Add configurable batch and sequence partitioning within GTP_remat

[Mellonta]

Summary

Adds --tensor-parallel-num-sequence-shards to configure batch and sequence partitioning within GTP_remat groups. By default, GTP_remat ranks process independent microbatches. Increasing the sequence-shard count partitions each sequence across more ranks and reduces the number of independent batch partitions.

Weight sharding remains controlled by --tensor-parallel-num-weight-shards. The new sequence-shard flag includes TP when sequence parallelism (SP) is enabled:

SP degree = TP with SP enabled, otherwise 1

gtp_remat_num_sequence_shards = tensor_parallel_num_sequence_shards / SP degree
gtp_remat_num_batch_shards = GTP_remat_size / gtp_remat_num_sequence_shards
batch_parallel_size = DP × gtp_remat_num_batch_shards

Both divisions must be exact. Leaving the flag unset preserves existing behavior: the sequence-shard count defaults to the SP degree, so gtp_remat_num_sequence_shards=1 and gtp_remat_num_batch_shards=GTP_remat_size. Sequence partitioning within GTP_remat requires SP when TP > 1.

For example, weight shards=64, TP=2 with SP enabled, and sequence shards=16 produce 4 independent batch partitions, each with 2 SP partitions × 8 sequence partitions within GTP_remat, in each TP × GTP_remat weight-sharding group.

--context-parallel-size adds further sequence partitions across weight replicas. The resolved runtime CP size excludes TP/SP:

runtime_CP = configured_CP × gtp_remat_num_sequence_shards

This allows sequence partitioning to span more ranks than a single weight-sharding group.

The implementation introduces GTPParallelLayout, which describes two views of the same ranks:

• Weight layout: TP × GTP_remat shards weights; configured CP × DP replicates those shards. These groups drive weight materialization, gradient reduction, and optimizer-state sharding.
• Batch and sequence layout: DP × gtp_remat_num_batch_shards distributes independent microbatches, while configured CP × gtp_remat_num_sequence_shards partitions each sequence, in addition to TP/SP. These groups drive data distribution, CP communication, and sample accounting.

Argument validation resolves the runtime CP size, and process-group construction builds the existing groups from these views. Global batch size is calculated as:

global_batch_size = micro_batch_size × batch_parallel_size × num_microbatches

Training progress and sample-based schedules use the same independent-sample count.

For example, on 128 GPUs with PP=1:

--tensor-model-parallel-size 1 \
--tensor-parallel-num-weight-shards 64 \
--tensor-parallel-num-sequence-shards 64 \
--context-parallel-size 2

This gives DP=1, with weights sharded across GTP_remat groups of 64 ranks and two replicas of each weight shard. Each sequence is partitioned across 128 ranks. The batch dimension is unpartitioned (batch_parallel_size=1), so all 128 ranks process different sequence partitions of the same microbatch.

This placement example assumes 64 GPUs per NVLink domain. The current implementation supports rank order tp-cp-ep-dp-pp, with TP as the innermost axis and PP as the outermost axis. Place each 64-rank weight-sharding group entirely within one NVLink domain. CP then spans both domains, while weight all-gathers and gradient reduce-scatters stay within each domain.

Comparing with PR #6886

PR #6886 folds the entire CP axis into the weight-sharding group, so its folded weight group is at least as large as its CP degree. This PR allows 128 sequence partitions with 64-rank weight groups, as in the example above.

PR #6886 additionally preserves optimizer-state sharding across CP for parameters not sharded by GTP, such as norms and biases. This implementation uses one dense weight-replica group, so those optimizer states remain replicated within each GTP group, although they can still be sharded across weight replicas. This trades some memory efficiency for simpler ownership and reduction rules.

The flag values also differ: for TP=1, #6886's CP=64 / GTP_remat=1 corresponds here to weight shards=64, sequence shards=64, and CP=1.

The changes are deliberately minimal: resolve weight and token layouts during process-group construction, update sample accounting, and reuse existing GTP, DDP, and optimizer paths. This covers the target static-CP use case without adding CP-specific parameter flags, buffer routing, or optimizer checkpoint keys, reducing the number of interacting special cases to validate and maintain.

Test

Layout, process-group, and sample-accounting tests:

uv run python -m torch.distributed.run --nproc-per-node 8 -m pytest -q \
  tests/unit_tests/test_gtp_parallel_layout.py \
  tests/unit_tests/test_process_groups_config.py \
  tests/unit_tests/test_eval_batch_size.py \
  tests/unit_tests/test_training.py::TestPackedSampleAccounting::test_static_cp_sample_count_and_scheduler

GPU coverage includes attention gradients, Adam updates, optimizer groups, and model/optimizer checkpoint loading after changing CP. Run both process counts to cover all parameterized layouts:

for n in 4 8; do
  uv run python -m torch.distributed.run --nproc-per-node "$n" -m pytest -q \
    tests/unit_tests/generalized_tensor_parallel/test_gtp_cp_overlap.py
done

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

[Mellonta]

/ok to test dcd3bb7

[Mellonta]

/ok to test 0cf852d

[deepakn94]

What is "sample_parallel_size" now? :O

[Mellonta]

This is a helper function to replace args.data_parallel_size * args.gtp_weight_remat_size.
Without folding: DP × GTP (old codes)
With folding: DP × GTP / gcd(CP, GTP) because some GTP ranks also serve as CP ranks

[deepakn94]

replica_size or num_weight_replicas?

[Mellonta]

/ok to test 4cad561

[fanshiqing]

/review

[svcnvidia-nemo-ci]

Apply sequence-shard defaults before resolving the layout

Failure: In validate_args, an unset tensor_parallel_num_sequence_shards is immediately resolved to 1 (or TP with SP). The later programmatic-defaults loop therefore treats it as explicitly supplied and ignores args_defaults for the new flag.

For example, with world=8, TP=PP=1, weight shards=2, configured CP=1, and ckpt_format="torch_dist", calling validate_args(args, {"tensor_parallel_num_sequence_shards": 2}) leaves sequence shards=1, runtime CP=1, and batch parallel size=8. Expected: sequence shards=2, runtime CP=2, and batch parallel size=4. Bootstrap consequently constructs a different token layout from the caller's requested default.

Correction: Apply this particular default while the argument still retains its unset state, before the sequence-shard resolver. Keep the invariant that explicit CLI values win and all layout arithmetic uses the effective sequence-shard value. Suggested insertion immediately before line 450:

if args.tensor_parallel_num_sequence_shards is None:
    args.tensor_parallel_num_sequence_shards = defaults.get(
        "tensor_parallel_num_sequence_shards"
    )

This is a source-reasoned suggestion; verify it before adopting it.

Preservation and acceptance — proposed checks, not executed:

• Extend the argument-validation tests with the setup above. Assert sequence shards=2, GTP sequence shards=2, runtime CP=2, data parallel size=4, and batch parallel size=4.
• With the same setup but an explicit sequence-shard value of 1 and a programmatic default of 2, retain sequence shards=1, runtime CP=1, and batch parallel size=8.
• With neither an explicit value nor a programmatic default, retain the existing TP/SP defaults and the runtime-CP/batch-size assertions in test_sequence_shard_cli_resolves_runtime_cp.

Completed the light review across all changed files. Tests could not be executed in this read-only review.

Review generated by codex

[zhongbozhu]

Nit: hybrid_context_parallel is an unfortunate config that's going to retire (and therefore shouldn't be added in the first place :) ) #7047 , we can remove hybrid_context_parallel checks

doesn't seem to matter

[fanshiqing]

@Mellonta Can we report this incompatibility in advance (e.g., in the argumentts.py around GTP's argument check) instead?

cc @kvareddy

[Phlip79]

Please move to megatron/core/tensor_parallel/ with the other gtp files.

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38375109130

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38375470909