Fix reduce_amax NotImplementedError on FP8 weights (NVBug 6360175) (#1824)

### What does this PR do?

Type of change: Bug fix

Fixes [NVBug 6360175](https://nvbugspro.nvidia.com/bug/6360175) /
OMNIML-5265: quantizing a model whose weights are stored natively in FP8
(e.g. DeepSeek-V3 in `float8_e4m3fn`) crashes during `mtq.quantize`
calibration with:

```
File ".../modelopt/torch/quantization/utils/core_utils.py", line 162, in reduce_amax
    max_val = torch.max(input)
NotImplementedError: "max_all_cuda" not implemented for 'Float8_e4m3fn'
```

**Root cause:** FP8 dtypes (`float8_e4m3fn` / `float8_e5m2`) implement
no full-tensor reduction kernel (`max_all_cuda`/`min_all_cuda`), nor
`amax`/`amin`, `abs`, or elementwise `maximum`. `reduce_amax` called
these directly on the FP8 weight tensor.

**Fix:** Upcast FP8 inputs to the default float dtype
(`torch.get_default_dtype()`) at the top of `reduce_amax`, before any
reduction. The upcast is **lossless** (any default float dtype
represents every FP8 value exactly) and only affects the FP8 path — the
common (fp16/bf16/fp32) path is untouched. Placing the upcast at the top
covers all branches (`torch.max`/`min`, `torch.amax`/`amin`,
`torch.abs`), not just the line in the traceback.

### Usage

No API change. Quantization of natively-FP8 checkpoints (e.g.
DeepSeek-V3 NVFP4 PTQ) now runs through calibration instead of raising.

### Testing

- New CPU regression test `test_reduce_amax_fp8` in
`tests/unit/torch/quantization/test_utils.py` covering both FP8 dtypes
(`float8_e4m3fn`, `float8_e5m2`) across all axis modes (`None`, `0`,
`1`, `(0, 1)`); asserts results equal the float reference and the output
dtype is the default float dtype. CPU reproduces the original error (no
FP8 reduction kernel there either), so the test is GPU-free.
- `pre-commit run --files ...` passes (ruff, mypy, bandit, license, rst
checks).

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update Changelog?: ✅ (0.45 Bug Fixes)
- Did you get Claude approval on this PR?: ❌ (not yet)

### Additional Information

NVBug 6360175 is tagged `Committed_ModelOpt_0.45.0` (regression); the
changelog entry is under 0.45 and this will be cherry-picked to
`release/0.45` after merge.

Supersedes #1823, which got a stuck head ref (frozen at the original
commit, no sync on force-push) after the repo move
`TensorRT-Model-Optimizer` → `Model-Optimizer`; it could not be
re-synced or reopened, so this PR replaces it from the same branch.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

CHANGELOG.rst

@@ -125,6 +125,7 @@ Changelog
 - Fix the GPT-OSS MXFP4 → NVFP4 PTQ path in ``examples/llm_ptq/hf_ptq.py`` (used with ``--cast_mxfp4_to_nvfp4``). ``get_model`` now loads native MXFP4 checkpoints (``openai/gpt-oss-*``) dequantized to BF16 ``GptOssExperts`` via ``Mxfp4Config(dequantize=True)`` on a sequential device map. This fixes a CUDA illegal-memory access during the multi-GPU dequant load and the ``NotImplementedError`` for experts type ``Mxfp4GptOssExperts`` during unified HF export (the packed-kernel experts wrapper, used when the optional ``kernels`` package is installed, is unsupported by export); ``kernels`` is no longer required. The ``--cast_mxfp4_to_nvfp4`` step now also resolves a HF Hub ID ``--pyt_ckpt_path`` to its local snapshot directory instead of failing with ``FileNotFoundError``.
 - Fix ``_QuantGptOssExperts`` / ``_QuantLlama4TextExperts`` static-block NVFP4 weight calibration raising ``ValueError: Input shape has changed`` during the calibration forward. These experts quantize their weights transposed (``_transposed_quantize``); ``iter_weights_for_calibration`` now yields the same transposed view so weight-only calibration and the forward agree on the block-quant shape (and the export ``_amax`` orientation).
 - Fix unified HF checkpoint export for Llama4 MoE models. The uncalibrated-experts input-quantizer ``amax`` fallback in ``_export_transformers_checkpoint`` special-cased only ``QuantGptOssExperts``; ``QuantLlama4TextExperts`` uses the same fused ``gate_up_proj`` / ``down_proj`` layout and is now handled by the same branch, fixing the export failure.
+- Fix ``NotImplementedError: "max_all_cuda" not implemented for 'Float8_e4m3fn'`` during quantization calibration of models with natively FP8 (``float8_e4m3fn`` / ``float8_e5m2``) weights, such as DeepSeek-V3. FP8 dtypes implement no reduction (``max``/``amax``), ``abs``, or elementwise ``maximum`` kernels, so ``reduce_amax`` now upcasts FP8 inputs to the default float dtype before reducing; the upcast is lossless and only affects the FP8 path.
 
 0.44 (2026-05-14)
 ^^^^^^^^^^^^^^^^^

modelopt/torch/quantization/utils/core_utils.py

@@ -33,6 +33,10 @@ from modelopt.torch.utils import get_unwrapped_name, print_rank_0
 if TYPE_CHECKING:
     from collections.abc import Generator
 
+# FP8 dtypes do not implement reduction kernels (e.g. ``max_all_cuda``), ``abs``, or
+# elementwise ``maximum``, so tensors of these dtypes must be upcast before amax reduction.
+_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2)
+
 
 def reduce_block_amax(input_tensor: torch.Tensor, block_sizes: dict):
     """Computes the amax of the input tensor using block-based reduction for each dimension.
@@ -157,6 +161,10 @@ def reduce_amax(input, axis=None, keepdims=True, squeeze_scalar=True):
     Returns:
         The reduced tensor.
     """
+    # FP8 dtypes lack reduction/abs kernels (e.g. ``max_all_cuda``); upcast to the default
+    # float dtype, which represents every FP8 value exactly so the amax is computed losslessly.
+    if input.dtype in _FP8_DTYPES:
+        input = input.to(torch.get_default_dtype())
     # A memory-efficient implementation that avoids copying input tensor
     if axis is None:
         max_val = torch.max(input)

tests/unit/torch/quantization/test_utils.py

@@ -18,6 +18,7 @@ import torch
 
 from modelopt.torch.quantization.utils import (
     convert_quantization_axis_to_reduce_axis,
+    reduce_amax,
     reduce_block_amax,
 )
 from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
@@ -58,6 +59,25 @@ def test_reduce_block_amax(block_sizes, test_input, expected_scales):
     torch.allclose(scales, expected_scales)
 
 
+@pytest.mark.parametrize("fp8_dtype", [torch.float8_e4m3fn, torch.float8_e5m2])
+@pytest.mark.parametrize("axis", [None, 0, 1, (0, 1)])
+def test_reduce_amax_fp8(fp8_dtype, axis):
+    """FP8 tensors have no reduction/abs kernels; reduce_amax must upcast them.
+
+    Regression test for ``NotImplementedError: "max_all_cuda" not implemented for
+    'Float8_e4m3fn'`` when calibrating models with natively FP8 weights (e.g. DeepSeek-V3).
+    """
+    # Values chosen to be exactly representable in both FP8 formats so the upcast is lossless.
+    ref = torch.tensor([[1.0, -3.0, 2.0], [0.5, -0.25, 4.0]])
+    x_fp8 = ref.to(fp8_dtype)
+
+    out = reduce_amax(x_fp8, axis=axis)
+    expected = reduce_amax(ref, axis=axis)
+
+    assert out.dtype == torch.get_default_dtype()
+    assert torch.equal(out, expected)
+
+
 @pytest.mark.parametrize(
     ("shape", "quant_axis", "expected_reduce_axis"),
     [