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

yao-matrix · SunMarc · web-flow · commit c8c1e525ed87 · 2025-06-11T12:16:25.000Z
* since 1.11.0, torchao.prototype.low_bit_optim is promoted to
torchao.optim

Signed-off-by: YAO Matrix &lt;matrix.yao@intel.com&gt;

* fix review comments

Signed-off-by: YAO Matrix &lt;matrix.yao@intel.com&gt;

---------

Signed-off-by: YAO Matrix &lt;matrix.yao@intel.com&gt;
Co-authored-by: Marc Sun &lt;57196510+SunMarc@users.noreply.github.com&gt;
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
@@ -5368,19 +5368,19 @@ def hp_name(trial):
             )
         )
     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,
             )
         )

Original file line number	Diff line number	Diff line change
`@@ -5368,19 +5368,19 @@ def hp_name(trial):`
`5368`	`5368`	`)`
`5369`	`5369`	`)`
`5370`	`5370`	`if is_torchao_available():`
`5371`		`- import torchao`
	`5371`	`+ from torchao.optim import AdamW4bit, AdamW8bit`
`5372`	`5372`
`5373`	`5373`	`optim_test_params.append(`
`5374`	`5374`	`(`
`5375`	`5375`	`OptimizerNames.ADAMW_TORCH_4BIT,`
`5376`		`- torchao.prototype.low_bit_optim.AdamW4bit,`
	`5376`	`+ AdamW4bit,`
`5377`	`5377`	`default_adam_kwargs,`
`5378`	`5378`	`)`
`5379`	`5379`	`)`
`5380`	`5380`	`optim_test_params.append(`
`5381`	`5381`	`(`
`5382`	`5382`	`TrainingArguments(optim=OptimizerNames.ADAMW_TORCH_8BIT, output_dir="None"),`
`5383`		`- torchao.prototype.low_bit_optim.AdamW8bit,`
	`5383`	`+ AdamW8bit,`
`5384`	`5384`	`default_adam_kwargs,`
`5385`	`5385`	`)`
`5386`	`5386`	`)`