Update train.py

train.py

@@ -15,7 +15,7 @@ random.seed(random_seed)
 
 dataset = load_dataset("Crystalcareai/Self-Discover-MM-Instruct-openai", split="train_sft")
 
-n_ahead_talk_global = 3
+n_ahead_talk_global = 4
 n_passes_global = 2
 n_ahead_global = 2
 n_examples = 0
@@ -64,7 +64,8 @@ def model_init(params):
     )
     print("Loaded model")
    
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id,truncation=True,padding="left")
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
+    tokenizer.padding_side = 'left'  # Adjust padding side to 'left' to avoid batch generation issues with Flash Attention
     tokenizer.pad_token_id = tokenizer.eos_token_id
 
     special_tokens_to_add = []
@@ -96,15 +97,15 @@ def model_init(params):
     model.train()
     return model
 
-max_seq_length = 1024
+max_seq_length = 2048
 run_id = int(time.time())
 training_args = TrainingArguments(
     output_dir="./out",
-    num_train_epochs=1.5,
+    num_train_epochs=3,
     per_device_train_batch_size=1,
     gradient_checkpointing=False,
-    gradient_accumulation_steps=8,
-    optim="lion_32bit",
+    gradient_accumulation_steps=16,
+    optim="adamw_torch_fused",
     logging_steps=1,
     save_strategy="steps",
     save_steps=300,
@@ -114,8 +115,8 @@ training_args = TrainingArguments(
     # beta1=0.9,
     # beta2=0.95,
     # auto_find_batch_size=True
-    learning_rate=3e-07,
-    max_grad_norm=0.3,  # Gradient clipping with a maximum gradient norm of 0.3
+    learning_rate=2e-07,
+    max_grad_norm=1.0,  # Gradient clipping with a maximum gradient norm of 0.3
     warmup_steps=10,
     lr_scheduler_type="cosine",
     push_to_hub=False,
@@ -147,4 +148,4 @@ trainer = SFTTrainer(
     max_seq_length=max_seq_length,
 )
 
-trainer.train()
+trainer.train()