pytorch - 💡(How to fix) Fix Dynamo is leaking internal state during optimizer tracing [1 comments, 2 participants]

import torch
from copy import deepcopy

m = torch.nn.Linear(2, 2)
opt = torch.optim.AdamW(m.parameters(), lr=0.01)
m(torch.randn(2, 2)).sum().backward()
opt.step()

ORIGINAL_LR = opt.param_groups[0]["lr"]


def fn():
    # 1. Mutate: registers the param_group VT for SideEffects flush.
    opt.param_groups[0]["lr"] = 5.0
    # 2. Graph break: dynamo flushes tracked VT mutations onto the live dict.
    sd = deepcopy(opt.state_dict())
    # 3. Undo: restore original. Visible to eager Python from this point on.
    opt.param_groups[0]["lr"] = ORIGINAL_LR
    return sd


print(f"Before: capturable = {[g['capturable'] for g in opt.param_groups]}")
print(f"        lr         = {[g['lr'] for g in opt.param_groups]}")
torch._dynamo.optimize("eager_noexcept", nopython=False)(fn)()
print(f"After:  capturable = {[g['capturable'] for g in opt.param_groups]}")
print(f"        lr         = {[g['lr'] for g in opt.param_groups]}")

---

Before: capturable = [False]
        lr         = [0.01]
After:  capturable = [True]
        lr         = [0.01]