Project-Chronos 0.1.1__py3-none-any.whl

train_chronos_grpo.py

@@ -0,0 +1,88 @@
+"""train_chronos_grpo.py — Stage 5 (GRPO) entry point."""
+import argparse
+import os
+
+import chronos.deps  # noqa
+import torch
+from transformers import AutoTokenizer
+
+from chronos.backend import resolve_training_device
+from chronos.model.config import ChronosConfig
+from chronos.model.model_chronos import ChronosForCausalLM
+from chronos.trainer.grpo_trainer import ChronosGRPOTrainer, load_grpo_prompts
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--data_path", required=True)
+    p.add_argument("--from_weight", default="orpo")
+    p.add_argument("--save_dir", default="out")
+    p.add_argument("--steps", type=int, default=None)
+    p.add_argument("--epochs", type=int, default=1)
+    p.add_argument("--max_seq_len", type=int, default=96)
+    p.add_argument("--max_gen_len", type=int, default=24)
+    p.add_argument("--num_generations", type=int, default=4)
+    p.add_argument("--temperature", type=float, default=1.0)
+    p.add_argument("--learning_rate", type=float, default=5e-6)
+    p.add_argument("--accumulation_steps", type=int, default=1)
+    p.add_argument("--grad_clip", type=float, default=1.0)
+    p.add_argument("--log_interval", type=int, default=5)
+    p.add_argument("--save_interval", type=int, default=10000)
+    p.add_argument("--device", default="auto")
+    p.add_argument("--dtype", default="float16")
+    p.add_argument("--hidden_size", type=int, default=256)
+    p.add_argument("--num_hidden_layers", type=int, default=4)
+    p.add_argument("--num_experts", type=int, default=4)
+    p.add_argument("--beta", type=float, default=0.04, help="KL penalty weight")
+    p.add_argument("--lambda_router_anchor", type=float, default=0.1)
+    p.add_argument("--reward", type=str, default="toy",
+                   help="Reward spec: 'toy' (default) or 'lm:/path/to/reward-model'.")
+    args = p.parse_args()
+    selected_backend, resolved_device = resolve_training_device(args.device)
+    args.device = resolved_device
+
+    cfg = ChronosConfig(
+        hidden_size=args.hidden_size,
+        num_hidden_layers=args.num_hidden_layers,
+        num_experts=args.num_experts,
+        num_experts_per_tok=1,
+        max_position_embeddings=args.max_seq_len + args.max_gen_len,
+        flash_attn=False,
+        lambda_router_anchor=args.lambda_router_anchor,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(chronos.deps.get_tokenizer_path())
+    model = ChronosForCausalLM(cfg).to(args.device)
+
+    ckp_path = os.path.join(args.save_dir, f"{args.from_weight}_{cfg.hidden_size}_moe.pth")
+    if os.path.exists(ckp_path):
+        state = torch.load(ckp_path, map_location=args.device)
+        model.load_state_dict(state, strict=False)
+        print(f"[GRPO] Loaded {ckp_path}")
+    else:
+        print(f"[GRPO] No checkpoint at {ckp_path} — training from random init.")
+    print(f"[GRPO] Training backend: {selected_backend}  device={args.device}")
+
+    prompts = load_grpo_prompts(args.data_path, max_prompts=args.steps)
+    from chronos.trainer.reward import build_reward_fn
+    reward_fn = build_reward_fn(args.reward)
+    trainer = ChronosGRPOTrainer(model, cfg, args, tokenizer, reward_fn=reward_fn)
+
+    if cfg.lambda_router_anchor > 0 and prompts:
+        # Build a calibration batch from the first prompt
+        ptext = prompts[0]
+        calib = tokenizer(ptext, return_tensors="pt", truncation=True,
+                          max_length=args.max_seq_len).input_ids
+        trainer.set_calibration_batch(calib)
+        print("[GRPO] Router anchor reference captured.")
+
+    iters = len(prompts) if args.steps is None else min(args.steps, len(prompts))
+    for epoch in range(args.epochs):
+        trainer.train_epoch(epoch, prompts, iters, max_steps=args.steps)
+        if args.steps is not None:
+            break
+    trainer._save(epoch=args.epochs - 1, step=iters)
+    print(f"[GRPO] Done. Checkpoint: {args.save_dir}/grpo_{cfg.hidden_size}_moe.pth")
+
+
+if __name__ == "__main__":
+    main()

train_chronos_orpo.py

@@ -0,0 +1,81 @@
+"""train_chronos_orpo.py — Stage 4 (ORPO) entry point."""
+import argparse
+import os
+
+import chronos.deps  # noqa
+import torch
+from transformers import AutoTokenizer
+
+from chronos.backend import resolve_training_device
+from chronos.model.config import ChronosConfig
+from chronos.model.model_chronos import ChronosForCausalLM
+from chronos.trainer.orpo_trainer import ChronosORPOTrainer, build_orpo_loader
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--data_path", required=True)
+    p.add_argument("--from_weight", default="sft")
+    p.add_argument("--save_dir", default="out")
+    p.add_argument("--steps", type=int, default=None)
+    p.add_argument("--epochs", type=int, default=1)
+    p.add_argument("--batch_size", type=int, default=2)
+    p.add_argument("--max_seq_len", type=int, default=128)
+    p.add_argument("--learning_rate", type=float, default=5e-6)
+    p.add_argument("--accumulation_steps", type=int, default=1)
+    p.add_argument("--grad_clip", type=float, default=1.0)
+    p.add_argument("--log_interval", type=int, default=10)
+    p.add_argument("--save_interval", type=int, default=10000)
+    p.add_argument("--device", default="auto")
+    p.add_argument("--dtype", default="float16")
+    p.add_argument("--hidden_size", type=int, default=256)
+    p.add_argument("--num_hidden_layers", type=int, default=4)
+    p.add_argument("--num_experts", type=int, default=4)
+    p.add_argument("--beta", type=float, default=0.1)
+    p.add_argument("--lambda_or", type=float, default=0.1)
+    p.add_argument("--lambda_router_anchor", type=float, default=0.1)
+    args = p.parse_args()
+    selected_backend, resolved_device = resolve_training_device(args.device)
+    args.device = resolved_device
+
+    cfg = ChronosConfig(
+        hidden_size=args.hidden_size,
+        num_hidden_layers=args.num_hidden_layers,
+        num_experts=args.num_experts,
+        num_experts_per_tok=1,
+        max_position_embeddings=args.max_seq_len,
+        flash_attn=False,
+        lambda_router_anchor=args.lambda_router_anchor,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(chronos.deps.get_tokenizer_path())
+    model = ChronosForCausalLM(cfg).to(args.device)
+
+    ckp_path = os.path.join(args.save_dir, f"{args.from_weight}_{cfg.hidden_size}_moe.pth")
+    if os.path.exists(ckp_path):
+        state = torch.load(ckp_path, map_location=args.device)
+        model.load_state_dict(state, strict=False)
+        print(f"[ORPO] Loaded {ckp_path}")
+    else:
+        print(f"[ORPO] No checkpoint at {ckp_path} — training from random init.")
+    print(f"[ORPO] Training backend: {selected_backend}  device={args.device}")
+
+    loader = build_orpo_loader(args.data_path, tokenizer, args.max_seq_len, args.batch_size)
+    trainer = ChronosORPOTrainer(model, cfg, args, tokenizer)
+
+    if cfg.lambda_router_anchor > 0:
+        first = next(iter(loader))
+        calib = torch.cat([first["x_chosen"], first["x_rejected"]], dim=0)
+        trainer.set_calibration_batch(calib)
+        print("[ORPO] Router anchor reference captured.")
+
+    iters = len(loader) if args.steps is None else min(args.steps, len(loader))
+    for epoch in range(args.epochs):
+        trainer.train_epoch(epoch, loader, iters, max_steps=args.steps)
+        if args.steps is not None:
+            break
+    trainer._save(epoch=args.epochs - 1, step=iters)
+    print(f"[ORPO] Done. Checkpoint: {args.save_dir}/orpo_{cfg.hidden_size}_moe.pth")
+
+
+if __name__ == "__main__":
+    main()

train_chronos_sft.py

@@ -0,0 +1,85 @@
+"""train_chronos_sft.py — Stage 2 (SFT) entry point."""
+import argparse
+import os
+import sys
+
+import chronos.deps  # noqa
+import torch
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer
+
+from chronos.backend import resolve_training_device
+from chronos.model.config import ChronosConfig
+from chronos.model.model_chronos import ChronosForCausalLM
+from chronos.trainer.sft_trainer import ChronosSFTTrainer, build_sft_loader
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--data_path", required=True)
+    p.add_argument("--from_weight", default="chronos",
+                   help="Pretrained checkpoint stem (loaded from save_dir/<stem>_<H>_moe.pth).")
+    p.add_argument("--save_dir", default="out")
+    p.add_argument("--steps", type=int, default=None,
+                   help="If set, stop after N steps (smoke runs).")
+    p.add_argument("--epochs", type=int, default=1)
+    p.add_argument("--batch_size", type=int, default=2)
+    p.add_argument("--max_seq_len", type=int, default=128)
+    p.add_argument("--learning_rate", type=float, default=1e-4)
+    p.add_argument("--accumulation_steps", type=int, default=1)
+    p.add_argument("--grad_clip", type=float, default=1.0)
+    p.add_argument("--log_interval", type=int, default=10)
+    p.add_argument("--save_interval", type=int, default=10000)
+    p.add_argument("--device", default="auto")
+    p.add_argument("--dtype", default="float16")
+    p.add_argument("--hidden_size", type=int, default=256)
+    p.add_argument("--num_hidden_layers", type=int, default=4)
+    p.add_argument("--num_experts", type=int, default=4)
+    p.add_argument("--lambda_router_anchor", type=float, default=0.01,
+                   help="SFT default: weak anchor.")
+    args = p.parse_args()
+    selected_backend, resolved_device = resolve_training_device(args.device)
+    args.device = resolved_device
+
+    cfg = ChronosConfig(
+        hidden_size=args.hidden_size,
+        num_hidden_layers=args.num_hidden_layers,
+        num_experts=args.num_experts,
+        num_experts_per_tok=1,
+        max_position_embeddings=args.max_seq_len,
+        flash_attn=False,
+        lambda_router_anchor=args.lambda_router_anchor,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(chronos.deps.get_tokenizer_path())
+    model = ChronosForCausalLM(cfg).to(args.device)
+
+    # Load pretrain checkpoint if present
+    ckp_path = os.path.join(args.save_dir, f"{args.from_weight}_{cfg.hidden_size}_moe.pth")
+    if os.path.exists(ckp_path):
+        state = torch.load(ckp_path, map_location=args.device)
+        model.load_state_dict(state, strict=False)
+        print(f"[SFT] Loaded pretrain weights from {ckp_path}")
+    else:
+        print(f"[SFT] No pretrain checkpoint at {ckp_path} — training from random init.")
+    print(f"[SFT] Training backend: {selected_backend}  device={args.device}")
+
+    loader = build_sft_loader(args.data_path, tokenizer, args.max_seq_len, args.batch_size)
+    trainer = ChronosSFTTrainer(model, cfg, args, tokenizer)
+
+    # Capture router reference on the first batch BEFORE training
+    if cfg.lambda_router_anchor > 0:
+        first_batch = next(iter(loader))
+        trainer.set_calibration_batch(first_batch[0])
+        print("[SFT] Router anchor reference captured.")
+
+    iters = len(loader) if args.steps is None else min(args.steps, len(loader))
+    for epoch in range(args.epochs):
+        trainer.train_epoch(epoch, loader, iters, max_steps=args.steps)
+        if args.steps is not None:
+            break
+    trainer._save(epoch=args.epochs - 1, step=iters)
+    print(f"[SFT] Done. Checkpoint: {args.save_dir}/sft_{cfg.hidden_size}_moe.pth")
+
+
+if __name__ == "__main__":
+    main()

ui/__init__.py

@@ -0,0 +1,3 @@
+"""
+ui/__init__.py
+"""

ui/estimator.py

@@ -0,0 +1,204 @@
+"""
+ui/estimator.py
+
+Parameter-count and memory-footprint estimator for the Architecture
+Designer tab. Pure-math functions; no torch dependency.
+
+The numbers match an actual ``ChronosForCausalLM`` state_dict to within
+~1% for standard configs — the mismatch comes from bias terms and RMSNorm
+scale vectors which are tiny relative to the weight matrices we count.
+
+IMPORTANT: MiniMindConfig auto-derives ``intermediate_size`` from
+``hidden_size`` as ``ceil(H·π/64)·64`` unless the user explicitly sets
+it, and ``moe_intermediate_size`` defaults to ``intermediate_size``.
+The estimator replicates that auto-derivation here so the UI preview
+matches what the trainer actually builds. Callers can still override
+either value by passing it into ``ArchConfig`` explicitly.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+
+DTYPE_BYTES = {
+    "fp32": 4, "float32": 4,
+    "bf16": 2, "bfloat16": 2,
+    "fp16": 2, "float16": 2,
+    "int8": 1,
+    "nf4": 0.5,
+}
+
+
+def _minimind_intermediate(hidden_size: int) -> int:
+    """Replicates MiniMindConfig's auto-derivation rule."""
+    return int(math.ceil(hidden_size * math.pi / 64) * 64)
+
+
+@dataclass
+class ArchConfig:
+    hidden_size: int = 256
+    num_hidden_layers: int = 4
+    num_attention_heads: int = 8
+    num_key_value_heads: int = 4
+    vocab_size: int = 6400
+    # intermediate_size / moe_intermediate_size: when None, auto-derive from
+    # hidden_size via the MiniMind rule so the estimate matches the trainer.
+    intermediate_size: Optional[int] = None
+    moe_intermediate_size: Optional[int] = None
+    num_experts: int = 4
+    num_experts_per_tok: int = 1
+    num_shared_experts: int = 1
+    max_position_embeddings: int = 2048
+    lookahead_steps: int = 2
+    kv_latent_dim: int = 64
+    rope_dim: int = 32
+    sliding_window_size: int = 2048
+    use_hybrid_attention: bool = True
+    tie_word_embeddings: bool = True
+    dtype: str = "fp16"
+
+    def __post_init__(self):
+        auto = _minimind_intermediate(self.hidden_size)
+        if self.intermediate_size is None or self.intermediate_size <= 0:
+            self.intermediate_size = auto
+        if self.moe_intermediate_size is None or self.moe_intermediate_size <= 0:
+            self.moe_intermediate_size = auto
+
+
+def _attn_params(cfg: ArchConfig) -> int:
+    """Params per attention block (MLA/SW hybrid has similar Q/K/V/O sizes)."""
+    h = cfg.hidden_size
+    n_h = cfg.num_attention_heads
+    n_kv = cfg.num_key_value_heads
+    head_dim = h // n_h
+    q = h * (n_h * head_dim)
+    k = h * (n_kv * head_dim)
+    v = h * (n_kv * head_dim)
+    o = (n_h * head_dim) * h
+    if cfg.use_hybrid_attention:
+        # MLA adds latent KV projections; negligible vs full KV projections.
+        q += h * cfg.kv_latent_dim
+        k += cfg.kv_latent_dim * (n_kv * head_dim)
+        v += cfg.kv_latent_dim * (n_kv * head_dim)
+    return q + k + v + o
+
+
+def _expert_params(cfg: ArchConfig) -> int:
+    """Params per single expert FeedForward (SwiGLU: gate+up+down)."""
+    h = cfg.hidden_size
+    d = cfg.moe_intermediate_size
+    return 3 * h * d
+
+
+def _per_layer_params(cfg: ArchConfig) -> int:
+    """Attention + routing + all experts + shared experts + 2×RMSNorm."""
+    attn = _attn_params(cfg)
+    router = cfg.hidden_size * cfg.num_experts
+    experts = cfg.num_experts * _expert_params(cfg)
+    shared = cfg.num_shared_experts * _expert_params(cfg)
+    rmsnorm = 2 * cfg.hidden_size  # 2 norms per block
+    return attn + router + experts + shared + rmsnorm
+
+
+def _active_expert_slots(cfg: ArchConfig) -> int:
+    """Experts whose weights are *used* on a given token: top-k + shared."""
+    return cfg.num_experts_per_tok + cfg.num_shared_experts
+
+
+def total_params(cfg: ArchConfig) -> int:
+    """Full model parameter count."""
+    embed = cfg.vocab_size * cfg.hidden_size
+    layers = cfg.num_hidden_layers * _per_layer_params(cfg)
+    lookahead = (
+        cfg.hidden_size * (cfg.hidden_size // 4)
+        + (cfg.hidden_size // 4) * cfg.num_experts * (cfg.lookahead_steps + 1)
+    )
+    lm_head = 0 if cfg.tie_word_embeddings else embed
+    return int(embed + layers + lookahead + lm_head + cfg.hidden_size)  # + final norm
+
+
+def active_params(cfg: ArchConfig) -> int:
+    """Params touched by one token: embed + L·(attn + norm + (topK+shared)·expert)."""
+    embed = cfg.vocab_size * cfg.hidden_size
+    per_layer_active = (
+        _attn_params(cfg)
+        + 2 * cfg.hidden_size
+        + cfg.hidden_size * cfg.num_experts  # router always active
+        + _active_expert_slots(cfg) * _expert_params(cfg)
+    )
+    lm_head_active = 0 if cfg.tie_word_embeddings else embed
+    return int(embed + cfg.num_hidden_layers * per_layer_active + lm_head_active)
+
+
+def kv_cache_bytes(cfg: ArchConfig, seq_len: int = 1024) -> int:
+    """Approximate KV-cache bytes for a given sequence length."""
+    h = cfg.hidden_size
+    n_h = cfg.num_attention_heads
+    n_kv = cfg.num_key_value_heads
+    head_dim = h // n_h
+    bytes_per = DTYPE_BYTES.get(cfg.dtype, 2)
+    if cfg.use_hybrid_attention:
+        # Half the layers use MLA (compressed), half use sliding window.
+        # MLA cache ≈ kv_latent_dim per position; SW cache capped at window.
+        mla_layers = cfg.num_hidden_layers // 2
+        sw_layers = cfg.num_hidden_layers - mla_layers
+        mla_bytes = mla_layers * seq_len * cfg.kv_latent_dim * bytes_per
+        sw_bytes = sw_layers * min(seq_len, cfg.sliding_window_size) * 2 * n_kv * head_dim * bytes_per
+        return int(mla_bytes + sw_bytes)
+    return int(cfg.num_hidden_layers * seq_len * 2 * n_kv * head_dim * bytes_per)
+
+
+def memory_footprint(cfg: ArchConfig, seq_len: int = 1024) -> Dict[str, float]:
+    """Return a dict of byte sizes for display in the UI."""
+    bytes_per = DTYPE_BYTES.get(cfg.dtype, 2)
+    total_b = total_params(cfg) * bytes_per
+    active_b = active_params(cfg) * bytes_per
+    kv_b = kv_cache_bytes(cfg, seq_len=seq_len)
+    cold_b = total_b - active_b
+    return {
+        "total_bytes": total_b,
+        "active_bytes": active_b,
+        "cold_bytes": cold_b,          # candidates for SSD offload
+        "kv_cache_bytes": kv_b,
+        "vram_estimate_bytes": active_b + kv_b,
+        "ssd_estimate_bytes": cold_b,
+    }
+
+
+def estimated_decode_tps(cfg: ArchConfig, hit_rate: float = 0.85,
+                         ssd_gbps: float = 3.0, vram_gbps: float = 600.0) -> float:
+    """Order-of-magnitude decode throughput estimator.
+
+    Model: tps ≈ 1 / (t_compute + t_miss).
+      t_compute ≈ active_params · flop_per_param / compute_gflops  (simplified
+                  to a linear-in-active-params term with 1e-10 coefficient
+                  so the output lives in a sensible range for toy models).
+      t_miss    ≈ (1 − hit_rate) · expert_bytes_per_layer · num_layers / ssd_gbps.
+    Conservative; intended as a UI hint, not a benchmark.
+    """
+    bytes_per = DTYPE_BYTES.get(cfg.dtype, 2)
+    active_b = active_params(cfg) * bytes_per
+    t_compute = max(active_b / (vram_gbps * 1e9), 1e-6)
+    expert_b = _expert_params(cfg) * bytes_per
+    t_miss = (1.0 - hit_rate) * expert_b * cfg.num_hidden_layers / (ssd_gbps * 1e9)
+    return 1.0 / (t_compute + t_miss)
+
+
+def fmt_bytes(n: int) -> str:
+    for unit in ("B", "KB", "MB", "GB", "TB"):
+        if abs(n) < 1024:
+            return f"{n:.2f} {unit}"
+        n /= 1024
+    return f"{n:.2f} PB"
+
+
+def fmt_params(n: int) -> str:
+    if n >= 1e9:
+        return f"{n/1e9:.2f} B"
+    if n >= 1e6:
+        return f"{n/1e6:.2f} M"
+    if n >= 1e3:
+        return f"{n/1e3:.2f} K"
+    return str(n)

ui/gradio_compat.py

@@ -0,0 +1,72 @@
+"""Optional Gradio compatibility layer.
+
+The core smoke tests exercise helper functions from the UI modules in minimal
+CI environments where the optional WebUI dependency may not be installed. In a
+real WebUI install this module simply re-exports gradio. Without gradio it
+provides a tiny no-op component surface so imports and build_app() smoke tests
+still validate the wiring that does not require a browser runtime.
+"""
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+
+try:  # pragma: no cover - exercised when gradio is installed
+    import gradio as gr  # type: ignore
+    HAS_GRADIO = True
+except ModuleNotFoundError:  # pragma: no cover - exercised in slim CI
+    HAS_GRADIO = False
+
+    class _Component:
+        def __init__(self, *args, **kwargs):
+            self.args = args
+            self.kwargs = kwargs
+            self.value = kwargs.get("value")
+            self.label = kwargs.get("label")
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc, tb):
+            return False
+
+        def click(self, *args, **kwargs):
+            return self
+
+        def change(self, *args, **kwargs):
+            return self
+
+        def tick(self, *args, **kwargs):
+            return self
+
+        def launch(self, **kwargs):
+            self.launch_kwargs = kwargs
+            return self
+
+    def _update(**kwargs):
+        return SimpleNamespace(**kwargs)
+
+    gr = SimpleNamespace(
+        Accordion=_Component,
+        BarPlot=_Component,
+        Blocks=_Component,
+        Button=_Component,
+        Checkbox=_Component,
+        Column=_Component,
+        Dropdown=_Component,
+        File=_Component,
+        HTML=_Component,
+        JSON=_Component,
+        Markdown=_Component,
+        Number=_Component,
+        Plot=_Component,
+        Radio=_Component,
+        Row=_Component,
+        Slider=_Component,
+        State=_Component,
+        Tab=_Component,
+        Textbox=_Component,
+        Timer=_Component,
+        update=_update,
+    )
+