from __future__ import annotations import argparse import io import json import math import os import queue import random import threading import time from dataclasses import dataclass from pathlib import Path from typing import Any import numpy as np import torch from datasets import load_dataset from huggingface_hub import HfApi from safetensors.torch import load_file, save_file from torch.nn import functional as F from cascade_model import CascadeConfig, CascadeForCausalLM @dataclass(frozen=True) class Source: repo: str config: str | None weight: float kind: str SOURCES = ( Source("HuggingFaceFW/fineweb-edu", "sample-10BT", 0.58, "text"), Source("HuggingFaceTB/finemath", "finemath-4plus", 0.22, "text"), Source("HuggingFaceTB/smollm-corpus", "cosmopedia-v2", 0.10, "text"), Source("open-r1/OpenR1-Math-220k", "default", 0.07, "math"), Source("open-r1/codeforces-cots", "solutions_w_editorials_py_decontaminated", 0.03, "code"), ) def arguments() -> argparse.Namespace: p = argparse.ArgumentParser() p.add_argument("--repo", default="Asilarkness/cascade-1b-logic") p.add_argument("--token-file", type=Path, default=Path("/marimo/hf.txt")) p.add_argument("--output", type=Path, default=Path("/marimo/storage/cascade-1b")) p.add_argument("--batch-size", type=int, default=128) p.add_argument("--sequence-length", type=int, default=1024) p.add_argument("--max-patches", type=int, default=192) p.add_argument("--patch-rate", type=float, default=0.125) p.add_argument("--max-steps", type=int, default=100000) p.add_argument("--warmup-steps", type=int, default=1000) p.add_argument("--adam-lr", type=float, default=3e-4) p.add_argument("--muon-lr", type=float, default=0.02) p.add_argument("--weight-decay", type=float, default=0.1) p.add_argument("--log-every", type=int, default=10) p.add_argument("--eval-every", type=int, default=1000) p.add_argument("--checkpoint-every", type=int, default=1000) p.add_argument("--first-checkpoint", type=int, default=20) p.add_argument("--state-upload-every", type=int, default=100) p.add_argument("--bootstrap-mib", type=int, default=8) p.add_argument("--validation-batches", type=int, default=4) p.add_argument("--seed", type=int, default=3407) p.add_argument("--compile", action="store_true") p.add_argument("--compile-blocks", action="store_true") p.add_argument("--resume", action="store_true") p.add_argument("--probe-only", action="store_true") p.add_argument("--probe-batches", default="64,96,128,160,192,224,256") p.add_argument("--no-publish", action="store_true") return p.parse_args() def flatten_text(value: Any) -> list[str]: if isinstance(value, str): return [value] if isinstance(value, dict): if isinstance(value.get("content"), str): return [value["content"]] out: list[str] = [] for item in value.values(): out.extend(flatten_text(item)) return out if isinstance(value, (list, tuple)): out = [] for item in value: out.extend(flatten_text(item)) return out return [] def document(row: dict[str, Any], kind: str) -> str: if kind == "text": return str(row.get("text") or "") if kind == "math": problem = str(row.get("problem") or "") solution = str(row.get("solution") or "") if not solution: values = flatten_text(row.get("generations", [])) solution = values[0] if values else "" answer = str(row.get("answer") or "") return f"Problem:\n{problem}\n\nSolution:\n{solution}\n\nAnswer:\n{answer}" keys = ( "description", "input_format", "output_format", "editorial", "prompt", "generation", "messages", "accepted_solutions", "solutions", "solution", "code", ) values: list[str] = [] for key in keys: values.extend(flatten_text(row.get(key))) return "\n\n".join(dict.fromkeys(x for x in values if len(x.strip()) > 20)) class SourceReader: def __init__(self, source: Source, token: str, seed: int) -> None: self.source = source self.token = token self.seed = seed self.queue: queue.Queue[str] = queue.Queue(32) self.errors: queue.Queue[str] = queue.Queue(8) self.thread = threading.Thread(target=self.run, daemon=True) self.thread.start() def run(self) -> None: attempt = 0 while True: try: data = load_dataset( self.source.repo, self.source.config, split="train", streaming=True, token=self.token, ).shuffle(seed=self.seed + attempt, buffer_size=64) for row in data: text = document(row, self.source.kind).strip() if len(text) >= 80: self.queue.put(text[:4_000_000]) attempt += 1 except Exception as error: if not self.errors.full(): self.errors.put(f"{type(error).__name__}: {str(error)[:240]}") attempt += 1 time.sleep(min(30, 2**min(attempt, 5))) class StreamPacker: def __init__(self, batch_size: int, sequence_length: int, token: str, seed: int) -> None: self.batch_size = batch_size self.sequence_length = sequence_length self.rng = random.Random(seed) self.readers = [SourceReader(source, token, seed + i * 1009) for i, source in enumerate(SOURCES)] self.buffer = bytearray() self.offset = 0 self.documents = {source.repo + ":" + str(source.config): 0 for source in SOURCES} self.bytes = {source.repo + ":" + str(source.config): 0 for source in SOURCES} def wait_essential(self, timeout: float = 300) -> None: deadline = time.time() + timeout while time.time() < deadline: if not self.readers[0].queue.empty() and not self.readers[1].queue.empty(): return time.sleep(0.1) errors = {r.source.repo: list(r.errors.queue) for r in self.readers[:2]} raise TimeoutError(json.dumps(errors)) def get_document(self) -> bytes: deadline = time.time() + 300 while time.time() < deadline: ready = [(reader, reader.source.weight) for reader in self.readers if not reader.queue.empty()] if ready: reader = min( (x[0] for x in ready), key=lambda item: self.bytes[item.source.repo + ":" + str(item.source.config)] / item.source.weight, ) try: text = reader.queue.get_nowait() except queue.Empty: continue key = reader.source.repo + ":" + str(reader.source.config) self.documents[key] += 1 payload = (text + "\n\n<|endoftext|>\n\n").encode("utf-8", errors="replace")[:65536] self.bytes[key] += len(payload) return payload time.sleep(0.05) errors = {r.source.repo: list(r.errors.queue) for r in self.readers} raise TimeoutError(json.dumps(errors)) def next(self) -> torch.Tensor: needed = self.batch_size * (self.sequence_length + 1) while len(self.buffer) - self.offset < needed: self.buffer.extend(self.get_document()) chunk = np.frombuffer(self.buffer, dtype=np.uint8, count=needed, offset=self.offset).copy() self.offset += needed if self.offset >= 16_000_000: del self.buffer[: self.offset] self.offset = 0 return torch.from_numpy(chunk).reshape(self.batch_size, self.sequence_length + 1).pin_memory() def entropy_bootstrap(packer: StreamPacker, target_mib: int, patch_rate: float) -> tuple[torch.Tensor, float, list[torch.Tensor]]: counts = np.zeros(256 * 256, dtype=np.int64) batches: list[torch.Tensor] = [] size = 0 target = target_mib * 2**20 while size < target: batch = packer.next() batches.append(batch) values = batch.numpy().reshape(-1).astype(np.int64) counts += np.bincount(values[:-1] * 256 + values[1:], minlength=256 * 256) size += values.size print(json.dumps({"bootstrap_mib": round(size / 2**20, 2)}), flush=True) matrix = counts.reshape(256, 256).astype(np.float64) rows = matrix.sum(1, keepdims=True) surprise = -np.log((matrix + 0.1) / (rows + 25.6)) flat_count = matrix.reshape(-1) flat_score = surprise.reshape(-1) order = np.argsort(flat_score)[::-1] cumulative = np.cumsum(flat_count[order]) index = min(np.searchsorted(cumulative, flat_count.sum() * patch_rate), len(order) - 1) threshold = float(flat_score[order[index]]) return torch.from_numpy(surprise).float(), threshold, batches class Optimizers: def __init__(self, model: CascadeForCausalLM, args: argparse.Namespace) -> None: matrices: list[torch.nn.Parameter] = [] other: list[torch.nn.Parameter] = [] seen: set[int] = set() for name, parameter in model.named_parameters(): if id(parameter) in seen: continue seen.add(id(parameter)) excluded = "embed" in name or "head" in name (matrices if parameter.ndim == 2 and not excluded else other).append(parameter) self.muon = torch.optim.Muon( matrices, lr=args.muon_lr, weight_decay=args.weight_decay, momentum=0.95, nesterov=True, ns_steps=5, adjust_lr_fn="match_rms_adamw", ) self.adam = torch.optim.AdamW( other, lr=args.adam_lr, betas=(0.9, 0.95), eps=1e-8, weight_decay=args.weight_decay, fused=True, ) self.base_muon = args.muon_lr self.base_adam = args.adam_lr def zero_grad(self) -> None: self.muon.zero_grad(set_to_none=True) self.adam.zero_grad(set_to_none=True) def step(self) -> None: self.muon.step() self.adam.step() def set_scale(self, scale: float) -> None: for group in self.muon.param_groups: group["lr"] = self.base_muon * scale for group in self.adam.param_groups: group["lr"] = self.base_adam * scale def lr_scale(step: int, args: argparse.Namespace) -> float: if step < args.warmup_steps: return (step + 1) / max(1, args.warmup_steps) progress = (step - args.warmup_steps) / max(1, args.max_steps - args.warmup_steps) return 0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(1.0, progress))) class Hub: def __init__(self, repo: str, token: str, output: Path) -> None: self.repo = repo self.api = HfApi(token=token) self.output = output self.thread: threading.Thread | None = None def upload_sources(self, files: list[Path]) -> None: for path in files: if path.exists(): target = "README.md" if path.name == "MODEL_CARD.md" else path.name self.api.upload_file(path_or_fileobj=path, path_in_repo=target, repo_id=self.repo) def state(self, values: dict[str, Any]) -> None: payload = io.BytesIO(json.dumps(values, indent=2).encode()) self.api.upload_file(path_or_fileobj=payload, path_in_repo="training_state.json", repo_id=self.repo) def checkpoint(self, step: int, model: CascadeForCausalLM, config: CascadeConfig, metrics: dict[str, Any]) -> bool: if self.thread is not None and self.thread.is_alive(): return False directory = self.output / "checkpoint-latest" directory.mkdir(parents=True, exist_ok=True) state = { key: value.detach().cpu().contiguous() for key, value in model.state_dict().items() if key != "head.weight" } config.save(directory / "config.json") (directory / "training_state.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8") def save_and_upload() -> None: save_file(state, directory / "model.safetensors", metadata={"format": "pt", "step": str(step)}) for attempt in range(5): try: self.api.upload_folder( folder_path=directory, path_in_repo=f"checkpoints/step-{step:08d}", repo_id=self.repo, commit_message=f"checkpoint step {step}", ) return except Exception: if attempt == 4: raise time.sleep(15 * (attempt + 1)) self.thread = threading.Thread(target=save_and_upload, daemon=True) self.thread.start() return True def wait(self) -> None: if self.thread is not None: self.thread.join() @torch.no_grad() def evaluate(model: CascadeForCausalLM, batches: list[torch.Tensor], routes: int) -> float: model.eval() losses = [] for route in range(routes): for batch in batches: raw = batch.cuda(non_blocking=True).long() logits, _, _ = model(raw[:, :-1], route) losses.append(float(F.cross_entropy(logits.flatten(0, 1), raw[:, 1:].flatten()))) model.train() return sum(losses) / len(losses) def load_resume(model: CascadeForCausalLM, output: Path) -> int: directory = output / "checkpoint-latest" weights = directory / "model.safetensors" state_path = directory / "training_state.json" if not weights.exists() or not state_path.exists(): return 0 missing, unexpected = model.load_state_dict(load_file(weights), strict=False) if unexpected or any(x != "head.weight" for x in missing): raise RuntimeError({"missing": missing, "unexpected": unexpected}) return int(json.loads(state_path.read_text())["step"]) def probe(args: argparse.Namespace) -> None: config = CascadeConfig(sequence_length=args.sequence_length) model = CascadeForCausalLM(config).cuda().to(torch.bfloat16) print(json.dumps(model.parameter_counts()), flush=True) for batch_size in map(int, args.probe_batches.split(",")): try: torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() x = torch.randint(0, 256, (batch_size, args.sequence_length), device="cuda") started = time.perf_counter() logits, auxiliary, _ = model(x, 0) loss = F.cross_entropy(logits.flatten(0, 1), x.flatten()) + config.patch_aux_weight * auxiliary loss.backward() torch.cuda.synchronize() elapsed = time.perf_counter() - started print(json.dumps({"batch": batch_size, "seconds": elapsed, "tokens_s": batch_size * args.sequence_length / elapsed, "vram_gib": torch.cuda.max_memory_allocated() / 2**30}), flush=True) model.zero_grad(set_to_none=True) del x, logits, loss, auxiliary except torch.cuda.OutOfMemoryError: print(json.dumps({"batch": batch_size, "oom": True}), flush=True) model.zero_grad(set_to_none=True) torch.cuda.empty_cache() break def train(args: argparse.Namespace) -> None: token = args.token_file.read_text().strip() args.output.mkdir(parents=True, exist_ok=True) torch.manual_seed(args.seed) random.seed(args.seed) np.random.seed(args.seed) torch.set_float32_matmul_precision("high") packer = StreamPacker(args.batch_size, args.sequence_length, token, args.seed) packer.wait_essential() surprise, threshold, bootstrap = entropy_bootstrap(packer, args.bootstrap_mib, args.patch_rate) validation = [x[: min(8, args.batch_size)] for x in bootstrap[: args.validation_batches]] warm = bootstrap[args.validation_batches :] config = CascadeConfig( sequence_length=args.sequence_length, max_patches=args.max_patches, patch_rate=args.patch_rate, entropy_threshold=threshold, ) model = CascadeForCausalLM(config, surprise).cuda().to(torch.bfloat16) start_step = load_resume(model, args.output) if args.resume else 0 if start_step: config.entropy_threshold = float(model.entropy_threshold) counts = model.parameter_counts() hub = None if args.no_publish else Hub(args.repo, token, args.output) here = Path(__file__).resolve().parent if hub is not None: hub.upload_sources([here / "cascade_model.py", Path(__file__).resolve(), here / "MODEL_CARD.md"]) config.save(args.output / "config.json") if hub is not None: hub.api.upload_file(path_or_fileobj=args.output / "config.json", path_in_repo="config.json", repo_id=args.repo) optimizers = Optimizers(model, args) initial_validation = evaluate(model, validation, config.route_paths) print(json.dumps({"parameters": counts, "threshold": threshold, "initial_validation": initial_validation, "sources": packer.documents}), flush=True) if args.compile_blocks: model.compile_blocks() compiled = torch.compile(model, mode="max-autotune-no-cudagraphs") if args.compile else model history = args.output / "training.jsonl" tokens = start_step * args.batch_size * args.sequence_length started = time.perf_counter() last_time = started last_tokens = tokens last_metrics: dict[str, Any] = {"step": start_step, "tokens": tokens, "validation_loss": initial_validation} for step in range(start_step + 1, args.max_steps + 1): batch = warm.pop(0) if warm else packer.next() raw = batch.cuda(non_blocking=True).long() x, y = raw[:, :-1], raw[:, 1:] route = (step - 1) % config.route_paths scale = lr_scale(step - 1, args) optimizers.set_scale(scale) optimizers.zero_grad() logits, auxiliary, stats = compiled(x, route) language = F.cross_entropy(logits.flatten(0, 1), y.flatten()) loss = language + config.patch_aux_weight * auxiliary loss.backward() norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizers.step() tokens += x.numel() if step == 1 or step % args.log_every == 0: torch.cuda.synchronize() now = time.perf_counter() metrics = { "step": step, "tokens": tokens, "parameters": counts, "batch_size": args.batch_size, "sequence_length": args.sequence_length, "loss": float(language.detach()), "auxiliary_loss": float(auxiliary.detach()), "gradient_norm": float(norm.detach()), "route": route, "lr_scale": scale, "tokens_per_second": (tokens - last_tokens) / (now - last_time), "elapsed_seconds": now - started, "average_patch_bytes": x.numel() / float(stats["patches"]), "patch_coalesced_rate": float(stats["overflow"]) / x.shape[0], "gate": float(stats["gate"].detach()), "peak_vram_gib": torch.cuda.max_memory_allocated() / 2**30, "documents": dict(packer.documents), "source_bytes": dict(packer.bytes), } last_metrics = metrics with history.open("a", encoding="utf-8") as file: file.write(json.dumps(metrics) + "\n") print(json.dumps(metrics), flush=True) last_time, last_tokens = now, tokens if step % args.eval_every == 0: last_metrics["validation_loss"] = evaluate(compiled, validation, config.route_paths) print(json.dumps({"step": step, "validation_loss": last_metrics["validation_loss"]}), flush=True) if hub is not None and (step == args.first_checkpoint or step % args.checkpoint_every == 0): last_metrics["checkpoint_started"] = hub.checkpoint(step, model, config, last_metrics) print(json.dumps({"step": step, "checkpoint_started": last_metrics["checkpoint_started"]}), flush=True) if hub is not None and step % args.state_upload_every == 0: hub.state(last_metrics) if hub is not None: hub.checkpoint(args.max_steps, model, config, last_metrics) hub.wait() hub.state(last_metrics) def main() -> None: args = arguments() if not torch.cuda.is_available(): raise RuntimeError("CUDA required") if args.probe_only: probe(args) else: train(args) if __name__ == "__main__": main()