[cuichenx]

Summary

Fold high-rank LinearWithFrozenWeight dgrad inputs into a 2D GEMM before multiplying by the frozen weight, then reshape the result back to the original leading dimensions.

This preserves the existing math while avoiding a slow PyTorch batched-matmul dispatch for frozen linear backward when grad_output has more than two dimensions and its leading dimensions are not foldable by PyTorch's high-rank matmul path.

Why this is needed

Megatron Bridge PEFT/LoRA freezes the base model weights and relies on Megatron-Core's frozen-weight linear path to propagate activation gradients through the frozen base model. With packed sequence training, the frozen linear backward can receive a 3D grad_output shaped like sequence/microbatch/hidden. The previous implementation called:

grad_output.matmul(weight)

directly on that high-rank tensor. For some high-rank, non-contiguous activation-gradient layouts, PyTorch dispatches this as aten::matmul -> aten::bmm, which is much slower than the equivalent flattened 2D GEMM. Full-parameter SFT does not use this frozen-weight path in the same way because base weights require gradients, so the issue shows up specifically in PEFT/frozen-weight training.

Foldable vs non-foldable layouts

The issue is layout-sensitive. PyTorch can already fold a high-rank matmul into aten::mm when the first N-1 dimensions are view-flattenable, using a stride predicate equivalent to:

stride[i] == stride[i + 1] * shape[i + 1]

For a 3D [S, B, H] @ [H, K] dgrad, the important leading-dimension check is:

stride[0] == stride[1] * B

When this holds, the old path already routes to aten::mm, so this patch is not expected to improve performance. When this does not hold, the old path can fall through to aten::bmm; this patch forces the intended 2D GEMM.

Benchmark evidence

Benchmarks were run with one-off external repro/profiling scripts; no standalone benchmark script is included in this PR.

Environment:

• Main container: nvcr.io/nvidian/nemo:26.06.rc1
• Additional container checks: nvcr.io/nvidia/nemo:26.02, nvcr.io/nvidia/nemo:25.11
• Hardware: H100 and GB200
• Dtype: BF16

Minimal frozen-linear forward/backward repro on H100:

• Shape: grad_output=[384, 8, 4096], weight=[4096, 4096]
• grad_output stride: [4096, 1572864, 1]
• Old path: 1.756 ms median, includes aten::bmm
• Flattened path: 0.323 ms median, uses aten::mm
• Speedup: 5.43x

Isolated dgrad dispatch sweep:

grad_output problem / layout	Representative stride	H100 speedup	GB200 speedup	Old-path dispatch
[384, 8, 4096], contiguous	[32768, 4096, 1]	1.0x no speedup	1.0x no speedup	aten::mm
[384, 8, 4096], non-contiguous but foldable leading dims	[8, 1, 3072]	1.0x no speedup	1.0x no speedup	aten::mm
[384, 8, 4096], strided hidden dim but foldable leading dims	[65536, 8192, 2]	1.0x no speedup	1.0x no speedup	aten::mm
[384, 8, 4096], non-contiguous leading dims	[4096, 1572864, 1]	9.1x	5.1x	aten::bmm
[384, 8, 4096], expanded middle dim	[4096, 0, 1]	9.9x	5.8x	aten::bmm
[384, 16, 4096], contiguous	[65536, 4096, 1]	1.0x no speedup	1.0x no speedup	aten::mm
[384, 16, 4096], non-contiguous leading dims	[4096, 1572864, 1]	4.7x	3.0x	aten::bmm
[768, 8, 4096], contiguous	[32768, 4096, 1]	1.0x no speedup	1.0x no speedup	aten::mm
[768, 8, 4096], non-contiguous leading dims	[4096, 3145728, 1]	9.4x	6.1x	aten::bmm
[384, 8, 8192], contiguous	[65536, 8192, 1]	1.0x no speedup	1.0x no speedup	aten::mm
[384, 8, 8192], non-contiguous leading dims	[8192, 3145728, 1]	19.4x	7.5x	aten::bmm

Older-container checks on H100 showed the same pattern:

Container	Reported PyTorch / CUDA	[384,8,4096] contiguous	[384,8,4096] transpose-leading	[384,8,4096] expanded-middle	[384,8,8192] contiguous	[384,8,8192] transpose-leading	[384,8,8192] expanded-middle
nvcr.io/nvidia/nemo:25.11	2.9.0a0+50eac811a6.nv25.09 / CUDA 13.0	1.0x no speedup	5.3x	5.6x	0.9x no speedup	10.5x	10.8x
nvcr.io/nvidia/nemo:26.02	2.10.0a0+b558c986e8.nv25.11 / CUDA 13.0	1.0x no speedup	5.2x	5.4x	1.0x no speedup	10.3x	10.8x

Llama 70B LoRA recipe check

I also ran the H100 BF16 Llama 70B LoRA performance recipe before and after this patch:

• Recipe: llama3_70b_lora_config_h100
• Container: nvcr.io/nvidian/nemo:26.06.rc1
• Hardware: H100
• Precision: BF16
• Parallelism/batch: 8 GPUs, TP=1, PP=4, VP=20, CP=1, MBS=1, GBS=32
• Steps: 50
• LoRA target modules: ["linear_qkv"]
• Checkpoint: existing torch-dist Llama 70B-family checkpoint with matching architecture dimensions, used only to satisfy PEFT checkpoint loading for the perf recipe

Results:

Run	Median iter 2-50	Mean iter 2-50	Last 15 median	Speedup vs baseline
Baseline	12,746.0 ms	12,745.7 ms	12,743.1 ms	-
Patched	12,745.1 ms	12,744.4 ms	12,743.3 ms	1.00x

This recipe does not show a meaningful end-to-end speedup. That is expected: with MBS=1, the main LoRA-targeted frozen QKV dgrad shape is effectively:

grad_output [4096, 1, 10240]
weight      [10240, 8192]

A normal contiguous [S, 1, H] tensor satisfies the foldability predicate, so the old PyTorch path already uses aten::mm. This negative result is important because it confirms the patch is not a blanket end-to-end speedup for every LoRA recipe; it targets the non-foldable high-rank layouts that currently route through aten::bmm.

Alternative considered

An alternative fix would be to call grad_output.contiguous() before the high-rank matmul. That can avoid the bad bmm route for reproduced non-foldable layouts because the copied tensor becomes foldable.

The explicit reshape(-1, hidden).matmul(weight).reshape(...) form is preferred because it directly expresses the intended 2D GEMM, avoids relying on high-rank matmul dispatch heuristics, and lets reshape stay as a view for already-foldable layouts. A blanket contiguous() would force a copy for every non-contiguous grad_output, including non-contiguous-but-foldable cases that already route to aten::mm.

Model and shape impact

The code change is model-family agnostic, but the size of the performance win is shape and layout dependent.

Any model that uses LinearWithFrozenWeight with rank > 2 activation gradients can take this path. That includes PEFT/frozen-weight training for dense or MoE model families used through Megatron Bridge. However, whether the old high-rank matmul becomes pathological depends on the GEMM shape and runtime stack: hidden size, output size, packed sequence/token shape, dtype, tensor strides, GPU architecture, PyTorch version, and cuBLASLt kernel selection all matter.

The patch is still safe for all such models because it only folds leading dimensions for an equivalent matrix multiply and then restores the original leading shape.

Normal trainable linear layers and full-parameter SFT should not be affected because they do not use this frozen-weight autograd function for the base model dgrad.

Testing

• Added a unit test comparing 3D LinearWithFrozenWeight forward and input gradient against torch.nn.functional.linear.
• Ran a non-install local syntax check:
  • python -m py_compile megatron/core/tensor_parallel/layers.py tests/unit_tests/tensor_parallel/test_layers.py
• Ran external H100/GB200 microbenchmarks for foldable and non-foldable high-rank dgrad layouts.
• Ran the H100 BF16 Llama 70B LoRA performance recipe before/after; result was ~1.00x because that recipe's MBS=1 QKV dgrad layout is foldable.

[copy-pr-bot]

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.

[greptile-apps]

_{Reviews (1): Last reviewed commit: "test: add frozen linear dgrad repro scri..." | Re-trigger Greptile}

[cuichenx]

/claude strict-review

[greptile-apps]

_{Reviews (2): Last reviewed commit: "test: remove frozen linear dgrad repro s..." | Re-trigger Greptile}

[claude]

Strict review passed — no significant issues found. LGTM.

CRITICAL: 0 | IMPORTANT: 0 | SUGGESTION: 0

The reshape-matmul-reshape pattern in the frozen-weight backward is correct:

• grad_output.reshape(-1, out_features) @ weight[out_features, in_features] → [N, in_features], reshaped back to [*leading_dims, in_features] via weight.size(1).
• Mathematically identical to the original grad_output.matmul(weight) — the matmul is row-independent so flattening leading dims is safe.
• reshape(-1, ...) stays as a zero-copy view for already-foldable layouts, only materializing a copy for the non-contiguous cases that would otherwise fall through to the slow aten::bmm dispatch.
• Consistent with the existing prepare_input_tensors_for_wgrad_compute pattern in megatron/core/utils.py which does the same 3D→2D flatten for the weight gradient.
• The subsequent all_reduce at L385 is element-wise and unaffected by the output tensor's contiguity.
• Test validates both forward and input-gradient correctness for 3D inputs against torch.nn.functional.linear.

Clean, well-scoped change with thorough profiling evidence.

[greptile-apps]

_{Reviews (3): Last reviewed commit: "chore: fix lint formatting and guidance" | Re-trigger Greptile}

[cuichenx]

Hi @NVIDIA/core-adlr @NVIDIA/core-nemo Could you please review this PR?

The benchmarking results can be reproduced with this script

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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