Upload train_math.py

train_math.py

@@ -18,7 +18,7 @@ LR = 2e-4
 PER_DEVICE_BATCH = 4
 GRADIENT_ACCUMULATION = 32
 MAX_SEQ_LENGTH = 4096
-WARMUP_RATIO = 0.1
+WARMUP_STEPS = 500
 LORA_R = 32
 LORA_ALPHA = 16
 LORA_DROPOUT = 0.05
@@ -98,7 +98,7 @@ def main():
         per_device_train_batch_size=PER_DEVICE_BATCH,
         gradient_accumulation_steps=GRADIENT_ACCUMULATION,
         learning_rate=LR, bf16=True,
-        lr_scheduler_type="cosine", warmup_ratio=WARMUP_RATIO,
+        lr_scheduler_type="cosine", warmup_steps=WARMUP_STEPS,
         logging_steps=10, save_strategy="epoch", save_total_limit=2,
         gradient_checkpointing=True, push_to_hub=True, hub_model_id=HUB_MODEL_ID,
         hub_private_repo=False, report_to="trackio", disable_tqdm=True,
@@ -107,7 +107,7 @@ def main():
 
     print("Initializing SFTTrainer...")
     trainer = SFTTrainer(
-        model=model, tokenizer=tokenizer, train_dataset=train_dataset,
+        model=model, processing_class=tokenizer, train_dataset=train_dataset,
         peft_config=peft_config, args=training_args,
         max_seq_length=MAX_SEQ_LENGTH,
     )