From c8c1e525ed871eabf053cd48f26cf83edb90c00b Mon Sep 17 00:00:00 2001 From: Yao Matrix Date: Wed, 11 Jun 2025 20:16:25 +0800 Subject: [PATCH] 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 * fix review comments Signed-off-by: YAO Matrix --------- Signed-off-by: YAO Matrix Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com> --- tests/trainer/test_trainer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py index 65a16066c9..7b11ecc3ec 100644 --- a/tests/trainer/test_trainer.py +++ b/tests/trainer/test_trainer.py @@ -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, ) )