feat(megatron): support GTP weight rematerialization for training + refit

[shanmugamr1992]

What this does

Adds support for Megatron-Core Generalized Tensor Parallelism (GTP) in NeMo-RL, including the refit path, so a GRPO run can train with GTP-sharded weights and still refit into a whole-weight inference model.

GTP shards each weight along dim 0 across a gtp_remat process group carved out of the DP axis, and all-gathers ("rematerializes") it on demand in both forward and backward. The world decomposition becomes world = TP x GTP x CP x DP x PP. It is exercised against Megatron-LM#6133, which added the refit support on the MCore side and has since merged.

User-facing config

Two new optional policy.megatron_cfg keys, both defaulting to null (= GTP off, i.e. no behaviour change for existing configs):

policy:
  megatron_cfg:
    tensor_parallel_num_weight_shards: null         # dense axis
    expert_tensor_parallel_num_weight_shards: null  # expert axis

These mirror MCore's own ModelParallelConfig names. setup.py resolves them via MCore's resolve_tensor_parallel_weight_shards() into the derived gtp_weight_remat_size = num_weight_shards // tp_size. Following config-conventions, the keys are NotRequired on the (still v1 TypedDict) MegatronConfig with the default documented in the exemplar YAML only -- no .get(key, default) at any call site.

Why the enablement is a shim in NeMo-RL

Megatron-Bridge does not yet forward gtp_remat_size / expert_gtp_remat_size to parallel_state.initialize_model_parallel (src/megatron/bridge/training/initialize.py:805). Rather than take a cross-repo dependency on an unreleased Bridge change, setup_model_and_optimizer wraps Bridge's init in a narrowly-scoped context manager that monkeypatches parallel_state.initialize_model_parallel to inject the two sizes, then restores the original in a finally.

The patch uses functools.wraps deliberately: Bridge probes inspect.signature(parallel_state.initialize_model_parallel).parameters a few lines earlier (initialize.py:794), and an unwrapped wrapper would change what that probe sees.

It is marked with a TODO to delete once Bridge forwards these natively. Reviewers should feel free to push back if you would rather this wait for Bridge.

Inference deliberately pins GTP off (build_inference_model sets the shard counts equal to their TP degree, gtp_weight_remat_size = 1): refit reassembles the training model's dim-0 shards into whole inference weights, so the inference model never shards.

Bug found along the way: the refit path was being skipped silently

This is the part most worth reviewing, and it is a pre-existing bug, not one this feature introduced.

dedicated_inference_megatron_cfg() decides whether colocated generation runs directly on the shared training model, or whether the worker must build a second, resharded inference model (and therefore refit). Its layout comparison did not include the weight-shard counts. So a TP1 x GTP2 training model compared equal to a TP1 inference model, the function returned None, no dedicated inference model was built, and no refit ever happened.

The failure mode is silent, which is what makes it nasty: generation still produces tokens, the metrics still look fine, the test still exits 0. The new functional test would have passed while exercising exactly none of the code it was written for.

Fixed by adding tensor_parallel_num_weight_shards and expert_tensor_parallel_num_weight_shards to layout_keys, normalizing null -> TP degree before comparing so that null and an explicit TP-equal value are treated as the same unsharded layout (otherwise writing the default out longhand would cost you a pointless dedicated inference model and a full refit every wake).

Pinned by a new CPU-only unit test, tests/unit/models/generation/test_megatron_generation_layout.py (8 assertions, no GPU, runs in the normal unit suite) so this cannot silently regress again.

The functional test

tests/functional/grpo_megatron_generation_gtp.sh -- GRPO on Qwen/Qwen2.5-0.5B, cluster.gpus_per_node=2, training TP=1 + tensor_parallel_num_weight_shards=2 (GTP=2), megatron generation backend, colocated, refit_backend=nccl, 2 steps.

Asserts max(train/token_mult_prob_error) < 1.05, plus two anti-vacuity greps on the run log, added precisely because of the bug above:

• [gtp] enabling GTP weight rematerialization (gtp_remat_size=2 -- proves GTP actually got turned on
• [colocated-reshard] building dedicated inference model -- proves the reshard/refit path actually ran

⚠️ Merge blockers, all outside this repo

The feature is code-complete here, but the configuration it enables cannot run on the pinned dependencies. Three things must land first; none of them belong in this PR.

1. TransformerEngine >= 2.19.0.dev0

_GTP_TE_MIN_VERSION = Version("2.19.0.dev0") is declared in megatron/core/tensor_parallel/generalized_tensor_parallelism.py. NeMo-RL pins TE v2.14.1 / release_v2.15. This floor predates PR 6133 -- it is present at the image's own mcore pin.

Below the floor MCore does not raise. It catches its own ImportError and rebinds every TE name to a stub, including QuantizedTensor = MagicMock() -- an instance, not a type -- and sets HAVE_TE = False -> HAVE_GTP = False. Nothing downstream re-checks, because MCore's layer constructors gate GTP on the process-group size alone. The result was a run dying six frames deep on:

TypeError: isinstance() arg 2 must be a type, a tuple of types, or a union
  at generalized_tensor_parallelism.py:702

which says nothing whatsoever about TransformerEngine.

Two mitigations, both self-retiring, are in this PR:

1. HAVE_GTP preflight in the shim (setup.py) turns that into a named error stating the actual requirement and how to opt out.
2. megatron_gtp_supported() capability guard in L1_Functional_Tests_Megatron_4.sh skips the test rather than leaving L1 permanently red. It asks MCore directly (from megatron.core.tensor_parallel.gtp_api import HAVE_GTP) instead of hardcoding a version, so the moment the TE pin is bumped, the test re-enables itself with no edit to this branch.

Bumping TE needs a container rebuild and affects every backend, so it is deliberately not in this PR. TE release_v2.19 was built into the test venvs out-of-band to get the validation below.

2. NVIDIA/Megatron-LM#6940 (draft) -- four GTP bugs

Found while validating this branch. All four are megatron-core bugs, none are NeMo-RL's:

• TE sizes its bias off the pre-sharded out_features _gtp_pre_init passes it, so the bias comes out 1/gtp_remat too small -- but a bias is GTP-replicated. Breaks the GEMM epilogue and declares a wrong global shape to dist-checkpointing.
• GPTModel passed only tp_group (not pg_collection) to LanguageModelEmbedding and the output layer, so both resolved their GTP axis from the global MPU groups instead of the model's own. TELMHeadColumnParallelLinear did not accept pg_collection at all.
• build_inference_pg_collection left gtp_remat unset, and resolve_gtp_remat_group falls back to the MPU globals for an absent field -- so the dedicated inference model came out GTP-sharded from the training layout.
• ColumnParallelLinear's supplied-weight shape guard compared a GTP shard's local shape against the logical one. With tied embeddings, shared_embedding_or_output_weight() hands the output layer exactly that shard. This one breaks any GTP run with share_embeddings_and_output_weights=True, independent of NeMo-RL.

The last two both surfaced as supplied weight's shape is (75968, 896), not (151936, 896) as expected.

3. NVIDIA-NeMo/Megatron-Bridge#5862 (draft) -- checkpoint replica group

Bridge derives sharded-checkpoint replica_id from pg_collection.dp_cp, which is the GTP replicate group and excludes the gtp_remat axis. Every gtp_remat peer then claims replica_id 0 for the same non-GTP-sharded tensor (norms, biases, _extra_state), and dist-checkpointing rejects it with Invalid access pattern for ShardedTensor(...). No-op when GTP is off.

What was validated on hardware

OCI HSG GB200 cluster (aarch64, 4 GPUs/node), with TE release_v2.19 installed into the megatron venvs and the two draft fix PRs above applied.

With all of the above in place the test runs end-to-end: GTP group built, dedicated inference model built, reshard + NCCL refit, two full GRPO steps, both anti-vacuity greps hit. Before the megatron-core fixes it could not even construct the models.

Earlier attempts, kept for the record:

#	Result	Finding
a1	rc=1, 161s	uv re-resolved and failed to build torch-memory-saver (TMS_CUDA_MAJOR env var must be set). Test-harness issue, not product: mcore is an editable uv workspace member, so checking out the revision under test staled the image's uv.lock. Fixed in the harness with UV_FROZEN=1.
a2	rc=1, 83s	The isinstance() TypeError above. Crucially, both ranks first printed [gtp] enabling GTP weight rematerialization (gtp_remat_size=2 expert_gtp_remat_size=1) -- so Bridge forwarded the sizes and MCore built the gtp_remat group. The shim works; only the TE-gated kernels fail.
a3	rc=1, 75s	Confirms the new HAVE_GTP preflight fires cleanly on both ranks with the actionable message instead of the cryptic one.

No assertion was loosened, no threshold moved, no test skipped or shortened to manufacture a green run.

One open question the test does not answer

The end-to-end run's train/token_mult_prob_error measures 1.1011, against the < 1.05 threshold this test inherited by copy from its non-GTP sibling. The value is bit-identical across repeated runs, so it is a deterministic train-vs-generation logprob disagreement rather than noise.

I could not attribute it. The non-GTP control cannot exercise this code path -- with GTP off, colocated generation takes the shared-training-model branch instead of building a dedicated inference model, and that branch hits an unrelated CUDA-graph capture error -- so there is no measured baseline for the 1.05 gate on this configuration. It was never validated here; it was copied. Several other non-vanilla Megatron generation variants sit similarly above their own 1.05-class gates.

Flagging it rather than moving it: whether the GTP path adds real numerical error, or the inherited threshold is simply wrong for it, is a call for a reviewer who owns this path. If the answer is "the threshold is wrong", say so on this PR and I will change it with that rationale recorded.

Review notes

• nemo_rl/models/generation/megatron/config.py is in pyrefly.toml; nemo_rl/models/megatron/setup.py is not (pre-existing).
• The new YAML keys are not in tests/unit/reference_configs/, so no reference-config update was needed.
• uvx ruff format --diff and uvx ruff check are clean; bash -n clean on both shell scripts.
• Opened as a draft because CI cannot demonstrate the feature working until the TE pin moves and the two upstream fix PRs merge.

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

[shanmugamr1992]

TransformerEngine 2.19 run on GB200 — GTP now activates, but a pre-existing upstream bug blocks the test

Follow-up to the "TE too old" note in the description. I upgraded TransformerEngine on the
OCI HSG GB200 cluster (aarch64, GB200 NVL72, 2 GPUs used) and re-ran. HAVE_GTP is now
True and GTP genuinely engages, but the functional test still fails — on a different,
pre-existing megatron-core bug that is not caused by NVIDIA/Megatron-LM#6133.

Result

megatron-core	a0fa9c98913c92435b75471539eee7d4b8775784 (PR 6133 head)
TransformerEngine	2.19.0+0bf88ec4 (built from release_v2.19)
HAVE_GTP	True in all 18 megatron worker venvs
GTP activation	✅ both workers log gtp_remat_size=2, each rank holds 247,038,336 params (exactly half of Qwen2.5-0.5B)
grpo_megatron_generation_gtp.sh	❌ fails at checkpoint load, before generation/refit

So the TE floor is cleared and the NeMo-RL shim in this PR does its job — the parallel state
comes up as (tensor, gtp_remat, pipeline) and the weights are sharded. The run dies later.

Two fixes were needed to get TE 2.19 built at all

Both are worth knowing for anyone who bumps TE in this repo; neither is obvious from the failure.

1. uv pip install inside this workspace silently ignores the ref you ask for.
pyproject.toml [tool.uv] override-dependencies pins
transformer-engine[pytorch,core_cu13] @ git+…@release_v2.15 "across all extras". That
override also applies to ad-hoc uv pip install, so asking for release_v2.19 gets you
2.15 — with rc=0 and no warning. One run spent 21 minutes faithfully recompiling the
version that was already installed. The tell was in the log:

Built `transformer-engine @ git+…@42b840051647eef89761a16dfdff87e82bb253ab`

42b84005 is release_v2.15. The fix is uv pip install --no-config, run from a directory
outside the workspace, with a bare spec — the [pytorch]/[core_cu13] extras have to be
dropped too, since they exist only on the PyPI package and not on the git source.

2. TE 2.19 needs nvidia-cudnn-frontend>=1.25.0, and --no-build-isolation won't install it.
--no-build-isolation is required here so TE links the venv's existing torch instead of
downloading its own, but it also means none of TE's declared build requirements get installed.
TE 2.19 added nvidia-cudnn-frontend>=1.25.0 to build_tools/pytorch.py setup_requires and
resolves the header directory via importlib.metadata.distribution("nvidia-cudnn-frontend").
The venvs carry 1.23.0 from the TE 2.15 era, so the build finds a header, proceeds, and
dies ~25 minutes in:

fused_attn_f16_arbitrary_seqlen.cu(262): error: class "cudnn_frontend::graph::Tensor_attributes"
    has no member "set_ragged_offset_multiplier"
fused_attn_fp8.cu(355): error: class "cudnn_frontend::graph::SDPA_attributes"
    has no member "set_cu_seq_len_q"

Installing nvidia-cudnn-frontend>=1.25.0 (resolved to 1.27.0) first fixes it. Build then takes
~30 min for sm_100a; the wheel caches, so the other 17 venvs installed in 1–25 s each.

The remaining blocker: GTP-sharded models can't load a torch_dist checkpoint

Correction (this comment was edited). The first version blamed
make_tp_sharded_tensor_for_checkpoint for never folding the GTP axis into the checkpoint
global shape. That was wrong: megatron/core/utils.py is GTP-aware and composes
tp_rank * gtp_remat_size + gtp_rank over tp_size * gtp_remat_size. The two real causes are
below; both are about tensors GTP replicates rather than shards.

CheckpointingException: Global shape mismatch for loaded (torch.Size([24, 1152]))
  and expected ((24, 576)) tensor for key decoder.layers.self_attention.linear_qkv.bias

CheckpointingException: Invalid sharding pattern validation. Errors: Invalid access pattern
  for ShardedTensor(key='decoder.final_layernorm.weight', local_shape=(896,),
  global_shape=(896,), global_offset=(0,), ...)

1. The bias is physically sized as a GTP shard, but everything treats it as replicated.
_gtp_pre_init in megatron/core/extensions/transformer_engine.py deliberately pre-shards
out_features so the plain TE constructor allocates only this rank's weight shard. TE sizes the
bias off that same out_features, so the bias also comes out 1/gtp_remat too small - 576 instead
of 1152 for Qwen2.5-0.5B's linear_qkv. Nothing else agrees with that:
attach_gtp_to_presharded_module wraps only weight_names, so is_gtp_param(bias) is False;
mcore's own ColumnParallelLinear allocates its bias at the full per-TP size; and
make_sharded_tensors_for_checkpoint_with_gtp_remat has an explicit branch commented
"Non-GTPShardedParam under a GTP-active module (e.g. bias): GTP-replicated". The checkpoint
therefore declares a 576-wide global shape for a tensor that is 1152 wide on disk. This is also a
latent runtime bug, not only a checkpoint one: the all-gathered weight produces a full-width GEMM
output that a half-width bias cannot be added to. It goes unnoticed with Llama-style models because
they have no linear biases at all.

2. GTP-replicated tensors all claim replica_id 0. replica_id for a replicated tensor comes
from the dp_cp group, which deliberately excludes the gtp_remat axis (it is the replicate
group - get_data_parallel_group(..., with_gtp_remat=False) in process_groups_config.py). Every
gtp_remat peer thus reports the same DP rank, so both ranks claim to be the main replica of
final_layernorm.weight and _validate_sharding_for_key rejects it.
make_sharded_tensors_for_checkpoint_with_gtp_remat already folds gtp_rank into replica_id for
the replicated entries of a GTP-active module - but a module holding no GTP param at all
(final_layernorm, attention _extra_state, and that function's own non-GTP fast path) never
reaches those branches.

Neither is PR 6133's bug. git log -L 2700,2810 on
generalized_tensor_parallelism.py attributes the code to c5ff22b7f
([feat] Generalized Tensor Parallelism (GTP), #4967), last modified by 9cce740fa (#6775) and
2dc1b2bb7 (#6781) - all on main. PR 6133's commits (8d82590f8 "Reconcile GTP refit with
current main", a33186b8f, 4d7b88732, ...) add GTP refit/resharding and don't touch this path.

Fix in progress — and the two fixes belong in different repos.

Bug 1 (bias sizing) is a Megatron-LM bug. Fixed on
shanmugamr1992/Megatron-LM@mcore-6133-fix-gtp-bias,
branched off PR 6133's head a0fa9c989: _gtp_attach_post_init now re-allocates the GEMM biases
at their logical length after TE construction, right where it already restores
module.out_features. Only names matching bias\d* are touched, so a layer_norm_bias (sized
off in_features) is left alone.

Bug 2 (replica_id) is a Megatron-Bridge bug, not a Megatron-LM one. My first attempt folded
gtp_rank into replica_id inside make_sharded_tensor_for_checkpoint, which was wrong: mcore has
a deliberate regression test,
tests/unit_tests/generalized_tensor_parallel/test_gtp_dcp.py::_worker_replicated_param_needs_gtp_inclusive_dp_cp,
that asserts passing the gtp-excluded pg.dp_cp must produce more than one main replica. The
group choice is the caller's responsibility by design, and Megatron-LM's own training loop already
makes it (megatron/training/training.py: dp_cp_group = getattr(ckpt_pgc, "dp_cp_gtp_remat", None)).
Megatron-Bridge has its own checkpoint entrypoints — the ones NeMo-RL imports — and they pass
pg_collection.dp_cp. Fixed on
shanmugamr1992/Megatron-Bridge@mcore-6133-fix-gtp-dcp:
save_checkpoint, _load_checkpoint_from_path, and the three FullyParallel{Save,Load}StrategyWrapper
sites now use dp_cp_gtp_remat when present. The wrappers matter independently — they re-elect one
writer per shard within the group they are given, so a gtp-excluded group would reintroduce the
duplicate even after the metadata is corrected.

I will report the functional-test result against both branches here, and only propose them upstream
once the run is green.

What this means for reviewing this PR

• The NeMo-RL side of the change is exercised as far as it can be: GTP config plumbing,
  the initialize_model_parallel shim, and the layout/refit decision all work, and the
  HAVE_GTP preflight in setup.py correctly stopped the earlier runs with an actionable
  message instead of a confusing crash.
• tests/unit/models/generation/test_megatron_generation_layout.py (CPU-only) still covers
  the silent refit-skip bug and needs no GPU.
• grpo_megatron_generation_gtp.sh cannot pass until the two upstream checkpoint bugs are
  fixed; candidate fixes are linked above and are being re-run now. The test stays gated behind
  the megatron_gtp_supported() capability check in L1_Functional_Tests_Megatron_4.sh, so it
  skips rather than fails on images without GTP.
• Landing GTP for real needs a TE bump in pyproject.toml plus an image rebuild.
  Everything above was done by rebuilding TE inside a running container, which is fine for
  investigation but is not a shippable configuration.

No assertion, threshold, step count, or batch size was loosened at any point, and the two
anti-vacuity grep guards in the test script are still in place — I'd rather report a red
test than a green one that proves nothing.

[shanmugamr1992]

/ok to test 95fb64d