diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py index e394859002..70b59ce53d 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py @@ -134,6 +134,9 @@ def _get_scale_format(args, name, weight_block_size): if ".experts." not in name: # Non-MoE linear weights: ue8m0 when deepgemm is enabled + dense_backend = getattr(args, "sglang_fp8_gemm_runner_backend", "auto") + if dense_backend not in ("auto", "deep_gemm"): + return None return "ue8m0" # MoE expert weights: only ue8m0 when runner is deep_gemm diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_mxfp8.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_mxfp8.py index e6e86e4457..829cb14f96 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_mxfp8.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_mxfp8.py @@ -94,6 +94,12 @@ def quantize_params_mxfp8(args, megatron_name, converted_named_params, quantizat "self_attention.linear_kv_up_proj.weight", "self_attention.wq_b.weight", "self_attention.wk.weight", + # DeepSeek V4 attention + "self_attention.wq_a.weight", + "self_attention.wkv.weight", + "self_attention.wo_b.weight", + "self_attention.indexer.linear_wq_b.weight", + "self_attention.indexer.linear_wk.weight", ]: quantize_named_params = [] for converted_name, param in converted_named_params: diff --git a/miles_plugins/models/deepseek_v4/ops/compressor.py b/miles_plugins/models/deepseek_v4/ops/compressor.py index 55a93f2949..a7ef50e71b 100644 --- a/miles_plugins/models/deepseek_v4/ops/compressor.py +++ b/miles_plugins/models/deepseek_v4/ops/compressor.py @@ -1,3 +1,5 @@ +import os + import einops import torch import torch.nn as nn @@ -6,7 +8,7 @@ from miles_plugins.models.deepseek_v4.ops.cp_utils import all_gather_cp, get_freqs_cis_for_cp from miles_plugins.models.deepseek_v4.ops.kernel.precision_aligned_ops import linear_bf16_fp32 -from miles_plugins.models.deepseek_v4.ops.qat import fp8_simulate_qat +from miles_plugins.models.deepseek_v4.ops.qat import fp8_simulate_qat, mxfp4_simulate_qat from miles_plugins.models.deepseek_v4.ops.rope import apply_rotary_emb, wrapped_precompute_freqs_cis from miles_plugins.models.deepseek_v4.ops.utils import rotate_activation @@ -77,6 +79,7 @@ def __init__( self.rotate = rotate coff = 1 + self.overlap self.use_fp8_qat = config.fp8 is not None + self.use_mxfp4_qat = os.environ.get("USE_MXFP4_QAT", "0") == "1" self.cp_group = cp_group self.cp_size = cp_group.size() if cp_group is not None else 1 @@ -155,7 +158,9 @@ def forward_raw(self, x: torch.Tensor) -> torch.Tensor: if self.rotate: kv = rotate_activation(kv) - if self.use_fp8_qat: + if self.use_mxfp4_qat: + kv = mxfp4_simulate_qat(kv, 32) + elif self.use_fp8_qat: kv = fp8_simulate_qat(kv, 128) else: if self.use_fp8_qat: diff --git a/miles_plugins/models/deepseek_v4/ops/kernel/act_quant.py b/miles_plugins/models/deepseek_v4/ops/kernel/act_quant.py index a26e90e877..e9a75b0108 100644 --- a/miles_plugins/models/deepseek_v4/ops/kernel/act_quant.py +++ b/miles_plugins/models/deepseek_v4/ops/kernel/act_quant.py @@ -164,3 +164,86 @@ def act_quant( x.copy_(y) return x return y, s + + +@tilelang.jit(pass_configs=pass_configs) +def fp4_quant_kernel(N, block_size=32, in_dtype=BF16, scale_dtype=FE8M0, inplace=False): + """Block-wise FP4 (E2M1) quantization with power-of-2 (UE8M0) scales. + + ``inplace=True`` does a fused quant+dequant back to ``in_dtype`` (the MXFP4 + fake-quant used for QAT). ``inplace=False`` emits packed E2M1 plus the UE8M0 + scale. The pow2 scale is computed via the same IEEE-754 bit tricks as the FP8 + path (``fast_round_scale``); the inner cast goes through ``FP4`` explicitly so + the value is rounded onto the E2M1 grid regardless of the output storage dtype. + """ + M = T.symbolic("M") + fp4_max = 6.0 + fp4_max_inv = 1.0 / fp4_max + blk_m = 32 + group_size = block_size + compute_dtype = FP32 + out_dtype = in_dtype if inplace else FP4 + + @T.prim_func + def fp4_quant_kernel_( + X: T.Tensor[(M, N), in_dtype], + Y: T.Tensor[(M, N), out_dtype], + S: T.Tensor[(M, T.ceildiv(N, group_size)), scale_dtype], + ): + with T.Kernel(T.ceildiv(M, blk_m), T.ceildiv(N, group_size), threads=128) as ( + pid_m, + pid_n, + ): + x_shared = T.alloc_shared((blk_m, group_size), in_dtype) + x_local = T.alloc_fragment((blk_m, group_size), in_dtype) + amax_local = T.alloc_fragment((blk_m,), compute_dtype) + s_local = T.alloc_fragment((blk_m,), compute_dtype) + y_local = T.alloc_fragment((blk_m, group_size), out_dtype) + y_shared = T.alloc_shared((blk_m, group_size), out_dtype) + + for _ in T.Pipelined(1, num_stages=2): + T.copy(X[pid_m * blk_m, pid_n * group_size], x_shared) + T.copy(x_shared, x_local) + T.reduce_absmax(x_local, amax_local, dim=1) + for i in T.Parallel(blk_m): + amax_local[i] = T.max(amax_local[i], 6 * (2**-126)) + s_local[i] = fast_round_scale(amax_local[i], fp4_max_inv) + if inplace: + for i, j in T.Parallel(blk_m, group_size): + y_local[i, j] = T.Cast( + out_dtype, + T.Cast(compute_dtype, T.Cast(FP4, T.clamp(x_local[i, j] / s_local[i], -fp4_max, fp4_max))) + * s_local[i], + ) + else: + for i, j in T.Parallel(blk_m, group_size): + y_local[i, j] = T.clamp(x_local[i, j] / s_local[i], -fp4_max, fp4_max) + for i in T.Parallel(blk_m): + S[pid_m * blk_m + i, pid_n] = T.Cast(scale_dtype, s_local[i]) + T.copy(y_local, y_shared) + T.copy(y_shared, Y[pid_m * blk_m, pid_n * group_size]) + + return fp4_quant_kernel_ + + +def fp4_act_quant( + x: torch.Tensor, + block_size: int = 32, + inplace: bool = False, +) -> torch.Tensor: + """Block-wise FP4 (E2M1, ``1 x block_size`` UE8M0) quantization. + + ``inplace=True`` does the fused quant+dequant back to BF16 (the MXFP4 fake-quant + used by QAT) and returns ``x``; ``inplace=False`` returns ``(packed_e2m1, scale)``. + """ + N = x.size(-1) + assert N % block_size == 0 + z = x.contiguous() + y = torch.empty_like(z) if inplace else z.new_empty(*z.shape[:-1], N // 2, dtype=torch.float4_e2m1fn_x2) + s = z.new_empty(*z.size()[:-1], N // block_size, dtype=torch.float8_e8m0fnu) + kernel = fp4_quant_kernel(N, block_size, inplace=inplace) + kernel(z.view(-1, N), y.view(-1, y.size(-1)), s.view(-1, N // block_size)) + if inplace: + x.copy_(y) + return x + return y, s diff --git a/miles_plugins/models/deepseek_v4/ops/qat.py b/miles_plugins/models/deepseek_v4/ops/qat.py index 60cadea9b3..8db72ec9f7 100644 --- a/miles_plugins/models/deepseek_v4/ops/qat.py +++ b/miles_plugins/models/deepseek_v4/ops/qat.py @@ -2,7 +2,7 @@ from tile_kernels.quant import per_token_cast_back -from .kernel.act_quant import act_quant +from .kernel.act_quant import act_quant, fp4_act_quant def fp8_simulate(x: torch.Tensor, block_size: int): @@ -13,7 +13,10 @@ def fp8_simulate(x: torch.Tensor, block_size: int): the rest of the DeepSeek stack. """ x_c = x.contiguous() - y, scale = act_quant(x_c, block_size, "ue8m0") + # Force fp32 scale storage: act_quant's auto dtype picks float8_e8m0fnu on + # Blackwell+DeepGEMM, but per_token_cast_back below only accepts int32/fp32 + # scales. scale_fmt="ue8m0" still rounds the values to powers of two. + y, scale = act_quant(x_c, block_size, "ue8m0", scale_dtype=torch.float32) N = x_c.size(-1) y_flat = y.view(-1, N) @@ -34,3 +37,27 @@ def backward(ctx, grad_kv): fp8_simulate_qat = DeepSeekV4LinearQATFunc.apply + + +def mxfp4_simulate(x: torch.Tensor, block_size: int = 32): + """Simulate MXFP4 (E2M1) cast + dequant with ``1 x block_size`` UE8M0 (pow2) scaling. + + Mirrors :func:`fp8_simulate` but for FP4: a single fused on-device kernel + (``fp4_act_quant`` with ``inplace=True``) rounds each value onto the E2M1 grid + and dequantizes back to ``x.dtype``. We operate on ``x.clone()`` so the op is + non-mutating (safe inside an autograd Function / straight-through estimator). + """ + return fp4_act_quant(x.clone(), block_size, inplace=True).to(x.dtype) + + +class DeepSeekV4MXFP4LinearQATFunc(torch.autograd.Function): + @staticmethod + def forward(ctx, kv, block_size=32): + return mxfp4_simulate(kv, block_size) + + @staticmethod + def backward(ctx, grad_kv): + return grad_kv, None + + +mxfp4_simulate_qat = DeepSeekV4MXFP4LinearQATFunc.apply diff --git a/miles_plugins/models/deepseek_v4/ops/v4_indexer.py b/miles_plugins/models/deepseek_v4/ops/v4_indexer.py index 45ba7c6030..d31d7ce2bc 100644 --- a/miles_plugins/models/deepseek_v4/ops/v4_indexer.py +++ b/miles_plugins/models/deepseek_v4/ops/v4_indexer.py @@ -1,3 +1,5 @@ +import os + import einops import torch from megatron.core import parallel_state @@ -13,7 +15,7 @@ _make_causal_cu_seqlens, batched_indexer_fwd, ) -from miles_plugins.models.deepseek_v4.ops.qat import fp8_simulate_qat +from miles_plugins.models.deepseek_v4.ops.qat import fp8_simulate_qat, mxfp4_simulate_qat from miles_plugins.models.deepseek_v4.ops.rope import apply_rotary_emb, wrapped_precompute_freqs_cis from miles_plugins.models.deepseek_v4.ops.utils import rotate_activation @@ -32,6 +34,7 @@ def __init__(self, config: TransformerConfig, pg_collection=None): self.rope_head_dim = config.qk_pos_emb_head_dim self.compress_ratio = 4 self.use_fp8_qat = config.fp8 is not None + self.use_mxfp4_qat = os.environ.get("USE_MXFP4_QAT", "0") == "1" if pg_collection is None: pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=["tp", "cp"]) @@ -106,7 +109,9 @@ def forward(self, x: torch.Tensor, qr: torch.Tensor, mask=None, packed_seq_param q = einops.rearrange(q, "b s ... -> s b ...") q = rotate_activation(q) - if self.use_fp8_qat: + if self.use_mxfp4_qat: + q = mxfp4_simulate_qat(q, 32) + elif self.use_fp8_qat: q = fp8_simulate_qat(q, 128) k = self.compressor(x) @@ -127,6 +132,10 @@ def forward(self, x: torch.Tensor, qr: torch.Tensor, mask=None, packed_seq_param cu_ks = cu_ks[cp_rank * seqlen : (cp_rank + 1) * seqlen] cu_ke = cu_ke[cp_rank * seqlen : (cp_rank + 1) * seqlen] index_scores = batched_indexer_fwd(q, k, weights.float(), cu_ks, cu_ke) + if self.use_mxfp4_qat: + # QAT: index scores are cast FP32 -> BF16 (and back) to match the + # BF16 top-k selector used at inference (~2x faster top-k). + index_scores = index_scores.to(torch.bfloat16).to(torch.float32) topk_count = min(self.index_topk, index_scores.size(-1)) topk_indices = index_scores.topk(topk_count, dim=-1)[1] diff --git a/scripts/run_deepseek_v4.py b/scripts/run_deepseek_v4.py index 88f94e3a16..8e01b03ca3 100644 --- a/scripts/run_deepseek_v4.py +++ b/scripts/run_deepseek_v4.py @@ -56,6 +56,19 @@ _PRO_MODEL_NAMES = ("DeepSeek-V4-Pro-FP8",) _BLACKWELL_HARDWARE = ("B200", "B300", "GB200", "GB300") +_DSV4_TE_PRECISION_CONFIG = """ +configs: + bf16: + transformer_engine_config_type: "TEQuantizationParams" + training_recipe: {} +matchers: + dsa_indexer_weights_proj_bf16: + type: "glob" + enabled: true + pattern: "*.self_attention.indexer.linear_weights_proj" + config: "bf16" +""".strip() + @dataclass class ScriptArgs(U.ExecuteTrainConfig): @@ -103,7 +116,14 @@ class ScriptArgs(U.ExecuteTrainConfig): # precision configs enable_r3: bool = True train_deterministic: bool = True - fp8_training: bool | None = None + # Precision recipe for BOTH trainer (Megatron/TE) and rollout (sglang checkpoint): + # fp8 - blockwise FP8 128x128 (Hopper: fp32 scales; Blackwell: pow2 scales, MXFP8-emulated) + # mxfp8 - MXFP8 1x32 ue8m0 (Blackwell only; rollout serves the BF16->MXFP8 converted ckpt) + # mxfp4_qat - MXFP4 QAT on top of blockwise FP8 (Blackwell only): sets USE_MXFP4_QAT=1 so the + # model MXFP4 fake-quantizes the indexer q / compressor kv + casts index scores to + # BF16, and MoE expert weights are MXFP4 fake-quantized before TE's FP8 cast. + # bf16 - BF16 training; rollout serves the source FP8 checkpoint + recipe: Literal["fp8", "mxfp8", "bf16", "mxfp4_qat"] = "fp8" enable_mis: bool = False # pass any extra sglang/miles/megatron args through `--extra-args '--your-arg'` @@ -116,6 +136,11 @@ def __post_init__(self): self.model_local_dir = self.model_dir if self.model_name in _PRO_MODEL_NAMES: self.enable_r3 = False + if self.hardware in ("H100", "H200"): + assert self.recipe not in ( + "mxfp8", + "mxfp4_qat", + ), "MXFP8 / MXFP4-QAT require Blackwell (no native MXFP4/MXFP8 on H100/H200)" assert self.rollout_num_nodes >= 0 assert self.rollout_num_nodes < self.num_nodes self.colocate = self.rollout_num_nodes == 0 @@ -140,6 +165,10 @@ def bf16_name(self): return "DeepSeek-V4-Pro-BF16" return f"{self.model_name}-bf16" + @property + def mxfp8_name(self): + return f"{self.model_name}-MXFP8" + def _is_blackwell(args: ScriptArgs) -> bool: if args.hardware != "auto": @@ -153,12 +182,6 @@ def _is_blackwell(args: ScriptArgs) -> bool: return major >= 10 -def _resolve_precision_defaults(args: ScriptArgs): - if args.fp8_training is None: - args.fp8_training = not _is_blackwell(args) - print(f"[precision] fp8_training auto -> {args.fp8_training}") - - def _download_dataset(args: ScriptArgs): """Download the task-specific dataset(s).""" match args.task: @@ -223,6 +246,28 @@ def prepare_single(args: ScriptArgs): _prepare_single(args) +def _prepare_mxfp8(args: ScriptArgs): + """BF16 -> MXFP8 conversion for sglang rollout (Blackwell only). + + head/wo_a/ffn.gate/compressor/norms/embed and the DSA indexer weights_proj + are all kept BF16 by SKIP_WEIGHT_SUBSTRINGS in tools/convert_hf_to_mxfp8.py. + """ + if args.recipe != "mxfp8": + return + U.exec_command( + f"python tools/convert_hf_to_mxfp8.py " + f"--model-dir {args.model_dir}/{args.bf16_name} " + f"--save-dir {args.model_dir}/{args.mxfp8_name} " + ) + + +@app.command() +@U.dataclass_cli +def prepare_mxfp8(args: ScriptArgs): + """BF16 -> MXFP8 conversion (needs prepare-single done first). One node.""" + _prepare_mxfp8(args) + + def _prepare_spmd(args: ScriptArgs): is_4layer = args.model_name == "DeepSeek-V4-Flash-FP8-4layer" actor_num_nodes = args.actor_num_nodes @@ -291,9 +336,10 @@ def _prepare_cp(args: ScriptArgs): path_dst=f"{args.model_local_dir}/{args.torch_dist_name}", num_nodes=args.num_nodes, ) + hf_name = args.mxfp8_name if args.recipe == "mxfp8" else args.model_name U.rsync_simple( - path_src=f"{args.model_dir}/{args.model_name}", - path_dst=f"{args.model_local_dir}/{args.model_name}", + path_src=f"{args.model_dir}/{hf_name}", + path_dst=f"{args.model_local_dir}/{hf_name}", num_nodes=args.num_nodes, ) @@ -365,7 +411,13 @@ def _get_parallel_config(args: ScriptArgs) -> str: def _train(args: ScriptArgs): - _resolve_precision_defaults(args) + if args.recipe == "mxfp8": + assert _is_blackwell(args), "MXFP8 requires Blackwell (B200/B300/GB200/GB300)" + mxfp8_checkpoint = f"{args.model_local_dir}/{args.mxfp8_name}" + if args.hf_checkpoint != mxfp8_checkpoint: + print(f"[precision] recipe=mxfp8: hf_checkpoint {args.hf_checkpoint} -> {mxfp8_checkpoint}") + args.hf_checkpoint = mxfp8_checkpoint + print(f"[precision] recipe={args.recipe}") print( f"running on {args.num_nodes} nodes " f"({args.actor_num_nodes} actor nodes x {args.actor_num_gpus_per_node} GPUs/node, " @@ -472,8 +524,21 @@ def _train(args: ScriptArgs): sglang_dp_size = sglang_world_size sglang_ep_size = sglang_world_size sglang_a2a_backend = None + + match args.recipe: + case "mxfp8": + sglang_fp8_gemm_backend = "flashinfer_cutlass" + sglang_moe_runner_backend = "flashinfer_trtllm_routed" + case "mxfp4_qat": + sglang_fp8_gemm_backend = "flashinfer_mxfp4" + sglang_moe_runner_backend = "auto" + case _: # fp8 / bf16 + sglang_fp8_gemm_backend = "auto" + sglang_moe_runner_backend = "auto" sglang_args = ( f"--rollout-num-gpus-per-engine {sglang_world_size} " + f"--sglang-fp8-gemm-backend {sglang_fp8_gemm_backend} " + f"--sglang-moe-runner-backend {sglang_moe_runner_backend} " f"--sglang-tp-size {sglang_tp_size} " f"--sglang-dp-size {sglang_dp_size} " "--sglang-enable-dp-attention " @@ -563,11 +628,43 @@ def _train(args: ScriptArgs): "CUBLAS_WORKSPACE_CONFIG": ":4096:8", } - if args.fp8_training: - misc_args += "--transformer-impl transformer_engine " "--bf16 " "--fp8-format e4m3 " "--fp8-recipe blockwise " - extra_env_vars |= { - "NVTE_FP8_BLOCK_SCALING_FP32_SCALES": "1", - } + match args.recipe: + case "mxfp8": + misc_args += "--transformer-impl transformer_engine " "--bf16 " "--fp8-format e4m3 " "--fp8-recipe mxfp8 " + case "fp8": + misc_args += ( + "--transformer-impl transformer_engine " "--bf16 " "--fp8-format e4m3 " "--fp8-recipe blockwise " + ) + if not _is_blackwell(args): + extra_env_vars |= { + "NVTE_FP8_BLOCK_SCALING_FP32_SCALES": "1", + } + else: + # actor_group.py unconditionally injects NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1 + # into train actors; on Blackwell the TE blockwise recipe is emulated with + # MXFP8 and requires power-of-two scales, so force it back to 0 here. + misc_args += """--train-env-vars '{"NVTE_FP8_BLOCK_SCALING_FP32_SCALES":"0"}' """ + case "mxfp4_qat": + # MXFP4 QAT builds on the blockwise FP8 (128x128) framework. USE_MXFP4_QAT=1 + # turns on the in-model MXFP4 fake-quant (indexer q / compressor kv), the + # FP32->BF16 index-score cast, and the MoE expert-weight MXFP4 fake-quant in + # TE's _get_weight_tensors. Because FP4->FP8 is lossless, the existing + # blockwise FP8 GEMM runs unchanged. Blackwell-only (asserted above), so + # force power-of-two scales like the Blackwell fp8 path. + misc_args += ( + "--transformer-impl transformer_engine " "--bf16 " "--fp8-format e4m3 " "--fp8-recipe blockwise " + ) + misc_args += ( + """--train-env-vars '{"USE_MXFP4_QAT":"1","NVTE_FP8_BLOCK_SCALING_FP32_SCALES":"0"}' """ + ) + case "bf16": + pass + + if args.recipe in ("mxfp8", "fp8", "mxfp4_qat") and "--te-precision-config-file" not in args.extra_args: + misc_args += ( + f"--te-precision-config-file " + f"{U.save_to_temp_file(_DSV4_TE_PRECISION_CONFIG, 'yaml')} " + ) train_args = ( f"{ckpt_args} " @@ -611,6 +708,13 @@ def full_train(args: ScriptArgs): else: print(f"[full_train] Skipping FP8->BF16 cast: {bf16_sentinel} already exists.") + if args.recipe == "mxfp8": + mxfp8_sentinel = Path(f"{args.model_dir}/{args.mxfp8_name}") / "model.safetensors.index.json" + if not mxfp8_sentinel.exists(): + _prepare_mxfp8(args) + else: + print(f"[full_train] Skipping BF16->MXFP8 conversion: {mxfp8_sentinel} already exists.") + torch_dist_dir = Path(f"{args.model_dir}/{args.torch_dist_name}") torch_dist_sentinel = torch_dist_dir / "latest_checkpointed_iteration.txt" if not torch_dist_sentinel.exists(): diff --git a/tools/convert_hf_to_fp8.py b/tools/convert_hf_to_fp8.py index 4043025e54..00e4ea7673 100644 --- a/tools/convert_hf_to_fp8.py +++ b/tools/convert_hf_to_fp8.py @@ -137,6 +137,10 @@ def process_file(input_path, output_path, filename, strategy, block_size, result and "lm_head" not in key and "eh_proj" not in key and "weights_proj" not in key + and "head." not in key + and "wo_a" not in key + and "ffn.gate." not in key + and "compressor." not in key and "vision_tower" not in key and "mm_projector" not in key ): diff --git a/tools/convert_hf_to_mxfp8.py b/tools/convert_hf_to_mxfp8.py index d99db93f40..addc03939a 100644 --- a/tools/convert_hf_to_mxfp8.py +++ b/tools/convert_hf_to_mxfp8.py @@ -45,6 +45,10 @@ "lm_head", "eh_proj", "weights_proj", + "head.", + "wo_a", + "ffn.gate.", + "compressor.", ) SOURCE_FP8_BLOCK_SIZE = [128, 128]