init

README.md

@@ -0,0 +1,52 @@
+# Tulu Laptop Finetune + W&B
+
+Minimal setup to finetune a laptop-friendly Tulu checkpoint with QLoRA and track runs in Weights & Biases.
+
+## Prereqs
+- Recent NVIDIA GPU with CUDA for 4-bit (bitsandbytes) set `--use_4bit true`. On CPU/MPS (default), set `--use_4bit false`, but expect much slower/limited runs.
+- Conda (Miniconda/Anaconda).
+- A Weights & Biases account + API key.
+
+## Setup
+1) Create the env (Conda)
+```bash
+conda env create -f environment.yml
+conda activate tulu-train
+```
+2) Add secrets (keep `.env` out of git)
+```bash
+cp .env.example .env
+# Edit .env with your WANDB_API_KEY / project / entity
+```
+3) Verify packages (optional if you prefer pip)
+```bash
+pip install -r requirements.txt
+```
+
+## Run a quick finetune
+The defaults use `allenai/tulu-2-7b` with a small instruction dataset (`mlabonne/guanaco-llama2-1k`) and 4-bit QLoRA. This keeps memory needs closer to laptop GPUs.
+```bash
+python train_tulu.py \
+  --output_dir outputs/tulu-lora \
+  --max_seq_length 512 \
+  --per_device_batch_size 1 \
+  --gradient_accumulation_steps 16
+```
+
+Key flags:
+- `--use_4bit false` if bitsandbytes/CUDA are unavailable (will be slower and need more RAM).
+- `--dataset_name` to try another instruction set (any HF dataset with `instruction/input/output` fields).
+- `--model_name` if you want a different Tulu variant (e.g., `allenai/tulu-2-dpo-7b`).
+
+## How W&B is used
+- `train_tulu.py` loads `.env`, logs into W&B, and reports through `Trainer(report_to=["wandb"])`.
+- Ensure `WANDB_API_KEY`, `WANDB_PROJECT`, and (optionally) `WANDB_ENTITY` are set in `.env`.
+- Each run captures hyperparameters and metrics; check the W&B UI for live loss curves and checkpoints.
+
+## Output
+- Finetuned adapters + tokenizer are written to `outputs/tulu-lora` (configurable via `--output_dir`). Push this to the Hub with `huggingface-cli upload` if desired.
+
+## Troubleshooting
+- OOM? Reduce `max_seq_length`, increase `gradient_accumulation_steps`, or switch to a smaller dataset.
+- bitsandbytes import errors on macOS/CPU: run with `--use_4bit false` or use a Linux+CUDA machine.
+- bitsandbytes install error? We pin to `0.42.0`, the latest widely distributed wheel. If you cannot install it (CPU-only/MPS), remove it from `requirements.txt` and set `--use_4bit false`.

environment.yml

@@ -0,0 +1,27 @@
+name: deeai
+channels:
+  # Use conda-forge for up-to-date builds of Python and libs.
+  - conda-forge
+dependencies:
+  # Base interpreter; Python 3.10 has broad wheel support across ML libs.
+  - python=3.10
+  # Core tooling and a clean pip inside the env.
+  - pip
+  - pip:
+      # Core model + tokenizer stack.
+      - transformers>=4.44
+      - datasets>=2.19
+      # Parameter-efficient finetuning (LoRA).
+      - peft>=0.11
+      # Multi-GPU/accelerator launcher + config helper.
+      - accelerate>=0.33
+      # 4-bit quantization backend for laptop-friendly training (CUDA required).
+      # 0.42 is the latest widely available pip release.
+      - bitsandbytes==0.42.0
+      # Logging + experiment tracking.
+      - wandb>=0.17
+      # Env loader so secrets stay in .env, not code.
+      - python-dotenv>=1.0
+      # Optional: small utilities.
+      - tqdm>=4.66
+      - scipy>=1.11

requirements.txt

@@ -0,0 +1,12 @@
+# Core model stack
+transformers>=4.44
+datasets>=2.19
+peft>=0.11
+accelerate>=0.33
+bitsandbytes==0.42.0  # CUDA-only; required for 4-bit QLoRA
+
+# Tracking and utilities
+wandb>=0.17
+python-dotenv>=1.0
+tqdm>=4.66
+scipy>=1.11

train_tulu.py

@@ -0,0 +1,178 @@
+"""
+Minimal QLoRA finetune for a laptop-friendly Tulu checkpoint with W&B logging.
+
+Defaults aim to run on a single consumer GPU using 4-bit quantization.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from dataclasses import dataclass
+from typing import Dict, List
+
+import torch
+import wandb
+from datasets import load_dataset
+from dotenv import load_dotenv
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    DataCollatorForLanguageModeling,
+    Trainer,
+    TrainingArguments,
+)
+
+
+@dataclass
+class ScriptConfig:
+    model_name: str = "allenai/tulu-2-7b"
+    dataset_name: str = "mlabonne/guanaco-llama2-1k"  # small, instruction-style set
+    output_dir: str = "outputs/tulu-lora"
+    max_seq_length: int = 512
+    per_device_batch_size: int = 1
+    gradient_accumulation_steps: int = 16
+    num_train_epochs: int = 1
+    learning_rate: float = 2e-4
+    warmup_ratio: float = 0.03
+    logging_steps: int = 10
+    save_steps: int = 200
+    use_4bit: bool = True
+
+
+def format_chat(example: Dict[str, str]) -> str:
+    """Simple instruction->response template that fits Tulu-style tuning."""
+    user_input = example.get("input") or "N/A"
+    return (
+        f"### Instruction:\n{example['instruction']}\n\n"
+        f"### Input:\n{user_input}\n\n"
+        f"### Response:\n{example['output']}"
+    )
+
+
+def tokenize_example(example: Dict[str, str], tokenizer, max_seq_length: int):
+    prompt = format_chat(example)
+    # We build labels that are the same as input_ids for causal LM.
+    tokenized = tokenizer(
+        prompt,
+        truncation=True,
+        max_length=max_seq_length,
+        padding="max_length",
+    )
+    tokenized["labels"] = tokenized["input_ids"].copy()
+    return tokenized
+
+
+def load_model_and_tokenizer(cfg: ScriptConfig):
+    quantization_config = None
+    if cfg.use_4bit:
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.bfloat16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+        )
+
+    tokenizer = AutoTokenizer.from_pretrained(cfg.model_name, use_fast=False)
+    tokenizer.padding_side = "right"
+    tokenizer.pad_token = tokenizer.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(
+        cfg.model_name,
+        quantization_config=quantization_config,
+        device_map="auto",
+    )
+    if cfg.use_4bit:
+        model = prepare_model_for_kbit_training(model)
+
+    lora_cfg = LoraConfig(
+        r=64,
+        lora_alpha=16,
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+        lora_dropout=0.05,
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+    model = get_peft_model(model, lora_cfg)
+    return model, tokenizer
+
+
+def init_wandb(cfg: ScriptConfig):
+    project = os.getenv("WANDB_PROJECT", "tulu-laptop-run")
+    entity = os.getenv("WANDB_ENTITY")
+    api_key = os.getenv("WANDB_API_KEY")
+    if not api_key:
+        raise RuntimeError("WANDB_API_KEY is missing. Put it in your .env before running.")
+    wandb.login(key=api_key)
+    wandb.init(project=project, entity=entity, config=vars(cfg))
+
+
+def parse_args() -> ScriptConfig:
+    parser = argparse.ArgumentParser(description="Finetune Tulu with QLoRA + W&B")
+    parser.add_argument("--model_name", default=ScriptConfig.model_name)
+    parser.add_argument("--dataset_name", default=ScriptConfig.dataset_name)
+    parser.add_argument("--output_dir", default=ScriptConfig.output_dir)
+    parser.add_argument("--max_seq_length", type=int, default=ScriptConfig.max_seq_length)
+    parser.add_argument("--per_device_batch_size", type=int, default=ScriptConfig.per_device_batch_size)
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=ScriptConfig.gradient_accumulation_steps)
+    parser.add_argument("--num_train_epochs", type=float, default=ScriptConfig.num_train_epochs)
+    parser.add_argument("--learning_rate", type=float, default=ScriptConfig.learning_rate)
+    parser.add_argument("--warmup_ratio", type=float, default=ScriptConfig.warmup_ratio)
+    parser.add_argument("--logging_steps", type=int, default=ScriptConfig.logging_steps)
+    parser.add_argument("--save_steps", type=int, default=ScriptConfig.save_steps)
+    parser.add_argument("--use_4bit", action=argparse.BooleanOptionalAction, default=False)
+    args = parser.parse_args()
+    return ScriptConfig(**vars(args))
+
+
+def main():
+    load_dotenv()
+    cfg = parse_args()
+
+    init_wandb(cfg)
+    model, tokenizer = load_model_and_tokenizer(cfg)
+
+    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
+    use_fp16 = torch.cuda.is_available() and not use_bf16
+
+    raw_dataset = load_dataset(cfg.dataset_name)
+    tokenized = raw_dataset["train"].map(
+        lambda ex: tokenize_example(ex, tokenizer, cfg.max_seq_length),
+        remove_columns=raw_dataset["train"].column_names,
+    )
+
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+    training_args = TrainingArguments(
+        output_dir=cfg.output_dir,
+        per_device_train_batch_size=cfg.per_device_batch_size,
+        gradient_accumulation_steps=cfg.gradient_accumulation_steps,
+        num_train_epochs=cfg.num_train_epochs,
+        learning_rate=cfg.learning_rate,
+        warmup_ratio=cfg.warmup_ratio,
+        logging_steps=cfg.logging_steps,
+        save_steps=cfg.save_steps,
+        bf16=use_bf16,
+        fp16=use_fp16,
+        report_to=["wandb"],
+        optim="paged_adamw_32bit",
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+    )
+
+    trainer.train()
+    trainer.save_model(cfg.output_dir)
+    tokenizer.save_pretrained(cfg.output_dir)
+    wandb.finish()
+
+
+if __name__ == "__main__":
+    main()