Project-Chronos 0.1.1__py3-none-any.whl

chronos/data/flexible_dataset.py

@@ -0,0 +1,309 @@
+"""
+chronos/data/flexible_dataset.py
+
+Streaming JSONL loaders with byte-offset indexing. Three variants cover
+the six Chronos training stages without ever loading the file into RAM:
+
+  FlexibleDataset       — pretrain / language modeling. Auto-detects
+                           common text-field schemas (text, content,
+                           instruction+output, conversations, messages).
+  StreamingSFTDataset   — supervised fine-tuning. Uses the tokenizer's
+                           chat template; only assistant tokens contribute
+                           to the loss (label = -100 elsewhere). Drop-in
+                           for minimind's SFTDataset.
+  StreamingDPODataset   — preference pairs. Returns the dict shape DPO /
+                           ORPO trainers expect.
+
+Each dataset stores only `offsets: list[int]` (8 bytes/sample) plus the
+file path. `__getitem__` does seek+readline+json.loads on demand, so a
+multi-GB corpus uses < 1 GB of resident memory.
+
+The file handle is opened lazily per process (re-opens after fork so
+DataLoader num_workers > 0 still works) and protected by a thread lock
+so num_workers = 0 paths with concurrent access are safe.
+"""
+from __future__ import annotations
+
+import json
+import os
+import threading
+
+import torch
+from torch.utils.data import Dataset
+
+
+# ── Field-name auto-detection (FlexibleDataset only) ──────────────────
+
+_TEXT_KEYS = ("text", "content", "story", "passage", "document", "article")
+_PAIR_KEYS = (
+    ("instruction", "output"),
+    ("instruction", "response"),
+    ("prompt", "response"),
+    ("prompt", "completion"),
+    ("input", "output"),
+    ("question", "answer"),
+    ("query", "answer"),
+)
+
+
+def _extract_text(record: dict) -> str:
+    for k in _TEXT_KEYS:
+        if k in record and record[k]:
+            return str(record[k])
+    for k1, k2 in _PAIR_KEYS:
+        if k1 in record and k2 in record:
+            parts = [str(record[k1]).strip(), str(record[k2]).strip()]
+            return "\n".join(p for p in parts if p)
+    if "conversations" in record and isinstance(record["conversations"], list):
+        return " ".join(
+            str(c.get("value", c.get("content", "")))
+            for c in record["conversations"]
+        )
+    if "messages" in record and isinstance(record["messages"], list):
+        return " ".join(str(m.get("content", "")) for m in record["messages"])
+    parts = [str(v) for v in record.values()
+             if isinstance(v, (str, int, float)) and str(v).strip()]
+    return " ".join(parts)
+
+
+# ── Shared streaming infrastructure ───────────────────────────────────
+
+class _StreamingJSONLBase(Dataset):
+    """Base class providing byte-offset index + lazy per-process file handle."""
+
+    def __init__(self, data_path: str):
+        self.data_path = os.path.abspath(data_path)
+        self.offsets: list[int] = []
+        with open(self.data_path, "rb") as f:
+            pos = 0
+            while True:
+                line = f.readline()
+                if not line:
+                    break
+                if line.strip():
+                    self.offsets.append(pos)
+                pos = f.tell()
+        if not self.offsets:
+            raise ValueError(f"No JSONL records in {data_path}")
+        self._fh = None
+        self._fh_lock = threading.Lock()
+
+    def _get_fh(self):
+        pid = os.getpid()
+        if self._fh is None or getattr(self._fh, "_chronos_pid", None) != pid:
+            if self._fh is not None:
+                try:
+                    self._fh.close()
+                except Exception:
+                    pass
+            self._fh = open(self.data_path, "rb")
+            self._fh._chronos_pid = pid
+        return self._fh
+
+    def _read_record(self, index: int) -> dict:
+        offset = self.offsets[index]
+        with self._fh_lock:
+            fh = self._get_fh()
+            fh.seek(offset)
+            raw = fh.readline()
+        return json.loads(raw.decode("utf-8"))
+
+    def __len__(self):
+        return len(self.offsets)
+
+    def __del__(self):
+        try:
+            if self._fh is not None:
+                self._fh.close()
+        except Exception:
+            pass
+
+
+# ── Pretrain / generic LM ────────────────────────────────────────────
+
+class FlexibleDataset(_StreamingJSONLBase):
+    """Pretrain-style streaming dataset; auto-detects the text field."""
+
+    def __init__(self, data_path: str, tokenizer, max_length: int = 512):
+        super().__init__(data_path)
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        first = self._read_record(0)
+        detected = _extract_text(first)
+        field_hint = (
+            next((k for k in _TEXT_KEYS if k in first), None)
+            or next((f"{k1}+{k2}" for k1, k2 in _PAIR_KEYS
+                     if k1 in first and k2 in first), None)
+            or ("conversations" if "conversations" in first else "auto")
+        )
+        print(f"[FlexibleDataset] {len(self.offsets)} records (streaming), "
+              f"detected format: '{field_hint}', "
+              f"sample preview: {detected[:80]!r}")
+
+    def __getitem__(self, index):
+        record = self._read_record(index)
+        text = _extract_text(record)
+        tok = self.tokenizer(
+            text, add_special_tokens=False,
+            max_length=self.max_length - 2, truncation=True,
+        )
+        tokens = (
+            [self.tokenizer.bos_token_id]
+            + tok.input_ids
+            + [self.tokenizer.eos_token_id]
+        )
+        pad_id = self.tokenizer.pad_token_id or 0
+        tokens = tokens + [pad_id] * (self.max_length - len(tokens))
+        input_ids = torch.tensor(tokens, dtype=torch.long)
+        labels = input_ids.clone()
+        labels[input_ids == pad_id] = -100
+        return input_ids, labels
+
+
+# ── SFT (conversations → assistant-token loss mask) ─────────────────
+
+class StreamingSFTDataset(_StreamingJSONLBase):
+    """Streaming counterpart to minimind's SFTDataset.
+
+    Reads `conversations: [{role, content, ...}, ...]` records, applies
+    the tokenizer's chat template, and masks every position outside an
+    assistant turn with -100 so the loss only fires on assistant tokens.
+    """
+
+    def __init__(self, data_path: str, tokenizer, max_length: int = 1024):
+        super().__init__(data_path)
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        # Sentinels marking the assistant turn boundary in encoded form.
+        # Same formulation minimind uses (`bos`+'assistant\n' / `eos`+'\n').
+        self.bos_id = tokenizer(
+            f'{tokenizer.bos_token}assistant\n', add_special_tokens=False
+        ).input_ids
+        self.eos_id = tokenizer(
+            f'{tokenizer.eos_token}\n', add_special_tokens=False
+        ).input_ids
+        print(f"[StreamingSFTDataset] {len(self.offsets)} records (streaming)")
+
+    def _build_prompt(self, conversations) -> str:
+        messages = []
+        tools = None
+        for m in conversations:
+            m = dict(m)
+            if m.get("role") == "system" and m.get("tools"):
+                tools = json.loads(m["tools"]) if isinstance(m["tools"], str) else m["tools"]
+            if m.get("tool_calls") and isinstance(m["tool_calls"], str):
+                m["tool_calls"] = json.loads(m["tool_calls"])
+            messages.append(m)
+        return self.tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=False, tools=tools,
+        )
+
+    def _label_mask(self, input_ids: list[int]) -> list[int]:
+        """-100 outside every <bos>assistant\\n ... <eos>\\n span."""
+        labels = [-100] * len(input_ids)
+        i = 0
+        nb = len(self.bos_id)
+        ne = len(self.eos_id)
+        while i < len(input_ids):
+            if input_ids[i:i + nb] == self.bos_id:
+                start = i + nb
+                end = start
+                while end < len(input_ids):
+                    if input_ids[end:end + ne] == self.eos_id:
+                        break
+                    end += 1
+                for j in range(start, min(end + ne, self.max_length)):
+                    labels[j] = input_ids[j]
+                i = end + ne if end < len(input_ids) else len(input_ids)
+            else:
+                i += 1
+        return labels
+
+    def __getitem__(self, index):
+        record = self._read_record(index)
+        convs = record.get("conversations")
+        if not convs:
+            # Graceful fall-through for mixed corpora — treat as pretrain.
+            text = _extract_text(record)
+            ids = self.tokenizer(text, truncation=True,
+                                 max_length=self.max_length).input_ids
+        else:
+            prompt = self._build_prompt(convs)
+            ids = self.tokenizer(prompt).input_ids[:self.max_length]
+        pad_id = self.tokenizer.pad_token_id or 0
+        ids = ids + [pad_id] * (self.max_length - len(ids))
+        labels = self._label_mask(ids)
+        return (
+            torch.tensor(ids, dtype=torch.long),
+            torch.tensor(labels, dtype=torch.long),
+        )
+
+
+# ── DPO / ORPO (chosen + rejected pair) ─────────────────────────────
+
+class StreamingDPODataset(_StreamingJSONLBase):
+    """Streaming counterpart to minimind's DPODataset.
+
+    Returns the dict shape the DPO/ORPO trainers expect:
+      {x_chosen, y_chosen, mask_chosen, x_rejected, y_rejected, mask_rejected}
+    """
+
+    def __init__(self, data_path: str, tokenizer, max_length: int = 4096):
+        super().__init__(data_path)
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.padding = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
+        self.bos_id = tokenizer(
+            f'{tokenizer.bos_token}assistant\n', add_special_tokens=False
+        ).input_ids
+        self.eos_id = tokenizer(
+            f'{tokenizer.eos_token}\n', add_special_tokens=False
+        ).input_ids
+        print(f"[StreamingDPODataset] {len(self.offsets)} pairs (streaming)")
+
+    def _loss_mask(self, ids: list[int]) -> list[int]:
+        mask = [0] * len(ids)
+        i = 0
+        nb = len(self.bos_id); ne = len(self.eos_id)
+        while i < len(ids):
+            if ids[i:i + nb] == self.bos_id:
+                start = i + nb
+                end = start
+                while end < len(ids):
+                    if ids[end:end + ne] == self.eos_id:
+                        break
+                    end += 1
+                for j in range(start, min(end + ne, self.max_length)):
+                    mask[j] = 1
+                i = end + ne if end < len(ids) else len(ids)
+            else:
+                i += 1
+        return mask
+
+    def __getitem__(self, index):
+        rec = self._read_record(index)
+        chosen = rec["chosen"]; rejected = rec["rejected"]
+
+        chosen_prompt = self.tokenizer.apply_chat_template(
+            chosen, tokenize=False, add_generation_prompt=False,
+        )
+        rejected_prompt = self.tokenizer.apply_chat_template(
+            rejected, tokenize=False, add_generation_prompt=False,
+        )
+
+        c_enc = self.tokenizer(chosen_prompt, truncation=True,
+                               max_length=self.max_length, padding="max_length")
+        r_enc = self.tokenizer(rejected_prompt, truncation=True,
+                               max_length=self.max_length, padding="max_length")
+        c_ids = c_enc["input_ids"]
+        r_ids = r_enc["input_ids"]
+        c_mask = self._loss_mask(c_ids)
+        r_mask = self._loss_mask(r_ids)
+        return {
+            "x_chosen":     torch.tensor(c_ids[:-1], dtype=torch.long),
+            "y_chosen":     torch.tensor(c_ids[1:],  dtype=torch.long),
+            "mask_chosen":  torch.tensor(c_mask[1:], dtype=torch.long),
+            "x_rejected":   torch.tensor(r_ids[:-1], dtype=torch.long),
+            "y_rejected":   torch.tensor(r_ids[1:],  dtype=torch.long),
+            "mask_rejected":torch.tensor(r_mask[1:], dtype=torch.long),
+        }

chronos/deps.py

@@ -0,0 +1,72 @@
+"""
+chronos/deps.py — Automatic dependency bootstrap for minimind-master kernel.
+
+When minimind-master is not found locally, this module downloads it from
+GitHub and makes it importable. Called automatically on first import of
+any chronos.model module.
+
+Legal notice: minimind is licensed under Apache-2.0.
+See THIRD_PARTY_NOTICES.md for full attribution.
+"""
+import sys
+import os
+import subprocess
+import importlib
+from pathlib import Path
+
+MINIMIND_REPO = "https://github.com/jingyaogong/minimind"
+MINIMIND_COMMIT = "master"  # pin to a tag/commit for reproducibility
+
+# Candidate local paths (in priority order)
+_CANDIDATE_PATHS = [
+    Path(__file__).parent.parent.parent / "minimind-master",   # sibling dir
+    Path(__file__).parent.parent / "minimind-master",
+    Path.home() / ".cache" / "chronos" / "minimind-master",
+]
+
+
+def _find_local() -> Path | None:
+    for p in _CANDIDATE_PATHS:
+        if (p / "model" / "model_minimind.py").exists():
+            return p
+    return None
+
+
+def _download(target: Path) -> Path:
+    target.parent.mkdir(parents=True, exist_ok=True)
+    print(f"[Chronos] minimind not found locally. Downloading from {MINIMIND_REPO} ...")
+    print(f"[Chronos] Target: {target}")
+    print(f"[Chronos] License: Apache-2.0  |  Copyright jingyaogong")
+    subprocess.run(
+        ["git", "clone", "--depth", "1", "--branch", MINIMIND_COMMIT,
+         MINIMIND_REPO, str(target)],
+        check=True,
+    )
+    print(f"[Chronos] minimind downloaded to {target}")
+    return target
+
+
+def ensure_minimind() -> Path:
+    """
+    Ensure minimind-master is available and on sys.path.
+    Returns the path to the minimind root directory.
+    """
+    local = _find_local()
+    if local is None:
+        cache_dir = Path.home() / ".cache" / "chronos" / "minimind-master"
+        local = _download(cache_dir)
+
+    root = str(local)
+    if root not in sys.path:
+        sys.path.insert(0, root)
+    return local
+
+
+def get_tokenizer_path() -> str:
+    """Return the path to the minimind tokenizer directory."""
+    root = ensure_minimind()
+    return str(root / "model")
+
+
+# Run on import
+_minimind_root = ensure_minimind()

chronos/eval/__init__.py

@@ -0,0 +1,2 @@
+from .io_profiler import evaluate_lookahead_accuracy, estimate_cache_hit_rate
+from .benchmark import compute_perplexity, measure_throughput, kv_cache_memory_bytes

chronos/eval/benchmark.py

@@ -0,0 +1,159 @@
+"""
+chronos/eval/benchmark.py
+
+Phase 3 end-to-end evaluation:
+- Perplexity on held-out text (accuracy proxy)
+- Throughput (tokens/s) with and without async prefetch
+- KV cache memory usage comparison: standard vs hybrid attention
+- Soft gating fallback rate
+
+Usage:
+    python -m chronos.eval.benchmark \
+        --model_path ./out/chronos_512_moe.pth \
+        --data_path  ./dataset/eval.jsonl \
+        --device     cpu \
+        --max_new_tokens 64
+"""
+import sys
+import chronos.deps  # ensure minimind on sys.path
+
+import argparse
+import math
+import time
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer
+
+from chronos.model.config import ChronosConfig
+from chronos.model.model_chronos import ChronosForCausalLM
+from chronos.model.hybrid_attention import MLAAttention, SlidingWindowAttention
+
+
+def compute_perplexity(model, dataloader, device, max_batches=50):
+    model.eval()
+    total_loss, total_tokens = 0.0, 0
+    with torch.no_grad():
+        for i, (input_ids, labels) in enumerate(dataloader):
+            if i >= max_batches:
+                break
+            input_ids, labels = input_ids.to(device), labels.to(device)
+            out, _ = model(input_ids, labels=labels)
+            n_tokens = (labels != -100).sum().item()
+            total_loss += out.loss.item() * n_tokens
+            total_tokens += n_tokens
+    ppl = math.exp(total_loss / max(total_tokens, 1))
+    return ppl
+
+
+def measure_throughput(model, input_ids, device, max_new_tokens=64):
+    model.eval()
+    input_ids = input_ids.to(device)
+    t0 = time.monotonic()
+    with torch.no_grad():
+        out, lp = model(input_ids, use_cache=True)
+        past_kv = out.past_key_values
+        for _ in range(max_new_tokens):
+            x = torch.randint(0, model.config.vocab_size, (1, 1), device=device)
+            out2, _ = model(x, past_key_values=past_kv, use_cache=True)
+            past_kv = out2.past_key_values
+    elapsed = time.monotonic() - t0
+    return max_new_tokens / elapsed
+
+
+def kv_cache_memory_bytes(past_key_values, model):
+    """Estimate total KV cache memory in bytes (fp32)."""
+    total = 0
+    for i, kv in enumerate(past_key_values):
+        if kv is None:
+            continue
+        for t in kv:
+            total += t.numel() * t.element_size()
+    return total
+
+
+def run_benchmark(args):
+    config = ChronosConfig(
+        hidden_size=args.hidden_size,
+        num_hidden_layers=args.num_hidden_layers,
+        num_experts=args.num_experts,
+        use_moe=True,
+        use_hybrid_attention=True,
+        kv_latent_dim=args.kv_latent_dim,
+        sliding_window_size=args.sliding_window_size,
+        vram_budget_gb=args.vram_budget_gb,
+    )
+    model = ChronosForCausalLM(config).to(args.device)
+
+    if args.model_path:
+        weights = torch.load(args.model_path, map_location=args.device)
+        model.load_state_dict(weights, strict=False)
+        print(f"Loaded: {args.model_path}")
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        chronos.deps.get_tokenizer_path()
+    )
+
+    from dataset.lm_dataset import PretrainDataset
+    dataset = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
+    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False)
+
+    print("\n=== Phase 3 Benchmark ===")
+
+    # 1. Perplexity
+    ppl = compute_perplexity(model, loader, args.device)
+    print(f"Perplexity: {ppl:.2f}")
+
+    # 2. Throughput
+    sample_ids = torch.randint(0, config.vocab_size, (1, 32))
+    tps = measure_throughput(model, sample_ids, args.device, args.max_new_tokens)
+    print(f"Throughput: {tps:.1f} tokens/s")
+
+    # 3. KV cache memory after N decode steps
+    model.eval()
+    with torch.no_grad():
+        x = torch.randint(0, config.vocab_size, (1, 32), device=args.device)
+        out, _ = model(x, use_cache=True)
+        past_kv = out.past_key_values
+        for _ in range(args.max_new_tokens):
+            xn = torch.randint(0, config.vocab_size, (1, 1), device=args.device)
+            out2, _ = model(xn, past_key_values=past_kv, use_cache=True)
+            past_kv = out2.past_key_values
+
+    kv_bytes = kv_cache_memory_bytes(past_kv, model)
+    print(f"KV cache after {args.max_new_tokens} steps: {kv_bytes/1024:.1f} KB")
+
+    # 4. Layer-wise cache sizes
+    print("\nLayer-wise KV cache:")
+    for i, kv in enumerate(past_kv):
+        if kv is not None:
+            layer_type = type(model.model.layers[i].self_attn).__name__
+            seq_len = kv[0].shape[1]
+            print(f"  Layer {i:2d} ({layer_type:25s}): seq_len={seq_len:4d}  "
+                  f"bytes={sum(t.numel()*t.element_size() for t in kv)//1024}KB")
+
+    print(f"\n=== Summary ===")
+    print(f"  PPL={ppl:.2f}  TPS={tps:.1f}  KV={kv_bytes//1024}KB  "
+          f"hybrid_attn=True  window={config.sliding_window_size}")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, default=None)
+    parser.add_argument("--data_path", type=str, required=True)
+    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument("--max_seq_len", type=int, default=256)
+    parser.add_argument("--max_new_tokens", type=int, default=64)
+    parser.add_argument("--hidden_size", type=int, default=512)
+    parser.add_argument("--num_hidden_layers", type=int, default=8)
+    parser.add_argument("--num_experts", type=int, default=4)
+    parser.add_argument("--kv_latent_dim", type=int, default=64)
+    parser.add_argument("--sliding_window_size", type=int, default=2048)
+    parser.add_argument("--vram_budget_gb", type=float, default=4.0)
+    args = parser.parse_args()
+    run_benchmark(args)
+
+
+if __name__ == "__main__":
+    main()