Update

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>

Author: cjluo-nv

examples/llm_ptq/hf_ptq.py

@@ -53,6 +53,7 @@
     export_hf_checkpoint,
     export_tensorrt_llm_checkpoint,
     get_model_type,
+    save_expert_token_count_table,
 )
 from modelopt.torch.export.model_utils import get_language_model_from_vl, is_multimodal_model
 from modelopt.torch.quantization.config import _default_disabled_quantizer_cfg, need_calibration
@@ -727,6 +728,7 @@ def post_quantize(
 
     if args.verbose:
         mtq.print_quant_summary(full_model)
+        save_expert_token_count_table(full_model, args.export_path)
 
     # Run some samples
     torch.cuda.empty_cache()

modelopt/torch/export/__init__.py

@@ -19,6 +19,7 @@
 from .model_config import *
 from .model_config_export import *
 from .model_utils import *
+from .moe_utils import *
 from .plugins import *
 from .transformer_engine import *
 from .unified_export_hf import *

modelopt/torch/export/unified_export_hf.py

@@ -76,7 +76,6 @@
     QUANTIZATION_W4A8_NVFP4_FP8,
 )
 from .model_utils import get_language_model_from_vl, is_multimodal_model
-from .moe_utils import save_expert_token_count_table
 from .plugins import export_spec_ckpt_config, export_spec_ckpt_state_dict, spec_opt_only
 from .quant_utils import (
     fuse_prequant_layernorm,
@@ -1004,8 +1003,6 @@ def export_hf_checkpoint(
     try:
         post_state_dict, hf_quant_config = _export_transformers_checkpoint(model, dtype)
 
-        save_expert_token_count_table(model, export_dir)
-
         if hf_quant_config is not None:
             # Save hf_quant_config.json for backward compatibility
             with open(f"{export_dir}/hf_quant_config.json", "w") as file: