Repository: pytorch
Issue: [TestOptimRenewed] test_set_default_dtype_works_with_foreach errors with dynamo 
Labels: module: optimizer,triaged,oncall: pt2,module: dynamo,dynamo-must-fix

### 🐛 Describe the bug

Running the test with `PYTORCH_TEST_WITH_DYNAMO=1 python test/test_optim.py -k test_set_default_dtype_works_with_foreach` will fail the first test it runs. It doesn't matter which test, they will all individually fail (e.g.f through running `PYTORCH_TEST_WITH_DYNAMO=1 python test/test_optim.py -k test_set_default_dtype_works_with_foreach_AdamW`. If multiple tests are run together, only the first will fail.

Additional observations:
* This test calls into the same helper (`_test_derived_optimizers`) as `test_foreach_matches_forloop`, EXCEPT it sets the default device before and after the call.
* There is a graph break within a for loop in `_test_derived_optimizers`, which should be okay.
* Thus, I believe this likely has to do with https://github.com/pytorch/pytorch/issues/112787, which was fixed

### Error logs

```
(pytorch-3.10) [janeyx@devgpu023.odn1 ~/local/pytorch (ec48f6d8|REBASE-i|detached HEAD)]$  PYTORCH_TEST_WITH_DYNAMO=1 python test/test_optim.py -k test_set_default_dtype_works_with_foreach_AdamW
/home/janeyx/.conda/envs/pytorch-3.10/lib/python3.10/site-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.17.3 and <1.25.0 is required for this version of SciPy (detected version 1.26.0
  warnings.warn(f"A NumPy version >={np_minversion} and <{np_maxversion}"
sF
======================================================================
FAIL: test_set_default_dtype_works_with_foreach_AdamW_cuda_float64 (__main__.TestOptimRenewedCUDA)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/data/users/janeyx/pytorch/torch/_dynamo/symbolic_convert.py", line 463, in wrapper
    return inner_fn(self, inst)
  File "/data/users/janeyx/pytorch/torch/_dynamo/symbolic_convert.py", line 1227, in CALL_FUNCTION_KW
    self.call_function(fn, args, kwargs)
  File "/data/users/janeyx/pytorch/torch/_dynamo/symbolic_convert.py", line 645, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/data/users/janeyx/pytorch/torch/_dynamo/variables/functions.py", line 298, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/janeyx/pytorch/torch/_dynamo/variables/functions.py", line 252, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/janeyx/pytorch/torch/_dynamo/variables/functions.py", line 81, in call_function
    return tx.inline_user_function_return(
  File "/data/users/janeyx/pytorch/torch/_dynamo/symbolic_convert.py", line 651, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/data/users/janeyx/pytorch/torch/_dynamo/symbolic_convert.py", line 2231, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/data/users/janeyx/pytorch/torch/_dynamo/symbolic_convert.py", line 2346, in inline_call_
    tracer.run()
  File "/data/users/janeyx/pytorch/torch/_dynamo/symbolic_convert.py", line 781, in run
    and self.step()
  File "/data/users/janeyx/pytorch/torch/_dynamo/symbolic_convert.py", line 744, in step
    getattr(self, inst.opname)(inst)
  File "/data/users/janeyx/pytorch/torch/_dynamo/symbolic_convert.py", line 463, in wrapper
    return inner_fn(self, inst)
  File "/data/users/janeyx/pytorch/torch/_dynamo/symbolic_convert.py", line 1176, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/data/users/janeyx/pytorch/torch/_dynamo/symbolic_convert.py", line 645, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/data/users/janeyx/pytorch/torch/_dynamo/variables/builtin.py", line 648, in call_function
    result = handler(tx, *args, **kwargs)
  File "/data/users/janeyx/pytorch/torch/_dynamo/variables/builtin.py", line 1479, in call_deepcopy
    unimplemented(f"copy.deepcopy {repr(x)}")
  File "/data/users/janeyx/pytorch/torch/_dynamo/exc.py", line 193, in unimplemented
    raise Unsupported(msg)
torch._dynamo.exc.Unsupported: copy.deepcopy ConstDictVariable()

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/data/users/janeyx/pytorch/torch/_dynamo/convert_frame.py", line 150, in _fn
    return fn(*args, **kwargs)
  File "/data/users/janeyx/pytorch/torch/_dynamo/convert_frame.py", line 484, in transform
    tracer.run()
  File "/data/users/janeyx/pytorch/torch/_dynamo/symbolic_convert.py", line 2096, in run
    super().run()
  File "/data/users/janeyx/pytorch/torch/_dynamo/symbolic_convert.py", line 781, in run
    and self.step()
  File "/data/users/janeyx/pytorch/torch/_dynamo/symbolic_convert.py", line 744, in step
    getattr(self, inst.opname)(inst)
  File "/data/users/janeyx/pytorch/torch/_dynamo/symbolic_convert.py", line 471, in wrapper
    raise exc.SkipFrame(msg) from excp
torch._dynamo.exc.SkipFrame: Skipping frame because there is a graph break in a for/while loop
<FrameSummary file /data/users/janeyx/pytorch/test/test_optim.py, line 133 in test_set_default_dtype_works_with_foreach>

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/data/users/janeyx/pytorch/torch/testing/_internal/common_utils.py", line 2674, in wrapper
    method(*args, **kwargs)
  File "/data/users/janeyx/pytorch/torch/testing/_internal/common_utils.py", line 2674, in wrapper
    method(*args, **kwargs)
  File "/data/users/janeyx/pytorch/torch/testing/_internal/common_device_type.py", line 416, in instantiated_test
    result = test(self, **param_kwargs)
  File "/data/users/janeyx/pytorch/torch/testing/_internal/common_optimizers.py", line 169, in test_wrapper
    return test(*args, **kwargs)
  File "/data/users/janeyx/pytorch/torch/testing/_internal/common_device_type.py", line 1110, in only_fn
    return fn(slf, *args, **kwargs)
  File "/data/users/janeyx/pytorch/torch/_dynamo/eval_frame.py", line 597, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state)
  File "/data/users/janeyx/pytorch/torch/_dynamo/convert_frame.py", line 736, in _convert_frame
    result = inner_convert(frame, cache_entry, hooks, frame_state)
  File "/data/users/janeyx/pytorch/torch/_dynamo/convert_frame.py", line 381, in _convert_frame_assert
    return _compile(
  File "/data/users/janeyx/pytorch/torch/_dynamo/convert_frame.py", line 638, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/data/users/janeyx/pytorch/torch/_dynamo/utils.py", line 244, in time_wrapper
    r = func(*args, **kwargs)
  File "/data/users/janeyx/pytorch/torch/_dynamo/convert_frame.py", line 519, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/data/users/janeyx/pytorch/torch/_dynamo/bytecode_transformation.py", line 1033, in transform_code_object
    transformations(instructions, code_options)
  File "/data/users/janeyx/pytorch/torch/_dynamo/convert_frame.py", line 163, in _fn
    assert (
AssertionError: Global state changed while dynamo tracing, please report a bug


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True


To execute this test, run the following from the base repo dir:
    PYTORCH_TEST_WITH_DYNAMO=1 python test/test_optim.py -k test_set_default_dtype_works_with_foreach_AdamW_cuda_float64

This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0

----------------------------------------------------------------------
Ran 2 tests in 18.870s

FAILED (failures=1, skipped=1)
```

### Minified repro

`PYTORCH_TEST_WITH_DYNAMO=1 python test/test_optim.py -k test_set_default_dtype_works_with_foreach_AdamW`

### Versions

on main

cc @vincentqb @jbschlosser @albanD @crcrpar @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng