[feature] Integrate Megatron-Core Generalized Tensor Parallelism (GTP) #5434

User problem

Megatron Bridge main now pins Megatron-Core 6513e3e23, which contains the experimental Generalized Tensor Parallelism (GTP) implementation introduced by Megatron-LM #4967. However, Bridge does not currently provide functional end-to-end GTP enablement through its training configuration and recipe path.

Bridge providers inherit MCore's user-facing tensor_parallel_num_weight_shards and expert_tensor_parallel_num_weight_shards fields, but the rest of the runtime contract is not wired:

• training/initialize.py::_initialize_distributed() does not pass the derived gtp_weight_remat_size / expert_gtp_weight_remat_size into parallel_state.initialize_model_parallel().
• The Bridge model-construction path does not invoke MCore's configure_gtp_remat_from_recipe() and classify_gtp_remat_chains() lifecycle around model construction.
• Bridge batch-size, consumed-sample, LR-scheduler, logging, and evaluation accounting do not yet reflect MCore's DP × GTP_remat data-distribution contract.
• The decentralized HyperCommGrid path currently has no GTP axes. Existing MIMO compatibility aliases collapse GTP-inclusive groups to non-GTP groups, which is not GTP enablement.
• There is no supported Bridge recipe, documentation, or functional validation for dense GTP or MoE expert-GTP.

As a result, users cannot safely enable the MCore feature by setting the inherited config fields alone.

Desired outcome

Add first-class Megatron Bridge integration for MCore GTP while keeping MCore as the owner of the implementation.

Runtime and configuration

• Expose the MCore public knobs through normal Bridge recipe overrides:
  • tensor_parallel_num_weight_shards
  • expert_tensor_parallel_num_weight_shards
• Preserve MCore's derivation and validation of gtp_weight_remat_size / expert_gtp_weight_remat_size, including divisibility constraints.
• Pass the derived dense and expert GTP-remat sizes into standard model-parallel process-group initialization.
• Configure GTP before model construction and classify prefetch chains after all model chunks are built, including VPP/multi-chunk models.
• Reset GTP global state when Bridge rebuilds models in-process.
• Correct data distribution, global-batch, consumed-sample, LR-scheduler, logging, and evaluation microbatch accounting for the GTP axes.
• Integrate with the existing Bridge DDP/finalize-grad, optimizer, and torch_dist checkpoint paths without introducing a separate checkpoint format.
• Validate required TransformerEngine and precision combinations, including MXFP8/NVFP4 param-gather requirements and high-priority GTP stream groups.
• For decentralized process groups, either implement real GTP axes or fail early with a clear unsupported-configuration error; do not silently collapse the topology.

Recipes and documentation

• Add at least one supported Bridge training recipe that demonstrates GTP on current hardware.
• Document topology arithmetic, memory/communication tradeoffs, required precision flags, stream-priority settings, and interactions with TP/SP/EP/DDP, activation recompute, and CUDA graphs.
• Clearly mark the feature experimental while the upstream MCore API remains experimental.

Validation / acceptance criteria

• Config tests cover dense/expert weight-shard derivation and invalid divisibility.
• A targeted distributed BF16 smoke demonstrates finite training with GTP active and confirms the expected process-group sizes.
• A focused MoE smoke covers expert-GTP initialization and gradients.
• GTP-off remains equivalent to the existing TP/DP path.
• Save/load works through the standard distributed checkpoint path, including a supported topology-reshard case.
• Batch-size, scheduler increments, consumed samples, and evaluation counts are correct with GTP enabled.
• Unsupported process-group/precision combinations fail before training with actionable messages.
• Documentation and a representative recipe are included.

Alternatives considered

• Use Megatron-LM's launcher directly. This exercises GTP but bypasses Bridge recipes, data, checkpoint orchestration, and model-support workflows.
• Add recipe-local calls into MCore's GTP API. This would be fragile because process groups, model lifecycle, sample accounting, and checkpoint behavior span the shared Bridge training runtime.
• Use TP, Megatron FSDP, or activation recompute instead. These can solve related memory problems but do not provide GTP's per-weight rematerialization and orthogonal TP/GTP topology.

Affected area

area:training

Urgency / use case

Important but not blocking

Extra context

• Upstream implementation: NVIDIA/Megatron-LM#4967
• Upstream documentation: Generalized Tensor Parallelism
• Related upstream follow-ups: MTP support #6078, documentation #6236
• Bridge main MCore pin checked on 2026-08-07: 6513e3e23d6b5eda6a1c934990b15e804237732b
• No existing Megatron Bridge GTP tracking issue or implementation PR was found before filing this request.