Create sft.py

sft.py

@@ -0,0 +1,249 @@
+"""
+© SupraLabs 2026 - SFT script for Supra-50M on alpaca-cleaned
+No TRL. Uses HuggingFace Trainer with prompt-masked cross-entropy loss.
+"""
+
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+print("[*] Loading libraries...")
+import torch
+import numpy as np
+from dataclasses import dataclass
+from typing import Optional
+from datasets import load_dataset
+from transformers import (
+    AutoModelForCausalLM,
+    Trainer,
+    TrainingArguments,
+    PreTrainedTokenizerBase,
+    PreTrainedTokenizerFast
+)
+from torch.utils.data import Dataset
+
+# ── Config ────────────────────────────────────────────────────────────────────
+
+MODEL_ID = "./Chimera-FINAL"
+OUTPUT_DIR      = "./Supra-50M-SFT"
+MAX_LENGTH      = 512       # alpaca samples are short, 512 is plenty
+IGNORE_INDEX    = -100      # standard label mask value for cross-entropy
+
+# Conservative hyperparameters — small model, don't nuke the pretraining
+LEARNING_RATE   = 3e-4
+EPOCHS          = 4
+BATCH_SIZE      = 8
+GRAD_ACCUM      = 2         # effective batch size = 16
+WARMUP_RATIO    = 0.1
+WEIGHT_DECAY    = 0.0
+MAX_GRAD_NORM   = 1.0
+
+# ── Alpaca prompt template ────────────────────────────────────────────────────
+
+PROMPT_WITH_INPUT = (
+    "Below is an instruction that describes a task, paired with an input "
+    "that provides further context. Write a response that appropriately "
+    "completes the request.\n\n"
+    "### Instruction:\n{instruction}\n\n"
+    "### Input:\n{input}\n\n"
+    "### Response:\n"
+)
+
+PROMPT_WITHOUT_INPUT = (
+    "Below is an instruction that describes a task. Write a response that "
+    "appropriately completes the request.\n\n"
+    "### Instruction:\n{instruction}\n\n"
+    "### Response:\n"
+)
+
+def build_prompt(sample: dict) -> tuple[str, str]:
+    """Returns (prompt, response) — kept separate so we can mask the prompt."""
+    instruction = sample["instruction"].strip()
+    inp         = sample.get("input", "").strip()
+    output      = sample["output"].strip()
+
+    if inp:
+        prompt = PROMPT_WITH_INPUT.format(instruction=instruction, input=inp)
+    else:
+        prompt = PROMPT_WITHOUT_INPUT.format(instruction=instruction)
+
+    return prompt, output
+
+
+# ── Dataset ───────────────────────────────────────────────────────────────────
+
+class AlpacaDataset(Dataset):
+    """
+    Tokenizes each sample and masks the prompt portion of the labels so the
+    model only computes loss on the response tokens — not on the instruction.
+    """
+
+    def __init__(self, hf_dataset, tokenizer: PreTrainedTokenizerBase, max_length: int):
+        self.tokenizer  = tokenizer
+        self.max_length = max_length
+        self.samples    = hf_dataset
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        prompt, response = build_prompt(self.samples[idx])
+
+        # Tokenize prompt and response separately so we know the prompt length
+        prompt_ids = self.tokenizer.encode(prompt, add_special_tokens=False)
+        prompt_ids = [self.tokenizer.bos_token_id] + prompt_ids   # explizit
+        response_ids = self.tokenizer.encode(response, add_special_tokens=False) + [self.tokenizer.eos_token_id]
+
+        input_ids = prompt_ids + response_ids
+
+        # Truncate to max_length
+        input_ids = input_ids[:self.max_length]
+
+        # Labels: mask prompt tokens with IGNORE_INDEX
+        prompt_len = min(len(prompt_ids), len(input_ids))
+        labels     = [IGNORE_INDEX] * prompt_len + input_ids[prompt_len:]
+
+        # Sanity: both must be the same length after truncation
+        assert len(input_ids) == len(labels)
+
+        return {
+            "input_ids":      torch.tensor(input_ids, dtype=torch.long),
+            "labels":         torch.tensor(labels,    dtype=torch.long),
+        }
+
+
+# ── Collator ──────────────────────────────────────────────────────────────────
+
+@dataclass
+class PaddingCollator:
+    """
+    Right-pads input_ids and labels to the longest sequence in the batch.
+    Labels are padded with IGNORE_INDEX so padding never contributes to loss.
+    """
+    tokenizer:  PreTrainedTokenizerBase
+    max_length: int
+
+    def __call__(self, batch):
+        max_len = max(len(x["input_ids"]) for x in batch)
+        max_len = min(max_len, self.max_length)
+
+        input_ids_padded = []
+        labels_padded    = []
+        attention_masks  = []
+
+        for item in batch:
+            ids    = item["input_ids"][:max_len]
+            lbls   = item["labels"][:max_len]
+            pad_n  = max_len - len(ids)
+
+            input_ids_padded.append(
+                torch.cat([ids, torch.full((pad_n,), self.tokenizer.pad_token_id, dtype=torch.long)])
+            )
+            labels_padded.append(
+                torch.cat([lbls, torch.full((pad_n,), IGNORE_INDEX, dtype=torch.long)])
+            )
+            attention_masks.append(
+                torch.cat([torch.ones(len(ids), dtype=torch.long),
+                           torch.zeros(pad_n, dtype=torch.long)])
+            )
+
+        return {
+            "input_ids":      torch.stack(input_ids_padded),
+            "labels":         torch.stack(labels_padded),
+            "attention_mask": torch.stack(attention_masks),
+        }
+
+
+# ── Main ──────────────────────────────────────────────────────────────────────
+
+def main():
+    # Load tokenizer + model from Hub
+    print(f"[*] Loading tokenizer from {MODEL_ID}...")
+    from tokenizers import ByteLevelBPETokenizer
+
+    fast_tokenizer = ByteLevelBPETokenizer(
+        "custom_llama_tokenizer-vocab.json",
+        "custom_llama_tokenizer-merges.txt"
+    )
+    tokenizer = PreTrainedTokenizerFast(
+        tokenizer_object=fast_tokenizer,
+        bos_token="<s>",
+        eos_token="</s>",
+        unk_token="<unk>",
+        pad_token="<pad>",
+    )
+
+    print(f"[*] Loading model from {MODEL_ID}...")
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        dtype=torch.bfloat16,
+        device_map="auto",
+    )
+
+    print(f"[+] Model loaded — {model.num_parameters():,} parameters")
+
+    # Load alpaca-cleaned (≈52k instruction-tuning pairs)
+    print("[*] Loading alpaca-cleaned dataset...")
+    raw = load_dataset("yahma/alpaca-cleaned", split="train")
+    print(f"[+] Dataset: {len(raw):,} samples")
+
+    # Optional: quick sanity-check split (comment out for full training)
+    # raw = raw.select(range(1000))
+
+    split = raw.train_test_split(test_size=0.01, seed=42)
+    train_dataset = AlpacaDataset(split["train"], tokenizer, MAX_LENGTH)
+    eval_dataset  = AlpacaDataset(split["test"],  tokenizer, MAX_LENGTH)
+    collator      = PaddingCollator(tokenizer=tokenizer, max_length=MAX_LENGTH)
+
+    print(f"[+] Dataset ready: {len(train_dataset):,} samples")
+    print(f"[+] Example prompt preview:\n{build_prompt(raw[0])[0][:800]}...")
+
+    # Training arguments
+    training_args = TrainingArguments(
+        output_dir=OUTPUT_DIR,
+        num_train_epochs=EPOCHS,
+        per_device_train_batch_size=BATCH_SIZE,
+        gradient_accumulation_steps=GRAD_ACCUM,
+        learning_rate=LEARNING_RATE,
+        lr_scheduler_type="cosine",
+        warmup_ratio=WARMUP_RATIO,
+        weight_decay=WEIGHT_DECAY,
+        max_grad_norm=MAX_GRAD_NORM,
+        bf16=True,
+        fp16=False,
+        logging_steps=50,
+        save_total_limit=2,
+        report_to="none",
+        dataloader_num_workers=8,
+        dataloader_pin_memory=True,
+        optim="adamw_torch_fused",
+        adam_beta1=0.9,
+        adam_beta2=0.999,
+        push_to_hub=False,
+        seed=42,
+        data_seed=42,
+        eval_strategy="epoch",
+        save_strategy="epoch",
+        load_best_model_at_end=True,
+        metric_for_best_model="eval_loss",
+        greater_is_better=False,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        data_collator=collator,
+    )
+
+    print("[*] Starting SFT...")
+    trainer.train()
+
+    print(f"[*] Saving final model to {OUTPUT_DIR}-FINAL ...")
+    trainer.save_model(f"{OUTPUT_DIR}-FINAL")
+    tokenizer.save_pretrained(f"{OUTPUT_DIR}-FINAL")
+    print("[+] Done.")
+
+
+if __name__ == "__main__":
+    main()