install debug expert token counters on nvfp4 moe test script

[vkuzo]

Summary:

Add expert counters on HF olmoe model to get token-expert counts.
This helps understand how much calibration is needed to get
a reasonable amount for tokens per expert for PTQ algorithms.

For this specific model, seems that calibrating on c4 will quickly hit
all the experts:

(pt_nightly) dev@gpu-dev-fef610c2:~/ao (20260420_gptq_nvfp4)$ time python scripts/prototype/test_nvfp4_moe.py --recipe=bf16 --calibrate_on_c4 --num_calibration_samples 1
...
=== Global expert utilization summary ===
experts with <= 0 tokens: 4/1024 (0.4%)
experts with <= 10 tokens: 60/1024 (5.9%)
experts with <= 20 tokens: 136/1024 (13.3%)
experts with <= 30 tokens: 220/1024 (21.5%)
experts with <= 40 tokens: 322/1024 (31.4%)
experts with <= 50 tokens: 404/1024 (39.5%)
experts with <= 60 tokens: 507/1024 (49.5%)
experts with <= 70 tokens: 612/1024 (59.8%)
experts with <= 80 tokens: 706/1024 (68.9%)
experts with <= 90 tokens: 799/1024 (78.0%)
experts with <= 100 tokens: 878/1024 (85.7%)
experts with <= 110 tokens: 923/1024 (90.1%)
experts with <= 120 tokens: 958/1024 (93.6%)

// min number for at least one token per expert
(pt_nightly) dev@gpu-dev-fef610c2:~/ao (20260420_gptq_nvfp4)$ time python scripts/prototype/test_nvfp4_moe.py --recipe=bf16 --calibrate_on_c4 --num_calibration_samples 3
...
=== Global expert utilization summary ===
experts with <= 0 tokens: 0/1024 (0.0%)
experts with <= 10 tokens: 1/1024 (0.1%)
experts with <= 20 tokens: 9/1024 (0.9%)
experts with <= 30 tokens: 22/1024 (2.1%)
experts with <= 40 tokens: 32/1024 (3.1%)
experts with <= 50 tokens: 49/1024 (4.8%)
experts with <= 60 tokens: 60/1024 (5.9%)
experts with <= 70 tokens: 69/1024 (6.7%)
experts with <= 80 tokens: 81/1024 (7.9%)
experts with <= 90 tokens: 88/1024 (8.6%)
experts with <= 100 tokens: 115/1024 (11.2%)
experts with <= 110 tokens: 130/1024 (12.7%)
experts with <= 120 tokens: 149/1024 (14.6%)

// resonable default
(pt_nightly) dev@gpu-dev-fef610c2:~/ao (20260420_gptq_nvfp4)$ time python scripts/prototype/test_nvfp4_moe.py --recipe=bf16 --calibrate_on_c4 --num_calibration_samples 128
...
=== Global expert utilization summary ===
experts with <= 0 tokens: 0/1024 (0.0%)
experts with <= 10 tokens: 0/1024 (0.0%)
experts with <= 20 tokens: 0/1024 (0.0%)
experts with <= 30 tokens: 0/1024 (0.0%)
experts with <= 40 tokens: 0/1024 (0.0%)
experts with <= 50 tokens: 0/1024 (0.0%)
experts with <= 60 tokens: 0/1024 (0.0%)
experts with <= 70 tokens: 0/1024 (0.0%)
experts with <= 80 tokens: 0/1024 (0.0%)
experts with <= 90 tokens: 0/1024 (0.0%)
experts with <= 100 tokens: 0/1024 (0.0%)
experts with <= 110 tokens: 0/1024 (0.0%)
experts with <= 120 tokens: 0/1024 (0.0%)

90% clauded

Test Plan:

[vkuzo]

Stack from ghstack (oldest at bottom):

• extend gptq example script with olmoe model #4329
• extend GPTQ coverage to grouped_mm #4328
• nvfp4 gptq for bmm #4327
• nvfp4 gptq: prep for bmm and grouped_mm #4323
• -> install debug expert token counters on nvfp4 moe test script #4322

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4322

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Rolling out OSDC (ARC) runners on pull & trunk workflows in PyTorch main

❌ 1 New Failure

As of commit f8d1861 with merge base 67a78e5 ():

NEW FAILURE - The following job has failed:

• Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio mslk --index-url https://download.... / linux-job (gh)
  RuntimeError: Command docker exec -t 703808054d5b07cf875f3603a7d8043aef041013322b7ddaca6e549a64dc917e /exec failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.