tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

feat: add complete optimizer docs (#3017) [skip ci]

Browse Source 
* feat: add complete optimizer docs

* fix: deprecate old torchao adamw low bit
This commit is contained in:
NanoCode012
2025-08-06 19:01:51 +07:00
committed by GitHub
parent 70faea331f
commit e3177c3210
3 changed files with 114 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/axolotl/core/builders/base.py
View File

@@ -29,7 +29,6 @@ from transformers import (
TrainerCallback,
)
from transformers.trainer_pt_utils import AcceleratorConfig
from transformers.training_args import OptimizerNames
from axolotl.integrations.base import PluginManager
from axolotl.monkeypatch.trainer.lr import patch_trainer_get_lr
@@ -284,21 +283,6 @@ class TrainerBuilderBase(abc.ABC):
optimizer_kwargs["foreach"] = False
optimizer_cls = AdamW
optimizer_kwargs.update(adam_kwargs)
elif self.cfg.optimizer == "ao_adamw_4bit":
# TODO remove 20250401
from torchao.prototype.low_bit_optim import AdamW4bit
optimizer_cls = AdamW4bit
optimizer_kwargs.update(adam_kwargs)
LOG.warning(
f"`ao_adamw_4bit` will be deprecated soon. Please use `{OptimizerNames.ADAMW_TORCH_4BIT}` instead."
)
elif self.cfg.optimizer == "ao_adamw_8bit":
from torchao.prototype.low_bit_optim import AdamW8bit
optimizer_cls = AdamW8bit
optimizer_kwargs.update(adam_kwargs)
elif self.cfg.optimizer == "ao_adamw_fp8":
from torchao.prototype.low_bit_optim import AdamWFp8