[PyTorch] Fix deferred initialization in fusible ops

[denera]

Description

Ops replace their params when materializing them on the first fuser forward pass, which invalidated the fuser's cached params and ops.Linear's parameter aliases. Backward then failed with a meta-vs-cuda device mismatch.

• Re-cache the basic ops params after first materialization
• Run first forward initialization on the top-level ops, before the cached params are used to pick which ops require backward
• Sync te.ops.Linear weight/bias aliases with its basic ops
• Add deferred initialization tests for the affected ops

Fixes #3322

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

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]

Greptile Summary

The PR fixes deferred initialization for PyTorch fusible operations by initializing top-level operations before backward selection, refreshing cached basic-operation parameters after materialization, and synchronizing Linear parameter aliases. It also adds CUDA forward/backward coverage for deferred initialization across basic, grouped, composite, sequential, and quantized operations.

• Preserve the original top-level operation list so composite initialization hooks run before fusion planning.
• Rebuild flattened parameter caches after meta-device parameters materialize.
• Synchronize Linear.weight and Linear.bias with their underlying basic operations.
• Test materialization, gradients, device placement, alias identity, and quantized weights.

Confidence Score: 5/5

The PR appears safe to merge, with no concrete blocking or non-blocking defects identified in the changed paths.

First-forward initialization now reaches nested basic operations, synchronizes wrapper aliases, and refreshes the parameter cache before backward eligibility and fusion planning consume it.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fuser.py	Defers parameter caching until first-forward initialization and orders recipe reset, top-level materialization, cache refresh, and fusion planning consistently.
transformer_engine/pytorch/ops/linear.py	Re-registers Linear parameter aliases after its underlying basic operations replace meta-device parameters.
tests/pytorch/test_fusible_ops.py	Adds deferred-initialization forward/backward tests covering device placement, gradients, aliases, grouped parameters, composition, and quantized weights.

Sequence Diagram

sequenceDiagram
  participant User
  participant Fuser as OperationFuser
  participant Basic as Basic operations
  participant Top as Top-level operations
  participant Autograd
  User->>Fuser: First forward
  Fuser->>Basic: reset_recipe_state(recipe)
  Fuser->>Top: pre_first_fuser_forward()
  Top->>Basic: materialize meta parameters
  Top->>Top: synchronize parameter aliases
  Fuser->>Basic: cache materialized parameters
  Fuser->>Fuser: select and build forward/backward fusions
  Fuser->>Autograd: execute with refreshed parameters
  Autograd-->>User: outputs and parameter gradients

Loading

_{Reviews (1): Last reviewed commit: "[PyTorch] Fix deferred initialization in..." | Re-trigger Greptile}

[denera]

/te-ci pytorch