[PyTorch] [torch.compile] transformer_engine.pytorch.autocast suport inside torch.compile

[pggPL]

Description

Enable torch.compile(fullgraph=True) for FP8 autocast by moving compile-visible mutable state off class attributes, avoiding tracing through support checks, and adding test.

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

Please list the changes introduced in this PR:

• Move mutable FP8 autocast state from direct cls attribute writes to a dataclass-backed singleton object, because torch.compile does not support writes directly to class attributes.
• Replace lru_cache-based support checks with explicit module-level caches and mark the wrapper functions with @torch.compiler.assume_constant_result so torch.compile does not trace into check_*_support().
• Add torch.compile coverage for FP8 autocast using a custom test module; the test is more involved because there is currently no simple TE layer that supports both FP8 and torch.compile.
• Make DelayedScaling explicitly unsupported under torch.compile and raise a clear error for that case.

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

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

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

Please list the changes introduced in this PR:

• Change A
• Change B

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 enables torch.compile(fullgraph=True) for FP8 autocast by migrating mutable global state from class-level attributes to a FP8GlobalState dataclass instance (quantization_state) and replacing lru_cache-based support checks with @torch.compiler.assume_constant_result-decorated wrappers. The mechanical refactor across modules is clean and consistent, and the ToyLinear + opaque-type test infrastructure is a thoughtful solution to the difficulty of testing FP8 + compile together.

Several P1 concerns surfaced in prior review rounds remain open — most critically, autocast_enter unconditionally calls get_default_fp8_recipe() when no recipe is supplied (even when enabled=False), which will AssertionError under torch.compile; reset() swaps the singleton object rather than mutating it in-place, invalidating compiled-graph guards on quantization_state; and the null-guard on skip_fp8_weight_update_tensor at graph replay time is still missing in graph.py.

Confidence Score: 4/5

• The core state-migration refactor is correct and the torch.compile guard approach is sound, but several P1 issues from prior review rounds remain unaddressed and should be fixed before merge.
• Score reflects that the mechanical refactor across 10 files is clean and consistent, but open issues from prior rounds (unconditional get_default_fp8_recipe() crash path, reset() breaking compiled-graph guards, missing null guard in graph replay, assert strippable under -O) are confirmed present in the submitted code and represent real defect paths in the compile-enabled flow this PR is trying to unlock.
• transformer_engine/pytorch/quantization.py (autocast_enter recipe path, reset(), fp8_graph_capturing assert) and transformer_engine/pytorch/graph.py (null guard on skip_fp8_weight_update_tensor at line 841)

Important Files Changed

Filename	Overview
transformer_engine/pytorch/quantization.py	Core of the PR: replaces class-level attribute writes with a FP8GlobalState dataclass singleton and adds @torch.compiler.assume_constant_result to support checks. Several open P1 concerns from prior review rounds remain unresolved (unconditional get_default_fp8_recipe() in autocast_enter, reset() swapping the singleton object, assert strippable under -O).
tests/pytorch/test_torch_compile.py	New test file with sophisticated ToyLinear + opaque-type infrastructure for testing torch.compile; however test_autocast_sanity exercises only torch.nn.Linear (not TE ops), no DelayedScaling-under-compile error test, and private torch._opaque_base/torch._library.opaque_object APIs are used without stability guards.
transformer_engine/pytorch/graph.py	Moves skip_fp8_weight_update_tensor initialization inline with a null check at setup time; the inner call site on line 841 still has no null guard and will crash with AttributeError if reset() is called between graph capture and replay.
transformer_engine/pytorch/module/base.py	Correctly redirects amax buffer writes to quantization_state; update order of global_amax_history_buffer and global_amax_buffer is reversed compared to ops/op.py (cosmetic inconsistency, not a correctness bug).

Class Diagram

%%{init: {'theme': 'neutral'}}%%
classDiagram
    class FP8GlobalState {
        +bool fp8_enabled
        +bool fp8_calibration
        +Recipe fp8_recipe
        +dist_group_type fp8_distributed_group
        +bool fp8_parameters
        +bool high_precision_init_val
        +bool is_first_fp8_module
        +bool fp8_graph_capturing
        +int autocast_depth
        +Dict global_amax_buffer
        +Dict global_amax_history_buffer
        +Dict global_scale_buffer
        +list fp8_tensors_recompute_buffer
        +Dict autocast_arguments
        +Tensor skip_fp8_weight_update_tensor
    }

    class FP8GlobalStateManager {
        +FP8GlobalState quantization_state
        +reset() void
        +autocast_enter() void
        +autocast_exit() void
        +is_fp8_enabled() bool
        +fp8_graph_capturing() bool
        +is_first_fp8_module() bool
        +get_fp8_recipe() Recipe
        +get_autocast_state() tuple
        +set_autocast_state() void
        +reduce_and_update_fp8_tensors() void
    }

    class ModuleLevelCaches {
        +_FP8_SUPPORT: Optional[Tuple]
        +_MXFP8_SUPPORT: Optional[Tuple]
        +_NVFP4_SUPPORT: Optional[Tuple]
        +_FP8_BLOCK_SCALING_SUPPORT: Optional[Tuple]
    }

    class SupportChecks {
        +check_fp8_support() [assume_constant_result]
        +check_mxfp8_support() [assume_constant_result]
        +check_nvfp4_support() [assume_constant_result]
        +check_fp8_block_scaling_support() [assume_constant_result]
    }

    FP8GlobalStateManager --> FP8GlobalState : quantization_state (singleton)
    SupportChecks --> ModuleLevelCaches : reads/writes
    FP8GlobalStateManager --> SupportChecks : calls

Loading

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

[pggPL]

/te-ci pytorch L1

[pggPL]

/te-ci pytorch

[pggPL]

/te-ci pytorch

[ptrendx]

/te-ci pytorch