[feat] 2/n Adding NVFP4 Humans& Quantization

.claude/docs/backends/megatron.md

@@ -24,6 +24,53 @@ Key strategies:
 
 Note: Sequence parallelism is auto-enabled when `tensor_model_parallel_size > 1` — there is no separate config field for it.
 
+## Expert MXFP8
+
+Set `trainer.policy.model.expert_mxfp8.enabled=true` on SM100/SM103 to run
+routed expert GEMMs with Transformer Engine MXFP8 and configure vLLM for
+expert-only online MXFP8. Checkpoints and weight sync remain high precision.
+Set `trainer.policy.model.expert_mxfp8.persistent=true` with
+`trainer.policy.megatron_config.ddp_config.fp8_param_gather=true` to keep
+routed-expert primary parameters in MXFP8 between optimizer steps. Set
+`generator.inference_engine.serialized_weight_sync_mode=serialized_mxfp8` to send
+routed experts to vLLM as MXFP8 data and E8M0 scales. Unmerged LoRA is not
+supported. Shared model layouts, MXFP8 format handling, and Megatron/vLLM
+adapters live in `skyrl/backends/skyrl_train/quantization/`.
+
+## Expert NVFP4
+
+Set `trainer.policy.model.expert_nvfp4.enabled=true` to quantize routed
+experts with Transformer Engine NVFP4. Serialized rollout updates require
+`generator.inference_engine.serialized_weight_sync_mode=serialized_nvfp4`
+and use ModelOpt NVFP4 loading in vLLM. Set
+`expert_nvfp4.row_scaled_activation=true` for dynamically quantized per-token
+W4A4 activations; otherwise rollout uses W4A16.
+
+Row-scaled activations, scoped Four-Over-Six, and dequantized backward are
+configured under `trainer.policy.model.expert_nvfp4`. These options require a
+Transformer Engine build that exposes the corresponding NVFP4 recipe fields.
+
+The Humans&-style RL recipe keeps BF16 primary parameters and uses all of
+these settings together:
+
+```text
+trainer.policy.model.expert_nvfp4.enabled=true
+trainer.policy.model.expert_nvfp4.training=true
+trainer.policy.model.expert_nvfp4.persistent=false
+trainer.policy.model.expert_nvfp4.backward_override=dequantized
+trainer.policy.model.expert_nvfp4.row_scaled_activation=true
+trainer.policy.model.expert_nvfp4.disable_rht=true
+trainer.policy.model.expert_nvfp4.disable_stochastic_rounding=true
+trainer.policy.model.expert_nvfp4.disable_2d_quantization=true
+trainer.policy.model.expert_nvfp4.four_over_six_scope=all
+trainer.policy.model.expert_nvfp4.four_over_six_e4m3_use_256_scope=all
+trainer.policy.model.expert_nvfp4.four_over_six_error_mode=MSE
+trainer.policy.model.expert_nvfp4.four_over_six_error_use_fast_math=true
+trainer.policy.model.expert_nvfp4.disable_fp4_quant_fast_math=true
+trainer.policy.model.expert_nvfp4.high_precision_last_layers=8
+generator.inference_engine.serialized_weight_sync_mode=serialized_nvfp4
+```
+
 ## Test Requirements
 
 Megatron GPU tests need: `NVTE_FLASH_ATTN=0`

.claude/docs/weight_sync.md

@@ -16,6 +16,9 @@ skyrl/backends/skyrl_train/weight_sync/
 └── weight_extractor_utils.py
 ```
 
+Quantization policies and serialized formats live in
+`skyrl/backends/skyrl_train/quantization/`.
+
 vLLM worker-extension class (loaded via `--worker-extension-cls`):
 
 - `skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py` — `NewInferenceWorkerWrap`. Three-phase chunked lifecycle.

docs/content/docs/examples/quantized_rollouts.mdx

@@ -31,6 +31,42 @@ generator.inference_engine.engine_init_kwargs.quantization=fp8
 
 This uses vLLM's [online dynamic FP8 quantization](https://docs.vllm.ai/en/latest/features/quantization/fp8.html), so no calibration data or pre-quantized checkpoint is required.
 
+## Expert-only MXFP8
+
+On Blackwell GPUs, the Megatron backend can run routed MoE expert GEMMs in
+MXFP8 while vLLM serves the same experts with online MXFP8 quantization:
+
+```bash
+trainer.strategy=megatron \
+trainer.policy.model.expert_mxfp8.enabled=true \
+trainer.policy.model.expert_mxfp8.persistent=true \
+trainer.policy.megatron_config.ddp_config.fp8_param_gather=true \
+generator.inference_engine.serialized_weight_sync_mode=serialized_mxfp8
+```
+
+Persistent expert weights are quantized after each optimizer step and reused by
+Transformer Engine until the next update. Weight sync sends routed experts as
+MXFP8 data with E8M0 scales, while checkpoints remain high precision. Attention,
+dense layers, shared experts, the router, and the LM head remain high precision.
+
+This mode requires SM100 or SM103 and does not support unmerged LoRA. Set
+`training=false` or `rollout=false` under `expert_mxfp8` to enable only one
+side for benchmarking. Set `persistent=false` to quantize expert weights during
+each forward pass. Leave `serialized_weight_sync_mode=null` to let vLLM
+quantize experts after each weight sync. `fp8_weight_sync_mode` remains a
+deprecated alias.
+
+Run the Modal benchmark with identical BF16 and MXFP8 workloads:
+
+```bash
+uv run --isolated --with modal modal run --detach \
+  examples/train/megatron/modal_expert_mxfp8_benchmark.py \
+  --mode both --steps 10
+```
+
+Use `--mode bf16` or `--mode mxfp8` for one side. Set `MODAL_GPU` to override
+the default `B200:8` allocation.
+
 ## Enabling off-policy correction (TIS)
 
 To apply TIS, we need the inference engine to return the rollout logprobs for the generated tokens, and we configure the correction on the policy loss: