Support adamw_torch_8bit (#34993)

* var * more * test
2025-01-21 21:17:49 +08:00
parent f82b19cb6f
commit dc10f7906a
3 changed files with 19 additions and 3 deletions
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -5017,6 +5017,13 @@ if is_torch_available():
                default_adam_kwargs,
            )
        )
+        optim_test_params.append(
+            (
+                TrainingArguments(optim=OptimizerNames.ADAMW_TORCH_8BIT, output_dir="None"),
+                torchao.prototype.low_bit_optim.AdamW8bit,
+                default_adam_kwargs,
+            )
+        )


@require_torch