[torch.compile][PyTorch] Prepare linear for torch compile

[pggPL]

Description

Refactor of transformer_engine/pytorch/module/linear.py to lift the
Linear module into a shape that can be wrapped in a
torch.library.custom_op (and matching backward op) in a follow-up PR.

Type of change

• Code refactoring

Changes

• Pack forward/backward state into LinearFwdArgs / LinearBwdArgs
  dataclasses. _linear_forward_impl, _linear_setup_ctx,
  _linear_backward and the _Linear.{forward,backward} autograd
  methods all take a single structured argument instead of 25+
  positional ones. A custom op requires a fully-declared signature on
  both sides; the previous pattern of writing arbitrary
  ctx.something = ... attributes scattered throughout forward made
  it impossible to tell from the call site what state backward
  actually consumes. The dataclasses make the read/write contract
  explicit and grep-able. Concretely, things that used to be re-queried
  from tensor objects (input.requires_grad, weight.requires_grad,
  bias.requires_grad) are now captured up front as
  input_requires_grad / weight_requires_grad /
  bias_requires_grad and consumed as requires_dgrad /
  requires_wgrad in backward — backward no longer has to assume the
  Python tensor objects survive the op boundary.
• Move prepare_for_saving / ctx.save_for_backward to the autograd
  boundary. _linear_forward_impl returns the raw tensors it
  produced (along with tensors_to_save_from_forward aliases) and
  _linear_setup_ctx returns the raw merged tensor list it wants
  saved; _Linear.forward is the only place that actually calls
  prepare_for_saving(*tensors_to_save_from_setup) and
  ctx.save_for_backward(...). This shape fell out of the compile-path
  experiments: under torch.library.register_autograd, the
  setup_context callback is the only legal place to call
  save_for_backward, and the helper has to hand back tensors rather
  than mutate the autograd ctx itself. Same contract is now used in
  eager so a single helper serves both modes.
• Deduplicate saved tensors that alias forward inputs. Save-slots
  that would alias inp / weight / bias are emitted as None and
  reconstructed in _linear_setup_ctx from the original refs. An
  opaque custom op cannot return aliases of its inputs (the tracer
  has no way to reason about the aliasing).
• Minimize the ctx_attrs blob plumbed from forward to backward
  setup. Anything that can be re-derived from LinearFwdArgs
  (weight_quantizer, is_fsdp2, owns_input) is recomputed in
  _linear_setup_ctx. The compile path needs the fake forward impl
  to return a structurally identical ctx_attrs, so a smaller surface
  is a smaller cross-impl contract.

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

This PR is a large refactor of transformer_engine/pytorch/module/linear.py in preparation for torch.compile support. It packs all forward/backward state into LinearFwdArgs / LinearBwdArgs dataclasses, moves prepare_for_saving / ctx.save_for_backward to the autograd boundary, deduplicates saved tensors that alias forward inputs via an alias map, and shrinks the ctx_attrs blob crossing the forward/backward boundary.

• LinearFwdArgs / LinearBwdArgs dataclasses replace the ~25-element positional-argument tuples, making the forward→backward data contract explicit and static; the _Linear.forward signature shrinks to (weight, inp, bias, fwd_args).
• Saved-tensor deduplication uses a saved_tensor_aliases tuple to mark which save slots alias forward inputs; _linear_setup_ctx reconstructs those refs instead of saving them, which is required for the opaque-custom-op tracer.
• ctx.backward_objects stores the populated LinearBwdArgs on the autograd ctx so backward can retrieve it; the reference is nulled after use to prevent tensor lifetime extension under retain_graph.

Confidence Score: 4/5

The refactor is internally consistent and the new dataclass contract is well-structured, but the backward path has a known crash when the same graph node is traversed a second time.

The forward→backward data handoff via LinearFwdArgs/LinearBwdArgs is clean and the alias-deduplication logic is correct. However, ctx.backward_objects is nulled after the first backward call; a second backward on the same node (retain_graph=True) immediately dereferences None and raises AttributeError, which is a real behavioral regression from the prior flat-ctx approach.

transformer_engine/pytorch/module/linear.py — specifically the _Linear.backward method and the ctx.backward_objects lifecycle under retain_graph.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/linear.py	Core refactor file: introduces LinearFwdArgs/LinearBwdArgs dataclasses, restructures forward/backward plumbing, and changes weight_quantizer derivation in _linear_setup_ctx to use weight._quantizer for QuantizedTensor weights. The retain_graph second-backward crash (ctx.backward_objects=None after first backward) remains present.
tests/pytorch/test_backward_override.py	Test helper updated to look up backward state via grad_fn.backward_objects (LinearBwdArgs) with fallback to grad_fn; getattr default does not trigger when backward_objects exists but is None, so a post-backward call yields a confusing error instead of the fallback.

Sequence Diagram

sequenceDiagram
    participant L as Linear.forward
    participant FA as LinearFwdArgs
    participant FI as _linear_forward_impl
    participant SC as _linear_setup_ctx
    participant BA as LinearBwdArgs
    participant CTX as autograd ctx
    participant BW as _Linear.backward
    participant BI as _linear_backward

    L->>FA: construct(weight, inp, bias, ...)
    L->>FI: _linear_forward_impl(fwd_args)
    FI-->>L: out, new_wks, tensors_to_save_from_forward, ctx_attrs
    L->>BA: LinearBwdArgs()
    L->>SC: _linear_setup_ctx(bwd_args, fwd_args, out, ctx_attrs, tensors)
    SC-->>BA: populate all backward fields
    SC-->>L: (saved_inputmat, wt_save, saved_weight, saved_bias)
    L->>CTX: "prepare_for_saving(*tensors)"
    L->>CTX: ctx.save_for_backward(...)
    L->>CTX: "ctx.backward_objects = bwd_args"

    BW->>CTX: "bwd_args = ctx.backward_objects"
    BW->>BA: "bwd_args.grad_output = grad_output"
    BW->>BA: bwd_args.setup_saved_tensors(ctx)
    BW->>BI: _linear_backward(bwd_args)
    BI-->>BW: (wgrad, dgrad, grad_bias)
    BW->>CTX: "ctx.backward_objects = None"

Loading

_{Reviews (3): Last reviewed commit: "Merge branch 'main' into prepare_linear_..." | Re-trigger Greptile}

[pggPL]

/te-ci pytorch L1

[timmoon10]

LGTM. Most of my comments are just me talking through the design.

[timmoon10]

If we destroy the cached state, we should also mark backward with function.once_differentiable.

[timmoon10]

This design is nicer than exposing a long list of positional args or passing in a tuple of args, but it is less nice than exposing kwargs. However, it has some advantages for this case:

• torch.compile infrastructure will only need to handle one non-tensor arg to the autograd function.
• Autograd functions do some processing on each arg, so minimizing the number of non-tensor args reduces CPU overhead.
• _Linear.forward calls _linear_forward_impl and it doesn't need to be aware of the impl specifics.

[timmoon10]

In typical usage, LinearBwdArgs is somewhat redundant with the autograd context class. But we need it so that the forward-context-saving and backward are less entangled with autograd, which will allow code reuse when we implement an alternate torch.compile code path.

[pggPL]

/te-ci pytorch L1