[torch.compile] Bunch of small changes needed for enabling torch.compile

[pggPL]

Description

Small standalone fixes extracted from a larger torch.compile branch, going directly from main. Two independent changes: making Userbuffers pybind11 queries compile-friendly, and freeing quantized grad_output early for column-parallel SP. Plus a custom-recipe caching fix, a split-accumulator refactor, and a CI test hook-up.

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

Changes

1. Userbuffers pybind11 queries under torch.compile

• is_fp8_ubuf() / with_cublasmp() are compile-time constants but graph-break when traced. At the nn.Module.forward boundary (where no UB communicator object is in hand yet) they go through get_ub_is_fp8(name, use_fp8), wrapped in torch.compiler.assume_constant_result — only plain (str, bool) args are baked, so guards are well-defined and don't rely on pybind-object identity.
• In the hot forward/backward implementation paths the UB communicator is already fetched, so those call ub_obj.is_fp8_ubuf() / ub_obj.with_cublasmp() directly — no wrapper, no string concatenation, no redundant registry lookup. Eager speed is preserved.

2. Free quantized grad_output early for column-parallel SP

• Row-parallel SP already called clear_tensor_data(grad_output) on the gathered tensor. Column-parallel SP quantizes grad_output to a Float8TensorStorage (an internal tensor) but never freed it. Under torch.compile reduce-overhead this left live pool tensors at recording end ("Detected N tensor(s) in the cudagraph pool not tracked as outputs"). The free now covers row-SP and column-SP-FP8 (column-SP non-FP8 is a no-op view, so it's excluded).

3. Replace fp8_recipe in LinearBwdArgs with pre-resolved split-accumulator booleans

• LinearBwdArgs no longer carries the recipe object (which holds process-group references and is compile-unfriendly). dgrad_use_split_accumulator / wgrad_use_split_accumulator are resolved once in Linear.forward (reusing the existing get_fp8_recipe() call) and threaded through as plain booleans.

4. Custom-recipe quantizer caching fix

• CustomRecipeState early-exit was missing an identity check, so quantizers were rebuilt on every forward even when the recipe was unchanged. Added if recipe_state.recipe is recipe: return.

5. Test hook-up

• Added test_torch_compile.py to L0_pytorch_unittest.

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 bundles five targeted torch.compile-readiness fixes extracted from a larger branch. Each change is independent and non-breaking.

• UB pybind11 graph-break fix: get_ub_is_fp8() wraps is_fp8_ubuf() with @torch.compiler.assume_constant_result so the module boundary calls are folded as compile-time constants; destroy_ub() now calls torch.compiler.reset() to invalidate those constants on teardown.
• Column-parallel SP FP8 pool-tensor leak: clear_tensor_data(grad_output) is extended to cover the column-parallel SP FP8 path (in addition to row-SP) in both linear.py and layernorm_linear.py — the quantized Float8TensorStorage is freed immediately after the wgrad GEMM, after which it is no longer referenced.
• LinearBwdArgs recipe-object removal: fp8_recipe is replaced by two plain booleans (dgrad_use_split_accumulator, wgrad_use_split_accumulator) resolved once at forward time, eliminating a compile-unfriendly process-group reference from the autograd context.
• CustomRecipeState rebuild fix: Adds an is identity guard so quantizers are not re-built on every forward when the custom recipe object is unchanged.
• CI hook-up: test_torch_compile.py is added to the L0 test suite; the test's qfactory is corrected to dispatch on QuantizerRole.tensor_type (the previous string-keyed dict would have raised KeyError at runtime).

Confidence Score: 5/5

All changes are well-scoped and non-breaking; the column-SP FP8 tensor free is correctly placed after the wgrad GEMM and the freed variable is not accessed again.

Each change is narrowly targeted: the split-accumulator refactor preserves the same defaults for non-FP8 and reproduces the same hasattr guards for FP8 recipes; the grad_output free only triggers after the GEMM that consumed it; the CustomRecipeState identity check uses is which is appropriate for a mutable recipe object; and torch.compiler.reset() in destroy_ub() correctly handles stale cached constants. No correctness regressions were found across the traced code paths.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/base.py	Adds get_ub_is_fp8() wrapper with @assume_constant_result, adds identity check to CustomRecipeState early-exit, and calls torch.compiler.reset() in destroy_ub() to invalidate cached constants after UB teardown.
transformer_engine/pytorch/module/linear.py	Replaces fp8_recipe in LinearBwdArgs with pre-resolved split-accumulator booleans, extends grad_output early-free to column-parallel SP FP8 (after wgrad GEMM), and switches UB fp8 queries to get_ub_is_fp8().
transformer_engine/pytorch/module/layernorm_linear.py	Applies the same column-parallel SP FP8 grad_output early-free as linear.py, and switches UB fp8 queries to get_ub_is_fp8(). ctx.fp8_recipe remains for split-accumulator (not in scope of this PR).
transformer_engine/pytorch/module/layernorm_mlp.py	Single-line change: switches get_ub(...).is_fp8_ubuf() to get_ub_is_fp8() in the fc2_fprop path for compile-time constant folding.
tests/pytorch/test_torch_compile.py	Adds get_quantizer_roles() to ToyLinear module and fixes qfactory to dispatch on QuantizerRole.tensor_type; previous string keys like 'linear_input' would have caused a KeyError.
qa/L0_pytorch_unittest/test.sh	Adds test_torch_compile.py to the L0 CI test suite.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant M as Linear.forward()
    participant GSM as FP8GlobalStateManager
    participant FA as LinearFwdArgs
    participant BA as LinearBwdArgs
    participant BW as _linear_backward()

    M->>GSM: get_fp8_recipe()
    GSM-->>M: _recipe
    M->>M: resolve dgrad/wgrad split-accumulator bools
    M->>FA: fwd_args(dgrad_use_split_accumulator, wgrad_use_split_accumulator)
    FA->>BA: _linear_setup_ctx copies plain bools
    Note over FA,BA: No recipe object stored

    BW->>BW: "use_split_accumulator = bwd_args.dgrad_use_split_accumulator"
    BW->>BW: dgrad GEMM
    BW->>BW: "use_split_accumulator = bwd_args.wgrad_use_split_accumulator"
    BW->>BW: wgrad GEMM
    BW->>BW: clear_tensor_data(grad_output) [row-SP or col-SP+FP8]

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant M as Linear.forward()
    participant GSM as FP8GlobalStateManager
    participant FA as LinearFwdArgs
    participant BA as LinearBwdArgs
    participant BW as _linear_backward()

    M->>GSM: get_fp8_recipe()
    GSM-->>M: _recipe
    M->>M: resolve dgrad/wgrad split-accumulator bools
    M->>FA: fwd_args(dgrad_use_split_accumulator, wgrad_use_split_accumulator)
    FA->>BA: _linear_setup_ctx copies plain bools
    Note over FA,BA: No recipe object stored

    BW->>BW: "use_split_accumulator = bwd_args.dgrad_use_split_accumulator"
    BW->>BW: dgrad GEMM
    BW->>BW: "use_split_accumulator = bwd_args.wgrad_use_split_accumulator"
    BW->>BW: wgrad GEMM
    BW->>BW: clear_tensor_data(grad_output) [row-SP or col-SP+FP8]

Loading

_{Reviews (3): Last reviewed commit: "Provide explicit QuantizerRoles in torch..." | Re-trigger Greptile}

[pggPL]

/te-ci pytorch L1

[IvanYashchuk]

The current helper call sites are all inside @no_torch_dynamo() forwards and the added test_torch_compile.py coverage does not exercise user buffers or it's done implicitly in the test?

Is it possible avoid a process-wide compiler reset on UB teardown, or add a targeted compiled UB test that proves the stale-constant case and justifies this global invalidation?

[IvanYashchuk]

That file only compiles a local ToyLinear helper and torch.nn.Linear under te.autocast. It does not instantiate changed in this PR te.Linear, LayerNormLinear, or LayerNormMLP, and it has no UB, sequence_parallel/parallel_mode.

What tests would fail without changes to layernorm_linear, layernorm_mlp files?

[pggPL]

I fix the issue that the test was not connected to the CI.
Currently it tests only if te.autocast() can be traced inside torch.compile.

This is first of series of PRs and I change here only small things to make next PRs cleaner.