Skip to content

Optimize torch.compile's recompilation limit #6142

Description

@nzmora-nvidia

System Info

TLDR; when using cuda-graphs we compile once for each batch size and this triggers a recompilation of the kernels. Once the same frame has been recompiled that many times, Dynamo stops compiling it further and falls back to normal eager execution to avoid compile thrashing - and we will lose perf.

[rank1]:W0716 01:10:47.676000 567101 torch/_dynamo/convert_frame.py:984] [0/8] torch._dynamo hit config.recompile_limit (8)

Details:

The debug trace pasted above is emitted by PyTorch Dynamo when we run with a compile-based backend (torch-compile, torch-opt, etc.).
PyTorch Dynamo/​torch.compile captured the model’s forward once and produced a compiled graph that is guarded by the exact tensor shapes it first saw. The next time we call the model (we send requests with different batch / sequence-length combinations) those shape guards are checked.
When an input violates a guard (e.g. input_ids dimension was expected 1, got 128) the compiled graph cannot be reused, so Dynamo recompiles the function. After the 8-th unique shape it reached the limit and issued the warning.
==>
Subsequent calls to that function will now run in regular eager mode on that rank, so we lose the speed-ups from torch.compile for those inputs.

Task

We need to raise the limit so it's high enough to prevent eager mode.
This is the backend code: tensorrt_llm/_torch/auto_deploy/compile/backends/torch_opt.py

   import torch._dynamo as dynamo
   dynamo.config.recompile_limit = 32

Suggestions from @lucaslie

+        # TODO: every graph capture will trigger a recompile, hence we need a higher limit
+        # TODO: refine this fix:
+        #   1. rewrite CapturedGraph utility to always take cudagraph batch sizes as input
+        #       a. max batch size can be inferred from cudagraph batch sizes
+        #       b. move heuristic to batch sizes from max_batch_size to TorchCudagraphCompiler class
+        #   2. Set the recompile limit here to be max(len(cuda_graph_batch_sizes), torch._dynamo.config.recompile_limit)
+        #   3. Consider moving cache_size_limit config to here or TorchCompile backend as well
+        #      --> right now, it's hard-coded into build_and_run_ad.py

Who can help?

No response

Information

  • The official example scripts
  • My own modified scripts

Tasks

  • An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
  • My own task or dataset (give details below)

Reproduction

d

Expected behavior

d

actual behavior

d

additional notes

TLDR; when using cuda-graphs we compile once for each batch size and this triggers a recompilation of the kernels. Once the same frame has been recompiled that many times, Dynamo stops compiling it further and falls back to normal eager execution to avoid compile thrashing - and we will lose perf.

[rank1]:W0716 01:10:47.676000 567101 torch/_dynamo/convert_frame.py:984] [0/8] torch._dynamo hit config.recompile_limit (8)

Details:

The debug trace pasted above is emitted by PyTorch Dynamo when we run with a compile-based backend (torch-compile, torch-opt, etc.).
PyTorch Dynamo/​torch.compile captured the model’s forward once and produced a compiled graph that is guarded by the exact tensor shapes it first saw. The next time we call the model (we send requests with different batch / sequence-length combinations) those shape guards are checked.
When an input violates a guard (e.g. input_ids dimension was expected 1, got 128) the compiled graph cannot be reused, so Dynamo recompiles the function. After the 8-th unique shape it reached the limit and issued the warning.
==>
Subsequent calls to that function will now run in regular eager mode on that rank, so we lose the speed-ups from torch.compile for those inputs.

Task

We need to raise the limit so it's high enough to prevent eager mode.
This is the backend code: tensorrt_llm/_torch/auto_deploy/compile/backends/torch_opt.py

   import torch._dynamo as dynamo
   dynamo.config.recompile_limit = 32

Suggestions from @lucaslie

+        # TODO: every graph capture will trigger a recompile, hence we need a higher limit
+        # TODO: refine this fix:
+        #   1. rewrite CapturedGraph utility to always take cudagraph batch sizes as input
+        #       a. max batch size can be inferred from cudagraph batch sizes
+        #       b. move heuristic to batch sizes from max_batch_size to TorchCudagraphCompiler class
+        #   2. Set the recompile limit here to be max(len(cuda_graph_batch_sizes), torch._dynamo.config.recompile_limit)
+        #   3. Consider moving cache_size_limit config to here or TorchCompile backend as well
+        #      --> right now, it's hard-coded into build_and_run_ad.py

Metadata

Metadata

Labels

AutoDeploy<NV> AutoDeploy BackendInvestigatingPerformanceTRTLLM model inference speed, throughput, efficiency. Latency, benchmarks, regressions, opts.triagedIssue has been triaged by maintainers

Type

Projects

Status
Done

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions