Rename train.py to train-h100-sharegpt-sft.py

train.py → train-h100-sharegpt-sft.py

@@ -44,9 +44,11 @@ def model_init(params):
     model_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,14 +60,13 @@ def model_init(params):
         use_complex_talk_head=True,
         use_weighted_talk_head=True,
         trust_remote_code=True,  
-        device_map="auto",
+        # device_map="auto",
         # load_in_4bit=True,
         # attn_implementation="flash_attention_2",
     )
     print("Loaded model")
    
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
-    tokenizer.padding_side = 'left'  # Adjust padding side to 'left' to avoid batch generation issues with Flash Attention
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id,truncation=True,padding_side="right")
     tokenizer.pad_token_id = tokenizer.eos_token_id
 
     special_tokens_to_add = []
@@ -97,42 +98,42 @@ def model_init(params):
     model.train()
     return model
 
-max_seq_length = 2048
+max_seq_length = 1024
 run_id = int(time.time())
 training_args = TrainingArguments(
     output_dir="./out",
     num_train_epochs=3,
     per_device_train_batch_size=1,
     gradient_checkpointing=False,
-    gradient_accumulation_steps=16,
-    optim="adamw_torch_fused",
+    gradient_accumulation_steps=6,
+    optim="lion_32bit",
     logging_steps=1,
     save_strategy="steps",
-    save_steps=300,
+    save_steps=25,
     bf16=True,
+    # fp16=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=10,
+    learning_rate=6e-05,
+    max_grad_norm=0.3,  # Gradient clipping with a maximum gradient norm of 0.3
+    warmup_ratio=0.06,
     lr_scheduler_type="cosine",
     push_to_hub=False,
     report_to="wandb"
     
 )
 
-# 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,
-# )
+peft_config = LoraConfig(
+          r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
+    target_modules = ["q_proj", "k_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
@@ -143,7 +144,7 @@ trainer = SFTTrainer(
     args=training_args,
     train_dataset=dataset,
     model=model,
-    # peft_config=peft_config,
+    peft_config=peft_config,
     tokenizer=tokenizer,  
     max_seq_length=max_seq_length,
 )