SIGN IN SIGN UP

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