Fix Llama4 example

[yiliu30]

No description provided.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 9 New Failures, 1 Cancelled Job

As of commit c6fdc2c with merge base df7bf37 ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test (CPU 2.6, linux.4xlarge, torch==2.6.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  test/quantization/test_moe_quant.py::TestMoEQuantCompile::test_int8wo_base_cpu_1_multiple_tokens
• Run Regression Tests / test (CPU 2.7, linux.4xlarge, torch==2.7.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  test/quantization/test_moe_quant.py::TestMoEQuantCompile::test_int8wo_base_cpu_1_multiple_tokens
• Run Regression Tests / test (CPU 2.8, linux.4xlarge, torch==2.8.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  test/quantization/test_moe_quant.py::TestMoEQuantCompile::test_int8wo_base_cpu_1_multiple_tokens
• Run Regression Tests / test (CUDA 2.6, linux.g5.12xlarge.nvidia.gpu, torch==2.6.0, cuda, 12.6) / linux-job (gh)
  test/quantization/test_moe_quant.py::TestMoEQuantCompile::test_int8wo_fake_dim_1_multiple_tokens
• Run Regression Tests / test (CUDA 2.7, linux.g5.12xlarge.nvidia.gpu, torch==2.7.0, cuda, 12.6) / linux-job (gh)
  test/quantization/test_moe_quant.py::TestMoEQuantCompile::test_int8wo_fake_dim_1_multiple_tokens
• Run Regression Tests / test (CUDA 2.8, linux.g5.12xlarge.nvidia.gpu, torch==2.8.0, cuda, 12.6) / linux-job (gh)
  test/quantization/test_moe_quant.py::TestMoEQuantCompile::test_int8wo_fake_dim_1_multiple_tokens
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh)
  test/quantization/test_moe_quant.py::TestMoEQuantCompile::test_int8wo_base_cpu_1_multiple_tokens
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  test/quantization/test_moe_quant.py::TestMoEQuantCompile::test_int8wo_fake_dim_1_multiple_tokens
• Run TorchAO Experimental Tests / test-cpu-ops (macos-14) (gh)
  torchao/experimental/tests/test_int8_dynamic_activation_intx_weight.py::TestInt8DynamicActivationIntxWeight::test_moe_quant_intx

CANCELLED JOB - The following job was cancelled. Please retry:

• Run TorchAO Experimental Tests / test-cpu-ops (linux.arm64.2xlarge) (gh)
  ##[error]The operation was canceled.

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

[yiliu30]

The original router returns router_scores and router_logits.
https://github.com/huggingface/transformers/blob/19ffe0219dae122203f9726669f88ef1c6ea3bb4/src/transformers/models/llama4/modeling_llama4.py#L143

[jerryzh168]

router here is a nn.Linear, not Llama4Router actually, see L21?

[yiliu30]

This rounter was replaced with the original rounter.

ao/torchao/prototype/moe_quant/llama4_quant.py

Lines 48 to 53 in 9d01b43

	router = module.router
	up_proj = module.experts.gate_up_proj
	w1, w3 = up_proj.permute(0, 2, 1).chunk(2, dim=1)
	w2 = module.experts.down_proj.permute(0, 2, 1)
	new_mod.router = router

[jerryzh168]

so this module does not run by itself? seems quite confusing

[yiliu30]

The MOEFeedForwardAOQuantizable was used, but it seems that its rounter shouldn’t be quantized in order to preserve accuracy. Could you confirm that? @HDCharles

[yiliu30]

Hi @liangel-02 @andrewor14 Looks like the failed CI checks are not related to this PR. Could you help retrigger them?