Accumulate delayed wgrads into shared (tied) parameters by 0z5a · Pull Request #3552 · NVIDIA/TransformerEngine

0z5a · 2026-09-21T04:51:09Z

Delayed backward_dw() calls from modules sharing a parameter must add their contributions to the existing gradient. Assignment drops earlier tied-weight and microbatch contributions. Apply the same None-versus-add contract to all four implementations: the base path used by te.Linear, te.LayerNormMLP, te.GroupedLinear, and te.ops.GroupedLinear.

LayerNormMLP accumulates both FC weight gradients and nonempty deferred bias gradients. Biases computed in the main backward, including FP8 block-scaling biases, are preserved without recomputing them from saved activations. Grouped paths cover discrete and packed weights, with deferred discrete bias accumulation in the module path. Fused main_grad accumulation and reduction hooks retain their existing behavior.

Merge current main (5e994640) and resolve the LayerNormMLP bias conflict while preserving the accumulation fix and upstream changes.

The public regression file is now launched by qa/L0_pytorch_unittest/test.sh. Its 50 parametrized GPU cases cover tied and non-tied Linear, an independent nn.Linear reference, both drain orders, multiple microbatches, two reset/step iterations with both set_to_none settings, both LayerNormMLP FCs and biases, FP8 block-scaling bias preservation, and discrete/packed grouped weights. Tests keep their sizes local and spell out execution and drain steps; only the two-Linear factory is shared. Packed and FP8 cases check their existing runtime support predicates.

CUDA Graph follow-up at dbe1326e7a2ddfca0ee40b19142f84404eaa9399:

The two deferred-weight-gradient capture paths now clear ordinary-allocator .grad targets while capturing, retain the fresh graph-pool contributions, and clone or add those contributions into caller-owned gradients before delayed hooks during replay. zero_grad(set_to_none=True) can release the caller's gradients without invalidating a captured accumulation address. Reset releases the retained graph gradients.

tests/pytorch/test_cuda_graphs.py adds 26 parametrized cases covering both capture schedules, both gradient reset modes, FP32/BF16, bias, retained gradients, tied weights, multiple microbatches and fused main_grad, compared with independent PyTorch references.

Validation for this follow-up:

Private CPU controls execute the actual four method bodies and production WeightGradStore using real CPU Parameters and independent PyTorch autograd results. All 36 accumulation controls passed; the previous review head fails 32 of them in the three remaining implementations. Another nine controls passed for fused-main-grad preservation and hook invocation.
Private import-free definition collection confirms all 50 parametrizations. This is a test-definition check, not native TE execution.
The 36 accumulation and nine preservation controls passed again for this follow-up. Eight additional CPU source controls passed for the actual capture/replay/reset helpers with real CPU Parameters; graph capture and replay are modeled, so these do not validate CUDA allocation or native TE kernels.
Scoped pre-commit checks, Python syntax and diff checks passed. Targeted pylint for transformer_engine/pytorch/graph.py passed (10.00/10).
The reviewer reported 24 failures among the 50 eager cases on SM100. Those failures have not been rerun or diagnosed from per-case logs. Native SM100, CUDA Graph and FP8 validation is deferred until SM100 access is available; no GPU pass is claimed for this follow-up. The earlier eight-case Linear GPU results apply only to the prior revision.

Microbatch accumulation can differ in BF16 rounding order; those cases use a 1e-2 tolerance. Single-microbatch Linear comparisons remain exact.

greptile-apps · 2026-09-21T04:53:45Z

Confidence Score: 5/5

[Medium impact] The change since the previous review appears safe to merge.

Summary

The PR makes delayed gradients add to existing gradients for shared parameters and microbatches.

Delayed gradients from tied layers and microbatches add together.
Grouped layers add delayed weight results across uses.
CUDA Graph replay adds delayed gradients to caller-owned buffers.

Since the previous review, only test_non_tied_linear_gradients changed. Its new zip loop preserves both gradient comparisons. No new issues were found.

_{Reviews (11) · Last reviewed commit: "Iterate the non-tied test's module pairs..." · Reviewed by Greptile}

wujingyue · 2026-09-21T06:31:19Z

+                # module overwrite the contributions of the earlier ones, so
+                # accumulate instead: initialise on the first contribution of the
+                # step and add on every later one.
+                if weight_tensor.grad is None:


Is it safe? There's also zero_grad(set_to_none=False)

Yes. set_to_none=False leaves a zeroed tensor where set_to_none=True leaves None, so the first contribution of the step takes the add path and lands on the same value; nothing in the branch needs grad to be None to start a step. Said so in the comment (7ed69a4).

It also turned out not to be pinned: the test ran both settings but only asserted second != first, which cannot show pollution, because a first-step value carried into the second still differs from the first. It now compares the second step's grad against a fresh model on the same weights and input, bit-exactly.

wujingyue

Core changes look good! Tests need a lot of cleanup in addition to my comments.

wujingyue · 2026-09-21T14:11:48Z

+SIZE = 16
+BATCH = 2
+DTYPE = torch.bfloat16
+DEVICE = "cuda"


Should go with the test rather than global.

Moved (0f419b2): the global is gone. _build and _x take device="cuda" as a default, the micro-batch test uses _x(4), and the oracle takes the device from model[0].weight.device.

wujingyue · 2026-09-21T14:11:53Z

+# differ by a BF16 ULP. That is not something the fix promises to remove, so the
+# tolerance is declared here and applies to the multi-microbatch case only;
+# every single-microbatch comparison below stays bit-exact.
+MULTI_MICROBATCH_RTOL = 1e-2


Moved into test_multiple_microbatches_accumulate, the only test that needs it.

wujingyue · 2026-09-21T14:12:32Z

+MULTI_MICROBATCH_ATOL = 1e-2
+
+
+def _tie(model, bias):


Suggested change

def _tie(model, bias):

def _tie(model, *, bias):

otherwise it's hard to understand what _tie(model, True) means

Done: _tie(model, *, bias).

wujingyue · 2026-09-21T14:12:38Z

+    return model
+
+
+def _build(delay, tied, bias=False):


Done: _build(delay, *, bias=False).

wujingyue · 2026-09-21T14:13:13Z

+            fuse_wgrad_accumulation=False,
+        ),
+    )
+    return _tie(model, bias) if tied else model


This should be separated out of _build so SRP

Done: _build now only builds. Tying happens at the call sites, where each test says whether it wants a tied pair or two independent ones.

wujingyue · 2026-09-21T14:13:30Z

+    return _tie(model, bias) if tied else model
+
+
+def _run(model, x, microbatches=1, delay=False, order=(0, 1)):


What's order?

Renamed to backward_dw_order and documented: it is the module order used to drain the delayed wgrad queues, which decides which tied module pops its delayed wgrad first.

greptile-apps · 2026-09-24T01:13:35Z

Want your agent to iterate on Greptile's feedback? Try greploops.

backward_dw() assigned the popped delayed wgrad to Parameter.grad. When two modules share one Parameter, each pops its own contribution and the later one silently replaced the earlier, so the shared gradient came out too small -- on the issue's own 16x16 / batch=2 BF16 repro it was 0.6015x the regular path, with 256/256 elements differing. The bias takes the same pop path and hits the same bug, with an extra twist: the second module's bgrad is numerically ~0, so it wiped out a contribution the first module had already written correctly (1.0984 delivered vs 33.0859 expected). Initialise on the first contribution of the step and add on every later one, for both the weight and the bias. dtype/shape/buffer-alias contracts are unchanged, and the fused wgrad-accumulation branch is untouched. tests/pytorch/test_tied_delayed_wgrad.py covers the matrix: tied and non-tied, both backward_dw() orders, multiple microbatches, two optimizer steps with set_to_none True/False, the bias case, and parity with a plain nn.Linear oracle. It uses only the public te.Linear API so it cannot be satisfied by a shared implementation detail of the fix. Without the fix it fails 5 and passes 3; with it, 8 pass. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

wujingyue · 2026-09-24T05:58:22Z

+    return model
+
+
+def _build(delay, *, bias=False, device="cuda"):


Suggested change

def _build(delay, *, bias=False, device="cuda"):

def _build(*, delay, bias, device="cuda"):

Went with your follow-up suggestion in the next thread instead; see the reply there.

wujingyue · 2026-09-24T06:00:36Z

+    return model
+
+
+def _build(delay, *, bias=False, device="cuda"):


Suggested change

def _build(delay, *, bias=False, device="cuda"):

def _build_two_linears(delay, *, bias=False, device="cuda"):

Sorry for the very late replies on this round. Applied in c5cf151: the helper is now _build_two_linears(delay, *, bias=False, device="cuda") and only builds the two layers.

wujingyue · 2026-09-24T06:01:00Z

+SIZE = 16
+BATCH = 2
+DTYPE = torch.bfloat16


https://testing.googleblog.com/2019/12/testing-on-toilet-tests-too-dry-make.html

Removed in c5cf151. The SIZE/BATCH/DTYPE constants are gone, and each test spells out its own sizes and torch.bfloat16.

wujingyue · 2026-09-24T06:03:24Z

+    ``backward_dw_order`` is the module order used for draining, which decides which tied
+    module pops its delayed wgrad first.
+    """
+    for xb in torch.chunk(x, microbatches, dim=0):


xb?

https://google.github.io/styleguide/pyguide.html#316-naming

Renamed to microbatch (c5cf151).

wujingyue · 2026-09-24T06:04:12Z

+                model[i].backward_dw()
+
+
+def _x(batch=BATCH, device="cuda"):


https://testing.googleblog.com/2019/12/testing-on-toilet-tests-too-dry-make.html

Removed in c5cf151. Each test builds its input with an inline torch.randn(...).

wujingyue · 2026-09-24T06:04:25Z

+
+
+@pytest.fixture(autouse=True)
+def _seed():


https://testing.googleblog.com/2019/12/testing-on-toilet-tests-too-dry-make.html

Removed in c5cf151. Each test calls torch.manual_seed(9) itself.

wujingyue · 2026-09-24T06:04:39Z

+    xd = _x()
+    _run(delayed, xd, delay=True)
+    _run(regular, xd.detach().clone().requires_grad_(True))
+    for i in (0, 1):


Done in d14c2d2: test_non_tied_linear_gradients now iterates zip(delayed, regular) instead of indexing both models.

Review question: is the `grad is None` first-contribution test safe when the optimizer was built with `zero_grad(set_to_none=False)`? It is: that setting leaves a zeroed tensor where `set_to_none=True` leaves `None`, so the first contribution takes the add path and lands on the same value, and nothing in the branch depends on `grad` being `None`. Say so where the branch is, and pin it in the test. `test_two_optimizer_steps_do_not_pollute_each_other` already ran both settings, but `second != first` cannot show pollution - a first-step value carried into the second still differs from the first. It now compares the second step's gradient against a fresh model on the same weights and input, bit-exactly. Verified with the new test over an out-of-tree copy of an installed TransformerEngine: 5 failed, 3 passed against the unfixed `backward_dw` (the same five this PR describes), and 8 passed with this PR's `backward_dw` in place. Signed-off-by: 0z5a <192209249+0z5a@users.noreply.github.com> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

- `_tie(model, *, bias)` and `_build(delay, *, bias=False, device="cuda")`: the boolean arguments are keyword-only, so a call site says what it means. - `_build` only builds. Tying moved out to the call sites, which is where a test decides whether it wants a tied pair or two independent ones. - `order` is now `backward_dw_order`, and the docstring says what it orders. - The module-level `DEVICE` and tolerance globals are gone: the device is a default on the two helpers that need it, and the multi-microbatch tolerance is local to the test that needs it. Signed-off-by: 0z5a <192209249+0z5a@users.noreply.github.com> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

denera

Thank you for the PR!

In addition to @wujingyue's review comments, I would request two additional changes, one trivial but the other a critical omission.

[TRIVIAL] The new test will not run in our CI. I would recommend moving the cases into test_numerics.py OR adding test_tied_delayed_wgrad.py into the qa/L0_pytorch_unittest/test.sh script (this is what our CI pipeline actually launches).
[CRITICAL] Your fix covers only one of the four delayed-wgrad implementations. te.LayerNormMLP.backward_dw, te.GroupedLinear.backward_dw and te.ops.GroupedLinear.backward_dw all remain unchanged and still drop all but the last wgrad contribution. These should be addressed in this PR as well.

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

denera

@0z5a I just manually tested this PR and I'm seeing 24 of the 50 test cases added in this PR failing on SM100. Could you please make sure that all the tests pass locally on your end before we launch the CI pipeline? Thanks!

Additionally, there's a CUDA-graph hazard here whenever optimizer.zero_grad() is invoked.

make_graphed_callables() does a few eager warmup passes before capture. After these warmups, param.grad holds a tensor G from the normal caching allocator, not the graph's private pool.
Graph capture on backward_dw() records kernels, not Python assignments. It captures the WGRAD GEMM writing into a fresh buffer S from the graph's pool, and the G.add_(S) for the accumulation with the raw address for tensor G baked in at capture time (NOTE: This only affects fuse_wgrad_accumulation=False).
optimizer.zero_grad() by default sets param.grad = None. This drops the last reference to G and the normal caching allocator frees the block.
Next time the graph runs, the captured raw address for tensor G is no longer valid. If the caching allocator has handed that address block to something else, the WGRAD accumulation corrupts it on graph replay.

Fixing this requires setting the param.grad = None for the parameters of modules that require backward_dw() right before the two bwd_dw_graph captures in lines 809 and 1008 for transformer_engine/pytorch/graph.py.

I would also recommend adding a test for backward_dw() capture under test_cuda_graphs.py. Something simple like te.Linear with delay_wgrad_compute=True graphed and compared against eager would reveal the above graph capture issue in the future.

Capture deferred contributions in the graph pool and accumulate them into caller-owned gradients before delayed hooks. Cover both capture orders, default and explicit gradient resets, tied weights, microbatches, bias, retained gradients, and fused main_grad with independent PyTorch references. Native GPU validation is pending on the coordinated test machine. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

With delay_wgrad_compute=True and fuse_wgrad_accumulation=False, GroupedLinear returned the still-empty wgrad buffer to autograd from the main backward pass. Autograd stored it as (or added it into) .grad before backward_dw() had written it, so accumulating the delayed result on top read uninitialized memory. Return None instead, as Linear does, and let backward_dw() own .grad. The tied-gradient tests compared multi-term BF16 sums from the delayed and regular paths with a fixed 1e-2 tolerance. The two paths add the same terms in a different order, so allow a few BF16 roundings at the gradient's magnitude. The graphed microbatch test compared against a PyTorch FP32 reference bitwise, but TE runs FP32 GEMMs in TF32; compare graph replay against the same TE modules run eagerly instead. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

0z5a · 2026-10-27T07:40:08Z

@denera Sorry for the slow follow-up, and thank you for running the tests. You were right: the failures are not specific to SM100. Applying this PR's head (dbe1326e) on current main reproduces the same 24/50 failures in test_tied_delayed_wgrad.py on an H100-class GPU (GH200, sm_90). It also fails the 8 test_graphed_delayed_wgrad_accumulates_microbatches cases. The rest of test_cuda_graphs.py is unchanged (438 passed, 423 skipped on both main and the PR).

There were three causes:

GroupedLinear handed autograd an unfilled buffer. With delay_wgrad_compute=True and fuse_wgrad_accumulation=False, the module backward returned the still-empty torch.empty wgrad buffer to autograd. .grad therefore held uninitialized memory, or the sum of two such buffers for tied weights, before backward_dw() ran. Accumulating into it produced values around 5e37. Linear already returns None when the wgrad is delayed. GroupedLinear now does the same, and backward_dw() owns .grad.
Tolerances for BF16 sums that add the same terms in a different order. With two microbatches the delayed path computes ((a1+b1)+a2)+b2, while autograd computes (a1+b1)+(a2+b2). The LayerNormMLP and te.ops.GroupedLinear mismatches were all a single BF16 ulp at the gradient's magnitude: 0.0625 and 0.125 for values of 8–16, and 0.031 for values of 4–8. The comparison now allows a few BF16 roundings at that magnitude. As a negative control, with the new tests, main still fails the same 49 of the 50 cases it failed before on sm_90 (and all 33 non-skipped tied cases on sm_89).
The graphed microbatch test compared TE FP32 against PyTorch FP32 bitwise. TE runs FP32 GEMMs in TF32 (cublaslt_gemm.cu). Once SGD moves the weights off TF32-representable values, every element differs by about 1.8e-3 relative. With NVIDIA_TF32_OVERRIDE=0, the original test passes 8/8. The test now compares graph replay against the same TE modules run eagerly, bitwise. It keeps the PyTorch reference with a TF32 tolerance.

Results on the new head 2506cec2 (main fac8d4a9 merged in, no conflicts):

GPU	`test_tied_delayed_wgrad.py`	`test_cuda_graphs.py`
GH200 (sm_90)	50/50 passed	464 passed, 423 skipped (incl. all 26 new delayed-wgrad graph cases)
L40S (sm_89)	34 passed, 16 skipped (packed grouped weights / FP8 block scaling unavailable on sm_89)	320 passed, 333 skipped (incl. all 26 new delayed-wgrad graph cases)

The CUDA-graph hazard you described (zero_grad() freeing the captured .grad) is handled since dbe1326e: test_make_graphed_callables_with_delayed_wgrad covers it with allocator-pressure guards and passes 16/16.

I don't have SM100 access right now, so I'd appreciate it if you could rerun there before CI. pylint (3.3.1, repo rc) is clean on the touched files.

Pair each delayed Linear with its regular counterpart directly instead of indexing both Sequentials, as suggested in review. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

0z5a · 2026-10-28T16:16:08Z

@denera Follow-up on Blackwell. I ran the tests at the current head (d14c2d28) on a B300 (sm_103). TE was built from main fac8d4a9 for 103a, with this PR's Python files applied on top:

TE Python package	Tests	Result
main	`test_tied_delayed_wgrad.py`	49 of 50 failed (the bug these tests target)
main + this PR	`test_tied_delayed_wgrad.py` + `test_cuda_graphs.py`	752 passed, 0 failed, 1121 skipped

The 24/50 failures you saw on SM100 do not occur on sm_103 with the current head. The three graph tests this PR adds also pass, 26 cases in all:

test_make_graphed_callables_with_delayed_wgrad, which graphs backward_dw() and compares against eager, as you suggested
test_graphed_delayed_wgrad_accumulates_microbatches
test_graphed_delayed_wgrad_preserves_main_grad

0z5a requested a review from ksivaman as a code owner September 21, 2026 04:51

github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Sep 21, 2026

0z5a mentioned this pull request Sep 21, 2026

Tied weights do not accumulate with delayed wgrad #3437

Open

greptile-apps Bot reviewed Sep 21, 2026

View reviewed changes

Comment thread tests/pytorch/test_tied_delayed_wgrad.py Outdated

wujingyue reviewed Sep 21, 2026

View reviewed changes

0z5a force-pushed the tied-delayed-wgrad-accumulate branch from 0f419b2 to 6681062 Compare September 24, 2026 01:11

wujingyue suggested changes Sep 24, 2026

View reviewed changes

0z5a force-pushed the tied-delayed-wgrad-accumulate branch from 6681062 to a1ac680 Compare September 24, 2026 06:05

0z5a added 2 commits September 24, 2026 15:07

0z5a force-pushed the tied-delayed-wgrad-accumulate branch from a1ac680 to 557a35b Compare September 24, 2026 07:07

denera self-assigned this Oct 2, 2026

denera self-requested a review October 2, 2026 16:21

denera requested changes Oct 2, 2026

View reviewed changes

fix(pytorch): accumulate all delayed tied-parameter gradients

c5cf151

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

0z5a requested a review from timmoon10 as a code owner October 7, 2026 20:15

Merge main and preserve delayed tied-gradient accumulation

67a310d

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

denera requested changes Oct 8, 2026

View reviewed changes

0z5a added 3 commits October 8, 2026 22:09

Merge main into tied-delayed-wgrad-accumulate

10a1c98

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

Iterate the non-tied test's module pairs with zip

d14c2d2

Pair each delayed Linear with its regular counterpart directly instead of indexing both Sequentials, as suggested in review. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>

		return _tie(model, bias) if tied else model


		def _run(model, x, microbatches=1, delay=False, order=(0, 1)):

		return model


		def _build(delay, *, bias=False, device="cuda"):

	def _build(delay, *, bias=False, device="cuda"):
	def _build(*, delay, bias, device="cuda"):

	def _build(delay, *, bias=False, device="cuda"):
	def _build_two_linears(delay, *, bias=False, device="cuda"):

Conversation

0z5a commented Sep 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

greptile-apps Bot commented Sep 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Summary

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wujingyue left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot commented Sep 24, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

denera left a comment

Choose a reason for hiding this comment

Uh oh!

denera left a comment • 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!

0z5a commented Oct 27, 2026

Uh oh!

0z5a commented Sep 21, 2026 •

edited

Loading

greptile-apps Bot commented Sep 21, 2026 •

edited

Loading

denera left a comment •

edited

Loading