FP8 RL Training and Rollout Weight Synchronization

.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`.
 

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`.

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 \
+  $@

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 \
+  $@