bnb.optim.AdamW performance differs from torch.optim.AdamW despite being called with the same hyperparameters

[inkitori]

Hi, I noticed that when I swap out torch.optim.AdamW with bnb.optim.AdamW (both in 32 bit precision) and I fine tune a model loaded in bf16, the overall performance of the model trained using bnb's AdamW is higher than with Torch's. Furthermore, if I fine tune a model loaded in fp16, torch's AdamW leads to NaN values after the very first optim.step(), while bnb's AdamW trains perfectly fine without any NaNs. Is there a reason for this performance discrepancy, especially in the latter case? Does bnb internally use a different dtype than torch, or something else? Thank you

[matthewdouglas]

Interesting, thanks for asking about this! For the most part, we're using fp32 for all of the internal computations; we also cast the parameter value to fp32 when updating. Can you share what version of PyTorch you're using, and any example code to reproduce?

[inkitori]

I'm currently on torch==2.8.0.

from transformers import Qwen2_5OmniProcessor, Qwen2_5OmniThinkerForConditionalGeneration
model = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(model_id, torch_dtype="float16")

optim = torch.optim.AdamW(model.parameters(), lr=training_config['learning_rate'])
# optim = bnb.optim.AdamW(model.parameters(), lr=training_config['learning_rate'])

The rest of my code is just a regular torch training loop (no trainer/trl). I calculate my loss by passing labels directly into model.forward. Should I include more details?