pytorch - 💡(How to fix) Fix [inductor] Scheduler phases scale poorly on lowered aten.lstm/aten.gru graphs

Problem

When torch._dynamo.config.allow_rnn = True is set, dynamo traces through nn.LSTM / nn.GRU successfully and produces a small FX graph (12–32 nodes). Inductor then decomposes the single aten.lstm op into a much larger lowered graph. The scheduler phases operating on this expanded graph take a long time — 7s for a 2-layer LSTM, scaling up to 80s+ for 6-layer models, making compilation slow for common RNN configurations.

This is a one-time compilation cost (cached runs are fast), but it's high enough that users hitting #158007 who try allow_rnn=True as suggested by the error message will likely give up thinking it's hung.

Measured data

All measurements from isolated subprocesses with TORCHINDUCTOR_FORCE_DISABLE_CACHES=1 on PyTorch 2.13.0a0+git795be92, NVIDIA H200.

Lowered graph expansion — inductor decomposes the aten.lstm op into many buffers:

12 FX nodes become 67 lowered buffers for a 1-layer LSTM with h=64. Adding a second layer doubles this to 134.

Where the time goes — phase breakdown from PyTorch's own compile_times():

The dominant cost is the scheduler: Scheduler.__init__ + Scheduler.fused_nodes + Scheduler.codegen together account for ~55% of compile_fx time and roughly double when the buffer count doubles (1L→2L).

For larger models (6-layer, h=256), earlier wall-clock measurements showed compile_fx taking 80+ seconds.

Reproduction

import torch
import torch.nn as nn
import time
import os

os.environ["TORCHINDUCTOR_FORCE_DISABLE_CACHES"] = "1"
torch._dynamo.config.allow_rnn = True

model = nn.LSTM(32, 64, num_layers=2, batch_first=True).cuda().eval()
x = torch.randn(2, 16, 32, device="cuda")

def timing_backend(gm, example_inputs):
    from torch._inductor.compile_fx import compile_fx
    print(f"FX graph: {len(list(gm.graph.nodes))} nodes")
    t0 = time.time()
    result = compile_fx(gm, example_inputs)
    print(f"compile_fx: {time.time() - t0:.1f}s")
    return result

compiled = torch.compile(model, fullgraph=True, backend=timing_backend)
with torch.no_grad():
    out = compiled(x)

from torch._dynamo.utils import compile_times
print(compile_times(repr="str"))

Expected output: compile_fx: ~7s, with Scheduler.__init__ and Scheduler.codegen as the top phases.

Context

• Dynamo tracing works correctly with allow_rnn=True — compiled outputs match eager within float32 precision
• torch.export is unaffected since it doesn't use inductor
• Cached compilations are fast (~0.05s), so this is purely a cold-compile cost
• The scheduler scaling is likely not LSTM-specific — any op that decomposes into 100+ lowered buffers would likely hit similar costs

Related

• #158007 — torch.compile cannot compile LSTM (blocked by allow_rnn=False default)
• #141184 — (separate dynamo bug, 0 graphs compiled)
• https://github.com/pytorch/pytorch/pull/163207

Environment

PyTorch: 2.13.0a0+git795be92
CUDA: 13.0
GPU: NVIDIA H200
OS: RHEL 9 (kernel 5.14.0-615.el9.x86_64)
Python: 3.12

cc @jerryzh168 @mikaylagawarecki @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo @oulgen @jamesjwu @aorenste @anijain2305 @laithsakka @masnesral @aditvenk @chenyang78