torch.compile Not Applied to FlexAttention with SAC `selective_ac_option=op`

[fegin]

Bug description

Description

When using the SAC with the option selective_ac_option=op, torch.compile is not applied to FlexAttention. This occurs despite TorchTitan explicitly adding calls to torch.compile(flex_attention).

As a result, the following warning is emitted and results in a very low performance:

[rank0]:
[rank0]:SOLUTION: Use torch.compile(flex_attention)(...)
[rank0]:
[rank0]:If you want to debug your score_mod/mask_mod, you can set:
[rank0]:torch.nn.attention.flex_attention._FLEX_ATTENTION_DISABLE_COMPILE_DEBUG = True
[rank0]:
[rank0]:This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results.

Steps to Reproduce

1. Use the default DeepSeek 16B configuration, or
2. Use the debug model with the flag --model.flavor=debugmodel_flex_attn.

Additional Information

• This issue does not occur when:
  • No Activation Checkpointing (AC) is used,
  • Full Activation Checkpointing is used,
  • SAC is used with selective_ac_option=2.

Versions

Nightly

[fegin]

cc., @tianyu-l, @soulitzer, @drisspg

[soulitzer]

This is expected if we're running eager SAC because TorchDispatchMode is not supported w/ torch.compile today, causing the compilation of fn to be skipped.

[fegin]

The expected behavior is torch.compile(SAC(flex_attention)), whereas TorchTitan currently does SAC(torch.compile(flex_attention)).

The main difference between flex_attention and other torch.compile use cases in TorchTitan is that flex_attention is a single operation inside an Attention module, which itself resides within a TransformerBlock. Since we also apply SAC at the granularity of the TransformerBlock, torch.compile(SAC(TransformerBlock)) is feasible. However, achieving torch.compile(SAC(flex_attention)) is challenging unless SAC is applied at a very fine-grained level.

At this point, I can only think of erroring out when FlexAttention and SAC are used together. We can do layer SAC (it is actually a full_ac at TransformerBlock level) or full AC.

cc., @tianyu-l

[tianyu-l]

This is expected if we're running eager SAC because TorchDispatchMode is not supported w/ torch.compile today, causing the compilation of fn to be skipped.

@xmfan @bdhirsh is this something fundamentally difficult?

@soulitzer would AC2 help solve this issue?

At this point, I can only think of erroring out when FlexAttention and SAC are used together. We can do layer SAC (it is actually a full_ac at TransformerBlock level) or full AC.

This will be quite inefficient for the layers that do full AC -- the blocking a2a and d2h sync will all appear again in the backward.

[soulitzer]

Even if TorchDispatchMode were supported by compile, SAC wrapping compile wouldn't compose because compile duplicated the nodes already and the forward and backward graphs are opaque to SAC.

AC2 also uses TorchDispatchMode (even without SAC!), so it should also run into this issue.

@fegin is trying out temporarily disabling SAC for just the compile(flex_attn op, I think that could work, although we're relying on an private torch-dispatch API to do this.

[fegin]

Using _pop_mode_temporarily doesn't work as this API doesn't update _is_in_non_infra_torch_dispatch_mode global variable. If we do _pop_mode_temporarily plus manually changing _is_in_non_infra_torch_dispatch_mode value, then FlexAttention can be compiled. Need a more formal guideline how we can achieve this.

cc., @soulitzer @bdhirsh

[fegin]

Tried torch.compile(SAC(Flex)) and we got an error message like this:

[rank0]:[rank0]:     return wrapper()
[rank0]:[rank0]:   File "/data/users/chienchin/mywork/pytorch/torch/_ops.py", line 531, in wrapper
[rank0]:[rank0]:     return self.dispatch(
[rank0]:[rank0]:   File "/data/users/chienchin/mywork/pytorch/torch/_ops.py", line 422, in dispatch
[rank0]:[rank0]:     raise NotImplementedError(
[rank0]:[rank0]: torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
[rank0]:[rank0]: NotImplementedError: There was no rule registered for HigherOrderOperator triton_kernel_wrapper_mutation and mode <torch.utils.checkpoint._CachingTorchDispatchMode object
at 0x7f1f5520a6b0>.Hint: set <torch.utils.checkpoint._CachingTorchDispatchMode object at 0x7f1f5520a6b0>'s supports_higher_order_operators to True. This causes all higher order operators to pass through <torch.utils.checkpoint._CachingTorchDispatchMode object at 0x7f1f5520a6b0>'s __torch_dispatch__, so handle them accordingly by adding support for HigerOrderOperators (in this case, triton_kernel_wrapper_mutation) in <torch.utils.checkpoint._CachingTorchDispatchMode
object at 0x7f1f5520a6b0>.__torch_dispatch__ or returning NotImplemented when not supported.

cc., @soulitzer @drisspg

[fegin]

pytorch/pytorch#161541 will fix the compilation error.

#1645 adds the warning when --compile.enable is not set. This PR depends on the above one.