From bf6bd99f7d89d351365f758f8ccb6c6dd0e84f75 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Mon, 25 May 2026 03:03:02 -0700 Subject: [PATCH 01/37] feat: add flashinfer nvfp4 quantization --- .../megatron_to_hf/processors/__init__.py | 11 +- .../processors/quantizer_nvfp4.py | 72 ++-- miles/utils/nvfp4.py | 199 +++++++++++ scripts/models/glm5-744B-A40B_5layer.sh | 12 + scripts/run_qwen3_30b_a3b.py | 136 +++++++- .../test_glm5_744b_a40b_5layer_nvfp4.py | 309 ++++++++++++++++++ tests/fast-gpu/test_nvfp4_quantizer.py | 203 +++++++++++- tools/convert_hf_to_nvfp4.py | 92 +++--- 8 files changed, 929 insertions(+), 105 deletions(-) create mode 100644 miles/utils/nvfp4.py create mode 100644 scripts/models/glm5-744B-A40B_5layer.sh create mode 100644 tests/e2e/megatron/test_glm5_744b_a40b_5layer_nvfp4.py diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py b/miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py index fba2c4e165a..cb1a7d613d0 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py @@ -2,12 +2,14 @@ from .quantizer_compressed_tensors import quantize_params_compressed_tensors from .quantizer_fp8 import quantize_params_fp8 from .quantizer_mxfp8 import quantize_params_mxfp8 +from .quantizer_nvfp4 import is_nvfp4_quantization_config, quantize_params_nvfp4 __all__ = [ "remove_padding", "quantize_param", "quantize_params_fp8", "quantize_params_mxfp8", + "quantize_params_nvfp4", "quantize_params_compressed_tensors", ] @@ -15,10 +17,13 @@ def quantize_params(args, megatron_name, converted_named_params, quantization_config): if quantization_config is None: return converted_named_params - elif quantization_config["quant_method"] == "fp8": + elif quantization_config.get("quant_method") == "fp8": return quantize_params_fp8(args, megatron_name, converted_named_params, quantization_config) - elif quantization_config["quant_method"] == "mxfp8": + elif quantization_config.get("quant_method") == "mxfp8": return quantize_params_mxfp8(args, megatron_name, converted_named_params, quantization_config) - elif quantization_config["quant_method"] == "compressed-tensors": + elif is_nvfp4_quantization_config(quantization_config): + return quantize_params_nvfp4(args, megatron_name, converted_named_params, quantization_config) + elif quantization_config.get("quant_method") == "compressed-tensors": # only int4 at the moment. return quantize_params_compressed_tensors(converted_named_params, quantization_config) + return converted_named_params diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py index 6432d5ebed4..5af1d08614f 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py @@ -2,9 +2,7 @@ import torch -FP4_E2M1_MAX = 6.0 -FP8_E4M3_MAX = 448.0 -NVFP4_GROUP_SIZE = 16 +from miles.utils.nvfp4 import NVFP4_GROUP_SIZE, nvfp4_4over6_enabled, nvfp4_global_decode_scale_te, nvfp4_quantize_1d GATED_PAIR_SUFFIXES = { ".gate_proj.weight": "gate", @@ -14,6 +12,23 @@ } +def is_nvfp4_quantization_config(quantization_config) -> bool: + if quantization_config is None: + return False + return quantization_config.get("quant_algo") == "NVFP4" or quantization_config.get("quant_method") == "nvfp4" + + +def fp4_param_gather_enabled(args) -> bool: + if args is None: + return False + return bool(getattr(args, "fp4_param", False) or getattr(args, "fp4_param_gather", False)) + + +def assert_no_fp4_param_gather(args) -> None: + if fp4_param_gather_enabled(args): + raise NotImplementedError("fp4-param-gather is unsupported for Miles NVFP4 checkpoint export.") + + def _get_ignore_rules(quantization_config) -> list[str]: ignore_rules = quantization_config.get("ignore", []) or [] if isinstance(ignore_rules, str): @@ -37,7 +52,8 @@ def _is_ignored(name: str, ignore_rules: list[str]) -> bool: def quantize_params_nvfp4(args, megatron_name, converted_named_params, quantization_config): assert quantization_config is not None - assert quantization_config.get("quant_algo") == "NVFP4" or quantization_config.get("quant_method") == "nvfp4" + assert is_nvfp4_quantization_config(quantization_config) + assert_no_fp4_param_gather(args) if getattr(args, "extra_high_precision_layers_megatron", False): for layer_name in getattr(args, "extra_high_precision_layers_megatron", ()): @@ -162,37 +178,16 @@ def _split_gated_pair_name(name: str): return None, None -def _nvfp4_global_decode_scale_te(global_amax: torch.Tensor) -> torch.Tensor: - fp4_max = torch.tensor(FP4_E2M1_MAX, device=global_amax.device, dtype=torch.float32) - fp8_max = torch.tensor(FP8_E4M3_MAX, device=global_amax.device, dtype=torch.float32) - global_encode_scale = torch.div(fp8_max * fp4_max, global_amax.to(torch.float32)) - global_encode_scale = torch.min( - global_encode_scale, - torch.tensor( - torch.finfo(torch.float32).max, - device=global_encode_scale.device, - dtype=torch.float32, - ), - ) - if global_encode_scale.numel() == 1: - if global_encode_scale == torch.tensor(0.0, device=global_amax.device, dtype=torch.float32): - global_encode_scale = torch.tensor(1.0, device=global_amax.device, dtype=torch.float32) - else: - global_encode_scale = torch.where( - global_encode_scale == 0.0, - torch.ones_like(global_encode_scale), - global_encode_scale, - ) - return torch.div(1.0, global_encode_scale) +_nvfp4_global_decode_scale_te = nvfp4_global_decode_scale_te def _quantize_nvfp4_1d( weight: torch.Tensor, global_amax: torch.Tensor | None = None, + use_4over6: bool | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ - NVFP4 1D quantization (tile shape = 1x16), adapted from - TransformerEngine NVFP4QuantizerRef._quantize_blockwise_reference. + NVFP4 1D quantization (tile shape = 1x16). Returns: qweight: uint8 packed fp4, shape (M, K // 2) @@ -200,7 +195,7 @@ def _quantize_nvfp4_1d( global_scale: float32 scalar tensor """ weight = weight.contiguous() - m, n = weight.shape + _, n = weight.shape if n % NVFP4_GROUP_SIZE != 0: raise ValueError(f"NVFP4 requires K divisible by {NVFP4_GROUP_SIZE}, got {n}.") @@ -209,25 +204,18 @@ def _quantize_nvfp4_1d( else: global_amax = global_amax.to(device=weight.device, dtype=torch.float32) - from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import NVFP4QuantizerRef - - qweight, block_scale = NVFP4QuantizerRef._quantize_blockwise_reference( - weight, - global_amax, - NVFP4_GROUP_SIZE, - 1, - pow_2_scales=False, - eps=0.0, - ) - return qweight, block_scale, _nvfp4_global_decode_scale_te(global_amax) + return nvfp4_quantize_1d(weight, global_amax, use_4over6=use_4over6) def quantize_nvfp4( weight: torch.Tensor, global_amax: torch.Tensor | None = None, + use_4over6: bool | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if use_4over6 is None: + use_4over6 = nvfp4_4over6_enabled() if weight.dim() == 2: - return _quantize_nvfp4_1d(weight, global_amax=global_amax) + return _quantize_nvfp4_1d(weight, global_amax=global_amax, use_4over6=use_4over6) if weight.dim() == 3: if global_amax is not None: raise ValueError("global_amax override is only supported for 2D weights.") @@ -235,7 +223,7 @@ def quantize_nvfp4( block_scales = [] global_scales = [] for idx in range(weight.shape[0]): - qweight, block_scale, global_scale = _quantize_nvfp4_1d(weight[idx]) + qweight, block_scale, global_scale = _quantize_nvfp4_1d(weight[idx], use_4over6=use_4over6) qweights.append(qweight) block_scales.append(block_scale) global_scales.append(global_scale) diff --git a/miles/utils/nvfp4.py b/miles/utils/nvfp4.py new file mode 100644 index 00000000000..002c368a3f4 --- /dev/null +++ b/miles/utils/nvfp4.py @@ -0,0 +1,199 @@ +import logging +import os +from contextlib import contextmanager + +import torch + +FP4_E2M1_MAX = 6.0 +FP8_E4M3_MAX = 448.0 +NVFP4_GROUP_SIZE = 16 + +logger = logging.getLogger(__name__) +FLASHINFER_NVFP4_ENV_KEYS = ( + "FLASHINFER_NVFP4_4OVER6", + "FLASHINFER_NVFP4_4OVER6_E4M3_USE_256", + "FLASHINFER_NVFP4_4OVER6_ERR_MODE", + "FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH", + "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH", +) + + +def str_to_bool(value) -> bool: + if isinstance(value, bool): + return value + if value is None: + return False + return str(value).strip().lower() in ("1", "true", "yes", "on") + + +def nvfp4_4over6_weight_scope_enabled(value) -> bool: + if isinstance(value, str): + value = value.strip().lower() + if value in ("weights", "all"): + return True + if value in ("none", "activations"): + return False + return str_to_bool(value) + + +def nvfp4_4over6_enabled() -> bool: + return nvfp4_4over6_weight_scope_enabled(os.getenv("NVTE_NVFP4_4OVER6")) + + +def nvfp4_4over6_weight_scope(use_4over6: bool | None = None) -> str: + if use_4over6 is None: + use_4over6 = nvfp4_4over6_enabled() + return "weights" if use_4over6 else "none" + + +def nvfp4_weight_e4m3_max(use_4over6: bool) -> int: + if use_4over6 and nvfp4_4over6_weight_scope_enabled(os.getenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all")): + return 256 + return int(FP8_E4M3_MAX) + + +def nvfp4_4over6_err_mode() -> str: + err_mode = os.getenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MAE").strip().upper() + if err_mode not in ("MAE", "MSE"): + raise ValueError("NVTE_NVFP4_4OVER6_ERR_MODE must be one of: 'MAE', 'MSE'.") + return err_mode + + +def nvfp4_global_encode_scale_te( + global_amax: torch.Tensor, + nvfp4_e4m3_max: int = int(FP8_E4M3_MAX), +) -> torch.Tensor: + fp4_max = torch.tensor(FP4_E2M1_MAX, device=global_amax.device, dtype=torch.float32) + fp8_max = torch.tensor(float(nvfp4_e4m3_max), device=global_amax.device, dtype=torch.float32) + global_encode_scale = torch.div(fp8_max * fp4_max, global_amax.to(torch.float32)) + global_encode_scale = torch.min( + global_encode_scale, + torch.tensor( + torch.finfo(torch.float32).max, + device=global_encode_scale.device, + dtype=torch.float32, + ), + ) + if global_encode_scale.numel() == 1: + if global_encode_scale == torch.tensor(0.0, device=global_amax.device, dtype=torch.float32): + global_encode_scale = torch.tensor(1.0, device=global_amax.device, dtype=torch.float32) + else: + global_encode_scale = torch.where( + global_encode_scale == 0.0, + torch.ones_like(global_encode_scale), + global_encode_scale, + ) + return global_encode_scale + + +def nvfp4_global_decode_scale_te( + global_amax: torch.Tensor, + nvfp4_e4m3_max: int = int(FP8_E4M3_MAX), +) -> torch.Tensor: + return torch.div(1.0, nvfp4_global_encode_scale_te(global_amax, nvfp4_e4m3_max)) + + +def sync_flashinfer_nvfp4_env_from_nvte(use_4over6: bool | None = None) -> dict[str, str]: + if use_4over6 is None: + use_4over6 = nvfp4_4over6_enabled() + + flashinfer_env = { + "FLASHINFER_NVFP4_4OVER6": "1" if use_4over6 else "0", + "FLASHINFER_NVFP4_4OVER6_E4M3_USE_256": ( + "1" + if use_4over6 and nvfp4_4over6_weight_scope_enabled(os.getenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all")) + else "0" + ), + "FLASHINFER_NVFP4_4OVER6_ERR_MODE": nvfp4_4over6_err_mode(), + "FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH": ( + "1" if str_to_bool(os.getenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", "0")) else "0" + ), + } + os.environ.update(flashinfer_env) + os.environ.setdefault("TRTLLM_DISABLE_FP4_QUANT_FAST_MATH", "1") + return {**flashinfer_env, "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH": os.environ["TRTLLM_DISABLE_FP4_QUANT_FAST_MATH"]} + + +@contextmanager +def flashinfer_nvfp4_env_from_nvte(use_4over6: bool | None = None): + original_env = {key: os.environ.get(key) for key in FLASHINFER_NVFP4_ENV_KEYS} + sync_flashinfer_nvfp4_env_from_nvte(use_4over6) + try: + yield + finally: + for key, value in original_env.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +try: + from flashinfer import nvfp4_quantize as _flashinfer_nvfp4_quantize + from flashinfer.tllm_enums import SfLayout +except ImportError: + _flashinfer_nvfp4_quantize = None + SfLayout = None + logger.warning("FlashInfer nvfp4_quantize not available; falling back to TransformerEngine reference.") + + +def _te_nvfp4_quantize_1d( + weight: torch.Tensor, + global_amax: torch.Tensor, + use_4over6: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + try: + from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef + except ImportError: + from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import NVFP4QuantizerRef + + nvfp4_e4m3_max = nvfp4_weight_e4m3_max(use_4over6) + try: + qweight, block_scale = NVFP4QuantizerRef._quantize_blockwise_reference( + weight, + global_amax, + NVFP4_GROUP_SIZE, + 1, + pow_2_scales=False, + nvfp4_use_4over6=use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode(), + eps=0.0, + ) + except TypeError: + if use_4over6: + raise + qweight, block_scale = NVFP4QuantizerRef._quantize_blockwise_reference( + weight, + global_amax, + NVFP4_GROUP_SIZE, + 1, + pow_2_scales=False, + eps=0.0, + ) + return qweight, block_scale, nvfp4_global_decode_scale_te(global_amax, nvfp4_e4m3_max) + + +def nvfp4_quantize_1d( + weight: torch.Tensor, + global_amax: torch.Tensor, + use_4over6: bool | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if use_4over6 is None: + use_4over6 = nvfp4_4over6_enabled() + + if _flashinfer_nvfp4_quantize is None or weight.dtype == torch.float32: + return _te_nvfp4_quantize_1d(weight, global_amax, use_4over6) + + nvfp4_e4m3_max = nvfp4_weight_e4m3_max(use_4over6) + global_encode_scale = nvfp4_global_encode_scale_te(global_amax, nvfp4_e4m3_max) + with flashinfer_nvfp4_env_from_nvte(use_4over6): + qweight, block_scale = _flashinfer_nvfp4_quantize( + weight, + global_encode_scale.reshape(1).contiguous(), + sfLayout=SfLayout.layout_linear, + do_shuffle=False, + sf_vec_size=NVFP4_GROUP_SIZE, + backend="cuda", + ) + return qweight, block_scale.view(torch.float8_e4m3fn), nvfp4_global_decode_scale_te(global_amax, nvfp4_e4m3_max) diff --git a/scripts/models/glm5-744B-A40B_5layer.sh b/scripts/models/glm5-744B-A40B_5layer.sh new file mode 100644 index 00000000000..e6e93cd0890 --- /dev/null +++ b/scripts/models/glm5-744B-A40B_5layer.sh @@ -0,0 +1,12 @@ +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]:-$0}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/glm5-744B-A40B.sh" + +# Override for 5-layer pruned model (first 5 layers: 3 dense + 2 MoE) +N_MOE_LAYERS=2 + +for ((i=0; i<${#MODEL_ARGS[@]}; i++)); do + case "${MODEL_ARGS[$i]}" in + --num-layers) MODEL_ARGS[$((i+1))]=$((N_DENSE_LAYERS + N_MOE_LAYERS)) ;; + --moe-layer-freq) MODEL_ARGS[$((i+1))]="[0]*${N_DENSE_LAYERS}+[1]*${N_MOE_LAYERS}" ;; + esac +done diff --git a/scripts/run_qwen3_30b_a3b.py b/scripts/run_qwen3_30b_a3b.py index cb9e225f7c2..d9c9609bcf7 100644 --- a/scripts/run_qwen3_30b_a3b.py +++ b/scripts/run_qwen3_30b_a3b.py @@ -1,3 +1,5 @@ +import os +import shlex from dataclasses import dataclass from typing import Literal @@ -5,6 +7,20 @@ import miles.utils.external_utils.command_utils as U +BLACKWELL_HARDWARE = ("B200", "B300", "GB200", "GB300") +FP4_ENV_MARKERS = ["NVTE", "FLASHINFER", "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH"] + + +def fp4_env_vars() -> dict[str, str]: + return {key: value for key, value in os.environ.items() if any(marker in key for marker in FP4_ENV_MARKERS)} + + +def env_prefix(env_vars: dict[str, str]) -> str: + if not env_vars: + return "" + assignments = (f"{key}={shlex.quote(value)}" for key, value in sorted(env_vars.items())) + return " ".join(assignments) + " " + @dataclass class ScriptArgs(U.ExecuteTrainConfig): @@ -13,6 +29,9 @@ class ScriptArgs(U.ExecuteTrainConfig): model_name: str = "Qwen3-30B-A3B" megatron_model_type: str = "qwen3-30B-A3B" num_gpus_per_node: int | None = None + actor_num_gpus_per_node: int | None = None + rollout_num_gpus: int | None = None + no_colocate: bool = False hardware: Literal["H100", "B200", "B300", "GB200", "GB300"] = "H100" enable_eval: bool = True extra_args: str = "" @@ -22,9 +41,11 @@ class ScriptArgs(U.ExecuteTrainConfig): rollout_fp8: bool = False rollout_mxfp8: bool = False rollout_int4: bool = False + rollout_nvfp4: bool = False rollout_attn_fp8: bool = False train_fp8: bool = False train_mxfp8: bool = False + train_nvfp4: bool = False enable_megatron_bridge: bool = False enable_mis: bool = False # TODO improve, should be able to override more easily @@ -32,16 +53,38 @@ class ScriptArgs(U.ExecuteTrainConfig): def __post_init__(self): self.num_gpus_per_node = self.num_gpus_per_node or U.NUM_GPUS_OF_HARDWARE[self.hardware] + self.no_colocate = self.no_colocate or self.rollout_nvfp4 + if self.no_colocate: + self.actor_num_gpus_per_node = self.actor_num_gpus_per_node or self.num_gpus_per_node // 2 + self.rollout_num_gpus = self.rollout_num_gpus or self.num_gpus_per_node - self.actor_num_gpus_per_node + assert self.actor_num_gpus_per_node > 0, "actor_num_gpus_per_node must be positive" + assert self.rollout_num_gpus > 0, "rollout_num_gpus must be positive" + assert ( + self.actor_num_gpus_per_node + self.rollout_num_gpus <= self.num_gpus_per_node + ), "actor and rollout GPU allocations cannot exceed num_gpus_per_node" + else: + self.actor_num_gpus_per_node = self.actor_num_gpus_per_node or self.num_gpus_per_node + self.rollout_num_gpus = self.rollout_num_gpus or self.num_gpus_per_node if self.rollout_int4: assert not self.rollout_fp8, "rollout_int4 and rollout_fp8 cannot be enabled at the same time" assert not self.rollout_mxfp8, "rollout_int4 and rollout_mxfp8 cannot be enabled at the same time" + assert not self.rollout_nvfp4, "rollout_int4 and rollout_nvfp4 cannot be enabled at the same time" if self.rollout_mxfp8: assert not self.rollout_fp8, "rollout_mxfp8 and rollout_fp8 cannot be enabled at the same time" - assert self.hardware in ("B200", "B300", "GB200", "GB300"), "rollout_mxfp8 only supports Blackwell GPUs" + assert not self.rollout_nvfp4, "rollout_mxfp8 and rollout_nvfp4 cannot be enabled at the same time" + assert self.hardware in BLACKWELL_HARDWARE, "rollout_mxfp8 only supports Blackwell GPUs" + if self.rollout_nvfp4: + assert not self.rollout_fp8, "rollout_nvfp4 and rollout_fp8 cannot be enabled at the same time" + assert self.hardware in BLACKWELL_HARDWARE, "rollout_nvfp4 only supports Blackwell GPUs" if self.train_mxfp8: assert not self.train_fp8, "train_mxfp8 and train_fp8 cannot be enabled at the same time" - assert self.hardware in ("B200", "B300", "GB200", "GB300"), "train_mxfp8 only supports Blackwell GPUs" + assert not self.train_nvfp4, "train_mxfp8 and train_nvfp4 cannot be enabled at the same time" + assert self.hardware in BLACKWELL_HARDWARE, "train_mxfp8 only supports Blackwell GPUs" assert self.rollout_mxfp8, "train_mxfp8 requires rollout_mxfp8 to be enabled" + if self.train_nvfp4: + assert not self.train_fp8, "train_nvfp4 and train_fp8 cannot be enabled at the same time" + assert self.hardware in BLACKWELL_HARDWARE, "train_nvfp4 only supports Blackwell GPUs" + assert self.rollout_nvfp4, "train_nvfp4 requires rollout_nvfp4 to be enabled" def prepare(args: ScriptArgs): @@ -60,6 +103,21 @@ def prepare(args: ScriptArgs): f"{args.extra_args} " ) + if args.rollout_nvfp4 or args.train_nvfp4: + nvfp4_env_prefix = env_prefix( + { + "NVTE_USE_FAST_MATH": "0", + "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH": "1", + **fp4_env_vars(), + } + ) + U.exec_command( + f"{nvfp4_env_prefix}" + f"python tools/convert_hf_to_nvfp4.py --model-dir {args.model_dir}/{args.model_name} " + f"--save-dir {args.model_dir}/{args.model_name}-NVFP4 " + f"{args.extra_args} " + ) + if args.rollout_int4: U.exec_command( f"python tools/convert_hf_to_int4_direct.py --model-dir {args.model_dir}/{args.model_name} --save-dir {args.model_dir}/{args.model_name}-INT4" @@ -90,6 +148,8 @@ def execute(args: ScriptArgs): hf_checkpoint = f"{args.model_dir}/{args.model_name}-FP8" elif args.train_mxfp8: hf_checkpoint = f"{args.model_dir}/{args.model_name}-MXFP8" + elif args.rollout_nvfp4 or args.train_nvfp4: + hf_checkpoint = f"{args.model_dir}/{args.model_name}-NVFP4" elif args.rollout_int4: hf_checkpoint = f"{args.model_dir}/{args.model_name}-INT4" else: @@ -167,12 +227,16 @@ def execute(args: ScriptArgs): # need to comment this when using model with MLA "--attention-backend flash " f"--actor-num-nodes {args.num_nodes} " - f"--actor-num-gpus-per-node {args.num_gpus_per_node} " f"--num-gpus-per-node {args.num_gpus_per_node} " - "--colocate " "--use-fault-tolerance " f"--dump-details {args.output_dir}/{args.run_id}/dump_details " ) + if args.no_colocate: + misc_args += ( + f"--actor-num-gpus-per-node {args.actor_num_gpus_per_node} " f"--rollout-num-gpus {args.rollout_num_gpus} " + ) + else: + misc_args += f"--actor-num-gpus-per-node {args.num_gpus_per_node} " "--colocate " misc_env_vars = {} if args.rollout_int4: @@ -205,6 +269,48 @@ def execute(args: ScriptArgs): misc_env_vars |= { "NVTE_FP8_BLOCK_SCALING_FP32_SCALES": "1", } + elif args.train_nvfp4: + match args.hardware: + case "B200" | "B300" | "GB200" | "GB300": + misc_args += ( + "--transformer-impl transformer_engine " "--bf16 " "--fp4-format e2m1 " "--fp4-recipe nvfp4 " + ) + misc_env_vars |= { + "NVTE_NVFP4_DISABLE_2D_QUANTIZATION": "1", + "NVTE_NVFP4_DISABLE_RHT": "1", + "NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING": "1", + "NVTE_NVFP4_ROW_SCALED_ACTIVATION": "1", + "NVTE_BACKWARD_OVERRIDE": "dequantized", + "NVTE_USE_FAST_MATH": "0", + } + optimizer_args += "--optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer " + te_precision_config_text = """ +configs: + nvfp4: + transformer_engine_config_type: "TEQuantizationParams" + training_recipe: + fp4_quantization_recipe: "nvfp4" + bf16: + transformer_engine_config_type: "TEQuantizationParams" + training_recipe: {} +matchers: + routed_experts_fc1_nvfp4: + type: "glob" + enabled: true + pattern: "*.mlp.experts.linear_fc1" + config: "nvfp4" + routed_experts_fc2_nvfp4: + type: "glob" + enabled: true + pattern: "*.mlp.experts.linear_fc2" + config: "nvfp4" + default_bf16: + type: "glob" + enabled: true + pattern: "*" + config: "bf16" +""".strip() + misc_args += f"--te-precision-config-file {U.save_to_temp_file(te_precision_config_text, 'yaml')} " if args.enable_megatron_bridge: misc_args += "--megatron-to-hf-mode bridge " @@ -233,9 +339,12 @@ def execute(args: ScriptArgs): "--sequence-parallel " "--pipeline-model-parallel-size 1 " "--context-parallel-size 1 " - f"--expert-model-parallel-size {args.num_gpus_per_node if args.train_mxfp8 else 4} " "--expert-tensor-parallel-size 1 " ) + if args.no_colocate: + perf_args += f"--expert-model-parallel-size {args.actor_num_gpus_per_node} " + else: + perf_args += f"--expert-model-parallel-size {args.num_gpus_per_node if args.train_mxfp8 else 4} " sglang_args = "--sglang-mem-fraction-static 0.7 " "--sglang-attention-backend trtllm_mha " if args.rollout_fp8: sglang_world_size = 2 @@ -266,6 +375,21 @@ def execute(args: ScriptArgs): f"--sglang-chunked-prefill-size {sglang_world_size * sglang_decode_max_bs} " f"--sglang-cuda-graph-max-bs {sglang_decode_max_bs} " ) + elif args.rollout_nvfp4: + sglang_world_size = 2 + sglang_decode_max_bs = 256 + sglang_args += ( + f"--rollout-num-gpus-per-engine {sglang_world_size} " + "--sglang-moe-runner-backend flashinfer_trtllm_routed " + f"--sglang-tp-size {sglang_world_size} " + f"--sglang-ep-size {sglang_world_size} " + f"--sglang-cuda-graph-max-bs {sglang_decode_max_bs} " + "--sglang-kv-cache-dtype bf16 " + ) + misc_env_vars |= { + "SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION": "1", + "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH": "1", + } else: sglang_args += "--rollout-num-gpus-per-engine 4 " "--sglang-cuda-graph-max-bs 512 " case _: @@ -306,6 +430,8 @@ def execute(args: ScriptArgs): f"{args.extra_args} " ) + misc_env_vars |= fp4_env_vars() + U.execute_train( train_args=train_args, num_gpus_per_node=args.num_gpus_per_node, diff --git a/tests/e2e/megatron/test_glm5_744b_a40b_5layer_nvfp4.py b/tests/e2e/megatron/test_glm5_744b_a40b_5layer_nvfp4.py new file mode 100644 index 00000000000..8a10fa20b32 --- /dev/null +++ b/tests/e2e/megatron/test_glm5_744b_a40b_5layer_nvfp4.py @@ -0,0 +1,309 @@ +import json +import os +from pathlib import Path + +from tests.ci.ci_register import register_cuda_ci + +import miles.utils.external_utils.command_utils as U + +register_cuda_ci( + est_time=3600, + suite="stage-c-8-gpu-h100", + labels=["model-scripts"], + disabled="Requires Blackwell/B200 CI runner for NVFP4.", +) + +MODEL_ORG = "Pinaster" +MODEL_NAME = "GLM-5_5layer" +MODEL_TYPE = "glm5-744B-A40B_5layer" +NUM_GPUS = 8 +ACTOR_NUM_GPUS = 4 +ROLLOUT_NUM_GPUS = 4 +ROLLOUT_GPUS_PER_ENGINE = 2 +NUM_LAYERS_AT_START_IN_BF16 = 1 +NUM_LAYERS_AT_END_IN_BF16 = 1 +RUN_ID = U.create_run_id() + +MODEL_DIR = "/root/models" +DATA_DIR = "/root/datasets" +MEGATRON_PATH = "/root/TransformerEngine:/root/Megatron-LM" + +EXTRA_HIGH_PRECISION_LAYERS_HF = (".shared_experts.",) +EXTRA_HIGH_PRECISION_LAYERS_MEGATRON = ( + ".shared_experts.linear_fc1", + ".shared_experts.linear_fc2", +) + +NVFP4_ENV = { + "NVTE_NVFP4_DISABLE_2D_QUANTIZATION": "1", + "NVTE_NVFP4_DISABLE_RHT": "1", + "NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING": "1", + "NVTE_NVFP4_ROW_SCALED_ACTIVATION": "1", + "NVTE_BACKWARD_OVERRIDE": "dequantized", + "NVTE_USE_FAST_MATH": "0", + "NVTE_NVFP4_4OVER6": "all", + "FLASHINFER_NVFP4_4OVER6": "1", + "NVTE_NVFP4_4OVER6_E4M3_USE_256": "all", + "FLASHINFER_NVFP4_4OVER6_E4M3_USE_256": "1", + "NVTE_NVFP4_4OVER6_ERR_MODE": "MSE", + "FLASHINFER_NVFP4_4OVER6_ERR_MODE": "MSE", + "NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH": "0", + "FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH": "0", + "SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION": "1", + "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH": "1", +} + +GLM5_ENV = { + "SGLANG_NSA_FORCE_MLA": "1", + "INDEXER_ROPE_NEOX_STYLE": "0", + "NVSHMEM_DISABLE_NCCL": "1", +} + +TE_PRECISION_CONFIG = """ +configs: + nvfp4: + transformer_engine_config_type: "TEQuantizationParams" + training_recipe: + fp4_quantization_recipe: "nvfp4" + bf16: + transformer_engine_config_type: "TEQuantizationParams" + training_recipe: {} +matchers: + routed_experts_fc1_nvfp4: + type: "glob" + enabled: true + pattern: "*.mlp.experts.linear_fc1" + config: "nvfp4" + routed_experts_fc2_nvfp4: + type: "glob" + enabled: true + pattern: "*.mlp.experts.linear_fc2" + config: "nvfp4" + shared_experts_fc1_bf16: + type: "glob" + enabled: true + pattern: "*.mlp.shared_experts.linear_fc1" + config: "bf16" + shared_experts_fc2_bf16: + type: "glob" + enabled: true + pattern: "*.mlp.shared_experts.linear_fc2" + config: "bf16" + default_bf16: + type: "glob" + enabled: true + pattern: "*" + config: "bf16" +""".strip() + + +def _extra_high_precision_layers_hf_args() -> str: + return "--extra-high-precision-layers-hf " + " ".join(EXTRA_HIGH_PRECISION_LAYERS_HF) + " " + + +def _extra_high_precision_layers_megatron_args() -> str: + return "--extra-high-precision-layers-megatron " + " ".join(EXTRA_HIGH_PRECISION_LAYERS_MEGATRON) + " " + + +def _validate_glm_checkpoint(): + config_path = Path(MODEL_DIR) / MODEL_NAME / "config.json" + if not config_path.exists(): + raise FileNotFoundError(f"{config_path} not found") + + with open(config_path) as f: + config = json.load(f) + + if ( + config.get("model_type") != "glm_moe_dsa" + or config.get("architectures") != ["GlmMoeDsaForCausalLM"] + or config.get("num_hidden_layers") != 5 + ): + raise RuntimeError( + f"{config_path} must use native GLM-5 5-layer config with " + f"model_type=glm_moe_dsa, architectures=[GlmMoeDsaForCausalLM], " + "and num_hidden_layers=5" + ) + if "auto_map" in config: + raise RuntimeError(f"{config_path} must not contain auto_map. Try updating the checkpoint.") + + +def prepare(): + os.environ.update(NVFP4_ENV) + U.exec_command(f"mkdir -p {MODEL_DIR} {DATA_DIR}") + U.exec_command(f"hf download {MODEL_ORG}/{MODEL_NAME} --local-dir {MODEL_DIR}/{MODEL_NAME}") + U.hf_download_dataset("zhuzilin/dapo-math-17k", data_dir=DATA_DIR) + + _validate_glm_checkpoint() + U.exec_command(f"rm -rf {MODEL_DIR}/{MODEL_NAME}-NVFP4 {MODEL_DIR}/{MODEL_NAME}_torch_dist") + + U.exec_command( + f"python tools/convert_hf_to_nvfp4.py " + f"--model-dir {MODEL_DIR}/{MODEL_NAME} " + f"--save-dir {MODEL_DIR}/{MODEL_NAME}-NVFP4 " + f"--num-layers-at-start-in-bf16 {NUM_LAYERS_AT_START_IN_BF16} " + f"--num-layers-at-end-in-bf16 {NUM_LAYERS_AT_END_IN_BF16} " + f"{_extra_high_precision_layers_hf_args()}" + ) + + U.convert_checkpoint( + model_name=MODEL_NAME, + megatron_model_type=MODEL_TYPE, + num_gpus_per_node=ACTOR_NUM_GPUS, + extra_args=( + "--tensor-model-parallel-size 1 " + "--expert-tensor-parallel-size 1 " + "--pipeline-model-parallel-size 1 " + "--expert-model-parallel-size 1 " + ), + dir_dst=MODEL_DIR, + hf_checkpoint=f"{MODEL_DIR}/{MODEL_NAME}", + megatron_path=MEGATRON_PATH, + ) + + +def execute(): + os.environ.update(NVFP4_ENV) + os.environ.update(GLM5_ENV) + os.environ.setdefault("RAY_TMPDIR", "/tmp/ray") + te_precision_config_path = U.save_to_temp_file(TE_PRECISION_CONFIG, "yaml") + load_save_path = f"/root/shared_data/{RUN_ID}/checkpoints" + + ckpt_args = ( + f"--hf-checkpoint {MODEL_DIR}/{MODEL_NAME}-NVFP4/ " + f"--ref-load {MODEL_DIR}/{MODEL_NAME}_torch_dist " + f"--load {load_save_path} " + f"--save {load_save_path} " + "--save-interval 2 " + "--save-retain-interval 2 " + ) + + rollout_args = ( + f"--prompt-data {DATA_DIR}/dapo-math-17k/dapo-math-17k.jsonl " + "--input-key prompt " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type deepscaler " + "--num-rollout 2 " + "--rollout-batch-size 32 " + "--n-samples-per-prompt 8 " + "--rollout-max-response-len 100 " + "--rollout-temperature 1 " + "--global-batch-size 256 " + "--balance-data " + ) + + perf_args = ( + f"--tensor-model-parallel-size {ACTOR_NUM_GPUS} " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " + f"--expert-model-parallel-size {ACTOR_NUM_GPUS} " + "--expert-tensor-parallel-size 1 " + "--recompute-granularity full " + "--recompute-method uniform " + "--recompute-num-layers 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 32768 " + "--data-pad-size-multiplier 4096 " + "--log-probs-chunk-size 1024 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--use-kl-loss " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + "--optimizer-cpu-offload " + "--overlap-cpu-optimizer-d2h-h2d " + "--use-precision-aware-optimizer " + ) + + sglang_args = ( + "--sglang-mem-fraction-static 0.7 " + "--sglang-attention-backend nsa " + "--sglang-nsa-decode-backend flashmla_sparse " + "--sglang-nsa-prefill-backend flashmla_sparse " + "--sglang-kv-cache-dtype bf16 " + "--sglang-page-size 64 " + f"--rollout-num-gpus-per-engine {ROLLOUT_GPUS_PER_ENGINE} " + "--sglang-moe-runner-backend flashinfer_trtllm_routed " + f"--sglang-tp-size {ROLLOUT_GPUS_PER_ENGINE} " + f"--sglang-ep-size {ROLLOUT_GPUS_PER_ENGINE} " + "--sglang-cuda-graph-max-bs 256 " + "--sglang-watchdog-timeout 3600 " + ) + + ci_args = "--ci-test --ci-disable-logprobs-checker " + + mixed_precision_args = ( + "--transformer-impl transformer_engine " + "--bf16 " + "--fp4-format e2m1 " + "--fp4-recipe nvfp4 " + "--first-last-layers-bf16 " + f"--num-layers-at-start-in-bf16 {NUM_LAYERS_AT_START_IN_BF16} " + f"--num-layers-at-end-in-bf16 {NUM_LAYERS_AT_END_IN_BF16} " + f"{_extra_high_precision_layers_hf_args()}" + f"{_extra_high_precision_layers_megatron_args()}" + f"--te-precision-config-file {te_precision_config_path} " + ) + + misc_args = ( + "--use-rollout-routing-replay " + "--use-miles-router " + "--sglang-disable-shared-experts-fusion " + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--allgather-cp " + f"--update-weight-buffer-size {2 * 1024 ** 3} " + "--actor-num-nodes 1 " + f"--actor-num-gpus-per-node {ACTOR_NUM_GPUS} " + f"--num-gpus-per-node {NUM_GPUS} " + f"--rollout-num-gpus {ROLLOUT_NUM_GPUS} " + "--use-fault-tolerance " + f"--dump-details /root/shared_data/{RUN_ID}/dump_details " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__, run_id=RUN_ID)} " + f"{perf_args} " + f"{sglang_args} " + f"{ci_args} " + f"{mixed_precision_args} " + f"{misc_args} " + ) + + U.execute_train( + train_args=train_args, + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + megatron_path=MEGATRON_PATH, + extra_env_vars={**NVFP4_ENV, **GLM5_ENV}, + ) + + +if __name__ == "__main__": + prepare() + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute() diff --git a/tests/fast-gpu/test_nvfp4_quantizer.py b/tests/fast-gpu/test_nvfp4_quantizer.py index 0855591a819..7b49b6baef7 100644 --- a/tests/fast-gpu/test_nvfp4_quantizer.py +++ b/tests/fast-gpu/test_nvfp4_quantizer.py @@ -3,20 +3,31 @@ register_cuda_ci(est_time=60, suite="stage-b-2-gpu-h200", labels=[]) +import json +import os + import pytest import torch +from tools.convert_hf_to_nvfp4 import _update_quantization_config as tool_update_quantization_config +from tools.convert_hf_to_nvfp4 import _write_hf_quant_config as tool_write_hf_quant_config from tools.convert_hf_to_nvfp4 import quantize_nvfp4 as tool_quantize_nvfp4 from tools.convert_hf_to_nvfp4 import should_quantize as tool_should_quantize_nvfp4 -from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef -from miles.backends.megatron_utils.megatron_to_hf.processors.quantizer_nvfp4 import ( - NVFP4_GROUP_SIZE, - _nvfp4_global_decode_scale_te, -) +from miles.backends.megatron_utils.megatron_to_hf.processors import quantize_params from miles.backends.megatron_utils.megatron_to_hf.processors.quantizer_nvfp4 import ( quantize_nvfp4 as processor_quantize_nvfp4, ) from miles.backends.megatron_utils.megatron_to_hf.processors.quantizer_nvfp4 import quantize_params_nvfp4 +from miles.utils.nvfp4 import ( + NVFP4_GROUP_SIZE, + flashinfer_nvfp4_env_from_nvte, + nvfp4_4over6_err_mode, + nvfp4_4over6_weight_scope, + nvfp4_global_decode_scale_te, + nvfp4_weight_e4m3_max, + sync_flashinfer_nvfp4_env_from_nvte, +) NVFP4_SHAPES = [ (1, 64), @@ -33,6 +44,23 @@ (2048, 7168), (128, 16384), ] +NVFP4_ENV_KEYS = ( + "NVTE_NVFP4_4OVER6", + "NVTE_NVFP4_4OVER6_E4M3_USE_256", + "NVTE_NVFP4_4OVER6_ERR_MODE", + "NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", + "FLASHINFER_NVFP4_4OVER6", + "FLASHINFER_NVFP4_4OVER6_E4M3_USE_256", + "FLASHINFER_NVFP4_4OVER6_ERR_MODE", + "FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH", + "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH", +) + + +@pytest.fixture(autouse=True) +def clean_nvfp4_env(monkeypatch): + for key in NVFP4_ENV_KEYS: + monkeypatch.delenv(key, raising=False) def _make_weight(init_data: str, dtype: torch.dtype, shape: tuple[int, int], device: str) -> torch.Tensor: @@ -54,18 +82,47 @@ def _make_weight(init_data: str, dtype: torch.dtype, shape: tuple[int, int], dev raise ValueError(f"Unknown init_data: {init_data}") -def _te_nvfp4_reference(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +def _te_nvfp4_reference( + weight: torch.Tensor, + use_4over6: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: weight = weight.contiguous() global_amax = torch.max(torch.abs(weight.to(torch.float32))) + return _te_nvfp4_reference_with_global_amax(weight, global_amax, use_4over6=use_4over6) + + +def _te_nvfp4_reference_with_global_amax( + weight: torch.Tensor, + global_amax: torch.Tensor, + use_4over6: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + weight = weight.contiguous() + nvfp4_e4m3_max = nvfp4_weight_e4m3_max(use_4over6) qweight, block_scale = NVFP4QuantizerRef._quantize_blockwise_reference( weight, global_amax, NVFP4_GROUP_SIZE, 1, pow_2_scales=False, + nvfp4_use_4over6=use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode(), eps=0.0, ) - return qweight, block_scale, _nvfp4_global_decode_scale_te(global_amax) + return qweight, block_scale, nvfp4_global_decode_scale_te(global_amax, nvfp4_e4m3_max) + + +def test_nvfp4_dispatch_accepts_quant_algo_without_quant_method(): + converted_named_params = [("model.embed_tokens.weight", torch.zeros((1, 1), dtype=torch.bfloat16))] + + out = quantize_params( + args=None, + megatron_name="embedding.word_embeddings.weight", + converted_named_params=converted_named_params, + quantization_config={"quant_algo": "NVFP4"}, + ) + + assert out is converted_named_params def test_nvfp4_quantize_params_requires_complete_gated_pair(): @@ -99,6 +156,22 @@ def test_nvfp4_quantize_params_respects_extra_high_precision_layers_megatron(): assert out is converted_named_params +def test_nvfp4_quantize_params_rejects_fp4_param_gather(): + weight = torch.randn((4, NVFP4_GROUP_SIZE), dtype=torch.bfloat16) + args = type("Args", (), {"fp4_param": True})() + + with pytest.raises(NotImplementedError, match="fp4-param-gather is unsupported"): + quantize_params_nvfp4( + args=args, + megatron_name="decoder.layers.0.mlp.experts.linear_fc1.weight0", + converted_named_params=[ + ("model.layers.0.mlp.experts.0.gate_proj.weight", weight), + ("model.layers.0.mlp.experts.0.up_proj.weight", weight), + ], + quantization_config={"quant_method": "nvfp4"}, + ) + + @pytest.mark.parametrize("layer_idx", [0, 3]) def test_nvfp4_quantize_params_respects_first_last_layers_bf16(layer_idx): weight = torch.randn((4, NVFP4_GROUP_SIZE), dtype=torch.bfloat16) @@ -142,6 +215,57 @@ def test_nvfp4_hf_should_quantize_respects_extra_high_precision_layers_hf(): ) +def test_nvfp4_converter_records_4over6_mode(tmp_path, monkeypatch): + monkeypatch.setenv("NVTE_NVFP4_4OVER6", "weights") + cfg = {} + tool_update_quantization_config(cfg, ignore_list=["model.layers.0"]) + assert cfg["quantization_config"]["nvfp4_4over6"] == "weights" + + input_dir = tmp_path / "input" + output_dir = tmp_path / "output" + input_dir.mkdir() + output_dir.mkdir() + tool_write_hf_quant_config(str(output_dir), ignore_list=["model.layers.0"], input_path=str(input_dir)) + + hf_quant_config = json.loads((output_dir / "hf_quant_config.json").read_text()) + assert hf_quant_config["quantization"]["nvfp4_4over6"] == "weights" + + +def test_nvfp4_flashinfer_env_syncs_from_nvte(monkeypatch): + monkeypatch.setenv("NVTE_NVFP4_4OVER6", "all") + monkeypatch.setenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all") + monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MSE") + monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", "0") + monkeypatch.delenv("TRTLLM_DISABLE_FP4_QUANT_FAST_MATH", raising=False) + + synced = sync_flashinfer_nvfp4_env_from_nvte() + + assert synced["FLASHINFER_NVFP4_4OVER6"] == "1" + assert synced["FLASHINFER_NVFP4_4OVER6_E4M3_USE_256"] == "1" + assert synced["FLASHINFER_NVFP4_4OVER6_ERR_MODE"] == "MSE" + assert synced["FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH"] == "0" + assert synced["TRTLLM_DISABLE_FP4_QUANT_FAST_MATH"] == "1" + + +def test_nvfp4_flashinfer_env_context_restores_previous_values(monkeypatch): + monkeypatch.setenv("NVTE_NVFP4_4OVER6", "all") + monkeypatch.setenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all") + monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MSE") + monkeypatch.setenv("FLASHINFER_NVFP4_4OVER6", "old") + monkeypatch.setenv("TRTLLM_DISABLE_FP4_QUANT_FAST_MATH", "0") + + with flashinfer_nvfp4_env_from_nvte(): + assert os.environ["FLASHINFER_NVFP4_4OVER6"] == "1" + assert os.environ["FLASHINFER_NVFP4_4OVER6_E4M3_USE_256"] == "1" + assert os.environ["FLASHINFER_NVFP4_4OVER6_ERR_MODE"] == "MSE" + assert os.environ["TRTLLM_DISABLE_FP4_QUANT_FAST_MATH"] == "0" + + assert os.environ["FLASHINFER_NVFP4_4OVER6"] == "old" + assert "FLASHINFER_NVFP4_4OVER6_E4M3_USE_256" not in os.environ + assert "FLASHINFER_NVFP4_4OVER6_ERR_MODE" not in os.environ + assert os.environ["TRTLLM_DISABLE_FP4_QUANT_FAST_MATH"] == "0" + + @pytest.mark.parametrize( "quantize_fn", [processor_quantize_nvfp4, tool_quantize_nvfp4], @@ -163,6 +287,71 @@ def test_nvfp4_quantize_matches_te_reference_bitwise(quantize_fn, shape, dtype, torch.testing.assert_close(global_scale, global_scale_ref, rtol=0, atol=0) +@pytest.mark.parametrize("use_4over6", [False, True], ids=lambda value: nvfp4_4over6_weight_scope(value)) +@pytest.mark.parametrize("err_mode", ["MAE", "MSE"]) +def test_nvfp4_quantize_matches_te_reference_with_4over6_modes(monkeypatch, use_4over6, err_mode): + device = "cuda" + torch.manual_seed(42) + monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_MODE", err_mode) + monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", "0") + if use_4over6: + monkeypatch.setenv("NVTE_NVFP4_4OVER6", "weights") + monkeypatch.setenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all") + else: + monkeypatch.delenv("NVTE_NVFP4_4OVER6", raising=False) + + weight = _make_weight("random", torch.bfloat16, (128, 1024), device) + qweight, block_scale, global_scale = processor_quantize_nvfp4(weight) + qweight_ref, block_scale_ref, global_scale_ref = _te_nvfp4_reference(weight, use_4over6=use_4over6) + + torch.testing.assert_close(qweight, qweight_ref, rtol=0, atol=0) + torch.testing.assert_close(block_scale.view(torch.uint8), block_scale_ref.view(torch.uint8), rtol=0, atol=0) + torch.testing.assert_close(global_scale, global_scale_ref, rtol=0, atol=0) + + +def test_nvfp4_quantize_params_reads_4over6_from_env(monkeypatch): + device = "cuda" + torch.manual_seed(42) + monkeypatch.setenv("NVTE_NVFP4_4OVER6", "weights") + monkeypatch.setenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all") + + gate = _make_weight("random", torch.bfloat16, (4, NVFP4_GROUP_SIZE), device) + up = _make_weight("random", torch.bfloat16, (4, NVFP4_GROUP_SIZE), device) + shared_amax = torch.max(gate.abs().max().to(torch.float32), up.abs().max().to(torch.float32)) + converted_named_params = [ + ("model.layers.0.mlp.experts.0.gate_proj.weight", gate), + ("model.layers.0.mlp.experts.0.up_proj.weight", up), + ] + + out = dict( + quantize_params_nvfp4( + args=None, + megatron_name="decoder.layers.0.mlp.experts.linear_fc1.weight0", + converted_named_params=converted_named_params, + quantization_config={"quant_method": "nvfp4"}, + ) + ) + qweight_ref, block_scale_ref, global_scale_ref = _te_nvfp4_reference_with_global_amax( + gate, + shared_amax, + use_4over6=True, + ) + + torch.testing.assert_close(out["model.layers.0.mlp.experts.0.gate_proj.weight"], qweight_ref, rtol=0, atol=0) + torch.testing.assert_close( + out["model.layers.0.mlp.experts.0.gate_proj.weight_scale"].view(torch.uint8), + block_scale_ref.view(torch.uint8), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + out["model.layers.0.mlp.experts.0.gate_proj.weight_scale_2"], + global_scale_ref, + rtol=0, + atol=0, + ) + + if __name__ == "__main__": import sys diff --git a/tools/convert_hf_to_nvfp4.py b/tools/convert_hf_to_nvfp4.py index f8d1548ef60..1b72b6442cb 100644 --- a/tools/convert_hf_to_nvfp4.py +++ b/tools/convert_hf_to_nvfp4.py @@ -10,14 +10,16 @@ Use --extra-high-precision-layers-hf to keep additional HF weight-name substrings unquantized. -This follows the NVFP4 reference quantization in Transformer Engine and uses -1D block scaling (NVTE_NVFP4_1D_SCALING, group size = 16). +This follows the NVFP4 reference quantization in Transformer Engine, uses the +FlashInfer CUDA quantizer when available, and keeps 4over6 controlled by the +NVTE_NVFP4_* environment variables. """ import argparse import gc import json import os +import re import shutil import safetensors @@ -25,9 +27,14 @@ import torch from tqdm import tqdm -FP4_E2M1_MAX = 6.0 -FP8_E4M3_MAX = 448.0 -NVFP4_GROUP_SIZE = 16 +from miles.utils.nvfp4 import ( + NVFP4_GROUP_SIZE, + nvfp4_4over6_enabled, + nvfp4_4over6_weight_scope, + nvfp4_global_decode_scale_te, + nvfp4_quantize_1d, +) + DEFAULT_KV_CACHE_SCHEME = {"dynamic": False, "num_bits": 8, "type": "float"} DEFAULT_KV_CACHE_QUANT_ALGO = "FP8" @@ -100,37 +107,16 @@ def should_quantize( return True -def _nvfp4_global_decode_scale_te(global_amax: torch.Tensor) -> torch.Tensor: - fp4_max = torch.tensor(FP4_E2M1_MAX, device=global_amax.device, dtype=torch.float32) - fp8_max = torch.tensor(FP8_E4M3_MAX, device=global_amax.device, dtype=torch.float32) - global_encode_scale = torch.div(fp8_max * fp4_max, global_amax.to(torch.float32)) - global_encode_scale = torch.min( - global_encode_scale, - torch.tensor( - torch.finfo(torch.float32).max, - device=global_encode_scale.device, - dtype=torch.float32, - ), - ) - if global_encode_scale.numel() == 1: - if global_encode_scale == torch.tensor(0.0, device=global_amax.device, dtype=torch.float32): - global_encode_scale = torch.tensor(1.0, device=global_amax.device, dtype=torch.float32) - else: - global_encode_scale = torch.where( - global_encode_scale == 0.0, - torch.ones_like(global_encode_scale), - global_encode_scale, - ) - return torch.div(1.0, global_encode_scale) +_nvfp4_global_decode_scale_te = nvfp4_global_decode_scale_te def _quantize_nvfp4_1d( weight: torch.Tensor, global_amax: torch.Tensor | None = None, + use_4over6: bool | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ - NVFP4 1D quantization (tile shape = 1x16), adapted from - TransformerEngine NVFP4QuantizerRef._quantize_blockwise_reference. + NVFP4 1D quantization (tile shape = 1x16). Returns: qweight: uint8 packed fp4, shape (M, K // 2) @@ -138,7 +124,7 @@ def _quantize_nvfp4_1d( global_scale: float32 scalar tensor """ weight = weight.contiguous() - m, n = weight.shape + _, n = weight.shape if n % NVFP4_GROUP_SIZE != 0: raise ValueError(f"NVFP4 requires K divisible by {NVFP4_GROUP_SIZE}, got {n}.") @@ -147,25 +133,18 @@ def _quantize_nvfp4_1d( else: global_amax = global_amax.to(device=weight.device, dtype=torch.float32) - from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import NVFP4QuantizerRef - - qweight, block_scale = NVFP4QuantizerRef._quantize_blockwise_reference( - weight, - global_amax, - NVFP4_GROUP_SIZE, - 1, - pow_2_scales=False, - eps=0.0, - ) - return qweight, block_scale, _nvfp4_global_decode_scale_te(global_amax) + return nvfp4_quantize_1d(weight, global_amax, use_4over6=use_4over6) def quantize_nvfp4( weight: torch.Tensor, global_amax: torch.Tensor | None = None, + use_4over6: bool | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if use_4over6 is None: + use_4over6 = nvfp4_4over6_enabled() if weight.dim() == 2: - return _quantize_nvfp4_1d(weight, global_amax=global_amax) + return _quantize_nvfp4_1d(weight, global_amax=global_amax, use_4over6=use_4over6) if weight.dim() == 3: if global_amax is not None: raise ValueError("global_amax override is only supported for 2D weights.") @@ -173,7 +152,7 @@ def quantize_nvfp4( block_scales = [] global_scales = [] for idx in range(weight.shape[0]): - qweight, block_scale, global_scale = _quantize_nvfp4_1d(weight[idx]) + qweight, block_scale, global_scale = _quantize_nvfp4_1d(weight[idx], use_4over6=use_4over6) qweights.append(qweight) block_scales.append(block_scale) global_scales.append(global_scale) @@ -198,7 +177,7 @@ def add_result(self, filename: str, q_weights: dict[str, torch.Tensor], module_n self.modules_to_not_convert.extend(module_names) -def _update_quantization_config(cfg: dict, ignore_list: list[str]) -> None: +def _update_quantization_config(cfg: dict, ignore_list: list[str], use_4over6: bool | None = None) -> None: quant_cfg = cfg.get("quantization_config") if not isinstance(quant_cfg, dict): quant_cfg = {} @@ -206,6 +185,7 @@ def _update_quantization_config(cfg: dict, ignore_list: list[str]) -> None: quant_cfg["quant_algo"] = "NVFP4" quant_cfg["quant_method"] = "modelopt" quant_cfg["group_size"] = NVFP4_GROUP_SIZE + quant_cfg["nvfp4_4over6"] = nvfp4_4over6_weight_scope(use_4over6) quant_cfg["ignore"] = ignore_list quant_cfg.setdefault("kv_cache_scheme", DEFAULT_KV_CACHE_SCHEME) @@ -227,7 +207,12 @@ def _update_quantization_config(cfg: dict, ignore_list: list[str]) -> None: cfg["quantization_config"] = quant_cfg -def _write_hf_quant_config(output_path: str, ignore_list: list[str], input_path: str) -> None: +def _write_hf_quant_config( + output_path: str, + ignore_list: list[str], + input_path: str, + use_4over6: bool | None = None, +) -> None: hf_quant_path = os.path.join(input_path, "hf_quant_config.json") if os.path.exists(hf_quant_path): with open(hf_quant_path) as f: @@ -242,6 +227,7 @@ def _write_hf_quant_config(output_path: str, ignore_list: list[str], input_path: quant_section["quant_algo"] = "NVFP4" quant_section["kv_cache_quant_algo"] = DEFAULT_KV_CACHE_QUANT_ALGO quant_section["group_size"] = NVFP4_GROUP_SIZE + quant_section["nvfp4_4over6"] = nvfp4_4over6_weight_scope(use_4over6) quant_section["exclude_modules"] = ignore_list hf_quant_cfg["quantization"] = quant_section @@ -258,6 +244,9 @@ def _augment_ignore_list(ignore_list: list[str]) -> list[str]: if name.endswith(suffix): extra.add(name[: -len(suffix)] + ".qkv_proj") break + match = re.match(r"(.*\.mlp\.experts)\.\d+\.(gate_proj|up_proj|down_proj)$", name) + if match: + extra.add(match.group(1)) ignore_set.update(extra) return sorted(ignore_set) @@ -315,6 +304,7 @@ def process_file( num_layers_at_end_in_bf16: int, extra_high_precision_layers_hf: tuple[str, ...], shared_global_amax: dict[str, torch.Tensor], + use_4over6: bool, ) -> None: if not filename.endswith(".safetensors"): return @@ -341,7 +331,11 @@ def process_file( if should_quantize(key, tensor, skip_weight_substrings=dynamic_skip_substrings): base, _role = _split_gated_pair_name(key) global_amax = shared_global_amax.get(base) if base else None - qweight, block_scale, weight_scale_2 = quantize_nvfp4(tensor, global_amax=global_amax) + qweight, block_scale, weight_scale_2 = quantize_nvfp4( + tensor, + global_amax=global_amax, + use_4over6=use_4over6, + ) q_weights[key] = qweight q_weights[key.replace(".weight", ".weight_scale")] = block_scale q_weights[key.replace(".weight", ".weight_scale_2")] = weight_scale_2 @@ -385,6 +379,7 @@ def convert_nvfp4( *extra_high_precision_layers_hf, *sorted(dynamic_skip_layer_prefixes), ) + use_4over6 = nvfp4_4over6_enabled() shared_global_amax = _collect_shared_global_amax( input_path=input_path, @@ -405,6 +400,7 @@ def convert_nvfp4( num_layers_at_end_in_bf16, extra_high_precision_layers_hf, shared_global_amax, + use_4over6, ) gc.collect() if torch.cuda.is_available(): @@ -415,10 +411,10 @@ def convert_nvfp4( config_path = os.path.join(input_path, "config.json") if os.path.exists(config_path): cfg = json.load(open(config_path)) - _update_quantization_config(cfg, ignore_list) + _update_quantization_config(cfg, ignore_list, use_4over6=use_4over6) json.dump(cfg, open(os.path.join(output_path, "config.json"), "w"), indent=2) - _write_hf_quant_config(output_path, ignore_list, input_path) + _write_hf_quant_config(output_path, ignore_list, input_path, use_4over6=use_4over6) index_dict = { "weight_map": result_collector.weight_map, From 9443aee73db8b9b47bb27cfcf8db6fd1e413f49a Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Wed, 27 May 2026 01:07:11 -0700 Subject: [PATCH 02/37] Skip checkpoint saving in GLM5 NVFP4 e2e --- .../e2e/megatron/test_glm5_744b_a40b_5layer_nvfp4.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/tests/e2e/megatron/test_glm5_744b_a40b_5layer_nvfp4.py b/tests/e2e/megatron/test_glm5_744b_a40b_5layer_nvfp4.py index 8a10fa20b32..2258154f49e 100644 --- a/tests/e2e/megatron/test_glm5_744b_a40b_5layer_nvfp4.py +++ b/tests/e2e/megatron/test_glm5_744b_a40b_5layer_nvfp4.py @@ -166,16 +166,8 @@ def execute(): os.environ.update(GLM5_ENV) os.environ.setdefault("RAY_TMPDIR", "/tmp/ray") te_precision_config_path = U.save_to_temp_file(TE_PRECISION_CONFIG, "yaml") - load_save_path = f"/root/shared_data/{RUN_ID}/checkpoints" - - ckpt_args = ( - f"--hf-checkpoint {MODEL_DIR}/{MODEL_NAME}-NVFP4/ " - f"--ref-load {MODEL_DIR}/{MODEL_NAME}_torch_dist " - f"--load {load_save_path} " - f"--save {load_save_path} " - "--save-interval 2 " - "--save-retain-interval 2 " - ) + + ckpt_args = f"--hf-checkpoint {MODEL_DIR}/{MODEL_NAME}-NVFP4/ " f"--ref-load {MODEL_DIR}/{MODEL_NAME}_torch_dist " rollout_args = ( f"--prompt-data {DATA_DIR}/dapo-math-17k/dapo-math-17k.jsonl " From ff7598d61cfef391cd12951317a170caf73c07ef Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Sat, 30 May 2026 16:34:36 -0700 Subject: [PATCH 03/37] refactor: simplify nvfp4 env handling --- .../processors/quantizer_nvfp4.py | 20 +-- miles/utils/nvfp4.py | 47 +++-- tests/fast-gpu/test_nvfp4_quantizer.py | 162 +----------------- tools/convert_hf_to_nvfp4.py | 27 +-- 4 files changed, 37 insertions(+), 219 deletions(-) diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py index 5af1d08614f..9bad6070d3e 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py @@ -2,7 +2,7 @@ import torch -from miles.utils.nvfp4 import NVFP4_GROUP_SIZE, nvfp4_4over6_enabled, nvfp4_global_decode_scale_te, nvfp4_quantize_1d +from miles.utils.nvfp4 import NVFP4_GROUP_SIZE, nvfp4_global_decode_scale_te, nvfp4_quantize_1d GATED_PAIR_SUFFIXES = { ".gate_proj.weight": "gate", @@ -18,14 +18,8 @@ def is_nvfp4_quantization_config(quantization_config) -> bool: return quantization_config.get("quant_algo") == "NVFP4" or quantization_config.get("quant_method") == "nvfp4" -def fp4_param_gather_enabled(args) -> bool: - if args is None: - return False - return bool(getattr(args, "fp4_param", False) or getattr(args, "fp4_param_gather", False)) - - def assert_no_fp4_param_gather(args) -> None: - if fp4_param_gather_enabled(args): + if args is not None and bool(getattr(args, "fp4_param", False) or getattr(args, "fp4_param_gather", False)): raise NotImplementedError("fp4-param-gather is unsupported for Miles NVFP4 checkpoint export.") @@ -184,7 +178,6 @@ def _split_gated_pair_name(name: str): def _quantize_nvfp4_1d( weight: torch.Tensor, global_amax: torch.Tensor | None = None, - use_4over6: bool | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ NVFP4 1D quantization (tile shape = 1x16). @@ -204,18 +197,15 @@ def _quantize_nvfp4_1d( else: global_amax = global_amax.to(device=weight.device, dtype=torch.float32) - return nvfp4_quantize_1d(weight, global_amax, use_4over6=use_4over6) + return nvfp4_quantize_1d(weight, global_amax) def quantize_nvfp4( weight: torch.Tensor, global_amax: torch.Tensor | None = None, - use_4over6: bool | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - if use_4over6 is None: - use_4over6 = nvfp4_4over6_enabled() if weight.dim() == 2: - return _quantize_nvfp4_1d(weight, global_amax=global_amax, use_4over6=use_4over6) + return _quantize_nvfp4_1d(weight, global_amax=global_amax) if weight.dim() == 3: if global_amax is not None: raise ValueError("global_amax override is only supported for 2D weights.") @@ -223,7 +213,7 @@ def quantize_nvfp4( block_scales = [] global_scales = [] for idx in range(weight.shape[0]): - qweight, block_scale, global_scale = _quantize_nvfp4_1d(weight[idx], use_4over6=use_4over6) + qweight, block_scale, global_scale = _quantize_nvfp4_1d(weight[idx]) qweights.append(qweight) block_scales.append(block_scale) global_scales.append(global_scale) diff --git a/miles/utils/nvfp4.py b/miles/utils/nvfp4.py index 002c368a3f4..1a60d9f11c3 100644 --- a/miles/utils/nvfp4.py +++ b/miles/utils/nvfp4.py @@ -40,14 +40,14 @@ def nvfp4_4over6_enabled() -> bool: return nvfp4_4over6_weight_scope_enabled(os.getenv("NVTE_NVFP4_4OVER6")) -def nvfp4_4over6_weight_scope(use_4over6: bool | None = None) -> str: - if use_4over6 is None: - use_4over6 = nvfp4_4over6_enabled() - return "weights" if use_4over6 else "none" +def nvfp4_4over6_weight_scope() -> str: + return "weights" if nvfp4_4over6_enabled() else "none" -def nvfp4_weight_e4m3_max(use_4over6: bool) -> int: - if use_4over6 and nvfp4_4over6_weight_scope_enabled(os.getenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all")): +def nvfp4_weight_e4m3_max() -> int: + if nvfp4_4over6_enabled() and nvfp4_4over6_weight_scope_enabled( + os.getenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all") + ): return 256 return int(FP8_E4M3_MAX) @@ -93,15 +93,14 @@ def nvfp4_global_decode_scale_te( return torch.div(1.0, nvfp4_global_encode_scale_te(global_amax, nvfp4_e4m3_max)) -def sync_flashinfer_nvfp4_env_from_nvte(use_4over6: bool | None = None) -> dict[str, str]: - if use_4over6 is None: - use_4over6 = nvfp4_4over6_enabled() - +def sync_flashinfer_nvfp4_env_from_nvte() -> dict[str, str]: + weight_4over6_enabled = nvfp4_4over6_enabled() flashinfer_env = { - "FLASHINFER_NVFP4_4OVER6": "1" if use_4over6 else "0", + "FLASHINFER_NVFP4_4OVER6": "1" if weight_4over6_enabled else "0", "FLASHINFER_NVFP4_4OVER6_E4M3_USE_256": ( "1" - if use_4over6 and nvfp4_4over6_weight_scope_enabled(os.getenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all")) + if weight_4over6_enabled + and nvfp4_4over6_weight_scope_enabled(os.getenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all")) else "0" ), "FLASHINFER_NVFP4_4OVER6_ERR_MODE": nvfp4_4over6_err_mode(), @@ -115,10 +114,10 @@ def sync_flashinfer_nvfp4_env_from_nvte(use_4over6: bool | None = None) -> dict[ @contextmanager -def flashinfer_nvfp4_env_from_nvte(use_4over6: bool | None = None): +def flashinfer_nvfp4_env_from_nvte(): original_env = {key: os.environ.get(key) for key in FLASHINFER_NVFP4_ENV_KEYS} - sync_flashinfer_nvfp4_env_from_nvte(use_4over6) try: + sync_flashinfer_nvfp4_env_from_nvte() yield finally: for key, value in original_env.items(): @@ -140,14 +139,14 @@ def flashinfer_nvfp4_env_from_nvte(use_4over6: bool | None = None): def _te_nvfp4_quantize_1d( weight: torch.Tensor, global_amax: torch.Tensor, - use_4over6: bool, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: try: from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef except ImportError: from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import NVFP4QuantizerRef - nvfp4_e4m3_max = nvfp4_weight_e4m3_max(use_4over6) + weight_4over6_enabled = nvfp4_4over6_enabled() + nvfp4_e4m3_max = nvfp4_weight_e4m3_max() try: qweight, block_scale = NVFP4QuantizerRef._quantize_blockwise_reference( weight, @@ -155,13 +154,13 @@ def _te_nvfp4_quantize_1d( NVFP4_GROUP_SIZE, 1, pow_2_scales=False, - nvfp4_use_4over6=use_4over6, + nvfp4_use_4over6=weight_4over6_enabled, nvfp4_e4m3_max=nvfp4_e4m3_max, nvfp4_4over6_err_mode=nvfp4_4over6_err_mode(), eps=0.0, ) except TypeError: - if use_4over6: + if weight_4over6_enabled: raise qweight, block_scale = NVFP4QuantizerRef._quantize_blockwise_reference( weight, @@ -177,17 +176,13 @@ def _te_nvfp4_quantize_1d( def nvfp4_quantize_1d( weight: torch.Tensor, global_amax: torch.Tensor, - use_4over6: bool | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - if use_4over6 is None: - use_4over6 = nvfp4_4over6_enabled() - - if _flashinfer_nvfp4_quantize is None or weight.dtype == torch.float32: - return _te_nvfp4_quantize_1d(weight, global_amax, use_4over6) + if _flashinfer_nvfp4_quantize is None or weight.dtype == torch.float32 or not weight.is_cuda: + return _te_nvfp4_quantize_1d(weight, global_amax) - nvfp4_e4m3_max = nvfp4_weight_e4m3_max(use_4over6) + nvfp4_e4m3_max = nvfp4_weight_e4m3_max() global_encode_scale = nvfp4_global_encode_scale_te(global_amax, nvfp4_e4m3_max) - with flashinfer_nvfp4_env_from_nvte(use_4over6): + with flashinfer_nvfp4_env_from_nvte(): qweight, block_scale = _flashinfer_nvfp4_quantize( weight, global_encode_scale.reshape(1).contiguous(), diff --git a/tests/fast-gpu/test_nvfp4_quantizer.py b/tests/fast-gpu/test_nvfp4_quantizer.py index 7b49b6baef7..2311615ca04 100644 --- a/tests/fast-gpu/test_nvfp4_quantizer.py +++ b/tests/fast-gpu/test_nvfp4_quantizer.py @@ -3,30 +3,21 @@ register_cuda_ci(est_time=60, suite="stage-b-2-gpu-h200", labels=[]) -import json -import os - import pytest import torch -from tools.convert_hf_to_nvfp4 import _update_quantization_config as tool_update_quantization_config -from tools.convert_hf_to_nvfp4 import _write_hf_quant_config as tool_write_hf_quant_config from tools.convert_hf_to_nvfp4 import quantize_nvfp4 as tool_quantize_nvfp4 from tools.convert_hf_to_nvfp4 import should_quantize as tool_should_quantize_nvfp4 from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef -from miles.backends.megatron_utils.megatron_to_hf.processors import quantize_params from miles.backends.megatron_utils.megatron_to_hf.processors.quantizer_nvfp4 import ( quantize_nvfp4 as processor_quantize_nvfp4, ) from miles.backends.megatron_utils.megatron_to_hf.processors.quantizer_nvfp4 import quantize_params_nvfp4 from miles.utils.nvfp4 import ( NVFP4_GROUP_SIZE, - flashinfer_nvfp4_env_from_nvte, nvfp4_4over6_err_mode, - nvfp4_4over6_weight_scope, nvfp4_global_decode_scale_te, nvfp4_weight_e4m3_max, - sync_flashinfer_nvfp4_env_from_nvte, ) NVFP4_SHAPES = [ @@ -84,27 +75,25 @@ def _make_weight(init_data: str, dtype: torch.dtype, shape: tuple[int, int], dev def _te_nvfp4_reference( weight: torch.Tensor, - use_4over6: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: weight = weight.contiguous() global_amax = torch.max(torch.abs(weight.to(torch.float32))) - return _te_nvfp4_reference_with_global_amax(weight, global_amax, use_4over6=use_4over6) + return _te_nvfp4_reference_with_global_amax(weight, global_amax) def _te_nvfp4_reference_with_global_amax( weight: torch.Tensor, global_amax: torch.Tensor, - use_4over6: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: weight = weight.contiguous() - nvfp4_e4m3_max = nvfp4_weight_e4m3_max(use_4over6) + nvfp4_e4m3_max = nvfp4_weight_e4m3_max() qweight, block_scale = NVFP4QuantizerRef._quantize_blockwise_reference( weight, global_amax, NVFP4_GROUP_SIZE, 1, pow_2_scales=False, - nvfp4_use_4over6=use_4over6, + nvfp4_use_4over6=False, nvfp4_e4m3_max=nvfp4_e4m3_max, nvfp4_4over6_err_mode=nvfp4_4over6_err_mode(), eps=0.0, @@ -112,19 +101,6 @@ def _te_nvfp4_reference_with_global_amax( return qweight, block_scale, nvfp4_global_decode_scale_te(global_amax, nvfp4_e4m3_max) -def test_nvfp4_dispatch_accepts_quant_algo_without_quant_method(): - converted_named_params = [("model.embed_tokens.weight", torch.zeros((1, 1), dtype=torch.bfloat16))] - - out = quantize_params( - args=None, - megatron_name="embedding.word_embeddings.weight", - converted_named_params=converted_named_params, - quantization_config={"quant_algo": "NVFP4"}, - ) - - assert out is converted_named_params - - def test_nvfp4_quantize_params_requires_complete_gated_pair(): weight = torch.randn((4, NVFP4_GROUP_SIZE), dtype=torch.float32) with pytest.raises(ValueError, match="requires gate/up tensors to be quantized together"): @@ -156,22 +132,6 @@ def test_nvfp4_quantize_params_respects_extra_high_precision_layers_megatron(): assert out is converted_named_params -def test_nvfp4_quantize_params_rejects_fp4_param_gather(): - weight = torch.randn((4, NVFP4_GROUP_SIZE), dtype=torch.bfloat16) - args = type("Args", (), {"fp4_param": True})() - - with pytest.raises(NotImplementedError, match="fp4-param-gather is unsupported"): - quantize_params_nvfp4( - args=args, - megatron_name="decoder.layers.0.mlp.experts.linear_fc1.weight0", - converted_named_params=[ - ("model.layers.0.mlp.experts.0.gate_proj.weight", weight), - ("model.layers.0.mlp.experts.0.up_proj.weight", weight), - ], - quantization_config={"quant_method": "nvfp4"}, - ) - - @pytest.mark.parametrize("layer_idx", [0, 3]) def test_nvfp4_quantize_params_respects_first_last_layers_bf16(layer_idx): weight = torch.randn((4, NVFP4_GROUP_SIZE), dtype=torch.bfloat16) @@ -215,57 +175,6 @@ def test_nvfp4_hf_should_quantize_respects_extra_high_precision_layers_hf(): ) -def test_nvfp4_converter_records_4over6_mode(tmp_path, monkeypatch): - monkeypatch.setenv("NVTE_NVFP4_4OVER6", "weights") - cfg = {} - tool_update_quantization_config(cfg, ignore_list=["model.layers.0"]) - assert cfg["quantization_config"]["nvfp4_4over6"] == "weights" - - input_dir = tmp_path / "input" - output_dir = tmp_path / "output" - input_dir.mkdir() - output_dir.mkdir() - tool_write_hf_quant_config(str(output_dir), ignore_list=["model.layers.0"], input_path=str(input_dir)) - - hf_quant_config = json.loads((output_dir / "hf_quant_config.json").read_text()) - assert hf_quant_config["quantization"]["nvfp4_4over6"] == "weights" - - -def test_nvfp4_flashinfer_env_syncs_from_nvte(monkeypatch): - monkeypatch.setenv("NVTE_NVFP4_4OVER6", "all") - monkeypatch.setenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all") - monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MSE") - monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", "0") - monkeypatch.delenv("TRTLLM_DISABLE_FP4_QUANT_FAST_MATH", raising=False) - - synced = sync_flashinfer_nvfp4_env_from_nvte() - - assert synced["FLASHINFER_NVFP4_4OVER6"] == "1" - assert synced["FLASHINFER_NVFP4_4OVER6_E4M3_USE_256"] == "1" - assert synced["FLASHINFER_NVFP4_4OVER6_ERR_MODE"] == "MSE" - assert synced["FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH"] == "0" - assert synced["TRTLLM_DISABLE_FP4_QUANT_FAST_MATH"] == "1" - - -def test_nvfp4_flashinfer_env_context_restores_previous_values(monkeypatch): - monkeypatch.setenv("NVTE_NVFP4_4OVER6", "all") - monkeypatch.setenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all") - monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MSE") - monkeypatch.setenv("FLASHINFER_NVFP4_4OVER6", "old") - monkeypatch.setenv("TRTLLM_DISABLE_FP4_QUANT_FAST_MATH", "0") - - with flashinfer_nvfp4_env_from_nvte(): - assert os.environ["FLASHINFER_NVFP4_4OVER6"] == "1" - assert os.environ["FLASHINFER_NVFP4_4OVER6_E4M3_USE_256"] == "1" - assert os.environ["FLASHINFER_NVFP4_4OVER6_ERR_MODE"] == "MSE" - assert os.environ["TRTLLM_DISABLE_FP4_QUANT_FAST_MATH"] == "0" - - assert os.environ["FLASHINFER_NVFP4_4OVER6"] == "old" - assert "FLASHINFER_NVFP4_4OVER6_E4M3_USE_256" not in os.environ - assert "FLASHINFER_NVFP4_4OVER6_ERR_MODE" not in os.environ - assert os.environ["TRTLLM_DISABLE_FP4_QUANT_FAST_MATH"] == "0" - - @pytest.mark.parametrize( "quantize_fn", [processor_quantize_nvfp4, tool_quantize_nvfp4], @@ -287,71 +196,6 @@ def test_nvfp4_quantize_matches_te_reference_bitwise(quantize_fn, shape, dtype, torch.testing.assert_close(global_scale, global_scale_ref, rtol=0, atol=0) -@pytest.mark.parametrize("use_4over6", [False, True], ids=lambda value: nvfp4_4over6_weight_scope(value)) -@pytest.mark.parametrize("err_mode", ["MAE", "MSE"]) -def test_nvfp4_quantize_matches_te_reference_with_4over6_modes(monkeypatch, use_4over6, err_mode): - device = "cuda" - torch.manual_seed(42) - monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_MODE", err_mode) - monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", "0") - if use_4over6: - monkeypatch.setenv("NVTE_NVFP4_4OVER6", "weights") - monkeypatch.setenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all") - else: - monkeypatch.delenv("NVTE_NVFP4_4OVER6", raising=False) - - weight = _make_weight("random", torch.bfloat16, (128, 1024), device) - qweight, block_scale, global_scale = processor_quantize_nvfp4(weight) - qweight_ref, block_scale_ref, global_scale_ref = _te_nvfp4_reference(weight, use_4over6=use_4over6) - - torch.testing.assert_close(qweight, qweight_ref, rtol=0, atol=0) - torch.testing.assert_close(block_scale.view(torch.uint8), block_scale_ref.view(torch.uint8), rtol=0, atol=0) - torch.testing.assert_close(global_scale, global_scale_ref, rtol=0, atol=0) - - -def test_nvfp4_quantize_params_reads_4over6_from_env(monkeypatch): - device = "cuda" - torch.manual_seed(42) - monkeypatch.setenv("NVTE_NVFP4_4OVER6", "weights") - monkeypatch.setenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all") - - gate = _make_weight("random", torch.bfloat16, (4, NVFP4_GROUP_SIZE), device) - up = _make_weight("random", torch.bfloat16, (4, NVFP4_GROUP_SIZE), device) - shared_amax = torch.max(gate.abs().max().to(torch.float32), up.abs().max().to(torch.float32)) - converted_named_params = [ - ("model.layers.0.mlp.experts.0.gate_proj.weight", gate), - ("model.layers.0.mlp.experts.0.up_proj.weight", up), - ] - - out = dict( - quantize_params_nvfp4( - args=None, - megatron_name="decoder.layers.0.mlp.experts.linear_fc1.weight0", - converted_named_params=converted_named_params, - quantization_config={"quant_method": "nvfp4"}, - ) - ) - qweight_ref, block_scale_ref, global_scale_ref = _te_nvfp4_reference_with_global_amax( - gate, - shared_amax, - use_4over6=True, - ) - - torch.testing.assert_close(out["model.layers.0.mlp.experts.0.gate_proj.weight"], qweight_ref, rtol=0, atol=0) - torch.testing.assert_close( - out["model.layers.0.mlp.experts.0.gate_proj.weight_scale"].view(torch.uint8), - block_scale_ref.view(torch.uint8), - rtol=0, - atol=0, - ) - torch.testing.assert_close( - out["model.layers.0.mlp.experts.0.gate_proj.weight_scale_2"], - global_scale_ref, - rtol=0, - atol=0, - ) - - if __name__ == "__main__": import sys diff --git a/tools/convert_hf_to_nvfp4.py b/tools/convert_hf_to_nvfp4.py index 1b72b6442cb..12904887f6f 100644 --- a/tools/convert_hf_to_nvfp4.py +++ b/tools/convert_hf_to_nvfp4.py @@ -29,7 +29,6 @@ from miles.utils.nvfp4 import ( NVFP4_GROUP_SIZE, - nvfp4_4over6_enabled, nvfp4_4over6_weight_scope, nvfp4_global_decode_scale_te, nvfp4_quantize_1d, @@ -113,7 +112,6 @@ def should_quantize( def _quantize_nvfp4_1d( weight: torch.Tensor, global_amax: torch.Tensor | None = None, - use_4over6: bool | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ NVFP4 1D quantization (tile shape = 1x16). @@ -133,18 +131,15 @@ def _quantize_nvfp4_1d( else: global_amax = global_amax.to(device=weight.device, dtype=torch.float32) - return nvfp4_quantize_1d(weight, global_amax, use_4over6=use_4over6) + return nvfp4_quantize_1d(weight, global_amax) def quantize_nvfp4( weight: torch.Tensor, global_amax: torch.Tensor | None = None, - use_4over6: bool | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - if use_4over6 is None: - use_4over6 = nvfp4_4over6_enabled() if weight.dim() == 2: - return _quantize_nvfp4_1d(weight, global_amax=global_amax, use_4over6=use_4over6) + return _quantize_nvfp4_1d(weight, global_amax=global_amax) if weight.dim() == 3: if global_amax is not None: raise ValueError("global_amax override is only supported for 2D weights.") @@ -152,7 +147,7 @@ def quantize_nvfp4( block_scales = [] global_scales = [] for idx in range(weight.shape[0]): - qweight, block_scale, global_scale = _quantize_nvfp4_1d(weight[idx], use_4over6=use_4over6) + qweight, block_scale, global_scale = _quantize_nvfp4_1d(weight[idx]) qweights.append(qweight) block_scales.append(block_scale) global_scales.append(global_scale) @@ -177,7 +172,7 @@ def add_result(self, filename: str, q_weights: dict[str, torch.Tensor], module_n self.modules_to_not_convert.extend(module_names) -def _update_quantization_config(cfg: dict, ignore_list: list[str], use_4over6: bool | None = None) -> None: +def _update_quantization_config(cfg: dict, ignore_list: list[str]) -> None: quant_cfg = cfg.get("quantization_config") if not isinstance(quant_cfg, dict): quant_cfg = {} @@ -185,7 +180,7 @@ def _update_quantization_config(cfg: dict, ignore_list: list[str], use_4over6: b quant_cfg["quant_algo"] = "NVFP4" quant_cfg["quant_method"] = "modelopt" quant_cfg["group_size"] = NVFP4_GROUP_SIZE - quant_cfg["nvfp4_4over6"] = nvfp4_4over6_weight_scope(use_4over6) + quant_cfg["nvfp4_4over6"] = nvfp4_4over6_weight_scope() quant_cfg["ignore"] = ignore_list quant_cfg.setdefault("kv_cache_scheme", DEFAULT_KV_CACHE_SCHEME) @@ -211,7 +206,6 @@ def _write_hf_quant_config( output_path: str, ignore_list: list[str], input_path: str, - use_4over6: bool | None = None, ) -> None: hf_quant_path = os.path.join(input_path, "hf_quant_config.json") if os.path.exists(hf_quant_path): @@ -227,7 +221,7 @@ def _write_hf_quant_config( quant_section["quant_algo"] = "NVFP4" quant_section["kv_cache_quant_algo"] = DEFAULT_KV_CACHE_QUANT_ALGO quant_section["group_size"] = NVFP4_GROUP_SIZE - quant_section["nvfp4_4over6"] = nvfp4_4over6_weight_scope(use_4over6) + quant_section["nvfp4_4over6"] = nvfp4_4over6_weight_scope() quant_section["exclude_modules"] = ignore_list hf_quant_cfg["quantization"] = quant_section @@ -304,7 +298,6 @@ def process_file( num_layers_at_end_in_bf16: int, extra_high_precision_layers_hf: tuple[str, ...], shared_global_amax: dict[str, torch.Tensor], - use_4over6: bool, ) -> None: if not filename.endswith(".safetensors"): return @@ -334,7 +327,6 @@ def process_file( qweight, block_scale, weight_scale_2 = quantize_nvfp4( tensor, global_amax=global_amax, - use_4over6=use_4over6, ) q_weights[key] = qweight q_weights[key.replace(".weight", ".weight_scale")] = block_scale @@ -379,8 +371,6 @@ def convert_nvfp4( *extra_high_precision_layers_hf, *sorted(dynamic_skip_layer_prefixes), ) - use_4over6 = nvfp4_4over6_enabled() - shared_global_amax = _collect_shared_global_amax( input_path=input_path, safetensors_files=safetensors_files, @@ -400,7 +390,6 @@ def convert_nvfp4( num_layers_at_end_in_bf16, extra_high_precision_layers_hf, shared_global_amax, - use_4over6, ) gc.collect() if torch.cuda.is_available(): @@ -411,10 +400,10 @@ def convert_nvfp4( config_path = os.path.join(input_path, "config.json") if os.path.exists(config_path): cfg = json.load(open(config_path)) - _update_quantization_config(cfg, ignore_list, use_4over6=use_4over6) + _update_quantization_config(cfg, ignore_list) json.dump(cfg, open(os.path.join(output_path, "config.json"), "w"), indent=2) - _write_hf_quant_config(output_path, ignore_list, input_path, use_4over6=use_4over6) + _write_hf_quant_config(output_path, ignore_list, input_path) index_dict = { "weight_map": result_collector.weight_map, From ac889b47793b3801266fee02cb30087c12891440 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Sat, 30 May 2026 16:41:36 -0700 Subject: [PATCH 04/37] refactor: inline nvfp4 env helpers --- .../megatron_to_hf/processors/__init__.py | 4 +- .../processors/quantizer_nvfp4.py | 8 +-- miles/utils/nvfp4.py | 53 ++++--------------- tests/fast-gpu/test_nvfp4_quantizer.py | 11 ++-- tools/convert_hf_to_nvfp4.py | 15 +++--- 5 files changed, 25 insertions(+), 66 deletions(-) diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py b/miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py index cb1a7d613d0..42aed336921 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py @@ -2,7 +2,7 @@ from .quantizer_compressed_tensors import quantize_params_compressed_tensors from .quantizer_fp8 import quantize_params_fp8 from .quantizer_mxfp8 import quantize_params_mxfp8 -from .quantizer_nvfp4 import is_nvfp4_quantization_config, quantize_params_nvfp4 +from .quantizer_nvfp4 import quantize_params_nvfp4 __all__ = [ "remove_padding", @@ -21,7 +21,7 @@ def quantize_params(args, megatron_name, converted_named_params, quantization_co return quantize_params_fp8(args, megatron_name, converted_named_params, quantization_config) elif quantization_config.get("quant_method") == "mxfp8": return quantize_params_mxfp8(args, megatron_name, converted_named_params, quantization_config) - elif is_nvfp4_quantization_config(quantization_config): + elif quantization_config.get("quant_algo") == "NVFP4" or quantization_config.get("quant_method") == "nvfp4": return quantize_params_nvfp4(args, megatron_name, converted_named_params, quantization_config) elif quantization_config.get("quant_method") == "compressed-tensors": # only int4 at the moment. diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py index 9bad6070d3e..2f0469c26dd 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py @@ -12,12 +12,6 @@ } -def is_nvfp4_quantization_config(quantization_config) -> bool: - if quantization_config is None: - return False - return quantization_config.get("quant_algo") == "NVFP4" or quantization_config.get("quant_method") == "nvfp4" - - def assert_no_fp4_param_gather(args) -> None: if args is not None and bool(getattr(args, "fp4_param", False) or getattr(args, "fp4_param_gather", False)): raise NotImplementedError("fp4-param-gather is unsupported for Miles NVFP4 checkpoint export.") @@ -46,7 +40,7 @@ def _is_ignored(name: str, ignore_rules: list[str]) -> bool: def quantize_params_nvfp4(args, megatron_name, converted_named_params, quantization_config): assert quantization_config is not None - assert is_nvfp4_quantization_config(quantization_config) + assert quantization_config.get("quant_algo") == "NVFP4" or quantization_config.get("quant_method") == "nvfp4" assert_no_fp4_param_gather(args) if getattr(args, "extra_high_precision_layers_megatron", False): diff --git a/miles/utils/nvfp4.py b/miles/utils/nvfp4.py index 1a60d9f11c3..8a94799d76c 100644 --- a/miles/utils/nvfp4.py +++ b/miles/utils/nvfp4.py @@ -18,47 +18,14 @@ ) -def str_to_bool(value) -> bool: - if isinstance(value, bool): - return value - if value is None: - return False - return str(value).strip().lower() in ("1", "true", "yes", "on") - - -def nvfp4_4over6_weight_scope_enabled(value) -> bool: - if isinstance(value, str): - value = value.strip().lower() - if value in ("weights", "all"): - return True - if value in ("none", "activations"): - return False - return str_to_bool(value) - - -def nvfp4_4over6_enabled() -> bool: - return nvfp4_4over6_weight_scope_enabled(os.getenv("NVTE_NVFP4_4OVER6")) - - -def nvfp4_4over6_weight_scope() -> str: - return "weights" if nvfp4_4over6_enabled() else "none" - - def nvfp4_weight_e4m3_max() -> int: - if nvfp4_4over6_enabled() and nvfp4_4over6_weight_scope_enabled( - os.getenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all") - ): + if os.getenv("NVTE_NVFP4_4OVER6", "").strip().lower() in ("weights", "all") and os.getenv( + "NVTE_NVFP4_4OVER6_E4M3_USE_256", "all" + ).strip().lower() in ("weights", "all"): return 256 return int(FP8_E4M3_MAX) -def nvfp4_4over6_err_mode() -> str: - err_mode = os.getenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MAE").strip().upper() - if err_mode not in ("MAE", "MSE"): - raise ValueError("NVTE_NVFP4_4OVER6_ERR_MODE must be one of: 'MAE', 'MSE'.") - return err_mode - - def nvfp4_global_encode_scale_te( global_amax: torch.Tensor, nvfp4_e4m3_max: int = int(FP8_E4M3_MAX), @@ -94,18 +61,20 @@ def nvfp4_global_decode_scale_te( def sync_flashinfer_nvfp4_env_from_nvte() -> dict[str, str]: - weight_4over6_enabled = nvfp4_4over6_enabled() + weight_4over6_enabled = os.getenv("NVTE_NVFP4_4OVER6", "").strip().lower() in ("weights", "all") flashinfer_env = { "FLASHINFER_NVFP4_4OVER6": "1" if weight_4over6_enabled else "0", "FLASHINFER_NVFP4_4OVER6_E4M3_USE_256": ( "1" if weight_4over6_enabled - and nvfp4_4over6_weight_scope_enabled(os.getenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all")) + and os.getenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all").strip().lower() in ("weights", "all") else "0" ), - "FLASHINFER_NVFP4_4OVER6_ERR_MODE": nvfp4_4over6_err_mode(), + "FLASHINFER_NVFP4_4OVER6_ERR_MODE": os.getenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MAE").strip().upper(), "FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH": ( - "1" if str_to_bool(os.getenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", "0")) else "0" + "1" + if os.getenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", "0").strip().lower() in ("1", "true", "yes", "on") + else "0" ), } os.environ.update(flashinfer_env) @@ -145,7 +114,7 @@ def _te_nvfp4_quantize_1d( except ImportError: from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import NVFP4QuantizerRef - weight_4over6_enabled = nvfp4_4over6_enabled() + weight_4over6_enabled = os.getenv("NVTE_NVFP4_4OVER6", "").strip().lower() in ("weights", "all") nvfp4_e4m3_max = nvfp4_weight_e4m3_max() try: qweight, block_scale = NVFP4QuantizerRef._quantize_blockwise_reference( @@ -156,7 +125,7 @@ def _te_nvfp4_quantize_1d( pow_2_scales=False, nvfp4_use_4over6=weight_4over6_enabled, nvfp4_e4m3_max=nvfp4_e4m3_max, - nvfp4_4over6_err_mode=nvfp4_4over6_err_mode(), + nvfp4_4over6_err_mode=os.getenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MAE").strip().upper(), eps=0.0, ) except TypeError: diff --git a/tests/fast-gpu/test_nvfp4_quantizer.py b/tests/fast-gpu/test_nvfp4_quantizer.py index 2311615ca04..1cbe1bb98c8 100644 --- a/tests/fast-gpu/test_nvfp4_quantizer.py +++ b/tests/fast-gpu/test_nvfp4_quantizer.py @@ -3,6 +3,8 @@ register_cuda_ci(est_time=60, suite="stage-b-2-gpu-h200", labels=[]) +import os + import pytest import torch from tools.convert_hf_to_nvfp4 import quantize_nvfp4 as tool_quantize_nvfp4 @@ -13,12 +15,7 @@ quantize_nvfp4 as processor_quantize_nvfp4, ) from miles.backends.megatron_utils.megatron_to_hf.processors.quantizer_nvfp4 import quantize_params_nvfp4 -from miles.utils.nvfp4 import ( - NVFP4_GROUP_SIZE, - nvfp4_4over6_err_mode, - nvfp4_global_decode_scale_te, - nvfp4_weight_e4m3_max, -) +from miles.utils.nvfp4 import NVFP4_GROUP_SIZE, nvfp4_global_decode_scale_te, nvfp4_weight_e4m3_max NVFP4_SHAPES = [ (1, 64), @@ -95,7 +92,7 @@ def _te_nvfp4_reference_with_global_amax( pow_2_scales=False, nvfp4_use_4over6=False, nvfp4_e4m3_max=nvfp4_e4m3_max, - nvfp4_4over6_err_mode=nvfp4_4over6_err_mode(), + nvfp4_4over6_err_mode=os.getenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MAE").strip().upper(), eps=0.0, ) return qweight, block_scale, nvfp4_global_decode_scale_te(global_amax, nvfp4_e4m3_max) diff --git a/tools/convert_hf_to_nvfp4.py b/tools/convert_hf_to_nvfp4.py index 12904887f6f..06121ed5ede 100644 --- a/tools/convert_hf_to_nvfp4.py +++ b/tools/convert_hf_to_nvfp4.py @@ -27,12 +27,7 @@ import torch from tqdm import tqdm -from miles.utils.nvfp4 import ( - NVFP4_GROUP_SIZE, - nvfp4_4over6_weight_scope, - nvfp4_global_decode_scale_te, - nvfp4_quantize_1d, -) +from miles.utils.nvfp4 import NVFP4_GROUP_SIZE, nvfp4_global_decode_scale_te, nvfp4_quantize_1d DEFAULT_KV_CACHE_SCHEME = {"dynamic": False, "num_bits": 8, "type": "float"} DEFAULT_KV_CACHE_QUANT_ALGO = "FP8" @@ -180,7 +175,9 @@ def _update_quantization_config(cfg: dict, ignore_list: list[str]) -> None: quant_cfg["quant_algo"] = "NVFP4" quant_cfg["quant_method"] = "modelopt" quant_cfg["group_size"] = NVFP4_GROUP_SIZE - quant_cfg["nvfp4_4over6"] = nvfp4_4over6_weight_scope() + quant_cfg["nvfp4_4over6"] = ( + "weights" if os.getenv("NVTE_NVFP4_4OVER6", "").strip().lower() in ("weights", "all") else "none" + ) quant_cfg["ignore"] = ignore_list quant_cfg.setdefault("kv_cache_scheme", DEFAULT_KV_CACHE_SCHEME) @@ -221,7 +218,9 @@ def _write_hf_quant_config( quant_section["quant_algo"] = "NVFP4" quant_section["kv_cache_quant_algo"] = DEFAULT_KV_CACHE_QUANT_ALGO quant_section["group_size"] = NVFP4_GROUP_SIZE - quant_section["nvfp4_4over6"] = nvfp4_4over6_weight_scope() + quant_section["nvfp4_4over6"] = ( + "weights" if os.getenv("NVTE_NVFP4_4OVER6", "").strip().lower() in ("weights", "all") else "none" + ) quant_section["exclude_modules"] = ignore_list hf_quant_cfg["quantization"] = quant_section From 99dd0893804b20fd56f9ed3f6c352361d3b90a5c Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Sat, 30 May 2026 16:50:49 -0700 Subject: [PATCH 05/37] refactor: inline nvfp4 cleanup guards --- .../megatron_to_hf/processors/quantizer_nvfp4.py | 8 ++------ scripts/run_qwen3_30b_a3b.py | 9 ++++----- tools/convert_hf_to_nvfp4.py | 12 +++--------- 3 files changed, 9 insertions(+), 20 deletions(-) diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py index 2f0469c26dd..be46f57ce73 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py @@ -12,11 +12,6 @@ } -def assert_no_fp4_param_gather(args) -> None: - if args is not None and bool(getattr(args, "fp4_param", False) or getattr(args, "fp4_param_gather", False)): - raise NotImplementedError("fp4-param-gather is unsupported for Miles NVFP4 checkpoint export.") - - def _get_ignore_rules(quantization_config) -> list[str]: ignore_rules = quantization_config.get("ignore", []) or [] if isinstance(ignore_rules, str): @@ -41,7 +36,8 @@ def _is_ignored(name: str, ignore_rules: list[str]) -> bool: def quantize_params_nvfp4(args, megatron_name, converted_named_params, quantization_config): assert quantization_config is not None assert quantization_config.get("quant_algo") == "NVFP4" or quantization_config.get("quant_method") == "nvfp4" - assert_no_fp4_param_gather(args) + if args is not None and bool(getattr(args, "fp4_param", False) or getattr(args, "fp4_param_gather", False)): + raise NotImplementedError("fp4-param-gather is unsupported for Miles NVFP4 checkpoint export.") if getattr(args, "extra_high_precision_layers_megatron", False): for layer_name in getattr(args, "extra_high_precision_layers_megatron", ()): diff --git a/scripts/run_qwen3_30b_a3b.py b/scripts/run_qwen3_30b_a3b.py index d9c9609bcf7..d2a51964308 100644 --- a/scripts/run_qwen3_30b_a3b.py +++ b/scripts/run_qwen3_30b_a3b.py @@ -7,7 +7,6 @@ import miles.utils.external_utils.command_utils as U -BLACKWELL_HARDWARE = ("B200", "B300", "GB200", "GB300") FP4_ENV_MARKERS = ["NVTE", "FLASHINFER", "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH"] @@ -72,18 +71,18 @@ def __post_init__(self): if self.rollout_mxfp8: assert not self.rollout_fp8, "rollout_mxfp8 and rollout_fp8 cannot be enabled at the same time" assert not self.rollout_nvfp4, "rollout_mxfp8 and rollout_nvfp4 cannot be enabled at the same time" - assert self.hardware in BLACKWELL_HARDWARE, "rollout_mxfp8 only supports Blackwell GPUs" + assert self.hardware in ("B200", "B300", "GB200", "GB300"), "rollout_mxfp8 only supports Blackwell GPUs" if self.rollout_nvfp4: assert not self.rollout_fp8, "rollout_nvfp4 and rollout_fp8 cannot be enabled at the same time" - assert self.hardware in BLACKWELL_HARDWARE, "rollout_nvfp4 only supports Blackwell GPUs" + assert self.hardware in ("B200", "B300", "GB200", "GB300"), "rollout_nvfp4 only supports Blackwell GPUs" if self.train_mxfp8: assert not self.train_fp8, "train_mxfp8 and train_fp8 cannot be enabled at the same time" assert not self.train_nvfp4, "train_mxfp8 and train_nvfp4 cannot be enabled at the same time" - assert self.hardware in BLACKWELL_HARDWARE, "train_mxfp8 only supports Blackwell GPUs" + assert self.hardware in ("B200", "B300", "GB200", "GB300"), "train_mxfp8 only supports Blackwell GPUs" assert self.rollout_mxfp8, "train_mxfp8 requires rollout_mxfp8 to be enabled" if self.train_nvfp4: assert not self.train_fp8, "train_nvfp4 and train_fp8 cannot be enabled at the same time" - assert self.hardware in BLACKWELL_HARDWARE, "train_nvfp4 only supports Blackwell GPUs" + assert self.hardware in ("B200", "B300", "GB200", "GB300"), "train_nvfp4 only supports Blackwell GPUs" assert self.rollout_nvfp4, "train_nvfp4 requires rollout_nvfp4 to be enabled" diff --git a/tools/convert_hf_to_nvfp4.py b/tools/convert_hf_to_nvfp4.py index 06121ed5ede..54e0010d452 100644 --- a/tools/convert_hf_to_nvfp4.py +++ b/tools/convert_hf_to_nvfp4.py @@ -199,11 +199,7 @@ def _update_quantization_config(cfg: dict, ignore_list: list[str]) -> None: cfg["quantization_config"] = quant_cfg -def _write_hf_quant_config( - output_path: str, - ignore_list: list[str], - input_path: str, -) -> None: +def _write_hf_quant_config(output_path: str, ignore_list: list[str], input_path: str) -> None: hf_quant_path = os.path.join(input_path, "hf_quant_config.json") if os.path.exists(hf_quant_path): with open(hf_quant_path) as f: @@ -323,10 +319,7 @@ def process_file( if should_quantize(key, tensor, skip_weight_substrings=dynamic_skip_substrings): base, _role = _split_gated_pair_name(key) global_amax = shared_global_amax.get(base) if base else None - qweight, block_scale, weight_scale_2 = quantize_nvfp4( - tensor, - global_amax=global_amax, - ) + qweight, block_scale, weight_scale_2 = quantize_nvfp4(tensor, global_amax=global_amax) q_weights[key] = qweight q_weights[key.replace(".weight", ".weight_scale")] = block_scale q_weights[key.replace(".weight", ".weight_scale_2")] = weight_scale_2 @@ -370,6 +363,7 @@ def convert_nvfp4( *extra_high_precision_layers_hf, *sorted(dynamic_skip_layer_prefixes), ) + shared_global_amax = _collect_shared_global_amax( input_path=input_path, safetensors_files=safetensors_files, From 13b732e6baec5774fff958bf5aeeaa33b502b825 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Sat, 30 May 2026 17:15:18 -0700 Subject: [PATCH 06/37] Minor clean up doc --- tools/convert_hf_to_nvfp4.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tools/convert_hf_to_nvfp4.py b/tools/convert_hf_to_nvfp4.py index 54e0010d452..d829724f30a 100644 --- a/tools/convert_hf_to_nvfp4.py +++ b/tools/convert_hf_to_nvfp4.py @@ -10,9 +10,8 @@ Use --extra-high-precision-layers-hf to keep additional HF weight-name substrings unquantized. -This follows the NVFP4 reference quantization in Transformer Engine, uses the -FlashInfer CUDA quantizer when available, and keeps 4over6 controlled by the -NVTE_NVFP4_* environment variables. +This follows the NVFP4 reference quantization in Transformer Engine and uses +1D block scaling (NVTE_NVFP4_1D_SCALING, group size = 16). """ import argparse From 85abaa6f6a532aee63d948b9a8e265e3a824eea6 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Sat, 30 May 2026 17:17:00 -0700 Subject: [PATCH 07/37] refactor: remove nvfp4 decode alias --- .../megatron_to_hf/processors/quantizer_nvfp4.py | 5 +---- tools/convert_hf_to_nvfp4.py | 5 +---- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py index be46f57ce73..fadc47bceb1 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py @@ -2,7 +2,7 @@ import torch -from miles.utils.nvfp4 import NVFP4_GROUP_SIZE, nvfp4_global_decode_scale_te, nvfp4_quantize_1d +from miles.utils.nvfp4 import NVFP4_GROUP_SIZE, nvfp4_quantize_1d GATED_PAIR_SUFFIXES = { ".gate_proj.weight": "gate", @@ -162,9 +162,6 @@ def _split_gated_pair_name(name: str): return None, None -_nvfp4_global_decode_scale_te = nvfp4_global_decode_scale_te - - def _quantize_nvfp4_1d( weight: torch.Tensor, global_amax: torch.Tensor | None = None, diff --git a/tools/convert_hf_to_nvfp4.py b/tools/convert_hf_to_nvfp4.py index d829724f30a..cd6c4641d82 100644 --- a/tools/convert_hf_to_nvfp4.py +++ b/tools/convert_hf_to_nvfp4.py @@ -26,7 +26,7 @@ import torch from tqdm import tqdm -from miles.utils.nvfp4 import NVFP4_GROUP_SIZE, nvfp4_global_decode_scale_te, nvfp4_quantize_1d +from miles.utils.nvfp4 import NVFP4_GROUP_SIZE, nvfp4_quantize_1d DEFAULT_KV_CACHE_SCHEME = {"dynamic": False, "num_bits": 8, "type": "float"} DEFAULT_KV_CACHE_QUANT_ALGO = "FP8" @@ -100,9 +100,6 @@ def should_quantize( return True -_nvfp4_global_decode_scale_te = nvfp4_global_decode_scale_te - - def _quantize_nvfp4_1d( weight: torch.Tensor, global_amax: torch.Tensor | None = None, From 7ce9115aa0cdaaf81f3804b6423ed5b86db83c0e Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Sat, 30 May 2026 17:21:47 -0700 Subject: [PATCH 08/37] refactor: gate fp4 env forwarding --- scripts/run_qwen3_30b_a3b.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/scripts/run_qwen3_30b_a3b.py b/scripts/run_qwen3_30b_a3b.py index d2a51964308..83b848b91f5 100644 --- a/scripts/run_qwen3_30b_a3b.py +++ b/scripts/run_qwen3_30b_a3b.py @@ -7,11 +7,10 @@ import miles.utils.external_utils.command_utils as U -FP4_ENV_MARKERS = ["NVTE", "FLASHINFER", "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH"] - def fp4_env_vars() -> dict[str, str]: - return {key: value for key, value in os.environ.items() if any(marker in key for marker in FP4_ENV_MARKERS)} + fp4_env_markers = ["NVTE", "FLASHINFER", "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH"] + return {key: value for key, value in os.environ.items() if any(marker in key for marker in fp4_env_markers)} def env_prefix(env_vars: dict[str, str]) -> str: @@ -429,7 +428,8 @@ def execute(args: ScriptArgs): f"{args.extra_args} " ) - misc_env_vars |= fp4_env_vars() + if args.rollout_nvfp4 or args.train_nvfp4: + misc_env_vars |= fp4_env_vars() U.execute_train( train_args=train_args, From a5f3543df1aa7368db9e1c0551374057111457d9 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Sat, 30 May 2026 17:23:41 -0700 Subject: [PATCH 09/37] Minor clean up script --- scripts/run_qwen3_30b_a3b.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/run_qwen3_30b_a3b.py b/scripts/run_qwen3_30b_a3b.py index 83b848b91f5..fe904292bbc 100644 --- a/scripts/run_qwen3_30b_a3b.py +++ b/scripts/run_qwen3_30b_a3b.py @@ -278,7 +278,7 @@ def execute(args: ScriptArgs): "NVTE_NVFP4_DISABLE_RHT": "1", "NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING": "1", "NVTE_NVFP4_ROW_SCALED_ACTIVATION": "1", - "NVTE_BACKWARD_OVERRIDE": "dequantized", + "NVTE_BACKWARD_OVERRIDE": "high_precision", "NVTE_USE_FAST_MATH": "0", } optimizer_args += "--optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer " From 77edb7eb520e279c7700cfa34ea46cc6a5ca2105 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Sat, 30 May 2026 17:25:44 -0700 Subject: [PATCH 10/37] refactor: drop nvfp4 4over6 metadata --- tools/convert_hf_to_nvfp4.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tools/convert_hf_to_nvfp4.py b/tools/convert_hf_to_nvfp4.py index cd6c4641d82..a7735f0d874 100644 --- a/tools/convert_hf_to_nvfp4.py +++ b/tools/convert_hf_to_nvfp4.py @@ -171,9 +171,6 @@ def _update_quantization_config(cfg: dict, ignore_list: list[str]) -> None: quant_cfg["quant_algo"] = "NVFP4" quant_cfg["quant_method"] = "modelopt" quant_cfg["group_size"] = NVFP4_GROUP_SIZE - quant_cfg["nvfp4_4over6"] = ( - "weights" if os.getenv("NVTE_NVFP4_4OVER6", "").strip().lower() in ("weights", "all") else "none" - ) quant_cfg["ignore"] = ignore_list quant_cfg.setdefault("kv_cache_scheme", DEFAULT_KV_CACHE_SCHEME) @@ -210,9 +207,6 @@ def _write_hf_quant_config(output_path: str, ignore_list: list[str], input_path: quant_section["quant_algo"] = "NVFP4" quant_section["kv_cache_quant_algo"] = DEFAULT_KV_CACHE_QUANT_ALGO quant_section["group_size"] = NVFP4_GROUP_SIZE - quant_section["nvfp4_4over6"] = ( - "weights" if os.getenv("NVTE_NVFP4_4OVER6", "").strip().lower() in ("weights", "all") else "none" - ) quant_section["exclude_modules"] = ignore_list hf_quant_cfg["quantization"] = quant_section From 1c8e7c83b20165ddb2e727d3f4795330e2a186f6 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Sat, 30 May 2026 17:27:14 -0700 Subject: [PATCH 11/37] refactor: localize flashinfer env keys --- miles/utils/nvfp4.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/miles/utils/nvfp4.py b/miles/utils/nvfp4.py index 8a94799d76c..fb1b35c4d5a 100644 --- a/miles/utils/nvfp4.py +++ b/miles/utils/nvfp4.py @@ -9,13 +9,6 @@ NVFP4_GROUP_SIZE = 16 logger = logging.getLogger(__name__) -FLASHINFER_NVFP4_ENV_KEYS = ( - "FLASHINFER_NVFP4_4OVER6", - "FLASHINFER_NVFP4_4OVER6_E4M3_USE_256", - "FLASHINFER_NVFP4_4OVER6_ERR_MODE", - "FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH", - "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH", -) def nvfp4_weight_e4m3_max() -> int: @@ -84,7 +77,14 @@ def sync_flashinfer_nvfp4_env_from_nvte() -> dict[str, str]: @contextmanager def flashinfer_nvfp4_env_from_nvte(): - original_env = {key: os.environ.get(key) for key in FLASHINFER_NVFP4_ENV_KEYS} + flashinfer_nvfp4_env_keys = ( + "FLASHINFER_NVFP4_4OVER6", + "FLASHINFER_NVFP4_4OVER6_E4M3_USE_256", + "FLASHINFER_NVFP4_4OVER6_ERR_MODE", + "FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH", + "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH", + ) + original_env = {key: os.environ.get(key) for key in flashinfer_nvfp4_env_keys} try: sync_flashinfer_nvfp4_env_from_nvte() yield From df17e07c5e833274b5d7b7d83cc6f69e73edc125 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Sat, 30 May 2026 17:32:35 -0700 Subject: [PATCH 12/37] test: cover nvfp4 4over6 bitwise path --- tests/fast-gpu/test_nvfp4_quantizer.py | 29 +++++++++----------------- 1 file changed, 10 insertions(+), 19 deletions(-) diff --git a/tests/fast-gpu/test_nvfp4_quantizer.py b/tests/fast-gpu/test_nvfp4_quantizer.py index 1cbe1bb98c8..62cb42a5a6f 100644 --- a/tests/fast-gpu/test_nvfp4_quantizer.py +++ b/tests/fast-gpu/test_nvfp4_quantizer.py @@ -32,23 +32,6 @@ (2048, 7168), (128, 16384), ] -NVFP4_ENV_KEYS = ( - "NVTE_NVFP4_4OVER6", - "NVTE_NVFP4_4OVER6_E4M3_USE_256", - "NVTE_NVFP4_4OVER6_ERR_MODE", - "NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", - "FLASHINFER_NVFP4_4OVER6", - "FLASHINFER_NVFP4_4OVER6_E4M3_USE_256", - "FLASHINFER_NVFP4_4OVER6_ERR_MODE", - "FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH", - "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH", -) - - -@pytest.fixture(autouse=True) -def clean_nvfp4_env(monkeypatch): - for key in NVFP4_ENV_KEYS: - monkeypatch.delenv(key, raising=False) def _make_weight(init_data: str, dtype: torch.dtype, shape: tuple[int, int], device: str) -> torch.Tensor: @@ -90,7 +73,7 @@ def _te_nvfp4_reference_with_global_amax( NVFP4_GROUP_SIZE, 1, pow_2_scales=False, - nvfp4_use_4over6=False, + nvfp4_use_4over6=os.getenv("NVTE_NVFP4_4OVER6", "").strip().lower() in ("weights", "all"), nvfp4_e4m3_max=nvfp4_e4m3_max, nvfp4_4over6_err_mode=os.getenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MAE").strip().upper(), eps=0.0, @@ -180,9 +163,17 @@ def test_nvfp4_hf_should_quantize_respects_extra_high_precision_layers_hf(): @pytest.mark.parametrize("shape", NVFP4_SHAPES) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=str) @pytest.mark.parametrize("init_data", ["random", "boundary", "zeros", "maxes"]) -def test_nvfp4_quantize_matches_te_reference_bitwise(quantize_fn, shape, dtype, init_data): +@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) +def test_nvfp4_quantize_matches_te_reference_bitwise(quantize_fn, shape, dtype, init_data, use_4over6, monkeypatch): device = "cuda" torch.manual_seed(42) + monkeypatch.delenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", raising=False) + monkeypatch.delenv("NVTE_NVFP4_4OVER6_ERR_MODE", raising=False) + monkeypatch.delenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", raising=False) + if use_4over6: + monkeypatch.setenv("NVTE_NVFP4_4OVER6", "all") + else: + monkeypatch.delenv("NVTE_NVFP4_4OVER6", raising=False) weight = _make_weight(init_data, dtype, shape, device) qweight, block_scale, global_scale = quantize_fn(weight) From 9df81cbc40ac171b6e8901963ddb212b3ef64e02 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Sat, 30 May 2026 17:34:37 -0700 Subject: [PATCH 13/37] test: collapse nvfp4 reference helper --- tests/fast-gpu/test_nvfp4_quantizer.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/tests/fast-gpu/test_nvfp4_quantizer.py b/tests/fast-gpu/test_nvfp4_quantizer.py index 62cb42a5a6f..a15f8d981ae 100644 --- a/tests/fast-gpu/test_nvfp4_quantizer.py +++ b/tests/fast-gpu/test_nvfp4_quantizer.py @@ -58,14 +58,6 @@ def _te_nvfp4_reference( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: weight = weight.contiguous() global_amax = torch.max(torch.abs(weight.to(torch.float32))) - return _te_nvfp4_reference_with_global_amax(weight, global_amax) - - -def _te_nvfp4_reference_with_global_amax( - weight: torch.Tensor, - global_amax: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - weight = weight.contiguous() nvfp4_e4m3_max = nvfp4_weight_e4m3_max() qweight, block_scale = NVFP4QuantizerRef._quantize_blockwise_reference( weight, From 31506aeee906870661df4195f370754bfb2e1381 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Sat, 30 May 2026 19:22:51 -0700 Subject: [PATCH 14/37] Clean up NVFP4 bitwise test env setup --- tests/fast-gpu/test_nvfp4_quantizer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/fast-gpu/test_nvfp4_quantizer.py b/tests/fast-gpu/test_nvfp4_quantizer.py index a15f8d981ae..d56336ee2d9 100644 --- a/tests/fast-gpu/test_nvfp4_quantizer.py +++ b/tests/fast-gpu/test_nvfp4_quantizer.py @@ -159,11 +159,9 @@ def test_nvfp4_hf_should_quantize_respects_extra_high_precision_layers_hf(): def test_nvfp4_quantize_matches_te_reference_bitwise(quantize_fn, shape, dtype, init_data, use_4over6, monkeypatch): device = "cuda" torch.manual_seed(42) - monkeypatch.delenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", raising=False) - monkeypatch.delenv("NVTE_NVFP4_4OVER6_ERR_MODE", raising=False) - monkeypatch.delenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", raising=False) if use_4over6: monkeypatch.setenv("NVTE_NVFP4_4OVER6", "all") + monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MSE") else: monkeypatch.delenv("NVTE_NVFP4_4OVER6", raising=False) From c92b62b073d1440d8d235e3ac413da3a4fee989f Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Mon, 1 Jun 2026 15:48:08 -0700 Subject: [PATCH 15/37] Use TE direct NVFP4 quantizer in Miles --- miles/utils/nvfp4.py | 150 ++++++++----------------- tests/fast-gpu/test_nvfp4_quantizer.py | 78 ++++++++++++- tools/convert_hf_to_nvfp4.py | 4 +- 3 files changed, 121 insertions(+), 111 deletions(-) diff --git a/miles/utils/nvfp4.py b/miles/utils/nvfp4.py index fb1b35c4d5a..13b5bc1f433 100644 --- a/miles/utils/nvfp4.py +++ b/miles/utils/nvfp4.py @@ -1,14 +1,12 @@ -import logging import os -from contextlib import contextmanager import torch +from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer FP4_E2M1_MAX = 6.0 FP8_E4M3_MAX = 448.0 NVFP4_GROUP_SIZE = 16 - -logger = logging.getLogger(__name__) +TE_NVFP4_ROW_ALIGNMENT = 16 def nvfp4_weight_e4m3_max() -> int: @@ -53,111 +51,53 @@ def nvfp4_global_decode_scale_te( return torch.div(1.0, nvfp4_global_encode_scale_te(global_amax, nvfp4_e4m3_max)) -def sync_flashinfer_nvfp4_env_from_nvte() -> dict[str, str]: - weight_4over6_enabled = os.getenv("NVTE_NVFP4_4OVER6", "").strip().lower() in ("weights", "all") - flashinfer_env = { - "FLASHINFER_NVFP4_4OVER6": "1" if weight_4over6_enabled else "0", - "FLASHINFER_NVFP4_4OVER6_E4M3_USE_256": ( - "1" - if weight_4over6_enabled - and os.getenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all").strip().lower() in ("weights", "all") - else "0" - ), - "FLASHINFER_NVFP4_4OVER6_ERR_MODE": os.getenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MAE").strip().upper(), - "FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH": ( - "1" - if os.getenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", "0").strip().lower() in ("1", "true", "yes", "on") - else "0" - ), - } - os.environ.update(flashinfer_env) - os.environ.setdefault("TRTLLM_DISABLE_FP4_QUANT_FAST_MATH", "1") - return {**flashinfer_env, "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH": os.environ["TRTLLM_DISABLE_FP4_QUANT_FAST_MATH"]} - - -@contextmanager -def flashinfer_nvfp4_env_from_nvte(): - flashinfer_nvfp4_env_keys = ( - "FLASHINFER_NVFP4_4OVER6", - "FLASHINFER_NVFP4_4OVER6_E4M3_USE_256", - "FLASHINFER_NVFP4_4OVER6_ERR_MODE", - "FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH", - "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH", - ) - original_env = {key: os.environ.get(key) for key in flashinfer_nvfp4_env_keys} - try: - sync_flashinfer_nvfp4_env_from_nvte() - yield - finally: - for key, value in original_env.items(): - if value is None: - os.environ.pop(key, None) - else: - os.environ[key] = value - - -try: - from flashinfer import nvfp4_quantize as _flashinfer_nvfp4_quantize - from flashinfer.tllm_enums import SfLayout -except ImportError: - _flashinfer_nvfp4_quantize = None - SfLayout = None - logger.warning("FlashInfer nvfp4_quantize not available; falling back to TransformerEngine reference.") - - -def _te_nvfp4_quantize_1d( - weight: torch.Tensor, - global_amax: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - try: - from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef - except ImportError: - from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import NVFP4QuantizerRef +def _nvfp4_4over6_enabled() -> bool: + return os.getenv("NVTE_NVFP4_4OVER6", "").strip().lower() in ("weights", "all") - weight_4over6_enabled = os.getenv("NVTE_NVFP4_4OVER6", "").strip().lower() in ("weights", "all") - nvfp4_e4m3_max = nvfp4_weight_e4m3_max() - try: - qweight, block_scale = NVFP4QuantizerRef._quantize_blockwise_reference( - weight, - global_amax, - NVFP4_GROUP_SIZE, - 1, - pow_2_scales=False, - nvfp4_use_4over6=weight_4over6_enabled, - nvfp4_e4m3_max=nvfp4_e4m3_max, - nvfp4_4over6_err_mode=os.getenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MAE").strip().upper(), - eps=0.0, - ) - except TypeError: - if weight_4over6_enabled: - raise - qweight, block_scale = NVFP4QuantizerRef._quantize_blockwise_reference( - weight, - global_amax, - NVFP4_GROUP_SIZE, - 1, - pow_2_scales=False, - eps=0.0, - ) - return qweight, block_scale, nvfp4_global_decode_scale_te(global_amax, nvfp4_e4m3_max) + +def _pad_rows_for_te_quantizer(weight: torch.Tensor) -> torch.Tensor: + pad_rows = (-weight.shape[0]) % TE_NVFP4_ROW_ALIGNMENT + if pad_rows == 0: + return weight + padding = torch.zeros((pad_rows, weight.shape[1]), device=weight.device, dtype=weight.dtype) + return torch.cat((weight, padding), dim=0) def nvfp4_quantize_1d( weight: torch.Tensor, - global_amax: torch.Tensor, + global_amax: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - if _flashinfer_nvfp4_quantize is None or weight.dtype == torch.float32 or not weight.is_cuda: - return _te_nvfp4_quantize_1d(weight, global_amax) - + weight = weight.contiguous() + num_rows, num_cols = weight.shape + row_scaled_nvfp4 = global_amax is not None and global_amax.ndim > 0 and global_amax.numel() == num_rows nvfp4_e4m3_max = nvfp4_weight_e4m3_max() - global_encode_scale = nvfp4_global_encode_scale_te(global_amax, nvfp4_e4m3_max) - with flashinfer_nvfp4_env_from_nvte(): - qweight, block_scale = _flashinfer_nvfp4_quantize( - weight, - global_encode_scale.reshape(1).contiguous(), - sfLayout=SfLayout.layout_linear, - do_shuffle=False, - sf_vec_size=NVFP4_GROUP_SIZE, - backend="cuda", - ) - return qweight, block_scale.view(torch.float8_e4m3fn), nvfp4_global_decode_scale_te(global_amax, nvfp4_e4m3_max) + + quantizer = NVFP4Quantizer( + rowwise=True, + columnwise=False, + with_amax_reduction=False, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=False, + stochastic_rounding=False, + row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=_nvfp4_4over6_enabled(), + nvfp4_e4m3_max=nvfp4_e4m3_max, + nvfp4_4over6_err_mode=os.getenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MAE").strip().upper(), + with_random_sign_mask=False, + ) + + quant_input = weight + if global_amax is not None and not row_scaled_nvfp4: + amax_row = torch.zeros((1, num_cols), device=weight.device, dtype=weight.dtype) + amax_row[0, 0] = global_amax.to(device=weight.device, dtype=weight.dtype).reshape(()) + quant_input = torch.cat((quant_input, amax_row), dim=0) + + quantized = quantizer.quantize(_pad_rows_for_te_quantizer(quant_input)) + qweight = quantized._rowwise_data[:num_rows, : num_cols // 2].contiguous() + block_scale = quantized._rowwise_scale_inv[:num_rows, : num_cols // NVFP4_GROUP_SIZE].contiguous() + if row_scaled_nvfp4: + amax = quantized._amax_rowwise[:num_rows].contiguous() + else: + amax = quantized._amax_rowwise.reshape(-1)[0] + return qweight, block_scale.view(torch.float8_e4m3fn), nvfp4_global_decode_scale_te(amax, nvfp4_e4m3_max) diff --git a/tests/fast-gpu/test_nvfp4_quantizer.py b/tests/fast-gpu/test_nvfp4_quantizer.py index d56336ee2d9..b208f929a11 100644 --- a/tests/fast-gpu/test_nvfp4_quantizer.py +++ b/tests/fast-gpu/test_nvfp4_quantizer.py @@ -55,9 +55,10 @@ def _make_weight(init_data: str, dtype: torch.dtype, shape: tuple[int, int], dev def _te_nvfp4_reference( weight: torch.Tensor, + global_amax: torch.Tensor, + row_scaled_nvfp4: bool, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: weight = weight.contiguous() - global_amax = torch.max(torch.abs(weight.to(torch.float32))) nvfp4_e4m3_max = nvfp4_weight_e4m3_max() qweight, block_scale = NVFP4QuantizerRef._quantize_blockwise_reference( weight, @@ -65,6 +66,7 @@ def _te_nvfp4_reference( NVFP4_GROUP_SIZE, 1, pow_2_scales=False, + row_scaled_nvfp4=row_scaled_nvfp4, nvfp4_use_4over6=os.getenv("NVTE_NVFP4_4OVER6", "").strip().lower() in ("weights", "all"), nvfp4_e4m3_max=nvfp4_e4m3_max, nvfp4_4over6_err_mode=os.getenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MAE").strip().upper(), @@ -73,6 +75,66 @@ def _te_nvfp4_reference( return qweight, block_scale, nvfp4_global_decode_scale_te(global_amax, nvfp4_e4m3_max) +def _global_amax(weight: torch.Tensor, global_scale_mode: str) -> torch.Tensor: + if global_scale_mode == "per_tensor": + return torch.max(torch.abs(weight.to(torch.float32))) + if global_scale_mode == "per_token": + return torch.max(torch.abs(weight.to(torch.float32)), dim=1).values + raise ValueError(f"Unknown global_scale_mode: {global_scale_mode}") + + +def test_nvfp4_quantize_uses_te_direct_rowwise_quantizer(monkeypatch): + import miles.utils.nvfp4 as nvfp4_utils + + calls = [] + + class FakeQuantizedTensor: + def __init__(self, tensor: torch.Tensor): + self._rowwise_data = torch.arange(tensor.shape[0] * (tensor.shape[1] // 2), dtype=torch.uint8).reshape( + tensor.shape[0], tensor.shape[1] // 2 + ) + self._rowwise_scale_inv = torch.arange( + tensor.shape[0] * (tensor.shape[1] // NVFP4_GROUP_SIZE), dtype=torch.uint8 + ).reshape(tensor.shape[0], tensor.shape[1] // NVFP4_GROUP_SIZE) + self._amax_rowwise = torch.tensor([2.0], dtype=torch.float32) + + class FakeQuantizer: + def __init__(self, **kwargs): + calls.append(kwargs) + + def quantize(self, tensor: torch.Tensor) -> FakeQuantizedTensor: + assert tensor.shape == (16, NVFP4_GROUP_SIZE) + return FakeQuantizedTensor(tensor) + + monkeypatch.setattr(nvfp4_utils, "NVFP4Quantizer", FakeQuantizer) + + qweight, block_scale, global_scale = nvfp4_utils.nvfp4_quantize_1d( + torch.ones((3, NVFP4_GROUP_SIZE), dtype=torch.float32), + torch.tensor(2.0, dtype=torch.float32), + ) + + assert calls == [ + { + "rowwise": True, + "columnwise": False, + "with_amax_reduction": False, + "with_rht": False, + "with_post_rht_amax": False, + "with_2d_quantization": False, + "stochastic_rounding": False, + "row_scaled_nvfp4": False, + "nvfp4_use_4over6": False, + "nvfp4_e4m3_max": 448, + "nvfp4_4over6_err_mode": "MAE", + "with_random_sign_mask": False, + } + ] + assert qweight.shape == (3, NVFP4_GROUP_SIZE // 2) + assert block_scale.shape == (3, 1) + assert block_scale.dtype == torch.float8_e4m3fn + torch.testing.assert_close(global_scale, nvfp4_global_decode_scale_te(torch.tensor(2.0), 448), rtol=0, atol=0) + + def test_nvfp4_quantize_params_requires_complete_gated_pair(): weight = torch.randn((4, NVFP4_GROUP_SIZE), dtype=torch.float32) with pytest.raises(ValueError, match="requires gate/up tensors to be quantized together"): @@ -156,7 +218,10 @@ def test_nvfp4_hf_should_quantize_respects_extra_high_precision_layers_hf(): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=str) @pytest.mark.parametrize("init_data", ["random", "boundary", "zeros", "maxes"]) @pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) -def test_nvfp4_quantize_matches_te_reference_bitwise(quantize_fn, shape, dtype, init_data, use_4over6, monkeypatch): +@pytest.mark.parametrize("global_scale_mode", ["per_tensor", "per_token"]) +def test_nvfp4_quantize_matches_te_reference_bitwise( + quantize_fn, shape, dtype, init_data, use_4over6, global_scale_mode, monkeypatch +): device = "cuda" torch.manual_seed(42) if use_4over6: @@ -166,8 +231,13 @@ def test_nvfp4_quantize_matches_te_reference_bitwise(quantize_fn, shape, dtype, monkeypatch.delenv("NVTE_NVFP4_4OVER6", raising=False) weight = _make_weight(init_data, dtype, shape, device) - qweight, block_scale, global_scale = quantize_fn(weight) - qweight_ref, block_scale_ref, global_scale_ref = _te_nvfp4_reference(weight) + global_amax = _global_amax(weight, global_scale_mode) + qweight, block_scale, global_scale = quantize_fn(weight, global_amax=global_amax) + qweight_ref, block_scale_ref, global_scale_ref = _te_nvfp4_reference( + weight, + global_amax, + row_scaled_nvfp4=global_scale_mode == "per_token", + ) torch.testing.assert_close(qweight, qweight_ref, rtol=0, atol=0) torch.testing.assert_close(block_scale.view(torch.uint8), block_scale_ref.view(torch.uint8), rtol=0, atol=0) diff --git a/tools/convert_hf_to_nvfp4.py b/tools/convert_hf_to_nvfp4.py index a7735f0d874..f420dc05c21 100644 --- a/tools/convert_hf_to_nvfp4.py +++ b/tools/convert_hf_to_nvfp4.py @@ -10,8 +10,8 @@ Use --extra-high-precision-layers-hf to keep additional HF weight-name substrings unquantized. -This follows the NVFP4 reference quantization in Transformer Engine and uses -1D block scaling (NVTE_NVFP4_1D_SCALING, group size = 16). +This follows Transformer Engine NVFP4 quantization and uses 1D block scaling +(NVTE_NVFP4_1D_SCALING, group size = 16). """ import argparse From e7c0df0ed1afecf62ef4c103a624612e22e6d6c3 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Mon, 1 Jun 2026 16:07:21 -0700 Subject: [PATCH 16/37] Clean up NVFP4 amax handling Signed-off-by: Ziang Li --- .../processors/quantizer_nvfp4.py | 24 ++++++++--------- miles/utils/nvfp4.py | 16 +++++------ tests/fast-gpu/test_nvfp4_quantizer.py | 27 ++++++++++--------- tools/convert_hf_to_nvfp4.py | 24 ++++++++--------- 4 files changed, 46 insertions(+), 45 deletions(-) diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py index fadc47bceb1..864733f8b81 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py @@ -129,8 +129,10 @@ def _quantize_moe_params(converted_named_params, ignore_rules): quantize_named_params.append((converted_name, param)) continue base, _role = _split_gated_pair_name(converted_name) - global_amax = shared_global_amax.get(base) if base else None - qweight, block_scale, weight_scale_2 = quantize_nvfp4(param, global_amax=global_amax) + qweight, block_scale, weight_scale_2 = quantize_nvfp4( + param, + shared_global_amax=shared_global_amax.get(base) if base else None, + ) quantize_named_params.append((converted_name, qweight)) quantize_named_params.append((converted_name.replace(".weight", ".weight_scale"), block_scale)) quantize_named_params.append((converted_name.replace(".weight", ".weight_scale_2"), weight_scale_2)) @@ -164,7 +166,7 @@ def _split_gated_pair_name(name: str): def _quantize_nvfp4_1d( weight: torch.Tensor, - global_amax: torch.Tensor | None = None, + shared_global_amax: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ NVFP4 1D quantization (tile shape = 1x16). @@ -179,23 +181,21 @@ def _quantize_nvfp4_1d( if n % NVFP4_GROUP_SIZE != 0: raise ValueError(f"NVFP4 requires K divisible by {NVFP4_GROUP_SIZE}, got {n}.") - if global_amax is None: - global_amax = torch.max(torch.abs(weight.to(torch.float32))) - else: - global_amax = global_amax.to(device=weight.device, dtype=torch.float32) + if shared_global_amax is not None: + shared_global_amax = shared_global_amax.to(device=weight.device, dtype=torch.float32) - return nvfp4_quantize_1d(weight, global_amax) + return nvfp4_quantize_1d(weight, shared_global_amax) def quantize_nvfp4( weight: torch.Tensor, - global_amax: torch.Tensor | None = None, + shared_global_amax: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if weight.dim() == 2: - return _quantize_nvfp4_1d(weight, global_amax=global_amax) + return _quantize_nvfp4_1d(weight, shared_global_amax=shared_global_amax) if weight.dim() == 3: - if global_amax is not None: - raise ValueError("global_amax override is only supported for 2D weights.") + if shared_global_amax is not None: + raise ValueError("shared_global_amax override is only supported for 2D weights.") qweights = [] block_scales = [] global_scales = [] diff --git a/miles/utils/nvfp4.py b/miles/utils/nvfp4.py index 13b5bc1f433..a10445ffd8b 100644 --- a/miles/utils/nvfp4.py +++ b/miles/utils/nvfp4.py @@ -65,12 +65,13 @@ def _pad_rows_for_te_quantizer(weight: torch.Tensor) -> torch.Tensor: def nvfp4_quantize_1d( weight: torch.Tensor, - global_amax: torch.Tensor | None = None, + shared_global_amax: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: weight = weight.contiguous() num_rows, num_cols = weight.shape - row_scaled_nvfp4 = global_amax is not None and global_amax.ndim > 0 and global_amax.numel() == num_rows nvfp4_e4m3_max = nvfp4_weight_e4m3_max() + if shared_global_amax is not None and shared_global_amax.numel() != 1: + raise ValueError("shared_global_amax must be a scalar tensor.") quantizer = NVFP4Quantizer( rowwise=True, @@ -80,7 +81,7 @@ def nvfp4_quantize_1d( with_post_rht_amax=False, with_2d_quantization=False, stochastic_rounding=False, - row_scaled_nvfp4=row_scaled_nvfp4, + row_scaled_nvfp4=False, nvfp4_use_4over6=_nvfp4_4over6_enabled(), nvfp4_e4m3_max=nvfp4_e4m3_max, nvfp4_4over6_err_mode=os.getenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MAE").strip().upper(), @@ -88,16 +89,13 @@ def nvfp4_quantize_1d( ) quant_input = weight - if global_amax is not None and not row_scaled_nvfp4: + if shared_global_amax is not None: amax_row = torch.zeros((1, num_cols), device=weight.device, dtype=weight.dtype) - amax_row[0, 0] = global_amax.to(device=weight.device, dtype=weight.dtype).reshape(()) + amax_row[0, 0] = shared_global_amax.to(device=weight.device, dtype=weight.dtype).reshape(()) quant_input = torch.cat((quant_input, amax_row), dim=0) quantized = quantizer.quantize(_pad_rows_for_te_quantizer(quant_input)) qweight = quantized._rowwise_data[:num_rows, : num_cols // 2].contiguous() block_scale = quantized._rowwise_scale_inv[:num_rows, : num_cols // NVFP4_GROUP_SIZE].contiguous() - if row_scaled_nvfp4: - amax = quantized._amax_rowwise[:num_rows].contiguous() - else: - amax = quantized._amax_rowwise.reshape(-1)[0] + amax = quantized._amax_rowwise.reshape(-1)[0] return qweight, block_scale.view(torch.float8_e4m3fn), nvfp4_global_decode_scale_te(amax, nvfp4_e4m3_max) diff --git a/tests/fast-gpu/test_nvfp4_quantizer.py b/tests/fast-gpu/test_nvfp4_quantizer.py index b208f929a11..fcaf19fde8a 100644 --- a/tests/fast-gpu/test_nvfp4_quantizer.py +++ b/tests/fast-gpu/test_nvfp4_quantizer.py @@ -75,12 +75,12 @@ def _te_nvfp4_reference( return qweight, block_scale, nvfp4_global_decode_scale_te(global_amax, nvfp4_e4m3_max) -def _global_amax(weight: torch.Tensor, global_scale_mode: str) -> torch.Tensor: - if global_scale_mode == "per_tensor": +def _shared_scalar_amax(weight: torch.Tensor, shared_amax_mode: str) -> torch.Tensor | None: + if shared_amax_mode == "te_generated": + return None + if shared_amax_mode == "shared_scalar": return torch.max(torch.abs(weight.to(torch.float32))) - if global_scale_mode == "per_token": - return torch.max(torch.abs(weight.to(torch.float32)), dim=1).values - raise ValueError(f"Unknown global_scale_mode: {global_scale_mode}") + raise ValueError(f"Unknown shared_amax_mode: {shared_amax_mode}") def test_nvfp4_quantize_uses_te_direct_rowwise_quantizer(monkeypatch): @@ -110,7 +110,7 @@ def quantize(self, tensor: torch.Tensor) -> FakeQuantizedTensor: qweight, block_scale, global_scale = nvfp4_utils.nvfp4_quantize_1d( torch.ones((3, NVFP4_GROUP_SIZE), dtype=torch.float32), - torch.tensor(2.0, dtype=torch.float32), + shared_global_amax=torch.tensor(2.0, dtype=torch.float32), ) assert calls == [ @@ -218,9 +218,9 @@ def test_nvfp4_hf_should_quantize_respects_extra_high_precision_layers_hf(): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=str) @pytest.mark.parametrize("init_data", ["random", "boundary", "zeros", "maxes"]) @pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) -@pytest.mark.parametrize("global_scale_mode", ["per_tensor", "per_token"]) +@pytest.mark.parametrize("shared_amax_mode", ["te_generated", "shared_scalar"]) def test_nvfp4_quantize_matches_te_reference_bitwise( - quantize_fn, shape, dtype, init_data, use_4over6, global_scale_mode, monkeypatch + quantize_fn, shape, dtype, init_data, use_4over6, shared_amax_mode, monkeypatch ): device = "cuda" torch.manual_seed(42) @@ -231,12 +231,15 @@ def test_nvfp4_quantize_matches_te_reference_bitwise( monkeypatch.delenv("NVTE_NVFP4_4OVER6", raising=False) weight = _make_weight(init_data, dtype, shape, device) - global_amax = _global_amax(weight, global_scale_mode) - qweight, block_scale, global_scale = quantize_fn(weight, global_amax=global_amax) + shared_global_amax = _shared_scalar_amax(weight, shared_amax_mode) + reference_amax = ( + shared_global_amax if shared_global_amax is not None else torch.max(torch.abs(weight.to(torch.float32))) + ) + qweight, block_scale, global_scale = quantize_fn(weight, shared_global_amax=shared_global_amax) qweight_ref, block_scale_ref, global_scale_ref = _te_nvfp4_reference( weight, - global_amax, - row_scaled_nvfp4=global_scale_mode == "per_token", + reference_amax, + row_scaled_nvfp4=False, ) torch.testing.assert_close(qweight, qweight_ref, rtol=0, atol=0) diff --git a/tools/convert_hf_to_nvfp4.py b/tools/convert_hf_to_nvfp4.py index f420dc05c21..9469a6c7014 100644 --- a/tools/convert_hf_to_nvfp4.py +++ b/tools/convert_hf_to_nvfp4.py @@ -102,7 +102,7 @@ def should_quantize( def _quantize_nvfp4_1d( weight: torch.Tensor, - global_amax: torch.Tensor | None = None, + shared_global_amax: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ NVFP4 1D quantization (tile shape = 1x16). @@ -117,23 +117,21 @@ def _quantize_nvfp4_1d( if n % NVFP4_GROUP_SIZE != 0: raise ValueError(f"NVFP4 requires K divisible by {NVFP4_GROUP_SIZE}, got {n}.") - if global_amax is None: - global_amax = torch.max(torch.abs(weight.to(torch.float32))) - else: - global_amax = global_amax.to(device=weight.device, dtype=torch.float32) + if shared_global_amax is not None: + shared_global_amax = shared_global_amax.to(device=weight.device, dtype=torch.float32) - return nvfp4_quantize_1d(weight, global_amax) + return nvfp4_quantize_1d(weight, shared_global_amax) def quantize_nvfp4( weight: torch.Tensor, - global_amax: torch.Tensor | None = None, + shared_global_amax: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if weight.dim() == 2: - return _quantize_nvfp4_1d(weight, global_amax=global_amax) + return _quantize_nvfp4_1d(weight, shared_global_amax=shared_global_amax) if weight.dim() == 3: - if global_amax is not None: - raise ValueError("global_amax override is only supported for 2D weights.") + if shared_global_amax is not None: + raise ValueError("shared_global_amax override is only supported for 2D weights.") qweights = [] block_scales = [] global_scales = [] @@ -308,8 +306,10 @@ def process_file( tensor = f.get_tensor(key) if should_quantize(key, tensor, skip_weight_substrings=dynamic_skip_substrings): base, _role = _split_gated_pair_name(key) - global_amax = shared_global_amax.get(base) if base else None - qweight, block_scale, weight_scale_2 = quantize_nvfp4(tensor, global_amax=global_amax) + qweight, block_scale, weight_scale_2 = quantize_nvfp4( + tensor, + shared_global_amax=shared_global_amax.get(base) if base else None, + ) q_weights[key] = qweight q_weights[key.replace(".weight", ".weight_scale")] = block_scale q_weights[key.replace(".weight", ".weight_scale_2")] = weight_scale_2 From 0d308ae3ded68883f49bf6b3490cd60fcac26026 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Mon, 1 Jun 2026 16:17:16 -0700 Subject: [PATCH 17/37] Use paired NVFP4 quantization for gated weights Signed-off-by: Ziang Li --- .../processors/quantizer_nvfp4.py | 34 ++-- miles/utils/nvfp4.py | 42 +++-- tests/fast-gpu/test_nvfp4_quantizer.py | 167 ++++++++++++++++-- tools/convert_hf_to_nvfp4.py | 125 ++++++++----- 4 files changed, 278 insertions(+), 90 deletions(-) diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py index 864733f8b81..34891cb5569 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py @@ -2,7 +2,7 @@ import torch -from miles.utils.nvfp4 import NVFP4_GROUP_SIZE, nvfp4_quantize_1d +from miles.utils.nvfp4 import NVFP4_GROUP_SIZE, nvfp4_quantize_1d, nvfp4_quantize_1d_pair GATED_PAIR_SUFFIXES = { ".gate_proj.weight": "gate", @@ -97,7 +97,6 @@ def quantize_params_nvfp4(args, megatron_name, converted_named_params, quantizat def _quantize_moe_params(converted_named_params, ignore_rules): - shared_global_amax = {} gated_candidates = {} for converted_name, param in converted_named_params: base, role = _split_gated_pair_name(converted_name) @@ -110,8 +109,9 @@ def _quantize_moe_params(converted_named_params, ignore_rules): f"NVFP4 requires a single complete gate/up pair per conversion batch; " f"found duplicate {role} tensor for {base}." ) - roles[role] = param + roles[role] = (converted_name, param) + paired_outputs = {} for base, roles in gated_candidates.items(): if set(roles) != {"gate", "up"}: present = ", ".join(sorted(roles)) @@ -119,20 +119,21 @@ def _quantize_moe_params(converted_named_params, ignore_rules): f"NVFP4 requires gate/up tensors to be quantized together so they can share " f"one global amax; found only {{{present}}} for {base}." ) - gate_amax = roles["gate"].abs().max().to(torch.float32) - up_amax = roles["up"].abs().max().to(torch.float32) - shared_global_amax[base] = torch.max(gate_amax, up_amax) + gate_name, gate_weight = roles["gate"] + up_name, up_weight = roles["up"] + gate_output, up_output = nvfp4_quantize_1d_pair(gate_weight, up_weight) + paired_outputs[gate_name] = gate_output + paired_outputs[up_name] = up_output quantize_named_params = [] for converted_name, param in converted_named_params: if not _should_quantize_param(converted_name, param, ignore_rules): quantize_named_params.append((converted_name, param)) continue - base, _role = _split_gated_pair_name(converted_name) - qweight, block_scale, weight_scale_2 = quantize_nvfp4( - param, - shared_global_amax=shared_global_amax.get(base) if base else None, - ) + if converted_name in paired_outputs: + qweight, block_scale, weight_scale_2 = paired_outputs[converted_name] + else: + qweight, block_scale, weight_scale_2 = quantize_nvfp4(param) quantize_named_params.append((converted_name, qweight)) quantize_named_params.append((converted_name.replace(".weight", ".weight_scale"), block_scale)) quantize_named_params.append((converted_name.replace(".weight", ".weight_scale_2"), weight_scale_2)) @@ -166,7 +167,6 @@ def _split_gated_pair_name(name: str): def _quantize_nvfp4_1d( weight: torch.Tensor, - shared_global_amax: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ NVFP4 1D quantization (tile shape = 1x16). @@ -181,21 +181,15 @@ def _quantize_nvfp4_1d( if n % NVFP4_GROUP_SIZE != 0: raise ValueError(f"NVFP4 requires K divisible by {NVFP4_GROUP_SIZE}, got {n}.") - if shared_global_amax is not None: - shared_global_amax = shared_global_amax.to(device=weight.device, dtype=torch.float32) - - return nvfp4_quantize_1d(weight, shared_global_amax) + return nvfp4_quantize_1d(weight) def quantize_nvfp4( weight: torch.Tensor, - shared_global_amax: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if weight.dim() == 2: - return _quantize_nvfp4_1d(weight, shared_global_amax=shared_global_amax) + return _quantize_nvfp4_1d(weight) if weight.dim() == 3: - if shared_global_amax is not None: - raise ValueError("shared_global_amax override is only supported for 2D weights.") qweights = [] block_scales = [] global_scales = [] diff --git a/miles/utils/nvfp4.py b/miles/utils/nvfp4.py index a10445ffd8b..1d8429957e2 100644 --- a/miles/utils/nvfp4.py +++ b/miles/utils/nvfp4.py @@ -65,13 +65,10 @@ def _pad_rows_for_te_quantizer(weight: torch.Tensor) -> torch.Tensor: def nvfp4_quantize_1d( weight: torch.Tensor, - shared_global_amax: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: weight = weight.contiguous() num_rows, num_cols = weight.shape nvfp4_e4m3_max = nvfp4_weight_e4m3_max() - if shared_global_amax is not None and shared_global_amax.numel() != 1: - raise ValueError("shared_global_amax must be a scalar tensor.") quantizer = NVFP4Quantizer( rowwise=True, @@ -88,14 +85,39 @@ def nvfp4_quantize_1d( with_random_sign_mask=False, ) - quant_input = weight - if shared_global_amax is not None: - amax_row = torch.zeros((1, num_cols), device=weight.device, dtype=weight.dtype) - amax_row[0, 0] = shared_global_amax.to(device=weight.device, dtype=weight.dtype).reshape(()) - quant_input = torch.cat((quant_input, amax_row), dim=0) - - quantized = quantizer.quantize(_pad_rows_for_te_quantizer(quant_input)) + quantized = quantizer.quantize(_pad_rows_for_te_quantizer(weight)) qweight = quantized._rowwise_data[:num_rows, : num_cols // 2].contiguous() block_scale = quantized._rowwise_scale_inv[:num_rows, : num_cols // NVFP4_GROUP_SIZE].contiguous() amax = quantized._amax_rowwise.reshape(-1)[0] return qweight, block_scale.view(torch.float8_e4m3fn), nvfp4_global_decode_scale_te(amax, nvfp4_e4m3_max) + + +def nvfp4_quantize_1d_pair( + first: torch.Tensor, + second: torch.Tensor, +) -> tuple[ + tuple[torch.Tensor, torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor, torch.Tensor], +]: + if first.dim() != 2 or second.dim() != 2: + raise ValueError("nvfp4_quantize_1d_pair expects two 2D tensors.") + if first.shape[1] != second.shape[1]: + raise ValueError( + f"NVFP4 paired quantization requires matching K dimensions, got {first.shape[1]} and {second.shape[1]}." + ) + + first_rows = first.shape[0] + combined_qweight, combined_block_scale, global_scale = nvfp4_quantize_1d( + torch.cat((first.contiguous(), second.contiguous()), dim=0) + ) + first_result = ( + combined_qweight[:first_rows].contiguous(), + combined_block_scale[:first_rows].contiguous(), + global_scale, + ) + second_result = ( + combined_qweight[first_rows:].contiguous(), + combined_block_scale[first_rows:].contiguous(), + global_scale, + ) + return first_result, second_result diff --git a/tests/fast-gpu/test_nvfp4_quantizer.py b/tests/fast-gpu/test_nvfp4_quantizer.py index fcaf19fde8a..2db157389d6 100644 --- a/tests/fast-gpu/test_nvfp4_quantizer.py +++ b/tests/fast-gpu/test_nvfp4_quantizer.py @@ -6,7 +6,10 @@ import os import pytest +import safetensors +import safetensors.torch import torch +from tools.convert_hf_to_nvfp4 import convert_nvfp4 from tools.convert_hf_to_nvfp4 import quantize_nvfp4 as tool_quantize_nvfp4 from tools.convert_hf_to_nvfp4 import should_quantize as tool_should_quantize_nvfp4 from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef @@ -15,7 +18,12 @@ quantize_nvfp4 as processor_quantize_nvfp4, ) from miles.backends.megatron_utils.megatron_to_hf.processors.quantizer_nvfp4 import quantize_params_nvfp4 -from miles.utils.nvfp4 import NVFP4_GROUP_SIZE, nvfp4_global_decode_scale_te, nvfp4_weight_e4m3_max +from miles.utils.nvfp4 import ( + NVFP4_GROUP_SIZE, + nvfp4_global_decode_scale_te, + nvfp4_quantize_1d_pair, + nvfp4_weight_e4m3_max, +) NVFP4_SHAPES = [ (1, 64), @@ -75,14 +83,6 @@ def _te_nvfp4_reference( return qweight, block_scale, nvfp4_global_decode_scale_te(global_amax, nvfp4_e4m3_max) -def _shared_scalar_amax(weight: torch.Tensor, shared_amax_mode: str) -> torch.Tensor | None: - if shared_amax_mode == "te_generated": - return None - if shared_amax_mode == "shared_scalar": - return torch.max(torch.abs(weight.to(torch.float32))) - raise ValueError(f"Unknown shared_amax_mode: {shared_amax_mode}") - - def test_nvfp4_quantize_uses_te_direct_rowwise_quantizer(monkeypatch): import miles.utils.nvfp4 as nvfp4_utils @@ -110,7 +110,6 @@ def quantize(self, tensor: torch.Tensor) -> FakeQuantizedTensor: qweight, block_scale, global_scale = nvfp4_utils.nvfp4_quantize_1d( torch.ones((3, NVFP4_GROUP_SIZE), dtype=torch.float32), - shared_global_amax=torch.tensor(2.0, dtype=torch.float32), ) assert calls == [ @@ -135,6 +134,51 @@ def quantize(self, tensor: torch.Tensor) -> FakeQuantizedTensor: torch.testing.assert_close(global_scale, nvfp4_global_decode_scale_te(torch.tensor(2.0), 448), rtol=0, atol=0) +def test_nvfp4_quantize_pair_concats_before_te_quantizer(monkeypatch): + import miles.utils.nvfp4 as nvfp4_utils + + quantized_input = None + + class FakeQuantizedTensor: + def __init__(self, tensor: torch.Tensor): + self._rowwise_data = torch.arange(tensor.shape[0] * (tensor.shape[1] // 2), dtype=torch.uint8).reshape( + tensor.shape[0], tensor.shape[1] // 2 + ) + self._rowwise_scale_inv = torch.arange( + tensor.shape[0] * (tensor.shape[1] // NVFP4_GROUP_SIZE), dtype=torch.uint8 + ).reshape(tensor.shape[0], tensor.shape[1] // NVFP4_GROUP_SIZE) + self._amax_rowwise = torch.tensor([7.0], dtype=torch.float32) + + class FakeQuantizer: + def __init__(self, **kwargs): + pass + + def quantize(self, tensor: torch.Tensor) -> FakeQuantizedTensor: + nonlocal quantized_input + quantized_input = tensor.clone() + return FakeQuantizedTensor(tensor) + + monkeypatch.setattr(nvfp4_utils, "NVFP4Quantizer", FakeQuantizer) + + gate = torch.ones((3, NVFP4_GROUP_SIZE), dtype=torch.float32) + up = torch.full((5, NVFP4_GROUP_SIZE), 2.0, dtype=torch.float32) + (gate_qweight, gate_block_scale, gate_global_scale), ( + up_qweight, + up_block_scale, + up_global_scale, + ) = nvfp4_utils.nvfp4_quantize_1d_pair(gate, up) + + assert quantized_input.shape == (16, NVFP4_GROUP_SIZE) + torch.testing.assert_close(quantized_input[:3], gate) + torch.testing.assert_close(quantized_input[3:8], up) + torch.testing.assert_close(quantized_input[8:], torch.zeros_like(quantized_input[8:])) + assert gate_qweight.shape == (3, NVFP4_GROUP_SIZE // 2) + assert up_qweight.shape == (5, NVFP4_GROUP_SIZE // 2) + assert gate_block_scale.shape == (3, 1) + assert up_block_scale.shape == (5, 1) + torch.testing.assert_close(gate_global_scale, up_global_scale, rtol=0, atol=0) + + def test_nvfp4_quantize_params_requires_complete_gated_pair(): weight = torch.randn((4, NVFP4_GROUP_SIZE), dtype=torch.float32) with pytest.raises(ValueError, match="requires gate/up tensors to be quantized together"): @@ -209,6 +253,59 @@ def test_nvfp4_hf_should_quantize_respects_extra_high_precision_layers_hf(): ) +def test_nvfp4_hf_converter_quantizes_cross_shard_gated_pair_together(tmp_path, monkeypatch): + monkeypatch.delenv("NVTE_NVFP4_4OVER6", raising=False) + model_dir = tmp_path / "model" + save_dir = tmp_path / "converted" + model_dir.mkdir() + (model_dir / "config.json").write_text('{"num_hidden_layers": 1}') + + gate_key = "model.layers.0.mlp.experts.0.gate_proj.weight" + up_key = "model.layers.0.mlp.experts.0.up_proj.weight" + gate = torch.randn((3, 128), dtype=torch.bfloat16) + up = torch.randn((5, 128), dtype=torch.bfloat16) + safetensors.torch.save_file({gate_key: gate}, model_dir / "gate.safetensors", metadata={"format": "pt"}) + safetensors.torch.save_file({up_key: up}, model_dir / "up.safetensors", metadata={"format": "pt"}) + + convert_nvfp4(str(model_dir), str(save_dir), device="cuda") + + (gate_qweight, gate_block_scale, gate_global_scale), ( + up_qweight, + up_block_scale, + up_global_scale, + ) = nvfp4_quantize_1d_pair(gate.cuda(), up.cuda()) + + with safetensors.safe_open(save_dir / "gate.safetensors", framework="pt", device="cuda") as f: + torch.testing.assert_close(f.get_tensor(gate_key), gate_qweight, rtol=0, atol=0) + torch.testing.assert_close( + f.get_tensor(gate_key.replace(".weight", ".weight_scale")).view(torch.uint8), + gate_block_scale.view(torch.uint8), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + f.get_tensor(gate_key.replace(".weight", ".weight_scale_2")), + gate_global_scale, + rtol=0, + atol=0, + ) + + with safetensors.safe_open(save_dir / "up.safetensors", framework="pt", device="cuda") as f: + torch.testing.assert_close(f.get_tensor(up_key), up_qweight, rtol=0, atol=0) + torch.testing.assert_close( + f.get_tensor(up_key.replace(".weight", ".weight_scale")).view(torch.uint8), + up_block_scale.view(torch.uint8), + rtol=0, + atol=0, + ) + torch.testing.assert_close( + f.get_tensor(up_key.replace(".weight", ".weight_scale_2")), + up_global_scale, + rtol=0, + atol=0, + ) + + @pytest.mark.parametrize( "quantize_fn", [processor_quantize_nvfp4, tool_quantize_nvfp4], @@ -218,10 +315,7 @@ def test_nvfp4_hf_should_quantize_respects_extra_high_precision_layers_hf(): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=str) @pytest.mark.parametrize("init_data", ["random", "boundary", "zeros", "maxes"]) @pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) -@pytest.mark.parametrize("shared_amax_mode", ["te_generated", "shared_scalar"]) -def test_nvfp4_quantize_matches_te_reference_bitwise( - quantize_fn, shape, dtype, init_data, use_4over6, shared_amax_mode, monkeypatch -): +def test_nvfp4_quantize_matches_te_reference_bitwise(quantize_fn, shape, dtype, init_data, use_4over6, monkeypatch): device = "cuda" torch.manual_seed(42) if use_4over6: @@ -231,11 +325,8 @@ def test_nvfp4_quantize_matches_te_reference_bitwise( monkeypatch.delenv("NVTE_NVFP4_4OVER6", raising=False) weight = _make_weight(init_data, dtype, shape, device) - shared_global_amax = _shared_scalar_amax(weight, shared_amax_mode) - reference_amax = ( - shared_global_amax if shared_global_amax is not None else torch.max(torch.abs(weight.to(torch.float32))) - ) - qweight, block_scale, global_scale = quantize_fn(weight, shared_global_amax=shared_global_amax) + reference_amax = torch.max(torch.abs(weight.to(torch.float32))) + qweight, block_scale, global_scale = quantize_fn(weight) qweight_ref, block_scale_ref, global_scale_ref = _te_nvfp4_reference( weight, reference_amax, @@ -247,6 +338,44 @@ def test_nvfp4_quantize_matches_te_reference_bitwise( torch.testing.assert_close(global_scale, global_scale_ref, rtol=0, atol=0) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=str) +@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) +def test_nvfp4_quantize_pair_matches_te_reference_bitwise(dtype, use_4over6, monkeypatch): + device = "cuda" + torch.manual_seed(42) + if use_4over6: + monkeypatch.setenv("NVTE_NVFP4_4OVER6", "all") + monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MSE") + else: + monkeypatch.delenv("NVTE_NVFP4_4OVER6", raising=False) + + gate = _make_weight("random", dtype, (3, 128), device) + up = _make_weight("boundary", dtype, (5, 128), device) + (gate_qweight, gate_block_scale, gate_global_scale), ( + up_qweight, + up_block_scale, + up_global_scale, + ) = nvfp4_quantize_1d_pair(gate, up) + + combined = torch.cat((gate, up), dim=0) + qweight_ref, block_scale_ref, global_scale_ref = _te_nvfp4_reference( + combined, + torch.max(torch.abs(combined.to(torch.float32))), + row_scaled_nvfp4=False, + ) + + torch.testing.assert_close(gate_qweight, qweight_ref[: gate.shape[0]], rtol=0, atol=0) + torch.testing.assert_close(up_qweight, qweight_ref[gate.shape[0] :], rtol=0, atol=0) + torch.testing.assert_close( + gate_block_scale.view(torch.uint8), block_scale_ref[: gate.shape[0]].view(torch.uint8), rtol=0, atol=0 + ) + torch.testing.assert_close( + up_block_scale.view(torch.uint8), block_scale_ref[gate.shape[0] :].view(torch.uint8), rtol=0, atol=0 + ) + torch.testing.assert_close(gate_global_scale, global_scale_ref, rtol=0, atol=0) + torch.testing.assert_close(up_global_scale, global_scale_ref, rtol=0, atol=0) + + if __name__ == "__main__": import sys diff --git a/tools/convert_hf_to_nvfp4.py b/tools/convert_hf_to_nvfp4.py index 9469a6c7014..ac8289581cd 100644 --- a/tools/convert_hf_to_nvfp4.py +++ b/tools/convert_hf_to_nvfp4.py @@ -26,7 +26,7 @@ import torch from tqdm import tqdm -from miles.utils.nvfp4 import NVFP4_GROUP_SIZE, nvfp4_quantize_1d +from miles.utils.nvfp4 import NVFP4_GROUP_SIZE, nvfp4_quantize_1d, nvfp4_quantize_1d_pair DEFAULT_KV_CACHE_SCHEME = {"dynamic": False, "num_bits": 8, "type": "float"} DEFAULT_KV_CACHE_QUANT_ALGO = "FP8" @@ -102,7 +102,6 @@ def should_quantize( def _quantize_nvfp4_1d( weight: torch.Tensor, - shared_global_amax: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ NVFP4 1D quantization (tile shape = 1x16). @@ -117,21 +116,15 @@ def _quantize_nvfp4_1d( if n % NVFP4_GROUP_SIZE != 0: raise ValueError(f"NVFP4 requires K divisible by {NVFP4_GROUP_SIZE}, got {n}.") - if shared_global_amax is not None: - shared_global_amax = shared_global_amax.to(device=weight.device, dtype=torch.float32) - - return nvfp4_quantize_1d(weight, shared_global_amax) + return nvfp4_quantize_1d(weight) def quantize_nvfp4( weight: torch.Tensor, - shared_global_amax: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if weight.dim() == 2: - return _quantize_nvfp4_1d(weight, shared_global_amax=shared_global_amax) + return _quantize_nvfp4_1d(weight) if weight.dim() == 3: - if shared_global_amax is not None: - raise ValueError("shared_global_amax override is only supported for 2D weights.") qweights = [] block_scales = [] global_scales = [] @@ -235,16 +228,14 @@ def _split_gated_pair_name(name: str) -> tuple[str | None, str | None]: return None, None -def _collect_shared_global_amax( +def _collect_gated_pair_locations( *, input_path: str, safetensors_files: list[str], device: str, skip_weight_substrings: tuple[str, ...], -) -> dict[str, torch.Tensor]: - """Collect shared gate/up amax across all shards to keep w1/w3 scales equal.""" - gate_amax: dict[str, torch.Tensor] = {} - up_amax: dict[str, torch.Tensor] = {} +) -> dict[str, dict[str, tuple[str, str]]]: + gated_pairs: dict[str, dict[str, tuple[str, str]]] = {} for filename in safetensors_files: with safetensors.safe_open(os.path.join(input_path, filename), framework="pt", device=device) as f: for key in f.keys(): @@ -254,20 +245,33 @@ def _collect_shared_global_amax( base, role = _split_gated_pair_name(key) if base is None or role is None: continue - amax = tensor.abs().max().to(torch.float32) - if role == "gate": - prev = gate_amax.get(base) - gate_amax[base] = amax if prev is None else torch.max(prev, amax) - elif role == "up": - prev = up_amax.get(base) - up_amax[base] = amax if prev is None else torch.max(prev, amax) - else: - continue + roles = gated_pairs.setdefault(base, {}) + if role in roles: + raise ValueError( + f"NVFP4 requires a single complete gate/up pair per converted checkpoint; " + f"found duplicate {role} tensor for {base}." + ) + roles[role] = (filename, key) + return {base: roles for base, roles in gated_pairs.items() if set(roles) == {"gate", "up"}} + + +def _nvfp4_quantized_entries( + key: str, + qweight: torch.Tensor, + block_scale: torch.Tensor, + weight_scale_2: torch.Tensor, +) -> dict[str, torch.Tensor]: + return { + key: qweight, + key.replace(".weight", ".weight_scale"): block_scale, + key.replace(".weight", ".weight_scale_2"): weight_scale_2, + key.replace(".weight", ".input_scale"): torch.ones_like(weight_scale_2, dtype=torch.float32), + } + - shared_global_amax: dict[str, torch.Tensor] = {} - for base in gate_amax.keys() & up_amax.keys(): - shared_global_amax[base] = torch.max(gate_amax[base], up_amax[base]) - return shared_global_amax +def _load_safetensors_tensor(input_path: str, filename: str, key: str, device: str) -> torch.Tensor: + with safetensors.safe_open(os.path.join(input_path, filename), framework="pt", device=device) as f: + return f.get_tensor(key) def process_file( @@ -280,7 +284,9 @@ def process_file( num_layers_at_start_in_bf16: int, num_layers_at_end_in_bf16: int, extra_high_precision_layers_hf: tuple[str, ...], - shared_global_amax: dict[str, torch.Tensor], + gated_pair_locations: dict[str, dict[str, tuple[str, str]]], + processed_gated_pairs: set[str], + deferred_quantized_entries: dict[str, dict[str, dict[str, torch.Tensor]]], ) -> None: if not filename.endswith(".safetensors"): return @@ -303,19 +309,46 @@ def process_file( with safetensors.safe_open(os.path.join(input_path, filename), framework="pt", device=device) as f: for key in f.keys(): + if key in q_weights: + continue + deferred_entries = deferred_quantized_entries.get(filename, {}).pop(key, None) + if deferred_entries is not None: + q_weights.update(deferred_entries) + continue + tensor = f.get_tensor(key) if should_quantize(key, tensor, skip_weight_substrings=dynamic_skip_substrings): base, _role = _split_gated_pair_name(key) - qweight, block_scale, weight_scale_2 = quantize_nvfp4( - tensor, - shared_global_amax=shared_global_amax.get(base) if base else None, - ) - q_weights[key] = qweight - q_weights[key.replace(".weight", ".weight_scale")] = block_scale - q_weights[key.replace(".weight", ".weight_scale_2")] = weight_scale_2 - q_weights[key.replace(".weight", ".input_scale")] = torch.ones_like( - weight_scale_2, dtype=torch.float32 - ) + if base in gated_pair_locations: + if base in processed_gated_pairs: + raise ValueError(f"Missing deferred NVFP4 output for already processed gated pair {base}.") + + pair = gated_pair_locations[base] + gate_filename, gate_key = pair["gate"] + up_filename, up_key = pair["up"] + gate_weight = ( + tensor + if key == gate_key + else _load_safetensors_tensor(input_path, gate_filename, gate_key, device) + ) + up_weight = ( + tensor if key == up_key else _load_safetensors_tensor(input_path, up_filename, up_key, device) + ) + gate_output, up_output = nvfp4_quantize_1d_pair(gate_weight, up_weight) + for target_filename, target_key, output in ( + (gate_filename, gate_key, gate_output), + (up_filename, up_key, up_output), + ): + entries = _nvfp4_quantized_entries(target_key, *output) + if target_filename == filename: + q_weights.update(entries) + else: + deferred_quantized_entries.setdefault(target_filename, {})[target_key] = entries + processed_gated_pairs.add(base) + continue + + qweight, block_scale, weight_scale_2 = quantize_nvfp4(tensor) + q_weights.update(_nvfp4_quantized_entries(key, qweight, block_scale, weight_scale_2)) else: if key.endswith(".weight"): modules_to_not_convert.append(key.replace(".weight", "")) @@ -354,12 +387,14 @@ def convert_nvfp4( *sorted(dynamic_skip_layer_prefixes), ) - shared_global_amax = _collect_shared_global_amax( + gated_pair_locations = _collect_gated_pair_locations( input_path=input_path, safetensors_files=safetensors_files, device=device, skip_weight_substrings=dynamic_skip_substrings, ) + processed_gated_pairs: set[str] = set() + deferred_quantized_entries: dict[str, dict[str, dict[str, torch.Tensor]]] = {} result_collector = ConversionResult() for filename in tqdm(safetensors_files, desc="Processing files"): process_file( @@ -372,12 +407,20 @@ def convert_nvfp4( num_layers_at_start_in_bf16, num_layers_at_end_in_bf16, extra_high_precision_layers_hf, - shared_global_amax, + gated_pair_locations, + processed_gated_pairs, + deferred_quantized_entries, ) gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() + remaining_deferred_entries = { + filename: sorted(entries) for filename, entries in deferred_quantized_entries.items() if entries + } + if remaining_deferred_entries: + raise RuntimeError(f"Unwritten deferred NVFP4 gated-pair entries: {remaining_deferred_entries}") + ignore_list = _augment_ignore_list(result_collector.modules_to_not_convert) config_path = os.path.join(input_path, "config.json") From 49ddb8620999e346302a422e56d8ce847b6d3284 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Mon, 1 Jun 2026 16:52:27 -0700 Subject: [PATCH 18/37] Avoid shared storage for paired NVFP4 scales Signed-off-by: Ziang Li --- miles/utils/nvfp4.py | 4 ++-- tests/fast-gpu/test_nvfp4_quantizer.py | 28 ++++++++++++++++++++++++++ 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/miles/utils/nvfp4.py b/miles/utils/nvfp4.py index 1d8429957e2..624b7aaea3d 100644 --- a/miles/utils/nvfp4.py +++ b/miles/utils/nvfp4.py @@ -113,11 +113,11 @@ def nvfp4_quantize_1d_pair( first_result = ( combined_qweight[:first_rows].contiguous(), combined_block_scale[:first_rows].contiguous(), - global_scale, + global_scale.clone(), ) second_result = ( combined_qweight[first_rows:].contiguous(), combined_block_scale[first_rows:].contiguous(), - global_scale, + global_scale.clone(), ) return first_result, second_result diff --git a/tests/fast-gpu/test_nvfp4_quantizer.py b/tests/fast-gpu/test_nvfp4_quantizer.py index 2db157389d6..194cfa8abd6 100644 --- a/tests/fast-gpu/test_nvfp4_quantizer.py +++ b/tests/fast-gpu/test_nvfp4_quantizer.py @@ -306,6 +306,34 @@ def test_nvfp4_hf_converter_quantizes_cross_shard_gated_pair_together(tmp_path, ) +def test_nvfp4_hf_converter_quantizes_same_shard_gated_pair_together(tmp_path, monkeypatch): + monkeypatch.delenv("NVTE_NVFP4_4OVER6", raising=False) + model_dir = tmp_path / "model" + save_dir = tmp_path / "converted" + model_dir.mkdir() + (model_dir / "config.json").write_text('{"num_hidden_layers": 1}') + + gate_key = "model.layers.0.mlp.experts.0.gate_proj.weight" + up_key = "model.layers.0.mlp.experts.0.up_proj.weight" + gate = torch.randn((3, 128), dtype=torch.bfloat16) + up = torch.randn((5, 128), dtype=torch.bfloat16) + safetensors.torch.save_file( + { + gate_key: gate, + up_key: up, + }, + model_dir / "model.safetensors", + metadata={"format": "pt"}, + ) + + convert_nvfp4(str(model_dir), str(save_dir), device="cuda") + + with safetensors.safe_open(save_dir / "model.safetensors", framework="pt", device="cuda") as f: + gate_global_scale = f.get_tensor(gate_key.replace(".weight", ".weight_scale_2")) + up_global_scale = f.get_tensor(up_key.replace(".weight", ".weight_scale_2")) + torch.testing.assert_close(gate_global_scale, up_global_scale, rtol=0, atol=0) + + @pytest.mark.parametrize( "quantize_fn", [processor_quantize_nvfp4, tool_quantize_nvfp4], From 60c104dba864deb405c58382e2f569c90be95272 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Mon, 1 Jun 2026 17:11:04 -0700 Subject: [PATCH 19/37] test: broaden paired NVFP4 quantizer coverage Signed-off-by: Ziang Li --- tests/fast-gpu/test_nvfp4_quantizer.py | 104 ++----------------------- 1 file changed, 5 insertions(+), 99 deletions(-) diff --git a/tests/fast-gpu/test_nvfp4_quantizer.py b/tests/fast-gpu/test_nvfp4_quantizer.py index 194cfa8abd6..a9b3b345d61 100644 --- a/tests/fast-gpu/test_nvfp4_quantizer.py +++ b/tests/fast-gpu/test_nvfp4_quantizer.py @@ -83,102 +83,6 @@ def _te_nvfp4_reference( return qweight, block_scale, nvfp4_global_decode_scale_te(global_amax, nvfp4_e4m3_max) -def test_nvfp4_quantize_uses_te_direct_rowwise_quantizer(monkeypatch): - import miles.utils.nvfp4 as nvfp4_utils - - calls = [] - - class FakeQuantizedTensor: - def __init__(self, tensor: torch.Tensor): - self._rowwise_data = torch.arange(tensor.shape[0] * (tensor.shape[1] // 2), dtype=torch.uint8).reshape( - tensor.shape[0], tensor.shape[1] // 2 - ) - self._rowwise_scale_inv = torch.arange( - tensor.shape[0] * (tensor.shape[1] // NVFP4_GROUP_SIZE), dtype=torch.uint8 - ).reshape(tensor.shape[0], tensor.shape[1] // NVFP4_GROUP_SIZE) - self._amax_rowwise = torch.tensor([2.0], dtype=torch.float32) - - class FakeQuantizer: - def __init__(self, **kwargs): - calls.append(kwargs) - - def quantize(self, tensor: torch.Tensor) -> FakeQuantizedTensor: - assert tensor.shape == (16, NVFP4_GROUP_SIZE) - return FakeQuantizedTensor(tensor) - - monkeypatch.setattr(nvfp4_utils, "NVFP4Quantizer", FakeQuantizer) - - qweight, block_scale, global_scale = nvfp4_utils.nvfp4_quantize_1d( - torch.ones((3, NVFP4_GROUP_SIZE), dtype=torch.float32), - ) - - assert calls == [ - { - "rowwise": True, - "columnwise": False, - "with_amax_reduction": False, - "with_rht": False, - "with_post_rht_amax": False, - "with_2d_quantization": False, - "stochastic_rounding": False, - "row_scaled_nvfp4": False, - "nvfp4_use_4over6": False, - "nvfp4_e4m3_max": 448, - "nvfp4_4over6_err_mode": "MAE", - "with_random_sign_mask": False, - } - ] - assert qweight.shape == (3, NVFP4_GROUP_SIZE // 2) - assert block_scale.shape == (3, 1) - assert block_scale.dtype == torch.float8_e4m3fn - torch.testing.assert_close(global_scale, nvfp4_global_decode_scale_te(torch.tensor(2.0), 448), rtol=0, atol=0) - - -def test_nvfp4_quantize_pair_concats_before_te_quantizer(monkeypatch): - import miles.utils.nvfp4 as nvfp4_utils - - quantized_input = None - - class FakeQuantizedTensor: - def __init__(self, tensor: torch.Tensor): - self._rowwise_data = torch.arange(tensor.shape[0] * (tensor.shape[1] // 2), dtype=torch.uint8).reshape( - tensor.shape[0], tensor.shape[1] // 2 - ) - self._rowwise_scale_inv = torch.arange( - tensor.shape[0] * (tensor.shape[1] // NVFP4_GROUP_SIZE), dtype=torch.uint8 - ).reshape(tensor.shape[0], tensor.shape[1] // NVFP4_GROUP_SIZE) - self._amax_rowwise = torch.tensor([7.0], dtype=torch.float32) - - class FakeQuantizer: - def __init__(self, **kwargs): - pass - - def quantize(self, tensor: torch.Tensor) -> FakeQuantizedTensor: - nonlocal quantized_input - quantized_input = tensor.clone() - return FakeQuantizedTensor(tensor) - - monkeypatch.setattr(nvfp4_utils, "NVFP4Quantizer", FakeQuantizer) - - gate = torch.ones((3, NVFP4_GROUP_SIZE), dtype=torch.float32) - up = torch.full((5, NVFP4_GROUP_SIZE), 2.0, dtype=torch.float32) - (gate_qweight, gate_block_scale, gate_global_scale), ( - up_qweight, - up_block_scale, - up_global_scale, - ) = nvfp4_utils.nvfp4_quantize_1d_pair(gate, up) - - assert quantized_input.shape == (16, NVFP4_GROUP_SIZE) - torch.testing.assert_close(quantized_input[:3], gate) - torch.testing.assert_close(quantized_input[3:8], up) - torch.testing.assert_close(quantized_input[8:], torch.zeros_like(quantized_input[8:])) - assert gate_qweight.shape == (3, NVFP4_GROUP_SIZE // 2) - assert up_qweight.shape == (5, NVFP4_GROUP_SIZE // 2) - assert gate_block_scale.shape == (3, 1) - assert up_block_scale.shape == (5, 1) - torch.testing.assert_close(gate_global_scale, up_global_scale, rtol=0, atol=0) - - def test_nvfp4_quantize_params_requires_complete_gated_pair(): weight = torch.randn((4, NVFP4_GROUP_SIZE), dtype=torch.float32) with pytest.raises(ValueError, match="requires gate/up tensors to be quantized together"): @@ -366,9 +270,11 @@ def test_nvfp4_quantize_matches_te_reference_bitwise(quantize_fn, shape, dtype, torch.testing.assert_close(global_scale, global_scale_ref, rtol=0, atol=0) +@pytest.mark.parametrize("shape", NVFP4_SHAPES) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=str) +@pytest.mark.parametrize("init_data", ["random", "boundary", "zeros", "maxes"]) @pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) -def test_nvfp4_quantize_pair_matches_te_reference_bitwise(dtype, use_4over6, monkeypatch): +def test_nvfp4_quantize_pair_matches_te_reference_bitwise(shape, dtype, init_data, use_4over6, monkeypatch): device = "cuda" torch.manual_seed(42) if use_4over6: @@ -377,8 +283,8 @@ def test_nvfp4_quantize_pair_matches_te_reference_bitwise(dtype, use_4over6, mon else: monkeypatch.delenv("NVTE_NVFP4_4OVER6", raising=False) - gate = _make_weight("random", dtype, (3, 128), device) - up = _make_weight("boundary", dtype, (5, 128), device) + gate = _make_weight(init_data, dtype, shape, device) + up = _make_weight(init_data, dtype, shape, device) (gate_qweight, gate_block_scale, gate_global_scale), ( up_qweight, up_block_scale, From 28803c2cb61ef363567a9c5c6a846c38b05a2d6d Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Thu, 18 Jun 2026 15:04:39 -0700 Subject: [PATCH 20/37] Clean up --- scripts/run_qwen3_30b_a3b.py | 42 +++++++++++++++--------------------- 1 file changed, 17 insertions(+), 25 deletions(-) diff --git a/scripts/run_qwen3_30b_a3b.py b/scripts/run_qwen3_30b_a3b.py index fe904292bbc..ca04f949842 100644 --- a/scripts/run_qwen3_30b_a3b.py +++ b/scripts/run_qwen3_30b_a3b.py @@ -1,5 +1,4 @@ import os -import shlex from dataclasses import dataclass from typing import Literal @@ -8,18 +7,6 @@ import miles.utils.external_utils.command_utils as U -def fp4_env_vars() -> dict[str, str]: - fp4_env_markers = ["NVTE", "FLASHINFER", "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH"] - return {key: value for key, value in os.environ.items() if any(marker in key for marker in fp4_env_markers)} - - -def env_prefix(env_vars: dict[str, str]) -> str: - if not env_vars: - return "" - assignments = (f"{key}={shlex.quote(value)}" for key, value in sorted(env_vars.items())) - return " ".join(assignments) + " " - - @dataclass class ScriptArgs(U.ExecuteTrainConfig): mode: Literal["normal", "debug_minimal"] = "normal" @@ -101,14 +88,17 @@ def prepare(args: ScriptArgs): f"{args.extra_args} " ) - if args.rollout_nvfp4 or args.train_nvfp4: - nvfp4_env_prefix = env_prefix( - { - "NVTE_USE_FAST_MATH": "0", - "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH": "1", - **fp4_env_vars(), - } - ) + if args.rollout_nvfp4: + nvfp4_env_vars = { + "NVTE_USE_FAST_MATH": "0", + "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH": "1", + **{ + key: value + for key, value in os.environ.items() + if "NVTE" in key or "FLASHINFER" in key or key == "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH" + }, + } + nvfp4_env_prefix = " ".join(f"{key}={value}" for key, value in nvfp4_env_vars.items()) + " " U.exec_command( f"{nvfp4_env_prefix}" f"python tools/convert_hf_to_nvfp4.py --model-dir {args.model_dir}/{args.model_name} " @@ -146,7 +136,7 @@ def execute(args: ScriptArgs): hf_checkpoint = f"{args.model_dir}/{args.model_name}-FP8" elif args.train_mxfp8: hf_checkpoint = f"{args.model_dir}/{args.model_name}-MXFP8" - elif args.rollout_nvfp4 or args.train_nvfp4: + elif args.rollout_nvfp4: hf_checkpoint = f"{args.model_dir}/{args.model_name}-NVFP4" elif args.rollout_int4: hf_checkpoint = f"{args.model_dir}/{args.model_name}-INT4" @@ -388,6 +378,11 @@ def execute(args: ScriptArgs): "SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION": "1", "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH": "1", } + misc_env_vars |= { + key: value + for key, value in os.environ.items() + if "NVTE" in key or "FLASHINFER" in key or key == "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH" + } else: sglang_args += "--rollout-num-gpus-per-engine 4 " "--sglang-cuda-graph-max-bs 512 " case _: @@ -428,9 +423,6 @@ def execute(args: ScriptArgs): f"{args.extra_args} " ) - if args.rollout_nvfp4 or args.train_nvfp4: - misc_env_vars |= fp4_env_vars() - U.execute_train( train_args=train_args, num_gpus_per_node=args.num_gpus_per_node, From 7d5550406274c5ac2dadc7c42d35b24263d817d0 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Thu, 18 Jun 2026 15:07:34 -0700 Subject: [PATCH 21/37] Also process FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH --- scripts/run_qwen3_30b_a3b.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/scripts/run_qwen3_30b_a3b.py b/scripts/run_qwen3_30b_a3b.py index ca04f949842..0edb9426a87 100644 --- a/scripts/run_qwen3_30b_a3b.py +++ b/scripts/run_qwen3_30b_a3b.py @@ -92,6 +92,7 @@ def prepare(args: ScriptArgs): nvfp4_env_vars = { "NVTE_USE_FAST_MATH": "0", "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH": "1", + "FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH": "1", **{ key: value for key, value in os.environ.items() @@ -377,6 +378,7 @@ def execute(args: ScriptArgs): misc_env_vars |= { "SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION": "1", "TRTLLM_DISABLE_FP4_QUANT_FAST_MATH": "1", + "FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH": "1", } misc_env_vars |= { key: value From 48683057eac2a2de585f0ea07e5708675673ac92 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Thu, 18 Jun 2026 15:11:31 -0700 Subject: [PATCH 22/37] Clean up --- .../megatron_utils/megatron_to_hf/processors/__init__.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py b/miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py index 42aed336921..1bb1b21a6fe 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/__init__.py @@ -17,13 +17,12 @@ def quantize_params(args, megatron_name, converted_named_params, quantization_config): if quantization_config is None: return converted_named_params - elif quantization_config.get("quant_method") == "fp8": + elif quantization_config["quant_method"] == "fp8": return quantize_params_fp8(args, megatron_name, converted_named_params, quantization_config) - elif quantization_config.get("quant_method") == "mxfp8": + elif quantization_config["quant_method"] == "mxfp8": return quantize_params_mxfp8(args, megatron_name, converted_named_params, quantization_config) - elif quantization_config.get("quant_algo") == "NVFP4" or quantization_config.get("quant_method") == "nvfp4": + elif quantization_config.get("quant_algo") == "NVFP4" or quantization_config["quant_method"] == "nvfp4": return quantize_params_nvfp4(args, megatron_name, converted_named_params, quantization_config) - elif quantization_config.get("quant_method") == "compressed-tensors": + elif quantization_config["quant_method"] == "compressed-tensors": # only int4 at the moment. return quantize_params_compressed_tensors(converted_named_params, quantization_config) - return converted_named_params From 93591f3ed4e36cfdcbce5a9bff9cb8990da90a6a Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Thu, 18 Jun 2026 15:23:49 -0700 Subject: [PATCH 23/37] Clean up assertions --- scripts/run_qwen3_30b_a3b.py | 34 +++++++++------------------------- 1 file changed, 9 insertions(+), 25 deletions(-) diff --git a/scripts/run_qwen3_30b_a3b.py b/scripts/run_qwen3_30b_a3b.py index 0edb9426a87..fba5926f4f6 100644 --- a/scripts/run_qwen3_30b_a3b.py +++ b/scripts/run_qwen3_30b_a3b.py @@ -42,34 +42,18 @@ def __post_init__(self): if self.no_colocate: self.actor_num_gpus_per_node = self.actor_num_gpus_per_node or self.num_gpus_per_node // 2 self.rollout_num_gpus = self.rollout_num_gpus or self.num_gpus_per_node - self.actor_num_gpus_per_node - assert self.actor_num_gpus_per_node > 0, "actor_num_gpus_per_node must be positive" - assert self.rollout_num_gpus > 0, "rollout_num_gpus must be positive" - assert ( - self.actor_num_gpus_per_node + self.rollout_num_gpus <= self.num_gpus_per_node - ), "actor and rollout GPU allocations cannot exceed num_gpus_per_node" else: self.actor_num_gpus_per_node = self.actor_num_gpus_per_node or self.num_gpus_per_node self.rollout_num_gpus = self.rollout_num_gpus or self.num_gpus_per_node - if self.rollout_int4: - assert not self.rollout_fp8, "rollout_int4 and rollout_fp8 cannot be enabled at the same time" - assert not self.rollout_mxfp8, "rollout_int4 and rollout_mxfp8 cannot be enabled at the same time" - assert not self.rollout_nvfp4, "rollout_int4 and rollout_nvfp4 cannot be enabled at the same time" - if self.rollout_mxfp8: - assert not self.rollout_fp8, "rollout_mxfp8 and rollout_fp8 cannot be enabled at the same time" - assert not self.rollout_nvfp4, "rollout_mxfp8 and rollout_nvfp4 cannot be enabled at the same time" - assert self.hardware in ("B200", "B300", "GB200", "GB300"), "rollout_mxfp8 only supports Blackwell GPUs" - if self.rollout_nvfp4: - assert not self.rollout_fp8, "rollout_nvfp4 and rollout_fp8 cannot be enabled at the same time" - assert self.hardware in ("B200", "B300", "GB200", "GB300"), "rollout_nvfp4 only supports Blackwell GPUs" - if self.train_mxfp8: - assert not self.train_fp8, "train_mxfp8 and train_fp8 cannot be enabled at the same time" - assert not self.train_nvfp4, "train_mxfp8 and train_nvfp4 cannot be enabled at the same time" - assert self.hardware in ("B200", "B300", "GB200", "GB300"), "train_mxfp8 only supports Blackwell GPUs" - assert self.rollout_mxfp8, "train_mxfp8 requires rollout_mxfp8 to be enabled" - if self.train_nvfp4: - assert not self.train_fp8, "train_nvfp4 and train_fp8 cannot be enabled at the same time" - assert self.hardware in ("B200", "B300", "GB200", "GB300"), "train_nvfp4 only supports Blackwell GPUs" - assert self.rollout_nvfp4, "train_nvfp4 requires rollout_nvfp4 to be enabled" + + assert ( + sum((self.rollout_fp8, self.rollout_mxfp8, self.rollout_int4, self.rollout_nvfp4)) <= 1 + ), "only one rollout precision mode can be enabled" + assert ( + sum((self.train_fp8, self.train_mxfp8, self.train_nvfp4)) <= 1 + ), "only one train precision mode can be enabled" + if any((self.rollout_mxfp8, self.rollout_nvfp4, self.train_mxfp8, self.train_nvfp4)): + assert self.hardware in ("B200", "B300", "GB200", "GB300"), "mxfp8 and nvfp4 only support Blackwell GPUs" def prepare(args: ScriptArgs): From 8bcc4d317bba9a2957172491e0ce7be586b6baa6 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Mon, 22 Jun 2026 13:48:25 -0700 Subject: [PATCH 24/37] test: move nvfp4 glm e2e to glm5.2 --- scripts/models/glm5-744B-A40B_5layer.sh | 12 ------ ... => test_glm5_2_744b_a40b_5layer_nvfp4.py} | 40 ++++++++++++------- 2 files changed, 25 insertions(+), 27 deletions(-) delete mode 100644 scripts/models/glm5-744B-A40B_5layer.sh rename tests/e2e/megatron/{test_glm5_744b_a40b_5layer_nvfp4.py => test_glm5_2_744b_a40b_5layer_nvfp4.py} (89%) diff --git a/scripts/models/glm5-744B-A40B_5layer.sh b/scripts/models/glm5-744B-A40B_5layer.sh deleted file mode 100644 index e6e93cd0890..00000000000 --- a/scripts/models/glm5-744B-A40B_5layer.sh +++ /dev/null @@ -1,12 +0,0 @@ -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]:-$0}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/glm5-744B-A40B.sh" - -# Override for 5-layer pruned model (first 5 layers: 3 dense + 2 MoE) -N_MOE_LAYERS=2 - -for ((i=0; i<${#MODEL_ARGS[@]}; i++)); do - case "${MODEL_ARGS[$i]}" in - --num-layers) MODEL_ARGS[$((i+1))]=$((N_DENSE_LAYERS + N_MOE_LAYERS)) ;; - --moe-layer-freq) MODEL_ARGS[$((i+1))]="[0]*${N_DENSE_LAYERS}+[1]*${N_MOE_LAYERS}" ;; - esac -done diff --git a/tests/e2e/megatron/test_glm5_744b_a40b_5layer_nvfp4.py b/tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py similarity index 89% rename from tests/e2e/megatron/test_glm5_744b_a40b_5layer_nvfp4.py rename to tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py index 2258154f49e..e48b09a6828 100644 --- a/tests/e2e/megatron/test_glm5_744b_a40b_5layer_nvfp4.py +++ b/tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py @@ -14,8 +14,8 @@ ) MODEL_ORG = "Pinaster" -MODEL_NAME = "GLM-5_5layer" -MODEL_TYPE = "glm5-744B-A40B_5layer" +MODEL_NAME = "GLM-5.2_5layer" +MODEL_TYPE = "glm5.2-744B-A40B_5layer" NUM_GPUS = 8 ACTOR_NUM_GPUS = 4 ROLLOUT_NUM_GPUS = 4 @@ -54,6 +54,7 @@ } GLM5_ENV = { + "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256", "SGLANG_NSA_FORCE_MLA": "1", "INDEXER_ROPE_NEOX_STYLE": "0", "NVSHMEM_DISABLE_NCCL": "1", @@ -119,7 +120,7 @@ def _validate_glm_checkpoint(): or config.get("num_hidden_layers") != 5 ): raise RuntimeError( - f"{config_path} must use native GLM-5 5-layer config with " + f"{config_path} must use native GLM-5.2 5-layer config with " f"model_type=glm_moe_dsa, architectures=[GlmMoeDsaForCausalLM], " "and num_hidden_layers=5" ) @@ -148,7 +149,7 @@ def prepare(): U.convert_checkpoint( model_name=MODEL_NAME, megatron_model_type=MODEL_TYPE, - num_gpus_per_node=ACTOR_NUM_GPUS, + num_gpus_per_node=1, extra_args=( "--tensor-model-parallel-size 1 " "--expert-tensor-parallel-size 1 " @@ -177,12 +178,11 @@ def execute(): "--rollout-shuffle " "--rm-type deepscaler " "--num-rollout 2 " - "--rollout-batch-size 32 " + "--rollout-batch-size 8 " "--n-samples-per-prompt 8 " "--rollout-max-response-len 100 " "--rollout-temperature 1 " - "--global-batch-size 256 " - "--balance-data " + "--global-batch-size 64 " ) perf_args = ( @@ -196,19 +196,22 @@ def execute(): "--recompute-method uniform " "--recompute-num-layers 1 " "--use-dynamic-batch-size " - "--max-tokens-per-gpu 32768 " - "--data-pad-size-multiplier 4096 " - "--log-probs-chunk-size 1024 " + "--max-tokens-per-gpu 2048 " + "--data-pad-size-multiplier 1024 " + "--log-probs-chunk-size 16384 " ) grpo_args = ( "--advantage-estimator grpo " - "--use-kl-loss " "--kl-loss-coef 0.00 " "--kl-loss-type low_var_kl " + "--kl-coef 0.00 " "--entropy-coef 0.00 " "--eps-clip 0.2 " "--eps-clip-high 0.28 " + "--use-tis " + "--tis-clip-low 0.5 " + "--tis-clip 2.0 " ) optimizer_args = ( @@ -225,20 +228,25 @@ def execute(): sglang_args = ( "--sglang-mem-fraction-static 0.7 " + "--sglang-enable-dp-attention " "--sglang-attention-backend nsa " - "--sglang-nsa-decode-backend flashmla_sparse " + "--sglang-nsa-decode-backend flashmla_kv " "--sglang-nsa-prefill-backend flashmla_sparse " - "--sglang-kv-cache-dtype bf16 " + "--sglang-kv-cache-dtype fp8_e4m3 " "--sglang-page-size 64 " f"--rollout-num-gpus-per-engine {ROLLOUT_GPUS_PER_ENGINE} " "--sglang-moe-runner-backend flashinfer_trtllm_routed " - f"--sglang-tp-size {ROLLOUT_GPUS_PER_ENGINE} " f"--sglang-ep-size {ROLLOUT_GPUS_PER_ENGINE} " + f"--sglang-dp-size {ROLLOUT_GPUS_PER_ENGINE} " + "--sglang-moe-dense-tp-size 1 " + "--sglang-enable-dp-lm-head " "--sglang-cuda-graph-max-bs 256 " + "--sglang-max-running-requests 512 " + f"--sglang-chunked-prefill-size {2048 * ROLLOUT_GPUS_PER_ENGINE} " "--sglang-watchdog-timeout 3600 " ) - ci_args = "--ci-test --ci-disable-logprobs-checker " + ci_args = "--ci-test --ci-disable-logprobs-checker --disable-weights-backuper " mixed_precision_args = ( "--transformer-impl transformer_engine " @@ -269,6 +277,8 @@ def execute(): f"--num-gpus-per-node {NUM_GPUS} " f"--rollout-num-gpus {ROLLOUT_NUM_GPUS} " "--use-fault-tolerance " + "--moe-enable-deepep " + "--moe-token-dispatcher-type flex " f"--dump-details /root/shared_data/{RUN_ID}/dump_details " ) From e45a27e693ffed8fb24a65280e797d8e91fad549 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Mon, 22 Jun 2026 14:09:08 -0700 Subject: [PATCH 25/37] test: log kl loss in glm5.2 nvfp4 e2e --- tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py b/tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py index e48b09a6828..6bc7a6f20ca 100644 --- a/tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py +++ b/tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py @@ -203,6 +203,7 @@ def execute(): grpo_args = ( "--advantage-estimator grpo " + "--use-kl-loss " "--kl-loss-coef 0.00 " "--kl-loss-type low_var_kl " "--kl-coef 0.00 " @@ -246,7 +247,7 @@ def execute(): "--sglang-watchdog-timeout 3600 " ) - ci_args = "--ci-test --ci-disable-logprobs-checker --disable-weights-backuper " + ci_args = "--ci-test --ci-disable-logprobs-checker " mixed_precision_args = ( "--transformer-impl transformer_engine " From 156900b0733b8334cc5061bb381475c2ea8669cb Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Mon, 22 Jun 2026 14:42:50 -0700 Subject: [PATCH 26/37] Optimize weight update time by avoiding concat materialization --- miles/utils/nvfp4.py | 27 +++++++++++++++++++++++--- tests/fast-gpu/test_nvfp4_quantizer.py | 16 +++++++++++++++ 2 files changed, 40 insertions(+), 3 deletions(-) diff --git a/miles/utils/nvfp4.py b/miles/utils/nvfp4.py index 624b7aaea3d..6b3c6c3a1f2 100644 --- a/miles/utils/nvfp4.py +++ b/miles/utils/nvfp4.py @@ -107,9 +107,10 @@ def nvfp4_quantize_1d_pair( ) first_rows = first.shape[0] - combined_qweight, combined_block_scale, global_scale = nvfp4_quantize_1d( - torch.cat((first.contiguous(), second.contiguous()), dim=0) - ) + combined = _contiguous_pair_view(first, second) + if combined is None: + combined = torch.cat((first.contiguous(), second.contiguous()), dim=0) + combined_qweight, combined_block_scale, global_scale = nvfp4_quantize_1d(combined) first_result = ( combined_qweight[:first_rows].contiguous(), combined_block_scale[:first_rows].contiguous(), @@ -121,3 +122,23 @@ def nvfp4_quantize_1d_pair( global_scale.clone(), ) return first_result, second_result + + +def _contiguous_pair_view(first: torch.Tensor, second: torch.Tensor) -> torch.Tensor | None: + if not first.is_contiguous() or not second.is_contiguous(): + return None + if first.device != second.device or first.dtype != second.dtype or first.stride() != second.stride(): + return None + if first.untyped_storage().data_ptr() != second.untyped_storage().data_ptr(): + return None + if first.storage_offset() + first.numel() != second.storage_offset(): + return None + + try: + return first.as_strided( + (first.shape[0] + second.shape[0], first.shape[1]), + first.stride(), + first.storage_offset(), + ) + except RuntimeError: + return None diff --git a/tests/fast-gpu/test_nvfp4_quantizer.py b/tests/fast-gpu/test_nvfp4_quantizer.py index a9b3b345d61..d91bf5e45a3 100644 --- a/tests/fast-gpu/test_nvfp4_quantizer.py +++ b/tests/fast-gpu/test_nvfp4_quantizer.py @@ -238,6 +238,22 @@ def test_nvfp4_hf_converter_quantizes_same_shard_gated_pair_together(tmp_path, m torch.testing.assert_close(gate_global_scale, up_global_scale, rtol=0, atol=0) +def test_nvfp4_quantize_pair_reuses_adjacent_storage(monkeypatch): + base = torch.randn((32, 64), dtype=torch.bfloat16, device="cuda") + gate, up = base.chunk(2, dim=0) + + def fail_cat(*args, **kwargs): + raise AssertionError("adjacent gate/up pair should not be materialized with torch.cat") + + monkeypatch.setattr(torch, "cat", fail_cat) + (gate_qweight, gate_block_scale, _), (up_qweight, up_block_scale, _) = nvfp4_quantize_1d_pair(gate, up) + + assert gate_qweight.shape == (16, 32) + assert up_qweight.shape == (16, 32) + assert gate_block_scale.shape == (16, 4) + assert up_block_scale.shape == (16, 4) + + @pytest.mark.parametrize( "quantize_fn", [processor_quantize_nvfp4, tool_quantize_nvfp4], From ab620da9e004e4434e97aca63c21554542086f13 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Mon, 22 Jun 2026 21:41:02 -0700 Subject: [PATCH 27/37] Drop activation scale during weight update --- .../processors/quantizer_nvfp4.py | 3 -- tests/fast-gpu/test_nvfp4_quantizer.py | 29 +++++++++++++++++++ 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py index 34891cb5569..0b7a6ef2407 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py @@ -137,9 +137,6 @@ def _quantize_moe_params(converted_named_params, ignore_rules): quantize_named_params.append((converted_name, qweight)) quantize_named_params.append((converted_name.replace(".weight", ".weight_scale"), block_scale)) quantize_named_params.append((converted_name.replace(".weight", ".weight_scale_2"), weight_scale_2)) - quantize_named_params.append( - (converted_name.replace(".weight", ".input_scale"), torch.ones_like(weight_scale_2, dtype=torch.float32)) - ) return quantize_named_params diff --git a/tests/fast-gpu/test_nvfp4_quantizer.py b/tests/fast-gpu/test_nvfp4_quantizer.py index d91bf5e45a3..17f136ce750 100644 --- a/tests/fast-gpu/test_nvfp4_quantizer.py +++ b/tests/fast-gpu/test_nvfp4_quantizer.py @@ -142,6 +142,35 @@ def test_nvfp4_quantize_params_respects_first_last_layers_bf16(layer_idx): assert out is converted_named_params +def test_nvfp4_quantize_params_omits_static_input_scale(monkeypatch): + weight = torch.randn((4, NVFP4_GROUP_SIZE), dtype=torch.bfloat16) + qweight = torch.empty((4, NVFP4_GROUP_SIZE // 2), dtype=torch.uint8) + block_scale = torch.empty((4, 1), dtype=torch.float8_e4m3fn) + global_scale = torch.ones((), dtype=torch.float32) + + def fake_quantize_1d_pair(_gate, _up): + return (qweight, block_scale, global_scale), (qweight, block_scale, global_scale) + + monkeypatch.setattr( + "miles.backends.megatron_utils.megatron_to_hf.processors.quantizer_nvfp4.nvfp4_quantize_1d_pair", + fake_quantize_1d_pair, + ) + + out = quantize_params_nvfp4( + args=None, + megatron_name="decoder.layers.0.mlp.experts.linear_fc1.weight0", + converted_named_params=[ + ("model.layers.0.mlp.experts.0.gate_proj.weight", weight), + ("model.layers.0.mlp.experts.0.up_proj.weight", weight), + ], + quantization_config={"quant_method": "nvfp4"}, + ) + + names = [name for name, _ in out] + assert "model.layers.0.mlp.experts.0.gate_proj.input_scale" not in names + assert "model.layers.0.mlp.experts.0.up_proj.input_scale" not in names + + def test_nvfp4_hf_should_quantize_respects_extra_high_precision_layers_hf(): weight = torch.randn((4, NVFP4_GROUP_SIZE), dtype=torch.bfloat16) From 1eb97cb32d675051f4d8a61d0eedae8ac8f26733 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Thu, 25 Jun 2026 16:04:12 -0700 Subject: [PATCH 28/37] Add topk backend --- tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py b/tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py index 6bc7a6f20ca..8be1250f1a4 100644 --- a/tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py +++ b/tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py @@ -55,7 +55,9 @@ GLM5_ENV = { "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256", - "SGLANG_NSA_FORCE_MLA": "1", + "SGLANG_DSA_FUSE_TOPK": "1", +"SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD": "0", + "SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK": "large", "INDEXER_ROPE_NEOX_STYLE": "0", "NVSHMEM_DISABLE_NCCL": "1", } @@ -233,6 +235,7 @@ def execute(): "--sglang-attention-backend nsa " "--sglang-nsa-decode-backend flashmla_kv " "--sglang-nsa-prefill-backend flashmla_sparse " + "--sglang-dsa-topk-backend flashinfer " "--sglang-kv-cache-dtype fp8_e4m3 " "--sglang-page-size 64 " f"--rollout-num-gpus-per-engine {ROLLOUT_GPUS_PER_ENGINE} " @@ -272,6 +275,7 @@ def execute(): "--attention-softmax-in-fp32 " "--attention-backend flash " "--allgather-cp " + "--miles-dsa-topk-backend flashinfer " f"--update-weight-buffer-size {2 * 1024 ** 3} " "--actor-num-nodes 1 " f"--actor-num-gpus-per-node {ACTOR_NUM_GPUS} " From bcfd5be2fe980cb401ec5a274680bf3498b644a1 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Thu, 25 Jun 2026 16:04:35 -0700 Subject: [PATCH 29/37] Opt out unimplemented weight checker --- miles/utils/arguments.py | 11 ++++++++++- .../megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py | 4 ++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 344a59141dc..20e26eb74a9 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1749,6 +1749,10 @@ def add_ci_arguments(parser): "--ci-disable-logprobs-checker", action="store_true", ) + parser.add_argument( + "--ci-disable-weight-update-checker", + action="store_true", + ) parser.add_argument( "--ci-metric-checker-key", type=str, @@ -2275,7 +2279,12 @@ def miles_validate_args(args): "debug_rollout_only and debug_train_only cannot be set at the same time, " "please set only one of them." ) - if args.ci_test and not args.debug_rollout_only and not args.debug_train_only: + if ( + args.ci_test + and not args.debug_rollout_only + and not args.debug_train_only + and not args.ci_disable_weight_update_checker + ): args.check_weight_update_equal = True # always true on offload for colocate at the moment. diff --git a/tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py b/tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py index 8be1250f1a4..b96eafad868 100644 --- a/tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py +++ b/tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4.py @@ -56,7 +56,7 @@ GLM5_ENV = { "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "256", "SGLANG_DSA_FUSE_TOPK": "1", -"SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD": "0", + "SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD": "0", "SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK": "large", "INDEXER_ROPE_NEOX_STYLE": "0", "NVSHMEM_DISABLE_NCCL": "1", @@ -250,7 +250,7 @@ def execute(): "--sglang-watchdog-timeout 3600 " ) - ci_args = "--ci-test --ci-disable-logprobs-checker " + ci_args = "--ci-test --ci-disable-logprobs-checker --ci-disable-weight-update-checker " mixed_precision_args = ( "--transformer-impl transformer_engine " From d6828729ef521f2788a0ca2192c146c43ad97677 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Thu, 9 Jul 2026 13:27:29 -0700 Subject: [PATCH 30/37] Bump TransformerEngine docker pins to 2.17 --- docker/Dockerfile | 8 ++++---- .../patch/cu13/te_fa2_sm103_whitelist.patch | 20 ++++++++++++------- 2 files changed, 17 insertions(+), 11 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index b259542f558..1493cdd5bf2 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -94,16 +94,16 @@ RUN pip install causal-conv1d==1.6.1 mamba-ssm==2.3.1 --no-build-isolation # transformer_engine RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \ - pip install --no-deps transformer_engine==2.12.0 && \ - pip install transformer_engine_cu13==2.12.0 && \ + pip install --no-deps transformer_engine==2.17.0 && \ + pip install transformer_engine_cu13==2.17.0 && \ if ls /tmp/wheels/transformer_engine_torch-*.whl 2>/dev/null | grep -q .; then \ pip install /tmp/wheels/transformer_engine_torch-*.whl; \ else \ pip install nvidia-mathdx==25.6.0 && \ - pip -v install --no-build-isolation transformer_engine_torch==2.12.0; \ + pip -v install --no-build-isolation transformer_engine_torch==2.17.0; \ fi; \ else \ - pip -v install --no-build-isolation "transformer_engine[pytorch]==2.10.0"; \ + pip -v install --no-build-isolation "transformer_engine[pytorch]==2.17.0"; \ fi # TE patches (cu13): B300/GB300 sm103 FA2 whitelist fix diff --git a/docker/patch/cu13/te_fa2_sm103_whitelist.patch b/docker/patch/cu13/te_fa2_sm103_whitelist.patch index 203fe27e65a..c2313d8816d 100644 --- a/docker/patch/cu13/te_fa2_sm103_whitelist.patch +++ b/docker/patch/cu13/te_fa2_sm103_whitelist.patch @@ -1,11 +1,17 @@ --- a/pytorch/attention/dot_product_attention/utils.py +++ b/pytorch/attention/dot_product_attention/utils.py -@@ -629,7 +629,7 @@ - or head_dim_qk % 8 != 0 - or ( - head_dim_qk > 192 -- and device_compute_capability not in ((8, 0), (9, 0), (10, 0), (12, 0)) -+ and device_compute_capability not in ((8, 0), (9, 0), (10, 0), (10, 3), (12, 0)) +@@ -848,16 +848,11 @@ + and ( + fa2_padded_head_dim > 256 + or fa2_padded_head_dim % 8 != 0 +- or ( +- fa2_padded_head_dim > 192 +- and device_compute_capability not in ((8, 0), (9, 0), (10, 0), (12, 0)) +- ) ) ): - if FlashAttentionUtils.is_installed: + logger.debug( + "Disabling FlashAttention 2 due to unsupported head_dim_qk and head_dim_v. " + "Supported after padding: padded head_dim %%8 = 0, padded head_dim <= 256 " +- "(>192 requires sm80/90/100+). " + "Found: head_dim_qk = %s, head_dim_v = %s, padded head_dim = %s, on sm%s.", From 7ccbf1d22000f3e1ae29d0aedc4d7a7ed55dcc61 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Thu, 9 Jul 2026 13:48:21 -0700 Subject: [PATCH 31/37] Add TransformerEngine dequantized override patch --- docker/Dockerfile | 2 +- .../te_dequantized_backward_override.patch | 32 +++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) create mode 100644 docker/patch/cu13/te_dequantized_backward_override.patch diff --git a/docker/Dockerfile b/docker/Dockerfile index 1493cdd5bf2..9a720800dd5 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -106,7 +106,7 @@ RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \ pip -v install --no-build-isolation "transformer_engine[pytorch]==2.17.0"; \ fi -# TE patches (cu13): B300/GB300 sm103 FA2 whitelist fix +# TE patches (cu13): B300/GB300 FA2 and backward override fixes COPY docker/patch/ /tmp/patches/ RUN if [ "${ENABLE_CUDA_13}" = "1" ] && [ -d /tmp/patches/cu13 ]; then \ TE_DIR=$(python -c 'import transformer_engine; print(transformer_engine.__path__[0])') && \ diff --git a/docker/patch/cu13/te_dequantized_backward_override.patch b/docker/patch/cu13/te_dequantized_backward_override.patch new file mode 100644 index 00000000000..593577b6d08 --- /dev/null +++ b/docker/patch/cu13/te_dequantized_backward_override.patch @@ -0,0 +1,32 @@ +--- a/pytorch/module/grouped_linear.py ++++ b/pytorch/module/grouped_linear.py +@@ -428,10 +428,12 @@ + if fp8: + backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + else: + backward_override = None + if backward_override == "high_precision": + save_original_input = True +- ++ elif backward_override == "dequantized": ++ save_original_input = False ++ + num_gemms = len(m_splits) + weights = weights_and_biases[:num_gemms] + biases = weights_and_biases[num_gemms:] +--- a/pytorch/module/linear.py ++++ b/pytorch/module/linear.py +@@ -282,10 +282,12 @@ + save_original_input = args.save_original_input + debug = args.debug + backward_override = args.backward_override + is_fsdp2 = args.is_fsdp2 + if backward_override == "high_precision": + save_original_input = True +- ++ elif backward_override == "dequantized": ++ save_original_input = False ++ + # NVTX label for profiling + nvtx_label = "transformer_engine._Linear.forward" + if ub_name is not None: From f9da769c569fffc4802628216ae8113bd6f71e78 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Thu, 9 Jul 2026 13:49:51 -0700 Subject: [PATCH 32/37] Document temporary TransformerEngine patch --- docker/Dockerfile | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docker/Dockerfile b/docker/Dockerfile index 9a720800dd5..7cceb9ae3da 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -107,6 +107,8 @@ RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \ fi # TE patches (cu13): B300/GB300 FA2 and backward override fixes +# te_dequantized_backward_override.patch is a hot fix from +# https://github.com/NVIDIA/TransformerEngine/pull/3141; drop it after TE v2.18. COPY docker/patch/ /tmp/patches/ RUN if [ "${ENABLE_CUDA_13}" = "1" ] && [ -d /tmp/patches/cu13 ]; then \ TE_DIR=$(python -c 'import transformer_engine; print(transformer_engine.__path__[0])') && \ From 0067637b4eaca2a2c8e50d3296cc7e0dfb0cf282 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Thu, 23 Jul 2026 17:24:58 -0700 Subject: [PATCH 33/37] Bump TransformerEngine to 2.17.0 (cu13 and cu12 paths) - transformer_engine / transformer_engine_cu13 / transformer_engine_torch 2.12.0 -> 2.17.0 - update te_fa2_sm103_whitelist.patch for TE 2.17 source layout - add te_dequantized_backward_override.patch (hot fix from NVIDIA/TransformerEngine#3141; drop after TE v2.18) --- docker/Dockerfile | 12 +++++----- .../te_dequantized_backward_override.patch | 22 +++++++++++++++++++ .../patch/cu13/te_fa2_sm103_whitelist.patch | 22 +++++++++++++------ 3 files changed, 44 insertions(+), 12 deletions(-) create mode 100644 docker/patch/cu13/te_dequantized_backward_override.patch diff --git a/docker/Dockerfile b/docker/Dockerfile index d2a3b038fdc..8a31f23aa97 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -94,19 +94,21 @@ RUN pip install causal-conv1d==1.6.1 mamba-ssm==2.3.1 --no-build-isolation # transformer_engine RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \ - pip install --no-deps transformer_engine==2.12.0 && \ - pip install transformer_engine_cu13==2.12.0 && \ + pip install --no-deps transformer_engine==2.17.0 && \ + pip install transformer_engine_cu13==2.17.0 && \ if ls /tmp/wheels/transformer_engine_torch-*.whl 2>/dev/null | grep -q .; then \ pip install /tmp/wheels/transformer_engine_torch-*.whl; \ else \ pip install nvidia-mathdx==25.6.0 && \ - pip -v install --no-build-isolation transformer_engine_torch==2.12.0; \ + pip -v install --no-build-isolation transformer_engine_torch==2.17.0; \ fi; \ else \ - pip -v install --no-build-isolation "transformer_engine[pytorch]==2.10.0"; \ + pip -v install --no-build-isolation "transformer_engine[pytorch]==2.17.0"; \ fi -# TE patches (cu13): B300/GB300 sm103 FA2 whitelist fix +# TE patches (cu13): B300/GB300 FA2 and backward override fixes +# te_dequantized_backward_override.patch is a hot fix from +# https://github.com/NVIDIA/TransformerEngine/pull/3141; drop it after TE v2.18. COPY docker/patch/ /tmp/patches/ RUN if [ "${ENABLE_CUDA_13}" = "1" ] && [ -d /tmp/patches/cu13 ]; then \ TE_DIR=$(python -c 'import transformer_engine; print(transformer_engine.__path__[0])') && \ diff --git a/docker/patch/cu13/te_dequantized_backward_override.patch b/docker/patch/cu13/te_dequantized_backward_override.patch new file mode 100644 index 00000000000..31c81da4904 --- /dev/null +++ b/docker/patch/cu13/te_dequantized_backward_override.patch @@ -0,0 +1,22 @@ +--- a/pytorch/module/grouped_linear.py ++++ b/pytorch/module/grouped_linear.py +@@ -431,6 +431,8 @@ + backward_override = None + if backward_override == "high_precision": + save_original_input = True ++ elif backward_override == "dequantized": ++ save_original_input = False + + num_gemms = len(m_splits) + weights = weights_and_biases[:num_gemms] +--- a/pytorch/module/linear.py ++++ b/pytorch/module/linear.py +@@ -285,6 +285,8 @@ + is_fsdp2 = args.is_fsdp2 + if backward_override == "high_precision": + save_original_input = True ++ elif backward_override == "dequantized": ++ save_original_input = False + + # NVTX label for profiling + nvtx_label = "transformer_engine._Linear.forward" diff --git a/docker/patch/cu13/te_fa2_sm103_whitelist.patch b/docker/patch/cu13/te_fa2_sm103_whitelist.patch index 203fe27e65a..259d4879407 100644 --- a/docker/patch/cu13/te_fa2_sm103_whitelist.patch +++ b/docker/patch/cu13/te_fa2_sm103_whitelist.patch @@ -1,11 +1,19 @@ --- a/pytorch/attention/dot_product_attention/utils.py +++ b/pytorch/attention/dot_product_attention/utils.py -@@ -629,7 +629,7 @@ - or head_dim_qk % 8 != 0 - or ( - head_dim_qk > 192 -- and device_compute_capability not in ((8, 0), (9, 0), (10, 0), (12, 0)) -+ and device_compute_capability not in ((8, 0), (9, 0), (10, 0), (10, 3), (12, 0)) +@@ -848,16 +848,11 @@ + and ( + fa2_padded_head_dim > 256 + or fa2_padded_head_dim % 8 != 0 +- or ( +- fa2_padded_head_dim > 192 +- and device_compute_capability not in ((8, 0), (9, 0), (10, 0), (12, 0)) +- ) ) ): - if FlashAttentionUtils.is_installed: + logger.debug( + "Disabling FlashAttention 2 due to unsupported head_dim_qk and head_dim_v. " + "Supported after padding: padded head_dim %%8 = 0, padded head_dim <= 256 " +- "(>192 requires sm80/90/100+). " + "Found: head_dim_qk = %s, head_dim_v = %s, padded head_dim = %s, on sm%s.", + head_dim_qk, + head_dim_v, From cdbafdc5cd9454cf966c4e38cc636f52de872ccd Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Thu, 23 Jul 2026 22:33:29 -0700 Subject: [PATCH 34/37] Allow tensor parallel size below 4 for small actor GPU counts --- scripts/run_qwen3_30b_a3b.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/run_qwen3_30b_a3b.py b/scripts/run_qwen3_30b_a3b.py index fba5926f4f6..3083e20bce5 100644 --- a/scripts/run_qwen3_30b_a3b.py +++ b/scripts/run_qwen3_30b_a3b.py @@ -308,7 +308,7 @@ def execute(args: ScriptArgs): ) case ("B200" | "B300" | "GB200" | "GB300", 1 | 2 | 4): perf_args += ( - "--tensor-model-parallel-size 4 " + f"--tensor-model-parallel-size {min(4, args.actor_num_gpus_per_node)} " "--sequence-parallel " "--pipeline-model-parallel-size 1 " "--context-parallel-size 1 " From 11c94d2653b453968d2126ea8dae27b4428e20b6 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Thu, 23 Jul 2026 23:05:29 -0700 Subject: [PATCH 35/37] Lazy-import NVFP4Quantizer so nvfp4 utils import on CPU-only envs --- miles/utils/nvfp4.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/miles/utils/nvfp4.py b/miles/utils/nvfp4.py index 6b3c6c3a1f2..69df0b87540 100644 --- a/miles/utils/nvfp4.py +++ b/miles/utils/nvfp4.py @@ -1,7 +1,6 @@ import os import torch -from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer FP4_E2M1_MAX = 6.0 FP8_E4M3_MAX = 448.0 @@ -66,6 +65,8 @@ def _pad_rows_for_te_quantizer(weight: torch.Tensor) -> torch.Tensor: def nvfp4_quantize_1d( weight: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + weight = weight.contiguous() num_rows, num_cols = weight.shape nvfp4_e4m3_max = nvfp4_weight_e4m3_max() From 97431446489c4eac6afb32e03a3c4dc9e7aa4d7a Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Thu, 23 Jul 2026 23:39:21 -0700 Subject: [PATCH 36/37] Fix torch_memory_saver LD_PRELOAD path for CUDA-suffixed binaries torch_memory_saver 0.0.9.post1 ships CUDA-major-suffixed preload binaries (torch_memory_saver_hook_mode_preload_cu13.abi3.so), but actor_factory hardcoded the old unsuffixed name and asserted it exists, breaking every offload_train megatron run. Use the package's get_binary_path_from_package helper to resolve the right variant. --- miles/ray/train/actor_factory.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/miles/ray/train/actor_factory.py b/miles/ray/train/actor_factory.py index 6c8245daec1..43b54227979 100644 --- a/miles/ray/train/actor_factory.py +++ b/miles/ray/train/actor_factory.py @@ -41,13 +41,9 @@ def allocate_gpus_for_actor( env_vars["DUMPER_SOURCE_PATCHER_CONFIG"] = source_patcher_config if args.offload_train and args.train_backend == "megatron": - import torch_memory_saver + from torch_memory_saver.utils import get_binary_path_from_package - dynlib_path = os.path.join( - os.path.dirname(os.path.dirname(torch_memory_saver.__file__)), - "torch_memory_saver_hook_mode_preload.abi3.so", - ) - assert os.path.exists(dynlib_path), f"LD_PRELOAD so file {dynlib_path} does not exist." + dynlib_path = str(get_binary_path_from_package("torch_memory_saver_hook_mode_preload")) env_vars["LD_PRELOAD"] = dynlib_path env_vars["TMS_INIT_ENABLE"] = "1" From e22a54818c122009151e93e345613629005824b2 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Fri, 24 Jul 2026 01:56:08 -0700 Subject: [PATCH 37/37] Disable NVFP4 quantizer test in CI until a Blackwell runner exists The NVFP4 quantize kernels are gated to sm_100+ (Blackwell); the test is registered on stage-b-2-gpu-h200 (Hopper) where tex.quantize hits NVTE_DEVICE_ERROR("sm_100 or higher is required"). Mark disabled, matching the GLM5.2 NVFP4 e2e test, until miles CI has a B-card suite. --- tests/fast-gpu/test_nvfp4_quantizer.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/fast-gpu/test_nvfp4_quantizer.py b/tests/fast-gpu/test_nvfp4_quantizer.py index 17f136ce750..66955d5635c 100644 --- a/tests/fast-gpu/test_nvfp4_quantizer.py +++ b/tests/fast-gpu/test_nvfp4_quantizer.py @@ -1,6 +1,11 @@ from tests.ci.ci_register import register_cuda_ci -register_cuda_ci(est_time=60, suite="stage-b-2-gpu-h200", labels=[]) +register_cuda_ci( + est_time=60, + suite="stage-b-2-gpu-h200", + labels=[], + disabled="Requires Blackwell/B200 CI runner for NVFP4.", +) import os