[pggPL]

Description

Restructure LayerNormLinear and LayerNormMLP the same way Linear was restructured for torch.compile, so that the follow-up PR registering them as torch.compile custom ops is a pure wiring change.

No functional change intended.

Type of change

• Code refactoring

Changes

• LayerNormLinearFwdArgs / LayerNormLinearBwdArgs and LayerNormMLPFwdArgs / LayerNormMLPBwdArgs dataclasses replace the positional non_tensor_args tuple and the ad-hoc ctx.<attr> plumbing.
• The autograd Functions become thin wrappers around module-level _layernorm_{linear,mlp}_forward_impl / _setup_ctx / _backward_impl that take one dataclass argument (the shape the custom-op framework expects). Saved tensors that alias forward inputs or outputs are deduplicated by name (saved_tensor_aliases) and rebuilt in setup_ctx.
• LayerNormMLP keeps activation checkpointing: the forward saves only its inputs, and _layernorm_mlp_recompute re-runs the forward impl in backward with recompute_for_bwd set, filling the backward args from the recomputed state.
• Weight workspaces move from positional autograd inputs into the args bag and are returned as extra outputs (as in Linear). Backward-only recipe knobs (dgrad/wgrad_use_split_accumulator, backward_override, debug UB flags, delayed wgrad_store) are resolved in Module.forward.
• check_fp8_reduce_and_update moves to module/_common.py and is shared with Linear.

Testing (RTX Ada, single GPU)

• tests/pytorch/test_sanity.py -k "layernorm_linear or layernorm_mlp": 7344 passed
• tests/pytorch/test_numerics.py -k "layernorm_linear or layernorm_mlp": 738 passed
• tests/pytorch/test_backward_override.py: 1316 passed
• tests/pytorch/test_cuda_graphs.py -k "layernorm or mlp": 105 passed
• test_gpt_full_activation_recompute (checkpoint recompute path): same pass/fail pattern as main (known Ada flakiness)
• pylint on the touched files: 10/10

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
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; no new actionable correctness, security, or repository-rule violations remain.

Summary

This PR restructures LayerNormLinear and LayerNormMLP into module-level forward, context-setup, and backward implementations suitable for subsequent torch.compile custom-op registration.

• Replaces positional and ad hoc autograd state with typed forward and backward argument dataclasses.
• Deduplicates saved tensor aliases and reconstructs them during autograd context setup.
• Preserves LayerNormMLP activation recomputation and moves backward-only configuration into module forward methods.
• Shares FP8 reduce-and-update bookkeeping with Linear through _common.py.
• Follow-up changes correctly handle MXFP8 grad-output gathering for cuBLASMp weight-gradient paths.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    M[Module.forward] --> A[Forward argument dataclass]
    A --> F[Forward implementation]
    F --> O[Outputs and weight workspace]
    F --> S[Saved tensors and context attributes]
    O --> C[Autograd setup_ctx]
    S --> C
    C --> B[Backward argument dataclass]
    B --> R{LayerNormMLP recomputation?}
    R -- No --> I[Backward implementation]
    R -- Yes --> RF[Re-run forward implementation]
    RF --> I
    I --> G[Input and parameter gradients]

Loading

_{Reviews (3) · Last reviewed commit: "Merge main into layernorm_compile_refact..."}

[pggPL]

CPU-overhead check vs. base d16c754 (RTX Ada, eager, tiny shapes [8, 2, 256] so the run is host-bound; single thread, pinned core; 6 interleaved rounds × 9 trials × 300 iters, median CPU time per fwd+bwd step):

module	recipe	base	PR	Δ
LayerNormLinear	bf16	209 µs	217 µs	+8 µs (+3.7%)
LayerNormLinear	fp8 (delayed)	360 µs	365 µs	+5 µs (+1.3%)
LayerNormMLP	bf16	387 µs	399 µs	+12 µs (+3.2%)
LayerNormMLP	fp8 (delayed)	469 µs	479 µs	+10 µs (+2.2%)

[negvet]

LGTM, pending ci and merge confix fix

Minor questions:

• Is eager CPU overhead comparable to Linear?
• I guess the compiled version built on top of this PR is semi-ready? Do we have measurements of its expected benefits?

In #3343, which enables recipe update at runtime, I enforce recompilation with the recipe update. I wonder how much time it takes to re-compile some representable model e.g. DSv3.

[pggPL]

/te-ci pytorch

[pggPL]

Is eager CPU overhead comparable to Linear?

Yes

I guess the compiled version built on top of this PR is semi-ready? Do we have measurements of its expected benefits?

It is not, but we aim to reducing cpu overhead of non-te parts of the model. My PRs will allow TE not to do graph break which adds addidtional cpu overhead and prevents as from using cuda graphs backend.

In #3343, which enables recipe update at runtime, I enforce recompilation with the recipe update. I wonder how much time it takes to re-compile some representable model e.g. DSv3.

I haven't measure it, but it is one-time cost and everything is cached, so if we change recipe more than once the cost should be neglible.