feat(training): wire generalized tensor parallelism

[yaoyu-33]

Summary

• forward the dense and expert GTP rematerialization sizes derived from the existing weight-shard configuration into standard MCore process-group initialization
• configure GTP before model construction and classify all model/VPP chunks after distributed wrapping
• use the GTP-inclusive data-distribution groups for dataloaders, batch/sample accounting, training/evaluation loss reductions, and packed-sequence statistics
• reject decentralized process groups and collective averaging early when GTP is active, and report the Transformer Engine requirement actionably

This is Stage 1 of #5434. It intentionally keeps GTP implementation ownership in Megatron Core.

Validation

Pinned revisions:

• Bridge base: 064b77054eed990778d7ab4f092083449445b764
• Megatron-LM: 14346b65a2d0790e451919858f7771078105c5f0
• Transformer Engine main: c280e7cb192f57a182a6aab64d3d6bd40d933d92
• runtime: PyTorch 2.13.0a0+8145d630e8.nv26.6, Megatron Core 0.20.0+14346b65a, Transformer Engine 2.19.0.dev0

Targeted results:

• 105 passed across GTP, distributed initialization, setup, process-group wiring, and decentralized-group unit tests
• 2-GPU BF16 Bridge smoke: three finite steps with GTP=2, classified GTP parameters, full data-distribution size 2, replicate-DP size 1, consumed samples 2, 4, 6, and no skipped/NaN iterations
• GTP-off control: finite two-rank DP training with no GTP parameters
• uv run --active --no-sync pre-commit run --all-files

Reproducible Bridge/MLM scalar comparison

Both runs used a 2-layer GPT (hidden size 128, FFN size 256, 4 heads, sequence length 64, vocabulary size 128), TP1/PP1/CP1/GTP2, BF16, mock data, micro/global batch sizes 1/2, seed 1234, dropout 0, untied embeddings/output weights, no linear bias, attention softmax in FP32, weight decay 0.01, clip grad 1.0, and the same three-step cosine schedule (3e-4, 1.65e-4, 3e-5; minimum 3e-5; one warmup step). The optimizer was Adam with betas 0.9/0.999 and epsilon 1e-8, without distributed optimizer.

Bridge command:

uv run --active --no-sync python -m torch.distributed.run --nproc_per_node=2 --master_port=29555 \
  -m pytest -q -s \
  tests/functional_tests/test_groups/training/test_pretrain.py::TestPretrain::test_pretrain_with_generalized_tensor_parallelism

Megatron-LM command:

uv run --active --no-sync python -m torch.distributed.run --nproc_per_node=2 --master_port=29556 \
  3rdparty/Megatron-LM/pretrain_gpt.py \
  --mock-data --tokenizer-type NullTokenizer --vocab-size 128 \
  --num-layers 2 --hidden-size 128 --ffn-hidden-size 256 --num-attention-heads 4 \
  --seq-length 64 --max-position-embeddings 64 \
  --micro-batch-size 1 --global-batch-size 2 --train-iters 3 \
  --eval-interval 100 --eval-iters 0 \
  --lr 3e-4 --min-lr 3e-5 --lr-warmup-iters 1 --lr-decay-iters 3 --lr-decay-style cosine \
  --weight-decay 0.01 --weight-decay-incr-style constant --clip-grad 1.0 \
  --bf16 --disable-bias-linear --untie-embeddings-and-output-weights \
  --no-gradient-accumulation-fusion --attention-dropout 0 --hidden-dropout 0 \
  --attention-softmax-in-fp32 \
  --tensor-model-parallel-size 1 --pipeline-model-parallel-size 1 \
  --tensor-parallel-num-weight-shards 2 \
  --seed 1234 --log-interval 1 --dataloader-type single --num-workers 0

GTP is enabled from the same source setting in both frontends: Bridge sets tensor_model_parallel_size=1 and tensor_parallel_num_weight_shards=2 on its model config; MLM uses the corresponding two CLI flags. MCore derives gtp_weight_remat_size = tensor_parallel_num_weight_shards / tensor_model_parallel_size = 2; the derived field is not set directly.

Step	Samples	Bridge loss	MLM loss	Bridge grad	MLM grad	LR (both)
1	2	4.860444	4.860444	2.126	2.126	3.00e-4
2	4	4.827285	4.843864	2.480	2.480	1.65e-4
3	6	4.781815	4.781815	2.413	2.413	3.00e-5

The step-2 absolute loss difference is 0.016579, or 0.342268% relative to MLM. It reproduced across reruns. Matching dropout and FP32-softmax settings aligned step 1 and the displayed grad norms but did not remove step 2. Both frontends reduce token-summed loss and token counts over the GTP-inclusive data-distribution group. Disabling cross-entropy fusion did not remove the Bridge value, and a matched SGD control retained the same step-2 difference, so it is not explained solely by Adam state. A GTP-off run changes data topology and is not a clean isolation control.

These are rounded, reported scalar observations only. No per-rank inputs, logits, tensors, parameter deltas, per-parameter gradients, or optimizer slots were captured, so this does not establish bitwise/tensor parity or rule out transient state drift. The exact source of the step-2 discrepancy remains unresolved. The checked-in functional test therefore asserts the intended GTP topology, nonempty GTP parameter classification, progress, and finite results; it is deliberately not a numerical-parity assertion.

Deferred

• decentralized HyperCommGrid GTP support
• broader MoE/expert-GTP, PP/CP/VPP, checkpoint-topology, recipes/docs, and performance coverage
• upstream full-GPT limitations for bias-enabled linears and tied embedding/output weights; the correctness smoke follows upstream GTP coverage by disabling bias and using untied embeddings

2026-08-17 latest-main refresh

• merged current main (46c42552bcb7ae6be22d7e117de0776775371d48) normally through DCO-signed merge commits, with final head fdc651cab1362d760fd6027c8fc7ee55c4df7764
• the substantive setup conflict was resolved as provider finalize -> configure GTP -> build distributed model -> classify GTP chains, preserving current-main provider finalization and the Stage 1 GTP lifecycle; a later main-only chat-processing commit merged without conflicts or GTP-path overlap
• final-tree validation: 106 focused unit tests passed; the checked-in 2-GPU BF16 GTP functional passed; a 2-GPU GTP-off control passed; and all all-files pre-commit hooks passed
• final-tree Bridge retained loss (4.860444, 4.827285, 4.781815), grad (2.126, 2.480, 2.413), LR (3e-4, 1.65e-4, 3e-5), and consumed samples (2, 4, 6)
• the exact MLM comparison was rerun on the final tree and retained loss (4.860444, 4.843864, 4.781815) with the same displayed grad/LR/sample sequence; the step-2 loss difference remains 0.016579 (0.342268%) and unexplained, so this remains scalar smoke evidence rather than bitwise or tensor-parity evidence

[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 008670b

[yaoyu-33]

/ok to test 40b50b4

[yaoyu-33]

/ok to test 8edbc78

[yaoyu-33]

/ok to test 8edbc78

[yaoyu-33]

/ok to test 152fced

[yaoyu-33]

/ok to test fdc651c

[yaoyu-33]

/ok to test 54bbc27

[cuichenx]

Looks good.
Several question that cross my mind:

• Do we have GTP on/off training experiments and verify loss identical?
• Is GTP supported for conversion? if not should we explicitly disallow it.