Repository navigation
Bug when using optimizer LAMB 32bits #1350
Description
System Info
wsl Ubuntu22.04, Python3.10, bnb 0.43.1
Reproduction
In order to reproduce the issue you can do this :
import bitsandbytes as bnb
import torch
import torch.nn as nn
model = nn.Linear(10,2).cuda()
model.train()
# We create an optimizer
optimizer = bnb.optim.LAMB(model.parameters())
# We create dummy input / output
input = torch.rand(size=(10,10)).cuda()
target = torch.zeros(10).cuda()
# We compute prediction / loss
optimizer.zero_grad()
prediction = model(input)
loss =nn.CrossEntropyLoss()(prediction, target.long())
loss.backward()
optimizer.step()It will result in something like that :
---------------------------------------------------------------------------
KeyError Traceback (most recent call last)
Cell In[12], [line 19](vscode-notebook-cell:?execution_count=12&line=19)
[16](vscode-notebook-cell:?execution_count=12&line=16) loss =nn.CrossEntropyLoss()(prediction, target.long())
[18](vscode-notebook-cell:?execution_count=12&line=18) loss.backward()
---> [19](vscode-notebook-cell:?execution_count=12&line=19) optimizer.step()
File /home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/torch/optim/optimizer.py:391, in Optimizer.profile_hook_step.<locals>.wrapper(*args, **kwargs)
[386](https://vscode-remote+wsl-002bwsl4datascience.vscode-resource.vscode-cdn.net/home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/torch/optim/optimizer.py:386) else:
[387](https://vscode-remote+wsl-002bwsl4datascience.vscode-resource.vscode-cdn.net/home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/torch/optim/optimizer.py:387) raise RuntimeError(
[388](https://vscode-remote+wsl-002bwsl4datascience.vscode-resource.vscode-cdn.net/home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/torch/optim/optimizer.py:388) f"{func} must return None or a tuple of (new_args, new_kwargs), but got {result}."
[389](https://vscode-remote+wsl-002bwsl4datascience.vscode-resource.vscode-cdn.net/home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/torch/optim/optimizer.py:389) )
--> [391](https://vscode-remote+wsl-002bwsl4datascience.vscode-resource.vscode-cdn.net/home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/torch/optim/optimizer.py:391) out = func(*args, **kwargs)
[392](https://vscode-remote+wsl-002bwsl4datascience.vscode-resource.vscode-cdn.net/home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/torch/optim/optimizer.py:392) self._optimizer_step_code()
[394](https://vscode-remote+wsl-002bwsl4datascience.vscode-resource.vscode-cdn.net/home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/torch/optim/optimizer.py:394) # call optimizer step post hooks
File /home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/torch/utils/_contextlib.py:115, in context_decorator.<locals>.decorate_context(*args, **kwargs)
[112](https://vscode-remote+wsl-002bwsl4datascience.vscode-resource.vscode-cdn.net/home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/torch/utils/_contextlib.py:112) @functools.wraps(func)
[113](https://vscode-remote+wsl-002bwsl4datascience.vscode-resource.vscode-cdn.net/home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/torch/utils/_contextlib.py:113) def decorate_context(*args, **kwargs):
[114](https://vscode-remote+wsl-002bwsl4datascience.vscode-resource.vscode-cdn.net/home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/torch/utils/_contextlib.py:114) with ctx_factory():
--> [115](https://vscode-remote+wsl-002bwsl4datascience.vscode-resource.vscode-cdn.net/home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/torch/utils/_contextlib.py:115) return func(*args, **kwargs)
File /home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/bitsandbytes/optim/optimizer.py:287, in Optimizer8bit.step(self, closure)
[284](https://vscode-remote+wsl-002bwsl4datascience.vscode-resource.vscode-cdn.net/home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/bitsandbytes/optim/optimizer.py:284) self.init_state(group, p, gindex, pindex)
[286](https://vscode-remote+wsl-002bwsl4datascience.vscode-resource.vscode-cdn.net/home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/bitsandbytes/optim/optimizer.py:286) self.prefetch_state(p)
...
-> [1584](https://vscode-remote+wsl-002bwsl4datascience.vscode-resource.vscode-cdn.net/home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/bitsandbytes/functional.py:1584) optim_func = str2optimizer32bit[optimizer_name][0]
[1585](https://vscode-remote+wsl-002bwsl4datascience.vscode-resource.vscode-cdn.net/home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/bitsandbytes/functional.py:1585) elif g.dtype == torch.float16:
[1586](https://vscode-remote+wsl-002bwsl4datascience.vscode-resource.vscode-cdn.net/home/default/miniconda/envs/domf_iris2/lib/python3.10/site-packages/bitsandbytes/functional.py:1586) optim_func = str2optimizer32bit[optimizer_name][1]
KeyError: 'lamb'Expected behavior
The optimizer should be working.
Activity
TimDettmers commentedon Feb 16, 2026
Closing as duplicate of #1810. The
KeyErrorinstr2optimizer32bitfor the LAMB optimizer is the same root cause — LAMB (and LARS) entries were missing from thestr2optimizer32bitdispatch dictionary. PR #1855 addresses this by adding the missing entries. Once that PR is merged, LAMB 32-bit should work correctly.If the issue persists after the fix is released, please reopen or file a new issue.