[pggPL]

Description

Add full-graph torch.compile support for LayerNormLinear and LayerNormMLP, including backward, dynamic batch dimensions and normalization outputs.

Targets main. This branch includes the NVFP4 packed-buffer, TensorSpec and Bias correctness fixes from NVIDIA/TransformerEngine#3628.

Changes

• Describe forward outputs, saved tensors, aliases and backward gradients for both LayerNorm modules; retain eager fallbacks for unsupported configurations.
• Share output-gradient setup and quantizer validation with Linear. Adapt the newly inherited fused-attention backward container to the same gradient setup contract, preserving its existing first-output gradient assignment.
• Validate both MLP GEMMs, including gated activations and sequence-parallel dimensions.
• Describe LayerNorm/RMSNorm output, mean and reciprocal standard deviation through dynamo.fake.fake_norm, used by both LayerNorm modules. Share the fused-quantization layout predicate with activation output metadata, including the RHT-fusion disable flag and output device; remove duplicate layout helpers. Preserve MXFP8/cuDNN normalization layouts and gathered LayerNormMLP output shapes.
• Parameterize module compile tests across recipes, shapes and modes, with LayerNorm-specific outputs, gradients and eager fallbacks. Retain the non-matrix Linear regressions and expanded cold CUDA Graph recipe coverage inherited from [PyTorch] Fix NVFP4 metadata in compiled Linear operations #3628.

Validation

B100/SM100, PyTorch 2.15.0a0+875d815502.nvinternal.main, CUDA 13.4; existing October 7 native libraries reused, no fresh native build.

• Focused LayerNorm compile selection: 180 passed, 4 expected skips for unsupported LayerNormMLP backward overrides. Includes full-graph forward/backward, non-matrix inputs, normalization output gradients, variants and eager fallbacks.
• The selection includes 60 normalization scale-layout cases: both modules, LayerNorm/RMSNorm, MXFP8 and NVFP4, cuDNN settings, shapes, returned norm output, and NVFP4 RHT fusion enabled/disabled.
• The inherited cold CUDA Graph regression is identical to the version validated in [PyTorch] Fix NVFP4 metadata in compiled Linear operations #3628: 48 passed with an independent eager model and no implicit recipe construction for unquantized inputs.
• Pre-commit, license, syntax and whitespace checks passed. Full PyTorch Python lint retains an unchanged C0411 import-order diagnostic in attention/fused_mla_q_uproj.py in the GPU container; no unrelated source cleanup is included. Pylint on all changed production files passed.

No complete compile-suite rerun or distributed TP/SP/Userbuffers validation. The gathered-output test emulates the collective and does not establish multi-GPU correctness.

[greptile-apps]

Confidence Score: 4/5

[Medium impact] The changes affect compiled PyTorch module workflows and shared backward paths.

The PR still needs the backward compatibility fix before merging.

Findings

1.  Backward API Compatibility Break ▶

Summary

Adds full-graph torch.compile support for LayerNormLinear and LayerNormMLP.

• LayerNormLinear runs as a compiled operation with backward support.
• LayerNormMLP runs as a compiled operation across both matrix steps.
• Each compiled operation assigns gradients for all its public outputs.
• Fake tensor descriptions now match normalization and activation outputs.

No new findings were identified in the changes since the previous review.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[LayerNorm module call] --> B{Compile settings supported?}
  B -->|No| C[Eager fallback]
  B -->|Yes| D[Registered forward operation]
  D --> E[Fake outputs describe shapes and buffers]
  D --> F[Real forward saves backward state]
  F --> G[Output gradients]
  G --> H[setup_grad_outputs]
  H --> I[Registered backward operation]
  I --> J[Input and parameter gradients]

Loading

_{Reviews (14) · Last reviewed commit: "Remove unused LayerNorm compile metadata" · Reviewed by Greptile}

[greptile-apps]

Backward API Compatibility Break

register_custom_op is exported from transformer_engine.pytorch.dynamo, but backward now unconditionally calls setup_grad_outputs. Existing integrations that follow the previous grad_output or grad_outputs field convention can therefore register and compile successfully, then fail with AttributeError during backward. Please preserve the field-based fallback or provide a compatibility default while callers migrate.

[greptile-apps]

Comments Outside Diff

These findings could not be posted inline.

•  Quantized fake tensors fail transformer_engine/pytorch/dynamo/tensor_spec.py:102 ▶

  When a quantized TensorSpec does not explicitly set with_gemm_swizzled_scales, these lines omit that value from the metadata passed to storage reconstruction. BasicLinear’s fake quantized forward and dgrad paths create such specs, but the MXFP8 and NVFP4 storage constructors require this argument. Tracing therefore fails while assembling the fake tensor, before the compiled operation can run. Please preserve the quantizer-derived default when the field is unspecified while still allowing fake paths to override a known native layout.