Update schedulefree.py

schedulefree.py

@@ -1,12 +1,12 @@
 import signal
 import sys
+import torch
 from datasets import load_dataset
-from transformers import TrainingArguments
+from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForCausalLM
 from trl import SFTTrainer
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from peft import LoraConfig
-from schedulefree import AdamWScheduleFree
+
+# Importing Sophia optimizer
+from sophia import SophiaG
 
 # Signal handler function
 def signal_handler(sig, frame):
@@ -16,88 +16,85 @@ def signal_handler(sig, frame):
 # Register signal handler
 signal.signal(signal.SIGINT, signal_handler)
 
-dataset = load_dataset("Crystalcareai/Orca-Reka")['train']
+# Load the dataset
+dataset = load_dataset("Crystalcareai/Orca-Reka", split="train")
+model_id = "./outkannn"
+tokenizer_id = model_id
+tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
+tokenizer.padding_side = 'right'
+
+# Formatting function for the dataset
+alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
 
-def chatml_format(example):
-    """Format the dataset for training, accounting for empty columns."""
-    return {
-        "instruction": example['instruction'] if 'instruction' in example else " \n",
-        "input": example['input'] if 'input' in example else " \n",
-        "system": example['system'] if 'system' in example else " \n",
-        "output": example['output'] if 'output' in example else " \n",
-    }
+### Instruction:
+{}
 
-# Format dataset
-dataset = dataset.map(chatml_format, remove_columns=dataset.column_names)
+### Input:
+{}
 
-# Load model and tokenizer
+### Response:
+{}"""
+def formatting_prompts_func(examples):
+    instructions = examples["instruction"]
+    inputs = examples["input"]
+    outputs = examples["output"]
+    texts = []
+    EOS_TOKEN = tokenizer.eos_token
+    for instruction, input, output in zip(instructions, inputs, outputs):
+        text = alpaca_prompt.format(instruction, input, output) + EOS_TOKEN
+        texts.append(text)
+    return {"text": texts}
+
+# Process and map the formatting function
+dataset = dataset.map(formatting_prompts_func, batched=True)
+
+# Load model
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     device_map="auto",
-    attn_implementation="flash_attention_2",
     torch_dtype=torch.bfloat16,
-)
-tokenizer = AutoTokenizer.from_pretrained(model)
-tokenizer.padding_side = 'right'  # to prevent warnings
-
-peft_config = LoraConfig(
-    lora_alpha=16,
-    lora_dropout=0.05,
-    r=32,
-    bias="none",
-    target_modules=[
-        "0.w1",
-        "0.w2",
-        "0.w3",
-        "q_proj",
-        "v_proj",
-        "k_proj",
-        "o_proj"
-    ],
-    task_type="CAUSAL_LM",
-    use_dora=False,  # Enable Dora method
+    trust_remote_code=True
 )
 
+# Define training arguments
 args = TrainingArguments(
-    output_dir="./out",  # directory to save and repository id
-    num_train_epochs=3,  # number of training epochs
-    per_device_train_batch_size=4,  # batch size per device during training
-    gradient_checkpointing=True,  # use gradient checkpointing to save memory
-    optim="adamw_hf",
+    output_dir="./out",
+    num_train_epochs=3,
+    per_device_train_batch_size=4,
+    gradient_checkpointing=True,
     logging_steps=2,
     save_strategy="steps",
     save_steps=300,
-    bf16=True,  # use bfloat16 precision
-    tf32=True,  # use tf32 precision
-    ### peft specific arguments ###
-    learning_rate=2e-4,
-    max_grad_norm=0.3,
+    bf16=True,
+    tf32=True,
+    learning_rate=1e-4,
+    max_grad_norm=0.1,
     warmup_ratio=0.00,
-    lr_scheduler_type="constant",
-    report_to="wandb",
-    push_to_hub=False,
-    # push model to hub
+    lr_scheduler_type="cosine",
+    push_to_hub=False
 )
 
-max_seq_length = 2048  # max sequence length for model and packing of the dataset
+max_seq_length = 2048
 
-# Create the schedulefree optimizer
-optimizer = AdamWScheduleFree(model.parameters(), lr=args.learning_rate, beta=0.9)
+# Custom Trainer Class
+class CustomTrainer(SFTTrainer):
+    def create_optimizer(self):
+        # Override to use SophiaG optimizer
+        self.optimizer = SophiaG(self.model.parameters(), lr=self.args.learning_rate, betas=(0.965, 0.99), rho=0.01, weight_decay=0.1)
 
-trainer = SFTTrainer(
+# Trainer configuration
+trainer = CustomTrainer(
     model=model,
     args=args,
     train_dataset=dataset,
-    ### peft specific arguments ###
-    peft_config=peft_config,
     max_seq_length=max_seq_length,
     tokenizer=tokenizer,
+    dataset_text_field="output",
     packing=False,
-    optimizers=(optimizer, None),  # Pass the schedulefree optimizer
 )
 
-# start training, the model will be automatically saved to the hub and the output directory
+# Start training
 trainer.train()
 
-# save model
-trainer.save_model()
+# Save model
+trainer.save_model()