optimizer.load_state_dict not working

[saurabh-kataria]

I am unable to debug this. Torchao seems to have incompatibility with huggingface accelerate maybe or something else

My optimizer is from torchao.prototype.low_bit_optim import Adam4bit

line 173, in load_checkpoint
[rank3]: optimizer.load_state_dict(optimizer_state_dict)
[rank3]: File "/home/skatar6/.local/lib/python3.9/site-packages/torch/_compile.py", line 24, in inner
[rank3]: return torch._dynamo.disable(fn, recursive)(*args, **kwargs)
[rank3]: File "/home/skatar6/.local/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 451, in _fn
[rank3]: return fn(*args, **kwargs)
[rank3]: File "/home/skatar6/.local/lib/python3.9/site-packages/torch/optim/optimizer.py", line 777, in load_state_dict
[rank3]: state[param] = _cast(param, v, param_id=k, param_groups=state_dict['param_groups'])
[rank3]: File "/home/skatar6/.local/lib/python3.9/site-packages/torch/optim/optimizer.py", line 764, in _cast
[rank3]: return {k: _cast(param, v, param_id=param_id, param_groups=param_groups, key=k) for k, v in value.items()}
[rank3]: File "/home/skatar6/.local/lib/python3.9/site-packages/torch/optim/optimizer.py", line 764, in
[rank3]: return {k: _cast(param, v, param_id=param_id, param_groups=param_groups, key=k) for k, v in value.items()}
[rank3]: File "/home/skatar6/.local/lib/python3.9/site-packages/torch/optim/optimizer.py", line 762, in _cast
[rank3]: return Optimizer.process_value_according_to_param_policy(param, value, param_id, param_groups, key)
[rank3]: File "/home/skatar6/.local/lib/python3.9/site-packages/torch/optim/optimizer.py", line 644, in process_value_according_to_param_policy
[rank3]: return value.to(dtype=param.dtype, device=param.device)
[rank3]: File "/home/skatar6/anaconda3/envs/tmp4/lib/python3.9/site-packages/torchao/utils.py", line 377, in dispatch__torch_function
[rank3]: return func(*args, **kwargs)
[rank3]: File "/home/skatar6/anaconda3/envs/tmp4/lib/python3.9/site-packages/torchao/utils.py", line 392, in dispatch__torch_dispatch
[rank3]: kwarg_types = {k: type(arg) for k, arg in kwargs}
[rank3]: File "/home/skatar6/anaconda3/envs/tmp4/lib/python3.9/site-packages/torchao/utils.py", line 392, in
[rank3]: kwarg_types = {k: type(arg) for k, arg in kwargs}
[rank3]: ValueError: too many values to unpack (expected 2)

[supriyar]

cc @gau-nernst @msaroufim

[gau-nernst]

@saurabh-kataria Can you try with torchao nightly? I think it should be fixed by #1021

Btw the cryptic error is caused by #963 (comment) cc @jerryzh168

[saurabh-kataria]

It works now with my setup. Thanks a lot!

[saurabh-kataria]

FWIW the training stability of bnb optimizer (from bitsandbytes.optim import Adam8bit) is far better than torchao (from torchao.prototype.low_bit_optim import Adam4bit, Adam8bit) on my setup. Although I haven't evaluated the performance.

[gau-nernst]

It doesn't look like training instability (no loss spikes/noise), but rather a very huge difference in optimization trajectories. Do you have a small reproduction that I can run? I would like to look into this, since it looks very serious.

Also, which bnb and torchao version are you using? Latest bnb 0.44 changes default block size from 2048 to 256, while torchao 0.5 still uses the old default of 2048. You can override torchao block size by passing block_size=256 to the optimizer. For torchao 0.6, I have changed the default value to 256 to align with bnb. But I don't think this would account for this divergence.

[saurabh-kataria]

It is hard for me to share code right now. It is a generative time series GPT model.

I can share some other observations. The problem almost disappears when Adam8bit uses default optimizer hyper parameters (beta1=0.9, beta2=0.999, weight_decay=0) over my choices (beta1=0.9, beta2=0.95, weight_decay=0.01). And when I use AdamW8it with my hyper parameters, it works well and close to bnb Adam8bit. So I think for now, I will prefer AdamW versions from torchao. Regular Adam for now is not robust to hyper parameter choice.

[gau-nernst]

Adam != AdamW. @matthewdouglas Do you know if Adam from bnb is actually AdamW? I have a feeling it is (i.e. there is no true Adam in bnb)

These 2 look exactly the same

• Adam: https://github.com/bitsandbytes-foundation/bitsandbytes/blob/723e2162fd514deee66268106775f4bfbee25794/bitsandbytes/optim/adam.py#L77-L135
• AdamW: https://github.com/bitsandbytes-foundation/bitsandbytes/blob/723e2162fd514deee66268106775f4bfbee25794/bitsandbytes/optim/adamw.py#L69-L127

@saurabh-kataria Also, it seems like you use different batch size for bnb 8-bit Adam(W?) and torchao 8-bit AdamW (8192 vs 4096). I would expect the loss curves to match exactly if you use the same setup for both optimizers.

[matthewdouglas]

@gau-nernst You're right, they're both effectively AdamW. The default hyperparam for weight decay is 0 in the bnb Adam, but the implementation is the same. I can see how this creates confusion (along with how AMSgrad isn't actually implemented either). We can log an issue on our end for this but I'm not sure how soon we can get to it.

[saurabh-kataria]

I added two experiments to confirm that

1. torchao Adam8bit with my hyper parameters is not converging. Actually, even in theory Adam with weight decay is problematic. I use same batch size as bnb Adam8bit (which is actually AdamW as we realize) now as 8192. In theory, weight decay should not be supported with Adam but I guess user demands them, uses them, and it is also present in other libraries.
2. With torchao AdamW8bit, I get same learning curve a s bnb Adam8bit. (with my hyper params, it becomes bnb AdamW8bit). As a high-level user, my problem is solved I think. I'll just use torchao AdamW8bit or bnb Adam8bit and avoid torchao Adam8bit with weight decay for now in my setup.

[saurabh-kataria]

P.S. I found out torch Adam is also unstable like torchao Adam8bit on my setup so the issue is just with my choice of hyper parameters and other experiment specific stuff. No issues with torchao. Sorry for the confusion and thanks for your inputs.

[nighting0le01]

@saurabh-kataria were you able to use these low-bit optims with FSDP2? Dtensor