FP8 training examples

​

Files

• run-qwen3-4b-fp8.sh: example launch script with Qwen3‑4B in FP8.
• run-qwen3-30b-a3b-fp8-two-nodes.sh: example launch script for running Qwen3‑30B‑A3B in FP8 across two nodes.

​

Quick Start

1. Check if your training script is properly configured. For training tasks, we need to add these flags:
   --fp8-format e4m3
   --fp8-recipe blockwise
   # --fp8-param-gather # [optional] Currently incompatible with CPU Adam
   
   Then ensure the NVTE_FP8_BLOCK_SCALING_FP32_SCALES environment variable is enabled. Note that only Linear and GroupLinear layers in TransformerEngine use fp8 format. embedding and lm_head remain in their original precision. If --fp8-param-gather is not enabled, weights in TransformerEngine remain in bf16 format, only being cast to fp8 format during GEMM or GroupGEMM operations.
2. Convert your HuggingFace model weights to FP8 format. You can use tools/convert_hf_to_fp8.py to convert bf16 weights to fp8 format. Ensure that the --hf-checkpoint parameter points to a directory where the config.json contains the correct quantization_config. miles will automatically use FP8 quantization during weight updates.
3. Start FP8 training.
   cd miles
   
   # Qwen3‑4B FP8 training (single node)
   bash examples/infra_features/low_precision/run-qwen3-4b-fp8.sh
   
   # Qwen3‑30B‑A3B FP8 training (two nodes)
   bash examples/infra_features/low_precision/run-qwen3-30b-a3b-fp8-two-nodes.sh
   
   Following the above command will launch FP8 training.
4. Use the saved checkpoint for evaluation. Note that TransformerEngine does not specifically save FP8 quantized weights; the saved torch dist remains in original precision (usually bf16). If you want to evaluate under FP8, you need to convert the checkpoint from torch_dist to HuggingFace format, then convert to FP8 HuggingFace format.

​

Quick Explanation

1. Initialization: If FP8 recipe is enabled, layers will be built in FP8 context.
2. Training: During training, weights and activations are quantized online to nvfp8 format, and cuBLAS FP8 GEMM is called for various GEMM computations in forward and backward passes.
3. Weight updates: During RL weight updates, Megatron first dequantizes FP8 weights to bf16 format, then miles quantizes these bf16 weights to fp8 format and sends them to sglang. (This additional dequantization and quantization is not elegant, but we haven’t modified the interface yet for framework compatibility.)
4. Save checkpoint: Similar to weight updates, if checkpoints need to be saved from the training engine, they will also be dequantized back to bf16 and saved to torch_dist format checkpoints.

​

TODO

• FP8 weights (--fp8-param-gather) can provide memory savings benefits, but currently FP8 weights must be used with TransformerEngine’s FusedAdam, which conflicts with the commonly used Adam CPU offload technique in Megatron-LM.

​

INT4 Training Examples

​

Files

• run-qwen3‑30B‑A3B-int4.sh: Launch script for Qwen3‑30B‑A3B (INT4) on 8x H200 GPUs.
• run-qwen3-235B-A22B-int4.sh: Launch script for Qwen3-235B-A22B (INT4) on 64x H200 GPUs.
• run-kimi-k2-Thinking-int4.sh: Launch script for Kimi-k2-Thinking (INT4) on 256x H200 GPUs.

​

Quick Start

​

1. Convert HuggingFace Weights to INT4

python tools/convert_hf_to_hf_int4.py \
  --input-dir /path/to/your/original/models \
  --output-dir /path/to/your/save/models \
  --data-dir /path/to/your/wikitext

​

2. Start INT4 Training

• OPEN_TRAINING_INT4_FAKE_QAT_FLAG: Enables fake quantization operations for INT4 training.
• OPEN_TRAINING_INT4_GROUP_SIZE: Specifies the block size (group size) for model quantization.
  • Set to 128 for qwen3-30B-A3B and qwen3-235B-A22B-int4.
  • Set to 32 for kimi-k2-Thinking-int4.

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

# Qwen3‑30B‑A3B Int4 training
bash examples/infra_features/low_precision/run-qwen3‑30B‑A3B-int4.sh

# Qwen3-235B-A22B Int4 training (8 nodes)
bash examples/infra_features/low_precision/run-qwen3-235B-A22B-int4.sh

# Kimi-k2-Thinking Int4 training (32 nodes)
bash examples/infra_features/low_precision/run-kimi-k2-Thinking-int4.sh

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