Add self-distill + train job script

scripts/hf_job_selfdistill_train.py

@@ -0,0 +1,610 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "torch>=2.4",
+#     "transformers>=4.48",
+#     "huggingface_hub>=1.9",
+#     "rich",
+# ]
+# ///
+"""
+FastMTP Self-Distill + Train — HuggingFace Job Script
+
+Runs on HF A100 as a single job:
+  1. Load Gemma 4 E4B-it from HF Hub
+  2. Generate self-distillation responses from prompts
+  3. Train MTP head (frozen base, only head updates)
+  4. Upload MTP head checkpoint to HF
+
+Usage:
+    hf jobs uv run experiments/hf_job_selfdistill_train.py \
+        --flavor a100-large --with torch transformers huggingface_hub rich \
+        --env HF_TOKEN=hf_xxx \
+        --env MODEL_ID=google/gemma-4-E4B-it \
+        --env MTP_CONFIG=e4b \
+        --env MAX_PROMPTS=2000 \
+        --env TRAIN_EPOCHS=3 \
+        --detach
+"""
+from __future__ import annotations
+
+import json
+import math
+import os
+import sys
+import time
+from pathlib import Path
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# ---------------------------------------------------------------------------
+# Config from env
+# ---------------------------------------------------------------------------
+MODEL_ID = os.environ.get("MODEL_ID", "google/gemma-4-E4B-it")
+MTP_CONFIG_NAME = os.environ.get("MTP_CONFIG", "e4b")
+MAX_PROMPTS = int(os.environ.get("MAX_PROMPTS", "2000"))
+TRAIN_EPOCHS = int(os.environ.get("TRAIN_EPOCHS", "3"))
+BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "4"))
+MAX_SEQ_LEN = int(os.environ.get("MAX_SEQ_LEN", "2048"))
+LR = float(os.environ.get("LR", "5e-5"))
+K = int(os.environ.get("K", "3"))
+BETA = float(os.environ.get("BETA", "0.6"))
+HF_UPLOAD_REPO = os.environ.get("HF_UPLOAD_REPO", "Cytrex/fastmtp-training")
+PROMPTS_REPO = os.environ.get("PROMPTS_REPO", "Cytrex/fastmtp-prompts")
+LOG_EVERY = int(os.environ.get("LOG_EVERY", "25"))
+GEN_BATCH = int(os.environ.get("GEN_BATCH", "8"))
+GEN_MAX_TOKENS = int(os.environ.get("GEN_MAX_TOKENS", "2048"))
+
+print(f"=== FastMTP Self-Distill + Train ===")
+print(f"Model: {MODEL_ID}")
+print(f"MTP config: {MTP_CONFIG_NAME}")
+print(f"Max prompts: {MAX_PROMPTS}")
+print(f"Epochs: {TRAIN_EPOCHS}")
+print(f"Batch size: {BATCH_SIZE}")
+print(f"Upload to: {HF_UPLOAD_REPO}")
+print(f"Prompts from: {PROMPTS_REPO}")
+
+
+# =====================================================================
+# MTP HEAD (inlined from lib/mtp_head.py)
+# =====================================================================
+
+GEMMA4_CONFIGS = {
+    "e2b": {
+        "hidden_size": 1536,
+        "intermediate_size": 6144,
+        "num_attention_heads": 8,
+        "num_key_value_heads": 2,
+        "vocab_size": 262144,
+    },
+    "e4b": {
+        "hidden_size": 2560,
+        "intermediate_size": 10240,
+        "num_attention_heads": 8,
+        "num_key_value_heads": 2,
+        "vocab_size": 262144,
+    },
+    "27b": {
+        "hidden_size": 3584,
+        "intermediate_size": 14336,
+        "num_attention_heads": 32,
+        "num_key_value_heads": 8,
+        "vocab_size": 262144,
+    },
+    "31b": {
+        "hidden_size": 3584,
+        "intermediate_size": 14336,
+        "num_attention_heads": 32,
+        "num_key_value_heads": 8,
+        "vocab_size": 262144,
+    },
+}
+
+
+class MTPHead(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        vocab_size: int,
+        head_dim: Optional[int] = None,
+        rms_norm_eps: float = 1e-6,
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_attention_heads
+        self.num_kv_heads = num_key_value_heads
+        self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads
+        self.vocab_size = vocab_size
+
+        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=rms_norm_eps)
+        self.attn_norm = nn.RMSNorm(hidden_size, eps=rms_norm_eps)
+        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.ffn_norm = nn.RMSNorm(hidden_size, eps=rms_norm_eps)
+        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.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)
+
+    def _attention(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, _ = x.shape
+        residual = x
+        x = self.attn_norm(x)
+        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        if self.num_kv_heads < self.num_heads:
+            repeat = self.num_heads // self.num_kv_heads
+            k = k.repeat_interleave(repeat, dim=1)
+            v = v.repeat_interleave(repeat, dim=1)
+        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        out = out.transpose(1, 2).contiguous().view(B, T, self.num_heads * self.head_dim)
+        out = self.o_proj(out)
+        return residual + out
+
+    def _ffn(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = self.ffn_norm(x)
+        gate = F.silu(self.gate_proj(x))
+        up = self.up_proj(x)
+        out = self.down_proj(gate * up)
+        return residual + out
+
+    def forward(self, hidden_states: torch.Tensor, shifted_token_ids: torch.Tensor):
+        token_embed = self.embed_tokens(shifted_token_ids)
+        fused = torch.cat([hidden_states, token_embed], dim=-1)
+        fused = self.fusion_proj(fused)
+        fused = self.fusion_norm(fused)
+        x = self._attention(fused)
+        x = self._ffn(x)
+        logits = self.lm_head(x)
+        return logits, x
+
+    @classmethod
+    def from_base_model(cls, base_model: nn.Module, config: dict) -> "MTPHead":
+        head = cls(
+            hidden_size=config["hidden_size"],
+            intermediate_size=config["intermediate_size"],
+            num_attention_heads=config["num_attention_heads"],
+            num_key_value_heads=config["num_key_value_heads"],
+            vocab_size=config["vocab_size"],
+        )
+        if hasattr(base_model, "model") and hasattr(base_model.model, "language_model"):
+            lm_body = base_model.model.language_model
+        elif hasattr(base_model, "model"):
+            lm_body = base_model.model
+        else:
+            lm_body = base_model
+        if hasattr(lm_body, "embed_tokens"):
+            head.embed_tokens = lm_body.embed_tokens
+            head.embed_tokens.requires_grad_(False)
+        if hasattr(base_model, "lm_head"):
+            head.lm_head = base_model.lm_head
+            head.lm_head.requires_grad_(False)
+        elif hasattr(lm_body, "lm_head"):
+            head.lm_head = lm_body.lm_head
+            head.lm_head.requires_grad_(False)
+        return head
+
+    def trainable_params(self) -> list[nn.Parameter]:
+        tied_ids = set()
+        for module in (self.embed_tokens, self.lm_head):
+            for p in module.parameters():
+                tied_ids.add(id(p))
+        return [p for p in self.parameters() if id(p) not in tied_ids and p.requires_grad]
+
+    def n_trainable(self) -> int:
+        return sum(p.numel() for p in self.trainable_params())
+
+
+# =====================================================================
+# MTP TRAINING UTILS (inlined from lib/mtp_train.py)
+# =====================================================================
+
+def compute_alpha_weights(k: int = 3, beta: float = 0.6) -> list[float]:
+    raw = [beta ** (i - 1) for i in range(1, k + 1)]
+    total = sum(raw)
+    return [w / total for w in raw]
+
+
+def mtp_loss(draft_logits, target_ids, k=3, beta=0.6):
+    alphas = compute_alpha_weights(k=k, beta=beta)
+    total_loss = torch.tensor(0.0, device=draft_logits[0].device, requires_grad=True)
+    for i in range(k):
+        logits = draft_logits[i]
+        targets = target_ids[i]
+        B, T, V = logits.shape
+        ce = F.cross_entropy(logits.reshape(B * T, V), targets.reshape(B * T), ignore_index=-100)
+        total_loss = total_loss + alphas[i] * ce
+    return total_loss
+
+
+def _get_last_layer(model: nn.Module) -> nn.Module:
+    if hasattr(model, "model") and hasattr(model.model, "language_model"):
+        return model.model.language_model.layers[-1]
+    if hasattr(model, "model") and hasattr(model.model, "layers"):
+        return model.model.layers[-1]
+    raise ValueError("Cannot find transformer layers for hook registration")
+
+
+@torch.no_grad()
+def extract_hidden_states(model, input_ids, attention_mask=None):
+    captured = {}
+    def hook_fn(module, args, output):
+        if isinstance(output, tuple):
+            captured["hidden"] = output[0]
+        else:
+            captured["hidden"] = output
+    last_layer = _get_last_layer(model)
+    handle = last_layer.register_forward_hook(hook_fn)
+    try:
+        kwargs = {"input_ids": input_ids}
+        if attention_mask is not None:
+            kwargs["attention_mask"] = attention_mask
+        model(**kwargs)
+    finally:
+        handle.remove()
+    return captured["hidden"]
+
+
+def prepare_mtp_targets(input_ids, k=3, pad_id=-100):
+    if input_ids.dim() == 1:
+        input_ids = input_ids.unsqueeze(0)
+    B, S = input_ids.shape
+    valid_len = S - 2
+    targets = []
+    for i in range(1, k + 1):
+        shift = i + 1
+        avail = S - shift
+        if avail >= valid_len:
+            t = input_ids[:, shift: shift + valid_len]
+        else:
+            t_real = input_ids[:, shift:]
+            pad_len = valid_len - avail
+            pad = torch.full((B, pad_len), pad_id, dtype=input_ids.dtype, device=input_ids.device)
+            t = torch.cat([t_real, pad], dim=1)
+        targets.append(t)
+    return targets, valid_len
+
+
+# =====================================================================
+# SELF-DISTILL: Generate responses from the model
+# =====================================================================
+
+def generate_self_distill(model, tokenizer, prompts: list[str], batch_size: int = 8, max_tokens: int = 2048) -> list[dict]:
+    """Generate responses for all prompts, return list of {prompt, response}."""
+    print(f"\n=== Self-Distill: generating {len(prompts)} responses ===")
+    results = []
+    t0 = time.perf_counter()
+
+    for i in range(0, len(prompts), batch_size):
+        batch_prompts = prompts[i:i + batch_size]
+
+        # Format as chat
+        conversations = []
+        for p in batch_prompts:
+            conversations.append([{"role": "user", "content": p}])
+
+        # Tokenize with chat template
+        texts = [tokenizer.apply_chat_template(c, tokenize=False, add_generation_prompt=True) for c in conversations]
+        inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=max_tokens // 2)
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=max_tokens // 2,
+                temperature=0.6,
+                top_p=0.95,
+                do_sample=True,
+                pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
+            )
+
+        # Decode only the generated part
+        for j, output_ids in enumerate(outputs):
+            input_len = inputs["input_ids"][j].shape[0]
+            response_ids = output_ids[input_len:]
+            response = tokenizer.decode(response_ids, skip_special_tokens=True)
+            if response.strip():
+                results.append({"prompt": batch_prompts[j], "response": response.strip()})
+
+        if (i // batch_size + 1) % 10 == 0 or i + batch_size >= len(prompts):
+            elapsed = time.perf_counter() - t0
+            rate = len(results) / elapsed if elapsed > 0 else 0
+            print(f"  [{len(results)}/{len(prompts)}] rate={rate:.1f}/s elapsed={elapsed:.0f}s")
+
+    print(f"Self-distill complete: {len(results)} responses in {time.perf_counter() - t0:.0f}s")
+    return results
+
+
+# =====================================================================
+# LOAD PROMPTS FROM HF DATASET
+# =====================================================================
+
+def load_prompts_from_hf(repo_id: str, max_prompts: int) -> list[str]:
+    """Download prompts from HF dataset repo."""
+    from huggingface_hub import hf_hub_download, list_repo_files
+
+    print(f"\n=== Loading prompts from {repo_id} ===")
+    prompts = []
+
+    try:
+        files = list_repo_files(repo_id, repo_type="dataset")
+        jsonl_files = sorted([f for f in files if f.endswith(".jsonl")])
+        print(f"Found {len(jsonl_files)} JSONL files: {jsonl_files}")
+
+        for fname in jsonl_files:
+            local = hf_hub_download(repo_id, fname, repo_type="dataset")
+            with open(local, "r") as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    try:
+                        obj = json.loads(line)
+                        text = obj.get("prompt") or obj.get("text") or obj.get("instruction", "")
+                        if text and len(text) > 10:
+                            prompts.append(text)
+                    except json.JSONDecodeError:
+                        pass
+            if len(prompts) >= max_prompts:
+                break
+    except Exception as e:
+        print(f"Error loading from HF: {e}")
+        print("Falling back to built-in sample prompts")
+        prompts = _fallback_prompts()
+
+    prompts = prompts[:max_prompts]
+    print(f"Loaded {len(prompts)} prompts")
+    return prompts
+
+
+def _fallback_prompts() -> list[str]:
+    """Minimal set of diverse prompts for PoC if dataset not available."""
+    return [
+        "Explain the difference between TCP and UDP protocols.",
+        "Write a Python function that implements binary search on a sorted list.",
+        "What are the main causes of climate change and what can individuals do?",
+        "Erkläre den Unterschied zwischen funktionaler und objektorientierter Programmierung.",
+        "Solve the equation: 3x^2 + 5x - 2 = 0",
+        "Write a short story about a robot discovering emotions.",
+        "Describe the process of photosynthesis in detail.",
+        "Implement a simple linked list in Python with insert, delete, and search operations.",
+        "What are the key principles of quantum computing?",
+        "Explain how neural networks learn through backpropagation.",
+    ] * 10  # 100 prompts for minimal PoC
+
+
+# =====================================================================
+# TRAINING LOOP
+# =====================================================================
+
+def train_mtp(model, mtp_head, distill_data: list[dict], tokenizer, device: str):
+    """Train MTP head on self-distilled data."""
+    print(f"\n=== MTP Training: {len(distill_data)} samples, {TRAIN_EPOCHS} epochs ===")
+
+    # Freeze base model
+    for p in model.parameters():
+        p.requires_grad_(False)
+    model.eval()
+
+    base_dtype = next(model.parameters()).dtype
+    mtp_head = mtp_head.to(device=device, dtype=base_dtype)
+
+    optimizer = torch.optim.AdamW(mtp_head.trainable_params(), lr=LR, weight_decay=0.01)
+
+    # Tokenize all data
+    print("Tokenizing training data...")
+    all_tokens = []
+    for rec in distill_data:
+        text = f"<start_of_turn>user\n{rec['prompt']}<end_of_turn>\n<start_of_turn>model\n{rec['response']}<end_of_turn>"
+        tokens = tokenizer(text, truncation=True, max_length=MAX_SEQ_LEN, return_tensors="pt")["input_ids"].squeeze(0)
+        if len(tokens) >= K + 4:
+            all_tokens.append(tokens)
+    print(f"  {len(all_tokens)} valid samples after tokenization")
+
+    # Simple training loop
+    total_steps = (len(all_tokens) // BATCH_SIZE) * TRAIN_EPOCHS
+    warmup_steps = max(1, int(total_steps * 0.05))
+
+    scheduler = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer, max_lr=LR, total_steps=max(total_steps, 1),
+        pct_start=0.05, anneal_strategy="cos",
+    )
+
+    global_step = 0
+    best_loss = float("inf")
+    t_start = time.perf_counter()
+    checkpoint_dir = Path("/tmp/mtp_checkpoints")
+    checkpoint_dir.mkdir(exist_ok=True)
+
+    for epoch in range(1, TRAIN_EPOCHS + 1):
+        print(f"\n--- Epoch {epoch}/{TRAIN_EPOCHS} ---")
+        # Shuffle
+        import random
+        random.shuffle(all_tokens)
+
+        epoch_loss = 0.0
+        epoch_steps = 0
+
+        for batch_start in range(0, len(all_tokens), BATCH_SIZE):
+            batch_toks = all_tokens[batch_start:batch_start + BATCH_SIZE]
+            if not batch_toks:
+                continue
+
+            # Pad batch
+            max_len = max(t.shape[0] for t in batch_toks)
+            padded = torch.zeros(len(batch_toks), max_len, dtype=torch.long)
+            for i, t in enumerate(batch_toks):
+                padded[i, :t.shape[0]] = t
+            input_ids = padded.to(device)
+
+            # Extract hidden states
+            hidden = extract_hidden_states(model, input_ids)
+
+            # Prepare targets
+            targets, valid_len = prepare_mtp_targets(input_ids, k=K)
+            targets = [t.to(device) for t in targets]
+
+            # Run MTP head
+            h = hidden[:, :valid_len, :]
+            shifted_ids = input_ids[:, 1:1 + valid_len]
+
+            draft_logits = []
+            for _ in range(K):
+                logits, h = mtp_head(h, shifted_ids)
+                draft_logits.append(logits)
+
+            # Loss + backprop
+            optimizer.zero_grad()
+            loss = mtp_loss(draft_logits, targets, k=K, beta=BETA)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(mtp_head.trainable_params(), 1.0)
+            optimizer.step()
+            scheduler.step()
+
+            loss_val = loss.item()
+            epoch_loss += loss_val
+            epoch_steps += 1
+            global_step += 1
+
+            if global_step % LOG_EVERY == 0:
+                avg = epoch_loss / epoch_steps
+                lr_now = scheduler.get_last_lr()[0]
+                elapsed = time.perf_counter() - t_start
+                print(f"  step={global_step} loss={loss_val:.4f} avg={avg:.4f} lr={lr_now:.2e} {global_step/elapsed:.1f}step/s")
+
+        avg_epoch_loss = epoch_loss / max(epoch_steps, 1)
+        print(f"  Epoch {epoch} done — avg_loss={avg_epoch_loss:.4f}")
+
+        # Save checkpoint
+        ckpt = {
+            "mtp_head_state_dict": mtp_head.state_dict(),
+            "epoch": epoch,
+            "step": global_step,
+            "loss": avg_epoch_loss,
+            "k": K,
+            "beta": BETA,
+            "model_id": MODEL_ID,
+            "mtp_config": MTP_CONFIG_NAME,
+        }
+        torch.save(ckpt, checkpoint_dir / f"checkpoint_epoch{epoch}.pt")
+
+        if avg_epoch_loss < best_loss:
+            best_loss = avg_epoch_loss
+            torch.save(ckpt, checkpoint_dir / "best.pt")
+
+    torch.save(ckpt, checkpoint_dir / "final.pt")
+
+    # Save config
+    with open(checkpoint_dir / "train_config.json", "w") as f:
+        json.dump({
+            "model_id": MODEL_ID,
+            "mtp_config": MTP_CONFIG_NAME,
+            "k": K, "beta": BETA,
+            "lr": LR, "epochs": TRAIN_EPOCHS,
+            "batch_size": BATCH_SIZE,
+            "max_seq_len": MAX_SEQ_LEN,
+            "num_samples": len(all_tokens),
+            "best_loss": best_loss,
+        }, f, indent=2)
+
+    elapsed = time.perf_counter() - t_start
+    print(f"\nTraining complete. best_loss={best_loss:.4f} time={elapsed/60:.1f}min")
+    return checkpoint_dir
+
+
+# =====================================================================
+# UPLOAD RESULTS
+# =====================================================================
+
+def upload_results(checkpoint_dir: Path, distill_data: list[dict]):
+    """Upload checkpoints + distill data to HF."""
+    from huggingface_hub import HfApi
+
+    print(f"\n=== Uploading to {HF_UPLOAD_REPO} ===")
+    api = HfApi()
+
+    # Save distill data
+    distill_path = checkpoint_dir / "distill_data.jsonl"
+    with open(distill_path, "w") as f:
+        for rec in distill_data:
+            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
+
+    # Upload entire checkpoint dir
+    api.upload_folder(
+        folder_path=str(checkpoint_dir),
+        repo_id=HF_UPLOAD_REPO,
+        path_in_repo=f"mtp_head_{MTP_CONFIG_NAME}_k{K}",
+        repo_type="model",
+    )
+    print(f"Uploaded to {HF_UPLOAD_REPO}/mtp_head_{MTP_CONFIG_NAME}_k{K}")
+
+
+# =====================================================================
+# MAIN
+# =====================================================================
+
+def main():
+    t_total = time.perf_counter()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"\nDevice: {device}")
+    if device == "cuda":
+        print(f"GPU: {torch.cuda.get_device_name(0)}")
+        print(f"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB")
+
+    # Step 1: Load model
+    print(f"\n=== Loading model: {MODEL_ID} ===")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.bfloat16,
+        device_map=device,
+    )
+    model.eval()
+    print(f"Model loaded: {sum(p.numel() for p in model.parameters()) / 1e9:.2f}B params")
+
+    # Step 2: Load prompts
+    prompts = load_prompts_from_hf(PROMPTS_REPO, MAX_PROMPTS)
+
+    # Step 3: Self-distill
+    distill_data = generate_self_distill(
+        model, tokenizer, prompts,
+        batch_size=GEN_BATCH, max_tokens=GEN_MAX_TOKENS,
+    )
+
+    if not distill_data:
+        print("ERROR: No distill data generated!")
+        sys.exit(1)
+
+    # Step 4: Build MTP head
+    config = GEMMA4_CONFIGS[MTP_CONFIG_NAME]
+    mtp_head = MTPHead.from_base_model(model, config)
+    print(f"MTP head: {mtp_head.n_trainable() / 1e6:.1f}M trainable params")
+
+    # Step 5: Train
+    checkpoint_dir = train_mtp(model, mtp_head, distill_data, tokenizer, device)
+
+    # Step 6: Upload
+    upload_results(checkpoint_dir, distill_data)
+
+    elapsed = time.perf_counter() - t_total
+    print(f"\n=== JOB COMPLETE in {elapsed/60:.1f}min ===")
+
+
+if __name__ == "__main__":
+    main()