Update train-h100-sharegpt-sft.py

train-h100-sharegpt-sft.py

@@ -15,12 +15,12 @@ random_seed = 42
 torch.manual_seed(random_seed)
 random.seed(random_seed)
 
-dataset = load_dataset("HuggingFaceH4/orca-math-word-problems-200k", split="train_sft").select(range(3000))
+dataset = load_dataset("HuggingFaceH4/orca-math-word-problems-200k", split="train_sft").select(range(1000))
 
 
-n_ahead_talk_global = 2
+n_ahead_talk_global = 4
 n_passes_global = 1
-n_ahead_global = 8
+n_ahead_global = 4
 # n_examples = 1000
 # full_batch_size = 8
 
@@ -96,15 +96,15 @@ def model_init(params):
     model.train()
     return model
 
-max_seq_length = 8192
+max_seq_length = 1024
 run_id = int(time.time())
 training_args = TrainingArguments(
     output_dir="./out",
     num_train_epochs=1,
     per_device_train_batch_size=1,
     gradient_checkpointing=False,
-    gradient_accumulation_steps=1,
-    optim="lion_32bit",
+    gradient_accumulation_steps=8,
+    optim="adamw_torch_fused",
     logging_steps=1,
     save_strategy="steps",
     save_steps=100,
@@ -114,74 +114,74 @@ training_args = TrainingArguments(
     bf16=True,
     
     tf32=True,
-    learning_rate=1e-07,
+    learning_rate=2e-10,
     max_grad_norm=0,
     warmup_steps=20,
-    lr_scheduler_type="constant",
+    lr_scheduler_type="cosine",
     push_to_hub=False,
     report_to="wandb"
 )
 
 peft_config = LoraConfig(
     r = 8, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
-    target_modules =["up_proj", "down_proj", "gate_proj"],
+    target_modules =["q_proj", "v_proj"],
     lora_alpha = 32,
     lora_dropout = 0, # Supports any, but = 0 is optimized
     bias = "none", 
-    use_dora=False,
+    use_dora=True,
     task_type="CAUSAL_LM"
 )
 
 torch.autograd.set_detect_anomaly(True)
 
-class CustomSFTTrainer(SFTTrainer):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.beta = 0.9  # momentum factor
-        self.clip_factor = 1.0  # clipping factor
-        self.moving_avg = 0.0
+# class CustomSFTTrainer(SFTTrainer):
+#     def __init__(self, *args, **kwargs):
+#         super().__init__(*args, **kwargs)
+#         self.beta = 0.9  # momentum factor
+#         self.clip_factor = 1.0  # clipping factor
+#         self.moving_avg = 0.0
 
-    def training_step(self, model, inputs):
-        model.train()
-        inputs = self._prepare_inputs(inputs)
+#     def training_step(self, model, inputs):
+#         model.train()
+#         inputs = self._prepare_inputs(inputs)
 
-        outputs = model(**inputs)
-        loss = outputs.loss if isinstance(outputs, dict) else outputs[0]
+#         outputs = model(**inputs)
+#         loss = outputs.loss if isinstance(outputs, dict) else outputs[0]
 
-        if self.args.gradient_accumulation_steps > 1:
-            loss = loss / self.args.gradient_accumulation_steps
+#         if self.args.gradient_accumulation_steps > 1:
+#             loss = loss / self.args.gradient_accumulation_steps
 
-        loss.backward()
+#         loss.backward()
 
-        # Compute gradients and their norm
-        grad_norm = torch.sqrt(sum(p.grad.data.norm().to(model.device)**2 for p in model.parameters() if p.grad is not None))
+#         # Compute gradients and their norm
+#         grad_norm = torch.sqrt(sum(p.grad.data.norm().to(model.device)**2 for p in model.parameters() if p.grad is not None))
 
-        # Update moving average and apply gradient clipping
-        if self.state.global_step == 0:
-            self.moving_avg = grad_norm
-        else:
-            self.moving_avg = self.beta * self.moving_avg + (1 - self.beta) * grad_norm
+#         # Update moving average and apply gradient clipping
+#         if self.state.global_step == 0:
+#             self.moving_avg = grad_norm
+#         else:
+#             self.moving_avg = self.beta * self.moving_avg + (1 - self.beta) * grad_norm
 
-        if grad_norm > self.clip_factor * self.moving_avg:
-            clip_coef = (self.clip_factor * self.moving_avg / grad_norm).item()
-            for param in model.parameters():
-                if param.grad is not None:
-                    param.grad.data.mul_(clip_coef)
+#         if grad_norm > self.clip_factor * self.moving_avg:
+#             clip_coef = (self.clip_factor * self.moving_avg / grad_norm).item()
+#             for param in model.parameters():
+#                 if param.grad is not None:
+#                     param.grad.data.mul_(clip_coef)
 
-        if (self.state.global_step + 1) % self.args.gradient_accumulation_steps == 0:
-            self.optimizer.step()
-            self.lr_scheduler.step()
-            model.zero_grad()
-            self.state.global_step += 1
+#         if (self.state.global_step + 1) % self.args.gradient_accumulation_steps == 0:
+#             self.optimizer.step()
+#             self.lr_scheduler.step()
+#             model.zero_grad()
+#             self.state.global_step += 1
 
-        # Return the loss as a Tensor
-        return loss
+#         # Return the loss as a Tensor
+#         return loss
     
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 model = model_init(None)  
 
-trainer = CustomSFTTrainer(
+trainer = SFTTrainer(
     model=model,
     args=training_args,
     train_dataset=dataset,