[PyTorch] Torch.compile support

[pggPL]

Make TE layers compatible with torch.compile.

[sbhavani]

Related:

• How about the torch.compile in TransformerEngine ? #1241
• torch.compile graph breaks at forward #705

[cspades]

A few months ago, I also had trouble compiling TE fused RoPE in TorchTitan nightly:

        # Apply RoPE.
        # xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
        xq = apply_rotary_pos_emb_te(
            xq,
            # Needs to be real-valued for TE fused RoPE.
            freqs_cis.unsqueeze(1).unsqueeze(1),
            tensor_format="bshd",
            interleaved=False,
            fused=True,
        )
        xk = apply_rotary_pos_emb_te(
            xk,
            freqs_cis.unsqueeze(1).unsqueeze(1),
            tensor_format="bshd",
            interleaved=False,
            fused=True,
        )

@sbhavani I think I may or may not have talked to you about this in our many conversations around that time, but this blocked me from getting Llama 405B metrics with FSDP2 + Torch Compile with this fused RoPE (for parity with Megatron-LM's fused RoPE) so we had to math out the best-case improvement instead.

I tried this adding it as a custom op, but didn't have time to dig further.

[pstjohn]

Here's a simple repro

import torch
import transformer_engine.pytorch
from transformer_engine.pytorch.attention.rope import RotaryPositionEmbedding

layer = transformer_engine.pytorch.TransformerLayer(128, 128, 8)
rope_layer = RotaryPositionEmbedding(128 // 8)

inputs = torch.randn(8, 128).to("cuda")
freqs = rope_layer.forward(128)

inputs.to("cuda")
freqs.to("cuda")
layer.to("cuda")

layer = torch.compile(layer, fullgraph=True)
output = layer(inputs, rotary_pos_emb=freqs)

[ntadimeti]

related issue: #1241

[pggPL]

There is draft PR to add torch.compile support for te.Sequential #2608 - te.sequential will be the first TE layers will torch.compile support.

[ben-da6]

Is this not relatively easy to add, and I think would be a big win. All that is needed is to register some ops and remove the explicit no dynamo tags. For me I see no improvements using TE as we compile our whole graph, with all the graph breaks its not faster or more memory efficient.

[sbhavani]

@ben-da6 we would welcome a PR! It'd be great to see even in a small scale test where torch.compile could speed up

[pggPL]

We are actively working on that, this is much harder that it seems - there are multiple issues reported by us to PyTorch which will need to be closed before full torch.compile support will be merged. They are mostly related to new PyTorch OpaqueObject API, which allows to passing arbitrary objects as an inputs to torch custom operators. For example one merged PR here.

There are some parts of the TE which can be converted simply - like MoE related permutation ops which we will merge soon. Then we work on te.Linear and te.Sequential - which are more complex, but we already have some prototypes and we work towards exposing them as non-draft PRs. The attention is the most complex part, we have some issues there - for example, but not limited to, choice of the attention backend based on tensor dptr() parameters. It will be merged probably as the last part.

[omaralvarez]

Any updates on this?

[Abhishekghosh1998]

Any updates on this?

[pggPL]

We have draft PR #3053 which is actively developed and the reviews will start soon.

[honeyspoon]

I think the hf kernels framework is compatible with torch.compile
I think the best options would be to just expose the custom kernels in this repo there instead