Project-Chronos 0.1.1__py3-none-any.whl

chronos/trainer/optim_utils.py

@@ -0,0 +1,101 @@
+"""
+chronos/trainer/optim_utils.py
+
+Shared optimizer + LR-schedule helpers used by every Chronos training stage
+(pretrain, SFT, DPO, ORPO, GRPO, distill) and the WebUI Train tab.
+
+Two reasons to centralize this:
+
+1. Decay groups: AdamW with weight_decay applied uniformly to all params
+   destabilizes LayerNorm / RMSNorm scales and hurts embeddings. Standard
+   recipe is to exclude norms, biases, and embeddings from decay. Doing
+   this in one place means every trainer benefits from the same correct
+   recipe and we don't get drift.
+
+2. Warmup-then-cosine schedule. minimind's get_lr is pure cosine (effective
+   LR at step 0 is already 0.55·base_lr), which destabilizes the first few
+   hundred steps with small batches. Linear warmup over the first ~5% of
+   steps then cosine decay to 10% of base is the boring-and-correct choice.
+"""
+from __future__ import annotations
+
+import math
+from typing import Iterable, Tuple
+
+import torch
+import torch.nn as nn
+
+
+def _decay_groups(model: nn.Module) -> Tuple[list, list]:
+    """Split trainable parameters into (decay, no_decay) lists.
+
+    No-decay: any parameter whose owning module is a LayerNorm/RMSNorm,
+    any 1-D parameter (typically bias / norm scale), and any embedding
+    weight (regardless of dim — embeddings are looked-up rows, not
+    rotated subspaces, so decay would push *all* tokens toward zero).
+    """
+    decay, no_decay = [], []
+    seen = set()
+    for mod in model.modules():
+        is_norm = isinstance(mod, nn.LayerNorm) or mod.__class__.__name__.endswith("RMSNorm")
+        is_embed = isinstance(mod, nn.Embedding)
+        for p in mod.parameters(recurse=False):
+            if not p.requires_grad or id(p) in seen:
+                continue
+            seen.add(id(p))
+            if is_norm or is_embed or p.ndim < 2:
+                no_decay.append(p)
+            else:
+                decay.append(p)
+    return decay, no_decay
+
+
+def build_optimizer(
+    model: nn.Module,
+    lr: float,
+    *,
+    weight_decay: float = 0.01,
+    betas: Tuple[float, float] = (0.9, 0.95),
+    eps: float = 1e-8,
+) -> torch.optim.AdamW:
+    """AdamW with the standard transformer recipe applied to a Chronos model."""
+    decay, no_decay = _decay_groups(model)
+    return torch.optim.AdamW(
+        [
+            {"params": decay,    "weight_decay": float(weight_decay)},
+            {"params": no_decay, "weight_decay": 0.0},
+        ],
+        lr=float(lr), betas=betas, eps=eps,
+    )
+
+
+def get_lr(
+    current_step: int,
+    total_steps: int,
+    base_lr: float,
+    *,
+    warmup_steps: int | None = None,
+    min_lr_ratio: float = 0.1,
+) -> float:
+    """Linear warmup → cosine decay to `min_lr_ratio · base_lr`.
+
+    `current_step` is 1-indexed (matches every trainer's `enumerate(loader, start=1)`).
+    Falls back to a constant `base_lr` when `total_steps <= 0`.
+    """
+    if total_steps is None or total_steps <= 0:
+        return float(base_lr)
+    if warmup_steps is None:
+        warmup_steps = max(1, min(500, total_steps // 20))
+    step = max(1, int(current_step))
+    if step <= warmup_steps:
+        return float(base_lr) * (step / float(warmup_steps))
+    progress = (step - warmup_steps) / max(1, (total_steps - warmup_steps))
+    progress = min(1.0, max(0.0, progress))
+    cos = 0.5 * (1.0 + math.cos(math.pi * progress))
+    return float(base_lr) * (min_lr_ratio + (1.0 - min_lr_ratio) * cos)
+
+
+def apply_lr(optimizer: torch.optim.Optimizer, lr: float) -> None:
+    """Set the learning rate on every param-group in `optimizer`."""
+    for g in optimizer.param_groups:
+        g["lr"] = float(lr)

chronos/trainer/orpo_trainer.py

@@ -0,0 +1,189 @@
+"""
+chronos/trainer/orpo_trainer.py — Stage 4 ORPO trainer.
+
+ORPO (Hong et al. 2024) combines supervised NLL on the chosen response
+with an odds-ratio penalty on the rejected response, in a single objective
+that does NOT require a separate reference model. This matters for
+memory-constrained Chronos deployments: DPO needs two full models in
+memory; ORPO needs one.
+
+Loss:
+    L_NLL  = −mean_token log p_θ(y_w | x)       # supervised on chosen
+    L_OR   = −log σ( β · log [ odds(y_w) / odds(y_l) ] )
+             where odds(y) = p_θ(y) / (1 - p_θ(y))
+    L_ORPO = L_NLL + λ · L_OR
+
+Plus Chronos mixin (temporal / lookahead / aux) and router KL anchor.
+"""
+from __future__ import annotations
+
+import os
+import time
+from contextlib import nullcontext
+
+import torch
+import torch.nn.functional as F
+from torch import optim
+from torch.utils.data import DataLoader
+
+import chronos.deps  # noqa
+from trainer.trainer_utils import Logger, is_main_process  # type: ignore
+from chronos.data.flexible_dataset import StreamingDPODataset as DPODataset
+
+from chronos.model.config import ChronosConfig
+from chronos.model.model_chronos import ChronosForCausalLM
+from chronos.model.temporal_loss import lookahead_supervision_loss
+from chronos.trainer.optim_utils import get_lr
+from chronos.trainer.loss_mixin import (
+    chronos_loss_term,
+    collect_router_probs,
+    router_kl_anchor,
+    capture_reference_routing,
+)
+
+
+def _mean_logprob(logits: torch.Tensor, labels: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    """Mean log-prob of `labels` under `logits`, averaged over masked positions."""
+    log_probs = F.log_softmax(logits, dim=-1)
+    lp = log_probs.gather(dim=-1, index=labels.unsqueeze(-1)).squeeze(-1)  # [B, T]
+    lp = (lp * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1.0)
+    return lp
+
+
+class ChronosORPOTrainer:
+    def __init__(self, model: ChronosForCausalLM, config: ChronosConfig, args, tokenizer):
+        self.model = model
+        self.config = config
+        self.args = args
+        self.device = args.device
+        self.beta = float(getattr(args, "beta", 0.1))
+        self.lambda_or = float(getattr(args, "lambda_or", 0.1))
+
+        device_type = "cuda" if "cuda" in str(self.device) else "cpu"
+        dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+        self.autocast_ctx = nullcontext() if device_type == "cpu" \
+            else torch.cuda.amp.autocast(dtype=dtype)
+        self.scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == "float16"))
+        from chronos.trainer.optim_utils import build_optimizer
+        self.optimizer = build_optimizer(
+            model, lr=args.learning_rate,
+            weight_decay=float(getattr(args, "weight_decay", 0.01)),
+        )
+        self._router_ref = None
+
+    def set_calibration_batch(self, x: torch.Tensor):
+        self._router_ref = capture_reference_routing(self.model, x, self.device)
+
+    def train_step(self, batch, step, total_steps):
+        x_c = batch["x_chosen"].to(self.device)
+        x_r = batch["x_rejected"].to(self.device)
+        y_c = batch["y_chosen"].to(self.device)
+        y_r = batch["y_rejected"].to(self.device)
+        m_c = batch["mask_chosen"].to(self.device).float()
+        m_r = batch["mask_rejected"].to(self.device).float()
+        x = torch.cat([x_c, x_r], dim=0)
+        y = torch.cat([y_c, y_r], dim=0)
+
+        lr = get_lr(step, total_steps, self.args.learning_rate)
+        for pg in self.optimizer.param_groups:
+            pg["lr"] = lr
+
+        with self.autocast_ctx:
+            outputs, lookahead_probs = self.model(x)
+            B2 = x_c.shape[0]
+            logits_c = outputs.logits[:B2]
+            logits_r = outputs.logits[B2:]
+
+            lp_c = _mean_logprob(logits_c, y_c, m_c)   # [B]
+            lp_r = _mean_logprob(logits_r, y_r, m_r)
+
+            # NLL on chosen
+            nll = -lp_c.mean()
+
+            # Odds-ratio term (numerically stable via log-space)
+            # log odds(y) = lp - log(1 - exp(lp)). Use log(1 - exp(lp)) ≈
+            # log1p(-exp(lp)) with clamping to avoid NaN at lp ≈ 0.
+            def _log_odds(lp):
+                one_minus = (1.0 - lp.exp()).clamp_min(1e-6)
+                return lp - one_minus.log()
+            log_odds_ratio = _log_odds(lp_c) - _log_odds(lp_r)
+            l_or = -F.logsigmoid(self.beta * log_odds_ratio).mean()
+
+            base = nll + self.lambda_or * l_or
+
+            loss = chronos_loss_term(
+                self.model, base, lookahead_probs, self.config,
+                aux_loss=outputs.aux_loss,
+            )
+
+            anc_val = 0.0
+            lam = float(getattr(self.config, "lambda_router_anchor", 0.0))
+            if lam > 0.0 and self._router_ref is not None:
+                r4 = collect_router_probs(self.model)
+                if r4 is not None and r4.shape[1] > 1:
+                    cur = r4.mean(dim=2)
+                    ref = self._router_ref
+                    T = min(cur.shape[1], ref.shape[1])
+                    Bm = min(cur.shape[0], ref.shape[0])
+                    anc = router_kl_anchor(cur[:Bm, :T, :], ref[:Bm, :T, :], lam)
+                    loss = loss + anc
+                    anc_val = float(anc.item())
+
+            la_val = 0.0
+            r4 = collect_router_probs(self.model)
+            if (r4 is not None and r4.shape[1] > 1
+                    and lookahead_probs is not None
+                    and self.config.lookahead_steps > 0):
+                la_val = float(lookahead_supervision_loss(
+                    lookahead_probs, r4.mean(dim=2).detach(),
+                    self.config.lookahead_steps,
+                ).item())
+
+            loss = loss / self.args.accumulation_steps
+
+        self.scaler.scale(loss).backward()
+
+        if step % self.args.accumulation_steps == 0:
+            self.scaler.unscale_(self.optimizer)
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.grad_clip)
+            self.scaler.step(self.optimizer)
+            self.scaler.update()
+            self.optimizer.zero_grad(set_to_none=True)
+
+        return (
+            loss.item() * self.args.accumulation_steps,
+            float(nll.item()),
+            float(l_or.item()),
+            la_val,
+            anc_val,
+        )
+
+    def train_epoch(self, epoch, loader, iters, max_steps=None):
+        self.model.train()
+        total_steps = self.args.epochs * iters
+        for step, batch in enumerate(loader, start=1):
+            if max_steps is not None and step > max_steps:
+                break
+            loss, nll, lor, la, anc = self.train_step(
+                batch, epoch * iters + step, total_steps,
+            )
+            if step % self.args.log_interval == 0 or step == iters:
+                lr = self.optimizer.param_groups[-1]["lr"]
+                Logger(
+                    f"[ORPO] Epoch[{epoch+1}/{self.args.epochs}]({step}/{iters}) "
+                    f"loss:{loss:.4f} nll:{nll:.4f} or:{lor:.4f} "
+                    f"la:{la:.4f} anchor:{anc:.4f} lr:{lr:.2e}"
+                )
+
+    def _save(self, epoch, step):
+        os.makedirs(self.args.save_dir, exist_ok=True)
+        ckp = os.path.join(self.args.save_dir, f"orpo_{self.config.hidden_size}_moe.pth")
+        self.model.eval()
+        state = self.model.state_dict()
+        torch.save({k: v.half().cpu() for k, v in state.items()}, ckp)
+        self.model.train()
+
+
+def build_orpo_loader(data_path, tokenizer, max_seq_len, batch_size):
+    ds = DPODataset(data_path, tokenizer, max_length=max_seq_len)
+    return DataLoader(ds, batch_size=batch_size, shuffle=True, num_workers=0)

chronos/trainer/reward.py

@@ -0,0 +1,130 @@
+"""
+chronos/trainer/reward.py
+
+Pluggable reward models for GRPO. Two built-ins:
+
+- ``ToyReward`` — length + vocab-purity + repetition-penalty heuristic.
+  No external dependencies, used in smoke tests and when no real reward
+  model is configured.
+- ``LMRewardModel`` — wraps an HF AutoModel that exposes a ``.get_score(...)``
+  method (the minimind convention, see
+  ``minimind-master/trainer/trainer_utils.py:160``). Loads lazily so
+  importing this file never forces an HF network download.
+
+Any callable ``score(prompt: str, response_tokens: List[int], tokenizer) -> float``
+is accepted by ``ChronosGRPOTrainer``; the classes below simply package
+that protocol.
+"""
+from __future__ import annotations
+
+from typing import List, Optional
+
+
+TOY_VOCAB = {
+    "the", "quick", "brown", "fox", "machine", "learning", "model", "expert",
+    "router", "cluster", "memory", "cache", "prefetch", "latency", "throughput",
+    "tensor", "training", "inference", "decode", "weight", "gradient",
+    "softmax", "attention", "layer", "embedding", "vocab", "sequence",
+    "loss", "optimizer", "epoch", "batch", "step", "token", "logit",
+    "minimind", "chronos", "hybrid", "mixture", "of", "experts", "sparse",
+    "dense", "neural", "network", "transformer", "rope", "rotary", "kv",
+}
+
+
+class ToyReward:
+    """Zero-dependency heuristic reward: length + vocab-purity − repetition."""
+
+    def __init__(self, target_len: int = 24, vocab: Optional[set] = None):
+        self.target_len = target_len
+        self.vocab = vocab or TOY_VOCAB
+
+    def __call__(self, prompt: str, response_tokens: List[int], tokenizer) -> float:
+        return self.score(prompt, response_tokens, tokenizer)
+
+    def score(self, prompt: str, response_tokens: List[int], tokenizer) -> float:
+        text = tokenizer.decode(response_tokens, skip_special_tokens=True)
+        words = text.split()
+        if not words:
+            return -1.0
+        length_reward = min(len(words) / max(self.target_len, 1), 1.0)
+        on_vocab = sum(1 for w in words if w.lower() in self.vocab)
+        vocab_purity = on_vocab / len(words)
+        dupes = sum(1 for a, b in zip(words[:-1], words[1:]) if a == b)
+        repetition_penalty = -dupes / max(len(words), 1)
+        return length_reward + vocab_purity + repetition_penalty
+
+
+class LMRewardModel:
+    """HF-backed reward model.
+
+    Wraps any ``AutoModel`` that exposes ``get_score(tokenizer, messages) -> float``
+    (the reward-model ABI used by minimind and several open reward-model
+    checkpoints). The model is loaded lazily on the first call so that
+    importing this module is free.
+
+    Example::
+
+        reward = LMRewardModel("path/to/internlm2-1_8b-reward")
+        r = reward.score(prompt, response_tokens, tokenizer)
+
+    The score is clipped to [-3.0, 3.0] to keep GRPO advantages well-scaled,
+    matching the minimind convention.
+    """
+
+    def __init__(self, model_path: str, device: str = "cpu",
+                 dtype_str: str = "float16", score_clip: float = 3.0):
+        self.model_path = model_path
+        self.device = device
+        self.dtype_str = dtype_str
+        self.score_clip = float(score_clip)
+        self._model = None
+        self._tokenizer = None
+
+    def _load(self):
+        if self._model is not None:
+            return
+        import torch
+        from transformers import AutoTokenizer, AutoModel
+        dtype = torch.float16 if self.dtype_str == "float16" else torch.float32
+        self._tokenizer = AutoTokenizer.from_pretrained(
+            self.model_path, trust_remote_code=True,
+        )
+        self._model = AutoModel.from_pretrained(
+            self.model_path, torch_dtype=dtype, trust_remote_code=True,
+        ).to(self.device).eval()
+
+    def __call__(self, prompt: str, response_tokens: List[int], tokenizer) -> float:
+        return self.score(prompt, response_tokens, tokenizer)
+
+    def score(self, prompt: str, response_tokens: List[int], tokenizer) -> float:
+        self._load()
+        response = tokenizer.decode(response_tokens, skip_special_tokens=True)
+        messages = [
+            {"role": "user", "content": prompt},
+            {"role": "assistant", "content": response},
+        ]
+        s = self._model.get_score(self._tokenizer, messages)
+        return max(min(float(s), self.score_clip), -self.score_clip)
+
+
+def build_reward_fn(spec: Optional[str]) -> callable:
+    """Convert a CLI spec string into a reward callable.
+
+    Spec forms:
+        None | ""  | "toy"                 → ToyReward()
+        "lm:/path/to/reward-model"         → LMRewardModel(path)
+        "lm:/path/to/reward,device=cuda"   → LMRewardModel(path, device="cuda")
+    """
+    if not spec or spec == "toy":
+        return ToyReward()
+    if spec.startswith("lm:"):
+        body = spec[3:]
+        parts = body.split(",")
+        path = parts[0]
+        kwargs = {}
+        for kv in parts[1:]:
+            if "=" in kv:
+                k, v = kv.split("=", 1)
+                kwargs[k.strip()] = v.strip()
+        return LMRewardModel(path, **kwargs)
+    raise ValueError(f"unknown reward spec: {spec!r}")

chronos/trainer/sft_trainer.py

@@ -0,0 +1,189 @@
+"""
+chronos/trainer/sft_trainer.py
+
+Chronos-aware supervised fine-tuning trainer.
+
+Built around minimind's SFTDataset (jsonl with `conversations` field,
+tokenizer.apply_chat_template + per-assistant-token loss mask) so the data
+pipeline is identical to the one in minimind-master/trainer/train_full_sft.py.
+What we add:
+
+1. The model returns ``(outputs, lookahead_probs)``, which minimind's trainer
+   does not know about. We forward it explicitly.
+2. ``chronos_loss_term`` mixes in λ_balance·aux + λ_temporal·temporal +
+   λ_lookahead·lookahead on top of the base CE.
+3. Optional ``router_kl_anchor`` pulls the current router distribution back
+   toward the pretrain reference (default λ=0.01 for SFT — mild drift is
+   OK). The reference is captured once at stage start.
+"""
+from __future__ import annotations
+
+import os
+import time
+from contextlib import nullcontext
+
+import torch
+from torch import optim
+from torch.utils.data import DataLoader
+
+import chronos.deps  # ensure minimind on sys.path
+
+from trainer.trainer_utils import Logger, is_main_process  # type: ignore
+from chronos.data.flexible_dataset import StreamingSFTDataset
+
+from chronos.model.config import ChronosConfig
+from chronos.model.model_chronos import ChronosForCausalLM
+from chronos.model.temporal_loss import lookahead_supervision_loss
+from chronos.trainer.optim_utils import get_lr
+from chronos.trainer.loss_mixin import (
+    chronos_loss_term,
+    collect_router_probs,
+    router_kl_anchor,
+    capture_reference_routing,
+)
+
+
+class ChronosSFTTrainer:
+    """Stage 2 (SFT) trainer for Chronos."""
+
+    def __init__(self, model: ChronosForCausalLM, config: ChronosConfig, args, tokenizer):
+        self.model = model
+        self.config = config
+        self.args = args
+        self.tokenizer = tokenizer
+        self.device = args.device
+
+        device_type = "cuda" if "cuda" in str(self.device) else "cpu"
+        dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
+        self.autocast_ctx = (
+            nullcontext() if device_type == "cpu"
+            else torch.cuda.amp.autocast(dtype=dtype)
+        )
+        self.scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == "float16"))
+        from chronos.trainer.optim_utils import build_optimizer
+        self.optimizer = build_optimizer(
+            model, lr=args.learning_rate,
+            weight_decay=float(getattr(args, "weight_decay", 0.01)),
+        )
+
+        # Router anchor reference — captured at the start of a stage
+        self._router_ref = None
+        self._calib_batch = None
+
+    def set_calibration_batch(self, input_ids: torch.Tensor):
+        """Provide a representative input_ids batch; at stage start the
+        current (pre-SFT) model's router is frozen as the anchor teacher."""
+        self._calib_batch = input_ids
+        self._router_ref = capture_reference_routing(self.model, input_ids, self.device)
+
+    def train_step(self, input_ids, labels, step, total_steps):
+        input_ids = input_ids.to(self.device)
+        labels = labels.to(self.device)
+
+        lr = get_lr(step, total_steps, self.args.learning_rate)
+        for pg in self.optimizer.param_groups:
+            pg['lr'] = lr
+
+        with self.autocast_ctx:
+            outputs, lookahead_probs = self.model(input_ids, labels=labels)
+            ce_loss = outputs.loss
+            aux_loss = outputs.aux_loss
+
+            loss = chronos_loss_term(
+                self.model, ce_loss, lookahead_probs, self.config,
+                aux_loss=aux_loss,
+            )
+
+            # Router KL anchor
+            anchor_val = 0.0
+            lam = float(getattr(self.config, "lambda_router_anchor", 0.0))
+            if lam > 0.0 and self._router_ref is not None:
+                router_4d = collect_router_probs(self.model)
+                if router_4d is not None and router_4d.shape[1] > 1:
+                    cur = router_4d.mean(dim=2)
+                    # Shape-align: calibration batch shape may differ; crop
+                    # the shorter axis on both sides.
+                    ref = self._router_ref
+                    T = min(cur.shape[1], ref.shape[1])
+                    B = min(cur.shape[0], ref.shape[0])
+                    anchor = router_kl_anchor(
+                        cur[:B, :T, :], ref[:B, :T, :], lam,
+                    )
+                    loss = loss + anchor
+                    anchor_val = float(anchor.item())
+
+            # Standalone lookahead loss for logging
+            la_val = 0.0
+            router_4d = collect_router_probs(self.model)
+            if (router_4d is not None and router_4d.shape[1] > 1
+                    and lookahead_probs is not None
+                    and self.config.lookahead_steps > 0):
+                teacher = router_4d.mean(dim=2).detach()
+                la = lookahead_supervision_loss(
+                    lookahead_probs, teacher, self.config.lookahead_steps,
+                )
+                la_val = float(la.item())
+
+            loss = loss / self.args.accumulation_steps
+
+        self.scaler.scale(loss).backward()
+
+        if step % self.args.accumulation_steps == 0:
+            self.scaler.unscale_(self.optimizer)
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.grad_clip)
+            self.scaler.step(self.optimizer)
+            self.scaler.update()
+            self.optimizer.zero_grad(set_to_none=True)
+
+        return (
+            loss.item() * self.args.accumulation_steps,
+            ce_loss.item(),
+            float(aux_loss.item()) if aux_loss is not None else 0.0,
+            la_val,
+            anchor_val,
+        )
+
+    def train_epoch(self, epoch, loader, iters, start_step=0, wandb=None, max_steps=None):
+        self.model.train()
+        start_time = time.time()
+        total_steps = self.args.epochs * iters
+
+        for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
+            if max_steps is not None and step > max_steps:
+                break
+            loss, ce, aux, la, anc = self.train_step(
+                input_ids, labels, epoch * iters + step, total_steps,
+            )
+            if step % self.args.log_interval == 0 or step == iters:
+                elapsed = time.time() - start_time
+                eta = elapsed / max(step - start_step, 1) * (iters - step) // 60
+                lr = self.optimizer.param_groups[-1]['lr']
+                Logger(
+                    f'[SFT] Epoch[{epoch+1}/{self.args.epochs}]({step}/{iters}) '
+                    f'loss:{loss:.4f} ce:{ce:.4f} aux:{aux:.4f} '
+                    f'la:{la:.4f} anchor:{anc:.4f} lr:{lr:.2e} eta:{eta:.1f}min'
+                )
+                if wandb:
+                    wandb.log({
+                        "sft_loss": loss, "sft_ce": ce, "sft_aux": aux,
+                        "sft_lookahead": la, "sft_anchor": anc, "lr": lr,
+                    })
+
+            if (step % self.args.save_interval == 0 or step == iters) and is_main_process():
+                self._save(epoch, step)
+
+    def _save(self, epoch, step):
+        os.makedirs(self.args.save_dir, exist_ok=True)
+        ckp = os.path.join(
+            self.args.save_dir,
+            f'sft_{self.config.hidden_size}_moe.pth',
+        )
+        self.model.eval()
+        state = self.model.state_dict()
+        torch.save({k: v.half().cpu() for k, v in state.items()}, ckp)
+        self.model.train()
+
+
+def build_sft_loader(data_path, tokenizer, max_seq_len, batch_size):
+    ds = StreamingSFTDataset(data_path, tokenizer, max_length=max_seq_len)
+    return DataLoader(ds, batch_size=batch_size, shuffle=True, num_workers=0)

chronos/tuning/__init__.py

@@ -0,0 +1 @@
+from .chronos_auto_tuner import ChronosAutoTuner, ChronosSearchSpaceConfig