fix: make GTP weight rematerialization work end-to-end with tied embeddings and inference resharding by shanmugamr1992 · Pull Request #6940 · NVIDIA/Megatron-LM

shanmugamr1992 · 2026-08-28T03:22:29Z

Found while validating GTP weight rematerialization (#6133, "Gtp refit support")
against the NeMo-RL Megatron functional suites on GB200 (4 GPUs), running GRPO
with a colocated Megatron generation backend at TP=1, gtp_remat_size=2, on
Qwen2.5-0.5B (which has share_embeddings_and_output_weights=True).

Before these four commits the configuration could not start. After them it runs
two full GRPO steps end-to-end: build → generate → reshard/refit → train.

Each commit is independent and stands on its own; three of them are one-liners.

1. `fix: restore GTP-replicated bias to its logical size in TE layers`

_gtp_pre_init passes a pre-sharded out_features to TE's constructor so the
weight comes out GTP-sharded. TE also sizes its bias off that same
out_features, so the bias silently comes out 1/gtp_remat too small — but a
bias is GTP-replicated, not GTP-sharded (mcore's own ColumnParallelLinear
allocates it at the full per-TP size, attach_gtp_to_presharded_module never
wraps it, and make_sharded_tensors_for_checkpoint_with_gtp_remat treats a
non-GTP param under a GTP module as replicated).

Two consequences: the all-gathered weight produces a full-width GEMM output that
the shard-sized bias cannot be added to, and distributed checkpointing declares a
global shape 1/gtp_remat of the real one, so a plain checkpoint fails to load
with Global shape mismatch ... linear_qkv.bias.

_gtp_restore_presharded_bias re-allocates the affected biases at their logical
size after super().__init__. Biases are zero-initialized by both TE and mcore,
so nothing is lost. It matches only GEMM biases (bias, bias0..biasN) — never
layer_norm_bias, which is sized off in_features — and only when the observed
numel is exactly the shard size, so it is a no-op on every non-GTP path.

2. `fix: resolve the GTP axis from GPTModel's own process groups`

GPTModel passed only tp_group=self.pg_collection.tp to LanguageModelEmbedding
and to the output layer. resolve_gtp_remat_group falls back to the global MPU
groups when it is given no collection, so those two layers picked up whatever GTP
layout parallel_state happened to hold rather than the model's own. Passing
pg_collection through makes the model self-describing, per the repo guidance on
avoiding new implicit parallel_state reads in megatron/core.

TELMHeadColumnParallelLinear did not accept pg_collection at all, so the
output-layer half of this was unreachable for TE LM heads; it now forwards it to
TEColumnParallelLinear.

3. `fix: declare the GTP axes off on inference process-group collections`

build_inference_pg_collection left gtp_remat / expt_gtp_remat unset.
resolve_gtp_remat_group inspects vars(pg_collection) and falls back to the
MPU globals for a field that is absent, so an inference collection built beside
a GTP-enabled training model inherited the training GTP layout and produced a
GTP-sharded inference model. An inference layout never rematerializes weights, so
the axes are now declared explicitly None — present in vars(), fallback not
taken.

This is what made a colocated reshard fail with
supplied weight's shape is (75968, 896), not (151936, 896) as expected while
building the dedicated inference model.

4. `fix: accept a GTP-sharded weight in ColumnParallelLinear's shape guard`

With tied embeddings, shared_embedding_or_output_weight() hands the output
layer the embedding's GTP-sharded parameter, and _forward_impl all-gathers it to
the logical shape before the matmul (via gtp_remat_size). The
skip_weight_param_allocation guard compared the local shape against the
logical one and rejected it, so the same
supplied weight's shape is (75968, 896), not (151936, 896) aborted the training
forward — this one independently of NeMo-RL or resharding, in any GTP run with
share_embeddings_and_output_weights=True.

The guard now compares _unsharded_shape for a GTP param, so it still rejects a
genuinely wrong weight.

Validation

Rebased onto main at 5e2b1da45 and re-verified (each commit cherry-picked
clean; no conflicts with the merged #6133).

Validated as a set on GB200 (aarch64, 4 GPUs) with a GRPO functional test
configured TP=1, gtp_remat_size=2, NCCL refit, colocated Megatron generation:
without them the run aborts during model or inference-model construction; with
them it completes two full training steps.

One thing this does not fix

That run's train/token_mult_prob_error measures 1.1011, against a < 1.05
threshold the test inherited 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 (it
takes a different colocated-generation branch and hits an unrelated CUDA-graph
capture error), so there is no measured baseline for that threshold on this
configuration — it was never validated at 1.05 here. Several other non-vanilla
Megatron generation variants sit similarly above their own 1.05-class gates.

Flagging it rather than moving the threshold: whether the GTP path adds real
numerical error or the inherited threshold is simply wrong for it is a call for
whoever owns the GTP refit path.

Raised by an automated agent while testing #6133 against NeMo-RL. Draft — needs a
human review before merge.

copy-pr-bot · 2026-08-28T03:22:33Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

…e layout Forcing gtp_weight_remat_size=1 on the inference provider does not by itself keep the inference model off the training GTP layout: megatron's layers resolve the axis from the process-group collection, which falls back to the MPU globals unless build_inference_pg_collection declares it off (NVIDIA/Megatron-LM#6940). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> Signed-off-by: shanmugamr1992 <shanmugamr1992@gmail.com>

wdykas · 2026-08-28T18:28:11Z

I think this make sense to me. @fanshiqing can you take a look?

wdykas · 2026-09-01T02:05:06Z

/ok to test 95fb64d

fanshiqing · 2026-09-01T02:12:15Z

    return shard_out, gtp_ctx


+def _gtp_restore_presharded_bias(module, gtp_ctx, is_grouped=False):


This is a real bug and has also been fixed similarly in https://github.com/NVIDIA/Megatron-LM/pull/6986/changes#diff-781136995db953ab5d54d564383668584688622746d2e5576dd4e389ecd034a8R737

I prefer the fix in above MR which is simpler :P

Agreed — dropped this in favour of the #6986 fix, which is now on main. This hunk is no longer in the diff.

copy-pr-bot · 2026-09-01T15:36:07Z

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.

wdykas · 2026-09-01T16:16:12Z

/ok to test 5256309

_gtp_pre_init shrinks ``out_features`` to the local GTP shard before TE builds the module, so TE allocates every GEMM bias at 1/gtp_remat of its real length. _gtp_attach_post_init then restores ``module.out_features`` and hands the weight to attach_gtp_to_presharded_module, which only walks ``weight_names`` -- the bias is left at the shard size. A bias is GTP-*replicated*, not GTP-sharded: after the weight all-gather the GEMM produces a full-width output that a 1/gtp_remat-length bias cannot be added to, and dist-checkpointing sees a tensor whose local shape does not match its declared global shape. It goes unnoticed on bias-free models (Llama) and shows up immediately on Qwen2.5, whose qkv projection has a bias. Re-allocate those biases at their logical length after TE construction, copying the tensor-parallel attributes across. Only names matching ``bias\d*`` (the GEMM biases, including the grouped-linear ``bias0..biasN``) are considered, so a ``layer_norm_bias`` -- sized off in_features and unaffected by the out_features pre-shard -- is never touched, and a bias whose length is not the shard size is left alone. Signed-off-by: shanmugamr1992 <shanmugamr1992@gmail.com> (cherry picked from commit 490fbae)

With tied embeddings, GPTModel._postprocess passes shared_embedding_or_output_weight() -- the embedding's GTP weight-remat shard -- to the output layer. ColumnParallelLinear.forward then compared that shard against the unsharded expected shape and raised RuntimeError: supplied weight's shape is (75968, 896), not (151936, 896) as expected even though the shard is exactly what the layer wants: when gtp_remat_size > 1 both _forward_impl paths all-gather the supplied weight to its logical shape before the matmul, and save_for_backward keeps the sharded weight so wgrad reduces back into the shard. Compare the logical (gathered, pad-stripped) shape for a GTP param so the guard still rejects a genuinely wrong weight. Signed-off-by: shanmugamr1992 <shanmugamr1992@gmail.com> (cherry picked from commit 2e7d373)

Move replicated-bias restoration into the shared pre-sharded GTP attachment helper and resize bias Parameters in place, matching the simpler implementation from PR NVIDIA#6986. Skip the external-weight shape guard for GTP parameters while preserving it for ordinary weights, and add focused regression coverage for both behaviors. Signed-off-by: William Dykas <wdykas@nvidia.com>

Address review feedback on PR NVIDIA#6940: - Remove the two explanatory comments above the `pg_collection` arguments in GPTModel; the argument is self-explanatory alongside `tp_group`. - Remove the duplicate `TestGTPReplicatedBias`. `_gtp_restore_replicated_bias` and an equivalent, stronger multi-GPU test both landed on main in NVIDIA#6986, and this file already defines a class of the same name later on, which would have silently shadowed the copy added here. Signed-off-by: Shanmugam Ramasamy <shanmugamr@nvidia.com>

fanshiqing · 2026-09-03T00:02:36Z

LGTM, thanks for these fixes!

fanshiqing · 2026-09-03T00:03:01Z

/ok to test 636bfd6

nemo-automation-bot · 2026-09-14T17:29:01Z

🔄 Merge queue validation started!

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

nemo-automation-bot · 2026-09-15T00:14:25Z

🔄 Merge queue validation started!

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

…ddings and inference resharding (NVIDIA#6940) Signed-off-by: shanmugamr1992 <shanmugamr1992@gmail.com> Signed-off-by: William Dykas <wdykas@nvidia.com> Signed-off-by: Shanmugam Ramasamy <shanmugamr@nvidia.com> Co-authored-by: William Dykas <wdykas@nvidia.com> (cherry picked from commit 78d917e)

…ddings and inference resharding (NVIDIA#6940) Signed-off-by: shanmugamr1992 <shanmugamr1992@gmail.com> Signed-off-by: William Dykas <wdykas@nvidia.com> Signed-off-by: Shanmugam Ramasamy <shanmugamr@nvidia.com> Co-authored-by: William Dykas <wdykas@nvidia.com> (cherry picked from commit 78d917e) Signed-off-by: xuwenc <xuwenc@nvidia.com>

This was referenced Aug 28, 2026

fix(ckpt): derive checkpoint replica_id from the gtp_remat-inclusive DP group NVIDIA-NeMo/Megatron-Bridge#5862

Draft

feat(megatron): support GTP weight rematerialization for training + refit NVIDIA-NeMo/RL#3857

Draft

copy-pr-bot Bot temporarily deployed to test September 1, 2026 02:06 Inactive

copy-pr-bot Bot temporarily deployed to public September 1, 2026 02:08 Inactive

wdykas marked this pull request as ready for review September 1, 2026 02:10

wdykas requested review from a team as code owners September 1, 2026 02:10

github-actions Bot added the complexity: low label Sep 1, 2026

fanshiqing reviewed Sep 1, 2026

View reviewed changes

lhb8125 mentioned this pull request Sep 1, 2026

Megatron-LM Monthly Activity Report — August 2026 #7006

Open

fanshiqing reviewed Sep 1, 2026

View reviewed changes

Comment thread megatron/core/tensor_parallel/layers.py Resolved

fanshiqing requested changes Sep 1, 2026

View reviewed changes

shanmugamr1992 requested a review from a team as a code owner September 1, 2026 15:36

copy-pr-bot Bot temporarily deployed to test September 1, 2026 16:17 Inactive

copy-pr-bot Bot temporarily deployed to public September 1, 2026 16:22 Inactive

fanshiqing approved these changes Sep 2, 2026

View reviewed changes

fanshiqing reviewed Sep 2, 2026

View reviewed changes

Comment thread megatron/core/models/gpt/gpt_model.py Resolved

Comment thread megatron/core/models/gpt/gpt_model.py Resolved

Comment thread megatron/core/tensor_parallel/generalized_tensor_parallelism.py Resolved

shanmugamr1992 and others added 3 commits September 2, 2026 16:39

shanmugamr1992 force-pushed the mcore-6133-fix branch from 5256309 to 636bfd6 Compare September 2, 2026 23:48

fanshiqing enabled auto-merge September 3, 2026 00:03

copy-pr-bot Bot temporarily deployed to test September 3, 2026 00:04 Inactive

copy-pr-bot Bot temporarily deployed to public September 3, 2026 00:07 Inactive

kvareddy approved these changes Sep 3, 2026

View reviewed changes

lhb8125 mentioned this pull request Sep 4, 2026

Megatron-LM Monthly Activity Report — August 2026 (Final, through August 31, 2026) #6921

Open

yaoyu-33 approved these changes Sep 4, 2026

View reviewed changes

fanshiqing added this pull request to the merge queue Sep 14, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Sep 14, 2026

jaredcasper approved these changes Sep 14, 2026

View reviewed changes

svcnvidia-nemo-ci added the Approved All necessary approvals have been made label Sep 14, 2026

wdykas added this pull request to the merge queue Sep 14, 2026

Merged via the queue into NVIDIA:main with commit 78d917e Sep 15, 2026
108 of 110 checks passed

wdykas deleted the mcore-6133-fix branch September 15, 2026 00:14

wdykas mentioned this pull request Sep 15, 2026

fix(ckpt): derive checkpoint replica_id from the gtp_remat-inclusive DP group NVIDIA-NeMo/Megatron-Bridge#6102

Merged

lhb8125 mentioned this pull request Sep 21, 2026

Megatron-LM Monthly Activity Report — September 2026 (Final, through September 30, 2026) #7100

Open

lhb8125 mentioned this pull request Oct 1, 2026

Megatron-LM Monthly Activity Report — September 2026 #7773

Open

		return shard_out, gtp_ctx


		def _gtp_restore_presharded_bias(module, gtp_ctx, is_grouped=False):

Conversation

shanmugamr1992 commented Aug 28, 2026

1. fix: restore GTP-replicated bias to its logical size in TE layers

2. fix: resolve the GTP axis from GPTModel's own process groups

3. fix: declare the GTP axes off on inference process-group collections

4. fix: accept a GTP-sharded weight in ColumnParallelLinear's shape guard

Validation

One thing this does not fix

Uh oh!

copy-pr-bot Bot commented Aug 28, 2026

Uh oh!

wdykas commented Aug 28, 2026

Uh oh!

wdykas commented Sep 1, 2026

Uh oh!

fanshiqing Sep 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

fanshiqing Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

shanmugamr1992 Sep 2, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

copy-pr-bot Bot commented Sep 1, 2026

Uh oh!

wdykas commented Sep 1, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

fanshiqing commented Sep 3, 2026

Uh oh!

fanshiqing commented Sep 3, 2026

Uh oh!

nemo-automation-bot Bot commented Sep 14, 2026

Uh oh!

Uh oh!

nemo-automation-bot Bot commented Sep 15, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

1. `fix: restore GTP-replicated bias to its logical size in TE layers`

2. `fix: resolve the GTP axis from GPTModel's own process groups`

3. `fix: declare the GTP axes off on inference process-group collections`

4. `fix: accept a GTP-sharded weight in ColumnParallelLinear's shape guard`

fanshiqing Sep 1, 2026 •

edited

Loading