Upload distill_sharded.py with huggingface_hub

distill_sharded.py

@@ -0,0 +1,807 @@
+#!/usr/bin/env python3
+"""
+Single-process KL distillation with a sharded frozen teacher and one trainable
+student GPU.
+
+This is a derivative of distill.py tailored for large-teacher / smaller-student
+ setups where replicating the teacher per process is wasteful or infeasible.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+import logging
+import random
+import re
+import shutil
+import time
+import tomllib
+from collections import OrderedDict
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint_utils
+from torch.optim import AdamW
+
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(message)s",
+    datefmt="%H:%M:%S",
+)
+log = logging.getLogger("distill_sharded")
+
+
+REQUIRED_SECTIONS = ("model", "data", "train", "eval", "log", "init")
+REQUIRED_KEYS = {
+    "model": ("teacher", "student", "tokenizer", "student_device", "teacher_devices", "teacher_max_memory_gb"),
+    "data": ("min_chars", "max_seq_len", "kl_start_pos", "seed", "shuffle_buffer"),
+    "train": (
+        "seed",
+        "lr",
+        "schedule",
+        "warmup_steps",
+        "weight_decay",
+        "grad_clip",
+        "betas",
+        "eps",
+        "samples_per_step",
+        "max_steps",
+        "grad_checkpointing",
+        "attn_implementation",
+        "student_dtype",
+        "teacher_dtype",
+        "kl_chunk_size",
+        "micro_batch_size",
+        "new_layer_lr_mul",
+    ),
+    "eval": ("every_steps", "samples", "seed", "cache_path"),
+    "log": ("wandb", "wandb_project", "wandb_run", "log_every", "output_dir", "experiment_log"),
+    "init": ("zero_layers", "target_num_layers"),
+}
+
+DTYPE_MAP = {
+    "float32": torch.float32,
+    "bfloat16": torch.bfloat16,
+}
+
+
+def parse_dtype(s: str) -> torch.dtype:
+    if s not in DTYPE_MAP:
+        raise ValueError(f"unknown dtype {s!r}; must be one of {list(DTYPE_MAP)}")
+    return DTYPE_MAP[s]
+
+
+def load_config(path: str) -> dict:
+    with open(path, "rb") as f:
+        cfg = tomllib.load(f)
+    for sec in REQUIRED_SECTIONS:
+        if sec not in cfg:
+            raise KeyError(f"config missing required section [{sec}]")
+        for key in REQUIRED_KEYS[sec]:
+            if key not in cfg[sec]:
+                raise KeyError(f"config missing required key [{sec}].{key}")
+    return cfg
+
+
+def get_inner_with_layers(model):
+    seen = set()
+    stack = [model]
+    while stack:
+        m = stack.pop()
+        if id(m) in seen:
+            continue
+        seen.add(id(m))
+        if hasattr(m, "layers"):
+            return m
+        for attr in ("model", "language_model", "transformer", "base_model"):
+            child = getattr(m, attr, None)
+            if child is not None:
+                stack.append(child)
+    raise RuntimeError(f"Could not locate `.layers` inside {type(model).__name__}")
+
+
+def zero_layers(model, layer_indices):
+    inner = get_inner_with_layers(model)
+    layers = inner.layers
+    n = len(layers)
+    for idx in layer_indices:
+        if idx < 0 or idx >= n:
+            raise IndexError(f"layer {idx} out of range (0..{n - 1})")
+        with torch.no_grad():
+            for p in layers[idx].parameters():
+                p.zero_()
+    return n
+
+
+def _zero_output_projections(layer):
+    zeroed = []
+    with torch.no_grad():
+        if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "o_proj"):
+            layer.self_attn.o_proj.weight.zero_()
+            zeroed.append("self_attn.o_proj")
+        if hasattr(layer, "linear_attn") and hasattr(layer.linear_attn, "out_proj"):
+            layer.linear_attn.out_proj.weight.zero_()
+            zeroed.append("linear_attn.out_proj")
+        if hasattr(layer, "mlp") and hasattr(layer.mlp, "down_proj"):
+            layer.mlp.down_proj.weight.zero_()
+            zeroed.append("mlp.down_proj")
+    return zeroed
+
+
+def grow_layers(model, target_n):
+    inner = get_inner_with_layers(model)
+    cur_n = len(inner.layers)
+    if target_n == cur_n:
+        return cur_n, []
+    if target_n < cur_n:
+        raise ValueError(f"target_num_layers={target_n} < current {cur_n}; cannot shrink")
+
+    cfg = model.config
+    text_cfg = getattr(cfg, "text_config", cfg)
+    if not hasattr(text_cfg, "layer_types") or not text_cfg.layer_types:
+        raise RuntimeError("text config has no layer_types; cannot extend pattern")
+
+    period = getattr(text_cfg, "full_attention_interval", 4)
+    new_types = list(text_cfg.layer_types)
+    while len(new_types) < target_n:
+        new_types.append(new_types[len(new_types) % period])
+    text_cfg.layer_types = new_types
+    text_cfg.num_hidden_layers = target_n
+    if hasattr(cfg, "num_hidden_layers") and cfg is not text_cfg:
+        cfg.num_hidden_layers = target_n
+
+    layer_cls = type(inner.layers[0])
+    device = next(inner.parameters()).device
+    dtype = next(inner.parameters()).dtype
+
+    new_layer_zeroed = []
+    for i in range(cur_n, target_n):
+        new_layer = layer_cls(text_cfg, layer_idx=i)
+        new_layer.apply(model._init_weights)
+        new_layer.to(device=device, dtype=dtype)
+        zeroed = _zero_output_projections(new_layer)
+        new_layer_zeroed.append((i, zeroed))
+        inner.layers.append(new_layer)
+
+    return target_n, new_layer_zeroed
+
+
+def detect_model_kind(model_id: str) -> str:
+    from transformers import AutoConfig
+
+    cfg = AutoConfig.from_pretrained(model_id)
+    archs = list(getattr(cfg, "architectures", []) or [])
+    arch = archs[0] if archs else ""
+    if "ConditionalGeneration" in arch or "ImageText" in arch:
+        return "image_text"
+    return "causal_lm"
+
+
+def load_student(model_id: str, dtype: torch.dtype, grad_ckpt: bool, attn_impl: str):
+    kind = detect_model_kind(model_id)
+    if kind == "image_text":
+        from transformers import AutoModelForImageTextToText
+
+        model = AutoModelForImageTextToText.from_pretrained(
+            model_id,
+            dtype=dtype,
+            low_cpu_mem_usage=True,
+            attn_implementation=attn_impl,
+        )
+    else:
+        from transformers import AutoModelForCausalLM
+
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            dtype=dtype,
+            low_cpu_mem_usage=True,
+            attn_implementation=attn_impl,
+        )
+    model.config.use_cache = False
+    if grad_ckpt:
+        model.gradient_checkpointing_enable(
+            gradient_checkpointing_kwargs={"use_reentrant": False}
+        )
+    return model
+
+
+def load_teacher(model_id: str, dtype: torch.dtype, attn_impl: str, devices: list[int], max_mem_gb: int):
+    kind = detect_model_kind(model_id)
+    max_memory = {idx: f"{max_mem_gb}GiB" for idx in devices}
+    max_memory["cpu"] = "256GiB"
+    common = dict(
+        dtype=dtype,
+        low_cpu_mem_usage=True,
+        attn_implementation=attn_impl,
+        device_map="auto",
+        max_memory=max_memory,
+    )
+
+    if kind == "image_text":
+        from transformers import AutoModelForImageTextToText
+
+        model = AutoModelForImageTextToText.from_pretrained(model_id, **common)
+    else:
+        from transformers import AutoModelForCausalLM
+
+        model = AutoModelForCausalLM.from_pretrained(model_id, **common)
+    model.config.use_cache = False
+    model.eval()
+    for p in model.parameters():
+        p.requires_grad_(False)
+    return model
+
+
+def get_teacher_devices(model) -> tuple[torch.device, torch.device]:
+    device_map = getattr(model, "hf_device_map", None) or {}
+    ordered = OrderedDict()
+    for _, dev in device_map.items():
+        if isinstance(dev, int):
+            ordered.setdefault(f"cuda:{dev}", None)
+        elif isinstance(dev, str) and dev.startswith("cuda:"):
+            ordered.setdefault(dev, None)
+    if not ordered:
+        first = next(model.parameters()).device
+        return first, first
+    keys = list(ordered.keys())
+    return torch.device(keys[0]), torch.device(keys[-1])
+
+
+def teacher_forward(teacher, input_ids, attention_mask, out_device):
+    out = teacher(input_ids=input_ids, attention_mask=attention_mask)
+    logits = getattr(out, "logits", None)
+    if logits is None:
+        raise RuntimeError("teacher forward did not return .logits")
+    if logits.device != out_device:
+        logits = logits.to(out_device, non_blocking=True)
+    return logits
+
+
+class StreamingTextLoader:
+    def __init__(
+        self,
+        name,
+        text_field,
+        min_chars,
+        max_seq_len,
+        kl_start_pos,
+        tokenizer,
+        seed,
+        shuffle_buffer,
+    ):
+        from datasets import load_dataset
+
+        last_err = None
+        for attempt in range(8):
+            try:
+                ds = load_dataset(name, split="train", streaming=True)
+                break
+            except Exception as e:
+                last_err = e
+                wait = min(2 ** attempt, 30)
+                log.warning(
+                    f"load_dataset({name!r}) failed (attempt {attempt + 1}/8): "
+                    f"{type(e).__name__}: {e}; sleeping {wait}s"
+                )
+                time.sleep(wait)
+        else:
+            raise RuntimeError(f"load_dataset failed after 8 retries") from last_err
+        ds = ds.shuffle(seed=seed, buffer_size=shuffle_buffer)
+        self._ds = iter(ds)
+        self._text_field = text_field
+        self._min_chars = min_chars
+        self._max_seq_len = max_seq_len
+        self._min_tokens = kl_start_pos + 16
+        self._tokenizer = tokenizer
+        self._name = name
+
+    def next_sample(self):
+        scanned = 0
+        while scanned < 100:
+            try:
+                item = next(self._ds)
+            except StopIteration:
+                return None
+            scanned += 1
+            text = item.get(self._text_field, "") or ""
+            if len(text) < self._min_chars:
+                continue
+            ids = self._tokenizer(
+                text,
+                return_tensors="pt",
+                truncation=True,
+                max_length=self._max_seq_len,
+            ).input_ids.squeeze(0)
+            if ids.shape[0] < self._min_tokens:
+                continue
+            return ids
+        return None
+
+
+class MixedStreamingLoader:
+    def __init__(self, specs, tokenizer, min_chars, max_seq_len, kl_start_pos, seed, shuffle_buffer):
+        self._rng = random.Random(seed)
+        self._weights = []
+        self._loaders = []
+        for spec in specs:
+            self._weights.append(spec["weight"])
+            self._loaders.append(
+                StreamingTextLoader(
+                    name=spec["name"],
+                    text_field=spec["text_field"],
+                    min_chars=min_chars,
+                    max_seq_len=max_seq_len,
+                    kl_start_pos=kl_start_pos,
+                    tokenizer=tokenizer,
+                    seed=seed + len(self._loaders),
+                    shuffle_buffer=shuffle_buffer,
+                )
+            )
+
+    def next_batch(self, n):
+        out = []
+        while len(out) < n:
+            idx = self._rng.choices(range(len(self._loaders)), weights=self._weights, k=1)[0]
+            sample = self._loaders[idx].next_sample()
+            if sample is None:
+                continue
+            out.append(sample)
+        return out
+
+
+def collate_pad(token_lists, pad_id):
+    max_len = max(t.shape[0] for t in token_lists)
+    B = len(token_lists)
+    input_ids = torch.full((B, max_len), pad_id, dtype=torch.long)
+    attention_mask = torch.zeros((B, max_len), dtype=torch.long)
+    for i, t in enumerate(token_lists):
+        L = t.shape[0]
+        input_ids[i, :L] = t
+        attention_mask[i, :L] = 1
+    return input_ids, attention_mask
+
+
+def _kl_chunk_sum(s_chunk, t_chunk, m_chunk):
+    s = s_chunk.float()
+    t = t_chunk.float()
+    t_log_p = F.log_softmax(t, dim=-1)
+    s_log_p = F.log_softmax(s, dim=-1)
+    t_p = t_log_p.exp()
+    per_token = (t_p * (t_log_p - s_log_p)).sum(-1)
+    return (per_token * m_chunk).sum()
+
+
+def kl_loss_masked(student_logits, teacher_logits, attention_mask, start_pos, chunk_size):
+    s_full = student_logits[:, start_pos:, :]
+    t_full = teacher_logits[:, start_pos:, :].detach()
+    m_full = attention_mask[:, start_pos:].float()
+
+    T = s_full.shape[1]
+    if chunk_size <= 0 or chunk_size >= T:
+        return _kl_chunk_sum(s_full, t_full, m_full) / m_full.sum().clamp_min(1.0)
+
+    total_kl = torch.zeros((), device=s_full.device, dtype=torch.float32)
+    for i in range(0, T, chunk_size):
+        end = min(i + chunk_size, T)
+        s_c = s_full[:, i:end, :]
+        t_c = t_full[:, i:end, :]
+        m_c = m_full[:, i:end]
+        chunk_kl = checkpoint_utils.checkpoint(
+            _kl_chunk_sum, s_c, t_c, m_c, use_reentrant=False
+        )
+        total_kl = total_kl + chunk_kl
+    return total_kl / m_full.sum().clamp_min(1.0)
+
+
+def apply_trainable_masks(model, train_cfg):
+    trainable = train_cfg.get("trainable_patterns", [])
+    frozen = train_cfg.get("freeze_patterns", [])
+    if not trainable and not frozen:
+        return
+
+    trainable_re = [re.compile(p) for p in trainable]
+    frozen_re = [re.compile(p) for p in frozen]
+    for name, p in model.named_parameters():
+        keep = True
+        if trainable_re:
+            keep = any(r.search(name) for r in trainable_re)
+        if keep and frozen_re and any(r.search(name) for r in frozen_re):
+            keep = False
+        p.requires_grad_(keep)
+
+
+def make_optimizer(model, train_cfg, new_layer_indices=None):
+    base_lr = train_cfg["lr"]
+    mul = train_cfg["new_layer_lr_mul"]
+    common = dict(
+        weight_decay=train_cfg["weight_decay"],
+        betas=tuple(train_cfg["betas"]),
+        eps=train_cfg["eps"],
+    )
+
+    if not new_layer_indices or mul == 1.0:
+        return AdamW(
+            [p for p in model.parameters() if p.requires_grad],
+            lr=base_lr,
+            **common,
+        )
+
+    inner = get_inner_with_layers(model)
+    new_pids = set()
+    for idx in new_layer_indices:
+        for p in inner.layers[idx].parameters():
+            if p.requires_grad:
+                new_pids.add(id(p))
+
+    new_params = []
+    rest_params = []
+    for p in model.parameters():
+        if not p.requires_grad:
+            continue
+        (new_params if id(p) in new_pids else rest_params).append(p)
+
+    return AdamW(
+        [
+            {"params": rest_params, "lr": base_lr},
+            {"params": new_params, "lr": base_lr * mul},
+        ],
+        **common,
+    )
+
+
+def make_scheduler(optimizer, train_cfg):
+    schedule = train_cfg["schedule"]
+    warmup = train_cfg["warmup_steps"]
+    total = train_cfg["max_steps"]
+
+    if schedule == "constant":
+        from transformers import get_constant_schedule_with_warmup
+
+        return get_constant_schedule_with_warmup(optimizer, warmup)
+    if schedule == "cosine":
+        from transformers import get_cosine_schedule_with_warmup
+
+        return get_cosine_schedule_with_warmup(optimizer, warmup, total)
+    if schedule == "linear":
+        from transformers import get_linear_schedule_with_warmup
+
+        return get_linear_schedule_with_warmup(optimizer, warmup, total)
+    raise ValueError(f"unknown schedule: {schedule!r}")
+
+
+def build_dataset_specs(data_cfg):
+    if "datasets" in data_cfg:
+        names = data_cfg["datasets"]
+        text_fields = data_cfg.get("text_fields", [data_cfg.get("text_field", "text")] * len(names))
+        weights = data_cfg.get("dataset_weights", [1.0] * len(names))
+        if not (len(names) == len(text_fields) == len(weights)):
+            raise ValueError("datasets/text_fields/dataset_weights length mismatch")
+        return [
+            {"name": name, "text_field": field, "weight": weight}
+            for name, field, weight in zip(names, text_fields, weights)
+        ]
+    return [
+        {
+            "name": data_cfg["dataset"],
+            "text_field": data_cfg["text_field"],
+            "weight": 1.0,
+        }
+    ]
+
+
+def build_or_load_eval_cache(path, loader=None, samples=None):
+    path = Path(path)
+    if path.exists():
+        log.info(f"Loading eval cache from {path}")
+        raw = torch.load(path)
+        return [torch.tensor(x, dtype=torch.long) for x in raw]
+    if loader is None or samples is None:
+        raise ValueError("loader and samples are required when building a new eval cache")
+    path.parent.mkdir(parents=True, exist_ok=True)
+    log.info(f"Building eval cache at {path}")
+    batches = loader.next_batch(samples)
+    torch.save([x.tolist() for x in batches], path)
+    return batches
+
+
+def log_jsonl(path: Path, record: dict):
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("a") as f:
+        f.write(json.dumps(record, sort_keys=True) + "\n")
+
+
+@torch.no_grad()
+def evaluate(student, teacher, eval_batches, pad_id, kl_start_pos, kl_chunk_size, student_device, teacher_input_device):
+    student.eval()
+    total = 0.0
+    n = 0
+    for sample in eval_batches:
+        ids, mask = collate_pad([sample], pad_id)
+        teacher_ids = ids.to(teacher_input_device, non_blocking=True)
+        teacher_mask = mask.to(teacher_input_device, non_blocking=True)
+        student_ids = ids.to(student_device, non_blocking=True)
+        student_mask = mask.to(student_device, non_blocking=True)
+        t_logits = teacher_forward(teacher, teacher_ids, teacher_mask, student_device)
+        s_logits = student(input_ids=student_ids, attention_mask=student_mask).logits
+        loss = kl_loss_masked(
+            s_logits,
+            t_logits,
+            student_mask,
+            start_pos=kl_start_pos,
+            chunk_size=kl_chunk_size,
+        )
+        total += loss.item()
+        n += 1
+        del t_logits, s_logits, loss, teacher_ids, teacher_mask, student_ids, student_mask
+    student.train()
+    return total / max(n, 1)
+
+
+def save_best(student, tokenizer, output_dir, step, eval_kl):
+    out_dir = Path(output_dir) / "best"
+    if out_dir.exists():
+        shutil.rmtree(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    student.save_pretrained(out_dir, safe_serialization=True)
+    tokenizer.save_pretrained(out_dir)
+    with (out_dir / "best.json").open("w") as f:
+        json.dump({"step": step, "eval_kl": eval_kl}, f, indent=2)
+    log.info(f"saved best @ step {step}: eval_kl={eval_kl:.6f} -> {out_dir}")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", required=True)
+    args = parser.parse_args()
+
+    cfg = load_config(args.config)
+    torch.manual_seed(cfg["train"]["seed"])
+    random.seed(cfg["train"]["seed"])
+
+    student_device = torch.device(cfg["model"]["student_device"])
+    teacher_devices = list(cfg["model"]["teacher_devices"])
+
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(cfg["model"]["tokenizer"], trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    pad_id = tokenizer.pad_token_id
+
+    student = load_student(
+        cfg["model"]["student"],
+        parse_dtype(cfg["train"]["student_dtype"]),
+        grad_ckpt=cfg["train"]["grad_checkpointing"],
+        attn_impl=cfg["train"]["attn_implementation"],
+    )
+    student.to(student_device)
+    student.config.use_cache = False
+
+    target_n = cfg["init"]["target_num_layers"]
+    cur_n = len(get_inner_with_layers(student).layers)
+    new_layer_indices = []
+    if target_n != cur_n:
+        new_n, new_zeroed = grow_layers(student, target_n)
+        new_layer_indices = [idx for idx, _ in new_zeroed]
+        log.info(f"Grew student from {cur_n} -> {new_n} layers")
+        for idx, names in new_zeroed:
+            log.info(f"  layer {idx}: zeroed {names}")
+
+    zero_idx = cfg["init"]["zero_layers"]
+    if zero_idx:
+        n = zero_layers(student, zero_idx)
+        log.info(f"Zeroed student layers {zero_idx} (model has {n} layers)")
+
+    apply_trainable_masks(student, cfg["train"])
+    trainable_params = sum(p.numel() for p in student.parameters() if p.requires_grad)
+    total_params = sum(p.numel() for p in student.parameters())
+    if trainable_params == 0:
+        raise RuntimeError("No trainable parameters remain after applying trainable/freeze patterns")
+    log.info(f"Student params: total={total_params/1e9:.3f}B trainable={trainable_params/1e9:.3f}B")
+
+    teacher = load_teacher(
+        cfg["model"]["teacher"],
+        parse_dtype(cfg["train"]["teacher_dtype"]),
+        attn_impl=cfg["train"]["attn_implementation"],
+        devices=teacher_devices,
+        max_mem_gb=cfg["model"]["teacher_max_memory_gb"],
+    )
+    teacher_input_device, _ = get_teacher_devices(teacher)
+    log.info(f"Teacher input device: {teacher_input_device}")
+
+    optimizer = make_optimizer(student, cfg["train"], new_layer_indices=new_layer_indices)
+    scheduler = make_scheduler(optimizer, cfg["train"])
+
+    output_dir = Path(cfg["log"]["output_dir"])
+    output_dir.mkdir(parents=True, exist_ok=True)
+    shutil.copy2(args.config, output_dir / "config.snapshot.toml")
+    metrics_path = output_dir / "metrics.jsonl"
+    experiment_log = Path(cfg["log"]["experiment_log"])
+
+    use_wandb = cfg["log"]["wandb"]
+    if use_wandb:
+        import wandb
+
+        wandb.init(
+            project=cfg["log"]["wandb_project"],
+            name=cfg["log"]["wandb_run"],
+            config=cfg,
+        )
+
+    specs = build_dataset_specs(cfg["data"])
+    train_loader = MixedStreamingLoader(
+        specs=specs,
+        tokenizer=tokenizer,
+        min_chars=cfg["data"]["min_chars"],
+        max_seq_len=cfg["data"]["max_seq_len"],
+        kl_start_pos=cfg["data"]["kl_start_pos"],
+        seed=cfg["data"]["seed"],
+        shuffle_buffer=cfg["data"]["shuffle_buffer"],
+    )
+    eval_cache_path = Path(cfg["eval"]["cache_path"])
+    if eval_cache_path.exists():
+        eval_batches = build_or_load_eval_cache(eval_cache_path)
+    else:
+        eval_loader = MixedStreamingLoader(
+            specs=specs,
+            tokenizer=tokenizer,
+            min_chars=cfg["data"]["min_chars"],
+            max_seq_len=cfg["data"]["max_seq_len"],
+            kl_start_pos=cfg["data"]["kl_start_pos"],
+            seed=cfg["eval"]["seed"],
+            shuffle_buffer=cfg["data"]["shuffle_buffer"],
+        )
+        eval_batches = build_or_load_eval_cache(eval_cache_path, eval_loader, cfg["eval"]["samples"])
+    log.info(f"Eval samples: {len(eval_batches)}")
+
+    samples_per_step = cfg["train"]["samples_per_step"]
+    micro_batch_size = cfg["train"]["micro_batch_size"]
+    grad_clip = cfg["train"]["grad_clip"]
+    kl_start_pos = cfg["data"]["kl_start_pos"]
+    kl_chunk_size = cfg["train"]["kl_chunk_size"]
+    max_steps = cfg["train"]["max_steps"]
+    eval_every = cfg["eval"]["every_steps"]
+    log_every = cfg["log"]["log_every"]
+
+    student.train()
+    best_kl = float("inf")
+    global_step = 0
+    run_summary = {
+        "config": args.config,
+        "run_name": cfg["log"]["wandb_run"],
+        "student": cfg["model"]["student"],
+        "teacher": cfg["model"]["teacher"],
+        "start_time": int(time.time()),
+    }
+
+    while global_step < max_steps:
+        t0 = time.time()
+        batch = train_loader.next_batch(samples_per_step)
+        optimizer.zero_grad(set_to_none=True)
+        batch_n = len(batch)
+        kl_sum = 0.0
+
+        for mb_start in range(0, batch_n, micro_batch_size):
+            micro = batch[mb_start : mb_start + micro_batch_size]
+            mb_n = len(micro)
+            ids, mask = collate_pad(micro, pad_id)
+            teacher_ids = ids.to(teacher_input_device, non_blocking=True)
+            teacher_mask = mask.to(teacher_input_device, non_blocking=True)
+            student_ids = ids.to(student_device, non_blocking=True)
+            student_mask = mask.to(student_device, non_blocking=True)
+
+            with torch.no_grad():
+                t_logits = teacher_forward(teacher, teacher_ids, teacher_mask, student_device)
+            s_logits = student(input_ids=student_ids, attention_mask=student_mask).logits
+            loss = kl_loss_masked(
+                s_logits,
+                t_logits,
+                student_mask,
+                start_pos=kl_start_pos,
+                chunk_size=kl_chunk_size,
+            )
+            scaled = loss * (mb_n / batch_n)
+            scaled.backward()
+            kl_sum += loss.item() * mb_n
+            del teacher_ids, teacher_mask, student_ids, student_mask, t_logits, s_logits, loss, scaled
+
+        if grad_clip > 0:
+            torch.nn.utils.clip_grad_norm_(student.parameters(), grad_clip)
+        optimizer.step()
+        scheduler.step()
+        global_step += 1
+
+        elapsed = time.time() - t0
+        kl_avg = kl_sum / batch_n
+        lr_now = scheduler.get_last_lr()[0]
+        record = {
+            "step": global_step,
+            "train_kl": kl_avg,
+            "lr": lr_now,
+            "step_time_s": elapsed,
+        }
+        log_jsonl(metrics_path, record)
+
+        if global_step % log_every == 0:
+            log.info(
+                f"step {global_step}/{max_steps} | kl {kl_avg:.6f} | "
+                f"lr {lr_now:.2e} | {elapsed:.2f}s"
+            )
+            if use_wandb:
+                import wandb
+
+                wandb.log(
+                    {
+                        "train/kl": kl_avg,
+                        "train/lr": lr_now,
+                        "perf/step_time_s": elapsed,
+                    },
+                    step=global_step,
+                )
+
+        if global_step % eval_every == 0:
+            eval_kl = evaluate(
+                student,
+                teacher,
+                eval_batches,
+                pad_id,
+                kl_start_pos,
+                kl_chunk_size,
+                student_device,
+                teacher_input_device,
+            )
+            log.info(f"eval @ step {global_step}: kl={eval_kl:.6f} (best={best_kl:.6f})")
+            log_jsonl(metrics_path, {"step": global_step, "eval_kl": eval_kl})
+            if use_wandb:
+                import wandb
+
+                wandb.log({"eval/kl": eval_kl}, step=global_step)
+            if eval_kl < best_kl:
+                best_kl = eval_kl
+                save_best(student, tokenizer, output_dir, global_step, eval_kl)
+            student.train()
+
+        if global_step % 10 == 0:
+            gc.collect()
+            torch.cuda.empty_cache()
+
+    final_eval = evaluate(
+        student,
+        teacher,
+        eval_batches,
+        pad_id,
+        kl_start_pos,
+        kl_chunk_size,
+        student_device,
+        teacher_input_device,
+    )
+    log.info(f"final eval: kl={final_eval:.6f} (best={best_kl:.6f})")
+    if final_eval < best_kl:
+        best_kl = final_eval
+        save_best(student, tokenizer, output_dir, global_step, final_eval)
+
+    run_summary.update(
+        {
+            "end_time": int(time.time()),
+            "best_eval_kl": best_kl,
+            "final_eval_kl": final_eval,
+            "max_steps": max_steps,
+            "student_total_params": total_params,
+            "student_trainable_params": trainable_params,
+        }
+    )
+    log_jsonl(experiment_log, run_summary)
+
+    if use_wandb:
+        import wandb
+
+        wandb.log({"eval/final_kl": final_eval, "eval/best_kl": best_kl}, step=global_step)
+        wandb.finish()
+
+
+if __name__ == "__main__":
+    main()