class_path: torchao.prototype.low_bit_optim.AdamW8bit

scripts/pretrain-core-model.yaml

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