[PyTorch] Enable LayerNormLinear UB AG overlap under no_grad

[ravimajeti]

Description

LayerNormLinear's forward UserBuffers All-Gather overlap (ub_overlap_ag_fprop) was additionally gated on torch.is_grad_enabled(). A caller could explicitly request the overlap and still have it silently disabled whenever the forward ran under torch.no_grad().

This matters for activation checkpointing/recomputation: with reentrant checkpointing the initial forward runs under no_grad while the backward-time recompute runs with gradients enabled, so exactly one of the two identical forwards lost the overlapped communication/GEMM path — for a reason unconnected to the forward computation itself. Gradient mode governs autograd bookkeeping (what to save for backward); it is not a reason to suppress a requested forward communication overlap.

This PR removes is_grad_enabled from the forward-path selection only:

-    ub_overlap_ag_fprop = ub_overlap_ag_fprop and is_grad_enabled and not return_layernorm_output
+    ub_overlap_ag_fprop = ub_overlap_ag_fprop and not return_layernorm_output

Every other gradient-mode check in _layernorm_linear_forward_impl (backward_needs_input, backward quantizer setup, saved tensors) is untouched.

Scope: LayerNormLinear only. LayerNormMLP carries the same gate on its ub_overlap_ag and ub_overlap_rs; it is not changed here and is left for a follow-up so it can get its own test and validation. (Linear does not have the gate.)

Fixes #2902

Type of change

• Bug fix (non-breaking change which fixes an issue)

Changes

• transformer_engine/pytorch/module/layernorm_linear.py: decouple ub_overlap_ag_fprop from is_grad_enabled.
• tests/pytorch/distributed/run_layer_with_overlap.py: add --no-grad — runs the forward inside torch.no_grad(), skips backward, compares the output against the non-overlap reference, and asserts fill_userbuffers_buffer_for_all_gather() was actually invoked (so the test checks that the overlap path ran, not only that the numbers match).
• tests/pytorch/distributed/test_comm_gemm_overlap.py: add test_layernorm_linear_all_gather_overlap_no_grad, parametrized over BF16 (none) and FP8 (fp8_delayed_scaling).

Testing

All runs use the repository's own test files with the same invocations as qa/L1_pytorch_distributed_unittest/test.sh and qa/L0_pytorch_unittest/test.sh, on 4x H100 SXM (NVLink), TP=4, PyTorch 2.11.0+cu128, CUDA 12.8, built from this branch as rebased on main (ae34b34f).

Check	This branch	Only the one-line predicate change reverted
test_layernorm_linear_all_gather_overlap_no_grad[none] / [fp8_delayed_scaling] (added here)	2 passed	2 failed: AssertionError: Userbuffers AllGather overlap was not used under no_grad
Full tests/pytorch/distributed/test_comm_gemm_overlap.py (154 cases: grad-enabled LayerNormLinear via test_layers_with_overlap_bf16/fp8, Linear, multi-layer TransformerLayer, split/bulk overlaps, torch.compile)	55 passed, 89 skipped, 10 failed*	53 passed, 89 skipped, 12 failed* — the delta is exactly the 2 new cases
tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py	3 passed	—
tests/pytorch/test_numerics.py -k "recompute or checkpoint" (170 cases incl. all 96 test_gpt_full_activation_recompute variants: FP8 on/off, DelayedScaling/Float8CurrentScaling, reentrant/non-reentrant)	170 passed	—

* The 10 failures are all test_linear_with_overlap_compile[...] (te.Linear under torch.compile(fullgraph=True)) and are unrelated to this change: they fail identically on unmodified main (ae34b34f) in this environment. With PyTorch 2.11.0 the custom-op registration in dynamo/custom_op.py is disabled at import (cannot import name 'is_opaque_constant_type' from 'torch._library.opaque_object'), so te.Linear takes its torch._dynamo.disabled eager fallback and fullgraph=True raises Unsupported. They need a newer PyTorch build. The 89 skips are the mxfp8 / cuBLASMp hardware-and-build gates.

Lint: pre-commit run --files on the three changed files passes all hooks (black with the repo's settings, vermin, whitespace/EOF checks); pylint reports no messages for the changed source file.

Performance note (ad-hoc microbenchmark, not part of the test suite)

The fix makes the existing overlap path reachable under no_grad; it does not add a new one, so the speedup is whatever UB all-gather overlap already gives on the grad-enabled path. For reference, a forward-only single-LayerNormLinear microbenchmark (TP=4, hidden 4096, seq 8192, 200 timed iterations after warmup, overlap vs. ub_overlap_ag=False) measured +12–14 % in BF16 (A100-SXM4) and +21–22 % in FP8 (H100-SXM), each reproduced twice. As with the grad-enabled path, overlap is a net loss at small shapes (e.g. −12 to −17 % at hidden ≤ 2048, seq ≤ 2048), where the overlap setup cost exceeds the all-gather latency it hides.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation (no user-facing docs reference this internal gate)
• My changes generate no new warnings
• I have added tests that prove my fix is effective
• New and existing unit tests pass locally with my changes

🤖 Generated with Claude Code

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; the newly reachable no-grad overlap path uses existing forward-only machinery and is covered by path-sensitive numerical tests.

Summary

This PR enables the explicitly configured LayerNormLinear UserBuffers all-gather overlap during no-grad forwards while retaining the existing restrictions around returned layer-normalization outputs and backward-only state.

• Removes gradient mode from the forward overlap-selection predicate.
• Adds a no-grad execution mode to the distributed overlap harness and verifies that the UserBuffers path is actually invoked.
• Adds BF16 and delayed-scaling FP8 regression coverage against a non-overlap reference.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[LayerNormLinear forward] --> B{Column parallel and sequence parallel?}
    B -- No --> C[Use local normalized input]
    B -- Yes --> D{UB all-gather overlap requested and layernorm output not returned?}
    D -- Yes --> E[Fill UserBuffers local shard]
    E --> F[Fused all-gather and GEMM]
    D -- No --> G[Standard all-gather]
    G --> H[GEMM]
    F --> I[Forward output]
    H --> I
    I --> J{Gradient recording enabled?}
    J -- Yes --> K[Save backward state]
    J -- No --> L[Return without backward state]

Loading

_{Reviews (1) · Last reviewed commit: "Fix LayerNormLinear UB overlap under no-..."}