Enable torch.compile for RNN modules by default (#185902)
My agent generated the quoted summary below. It is included because it explains the root cause, implementation choice, CI baseline updates, and validation for #149909. > Dynamo previously graph-broke on RNN, GRU, and LSTM modules unless > `torch._dynamo.config.allow_rnn` was enabled manually. This made > `torch.compile(..., fullgraph=True)` fail immediately for these modules. > Enable the existing RNN tracing path by default so these modules can be > captured. > > Enabling that path exposes an LSTM decomposition bug on CPU. The mkldnn > one-layer LSTM decomposition was selected whenever only the input did not > require gradients. In normal module usage, the input may not require gradients > while parameters do, so AOTAutograd selected an inference-only mkldnn path and > then failed when it treated mkldnn's undefined workspace output as a tensor. > Gate the mkldnn decomposition on grad mode and every LSTM tensor input so > parameter-gradient cases use the differentiable Python fallback. > > Keeping RNN tracing opt-in would preserve the old graph break and require users > to patch `allow_rnn`, so it would not fix the reported compile behavior. The > default flip intentionally changes TorchBench capture metrics. Record the exact > CUDA inference/training and CPU inference graph-break, captured-call, and > eager-fallback values emitted by CI. The CPU-specific rows are required after > the baseline split on current main; accuracy continues to pass for every > affected model. > > Test Plan: > > ``` > python test/dynamo/test_modules.py NNModuleTests.test_rnn_modules_compile_by_default > python test/inductor/test_cpu_repro.py CPUReproTests.test_lstm_compile_default_grad_enabled > ``` > > ``` > python benchmarks/dynamo/check_graph_breaks.py --actual <(printf '%s\n' 'name,accuracy,graph_breaks,calls_captured,fallbacks_to_eager,unique_graphs' 'demucs,pass,0,69,0,1' 'doctr_reco_predictor,pass,3,53,0,2' 'tts_angular,pass,3,16,0,5') --expected benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv > python benchmarks/dynamo/check_graph_breaks.py --actual <(printf '%s\n' 'name,accuracy,graph_breaks,calls_captured,fallbacks_to_eager,unique_graphs' 'demucs,pass,6,93,0,2' 'tts_angular,pass,10,80,0,7') --expected benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_training.csv > python benchmarks/dynamo/check_graph_breaks.py --actual <(printf '%s\n' 'name,accuracy,graph_breaks,calls_captured,fallbacks_to_eager,unique_graphs' 'demucs,pass,0,69,0,1' 'doctr_reco_predictor,pass,3,51,0,1' 'tts_angular,pass,0,17,0,1') --expected benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_torchbench_inference.csv > ``` > > ``` > lintrunner -a > git diff --cached --check > git diff HEAD~ --check > ``` > > Benchmark Results: > > - With `allow_rnn=False`, fullgraph Inductor LSTM fails with `Unsupported: Attempted to wrap RNN, GRU, or LSTM` before a compiled runtime can run. > - With RNN tracing enabled, fullgraph Inductor LSTM ran at a median 294.27 us (11.61 us IQR, 1663 measurements, one thread). Fixes #149909 Pull Request resolved: https://github.com/pytorch/pytorch/pull/185902 Approved by: https://github.com/anijain2305, https://github.com/mlazos
J
Jason Ansel committed
a2ed3ae5b69bb2ec3427f3ca18a27ae1c100b229
Parent: 68d20d4
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com>
on 8/24/2026, 2:06:07 AM