prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  🐛 [Bug] Encountered TRT Error in Dynamo compile path
  ##  Bug Description
  When compiling GPT-2 with Dynamo compile, the following error is encountered:
  ```python
  [07/01/2023-00:07:12] [TRT] [E] 3: [executionContext.cpp::enqueueInternal::795] Error Code 3: API Usage Error (Parameter check failed at: runtime/api/executionContext.cpp::enqueueInternal::795, condition: bindings[x] || nullBindingOK
  )
  ```
  Additionally, this does not seem to cause failures in the Dynamo runtime (not caught by `pass_through_build_failures`), and it appears both with and without the experimental runtime.
  
  ## To Reproduce
  ```python
  model = GPT2Model.from_pretrained("gpt2").eval().cuda()
  input_ids      = torch.randint(0, 2, (1, 14), dtype=torch.int32).cuda()
  attention_mask = torch.randint(0, 2, (1, 14), dtype=torch.int32).cuda()
  traced = transformers_trace(model, input_names=["input_ids", "attention_mask"]).eval().cuda()
  fx_trt_model = torch_tensorrt.compile(traced, ir="dynamo_compile", inputs=[input_ids, attention_mask], debug=True, pass_through_build_failures=True, min_block_size=10)
  ```
  
  ## Expected behavior
  
  The model should not encounter TRT errors while compiling.
  
  ## Environment
  
   - Torch-TensorRT Version (e.g. 1.0.0): 2844630
   - PyTorch Version (e.g. 1.0): `2.1.0.dev20230620+cu118`
   - TensorRT Version: `8.6.1`
  
  ## Additional context
  
  This error no longer appears when we apply the `@fake_tensor_unsupported` flag to the backends, as was removed in #1955. It is unclear whether this is the direct _cause_ of the bug, however.
  
  Additionally, this bug seems to appear only sometimes, and not consistently.
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `py/torch_tensorrt/dynamo/backend/backends.py`
  Content:
  ```
  1 import logging
  2 from typing import Sequence
  3 import torch
  4 from functools import partial
  5 import torch._dynamo as td
  6 
  7 from torch_tensorrt.dynamo.backend._settings import CompilationSettings
  8 from torch_tensorrt.dynamo.backend.lowering._decompositions import (
  9     get_decompositions,
  10 )
  11 from torch_tensorrt.dynamo.backend.lowering._pre_aot_lowering import (
  12     pre_aot_substitutions,
  13 )
  14 from torch_tensorrt.dynamo.backend.lowering._partition import (
  15     partition,
  16     get_submod_inputs,
  17 )
  18 from torch_tensorrt.dynamo.backend.utils import parse_dynamo_kwargs
  19 from torch_tensorrt.dynamo.backend.conversion import convert_module
  20 
  21 from torch._functorch.aot_autograd import aot_module_simplified, make_boxed_compiler
  22 
  23 
  24 logger = logging.getLogger(__name__)
  25 
  26 
  27 @td.register_backend(name="torch_tensorrt")
  28 def torch_tensorrt_backend(
  29     gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor], **kwargs
  30 ):
  31     DEFAULT_BACKEND = aot_torch_tensorrt_aten_backend
  32 
  33     return DEFAULT_BACKEND(gm, sample_inputs, **kwargs)
  34 
  35 
  36 @td.register_backend(name="aot_torch_tensorrt_aten")
  37 def aot_torch_tensorrt_aten_backend(
  38     gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor], **kwargs
  39 ):
  40     settings = parse_dynamo_kwargs(kwargs)
  41 
  42     custom_backend = partial(
  43         _pretraced_backend,
  44         settings=settings,
  45     )
  46 
  47     # Perform Pre-AOT Lowering for Module-Level Replacement
  48     gm = pre_aot_substitutions(gm)
  49 
  50     # Invoke AOTAutograd to translate operators to aten
  51     return aot_module_simplified(
  52         gm,
  53         sample_inputs,
  54         fw_compiler=make_boxed_compiler(custom_backend),
  55         decompositions=get_decompositions(),
  56     )
  57 
  58 
  59 def _pretraced_backend(
  60     gm: torch.fx.GraphModule,
  61     sample_inputs: Sequence[torch.Tensor],
  62     settings: CompilationSettings = CompilationSettings(),
  63 ):
  64     """Helper function to manage translation of traced FX module to TRT engines
  65 
  66     Args:
  67         module: FX GraphModule to convert
  68         inputs: Inputs to the module
  69         settings: Compilation settings
  70     Returns:
  71         Compiled FX GraphModule
  72     """
  73     try:
  74         logger.debug("Post-AOT Autograd graph:\n" + str(gm.graph))
  75 
  76         trt_compiled = _compile_module(
  77             gm,
  78             sample_inputs,
  79             settings=settings,
  80         )
  81         return trt_compiled
  82     except:
  83         if not settings.pass_through_build_failures:
  84             logger.warning(
  85                 "TRT conversion failed on the subgraph. See trace above. "
  86                 + "Returning GraphModule forward instead.",
  87                 exc_info=True,
  88             )
  89             return gm.forward
  90         else:
  91             raise AssertionError(
  92                 "Halting compilation on build failure since "
  93                 + "pass_through_build_failures was specified as True. "
  94                 + "To return the default Torch implementation and avoid "
  95                 + "halting compilation on engine build failures, "
  96                 + "specify pass_through_build_failures=False."
  97             )
  98 
  99 
  100 def _compile_module(
  101     gm: torch.fx.GraphModule,
  102     sample_inputs: Sequence[torch.Tensor],
  103     settings: CompilationSettings = CompilationSettings(),
  104 ) -> torch.fx.GraphModule:
  105     """Compile a traced FX module
  106 
  107     Includes: Partitioning + Conversion Phases
  108 
  109     Args:
  110         module: FX GraphModule to convert
  111         inputs: Inputs to the module
  112         settings: Compilation settings
  113     Returns:
  114         Compiled FX GraphModule
  115     """
  116     # Partition module into components that can be TRT-accelerated
  117     partitioned_module = partition(
  118         gm,
  119         verbose=settings.debug,
  120         min_block_size=settings.min_block_size,
  121         torch_executed_ops=settings.torch_executed_ops,
  122     )
  123 
  124     # Iterate over all components that can be accelerated
  125     # Generate the corresponding TRT Module for those
  126     for name, _ in partitioned_module.named_children():
  127         submodule = getattr(partitioned_module, name)
  128 
  129         # Get submodule inputs
  130         submodule_inputs = get_submod_inputs(
  131             partitioned_module, submodule, sample_inputs
  132         )
  133 
  134         # Create TRT Module from submodule
  135         trt_mod = convert_module(
  136             submodule,
  137             submodule_inputs,
  138             settings=settings,
  139         )
  140 
  141         # Replace FX Module with TRT Module
  142         setattr(partitioned_module, name, trt_mod)
  143 
  144     return partitioned_module
  145
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/py/torch_tensorrt/dynamo/backend/backends.py b/py/torch_tensorrt/dynamo/backend/backends.py
  --- a/py/torch_tensorrt/dynamo/backend/backends.py
  +++ b/py/torch_tensorrt/dynamo/backend/backends.py
  @@ -121,6 +121,9 @@
           torch_executed_ops=settings.torch_executed_ops,
       )
   
  +    # Store TRT replicas of Torch subgraphs
  +    trt_modules = {}
  +
       # Iterate over all components that can be accelerated
       # Generate the corresponding TRT Module for those
       for name, _ in partitioned_module.named_children():
  @@ -138,7 +141,10 @@
               settings=settings,
           )
   
  -        # Replace FX Module with TRT Module
  +        trt_modules[name] = trt_mod
  +
  +    # Replace all FX Modules with TRT Modules
  +    for name, trt_mod in trt_modules.items():
           setattr(partitioned_module, name, trt_mod)
   
       return partitioned_module
  
verification_info: {"golden_diff": "diff --git a/py/torch_tensorrt/dynamo/backend/backends.py b/py/torch_tensorrt/dynamo/backend/backends.py\n--- a/py/torch_tensorrt/dynamo/backend/backends.py\n+++ b/py/torch_tensorrt/dynamo/backend/backends.py\n@@ -121,6 +121,9 @@\n         torch_executed_ops=settings.torch_executed_ops,\n     )\n \n+    # Store TRT replicas of Torch subgraphs\n+    trt_modules = {}\n+\n     # Iterate over all components that can be accelerated\n     # Generate the corresponding TRT Module for those\n     for name, _ in partitioned_module.named_children():\n@@ -138,7 +141,10 @@\n             settings=settings,\n         )\n \n-        # Replace FX Module with TRT Module\n+        trt_modules[name] = trt_mod\n+\n+    # Replace all FX Modules with TRT Modules\n+    for name, trt_mod in trt_modules.items():\n         setattr(partitioned_module, name, trt_mod)\n \n     return partitioned_module\n", "issue": "\ud83d\udc1b [Bug] Encountered TRT Error in Dynamo compile path\n##  Bug Description\r\nWhen compiling GPT-2 with Dynamo compile, the following error is encountered:\r\n```python\r\n[07/01/2023-00:07:12] [TRT] [E] 3: [executionContext.cpp::enqueueInternal::795] Error Code 3: API Usage Error (Parameter check failed at: runtime/api/executionContext.cpp::enqueueInternal::795, condition: bindings[x] || nullBindingOK\r\n)\r\n```\r\nAdditionally, this does not seem to cause failures in the Dynamo runtime (not caught by `pass_through_build_failures`), and it appears both with and without the experimental runtime.\r\n\r\n## To Reproduce\r\n```python\r\nmodel = GPT2Model.from_pretrained(\"gpt2\").eval().cuda()\r\ninput_ids      = torch.randint(0, 2, (1, 14), dtype=torch.int32).cuda()\r\nattention_mask = torch.randint(0, 2, (1, 14), dtype=torch.int32).cuda()\r\ntraced = transformers_trace(model, input_names=[\"input_ids\", \"attention_mask\"]).eval().cuda()\r\nfx_trt_model = torch_tensorrt.compile(traced, ir=\"dynamo_compile\", inputs=[input_ids, attention_mask], debug=True, pass_through_build_failures=True, min_block_size=10)\r\n```\r\n\r\n## Expected behavior\r\n\r\nThe model should not encounter TRT errors while compiling.\r\n\r\n## Environment\r\n\r\n - Torch-TensorRT Version (e.g. 1.0.0): 2844630\r\n - PyTorch Version (e.g. 1.0): `2.1.0.dev20230620+cu118`\r\n - TensorRT Version: `8.6.1`\r\n\r\n## Additional context\r\n\r\nThis error no longer appears when we apply the `@fake_tensor_unsupported` flag to the backends, as was removed in #1955. It is unclear whether this is the direct _cause_ of the bug, however.\r\n\r\nAdditionally, this bug seems to appear only sometimes, and not consistently.\n", "before_files": [{"content": "import logging\nfrom typing import Sequence\nimport torch\nfrom functools import partial\nimport torch._dynamo as td\n\nfrom torch_tensorrt.dynamo.backend._settings import CompilationSettings\nfrom torch_tensorrt.dynamo.backend.lowering._decompositions import (\n    get_decompositions,\n)\nfrom torch_tensorrt.dynamo.backend.lowering._pre_aot_lowering import (\n    pre_aot_substitutions,\n)\nfrom torch_tensorrt.dynamo.backend.lowering._partition import (\n    partition,\n    get_submod_inputs,\n)\nfrom torch_tensorrt.dynamo.backend.utils import parse_dynamo_kwargs\nfrom torch_tensorrt.dynamo.backend.conversion import convert_module\n\nfrom torch._functorch.aot_autograd import aot_module_simplified, make_boxed_compiler\n\n\nlogger = logging.getLogger(__name__)\n\n\n@td.register_backend(name=\"torch_tensorrt\")\ndef torch_tensorrt_backend(\n    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor], **kwargs\n):\n    DEFAULT_BACKEND = aot_torch_tensorrt_aten_backend\n\n    return DEFAULT_BACKEND(gm, sample_inputs, **kwargs)\n\n\n@td.register_backend(name=\"aot_torch_tensorrt_aten\")\ndef aot_torch_tensorrt_aten_backend(\n    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor], **kwargs\n):\n    settings = parse_dynamo_kwargs(kwargs)\n\n    custom_backend = partial(\n        _pretraced_backend,\n        settings=settings,\n    )\n\n    # Perform Pre-AOT Lowering for Module-Level Replacement\n    gm = pre_aot_substitutions(gm)\n\n    # Invoke AOTAutograd to translate operators to aten\n    return aot_module_simplified(\n        gm,\n        sample_inputs,\n        fw_compiler=make_boxed_compiler(custom_backend),\n        decompositions=get_decompositions(),\n    )\n\n\ndef _pretraced_backend(\n    gm: torch.fx.GraphModule,\n    sample_inputs: Sequence[torch.Tensor],\n    settings: CompilationSettings = CompilationSettings(),\n):\n    \"\"\"Helper function to manage translation of traced FX module to TRT engines\n\n    Args:\n        module: FX GraphModule to convert\n        inputs: Inputs to the module\n        settings: Compilation settings\n    Returns:\n        Compiled FX GraphModule\n    \"\"\"\n    try:\n        logger.debug(\"Post-AOT Autograd graph:\\n\" + str(gm.graph))\n\n        trt_compiled = _compile_module(\n            gm,\n            sample_inputs,\n            settings=settings,\n        )\n        return trt_compiled\n    except:\n        if not settings.pass_through_build_failures:\n            logger.warning(\n                \"TRT conversion failed on the subgraph. See trace above. \"\n                + \"Returning GraphModule forward instead.\",\n                exc_info=True,\n            )\n            return gm.forward\n        else:\n            raise AssertionError(\n                \"Halting compilation on build failure since \"\n                + \"pass_through_build_failures was specified as True. \"\n                + \"To return the default Torch implementation and avoid \"\n                + \"halting compilation on engine build failures, \"\n                + \"specify pass_through_build_failures=False.\"\n            )\n\n\ndef _compile_module(\n    gm: torch.fx.GraphModule,\n    sample_inputs: Sequence[torch.Tensor],\n    settings: CompilationSettings = CompilationSettings(),\n) -> torch.fx.GraphModule:\n    \"\"\"Compile a traced FX module\n\n    Includes: Partitioning + Conversion Phases\n\n    Args:\n        module: FX GraphModule to convert\n        inputs: Inputs to the module\n        settings: Compilation settings\n    Returns:\n        Compiled FX GraphModule\n    \"\"\"\n    # Partition module into components that can be TRT-accelerated\n    partitioned_module = partition(\n        gm,\n        verbose=settings.debug,\n        min_block_size=settings.min_block_size,\n        torch_executed_ops=settings.torch_executed_ops,\n    )\n\n    # Iterate over all components that can be accelerated\n    # Generate the corresponding TRT Module for those\n    for name, _ in partitioned_module.named_children():\n        submodule = getattr(partitioned_module, name)\n\n        # Get submodule inputs\n        submodule_inputs = get_submod_inputs(\n            partitioned_module, submodule, sample_inputs\n        )\n\n        # Create TRT Module from submodule\n        trt_mod = convert_module(\n            submodule,\n            submodule_inputs,\n            settings=settings,\n        )\n\n        # Replace FX Module with TRT Module\n        setattr(partitioned_module, name, trt_mod)\n\n    return partitioned_module\n", "path": "py/torch_tensorrt/dynamo/backend/backends.py"}], "after_files": [{"content": "import logging\nfrom typing import Sequence\nimport torch\nfrom functools import partial\nimport torch._dynamo as td\n\nfrom torch_tensorrt.dynamo.backend._settings import CompilationSettings\nfrom torch_tensorrt.dynamo.backend.lowering._decompositions import (\n    get_decompositions,\n)\nfrom torch_tensorrt.dynamo.backend.lowering._pre_aot_lowering import (\n    pre_aot_substitutions,\n)\nfrom torch_tensorrt.dynamo.backend.lowering._partition import (\n    partition,\n    get_submod_inputs,\n)\nfrom torch_tensorrt.dynamo.backend.utils import parse_dynamo_kwargs\nfrom torch_tensorrt.dynamo.backend.conversion import convert_module\n\nfrom torch._functorch.aot_autograd import aot_module_simplified, make_boxed_compiler\n\n\nlogger = logging.getLogger(__name__)\n\n\n@td.register_backend(name=\"torch_tensorrt\")\ndef torch_tensorrt_backend(\n    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor], **kwargs\n):\n    DEFAULT_BACKEND = aot_torch_tensorrt_aten_backend\n\n    return DEFAULT_BACKEND(gm, sample_inputs, **kwargs)\n\n\n@td.register_backend(name=\"aot_torch_tensorrt_aten\")\ndef aot_torch_tensorrt_aten_backend(\n    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor], **kwargs\n):\n    settings = parse_dynamo_kwargs(kwargs)\n\n    custom_backend = partial(\n        _pretraced_backend,\n        settings=settings,\n    )\n\n    # Perform Pre-AOT Lowering for Module-Level Replacement\n    gm = pre_aot_substitutions(gm)\n\n    # Invoke AOTAutograd to translate operators to aten\n    return aot_module_simplified(\n        gm,\n        sample_inputs,\n        fw_compiler=make_boxed_compiler(custom_backend),\n        decompositions=get_decompositions(),\n    )\n\n\ndef _pretraced_backend(\n    gm: torch.fx.GraphModule,\n    sample_inputs: Sequence[torch.Tensor],\n    settings: CompilationSettings = CompilationSettings(),\n):\n    \"\"\"Helper function to manage translation of traced FX module to TRT engines\n\n    Args:\n        module: FX GraphModule to convert\n        inputs: Inputs to the module\n        settings: Compilation settings\n    Returns:\n        Compiled FX GraphModule\n    \"\"\"\n    try:\n        logger.debug(\"Post-AOT Autograd graph:\\n\" + str(gm.graph))\n\n        trt_compiled = _compile_module(\n            gm,\n            sample_inputs,\n            settings=settings,\n        )\n        return trt_compiled\n    except:\n        if not settings.pass_through_build_failures:\n            logger.warning(\n                \"TRT conversion failed on the subgraph. See trace above. \"\n                + \"Returning GraphModule forward instead.\",\n                exc_info=True,\n            )\n            return gm.forward\n        else:\n            raise AssertionError(\n                \"Halting compilation on build failure since \"\n                + \"pass_through_build_failures was specified as True. \"\n                + \"To return the default Torch implementation and avoid \"\n                + \"halting compilation on engine build failures, \"\n                + \"specify pass_through_build_failures=False.\"\n            )\n\n\ndef _compile_module(\n    gm: torch.fx.GraphModule,\n    sample_inputs: Sequence[torch.Tensor],\n    settings: CompilationSettings = CompilationSettings(),\n) -> torch.fx.GraphModule:\n    \"\"\"Compile a traced FX module\n\n    Includes: Partitioning + Conversion Phases\n\n    Args:\n        module: FX GraphModule to convert\n        inputs: Inputs to the module\n        settings: Compilation settings\n    Returns:\n        Compiled FX GraphModule\n    \"\"\"\n    # Partition module into components that can be TRT-accelerated\n    partitioned_module = partition(\n        gm,\n        verbose=settings.debug,\n        min_block_size=settings.min_block_size,\n        torch_executed_ops=settings.torch_executed_ops,\n    )\n\n    # Store TRT replicas of Torch subgraphs\n    trt_modules = {}\n\n    # Iterate over all components that can be accelerated\n    # Generate the corresponding TRT Module for those\n    for name, _ in partitioned_module.named_children():\n        submodule = getattr(partitioned_module, name)\n\n        # Get submodule inputs\n        submodule_inputs = get_submod_inputs(\n            partitioned_module, submodule, sample_inputs\n        )\n\n        # Create TRT Module from submodule\n        trt_mod = convert_module(\n            submodule,\n            submodule_inputs,\n            settings=settings,\n        )\n\n        trt_modules[name] = trt_mod\n\n    # Replace all FX Modules with TRT Modules\n    for name, trt_mod in trt_modules.items():\n        setattr(partitioned_module, name, trt_mod)\n\n    return partitioned_module\n", "path": "py/torch_tensorrt/dynamo/backend/backends.py"}]}