> This page is for version 0.7.0 · 26.10. > For other versions, use one of these documentation indexes: > - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt > - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt > - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt > - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt > - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt > For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server. # 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`](#nemo_automodel-components-quantization-nvfp4-NVFP4GroupedMM) | Grouped GEMM over NVFP4-packed frozen weights with dequantization on the fly. | ### Functions | Name | Description | | -------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------- | | [`_check_last_dim`](#nemo_automodel-components-quantization-nvfp4-_check_last_dim) | Reject contraction dims that are not a whole number of 16-column blocks. | | [`dequantize_nvfp4`](#nemo_automodel-components-quantization-nvfp4-dequantize_nvfp4) | Unpack fp4 e2m1 packed-uint8 values and apply the per-16-column e4m3 scale and the row scale. | | [`quantize_nvfp4`](#nemo_automodel-components-quantization-nvfp4-quantize_nvfp4) | Quantize along the last dim to the packed NVFP4 layout used by `dequantize_nvfp4`. | | [`quantize_nvfp4_modelopt`](#nemo_automodel-components-quantization-nvfp4-quantize_nvfp4_modelopt) | Quantize a 2-D `[out, in]` matrix with one tensor scale, as a ModelOpt NVFP4 checkpoint stores it. | ### Data [`E2M1_MAX`](#nemo_automodel-components-quantization-nvfp4-E2M1_MAX) [`E4M3_MAX`](#nemo_automodel-components-quantization-nvfp4-E4M3_MAX) [`E4M3_MIN_SUBNORMAL`](#nemo_automodel-components-quantization-nvfp4-E4M3_MIN_SUBNORMAL) [`NVFP4_BLOCK_SIZE`](#nemo_automodel-components-quantization-nvfp4-NVFP4_BLOCK_SIZE) [`_DEQUANT_CHUNK_ELEMS`](#nemo_automodel-components-quantization-nvfp4-_DEQUANT_CHUNK_ELEMS) ### API ```python 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. ```python 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, ```python 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:** **`x`** `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. ```python nemo_automodel.components.quantization.nvfp4._check_last_dim( k: int ) -> None ``` Reject contraction dims that are not a whole number of 16-column blocks. ```python 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`. ```python 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 | None` — default: 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]`, ```python 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]`, ```python nemo_automodel.components.quantization.nvfp4.E2M1_MAX = 6.0 ``` ```python nemo_automodel.components.quantization.nvfp4.E4M3_MAX = 448.0 ``` ```python nemo_automodel.components.quantization.nvfp4.E4M3_MIN_SUBNORMAL = 2.0 ** -9 ``` ```python nemo_automodel.components.quantization.nvfp4.NVFP4_BLOCK_SIZE = 16 ``` ```python nemo_automodel.components.quantization.nvfp4._DEQUANT_CHUNK_ELEMS = 2 ** 27 ```