Repository: Lightning-AI/litgpt (language: python)

Issue:
grokadamw and 8bit bnb optimizers
Hi everyone,

I wondering how to use optimizers such as [grokadamw](https://github.com/cognitivecomputations/grokadamw) and AdamW8bit.

When I try to use them I get error:

```
TypeError: AdamW8bit.__init__() got an unexpected keyword argument 'fused'
``` 

'fused' is always passed as argument to optimizer.


Discussion / hints:
Hi there,

theoretically other optimizers that have a PyTorch optimizer-like usage should work, I think. Here's how the optimizer is instantiated:

https://github.com/Lightning-AI/litgpt/blob/be8b28dad46caee5dd62fedcd6051485e27c404b/litgpt/utils.py#L559-L567

I only tested it with Adam and SGD optimizers though, which have the "fused" parameter. I think this should be updated to support other optimizers as well (a good test case would be `torch.optim.RMSprop`, which doesn't have a `fused` parameter).
I think this line causes issue:

https://github.com/Lightning-AI/litgpt/blob/be8b28dad46caee5dd62fedcd6051485e27c404b/litgpt/pretrain.py#L212

It assumes that `'fused'` should be `True` if `cuda` is present.
Good catch. Tbh I only checked the lora and full finetuning codes which is probably why I missed this. Ideally this could be updated to only provide these extra_args the chosen optimizer has the `fused` parameter.
Simplest check can be done using:

```python
import inspect

'fused' in inspect.signature(grokadamw.GrokAdamW).parameters
# False

'fused' in inspect.signature(bnb.optim.AdamW8bit).parameters
# False

'fused' in inspect.signature(torch.optim.AdamW).parameters
# True
```

Basically check if 'fused' is in constructor method as parameter.
Yes, a `'fused' in inspect.signature(optimizer)` check sounds totally reasonable. Please feel free to add this via a PR if you have time.

Fix patch:
diff --git a/litgpt/utils.py b/litgpt/utils.py
index f1af9a8a1c..c3b4b9fa5c 100644
--- a/litgpt/utils.py
+++ b/litgpt/utils.py
@@ -557,13 +557,37 @@ def instantiate_bnb_optimizer(optimizer, model_parameters):
 
 
 def instantiate_torch_optimizer(optimizer, model_parameters, **kwargs):
+    # Special care taken where some optimizers do not have some parameters referenced in some of the code, for example "fused" in the pretrain.py script:
+    #   bnb.optim.AdamW8bit
+    #   grokadamw.GrokAdamW
+    #   torch.optim.RMSprop
+    
     if isinstance(optimizer, str):
-        optimizer_cls = getattr(torch.optim, optimizer)
+        if "." in optimizer:
+            class_module, class_name = optimizer.rsplit(".", 1)
+        else:
+            class_module, class_name = "torch.optim", optimizer
+
+        module = __import__(class_module, fromlist=[class_name])
+        optimizer_cls = getattr(module, class_name)
+
+        valid_params = set(inspect.signature(optimizer_cls).parameters)
+        kwargs = {key: value for key, value in dict(kwargs).items() if key in valid_params}
         optimizer = optimizer_cls(model_parameters, **kwargs)
-    else:
-        optimizer = dict(optimizer)  # copy
+    elif isinstance(optimizer, dict):
+        optimizer = dict(optimizer)
+        class_module, class_name = optimizer["class_path"].rsplit(".", 1)
+        module = __import__(class_module, fromlist=[class_name])
+        optimizer_cls = getattr(module, class_name)
+
+        valid_params = set(inspect.signature(optimizer_cls).parameters)
+        kwargs = {key: value for key, value in dict(kwargs).items() if key in valid_params}
+        
         optimizer["init_args"].update(kwargs)
         optimizer = instantiate_class(model_parameters, optimizer)
+    else:
+        raise ValueError(f'Unrecognized "optimizer" value: {optimizer}')
+
     return optimizer
 
 


Test patch:
diff --git a/tests/test_pretrain.py b/tests/test_pretrain.py
index 5669a84008..3b28894793 100644
--- a/tests/test_pretrain.py
+++ b/tests/test_pretrain.py
@@ -18,6 +18,27 @@
 from tests.conftest import RunIf
 
 
+@RunIf(min_cuda_gpus=1, standalone=True)
+@mock.patch("litgpt.pretrain.save_hyperparameters")
+def test_optimizer_args(_, tmp_path):
+    model_config = Config(block_size=2, n_layer=2, n_embd=4, n_head=2, padded_vocab_size=8)
+
+    dataset = torch.tensor([[0, 1, 2], [3, 4, 5], [0, 1, 2]])
+    dataloader = DataLoader(dataset)
+    pretrain.get_dataloaders = Mock(return_value=(dataloader, dataloader))
+
+    for i in ("AdamW", "SGD", "RMSprop"):
+        pretrain.setup(
+            "pythia-14m",
+            devices=1,
+            optimizer="RMSprop",
+            model_config=model_config,
+            out_dir=tmp_path,
+            train=TrainArgs(global_batch_size=2, max_tokens=16, save_interval=1, micro_batch_size=1, max_norm=1.0),
+            eval=EvalArgs(interval=1, max_iters=1, final_validation=False),
+        )
+
+
 @RunIf(min_cuda_gpus=2, standalone=True)
 # Set CUDA_VISIBLE_DEVICES for FSDP hybrid-shard, if fewer GPUs are used than are available
 @mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1"})