import os
from quark.torch import LLMTemplate, ModelQuantizer


# Configuration
ckpt_path = "Qwen/Qwen3.5-397B-A17B-FP8"
output_dir = "amd/Qwen3.5-397B-A17B-MXFP4"
quant_scheme = "mxfp4"
exclude_layers = ["lm_head", "model.visual.*", "mtp.*", "*mlp.gate", "*shared_expert_gate*", "*.linear_attn.*", "*.self_attn.*", "*.shared_expert.*"]

# Get quant config from template
template = LLMTemplate.get("qwen3_5_moe")
quant_config = template.get_config(scheme=quant_scheme, exclude_layers=exclude_layers)

# Quantize with File-to-file mode
quantizer = ModelQuantizer(quant_config)
quantizer.direct_quantize_checkpoint(
    pretrained_model_path=ckpt_path,
    save_path=output_dir,
)

lm_eval \
  --model vllm \
  --model_args pretrained=amd/Qwen3.5-397B-A17B-MXFP4,tensor_parallel_size=4,max_model_len=262144,gpu_memory_utilization=0.90,max_gen_toks=2048,trust_remote_code=True,reasoning_parser=qwen3 \
  --tasks gsm8k  --num_fewshot 5 \
  --batch_size auto