Low bit optimizers quality

[tsengalb99]

I saw that the quality numbers in https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim changed recently. Was there a bug in the AO low bit optimizer implementation before?

[msaroufim]

There was a bug with how we handled LR schedulers that make things go faster now #736

[tsengalb99]

So nothing functionally changed and the optimizers are faster -> you were able to test on a larger model?

[msaroufim]

Largest model we tried was on the order of 500M parameters, for 1B+ try it out and feel free share loss curves

[tsengalb99]

Did you test pretraining or just finetuning only? Thanks

[msaroufim]

This was a finetuning benchmark

Benchmark script for fine-tuning a timm model on resisc45 dataset is available at benchmarks/benchmark_low_bit_adam.py

We were hoping to expand to more fine-tuning benchmarks https://github.com/pytorch/torchtune here

And we're doing with fully quantized training here https://github.com/pytorch/ao/tree/main/torchao/prototype/quantized_training

If you're doing research in this space feel free to reach out, the focus on the team now is more on infra and making sure things run fast

[tsengalb99]

I tried using the AO 8 bit optimizer in nanoGPT by swapping out AdamW with Adam8bit here and keeping the embedding and non-decayed parameters (layernorms etc) in a regular torch AdamW instance. This strategy seems to work with the lpmm codebase, where I get close to FP32 performance, but with AO I get severely degraded performance. Is there anything special I need to do with the AO low bit optimizers? Also, I wasn't able to get the AO 4 bit optimizer to work out of the box. I had to disable the compile call here or compile would complain about the input to lerp being negative. I'm using torch 2.4.0 if that matters.

[msaroufim]

Yeah for sure try to use ao and PyTorch nightlies. In the meantime @gau-nernst might make sense to do convergence benchmarks with at least llama 8B

@tsengalb99 do you mind sharing more details on what you're observing too? What size of model? Loss curves?

[gau-nernst]

@tsengalb99 Yes, having more details about your training would be great to debug the issue.

Input to lerp being negative sounds like something is wrong with the training. I think we only use lerp here

ao/torchao/prototype/low_bit_optim/adam.py

Lines 145 to 146 in af68031

	new_exp_avg = exp_avg.lerp(grad, 1 - beta1)
	new_exp_avg_sq = exp_avg_sq.lerp(grad.square(), 1 - beta2)

Would you mind trying bnb AdamW8bit too? Our implementation should match bnb's one exactly. Also, I notice you are using ao Adam8bit. For a fair comparison, shouldn't AdamW8bit be used instead?

@msaroufim Regarding llama-8B, are you thinking pre-training or fine-tuning? Should be a drop-in replacement for torchtune and torchtitan. I think @awgu ran a small test before?

[msaroufim]

I was thinking of prioritizing larger scale finetuning experiments first and pretraining more of a hail mary

[tsengalb99]

Input to lerp being negative sounds like something is wrong with the training. I think we only use lerp here

Yes, that's where the error is being thrown

Adam8bit. For a fair comparison, shouldn't AdamW8bit be used instead?

Ah, I totally forgot to put AdamW and not Adam. That's probably the issue - thanks!

[msaroufim]

@tsengalb99 is this still an issue?

[tsengalb99]

No - it turns out the issue was Adam vs AdamW so not AO related. AO 4 bit matches lpmm on my pretraining experiment. Get Outlook for Android<https://aka.ms/AAb9ysg>

…

________________________________ From: Mark Saroufim ***@***.***> Sent: Monday, August 26, 2024 4:55:58 PM To: pytorch/ao ***@***.***> Cc: Albert Tseng ***@***.***>; Mention ***@***.***> Subject: Re: [pytorch/ao] Low bit optimizers quality (Issue #744) @tsengalb99<https://github.com/tsengalb99> is this still an issue? — Reply to this email directly, view it on GitHub<#744 (comment)>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/AH6WZSE3373IXUQ54ET6ENTZTO6A5AVCNFSM6AAAAABNBYW7ZKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDGMJRGMYDSMJWGQ>. You are receiving this because you were mentioned.Message ID: ***@***.***>