nemo_automodel.components.quantization.nvfp4

NVFP4 (fp4 e2m1 + fp8-e4m3 block scales + fp32 tensor scale) pack/unpack utilities for MoE expert weights.

The packed layout matches the ModelOpt NVFP4 checkpoint format (hf_quant_config.json: quant_algo: NVFP4, group_size: 16): two e2m1 values per uint8 byte (the even column in the low nibble, the following odd column in the high nibble), one float8_e4m3fn scale per 16 contiguous columns (weight_scale) and one float32 scale per tensor (weight_scale_2), so a value decodes as e2m1 * e4m3 * fp32. In-model storage keeps the per-tensor scale per output row, which lets fused gate/up projections carry the two tensor scales of their halves exactly: 0.5625 bytes per parameter plus 4 bytes per row.

NVFP4GroupedMM provides a grouped GEMM over packed weights that re-dequantizes in backward instead of saving the dequantized tensor, so frozen expert weights stay packed at steady state during LoRA training. No FP4 tensor cores are needed: the GEMM runs on the dequantized bf16 operand, on Hopper or Blackwell alike.

Module Contents

Classes

Name	Description
`NVFP4GroupedMM`	Grouped GEMM over NVFP4-packed frozen weights with dequantization on the fly.

Functions

Name	Description
`_check_last_dim`	Reject contraction dims that are not a whole number of 16-column blocks.
`dequantize_nvfp4`	Unpack fp4 e2m1 packed-uint8 values and apply the per-16-column e4m3 scale and the row scale.
`quantize_nvfp4`	Quantize along the last dim to the packed NVFP4 layout used by `dequantize_nvfp4`.
`quantize_nvfp4_modelopt`	Quantize a 2-D `[out, in]` matrix with one tensor scale, as a ModelOpt NVFP4 checkpoint stores it.

Data

_DEQUANT_CHUNK_ELEMS

API

class nemo_automodel.components.quantization.nvfp4.NVFP4GroupedMM()

Bases: Function

Grouped GEMM over NVFP4-packed frozen weights with dequantization on the fly.

Saves only the packed codes and scales for backward and re-dequantizes there, so the floating-point weight tensor is a transient in both passes instead of being kept alive by autograd. Weights are stored as [E, N, K] packed along K, the natural dequantization output and the operand the backward GEMM needs directly (grad_x = grad_out @ W). The forward needs [E, K, N], which torch._grouped_mm consumes as a transposed view (cuBLAS transB) without a contiguous copy.

No weight gradient is produced: the base weights are frozen under LoRA.

nemo_automodel.components.quantization.nvfp4.NVFP4GroupedMM.backward(
    ctx: torch.autograd.function.FunctionCtx,
    grad_out: torch.Tensor
) -> tuple[torch.Tensor, None, None, None, None]

staticmethod

Backpropagate through the activations while keeping the packed base frozen.

Parameters:

grad_out

torch.Tensor

Tensor of shape [tokens, out_dim] with the activation dtype, grouped by local expert like the forward input.

Returns: torch.Tensor

Tuple of (grad_x of shape [tokens, in_dim] in the activation dtype, None for packed, scales,

nemo_automodel.components.quantization.nvfp4.NVFP4GroupedMM.forward(
    ctx: torch.autograd.function.FunctionCtx,
    x: torch.Tensor,
    packed: torch.Tensor,
    scales: torch.Tensor,
    global_scale: torch.Tensor,
    offs: torch.Tensor
) -> torch.Tensor

staticmethod

Multiply activations grouped by expert by packed frozen weights.

Parameters:

torch.Tensor

Tensor of shape [tokens, in_dim], grouped contiguously by local expert.

packed

torch.Tensor

uint8 tensor of shape [local_experts, out_dim, in_dim // 2].

scales

torch.Tensor

float8_e4m3fn tensor of shape [local_experts, out_dim, in_dim // 16].

global_scale

torch.Tensor

float32 tensor of shape [local_experts, out_dim].

offs

torch.Tensor

Int32 tensor of shape [local_experts], holding cumulative token counts. All tensors must be on the same device.

Returns: torch.Tensor

Tensor of shape [tokens, out_dim] with the activation dtype.

nemo_automodel.components.quantization.nvfp4._check_last_dim(
    k: int
) -> None

Reject contraction dims that are not a whole number of 16-column blocks.

nemo_automodel.components.quantization.nvfp4.dequantize_nvfp4(
    packed: torch.Tensor,
    scales: torch.Tensor,
    global_scale: torch.Tensor,
    dtype: torch.dtype
) -> torch.Tensor

Unpack fp4 e2m1 packed-uint8 values and apply the per-16-column e4m3 scale and the row scale.

Each value decodes as e2m1 * float(e4m3 block scale) * row scale with fp32 intermediates and a single final rounding to dtype; the work runs in row chunks of _DEQUANT_CHUNK_ELEMS elements to bound the transients.

Parameters:

packed

torch.Tensor

uint8 (or int8) tensor of shape [..., K // 2] holding two e2m1 values per byte (even column in the low nibble).

scales

torch.Tensor

float8_e4m3fn tensor of shape [..., K // 16].

global_scale

torch.Tensor

float32 tensor of shape [...] (one value per row of packed), broadcast over the row’s blocks.

dtype

torch.dtype

Output dtype.

Returns: torch.Tensor

Dequantized tensor of shape [..., K] in dtype.

nemo_automodel.components.quantization.nvfp4.quantize_nvfp4(
    weight: torch.Tensor,
    global_scale: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Quantize along the last dim to the packed NVFP4 layout used by dequantize_nvfp4.

Follows the ModelOpt convention global = amax / (6 * 448), block_scale = e4m3(amax(block) / 6 / global) and code = round_to_e2m1(w / (block_scale * global)), with one global scale per row unless global_scale is given (see quantize_nvfp4_modelopt for the per-tensor checkpoint convention).

Parameters:

weight

torch.Tensor

Floating-point tensor of shape [..., K] with K divisible by 16.

global_scale

torch.Tensor | NoneDefaults to None

Optional float32 tensor broadcastable to weight.shape[:-1] with the row scales to use instead of the per-row amax rule.

Returns: torch.Tensor

Tuple of (uint8 packed codes [..., K // 2], float8_e4m3fn block scales [..., K // 16],

nemo_automodel.components.quantization.nvfp4.quantize_nvfp4_modelopt(
    weight: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Quantize a 2-D [out, in] matrix with one tensor scale, as a ModelOpt NVFP4 checkpoint stores it.

Returns: torch.Tensor

Tuple of (weight: uint8 [out, in // 2], weight_scale: float8_e4m3fn [out, in // 16],

nemo_automodel.components.quantization.nvfp4.E2M1_MAX = 6.0

nemo_automodel.components.quantization.nvfp4.E4M3_MAX = 448.0

nemo_automodel.components.quantization.nvfp4.E4M3_MIN_SUBNORMAL = 2.0 ** -9

nemo_automodel.components.quantization.nvfp4.NVFP4_BLOCK_SIZE = 16

nemo_automodel.components.quantization.nvfp4._DEQUANT_CHUNK_ELEMS = 2 ** 27