feat(train): custom start_preconditioning_step

tools/train/train.py

@@ -248,6 +248,10 @@ class TrainingArguments:
         default=1024,
         metadata={"help": "Chunked size for large layers with Distributed Shampoo."},
     )
+    start_preconditioning_step: int = field(
+        default=100,
+        metadata={"help": "Number of steps before starting to update preconditioner."},
+    )
     preconditioning_compute_steps: int = field(
         default=10, metadata={"help": "Number of steps to update preconditioner."}
     )
@@ -608,7 +612,7 @@ def main():
             beta2=training_args.beta2,
             diagonal_epsilon=1e-10,
             matrix_epsilon=1e-8,
-            start_preconditioning_step=training_args.warmup_steps,
+            start_preconditioning_step=training_args.start_preconditioning_step,
             preconditioning_compute_steps=training_args.preconditioning_compute_steps,
             statistics_compute_steps=1,
             best_effort_shape_interpretation=True,