Kimi-k2 calib+export

[jingyu-ml]

What does this PR do?

Type of change: new example

Overview: Support Kimi-k2 Calibration

Usage

python hf_ptq.py --pyt_ckpt_path unsloth/Kimi-K2-Thinking-BF16 --qformat nvfp4_mlp_experts_only --export_path <quantized_ckpt_path> --trust_remote_code

Testing

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes/No
• Did you write any new necessary tests?: Yes/No
• Did you add or update any necessary documentation?: Yes/No
• Did you update Changelog?: Yes/No

Additional Information

[realAsma]

We already have this config. See NVFP4_MLP_ONLY_CFG

Suggested change

	NVFP4_MLP_EXPERTS_ONLY_CFG = {
	"quant_cfg": {
	"*mlp.experts*weight_quantizer": {
	"num_bits": (2, 1),
	"block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
	"enable": True,
	"pass_through_bwd": True,
	},
	"*mlp.experts*input_quantizer": {
	"num_bits": (2, 1),
	"block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
	"enable": True,
	"pass_through_bwd": True,
	},
	**_default_disabled_quantizer_cfg,
	},
	"algorithm": "max",
	}

[jingyu-ml]

that will quantize mlp.shared_experts

[cjluo-nv]

this is great. Let's not creating more cfgs

[cjluo-nv]

Is it a recommendation or it's something we have to do? An alterantive is to up cast the in4 to BF16 during calibration like we did with DS.

[jingyu-ml]

But there’s no INT4 support in PyTorch, as we discussed. People have to use vLLM if they want INT4. Me and Zhiyu are looking into the vLLM calibration of this model

[cjluo-nv]

could you move these to the examples_utils.py like: https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/examples/llm_ptq/example_utils.py#L303

[cjluo-nv]

I'm still against adding more configs here. I think this MR we should just stick with MLP_only if we have to. People can tune the recipe themselves if they want to do experts only.

If you really like to add this config, let's name it experts_only for short. experts are always in MLP

[jingyu-ml]

Renamed it to nvfp4_experts_only. Let’s keep this config for now; once the YAML config system is released, we can avoid using these recipe dictionaries.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 74.57%. Comparing base (c6c9905) to head (647755c).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #655      +/-   ##
==========================================
- Coverage   74.58%   74.57%   -0.02%     
==========================================
  Files         183      183              
  Lines       18451    18452       +1     
==========================================
- Hits        13762    13760       -2     
- Misses       4689     4692       +3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.