from 1.11.0, torchao.prototype.low_bit_optim is promoted to torchao.optim (#38689)

* since 1.11.0, torchao.prototype.low_bit_optim is promoted to
torchao.optim

Signed-off-by: YAO Matrix <matrix.yao@intel.com>

* fix review comments

Signed-off-by: YAO Matrix <matrix.yao@intel.com>

---------

Signed-off-by: YAO Matrix <matrix.yao@intel.com>
Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>
This commit is contained in:
Yao Matrix
2025-06-11 20:16:25 +08:00
committed by GitHub
parent 56a7cf5546
commit c8c1e525ed

View File

@@ -5368,19 +5368,19 @@ if is_torch_available():
)
)
if is_torchao_available():
import torchao
from torchao.optim import AdamW4bit, AdamW8bit
optim_test_params.append(
(
OptimizerNames.ADAMW_TORCH_4BIT,
torchao.prototype.low_bit_optim.AdamW4bit,
AdamW4bit,
default_adam_kwargs,
)
)
optim_test_params.append(
(
TrainingArguments(optim=OptimizerNames.ADAMW_TORCH_8BIT, output_dir="None"),
torchao.prototype.low_bit_optim.AdamW8bit,
AdamW8bit,
default_adam_kwargs,
)
)