training/gpu_distill.py

#!/usr/bin/env python3
"""
Distill Gemini Flash summaries into Qwen3-0.6B.

Fine-tunes Qwen3-0.6B with LoRA to generate one-sentence summaries from
raw markdown text — distilling from 6,720 high-quality Gemini-generated
summaries. At inference time, feed any markdown text and get a summary
back. Runs on CPU for inference (~1-2s per summary).

Input:  raw embedded_text (markdown)
Output: one-sentence summary (Gemini-quality, Qwen-speed)

Usage:
    python3 gpu_distill.py --data-dir /workspace/data --output-dir /workspace/output
"""
import json
import os
import sys
import time
import datetime
import argparse
import math

sys.stdout.reconfigure(line_buffering=True)
sys.stderr.reconfigure(line_buffering=True)


def log(msg, level="INFO"):
    ts = datetime.datetime.now().strftime("%H:%M:%S")
    print(f"[{ts}] [{level}] {msg}", flush=True)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-dir", default="/workspace/data")
    parser.add_argument("--output-dir", default="/workspace/output")
    parser.add_argument("--epochs", type=int, default=5)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--lr", type=float, default=2e-4)
    parser.add_argument("--lora-rank", type=int, default=16)
    parser.add_argument("--lora-alpha", type=int, default=32)
    parser.add_argument("--model-name", default="Qwen/Qwen3-0.6B")
    parser.add_argument("--max-input-len", type=int, default=384, help="Max input tokens")
    parser.add_argument("--max-output-len", type=int, default=64, help="Max output tokens")
    parser.add_argument("--log-every", type=int, default=10)
    parser.add_argument("--sample-every", type=int, default=2)
    args = parser.parse_args()

    log("=" * 60)
    log("DISTILLATION: Markdown → Summary (LoRA fine-tune)")
    log("=" * 60)
    log(f"Config: epochs={args.epochs} batch={args.batch_size} lr={args.lr} "
        f"lora_rank={args.lora_rank} input_len={args.max_input_len} output_len={args.max_output_len}")

    # Auto-install missing deps (don't touch torch — use image's version)
    import subprocess as _sp
    for pkg in ["numpy", "transformers", "accelerate", "safetensors"]:
        try:
            __import__(pkg)
        except ImportError:
            log(f"Installing {pkg}...")
            _sp.run([sys.executable, "-m", "pip", "install", "--break-system-packages",
                     "-q", pkg], check=True)

    import numpy as np
    import torch
    import torch.nn as nn
    from torch.utils.data import Dataset, DataLoader
    from transformers import AutoTokenizer, AutoModelForCausalLM

    log(f"PyTorch {torch.__version__} | CUDA: {torch.cuda.is_available()}")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if device.type == "cuda":
        props = torch.cuda.get_device_properties(0)
        log(f"GPU: {torch.cuda.get_device_name()} | VRAM: {props.total_memory / 1024**3:.1f} GB")

    os.makedirs(args.output_dir, exist_ok=True)
    def vram_mb():
        return torch.cuda.memory_allocated() / 1024**2 if device.type == "cuda" else 0

    metrics = {
        "config": vars(args), "device": str(device),
        "gpu": torch.cuda.get_device_name() if device.type == "cuda" else "cpu",
        "method": "distillation", "steps": [], "epochs": [], "samples": [],
        "start_time": time.time(),
    }

    # ── Load data ──────────────────────────────────────────────────────
    log("Loading data...")
    t0 = time.time()

    # Load texts (embedded_text from clouderic.db) and summaries
    with open(os.path.join(args.data_dir, "texts.json")) as f:
        text_data = json.load(f)  # [{"id": str, "text": str}]
    with open(os.path.join(args.data_dir, "summaries.json")) as f:
        sum_data = json.load(f)  # [{"id": str, "summary": str}]

    sum_map = {s["id"]: s["summary"] for s in sum_data}
    pairs = [(t["text"], sum_map[t["id"]]) for t in text_data
             if t["id"] in sum_map and t["text"] and len(t["text"].strip()) > 20]
    log(f"Loaded {len(pairs)} (text, summary) pairs in {time.time()-t0:.1f}s")

    # Stats
    text_lens = [len(t) for t, _ in pairs]
    sum_lens = [len(s) for _, s in pairs]
    log(f"Text lengths: mean={np.mean(text_lens):.0f} median={np.median(text_lens):.0f} "
        f"max={max(text_lens)} chars")
    log(f"Summary lengths: mean={np.mean(sum_lens):.0f} median={np.median(sum_lens):.0f} "
        f"max={max(sum_lens)} chars")

    # ── Load model ─────────────────────────────────────────────────────
    log(f"Loading {args.model_name}...")
    t0 = time.time()
    tokenizer = AutoTokenizer.from_pretrained(args.model_name, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"  # for decoder-only models

    model = AutoModelForCausalLM.from_pretrained(
        args.model_name, torch_dtype=torch.float16, trust_remote_code=True,
    ).to(device)

    for param in model.parameters():
        param.requires_grad = False

    hidden_dim = model.config.hidden_size
    log(f"Model loaded in {time.time()-t0:.1f}s: hidden={hidden_dim} | VRAM: {vram_mb():.0f}MB")

    # ── LoRA ───────────────────────────────────────────────────────────
    class LoRALayer(nn.Module):
        def __init__(self, original_layer, rank, alpha):
            super().__init__()
            self.original = original_layer
            in_f, out_f = original_layer.in_features, original_layer.out_features
            self.lora_A = nn.Linear(in_f, rank, bias=False)
            self.lora_B = nn.Linear(rank, out_f, bias=False)
            self.scaling = alpha / rank
            nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B.weight)

        def forward(self, x):
            orig_out = self.original(x)
            lora_out = self.lora_B(self.lora_A(x.to(self.lora_A.weight.dtype)))
            return orig_out + lora_out.to(orig_out.dtype) * self.scaling

    lora_modules = []
    n_adapted = 0
    for name, module in model.named_modules():
        if hasattr(module, 'q_proj') and isinstance(module.q_proj, nn.Linear):
            lora_q = LoRALayer(module.q_proj, args.lora_rank, args.lora_alpha).to(device)
            module.q_proj = lora_q
            lora_modules.append(lora_q)
            n_adapted += 1
        if hasattr(module, 'v_proj') and isinstance(module.v_proj, nn.Linear):
            lora_v = LoRALayer(module.v_proj, args.lora_rank, args.lora_alpha).to(device)
            module.v_proj = lora_v
            lora_modules.append(lora_v)
            n_adapted += 1

    lora_params = []
    for lm in lora_modules:
        lora_params.extend(lm.lora_A.parameters())
        lora_params.extend(lm.lora_B.parameters())

    lora_total = sum(p.numel() for p in lora_params)
    log(f"LoRA applied to {n_adapted} layers | {lora_total:,} trainable params | VRAM: {vram_mb():.0f}MB")

    # ── Dataset ────────────────────────────────────────────────────────
    PROMPT_TEMPLATE = "Summarize in one sentence:\n{text}\n\nSummary:"

    class DistillDataset(Dataset):
        def __init__(self, pairs, tokenizer, max_input, max_output):
            self.items = []
            for text, summary in pairs:
                # Truncate text to fit
                prompt = PROMPT_TEMPLATE.format(text=text[:2000])
                # Tokenize prompt and summary separately
                prompt_enc = tokenizer(prompt, truncation=True, max_length=max_input,
                                       return_tensors="pt")
                summary_enc = tokenizer(summary, truncation=True, max_length=max_output,
                                         return_tensors="pt")

                # Concatenate: [prompt_tokens] [summary_tokens] [eos]
                input_ids = torch.cat([
                    prompt_enc["input_ids"].squeeze(0),
                    summary_enc["input_ids"].squeeze(0),
                    torch.tensor([tokenizer.eos_token_id]),
                ])

                # Labels: -100 for prompt, actual ids for summary+eos
                n_prompt = prompt_enc["input_ids"].shape[1]
                labels = input_ids.clone()
                labels[:n_prompt] = -100

                # Truncate total to max_input + max_output
                max_total = max_input + max_output
                if len(input_ids) > max_total:
                    input_ids = input_ids[:max_total]
                    labels = labels[:max_total]

                self.items.append((input_ids, labels))

        def __len__(self):
            return len(self.items)

        def __getitem__(self, idx):
            return self.items[idx]

    def collate_fn(batch):
        input_ids_list, labels_list = zip(*batch)
        max_len = max(ids.shape[0] for ids in input_ids_list)

        input_ids = torch.full((len(batch), max_len), tokenizer.pad_token_id, dtype=torch.long)
        labels = torch.full((len(batch), max_len), -100, dtype=torch.long)
        attention_mask = torch.zeros((len(batch), max_len), dtype=torch.long)

        for i, (ids, lab) in enumerate(zip(input_ids_list, labels_list)):
            # Right-align (pad on left for decoder-only)
            offset = max_len - ids.shape[0]
            input_ids[i, offset:] = ids
            labels[i, offset:] = lab
            attention_mask[i, offset:] = 1

        return input_ids, labels, attention_mask

    # Split
    n_val = max(int(len(pairs) * 0.1), 1)
    rng = np.random.RandomState(42)
    indices = rng.permutation(len(pairs))
    val_pairs = [pairs[i] for i in indices[:n_val]]
    train_pairs = [pairs[i] for i in indices[n_val:]]

    train_ds = DistillDataset(train_pairs, tokenizer, args.max_input_len, args.max_output_len)
    val_ds = DistillDataset(val_pairs, tokenizer, args.max_input_len, args.max_output_len)
    train_dl = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True,
                          drop_last=True, collate_fn=collate_fn)
    val_dl = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False, collate_fn=collate_fn)

    steps_per_epoch = len(train_dl)
    total_steps = steps_per_epoch * args.epochs
    log(f"Data: train={len(train_ds)} val={len(val_ds)} | {steps_per_epoch} steps/epoch, "
        f"{total_steps} total")

    # ── Training ───────────────────────────────────────────────────────
    optimizer = torch.optim.AdamW(lora_params, lr=args.lr, weight_decay=0.01)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps, eta_min=1e-6)
    scaler = torch.amp.GradScaler("cuda") if device.type == "cuda" else None
    best_val_loss = float("inf")
    global_step = 0

    log("")
    log("=" * 60)
    log("TRAINING START")
    log("=" * 60)
    train_start = time.time()

    for epoch in range(args.epochs):
        model.train()
        epoch_loss, epoch_tokens = 0.0, 0
        epoch_start = time.time()
        log(f"")
        log(f"── Epoch {epoch+1}/{args.epochs} ──")

        for step, (input_ids, labels, attn_mask) in enumerate(train_dl):
            step_start = time.time()
            input_ids = input_ids.to(device)
            labels = labels.to(device)
            attn_mask = attn_mask.to(device)

            optimizer.zero_grad()

            if scaler:
                with torch.amp.autocast("cuda"):
                    outputs = model(input_ids=input_ids, attention_mask=attn_mask, labels=labels)
                    loss = outputs.loss
                if torch.isnan(loss):
                    log(f"NaN at step {step+1}!", "ERROR")
                    break
                scaler.scale(loss).backward()
                scaler.unscale_(optimizer)
                grad_norm = torch.nn.utils.clip_grad_norm_(lora_params, 1.0).item()
                scaler.step(optimizer)
                scaler.update()
            else:
                outputs = model(input_ids=input_ids, attention_mask=attn_mask, labels=labels)
                loss = outputs.loss
                loss.backward()
                grad_norm = torch.nn.utils.clip_grad_norm_(lora_params, 1.0).item()
                optimizer.step()

            scheduler.step()

            n_tokens = (labels != -100).sum().item()
            step_time = time.time() - step_start
            tps = n_tokens / step_time if step_time > 0 else 0
            epoch_loss += loss.item() * n_tokens
            epoch_tokens += n_tokens
            global_step += 1

            metrics["steps"].append({
                "epoch": epoch+1, "step": step+1, "global_step": global_step,
                "loss": round(loss.item(), 4), "lr": scheduler.get_last_lr()[0],
                "grad_norm": round(grad_norm, 4), "vram_mb": round(vram_mb()),
                "tokens_per_sec": round(tps),
            })

            if step % args.log_every == 0:
                elapsed = time.time() - train_start
                eta = elapsed / global_step * (total_steps - global_step) if global_step > 0 else 0
                log(f"  step {step+1:>3}/{steps_per_epoch} | loss={loss.item():.4f} | "
                    f"lr={scheduler.get_last_lr()[0]:.1e} | grad={grad_norm:.3f} | "
                    f"VRAM={vram_mb():.0f}MB | {tps:.0f} tok/s | ETA={eta/60:.0f}m")

        if torch.isnan(loss):
            break

        avg_train = epoch_loss / max(epoch_tokens, 1)

        # Validation
        log(f"  Validating...")
        model.eval()
        val_loss, val_tokens = 0.0, 0
        with torch.no_grad():
            for input_ids, labels, attn_mask in val_dl:
                input_ids, labels, attn_mask = input_ids.to(device), labels.to(device), attn_mask.to(device)
                with torch.amp.autocast("cuda") if device.type == "cuda" else torch.no_grad():
                    outputs = model(input_ids=input_ids, attention_mask=attn_mask, labels=labels)
                n = (labels != -100).sum().item()
                val_loss += outputs.loss.item() * n
                val_tokens += n

        avg_val = val_loss / max(val_tokens, 1)
        epoch_time = time.time() - epoch_start
        is_best = avg_val < best_val_loss

        metrics["epochs"].append({
            "epoch": epoch+1, "train_loss": round(avg_train, 4),
            "val_loss": round(avg_val, 4), "time_s": round(epoch_time, 1), "best": is_best,
        })

        marker = " ★ NEW BEST" if is_best else ""
        log(f"  Epoch {epoch+1}/{args.epochs} DONE | train={avg_train:.4f} val={avg_val:.4f} | "
            f"{epoch_time:.0f}s{marker}")

        if device.type == "cuda":
            torch.cuda.empty_cache()

        if is_best:
            best_val_loss = avg_val
            lora_state = {}
            for name, module in model.named_modules():
                if isinstance(module, LoRALayer):
                    lora_state[name + ".lora_A"] = module.lora_A.state_dict()
                    lora_state[name + ".lora_B"] = module.lora_B.state_dict()
            torch.save({
                "epoch": epoch, "val_loss": avg_val,
                "lora_state": lora_state,
                "config": vars(args),
            }, os.path.join(args.output_dir, "best_distill.pt"))

        # Samples
        if (epoch + 1) % args.sample_every == 0 or epoch == args.epochs - 1 or is_best:
            try:
                log(f"  Generating samples...")
                model.eval()
                sample_rng = np.random.RandomState(epoch)
                sample_idx = sample_rng.choice(len(val_pairs), size=min(3, len(val_pairs)), replace=False)

                for si in sample_idx:
                    text, ref = val_pairs[si]
                    prompt = PROMPT_TEMPLATE.format(text=text[:1500])
                    inputs = tokenizer(prompt, return_tensors="pt", truncation=True,
                                       max_length=args.max_input_len).to(device)

                    with torch.no_grad():
                        gen = model.generate(
                            **inputs, max_new_tokens=args.max_output_len,
                            do_sample=False, temperature=1.0,
                            pad_token_id=tokenizer.pad_token_id,
                        )
                    gen_text = tokenizer.decode(gen[0][inputs["input_ids"].shape[1]:],
                                                skip_special_tokens=True)

                    del gen
                    if device.type == "cuda":
                        torch.cuda.empty_cache()

                    metrics["samples"].append({"epoch": epoch+1, "ref": ref[:200], "gen": gen_text[:200]})
                    log(f"    REF: {ref[:100]}")
                    log(f"    GEN: {gen_text[:100]}")
                    log(f"")
            except Exception as e:
                log(f"  Sample generation failed: {e}", "WARN")

            if device.type == "cuda":
                torch.cuda.empty_cache()

    # ── Summary ────────────────────────────────────────────────────────
    total_time = time.time() - train_start
    metrics["total_time_s"] = round(total_time, 1)
    metrics["best_val_loss"] = round(best_val_loss, 4)

    with open(os.path.join(args.output_dir, "training_metrics.json"), "w") as f:
        json.dump(metrics, f, indent=2)

    log("")
    log("=" * 60)
    log("TRAINING COMPLETE")
    log("=" * 60)
    log(f"Total time: {total_time/60:.1f} minutes")
    log(f"Best val loss: {best_val_loss:.4f}")
    log(f"")
    log("Epoch | Train Loss | Val Loss | Time | Best")
    log("-" * 50)
    for e in metrics["epochs"]:
        m = " ★" if e["best"] else ""
        log(f"  {e['epoch']:>3}  |   {e['train_loss']:.4f}  |  {e['val_loss']:.4f} | {e['time_s']:.0f}s{m}")


if __name__ == "__main__":
    main()