Update finetune_script.py

finetune_script.py

@@ -1,112 +1,138 @@
-from accelerate import Accelerator
-from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
-from transformers import DataCollatorForLanguageModeling
-from datasets import Dataset
-from trl import SFTTrainer
+# 0.2 Import Dependencies
 import os
-import pandas as pd
-
-# Initialize Accelerator for distributed training
-accelerator = Accelerator()
-
-if accelerator.is_local_main_process:
-    print(f"Running on {accelerator.device}")
-
-# ================================
-# Load Model and Tokenizer
-# ================================
-model_name = "unsloth/Meta-Llama-3.1-70B-bnb-4bit"
-
-# Use AutoTokenizer and AutoModelForCausalLM
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    load_in_4bit=True,
-    device_map="auto",  # Automatically shard across GPUs
-    torch_dtype="auto"  # Auto-adjust to fp16/bf16 based on mixed_precision
-)
-
-# ================================
-# LoRA (Low-Rank Adaptation)
-# ================================
-from peft import LoraConfig, get_peft_model
-
-# Apply LoRA for fine-tuning
-lora_config = LoraConfig(
-    r=16,
-    lora_alpha=16,
-    lora_dropout=0.1,
-    bias="none",
-    target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
-)
-model = get_peft_model(model, lora_config)
-
-# ================================
-# Load and Prepare Data
-# ================================
-# Load CSVs
-train_df = pd.read_csv("train.csv").fillna("")
-valid_df = pd.read_csv("valid.csv").fillna("")
-
-EOS_TOKEN = tokenizer.eos_token
-
-
-def format_prompts(df):
-    prompts = []
-    for _, row in df.iterrows():
-        inst, inp, out = row["instruction"], row["input"], row["output"]
-        prompt = f"### Instruction:\n{inst}\n\n### Input:\n{inp}\n\n### Response:\n{out}{EOS_TOKEN}"
-        prompts.append(prompt)
-    return prompts
-
-
-train_texts = format_prompts(train_df)
-valid_texts = format_prompts(valid_df)
-
-# Convert to Hugging Face datasets
-train_dataset = Dataset.from_dict({"text": train_texts})
-valid_dataset = Dataset.from_dict({"text": valid_texts})
-
-# Data collator
-data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
-
-# ================================
-# Training Arguments
-# ================================
-training_args = TrainingArguments(
-    output_dir="/workspace/outputs",
-    per_device_train_batch_size=1,
-    gradient_accumulation_steps=16,
-    evaluation_strategy="steps",
-    eval_steps=100,
-    save_steps=100,
-    save_total_limit=2,
-    logging_steps=10,
-    learning_rate=2e-4,
-    max_steps=500,
-    fp16=True,  # Use mixed precision for efficiency
-    report_to="none",
-)
-
-# ================================
-# Initialize Trainer
-# ================================
-trainer = SFTTrainer(
-    model=model,
-    tokenizer=tokenizer,
-    train_dataset=train_dataset,
-    eval_dataset=valid_dataset,
-    args=training_args,
-    data_collator=data_collator,
-)
-
-# ================================
-# Start Training
-# ================================
-if accelerator.is_local_main_process:
-    print("Starting training...")
-
-trainer.train()
-
-if accelerator.is_local_main_process:
-    print("Training completed successfully.")
+import torch
+from transformers import TextStreamer, TrainingArguments
+from datasets import load_dataset
+from trl import SFTTrainer
+from unsloth import FastLanguageModel, is_bfloat16_supported
+
+# 0.3 Import notebook_launcher from Accelerate
+from accelerate import notebook_launcher
+
+def train():
+    # 1. Configuration
+    max_seq_length = 2048
+    dtype = None
+    load_in_4bit = True
+
+    # Example Alpaca-style prompt template
+    alpaca_prompt = """Below is an instruction that describes a task. 
+    Write a response that appropriately completes the request.
+    ### Instruction:
+    {}
+    ### Input:
+    {}
+    ### Response:
+    {}
+    """
+
+    instruction = """This assistant is trained to code executive ranks ..."""
+    input = "In 2015 the company ..."
+    huggingface_model_name = "daresearch/Llama-3.1-70B-bnb-4bit-Exec-Labeling"
+
+    # 2. Before Training
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name="unsloth/Meta-Llama-3.1-70B-bnb-4bit",
+        max_seq_length=max_seq_length,
+        dtype=dtype,
+        load_in_4bit=load_in_4bit,
+        token=os.getenv("HF_TOKEN"),
+    )
+
+    # Quick inference test (optional)
+    FastLanguageModel.for_inference(model)
+    inputs = tokenizer([alpaca_prompt.format(instruction, input, "")], return_tensors="pt").to("cuda")
+    text_streamer = TextStreamer(tokenizer)
+    _ = model.generate(**inputs, streamer=text_streamer, max_new_tokens=1000)
+
+    # 3. Load and Format Data
+    EOS_TOKEN = tokenizer.eos_token
+    def formatting_prompts_func(examples):
+        instructions = examples["instruction"]
+        inputs = examples["input"]
+        outputs = examples["output"]
+        texts = []
+        for i, inp, out in zip(instructions, inputs, outputs):
+            text = alpaca_prompt.format(i, inp, out) + EOS_TOKEN
+            texts.append(text)
+        return {"text": texts}
+
+    train_dataset = load_dataset("csv", data_files="train.csv", split="train")
+    valid_dataset = load_dataset("csv", data_files="valid.csv", split="train")
+
+    train_dataset = train_dataset.map(formatting_prompts_func, batched=True)
+    valid_dataset = valid_dataset.map(formatting_prompts_func, batched=True)
+
+    # 4. Prepare LoRA Model
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r=16,
+        target_modules=[
+            "q_proj","k_proj","v_proj","o_proj",
+            "gate_proj","up_proj","down_proj",
+        ],
+        lora_alpha=16,
+        lora_dropout=0,
+        bias="none",
+        use_gradient_checkpointing="unsloth",
+        random_state=3407,
+        use_rslora=False,
+        loftq_config=None,
+    )
+
+    # 5. Training
+    trainer = SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=train_dataset,
+        eval_dataset=valid_dataset,
+        dataset_text_field="text",
+        max_seq_length=max_seq_length,
+        dataset_num_proc=2,
+        packing=False,
+        args=TrainingArguments(
+            per_device_train_batch_size=2,
+            gradient_accumulation_steps=4,
+            warmup_steps=5,
+            max_steps=100,
+            learning_rate=2e-4,
+            fp16=not is_bfloat16_supported(),
+            bf16=is_bfloat16_supported(),
+            logging_steps=1,
+            evaluation_strategy="steps",
+            eval_steps=10,
+            optim="adamw_8bit",
+            weight_decay=0.01,
+            lr_scheduler_type="linear",
+            seed=3407,
+            output_dir="outputs",
+        ),
+    )
+
+    trainer_stats = trainer.train()
+    print("Training complete.")
+
+    # 6. Evaluation
+    eval_stats = trainer.evaluate(eval_dataset=valid_dataset)
+    print(f"Validation Loss: {eval_stats['eval_loss']}")
+    if "eval_accuracy" in eval_stats:
+        print(f"Validation Accuracy: {eval_stats['eval_accuracy']}")
+
+    # 7. Saving & Uploading
+    FastLanguageModel.for_inference(model)
+    model.save_pretrained("lora_model")
+    tokenizer.save_pretrained("lora_model")
+    model.push_to_hub(huggingface_model_name, token=os.getenv("HF_TOKEN"))
+    tokenizer.push_to_hub(huggingface_model_name, token=os.getenv("HF_TOKEN"))
+
+    # Merge LoRA weights to create a 16-bit model
+    model.save_pretrained_merged("model", tokenizer, save_method="merged_16bit")
+    model.push_to_hub_merged(
+        huggingface_model_name,
+        tokenizer,
+        save_method="merged_16bit",
+        token=os.getenv("HF_TOKEN"),
+    )
+
+# 0.4 Launch training inside this same script/notebook using multiple GPUs
+notebook_launcher(train, num_processes=4)  # Adjust num_processes for your number of GPUs