Add option to disable applying side effects in dynamo (pytorch#167239)

There are two motivating use cases for this change: 1) export (when we trace pytree calls into a graph, we don't want to accidentally trace the side effect bytecode which will pollute the initial state) -> We want to warn about side effects and don't want to actually apply them 2) VLLM -> They want to detect side effects and error out. We implement this with two configs where one config controls whether we want to apply side effects (by default yes) and the warning level for side effects (warning for export and error for VLLM). We intentionally ignore input side effects, because they are captured in the graph and export would never trace the actual dynamo graph module when tracing the pytree calls). Pull Request resolved: pytorch#167239 Approved by: https://github.com/williamwen42, https://github.com/anijain2305

Author: tugsbayasgalan

test/dynamo/test_misc.py

@@ -19,6 +19,7 @@
 import os
 import pickle
 import random
+import re
 import sys
 import tempfile
 import threading
@@ -5635,6 +5636,115 @@ def f2(a, b):
  self.assertTrue(same(res11, res12))
  self.assertTrue(same(res21, res22))
 
+ def test_replay_side_effects_config(self):
+ # Test that replay_side_effects config controls mutation replay
+ def fn(x, lst):
+ lst.append(x + 1)
+ return x * 2
+
+ x = torch.tensor([5.0])
+
+ # Test with replay enabled (default)
+ lst_with_replay = []
+ opt_fn_with_replay = torch.compile(fn, backend="eager")
+ result1 = opt_fn_with_replay(x, lst_with_replay)
+ self.assertEqual(len(lst_with_replay), 1) # Mutation should be replayed
+ self.assertTrue(same(result1, x * 2))
+
+ torch._dynamo.reset()
+
+ # Test with replay disabled
+ lst_without_replay = []
+ with torch._dynamo.config.patch(
+ replay_side_effects=False, side_effect_replay_policy="warn"
+ ):
+ opt_fn_without_replay = torch.compile(fn, backend="eager")
+ result2 = opt_fn_without_replay(x, lst_without_replay)
+ self.assertEqual(
+ len(lst_without_replay), 0
+ ) # Mutation should NOT be replayed
+ self.assertTrue(same(result2, x * 2))
+
+ torch._dynamo.reset()
+ lst_without_replay = []
+ with torch._dynamo.config.patch(
+ replay_side_effects=False, side_effect_replay_policy="error"
+ ):
+ opt_fn_without_replay = torch.compile(fn, backend="eager")
+ with self.assertRaisesRegex(
+ RuntimeError,
+ re.escape(
+ "While compiling, we found certain side effects happened in the model.forward. Here are the list of potential sources you can double check: [\"L['lst']\"]"
+ ),
+ ):
+ _ = opt_fn_without_replay(x, lst_without_replay)
+
+ def test_replay_side_effects_model_attr(self):
+ class Bar(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.const = 4
+
+ def forward(self, x):
+ return x.cos()
+
+ class Foo(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.const = 4
+ self.tensor = None
+ self.bar = Bar()
+
+ def forward(self, x):
+ self.const = 5
+ self.tensor = x.sin()
+ res = self.bar(x)
+ return x.cos() + res.sum() + self.tensor
+
+ with torch._dynamo.config.patch(
+ replay_side_effects=False, side_effect_replay_policy="error"
+ ):
+ foo = Foo()
+ with self.assertRaisesRegex(
+ RuntimeError,
+ re.escape(
+ "While compiling, we found certain side effects happened in the model.forward. Here are the list of potential sources you can double check: [\"L['self']\"]"
+ ),
+ ):
+ torch.compile(foo, fullgraph=True)(torch.randn(4, 4))
+
+ with torch._dynamo.config.patch(
+ replay_side_effects=False, side_effect_replay_policy="silent"
+ ):
+ foo_v2_compile = Foo()
+ foo_v2_eager = Foo()
+ inp = torch.randn(4, 4)
+ res = torch.compile(foo_v2_compile, fullgraph=True)(torch.randn(4, 4))
+ self.assertEqual(foo_v2_compile.tensor, None)
+ self.assertEqual(foo_v2_compile.const, 4)
+ self.assertEqual(foo_v2_compile.bar.const, 4)
+ same(res, foo_v2_eager(inp))
+
+ def test_replay_side_effects_input_mut(self):
+ class Foo(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.const = 4
+ self.tensor = None
+
+ def forward(self, x):
+ x.add_(5)
+ return x.cos()
+
+ # This is ok because we actually capture the graph which
+ # has mutation. In export, we never retrace the actual
+ # gm so we won't see any mutation applied to inputs
+ with torch._dynamo.config.patch(
+ replay_side_effects=False, side_effect_replay_policy="error"
+ ):
+ foo = Foo()
+ torch.compile(foo, fullgraph=True)(torch.randn(4, 4))
+
  def test_list_append_return_none(self):
  def fn(x):
  alist = []

test/export/test_experimental.py

@@ -349,6 +349,18 @@ def generate(self, *, input_tensor, input_tensor2):
  res2 = p.generate(input_tensor=inp, input_tensor2=inp2)
  self.assertTrue(torch.allclose(res, res2))
 
+ def test_side_effect(self):
+ global_env = []
+
+ class Foo(torch.nn.Module):
+ def forward(self, x):
+ global_env.append(x)
+ return x.sin()
+
+ with torch._dynamo.config.patch(replay_side_effects=False):
+ _ = dynamo_graph_capture_for_export(Foo())(torch.randn(4, 4))
+ self.assertEqual(len(global_env), 0)
+
  def test_export_add_in_out_info(self):
  class Foo(torch.nn.Module):
  def forward(self, dct, lst, bleh):

torch/_dynamo/config.py

@@ -44,6 +44,19 @@
 # turn on/off DCE pass (deprecated: always true)
 dead_code_elimination = True
 
+# Enable or disable side effect replay after graph execution.
+# When False, mutations to Python objects (lists, dicts, attributes) won't be
+# replayed after the compiled graph runs. This can cause correctness issues
+# if your code depends on these mutations being visible. This should probably
+# never be False by default. At the moment, only export will need it.
+replay_side_effects = True
+
+# Configure side effect warning level
+# If `silent`, we silently allow side effects
+# If `warn`, we warn side effects
+# If `error`, we error on side effects
+side_effect_replay_policy = "silent"
+
 # disable (for a function) when cache reaches this size
 
 # controls the maximum number of cache entries with a guard on same ID_MATCH'd

torch/_dynamo/output_graph.py

@@ -1845,7 +1845,7 @@ def compile_subgraph(
  [create_instruction("DELETE_FAST", argval=graph_output_var)]
  )
 
- if self.export:
+ if torch._dynamo.config.side_effect_replay_policy in ["warn", "error"]:
  from torch.export._trace import _ExportModuleSpecTrackerDict
 
  potential_side_effects = []
@@ -1881,10 +1881,16 @@ def compile_subgraph(
  ]
 
  if side_effect_refs:
- warnings.warn(
- f"While exporting, we found certain side effects happened in the model.forward. "
- f"Here are the list of potential sources you can double check: {side_effect_refs}"
- )
+ if torch._dynamo.config.side_effect_replay_policy == "warn":
+ warnings.warn(
+ f"While compiling, we found certain side effects happened in the model.forward. "
+ f"Here are the list of potential sources you can double check: {side_effect_refs}"
+ )
+ else:
+ raise RuntimeError(
+ f"While compiling, we found certain side effects happened in the model.forward. "
+ f"Here are the list of potential sources you can double check: {side_effect_refs}"
+ )
 
  return all_stack_locals_metas
 
@@ -1930,7 +1936,8 @@ def codegen_suffix(
  assert self.backward_state_var is not None
  cg.append_output(cg.create_load(self.backward_state_var))
  cg.store_attr(name)
- self.side_effects.codegen_hooks(cg)
+ if config.replay_side_effects:
+ self.side_effects.codegen_hooks(cg)
 
  # TODO get debug_locals working for nested graph breaks
  # Return variables used for logging at the end
@@ -1945,7 +1952,8 @@ def codegen_suffix(
  self.codegen_cells(tx, cg)
 
  cg.restore_stack(stack_values, value_from_source=not tx.export)
- self.side_effects.codegen_update_mutated(cg)
+ if config.replay_side_effects:
+ self.side_effects.codegen_update_mutated(cg)
 
  def cleanup_graph(self) -> None:
  """

torch/_dynamo/variables/builder.py

@@ -359,6 +359,13 @@ class GraphArg:
  # stash a strong reference too.
  example_strong_ref: Optional[torch.Tensor] = None
 
+ def __setattr__(self, name, value):
+ # Use object.__setattr__ to bypass Dynamo's STORE_ATTR interception.
+ # This is needed because when PYTORCH_TEST_WITH_DYNAMO=1, even internal
+ # GraphArg creation can be traced, and with replay_side_effects=False,
+ # normal STORE_ATTR bytecode only records mutations without applying them.
+ object.__setattr__(self, name, value)
+
  @property
  def example(self):
  if isinstance(self._example, TensorWeakRef):

torch/export/_trace.py

@@ -140,6 +140,8 @@ class ExportDynamoConfig:
  capture_dynamic_output_shape_ops: bool = True
  capture_scalar_outputs: bool = True
  prefer_deferred_runtime_asserts_over_guards: bool = False
+ replay_side_effects: bool = False
+ side_effect_replay_policy: str = "warn"
 
 
 @dataclasses.dataclass