Put train.py

train.py

@@ -87,8 +87,6 @@ def main():
             [r["text"] for r in train_ds.select(range(min(2000, len(train_ds))))],
             truncation=False,
         )
-        import statistics
-
         lens = sorted(len(x) for x in probe["input_ids"])
         print(
             "token length stats: min=%d p50=%d p90=%d p95=%d p99=%d max=%d"
@@ -136,7 +134,7 @@ def main():
         gradient_accumulation_steps=args.grad_accum,
         learning_rate=args.lr,
         lr_scheduler_type="cosine",
-        warmup_ratio=0.03,
+        warmup_steps=0.03,
         bf16=use_bf16,
         fp16=not use_bf16,
         max_length=args.max_length,