Low Precision Training and Rollout#

Low precision in vime is primarily used to make rollout faster and more memory-efficient while keeping training numerically stable. For large MoE RL jobs, the recommended production path is:

BF16 training in Megatron + FP8 rollout/inference in vLLM

Megatron keeps the trainable checkpoint in BF16/torch_dist format. vLLM serves an FP8 Hugging Face checkpoint for rollout. During weight updates, vime uses the quantization config in --hf-checkpoint to quantize updated BF16 weights before sending them to vLLM.

Feature Maturity#

Feature	Status	Recommended Use
BF16 training + FP8 rollout/inference	Stable	Default path for large MoE RL recipes. Keeps training stable while reducing rollout memory and bandwidth.
FP8 attention KV cache in vLLM rollout	Requires model, attention backend and GPU support	Increase attention KV capacity with `--vllm-kv-cache-dtype fp8_e4m3`; hybrid recurrent state is configured separately.
INT4 rollout / INT4 QAT	Beta	Use when rollout memory/throughput pressure is high and the model path has been validated.
FP8 training + FP8 rollout	Experimental	Useful for research on training/inference mismatch and throughput, but still has optimizer and checkpointing caveats.

BF16 Training with FP8 Rollout#

This is the main production path in vime.

You can run FP8 rollout by setting --hf-checkpoint to a blockwise-quantized Hugging Face checkpoint. Convert a BF16 checkpoint with:

python tools/convert_hf_to_fp8.py \
    --model-dir $BF16_MODEL \
    --save-dir $FP8_MODEL \
    --strategy block --block-size 128 128 \
    --max-workers 4

Make sure the converted checkpoint’s config.json contains the correct quantization_config. vime uses that config during weight updates, so the training side can remain BF16 while rollout receives FP8 weights.

Example:

# Megatron training checkpoint remains BF16 / torch_dist.
--ref-load /path/to/model_torch_dist

# vLLM rollout checkpoint is FP8 Hugging Face.
--hf-checkpoint /path/to/model-fp8-hf

FP8 KV Cache for Rollout#

For long-context, multi-turn, or agentic workloads, KV cache capacity is often the bottleneck. Because vLLM arguments are passed through by adding --vllm-, you can enable FP8 KV cache directly:

--vllm-kv-cache-dtype fp8_e4m3

This is a rollout-side setting. It does not change Megatron training precision; it increases effective vLLM KV cache capacity and can allow longer contexts or higher concurrency, subject to the accuracy/performance behavior of your vLLM version and GPU stack.

For hybrid models such as Qwen3.8-27B and GLM-5.3-Flash, this flag applies only to attention-layer KV. Linear-attention recurrent state uses --vllm-mamba-ssm-cache-dtype, while vLLM manages attention KV and recurrent state through a shared paged-cache budget rather than independently sized pools. FP8 KV does not imply FP8 recurrent or convolution state. The builder leaves these overrides unset by default, and offers FP32/BF16 recurrent state for Qwen3.8-27B. See hybrid cache settings, architecture sources and memory derivation.

FP8 Training with FP8 Rollout#

vime also supports experimental FP8 training paths. We observed that FP8 training plus FP8 inference can improve inference throughput and reduce training/inference mismatch in some settings. More details are available in this blog.

FP8 Quick Start#

Convert your Hugging Face model weights to FP8 format using tools/convert_hf_to_fp8.py.
Add the FP8 training flags:

--fp8-format e4m3
--fp8-recipe blockwise
# --fp8-param-gather # optional; currently incompatible with CPU Adam

Ensure NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 is set. vime sets this to 1 by default for Ray actors.
Start an FP8 training example:

# Qwen3-4B FP8 training
bash scripts/low_precision/run-qwen3-4b-fp8.sh

# Qwen3-30B-A3B FP8 training (2 nodes)
bash scripts/low_precision/run-qwen3-30b-a3b-fp8.sh

Implementation Notes#

If an FP8 recipe is enabled, TransformerEngine layers are built in an FP8 context.
During training, weights and activations are quantized online to NVFP8 format, and cuBLAS FP8 GEMM is used for forward and backward GEMMs.
During RL weight updates, Megatron dequantizes FP8 weights to BF16, then vime quantizes the BF16 weights to FP8 and sends them to vLLM.
Checkpoints saved from the training engine are dequantized back to BF16 and saved as torch_dist.

Only Linear and GroupLinear layers in TransformerEngine use FP8. embedding and lm_head remain in their original precision. If --fp8-param-gather is not enabled, TransformerEngine weights remain stored in BF16 and are cast to FP8 only during GEMM or GroupGEMM.

Known Caveat#

--fp8-param-gather can save memory, but currently requires TransformerEngine FusedAdam, which conflicts with the CPU Adam offload path commonly used for large Megatron-LM RL jobs.

INT4 QAT Training#

INT4 STE (Straight-Through Estimator) training and INT4 inference can further reduce rollout memory and improve throughput. Treat this path as beta unless you have validated the target model and reward setup.

INT4 Quick Start#

Convert Hugging Face weights to INT4:

python tools/convert_hf_to_int4_direct.py \
  --model-dir /path/to/your/original/models \
  --save-dir /path/to/your/save/models

If you only need INT4 rollout, set --hf-checkpoint to the converted INT4 checkpoint.

Enable INT4 fake QAT:

RUNTIME_ENV_JSON="{
  \"env_vars\": {
    \"OPEN_TRAINING_INT4_FAKE_QAT_FLAG\": \"1\",
    \"OPEN_TRAINING_INT4_GROUP_SIZE\": \"128\"
  }
}"

OPEN_TRAINING_INT4_GROUP_SIZE should usually be:

128 for moonlight-16B-A3B, qwen3-30B-A3B, and qwen3-235B-A22B-int4;
32 for kimi-k2-Thinking-int4.

Launch an example:

# Moonlight-16B-A3B INT4 training
bash scripts/low_precision/run-moonlight-16B-A3B-int4.sh

# Qwen3-30B-A3B INT4 training
bash scripts/low_precision/run-qwen3-30B-A3B-int4.sh

# Qwen3-235B-A22B INT4 training (8 nodes)
bash scripts/low_precision/run-qwen3-235B-A22B-int4.sh

# Kimi-k2-Thinking INT4 training (32 nodes)
bash scripts/low_precision/run-kimi-k2-Thinking-int4.sh

For multi-node environments, start the Ray service according to your cluster configuration.

Numerics oracle for low-precision formats#

vime/utils/quant_ref.py provides pure-torch reference implementations (MXFP8 with UE8M0 scales, NVFP4 with E4M3 block scales, block-FP8 dequant) used as the deterministic oracle for all low-precision quantizers; the contract is pinned by tests/test_low_precision_quant_ref.py (CPU, python -m pytest tests/test_low_precision_quant_ref.py -x -q).

MXFP8 (Blackwell only)#

MXFP8 uses E4M3 values with UE8M0 (power-of-two, uint8) scales per 32-element group. Train with Megatron’s --fp8-recipe mxfp8 and convert/point --hf-checkpoint at an MXFP8 HF checkpoint (tools/convert_hf_to_mxfp8.py); live weight sync then re-quantizes exported weights via megatron_to_hf/processors/quantizer_mxfp8.py. Note the scale tensor is named weight_scale (vLLM ModelOptMxFp8Config contract), not weight_scale_inv. Use --extra-high-precision-layers-hf / --extra-high-precision-layers-megatron to keep selected layers in BF16. MXFP8 requires Blackwell (SM100+) and has no ROCm support.

Offline conversion: python tools/convert_hf_to_mxfp8.py --model-dir <HF ckpt> --save-dir <out> [--num-layers-at-start-in-bf16 N] [--num-layers-at-end-in-bf16 N] [--extra-high-precision-layers-hf ...], or --metadata-only for a weightless rollout schema. BF16 and block-FP8/UE8M0 [128,128] sources are supported; the latter is rescaled to [1,32] without dequantizing when scales are power-of-two.

NVFP4 fake QAT (Blackwell only)#

NVFP4 uses E2M1 values with E4M3 block scales (group 16) plus an FP32 global scale. Training-time fake QAT is enabled with --train-env-vars '{"OPEN_TRAINING_NVFP4_FAKE_QAT_FLAG": "1", "NVTE_USE_FAST_MATH": "0", ...}' and wraps TE grouped-expert weights via a straight-through estimator; the implementation is a portable pure-torch QDQ (backends/megatron_utils/kernels/nvfp4_qdq_torch.py) validated against the quant_ref oracle — miles’ SM10x-only CuTe-DSL kernel is not ported. TE’s NVTE_NVFP4_4OVER6* knobs are honored. NVFP4 has no ROCm support.

Offline conversion: python tools/convert_hf_to_nvfp4.py --model-dir <HF ckpt> --save-dir <out> [--num-layers-at-start-in-bf16 N] [--num-layers-at-end-in-bf16 N] [--extra-high-precision-layers-hf ...]. Only MoE routed-expert GEMMs in the main language decoder are quantized (gate/up share one global scale); MTP, shared experts and dense layers stay BF16. Live weight sync re-quantizes via megatron_to_hf/processors/quantizer_nvfp4.py in the same layout.

In-loop NVFP4 weight updates (NCCL/IPC) require the vLLM patch docker/patch/latest/vllm-nvfp4-inloop.patch (applied by the Dockerfile): the layerwise reload machinery restores layers in kernel format, so the patch makes ModelOptNvFp4FusedMoE recreate checkpoint-format parameters before an update. Verifier: tests/test_nvfp4_inloop_recreate.py (runs on CPU inside the image).

Low Precision Training and Rollout

Contents