FP8 RL Training and Rollout Weight Synchronization

[YJHMITWEB]

This PR adds end-to-end FP8 training and rollout support. It covers persistent Megatron FP8 parameters, serialized blockwise weight synchronization, mixed-dtype transport, and FP8 vLLM execution.

Changes

Change	Explanation
Serialized blockwise FP8 sync	fp8_weight_sync_mode=serialized_blockwise converts eligible weights to E4M3 and sends matching FP32 weight_scale_inv tensors in vLLM checkpoint format.
FP8 vLLM execution	Builds the rollout model with quantization=fp8, load_format=dummy, and a matching 128 x 128 block quantization configuration; the first full sync replaces the dummy weights.
Qwen3.5 mapping	Handles dense attention and MLP linears, GDN projections, shared experts, and routed MoE experts while keeping incompatible GDN and vision projections in BF16.
Mixed-dtype transport	Preserves mixed-dtype logical chunks over NCCL and partitions CUDA IPC buffers by their actual tensor dtype.
Persistent FP8 parameters	Keeps eligible Megatron primary parameters in E4M3, reducing the required HBM in Megatron.
Exact optimizer initialization	Initializes FP32 masters and hybrid CPU copies from unquantized checkpoint shards before the first optimizer step.
FP8 parameter gather	Exposes fp8_param_gather, which refreshes persistent FP8 parameters from updated FP32 masters.
FP8-safe packed sequences	Uses one TP/CP-aware alignment rule in collation, packed-sequence preprocessing, and MoE replay.

Scale Support

Without persistent fp8_param, serialized FP8 synchronization supports both FP32 and power-of-two block scales. Hopper supports both modes and defaults to FP32 scales; the current Blackwell path requires power-of-two scales.

Persistent fp8_param is supported only with FP32 scales. The main reason is that at a 1e-6 learning rate, 99.97-99.999% of sampled FP32 master values changed per update, no power-of-two block scale changed. The configuration therefore rejects power-of-two scales with fp8_param. Because Blackwell currently requires power-of-two scales, persistent FP8 parameters are not enabled there.

Configuration

<role> denotes policy or ref. All FP8 features are disabled by default, so existing BF16 configurations retain their current behavior.

Args	Default	Validated FP8 value	Behavior
generator.inference_engine.fp8_weight_sync_mode	null	serialized_blockwise	Sends rollout weights as E4M3 tensors with matching FP32 block scales.
generator.inference_engine.engine_init_kwargs.quantization	unset	fp8	Enables FP8 vLLM execution.
trainer.<role>.megatron_config.transformer_config_kwargs.fp8	unset	e4m3	Enables Transformer Engine FP8 compute for that Megatron model.
trainer.<role>.megatron_config.transformer_config_kwargs.fp8_recipe	unset	blockwise	Selects the Transformer Engine blockwise FP8 recipe.
trainer.<role>.megatron_config.transformer_config_kwargs.fp8_param	unset (false)	true	Stores Megatron primary parameters persistently in E4M3.
trainer.policy.megatron_config.ddp_config.fp8_param_gather	false	true with policy fp8_param	Refreshes policy FP8 parameters from updated FP32 optimizer masters after each step.

Environment variable	Runtime default	Validated value (Hopper)	Behavior
NVTE_FP8_BLOCK_SCALING_FP32_SCALES	1 when FP8 is active	1	1 uses FP32 block scales; 0 uses power-of-two scales. Persistent fp8_param requires 1.
NVTE_FP8_BLOCK_AMAX_EPSILON	0	1e-4	Applies the same non-negative block-amax floor to Transformer Engine and serialized rollout quantization.
VLLM_USE_DEEP_GEMM_E8M0	0 with serialized FP32 scales	0	Prevents vLLM from converting FP32 scales to E8M0 power-of-two scales.

serialized_blockwise requires Megatron and full-weight synchronization; adapter-only Megatron LoRA is rejected. On Blackwell, use power-of-two scales and keep fp8_param disabled.

Performance

vLLM Rollout Throughput

On Qwen/Qwen3.5-9B-Base, NVIDIA H100 80GB, TP1, vLLM 0.23.0, using DAPO recipe dataset.

Concurrent requests	BF16 output tok/s	FP8 output tok/s	Gain
1	140.5	198.6	+41.4%
16	1,940.3	2,599.4	+34.0%
32	3,418.9	4,361.0	+27.6%
128	7,897.7	8,989.2	+13.8%
256	9,223.3	10,791.0	+17.0%

This gain comes from FP8 vLLM execution and does not require persistent FP8 Megatron parameters.

End-to-End RL Step Timing

The end-to-end comparison uses Qwen/Qwen3.5-9B-Base on eight NVIDIA H100 80GB GPUs with Megatron TP2, eight TP1 vLLM engines, and a training batch size of 32.

Phase	BF16 mean (s)	FP8 mean (s)	FP8 time change	Speedup
Rollout generation	96.51	72.25	-25.1%	1.34x
Log-probability and reward forward passes	109.87	106.31	-3.2%	1.03x
Policy training	46.50	50.47	+8.5%	0.92x
Policy-to-vLLM weight sync	10.03	10.26	+2.3%	0.98x
End-to-end step	268.90	245.49	-8.7%	1.10x

Persistent FP8 Policy HBM

Megatron primary storage	HBM/GPU	Change at highest sample
BF16	59,676-63,638 MiB	Baseline
Persistent E4M3	56,850-57,810 MiB	-5,828 MiB (-5.69 GiB, -9.2%)

Quality Checks

The following WANDB runs log the metrics. All FP8 runs aligns well with the BF16 baseline.
Qwen3.5-4B runs: https://wandb.ai/sky-posttraining-uc-berkeley/qwen35_4b_dapo_mainclean_fp8align_fastbatch_bf16_fp8_20260630?nw=nwuserjinghanyao1

Qwen3.5-9B runs: https://wandb.ai/sky-posttraining-uc-berkeley/qwen35_9b_h100_tp2_bf16_fp8_prready_20260709?nw=nwuserjinghanyao1

Qwen3.5-35B-A3B runs: https://wandb.ai/sky-posttraining-uc-berkeley/qwen35_35b_a3b_fp8_amaxeps_20260705?nw=nwuserjinghanyao1

[gemini-code-assist]

Code Review

This pull request introduces blockwise FP8 serialization for rollout weight synchronization (specifically targeting Qwen3.5 layouts), adds chunking to vocab_parallel_entropy to bound temporary memory usage, and implements a colocated worker memory barrier with hard eviction to manage HBM footprint. It also updates sequence packing and alignment logic for FP8/TP/CP layouts and supports persistent FP8 parameter training. The review feedback highlights two critical issues: first, the bytes_per_token calculation for vocab entropy chunking does not account for batch dimensions, which could lead to OOMs; second, accessing megatron_config directly in prepare_runtime_environment will crash FSDP training runs due to an AttributeError, requiring defensive guards.