Re-apply unsloth following official guidiance

train_with_unsloth.py

@@ -1,204 +1,131 @@
 #! /usr/bin/env python3
-import unsloth
-from transformers import AutoModelForCausalLM
-from transformers import AutoTokenizer
+"""
+Fine-tune “SmolLM2-135M-Instruct” on the TaiwanChat dataset using Unsloth’s 4-bit quantization
++ LoRA adapters, with evaluation on a 1% hold-out every step, and push the merged model to Hugging Face.
+
+Steps:
+1. Load a 4-bit quantized base model via Unsloth’s FastLanguageModel.
+2. Attach LoRA adapters (r=16) and enable gradient checkpointing for memory savings.
+3. Load TaiwanChat, render ChatML, and split 99/1 train/validation.
+4. Configure SFTTrainer to mask user prompts (train_on_responses_only), run eval every step, log to W&B.
+5. Train for up to 60 steps.
+6. Merge base+LoRA weights into 16-bit safetensors and push to Hugging Face with `push_to_hub_merged`.
+"""
+
 from unsloth import FastLanguageModel
+from trl import SFTTrainer, SFTConfig
 from transformers import DataCollatorForLanguageModeling
-from transformers import TrainingArguments, Trainer
-from transformers import pipeline
+from unsloth.chat_templates import train_on_responses_only
 from datasets import load_dataset
-import torch
 import os
-import wandb
 from transformers.integrations import WandbCallback
-import math
-from transformers import EvalPrediction
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 PROJECT_NAME='SmolLM2-135M-Instruct-TaiwanChat'
-BASE_MODEL_ID="HuggingFaceTB/SmolLM2-135M-Instruct"
+BASE_MODEL_ID="unsloth/SmolLM2-135M-Instruct"
 DATASET_ID="yentinglin/TaiwanChat"
 N_SAMPLES=80000
-MAX_LEN=512
+MAX_LEN=2048
 
 # Tell wandb which project to use, and that you want to log your model
 os.environ["WANDB_PROJECT"]    = f"{PROJECT_NAME}_CLOUD"
 os.environ["WANDB_LOG_MODEL"]  = "end"
 
-# Detect GPU Type
-device_str='cpu'
-if torch.xpu.is_available():
-    device_str='xpu'
-elif torch.cuda.is_available():
-    device_str='cuda'
-print(f'Device is {device_str}')
-
 ## Load with Unsloth’s optimized API
 # 1) Load quantized model
 model, tokenizer = FastLanguageModel.from_pretrained(
     model_name     = BASE_MODEL_ID,
     max_seq_length = MAX_LEN,
-    dtype          = torch.float16,
     load_in_4bit   = True,
     full_finetuning= False,  # we will add LoRA adapters next
 )
 
 # 2) Prepare it for k‑bit training (sets up layer norms, disables caching, etc.)
-from peft import prepare_model_for_kbit_training
-model = prepare_model_for_kbit_training(model)  # :contentReference[oaicite:0]{index=0}
-
-# 3) Attach LoRA adapters on top of the quantized weights
-from peft import LoraConfig, get_peft_model, TaskType
-
-lora_config = LoraConfig(
-    r               = 8,                     # low‑rank dimension
-    lora_alpha      = 16,                    # scaling
-    target_modules  = ["q_proj", "v_proj"],  # apply to attention
-    bias            = "none",
-    task_type       = TaskType.CAUSAL_LM,
-    inference_mode  = False,
+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, # Supports any, but = 0 is optimized
+    bias = "none",    # Supports any, but = "none" is optimized
+    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
+    use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
+    random_state = 3407,
+    max_seq_length = MAX_LEN,
+    use_rslora = False,  # We support rank stabilized LoRA
+    loftq_config = None, # And LoftQ
 )
-model = get_peft_model(model, lora_config)  # :contentReference[oaicite:1]{index=1}
-
-# Now `model` has ~1–2% trainable parameters (the LoRA adapters), 
-# and Trainer will no longer throw the “purely quantized” error.
 
 # Prepare the TaiwanChat Dataset
 # 1) Load & split
-full_ds = load_dataset(DATASET_ID, split=f"train[:{N_SAMPLES}]")
-splits  = full_ds.train_test_split(test_size=0.1, seed=42)  # :contentReference[oaicite:2]{index=2}
-train_ds = splits["train"]
-val_ds   = splits["test"]
-
-# Preprocessing Function
-def preprocess_examples(examples):
-    chats = examples["messages"]
-    # 1) Render ChatML
-    text = tokenizer.apply_chat_template(
-        chats, tokenize=False, add_generation_prompt=True
-    )
-
-    # 2) Tokenize _and_ pad/truncate to MAX_LEN
-    toks = tokenizer(
-        text,
-        truncation=True,
-        padding="max_length",
-        max_length=MAX_LEN,
+dataset = load_dataset(DATASET_ID, split=f"train[:{N_SAMPLES}]")
+
+# turn list-of-messages → a single “text” string per example, using Unsloth’s ChatML template
+def fmt(examples):
+    texts = [
+        tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+        for msgs in examples["messages"]
+    ]
+    return {"text": texts}
+
+dataset = dataset.map(fmt, batched=True, remove_columns=["messages"])
+new_dataset = dataset.train_test_split(test_size = 0.01)
+
+training_args = SFTConfig(
+        fp16_full_eval = True,
+        per_device_eval_batch_size = 2,
+        eval_accumulation_steps = 4,
+        eval_strategy = "steps",
+        eval_steps = 1,
+        dataset_text_field="text",
+        output_dir=PROJECT_NAME,
+        max_seq_length = MAX_LEN,
+        per_device_train_batch_size = 2,
+        gradient_accumulation_steps = 4,
+        warmup_steps = 10,
+        max_steps = 60,
+        logging_steps = 1,
+        optim = "adamw_8bit",
+        seed = 3407,
+        # ─── W&B integration ───
+        logging_dir=f"{PROJECT_NAME}/logs",    # where to store TensorBoard/W&B logs
+        report_to=["wandb"],                   # enable W&B reporting
+        run_name=f"{PROJECT_NAME}_CLOUD",                 # name this run in your W&B project
+        push_to_hub=True,
+        gradient_checkpointing=True
     )
-    input_ids     = toks["input_ids"]
-    attention_mask= toks["attention_mask"]
-
-    # 3) Find where the assistant reply starts
-    role_id = tokenizer.convert_tokens_to_ids("<|im_start|>assistant")
-    if role_id in input_ids:
-        idx = input_ids.index(role_id)
-        start_of_reply = idx + 2
-    else:
-        start_of_reply = 0
-
-    # 4) Build labels: -100 before reply, then copy the rest
-    labels = [-100] * start_of_reply + input_ids[start_of_reply:]
-
-    # 5) Pad or truncate labels to EXACTLY len(input_ids)
-    if len(labels) < len(input_ids):
-        labels += [-100] * (len(input_ids) - len(labels))
-    else:
-        labels = labels[: len(input_ids)]
-
-    return {
-        "input_ids": input_ids,
-        "attention_mask": attention_mask,
-        "labels": labels,
-    }
-# Tokenization & Data Collator
-tokenized_train = train_ds.map(
-    preprocess_examples, batched=True, remove_columns=train_ds.column_names
-)
-tokenized_val = val_ds.map(
-    preprocess_examples, batched=True, remove_columns=val_ds.column_names
-)
-
-data_collator = DataCollatorForLanguageModeling(
-    tokenizer=tokenizer, mlm=False
-)
-
-training_args = TrainingArguments(
-    output_dir=PROJECT_NAME,
-    per_device_train_batch_size=2,
-    gradient_accumulation_steps    = 16,
-    learning_rate=5e-5,
-    num_train_epochs=3,
-    fp16=False if device_str == 'xpu' else True,
-    bf16=True if device_str == 'xpu' else False,
-    #evaluation_strategy    = "steps",     # run validation every eval_steps
-    #eval_steps             = 1000,
-    #load_best_model_at_end = True,
-    #metric_for_best_model  = "perplexity",
-    greater_is_better      = False,
-    logging_steps=1000,
-    save_steps=5000,
-
-    # ─── W&B integration ───
-    logging_dir=f"{PROJECT_NAME}/logs",    # where to store TensorBoard/W&B logs
-    report_to=["wandb"],                   # enable W&B reporting
-    run_name=f"{PROJECT_NAME}_CLOUD",                 # name this run in your W&B project
-
-    push_to_hub=True,
-    gradient_checkpointing=True,
-)
-
-# Enable gradient checkpointing on the model
-model.gradient_checkpointing_enable()
-
-# Define a metrics function
-def compute_metrics(p: EvalPrediction):
-    # p.predictions are logits: (batch, seq_len, vocab_size)
-    # p.label_ids are (batch, seq_len)
-    # The Trainer will automatically compute loss on eval_dataset
-    # We can pull that from p.metrics if available, 
-    # but simplest is to use returned "eval_loss" in Trainer.evaluate()
-    # Here we compute perplexity manually:
-    eval_loss = p.metrics["eval_loss"] if "eval_loss" in p.metrics else None
-    if eval_loss is None:
-        raise ValueError("eval_loss not found in metrics; ensure compute_metrics is called after evaluation.")
-    return {"perplexity": math.exp(eval_loss)}
-
 
 # Training with Trainer
-trainer = Trainer(
+trainer = SFTTrainer(
     model=model,
     args=training_args,
-    train_dataset=tokenized_train,
-    #eval_dataset=tokenized_val,
-    compute_metrics=compute_metrics,
-    data_collator=data_collator,
+    data_collator = DataCollatorForLanguageModeling(tokenizer = tokenizer, mlm=False),
+    tokenizer=tokenizer,
     callbacks=[WandbCallback],  # ensure the W&B callback is attached
+    train_dataset = new_dataset["train"],
+    eval_dataset = new_dataset["test"],
+)
+trainer = train_on_responses_only(trainer)
+trainer.train()
+
+model.push_to_hub_merged(
+    f'Luigi/{PROJECT_NAME}', 
+    tokenizer, 
+    save_method="merged_16bit", 
+    safe_serialization=None
 )
-trainer.train(resume_from_checkpoint=True)
 
-# Save Model & Tokenizer Locally
-trainer.save_model(PROJECT_NAME)
-trainer.push_to_hub(f'Luigi/{PROJECT_NAME}')
-tokenizer.save_pretrained(PROJECT_NAME)
 
-# 1) Load from local folder
-model_dir = PROJECT_NAME
-tokenizer = AutoTokenizer.from_pretrained(model_dir)
-model     = AutoModelForCausalLM.from_pretrained(model_dir)  # loads your fine‑tuned weights :contentReference[oaicite:2]{index=2}
 
-# ─── patch for Uns­loth’s fast-forward ───
-# Uns­loth expects a `max_seq_length` attribute on the model
-setattr(model, "max_seq_length", MAX_LEN)
+# 1. load merged model + tokenizer from your HF repo
+tokenizer = AutoTokenizer.from_pretrained(f'Luigi/{PROJECT_NAME}')  
+model     = AutoModelForCausalLM.from_pretrained(f'Luigi/{PROJECT_NAME}')  
 
-# Test Fine-tuned Model
-hf_device = 0 if device_str in ("cuda","xpu") else -1
+# 2. run text-generation
 gen = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    device=hf_device,                # or device=0 for GPU
-    max_new_tokens=512,        # customize as desired
+  "text-generation", model=model, tokenizer=tokenizer,
+  device_map="auto",    # or device=0 for a single GPU
 )
-
 prompt = "請問台北今天的天氣如何？"
-output = gen(prompt, do_sample=True, temperature=0.8)
-print(output[0]["generated_text"])
+print(gen(prompt, max_new_tokens=MAX_LEN)[0]["generated_text"])