Dynamo hangs/takes long time for optimizers

[anijain2305]

Repro

1. Get this PR - [reland][dynamo] use optimizers correctly in benchmarking pytorch#87492
2. Run this command -
   python benchmarks/dynamo/huggingface.py --training --accuracy --device cuda --backend=eager --float32 --only=MobileBertForMaskedLM

W/o optimizer, the above command takes few seconds to run.

[davidberard98]

on hf_T5_large (run these in the torchbench dir)

by default, these models run with the optimizer included.

python run.py hf_T5_large -t train -d cuda --torchdynamo inductor

^ this OOMs

python run.py hf_T5_large -t train -d cuda --torchdynamo inductor --torchinductor_cudagraph False

^ this takes ~30 minutes

... to remove the optimizer, go to torchbenchmark/util/framework/huggingface/model_factory.py and comment out self.optimizer.step(). Then re-run either of the above commands; this takes ~ a few minutes.

[davidberard98]

Using pytorch/pytorch#87492 and Animesh's repro instructions, I see in the debug logs ~2200 guards from param_groups[]

dumped logs in https://gist.github.com/davidberard98/d7f49b110c1bb796208a9ad6e2cabb60

[mlazos]

So I investigated more, it looks like Dynamo tracing just takes a really long time when there are 1000s of parameters (one potential area for speedup I noticed is that we iterate over the whole graph every time we add a tensor variable).

I have a more minimal repro here:

import torch
import torch._dynamo as torchdynamo
import torch._inductor

model = torch.nn.Sequential(*[torch.nn.Linear(10, 10, True, device="cuda:0") for _ in range(1000)])
optimizer = torch.optim.Adam(model.parameters(), capturable=True)
x = torch.randn(10, 10, dtype=torch.float, device="cuda:0")
y = model(x)
torch.sum(y).backward()

@torchdynamo.optimize("eager")
def repro():
    optimizer.step()


repro()

After discussing with Yanbo, one potential solution I see is to avoid tracing the large loops over the parameter list in the optimizer by extracting them into functions and disabling dynamo on them. Then use the *for_each ops which will transfer the loop into the FX graph directly and avoid tracing each parameter individually. I'm trying this out now