diff --git a/.claude/docs/weight_sync.md b/.claude/docs/weight_sync.md index 85b4baa25b..8f03928727 100644 --- a/.claude/docs/weight_sync.md +++ b/.claude/docs/weight_sync.md @@ -48,13 +48,16 @@ The weight sync implementation relies on the native vLLM weight sync APIs - `Wei - **Broadcast** (`BroadcastTransferStrategy`): NCCL collective. Used for **non-colocated** setups. Training and inference are on different GPUs; weights cross the wire over a dedicated process group. - **CUDA IPC** (`CudaIpcTransferStrategy`): Per-chunk packed buffer + one IPC handle per rank. Used for **colocated** setups (`colocate_all=true`). Both sides live on the same GPU; the receiver maps the sender's CUDA allocation directly. -- **Delta** (`DeltaTransferStrategy`): Weights travel as compressed XOR deltas against the base checkpoint, through a shared filesystem or object store instead of the network fabric. Selected with `generator.inference_engine.weight_sync_backend=delta`; intended for **non-colocated** setups where the two sides are not NCCL-reachable (separate clusters, PD-disaggregated serving). Not supported with LoRA (`validate_cfg` rejects it). +- **Delta** (`DeltaTransferStrategy`): Weights travel as compressed XOR deltas against the base checkpoint, through a shared filesystem or object store instead of the network fabric. Selected with `generator.inference_engine.weight_sync_backend=delta`; intended for **non-colocated** setups where the two sides are not NCCL-reachable (separate clusters, PD-disaggregated serving). Not supported with LoRA, nor with serialized FP8 weight sync (`fp8_weight_sync_mode=blockwise`) whose marker names and scale tensors the delta checkpoint format cannot represent; `validate_cfg` rejects both. - **Sharded RDT** (`sharded_rdt`): the inference workers **pull** the slices they consume from the trainer ranks over NIXL/RDMA, instead of the trainer pushing every tensor to every worker. Selected with `generator.inference_engine.weight_sync_backend=sharded_rdt`; non-colocated only (`placement.colocate_all=false`), Megatron or FSDP, and it forces - `distributed_executor_backend=ray` because the workers dial named trainer actors. See - the dedicated section below for the capabilities it declares. + `distributed_executor_backend=ray` because the workers dial named trainer actors. Not + supported with serialized FP8 weight sync (`fp8_weight_sync_mode=blockwise`): the weight + sources publish whole bridge tensors, never payload+scale pairs, so + `validate_inference_engine_cfg` rejects the combination. See the dedicated section below + for the capabilities it declares. Strategy choice is decided by the sender (`get_transfer_strategy_cls`). The init info is expanded per server via `for_servers()` / `to_api_payload()` and pushed to the servers through the HTTP control plane (`init_weight_update_communicator` → vLLM's native `/init_weight_transfer_engine`); the receive side is vLLM's native weight-transfer engine, driven by `NewInferenceWorkerWrap`. diff --git a/examples/train/fp8/README.md b/examples/train/fp8/README.md new file mode 100644 index 0000000000..30b11edf03 --- /dev/null +++ b/examples/train/fp8/README.md @@ -0,0 +1,48 @@ +# FP8 RL training + rollout examples + +DAPO on AIME with FP8 across the performance-critical parts of the stack: +trainer linear-layer GEMMs, rollout weights, and the weight transfer between +them. All scripts use `fp8_weight_sync_mode=blockwise`, which sends +the trainer-produced FP8 payloads and block scales directly to vLLM instead of +re-quantizing a BF16 export — keeping the rollout policy numerically identical +to the trained one. + +Prepare the dataset once: + +```bash +bash examples/train/algorithms/dapo/prepare_dapo_data.sh +``` + +| Script | Hardware | Recipe | FP8 params | +| --- | --- | --- | --- | +| `run_fp8_hopper_blockwise_qwen35_9b.sh` | 8×H100 | blockwise, FP32 scales | — | +| `run_fp8_hopper_blockwise_fp8param_qwen35_9b.sh` | 8×H100 | blockwise, FP32 scales | E4M3 primary weights (~39% less parameter HBM) | +| `run_fp8_hopper_blockwise_qwen35_35b_a3b.sh` | 2×8×H100 | blockwise, FP32 scales | — | +| `run_fp8_hopper_blockwise_fp8param_qwen35_35b_a3b.sh` | 2×8×H100 | blockwise, FP32 scales | E4M3 primary weights (~42% less parameter HBM) | +| `run_fp8_blackwell_mxfp8_qwen35_9b.sh` | 8×B200 | `auto` → native MXFP8 | not yet supported on MXFP8 | +| `run_fp8_blackwell_mxfp8_qwen35_35b_a3b.sh` | 8×B200 | `auto` → native MXFP8 | not yet supported on MXFP8 | + +Notes: + +- **Colocated vs. non-colocated.** Every script defaults to + `trainer.placement.colocate_all=true` (training and inference share GPUs). + Run with `COLOCATE_ALL=false` and split the GPUs between + `trainer.placement.policy_num_gpus_per_node` and the inference engines for a + disaggregated placement. +- **Qwen3.5 runs text-only.** All scripts set `language_model_only=true` on the policy, ref and + inference engine: Qwen3.5 otherwise loads through the VL bridge, which packs sequences inside + its own forward and is rejected together with SkyRL sample packing. +- **GDN kernels on Blackwell.** The Blackwell scripts `export FLA_TILELANG=0` so fla uses its + Triton GatedDeltaNet kernels; the TileLang packed backward aborts on B200 (it shows up as a + CUDA "misaligned address" in the first backward). Leave it unset on Hopper, where the Triton + backward is the broken one. +- **Recipe selection.** `fp8_recipe=auto` picks the architecture-native + recipe: `blockwise` (FP32 scales) on Hopper, `mxfp8` on Blackwell/SM100+. + The Hopper scripts pin `blockwise` explicitly; the Blackwell scripts use + `auto`. +- **FP8 configuration surface.** The scripts use the top-level + `megatron_config.fp8*` fields; the same keys under + `transformer_config_kwargs` override them if you need to. +- **KV cache.** FP8 KV cache for these hybrid-attention models is a separate + compatibility PR; once available, add + `generator.inference_engine.engine_init_kwargs.kv_cache_dtype=fp8_e4m3`. diff --git a/examples/train/fp8/run_fp8_blackwell_mxfp8_qwen35_35b_a3b.sh b/examples/train/fp8/run_fp8_blackwell_mxfp8_qwen35_35b_a3b.sh new file mode 100755 index 0000000000..9b82a94153 --- /dev/null +++ b/examples/train/fp8/run_fp8_blackwell_mxfp8_qwen35_35b_a3b.sh @@ -0,0 +1,144 @@ +set -x + +# Colocated DAPO with MXFP8 training + FP8 rollout for Qwen3.5-35B-A3B-Base (MoE). +# Hardware: 1 node of 8xB200 +# +# bash examples/train/algorithms/dapo/prepare_dapo_data.sh +# bash examples/train/fp8/run_fp8_blackwell_mxfp8_qwen35_35b_a3b.sh +# +# FP8 here covers trainer linear-layer GEMMs plus rollout weights via +# representation-preserving weight sync (fp8_weight_sync_mode=blockwise): +# vLLM receives the trainer-produced FP8 payloads and block scales instead of +# re-quantizing a BF16 export. +# +# fp8_param is not yet supported on the native MXFP8 path; primary training +# weights stay in BF16 while GEMMs and rollout weights run FP8. + +MODEL_NAME="Qwen/Qwen3.5-35B-A3B-Base" +DATA_DIR="$HOME/data/dapo" +TRAIN_FILE="$DATA_DIR/dapo-math-17k-cleaned.parquet" +TEST_FILE="$DATA_DIR/aime-2024-cleaned.parquet" +LOGGER="wandb" # change to "console" to print to stdout + +# Colocated by default: training and inference share the same GPUs. For a +# disaggregated (non-colocated) run, set COLOCATE_ALL=false and split the GPUs, +# e.g. trainer.placement.policy_num_gpus_per_node=4 with the remaining GPUs +# given to the inference engines via generator.inference_engine.num_engines. +COLOCATE_ALL=${COLOCATE_ALL:-true} + +NUM_NODES=1 +NUM_GPUS_PER_NODE=8 +NUM_INFERENCE_ENGINES=8 +INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=1 + +MEGATRON_TP=1 +MEGATRON_PP=1 +MEGATRON_CP=1 +MEGATRON_EP=8 +MEGATRON_ETP=1 + +# Qwen3.5 goes through the VL bridge (Qwen3VLModel), which packs sequences in its own +# forward and conflicts with SkyRL sample packing; language_model_only routes it to the +# native GPTModel + GDN THD packing path on both the trainer and vLLM. +LANGUAGE_MODEL_ONLY=true + +# ---- FP8: trainer GEMMs + rollout weight sync ---- +# fp8_recipe=auto resolves to TE's architecture-native recipe: MXFP8 on +# Blackwell (SM100+). Weight sync still transfers 128x128 blockwise FP8 with +# power-of-2 scales, which Blackwell DeepGEMM consumes as E8M0. +MEGATRON_FP8=e4m3 +MEGATRON_FP8_RECIPE=auto +MEGATRON_FP8_AMAX_COMPUTE_ALGO=most_recent +MEGATRON_TP_ONLY_AMAX_RED=false +FP8_WEIGHT_SYNC_MODE=blockwise +export NVTE_FP8_BLOCK_SCALING_FP32_SCALES=0 +# Pinned rather than left to the default so the contract is explicit on +# both ends: power-of-2 wire scales, which vLLM consumes as E8M0. +export VLLM_USE_DEEP_GEMM_E8M0=1 +# fla's default TileLang GDN backend aborts in the packed backward on Blackwell (surfaces as +# a CUDA "misaligned address" from the next Triton launch); force the Triton GDN kernels. +# Leave unset on Hopper, where the Triton GDN backward is the broken one: +# https://github.com/fla-org/flash-linear-attention/issues/640#issuecomment-4236520788 +export FLA_TILELANG=0 + +uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ + data.train_data="['$TRAIN_FILE']" \ + data.val_data="['$TEST_FILE']" \ + trainer.algorithm.advantage_estimator="grpo" \ + trainer.algorithm.policy_loss_type="regular" \ + trainer.algorithm.overlong_buffer_len=4096 \ + trainer.algorithm.overlong_buffer_penalty_factor=1.0 \ + trainer.algorithm.loss_reduction=token_mean \ + trainer.algorithm.use_kl_loss=false \ + trainer.algorithm.clip_ratio_c=10.0 \ + trainer.algorithm.eps_clip_low=0.2 \ + trainer.algorithm.eps_clip_high=0.28 \ + generator.apply_overlong_filtering=true \ + generator.sampling_params.temperature=1.0 \ + generator.sampling_params.top_p=1.0 \ + generator.sampling_params.max_generate_length=8192 \ + generator.sampling_params.logprobs=1 \ + generator.eval_sampling_params.temperature=1.0 \ + generator.eval_sampling_params.top_p=1.0 \ + generator.eval_sampling_params.max_generate_length=8192 \ + trainer.policy.model.path="$MODEL_NAME" \ + trainer.policy.language_model_only=$LANGUAGE_MODEL_ONLY \ + trainer.ref.language_model_only=$LANGUAGE_MODEL_ONLY \ + generator.inference_engine.language_model_only=$LANGUAGE_MODEL_ONLY \ + trainer.placement.colocate_all=$COLOCATE_ALL \ + trainer.strategy=megatron \ + trainer.placement.policy_num_nodes=$NUM_NODES \ + trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + trainer.placement.ref_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.ref.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.ref.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.ref.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.ref.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.ref.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.policy.megatron_config.fp8=$MEGATRON_FP8 \ + trainer.ref.megatron_config.fp8=$MEGATRON_FP8 \ + trainer.policy.megatron_config.fp8_recipe=$MEGATRON_FP8_RECIPE \ + trainer.ref.megatron_config.fp8_recipe=$MEGATRON_FP8_RECIPE \ + trainer.policy.megatron_config.fp8_amax_compute_algo=$MEGATRON_FP8_AMAX_COMPUTE_ALGO \ + trainer.ref.megatron_config.fp8_amax_compute_algo=$MEGATRON_FP8_AMAX_COMPUTE_ALGO \ + trainer.policy.megatron_config.transformer_config_kwargs.tp_only_amax_red=$MEGATRON_TP_ONLY_AMAX_RED \ + trainer.ref.megatron_config.transformer_config_kwargs.tp_only_amax_red=$MEGATRON_TP_ONLY_AMAX_RED \ + generator.inference_engine.fp8_weight_sync_mode=$FP8_WEIGHT_SYNC_MODE \ + generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ + generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.gpu_memory_utilization=0.7 \ + generator.batched=true \ + environment.env_class=aime \ + generator.n_samples_per_prompt=8 \ + generator.eval_n_samples_per_prompt=16 \ + trainer.epochs=20 \ + trainer.max_training_steps=400 \ + trainer.eval_batch_size=512 \ + trainer.eval_before_train=false \ + trainer.eval_interval=-1 \ + trainer.update_epochs_per_batch=1 \ + trainer.train_batch_size=32 \ + trainer.policy_mini_batch_size=32 \ + trainer.micro_forward_batch_size_per_gpu=2 \ + trainer.micro_train_batch_size_per_gpu=2 \ + trainer.max_prompt_length=2048 \ + trainer.policy.optimizer_config.lr=1e-6 \ + trainer.policy.optimizer_config.num_warmup_steps=0 \ + trainer.policy.optimizer_config.weight_decay=0.1 \ + trainer.policy.optimizer_config.max_grad_norm=1.0 \ + trainer.logger="$LOGGER" \ + trainer.project_name="skyrl_fp8" \ + trainer.run_name="fp8_blackwell_mxfp8_qwen35_35b_a3b" \ + trainer.ckpt_interval=-1 \ + trainer.hf_save_interval=-1 \ + trainer.resume_mode=null \ + trainer.max_ckpts_to_keep=3 \ + $@ diff --git a/examples/train/fp8/run_fp8_blackwell_mxfp8_qwen35_9b.sh b/examples/train/fp8/run_fp8_blackwell_mxfp8_qwen35_9b.sh new file mode 100755 index 0000000000..055b07b345 --- /dev/null +++ b/examples/train/fp8/run_fp8_blackwell_mxfp8_qwen35_9b.sh @@ -0,0 +1,144 @@ +set -x + +# Colocated DAPO with MXFP8 training + FP8 rollout for Qwen3.5-9B-Base. +# Hardware: 1 node of 8xB200 +# +# bash examples/train/algorithms/dapo/prepare_dapo_data.sh +# bash examples/train/fp8/run_fp8_blackwell_mxfp8_qwen35_9b.sh +# +# FP8 here covers trainer linear-layer GEMMs plus rollout weights via +# representation-preserving weight sync (fp8_weight_sync_mode=blockwise): +# vLLM receives the trainer-produced FP8 payloads and block scales instead of +# re-quantizing a BF16 export. +# +# fp8_param is not yet supported on the native MXFP8 path; primary training +# weights stay in BF16 while GEMMs and rollout weights run FP8. + +MODEL_NAME="Qwen/Qwen3.5-9B-Base" +DATA_DIR="$HOME/data/dapo" +TRAIN_FILE="$DATA_DIR/dapo-math-17k-cleaned.parquet" +TEST_FILE="$DATA_DIR/aime-2024-cleaned.parquet" +LOGGER="wandb" # change to "console" to print to stdout + +# Colocated by default: training and inference share the same GPUs. For a +# disaggregated (non-colocated) run, set COLOCATE_ALL=false and split the GPUs, +# e.g. trainer.placement.policy_num_gpus_per_node=4 with the remaining GPUs +# given to the inference engines via generator.inference_engine.num_engines. +COLOCATE_ALL=${COLOCATE_ALL:-true} + +NUM_NODES=1 +NUM_GPUS_PER_NODE=8 +NUM_INFERENCE_ENGINES=8 +INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=1 + +MEGATRON_TP=2 +MEGATRON_PP=1 +MEGATRON_CP=1 +MEGATRON_EP=1 +MEGATRON_ETP=1 + +# Qwen3.5 goes through the VL bridge (Qwen3VLModel), which packs sequences in its own +# forward and conflicts with SkyRL sample packing; language_model_only routes it to the +# native GPTModel + GDN THD packing path on both the trainer and vLLM. +LANGUAGE_MODEL_ONLY=true + +# ---- FP8: trainer GEMMs + rollout weight sync ---- +# fp8_recipe=auto resolves to TE's architecture-native recipe: MXFP8 on +# Blackwell (SM100+). Weight sync still transfers 128x128 blockwise FP8 with +# power-of-2 scales, which Blackwell DeepGEMM consumes as E8M0. +MEGATRON_FP8=e4m3 +MEGATRON_FP8_RECIPE=auto +MEGATRON_FP8_AMAX_COMPUTE_ALGO=most_recent +MEGATRON_TP_ONLY_AMAX_RED=false +FP8_WEIGHT_SYNC_MODE=blockwise +export NVTE_FP8_BLOCK_SCALING_FP32_SCALES=0 +# Pinned rather than left to the default so the contract is explicit on +# both ends: power-of-2 wire scales, which vLLM consumes as E8M0. +export VLLM_USE_DEEP_GEMM_E8M0=1 +# fla's default TileLang GDN backend aborts in the packed backward on Blackwell (surfaces as +# a CUDA "misaligned address" from the next Triton launch); force the Triton GDN kernels. +# Leave unset on Hopper, where the Triton GDN backward is the broken one: +# https://github.com/fla-org/flash-linear-attention/issues/640#issuecomment-4236520788 +export FLA_TILELANG=0 + +uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ + data.train_data="['$TRAIN_FILE']" \ + data.val_data="['$TEST_FILE']" \ + trainer.algorithm.advantage_estimator="grpo" \ + trainer.algorithm.policy_loss_type="regular" \ + trainer.algorithm.overlong_buffer_len=4096 \ + trainer.algorithm.overlong_buffer_penalty_factor=1.0 \ + trainer.algorithm.loss_reduction=token_mean \ + trainer.algorithm.use_kl_loss=false \ + trainer.algorithm.clip_ratio_c=10.0 \ + trainer.algorithm.eps_clip_low=0.2 \ + trainer.algorithm.eps_clip_high=0.28 \ + generator.apply_overlong_filtering=true \ + generator.sampling_params.temperature=1.0 \ + generator.sampling_params.top_p=1.0 \ + generator.sampling_params.max_generate_length=8192 \ + generator.sampling_params.logprobs=1 \ + generator.eval_sampling_params.temperature=1.0 \ + generator.eval_sampling_params.top_p=1.0 \ + generator.eval_sampling_params.max_generate_length=8192 \ + trainer.policy.model.path="$MODEL_NAME" \ + trainer.policy.language_model_only=$LANGUAGE_MODEL_ONLY \ + trainer.ref.language_model_only=$LANGUAGE_MODEL_ONLY \ + generator.inference_engine.language_model_only=$LANGUAGE_MODEL_ONLY \ + trainer.placement.colocate_all=$COLOCATE_ALL \ + trainer.strategy=megatron \ + trainer.placement.policy_num_nodes=$NUM_NODES \ + trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + trainer.placement.ref_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.ref.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.ref.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.ref.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.ref.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.ref.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.policy.megatron_config.fp8=$MEGATRON_FP8 \ + trainer.ref.megatron_config.fp8=$MEGATRON_FP8 \ + trainer.policy.megatron_config.fp8_recipe=$MEGATRON_FP8_RECIPE \ + trainer.ref.megatron_config.fp8_recipe=$MEGATRON_FP8_RECIPE \ + trainer.policy.megatron_config.fp8_amax_compute_algo=$MEGATRON_FP8_AMAX_COMPUTE_ALGO \ + trainer.ref.megatron_config.fp8_amax_compute_algo=$MEGATRON_FP8_AMAX_COMPUTE_ALGO \ + trainer.policy.megatron_config.transformer_config_kwargs.tp_only_amax_red=$MEGATRON_TP_ONLY_AMAX_RED \ + trainer.ref.megatron_config.transformer_config_kwargs.tp_only_amax_red=$MEGATRON_TP_ONLY_AMAX_RED \ + generator.inference_engine.fp8_weight_sync_mode=$FP8_WEIGHT_SYNC_MODE \ + generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ + generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.gpu_memory_utilization=0.7 \ + generator.batched=true \ + environment.env_class=aime \ + generator.n_samples_per_prompt=8 \ + generator.eval_n_samples_per_prompt=16 \ + trainer.epochs=20 \ + trainer.max_training_steps=400 \ + trainer.eval_batch_size=512 \ + trainer.eval_before_train=false \ + trainer.eval_interval=-1 \ + trainer.update_epochs_per_batch=1 \ + trainer.train_batch_size=32 \ + trainer.policy_mini_batch_size=32 \ + trainer.micro_forward_batch_size_per_gpu=2 \ + trainer.micro_train_batch_size_per_gpu=2 \ + trainer.max_prompt_length=2048 \ + trainer.policy.optimizer_config.lr=1e-6 \ + trainer.policy.optimizer_config.num_warmup_steps=0 \ + trainer.policy.optimizer_config.weight_decay=0.1 \ + trainer.policy.optimizer_config.max_grad_norm=1.0 \ + trainer.logger="$LOGGER" \ + trainer.project_name="skyrl_fp8" \ + trainer.run_name="fp8_blackwell_mxfp8_qwen35_9b" \ + trainer.ckpt_interval=-1 \ + trainer.hf_save_interval=-1 \ + trainer.resume_mode=null \ + trainer.max_ckpts_to_keep=3 \ + $@ diff --git a/examples/train/fp8/run_fp8_hopper_blockwise_fp8param_qwen35_35b_a3b.sh b/examples/train/fp8/run_fp8_hopper_blockwise_fp8param_qwen35_35b_a3b.sh new file mode 100755 index 0000000000..eefaafdc3e --- /dev/null +++ b/examples/train/fp8/run_fp8_hopper_blockwise_fp8param_qwen35_35b_a3b.sh @@ -0,0 +1,141 @@ +set -x + +# Colocated DAPO with blockwise FP8 + persistent FP8 parameters for Qwen3.5-35B-A3B-Base (MoE). +# Hardware: 2 nodes of 8xH100 +# +# bash examples/train/algorithms/dapo/prepare_dapo_data.sh +# bash examples/train/fp8/run_fp8_hopper_blockwise_fp8param_qwen35_35b_a3b.sh +# +# FP8 here covers trainer linear-layer GEMMs plus rollout weights via +# representation-preserving weight sync (fp8_weight_sync_mode=blockwise): +# vLLM receives the trainer-produced FP8 payloads and block scales instead of +# re-quantizing a BF16 export. +# +# fp8_param is most effective on MoE models: the routed experts hold most of +# the parameters and all quantize, cutting per-GPU parameter storage by ~42%. + +MODEL_NAME="Qwen/Qwen3.5-35B-A3B-Base" +DATA_DIR="$HOME/data/dapo" +TRAIN_FILE="$DATA_DIR/dapo-math-17k-cleaned.parquet" +TEST_FILE="$DATA_DIR/aime-2024-cleaned.parquet" +LOGGER="wandb" # change to "console" to print to stdout + +# Colocated by default: training and inference share the same GPUs. For a +# disaggregated (non-colocated) run, set COLOCATE_ALL=false and split the GPUs, +# e.g. trainer.placement.policy_num_gpus_per_node=4 with the remaining GPUs +# given to the inference engines via generator.inference_engine.num_engines. +COLOCATE_ALL=${COLOCATE_ALL:-true} + +NUM_NODES=2 +NUM_GPUS_PER_NODE=8 +NUM_INFERENCE_ENGINES=2 +INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=8 + +MEGATRON_TP=2 +MEGATRON_PP=1 +MEGATRON_CP=1 +MEGATRON_EP=8 +MEGATRON_ETP=1 + +# Qwen3.5 goes through the VL bridge (Qwen3VLModel), which packs sequences in its own +# forward and conflicts with SkyRL sample packing; language_model_only routes it to the +# native GPTModel + GDN THD packing path on both the trainer and vLLM. +LANGUAGE_MODEL_ONLY=true + +# ---- FP8: trainer GEMMs + rollout weight sync ---- +# Hopper blockwise FP8 uses exact FP32 block scales end to end. Both exported +# variables are defaulted and validated by SkyRL at startup; set for clarity. +MEGATRON_FP8=e4m3 +MEGATRON_FP8_RECIPE=blockwise +MEGATRON_FP8_PARAM=true +MEGATRON_FP8_PARAM_GATHER=true +MEGATRON_FP8_AMAX_COMPUTE_ALGO=most_recent +MEGATRON_TP_ONLY_AMAX_RED=false +FP8_WEIGHT_SYNC_MODE=blockwise +export NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 +export VLLM_USE_DEEP_GEMM_E8M0=0 + +uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ + data.train_data="['$TRAIN_FILE']" \ + data.val_data="['$TEST_FILE']" \ + trainer.algorithm.advantage_estimator="grpo" \ + trainer.algorithm.policy_loss_type="regular" \ + trainer.algorithm.overlong_buffer_len=4096 \ + trainer.algorithm.overlong_buffer_penalty_factor=1.0 \ + trainer.algorithm.loss_reduction=token_mean \ + trainer.algorithm.use_kl_loss=false \ + trainer.algorithm.clip_ratio_c=10.0 \ + trainer.algorithm.eps_clip_low=0.2 \ + trainer.algorithm.eps_clip_high=0.28 \ + generator.apply_overlong_filtering=true \ + generator.sampling_params.temperature=1.0 \ + generator.sampling_params.top_p=1.0 \ + generator.sampling_params.max_generate_length=8192 \ + generator.sampling_params.logprobs=1 \ + generator.eval_sampling_params.temperature=1.0 \ + generator.eval_sampling_params.top_p=1.0 \ + generator.eval_sampling_params.max_generate_length=8192 \ + trainer.policy.model.path="$MODEL_NAME" \ + trainer.policy.language_model_only=$LANGUAGE_MODEL_ONLY \ + trainer.ref.language_model_only=$LANGUAGE_MODEL_ONLY \ + generator.inference_engine.language_model_only=$LANGUAGE_MODEL_ONLY \ + trainer.placement.colocate_all=$COLOCATE_ALL \ + trainer.strategy=megatron \ + trainer.placement.policy_num_nodes=$NUM_NODES \ + trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + trainer.placement.ref_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.ref.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.ref.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.ref.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.ref.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.ref.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.policy.megatron_config.fp8=$MEGATRON_FP8 \ + trainer.ref.megatron_config.fp8=$MEGATRON_FP8 \ + trainer.policy.megatron_config.fp8_recipe=$MEGATRON_FP8_RECIPE \ + trainer.ref.megatron_config.fp8_recipe=$MEGATRON_FP8_RECIPE \ + trainer.policy.megatron_config.fp8_amax_compute_algo=$MEGATRON_FP8_AMAX_COMPUTE_ALGO \ + trainer.ref.megatron_config.fp8_amax_compute_algo=$MEGATRON_FP8_AMAX_COMPUTE_ALGO \ + trainer.policy.megatron_config.transformer_config_kwargs.tp_only_amax_red=$MEGATRON_TP_ONLY_AMAX_RED \ + trainer.ref.megatron_config.transformer_config_kwargs.tp_only_amax_red=$MEGATRON_TP_ONLY_AMAX_RED \ + trainer.policy.megatron_config.fp8_param=$MEGATRON_FP8_PARAM \ + trainer.ref.megatron_config.fp8_param=$MEGATRON_FP8_PARAM \ + trainer.policy.megatron_config.ddp_config.fp8_param_gather=$MEGATRON_FP8_PARAM_GATHER \ + generator.inference_engine.fp8_weight_sync_mode=$FP8_WEIGHT_SYNC_MODE \ + generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ + generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.gpu_memory_utilization=0.7 \ + generator.batched=true \ + environment.env_class=aime \ + generator.n_samples_per_prompt=8 \ + generator.eval_n_samples_per_prompt=16 \ + trainer.epochs=20 \ + trainer.max_training_steps=400 \ + trainer.eval_batch_size=512 \ + trainer.eval_before_train=false \ + trainer.eval_interval=-1 \ + trainer.update_epochs_per_batch=1 \ + trainer.train_batch_size=32 \ + trainer.policy_mini_batch_size=32 \ + trainer.micro_forward_batch_size_per_gpu=2 \ + trainer.micro_train_batch_size_per_gpu=2 \ + trainer.max_prompt_length=2048 \ + trainer.policy.optimizer_config.lr=1e-6 \ + trainer.policy.optimizer_config.num_warmup_steps=0 \ + trainer.policy.optimizer_config.weight_decay=0.1 \ + trainer.policy.optimizer_config.max_grad_norm=1.0 \ + trainer.logger="$LOGGER" \ + trainer.project_name="skyrl_fp8" \ + trainer.run_name="fp8_hopper_blockwise_fp8param_qwen35_35b_a3b" \ + trainer.ckpt_interval=-1 \ + trainer.hf_save_interval=-1 \ + trainer.resume_mode=null \ + trainer.max_ckpts_to_keep=3 \ + $@ diff --git a/examples/train/fp8/run_fp8_hopper_blockwise_fp8param_qwen35_9b.sh b/examples/train/fp8/run_fp8_hopper_blockwise_fp8param_qwen35_9b.sh new file mode 100755 index 0000000000..628454e702 --- /dev/null +++ b/examples/train/fp8/run_fp8_hopper_blockwise_fp8param_qwen35_9b.sh @@ -0,0 +1,142 @@ +set -x + +# Colocated DAPO with blockwise FP8 + persistent FP8 parameters for Qwen3.5-9B-Base. +# Hardware: 1 node of 8xH100 +# +# bash examples/train/algorithms/dapo/prepare_dapo_data.sh +# bash examples/train/fp8/run_fp8_hopper_blockwise_fp8param_qwen35_9b.sh +# +# FP8 here covers trainer linear-layer GEMMs plus rollout weights via +# representation-preserving weight sync (fp8_weight_sync_mode=blockwise): +# vLLM receives the trainer-produced FP8 payloads and block scales instead of +# re-quantizing a BF16 export. +# +# fp8_param stores the Megatron primary weights in E4M3 (the distributed +# optimizer keeps FP32 masters), cutting policy parameter storage by ~39% +# on this model. Requires the Hopper blockwise recipe with FP32 block scales. + +MODEL_NAME="Qwen/Qwen3.5-9B-Base" +DATA_DIR="$HOME/data/dapo" +TRAIN_FILE="$DATA_DIR/dapo-math-17k-cleaned.parquet" +TEST_FILE="$DATA_DIR/aime-2024-cleaned.parquet" +LOGGER="wandb" # change to "console" to print to stdout + +# Colocated by default: training and inference share the same GPUs. For a +# disaggregated (non-colocated) run, set COLOCATE_ALL=false and split the GPUs, +# e.g. trainer.placement.policy_num_gpus_per_node=4 with the remaining GPUs +# given to the inference engines via generator.inference_engine.num_engines. +COLOCATE_ALL=${COLOCATE_ALL:-true} + +NUM_NODES=1 +NUM_GPUS_PER_NODE=8 +NUM_INFERENCE_ENGINES=8 +INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=1 + +MEGATRON_TP=2 +MEGATRON_PP=1 +MEGATRON_CP=1 +MEGATRON_EP=1 +MEGATRON_ETP=1 + +# Qwen3.5 goes through the VL bridge (Qwen3VLModel), which packs sequences in its own +# forward and conflicts with SkyRL sample packing; language_model_only routes it to the +# native GPTModel + GDN THD packing path on both the trainer and vLLM. +LANGUAGE_MODEL_ONLY=true + +# ---- FP8: trainer GEMMs + rollout weight sync ---- +# Hopper blockwise FP8 uses exact FP32 block scales end to end. Both exported +# variables are defaulted and validated by SkyRL at startup; set for clarity. +MEGATRON_FP8=e4m3 +MEGATRON_FP8_RECIPE=blockwise +MEGATRON_FP8_PARAM=true +MEGATRON_FP8_PARAM_GATHER=true +MEGATRON_FP8_AMAX_COMPUTE_ALGO=most_recent +MEGATRON_TP_ONLY_AMAX_RED=false +FP8_WEIGHT_SYNC_MODE=blockwise +export NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 +export VLLM_USE_DEEP_GEMM_E8M0=0 + +uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ + data.train_data="['$TRAIN_FILE']" \ + data.val_data="['$TEST_FILE']" \ + trainer.algorithm.advantage_estimator="grpo" \ + trainer.algorithm.policy_loss_type="regular" \ + trainer.algorithm.overlong_buffer_len=4096 \ + trainer.algorithm.overlong_buffer_penalty_factor=1.0 \ + trainer.algorithm.loss_reduction=token_mean \ + trainer.algorithm.use_kl_loss=false \ + trainer.algorithm.clip_ratio_c=10.0 \ + trainer.algorithm.eps_clip_low=0.2 \ + trainer.algorithm.eps_clip_high=0.28 \ + generator.apply_overlong_filtering=true \ + generator.sampling_params.temperature=1.0 \ + generator.sampling_params.top_p=1.0 \ + generator.sampling_params.max_generate_length=8192 \ + generator.sampling_params.logprobs=1 \ + generator.eval_sampling_params.temperature=1.0 \ + generator.eval_sampling_params.top_p=1.0 \ + generator.eval_sampling_params.max_generate_length=8192 \ + trainer.policy.model.path="$MODEL_NAME" \ + trainer.policy.language_model_only=$LANGUAGE_MODEL_ONLY \ + trainer.ref.language_model_only=$LANGUAGE_MODEL_ONLY \ + generator.inference_engine.language_model_only=$LANGUAGE_MODEL_ONLY \ + trainer.placement.colocate_all=$COLOCATE_ALL \ + trainer.strategy=megatron \ + trainer.placement.policy_num_nodes=$NUM_NODES \ + trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + trainer.placement.ref_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.ref.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.ref.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.ref.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.ref.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.ref.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.policy.megatron_config.fp8=$MEGATRON_FP8 \ + trainer.ref.megatron_config.fp8=$MEGATRON_FP8 \ + trainer.policy.megatron_config.fp8_recipe=$MEGATRON_FP8_RECIPE \ + trainer.ref.megatron_config.fp8_recipe=$MEGATRON_FP8_RECIPE \ + trainer.policy.megatron_config.fp8_amax_compute_algo=$MEGATRON_FP8_AMAX_COMPUTE_ALGO \ + trainer.ref.megatron_config.fp8_amax_compute_algo=$MEGATRON_FP8_AMAX_COMPUTE_ALGO \ + trainer.policy.megatron_config.transformer_config_kwargs.tp_only_amax_red=$MEGATRON_TP_ONLY_AMAX_RED \ + trainer.ref.megatron_config.transformer_config_kwargs.tp_only_amax_red=$MEGATRON_TP_ONLY_AMAX_RED \ + trainer.policy.megatron_config.fp8_param=$MEGATRON_FP8_PARAM \ + trainer.ref.megatron_config.fp8_param=$MEGATRON_FP8_PARAM \ + trainer.policy.megatron_config.ddp_config.fp8_param_gather=$MEGATRON_FP8_PARAM_GATHER \ + generator.inference_engine.fp8_weight_sync_mode=$FP8_WEIGHT_SYNC_MODE \ + generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ + generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.gpu_memory_utilization=0.7 \ + generator.batched=true \ + environment.env_class=aime \ + generator.n_samples_per_prompt=8 \ + generator.eval_n_samples_per_prompt=16 \ + trainer.epochs=20 \ + trainer.max_training_steps=400 \ + trainer.eval_batch_size=512 \ + trainer.eval_before_train=false \ + trainer.eval_interval=-1 \ + trainer.update_epochs_per_batch=1 \ + trainer.train_batch_size=32 \ + trainer.policy_mini_batch_size=32 \ + trainer.micro_forward_batch_size_per_gpu=2 \ + trainer.micro_train_batch_size_per_gpu=2 \ + trainer.max_prompt_length=2048 \ + trainer.policy.optimizer_config.lr=1e-6 \ + trainer.policy.optimizer_config.num_warmup_steps=0 \ + trainer.policy.optimizer_config.weight_decay=0.1 \ + trainer.policy.optimizer_config.max_grad_norm=1.0 \ + trainer.logger="$LOGGER" \ + trainer.project_name="skyrl_fp8" \ + trainer.run_name="fp8_hopper_blockwise_fp8param_qwen35_9b" \ + trainer.ckpt_interval=-1 \ + trainer.hf_save_interval=-1 \ + trainer.resume_mode=null \ + trainer.max_ckpts_to_keep=3 \ + $@ diff --git a/examples/train/fp8/run_fp8_hopper_blockwise_qwen35_35b_a3b.sh b/examples/train/fp8/run_fp8_hopper_blockwise_qwen35_35b_a3b.sh new file mode 100755 index 0000000000..1e50115206 --- /dev/null +++ b/examples/train/fp8/run_fp8_hopper_blockwise_qwen35_35b_a3b.sh @@ -0,0 +1,133 @@ +set -x + +# Colocated DAPO with blockwise FP8 training + FP8 rollout for Qwen3.5-35B-A3B-Base (MoE). +# Hardware: 2 nodes of 8xH100 +# +# bash examples/train/algorithms/dapo/prepare_dapo_data.sh +# bash examples/train/fp8/run_fp8_hopper_blockwise_qwen35_35b_a3b.sh +# +# FP8 here covers trainer linear-layer GEMMs plus rollout weights via +# representation-preserving weight sync (fp8_weight_sync_mode=blockwise): +# vLLM receives the trainer-produced FP8 payloads and block scales instead of +# re-quantizing a BF16 export. + +MODEL_NAME="Qwen/Qwen3.5-35B-A3B-Base" +DATA_DIR="$HOME/data/dapo" +TRAIN_FILE="$DATA_DIR/dapo-math-17k-cleaned.parquet" +TEST_FILE="$DATA_DIR/aime-2024-cleaned.parquet" +LOGGER="wandb" # change to "console" to print to stdout + +# Colocated by default: training and inference share the same GPUs. For a +# disaggregated (non-colocated) run, set COLOCATE_ALL=false and split the GPUs, +# e.g. trainer.placement.policy_num_gpus_per_node=4 with the remaining GPUs +# given to the inference engines via generator.inference_engine.num_engines. +COLOCATE_ALL=${COLOCATE_ALL:-true} + +NUM_NODES=2 +NUM_GPUS_PER_NODE=8 +NUM_INFERENCE_ENGINES=2 +INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=8 + +MEGATRON_TP=2 +MEGATRON_PP=1 +MEGATRON_CP=1 +MEGATRON_EP=8 +MEGATRON_ETP=1 + +# Qwen3.5 goes through the VL bridge (Qwen3VLModel), which packs sequences in its own +# forward and conflicts with SkyRL sample packing; language_model_only routes it to the +# native GPTModel + GDN THD packing path on both the trainer and vLLM. +LANGUAGE_MODEL_ONLY=true + +# ---- FP8: trainer GEMMs + rollout weight sync ---- +# Hopper blockwise FP8 uses exact FP32 block scales end to end. Both exported +# variables are defaulted and validated by SkyRL at startup; set for clarity. +MEGATRON_FP8=e4m3 +MEGATRON_FP8_RECIPE=blockwise +MEGATRON_FP8_AMAX_COMPUTE_ALGO=most_recent +MEGATRON_TP_ONLY_AMAX_RED=false +FP8_WEIGHT_SYNC_MODE=blockwise +export NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 +export VLLM_USE_DEEP_GEMM_E8M0=0 + +uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ + data.train_data="['$TRAIN_FILE']" \ + data.val_data="['$TEST_FILE']" \ + trainer.algorithm.advantage_estimator="grpo" \ + trainer.algorithm.policy_loss_type="regular" \ + trainer.algorithm.overlong_buffer_len=4096 \ + trainer.algorithm.overlong_buffer_penalty_factor=1.0 \ + trainer.algorithm.loss_reduction=token_mean \ + trainer.algorithm.use_kl_loss=false \ + trainer.algorithm.clip_ratio_c=10.0 \ + trainer.algorithm.eps_clip_low=0.2 \ + trainer.algorithm.eps_clip_high=0.28 \ + generator.apply_overlong_filtering=true \ + generator.sampling_params.temperature=1.0 \ + generator.sampling_params.top_p=1.0 \ + generator.sampling_params.max_generate_length=8192 \ + generator.sampling_params.logprobs=1 \ + generator.eval_sampling_params.temperature=1.0 \ + generator.eval_sampling_params.top_p=1.0 \ + generator.eval_sampling_params.max_generate_length=8192 \ + trainer.policy.model.path="$MODEL_NAME" \ + trainer.policy.language_model_only=$LANGUAGE_MODEL_ONLY \ + trainer.ref.language_model_only=$LANGUAGE_MODEL_ONLY \ + generator.inference_engine.language_model_only=$LANGUAGE_MODEL_ONLY \ + trainer.placement.colocate_all=$COLOCATE_ALL \ + trainer.strategy=megatron \ + trainer.placement.policy_num_nodes=$NUM_NODES \ + trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + trainer.placement.ref_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.ref.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.ref.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.ref.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.ref.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.ref.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.policy.megatron_config.fp8=$MEGATRON_FP8 \ + trainer.ref.megatron_config.fp8=$MEGATRON_FP8 \ + trainer.policy.megatron_config.fp8_recipe=$MEGATRON_FP8_RECIPE \ + trainer.ref.megatron_config.fp8_recipe=$MEGATRON_FP8_RECIPE \ + trainer.policy.megatron_config.fp8_amax_compute_algo=$MEGATRON_FP8_AMAX_COMPUTE_ALGO \ + trainer.ref.megatron_config.fp8_amax_compute_algo=$MEGATRON_FP8_AMAX_COMPUTE_ALGO \ + trainer.policy.megatron_config.transformer_config_kwargs.tp_only_amax_red=$MEGATRON_TP_ONLY_AMAX_RED \ + trainer.ref.megatron_config.transformer_config_kwargs.tp_only_amax_red=$MEGATRON_TP_ONLY_AMAX_RED \ + generator.inference_engine.fp8_weight_sync_mode=$FP8_WEIGHT_SYNC_MODE \ + generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ + generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.gpu_memory_utilization=0.7 \ + generator.batched=true \ + environment.env_class=aime \ + generator.n_samples_per_prompt=8 \ + generator.eval_n_samples_per_prompt=16 \ + trainer.epochs=20 \ + trainer.max_training_steps=400 \ + trainer.eval_batch_size=512 \ + trainer.eval_before_train=false \ + trainer.eval_interval=-1 \ + trainer.update_epochs_per_batch=1 \ + trainer.train_batch_size=32 \ + trainer.policy_mini_batch_size=32 \ + trainer.micro_forward_batch_size_per_gpu=2 \ + trainer.micro_train_batch_size_per_gpu=2 \ + trainer.max_prompt_length=2048 \ + trainer.policy.optimizer_config.lr=1e-6 \ + trainer.policy.optimizer_config.num_warmup_steps=0 \ + trainer.policy.optimizer_config.weight_decay=0.1 \ + trainer.policy.optimizer_config.max_grad_norm=1.0 \ + trainer.logger="$LOGGER" \ + trainer.project_name="skyrl_fp8" \ + trainer.run_name="fp8_hopper_blockwise_qwen35_35b_a3b" \ + trainer.ckpt_interval=-1 \ + trainer.hf_save_interval=-1 \ + trainer.resume_mode=null \ + trainer.max_ckpts_to_keep=3 \ + $@ diff --git a/examples/train/fp8/run_fp8_hopper_blockwise_qwen35_9b.sh b/examples/train/fp8/run_fp8_hopper_blockwise_qwen35_9b.sh new file mode 100755 index 0000000000..fc66ba802b --- /dev/null +++ b/examples/train/fp8/run_fp8_hopper_blockwise_qwen35_9b.sh @@ -0,0 +1,133 @@ +set -x + +# Colocated DAPO with blockwise FP8 training + FP8 rollout for Qwen3.5-9B-Base. +# Hardware: 1 node of 8xH100 +# +# bash examples/train/algorithms/dapo/prepare_dapo_data.sh +# bash examples/train/fp8/run_fp8_hopper_blockwise_qwen35_9b.sh +# +# FP8 here covers trainer linear-layer GEMMs plus rollout weights via +# representation-preserving weight sync (fp8_weight_sync_mode=blockwise): +# vLLM receives the trainer-produced FP8 payloads and block scales instead of +# re-quantizing a BF16 export. + +MODEL_NAME="Qwen/Qwen3.5-9B-Base" +DATA_DIR="$HOME/data/dapo" +TRAIN_FILE="$DATA_DIR/dapo-math-17k-cleaned.parquet" +TEST_FILE="$DATA_DIR/aime-2024-cleaned.parquet" +LOGGER="wandb" # change to "console" to print to stdout + +# Colocated by default: training and inference share the same GPUs. For a +# disaggregated (non-colocated) run, set COLOCATE_ALL=false and split the GPUs, +# e.g. trainer.placement.policy_num_gpus_per_node=4 with the remaining GPUs +# given to the inference engines via generator.inference_engine.num_engines. +COLOCATE_ALL=${COLOCATE_ALL:-true} + +NUM_NODES=1 +NUM_GPUS_PER_NODE=8 +NUM_INFERENCE_ENGINES=8 +INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE=1 + +MEGATRON_TP=2 +MEGATRON_PP=1 +MEGATRON_CP=1 +MEGATRON_EP=1 +MEGATRON_ETP=1 + +# Qwen3.5 goes through the VL bridge (Qwen3VLModel), which packs sequences in its own +# forward and conflicts with SkyRL sample packing; language_model_only routes it to the +# native GPTModel + GDN THD packing path on both the trainer and vLLM. +LANGUAGE_MODEL_ONLY=true + +# ---- FP8: trainer GEMMs + rollout weight sync ---- +# Hopper blockwise FP8 uses exact FP32 block scales end to end. Both exported +# variables are defaulted and validated by SkyRL at startup; set for clarity. +MEGATRON_FP8=e4m3 +MEGATRON_FP8_RECIPE=blockwise +MEGATRON_FP8_AMAX_COMPUTE_ALGO=most_recent +MEGATRON_TP_ONLY_AMAX_RED=false +FP8_WEIGHT_SYNC_MODE=blockwise +export NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 +export VLLM_USE_DEEP_GEMM_E8M0=0 + +uv run --isolated --extra megatron -m examples.train.algorithms.dapo.main_dapo \ + data.train_data="['$TRAIN_FILE']" \ + data.val_data="['$TEST_FILE']" \ + trainer.algorithm.advantage_estimator="grpo" \ + trainer.algorithm.policy_loss_type="regular" \ + trainer.algorithm.overlong_buffer_len=4096 \ + trainer.algorithm.overlong_buffer_penalty_factor=1.0 \ + trainer.algorithm.loss_reduction=token_mean \ + trainer.algorithm.use_kl_loss=false \ + trainer.algorithm.clip_ratio_c=10.0 \ + trainer.algorithm.eps_clip_low=0.2 \ + trainer.algorithm.eps_clip_high=0.28 \ + generator.apply_overlong_filtering=true \ + generator.sampling_params.temperature=1.0 \ + generator.sampling_params.top_p=1.0 \ + generator.sampling_params.max_generate_length=8192 \ + generator.sampling_params.logprobs=1 \ + generator.eval_sampling_params.temperature=1.0 \ + generator.eval_sampling_params.top_p=1.0 \ + generator.eval_sampling_params.max_generate_length=8192 \ + trainer.policy.model.path="$MODEL_NAME" \ + trainer.policy.language_model_only=$LANGUAGE_MODEL_ONLY \ + trainer.ref.language_model_only=$LANGUAGE_MODEL_ONLY \ + generator.inference_engine.language_model_only=$LANGUAGE_MODEL_ONLY \ + trainer.placement.colocate_all=$COLOCATE_ALL \ + trainer.strategy=megatron \ + trainer.placement.policy_num_nodes=$NUM_NODES \ + trainer.placement.policy_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + trainer.placement.ref_num_gpus_per_node=$NUM_GPUS_PER_NODE \ + trainer.policy.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.policy.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.policy.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.policy.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.policy.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.ref.megatron_config.tensor_model_parallel_size=$MEGATRON_TP \ + trainer.ref.megatron_config.pipeline_model_parallel_size=$MEGATRON_PP \ + trainer.ref.megatron_config.context_parallel_size=$MEGATRON_CP \ + trainer.ref.megatron_config.expert_model_parallel_size=$MEGATRON_EP \ + trainer.ref.megatron_config.expert_tensor_parallel_size=$MEGATRON_ETP \ + trainer.policy.megatron_config.fp8=$MEGATRON_FP8 \ + trainer.ref.megatron_config.fp8=$MEGATRON_FP8 \ + trainer.policy.megatron_config.fp8_recipe=$MEGATRON_FP8_RECIPE \ + trainer.ref.megatron_config.fp8_recipe=$MEGATRON_FP8_RECIPE \ + trainer.policy.megatron_config.fp8_amax_compute_algo=$MEGATRON_FP8_AMAX_COMPUTE_ALGO \ + trainer.ref.megatron_config.fp8_amax_compute_algo=$MEGATRON_FP8_AMAX_COMPUTE_ALGO \ + trainer.policy.megatron_config.transformer_config_kwargs.tp_only_amax_red=$MEGATRON_TP_ONLY_AMAX_RED \ + trainer.ref.megatron_config.transformer_config_kwargs.tp_only_amax_red=$MEGATRON_TP_ONLY_AMAX_RED \ + generator.inference_engine.fp8_weight_sync_mode=$FP8_WEIGHT_SYNC_MODE \ + generator.inference_engine.num_engines=$NUM_INFERENCE_ENGINES \ + generator.inference_engine.tensor_parallel_size=$INFERENCE_ENGINE_TENSOR_PARALLEL_SIZE \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.gpu_memory_utilization=0.7 \ + generator.batched=true \ + environment.env_class=aime \ + generator.n_samples_per_prompt=8 \ + generator.eval_n_samples_per_prompt=16 \ + trainer.epochs=20 \ + trainer.max_training_steps=400 \ + trainer.eval_batch_size=512 \ + trainer.eval_before_train=false \ + trainer.eval_interval=-1 \ + trainer.update_epochs_per_batch=1 \ + trainer.train_batch_size=32 \ + trainer.policy_mini_batch_size=32 \ + trainer.micro_forward_batch_size_per_gpu=2 \ + trainer.micro_train_batch_size_per_gpu=2 \ + trainer.max_prompt_length=2048 \ + trainer.policy.optimizer_config.lr=1e-6 \ + trainer.policy.optimizer_config.num_warmup_steps=0 \ + trainer.policy.optimizer_config.weight_decay=0.1 \ + trainer.policy.optimizer_config.max_grad_norm=1.0 \ + trainer.logger="$LOGGER" \ + trainer.project_name="skyrl_fp8" \ + trainer.run_name="fp8_hopper_blockwise_qwen35_9b" \ + trainer.ckpt_interval=-1 \ + trainer.hf_save_interval=-1 \ + trainer.resume_mode=null \ + trainer.max_ckpts_to_keep=3 \ + $@ diff --git a/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py b/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py index 388f43482a..b3364cfc45 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py +++ b/skyrl/backends/skyrl_train/distributed/megatron/megatron_utils.py @@ -426,6 +426,7 @@ def preprocess_packed_seqs( pre_process: bool = True, sub_seq_lengths: Optional[list[list[int]]] = None, fp8_enabled: bool = False, + fp8_recipe: Optional[str] = None, ) -> tuple[torch.Tensor, PackedSeqParams]: """ Preprocess packed sequences. @@ -438,12 +439,11 @@ def preprocess_packed_seqs( per row. This is the historical SkyRL behavior used by the RL path and the existing SFT path without mini-batch packing. - ``sub_seq_lengths is not None``: each row may contain multiple - sub-sequences concatenated end-to-end. ``sub_seq_lengths[r]`` lists - the per-sub-sequence valid token counts for row ``r``. Tokens - ``input_ids[r, :sum(sub_seq_lengths[r])]`` are assumed to be the - concatenated sub-sequences in order; any trailing tokens in the row - are pad. ``cu_seqlens`` enumerates every sub-sequence across every - row. + sub-sequences. ``sub_seq_lengths[r]`` lists their valid token counts. + Each sub-sequence begins at the next ``align_size`` boundary, so internal + alignment padding may separate adjacent sub-sequences; any remaining + trailing tokens are pad. ``cu_seqlens`` enumerates every sub-sequence + across every row. CP splits sequence into CP*2 chunks, and each GPU gets 2 chunks (GPU0 gets first and last chunks, GPU1 gets second and second last chunks, @@ -453,7 +453,7 @@ def preprocess_packed_seqs( tp_size = mpu.get_tensor_model_parallel_world_size() cp_size = mpu.get_context_parallel_world_size() cp_rank = mpu.get_context_parallel_rank() - align_size = get_packed_seq_align_size(tp_size, cp_size, fp8_enabled=fp8_enabled) + align_size = get_packed_seq_align_size(tp_size, cp_size, fp8_enabled=fp8_enabled, fp8_recipe=fp8_recipe) batch_size = input_ids.shape[0] @@ -597,6 +597,7 @@ def remove_left_padding( position_ids: torch.Tensor, pre_process: bool = True, fp8_enabled: bool = False, + fp8_recipe: Optional[str] = None, ): """ Remove left padding from input_ids, attention_mask and position_ids @@ -610,7 +611,9 @@ def remove_left_padding( shape = list(input_ids.shape) # batch_size, seq_len,... seq_lens = attention_mask.sum(dim=1) seq_len = seq_lens.max().item() - align_size = get_unpacked_seq_align_size(mpu.get_tensor_model_parallel_world_size(), fp8_enabled=fp8_enabled) + align_size = get_unpacked_seq_align_size( + mpu.get_tensor_model_parallel_world_size(), fp8_enabled=fp8_enabled, fp8_recipe=fp8_recipe + ) pad_size = (align_size - seq_len % align_size) % align_size seq_len = seq_len + pad_size shape[1] = seq_len diff --git a/skyrl/backends/skyrl_train/distributed/megatron/packing_utils.py b/skyrl/backends/skyrl_train/distributed/megatron/packing_utils.py index f6b506361b..f66aa8e8b6 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/packing_utils.py +++ b/skyrl/backends/skyrl_train/distributed/megatron/packing_utils.py @@ -1,27 +1,43 @@ import math -from typing import Any +from typing import Any, Optional +from skyrl.backends.skyrl_train.distributed.megatron.quantization_utils import ( + is_mxfp8_recipe, +) -def is_fp8_enabled(fp8: Any) -> bool: - """Return whether a Megatron/TE fp8 config value enables FP8 execution.""" - if isinstance(fp8, str): - return fp8.strip().lower() not in {"", "0", "false", "none", "null", "no"} - return bool(fp8) +def _fp8_token_align(tp_size: int, cp_size: int, fp8_recipe: Any) -> int: + # MXFP8 quantizes sequence-parallel all-gather inputs in 1x32 tiles, so + # every rank's local shard must hold a multiple of 32 tokens: 32*tp*cp + # globally at any TP. Blockwise FP8 quantizes in 1x128 tiles, requiring + # 128-token local shards under sequence parallelism (a 128*tp*cp global + # segment when tp>1) and 16-token local slabs at TP=1. + if is_mxfp8_recipe(fp8_recipe): + return 32 * tp_size * cp_size + if tp_size > 1: + return 128 * tp_size * cp_size + return 16 * cp_size -def get_packed_seq_align_size(tp_size: int, cp_size: int, fp8_enabled: bool = False) -> int: - """Return global per-subsequence padding needed for TP/CP layout.""" + +def get_packed_seq_align_size( + tp_size: int, cp_size: int, fp8_enabled: bool = False, fp8_recipe: Optional[str] = None +) -> int: + """Return the global alignment unit for packed TP/CP/FP8 sequences.""" + if tp_size < 1 or cp_size < 1: + raise ValueError(f"tp_size and cp_size must be positive, got tp_size={tp_size}, cp_size={cp_size}") if cp_size > 1: layout_align = tp_size * cp_size * 2 else: layout_align = tp_size if not fp8_enabled: return layout_align - return math.lcm(layout_align, 16 * cp_size) + return math.lcm(layout_align, _fp8_token_align(tp_size, cp_size, fp8_recipe)) -def get_unpacked_seq_align_size(tp_size: int, fp8_enabled: bool = False) -> int: - """Return sequence padding needed when removing microbatch padding without CP.""" +def get_unpacked_seq_align_size(tp_size: int, fp8_enabled: bool = False, fp8_recipe: Optional[str] = None) -> int: + """Return the alignment unit for unpacked TP/FP8 sequences without CP.""" + if tp_size < 1: + raise ValueError(f"tp_size must be positive, got {tp_size}") if not fp8_enabled: return tp_size - return math.lcm(tp_size, 16) + return math.lcm(tp_size, _fp8_token_align(tp_size, 1, fp8_recipe)) diff --git a/skyrl/backends/skyrl_train/distributed/megatron/quantization_utils.py b/skyrl/backends/skyrl_train/distributed/megatron/quantization_utils.py new file mode 100644 index 0000000000..b158b8d373 --- /dev/null +++ b/skyrl/backends/skyrl_train/distributed/megatron/quantization_utils.py @@ -0,0 +1,97 @@ +"""Recipe/architecture helpers shared by FP8 packing, weight sync, and workers.""" + +from typing import Any, MutableMapping, Optional + +from loguru import logger + +AUTO_FP8_RECIPE = "auto" + + +def is_fp8_enabled(fp8: Any) -> bool: + """Return whether a Megatron/TE fp8 config value enables FP8 execution.""" + if isinstance(fp8, str): + return fp8.strip().lower() not in {"", "0", "false", "none", "null", "no", "off"} + return bool(fp8) + + +def is_mxfp8_recipe(fp8_recipe: Any) -> bool: + """Return whether a Megatron/TE fp8 recipe value selects MXFP8.""" + return isinstance(fp8_recipe, str) and fp8_recipe.strip().lower() == "mxfp8" + + +def has_visible_cuda_device() -> bool: + """Return whether this process can see a CUDA device.""" + import torch + + return torch.cuda.is_available() + + +def is_blackwell_or_newer() -> bool: + """Return whether the visible CUDA device is SM100+ (Blackwell or newer).""" + import torch + + if not has_visible_cuda_device(): + return False + major, _minor = torch.cuda.get_device_capability() + return major >= 10 + + +def resolve_auto_fp8_recipe(transformer_config_kwargs: Optional[MutableMapping[str, Any]]) -> Any: + """Resolve ``fp8_recipe="auto"`` to the architecture-native TE recipe. + + Mutates ``transformer_config_kwargs`` in place and returns the resolved + recipe; any explicitly configured recipe is returned unchanged. ``"auto"`` + selects the recipe each architecture supports natively: ``blockwise`` + (``Float8BlockScaling``, 1x128/128x128 tiles with FP32 scales) on Hopper, + and ``mxfp8`` (``MXFP8BlockScaling``, hardware 1x32 tiles with E8M0 + scales) on Blackwell. + + A process with no visible CUDA device cannot know the workers' + architecture, so it leaves ``"auto"`` in place instead of guessing — a + GPU-less Ray driver must not bake ``blockwise`` into the config that + Blackwell workers receive. Each Megatron worker resolves again locally + before the value reaches TE. + """ + recipe = transformer_config_kwargs.get("fp8_recipe") if transformer_config_kwargs else None + if not isinstance(recipe, str) or recipe.strip().lower() != AUTO_FP8_RECIPE: + if not is_mxfp8_recipe(recipe) and isinstance(recipe, str) and recipe.strip() and is_blackwell_or_newer(): + # TE emulates the blockwise recipe on Blackwell's MX datapath with + # power-of-2 scales; MoE experts that receive zero tokens then emit + # amax==0 grad blocks whose degenerate scales overflow the grad + # norm. The native recipe has no such mode. + logger.warning( + "fp8_recipe={} is emulated on SM100+; prefer fp8_recipe='mxfp8' (or 'auto').", + recipe, + ) + return recipe + if not has_visible_cuda_device(): + logger.info('fp8_recipe="auto" left unresolved: no CUDA device visible; each worker resolves it locally.') + return AUTO_FP8_RECIPE + resolved = "mxfp8" if is_blackwell_or_newer() else "blockwise" + transformer_config_kwargs["fp8_recipe"] = resolved + return resolved + + +def validate_concrete_fp8_recipe(transformer_config_kwargs: Optional[MutableMapping[str, Any]]) -> None: + """Reject FP8 recipe/device combinations Transformer Engine cannot run. + + Runs wherever the recipe becomes concrete: on the driver when it has a + CUDA device, and on every Megatron worker after its local resolution — + which covers GPU-less drivers that ship ``"auto"`` through unresolved. + """ + recipe = transformer_config_kwargs.get("fp8_recipe") if transformer_config_kwargs else None + if not is_mxfp8_recipe(recipe): + return + if has_visible_cuda_device() and not is_blackwell_or_newer(): + raise ValueError( + "fp8_recipe=mxfp8 requires SM100+ (Blackwell): TE's MXFP8BlockScaling has " + "no pre-Blackwell kernel path. Use fp8_recipe='blockwise' (or 'auto') on Hopper." + ) + if is_fp8_enabled(transformer_config_kwargs.get("fp8_param")): + raise ValueError( + "fp8_param=true is not supported with fp8_recipe=mxfp8: Transformer Engine " + "2.11 cannot re-point MXFP8Tensor storage (replace_raw_data raises " + "NotImplementedError), and Megatron's reuse_grad_buf_for_mxfp8_param_ag " + "fallback zeroes freshly loaded persistent params on the first param " + "all-gather. Use fp8_param=false with mxfp8." + ) diff --git a/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py b/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py index cb3e526c95..94e11b860e 100644 --- a/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py +++ b/skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py @@ -1,6 +1,7 @@ """Token-aligned metadata layout transforms shared by training features.""" from dataclasses import dataclass +from typing import Optional import torch @@ -43,6 +44,7 @@ def build_token_metadata_layout( *, packed: bool, fp8_enabled: bool, + fp8_recipe: Optional[str] = None, ) -> TokenMetadataLayout: """Compute the shared layout once for all replayed token metadata.""" import megatron.core.parallel_state as mpu @@ -53,7 +55,7 @@ def build_token_metadata_layout( tp_size = mpu.get_tensor_model_parallel_world_size() if not packed: - align_size = get_unpacked_seq_align_size(tp_size, fp8_enabled=fp8_enabled) + align_size = get_unpacked_seq_align_size(tp_size, fp8_enabled=fp8_enabled, fp8_recipe=fp8_recipe) max_sequence_length = max(sequence_lengths) aligned_sequence_length = max_sequence_length + (-max_sequence_length % align_size) return TokenMetadataLayout( @@ -63,7 +65,7 @@ def build_token_metadata_layout( ) cp_size = mpu.get_context_parallel_world_size() - align_size = get_packed_seq_align_size(tp_size, cp_size, fp8_enabled=fp8_enabled) + align_size = get_packed_seq_align_size(tp_size, cp_size, fp8_enabled=fp8_enabled, fp8_recipe=fp8_recipe) padded_sequence_lengths_tensor = sequence_lengths_tensor + (-sequence_lengths_tensor % align_size) padded_sequence_lengths = padded_sequence_lengths_tensor.tolist() cu_seqlens_padded = torch.cat( diff --git a/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py b/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py index 9be72ae9ab..fa9867d354 100644 --- a/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py +++ b/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py @@ -24,12 +24,19 @@ skyrl.backends.skyrl_train.inference_servers.new_inference_worker_wrap.NewInferenceWorkerWrap """ +from typing import Any + import torch from skyrl.backends.skyrl_train.inference_servers.layerwise_reload import ( LayerwiseReloadWorkerMixin, _empty_cuda_cache_rocm, ) +from skyrl.backends.skyrl_train.weight_sync.base import cuda_uuid_to_str +from skyrl.backends.skyrl_train.weight_sync.fp8 import ( + SKYRL_BATCHED_MOE_FP8_PREFIX, + batched_moe_wire_targets, +) try: from skyrl.backends.skyrl_train.weight_sync.delta_engine import ( @@ -58,6 +65,128 @@ VLLM_NEW_INFERENCE_WORKER_EXTENSION_CLS = f"{__name__}.NewInferenceWorkerWrap" +# checkpoint-name suffix -> (fused vLLM parameter suffix, FusedMoE shard id), +# derived from the registered ModelFp8Specs so per-model fused-loader knowledge +# lives in exactly one place (weight_sync/fp8/models). +_BATCHED_MOE_TARGETS = batched_moe_wire_targets() + + +def _map_hf_weight_name(model: torch.nn.Module, name: str) -> str: + """Apply a top-level vLLM model's HF-to-runtime prefix mapping.""" + mapper = getattr(model, "hf_to_vllm_mapper", None) + if mapper is None: + return name + mapped = mapper.apply_list([name]) + if len(mapped) != 1: + raise ValueError(f"Unable to map batched MoE checkpoint name {name!r}") + return mapped[0] + + +def _load_batched_moe_fp8_tensor( + model: torch.nn.Module, + params_dict: dict[str, torch.nn.Parameter], + wire_name: str, + loaded_weight: torch.Tensor, +) -> bool: + """Load one expert-batched FP8 weight/scale through FusedMoE's loader. + + Returns ``False`` for ordinary checkpoint tensors. Marked tensors are + required to resolve successfully so a protocol mismatch cannot silently + leave stale rollout weights behind. + """ + if not wire_name.startswith(SKYRL_BATCHED_MOE_FP8_PREFIX): + return False + if loaded_weight.ndim != 3: + raise ValueError( + f"Batched MoE wire tensor must be 3D, got name={wire_name!r}, shape={tuple(loaded_weight.shape)}" + ) + + checkpoint_name = wire_name.removeprefix(SKYRL_BATCHED_MOE_FP8_PREFIX) + mapped_name = _map_hf_weight_name(model, checkpoint_name) + target_name = None + shard_id = None + for checkpoint_suffix, ( + target_suffix, + candidate_shard_id, + ) in _BATCHED_MOE_TARGETS.items(): + if mapped_name.endswith(checkpoint_suffix): + target_name = mapped_name[: -len(checkpoint_suffix)] + target_suffix + shard_id = candidate_shard_id + break + if target_name is None or shard_id is None: + raise ValueError(f"Unsupported batched MoE wire tensor name {wire_name!r}") + if target_name not in params_dict: + # vLLM 0.26 turned FusedMoE into a factory returning a MoERunner whose + # RoutedExperts submodule registers the expert parameters, adding one + # segment to every runtime name; earlier engines register them on the + # experts module directly. + module_path, _, param_leaf = target_name.rpartition(".") + nested_name = f"{module_path}.routed_experts.{param_leaf}" + if nested_name not in params_dict: + raise ValueError( + f"Batched MoE target parameter was not found for wire tensor {wire_name!r}: " + f"tried {target_name!r} and {nested_name!r}" + ) + target_name = nested_name + + param = params_dict[target_name] + weight_loader = getattr(param, "weight_loader", None) + if weight_loader is None or not getattr(weight_loader, "supports_moe_loading", False): + # Layerwise reload wraps the loader with functools.wraps, which copies + # this marker from FusedMoE.weight_loader onto the wrapper. + raise ValueError(f"Parameter {target_name!r} does not expose a FusedMoE weight loader") + + if param.shape[0] == loaded_weight.shape[0]: + success = weight_loader( + param, + loaded_weight, + target_name, + shard_id=shard_id, + expert_id=0, + return_success=True, + ) + if not success: + raise ValueError(f"Fused loading failed for batched MoE tensor {wire_name!r}") + return True + + # Expert-parallel vLLM keeps only a subset locally. Retain the compact wire + # format, but let the loader map global expert IDs one view at a time. + loaded_any = False + for expert_id, expert_weight in enumerate(loaded_weight.unbind(0)): + loaded_any = ( + bool( + weight_loader( + param, + expert_weight, + target_name, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + ) + or loaded_any + ) + if not loaded_any: + raise ValueError(f"No local expert accepted batched MoE tensor {wire_name!r}") + return True + + +def _load_checkpoint_weights(model: torch.nn.Module, weights: list[tuple[str, torch.Tensor]]) -> Any: + """Load ordinary HF tensors plus SkyRL's compact batched-MoE tensors.""" + params_dict: dict[str, torch.nn.Parameter] | None = None + ordinary_weights: list[tuple[str, torch.Tensor]] = [] + for name, weight in weights: + if name.startswith(SKYRL_BATCHED_MOE_FP8_PREFIX): + if params_dict is None: + params_dict = dict(model.named_parameters()) + _load_batched_moe_fp8_tensor(model, params_dict, name, weight) + else: + ordinary_weights.append((name, weight)) + if ordinary_weights: + return model.load_weights(weights=ordinary_weights) + return set() + + class _LoadWeightsProxy: """Wraps a model, overriding only ``load_weights``. @@ -140,7 +269,7 @@ def update_weights_ipc(self, update_info: dict) -> None: handles = pickle.loads(base64.b64decode(pickled)) device_index = torch.cuda.current_device() - physical_gpu_id = str(torch.cuda.get_device_properties(device_index).uuid) + physical_gpu_id = cuda_uuid_to_str(torch.cuda.get_device_properties(device_index).uuid) if physical_gpu_id not in handles: raise ValueError(f"IPC handle not found for GPU UUID {physical_gpu_id}. " f"Available: {list(handles)}") func, args = handles[physical_gpu_id] @@ -163,7 +292,7 @@ def update_weights_ipc(self, update_info: dict) -> None: model = self.model_runner.model with set_current_vllm_config(self.vllm_config), torch.device(self.device): if self._skyrl_is_checkpoint_format: - model.load_weights(weights=weights) + _load_checkpoint_weights(model, weights) # vLLM's load only updates the main model; the spec-decode (MTP/Eagle) # drafter is a separate module and must be reloaded from the same # checkpoint-format weights (see spec_decode_utils). @@ -224,7 +353,7 @@ def update_weights_nccl(self, update_info: dict) -> None: def _load_weights(weights): weights = list(weights) - loaded = model.load_weights(weights=weights) + loaded = _load_checkpoint_weights(model, weights) _reload_spec_decode_drafter(self.model_runner, weights) return loaded diff --git a/skyrl/backends/skyrl_train/inference_servers/utils.py b/skyrl/backends/skyrl_train/inference_servers/utils.py index 8abc5e0578..e3f379e2ed 100644 --- a/skyrl/backends/skyrl_train/inference_servers/utils.py +++ b/skyrl/backends/skyrl_train/inference_servers/utils.py @@ -16,6 +16,12 @@ # validates below. This must run on the driver before the WeightTransferConfig # is constructed; the import above already triggers it, this is just explicit. from skyrl.backends.skyrl_train.weight_sync import get_transfer_strategy +from skyrl.backends.skyrl_train.weight_sync.fp8 import ( + BLOCKWISE_FP8, + get_serialized_fp8_quantization_config, + registered_fp8_spec_names, + resolve_fp8_spec, +) from skyrl.backends.skyrl_train.weight_sync.sharded_rdt import ( rdt_vllm_register, # noqa: F401,E402 ) @@ -28,6 +34,84 @@ logger = logging.getLogger(__name__) +def _serialized_fp8_ignored_layers(model_path: Optional[str]) -> list[str]: + if not model_path: + raise ValueError("A model path is required when FP8 weight sync is enabled") + try: + from transformers import AutoConfig + + hf_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) + except Exception as exc: + raise RuntimeError( + "Could not inspect the model config required to derive FP8 ignored layers: " f"model_path={model_path!r}" + ) from exc + spec = resolve_fp8_spec(hf_config) + if spec is None: + raise ValueError( + "FP8 weight sync has no registered model spec for this checkpoint layout " + f"(registered specs: {', '.join(registered_fp8_spec_names())}); model_path={model_path!r}" + ) + return spec.ignored_layers(hf_config) + + +def _set_or_validate(mapping: Dict[str, Any], key: str, expected: Any, *, context: str) -> None: + if key in mapping and mapping[key] != expected: + raise ValueError( + f"{context}.{key} must be {expected!r} when FP8 weight sync is enabled, " f"got {mapping[key]!r}" + ) + mapping[key] = copy.deepcopy(expected) + + +def _apply_serialized_fp8_weight_sync_defaults( + ie_cfg: InferenceEngineConfig, + engine_kwargs: Dict[str, Any], + model_path: Optional[str] = None, +) -> None: + """Configure vLLM for checkpoint-format blockwise FP8 weight reloads.""" + + mode = ie_cfg.fp8_weight_sync_mode + if mode is None: + return + if mode != BLOCKWISE_FP8: + raise ValueError(f"Unsupported fp8_weight_sync_mode={mode!r}. " f"Supported value: {BLOCKWISE_FP8!r}.") + + _set_or_validate(engine_kwargs, "quantization", "fp8", context="engine_init_kwargs") + # Build FP8 modules without a bootstrap checkpoint; the first full-weight + # sync replaces the dummy values. + _set_or_validate(engine_kwargs, "load_format", "dummy", context="engine_init_kwargs") + + hf_overrides_value = engine_kwargs.get("hf_overrides") + hf_overrides = {} if hf_overrides_value is None else copy.deepcopy(hf_overrides_value) + if not isinstance(hf_overrides, dict): + raise ValueError("engine_init_kwargs.hf_overrides must be a dict when FP8 weight sync is enabled") + + qcfg_value = hf_overrides.get("quantization_config") + qcfg = {} if qcfg_value is None else copy.deepcopy(qcfg_value) + if not isinstance(qcfg, dict): + raise ValueError( + "engine_init_kwargs.hf_overrides.quantization_config must be a dict " "when FP8 weight sync is enabled" + ) + + ignored_layers = _serialized_fp8_ignored_layers(model_path) + if ignored_layers: + logger.info( + "FP8 weight sync will leave %d vLLM modules unquantized " "to match the model's FP8 quantization spec.", + len(ignored_layers), + ) + + for key, value in get_serialized_fp8_quantization_config( + ignored_layers=ignored_layers, + ).items(): + _set_or_validate( + qcfg, + key, + value, + context="engine_init_kwargs.hf_overrides.quantization_config", + ) + hf_overrides["quantization_config"] = qcfg + engine_kwargs["hf_overrides"] = hf_overrides + + def _uses_lora_weight_sync(cfg: SkyRLTrainConfig) -> bool: """Return True when the trainer syncs LoRA adapters (not merged weights). @@ -174,6 +258,11 @@ def build_vllm_cli_args(cfg: SkyRLTrainConfig) -> Namespace: logger.info(f"vLLM speculative decoding enabled: speculative_config={spec_cfg}") engine_kwargs = get_config_as_dict(ie_cfg.engine_init_kwargs) + _apply_serialized_fp8_weight_sync_defaults( + ie_cfg, + engine_kwargs, + cfg.trainer.policy.model.path, + ) for key, value in engine_kwargs.items(): setattr(args, key, value) diff --git a/skyrl/backends/skyrl_train/weight_sync/base.py b/skyrl/backends/skyrl_train/weight_sync/base.py index fd8753c24b..6a73137cfa 100644 --- a/skyrl/backends/skyrl_train/weight_sync/base.py +++ b/skyrl/backends/skyrl_train/weight_sync/base.py @@ -2,7 +2,7 @@ from dataclasses import asdict, dataclass, field from functools import cached_property -from typing import Any, Dict, List +from typing import Any, Dict, Iterator, List import torch @@ -104,3 +104,53 @@ def total_numel(self) -> int: def total_size_bytes(self) -> int: """Calculate total memory footprint in bytes.""" return sum(t.numel() * t.element_size() for t in self.tensors) + + +def torch_dtype_name(dtype: torch.dtype) -> str: + """Return the dtype spelling expected by vLLM weight-transfer metadata.""" + return str(dtype).split(".")[-1] + + +def cuda_uuid_to_str(uuid: str | bytes) -> str: + """Normalize CUDA UUIDs identically on both sides of an IPC transfer.""" + return uuid.decode("ascii") if isinstance(uuid, bytes) else str(uuid) + + +def iter_single_dtype_chunks(chunk: WeightChunk) -> Iterator[WeightChunk]: + """Yield dtype-homogeneous subchunks in first-seen dtype order. + + CUDA IPC packs tensors into a typed buffer, while serialized FP8 chunks mix + FP8 weights, FP32 scales, and unquantized BF16 tensors. vLLM's NCCL path + byte-packs mixed dtypes and does not require this split. + """ + by_dtype: Dict[torch.dtype, Dict[str, list]] = {} + dtype_order: List[torch.dtype] = [] + + for name, tensor in zip(chunk.names, chunk.tensors): + dtype = tensor.dtype + if dtype not in by_dtype: + dtype_order.append(dtype) + by_dtype[dtype] = {"names": [], "dtypes": [], "shapes": [], "tensors": []} + group = by_dtype[dtype] + group["names"].append(name) + group["dtypes"].append(str(dtype)) + group["shapes"].append(list(tensor.shape)) + group["tensors"].append(tensor) + + for dtype in dtype_order: + group = by_dtype[dtype] + yield WeightChunk( + names=group["names"], + dtypes=group["dtypes"], + shapes=group["shapes"], + tensors=group["tensors"], + ) + + +def get_weight_chunk_metadata(chunk: WeightChunk) -> Dict[str, List]: + """Return vLLM metadata for the tensors in a transfer chunk.""" + return { + "names": list(chunk.names), + "dtype_names": [torch_dtype_name(tensor.dtype) for tensor in chunk.tensors], + "shapes": [list(tensor.shape) for tensor in chunk.tensors], + } diff --git a/skyrl/backends/skyrl_train/weight_sync/broadcast_strategy.py b/skyrl/backends/skyrl_train/weight_sync/broadcast_strategy.py index 7aa9c455c1..3cd917f5cb 100644 --- a/skyrl/backends/skyrl_train/weight_sync/broadcast_strategy.py +++ b/skyrl/backends/skyrl_train/weight_sync/broadcast_strategy.py @@ -18,7 +18,11 @@ import ray import torch -from skyrl.backends.skyrl_train.weight_sync.base import WeightChunk, WeightUpdateRequest +from skyrl.backends.skyrl_train.weight_sync.base import ( + WeightChunk, + WeightUpdateRequest, + get_weight_chunk_metadata, +) from skyrl.backends.skyrl_train.weight_sync.nccl_trainer_send import ( nccl_trainer_init, nccl_trainer_send_weights, @@ -130,21 +134,27 @@ async def send_chunks( self, chunks: Iterable[WeightChunk], weight_metadata: Optional[Dict[str, list]] = None, + derive_metadata_from_chunks: bool = False, **kwargs, ) -> None: """Send chunks via broadcast or vLLM native NCCL. Args: chunks: Iterable of WeightChunk objects to send. - weight_metadata: Pre-computed metadata dict with "names", "dtype_names", - "shapes". Avoids materializing all chunks to collect metadata. + weight_metadata: Complete metadata for the batched update path. + derive_metadata_from_chunks: Send each chunk with derived metadata. """ - await self._send_chunks_vllm_native(chunks, weight_metadata) + if derive_metadata_from_chunks: + if weight_metadata is not None: + raise ValueError("weight_metadata must be omitted when deriving metadata from chunks") + await self._send_serialized_fp8_chunks_vllm_native(chunks) + else: + await self._send_chunks_vllm_native(chunks, weight_metadata) async def _send_chunks_vllm_native( self, chunks: Iterable[WeightChunk], - weight_metadata: Optional[Dict[str, list]] = None, + weight_metadata: Optional[Dict[str, list]], ) -> None: """Batched path: one update_weights call + nccl_trainer_send_weights. @@ -153,10 +163,7 @@ async def _send_chunks_vllm_native( tensors to vLLM via the NCCL weight transfer engine. """ if weight_metadata is None: - raise ValueError( - "weight_metadata is required for vLLM native path. " - "Call weight_extractor.get_weight_metadata() and pass it to send_chunks." - ) + raise ValueError("weight_metadata is required unless derive_metadata_from_chunks=true") def weight_iterator() -> Iterator[Tuple[str, torch.Tensor]]: for chunk in chunks: @@ -178,7 +185,7 @@ def weight_iterator() -> Iterator[Tuple[str, torch.Tensor]]: update_info = dict(weight_metadata) update_task = asyncio.create_task(self._inference_client.update_weights_nccl(update_info)) - # Run in thread so the HTTP update_task can progress concurrently + # Run in a thread so the HTTP update task can progress concurrently. await asyncio.to_thread( nccl_trainer_send_weights, weight_iterator(), @@ -189,12 +196,52 @@ def weight_iterator() -> Iterator[Tuple[str, torch.Tensor]]: await self._inference_client.finish_weight_update() else: - # Non-rank-0 still needs to participate in the all-gather + # Non-rank-0 still needs to participate in extractor collectives. for _ in weight_iterator(): pass torch.distributed.barrier() + async def _send_serialized_fp8_chunks_vllm_native( + self, + chunks: Iterable[WeightChunk], + ) -> None: + """Send lazy mixed-dtype serialized-FP8 chunks through vLLM NCCL.""" + if torch.distributed.get_rank() == 0: + await self._inference_client.start_weight_update(is_checkpoint_format=True) + + for chunk in chunks: + if torch.distributed.get_rank() == 0: + await self._send_chunk_vllm_native(chunk) + + if torch.distributed.get_rank() == 0: + await self._inference_client.finish_weight_update() + + torch.distributed.barrier() + + async def _send_chunk_vllm_native(self, chunk: WeightChunk) -> None: + """Send one logical chunk as its own NCCL update round. + + Same wire protocol as the batched path (vendored + ``nccl_trainer_send_weights`` + ``BroadcastInitInfo.packed``), just one + round per chunk because serialized-FP8 names/shapes are only known once + the chunk is built. The update info carries only names/dtype_names/shapes: + vLLM 0.28.0 rejects ``packed`` there (it is fixed at init). The packed + producer linearizes by bytes, so mixed fp8/fp32/bf16 tensors in one + round are fine. + """ + update_info = get_weight_chunk_metadata(chunk) + update_task = asyncio.create_task(self._inference_client.update_weights_nccl(update_info)) + + # Let the receiver enter its collective while the trainer broadcasts. + await asyncio.to_thread( + nccl_trainer_send_weights, + iter(zip(chunk.names, chunk.tensors)), + self._model_update_group, + packed=self._init_info.packed, + ) + await update_task + def teardown(self) -> None: """Destroy the process group used for weight transfer.""" if self._model_update_group is not None and isinstance( diff --git a/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py b/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py index 18642cca92..5ec4f09f69 100644 --- a/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py +++ b/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py @@ -28,13 +28,18 @@ import torch from torch.multiprocessing.reductions import reduce_tensor -from skyrl.backends.skyrl_train.weight_sync.base import WeightChunk, WeightUpdateRequest +from skyrl.backends.skyrl_train.weight_sync.base import ( + WeightChunk, + WeightUpdateRequest, + cuda_uuid_to_str, + iter_single_dtype_chunks, + torch_dtype_name, +) from skyrl.backends.skyrl_train.weight_sync.transfer_strategy import ( WeightSyncInitInfo, WeightTransferSender, WeightTransferStrategy, ) -from skyrl.train.utils.utils import str_to_torch_dtype # IPC handle type: (rebuild_func, args) returned by reduce_tensor IpcHandle = Tuple[Callable[..., torch.Tensor], Tuple[Any, ...]] @@ -148,15 +153,15 @@ async def send_chunks( self, chunks: Iterable[WeightChunk], weight_metadata: Optional[Dict[str, list]] = None, + derive_metadata_from_chunks: bool = False, **kwargs, ) -> None: """Send chunks via CUDA IPC. Args: chunks: Iterable of WeightChunk objects to send. - weight_metadata: Unused for IPC (metadata is derived from chunks - directly to avoid ordering mismatches). Kept for interface - compatibility with the base class. + weight_metadata: Unused; IPC derives metadata from each tensor. + derive_metadata_from_chunks: Accepted for sender interface compatibility. """ await self._send_chunks_vllm_native(chunks, weight_metadata) @@ -186,77 +191,91 @@ async def _send_chunks_vllm_native( rank = torch.distributed.get_rank() world_size = torch.distributed.get_world_size() device = torch.cuda.current_device() - gpu_uuid = str(torch.cuda.get_device_properties(device).uuid) - dtype = str_to_torch_dtype(self._init_info.model_dtype_str) - dtype_name = self._init_info.model_dtype_str.split(".")[-1] - + gpu_uuid = cuda_uuid_to_str(torch.cuda.get_device_properties(device).uuid) if rank == 0: await self._inference_client.start_weight_update(is_checkpoint_format=True) torch.distributed.barrier() - for chunk in chunks: - # --- pack all tensors in this chunk into one contiguous buffer --- - # Chunk tensors share a single dtype by construction (see - # weight_extractor_utils.py), so offsets in element units are safe. - names: List[str] = [] - dtype_names: List[str] = [] - shapes: List[List[int]] = [] - sizes: List[int] = [] - - total_numel = sum(t.numel() for t in chunk.tensors) - packed_tensor = torch.empty( - total_numel, - device=device, - dtype=dtype, - requires_grad=False, - ) - - offset = 0 - for name, tensor, shape in zip(chunk.names, chunk.tensors, chunk.shapes): - size = tensor.numel() - packed_tensor[offset : offset + size].copy_(tensor.detach().reshape(-1)) - offset += size - names.append(name) - dtype_names.append(dtype_name) - shapes.append(list(shape) if not isinstance(shape, list) else shape) - sizes.append(size) - - # --- one IPC handle per rank for the packed buffer --- - ipc_handle: IpcHandle = reduce_tensor(packed_tensor) - local_handle_dict: Dict[str, IpcHandle] = {gpu_uuid: ipc_handle} - gathered: List[Optional[Dict[str, IpcHandle]]] = [None] * world_size - torch.distributed.all_gather_object(gathered, local_handle_dict) - - torch.distributed.barrier() - torch.cuda.synchronize() - - if rank == 0: - merged_handles: Dict[str, IpcHandle] = {} - for d in gathered: - if d is not None: - merged_handles.update(d) - - pickled = base64.b64encode(pickle.dumps(merged_handles)).decode("utf-8") - chunk_update_info: Dict[str, Any] = { - "names": names, - "dtype_names": dtype_names, - "shapes": shapes, - "sizes": sizes, - "ipc_handles_pickled": pickled, - } - await self._inference_client.update_weights_ipc(chunk_update_info) - - # Keep packed_tensor alive past the barrier so the receiver's - # rebuilt view has valid backing storage while it copies into - # the model. Post-barrier drops the local ref safely. - torch.distributed.barrier() - torch.cuda.ipc_collect() - torch.cuda.synchronize() + for logical_chunk in chunks: + for chunk in iter_single_dtype_chunks(logical_chunk): + await self._send_single_dtype_chunk_vllm_native( + chunk=chunk, + device=device, + gpu_uuid=gpu_uuid, + world_size=world_size, + rank=rank, + ) if rank == 0: await self._inference_client.finish_weight_update() torch.distributed.barrier() + async def _send_single_dtype_chunk_vllm_native( + self, + *, + chunk: WeightChunk, + device: int, + gpu_uuid: str, + world_size: int, + rank: int, + ) -> None: + dtype = chunk.tensors[0].dtype + dtype_name = torch_dtype_name(dtype) + if any(tensor.dtype != dtype for tensor in chunk.tensors): + raise ValueError("CUDA IPC packed chunks must contain a single tensor dtype") + + names: List[str] = [] + dtype_names: List[str] = [] + shapes: List[List[int]] = [] + sizes: List[int] = [] + + total_numel = sum(t.numel() for t in chunk.tensors) + packed_tensor = torch.empty( + total_numel, + device=device, + dtype=dtype, + requires_grad=False, + ) + + offset = 0 + for name, tensor in zip(chunk.names, chunk.tensors): + size = tensor.numel() + packed_tensor[offset : offset + size].copy_(tensor.detach().reshape(-1)) + offset += size + names.append(name) + dtype_names.append(dtype_name) + shapes.append(list(tensor.shape)) + sizes.append(size) + + ipc_handle: IpcHandle = reduce_tensor(packed_tensor) + local_handle_dict: Dict[str, IpcHandle] = {gpu_uuid: ipc_handle} + gathered: List[Optional[Dict[str, IpcHandle]]] = [None] * world_size + torch.distributed.all_gather_object(gathered, local_handle_dict) + + torch.distributed.barrier() + torch.cuda.synchronize() + + if rank == 0: + merged_handles: Dict[str, IpcHandle] = {} + for d in gathered: + if d is not None: + merged_handles.update(d) + + pickled = base64.b64encode(pickle.dumps(merged_handles)).decode("utf-8") + chunk_update_info: Dict[str, Any] = { + "names": names, + "dtype_names": dtype_names, + "shapes": shapes, + "sizes": sizes, + "ipc_handles_pickled": pickled, + } + await self._inference_client.update_weights_ipc(chunk_update_info) + + # Keep the backing tensor alive until the receiver copies from its IPC view. + torch.distributed.barrier() + torch.cuda.ipc_collect() + torch.cuda.synchronize() + def teardown(self) -> None: """No-op for CUDA IPC sender (no custom process group to clean up).""" pass diff --git a/skyrl/backends/skyrl_train/weight_sync/delta_strategy.py b/skyrl/backends/skyrl_train/weight_sync/delta_strategy.py index 4d4aae37ad..0dec7ff9d4 100644 --- a/skyrl/backends/skyrl_train/weight_sync/delta_strategy.py +++ b/skyrl/backends/skyrl_train/weight_sync/delta_strategy.py @@ -90,8 +90,17 @@ async def send_chunks( self, chunks: Iterable[WeightChunk], weight_metadata: Optional[Dict[str, list]] = None, + derive_metadata_from_chunks: bool = False, reset_prefix_cache: bool = False, + **kwargs, ) -> None: + if derive_metadata_from_chunks: + # Serialized-FP8 wire chunks carry marker names and scale tensors + # that the delta-checkpoint format cannot represent. + raise ValueError( + "Delta weight sync does not support serialized FP8 chunks; " + "disable fp8_weight_sync_mode or use a broadcast/CUDA-IPC strategy" + ) if torch.distributed.is_available() and torch.distributed.is_initialized(): rank = torch.distributed.get_rank() else: diff --git a/skyrl/backends/skyrl_train/weight_sync/fp8/__init__.py b/skyrl/backends/skyrl_train/weight_sync/fp8/__init__.py new file mode 100644 index 0000000000..918160ef94 --- /dev/null +++ b/skyrl/backends/skyrl_train/weight_sync/fp8/__init__.py @@ -0,0 +1,47 @@ +"""Serialized blockwise FP8 weight sync: quantization, wire format, model specs.""" + +from skyrl.backends.skyrl_train.weight_sync.fp8.models import ( + ModelFp8Spec, + MoeExpertSpec, + MoeProjection, + batched_moe_wire_targets, + register_fp8_spec, + registered_fp8_spec_names, + resolve_fp8_spec, +) +from skyrl.backends.skyrl_train.weight_sync.fp8.quantize import ( + batched_blockwise_cast_to_fp8, + blockwise_cast_to_fp8, + normalize_block_size, + use_power_2_scales_default, +) +from skyrl.backends.skyrl_train.weight_sync.fp8.vllm_format import ( + BLOCKWISE_FP8, + SKYRL_BATCHED_MOE_FP8_PREFIX, + SerializedFp8Config, + get_serialized_fp8_quantization_config, + iter_batched_moe_expert_fp8_tensors, + iter_serialized_fp8_tensors, + scale_name_for_weight, +) + +__all__ = [ + "BLOCKWISE_FP8", + "SKYRL_BATCHED_MOE_FP8_PREFIX", + "ModelFp8Spec", + "MoeExpertSpec", + "MoeProjection", + "SerializedFp8Config", + "batched_blockwise_cast_to_fp8", + "batched_moe_wire_targets", + "blockwise_cast_to_fp8", + "get_serialized_fp8_quantization_config", + "iter_batched_moe_expert_fp8_tensors", + "iter_serialized_fp8_tensors", + "normalize_block_size", + "register_fp8_spec", + "registered_fp8_spec_names", + "resolve_fp8_spec", + "scale_name_for_weight", + "use_power_2_scales_default", +] diff --git a/skyrl/backends/skyrl_train/weight_sync/fp8/models/README.md b/skyrl/backends/skyrl_train/weight_sync/fp8/models/README.md new file mode 100644 index 0000000000..809ebb7a55 --- /dev/null +++ b/skyrl/backends/skyrl_train/weight_sync/fp8/models/README.md @@ -0,0 +1,52 @@ +# Per-model specs for serialized FP8 weight sync + +Serialized FP8 weight sync needs four model-specific answers, grouped in a +`ModelFp8Spec` (`base.py`) and resolved once per checkpoint with +`resolve_fp8_spec(hf_config)`: + +| Field | Question it answers | +| --- | --- | +| `matches(hf_config)` | Does this spec support the checkpoint layout? | +| `should_quantize(name, shape)` | Should this exported HF weight be FP8 on the wire? (Linear weights yes; embeddings, norms, conv, router gates no.) | +| `ignored_layers(hf_config)` | Which vLLM module prefixes must stay unquantized? (Modules whose shards can't share a 128-block FP8 scheme.) | +| `moe_expert_spec(name)` | Is this a Megatron-Bridge *batched* expert tensor, and how does it map onto per-projection wire tensors? `None` for ordinary tensors. | + +Everything else — blockwise casting, wire naming, the vLLM quantization +config, the receiver's fused-MoE loading — is generic and lives outside this +package. The receiver-side table (which fused vLLM parameter each expert +projection loads into) is **derived** from the specs via +`batched_moe_wire_targets()`, so the mapping is declared exactly once. + +## Adding a new model + +1. Create `models/.py`. Implement the four callables (plain + functions; see `qwen35.py` — suffix tables usually suffice for + `should_quantize`). +2. If the model has routed MoE experts exported as batched 3D tensors, + declare one `MoeProjection(hf_name, vllm_param, shard_id)` per projection + and return `MoeExpertSpec(experts_base, projections, split_dim)` from + `moe_expert_spec` — `split_dim` is the dimension that concatenates fused + projections (e.g. `gate_up_proj` splits in half along dim 1); use `None` + when the tensor is a single projection. +3. Register it at module bottom and import the module in + `models/__init__.py` (the import is what registers the spec): + + ```python + MYMODEL_FP8_SPEC = register_fp8_spec( + ModelFp8Spec( + name="mymodel", + matches=is_mymodel_config, + should_quantize=is_quantizable_weight_shape, + ignored_layers=get_mymodel_fp8_ignored_layers, + moe_expert_spec=batched_moe_expert_spec, + moe_projections=(_MOE_GATE, _MOE_UP, _MOE_DOWN), + ) + ) + ``` + +4. Add tests mirroring `tests/backends/skyrl_train/weight_sync/ + test_serialized_fp8.py` (quantize filter, ignored layers, MoE mapping) + and, for real coverage, an FP8 row in the GPU CI logprobs-roundtrip test. + +No generic file changes are needed; checkpoints that match no registered +spec are rejected with the list of registered spec names. diff --git a/skyrl/backends/skyrl_train/weight_sync/fp8/models/__init__.py b/skyrl/backends/skyrl_train/weight_sync/fp8/models/__init__.py new file mode 100644 index 0000000000..c5b01b12cf --- /dev/null +++ b/skyrl/backends/skyrl_train/weight_sync/fp8/models/__init__.py @@ -0,0 +1,25 @@ +"""Per-model quantization specs for serialized FP8 weight sync.""" + +from skyrl.backends.skyrl_train.weight_sync.fp8.models.base import ( + ModelFp8Spec, + MoeExpertSpec, + MoeProjection, + batched_moe_wire_targets, + register_fp8_spec, + registered_fp8_spec_names, + resolve_fp8_spec, +) + +# Importing a model module registers its spec. +from skyrl.backends.skyrl_train.weight_sync.fp8.models.qwen35 import QWEN35_FP8_SPEC + +__all__ = [ + "ModelFp8Spec", + "MoeExpertSpec", + "MoeProjection", + "QWEN35_FP8_SPEC", + "batched_moe_wire_targets", + "register_fp8_spec", + "registered_fp8_spec_names", + "resolve_fp8_spec", +] diff --git a/skyrl/backends/skyrl_train/weight_sync/fp8/models/base.py b/skyrl/backends/skyrl_train/weight_sync/fp8/models/base.py new file mode 100644 index 0000000000..bc81f49597 --- /dev/null +++ b/skyrl/backends/skyrl_train/weight_sync/fp8/models/base.py @@ -0,0 +1,101 @@ +"""Generic per-model spec for serialized blockwise FP8 weight sync. + +``ModelFp8Spec`` groups everything the sync path must know about one model +family — which HF configs it matches, which weights quantize, which vLLM +modules stay unquantized, and how Megatron-Bridge's batched expert tensors +map onto wire projections. ``resolve_fp8_spec`` selects the spec for a +checkpoint; unsupported layouts resolve to ``None`` and callers reject them +explicitly. The vLLM-side fused-loader targets are derived from the same +projections via ``batched_moe_wire_targets``, so sender and receiver share +one source of truth instead of hardcoding the mapping twice. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Optional, Sequence + + +@dataclass(frozen=True) +class MoeProjection: + """One routed-expert projection and its vLLM fused-loader target.""" + + hf_name: str # projection name on the wire / in HF checkpoints, e.g. "gate_proj" + vllm_param: str # fused vLLM parameter it loads into, e.g. "w13_weight" + shard_id: str # FusedMoE weight_loader shard id, e.g. "w1" + + +@dataclass(frozen=True) +class MoeExpertSpec: + """A Megatron-Bridge batched expert tensor mapped onto wire projections. + + ``split_dim`` names the tensor dimension that concatenates the + projections (split evenly, in order); ``None`` means the tensor is a + single projection. + """ + + experts_base: str # checkpoint prefix ending in the experts module + projections: tuple[MoeProjection, ...] + split_dim: Optional[int] = None + + +@dataclass(frozen=True) +class ModelFp8Spec: + """Per-model policy for serialized blockwise FP8 weight sync.""" + + name: str + # hf_config -> does this spec support the checkpoint layout? + matches: Callable[[Any], bool] + # (hf_name, shape) -> serialize this exported weight as FP8? + should_quantize: Callable[[str, Sequence[int]], bool] + # hf_config -> vLLM module prefixes that must stay unquantized + ignored_layers: Callable[[Any], list[str]] + # batched expert tensor name -> MoeExpertSpec, or None if not one + moe_expert_spec: Callable[[str], Optional[MoeExpertSpec]] + # module segment holding routed experts in vLLM parameter names + moe_module: str = "experts" + # every projection the model emits, for receiver-side target derivation + moe_projections: tuple[MoeProjection, ...] = field(default=()) + + +_REGISTRY: list[ModelFp8Spec] = [] + + +def register_fp8_spec(spec: ModelFp8Spec) -> ModelFp8Spec: + """Register a model spec for ``resolve_fp8_spec`` lookup.""" + + if any(existing.name == spec.name for existing in _REGISTRY): + raise ValueError(f"An FP8 model spec named {spec.name!r} is already registered") + _REGISTRY.append(spec) + return spec + + +def registered_fp8_spec_names() -> tuple[str, ...]: + return tuple(spec.name for spec in _REGISTRY) + + +def resolve_fp8_spec(hf_config: Any) -> Optional[ModelFp8Spec]: + """Return the registered spec matching an HF config, or ``None``.""" + + for spec in _REGISTRY: + if spec.matches(hf_config): + return spec + return None + + +def batched_moe_wire_targets() -> dict[str, tuple[str, str]]: + """Receiver mapping: checkpoint suffix -> (fused vLLM suffix, shard id). + + Derived from every registered spec's projections so the vLLM worker + extension never re-encodes per-model fused-loader knowledge. + """ + + targets: dict[str, tuple[str, str]] = {} + for spec in _REGISTRY: + for proj in spec.moe_projections: + for weight_suffix, param_suffix in ((".weight", ""), (".weight_scale_inv", "_scale_inv")): + key = f".{spec.moe_module}.{proj.hf_name}{weight_suffix}" + value = (f".{spec.moe_module}.{proj.vllm_param}{param_suffix}", proj.shard_id) + if targets.setdefault(key, value) != value: + raise ValueError(f"Conflicting batched MoE wire target registered for suffix {key!r}") + return targets diff --git a/skyrl/backends/skyrl_train/weight_sync/fp8/models/qwen35.py b/skyrl/backends/skyrl_train/weight_sync/fp8/models/qwen35.py new file mode 100644 index 0000000000..5ac4564a01 --- /dev/null +++ b/skyrl/backends/skyrl_train/weight_sync/fp8/models/qwen35.py @@ -0,0 +1,144 @@ +"""Qwen3.5 ``ModelFp8Spec`` for serialized blockwise FP8 weight sync.""" + +from __future__ import annotations + +from typing import Any, Optional, Sequence + +from skyrl.backends.skyrl_train.weight_sync.fp8.models.base import ( + ModelFp8Spec, + MoeExpertSpec, + MoeProjection, + register_fp8_spec, +) + +_QWEN35_FP8_WEIGHT_SUFFIXES = ( + ".self_attn.q_proj.weight", + ".self_attn.k_proj.weight", + ".self_attn.v_proj.weight", + ".self_attn.o_proj.weight", + ".mlp.gate_proj.weight", + ".mlp.up_proj.weight", + ".mlp.down_proj.weight", + ".linear_attn.in_proj_qkv.weight", + ".linear_attn.in_proj_z.weight", + ".linear_attn.out_proj.weight", + # Shared-expert linears use FP8; router and shared-expert gates remain BF16. + ".mlp.shared_expert.gate_proj.weight", + ".mlp.shared_expert.up_proj.weight", + ".mlp.shared_expert.down_proj.weight", +) +# Megatron Bridge exports routed experts in batched tensors. Keep the expert +# dimension intact on the wire so the receiver can use vLLM's fused MoE loader. +_QWEN35_MOE_GATE_UP_SUFFIX = ".mlp.experts.gate_up_proj" +_QWEN35_MOE_DOWN_SUFFIX = ".mlp.experts.down_proj" +_QWEN35_UNQUANTIZED_LINEAR_SUFFIXES = ( + ".in_proj_b", + ".in_proj_a", +) +_QWEN35_LINEAR_ATTN_PREFIX_TEMPLATES = ( + "{model_prefix}.layers.{layer_idx}.linear_attn", + "{model_prefix}.language_model.layers.{layer_idx}.linear_attn", +) +# Vision attention output and both vision MLP linears carry dims (e.g. 4304) +# that stop being 128-divisible once vLLM TP-shards them, so all three must be +# ignored for the engine to build at TP>1. The weight-sync spec keeps the +# vision tower BF16 regardless. +_QWEN35_VISION_BLOCK_PREFIX_TEMPLATES = ( + "{model_prefix}.visual.blocks.{block_idx}.attn.proj", + "{model_prefix}.visual.blocks.{block_idx}.mlp.linear_fc1", + "{model_prefix}.visual.blocks.{block_idx}.mlp.linear_fc2", +) + +_MOE_GATE = MoeProjection(hf_name="gate_proj", vllm_param="w13_weight", shard_id="w1") +_MOE_UP = MoeProjection(hf_name="up_proj", vllm_param="w13_weight", shard_id="w3") +_MOE_DOWN = MoeProjection(hf_name="down_proj", vllm_param="w2_weight", shard_id="w2") + + +def is_qwen35_config(hf_config: Any) -> bool: + """Return whether an HF config uses the supported Qwen3.5 text layout.""" + + text_config = getattr(hf_config, "text_config", None) or getattr(hf_config, "language_config", None) or hf_config + model_type = str(getattr(text_config, "model_type", "") or getattr(hf_config, "model_type", "")) + return model_type in {"qwen3_5", "qwen3_5_text", "qwen3_5_moe", "qwen3_5_moe_text"} + + +def get_qwen35_fp8_ignored_layers(hf_config: Any, model_prefix: str = "model") -> list[str]: + """Return Qwen3.5 vLLM module prefixes excluded from serialized FP8. + + Serialized sync excludes GDN ``in_proj_a`` and ``in_proj_b``. vLLM requires + every shard of the fused module to share a quantization scheme, so both + prefixes are ignored for text-only and conditional-generation checkpoints. + """ + + text_config = getattr(hf_config, "text_config", None) or getattr(hf_config, "language_config", None) or hf_config + if not is_qwen35_config(hf_config): + return [] + + layer_types = list(getattr(text_config, "layer_types", []) or []) + ignored: list[str] = [] + for layer_idx, layer_type in enumerate(layer_types): + if layer_type != "linear_attention": + continue + layer_prefixes = [] + for template in _QWEN35_LINEAR_ATTN_PREFIX_TEMPLATES: + prefix = template.format(model_prefix=model_prefix, layer_idx=layer_idx) + if prefix not in layer_prefixes: + layer_prefixes.append(prefix) + + for layer_prefix in layer_prefixes: + for suffix in _QWEN35_UNQUANTIZED_LINEAR_SUFFIXES: + ignored.append(f"{layer_prefix}{suffix}") + + # vLLM instantiates the vision tower even for text-only runs + # (language_model_only only affects multimodal weight loading), and ignore + # matching requires each block's exact module prefix. + vision_config = getattr(hf_config, "vision_config", None) or getattr(hf_config, "visual_config", None) + vision_depth = 0 + if vision_config is not None: + for attr in ("depth", "num_hidden_layers", "num_layers"): + value = getattr(vision_config, attr, None) + if isinstance(value, int) and value > 0: + vision_depth = value + break + for block_idx in range(vision_depth): + for template in _QWEN35_VISION_BLOCK_PREFIX_TEMPLATES: + ignored.append(template.format(model_prefix=model_prefix, block_idx=block_idx)) + return ignored + + +def is_quantizable_weight_shape(name: str, shape: Sequence[int]) -> bool: + """Return whether an exported HF weight should be serialized as FP8. + + vLLM's FP8 config applies to Linear modules. HF checkpoints also contain 2D + embedding/output weights, so keep known non-Linear weight tables unquantized. + """ + + if not name.endswith(".weight") or len(shape) != 2: + return False + return name.endswith(_QWEN35_FP8_WEIGHT_SUFFIXES) + + +def batched_moe_expert_spec(name: str) -> Optional[MoeExpertSpec]: + """Map a Megatron Bridge batched Qwen3.5 MoE tensor name onto projections.""" + + if name.endswith(_QWEN35_MOE_GATE_UP_SUFFIX): + return MoeExpertSpec( + experts_base=name[: -len(".gate_up_proj")], + projections=(_MOE_GATE, _MOE_UP), + split_dim=1, + ) + if name.endswith(_QWEN35_MOE_DOWN_SUFFIX): + return MoeExpertSpec(experts_base=name[: -len(".down_proj")], projections=(_MOE_DOWN,)) + return None + + +QWEN35_FP8_SPEC = register_fp8_spec( + ModelFp8Spec( + name="qwen3.5", + matches=is_qwen35_config, + should_quantize=is_quantizable_weight_shape, + ignored_layers=get_qwen35_fp8_ignored_layers, + moe_expert_spec=batched_moe_expert_spec, + moe_projections=(_MOE_GATE, _MOE_UP, _MOE_DOWN), + ) +) diff --git a/skyrl/backends/skyrl_train/weight_sync/fp8/quantize.py b/skyrl/backends/skyrl_train/weight_sync/fp8/quantize.py new file mode 100644 index 0000000000..ec8531d34a --- /dev/null +++ b/skyrl/backends/skyrl_train/weight_sync/fp8/quantize.py @@ -0,0 +1,136 @@ +"""Blockwise FP8 quantization kernels and scale-mode helpers.""" + +from __future__ import annotations + +import os +from operator import index +from typing import Sequence + +import torch + + +def use_power_2_scales_default() -> bool: + """Return whether rollout weights use power-of-two block scales. + + The setting must match Transformer Engine. Hopper defaults to FP32 scales; + Blackwell launchers select power-of-two scales by setting + ``NVTE_FP8_BLOCK_SCALING_FP32_SCALES=0``. + """ + + scale_mode = os.getenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "1") + if scale_mode not in {"0", "1"}: + raise ValueError( + "NVTE_FP8_BLOCK_SCALING_FP32_SCALES must be '0' (power-of-2) " f"or '1' (FP32 scales), got {scale_mode!r}" + ) + return scale_mode == "0" + + +def normalize_block_size(block_size: Sequence[int]) -> tuple[int, int]: + try: + raw_values = tuple(block_size) + if any(isinstance(value, bool) for value in raw_values): + raise TypeError + values = tuple(index(value) for value in raw_values) + except (TypeError, ValueError) as exc: + raise ValueError(f"weight_block_size must contain exactly two positive integers, got {block_size!r}") from exc + if len(values) != 2 or any(value <= 0 for value in values): + raise ValueError(f"weight_block_size must contain exactly two positive integers, got {block_size!r}") + return values + + +def blockwise_cast_to_fp8( + weight: torch.Tensor, + block_size: Sequence[int], + power_2_scale: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a 2D tensor to vLLM's blockwise E4M3 checkpoint format. + + Returns ``weight_scale_inv`` such that + ``weight ~= qweight.float() * scale``. Power-of-two mode rounds scales up + to match Transformer Engine's UE8M0 rule. + """ + + if weight.ndim != 2: + raise ValueError(f"Blockwise FP8 expects a 2D tensor, got shape={tuple(weight.shape)}") + + block_m, block_n = normalize_block_size(block_size) + rows, cols = weight.shape + padded_rows = ((rows + block_m - 1) // block_m) * block_m + padded_cols = ((cols + block_n - 1) // block_n) * block_n + + fp8_info = torch.finfo(torch.float8_e4m3fn) + weight_fp32 = weight.detach().to(torch.float32) + if padded_rows != rows or padded_cols != cols: + padded = weight_fp32.new_zeros((padded_rows, padded_cols)) + padded[:rows, :cols].copy_(weight_fp32) + else: + padded = weight_fp32 + + blocks = padded.view(padded_rows // block_m, block_m, padded_cols // block_n, block_n) + blocks = blocks.permute(0, 2, 1, 3) + # Nonzero floor keeps all-zero blocks from degenerating the scale. + scale = blocks.abs().amax(dim=(2, 3)).clamp(min=1e-10) / fp8_info.max + if power_2_scale: + # Rounding up preserves range and matches TE's power-of-two scale rule. + scale = torch.pow(2.0, torch.ceil(torch.log2(scale))) + q_blocks = (blocks / scale[:, :, None, None]).clamp(min=fp8_info.min, max=fp8_info.max) + q_blocks = q_blocks.to(torch.float8_e4m3fn) + q_padded = q_blocks.permute(0, 2, 1, 3).contiguous().view(padded_rows, padded_cols) + q_weight = q_padded[:rows, :cols].contiguous() + return q_weight, scale.to(torch.float32).contiguous() + + +def batched_blockwise_cast_to_fp8( + weight: torch.Tensor, + block_size: Sequence[int], + power_2_scale: bool = False, + expert_batch_size: int = 8, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a 3D ``[experts, rows, cols]`` tensor blockwise. + + Quantizing several experts per operation avoids launching the full 2D + conversion pipeline once per expert, while bounded batches limit peak FP32 + workspace. + """ + + if weight.ndim != 3: + raise ValueError(f"Batched blockwise FP8 expects a 3D tensor, got shape={tuple(weight.shape)}") + if isinstance(expert_batch_size, bool) or not isinstance(expert_batch_size, int) or expert_batch_size <= 0: + raise ValueError(f"expert_batch_size must be a positive integer, got {expert_batch_size!r}") + + block_m, block_n = normalize_block_size(block_size) + num_experts, rows, cols = weight.shape + padded_rows = ((rows + block_m - 1) // block_m) * block_m + padded_cols = ((cols + block_n - 1) // block_n) * block_n + row_blocks = padded_rows // block_m + col_blocks = padded_cols // block_n + + fp8_info = torch.finfo(torch.float8_e4m3fn) + q_weight = torch.empty(weight.shape, dtype=torch.float8_e4m3fn, device=weight.device) + scales = torch.empty( + (num_experts, row_blocks, col_blocks), + dtype=torch.float32, + device=weight.device, + ) + + for start in range(0, num_experts, expert_batch_size): + end = min(start + expert_batch_size, num_experts) + weight_fp32 = weight[start:end].detach().to(torch.float32).contiguous() + if padded_rows != rows or padded_cols != cols: + padded = weight_fp32.new_zeros((end - start, padded_rows, padded_cols)) + padded[:, :rows, :cols].copy_(weight_fp32) + else: + padded = weight_fp32 + + blocks = padded.view(end - start, row_blocks, block_m, col_blocks, block_n) + blocks = blocks.permute(0, 1, 3, 2, 4) + scale = blocks.abs().amax(dim=(3, 4)).clamp(min=1e-10) / fp8_info.max + if power_2_scale: + scale = torch.pow(2.0, torch.ceil(torch.log2(scale))) + q_blocks = (blocks / scale[:, :, :, None, None]).clamp(min=fp8_info.min, max=fp8_info.max) + q_blocks = q_blocks.to(torch.float8_e4m3fn) + q_padded = q_blocks.permute(0, 1, 3, 2, 4).contiguous().view(end - start, padded_rows, padded_cols) + q_weight[start:end].copy_(q_padded[:, :rows, :cols]) + scales[start:end].copy_(scale) + + return q_weight, scales diff --git a/skyrl/backends/skyrl_train/weight_sync/fp8/vllm_format.py b/skyrl/backends/skyrl_train/weight_sync/fp8/vllm_format.py new file mode 100644 index 0000000000..c4b89ad7d4 --- /dev/null +++ b/skyrl/backends/skyrl_train/weight_sync/fp8/vllm_format.py @@ -0,0 +1,136 @@ +"""Serialized wire format for vLLM blockwise FP8 checkpoints.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Iterator, Sequence + +import torch + +from skyrl.backends.skyrl_train.weight_sync.fp8.models.base import ModelFp8Spec +from skyrl.backends.skyrl_train.weight_sync.fp8.quantize import ( + batched_blockwise_cast_to_fp8, + blockwise_cast_to_fp8, + normalize_block_size, + use_power_2_scales_default, +) + +BLOCKWISE_FP8 = "blockwise" +# Internal wire-format marker for Qwen3.5 MoE tensors that remain batched over +# experts. The receiver strips this marker and routes the tensor directly to +# vLLM's fused-MoE parameter loader instead of the ordinary HF-name loader. +SKYRL_BATCHED_MOE_FP8_PREFIX = "__skyrl_batched_moe_fp8__:" + + +@dataclass(frozen=True) +class SerializedFp8Config: + """Configuration for serialized FP8 rollout weight sync. + + ``spec`` is the per-model quantization policy, resolved once from the HF + config via ``resolve_fp8_spec``; the tensor iterators require it. + """ + + weight_block_size: tuple[int, int] = (128, 128) + power_2_scale: bool = field(default_factory=use_power_2_scales_default) + spec: ModelFp8Spec | None = None + + def __post_init__(self) -> None: + object.__setattr__(self, "weight_block_size", normalize_block_size(self.weight_block_size)) + if type(self.power_2_scale) is not bool: + raise ValueError(f"power_2_scale must be a bool, got {self.power_2_scale!r}") + + def require_spec(self) -> ModelFp8Spec: + if self.spec is None: + raise ValueError( + "SerializedFp8Config.spec is not set; resolve the model spec with " + "resolve_fp8_spec(hf_config) before serializing weights" + ) + return self.spec + + +def get_serialized_fp8_quantization_config( + weight_block_size: Sequence[int] = (128, 128), + ignored_layers: Sequence[str] | None = None, +) -> dict: + """Return vLLM's Hugging Face quantization config for serialized FP8.""" + + block_m, block_n = normalize_block_size(weight_block_size) + qconfig = { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [block_m, block_n], + } + if ignored_layers: + qconfig["ignored_layers"] = list(ignored_layers) + return qconfig + + +def scale_name_for_weight(name: str) -> str: + if not name.endswith(".weight"): + raise ValueError(f"FP8 scale can only be derived from .weight tensors: {name}") + return name[: -len(".weight")] + ".weight_scale_inv" + + +def iter_batched_moe_expert_fp8_tensors( + name: str, + tensor: torch.Tensor, + config: SerializedFp8Config, +) -> Iterator[tuple[str, torch.Tensor]]: + """Convert a batched expert tensor without expanding expert names. + + The old wire format emitted one weight and one scale tensor for every + expert/projection pair. Keeping the expert dimension intact reduces each + routed MoE layer from ``6 * num_experts`` tensors to six and lets vLLM use + its fused 3D loader. + """ + moe_spec = config.require_spec().moe_expert_spec(name) + if moe_spec is None: + raise ValueError(f"Not a batched MoE expert tensor: {name}") + if tensor.ndim != 3: + raise ValueError(f"Batched MoE expert tensor must be 3D, got shape={tuple(tensor.shape)}") + if moe_spec.split_dim is not None: + num_projections = len(moe_spec.projections) + if tensor.shape[moe_spec.split_dim] % num_projections != 0: + raise ValueError( + f"Batched MoE tensor dim {moe_spec.split_dim} must split evenly across " + f"{num_projections} projections, got shape={tuple(tensor.shape)}" + ) + projection_tensors = torch.chunk(tensor, num_projections, dim=moe_spec.split_dim) + else: + projection_tensors = (tensor,) + + for proj, projection_tensor in zip(moe_spec.projections, projection_tensors): + q_weight, scale = batched_blockwise_cast_to_fp8( + projection_tensor, + config.weight_block_size, + config.power_2_scale, + ) + weight_name = f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{moe_spec.experts_base}.{proj.hf_name}.weight" + yield weight_name, q_weight + yield scale_name_for_weight(weight_name), scale + + +def iter_serialized_fp8_tensors( + name: str, + tensor: torch.Tensor, + target_dtype: torch.dtype, + config: SerializedFp8Config, +) -> Iterator[tuple[str, torch.Tensor]]: + """Yield vLLM checkpoint tensors for one Megatron-exported weight.""" + + spec = config.require_spec() + if spec.moe_expert_spec(name) is not None: + yield from iter_batched_moe_expert_fp8_tensors(name, tensor, config) + return + + if tensor.ndim == 2 and spec.should_quantize(name, tuple(tensor.shape)): + q_weight, scale = blockwise_cast_to_fp8( + tensor, + config.weight_block_size, + config.power_2_scale, + ) + yield name, q_weight + yield scale_name_for_weight(name), scale + return + + yield name, tensor.to(dtype=target_dtype) diff --git a/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_strategy.py b/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_strategy.py index 9a84e1be9a..f94475b370 100644 --- a/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_strategy.py +++ b/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_strategy.py @@ -96,6 +96,16 @@ async def send( prefix cache, so the worker resets it (``handles_prefix_cache_reset`` stays False). """ + if weight_extractor.derives_metadata_from_chunks: + # Serialized FP8 splits each tensor into a quantized payload plus + # scales; the weight sources here publish whole bridge tensors, so + # the consumers would pull dequantized weights into modules vLLM + # built for FP8. Config validation rejects the combination too; this + # covers senders built outside the training entrypoint. + raise ValueError( + "sharded_rdt does not support serialized FP8 chunks; " + "disable fp8_weight_sync_mode or use the nccl weight sync backend" + ) del dtype, kwargs await self._sender.send(weight_extractor) diff --git a/skyrl/backends/skyrl_train/weight_sync/transfer_strategy.py b/skyrl/backends/skyrl_train/weight_sync/transfer_strategy.py index 3dbe2de01d..65adb5aeae 100644 --- a/skyrl/backends/skyrl_train/weight_sync/transfer_strategy.py +++ b/skyrl/backends/skyrl_train/weight_sync/transfer_strategy.py @@ -69,14 +69,20 @@ async def send( is what keeps ``get_weight_metadata`` (a whole-model gather on the Megatron extractor) off their critical path entirely. + An extractor whose metadata depends on chunk contents (serialized FP8) + reports ``derives_metadata_from_chunks``; for those, precomputing is not + just wasteful but unsupported, so the flag is forwarded instead. + Args: weight_extractor: The worker's extractor, already built. dtype: Inference dtype to convert to. **kwargs: Forwarded to :meth:`send_chunks`. """ + derive_metadata_from_chunks = weight_extractor.derives_metadata_from_chunks await self.send_chunks( weight_extractor.extract_weights(dtype), - weight_metadata=weight_extractor.get_weight_metadata(dtype), + weight_metadata=(None if derive_metadata_from_chunks else weight_extractor.get_weight_metadata(dtype)), + derive_metadata_from_chunks=derive_metadata_from_chunks, **kwargs, ) @@ -85,6 +91,7 @@ async def send_chunks( self, chunks: Iterable[WeightChunk], weight_metadata: Optional[Dict[str, list]] = None, + derive_metadata_from_chunks: bool = False, **kwargs, ) -> None: """Send chunks using this transfer strategy. @@ -95,8 +102,7 @@ async def send_chunks( Args: chunks: Iterable of WeightChunk objects to send. weight_metadata: Optional pre-computed metadata (names, dtype_names, shapes). - When provided, allows the sender to avoid materializing all chunks - to collect metadata upfront. + derive_metadata_from_chunks: Derive metadata from each transferred chunk. """ ... diff --git a/skyrl/backends/skyrl_train/weight_sync/weight_extractor.py b/skyrl/backends/skyrl_train/weight_sync/weight_extractor.py index 055a7903ce..eaac15d7c1 100644 --- a/skyrl/backends/skyrl_train/weight_sync/weight_extractor.py +++ b/skyrl/backends/skyrl_train/weight_sync/weight_extractor.py @@ -33,6 +33,17 @@ def extract_weights(self, dtype: torch.dtype) -> Iterator[WeightChunk]: """ ... + @property + def derives_metadata_from_chunks(self) -> bool: + """Whether metadata must be derived from the transferred chunks. + + True when metadata depends on chunk contents (e.g. serialized FP8, where + each tensor expands into quantized payload + scales), so + :meth:`get_weight_metadata` cannot describe the stream ahead of time. + Senders consult this to decide whether to precompute metadata. + """ + return False + @abstractmethod def get_weight_metadata(self, dtype: torch.dtype) -> Dict[str, List]: """Return weight metadata without materializing tensors. diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py index 4a1b4b4925..88614d18a3 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_model_wrapper.py @@ -32,7 +32,9 @@ vocab_parallel_entropy, vocab_parallel_entropy_packed_sequences, ) -from skyrl.backends.skyrl_train.distributed.megatron.packing_utils import is_fp8_enabled +from skyrl.backends.skyrl_train.distributed.megatron.quantization_utils import ( + is_fp8_enabled, +) from skyrl.backends.skyrl_train.distributed.megatron.token_metadata import ( build_token_metadata_layout, ) @@ -370,6 +372,7 @@ def forward_step(batch_iter, model): model_config = get_model_config(model) fp8_enabled = is_fp8_enabled(getattr(model_config, "fp8", None)) + fp8_recipe = getattr(model_config, "fp8_recipe", None) rollout_expert_indices = batch.pop("rollout_expert_indices", None) router_padding_mask = batch.pop("router_padding_mask", None) @@ -393,6 +396,7 @@ def forward_step(batch_iter, model): pre_process=mpu.is_pipeline_first_stage(ignore_virtual=True) or self.is_vlm, sub_seq_lengths=sub_seq_lengths, fp8_enabled=fp8_enabled, + fp8_recipe=fp8_recipe, ) batch["packed_seq_params"] = packed_seq_params batch["packed_targets"] = _build_packed_targets( @@ -407,6 +411,7 @@ def forward_step(batch_iter, model): position_ids, pre_process=mpu.is_pipeline_first_stage(ignore_virtual=True) or self.is_vlm, fp8_enabled=fp8_enabled, + fp8_recipe=fp8_recipe, ) packed_seq_params = None # Qwen-style VLMs recompute 3D mRoPE positions internally from @@ -421,6 +426,7 @@ def forward_step(batch_iter, model): attention_mask.device, packed=packed_seq_params is not None, fp8_enabled=fp8_enabled, + fp8_recipe=fp8_recipe, ) model_replay_kwargs = {} @@ -972,6 +978,7 @@ def forward_step(batch_iter, model): model_config = get_model_config(model) fp8_enabled = is_fp8_enabled(getattr(model_config, "fp8", None)) + fp8_recipe = getattr(model_config, "fp8_recipe", None) rollout_expert_indices = batch.pop("rollout_expert_indices", None) router_padding_mask = batch.pop("router_padding_mask", None) @@ -1003,6 +1010,7 @@ def forward_step(batch_iter, model): pre_process=mpu.is_pipeline_first_stage(ignore_virtual=True) or self.is_vlm, sub_seq_lengths=sub_seq_lengths, fp8_enabled=fp8_enabled, + fp8_recipe=fp8_recipe, ) batch["packed_seq_params"] = packed_seq_params batch["packed_targets"] = _build_packed_targets( @@ -1028,6 +1036,7 @@ def forward_step(batch_iter, model): position_ids, pre_process=mpu.is_pipeline_first_stage(ignore_virtual=True) or self.is_vlm, fp8_enabled=fp8_enabled, + fp8_recipe=fp8_recipe, ) packed_seq_params = None # Qwen-style VLMs recompute 3D mRoPE positions internally from @@ -1044,6 +1053,7 @@ def forward_step(batch_iter, model): attention_mask.device, packed=packed_seq_params is not None, fp8_enabled=fp8_enabled, + fp8_recipe=fp8_recipe, ) model_replay_kwargs = {} diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index b6e8cde782..f2aaed8765 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -1,3 +1,4 @@ +import gc import os import shutil from collections import defaultdict @@ -39,6 +40,10 @@ get_megatron_optimizer_param_scheduler, init_megatron_optim_config, ) +from skyrl.backends.skyrl_train.distributed.megatron.quantization_utils import ( + resolve_auto_fp8_recipe, + validate_concrete_fp8_recipe, +) from skyrl.backends.skyrl_train.inference_servers.remote_inference_client import ( SKYRL_LORA_ADAPTER_NAME, ) @@ -62,6 +67,13 @@ WeightChunk, WeightExtractor, ) +from skyrl.backends.skyrl_train.weight_sync.fp8 import ( + BLOCKWISE_FP8, + SerializedFp8Config, + iter_serialized_fp8_tensors, + registered_fp8_spec_names, + resolve_fp8_spec, +) from skyrl.backends.skyrl_train.workers.megatron.adapter_store import ( AdapterStore, LoraSignature, @@ -73,6 +85,10 @@ from skyrl.backends.skyrl_train.workers.megatron.model_bridges import ( maybe_force_qwen35_text_bridge, ) +from skyrl.backends.skyrl_train.workers.megatron.quantization.fp8_param import ( + initialize_fp8_param_optimizer_masters, + is_fp8_param_enabled, +) from skyrl.backends.skyrl_train.workers.worker import ( CriticWorkerBase, PolicyWorkerBase, @@ -120,12 +136,27 @@ def __init__( enable_bucketing: bool = False, bucket_size_threshold_GB: float = 1.0, training_dtype: torch.dtype = torch.bfloat16, + fp8_weight_sync_mode: Optional[str] = None, + hf_config=None, ): self.bridge = bridge self.actor_module = actor_module self.enable_bucketing = enable_bucketing self.bucket_size_threshold_GB = bucket_size_threshold_GB self.training_dtype = training_dtype + if fp8_weight_sync_mode is None: + self.serialized_fp8_config = None + elif fp8_weight_sync_mode == BLOCKWISE_FP8: + spec = resolve_fp8_spec(hf_config) if hf_config is not None else None + if spec is None: + raise ValueError( + "FP8 weight sync requires a registered model spec for the " + f"checkpoint layout (registered specs: {', '.join(registered_fp8_spec_names())}); " + "no spec matches the provided hf_config" + ) + self.serialized_fp8_config = SerializedFp8Config(spec=spec) + else: + raise ValueError(f"Unsupported fp8_weight_sync_mode={fp8_weight_sync_mode!r}") # Defer bucket init to first extract_weights call. # At __init__ time the model may be CPU-offloaded (colocate_all), @@ -226,6 +257,13 @@ def calculate_size_in_bytes(param, tp_size, ep_size): self.bucket_index_groups[-1].append(idx) curr_size += size + @property + def derives_metadata_from_chunks(self) -> bool: + """Serialized FP8 expands each tensor into payload + scales whose names + and shapes are only known once the chunk is built, so metadata has to + come off the stream rather than from :meth:`get_weight_metadata`.""" + return self.serialized_fp8_config is not None + def get_weight_metadata(self, dtype: torch.dtype) -> dict: """Return weight metadata without keeping tensors in memory. @@ -233,6 +271,11 @@ def get_weight_metadata(self, dtype: torch.dtype) -> dict: (tensors are discarded immediately). Result is cached for subsequent calls. TODO (aaron): find a better way to get all metadata without materializing tensors. """ + if self.serialized_fp8_config is not None: + raise RuntimeError( + "Serialized FP8 metadata depends on quantized tensor contents; " + "consume extract_weights() chunks instead." + ) if hasattr(self, "_weight_metadata_cache"): return self._weight_metadata_cache @@ -280,6 +323,18 @@ def _ensure_buckets_initialized(self): self._init_param_buckets() self._buckets_initialized = True + def _iter_sync_tensors( + self, + name: str, + tensor: torch.Tensor, + dtype: torch.dtype, + device: int, + ): + if self.serialized_fp8_config is not None: + tensor = tensor.to(device=device, non_blocking=True) + return iter_serialized_fp8_tensors(name, tensor, dtype, self.serialized_fp8_config) + return [(name, tensor.to(device=device, dtype=dtype, non_blocking=True))] + def extract_weights(self, dtype: torch.dtype): """Extract weights from Megatron model. @@ -301,14 +356,21 @@ def extract_weights(self, dtype: torch.dtype): ) for name, tensor in hf_params_generator: - tensor = tensor.to(device=device, dtype=dtype, non_blocking=True) + tensor_iter = self._iter_sync_tensors(name, tensor, dtype, device) - yield WeightChunk( - names=[name], - dtypes=[str(dtype)], - shapes=[list(tensor.shape)], - tensors=[tensor], - ) + names = [] + dtypes = [] + shapes = [] + tensors = [] + for out_name, out_tensor in tensor_iter: + out_tensor = out_tensor.contiguous() + names.append(out_name) + dtypes.append(str(out_tensor.dtype)) + shapes.append(list(out_tensor.shape)) + tensors.append(out_tensor) + + if tensors: + yield WeightChunk(names=names, dtypes=dtypes, shapes=shapes, tensors=tensors) else: # Build fresh tasks each sync so mapping objects have clean # PP-collective caches; reuse the pre-computed bucket structure. @@ -329,13 +391,14 @@ def extract_weights(self, dtype: torch.dtype): tensors = [] for name, tensor in hf_params_generator: - # Move to device and convert dtype - tensor = tensor.to(device=device, dtype=dtype, non_blocking=True) + tensor_iter = self._iter_sync_tensors(name, tensor, dtype, device) - names.append(name) - dtypes_list.append(str(dtype)) - shapes.append(list(tensor.shape)) - tensors.append(tensor) + for out_name, out_tensor in tensor_iter: + out_tensor = out_tensor.contiguous() + names.append(out_name) + dtypes_list.append(str(out_tensor.dtype)) + shapes.append(list(out_tensor.shape)) + tensors.append(out_tensor) # Yield one chunk containing all parameters in this bucket if tensors: @@ -371,7 +434,7 @@ def _maybe_setup_fake_int4_qat(self): rank0 = getattr(self, "_rank", 0) == 0 if fq.enabled: - from skyrl.backends.skyrl_train.workers.megatron.fake_int4_qat import ( + from skyrl.backends.skyrl_train.workers.megatron.quantization.fake_int4_qat import ( install_fake_int4_qat, ) @@ -440,11 +503,18 @@ def init_configs( if isinstance(transformer_config_kwargs, dict) else OmegaConf.to_container(transformer_config_kwargs, resolve=True) ) + # validate_megatron_cfg resolves fp8_recipe="auto" on the driver when it + # can see a GPU; a GPU-less driver ships "auto" through unresolved. The + # worker always has the target device visible, so resolve here and + # re-run the device/recipe validation the blind driver had to skip. + resolve_auto_fp8_recipe(transformer_config_kwargs) + validate_concrete_fp8_recipe(transformer_config_kwargs) if not self.cfg.gradient_checkpointing: for key in ("recompute_granularity", "recompute_method", "recompute_num_layers"): transformer_config_kwargs[key] = None + fp8_param_enabled = is_fp8_param_enabled(transformer_config_kwargs) bridge_source = bridge_weights_path or model_path if bridge_weights_path: logger.info( @@ -462,7 +532,11 @@ def init_configs( "(native GDN thd packing path; vision tower dropped)" ) - provider = bridge.to_megatron_provider() + # Defer persistent-FP8 checkpoint import until the bridge can expose + # unquantized converted shards for optimizer-master initialization. + provider = bridge.to_megatron_provider(load_weights=not fp8_param_enabled) + if fp8_param_enabled: + provider.perform_initialization = False if not enable_mtp and getattr(provider, "mtp_num_layers", None): logger.info(f"Disabling MTP for training (mtp_num_layers={provider.mtp_num_layers} -> None)") @@ -574,6 +648,8 @@ def init_configs( # the bridge weights path only under fake-INT4 QAT (INT4 model.path, BF16 # bridge weights); used so saved LoRA adapters reference the INT4 base. self._logical_model_path = model_path + self._deferred_fp8_param_weight_load = fp8_param_enabled + self._fp8_param_unquantized_state_dict = None # strategy.hf_config is the on-disk source-of-truth used by # save_hf_configs and must NOT carry runtime overrides like @@ -582,6 +658,40 @@ def init_configs( self.tokenizer = tokenizer self.enable_router_replay = megatron_config.moe_enable_routing_replay + def _load_deferred_fp8_param_weights(self, *, retain_unquantized_state: bool = False) -> None: + """Load deferred FP8 weights and optionally retain exact master tensors.""" + if not self._deferred_fp8_param_weight_load: + return + + original_export_dtype = self.bridge.export_weight_dtype + try: + # The FP8 export path retains converted, unquantized local shards. + # Only policy workers with an optimizer keep this additional copy. + if retain_unquantized_state: + self.bridge.export_weight_dtype = "fp8" + self.bridge.load_hf_weights(self.actor_module) + state_dict = getattr(self.bridge, "unquantized_state_dict", None) + if retain_unquantized_state and not state_dict: + raise RuntimeError( + "Megatron-Bridge did not capture unquantized checkpoint shards " + "for persistent FP8 optimizer initialization." + ) + self._fp8_param_unquantized_state_dict = state_dict if retain_unquantized_state else None + finally: + self.bridge.export_weight_dtype = original_export_dtype + + def _release_fp8_param_unquantized_state(self) -> None: + """Release temporary checkpoint shards after deferred weight import.""" + if not self._deferred_fp8_param_weight_load: + return + self._fp8_param_unquantized_state_dict = None + if self.bridge is not None: + # Clear both owners; either reference keeps the unquantized shard in HBM. + self.bridge.unquantized_state_dict = None + self._deferred_fp8_param_weight_load = False + gc.collect() + torch.cuda.empty_cache() + def configure_lora(self, lora_config, lora_type: Optional[str] = "lora"): if lora_config.target_modules == "all-linear": if lora_type == "lora": @@ -996,6 +1106,8 @@ def init_model(self, model_path, num_training_steps: int = 1e9): snapshot_download(model_path) # will be no-op if already downloaded torch.distributed.barrier() + self._load_deferred_fp8_param_weights(retain_unquantized_state=not self.cfg.policy.inference_only_init) + if self._rank == 0: print_model_size(self.actor_module[0]) @@ -1013,7 +1125,17 @@ def init_model(self, model_path, num_training_steps: int = 1e9): self.cfg.policy.optimizer_config, self.cfg.policy.megatron_config.optimizer_config_kwargs ) self.optimizer = get_megatron_optimizer(self.actor_module, optim_config) - + fp8_param_masters = initialize_fp8_param_optimizer_masters( + self.optimizer, + fp8_param=is_fp8_param_enabled(self.cfg.policy.megatron_config.transformer_config_kwargs), + fp8_param_gather=self.cfg.policy.megatron_config.ddp_config.fp8_param_gather, + state_dict=self._fp8_param_unquantized_state_dict, + ) + if fp8_param_masters: + logger.info( + "Initialized {} persistent-FP8 optimizer master shard group(s) from exact checkpoint shards.", + fp8_param_masters, + ) # create scheduler self.scheduler = get_megatron_optimizer_param_scheduler( optimizer=self.optimizer, @@ -1032,6 +1154,8 @@ def init_model(self, model_path, num_training_steps: int = 1e9): f"({n_local} head main params on rank {self._rank}; 0 is normal under DP sharding)" ) + self._release_fp8_param_unquantized_state() + # create worker model self.model = MegatronModelWrapper( config=self.cfg, @@ -1453,6 +1577,8 @@ async def init_weight_sync_state(self, inference_engine_client, inference_engine enable_bucketing=True, bucket_size_threshold_GB=inference_engine_cfg.weight_transfer_threshold_cuda_ipc_GB, training_dtype=torch.bfloat16 if self.cfg.bf16 else torch.float32, + fp8_weight_sync_mode=inference_engine_cfg.fp8_weight_sync_mode, + hf_config=self.strategy.hf_config, ) # super picks the strategy and creates the sender (for sharded_rdt that @@ -1770,6 +1896,9 @@ def init_model(self, model_path, num_training_steps: int = 1e9): snapshot_download(model_path) # will be no-op if already downloaded torch.distributed.barrier() + self._load_deferred_fp8_param_weights() + self._release_fp8_param_unquantized_state() + # load weights if self._rank == 0: print_model_size(self.actor_module[0]) diff --git a/skyrl/backends/skyrl_train/workers/megatron/quantization/__init__.py b/skyrl/backends/skyrl_train/workers/megatron/quantization/__init__.py new file mode 100644 index 0000000000..4b4ea18210 --- /dev/null +++ b/skyrl/backends/skyrl_train/workers/megatron/quantization/__init__.py @@ -0,0 +1 @@ +"""Quantization helpers for the Megatron training backend.""" diff --git a/skyrl/backends/skyrl_train/workers/megatron/fake_int4_qat.py b/skyrl/backends/skyrl_train/workers/megatron/quantization/fake_int4_qat.py similarity index 100% rename from skyrl/backends/skyrl_train/workers/megatron/fake_int4_qat.py rename to skyrl/backends/skyrl_train/workers/megatron/quantization/fake_int4_qat.py diff --git a/skyrl/backends/skyrl_train/workers/megatron/quantization/fp8_param.py b/skyrl/backends/skyrl_train/workers/megatron/quantization/fp8_param.py new file mode 100644 index 0000000000..cc90b6e1e7 --- /dev/null +++ b/skyrl/backends/skyrl_train/workers/megatron/quantization/fp8_param.py @@ -0,0 +1,150 @@ +"""Initialize optimizer state for persistent Transformer Engine FP8 parameters.""" + +from collections.abc import Mapping +from typing import Any + +import torch + +from skyrl.backends.skyrl_train.distributed.megatron.quantization_utils import ( + is_fp8_enabled, +) + + +def is_fp8_param_enabled(transformer_config_kwargs: Mapping[str, Any]) -> bool: + """Return whether dictionary config enables persistent FP8 parameters.""" + return is_fp8_enabled(transformer_config_kwargs.get("fp8_param", False)) + + +def _copy_model_shards_to_main_params( + megatron_optimizer: Any, + state_dict: Mapping[str, Any], +) -> int: + """Reload MCore FP32 master shards from converted checkpoint tensors. + + HybridDeviceOptimizer bypasses Megatron's state-dict-aware reload path. + Copying from ``state_dict`` preserves unquantized values for persistent FP8 + parameters and refreshes ordinary BF16 parameters. + """ + model_groups = getattr(megatron_optimizer, "model_float16_groups", None) + main_groups = getattr(megatron_optimizer, "shard_fp32_from_float16_groups", None) + model_fp32_groups = getattr(megatron_optimizer, "model_fp32_groups", None) + main_fp32_groups = getattr(megatron_optimizer, "shard_fp32_groups", None) + get_range = getattr(megatron_optimizer, "_get_model_param_range_map", None) + if model_groups is None or main_groups is None or not callable(get_range): + raise TypeError( + "Persistent FP8 parameter training with CPU optimizer offload requires " + "Megatron DistributedOptimizer shard metadata." + ) + + build_state_dict_map = getattr(megatron_optimizer, "_build_model_param_to_state_dict_param_map", None) + if not callable(build_state_dict_map): + raise TypeError( + "Persistent FP8 parameter training requires Megatron's " + "_build_model_param_to_state_dict_param_map() to reload exact checkpoint masters." + ) + state_dict_params = build_state_dict_map(state_dict) + + group_pairs = [(model_groups, main_groups)] + if model_fp32_groups is not None and main_fp32_groups is not None: + group_pairs.append((model_fp32_groups, main_fp32_groups)) + + copied = 0 + for grouped_model_params, grouped_main_params in group_pairs: + for model_group, main_group in zip(grouped_model_params, grouped_main_params): + for model_param, main_param in zip(model_group, main_group): + if main_param is None: + continue + + param_range = get_range(model_param)["param"] + if param_range.size != main_param.numel(): + raise RuntimeError( + "Persistent FP8 master-shard range does not match its FP32 master tensor: " + f"{param_range.size=} {main_param.numel()=}." + ) + + source_param = state_dict_params[model_param] + source = source_param.detach().reshape(-1)[param_range.start : param_range.end] + main_param.data.copy_(source.to(device=main_param.device, dtype=main_param.dtype)) + copied += 1 + return copied + + +def _sync_hybrid_device_optimizer_masters(hybrid_optimizer: Any) -> int: + """Refresh HybridDeviceOptimizer's secondary masters after checkpoint import. + + CPU copies are created before checkpoint import and are not refreshed by + MCore's public reload helper. Stale copies would overwrite weights on the + first optimizer step. + """ + copied = 0 + for cpu_param, gpu_param in getattr(hybrid_optimizer, "cpu_copys_map_gpu_param", {}).items(): + cpu_param.data.copy_( + gpu_param.detach().to(device=cpu_param.device, dtype=cpu_param.dtype), + non_blocking=False, + ) + copied += 1 + + for param, fp32_param in getattr(hybrid_optimizer, "param_to_fp32_param", {}).items(): + fp32_param.data.copy_( + param.detach().to(device=fp32_param.device, dtype=fp32_param.dtype), + non_blocking=False, + ) + copied += 1 + return copied + + +def _uses_hybrid_device_optimizer(megatron_optimizer: Any) -> bool: + """Identify MCore's CPU-offload optimizer without importing it eagerly.""" + optimizer = getattr(megatron_optimizer, "optimizer", None) + return hasattr(optimizer, "cpu_copys_map_gpu_param") and hasattr(optimizer, "param_to_fp32_param") + + +def initialize_fp8_param_optimizer_masters( + optimizer: Any, + *, + fp8_param: bool, + fp8_param_gather: bool, + state_dict: Mapping[str, Any] | None = None, +) -> int: + """Initialize optimizer masters from unquantized checkpoint shards. + + Persistent FP8 parameters cannot seed exact FP32 masters. For CPU offload, + refresh both distributed shards and HybridDeviceOptimizer's secondary copies. + """ + if not fp8_param: + return 0 + if not fp8_param_gather: + raise ValueError( + "Persistent FP8 parameters require ddp_config.fp8_param_gather=true " + "so updated FP32 master weights are requantized into FP8 compute weights." + ) + if state_dict is None: + raise ValueError( + "Persistent FP8 optimizer masters require Megatron-Bridge's exact unquantized checkpoint state." + ) + + optimizers = getattr(optimizer, "chained_optimizers", None) + if optimizers is None: + optimizers = [optimizer] + + initialized = 0 + with torch.no_grad(): + for megatron_optimizer in optimizers: + if _uses_hybrid_device_optimizer(megatron_optimizer): + copied = _copy_model_shards_to_main_params(megatron_optimizer, state_dict) + if copied == 0: + raise RuntimeError( + "Persistent FP8 parameter training did not find any model shards " + "to seed into CPU-offloaded optimizer masters." + ) + _sync_hybrid_device_optimizer_masters(megatron_optimizer.optimizer) + else: + reload_main_params = getattr(megatron_optimizer, "_copy_model_params_to_main_params", None) + if not callable(reload_main_params): + raise TypeError( + "Persistent FP8 parameter training requires a Megatron optimizer " + "with _copy_model_params_to_main_params()." + ) + reload_main_params(state_dict=state_dict) + initialized += 1 + return initialized diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index 864be3f2a9..b952a87fb3 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -142,7 +142,7 @@ class FakeInt4QatConfig(BaseConfig): BF16 masters, enabling this fake-quantizes the frozen expert GEMMs onto the same INT4 grid in the forward pass (straight-through backward), removing the train/infer weight mismatch. See - ``skyrl.backends.skyrl_train.workers.megatron.fake_int4_qat``. + ``skyrl.backends.skyrl_train.workers.megatron.quantization.fake_int4_qat``. """ enabled: bool = False @@ -255,6 +255,10 @@ class MegatronDDPConfig(BaseConfig): grad_reduce_in_fp32: bool = True overlap_grad_reduce: bool = False overlap_param_gather: bool = False + fp8_param_gather: bool = False + """Keep the DDP parameter all-gather in FP8. + Must be ``True`` when training with ``transformer_config_kwargs.fp8_param=true`` so persistent + FP8 params stay FP8 through the distributed-optimizer all-gather.""" average_in_collective: bool = True @@ -498,13 +502,37 @@ class MegatronConfig(BaseConfig): https://docs.skyrl.ai/docs/examples/megatron for the accepted names and per-field checks. ``use_precision_aware_optimizer=True`` can cause checkpointing to fail (https://github.com/nvidia/megatron-lm/issues/1820); leaving it ``False`` is recommended.""" + fp8: Optional[str] = None + """TransformerEngine FP8 compute format for linear-layer GEMMs, e.g. ``"e4m3"`` or + ``"hybrid"``. ``None`` (default) trains without FP8. Folded into + ``transformer_config_kwargs["fp8"]``; an explicit kwarg takes precedence.""" + fp8_recipe: Optional[str] = None + """TransformerEngine FP8 scaling recipe. ``"auto"`` resolves to the architecture-native + recipe — ``blockwise``/FP32 scales on Hopper, native MXFP8 on Blackwell/SM100+ (where TE + also requires every weight dim to be divisible by 32). Folded into + ``transformer_config_kwargs["fp8_recipe"]``; an explicit kwarg takes precedence.""" + fp8_param: Optional[bool] = None + """Store Megatron primary parameters in FP8 (E4M3). Supported with + ``fp8_recipe=blockwise`` + FP32 block scales only and requires + ``ddp_config.fp8_param_gather=true``. Folded into + ``transformer_config_kwargs["fp8_param"]``; an explicit kwarg takes precedence.""" + fp8_amax_compute_algo: Optional[str] = None + """TransformerEngine amax history reduction, e.g. ``"most_recent"`` or ``"max"``. Folded + into ``transformer_config_kwargs["fp8_amax_compute_algo"]``; an explicit kwarg takes + precedence.""" transformer_config_kwargs: Dict[str, Any] = field( default_factory=lambda: copy.deepcopy(DEFAULT_TRANSFORMER_CONFIG_KWARGS) ) """Pass-through kwargs for Megatron's ``TransformerConfig``: https://github.com/NVIDIA/Megatron-LM/blob/core_r0.13.0/megatron/core/transformer/transformer_config.py Also the place to put HuggingFace config overrides (e.g. ``rope_parameters``) for the Megatron - backend, where FSDP would use ``model_config_kwargs``.""" + backend, where FSDP would use ``model_config_kwargs``. + + FP8 training is configured through the top-level ``fp8``, ``fp8_recipe``, ``fp8_param``, and + ``fp8_amax_compute_algo`` fields above, which fold into these kwargs; setting the same keys + here directly overrides them. The block-scale env contract + (``NVTE_FP8_BLOCK_SCALING_FP32_SCALES``, ``VLLM_USE_DEEP_GEMM_E8M0``) is defaulted and + validated per architecture at startup and forwarded to all Ray actors.""" empty_cuda_cache: Optional[bool] = True """Manually empty torch's CUDA cache between the forward/backward pass and the optimizer step. This frees reserved-but-unallocated memory and can help avoid OOMs in the optimizer.""" @@ -567,6 +595,16 @@ def __post_init__(self): # doesn't have to repeat every default just to set one value. if self.transformer_config_kwargs is None: self.transformer_config_kwargs = {} + # The top-level FP8 fields fold into the TransformerConfig kwargs; explicitly + # configured kwargs take precedence. + for key, value in ( + ("fp8", self.fp8), + ("fp8_recipe", self.fp8_recipe), + ("fp8_param", self.fp8_param), + ("fp8_amax_compute_algo", self.fp8_amax_compute_algo), + ): + if value is not None: + self.transformer_config_kwargs.setdefault(key, value) for k, v in DEFAULT_TRANSFORMER_CONFIG_KWARGS.items(): self.transformer_config_kwargs.setdefault(k, copy.deepcopy(v)) if self.optimizer_config_kwargs is None: @@ -1106,6 +1144,14 @@ class InferenceEngineConfig(BaseConfig): """Should match the dtype used by the inference engine. Also used during full-weight sync, where policy weights are cast to this dtype before being sent to the inference engine. The LoRA-adapter sync path exports fp32 instead.""" + fp8_weight_sync_mode: Optional[str] = None + """Optional rollout weight format. ``"blockwise"`` sends FP8 checkpoint weights and + scales (one FP32 scale per 128x128 block) instead of ``model_dtype`` tensors, halving transfer + volume and letting vLLM serve FP8. Requires ``trainer.strategy="megatron"`` and a model with a + registered FP8 spec (see ``skyrl/backends/skyrl_train/weight_sync/fp8/models/README.md``). The + vLLM engine settings this needs (``quantization="fp8"``, ``load_format="dummy"``, and the + matching ``hf_overrides.quantization_config`` with per-model ignored layers) are applied + automatically; the first weight sync supplies real weights before any generation.""" run_engines_locally: bool = True """Launch inference servers during the training run in the current Ray cluster. When ``False``, point SkyRL at an external HTTP/vLLM deployment via ``external_proxy_url`` and/or diff --git a/skyrl/train/dataset/collators.py b/skyrl/train/dataset/collators.py index cb6e3e385b..4647461245 100644 --- a/skyrl/train/dataset/collators.py +++ b/skyrl/train/dataset/collators.py @@ -18,7 +18,7 @@ from __future__ import annotations -from typing import List +from typing import List, Optional import numpy as np import torch @@ -72,7 +72,8 @@ class PackedDataCollator: Flow: 1. Compute per-example sequence lengths. - 2. FFD-pack with ``bin_capacity = max_tokens_per_microbatch``, + 2. FFD-pack using each sequence's alignment-padded footprint and + ``bin_capacity = max(max_tokens_per_microbatch, align_size)``, ``min_bin_count = dp_size``, ``bin_count_multiple = dp_size``. 3. Round-robin assign bins to DP shards (this happens implicitly inside ``MeshDispatch.dispatch`` because the rows are laid out in shard-major @@ -96,6 +97,7 @@ def __init__( batch_size: int, micro_train_batch_size_per_gpu: int, fp8_enabled: bool = False, + fp8_recipe: Optional[str] = None, ): if max_tokens_per_microbatch is None: raise ValueError("PackedDataCollator requires max_tokens_per_microbatch to be set explicitly.") @@ -106,6 +108,7 @@ def __init__( self.dp_size = dp_size self.batch_size = batch_size self.fp8_enabled = fp8_enabled + self.fp8_recipe = fp8_recipe self._default_collator = DefaultCollator(tokenizer, micro_train_batch_size_per_gpu) self._tokenizer = tokenizer @@ -133,19 +136,26 @@ def __call__(self, examples: list, batch_size: int) -> TrainingInputBatch: pp_size = self.pp_size cp_size = self.cp_size # Each sub-seq's padded length must satisfy these divisibility - # constraints, which is why ``align_size`` carries all factors: + # constraints, which is why ``align_size`` carries all required factors: # - Sequence Parallelism (auto-on when tp>1) shards along the seq # dim, so each segment must be divisible by ``tp_size``. # - Context Parallelism splits each segment into ``2*cp_size`` equal # load-balanced causal chunks, so each segment must be divisible by # ``2*cp_size``. - # - When FP8 is enabled, Transformer Engine GEMMs require each CP - # rank's local token slab to be 16-aligned; globally this means - # ``16*cp_size``. + # - FP8: TE with fp8_recipe=blockwise needs 16-token local slabs at + # TP=1 and quantizes sequence-parallel all-gather inputs in + # 128-token blocks (so the global segment includes the TP and CP + # shard factors); fp8_recipe=mxfp8 quantizes in 1x32 tiles, so the + # TP=1 slab grows to 32. # This MUST stay in lockstep with the worker's preprocess_packed_seqs # (megatron_utils.py): if the divisors drift, the per-rank CP/SP # gather/scatter offsets silently corrupt loss/grads (no crash). - align_size = get_packed_seq_align_size(tp_size, cp_size, fp8_enabled=self.fp8_enabled) + align_size = get_packed_seq_align_size( + tp_size, cp_size, fp8_enabled=self.fp8_enabled, fp8_recipe=self.fp8_recipe + ) + + def _round_up(x: int, multiple: int) -> int: + return ((x + multiple - 1) // multiple) * multiple dp_size = self.dp_size @@ -179,14 +189,26 @@ def __call__(self, examples: list, batch_size: int) -> TrainingInputBatch: # same number of micro-batches. Forcing the global bin count to a # multiple of ``dp_size`` makes the per-DP-rank bin count (and thus # ``num_microbatches``) identical across ranks. + # Under FP8, pack each sequence's *aligned* footprint rather than its raw + # length: align_size is 128*tp*cp (vs tp*cp*2 without FP8), so per-sequence + # padding can otherwise push a bin far past max_tokens_per_microbatch and + # overflow the row budget. Also allow at least one alignment unit per bin, + # since a single padded sub-seq already costs align_size tokens. + # Non-FP8 keeps upstream's raw-length packing byte-for-byte. + if self.fp8_enabled: + packing_lengths = [_round_up(length, align_size) for length in seq_lengths] + packing_capacity = max(bin_capacity, align_size) + else: + packing_lengths = seq_lengths + packing_capacity = bin_capacity bin_count_multiple = dp_size packer = make_seq_packer( "first_fit_decreasing", - bin_capacity=bin_capacity, + bin_capacity=packing_capacity, min_bin_count=bin_count_multiple, bin_count_multiple=bin_count_multiple, ) - bins: List[List[int]] = packer.pack(seq_lengths) + bins: List[List[int]] = packer.pack(packing_lengths) # Assign bins to DP shards via round-robin (bin_idx % shards). # Concretely we want the resulting layout to be shard-major: @@ -206,9 +228,6 @@ def __call__(self, examples: list, batch_size: int) -> TrainingInputBatch: # 3. Compute packed-row lengths (with align_size padding per sub-seq) # and the global max packed length (for PP > 1 uniform padding). # ------------------------------------------------------------------ - def _round_up(x: int, m: int) -> int: - return ((x + m - 1) // m) * m - bin_packed_lengths: List[int] = [] bin_subseq_lengths: List[List[int]] = [] # one list per bin row for bin_indices in flat_bins: @@ -244,7 +263,8 @@ def _round_up(x: int, m: int) -> int: n_samples = len(examples) logger.info( f"sequence packing | packed {n_samples} samples into {num_bins} bins " - f"(~{num_bins // dp_size}/DP rank, bin_capacity={bin_capacity} tokens)" + f"(~{num_bins // dp_size}/DP rank, bin_capacity={packing_capacity}" + f"{' aligned' if self.fp8_enabled else ''} tokens)" ) # Fill NumPy buffers by slice, then convert once. @@ -270,8 +290,7 @@ def _round_up(x: int, m: int) -> int: if n_write > 0: loss_mask_np[row_idx, row_offset:write_end] = full_loss_masks[ex_idx][1 : 1 + n_write] - # Advance row_offset, padding sub-seq to the TP/CP layout - # multiple, plus FP8's 16-token local-rank multiple when active. + # Match the aligned footprint consumed by preprocess_packed_seqs. row_offset += _round_up(s, align_size) # Count response-token loss slots before normalization. The vectorized diff --git a/skyrl/train/sft_trainer.py b/skyrl/train/sft_trainer.py index 80d187ee33..1df1d28ab9 100644 --- a/skyrl/train/sft_trainer.py +++ b/skyrl/train/sft_trainer.py @@ -863,7 +863,7 @@ def _build_collator(self, tokenizer): from skyrl.train.dataset.collators import DefaultCollator, PackedDataCollator if self.sft_cfg.use_sequence_packing: - from skyrl.backends.skyrl_train.distributed.megatron.packing_utils import ( + from skyrl.backends.skyrl_train.distributed.megatron.quantization_utils import ( is_fp8_enabled, ) @@ -879,6 +879,7 @@ def _build_collator(self, tokenizer): batch_size=self.sft_cfg.batch_size, micro_train_batch_size_per_gpu=self.sft_cfg.micro_train_batch_size_per_gpu, fp8_enabled=is_fp8_enabled(transformer_config_kwargs.get("fp8")), + fp8_recipe=transformer_config_kwargs.get("fp8_recipe"), ) return DefaultCollator( tokenizer=tokenizer, diff --git a/skyrl/train/utils/utils.py b/skyrl/train/utils/utils.py index 6808dd4b81..d6a887a335 100644 --- a/skyrl/train/utils/utils.py +++ b/skyrl/train/utils/utils.py @@ -20,6 +20,16 @@ ) from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy +from skyrl.backends.skyrl_train.distributed.megatron.quantization_utils import ( + has_visible_cuda_device, + is_blackwell_or_newer, + is_fp8_enabled, + resolve_auto_fp8_recipe, + validate_concrete_fp8_recipe, +) +from skyrl.backends.skyrl_train.weight_sync.fp8 import ( + BLOCKWISE_FP8, +) from skyrl.env_vars import ( SKYRL_DUMP_INFRA_LOG_TO_STDOUT, SKYRL_LD_LIBRARY_PATH_EXPORT, @@ -202,6 +212,31 @@ def validate_megatron_cfg(cfg: SkyRLTrainConfig): assert ie_cfg.backend == "vllm", "only vllm is supported for with megatron" assert cfg.trainer.critic.model.path is None, "only GRPO training is currently supported for megatron" + policy_cfg = cfg.trainer.policy + policy_fp8_param = is_fp8_enabled(policy_cfg.megatron_config.transformer_config_kwargs.get("fp8_param")) + if ( + policy_fp8_param + and not policy_cfg.inference_only_init + and not policy_cfg.megatron_config.ddp_config.fp8_param_gather + ): + raise ValueError( + "Persistent policy fp8_param training requires " + "trainer.policy.megatron_config.ddp_config.fp8_param_gather=true" + ) + + # Resolve fp8_recipe="auto" to the architecture-native recipe (blockwise on + # Hopper, mxfp8 on Blackwell) before the config is shipped to Ray actors. + # A GPU-less driver leaves "auto" in place — guessing here would bake the + # wrong recipe into every worker's config — and each Megatron worker then + # resolves and re-validates locally against its own device. + for worker_cfg in (cfg.trainer.policy, cfg.trainer.ref): + megatron_config = getattr(worker_cfg, "megatron_config", None) + transformer_kwargs = getattr(megatron_config, "transformer_config_kwargs", None) + if not transformer_kwargs: + continue + resolve_auto_fp8_recipe(transformer_kwargs) + validate_concrete_fp8_recipe(transformer_kwargs) + if cfg.trainer.policy.megatron_config.moe_enable_routing_replay: assert ( cfg.generator.inference_engine.enable_return_routed_experts @@ -543,6 +578,32 @@ def validate_inference_engine_cfg(cfg: SkyRLTrainConfig): """ ie_cfg = cfg.generator.inference_engine + if ie_cfg.fp8_weight_sync_mode not in (None, BLOCKWISE_FP8): + raise ValueError( + f"Unsupported fp8_weight_sync_mode={ie_cfg.fp8_weight_sync_mode!r}; " f"expected {BLOCKWISE_FP8!r} or None" + ) + if ie_cfg.fp8_weight_sync_mode == BLOCKWISE_FP8: + if cfg.trainer.strategy != "megatron": + raise ValueError("blockwise FP8 weight sync requires trainer.strategy='megatron'") + if ie_cfg.weight_sync_backend in {"sharded_rdt", "delta"}: + # Neither backend can carry the quantized payload + scale pairs that + # blockwise FP8 sync is made of: the RDT weight sources export bridge + # tensors cast to the inference dtype, and the delta checkpoint format + # cannot represent the marker names and scale tensors. Both senders + # refuse at send time too, but vLLM is built with quantization="fp8" + # and load_format="dummy" long before the first sync, so the model is + # already loaded by then. + raise ValueError( + "blockwise FP8 weight sync is not supported with " + f"weight_sync_backend={ie_cfg.weight_sync_backend!r}; use 'nccl'" + ) + lora_cfg = cfg.trainer.policy.model.lora + if lora_cfg.rank > 0 and not cfg.trainer.policy.megatron_config.lora_config.merge_lora: + raise ValueError( + "blockwise FP8 weight sync requires full-weight updates; " + "Megatron LoRA with merge_lora=false syncs adapters only" + ) + if ie_cfg.enable_pd: assert ie_cfg.num_prefill > 0, "num_prefill must be > 0 when enable_pd=True" assert ( @@ -873,6 +934,85 @@ def prepare_runtime_environment(cfg: SkyRLTrainConfig) -> dict[str, str]: logger.info(f"Exporting SKYRL_* overrides to ray runtime env: {sorted(forwarded)}") env_vars.update(forwarded) + # Forward one block-scale contract to all Ray actors. Hopper defaults to FP32 + # scales; Blackwell (SM100+) defaults to power-of-two scales, the only mode TE + # supports for blockwise quantization there (it emulates Float8BlockScaling on + # the MX datapath). + serialized_fp8 = cfg.generator.inference_engine.fp8_weight_sync_mode == BLOCKWISE_FP8 + use_ref_model = cfg.trainer.algorithm.use_kl_loss or cfg.trainer.algorithm.use_kl_in_reward + policy_megatron_config = getattr(cfg.trainer.policy, "megatron_config", None) + ref_megatron_config = getattr(cfg.trainer.ref, "megatron_config", None) + policy_transformer_kwargs = getattr(policy_megatron_config, "transformer_config_kwargs", None) or {} + ref_transformer_kwargs = getattr(ref_megatron_config, "transformer_config_kwargs", None) or {} + policy_fp8_param = is_fp8_enabled(policy_transformer_kwargs.get("fp8_param")) + ref_fp8_param = use_ref_model and is_fp8_enabled(ref_transformer_kwargs.get("fp8_param")) + fp8_compute = is_fp8_enabled(policy_transformer_kwargs.get("fp8")) or ( + use_ref_model and is_fp8_enabled(ref_transformer_kwargs.get("fp8")) + ) + fp8_contract_enabled = serialized_fp8 or fp8_compute or policy_fp8_param or ref_fp8_param + + fp8_env_defaults: dict[str, str] = {} + configured_scale_mode = os.environ.get("NVTE_FP8_BLOCK_SCALING_FP32_SCALES") + if fp8_contract_enabled or configured_scale_mode is not None: + if configured_scale_mode is None and not has_visible_cuda_device(): + # The block-scale contract must be identical in every actor, so it is + # fixed here, before ray.init ships it in the runtime env — too early + # to probe the cluster. A GPU-less head cannot infer the workers' + # architecture, and defaulting to the Hopper contract would silently + # hand FP32 block scales to Blackwell workers, where TE emulates + # blockwise on the MX datapath and only supports power-of-2 scales. + raise ValueError( + "FP8 is enabled but this driver sees no CUDA device, so the block-scale " + "contract cannot be inferred from the workers' architecture. Export " + "NVTE_FP8_BLOCK_SCALING_FP32_SCALES explicitly: '0' (power-of-2 scales) " + "on Blackwell/SM100+, '1' (FP32 scales) on Hopper." + ) + scale_mode = configured_scale_mode or ("0" if is_blackwell_or_newer() else "1") + if scale_mode not in {"0", "1"}: + raise ValueError("NVTE_FP8_BLOCK_SCALING_FP32_SCALES must be '0' (power-of-2) " "or '1' (FP32 scales).") + + if scale_mode == "0" and (policy_fp8_param or ref_fp8_param): + raise ValueError( + "Persistent fp8_param requires FP32 block scales. Blackwell only supports " + "power-of-2 block scales, so use fp8_param=false on Blackwell." + ) + + if fp8_contract_enabled: + fp8_env_defaults["NVTE_FP8_BLOCK_SCALING_FP32_SCALES"] = scale_mode + if serialized_fp8 and scale_mode == "1": + e8m0_mode = os.environ.get("VLLM_USE_DEEP_GEMM_E8M0", "0") + if e8m0_mode != "0": + raise ValueError( + "FP32 block scales require VLLM_USE_DEEP_GEMM_E8M0=0 so vLLM " + "does not requantize them to power-of-2 scales." + ) + fp8_env_defaults["VLLM_USE_DEEP_GEMM_E8M0"] = e8m0_mode + elif serialized_fp8 and scale_mode == "0": + # The symmetric rule, and a property of the wire format rather than of + # this process's device: power-of-2 scales are exactly representable in + # E8M0, so vLLM's requantization is lossless. vLLM then picks the + # per-device form itself (UE8M0 on SM100/SM120, FP32-ceil-to-UE8M0 on + # Hopper), which is why this default must not be gated on the driver's + # architecture — a GPU-less head would drop it. SM100 DeepGEMM also + # rejects the alternative outright ("Unsupported architecture or + # scaling factor types" with VLLM_USE_DEEP_GEMM_E8M0=0). + e8m0_mode = os.environ.get("VLLM_USE_DEEP_GEMM_E8M0", "1") + if e8m0_mode != "1": + raise ValueError( + "Power-of-2 block scales require VLLM_USE_DEEP_GEMM_E8M0=1: they " + "requantize to E8M0 losslessly, and Blackwell DeepGEMM accepts no " + "other scale factor type. Unset the variable or set it to 1." + ) + fp8_env_defaults["VLLM_USE_DEEP_GEMM_E8M0"] = e8m0_mode + + for var_name in ( + "NVTE_FP8_BLOCK_SCALING_FP32_SCALES", + "VLLM_USE_DEEP_GEMM_E8M0", + ): + if value := os.environ.get(var_name, fp8_env_defaults.get(var_name)): + logger.info(f"Exporting `{var_name}` to ray runtime env: {value}") + env_vars[var_name] = value + return env_vars diff --git a/tests/backends/skyrl_train/distributed/test_packing_utils.py b/tests/backends/skyrl_train/distributed/test_packing_utils.py index eb99bd2d6d..bcfc8f910a 100644 --- a/tests/backends/skyrl_train/distributed/test_packing_utils.py +++ b/tests/backends/skyrl_train/distributed/test_packing_utils.py @@ -3,37 +3,48 @@ from skyrl.backends.skyrl_train.distributed.megatron.packing_utils import ( get_packed_seq_align_size, get_unpacked_seq_align_size, - is_fp8_enabled, ) -@pytest.mark.parametrize( - ("fp8", "expected"), - [ - (None, False), - ("", False), - ("false", False), - ("0", False), - (False, False), - ("hybrid", True), - ("e4m3", True), - (True, True), - ], -) -def test_is_fp8_enabled(fp8, expected): - assert is_fp8_enabled(fp8) is expected - - def test_packed_alignment_uses_layout_only_without_fp8(): assert get_packed_seq_align_size(tp_size=4, cp_size=1) == 4 assert get_packed_seq_align_size(tp_size=1, cp_size=2) == 4 def test_packed_alignment_adds_fp8_local_rank_multiple(): - assert get_packed_seq_align_size(tp_size=4, cp_size=1, fp8_enabled=True) == 16 + assert get_packed_seq_align_size(tp_size=4, cp_size=1, fp8_enabled=True) == 512 assert get_packed_seq_align_size(tp_size=1, cp_size=2, fp8_enabled=True) == 32 + assert get_packed_seq_align_size(tp_size=2, cp_size=1, fp8_enabled=True) == 256 + assert get_packed_seq_align_size(tp_size=2, cp_size=2, fp8_enabled=True) == 512 def test_unpacked_alignment_adds_fp8_multiple_only_when_enabled(): assert get_unpacked_seq_align_size(tp_size=4) == 4 - assert get_unpacked_seq_align_size(tp_size=4, fp8_enabled=True) == 16 + assert get_unpacked_seq_align_size(tp_size=1, fp8_enabled=True) == 16 + assert get_unpacked_seq_align_size(tp_size=2, fp8_enabled=True) == 256 + assert get_unpacked_seq_align_size(tp_size=4, fp8_enabled=True) == 512 + + +def test_mxfp8_recipe_aligns_to_32_token_local_shards(): + # 32*tp*cp at any TP, never the blockwise 128*tp*cp segments. + assert get_packed_seq_align_size(tp_size=1, cp_size=1, fp8_enabled=True, fp8_recipe="mxfp8") == 32 + assert get_packed_seq_align_size(tp_size=1, cp_size=2, fp8_enabled=True, fp8_recipe="mxfp8") == 64 + assert get_packed_seq_align_size(tp_size=2, cp_size=1, fp8_enabled=True, fp8_recipe="mxfp8") == 64 + assert get_packed_seq_align_size(tp_size=2, cp_size=2, fp8_enabled=True, fp8_recipe="mxfp8") == 128 + assert get_packed_seq_align_size(tp_size=4, cp_size=1, fp8_enabled=True, fp8_recipe="mxfp8") == 128 + assert get_unpacked_seq_align_size(tp_size=1, fp8_enabled=True, fp8_recipe="mxfp8") == 32 + assert get_unpacked_seq_align_size(tp_size=2, fp8_enabled=True, fp8_recipe="mxfp8") == 64 + # Non-mx recipes keep the blockwise constants. + assert get_packed_seq_align_size(tp_size=1, cp_size=1, fp8_enabled=True, fp8_recipe="blockwise") == 16 + assert get_unpacked_seq_align_size(tp_size=1, fp8_enabled=True, fp8_recipe=None) == 16 + + +@pytest.mark.parametrize(("tp_size", "cp_size"), [(0, 1), (1, 0), (-1, 1)]) +def test_packed_alignment_rejects_nonpositive_parallel_sizes(tp_size, cp_size): + with pytest.raises(ValueError, match="must be positive"): + get_packed_seq_align_size(tp_size, cp_size, fp8_enabled=True) + + +def test_unpacked_alignment_rejects_nonpositive_tp_size(): + with pytest.raises(ValueError, match="must be positive"): + get_unpacked_seq_align_size(0, fp8_enabled=True) diff --git a/tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_cp.py b/tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_cp.py index 836aa16d11..b3f87a561c 100644 --- a/tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_cp.py +++ b/tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_cp.py @@ -143,3 +143,47 @@ def test_short_seq_no_crash(self, tp_size, cp_size, real_tokens): assert result_ids.shape[1] % 16 == 0 assert packed_params.max_seqlen_q % (16 * cp_size) == 0 assert packed_params.qkv_format == "thd" + + def test_remove_left_padding_tp1_aligns_only_for_fp8(self): + """TP1 applies 16-token alignment only when FP8 is enabled.""" + from skyrl.backends.skyrl_train.distributed.megatron.megatron_utils import ( + remove_left_padding, + ) + + input_ids = torch.arange(7000).unsqueeze(0) + attention_mask = torch.zeros((1, 7000), dtype=torch.bool) + attention_mask[0, :6541] = True + position_ids = torch.arange(7000).unsqueeze(0) + + with patch("skyrl.backends.skyrl_train.distributed.megatron.megatron_utils.mpu") as mock_mpu: + mock_mpu.get_tensor_model_parallel_world_size.return_value = 1 + mock_mpu.get_context_parallel_world_size.return_value = 1 + + bf16_ids, bf16_mask, _ = remove_left_padding(input_ids, attention_mask, position_ids, fp8_enabled=False) + fp8_ids, fp8_mask, _ = remove_left_padding(input_ids, attention_mask, position_ids, fp8_enabled=True) + + assert bf16_ids.shape == (1, 6541) + assert int(bf16_mask.sum().item()) == 6541 + assert fp8_ids.shape == (1, 6544) + assert int(fp8_mask.sum().item()) == 6541 + + def test_remove_left_padding_tp_gt_1_fp8_uses_local_128_alignment(self): + """TP2 rounds 8,552 tokens to the 8,704-token global alignment.""" + from skyrl.backends.skyrl_train.distributed.megatron.megatron_utils import ( + remove_left_padding, + ) + + input_ids = torch.arange(9000).unsqueeze(0) + attention_mask = torch.zeros((1, 9000), dtype=torch.bool) + attention_mask[0, :8552] = True + position_ids = torch.arange(9000).unsqueeze(0) + + with patch("skyrl.backends.skyrl_train.distributed.megatron.megatron_utils.mpu") as mock_mpu: + mock_mpu.get_tensor_model_parallel_world_size.return_value = 2 + mock_mpu.get_context_parallel_world_size.return_value = 1 + + fp8_ids, fp8_mask, _ = remove_left_padding(input_ids, attention_mask, position_ids, fp8_enabled=True) + + assert fp8_ids.shape == (1, 8704) + assert int(fp8_mask.sum().item()) == 8552 + assert fp8_ids.shape[1] % (128 * 2) == 0 diff --git a/tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_multiseq.py b/tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_multiseq.py index 95c6880687..2f510092c2 100644 --- a/tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_multiseq.py +++ b/tests/backends/skyrl_train/distributed/test_preprocess_packed_seqs_multiseq.py @@ -187,24 +187,24 @@ def test_multiseq_with_tp_alignment(self): The intra-row offsets read by preprocess must match the collator's row layout, which advances ``row_offset += round_up(s, align_size)`` between sub-seqs. So with sub-seqs of length 3 and - 5 and tp_size=4, the collator places sub-seq 1 at row column 16 - (after the FP8/TP alignment pad gap), NOT row column 3. + 5 and tp_size=4, the collator places sub-seq 1 at row column + ``align_size`` after the FP8/TP alignment gap, not row column 3. """ from skyrl.backends.skyrl_train.distributed.megatron.megatron_utils import ( preprocess_packed_seqs, ) - seq_len = 48 batch_size = 1 - # Two sub-seqs of length 3 and 5; tp_size=4 still pads each to 16 - # for Transformer Engine FP8 compatibility. + # Each sequence uses a global footprint that leaves TP-local inputs + # aligned to 128 tokens. align_size = _get_align_size(tp_size=4, cp_size=1, fp8_enabled=True) + seq_len = 2 * align_size # Row layout mirrors what PackedDataCollator produces: - # row[0:3] = sub-seq 0 tokens - # row[3:16] = alignment pad (zero) - # row[16:21] = sub-seq 1 tokens - # row[21:32] = alignment pad (zero) + # row[0:3] = sub-seq 0 tokens + # row[3:align_size] = alignment pad (zero) + # row[align_size:align_size + 5] = sub-seq 1 tokens + # row[align_size + 5:2*align_size] = alignment pad (zero) input_ids = torch.zeros(batch_size, seq_len, dtype=torch.long) input_ids[0, :3] = torch.tensor([1, 2, 3]) input_ids[0, align_size : align_size + 5] = torch.tensor([10, 11, 12, 13, 14]) @@ -224,11 +224,11 @@ def test_multiseq_with_tp_alignment(self): fp8_enabled=True, ) - assert params.cu_seqlens_q.tolist() == [0, 16, 32] - assert packed.shape == (1, 32) + assert params.cu_seqlens_q.tolist() == [0, align_size, 2 * align_size] + assert packed.shape == (1, 2 * align_size) assert packed[0, :3].tolist() == [1, 2, 3] - assert packed[0, 3:16].tolist() == [0] * 13 - assert packed[0, 16:21].tolist() == [10, 11, 12, 13, 14] + assert packed[0, 3:align_size].tolist() == [0] * (align_size - 3) + assert packed[0, align_size : align_size + 5].tolist() == [10, 11, 12, 13, 14] def test_multiple_bin_rows(self): """Two bin rows, each with two sub-seqs, produce 4+1 cu_seqlens entries.""" diff --git a/tests/backends/skyrl_train/distributed/test_quantization_utils.py b/tests/backends/skyrl_train/distributed/test_quantization_utils.py new file mode 100644 index 0000000000..f54daedd82 --- /dev/null +++ b/tests/backends/skyrl_train/distributed/test_quantization_utils.py @@ -0,0 +1,111 @@ +import pytest + +from skyrl.backends.skyrl_train.distributed.megatron import quantization_utils +from skyrl.backends.skyrl_train.distributed.megatron.quantization_utils import ( + is_fp8_enabled, + is_mxfp8_recipe, + resolve_auto_fp8_recipe, + validate_concrete_fp8_recipe, +) + + +@pytest.mark.parametrize( + ("fp8", "expected"), + [ + (None, False), + ("", False), + ("false", False), + ("0", False), + (False, False), + ("hybrid", True), + ("e4m3", True), + (True, True), + ], +) +def test_is_fp8_enabled(fp8, expected): + assert is_fp8_enabled(fp8) is expected + + +@pytest.mark.parametrize( + ("recipe", "expected"), + [ + (None, False), + ("", False), + ("blockwise", False), + ("delayed", False), + ("mxfp8", True), + (" MXFP8 ", True), + ], +) +def test_is_mxfp8_recipe(recipe, expected): + assert is_mxfp8_recipe(recipe) is expected + + +def test_resolve_auto_fp8_recipe_picks_mxfp8_on_blackwell(monkeypatch): + monkeypatch.setattr(quantization_utils, "has_visible_cuda_device", lambda: True) + monkeypatch.setattr(quantization_utils, "is_blackwell_or_newer", lambda: True) + kwargs = {"fp8": "e4m3", "fp8_recipe": "auto"} + assert resolve_auto_fp8_recipe(kwargs) == "mxfp8" + assert kwargs["fp8_recipe"] == "mxfp8" + + +def test_resolve_auto_fp8_recipe_picks_blockwise_on_hopper(monkeypatch): + monkeypatch.setattr(quantization_utils, "has_visible_cuda_device", lambda: True) + monkeypatch.setattr(quantization_utils, "is_blackwell_or_newer", lambda: False) + kwargs = {"fp8": "e4m3", "fp8_recipe": "AUTO"} + assert resolve_auto_fp8_recipe(kwargs) == "blockwise" + assert kwargs["fp8_recipe"] == "blockwise" + + +def test_resolve_auto_fp8_recipe_defers_without_cuda(monkeypatch): + """A GPU-less driver must not guess the workers' architecture.""" + monkeypatch.setattr(quantization_utils, "has_visible_cuda_device", lambda: False) + kwargs = {"fp8": "e4m3", "fp8_recipe": "auto"} + assert resolve_auto_fp8_recipe(kwargs) == "auto" + assert kwargs["fp8_recipe"] == "auto" + + +def test_validate_concrete_fp8_recipe_ignores_non_mxfp8(): + validate_concrete_fp8_recipe({"fp8_recipe": "blockwise", "fp8_param": True}) + validate_concrete_fp8_recipe({"fp8_recipe": "auto"}) + validate_concrete_fp8_recipe({}) + validate_concrete_fp8_recipe(None) + + +def test_validate_concrete_fp8_recipe_rejects_mxfp8_before_blackwell(monkeypatch): + monkeypatch.setattr(quantization_utils, "has_visible_cuda_device", lambda: True) + monkeypatch.setattr(quantization_utils, "is_blackwell_or_newer", lambda: False) + with pytest.raises(ValueError, match="requires SM100"): + validate_concrete_fp8_recipe({"fp8_recipe": "mxfp8"}) + + +def test_validate_concrete_fp8_recipe_rejects_mxfp8_with_fp8_param(monkeypatch): + # Device-independent: must fire even on a GPU-less process. + monkeypatch.setattr(quantization_utils, "has_visible_cuda_device", lambda: False) + with pytest.raises(ValueError, match="fp8_param"): + validate_concrete_fp8_recipe({"fp8_recipe": "mxfp8", "fp8_param": True}) + + +def test_resolve_auto_fp8_recipe_passes_explicit_values_through(monkeypatch): + monkeypatch.setattr(quantization_utils, "is_blackwell_or_newer", lambda: False) + kwargs = {"fp8_recipe": "blockwise"} + assert resolve_auto_fp8_recipe(kwargs) == "blockwise" + assert kwargs["fp8_recipe"] == "blockwise" + assert resolve_auto_fp8_recipe({}) is None + assert resolve_auto_fp8_recipe(None) is None + + +def test_resolve_auto_fp8_recipe_warns_for_emulated_blockwise_on_blackwell(monkeypatch): + monkeypatch.setattr(quantization_utils, "is_blackwell_or_newer", lambda: True) + warnings = [] + monkeypatch.setattr( + quantization_utils.logger, "warning", lambda msg, *args, **kw: warnings.append(msg.format(*args)) + ) + kwargs = {"fp8_recipe": "blockwise"} + assert resolve_auto_fp8_recipe(kwargs) == "blockwise" + assert kwargs["fp8_recipe"] == "blockwise" + assert any("emulated" in w for w in warnings) + # The native recipe stays silent. + warnings.clear() + assert resolve_auto_fp8_recipe({"fp8_recipe": "mxfp8"}) == "mxfp8" + assert not warnings diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fake_int4_qat.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fake_int4_qat.py index d2b7d7121e..b8f46820d5 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fake_int4_qat.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_fake_int4_qat.py @@ -21,8 +21,10 @@ WorkerOutput, loss_fn_outputs_to_tensor, ) -from skyrl.backends.skyrl_train.workers.megatron import fake_int4_qat as fq_mod from skyrl.backends.skyrl_train.workers.megatron.megatron_worker import MegatronWorker +from skyrl.backends.skyrl_train.workers.megatron.quantization import ( + fake_int4_qat as fq_mod, +) from skyrl.train.config import SkyRLTrainConfig from skyrl.train.utils.utils import validate_cfg from tests.backends.skyrl_train.gpu.utils import ( @@ -47,7 +49,7 @@ def test_te_grouped_linear_still_exposes_get_weight_tensors(): assert hasattr(TEGroupedLinear, "_get_weight_tensors"), ( "TEGroupedLinear._get_weight_tensors is gone: the fake-INT4 QAT monkeypatch " - "(skyrl.backends.skyrl_train.workers.megatron.fake_int4_qat) no longer has an attach point. " + "(skyrl.backends.skyrl_train.workers.megatron.quantization.fake_int4_qat) no longer has an attach point. " "Find where the new TE version fetches weights in GroupedLinear.forward and re-target the patch." ) fwd_src = inspect.getsource(te_grouped_linear.GroupedLinear.forward) @@ -113,7 +115,7 @@ def _run_hook_effect_canary(): ) from megatron.core.transformer import TransformerConfig - from skyrl.backends.skyrl_train.workers.megatron.fake_int4_qat import ( + from skyrl.backends.skyrl_train.workers.megatron.quantization.fake_int4_qat import ( fake_int4_quantize_ste, install_fake_int4_qat, ) diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py index fe25b202de..2ce3e7d749 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py @@ -1,8 +1,17 @@ """ Run with: uv run --isolated --extra dev --extra megatron -- pytest -s tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_models.py + +The *_full_fp8 / *_fp8_param rows are Hopper-only (pytest.mark.h100): they run +blockwise FP8 on both Megatron (fp8=e4m3 + fp8_recipe=blockwise, plus +fp8_param=true persistent params for the fp8_param row) and vLLM +(quantization=fp8 fed by fp8_weight_sync_mode=blockwise), with FP32 +block scales (NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1, set by +_extra_env_vars_for_model). Select them with: -k "full_fp8 or fp8_param". """ +import os + import pytest import ray import torch @@ -12,6 +21,9 @@ WorkerOutput, loss_fn_outputs_to_tensor, ) +from skyrl.backends.skyrl_train.distributed.megatron.quantization_utils import ( + is_blackwell_or_newer, +) from skyrl.backends.skyrl_train.inference_servers.engine_utils import ( get_sampling_params_for_backend, ) @@ -86,15 +98,34 @@ def get_test_actor_config(model_name) -> SkyRLTrainConfig: return cfg -def _extra_env_vars_for_model(model_name: str) -> dict[str, str] | None: +def _extra_env_vars_for_model(model_name: str, fp8_mode: str | None = None) -> dict[str, str] | None: + env: dict[str, str] = {} # MLA models need cuDNN fused attention (the conftest globally sets # NVTE_FUSED_ATTN=0; re-enable it here so the fused backend is available). if "moonlight" in model_name.lower() or "glm-4" in model_name.lower(): - return {"NVTE_FUSED_ATTN": "1"} - return None + env["NVTE_FUSED_ATTN"] = "1" + if fp8_mode: + # Serialized-FP8 block-scale contract, mirroring what + # train/utils/utils.py pins in production (the test sets them + # explicitly because the fp8 fields are applied after + # get_test_actor_config's validate_cfg). Hopper: FP32 block scales + # end-to-end, and vLLM must not requantize wire scales to E8M0. + # Blackwell (SM100+): TE only supports power-of-2 block scales for + # blockwise quantization, and SM100 DeepGEMM only accepts E8M0 scale + # factors -- power-of-2 wire scales requantize to E8M0 losslessly. + if is_blackwell_or_newer(): + scale_mode, e8m0_mode = "0", "1" + else: + scale_mode, e8m0_mode = "1", "0" + env["NVTE_FP8_BLOCK_SCALING_FP32_SCALES"] = os.environ.get("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", scale_mode) + env["VLLM_USE_DEEP_GEMM_E8M0"] = os.environ.get("VLLM_USE_DEEP_GEMM_E8M0", e8m0_mode) + # fla's TileLang GDN backend aborts on Blackwell; fall back to Triton. + if "qwen3.5" in model_name.lower(): + env["FLA_TILELANG"] = os.environ.get("FLA_TILELANG", "0" if is_blackwell_or_newer() else "1") + return env or None -def _engine_overrides_for_model(model_name: str) -> dict: +def _engine_overrides_for_model(model_name: str, fp8_mode: str | None = None) -> dict: """Per-model overrides for vLLM engine init.""" overrides = {"engine_init_kwargs": {}, "gpu_memory_utilization": 0.9} if "Nemotron-3.5-Lightning" in model_name: @@ -108,6 +139,16 @@ def _engine_overrides_for_model(model_name: str) -> dict: # same GPU, so lower vLLM's pool footprint. if "qwen3.5-35b" in model_name.lower() and "tiny" not in model_name.lower(): overrides["gpu_memory_utilization"] = 0.5 + if fp8_mode: + # FP8 runs vLLM TP=1, so each rank holds the full ~35 GiB of FP8 + # weights; at gmu 0.5 on H100-80G the KV pool cannot cover the + # checkpoint's 262144 max_model_len. The test generates ~640 + # tokens per sequence. + overrides["engine_init_kwargs"]["max_model_len"] = 4096 + # GDN hybrid: one Mamba cache block per decode seq; the slim KV + # pool fits ~163 blocks, and the vLLM default max_num_seqs=1024 + # fails CUDA-graph capture. The test runs <= 80 concurrent seqs. + overrides["max_num_seqs"] = 128 if "glm-4.7-flash" in model_name.lower(): # GLM-4.7-Flash's 202k default context would size the KV pool far past # what is left next to the colocated Megatron policy shard. @@ -177,9 +218,9 @@ async def generate_with_vllm(generator, client, model_name, tokenizer, return_tr } ) training_input.metadata = {"response_length": num_actions} - return (response_mask, logprobs_t), training_input + return (response_mask, logprobs_t, generator_output), training_input else: - return (response_mask, logprobs_t) + return (response_mask, logprobs_t, generator_output) async def construct_training_input_from_generator_output(generator_output, tokenizer): @@ -195,10 +236,10 @@ async def construct_training_input_from_generator_output(generator_output, token @pytest.mark.asyncio @pytest.mark.megatron_models @pytest.mark.parametrize( - "tp,pp,cp,ep,etp,inference_tp,num_gpus,model_name,vllm_threshold,megatron_threshold", + "tp,pp,cp,ep,etp,inference_tp,num_gpus,model_name,vllm_threshold,megatron_threshold,fp8_mode", [ - pytest.param(2, 1, 1, 2, 1, 2, 4, "eatang/qwen3-moe-tiny-random", 1e-1, 2e-1, id="qwen3-moe_tp2_ep2"), - pytest.param(1, 2, 2, 1, None, 2, 4, "eatang/qwen3-moe-tiny-random", 1e-1, 2e-1, id="qwen3-moe_pp2_cp2"), + pytest.param(2, 1, 1, 2, 1, 2, 4, "eatang/qwen3-moe-tiny-random", 1e-1, 2e-1, None, id="qwen3-moe_tp2_ep2"), + pytest.param(1, 2, 2, 1, None, 2, 4, "eatang/qwen3-moe-tiny-random", 1e-1, 2e-1, None, id="qwen3-moe_pp2_cp2"), # GLM-4.7-Flash (~31B MoE, MLA) on 4xH100-80G. Mesh: TP=4 EP=4 ETP=1 # -> DP=1, vLLM TP=4 colocated on the same GPUs, same layout as the # other large-MoE entries below. @@ -213,6 +254,7 @@ async def construct_training_input_from_generator_output(generator_output, token "zai-org/GLM-4.7-Flash", 3e-1, 5e-2, + None, id="glm-4.7-flash_h100_tp4_ep4", marks=pytest.mark.h100, ), @@ -227,6 +269,7 @@ async def construct_training_input_from_generator_output(generator_output, token "eatang/qwen3.5-moe-tiny-random", 1e-1, 2e-1, + None, id="qwen3.5-moe_tp2_ep2", marks=pytest.mark.skip(reason="running into correctness issues for tiny qwen3.5"), ), @@ -244,6 +287,7 @@ async def construct_training_input_from_generator_output(generator_output, token "Qwen/Qwen3.5-0.8B", 1e-1, 5e-2, + None, id="qwen3.5-0.8b-dense_tp2", ), # Nemotron-3.5-Lightning (30B MoE, bf16) on 4xH100-80G. Same @@ -262,6 +306,7 @@ async def construct_training_input_from_generator_output(generator_output, token "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16", 5e-1, 5e-2, + None, id="nemotron3.5-lightning_tp4_ep4_h100", marks=pytest.mark.h100, ), @@ -280,18 +325,75 @@ async def construct_training_input_from_generator_output(generator_output, token "Qwen/Qwen3.5-35B-A3B", 3e-1, 5e-2, + None, id="qwen3.5-35b-a3b_h100_tp4_ep4", marks=pytest.mark.h100, ), + # Full-FP8 rows: blockwise FP8 Megatron compute + FP8 vLLM rollout fed + # by serialized blockwise weight sync; the fp8_param row additionally + # keeps persistent FP8 Megatron params with exact optimizer-master + # init from unquantized checkpoint shards. Hopper-only: the wire + # contract and fp8_param require FP32 block scales + # (NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1); Blackwell runs power-of-2 + # scales with fp8_param=false. Thresholds mirror the matching bf16 + # rows; tune as we accumulate measured diffs. + pytest.param( + 2, + 1, + 1, + 1, + None, + 2, + 2, + "Qwen/Qwen3.5-0.8B", + 1e-1, + 5e-2, + "full_fp8", + id="qwen3.5-0.8b-dense_tp2_full_fp8", + marks=pytest.mark.h100, + ), + pytest.param( + 2, + 1, + 1, + 1, + None, + 2, + 2, + "Qwen/Qwen3.5-0.8B", + 1e-1, + 5e-2, + "fp8_param", + id="qwen3.5-0.8b-dense_tp2_fp8_param", + marks=pytest.mark.h100, + ), + # TP=1 x 4 engines mirrors the production layout: Megatron TP/EP shards + # feed full-width vLLM ranks. Blockwise FP8 also builds at inference + # TP=2/4, since the vision blocks sit on the FP8 ignore list. + pytest.param( + 4, + 1, + 1, + 4, + 1, + 1, + 4, + "Qwen/Qwen3.5-35B-A3B", + 3e-1, + 5e-2, + "full_fp8", + id="qwen3.5-35b-a3b_h100_tp4_ep4_full_fp8", + marks=pytest.mark.h100, + ), ], ) async def test_logprobs_matching_roundtrip( - tp, pp, cp, ep, etp, inference_tp, num_gpus, model_name, vllm_threshold, megatron_threshold + tp, pp, cp, ep, etp, inference_tp, num_gpus, model_name, vllm_threshold, megatron_threshold, fp8_mode ): """ Check that logprob diff matches acrosss vllm and megatron. """ - with ray_init(extra_env_vars=_extra_env_vars_for_model(model_name)): + with ray_init(extra_env_vars=_extra_env_vars_for_model(model_name, fp8_mode)): cfg = get_test_actor_config(model_name=model_name) cfg.trainer.strategy = "megatron" cfg.generator.inference_engine.tensor_parallel_size = inference_tp @@ -304,10 +406,38 @@ async def test_logprobs_matching_roundtrip( cfg.generator.batched = False cfg.generator.max_turns = 1 + if fp8_mode: + # Megatron: blockwise FP8 compute; the fp8_param variant keeps + # persistent FP8 params (requires fp8_param_gather so updated FP32 + # masters requantize into the FP8 compute weights). + mcfg = cfg.trainer.policy.megatron_config + transformer_config_kwargs = dict(mcfg.transformer_config_kwargs or {}) + transformer_config_kwargs.update( + { + "fp8": "e4m3", + "fp8_recipe": "blockwise", + "fp8_amax_compute_algo": "most_recent", + "fp8_param": fp8_mode == "fp8_param", + } + ) + mcfg.transformer_config_kwargs = transformer_config_kwargs + if fp8_mode == "fp8_param": + mcfg.ddp_config.fp8_param_gather = True + # vLLM: FP8 rollout fed by serialized blockwise weight sync + # (_apply_serialized_fp8_weight_sync_defaults injects + # quantization=fp8, load_format=dummy and the blockwise + # quantization_config into the engine kwargs). + cfg.generator.inference_engine.fp8_weight_sync_mode = "blockwise" + # The validated FP8 production runs use the mp executor; with the + # ray executor, vLLM 0.23's ray_executor_v2 ignores + # VLLM_RAY_BUNDLE_INDICES, so multi-engine colocate (e.g. the 35B + # row's 4 x TP=1) stacks every engine's worker on GPU 0 and OOMs. + cfg.generator.inference_engine.distributed_executor_backend = "mp" + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) tokenizer.pad_token = tokenizer.eos_token - engine_overrides = _engine_overrides_for_model(model_name) + engine_overrides = _engine_overrides_for_model(model_name, fp8_mode) async with InferenceEngineState.create( cfg=cfg, model=model_name, @@ -317,9 +447,9 @@ async def test_logprobs_matching_roundtrip( sleep_level=2, # full sleep — this test explicitly syncs weights gpu_memory_utilization=engine_overrides["gpu_memory_utilization"], engine_init_kwargs=engine_overrides["engine_init_kwargs"], + max_num_seqs=engine_overrides.get("max_num_seqs"), ) as engines: client, pg = engines.client, engines.pg - await client.wake_up() generator = SkyRLGymGenerator( generator_cfg=cfg.generator, @@ -328,10 +458,6 @@ async def test_logprobs_matching_roundtrip( tokenizer=tokenizer, ) - (response_mask, logprobs_t), training_input = await generate_with_vllm( - generator, client, model_name, tokenizer, return_training_input=True - ) - await client.sleep() cfg.trainer.placement.policy_num_gpus_per_node = num_gpus cfg.trainer.policy.megatron_config.tensor_model_parallel_size = tp cfg.trainer.policy.megatron_config.pipeline_model_parallel_size = pp @@ -341,18 +467,59 @@ async def test_logprobs_matching_roundtrip( cfg.trainer.micro_forward_batch_size_per_gpu = 2 cfg.trainer.micro_train_batch_size_per_gpu = 2 - policy = init_worker_with_type( - "policy", - shared_pg=pg, - colocate_all=True, - num_gpus_per_node=num_gpus, - cfg=cfg, - ) - ray.get( - policy.async_run_ray_method( - "pass_through", "init_weight_sync_state", client, cfg.generator.inference_engine + policy = None + if fp8_mode: + # Serialized FP8 boots vLLM with load_format="dummy", so real + # weights must be synced from Megatron before generating + # (mirrors the trainer, which always syncs before the first + # rollout). Build the policy with the engines asleep, then + # run the same offload/wake/broadcast dance as the sync below. + await client.sleep() + policy = init_worker_with_type( + "policy", + shared_pg=pg, + colocate_all=True, + num_gpus_per_node=num_gpus, + cfg=cfg, + ) + ray.get( + policy.async_run_ray_method( + "pass_through", "init_weight_sync_state", client, cfg.generator.inference_engine + ) ) + policy.offload_to_cpu(offload_optimizer=True, offload_model=False) + await client.wake_up(tags=["weights"]) + with Timer("initial_sync_weights"): + ray.get( + policy.async_run_ray_method( + "pass_through", "broadcast_to_inference_engines", client, cfg.generator.inference_engine + ) + ) + policy.offload_to_cpu(offload_optimizer=False, offload_model=True) + await client.wake_up(tags=["kv_cache"]) + else: + await client.wake_up() + + (response_mask, logprobs_t, gen_out_1), training_input = await generate_with_vllm( + generator, client, model_name, tokenizer, return_training_input=True ) + await client.sleep() + + if policy is None: + policy = init_worker_with_type( + "policy", + shared_pg=pg, + colocate_all=True, + num_gpus_per_node=num_gpus, + cfg=cfg, + ) + ray.get( + policy.async_run_ray_method( + "pass_through", "init_weight_sync_state", client, cfg.generator.inference_engine + ) + ) + else: + policy.backload_to_gpu(backload_optimizer=False, backload_model=True) refs = policy.async_run_ray_method("mesh", "forward", data=training_input) results = ray.get(refs) @@ -387,33 +554,64 @@ async def test_logprobs_matching_roundtrip( policy.offload_to_cpu(offload_optimizer=False, offload_model=True) await client.wake_up(tags=["kv_cache"]) - response_mask_2, logprobs_t_2 = await generate_with_vllm( + response_mask_2, logprobs_t_2, gen_out_2 = await generate_with_vllm( generator, client, model_name, tokenizer, return_training_input=False ) - logprobs_t_valid = logprobs_t[response_mask.bool()] - logprobs_t_2_valid = logprobs_t_2[response_mask_2.bool()] - - # Pre- and post-sync are two independent sampled generations - # so truncate to the shorter sequence for the magnitude check. - if logprobs_t_valid.shape[0] != logprobs_t_2_valid.shape[0]: - min_len = min(logprobs_t_valid.shape[0], logprobs_t_2_valid.shape[0]) + if fp8_mode: + # In the FP8 flow both generations ran on identical synced + # weights, so compare logprobs only on each sequence's common + # prefix: once greedy decoding diverges at a near-tie token, + # later positions score different tokens and their diff is + # pure noise (measured up to ~0.14 mean on identical weights, + # vs ~1e-3 on common prefixes). + ids_1, lp_1 = gen_out_1["response_ids"], gen_out_1["rollout_logprobs"] + ids_2, lp_2 = gen_out_2["response_ids"], gen_out_2["rollout_logprobs"] + assert lp_1 is not None and lp_2 is not None, "resync check needs rollout logprobs" + diffs = [] + divergent = 0 + for s1, s2, l1, l2 in zip(ids_1, ids_2, lp_1, lp_2): + n = 0 + for a, b in zip(s1, s2): + if a != b: + break + n += 1 + if n < min(len(s1), len(s2)): + divergent += 1 + diffs.extend(abs(x - y) for x, y in zip(l1[:n], l2[:n])) + assert diffs, "no common-prefix tokens between pre/post-sync generations" + logprobs_diff = torch.tensor(diffs) print( - f"NOTE: pre/post-sync generation lengths differ " - f"({logprobs_t_valid.shape[0]} vs {logprobs_t_2_valid.shape[0]}); " - f"truncating to {min_len} for the magnitude check." + f"vLLM resync common-prefix logprob diff mean: {logprobs_diff.mean().item():.6f}, " + f"std: {logprobs_diff.std().item():.6f} over {len(diffs)} tokens " + f"({divergent}/{len(ids_1)} sequences diverged at a near-tie token)" ) - logprobs_t_valid = logprobs_t_valid[:min_len] - logprobs_t_2_valid = logprobs_t_2_valid[:min_len] + else: + logprobs_t_valid = logprobs_t[response_mask.bool()] + logprobs_t_2_valid = logprobs_t_2[response_mask_2.bool()] + + # Pre- and post-sync are two independent sampled generations + # so truncate to the shorter sequence for the magnitude check. + if logprobs_t_valid.shape[0] != logprobs_t_2_valid.shape[0]: + min_len = min(logprobs_t_valid.shape[0], logprobs_t_2_valid.shape[0]) + print( + f"NOTE: pre/post-sync generation lengths differ " + f"({logprobs_t_valid.shape[0]} vs {logprobs_t_2_valid.shape[0]}); " + f"truncating to {min_len} for the magnitude check." + ) + logprobs_t_valid = logprobs_t_valid[:min_len] + logprobs_t_2_valid = logprobs_t_2_valid[:min_len] - logprobs_diff = (logprobs_t_valid - logprobs_t_2_valid).abs() - print( - f"vLLM logprobs - mean: {logprobs_t_valid.mean().item():.6f}, std: {logprobs_t_valid.std().item():.6f}" - ) - print( - f"vLLM logprobs after sync - mean: {logprobs_t_2_valid.mean().item():.6f}, std: {logprobs_t_2_valid.std().item():.6f}" - ) - print(f"vLLM logprob diff mean: {logprobs_diff.mean().item():.6f}, std: {logprobs_diff.std().item():.6f}") + logprobs_diff = (logprobs_t_valid - logprobs_t_2_valid).abs() + print( + f"vLLM logprobs - mean: {logprobs_t_valid.mean().item():.6f}, std: {logprobs_t_valid.std().item():.6f}" + ) + print( + f"vLLM logprobs after sync - mean: {logprobs_t_2_valid.mean().item():.6f}, std: {logprobs_t_2_valid.std().item():.6f}" + ) + print( + f"vLLM logprob diff mean: {logprobs_diff.mean().item():.6f}, std: {logprobs_diff.std().item():.6f}" + ) assert ( logprobs_diff.mean().item() < vllm_threshold ), f"Logprob diff should be less than {vllm_threshold}, but is {logprobs_diff.mean().item():.6f}" diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_sft_packing_parity.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_sft_packing_parity.py index 5bff36356d..186e89ecc0 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_sft_packing_parity.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_sft_packing_parity.py @@ -22,6 +22,8 @@ from skyrl.backends.skyrl_train.distributed.dispatch import WorkerOutput from skyrl.backends.skyrl_train.distributed.megatron.packing_utils import ( get_packed_seq_align_size, +) +from skyrl.backends.skyrl_train.distributed.megatron.quantization_utils import ( is_fp8_enabled, ) from skyrl.backends.skyrl_train.training_batch import TensorList, TrainingInputBatch diff --git a/tests/backends/skyrl_train/inference_servers/test_batched_moe_fp8_reload.py b/tests/backends/skyrl_train/inference_servers/test_batched_moe_fp8_reload.py new file mode 100644 index 0000000000..bc0cc67d1e --- /dev/null +++ b/tests/backends/skyrl_train/inference_servers/test_batched_moe_fp8_reload.py @@ -0,0 +1,96 @@ +from types import SimpleNamespace + +import pytest +import torch + +from skyrl.backends.skyrl_train.inference_servers.new_inference_worker_wrap import ( + _load_batched_moe_fp8_tensor, +) +from skyrl.backends.skyrl_train.weight_sync.fp8 import ( + SKYRL_BATCHED_MOE_FP8_PREFIX, +) + + +def test_batched_moe_tensor_uses_one_full_expert_loader_call(): + calls = [] + + def weight_loader(param, loaded_weight, weight_name, *, shard_id, expert_id, return_success): + calls.append((param, loaded_weight, weight_name, shard_id, expert_id, return_success)) + return True + + weight_loader.supports_moe_loading = True + param = torch.nn.Parameter(torch.empty(3, 8, 4), requires_grad=False) + param.weight_loader = weight_loader + target_name = "model.layers.0.mlp.experts.w13_weight" + loaded_weight = torch.randn(3, 4, 4) + wire_name = f"{SKYRL_BATCHED_MOE_FP8_PREFIX}model.layers.0.mlp.experts.gate_proj.weight" + + loaded = _load_batched_moe_fp8_tensor( + SimpleNamespace(), + {target_name: param}, + wire_name, + loaded_weight, + ) + + assert loaded + assert len(calls) == 1 + assert calls[0][1] is loaded_weight + assert calls[0][2:] == (target_name, "w1", 0, True) + + +def test_batched_moe_scale_maps_to_fused_scale_parameter(): + calls = [] + + def weight_loader(param, loaded_weight, weight_name, *, shard_id, expert_id, return_success): + calls.append((weight_name, shard_id, tuple(loaded_weight.shape))) + return True + + weight_loader.supports_moe_loading = True + param = torch.nn.Parameter(torch.empty(2, 6, 3), requires_grad=False) + param.weight_loader = weight_loader + target_name = "language_model.model.layers.2.mlp.experts.w13_weight_scale_inv" + loaded_weight = torch.randn(2, 3, 3) + mapper = SimpleNamespace( + apply_list=lambda names: [names[0].replace("model.language_model.", "language_model.model.", 1)] + ) + model = SimpleNamespace(hf_to_vllm_mapper=mapper) + wire_name = f"{SKYRL_BATCHED_MOE_FP8_PREFIX}" "model.language_model.layers.2.mlp.experts.up_proj.weight_scale_inv" + + assert _load_batched_moe_fp8_tensor(model, {target_name: param}, wire_name, loaded_weight) + assert calls == [(target_name, "w3", (2, 3, 3))] + + +def test_batched_moe_resolves_routed_experts_nesting(): + """vLLM 0.26: expert params live on MoERunner's RoutedExperts submodule.""" + + calls = [] + + def weight_loader(param, loaded_weight, weight_name, *, shard_id, expert_id, return_success): + calls.append((weight_name, shard_id, expert_id)) + return True + + weight_loader.supports_moe_loading = True + param = torch.nn.Parameter(torch.empty(3, 8, 4), requires_grad=False) + param.weight_loader = weight_loader + nested_name = "language_model.model.layers.0.mlp.experts.routed_experts.w13_weight" + mapper = SimpleNamespace( + apply_list=lambda names: [names[0].replace("model.language_model.", "language_model.model.", 1)] + ) + model = SimpleNamespace(hf_to_vllm_mapper=mapper) + wire_name = f"{SKYRL_BATCHED_MOE_FP8_PREFIX}model.language_model.layers.0.mlp.experts.gate_proj.weight" + + loaded = _load_batched_moe_fp8_tensor(model, {nested_name: param}, wire_name, torch.randn(3, 4, 4)) + + assert loaded + assert calls == [(nested_name, "w1", 0)] + + +def test_batched_moe_missing_target_error_names_both_candidates(): + wire_name = f"{SKYRL_BATCHED_MOE_FP8_PREFIX}model.layers.0.mlp.experts.down_proj.weight" + + with pytest.raises(ValueError) as excinfo: + _load_batched_moe_fp8_tensor(SimpleNamespace(), {}, wire_name, torch.randn(2, 4, 4)) + + message = str(excinfo.value) + assert "model.layers.0.mlp.experts.w2_weight" in message + assert "model.layers.0.mlp.experts.routed_experts.w2_weight" in message diff --git a/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py b/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py index f8ca12b1b0..1a2c32bf06 100644 --- a/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py +++ b/tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py @@ -1,10 +1,12 @@ """Tests for build_vllm_cli_args on GPU-less hosts.""" from argparse import Namespace +from types import SimpleNamespace import pytest from skyrl.backends.skyrl_train.inference_servers.utils import ( + _apply_serialized_fp8_weight_sync_defaults, build_vllm_cli_args, get_pd_cli_args, get_pd_p2p_connector_name, @@ -13,6 +15,114 @@ from skyrl.train.config import SkyRLTrainConfig +def test_serialized_fp8_weight_sync_defaults_configure_vllm_checkpoint_fp8(monkeypatch): + import skyrl.backends.skyrl_train.inference_servers.utils as inference_utils + + monkeypatch.setattr(inference_utils, "_serialized_fp8_ignored_layers", lambda _model_path: []) + cfg = SkyRLTrainConfig() + ie_cfg = cfg.generator.inference_engine + ie_cfg.fp8_weight_sync_mode = "blockwise" + engine_kwargs = {"hf_overrides": {"rope_theta": 10000.0}} + + _apply_serialized_fp8_weight_sync_defaults(ie_cfg, engine_kwargs, model_path="qwen35-test") + + assert engine_kwargs["quantization"] == "fp8" + assert engine_kwargs["load_format"] == "dummy" + assert engine_kwargs["hf_overrides"]["rope_theta"] == 10000.0 + assert engine_kwargs["hf_overrides"]["quantization_config"] == { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [128, 128], + } + + +@pytest.mark.parametrize( + "engine_kwargs", + [ + {"quantization": "awq"}, + {"load_format": "safetensors"}, + {"hf_overrides": {"quantization_config": {"weight_block_size": [64, 128]}}}, + ], +) +def test_serialized_fp8_weight_sync_rejects_conflicting_vllm_settings(engine_kwargs, monkeypatch): + import skyrl.backends.skyrl_train.inference_servers.utils as inference_utils + + monkeypatch.setattr(inference_utils, "_serialized_fp8_ignored_layers", lambda _model_path: []) + cfg = SkyRLTrainConfig() + cfg.generator.inference_engine.fp8_weight_sync_mode = "blockwise" + + with pytest.raises(ValueError, match="FP8 weight sync"): + _apply_serialized_fp8_weight_sync_defaults( + cfg.generator.inference_engine, + engine_kwargs, + model_path="qwen35-test", + ) + + +@pytest.mark.parametrize( + "engine_kwargs", + [ + {"hf_overrides": []}, + {"hf_overrides": {"quantization_config": []}}, + ], +) +def test_serialized_fp8_weight_sync_rejects_non_mapping_overrides(engine_kwargs): + cfg = SkyRLTrainConfig() + cfg.generator.inference_engine.fp8_weight_sync_mode = "blockwise" + + with pytest.raises(ValueError, match="must be a dict"): + _apply_serialized_fp8_weight_sync_defaults( + cfg.generator.inference_engine, + engine_kwargs, + model_path="qwen35-test", + ) + + +def test_serialized_fp8_requires_model_path(): + cfg = SkyRLTrainConfig() + cfg.generator.inference_engine.fp8_weight_sync_mode = "blockwise" + + with pytest.raises(ValueError, match="model path is required"): + _apply_serialized_fp8_weight_sync_defaults(cfg.generator.inference_engine, {}) + + +def test_serialized_fp8_fails_when_model_config_cannot_be_inspected(monkeypatch): + import transformers + + cfg = SkyRLTrainConfig() + cfg.generator.inference_engine.fp8_weight_sync_mode = "blockwise" + + def fail_config_load(*_args, **_kwargs): + raise OSError("missing config") + + monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", fail_config_load) + with pytest.raises(RuntimeError, match="Could not inspect the model config"): + _apply_serialized_fp8_weight_sync_defaults( + cfg.generator.inference_engine, + {}, + model_path="missing-model", + ) + + +def test_serialized_fp8_rejects_unsupported_model_layout(monkeypatch): + import transformers + + cfg = SkyRLTrainConfig() + cfg.generator.inference_engine.fp8_weight_sync_mode = "blockwise" + monkeypatch.setattr( + transformers.AutoConfig, + "from_pretrained", + lambda *_args, **_kwargs: SimpleNamespace(model_type="llama"), + ) + + with pytest.raises(ValueError, match="no registered model spec"): + _apply_serialized_fp8_weight_sync_defaults( + cfg.generator.inference_engine, + {}, + model_path="unsupported-model", + ) + + @pytest.mark.vllm def test_build_vllm_cli_args_succeeds_on_gpu_less_host(monkeypatch): import vllm.platforms diff --git a/tests/backends/skyrl_train/weight_sync/fp8/__init__.py b/tests/backends/skyrl_train/weight_sync/fp8/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/backends/skyrl_train/weight_sync/fp8/test_serialized_fp8.py b/tests/backends/skyrl_train/weight_sync/fp8/test_serialized_fp8.py new file mode 100644 index 0000000000..98244b61f9 --- /dev/null +++ b/tests/backends/skyrl_train/weight_sync/fp8/test_serialized_fp8.py @@ -0,0 +1,315 @@ +from types import SimpleNamespace + +import pytest +import torch + +from skyrl.backends.skyrl_train.weight_sync.fp8 import ( + SKYRL_BATCHED_MOE_FP8_PREFIX, + SerializedFp8Config, + batched_blockwise_cast_to_fp8, + batched_moe_wire_targets, + blockwise_cast_to_fp8, + get_serialized_fp8_quantization_config, + iter_serialized_fp8_tensors, + resolve_fp8_spec, +) +from skyrl.backends.skyrl_train.weight_sync.fp8.models import QWEN35_FP8_SPEC + + +def _qwen35_config(**kwargs) -> SerializedFp8Config: + return SerializedFp8Config(spec=QWEN35_FP8_SPEC, **kwargs) + + +def test_blockwise_cast_to_fp8_emits_weight_and_fp32_scale(): + weight = torch.arange(257 * 129, dtype=torch.float32).reshape(257, 129) / 1000 + + q_weight, scale = blockwise_cast_to_fp8(weight, [128, 128]) + + assert q_weight.shape == weight.shape + assert q_weight.dtype == torch.float8_e4m3fn + assert scale.shape == (3, 2) + assert scale.dtype == torch.float32 + + +def test_blockwise_cast_defaults_to_exact_fp32_scales(): + torch.manual_seed(7) + weight = torch.randn(256, 256, dtype=torch.float32) + + default_weight, default_scale = blockwise_cast_to_fp8(weight, [128, 128]) + exact_weight, exact_scale = blockwise_cast_to_fp8(weight, [128, 128], power_2_scale=False) + + assert torch.equal(default_weight.view(torch.uint8), exact_weight.view(torch.uint8)) + assert torch.equal(default_scale, exact_scale) + + +def test_blockwise_cast_floors_scale_for_all_zero_blocks(monkeypatch): + monkeypatch.setenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "1") + config = SerializedFp8Config() + weight = torch.zeros((128, 128), dtype=torch.float32) + + q_weight, scale = blockwise_cast_to_fp8( + weight, + config.weight_block_size, + config.power_2_scale, + ) + + # All-zero blocks must not degenerate the scale to zero/inf. + assert scale.item() == pytest.approx(1e-10 / torch.finfo(torch.float8_e4m3fn).max) + assert torch.isfinite(scale).all() + assert q_weight.view(torch.uint8).eq(0).all() + + +def test_blockwise_cast_pow2_scales_match_te_ue8m0_rule(): + # TE's UE8M0 mode rounds dequantization scales up to powers of two. + torch.manual_seed(0) + weight = torch.randn(256, 384, dtype=torch.float32) + + _, pow2_scale = blockwise_cast_to_fp8(weight, [128, 128], power_2_scale=True) + _, exact_scale = blockwise_cast_to_fp8(weight, [128, 128], power_2_scale=False) + + log2 = torch.log2(pow2_scale) + assert torch.allclose(log2, log2.round(), atol=0.0) + expected = torch.pow(2.0, torch.ceil(torch.log2(exact_scale))) + assert torch.allclose(pow2_scale, expected) + assert torch.all(pow2_scale >= exact_scale) + + +def test_serialized_config_power_2_scale_follows_te_env(monkeypatch): + monkeypatch.delenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", raising=False) + assert SerializedFp8Config().power_2_scale is False + + monkeypatch.setenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "1") + assert SerializedFp8Config().power_2_scale is False + + monkeypatch.setenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "0") + assert SerializedFp8Config().power_2_scale is True + + monkeypatch.setenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "invalid") + with pytest.raises(ValueError, match="must be '0'.*or '1'"): + SerializedFp8Config() + + +@pytest.mark.parametrize("block_size", [[], [128], [128, 128, 128], [128, 0], [128, 1.5], [True, 128]]) +def test_blockwise_cast_rejects_invalid_block_size(block_size): + with pytest.raises(ValueError, match="exactly two positive integers"): + blockwise_cast_to_fp8(torch.ones((2, 2)), block_size) + + +def test_quantizable_weight_filter_keeps_embeddings_in_target_dtype(): + config = _qwen35_config() + linear = (256, 256) + embedding_tensor = torch.ones((32000, 256), dtype=torch.bfloat16) + + should_quantize = QWEN35_FP8_SPEC.should_quantize + assert should_quantize("model.layers.0.mlp.down_proj.weight", linear) + assert should_quantize("model.layers.0.linear_attn.in_proj_qkv.weight", linear) + assert should_quantize("model.layers.0.linear_attn.in_proj_z.weight", linear) + assert should_quantize("model.layers.0.linear_attn.out_proj.weight", linear) + assert not should_quantize("model.layers.0.linear_attn.conv1d.weight", linear) + assert not should_quantize("model.layers.0.linear_attn.in_proj_b.weight", linear) + assert not should_quantize("model.layers.0.linear_attn.in_proj_a.weight", linear) + assert not should_quantize("model.embed_tokens.weight", tuple(embedding_tensor.shape)) + + tensors = list( + iter_serialized_fp8_tensors( + "model.embed_tokens.weight", + embedding_tensor, + torch.bfloat16, + config, + ) + ) + assert [(name, tensor.dtype) for name, tensor in tensors] == [("model.embed_tokens.weight", torch.bfloat16)] + + +def test_vllm_serialized_fp8_quantization_config(): + assert get_serialized_fp8_quantization_config() == { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [128, 128], + } + assert get_serialized_fp8_quantization_config(ignored_layers=["model.layers.0.linear_attn.in_proj_b"]) == { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [128, 128], + "ignored_layers": ["model.layers.0.linear_attn.in_proj_b"], + } + + +def test_qwen35_fp8_ignored_layers_use_linear_attention_layers(): + hf_config = SimpleNamespace( + model_type="qwen3_5_text", + layer_types=[ + "linear_attention", + "full_attention", + "linear_attention", + ], + ) + + assert QWEN35_FP8_SPEC.ignored_layers(hf_config) == [ + "model.layers.0.linear_attn.in_proj_b", + "model.layers.0.linear_attn.in_proj_a", + "model.language_model.layers.0.linear_attn.in_proj_b", + "model.language_model.layers.0.linear_attn.in_proj_a", + "model.layers.2.linear_attn.in_proj_b", + "model.layers.2.linear_attn.in_proj_a", + "model.language_model.layers.2.linear_attn.in_proj_b", + "model.language_model.layers.2.linear_attn.in_proj_a", + ] + + +def test_qwen35_ignored_layers_include_only_checkpoint_vision_prefixes(): + hf_config = SimpleNamespace( + model_type="qwen3_5", + text_config=SimpleNamespace(model_type="qwen3_5_text", layer_types=[]), + vision_config=SimpleNamespace(depth=2), + ) + + assert QWEN35_FP8_SPEC.ignored_layers(hf_config) == [ + "model.visual.blocks.0.attn.proj", + "model.visual.blocks.0.mlp.linear_fc1", + "model.visual.blocks.0.mlp.linear_fc2", + "model.visual.blocks.1.attn.proj", + "model.visual.blocks.1.mlp.linear_fc1", + "model.visual.blocks.1.mlp.linear_fc2", + ] + + +def test_qwen35_ignored_layers_are_not_inferred_from_unrelated_hybrid_config(): + hf_config = SimpleNamespace(model_type="unrelated_hybrid", layer_types=["linear_attention"]) + + assert QWEN35_FP8_SPEC.ignored_layers(hf_config) == [] + + +def test_moe_batched_expert_spec_recognizes_and_splits_gate_up(): + base = "model.language_model.layers.5.mlp.experts" + gate_up = QWEN35_FP8_SPEC.moe_expert_spec(f"{base}.gate_up_proj") + assert gate_up.experts_base == base + assert [p.hf_name for p in gate_up.projections] == ["gate_proj", "up_proj"] + assert gate_up.split_dim == 1 + down = QWEN35_FP8_SPEC.moe_expert_spec(f"{base}.down_proj") + assert down.experts_base == base + assert [p.hf_name for p in down.projections] == ["down_proj"] + assert down.split_dim is None + assert QWEN35_FP8_SPEC.moe_expert_spec("model.layers.5.mlp.gate.weight") is None + assert QWEN35_FP8_SPEC.moe_expert_spec("model.layers.5.self_attn.q_proj.weight") is None + + +def test_moe_shared_expert_quantized_router_and_gate_bf16(): + lin = (256, 256) + assert QWEN35_FP8_SPEC.should_quantize("model.layers.3.mlp.shared_expert.gate_proj.weight", lin) + assert QWEN35_FP8_SPEC.should_quantize("model.layers.3.mlp.shared_expert.up_proj.weight", lin) + assert QWEN35_FP8_SPEC.should_quantize("model.layers.3.mlp.shared_expert.down_proj.weight", lin) + assert not QWEN35_FP8_SPEC.should_quantize("model.layers.3.mlp.gate.weight", (256, 8)) + assert not QWEN35_FP8_SPEC.should_quantize("model.layers.3.mlp.shared_expert_gate.weight", (256, 1)) + + +def test_batched_blockwise_cast_matches_independent_expert_casts(): + torch.manual_seed(11) + weight = torch.randn(5, 257, 129, dtype=torch.bfloat16) + + q_batched, scale_batched = batched_blockwise_cast_to_fp8( + weight, + [128, 128], + power_2_scale=True, + expert_batch_size=2, + ) + + for expert_id in range(weight.shape[0]): + q_expected, scale_expected = blockwise_cast_to_fp8( + weight[expert_id], + [128, 128], + power_2_scale=True, + ) + assert torch.equal(q_batched[expert_id].view(torch.uint8), q_expected.view(torch.uint8)) + assert torch.equal(scale_batched[expert_id], scale_expected) + + +def test_batched_moe_experts_remain_fused_with_pow2_scales(): + num_experts, moe_inter, hidden = 3, 128, 256 + config = _qwen35_config(weight_block_size=(128, 128), power_2_scale=True) + base = "model.language_model.layers.7.mlp.experts" + + torch.manual_seed(0) + gate_up = torch.randn(num_experts, 2 * moe_inter, hidden, dtype=torch.bfloat16) + emitted = list(iter_serialized_fp8_tensors(f"{base}.gate_up_proj", gate_up, torch.bfloat16, config)) + by_name = dict(emitted) + assert len(emitted) == 4 + for projection_idx, proj in enumerate(("gate_proj", "up_proj")): + weight_name = f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{base}.{proj}.weight" + w = by_name[weight_name] + s = by_name[f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{base}.{proj}.weight_scale_inv"] + assert w.dtype == torch.float8_e4m3fn and tuple(w.shape) == (num_experts, moe_inter, hidden) + assert s.dtype == torch.float32 and tuple(s.shape) == (num_experts, 1, 2) + log2 = torch.log2(s) + assert torch.allclose(log2, log2.round(), atol=0.0) + for expert_id in range(num_experts): + source = gate_up[expert_id, projection_idx * moe_inter : (projection_idx + 1) * moe_inter] + q_expected, scale_expected = blockwise_cast_to_fp8( + source, + config.weight_block_size, + config.power_2_scale, + ) + assert torch.equal(w[expert_id].view(torch.uint8), q_expected.view(torch.uint8)) + assert torch.equal(s[expert_id], scale_expected) + + down = torch.randn(num_experts, hidden, moe_inter, dtype=torch.bfloat16) + emitted_d = dict(iter_serialized_fp8_tensors(f"{base}.down_proj", down, torch.bfloat16, config)) + down_weight_name = f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{base}.down_proj.weight" + assert set(emitted_d) == { + down_weight_name, + f"{SKYRL_BATCHED_MOE_FP8_PREFIX}{base}.down_proj.weight_scale_inv", + } + assert emitted_d[down_weight_name].shape == down.shape + + +def test_batched_moe_gate_up_rejects_odd_output_dimension(): + name = "model.layers.0.mlp.experts.gate_up_proj" + tensor = torch.randn(2, 255, 128, dtype=torch.bfloat16) + + with pytest.raises(ValueError, match="must split evenly"): + list(iter_serialized_fp8_tensors(name, tensor, torch.bfloat16, _qwen35_config())) + + +def test_batched_moe_expert_iterator_rejects_non_3d_input(): + name = "model.layers.0.mlp.experts.down_proj" + with pytest.raises(ValueError, match="must be 3D"): + list( + iter_serialized_fp8_tensors( + name, + torch.ones((128, 128)), + torch.bfloat16, + _qwen35_config(), + ) + ) + + +def test_resolve_fp8_spec_matches_qwen35_and_rejects_unknown_layouts(): + assert resolve_fp8_spec(SimpleNamespace(model_type="qwen3_5_moe_text")) is QWEN35_FP8_SPEC + assert ( + resolve_fp8_spec(SimpleNamespace(model_type="qwen3_5", text_config=SimpleNamespace(model_type="qwen3_5_text"))) + is QWEN35_FP8_SPEC + ) + assert resolve_fp8_spec(SimpleNamespace(model_type="llama")) is None + + +def test_batched_moe_wire_targets_derive_from_registered_specs(): + assert batched_moe_wire_targets() == { + ".experts.gate_proj.weight": (".experts.w13_weight", "w1"), + ".experts.up_proj.weight": (".experts.w13_weight", "w3"), + ".experts.down_proj.weight": (".experts.w2_weight", "w2"), + ".experts.gate_proj.weight_scale_inv": (".experts.w13_weight_scale_inv", "w1"), + ".experts.up_proj.weight_scale_inv": (".experts.w13_weight_scale_inv", "w3"), + ".experts.down_proj.weight_scale_inv": (".experts.w2_weight_scale_inv", "w2"), + } + + +def test_iter_serialized_fp8_tensors_requires_model_spec(): + with pytest.raises(ValueError, match="spec is not set"): + list( + iter_serialized_fp8_tensors( + "model.layers.0.mlp.down_proj.weight", + torch.ones((256, 256), dtype=torch.bfloat16), + torch.bfloat16, + SerializedFp8Config(), + ) + ) diff --git a/tests/backends/skyrl_train/weight_sync/test_send_chunks_protocol.py b/tests/backends/skyrl_train/weight_sync/test_send_chunks_protocol.py new file mode 100644 index 0000000000..1c55cdbc90 --- /dev/null +++ b/tests/backends/skyrl_train/weight_sync/test_send_chunks_protocol.py @@ -0,0 +1,51 @@ +"""Every sender must implement the full ``send_chunks`` protocol. + +``megatron_worker`` passes ``derive_metadata_from_chunks`` to whichever sender +strategy is active, so a sender that omits it from its signature raises +``TypeError`` on every Megatron sync — even with FP8 off. +""" + +import asyncio +import inspect + +import pytest + +from skyrl.backends.skyrl_train.weight_sync.broadcast_strategy import ( + BroadcastWeightTransferSender, +) +from skyrl.backends.skyrl_train.weight_sync.cuda_ipc_strategy import ( + CudaIpcWeightTransferSender, +) +from skyrl.backends.skyrl_train.weight_sync.delta_strategy import ( + DeltaWeightTransferSender, +) + +SENDER_CLASSES = ( + BroadcastWeightTransferSender, + CudaIpcWeightTransferSender, + DeltaWeightTransferSender, +) + + +@pytest.mark.parametrize("sender_cls", SENDER_CLASSES, ids=lambda c: c.__name__) +def test_send_chunks_accepts_protocol_arguments(sender_cls): + signature = inspect.signature(sender_cls.send_chunks) + signature.bind_partial( + object(), + chunks=[], + weight_metadata=None, + derive_metadata_from_chunks=False, + ) + + +def test_delta_sender_rejects_serialized_fp8_chunks(): + """Delta checkpoints cannot represent serialized-FP8 wire chunks.""" + + with pytest.raises(ValueError, match="serialized FP8"): + asyncio.run( + DeltaWeightTransferSender.send_chunks( + object(), + chunks=[], + derive_metadata_from_chunks=True, + ) + ) diff --git a/tests/backends/skyrl_train/weight_sync/test_sharded_rdt_strategy.py b/tests/backends/skyrl_train/weight_sync/test_sharded_rdt_strategy.py index a1b0f4421a..1dcc03f633 100644 --- a/tests/backends/skyrl_train/weight_sync/test_sharded_rdt_strategy.py +++ b/tests/backends/skyrl_train/weight_sync/test_sharded_rdt_strategy.py @@ -38,9 +38,10 @@ class _Client: class _Extractor: """Records whether the metadata/chunk channels were touched at all.""" - def __init__(self) -> None: + def __init__(self, derives_metadata_from_chunks: bool = False) -> None: self.metadata_calls = 0 self.extract_calls = 0 + self.derives_metadata_from_chunks = derives_metadata_from_chunks def get_weight_metadata(self, dtype): self.metadata_calls += 1 @@ -166,6 +167,19 @@ async def test_the_push_backends_kwargs_are_accepted_and_ignored(self): await ShardedRdtWeightTransferSender(inner).send(_Extractor(), "torch.bfloat16", reset_prefix_cache=True) assert len(inner.sent) == 1 + @pytest.mark.asyncio + async def test_send_refuses_serialized_fp8_extractors(self): + """The weight sources publish whole bridge tensors, so an FP8 sync would + land dequantized weights in modules vLLM built for FP8 -- wrong rollout + weights rather than an error.""" + inner = _FakeRdtSender() + extractor = _Extractor(derives_metadata_from_chunks=True) + + with pytest.raises(ValueError, match="serialized FP8"): + await ShardedRdtWeightTransferSender(inner).send(extractor, "torch.bfloat16") + + assert inner.sent == [] + @pytest.mark.asyncio async def test_send_chunks_refuses_with_an_explanation(self): """There is no chunk stream to push; the error has to say what to call.""" diff --git a/tests/backends/skyrl_train/weight_sync/test_transfer_strategies.py b/tests/backends/skyrl_train/weight_sync/test_transfer_strategies.py index f73b810a05..f7ae0aef98 100644 --- a/tests/backends/skyrl_train/weight_sync/test_transfer_strategies.py +++ b/tests/backends/skyrl_train/weight_sync/test_transfer_strategies.py @@ -1,4 +1,7 @@ +import asyncio + import pytest +import torch from skyrl.backends.skyrl_train.weight_sync import ( BroadcastInitInfo, @@ -14,6 +17,10 @@ get_transfer_strategy, get_transfer_strategy_cls, ) +from skyrl.backends.skyrl_train.weight_sync.base import WeightChunk +from skyrl.backends.skyrl_train.weight_sync.broadcast_strategy import ( + BroadcastWeightTransferSender, +) from skyrl.train.config import InferenceEngineConfig from skyrl.train.config.config import DeltaWeightSyncConfig @@ -783,7 +790,7 @@ def _make_ie_cfg( ) def test_cuda_ipc_create_init_info(self): - """CudaIpcTransferStrategy.create_init_info should create CudaIpcInitInfo with model_dtype_str.""" + """Preserve model-dtype metadata in CUDA IPC initialization.""" ie_cfg = self._make_ie_cfg(model_dtype="torch.float32") init_info = CudaIpcTransferStrategy.create_init_info(ie_cfg) @@ -883,6 +890,149 @@ def test_mismatched_lengths_raises(self): ) +def test_broadcast_sender_preserves_mixed_dtype_logical_chunk(monkeypatch): + """vLLM's byte-packed NCCL path accepts one mixed-dtype logical update.""" + import skyrl.backends.skyrl_train.weight_sync.broadcast_strategy as broadcast_module + + class FakeInferenceClient: + def __init__(self): + self.events = [] + + async def start_weight_update(self, is_checkpoint_format): + self.events.append(("start", is_checkpoint_format)) + + async def finish_weight_update(self): + self.events.append(("finish",)) + + client = FakeInferenceClient() + sender = BroadcastWeightTransferSender( + init_info=BroadcastInitInfo( + master_addr="127.0.0.1", + master_port=12345, + rank_offset=1, + world_size=2, + override_existing_receiver=False, + ), + model_update_group=object(), + inference_client=client, + ) + sent_chunks = [] + + async def record_chunk(chunk): + sent_chunks.append((list(chunk.names), [tensor.dtype for tensor in chunk.tensors])) + + monkeypatch.setattr(sender, "_send_chunk_vllm_native", record_chunk) + monkeypatch.setattr(broadcast_module.torch.distributed, "get_rank", lambda: 0) + monkeypatch.setattr(broadcast_module.torch.distributed, "barrier", lambda: None) + + mixed_chunk = WeightChunk( + names=["w0", "scale", "w1", "norm"], + dtypes=["ignored"] * 4, + shapes=[[4], [1], [8], [2]], + tensors=[ + torch.empty((4,), dtype=torch.float8_e4m3fn), + torch.empty((1,), dtype=torch.float32), + torch.empty((8,), dtype=torch.float8_e4m3fn), + torch.empty((2,), dtype=torch.bfloat16), + ], + ) + + asyncio.run(sender.send_chunks(iter([mixed_chunk]), derive_metadata_from_chunks=True)) + + assert client.events == [("start", True), ("finish",)] + assert sent_chunks == [ + ( + ["w0", "scale", "w1", "norm"], + [torch.float8_e4m3fn, torch.float32, torch.float8_e4m3fn, torch.bfloat16], + ) + ] + + +def test_broadcast_send_chunk_uses_vendored_send_and_init_time_packed(monkeypatch): + """Per-chunk FP8 rounds share the batched path's wire protocol. + + vLLM 0.28.0 removed ``NCCLWeightTransferEngine.trainer_send_weights`` and + rejects ``packed`` on ``NCCLWeightTransferUpdateInfo``, so the per-chunk + send must go through the vendored ``nccl_trainer_send_weights`` with the + ``packed`` agreed at init, and the update payload must carry only metadata. + """ + import skyrl.backends.skyrl_train.weight_sync.broadcast_strategy as broadcast_module + + class FakeInferenceClient: + def __init__(self): + self.update_infos = [] + + async def update_weights_nccl(self, update_info): + self.update_infos.append(update_info) + + client = FakeInferenceClient() + group = object() + sender = BroadcastWeightTransferSender( + init_info=BroadcastInitInfo( + master_addr="127.0.0.1", + master_port=12345, + rank_offset=1, + world_size=2, + packed=False, + override_existing_receiver=False, + ), + model_update_group=group, + inference_client=client, + ) + sends = [] + + def fake_send(iterator, model_update_group, *, packed): + sends.append((list(iterator), model_update_group, packed)) + + monkeypatch.setattr(broadcast_module, "nccl_trainer_send_weights", fake_send) + + chunk = WeightChunk( + names=["w0", "scale"], + dtypes=["float8_e4m3fn", "float32"], + shapes=[[4], [1]], + tensors=[ + torch.empty((4,), dtype=torch.float8_e4m3fn), + torch.empty((1,), dtype=torch.float32), + ], + ) + + asyncio.run(sender._send_chunk_vllm_native(chunk)) + + assert client.update_infos == [ + {"names": ["w0", "scale"], "dtype_names": ["float8_e4m3fn", "float32"], "shapes": [[4], [1]]} + ] + assert len(sends) == 1 + names, sent_group, packed = sends[0] + assert [name for name, _ in names] == ["w0", "scale"] + assert sent_group is group + assert packed is False + + +def test_broadcast_sender_retains_precomputed_metadata_path(monkeypatch): + sender = BroadcastWeightTransferSender( + init_info=BroadcastInitInfo( + master_addr="127.0.0.1", + master_port=12345, + rank_offset=1, + world_size=2, + override_existing_receiver=False, + ), + model_update_group=object(), + inference_client=object(), + ) + calls = [] + + async def record_batched(chunks, weight_metadata): + calls.append((list(chunks), weight_metadata)) + + monkeypatch.setattr(sender, "_send_chunks_vllm_native", record_batched) + metadata = {"names": ["w"], "dtype_names": ["bfloat16"], "shapes": [[1]]} + + asyncio.run(sender.send_chunks(iter([]), weight_metadata=metadata)) + + assert calls == [([], metadata)] + + class TestCudaIpcWeightUpdateRequest: """Tests for CudaIpcWeightUpdateRequest.""" diff --git a/tests/backends/skyrl_train/weight_sync/test_weight_chunk.py b/tests/backends/skyrl_train/weight_sync/test_weight_chunk.py index a66c4f5159..6251d3273e 100644 --- a/tests/backends/skyrl_train/weight_sync/test_weight_chunk.py +++ b/tests/backends/skyrl_train/weight_sync/test_weight_chunk.py @@ -2,6 +2,11 @@ import torch from skyrl.backends.skyrl_train.weight_sync import WeightChunk +from skyrl.backends.skyrl_train.weight_sync.base import ( + cuda_uuid_to_str, + get_weight_chunk_metadata, + iter_single_dtype_chunks, +) class TestWeightChunk: @@ -95,3 +100,50 @@ def test_total_size_bytes_mixed_dtypes(self): ) assert chunk.total_size_bytes == 40 + 20 + + +def test_single_dtype_chunk_partition_preserves_first_seen_dtype_order(): + chunk = WeightChunk( + names=["w0", "s0", "w1", "b0"], + dtypes=["ignored"] * 4, + shapes=[[4], [1], [8], [2]], + tensors=[ + torch.empty((4,), dtype=torch.float8_e4m3fn), + torch.ones((1,), dtype=torch.float32), + torch.empty((8,), dtype=torch.float8_e4m3fn), + torch.ones((2,), dtype=torch.bfloat16), + ], + ) + + chunks = list(iter_single_dtype_chunks(chunk)) + + assert [subchunk.names for subchunk in chunks] == [["w0", "w1"], ["s0"], ["b0"]] + assert [[tensor.dtype for tensor in subchunk.tensors] for subchunk in chunks] == [ + [torch.float8_e4m3fn, torch.float8_e4m3fn], + [torch.float32], + [torch.bfloat16], + ] + + +def test_weight_chunk_metadata_uses_actual_tensor_dtypes(): + chunk = WeightChunk( + names=["w", "scale", "norm"], + dtypes=["stale"] * 3, + shapes=[[999], [999], [999]], + tensors=[ + torch.empty((4,), dtype=torch.float8_e4m3fn), + torch.empty((1,), dtype=torch.float32), + torch.empty((2,), dtype=torch.bfloat16), + ], + ) + + assert get_weight_chunk_metadata(chunk) == { + "names": ["w", "scale", "norm"], + "dtype_names": ["float8_e4m3fn", "float32", "bfloat16"], + "shapes": [[4], [1], [2]], + } + + +@pytest.mark.parametrize("uuid", ["GPU-123", b"GPU-123"]) +def test_cuda_uuid_to_str_normalizes_string_and_bytes(uuid): + assert cuda_uuid_to_str(uuid) == "GPU-123" diff --git a/tests/backends/skyrl_train/workers/megatron/quantization/__init__.py b/tests/backends/skyrl_train/workers/megatron/quantization/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/backends/skyrl_train/test_fake_int4_qat.py b/tests/backends/skyrl_train/workers/megatron/quantization/test_fake_int4_qat.py similarity index 98% rename from tests/backends/skyrl_train/test_fake_int4_qat.py rename to tests/backends/skyrl_train/workers/megatron/quantization/test_fake_int4_qat.py index df25a932b5..3d24efb71e 100644 --- a/tests/backends/skyrl_train/test_fake_int4_qat.py +++ b/tests/backends/skyrl_train/workers/megatron/quantization/test_fake_int4_qat.py @@ -13,7 +13,7 @@ import pytest import torch -from skyrl.backends.skyrl_train.workers.megatron.fake_int4_qat import ( +from skyrl.backends.skyrl_train.workers.megatron.quantization.fake_int4_qat import ( _FakeInt4QuantizeSTE, fake_int4_quantize_ste, ) diff --git a/tests/backends/skyrl_train/workers/megatron/quantization/test_fp8_param.py b/tests/backends/skyrl_train/workers/megatron/quantization/test_fp8_param.py new file mode 100644 index 0000000000..ae0ea86b8f --- /dev/null +++ b/tests/backends/skyrl_train/workers/megatron/quantization/test_fp8_param.py @@ -0,0 +1,171 @@ +import pytest +import torch + +from skyrl.backends.skyrl_train.workers.megatron.quantization.fp8_param import ( + initialize_fp8_param_optimizer_masters, + is_fp8_param_enabled, +) + + +class _FakeOptimizer: + def __init__(self): + self.calls = 0 + self.state_dict = None + + def _copy_model_params_to_main_params(self, state_dict=None): + self.calls += 1 + self.state_dict = state_dict + + +class _Range: + def __init__(self, start: int, end: int): + self.start = start + self.end = end + self.size = end - start + + +class _FakeQuantizedParam: + def __init__(self, values): + self._values = torch.tensor(values, dtype=torch.float32) + + def dequantize(self): + return self._values + + +class _FakeHybridDeviceOptimizer: + def __init__(self, cpu_masters, gpu_masters): + self.cpu_copys_map_gpu_param = dict(zip(cpu_masters, gpu_masters)) + self.param_to_fp32_param = {} + + +class _FakeHybridMegatronOptimizer: + def __init__(self): + self.model = _FakeQuantizedParam([2.0, 4.0, 6.0, 8.0]) + self.unquantized_model = torch.tensor([10.0, 20.0, 30.0, 40.0]) + self.gpu_master = torch.full((2,), -1.0) + self.gpu_master_unquantized = torch.full((2,), -3.0) + self.cpu_master = torch.full((2,), -2.0) + self.cpu_master_unquantized = torch.full((2,), -4.0) + self.optimizer = _FakeHybridDeviceOptimizer( + [self.cpu_master, self.cpu_master_unquantized], + [self.gpu_master, self.gpu_master_unquantized], + ) + self.model_float16_groups = [[self.model, self.unquantized_model]] + self.shard_fp32_from_float16_groups = [[self.gpu_master, self.gpu_master_unquantized]] + self.model_fp32_groups = [] + self.shard_fp32_groups = [] + self.loaded_quantized_model = torch.tensor([100.0, 200.0, 300.0, 400.0]) + self.loaded_unquantized_model = torch.tensor([50.0, 60.0, 70.0, 80.0]) + self.exact_state_dict = {"model": {"ignored": torch.tensor(0.0)}} + + def _is_distopt_quantized_param(self, param): + return param is self.model + + def _get_model_param_range_map(self, param): + assert param is self.model or param is self.unquantized_model + return {"param": _Range(1, 3)} + + def _build_model_param_to_state_dict_param_map(self, state_dict): + assert state_dict is self.exact_state_dict + return { + self.model: self.loaded_quantized_model, + self.unquantized_model: self.loaded_unquantized_model, + } + + def _copy_model_params_to_main_params(self): + raise AssertionError("Hybrid CPU-offload path must seed both master copies directly") + + +def test_fp8_param_enablement_reads_skyrl_transformer_config_mapping(): + assert is_fp8_param_enabled({"fp8_param": True}) + assert not is_fp8_param_enabled({"fp8_param": False}) + assert not is_fp8_param_enabled({"fp8_param": "false"}) + assert not is_fp8_param_enabled({}) + + +def test_fp8_param_master_initialization_reloads_each_chained_optimizer(): + first = _FakeOptimizer() + second = _FakeOptimizer() + chained = type("FakeChainedOptimizer", (), {"chained_optimizers": [first, second]})() + + initialized = initialize_fp8_param_optimizer_masters( + chained, + fp8_param=True, + fp8_param_gather=True, + state_dict={"model": {}}, + ) + + assert initialized == 2 + assert first.calls == 1 + assert second.calls == 1 + + +def test_fp8_param_cpu_offload_uses_exact_checkpoint_state_for_all_masters(): + optimizer = _FakeHybridMegatronOptimizer() + + initialized = initialize_fp8_param_optimizer_masters( + optimizer, + fp8_param=True, + fp8_param_gather=True, + state_dict=optimizer.exact_state_dict, + ) + + assert initialized == 1 + torch.testing.assert_close(optimizer.gpu_master, torch.tensor([200.0, 300.0])) + torch.testing.assert_close(optimizer.cpu_master, torch.tensor([200.0, 300.0])) + torch.testing.assert_close(optimizer.gpu_master_unquantized, torch.tensor([60.0, 70.0])) + torch.testing.assert_close(optimizer.cpu_master_unquantized, torch.tensor([60.0, 70.0])) + + +def test_fp8_param_master_initialization_passes_exact_state_to_normal_optimizer(): + optimizer = _FakeOptimizer() + state_dict = {"model": {"weight": torch.tensor([1.0])}} + + initialized = initialize_fp8_param_optimizer_masters( + optimizer, + fp8_param=True, + fp8_param_gather=True, + state_dict=state_dict, + ) + + assert initialized == 1 + assert optimizer.state_dict is state_dict + + +def test_fp8_param_master_initialization_requires_fp8_param_gather(): + optimizer = _FakeOptimizer() + + with pytest.raises(ValueError, match="fp8_param_gather=true"): + initialize_fp8_param_optimizer_masters( + optimizer, + fp8_param=True, + fp8_param_gather=False, + ) + + assert optimizer.calls == 0 + + +def test_fp8_param_master_initialization_requires_exact_checkpoint_state(): + optimizer = _FakeOptimizer() + + with pytest.raises(ValueError, match="exact unquantized checkpoint state"): + initialize_fp8_param_optimizer_masters( + optimizer, + fp8_param=True, + fp8_param_gather=True, + ) + + assert optimizer.calls == 0 + + +def test_non_persistent_fp8_does_not_touch_optimizer_masters(): + optimizer = _FakeOptimizer() + + initialized = initialize_fp8_param_optimizer_masters( + optimizer, + fp8_param=False, + fp8_param_gather=False, + ) + + assert initialized == 0 + assert optimizer.calls == 0 diff --git a/tests/train/test_config.py b/tests/train/test_config.py index c567b49ec5..2304647b0e 100644 --- a/tests/train/test_config.py +++ b/tests/train/test_config.py @@ -11,6 +11,7 @@ import pytest from omegaconf import OmegaConf +from skyrl.backends.skyrl_train.distributed.megatron import quantization_utils from skyrl.train.config.config import ( BaseConfig, DeltaWeightSyncConfig, @@ -20,7 +21,12 @@ build_nested_dataclass, overrides_dict_to_dotlist, ) -from skyrl.train.utils.utils import validate_cfg, validate_inference_engine_cfg +from skyrl.train.utils import utils as train_utils +from skyrl.train.utils.utils import ( + prepare_runtime_environment, + validate_cfg, + validate_inference_engine_cfg, +) from tests.train.util import example_dummy_config @@ -140,6 +146,11 @@ def test_cli_overrides_empty_args(): assert cfg.trainer.seed == 42 +def test_cli_overrides_fp8_param_gather(): + cfg = SkyRLTrainConfig.from_cli_overrides(["trainer.policy.megatron_config.ddp_config.fp8_param_gather=true"]) + assert cfg.trainer.policy.megatron_config.ddp_config.fp8_param_gather is True + + @pytest.mark.parametrize( ("field_name", "value"), [ @@ -154,6 +165,254 @@ def test_trainer_config_rejects_invalid_vocab_entropy_chunking(field_name, value TrainerConfig(**{field_name: value}) +def test_runtime_env_forwards_te_block_scale_mode(monkeypatch): + monkeypatch.setenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "1") + monkeypatch.setattr(train_utils, "peer_access_supported", lambda **_kwargs: True) + + env_vars = prepare_runtime_environment(example_dummy_config()) + + assert env_vars["NVTE_FP8_BLOCK_SCALING_FP32_SCALES"] == "1" + + +def test_runtime_env_supports_fsdp_without_megatron_configs(monkeypatch): + monkeypatch.delenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", raising=False) + monkeypatch.delenv("NVTE_FP8_BLOCK_AMAX_EPSILON", raising=False) + monkeypatch.delenv("VLLM_USE_DEEP_GEMM_E8M0", raising=False) + monkeypatch.setattr(train_utils, "peer_access_supported", lambda **_kwargs: True) + cfg = example_dummy_config() + cfg.trainer.strategy = "fsdp" + cfg.trainer.policy.megatron_config = None + cfg.trainer.ref.megatron_config = None + + env_vars = prepare_runtime_environment(cfg) + + assert "NVTE_FP8_BLOCK_SCALING_FP32_SCALES" not in env_vars + assert "VLLM_USE_DEEP_GEMM_E8M0" not in env_vars + + +def test_serialized_fp8_runtime_defaults_to_fp32_scales(monkeypatch): + monkeypatch.delenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", raising=False) + monkeypatch.delenv("VLLM_USE_DEEP_GEMM_E8M0", raising=False) + monkeypatch.setattr(train_utils, "has_visible_cuda_device", lambda: True) + monkeypatch.setattr(train_utils, "peer_access_supported", lambda **_kwargs: True) + monkeypatch.setattr(train_utils, "is_blackwell_or_newer", lambda: False) + cfg = example_dummy_config() + cfg.generator.inference_engine.fp8_weight_sync_mode = "blockwise" + + env_vars = prepare_runtime_environment(cfg) + + assert env_vars["NVTE_FP8_BLOCK_SCALING_FP32_SCALES"] == "1" + assert env_vars["VLLM_USE_DEEP_GEMM_E8M0"] == "0" + + +def test_serialized_fp8_runtime_defaults_to_pow2_scales_on_blackwell(monkeypatch): + monkeypatch.delenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", raising=False) + monkeypatch.delenv("VLLM_USE_DEEP_GEMM_E8M0", raising=False) + monkeypatch.setattr(train_utils, "has_visible_cuda_device", lambda: True) + monkeypatch.setattr(train_utils, "peer_access_supported", lambda **_kwargs: True) + monkeypatch.setattr(train_utils, "is_blackwell_or_newer", lambda: True) + cfg = example_dummy_config() + cfg.generator.inference_engine.fp8_weight_sync_mode = "blockwise" + + env_vars = prepare_runtime_environment(cfg) + + assert env_vars["NVTE_FP8_BLOCK_SCALING_FP32_SCALES"] == "0" + assert env_vars["VLLM_USE_DEEP_GEMM_E8M0"] == "1" + + +def test_serialized_fp8_pow2_scales_reject_disabled_e8m0_on_blackwell(monkeypatch): + monkeypatch.delenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", raising=False) + monkeypatch.setenv("VLLM_USE_DEEP_GEMM_E8M0", "0") + monkeypatch.setattr(train_utils, "has_visible_cuda_device", lambda: True) + monkeypatch.setattr(train_utils, "peer_access_supported", lambda **_kwargs: True) + monkeypatch.setattr(train_utils, "is_blackwell_or_newer", lambda: True) + cfg = example_dummy_config() + cfg.generator.inference_engine.fp8_weight_sync_mode = "blockwise" + + with pytest.raises(ValueError, match="VLLM_USE_DEEP_GEMM_E8M0=1"): + prepare_runtime_environment(cfg) + + +def test_serialized_fp8_requires_an_explicit_scale_mode_without_a_driver_gpu(monkeypatch): + """The contract is baked into the runtime env before ray.init, so a GPU-less + head cannot infer it from the workers; guessing Hopper would hand FP32 block + scales to Blackwell workers.""" + monkeypatch.delenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", raising=False) + monkeypatch.delenv("VLLM_USE_DEEP_GEMM_E8M0", raising=False) + monkeypatch.setattr(train_utils, "peer_access_supported", lambda **_kwargs: True) + monkeypatch.setattr(train_utils, "has_visible_cuda_device", lambda: False) + cfg = example_dummy_config() + cfg.generator.inference_engine.fp8_weight_sync_mode = "blockwise" + + with pytest.raises(ValueError, match="NVTE_FP8_BLOCK_SCALING_FP32_SCALES"): + prepare_runtime_environment(cfg) + + +def test_serialized_fp8_pow2_scales_set_e8m0_without_a_driver_gpu(monkeypatch): + """E8M0 follows the wire scale format, not the driver's device: vLLM picks the + per-device form itself, so the default must survive a GPU-less head.""" + monkeypatch.setenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "0") + monkeypatch.delenv("VLLM_USE_DEEP_GEMM_E8M0", raising=False) + monkeypatch.setattr(train_utils, "peer_access_supported", lambda **_kwargs: True) + monkeypatch.setattr(train_utils, "has_visible_cuda_device", lambda: False) + monkeypatch.setattr(train_utils, "is_blackwell_or_newer", lambda: False) + cfg = example_dummy_config() + cfg.generator.inference_engine.fp8_weight_sync_mode = "blockwise" + + env_vars = prepare_runtime_environment(cfg) + + assert env_vars["NVTE_FP8_BLOCK_SCALING_FP32_SCALES"] == "0" + assert env_vars["VLLM_USE_DEEP_GEMM_E8M0"] == "1" + + +@pytest.mark.parametrize("backend", ["sharded_rdt", "delta"]) +def test_serialized_fp8_weight_sync_rejects_backends_without_a_chunk_channel(backend): + """Neither backend carries payload + scale pairs, and both would otherwise + fail only at the first sync -- after vLLM has loaded as FP8.""" + cfg = _make_validated_test_config() + cfg.trainer.strategy = "megatron" + cfg.generator.inference_engine.fp8_weight_sync_mode = "blockwise" + cfg.generator.inference_engine.weight_sync_backend = backend + + with pytest.raises(ValueError, match=backend): + validate_inference_engine_cfg(cfg) + + +def test_serialized_fp8_weight_sync_requires_megatron(): + cfg = _make_validated_test_config() + cfg.trainer.strategy = "fsdp" + cfg.generator.inference_engine.fp8_weight_sync_mode = "blockwise" + + with pytest.raises(ValueError, match="requires trainer.strategy='megatron'"): + validate_inference_engine_cfg(cfg) + + +def test_serialized_fp8_weight_sync_rejects_adapter_only_megatron_lora(): + cfg = _make_validated_test_config() + cfg.trainer.strategy = "megatron" + cfg.generator.inference_engine.fp8_weight_sync_mode = "blockwise" + cfg.trainer.policy.model.lora.rank = 8 + cfg.trainer.policy.megatron_config.lora_config.merge_lora = False + + with pytest.raises(ValueError, match="requires full-weight updates"): + validate_inference_engine_cfg(cfg) + + +def test_megatron_fp8_compute_defaults_to_fp32_scales_without_serialized_sync(monkeypatch): + monkeypatch.delenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", raising=False) + monkeypatch.setattr(train_utils, "has_visible_cuda_device", lambda: True) + monkeypatch.setattr(train_utils, "peer_access_supported", lambda **_kwargs: True) + monkeypatch.setattr(train_utils, "is_blackwell_or_newer", lambda: False) + cfg = example_dummy_config() + cfg.trainer.policy.megatron_config.transformer_config_kwargs["fp8"] = "hybrid" + + env_vars = prepare_runtime_environment(cfg) + + assert env_vars["NVTE_FP8_BLOCK_SCALING_FP32_SCALES"] == "1" + + +def test_power_2_mode_rejects_persistent_fp8_without_serialized_sync(monkeypatch): + monkeypatch.setenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "0") + monkeypatch.setattr(train_utils, "peer_access_supported", lambda **_kwargs: True) + cfg = example_dummy_config() + cfg.trainer.policy.megatron_config.transformer_config_kwargs["fp8_param"] = True + + with pytest.raises(ValueError, match="fp8_param=false on Blackwell"): + prepare_runtime_environment(cfg) + + +def test_megatron_validation_requires_fp8_param_gather_for_training(): + cfg = _make_validated_test_config() + cfg.trainer.strategy = "megatron" + cfg.trainer.policy.megatron_config.transformer_config_kwargs["fp8_param"] = True + cfg.trainer.policy.megatron_config.ddp_config.fp8_param_gather = False + + with pytest.raises(ValueError, match="fp8_param_gather=true"): + train_utils.validate_megatron_cfg(cfg) + + +def test_megatron_top_level_fp8_fields_fold_into_transformer_config_kwargs(): + from skyrl.train.config.config import MegatronConfig + + cfg = MegatronConfig(fp8="e4m3", fp8_recipe="auto", fp8_param=True, fp8_amax_compute_algo="most_recent") + assert cfg.transformer_config_kwargs["fp8"] == "e4m3" + assert cfg.transformer_config_kwargs["fp8_recipe"] == "auto" + assert cfg.transformer_config_kwargs["fp8_param"] is True + assert cfg.transformer_config_kwargs["fp8_amax_compute_algo"] == "most_recent" + # Defaults stay off: no FP8 keys appear unless requested. + assert "fp8" not in MegatronConfig().transformer_config_kwargs + + +def test_megatron_explicit_transformer_config_kwargs_override_top_level_fp8_fields(): + from skyrl.train.config.config import MegatronConfig + + cfg = MegatronConfig(fp8="e4m3", fp8_recipe="blockwise", transformer_config_kwargs={"fp8_recipe": "mxfp8"}) + assert cfg.transformer_config_kwargs["fp8_recipe"] == "mxfp8" + assert cfg.transformer_config_kwargs["fp8"] == "e4m3" + + +def test_megatron_validation_allows_inference_only_fp8_param_without_gather(): + cfg = _make_validated_test_config() + cfg.trainer.strategy = "megatron" + cfg.trainer.policy.inference_only_init = True + cfg.trainer.policy.megatron_config.transformer_config_kwargs["fp8_param"] = True + cfg.trainer.policy.megatron_config.ddp_config.fp8_param_gather = False + + train_utils.validate_megatron_cfg(cfg) + + +@pytest.mark.parametrize(("blackwell", "expected_recipe"), [(True, "mxfp8"), (False, "blockwise")]) +def test_megatron_validation_resolves_auto_fp8_recipe(monkeypatch, blackwell, expected_recipe): + monkeypatch.setattr(quantization_utils, "has_visible_cuda_device", lambda: True) + monkeypatch.setattr(quantization_utils, "is_blackwell_or_newer", lambda: blackwell) + monkeypatch.setattr(train_utils, "is_blackwell_or_newer", lambda: blackwell) + cfg = _make_validated_test_config() + cfg.trainer.strategy = "megatron" + cfg.trainer.policy.megatron_config.transformer_config_kwargs["fp8"] = "e4m3" + cfg.trainer.policy.megatron_config.transformer_config_kwargs["fp8_recipe"] = "auto" + + train_utils.validate_megatron_cfg(cfg) + + assert cfg.trainer.policy.megatron_config.transformer_config_kwargs["fp8_recipe"] == expected_recipe + + +def test_megatron_validation_rejects_mxfp8_before_blackwell(monkeypatch): + monkeypatch.setattr(quantization_utils, "has_visible_cuda_device", lambda: True) + monkeypatch.setattr(quantization_utils, "is_blackwell_or_newer", lambda: False) + monkeypatch.setattr(train_utils, "is_blackwell_or_newer", lambda: False) + cfg = _make_validated_test_config() + cfg.trainer.strategy = "megatron" + cfg.trainer.policy.megatron_config.transformer_config_kwargs["fp8"] = "e4m3" + cfg.trainer.policy.megatron_config.transformer_config_kwargs["fp8_recipe"] = "mxfp8" + + with pytest.raises(ValueError, match="requires SM100"): + train_utils.validate_megatron_cfg(cfg) + + +def test_megatron_validation_rejects_mxfp8_with_fp8_param(monkeypatch): + monkeypatch.setattr(train_utils, "is_blackwell_or_newer", lambda: True) + cfg = _make_validated_test_config() + cfg.trainer.strategy = "megatron" + cfg.trainer.policy.megatron_config.transformer_config_kwargs["fp8"] = "e4m3" + cfg.trainer.policy.megatron_config.transformer_config_kwargs["fp8_recipe"] = "mxfp8" + cfg.trainer.policy.megatron_config.transformer_config_kwargs["fp8_param"] = True + cfg.trainer.policy.megatron_config.ddp_config.fp8_param_gather = True + + with pytest.raises(ValueError, match="not supported with fp8_recipe=mxfp8"): + train_utils.validate_megatron_cfg(cfg) + + +def test_serialized_fp8_fp32_scales_reject_vllm_e8m0(monkeypatch): + monkeypatch.setenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "1") + monkeypatch.setenv("VLLM_USE_DEEP_GEMM_E8M0", "1") + monkeypatch.setattr(train_utils, "peer_access_supported", lambda **_kwargs: True) + cfg = example_dummy_config() + cfg.generator.inference_engine.fp8_weight_sync_mode = "blockwise" + + with pytest.raises(ValueError, match="VLLM_USE_DEEP_GEMM_E8M0=0"): + prepare_runtime_environment(cfg) + + def test_cli_overrides_plus_prefix_rejected(): with pytest.raises(ValueError, match="The '\\+' prefix"): SkyRLTrainConfig.from_cli_overrides(["+new_field=value"]) diff --git a/tests/train/test_packing_round_trip.py b/tests/train/test_packing_round_trip.py index 749bad1944..bc37c41193 100644 --- a/tests/train/test_packing_round_trip.py +++ b/tests/train/test_packing_round_trip.py @@ -156,9 +156,8 @@ def test_mbs2_tp4_multisubseq_rows_pack_correctly(self): - ``mbs = 2`` so the THD offset stride matters per micro-batch row. - ``tp_size = 4`` so the TP-alignment padding inside rows kicks in. - Multi-subseq rows ``[(7, 5)]`` and ``[(3, 11)]`` to expose offset - mismatches in both directions (sub-seq 0 length 7 is NOT a multiple - of 4, sub-seq 1 length 5 is also not; row 1's sub-seq 0 length 3 is - shorter than its sub-seq 1 length 11). + mismatches in both directions (none of the lengths are alignment + multiples, and row 1's first sub-seq is shorter than its second). """ from skyrl.backends.skyrl_train.distributed.megatron.megatron_utils import ( preprocess_packed_seqs, @@ -178,17 +177,15 @@ def test_mbs2_tp4_multisubseq_rows_pack_correctly(self): ) # Sanity: collator layout matches what PackedDataCollator does. - # row 0: [101..107] + pad-to-16 + [201..205] + pad-to-16 => 32 slots - # row 1: [301..303] + pad-to-16 + [401..411] + pad-to-16 => 32 slots - assert sequences.shape == (2, 32) + assert sequences.shape == (2, 2 * align_size) assert sequences[0, :7].tolist() == row_0_sub_0 - assert sequences[0, 7:16].tolist() == [0] * 9 - assert sequences[0, 16:21].tolist() == row_0_sub_1 + assert sequences[0, 7:align_size].tolist() == [0] * (align_size - 7) + assert sequences[0, align_size : align_size + 5].tolist() == row_0_sub_1 assert sequences[1, :3].tolist() == row_1_sub_0 - assert sequences[1, 3:16].tolist() == [0] * 13 - assert sequences[1, 16:27].tolist() == row_1_sub_1 + assert sequences[1, 3:align_size].tolist() == [0] * (align_size - 3) + assert sequences[1, align_size : align_size + 11].tolist() == row_1_sub_1 # attention_mask is True ONLY at valid slots (NOT at TP-alignment gaps). - assert attention_mask[0, 7:16].any().item() is False + assert attention_mask[0, 7:align_size].any().item() is False assert attention_mask[0].sum().item() == 7 + 5 assert attention_mask[1].sum().item() == 3 + 11 @@ -208,19 +205,14 @@ def test_mbs2_tp4_multisubseq_rows_pack_correctly(self): ) # cu_seqlens_q (== cu_seqlens_q_padded for THD) enumerates 4 sub-seqs: - assert params.cu_seqlens_q.tolist() == [0, 16, 32, 48, 64] - # Packed slab is 64 tokens. Verify each sub-seq's *valid* tokens were - # read from the *correct intra-row offset* — i.e. preprocess respected - # the TP-alignment gap that the collator inserted. - assert packed.shape == (1, 64) - assert packed[0, :7].tolist() == row_0_sub_0 # row 0 sub 0 - assert packed[0, 7:16].tolist() == [0] * 9 # alignment pad - assert packed[0, 16:21].tolist() == row_0_sub_1 # row 0 sub 1 - assert packed[0, 21:32].tolist() == [0] * 11 # alignment pad - assert packed[0, 32:35].tolist() == row_1_sub_0 # row 1 sub 0 - assert packed[0, 35:48].tolist() == [0] * 13 # alignment pad - assert packed[0, 48:59].tolist() == row_1_sub_1 # row 1 sub 1 - assert packed[0, 59:64].tolist() == [0] * 5 # alignment pad + assert params.cu_seqlens_q.tolist() == [i * align_size for i in range(5)] + # Verify valid tokens are read from offsets produced by alignment padding. + assert packed.shape == (1, 4 * align_size) + assert packed[0, :7].tolist() == row_0_sub_0 + assert packed[0, 7:align_size].tolist() == [0] * (align_size - 7) + assert packed[0, align_size : align_size + 5].tolist() == row_0_sub_1 + assert packed[0, 2 * align_size : 2 * align_size + 3].tolist() == row_1_sub_0 + assert packed[0, 3 * align_size : 3 * align_size + 11].tolist() == row_1_sub_1 def test_mbs2_tp1_singlesubseq_rows_match_legacy_preprocess_path(self): """No regression in the legacy path: when each row has 1 sub-seq and tp_size=1.""" diff --git a/tests/train/test_sft_packing_collate.py b/tests/train/test_sft_packing_collate.py index b2a1a2e26b..7a5fbe1787 100644 --- a/tests/train/test_sft_packing_collate.py +++ b/tests/train/test_sft_packing_collate.py @@ -212,17 +212,18 @@ def test_tp_alignment_pads_each_sub_seq(self): subseq_lengths = batch["sub_seq_lengths"][0].tolist() assert sum(subseq_lengths) == 12 # raw, un-padded - def test_fp8_tp_alignment_pads_each_sub_seq_to_16(self): - """When FP8 is enabled, TP-only packed subseqs are also 16-aligned.""" + def test_fp8_tp_alignment_cost_prevents_overpacking(self): + """Packing uses each sequence's aligned FP8/TP footprint.""" collator = _make_collator(num_gpus=4, batch_size=2, max_length=128, tp=4, fp8="hybrid") examples = [ _make_example(7, 3), _make_example(5, 3), ] batch = collator(examples, batch_size=2) - assert batch["sequences"].shape[1] >= 32 - subseq_lengths = batch["sub_seq_lengths"][0].tolist() - assert sum(subseq_lengths) == 12 + # TP4 gives each sequence a 512-token footprint, so the 128-token budget + # expands to one footprint without packing both sequences together. + assert batch["sequences"].shape == (2, 512) + assert sorted(lengths.tolist() for lengths in batch["sub_seq_lengths"]) == [[5], [7]] def test_bf16_cp_alignment_does_not_apply_fp8_padding(self): """With CP>1 and BF16, keep only TP/CP layout padding.""" @@ -256,6 +257,38 @@ def test_fp8_cp_alignment_pads_each_sub_seq(self): padded = ((s + 31) // 32) * 32 assert (padded // 2) % 16 == 0 + def test_fp8_alignment_is_conditional(self): + examples = [ + _make_example(3, 1, base_token=100), + _make_example(3, 1, base_token=200), + ] + + bf16 = _make_collator(num_gpus=1, batch_size=2, max_length=64, fp8=False) + bf16_batch = bf16(examples, batch_size=2) + assert bf16_batch["sequences"].shape[1] == 6 + bf16_chunks = [bf16_batch["sequences"][0, :3].tolist(), bf16_batch["sequences"][0, 3:6].tolist()] + assert sorted(bf16_chunks) == [[100, 101, 102], [200, 201, 202]] + + fp8 = _make_collator(num_gpus=1, batch_size=2, max_length=64, fp8=True) + fp8_batch = fp8(examples, batch_size=2) + assert fp8_batch["sequences"].shape[1] == 32 + fp8_chunks = [fp8_batch["sequences"][0, :3].tolist(), fp8_batch["sequences"][0, 16:19].tolist()] + assert sorted(fp8_chunks) == [[100, 101, 102], [200, 201, 202]] + + def test_tp_gt_1_fp8_alignment_uses_local_128(self): + """TP>1 FP8 row layout must satisfy TE blockwise input all-gather.""" + examples = [ + _make_example(3, 1, base_token=100), + _make_example(3, 1, base_token=200), + ] + + fp8 = _make_collator(num_gpus=2, batch_size=2, max_length=512, tp=2, fp8=True) + fp8_batch = fp8(examples, batch_size=2) + + assert fp8_batch["sequences"].shape[1] == 512 + fp8_chunks = [fp8_batch["sequences"][0, :3].tolist(), fp8_batch["sequences"][0, 256:259].tolist()] + assert sorted(fp8_chunks) == [[100, 101, 102], [200, 201, 202]] + def test_eval_path_falls_back_to_super(self): """When batch_size != self.sft_cfg.batch_size (eval), no packing happens.""" collator = _make_collator(num_gpus=1, batch_size=4, max_length=64)