FigQuant GPU training test (with dtype fix)

figquant_gpu_fixed.py

@@ -0,0 +1,92 @@
+#!/usr/bin/env python3
+"""FigQuant training on GPU with the dtype fix applied."""
+import os, sys, subprocess, time, gc
+import numpy as np
+
+subprocess.check_call([sys.executable, "-m", "pip", "install", "-q",
+    "transformers", "accelerate", "datasets", "sentencepiece", "protobuf", "psutil", "numpy"])
+subprocess.check_call(["git", "clone", "https://github.com/ticketguy/littlefig.git", "/app/littlefig"])
+subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-e", "/app/littlefig[train]"])
+sys.path.insert(0, "/app/littlefig/src")
+
+import torch
+
+def log(msg): print(f"[GPU] {msg}", flush=True)
+
+log(f"PyTorch {torch.__version__}, CUDA={torch.cuda.is_available()}")
+if torch.cuda.is_available():
+    log(f"GPU: {torch.cuda.get_device_name()} ({torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB)")
+
+from little_fig.engine import FigModel
+from little_fig.engine.tier import TrainingTier
+from datasets import load_dataset
+from torch.utils.data import DataLoader
+
+MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+ds = load_dataset("tatsu-lab/alpaca", split="train").select(range(1000))
+log(f"Data: {len(ds)} examples")
+
+log("Loading FigQuant (lowram mode)...")
+gc.collect(); torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats()
+
+model = FigModel.from_pretrained(MODEL, lora_r=16, lora_alpha=32,
+    tier=TrainingTier.STREAMING_LORA, target_modules=["q_proj","k_proj","v_proj","o_proj"],
+    fast=False)  # lowram mode
+tok = model.tokenizer
+
+examples = [dict(r) for r in ds]
+def tok_fn(ex):
+    inst=ex.get("instruction",""); inp=ex.get("input","").strip(); out=ex.get("output","")
+    txt = f"### Instruction:\n{inst}\n\n### Input:\n{inp}\n\n### Response:\n{out}" if inp else \
+          f"### Instruction:\n{inst}\n\n### Response:\n{out}"
+    e = tok(txt, truncation=True, max_length=512, padding="max_length")
+    return {"input_ids": e["input_ids"], "labels": e["input_ids"].copy(), "attention_mask": e["attention_mask"]}
+tokenized = [tok_fn(ex) for ex in examples]
+
+class DS(torch.utils.data.Dataset):
+    def __init__(s, d): s.d = d
+    def __len__(s): return len(s.d)
+    def __getitem__(s, i): return {k: torch.tensor(v, dtype=torch.long) for k, v in s.d[i].items()}
+
+dl = DataLoader(DS(tokenized), batch_size=4, shuffle=True, drop_last=True)
+dev = torch.device("cuda"); model = model.to(dev)
+params = model.get_trainable_parameters()
+opt = torch.optim.AdamW(params, lr=2e-4, weight_decay=0.01)
+model.model.train()
+
+losses = []; gs = 0; al = 0.0
+torch.cuda.reset_peak_memory_stats()
+t0 = time.time()
+
+for batch in dl:
+    if gs >= 400: break  # 100 optimizer steps × 4 grad accum
+    batch = {k: v.to(dev) for k, v in batch.items()}
+    with torch.autocast("cuda", dtype=torch.float16):
+        loss = model(input_ids=batch["input_ids"], attention_mask=batch["attention_mask"],
+                     labels=batch["labels"]).loss / 4
+    loss.backward()
+    al += loss.item(); gs += 1
+    if gs % 4 == 0:
+        torch.nn.utils.clip_grad_norm_(params, 1.0)
+        opt.step(); opt.zero_grad()
+        s = gs // 4; losses.append(al); al = 0.0
+        if s % 20 == 0: log(f"  step={s} loss={losses[-1]:.4f}")
+
+tt = time.time() - t0
+peak = torch.cuda.max_memory_allocated() / 1e6
+
+log(f"\n{'='*50}")
+log(f"  FigQuant LoRA (lowram) on GPU — RESULTS")
+log(f"{'='*50}")
+log(f"  Final loss:  {losses[-1]:.4f}")
+log(f"  Time:        {tt:.0f}s")
+log(f"  GPU Memory:  {peak:.0f} MB")
+log(f"  Steps:       {len(losses)}")
+log(f"")
+log(f"  COMPARISON (same model, same data, same config):")
+log(f"  {'Method':>16} {'Loss':>8} {'Time':>7} {'GPU MB':>8}")
+log(f"  {'─'*44}")
+log(f"  {'FP16 LoRA':>16} {'0.2252':>8} {'1309s':>7} {'3585':>8}")
+log(f"  {'BnB NF4 QLoRA':>16} {'0.2399':>8} {'1423s':>7} {'2441':>8}")
+log(f"  {'FigQuant LoRA':>16} {losses[-1]:>8.4f} {tt:>6.0f}s {peak:>7.0f}")
+log(f"{'='*50}")