Fix compute_flops_per_token for the MoE configs it claims to support

[vineethsaivs]

What does this PR do?

compute_flops_per_token raises AttributeError on two of the three MoE families its own docstring names.

transformers 5.5.0
Mixtral    -> AttributeError: 'MixtralConfig' object has no attribute 'moe_intermediate_size'
Qwen2-MoE  -> AttributeError: 'Qwen2MoeConfig' object has no attribute 'num_local_experts'
Qwen3-MoE  -> 12,982,419,456

Two causes:

1. Expert count and expert width are spelled differently per family. Mixtral counts experts in num_local_experts and sizes them with intermediate_size; Qwen2-MoE counts them in num_experts and sizes them with moe_intermediate_size. The transformers-version switch picked num_local_experts on 5.1+, which Qwen2-MoE never declares. Now read whichever spelling the config carries.

2. The MoE-layer rule was off by one and ignored mlp_only_layers. transformers uses (layer_idx not in config.mlp_only_layers) and (layer_idx + 1) % config.decoder_sparse_step == 0; this used layer_idx % sparse_step == 0.

Qwen3-MoE is unchanged by this PR (12,982,419,456 before and after), so the fix is a no-op on the path that already worked.

Before submitting

• Did you read the contributor guideline, Pull Request section?
• Did you write any new necessary tests?

AI writing disclosure

• AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.

Test

Three new tests in TestComputeFlopsPerToken, plus the existing test_moe_active_vs_total_experts mirror expression updated to the corrected layer rule.

pytest tests/test_utils.py -k "moe_config_attribute_spellings or moe_expert_size_from_intermediate_size or moe_sparse_step_and_mlp_only_layers"
before: 3 failed
pytest tests/test_utils.py -k "ComputeFlopsPerToken or ComputeMfu"
after: 8 passed

---

Note

Low Risk
Changes only FLOPs/MFU estimation helpers and tests; training forward/backward paths are untouched, though reported MFU for some MoE configs may shift slightly where the layer rule was wrong.

Overview
Fixes compute_flops_per_token for Mixtral and Qwen2-MoE so MFU-style metrics no longer crash with AttributeError and layer counts match transformers.

MoE expert count and width are read from whichever config fields exist (num_local_experts / num_experts, moe_intermediate_size / intermediate_size) instead of assuming one naming scheme or a transformers version gate. Routed-expert FLOPs and which layers count as sparse now follow transformers: (layer_idx + 1) % decoder_sparse_step == 0, respect mlp_only_layers, and default sparse step to 1 when unset.

Tests add coverage for both config spellings, Mixtral expert sizing via intermediate_size, and sparse-step / MLP-only layering; the existing active-experts test is updated to the corrected layer rule.

^{Reviewed by Cursor Bugbot for commit 68e09c1. Bugbot is set up for automated code reviews on this repo. Configure here.}