torch.compile()'d optimizer.step() has too many arguments in C++
Labels: triaged, oncall: pt2, module: aotdispatch, module: cpu inductor

### 🐛 Describe the bug

When one torch.compile()'s the optimizer.step() for the following models and optimizers on CPU, inductor generates a C++ kernel that takes in too many arguments for ctypes to handle and hits a "ctypes: too many arguments" error (https://github.com/microsoft/nnfusion/issues/399). The compiled optimizer step kernel unrolls the for loop of updates so that there are `num params` copies of the same code. The more parameters there are, the more likely it is to hit the error. Since it is just unrolling for-loops and the param updates do not interfere, the fix here would be to break up the kernels so that each take fewer arguments to avoid the error.

Models that ran into a 'too many arguments' error and their incompatible optimizers:
```
BERT_pytorch: AdamW, Rprop, Adam, ASGD, Adagrad, Adamax, Adadelta, NAdam
Background_Matting: AdamW, Rprop, Adam, ASGD, Adagrad, Adamax, Adadelta, NAdam
Super_SloMo: AdamW, Adamax, Adam, ASGD, NAdam
attention_is_all_you_need_pytorch: AdamW, Rprop, Adam, ASGD, Adagrad, Adamax, Adadelta, NAdam
densenet121: AdamW, Rprop, Adam, SGD, ASGD, Adagrad, Adamax, RMSprop, Adadelta, NAdam
detectron2_fasterrcnn_r_101_c4: AdamW, Adamax, Adam, ASGD, NAdam
detectron2_fasterrcnn_r_101_dc5: AdamW, Adamax, Adam, ASGD, NAdam
detectron2_fasterrcnn_r_101_fpn: AdamW, Rprop, Adamax, Adam, Adagrad, ASGD, NAdam
detectron2_fasterrcnn_r_50_fpn: Adamax, NAdam, ASGD
detectron2_maskrcnn: Adamax, NAdam, ASGD
detectron2_maskrcnn_r_101_c4: AdamW, Adamax, Adam, ASGD, NAdam
detectron2_maskrcnn_r_101_fpn: AdamW, Rprop, Adamax, Adam, Adagrad, ASGD, NAdam
detectron2_maskrcnn_r_50_fpn: Adamax, NAdam, ASGD
doctr_det_predictor: AdamW, Rprop, Adam, ASGD, Adagrad, Adamax, Adadelta, NAdam
fambench_xlmr: AdamW, Rprop, Adam, SGD, ASGD, Adagrad, Adamax, RMSprop, Adadelta, NAdam
fastNLP_Bert: AdamW, Rprop, Adam, ASGD, Adagrad, Adamax, Adadelta, NAdam
hf_Bart: AdamW, Rprop, Adam, ASGD, Adagrad, Adamax, RMSprop, Adadelta, NAdam
hf_Bert: AdamW, Rprop, Adam, ASGD, Adagrad, Adamax, Adadelta, NAdam
hf_Bert_large: AdamW, Rprop, Adam, SGD, ASGD, Adagrad, Adamax, RMSprop, Adadelta, NAdam
hf_BigBird: AdamW, Rprop, Adam, ASGD, Adagrad, Adamax, Adadelta, NAdam
hf_DistilBert: AdamW, Adamax, Adam, ASGD, NAdam
hf_GPT2: AdamW, Rprop, Adam, ASGD, Adagrad, Adamax, Adadelta, NAdam
hf_GPT2_large: AdamW, Rprop, Adam, SGD, ASGD, Adagrad, Adamax, RMSprop, Adadelta, NAdam
hf_Longformer: AdamW, Rprop, Adam, ASGD, Adagrad, Adamax, RMSprop, Adadelta, NAdam
hf_Reformer: Adamax, NAdam
hf_T5: AdamW, Rprop, Adamax, Adam, Adagrad, ASGD, NAdam
hf_T5_base: AdamW, Rprop, Adam, ASGD, Adagrad, Adamax, RMSprop, Adadelta, NAdam
hf_T5_large: AdamW, Rprop, Adam, SGD, ASGD, Adagrad, Adamax, RMSprop, Adadelta, NAdam
mnasnet1_0: Rprop, Adadelta
mobilenet_v2: Rprop, Adadelta
mobilenet_v2_quantized_qat: Rprop, Adadelta
mobilenet_v3_large: Rprop, Adadelta
phlippe_densenet: Rprop, Adadelta
resnet152: SGD, Rprop, RMSprop, Adadelta
resnet50: Rprop, Adadelta
resnet50_quantized_qat: Rprop, Adadelta
resnext50_32x4d: Rprop, Adadelta
shufflenet_v2_x1_0: Rprop, Adadelta
timm_efficientnet: Rprop, RMSprop, Adadelta
timm_nfnet: RMSprop, Adadelta
timm_regnet: Rprop, RMSprop, Adadelta
timm_vision_transformer: Rprop, Adadelta
```

### Error logs

Short error is: `ctypes.ArgumentError: too many arguments (1940), maximum is 1024` for running torch.compile on BERT_pytorch's parameters on CPU with default AdamW.

<details>

```
Traceback (most recent call last):
  File "/fsx/users/janeyx/conda/envs/torchbenchmark/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/fsx/users/janeyx/conda/envs/torchbenchmark/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/scratch/janeyx/work/benchmark/userbenchmark/optim/__init__.py", line 414, in <module>
    run(sys.argv[1:])
  File "/scratch/janeyx/work/benchmark/userbenchmark/optim/__init__.py", line 405, in run
    results = run_benchmarks(args.optims, args.funcs, args.models, args.devices)
  File "/scratch/janeyx/work/benchmark/userbenchmark/optim/__init__.py", line 344, in run_benchmarks
    bm = run_model(mn, d, O, defaults, func_str)
  File "/scratch/janeyx/work/benchmark/userbenchmark/optim/__init__.py", line 320, in run_model
    raise e
  File "/scratch/janeyx/work/benchmark/userbenchmark/optim/__init__.py", line 317, in run_model
    ).blocked_autorange()
  File "/fsx/users/janeyx/conda/envs/torchbenchmark/lib/python3.10/site-packages/torch/utils/benchmark/utils/timer.py", line 394, in blocked_autorange
    number = self._estimate_block_size(min_run_time)
  File "/fsx/users/janeyx/conda/envs/torchbenchmark/lib/python3.10/site-packages/torch/utils/benchmark/utils/timer.py", line 311, in _estimate_block_size
    time_taken = self._timeit(number)
  File "/fsx/users/janeyx/conda/envs/torchbenchmark/lib/python3.10/site-packages/torch/utils/benchmark/utils/timer.py", line 256, in _timeit
    return max(self._timer.timeit(number), 1e-9)
  File "/fsx/users/janeyx/conda/envs/torchbenchmark/lib/python3.10/timeit.py", line 178, in timeit
    timing = self.inner(it, self.timer)
  File "<timeit-src>", line 6, in inner
  File "/scratch/janeyx/work/benchmark/userbenchmark/optim/__init__.py", line 268, in pt2_optimizer_step
    f()
  File "/fsx/users/janeyx/conda/envs/torchbenchmark/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 235, in _fn
    return fn(*args, **kwargs)
  File "/scratch/janeyx/work/benchmark/userbenchmark/optim/__init__.py", line 267, in f
    optimizer.step()
  File "/fsx/users/janeyx/conda/envs/torchbenchmark/lib/python3.10/site-packages/torch/optim/optimizer.py", line 33, in _use_grad
    ret = func(self, *args, **kwargs)
  File "/fsx/users/janeyx/conda/envs/torchbenchmark/lib/python3.10/site-packages/torch/optim/adamw.py", line 143, in step
    self._cuda_graph_capture_health_check()
  File "/fsx/users/janeyx/conda/envs/torchbenchmark/lib/python3.10/site-packages/torch/optim/adamw.py", line 171, in <resume in step>
    adamw(
  File "/fsx/users/janeyx/conda/envs/torchbenchmark/lib/python3.10/site-packages/torch/optim/adamw.py", line 264, in adamw
    def adamw(
  File "/fsx/users/janeyx/conda/envs/torchbenchmark/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 235, in _fn
    return fn(*args, **kwargs)
  File "/fsx/users/janeyx/conda/envs/torchbenchmark/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 3001, in forward
    return compiled_fn(full_args)
  File "/fsx/users/janeyx/conda/envs/torchbenchmark/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 1175, in g
    return f(*args)
  File "/fsx/users/janeyx/conda/envs/torchbenchmark/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 2056, in runtime_wrapper
    all_outs = call_func_with_args(
  File "/fsx/users/janeyx/conda/envs/torchbenchmark/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 1200, in call_func_with_args
    out = normalize_as_list(f(args))
  File "/tmp/torchinductor_janeyx/ch/cchgxveesattdyuylo65f32dgd33oz4dmmlxaimi3i6xyv2gbyh5.py", line 21373, in call
    kernel_cpp_0(c_void_p(arg0_1.data_ptr()), c_void_p(arg388_1.data_ptr()), c_void_p(arg194_1.data_ptr()), c_void_p(arg582_1.data_ptr()), c_void_p(arg776_1.data_ptr()), c_void_p(arg1_1.data_ptr()), c_void_p(arg389_1.data_ptr()), ...c_void_p(arg967_1.data_ptr()), c_void_p(arg968_1.data_ptr()), c_void_p(arg969_1.data_ptr()))
ctypes.ArgumentError: too many arguments (1940), maximum is 1024
```

</details>

### Minified repro

Minimal repro with BERT_pytorch's param shapes:

```
def pt2_optimizer_step(optimizer):
    @torch.compile()
    def f():
        optimizer.step()
    f()

params = [torch.rand(20005, 768, dtype=torch.float32, device='cpu') for _ in range(194)]
for p in params:
    p.grad = torch.rand_like(p)

o = torch.optim.AdamW(params)
pt2_optimizer_step(o)
```

Code reduced from https://github.com/pytorch/benchmark/pull/1473, look there if you want to pull params from models.

### Versions

2.0, nightlies

cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh