[PyTorch] [torch.compile] torch.compile support for LayerNormLinear and LayerNormMLP

[pggPL]

Description

Add torch.compile custom-op support for LayerNormLinear and LayerNormMLP, including backward, so supported configurations can run with fullgraph=True.

Changes

• Add forward and backward fake implementations covering output shapes, dtypes, saved tensors and aliases.
• Support dynamic batch dimensions, normalization outputs and supported FP8 recipes.
• Share output-gradient setup and quantizer validation across Linear and both LayerNorm modules. Keep unsupported configurations on the eager fallback path.
• Validate both MLP GEMMs before custom-op dispatch, including FC2 dimensions for gated activations and sequence parallelism.
• Parameterize existing compile tests across the three modules and cover LayerNorm-specific variants and gradients.

Validation

• Existing targeted GPU compile suite: 135 passed, 3 skipped, 1 xpassed.
• Manual checks confirmed that invalid FC2 dimensions are rejected during tracing for GELU and SwiGLU.
• Eight manual FP32-parameter/BF16-autocast cases passed output and gradient comparisons against eager, covering both LayerNorm modules, LayerNorm/RMSNorm and FP8 CurrentScaling on/off.
• Pylint on the latest touched production files: 10/10. Python syntax and whitespace checks passed.

Validation used a single RTX 5880 Ada GPU with PyTorch 2.14.0a0+b2c75dd062.nvinternal.main. Distributed TP/SP and Userbuffers execution under compile remain unvalidated.

[greptile-apps]

Confidence Score: 4/5

The PR is not yet safe to merge because the exported custom-op registration API still breaks existing backward argument containers that rely on grad_output or grad_outputs fields.

Findings

1.  Backward API Compatibility Break ▶

Summary

Adds full-graph torch.compile support for LayerNormLinear and LayerNormMLP, extending the existing custom-operation infrastructure used by Linear.

• Adds fake forward and backward implementations describing outputs, saved tensors, aliases, gradients, and quantized metadata.
• Adds compiled dispatch eligibility checks and eager fallback for unsupported configurations.
• Preserves NVFP4 normalization scale-layout metadata and validates both LayerNormMLP GEMMs.
• Expands eager-versus-compiled tests across modules, recipes, dynamic batches, normalization variants, activations, gradients, and CUDA-graph modes.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[LayerNormLinear or LayerNormMLP call] --> B{Compile configuration supported?}
    B -- No --> C[Eager fallback]
    B -- Yes --> D[Registered forward custom op]
    D --> E[Fake implementation supplies shapes and metadata]
    D --> F[Native forward implementation]
    F --> G[User outputs and saved tensors]
    G --> H[Autograd slices output gradients]
    H --> I[setup_grad_outputs]
    I --> J[Registered backward custom op]
    J --> K[Input and parameter gradients]

Loading

_{Reviews (3) · Last reviewed commit: "Preserve explicit quantized scale layout..."}

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