[yaoyu-33]

What does this PR do ?

Fixes Gated DeltaNet FLOPs accounting when a decoder-only attention layout is combined with Multi-Token Prediction (MTP), unblocking the Qwen3.5-VL L1 recipe test on main.

Changelog

• Generate and validate linear_attention_freq against decoder layers only; Hugging Face layer layouts do not include MTP layers.
• Account for every MTP layer using the final decoder layer's attention type, matching MCore's MTP block construction.
• Add a parameterized unit regression for integer/list layouts whose final decoder layer is either standard attention or GDN, with an exact FLOPs-delta assertion across two MTP depths.

GitHub Actions CI

Validated in the repository container on EOS:

• PASS: uv run python -m pytest tests/unit_tests/training/utils/test_flop_utils.py::TestGDNLayerFlops::test_gdn_mtp_reuses_last_decoder_attention_type -q (4 passed)
• PASS: uv run python -m pytest tests/unit_tests/training/utils/test_flop_utils.py::TestGDNLayerFlops -q (9 passed)
• PASS: mutation check: removing the MTP GDN-layer term makes both final-GDN cases fail (2 failed), then both pass after restoring the implementation.
• PASS: 2-GPU H100 TestQwen35VLFinetuneRecipes::test_sft_nothing_frozen smoke train (1 passed, all 10 iterations plus validation/test evaluation)
• PASS: uv run pre-commit run --all-files

Before your PR is "Ready for review"

Pre checks:

• Read and followed the contributor guidelines.
• Added a unit regression covering the broken integer/list plus MTP paths.
• No documentation update is needed; this restores existing internal FLOPs semantics.
• The change does not affect optional-install components or import guards.

Additional Information

• Regression introduced by [model] fix: preserve Qwen hybrid attention schedules #4790 (2f69b067748a0842521a527b59ace55e18df7f44): exact Hugging Face Qwen attention schedules became decoder-length lists.
• Last passing target job: https://github.com/NVIDIA-NeMo/Megatron-Bridge/actions/runs/29123303699/job/86496304191
• First failing target job at the bad commit: https://github.com/NVIDIA-NeMo/Megatron-Bridge/actions/runs/29135911454/job/86513808812
• test(model): fix AutoBridge and Qwen VL regressions #4821 reduced the functional-test list from 64 entries to 4, but Bridge FLOPs accounting still required 5 entries after adding mtp_num_layers=1.
• Current failing main job fixed by this PR: https://github.com/NVIDIA-NeMo/Megatron-Bridge/actions/runs/29176488641/job/86607551503

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

/ok to test 58631b4

[claude]

Line one
Line two

[claude]

Review of fix for GDN FLOP count on MTP layers.

The fix is correct. When the attention variant is gated_delta_net and mtp_num_layers is greater than zero, MTP layers were previously counted only as standard-attention layers (via num_standard_attn_layers = num_layers - num_gdn_layers), even though MTP construction reuses the final decoder layer spec. Adding last_layer_is_gdn * mtp_num_layers to num_gdn_layers mirrors the existing MoE handling (last_layer_is_moe at line 726) and keeps the GDN/standard split consistent. Using cfg.model.num_layers for the pattern length is also correct, since num_layers includes the MTP layers.

One test-coverage gap worth closing (also flagged inline): both parametrized values in test_gdn_mtp_repeats_last_decoder_attention resolve to pattern [1,1,1,0] for num_layers=4, so the last decoder layer is always standard-attention (last_layer_is_gdn == 0). The new last_layer_is_gdn * mtp_num_layers term is never exercised with a nonzero value. Add a case where the final decoder layer is GDN (e.g. linear_attention_freq=[1,1,0,1]) with expected_delta computed from the GDN per-layer cost — that is the exact path this PR fixes.

Suggested test cases:

• tests/unit_tests/training/utils/test_flop_utils.py::test_gdn_mtp_repeats_last_decoder_attention (existing; extend with a GDN-final-layer parametrization)
• No perf tests impacted.

[yaoyu-33]

/ok to test 683abde

[yaoyu-33]

Post-merge cross-calculator audit evidence

An independent AutoModel ↔ Megatron Bridge theoretical-FLOPs audit reproduced
the Qwen3.5-35B-A3B + MTP error fixed by this PR exactly.

Using identical parameters and GBS 1 on the pre-fix frozen Bridge revision
77a43e56a192c6e6eaedf62fb91df9da270d5a3f:

Calculator	TF/step
AutoModel	92.041686
Bridge before #4834	91.839705
Bridge relative difference	-0.219445%

The difference decomposes as follows:

• Correct full-attention MTP layer: 1.082331758592 TF
• Incorrect GDN substitute: 0.880199860224 TF
• Attention-type classification error: -0.202131898368 TF
• Existing Bridge-only partial MTP norm proxy: +0.000150994944 TF
• Net Bridge minus AutoModel before this PR: -0.201980903424 TF

This exactly matches the reported -0.219445%. The implementation in #4834
removes the 0.202131898368 TF architecture error by deriving the hybrid
pattern from decoder layers only and reusing the final decoder attention type
for MTP. The remaining approximately 0.000151 TF difference is only the
previously documented partial norm-counting convention, not a layer-layout
error.

Related evidence and follow-up:

• Full TFLOPs/MFU audit on MB-812
• AutoModel #3091 fixes the independent GPT-OSS head_dim issue found by the same audit.

No additional Bridge PR is needed because #4834 already contains both the
correct fix and focused integer/list plus MTP regression coverage.