class_path: torchao.prototype.low_bit_optim.AdamW8bit

scripts/pretrain-core-model.yaml

@@ -115,8 +115,8 @@ eval:
 
 optimizer:
   # class_path: torch.optim.AdamW
-  # class_path: torchao.prototype.low_bit_optim.AdamW8bit
-  class_path: torchao.prototype.low_bit_optim.AdamW4bit
+  class_path: torchao.prototype.low_bit_optim.AdamW8bit
+  # class_path: torchao.prototype.low_bit_optim.AdamW4bit
   init_args:
     # (type: float, default: 0.001)
     lr: 1e-4