Update train.py

train.py

@@ -17,7 +17,6 @@ dataset = load_dataset("HuggingFaceH4/deita-10k-v0-sft", split="train_sft")
 n_ahead_talk_global = 4
 n_passes_global = 2
 n_ahead_global = 12
-n_examples = 1_000
 full_batch_size = 8
 eval_and_logging_steps = 2
 save_steps = 100
@@ -64,7 +63,8 @@ def model_init(params):
     )
     print("Loaded model")
    
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id,padding=False,truncation=True)
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
+    tokenizer.padding_side = "right"
     tokenizer.pad_token_id = tokenizer.eos_token_id
 
     special_tokens_to_add = []
@@ -103,14 +103,14 @@ training_args = TrainingArguments(
     output_dir="./out",
     num_train_epochs=3,
     per_device_train_batch_size=1,
-    gradient_checkpointing=False,
+    gradient_accumulation_steps=global_gradient_accumulation_steps,
+    gradient_checkpointing=True,
     optim="adamw_bnb_8bit",
     logging_steps=2,
     save_strategy="steps",
     save_steps=300,
-
     bf16=True,
-    tf32=True,
+    tf32=False,
     learning_rate=2e-4,
     max_grad_norm=0.3,
     warmup_ratio=0.00,
@@ -139,4 +139,4 @@ trainer = SFTTrainer(
     tokenizer=tokenizer,  
 )
 
-trainer.train()
+trainer.train()