[user empathy] compile for transformers model

@mcarilli

🐛 Describe the bug

I was looking at the current integration of compile for transformers repo. One of my main goal was to understand if we can reduce the compile time by moving torch.compile to a repeated block, instead of the full model. In diffusion models, this strategy leads to 8x-10x reduction in compile time w/o sacrificing perf benefits compared to full model compilation (data in this PR). So, naturally I was curious about LLMs, because if true, it can significantly improve ux/adoption.

But I was hit with multiple issues. This issue tries to create a small test case to bring those issues forward.

import torch


class Cache:
    def __init__(self, num_layers):
        self.key_cache = []
        self.value_cache = []
        for i in range(num_layers):
            k = torch.zeros(4, 4, device="cuda")
            v = torch.zeros(4, 4, device="cuda")
            torch._dynamo.mark_static_address(k)
            torch._dynamo.mark_static_address(v)
            self.key_cache.append(k)
            self.value_cache.append(v)
        

    def get(self, idx):
        return self.key_value[idx], self.value_cache[idx]

    def update(self, new_k, new_v, idx):
        self.key_cache[idx].copy_(new_k)
        self.value_cache[idx].copy_(new_v)
        return self.key_cache[idx], self.value_cache[idx]

class Transformer(torch.nn.Module):
    def __init__(self, layer_idx):
        super().__init__()
        self.layer_idx = layer_idx

    def forward(self, x, cache):
        k = torch.sin(x)
        v = torch.cos(x)
        k, v = cache.update(k, v, self.layer_idx)
        return k + v

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.num_layers = 16
        self.layers = [Transformer(i) for i in range(self.num_layers)]
        self.cache = Cache(self.num_layers)

    def forward(self, x):
        for i in range(self.num_layers):
            x = self.layers[i](x, self.cache)
        return x
        
mod= Model().cuda() 
x = torch.randn(4, 4, device="cuda")
mod(x)

def full_model_compile():
    # This works
    # https://manifold.edge.x2p.facebook.net/v0/read/tree/logs/.tmpsKSTQ8/index.html?bucketName=tlparse_reports&apiKey=tlparse_reports-key&withPayload=1&timeoutMsec=10000
    opt_mod = torch.compile(mod, mode="reduce-overhead", fullgraph=True)
    opt_mod(x)

# full_model_compile()

def transformer_compile():
    for submod in mod.layers:
        submod.compile(fullgraph=True)
        # submod.compile(fullgraph=True, mode="reduce-overhead")
    mod(x)

torch.compiler.reset()
transformer_compile()

In the above repro, full model works. This is also the current integration point, but the compile time issue has been raised multiple times.

The function transformer_compile compiles just the transformer layers. There are the problems

Too many recompilations on self.layer_idx - last reason: 0/7: self.layer_idx == 7. The interesting thing in this case is the graph is same for all layer_ids, but the pre-graph bytecode to extract the tensor is different. This might be very hard to fix completely in dynamo.
cudagraph issue (cc @mcarilli @ezyang @eellison @penguinwu @BoyuanFeng @chauhang @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames @oulgen @jamesjwu @aorenste @laithsakka @masnesral). This issues exists in diffusers models.

RuntimeError: Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run. Stack trace: File "/home/anijain/local/pytorch/examples/hf_gen_example.py", line 34, in forward
    return k + v. To prevent overwriting, clone the tensor outside of torch.compile() or call torch.compiler.cudagraph_mark_step_begin() before each model invocation.

Recompilations because of mark_static_address - I attempted to remove the self.layer_idx guard by refactoring Transformer method to read the past_k and past_v values from the cache, and then directly do copy_ on them later on (instead of using self.layer_idx). Something like this


class TransformerFailedAttempt(torch.nn.Module):
    def __init__(self, layer_idx):
        super().__init__()
        self.layer_idx = layer_idx


    @torch.compile
    def _forward(self, x, past_k, past_v):
        k = torch.sin(x)
        v = torch.cos(x)
        k, v = cache.update_with_pastkv(k, v, past_k, past_v)
        return k + v


    def forward(self, x, cache):
        past_k, past_v = cache.get(self.layer_idx)
        return self._forward(x, past_k, past_v)

But this too failed because mark_static_address guards on the ID of the tensors, resulting in id guard failures on past_k and past_v. This is discussed at #156377

Note - I tried 3 models - Gemma3, Gemma2 and Qwen3 (ok effectively 2). All suffer from same exact issues, so I think the many models will suffer the same fate.

Error logs

cc @chauhang @penguinwu @zou3519 @BoyuanFeng @eellison @StrongerXi @ezyang

cc @gante for transformers

Versions

nightly

@anijain2305

Hi @anijain2305 (and PT team) 👋

The accelerate team has been exploring this topic, and it seems they have reached a working solution. More specifically, this is a PR containing benchmarks, using this function. It might not be as efficient in the decoding phase though (1 token at a time), due to its low computational intensity.

I'm looping in the accelerate team members related to this work: @IlyasMoutawwakil @SunMarc

CUDAGraph issue Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run. when compiling each layer vs compiling full model.

I have a smaller repro.

Example 1 (error when compiling each layer):

import torch

def f(x):
    return x + 1

def g(x):
    return x + 2

f = torch.compile(f, mode="reduce-overhead")
g = torch.compile(g, mode="reduce-overhead")

for _ in range(3):
    # torch.compiler.cudagraph_mark_step_begin() # uncomment to fix the error
    x = torch.randn(10, device="cuda")
    y = f(x)
    z = g(y) # "Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run"

new_id = torch._inductor.cudagraph_trees.get_manager(0).new_graph_id()

print(new_id)
# print: GraphID(id=2)

Example 2 (succeed when compiling full model):

import torch

def f(x):
    return x + 1

def g(x):
    return x + 2

def h(x):
    y = f(x)
    z = g(y)
    return z

h = torch.compile(h, mode="reduce-overhead")

for _ in range(3):
    # no need for torch.compiler.cudagraph_mark_step_begin() here
    x = torch.randn(10, device="cuda")
    z = h(x)

new_id = torch._inductor.cudagraph_trees.get_manager(0).new_graph_id()

print(new_id)
# print: GraphID(id=1)

Why compiling full model is fine but compiling each layer is bad? Since cg tree needs to decide whether it's 1) root->f->g, or 2)

root -> f
root -> g

By default, cg tree assumes 2). So when we start to execute g, it cleans all output memory of f to reuse the memory. Then, g accesses f's output and errors.

torch.compiler.cudagraph_mark_step_begin() gives extra information to compiler that it's 1). So when executing g, output memory of f is still kept and the error does not happen.

When compiling the full model, we only have 1 function h so it's fine.

Question: Would torch.compiler.cudagraph_mark_step_begin() fix the issue?

As we discussed, we may have a config to clone CUDAGraph outputs for avoiding the CUDAGraph issue. In the layerwise compilation setting, the cloned outputs is no longer a static tensor address and requires clone again as the inputs to the next layer. I.e.,

x -> f's cg -> static_address_y_as_output -> (clone) -> eager_y -> (clone) -> static_address_y_as_input -> g's cg.

Sunbelt Computer Software

PL/B Language Development and Support

Navigation Menu

Repository navigation

[user empathy] compile for `transformers` model #156520

🐛 Describe the bug

Error logs

Versions

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Sunbelt Computer Software

PL/B Language Development and Support

[user empathy] compile for transformers model #156520

Description

🐛 Describe the bug

Error logs

Versions

Activity

gante commented on Jun 23, 2025

BoyuanFeng commented on Jun 23, 2025

BoyuanFeng commented on Jun 23, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

[user empathy] compile for `transformers` model #156520