[pggPL]

Description

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

The single late fallback check relies on the Dynamo fix in pggPL/pytorch#4, related to pytorch/pytorch#144820. Without that fix, a graph break inside forward's try/finally can disable compilation for subsequent supported instances of the same module class.

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.

• Decide fallback eligibility once after preparing quantizers in Linear, LayerNormLinear and LayerNormMLP; reuse the prepared eager path and preserve end_forward cleanup.

• 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.

• Support compiled FP8 weight caching in both LayerNorm modules, following Linear: replace the cache on the first microbatch and reuse it on subsequent microbatches, including both MLP weights. GPU-controlled cache updates and fused wgrad accumulation retain eager fallbacks.

• Share the cache regression across Linear, LayerNormLinear and LayerNormMLP, checking independent eager outputs/gradients, cache refresh/reuse, optimizer updates, immediate/deferred backward and CUDA Graph replay.

• 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.

• Preserve the native RMSNorm MXFP8 2D scale layout in fake normalization metadata; cover metadata reconstruction and add CustomRecipe MXFP8 2D cases to the compiled normalization scale-layout matrix.

[greptile-apps]

Confidence Score: 4/5

[Medium impact] The changes reach PyTorch users compiling several Transformer Engine modules.

The PR still needs the previously reported backward API compatibility fix before merging.

Findings

1.  Backward API Compatibility Break ▶

Summary

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

• LayerNormLinear and LayerNormMLP run inside compiled forward and backward graphs.
• Compiled normalization metadata follows the native MXFP8 scale layouts.
• Compiled operations share symbolic shapes and output-gradient setup.

The previous and current commits have different merge bases with the target commit, so this review used the full PR diff.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Module forward] --> B[Prepare input and quantizers]
  B --> C{Supported compiled state?}
  C -->|Yes| D[Custom forward operation]
  C -->|No| E[Eager operation]
  D --> F[Save tensors and output metadata]
  E --> G[Finish forward cleanup]
  F --> G
  F --> H[Assign output gradients]
  H --> I[Custom backward operation]

Loading

_{Reviews (30) · Last reviewed commit: "[PyTorch] Add torch.compile support for ..." · Reviewed by Greptile}

[greptile-apps]

Comments Outside Diff

These findings could not be posted inline.

•  Recipe changes reuse unsupported plans transformer_engine/pytorch/ops/fuser.py:709 ▶

  The cache key now includes only recipe type, backward boundary, and backward_override. Fusion also depends on disable_rht, MXFP8 format, and enable_cutedsl_fused_grouped_mlp. A same-class recipe change returns the old plan without checking those settings. For example, a plan built with NVFP4 RHT stays selected after RHT is disabled, although the fused path requires it. Include the fusion-relevant recipe settings in the key.

[ptrendx]

Why is this checking specifically for RMSNorm? This sounds like the original logic had some bug where only RMSNorm was checked rather than any normalization.

[ptrendx]

This is a very specific function and takes very specific arguments. Could we generalize it a little bit?