mirror of https://github.com/NVIDIA/Model-Optimizer.git synced 2026-10-20 05:59:32 +08:00

Files

T

History

Zhiyu 5500999d0b Add Kimi-K3 NVFP4 experts and FP8-PB attention recipe (#2206 )

### What does this PR do?

Type of change: new example

Adds the calibration-free conversion pipeline and checkpoint-mirror PTQ
recipe used for `nvidia/Kimi-K3-NVFP4`:

- streams the 96-shard Kimi-K3 checkpoint without loading the 2.8T
model;
- casts the source MXFP4 routed experts to NVFP4 with expert
`input_scale=1.0`;
- quantizes the selected KDA and MLA attention weights to 128x128 block
FP8;
- leaves shared/latent experts, routers, convolutions, norms, the vision
tower, `lm_head`, and KV cache unquantized;
- emits mixed-precision Hugging Face metadata for deployment; and
- adds an exact recipe under
`modelopt_recipes/huggingface/models/moonshotai/Kimi-K3/` that directly
configures the streaming converter.

It also fixes `NVFP4QTensor.quantize()` probing CUDA/Blackwell
capability before checking whether the tensor is on CUDA and whether the
optional TensorRT-LLM fast path was requested. That probe broke the
converter's supported CPU path on hosts without a compatible GPU.

### Usage

```bash
python examples/kimi/kimi_k3/quantize_to_nvfp4.py \
    --source_ckpt /models/moonshotai/Kimi-K3 \
    --output_ckpt /models/Kimi-K3-NVFP4 \
    --recipe huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention \
    --jobs 8
```

The conversion requires no calibration dataset, forward pass, or GPU.
Multi-node shard conversion is also supported through `--rank`,
`--world_size`, and `--run_id`.

### Testing

```bash
uv run --frozen --extra dev python -m pytest -q \
    tests/unit/torch/quantization/test_nvfp4_tensor.py \
    tests/unit/recipe/test_kimi_k3_recipe.py \
    tests/unit/recipe/test_recipe_docs.py \
    tests/unit/torch/export/test_shard_cast_utils.py \
    tests/examples/kimi/test_kimi_k3_quantize_to_nvfp4.py
```

Result: 34 passed.

All pre-commit hooks pass for the changed files, including recipe
validation, Ruff, mypy, Bandit, YAML formatting, and markdownlint.

### 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](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅
- Did you get Claude approval on this PR?: N/A

### Additional Information

The resulting checkpoint and model card are available at
https://huggingface.co/nvidia/Kimi-K3-NVFP4.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Added calibration-free Kimi-K3 MXFP4-to-NVFP4 conversion with optional
FP8 attention quantization and distributed processing.
- Added NVFP4 activation calibration, weight-only quantization recipes,
grouped-expert quantization, compiled quantization options, SFT-masked
distillation, and MLflow tracking.

- **Documentation**
- Updated quantization terminology, recipe catalogs, checkpoint
guidance, and Kimi-K3 conversion instructions.

- **Bug Fixes**
- Improved CPU NVFP4 behavior, tied-weight export handling, EAGLE-3
training compatibility, and checkpoint export reliability.
- Removed deprecated configuration options and legacy evaluation
examples.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>

2026-08-28 02:33:14 +00:00

kimi_k3

Add Kimi-K3 NVFP4 experts and FP8-PB attention recipe (#2206 )

2026-08-28 02:33:14 +00:00

README.md

Add Kimi-K3 NVFP4 experts and FP8-PB attention recipe (#2206 )

2026-08-28 02:33:14 +00:00

README.md

Kimi post-training quantization

Kimi-K3: NVFP4 routed experts and block-FP8 attention

moonshotai/Kimi-K3 is released with its routed experts already packed as MXFP4. Loading the 2.8T-parameter model and running the normal in-memory examples/hf_ptq/hf_ptq.py flow would both discard that source representation and require impractical host memory. The Kimi-K3 converter therefore streams one safetensors shard at a time:

routed-expert w1, w2, and w3 weights are cast from MXFP4 to NVFP4;
expert activation input_scale is fixed to 1.0;
KDA and MLA projection weights are quantized to FP8 in 128x128 blocks;
attention activations are dynamically quantized by the inference runtime;
shared experts, latent experts, routers, convolution weights, norms, the vision tower, and lm_head remain BF16.

The exact quantization map is recorded in modelopt_recipes/huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention.yaml. Because the source checkpoint uses packed MXFP4 tensors rather than ordinary Hugging Face Linear.weight tensors, run the streaming converter instead of passing this recipe to hf_ptq.py:

python examples/kimi/kimi_k3/quantize_to_nvfp4.py \
    --source_ckpt /models/moonshotai/Kimi-K3 \
    --output_ckpt /models/Kimi-K3-NVFP4 \
    --recipe huggingface/models/moonshotai/Kimi-K3/ptq/nvfp4_experts-fp8_pb_attention \
    --jobs 8

The conversion is calibration-free: it does not run a forward pass, require a dataset, or require a GPU. For multi-node conversion, launch one process per node against shared storage and set --rank, --world_size, and a common --run_id. Use --help for the full set of conversion and synchronization options.