Update train.py

train.py

@@ -14,10 +14,11 @@ random.seed(random_seed)
 
 dataset = load_dataset("HuggingFaceH4/deita-10k-v0-sft", split="train_sft")
 
-n_ahead_talk_global = 4
+n_ahead_talk_global = 2
 n_passes_global = 2
-n_ahead_global = 12
-full_batch_size = 8
+n_ahead_global = 2
+n_examples = 0
+full_batch_size = 2
 eval_and_logging_steps = 2
 save_steps = 100
 
@@ -43,11 +44,11 @@ def model_init(params):
     optimize_lm_head_only_at_start = params.get("optimize_lm_head_only_at_start", False)
 
     model_id = "Crystalcareai/Quiet-Star-Custom"
-    tokenizer_id = "Crystalcareai/Quiet-Star-Custom"
+    tokenizer_id = model_id
     print("Loading model")
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
-        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+        torch_dtype=torch.float16 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,
@@ -58,13 +59,12 @@ def model_init(params):
         use_complex_think_head=False,
         use_complex_talk_head=True,
         use_weighted_talk_head=True,
-        trust_remote_code=True,
-        load_in_4bit=True,     
+        trust_remote_code=True,  
+        device_map="auto",
     )
     print("Loaded model")
    
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
-    tokenizer.padding_side = "right"
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id,padding=False,truncation=True)
     tokenizer.pad_token_id = tokenizer.eos_token_id
 
     special_tokens_to_add = []
@@ -103,31 +103,33 @@ training_args = TrainingArguments(
     output_dir="./out",
     num_train_epochs=3,
     per_device_train_batch_size=1,
-    gradient_accumulation_steps=global_gradient_accumulation_steps,
-    gradient_checkpointing=True,
-    optim="adamw_bnb_8bit",
-    logging_steps=2,
+    gradient_checkpointing=False,
+    gradient_accumulation_steps=4,
+    optim="adamw_torch_fused",
+    logging_steps=1,
     save_strategy="steps",
     save_steps=300,
     bf16=True,
     tf32=False,
-    learning_rate=2e-4,
-    max_grad_norm=0.3,
-    warmup_ratio=0.00,
-    lr_scheduler_type="constant",
+    # 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,
+    lr_scheduler_type="cosine",
     push_to_hub=False,
 )
 
-peft_config = LoraConfig(
-        lora_alpha=16,
-        lora_dropout=0.05,
-        r=32,
-        bias="none",
-        target_modules = ["q_proj", "k_proj", "v_proj", "o_proj","gate_proj", "up_proj", "down_proj",],
-        task_type="CAUSAL_LM",
-        use_dora=False,  # Enable Dora method
-)
+# peft_config = LoraConfig(
+#           r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
+#     target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
+#                       "gate_proj", "up_proj", "down_proj",],
+#     lora_alpha = 16,
+#     lora_dropout = 0, # Supports any, but = 0 is optimized
+#     bias = "none", # Enable Dora method
+#     use_dora=True,
+# )
 
+torch.autograd.set_detect_anomaly(True)
 model = model_init(None)  # Initialize the model
 tokenizer = model.tokenizer  
 
@@ -135,8 +137,8 @@ trainer = SFTTrainer(
     args=training_args,
     train_dataset=dataset,
     model=model,
-    peft_config=peft_config,
+    # peft_config=peft_config,
     tokenizer=tokenizer,  
 )
 
-trainer.train()
+trainer.train()