NVFP4 RL

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 quantize_params_nvfp4
 
 __all__ = [
     "remove_padding",
     "quantize_param",
     "quantize_params_fp8",
     "quantize_params_mxfp8",
+    "quantize_params_nvfp4",
     "quantize_params_compressed_tensors",
 ]
 
@@ -19,6 +21,8 @@ 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["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["quant_method"] == "nvfp4":
+        return quantize_params_nvfp4(args, megatron_name, converted_named_params, quantization_config)
     elif quantization_config["quant_method"] == "compressed-tensors":
         # only int4 at the moment.
         return quantize_params_compressed_tensors(converted_named_params, quantization_config)

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_quantize_1d, nvfp4_quantize_1d_pair
 
 GATED_PAIR_SUFFIXES = {
     ".gate_proj.weight": "gate",
@@ -38,6 +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"
+    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", ()):
@@ -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,33 +109,34 @@ 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))
             raise ValueError(
                 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)
-        global_amax = shared_global_amax.get(base) if base else None
-        qweight, block_scale, weight_scale_2 = quantize_nvfp4(param, global_amax=global_amax)
+        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))
-        quantize_named_params.append(
-            (converted_name.replace(".weight", ".input_scale"), torch.ones_like(weight_scale_2, dtype=torch.float32))
-        )
 
     return quantize_named_params
 
@@ -162,75 +162,31 @@ 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)
-
-
 def _quantize_nvfp4_1d(
     weight: torch.Tensor,
-    global_amax: torch.Tensor | 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)
       block_scale: float8_e4m3fn, shape (M, K // 16)
       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}.")
 
-    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)
-
-    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)
 
 
 def quantize_nvfp4(
     weight: torch.Tensor,
-    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)
     if weight.dim() == 3:
-        if global_amax is not None:
-            raise ValueError("global_amax override is only supported for 2D weights.")
         qweights = []
         block_scales = []
         global_scales = []

miles/utils/arguments.py

@@ -2105,6 +2105,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,
@@ -2725,7 +2729,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.

miles/utils/nvfp4.py

@@ -0,0 +1,145 @@
+import os
+
+import torch
+
+FP4_E2M1_MAX = 6.0
+FP8_E4M3_MAX = 448.0
+NVFP4_GROUP_SIZE = 16
+TE_NVFP4_ROW_ALIGNMENT = 16
+
+
+def nvfp4_weight_e4m3_max() -> int:
+    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_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 _nvfp4_4over6_enabled() -> bool:
+    return os.getenv("NVTE_NVFP4_4OVER6", "").strip().lower() in ("weights", "all")
+
+
+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,
+) -> 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()
+
+    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=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(),
+        with_random_sign_mask=False,
+    )
+
+    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 = _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(),
+        global_scale.clone(),
+    )
+    second_result = (
+        combined_qweight[first_rows:].contiguous(),
+        combined_block_scale[first_rows:].contiguous(),
+        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