perf(magi): avoid TE padding for magi

[HuiyingLi]

What does this PR do ?

Avoid Transformer Engine-specific per-document CP padding when the attention backend owns context-parallel dispatch, including the custom Magi backend. Add a two-layer Qwen3-MoE exact-geometry recipe that exercises packed THD CP with the production attention, router, expert, and dispatcher dimensions.

TE block-diagonal THD CP requires each packed document to be aligned to 2 * cp_size. Magi builds and dispatches its own arbitrary-mask layout, so retaining that TE-only alignment wastes attention work and reduces the useful label tokens that fit in a pack.

packing_cp_size=1 changes only dataset packing alignment. Runtime distributed.cp_size remains CP8 or CP16.

Changelog

• Select packer CP alignment from attention ownership: model-owned CP and custom Magi use the baseline CP1 alignment; TE continues using the configured CP size.
• Add qwen3_moe_2layer_magi_packed_cp8_32k.yaml, preserving Qwen3-30B-A3B head, expert, router, vocabulary, DeepEP, TE-expert, activation-checkpointing, backward, and optimizer paths while reducing depth to two layers for a practical gate.
• Add parameterized unit coverage for TE, model-owned CP, custom Magi, and non-custom Magi behavior.

Validation

• ruff format --check nemo_automodel/recipes/llm/train_ft.py tests/unit_tests/recipes/test_train_ft.py
• ruff check nemo_automodel/recipes/llm/train_ft.py tests/unit_tests/recipes/test_train_ft.py
• python tools/lint_example_yamls.py examples/llm_finetune/qwen/qwen3_moe_2layer_magi_packed_cp8_32k.yaml
• python -m pytest tests/unit_tests/recipes/test_train_ft.py -q — 96 passed, 7 skipped
• Single-node CP8 forward/backward/optimizer gate before multi-node CP16
• Fresh W&B runs, 30 optimizer steps each; steady window is steps 10–29

Qwen3-MoE case	Magi physical vs TE	Magi effective-label vs TE
2-layer exact geometry, CP8 / 64K	+9.81%	+18.29%
Actual 48-layer Qwen3-30B-A3B, CP8 / 32K	+63.17%	+88.27%
2-layer exact geometry, CP16 / 64K	-29.35%	-16.11%

W&B workspace: https://wandb.ai/Nemo-automodel/magi

Before your PR is "Ready for review"

Pre checks:

• Read and followed the contributor guidelines.
• Added focused unit coverage and an end-to-end benchmark recipe.
• Added the example recipe needed to document and reproduce the path.

Additional Information

• Base: main at 3ddef9b1.
• Head commit is DCO signed off.
• This custom-model path does not depend on HF Magi PR feat(magi): honor AttnMaskSpec on the HF attention backend #2622; that change is intentionally absent from this diff.

[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.

[HuiyingLi]

/ok to test bf999ba

[HuiyingLi]

/claude review

[claude]

LGTM