feat: support MTP with context and pipeline parallelism

[HuiyingLi]

What does this PR do ?

Support the generic MTP + context-parallel + pipeline-parallel training combination, including the memory-saving fused linear cross-entropy path.

This is PR 2/3 in the HY4 support stack and depends on #3763. It contains only reusable training/parallelism infrastructure; no HY4 model code is included.

Changelog

• Prepare MTP inputs and targets from the global packed sequence before CP sharding, including the PP path at CP1 so CP1 and CP2 share one loss contract.
• Carry authoritative CP-local, per-depth MTP targets through the last pipeline stage instead of shifting rank-local labels.
• Allow models that explicitly opt in to combine pipeline MTP with FusedLinearCrossEntropy.
• Discover checkpoint-compatible MTP blocks stored directly on an inner decoder for EP/FSDP/activation-checkpointing traversal.
• Add unit coverage for target extraction, packed boundaries, CP+PP capability gating, fused loss, and inner-decoder MTP traversal.

Validation

GPU interactive allocation: Slurm job 16995628, H100 80GB.

• ruff check and ruff format --check: passed.
• Relevant unit suites across loss, recipe, MoE parallelizer, shared traversal, PP+MTP, distributed helpers, and validation: 343 passed, 1 skipped.
• 2-GPU packed CP + fused-linear-CE + FSDP correctness test: 1 passed.

The validation command covered:

• tests/unit_tests/loss/test_mtp_lm_head_gather.py
• tests/unit_tests/moe/test_parallelizer.py
• tests/unit_tests/recipes/test_train_ft.py
• tests/unit_tests/shared/test_model_utils.py
• tests/unit_tests/models/nemotron_v3/test_nemotron_v3_pp_mtp.py
• tests/unit_tests/distributed/test_parallelizer_utils.py
• tests/unit_tests/test_validation.py
• tests/functional_tests/training/test_cp_flce_fsdp_correctness.py

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation? (No new recipe surface; capability remains model opt-in.)

Additional Information

Validation logs:

• logs/hy_v4_prs/pr2-mtp-cp-pp-scoped-interactive-16995628.log
• logs/hy_v4_prs/pr2-mtp-cp-pp-interactive-16995628.log

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