Self-distill + train script v2

selfdistill_train.py

@@ -0,0 +1,348 @@
+"""FastMTP: Self-Distill + Train in one job on HF A100.
+
+1. Load E4B base model
+2. Generate 5k responses (self-distillation)
+3. Train MTP head on those responses
+4. Upload checkpoint to HF
+"""
+import os, sys, json, time, random
+sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', buffering=1)
+sys.stderr = os.fdopen(sys.stderr.fileno(), 'w', buffering=1)
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from pathlib import Path
+from torch.utils.data import DataLoader
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from datasets import load_dataset
+from huggingface_hub import HfApi
+
+# ============================================================
+# Config
+# ============================================================
+MODEL_ID = "InfinimindCreations/gemma-4-E4B-it-uncensored"
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
+UPLOAD_REPO = "Cytrex/fastmtp-e4b-selfdistill"
+
+# Self-distill config
+N_DISTILL = 5000
+GEN_MAX_TOKENS = 256
+GEN_TEMPERATURE = 0.6
+GEN_TOP_K = 20
+GEN_TOP_P = 0.95
+
+# Training config
+K = 3
+BETA = 0.6
+LR = 5e-5
+BATCH = 2
+EPOCHS = 3
+MAX_SEQ = 512
+
+OUTPUT = "/tmp/mtp_checkpoint"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ============================================================
+# MTP Head
+# ============================================================
+class MTPHead(nn.Module):
+    def __init__(self, hidden_size, intermediate_size, num_attention_heads, num_key_value_heads, vocab_size):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_attention_heads
+        self.num_kv_heads = num_key_value_heads
+        self.head_dim = hidden_size // num_attention_heads
+
+        self.embed_tokens = nn.Embedding(vocab_size, hidden_size)
+        self.fusion_proj = nn.Linear(hidden_size * 2, hidden_size, bias=False)
+        self.fusion_norm = nn.RMSNorm(hidden_size, eps=1e-6)
+
+        self.q_proj = nn.Linear(hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, hidden_size, bias=False)
+        self.attn_norm = nn.RMSNorm(hidden_size, eps=1e-6)
+
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+        self.ffn_norm = nn.RMSNorm(hidden_size, eps=1e-6)
+
+        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)
+
+    def forward(self, hidden_states, shifted_token_ids):
+        tok_embed = self.embed_tokens(shifted_token_ids)
+        fused = self.fusion_proj(torch.cat([hidden_states, tok_embed], dim=-1))
+        fused = self.fusion_norm(fused)
+
+        B, T, _ = fused.shape
+        normed = self.attn_norm(fused)
+        q = self.q_proj(normed).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(normed).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(normed).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        if self.num_kv_heads < self.num_heads:
+            n_rep = self.num_heads // self.num_kv_heads
+            k = k.repeat_interleave(n_rep, dim=1)
+            v = v.repeat_interleave(n_rep, dim=1)
+        attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, -1)
+        x = fused + self.o_proj(attn_out)
+
+        normed = self.ffn_norm(x)
+        x = x + self.down_proj(F.silu(self.gate_proj(normed)) * self.up_proj(normed))
+
+        return self.lm_head(x), x
+
+    def trainable_params(self):
+        return [p for p in self.parameters() if p.requires_grad]
+
+# ============================================================
+# Loss
+# ============================================================
+def mtp_loss(draft_logits, target_ids, k=3, beta=0.6):
+    raw = [beta ** i for i in range(k)]
+    total = sum(raw)
+    alphas = [w / total for w in raw]
+    loss = torch.tensor(0.0, device=draft_logits[0].device)
+    for i in range(k):
+        ce = F.cross_entropy(
+            draft_logits[i].reshape(-1, draft_logits[i].size(-1)),
+            target_ids[i].reshape(-1),
+            ignore_index=0, reduction="mean",
+        )
+        loss = loss + alphas[i] * ce
+    return loss
+
+# ============================================================
+# Phase 1: Self-Distillation
+# ============================================================
+def generate_selfdistill(model, tokenizer, prompts, max_tokens=256):
+    """Generate responses from the model itself."""
+    print(f"\n=== PHASE 1: Self-Distill ({len(prompts)} prompts, max_tokens={max_tokens}) ===")
+    samples = []
+    t0 = time.time()
+
+    for i, prompt in enumerate(prompts):
+        input_text = f"<start_of_turn>user\n{prompt}<end_of_turn>\n<start_of_turn>model\n"
+        input_ids = tokenizer.encode(input_text, return_tensors="pt").to(DEVICE)
+
+        with torch.no_grad():
+            output = model.generate(
+                input_ids,
+                max_new_tokens=max_tokens,
+                do_sample=True,
+                temperature=GEN_TEMPERATURE,
+                top_k=GEN_TOP_K,
+                top_p=GEN_TOP_P,
+            )
+
+        response_ids = output[0][input_ids.shape[1]:]
+        response = tokenizer.decode(response_ids, skip_special_tokens=True).strip()
+
+        if len(response) > 20:
+            full_text = f"<start_of_turn>user\n{prompt}<end_of_turn>\n<start_of_turn>model\n{response}<end_of_turn>"
+            ids = tokenizer.encode(full_text, max_length=MAX_SEQ, truncation=True)
+            if len(ids) >= K + 4:
+                samples.append(torch.tensor(ids, dtype=torch.long))
+
+        if (i + 1) % 100 == 0:
+            elapsed = time.time() - t0
+            rate = (i + 1) / elapsed
+            eta = (len(prompts) - i - 1) / rate / 60
+            print(f"  [{i+1}/{len(prompts)}] {len(samples)} valid | {rate:.1f} prompts/s | ETA {eta:.1f}min")
+
+    elapsed = time.time() - t0
+    print(f"Self-distill done: {len(samples)} valid samples in {elapsed:.0f}s ({elapsed/60:.1f}min)")
+    return samples
+
+# ============================================================
+# Phase 2: Training
+# ============================================================
+def train_mtp(model, mtp_head, samples):
+    print(f"\n=== PHASE 2: Training ({len(samples)} samples, {EPOCHS} epochs) ===")
+
+    def collate(batch):
+        mx = max(len(s) for s in batch)
+        padded = torch.zeros(len(batch), mx, dtype=torch.long)
+        for i, s in enumerate(batch):
+            padded[i, :len(s)] = s
+        return padded
+
+    loader = DataLoader(samples, batch_size=BATCH, shuffle=True, collate_fn=collate, num_workers=0)
+    optimizer = torch.optim.AdamW(mtp_head.trainable_params(), lr=LR, betas=(0.9, 0.95), weight_decay=0.01)
+
+    # Freeze base model
+    for p in model.parameters():
+        p.requires_grad_(False)
+    model.eval()
+
+    total_steps = len(loader) * EPOCHS
+    print(f"Steps: {len(loader)}/epoch, {total_steps} total")
+
+    t0 = time.time()
+    best_loss = float("inf")
+
+    for epoch in range(EPOCHS):
+        epoch_loss = 0
+        for step, batch in enumerate(loader):
+            input_ids = batch.to(DEVICE)
+            B, S = input_ids.shape
+            valid_len = S - K - 1
+            if valid_len <= 0:
+                continue
+
+            with torch.no_grad():
+                outputs = model(input_ids=input_ids, output_hidden_states=True)
+                hidden = outputs.hidden_states[-1][:, :valid_len, :]
+
+            targets = []
+            for i in range(K):
+                shift = i + 2
+                t = input_ids[:, shift:shift + valid_len]
+                if t.shape[1] < valid_len:
+                    pad = torch.zeros(B, valid_len - t.shape[1], dtype=torch.long, device=DEVICE)
+                    t = torch.cat([t, pad], dim=1)
+                targets.append(t)
+
+            draft_logits = []
+            h = hidden
+            for i in range(K):
+                shifted_ids = input_ids[:, i + 1:i + 1 + valid_len]
+                if shifted_ids.shape[1] < valid_len:
+                    pad = torch.zeros(B, valid_len - shifted_ids.shape[1], dtype=torch.long, device=DEVICE)
+                    shifted_ids = torch.cat([shifted_ids, pad], dim=1)
+                logits, h = mtp_head(h, shifted_ids)
+                draft_logits.append(logits)
+
+            loss = mtp_loss(draft_logits, targets, K, BETA)
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(mtp_head.trainable_params(), 1.0)
+            optimizer.step()
+
+            epoch_loss += loss.item()
+            if (step + 1) % 50 == 0:
+                avg = epoch_loss / (step + 1)
+                elapsed = time.time() - t0
+                steps_done = epoch * len(loader) + step + 1
+                eta = (elapsed / steps_done) * (total_steps - steps_done) / 60
+                print(f"  E{epoch+1} S{step+1}/{len(loader)} | loss={loss.item():.4f} avg={avg:.4f} | {elapsed:.0f}s | ETA {eta:.0f}min")
+
+        avg_loss = epoch_loss / max(len(loader), 1)
+        print(f"Epoch {epoch+1}/{EPOCHS} | avg_loss={avg_loss:.4f} | {time.time()-t0:.0f}s")
+
+        os.makedirs(OUTPUT, exist_ok=True)
+        ckpt = {
+            "mtp_head_state_dict": {k: v.cpu() for k, v in mtp_head.state_dict().items()
+                                     if not k.startswith("embed_tokens") and not k.startswith("lm_head")},
+            "epoch": epoch + 1,
+            "loss": avg_loss,
+            "k": K, "beta": BETA,
+            "config": {"hidden_size": 2560, "intermediate_size": 10240, "num_attention_heads": 8, "num_key_value_heads": 2, "vocab_size": 262144},
+        }
+        torch.save(ckpt, f"{OUTPUT}/mtp_checkpoint_e{epoch+1}.pt")
+        if avg_loss < best_loss:
+            best_loss = avg_loss
+            torch.save(ckpt, f"{OUTPUT}/mtp_best.pt")
+            print(f"  New best: {best_loss:.4f}")
+
+    return best_loss
+
+# ============================================================
+# Main
+# ============================================================
+def main():
+    print(f"Device: {DEVICE}")
+    if DEVICE == "cuda":
+        print(f"GPU: {torch.cuda.get_device_name(0)}")
+
+    # Load model
+    print("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN, trust_remote_code=True)
+
+    print("Loading base model...")
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID, dtype=torch.bfloat16, device_map="auto",
+        token=HF_TOKEN, trust_remote_code=True,
+    )
+
+    # Load prompts from Magpie (only the prompts, not responses)
+    print("Loading prompts from Magpie-Pro-300K...")
+    ds = load_dataset("Magpie-Align/Magpie-Pro-300K-Filtered", split="train")
+    prompts = []
+    indices = list(range(len(ds)))
+    random.seed(42)
+    random.shuffle(indices)
+    for idx in indices:
+        if len(prompts) >= N_DISTILL:
+            break
+        conv = ds[idx]["conversations"]
+        if len(conv) >= 1 and conv[0]["from"] == "human" and len(conv[0]["value"]) > 10:
+            prompts.append(conv[0]["value"])
+    print(f"Loaded {len(prompts)} prompts")
+
+    # Phase 1: Self-distill
+    samples = generate_selfdistill(model, tokenizer, prompts, GEN_MAX_TOKENS)
+
+    if len(samples) < 100:
+        print(f"ERROR: Only {len(samples)} valid samples — not enough for training")
+        return
+
+    # Phase 2: Create MTP head and train
+    print("\nCreating MTP head...")
+    config = {"hidden_size": 2560, "intermediate_size": 10240, "num_attention_heads": 8, "num_key_value_heads": 2, "vocab_size": 262144}
+    mtp_head = MTPHead(**config)
+
+    # Tie embed + lm_head
+    if hasattr(model, 'model') and hasattr(model.model, 'language_model'):
+        embed_w = model.model.language_model.embed_tokens.weight
+    elif hasattr(model, 'model'):
+        embed_w = model.model.embed_tokens.weight
+    else:
+        raise RuntimeError("Cannot find embed_tokens")
+    lm_head_w = model.lm_head.weight
+
+    mtp_head.embed_tokens.weight = embed_w
+    mtp_head.lm_head.weight = lm_head_w
+    mtp_head.embed_tokens.weight.requires_grad = False
+    mtp_head.lm_head.weight.requires_grad = False
+
+    base_dtype = next(model.parameters()).dtype
+    mtp_head = mtp_head.to(device=DEVICE, dtype=base_dtype)
+    n_trainable = sum(p.numel() for p in mtp_head.trainable_params())
+    print(f"MTP head: {n_trainable:,} trainable params, dtype={base_dtype}")
+
+    best_loss = train_mtp(model, mtp_head, samples)
+
+    print(f"\n=== DONE === Best loss: {best_loss:.4f}")
+
+    # Upload
+    if HF_TOKEN:
+        print(f"\nUploading to {UPLOAD_REPO}...")
+        api = HfApi(token=HF_TOKEN)
+        try:
+            api.create_repo(UPLOAD_REPO, exist_ok=True)
+        except Exception as e:
+            print(f"Repo: {e}")
+
+        meta = {
+            "type": "fastmtp_head",
+            "base_model": MODEL_ID,
+            "method": "self-distillation",
+            "distill_samples": len(samples),
+            "k": K, "beta": BETA, "epochs": EPOCHS,
+            "best_loss": best_loss,
+            "trainable_params": n_trainable,
+            "reference": "arXiv:2509.18362",
+        }
+        with open(f"{OUTPUT}/mtp_config.json", "w") as f:
+            json.dump(meta, f, indent=2)
+
+        api.upload_folder(folder_path=OUTPUT, repo_id=UPLOAD_REPO,
+                          commit_message=f"FastMTP E4B self-distill — {EPOCHS}ep, {len(samples)} samples, loss={best_loss:.4f}")
+        print(f"Uploaded: https://huggingface.co/{UPLOAD_REPO}")
+
+
+if __name__ == "__main__":
+    main()