Update train.py

train.py

@@ -7,20 +7,19 @@ from transformers import TrainingArguments
 from trl import SFTTrainer
 from peft import LoraConfig
 
+
 import time
 random_seed = 42
 torch.manual_seed(random_seed)
 random.seed(random_seed)
 
-dataset = load_dataset("HuggingFaceH4/deita-10k-v0-sft", split="train_sft")
+dataset = load_dataset("Crystalcareai/Self-Discover-MM-Instruct-openai", split="train_sft")
 
-n_ahead_talk_global = 2
+n_ahead_talk_global = 3
 n_passes_global = 2
 n_ahead_global = 2
 n_examples = 0
-full_batch_size = 2
-eval_and_logging_steps = 2
-save_steps = 100
+
 
 
 def model_init(params):
@@ -39,7 +38,6 @@ def model_init(params):
     include_policy_loss = params.get("include_policy_loss", True)
     gumbel_detach = params.get("gumbel_detach", True)
     merged_talk_heads = params.get("merged_talk_heads", True)
-    gradient_accumulation_steps = params.get("gradient_accumulation_steps", global_gradient_accumulation_steps)
     residual_think_head = params.get("residual_think_head", False)
     optimize_lm_head_only_at_start = params.get("optimize_lm_head_only_at_start", False)
 
@@ -48,7 +46,7 @@ def model_init(params):
     print("Loading model")
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
         max_thoughts=n_ahead + n_ahead_talk + 1,
         merged_talk_heads=merged_talk_heads,
         merged_lm_and_talk_heads=False,
@@ -61,10 +59,12 @@ def model_init(params):
         use_weighted_talk_head=True,
         trust_remote_code=True,  
         device_map="auto",
+        # load_in_4bit=True,
+        # attn_implementation="flash_attention_2",
     )
     print("Loaded model")
    
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id,padding=False,truncation=True)
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id,truncation=True,padding="left")
     tokenizer.pad_token_id = tokenizer.eos_token_id
 
     special_tokens_to_add = []
@@ -76,6 +76,10 @@ def model_init(params):
         tokenizer.add_special_tokens({"additional_special_tokens": special_tokens_to_add})
         model.resize_token_embeddings(len(tokenizer))
     model.tokenizer = tokenizer
+    for name, module in model.named_modules():
+        if "embed" in name:
+            print(module, flush=True)
+   
     model.gumbel_detach = gumbel_detach
     model.include_policy_loss = include_policy_loss
     model.use_end_thought_token = use_end_thought_token
@@ -83,40 +87,40 @@ def model_init(params):
     model.n_ahead = n_ahead
     model.n_ahead_talk = n_ahead_talk
     model.n_passes = n_passes
-    model.n_tokens_print = gradient_accumulation_steps
-    model.gradient_accumulation_steps = gradient_accumulation_steps
     model.residual_think_head = residual_think_head
     model.optimize_lm_head_only_at_start = optimize_lm_head_only_at_start
     model.gumbel_temperature = gumbel_temperature
     model.original_mode = original
     model.config_params = params
     model.run_start = int(time.time())
-    model.kill_after = 100
     model.train()
     return model
 
-
-batch_size = full_batch_size // n_passes_global
-global_gradient_accumulation_steps = full_batch_size // batch_size
+max_seq_length = 1024
 run_id = int(time.time())
 training_args = TrainingArguments(
     output_dir="./out",
-    num_train_epochs=3,
+    num_train_epochs=1.5,
     per_device_train_batch_size=1,
     gradient_checkpointing=False,
-    gradient_accumulation_steps=4,
-    optim="adamw_torch_fused",
+    gradient_accumulation_steps=8,
+    optim="lion_32bit",
     logging_steps=1,
     save_strategy="steps",
     save_steps=300,
     bf16=True,
     tf32=False,
+    # epsilson=1e-05,
+    # beta1=0.9,
+    # beta2=0.95,
     # auto_find_batch_size=True
-    learning_rate=2e-07,
-    max_grad_norm=1.0,  # Gradient clipping with a maximum gradient norm of 0.3
-    warmup_steps=100,
+    learning_rate=3e-07,
+    max_grad_norm=0.3,  # Gradient clipping with a maximum gradient norm of 0.3
+    warmup_steps=10,
     lr_scheduler_type="cosine",
     push_to_hub=False,
+    report_to="wandb"
+    
 )
 
 # peft_config = LoraConfig(
@@ -131,14 +135,16 @@ training_args = TrainingArguments(
 
 torch.autograd.set_detect_anomaly(True)
 model = model_init(None)  # Initialize the model
-tokenizer = model.tokenizer  
 
+tokenizer = model.tokenizer  
+        
 trainer = SFTTrainer(
     args=training_args,
     train_dataset=dataset,
     model=model,
     # peft_config=peft_config,
     tokenizer=tokenizer,  
+    max_seq_length=max_seq_length,
 )
 
-trainer.train()
+trainer.train()