Upload 24 files

nanochat/checkpoint_manager.py

@@ -0,0 +1,285 @@
+"""
+Utilities for saving and loading model/optim/state checkpoints.
+"""
+import os
+import re
+import json
+import shutil
+import logging
+import torch
+from filelock import FileLock
+
+from nanochat.common import get_base_dir
+from nanochat.gpt import GPT, GPTConfig
+from nanochat.tokenizer import get_tokenizer
+from nanochat.common import setup_default_logging
+
+# Set up logging
+setup_default_logging()
+logger = logging.getLogger(__name__)
+def log0(message):
+    if int(os.environ.get('RANK', 0)) == 0:
+        logger.info(message)
+
+def _patch_missing_config_keys(model_config_kwargs):
+    """Add default values for new config keys missing in old checkpoints."""
+    # Old models were trained with full context (no sliding window)
+    if "window_pattern" not in model_config_kwargs:
+        model_config_kwargs["window_pattern"] = "L"
+        log0(f"Patching missing window_pattern in model config to 'L'")
+    # Checkpoints predating the quantum feed-forward network contain c_fc/c_proj
+    # tensors and must continue to instantiate the original dense MLP.
+    if "mlp_type" not in model_config_kwargs:
+        model_config_kwargs["mlp_type"] = "classical"
+        log0("Patching missing mlp_type in model config to 'classical'")
+    # Checkpoints predating the LFM2-style hybrid backbone were all-attention with
+    # legacy RoPE base and value embeddings enabled. These defaults reproduce them.
+    if "mixer_pattern" not in model_config_kwargs:
+        model_config_kwargs["mixer_pattern"] = "A"
+        log0("Patching missing mixer_pattern in model config to 'A'")
+    if "conv_kernel" not in model_config_kwargs:
+        model_config_kwargs["conv_kernel"] = 3
+    if "rope_theta" not in model_config_kwargs:
+        model_config_kwargs["rope_theta"] = 100000.0
+        log0("Patching missing rope_theta in model config to 100000.0")
+    if "ffn_pattern" not in model_config_kwargs:
+        model_config_kwargs["ffn_pattern"] = ""
+    if "use_value_embeddings" not in model_config_kwargs:
+        model_config_kwargs["use_value_embeddings"] = True
+
+def _patch_missing_keys(model_data, model_config):
+    """Add default values for new parameters that may be missing in old checkpoints."""
+    n_layer = model_config.n_layer
+    # resid_lambdas defaults to 1.0 (identity scaling)
+    if "resid_lambdas" not in model_data:
+        model_data["resid_lambdas"] = torch.ones(n_layer)
+        log0(f"Patching missing resid_lambdas in model data to 1.0")
+    # x0_lambdas defaults to 0.0 (disabled)
+    if "x0_lambdas" not in model_data:
+        model_data["x0_lambdas"] = torch.zeros(n_layer)
+        log0(f"Patching missing x0_lambdas in model data to 0.0")
+
+def save_checkpoint(checkpoint_dir, step, model_data, optimizer_data, meta_data, rank=0):
+    if rank == 0:
+        os.makedirs(checkpoint_dir, exist_ok=True)
+        # Save the model state parameters
+        model_path = os.path.join(checkpoint_dir, f"model_{step:06d}.pt")
+        torch.save(model_data, model_path)
+        logger.info(f"Saved model parameters to: {model_path}")
+        # Save the metadata dict as json
+        meta_path = os.path.join(checkpoint_dir, f"meta_{step:06d}.json")
+        with open(meta_path, "w", encoding="utf-8") as f:
+            json.dump(meta_data, f, indent=2)
+        logger.info(f"Saved metadata to: {meta_path}")
+    # Note that optimizer state is sharded across ranks, so each rank must save its own.
+    if optimizer_data is not None:
+        os.makedirs(checkpoint_dir, exist_ok=True)
+        optimizer_path = os.path.join(checkpoint_dir, f"optim_{step:06d}_rank{rank:d}.pt")
+        torch.save(optimizer_data, optimizer_path)
+        logger.info(f"Saved optimizer state to: {optimizer_path}")
+
+def load_checkpoint(checkpoint_dir, step, device, load_optimizer=False, rank=0):
+    # Load the model state
+    model_path = os.path.join(checkpoint_dir, f"model_{step:06d}.pt")
+    model_data = torch.load(model_path, map_location=device)
+    # Load the optimizer state if requested
+    optimizer_data = None
+    if load_optimizer:
+        optimizer_path = os.path.join(checkpoint_dir, f"optim_{step:06d}_rank{rank:d}.pt")
+        optimizer_data = torch.load(optimizer_path, map_location=device)
+    # Load the metadata
+    meta_path = os.path.join(checkpoint_dir, f"meta_{step:06d}.json")
+    with open(meta_path, "r", encoding="utf-8") as f:
+        meta_data = json.load(f)
+    return model_data, optimizer_data, meta_data
+
+
+def build_model(checkpoint_dir, step, device, phase):
+    """
+    A bunch of repetitive code to build a model from a given checkpoint.
+    Returns:
+    - base model - uncompiled, not wrapped in DDP
+    - tokenizer
+    - meta data saved during base model training
+    """
+    assert phase in ["train", "eval"], f"Invalid phase: {phase}"
+    model_data, optimizer_data, meta_data = load_checkpoint(checkpoint_dir, step, device, load_optimizer=False)
+    if device.type in {"cpu", "mps"}:
+        # Convert bfloat16 tensors to float for CPU inference
+        model_data = {
+            k: v.float() if v.dtype == torch.bfloat16 else v
+            for k, v in model_data.items()
+        }
+    # Hack: fix torch compile issue, which prepends all keys with _orig_mod.
+    model_data = {k.removeprefix("_orig_mod."): v for k, v in model_data.items()}
+    model_config_kwargs = meta_data["model_config"]
+    _patch_missing_config_keys(model_config_kwargs)
+    log0(f"Building model with config: {model_config_kwargs}")
+    model_config = GPTConfig(**model_config_kwargs)
+    _patch_missing_keys(model_data, model_config)
+    with torch.device("meta"):
+        model = GPT(model_config)
+    # Load the model state
+    model.to_empty(device=device)
+    model.init_weights() # note: this is dumb, but we need to init the rotary embeddings. TODO: fix model re-init
+    model.load_state_dict(model_data, strict=True, assign=True)
+    # Put the model in the right training phase / mode
+    if phase == "eval":
+        model.eval()
+    else:
+        model.train()
+    # Load the Tokenizer
+    tokenizer = get_tokenizer()
+    # Sanity check: compatibility between model and tokenizer
+    assert tokenizer.get_vocab_size() == model_config_kwargs["vocab_size"], f"Tokenizer vocab size {tokenizer.get_vocab_size()} does not match model config vocab size {model_config_kwargs['vocab_size']}"
+    return model, tokenizer, meta_data
+
+
+def find_largest_model(checkpoints_dir):
+    # attempt to guess the model tag: take the biggest model available
+    model_tags = [f for f in os.listdir(checkpoints_dir) if os.path.isdir(os.path.join(checkpoints_dir, f))]
+    if not model_tags:
+        raise FileNotFoundError(f"No checkpoints found in {checkpoints_dir}")
+    # 1) normally all model tags are of the form d<number>, try that first:
+    candidates = []
+    for model_tag in model_tags:
+        match = re.match(r"d(\d+)", model_tag)
+        if match:
+            model_depth = int(match.group(1))
+            candidates.append((model_depth, model_tag))
+    if candidates:
+        candidates.sort(key=lambda x: x[0], reverse=True)
+        return candidates[0][1]
+    # 2) if that failed, take the most recently updated model:
+    model_tags.sort(key=lambda x: os.path.getmtime(os.path.join(checkpoints_dir, x)), reverse=True)
+    return model_tags[0]
+
+
+def find_last_step(checkpoint_dir):
+    # Look into checkpoint_dir and find model_<step>.pt with the highest step
+    checkpoint_files = [f for f in os.listdir(checkpoint_dir) if re.search(r'model_(\d+)\.pt$', f)]
+    if not checkpoint_files:
+        raise FileNotFoundError(f"No checkpoints found in {checkpoint_dir}")
+    last_step = max(int(f.split("_")[-1].split(".")[0]) for f in checkpoint_files)
+    return last_step
+
+# -----------------------------------------------------------------------------
+# convenience functions that take into account nanochat's directory structure
+
+def load_model_from_dir(checkpoints_dir, device, phase, model_tag=None, step=None):
+    if model_tag is None:
+        # guess the model tag by defaulting to the largest model
+        model_tag = find_largest_model(checkpoints_dir)
+        log0(f"No model tag provided, guessing model tag: {model_tag}")
+    checkpoint_dir = os.path.join(checkpoints_dir, model_tag)
+    if step is None:
+        # guess the step by defaulting to the last step
+        step = find_last_step(checkpoint_dir)
+    assert step is not None, f"No checkpoints found in {checkpoint_dir}"
+    # build the model
+    log0(f"Loading model from {checkpoint_dir} with step {step}")
+    model, tokenizer, meta_data = build_model(checkpoint_dir, step, device, phase)
+    return model, tokenizer, meta_data
+
+def load_model(source, *args, **kwargs):
+    model_dir = {
+        "base": "base_checkpoints",
+        "sft": "chatsft_checkpoints",
+        "rl": "chatrl_checkpoints",
+    }[source]
+    base_dir = get_base_dir()
+    checkpoints_dir = os.path.join(base_dir, model_dir)
+    return load_model_from_dir(checkpoints_dir, *args, **kwargs)
+
+TOKENIZER_FILES = ("tokenizer.pkl", "token_bytes.pt")
+
+def download_hub_checkpoint(repo_id, revision=None, local_dir=None, token=None):
+    """
+    Download a nanochat-format checkpoint repo from the HuggingFace Hub and return the
+    local directory, ready to hand to build_model()/find_last_step().
+
+    The repo is expected to contain model_<step>.pt / meta_<step>.json (and optionally
+    optim_<step>_rank<r>.pt) at its root, i.e. exactly what save_checkpoint() writes.
+    If it also ships the training tokenizer, we install it into the local tokenizer dir
+    when there isn't one already -- a Hub snapshot without its matching tokenizer would
+    otherwise trip the vocab-size assert in build_model() with a confusing message.
+    """
+    try:
+        from huggingface_hub import snapshot_download
+    except ImportError as e:
+        raise ImportError(
+            "huggingface_hub is required to download checkpoints from the Hub. "
+            "Install it with: uv sync --extra gpu --extra distill"
+        ) from e
+    if local_dir is None:
+        local_dir = os.path.join("models", repo_id.split("/")[-1])
+    # Under torchrun only one rank should download; the others block and then reuse it.
+    os.makedirs(os.path.dirname(local_dir) or ".", exist_ok=True)
+    with FileLock(local_dir + ".lock"):
+        log0(f"Downloading checkpoint {repo_id} to {local_dir}")
+        snapshot_download(repo_id=repo_id, revision=revision, local_dir=local_dir, token=token)
+    # Sanity check that this actually looks like a nanochat checkpoint directory.
+    if not any(re.search(r"model_(\d+)\.pt$", f) for f in os.listdir(local_dir)):
+        raise FileNotFoundError(
+            f"{repo_id} downloaded to {local_dir} but contains no model_<step>.pt file. "
+            "This does not look like a nanochat checkpoint repo (as written by save_checkpoint)."
+        )
+    _install_hub_tokenizer(local_dir, repo_id)
+    return local_dir
+
+def _install_hub_tokenizer(local_dir, repo_id):
+    """Copy a tokenizer shipped alongside Hub weights into the local tokenizer dir, if absent."""
+    tokenizer_dir = os.path.join(get_base_dir(), "tokenizer")
+    if os.path.exists(os.path.join(tokenizer_dir, "tokenizer.pkl")):
+        return # a local tokenizer already exists, never clobber it
+    # the tokenizer may sit at the repo root or in a tokenizer/ subdirectory
+    for candidate in (local_dir, os.path.join(local_dir, "tokenizer")):
+        if os.path.exists(os.path.join(candidate, "tokenizer.pkl")):
+            os.makedirs(tokenizer_dir, exist_ok=True)
+            for filename in TOKENIZER_FILES:
+                src = os.path.join(candidate, filename)
+                if os.path.exists(src):
+                    shutil.copy2(src, os.path.join(tokenizer_dir, filename))
+            log0(f"Installed tokenizer from {candidate} into {tokenizer_dir}")
+            return
+    raise FileNotFoundError(
+        f"{repo_id} does not ship a tokenizer.pkl and none exists at {tokenizer_dir}. "
+        "The model cannot be loaded without the tokenizer it was trained with: either "
+        "train/copy one there (see scripts/tok_train.py), or pass --tokenizer-dir pointing "
+        "at a directory containing tokenizer.pkl and token_bytes.pt."
+    )
+
+def install_tokenizer_dir(tokenizer_dir):
+    """Copy an explicitly provided tokenizer into the local tokenizer dir (--tokenizer-dir)."""
+    dest = os.path.join(get_base_dir(), "tokenizer")
+    src_pickle = os.path.join(tokenizer_dir, "tokenizer.pkl")
+    if not os.path.exists(src_pickle):
+        raise FileNotFoundError(f"No tokenizer.pkl found in {tokenizer_dir}")
+    os.makedirs(dest, exist_ok=True)
+    for filename in TOKENIZER_FILES:
+        src = os.path.join(tokenizer_dir, filename)
+        if os.path.exists(src):
+            shutil.copy2(src, os.path.join(dest, filename))
+    log0(f"Installed tokenizer from {tokenizer_dir} into {dest}")
+
+def load_optimizer_state(source, device, rank, model_tag=None, step=None):
+    """Load just the optimizer shard for a given rank, without re-loading the model."""
+    model_dir = {
+        "base": "base_checkpoints",
+        "sft": "chatsft_checkpoints",
+        "rl": "chatrl_checkpoints",
+    }[source]
+    base_dir = get_base_dir()
+    checkpoints_dir = os.path.join(base_dir, model_dir)
+    if model_tag is None:
+        model_tag = find_largest_model(checkpoints_dir)
+    checkpoint_dir = os.path.join(checkpoints_dir, model_tag)
+    if step is None:
+        step = find_last_step(checkpoint_dir)
+    optimizer_path = os.path.join(checkpoint_dir, f"optim_{step:06d}_rank{rank:d}.pt")
+    if not os.path.exists(optimizer_path):
+        log0(f"Optimizer checkpoint not found: {optimizer_path}")
+        return None
+    log0(f"Loading optimizer state from {optimizer_path}")
+    optimizer_data = torch.load(optimizer_path, map_location=device)
+    return optimizer_data

nanochat/common.py

@@ -0,0 +1,331 @@
+"""
+Common utilities for nanochat.
+"""
+
+import os
+import re
+import logging
+import urllib.request
+import torch
+import torch.distributed as dist
+from filelock import FileLock
+
+# The dtype used for compute (matmuls, activations). Master weights stay fp32 for optimizer precision.
+# Linear layers cast their weights to this dtype in forward, replacing torch.amp.autocast.
+# Override with NANOCHAT_DTYPE env var: "bfloat16", "float16", "float32"
+_DTYPE_MAP = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}
+def _detect_compute_dtype():
+    env = os.environ.get("NANOCHAT_DTYPE")
+    if env is not None:
+        return _DTYPE_MAP[env], f"set via NANOCHAT_DTYPE={env}"
+    if torch.cuda.is_available():
+        # bf16 requires SM 80+ (Ampere: A100, A10, etc.)
+        # Older GPUs like V100 (SM 70) and T4 (SM 75) only have fp16 tensor cores
+        capability = torch.cuda.get_device_capability()
+        if capability >= (8, 0):
+            return torch.bfloat16, f"auto-detected: CUDA SM {capability[0]}{capability[1]} (bf16 supported)"
+        # fp16 training requires GradScaler (not yet implemented), so fall back to fp32.
+        # Users can still force fp16 via NANOCHAT_DTYPE=float16 if they know what they're doing.
+        return torch.float32, f"auto-detected: CUDA SM {capability[0]}{capability[1]} (pre-Ampere, bf16 not supported, using fp32)"
+    # Note: MPS on recent macOS also handles bf16 fine, opt in via NANOCHAT_DTYPE=bfloat16
+    return torch.float32, "auto-detected: no CUDA (CPU/MPS)"
+COMPUTE_DTYPE, COMPUTE_DTYPE_REASON = _detect_compute_dtype()
+
+class ColoredFormatter(logging.Formatter):
+    """Custom formatter that adds colors to log messages."""
+    # ANSI color codes
+    COLORS = {
+        'DEBUG': '\033[36m',    # Cyan
+        'INFO': '\033[32m',     # Green
+        'WARNING': '\033[33m',  # Yellow
+        'ERROR': '\033[31m',    # Red
+        'CRITICAL': '\033[35m', # Magenta
+    }
+    RESET = '\033[0m'
+    BOLD = '\033[1m'
+    def format(self, record):
+        # Add color to the level name
+        levelname = record.levelname
+        if levelname in self.COLORS:
+            record.levelname = f"{self.COLORS[levelname]}{self.BOLD}{levelname}{self.RESET}"
+        # Format the message
+        message = super().format(record)
+        # Add color to specific parts of the message
+        if levelname == 'INFO':
+            # Highlight numbers and percentages
+            message = re.sub(r'(\d+\.?\d*\s*(?:GB|MB|%|docs))', rf'{self.BOLD}\1{self.RESET}', message)
+            message = re.sub(r'(Shard \d+)', rf'{self.COLORS["INFO"]}{self.BOLD}\1{self.RESET}', message)
+        return message
+
+def setup_default_logging():
+    handler = logging.StreamHandler()
+    handler.setFormatter(ColoredFormatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
+    logging.basicConfig(
+        level=logging.INFO,
+        handlers=[handler]
+    )
+
+setup_default_logging()
+logger = logging.getLogger(__name__)
+
+def get_base_dir():
+    # co-locate nanochat intermediates with other cached data in ~/.cache (by default)
+    if os.environ.get("NANOCHAT_BASE_DIR"):
+        nanochat_dir = os.environ.get("NANOCHAT_BASE_DIR")
+    else:
+        home_dir = os.path.expanduser("~")
+        cache_dir = os.path.join(home_dir, ".cache")
+        nanochat_dir = os.path.join(cache_dir, "nanochat")
+    os.makedirs(nanochat_dir, exist_ok=True)
+    return nanochat_dir
+
+def download_file_with_lock(url, filename, postprocess_fn=None):
+    """
+    Downloads a file from a URL to a local path in the base directory.
+    Uses a lock file to prevent concurrent downloads among multiple ranks.
+    """
+    base_dir = get_base_dir()
+    file_path = os.path.join(base_dir, filename)
+    lock_path = file_path + ".lock"
+
+    if os.path.exists(file_path):
+        return file_path
+
+    with FileLock(lock_path):
+        # Only a single rank can acquire this lock
+        # All other ranks block until it is released
+
+        # Recheck after acquiring lock
+        if os.path.exists(file_path):
+            return file_path
+
+        # Download the content as bytes
+        print(f"Downloading {url}...")
+        with urllib.request.urlopen(url) as response:
+            content = response.read() # bytes
+
+        # Write to local file
+        with open(file_path, 'wb') as f:
+            f.write(content)
+        print(f"Downloaded to {file_path}")
+
+        # Run the postprocess function if provided
+        if postprocess_fn is not None:
+            postprocess_fn(file_path)
+
+    return file_path
+
+def print0(s="",**kwargs):
+    ddp_rank = int(os.environ.get('RANK', 0))
+    if ddp_rank == 0:
+        print(s, **kwargs)
+
+def print_banner():
+    # Cool DOS Rebel font ASCII banner made with https://manytools.org/hacker-tools/ascii-banner/
+    banner = """
+                                                       █████                █████
+                                                      ░░███                ░░███
+     ████████    ██████   ████████    ██████   ██████  ░███████    ██████  ███████
+    ░░███░░███  ░░░░░███ ░░███░░███  ███░░███ ███░░███ ░███░░███  ░░░░░███░░░███░
+     ░███ ░███   ███████  ░███ ░███ ░███ ░███░███ ░░░  ░███ ░███   ███████  ░███
+     ░███ ░███  ███░░███  ░███ ░███ ░███ ░███░███  ███ ░███ ░███  ███░░███  ░███ ███
+     ████ █████░░████████ ████ █████░░██████ ░░██████  ████ █████░░███████  ░░█████
+    ░░░░ ░░░░░  ░░░░░░░░ ░░░░ ░░░░░  ░░░░░░   ░░░░░░  ░░░░ ░░░░░  ░░░░░░░░   ░░░░░
+    """
+    print0(banner)
+
+def is_ddp_requested() -> bool:
+    """
+    True if launched by torchrun (env present), even before init.
+    Used to decide whether we *should* initialize a PG.
+    """
+    return all(k in os.environ for k in ("RANK", "LOCAL_RANK", "WORLD_SIZE"))
+
+def is_ddp_initialized() -> bool:
+    """
+    True if torch.distributed is available and the process group is initialized.
+    Used at cleanup to avoid destroying a non-existent PG.
+    """
+    return dist.is_available() and dist.is_initialized()
+
+def get_dist_info():
+    if is_ddp_requested():
+        # We rely on torchrun's env to decide if we SHOULD init.
+        # (Initialization itself happens in compute init.)
+        assert all(var in os.environ for var in ['RANK', 'LOCAL_RANK', 'WORLD_SIZE'])
+        ddp_rank = int(os.environ['RANK'])
+        ddp_local_rank = int(os.environ['LOCAL_RANK'])
+        ddp_world_size = int(os.environ['WORLD_SIZE'])
+        return True, ddp_rank, ddp_local_rank, ddp_world_size
+    else:
+        return False, 0, 0, 1
+
+def autodetect_device_type():
+    # prefer to use CUDA if available, otherwise use MPS, otherwise fallback on CPU
+    if torch.cuda.is_available():
+        device_type = "cuda"
+    elif torch.backends.mps.is_available():
+        device_type = "mps"
+    else:
+        device_type = "cpu"
+    print0(f"Autodetected device type: {device_type}")
+    return device_type
+
+def compute_init(device_type="cuda"): # cuda|cpu|mps
+    """Basic initialization that we keep doing over and over, so make common."""
+
+    assert device_type in ["cuda", "mps", "cpu"], "Invalid device type atm"
+    if device_type == "cuda":
+        assert torch.cuda.is_available(), "Your PyTorch installation is not configured for CUDA but device_type is 'cuda'"
+    if device_type == "mps":
+        assert torch.backends.mps.is_available(), "Your PyTorch installation is not configured for MPS but device_type is 'mps'"
+
+    # Reproducibility
+    # Note that we set the global seeds here, but most of the code uses explicit rng objects.
+    # The only place where global rng might be used is nn.Module initialization of the model weights.
+    torch.manual_seed(42)
+    if device_type == "cuda":
+        torch.cuda.manual_seed(42)
+    # skipping full reproducibility for now, possibly investigate slowdown later
+    # torch.use_deterministic_algorithms(True)
+
+    # Precision
+    if device_type == "cuda":
+        torch.set_float32_matmul_precision("high") # uses tf32 instead of fp32 for matmuls, see https://docs.pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html
+
+    # Distributed setup: Distributed Data Parallel (DDP), optional, and requires CUDA
+    is_ddp_requested, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info()
+    if is_ddp_requested and device_type == "cuda":
+        device = torch.device("cuda", ddp_local_rank)
+        torch.cuda.set_device(device)  # make "cuda" default to this device
+        dist.init_process_group(backend="nccl", device_id=device)
+        dist.barrier()
+    else:
+        device = torch.device(device_type) # mps|cpu
+
+    if ddp_rank == 0:
+        logger.info(f"Distributed world size: {ddp_world_size}")
+
+    return is_ddp_requested, ddp_rank, ddp_local_rank, ddp_world_size, device
+
+def compute_cleanup():
+    """Companion function to compute_init, to clean things up before script exit"""
+    if is_ddp_initialized():
+        dist.destroy_process_group()
+
+class DummyWandb:
+    """Useful if we wish to not use wandb but have all the same signatures"""
+    def __init__(self):
+        pass
+    def log(self, *args, **kwargs):
+        pass
+    def finish(self):
+        pass
+
+# hardcoded BF16 peak flops for various GPUs
+# inspired by torchtitan: https://github.com/pytorch/torchtitan/blob/main/torchtitan/tools/utils.py
+# and PR: https://github.com/karpathy/nanochat/pull/147
+def get_peak_flops(device_name: str) -> float:
+    name = device_name.lower()
+
+    # Table order matters: more specific patterns first.
+    _PEAK_FLOPS_TABLE = (
+        # NVIDIA Blackwell
+        (["gb200"], 2.5e15),
+        (["grace blackwell"], 2.5e15),
+        (["b200"], 2.25e15),
+        (["b100"], 1.8e15),
+        # NVIDIA Hopper
+        (["h200", "nvl"], 836e12),
+        (["h200", "pcie"], 836e12),
+        (["h200"], 989e12),
+        (["h100", "nvl"], 835e12),
+        (["h100", "pcie"], 756e12),
+        (["h100"], 989e12),
+        (["h800", "nvl"], 989e12),
+        (["h800"], 756e12),
+        # NVIDIA Ampere data center
+        (["a100"], 312e12),
+        (["a800"], 312e12),
+        (["a40"], 149.7e12),
+        (["a30"], 165e12),
+        # NVIDIA Ada data center
+        (["l40s"], 362e12),
+        (["l40-s"], 362e12),
+        (["l40 s"], 362e12),
+        (["l4"], 121e12),
+        # AMD CDNA accelerators
+        (["mi355"], 2.5e15),
+        (["mi325"], 1.3074e15),
+        (["mi300x"], 1.3074e15),
+        (["mi300a"], 980.6e12),
+        (["mi250x"], 383e12),
+        (["mi250"], 362.1e12),
+        # Consumer RTX
+        (["5090"], 209.5e12),
+        (["4090"], 165.2e12),
+        # 7,424 CUDA cores * 1.59 GHz max boost * 4 dense BF16 tensor FLOPs.
+        # Laptop boost varies with configured TGP, so actual peak may be lower.
+        (["3080 ti", "laptop"], 47.2e12),
+        (["3090"], 71e12),
+    )
+    for patterns, flops in _PEAK_FLOPS_TABLE:
+        if all(p in name for p in patterns):
+            return flops
+    if "data center gpu max 1550" in name:
+        # Ponte Vecchio (PVC) - dynamic based on compute units
+        max_comp_units = torch.xpu.get_device_properties("xpu").max_compute_units
+        return 512 * max_comp_units * 1300 * 10**6
+
+    # Unknown GPU - return inf so MFU shows as 0% rather than a wrong guess
+    logger.warning(f"Peak flops undefined for: {device_name}, MFU will show as 0%")
+    return float('inf')
+
+def get_peak_bandwidth(device_name: str) -> float:
+    """Peak HBM/GDDR memory bandwidth in bytes/sec. The decode phase of inference
+    is memory-bandwidth-bound, so this is the roofline for tokens/sec (see MBU)."""
+    name = device_name.lower()
+
+    # Table order matters: more specific patterns first.
+    _PEAK_BANDWIDTH_TABLE = (
+        # NVIDIA Blackwell (HBM3e)
+        (["gb200"], 8.0e12),
+        (["grace blackwell"], 8.0e12),
+        (["b200"], 8.0e12),
+        (["b100"], 8.0e12),
+        # NVIDIA Hopper
+        (["h200"], 4.8e12),
+        (["h100", "nvl"], 3.9e12),
+        (["h100", "pcie"], 2.0e12),
+        (["h100"], 3.35e12), # SXM
+        (["h800", "pcie"], 2.0e12),
+        (["h800"], 3.35e12), # SXM
+        # NVIDIA Ampere data center (A100 80GB; the 40GB variant is 1.6e12)
+        (["a100"], 2.0e12),
+        (["a800"], 2.0e12),
+        (["a40"], 696e9),
+        (["a30"], 933e9),
+        # NVIDIA Ada data center
+        (["l40s"], 864e9),
+        (["l40-s"], 864e9),
+        (["l40 s"], 864e9),
+        (["l4"], 300e9),
+        # AMD CDNA accelerators
+        (["mi355"], 8.0e12),
+        (["mi325"], 6.0e12),
+        (["mi300x"], 5.3e12),
+        (["mi300a"], 5.3e12),
+        (["mi250x"], 3.28e12),
+        (["mi250"], 3.28e12),
+        # Consumer RTX
+        (["5090"], 1.79e12),
+        (["4090"], 1.01e12),
+        (["3090"], 936e9),
+    )
+    for patterns, bandwidth in _PEAK_BANDWIDTH_TABLE:
+        if all(p in name for p in patterns):
+            return bandwidth
+
+    # Unknown GPU - return inf so MBU shows as 0% rather than a wrong guess
+    logger.warning(f"Peak bandwidth undefined for: {device_name}, MBU will show as 0%")
+    return float('inf')

nanochat/core_eval.py

@@ -0,0 +1,262 @@
+"""
+Functions for evaluating the CORE metric, as described in the DCLM paper.
+https://arxiv.org/abs/2406.11794
+
+TODOs:
+- All tasks ~match except for squad. We get 31% reference is 37%. Figure out why.
+"""
+import random
+
+from jinja2 import Template
+import torch
+import torch.distributed as dist
+
+# -----------------------------------------------------------------------------
+# Prompt rendering utilities
+
+def render_prompts_mc(item, continuation_delimiter, fewshot_examples=None):
+    """Render complete prompts for a multiple choice question"""
+    template_str = """
+{%- for example in fewshot_examples -%}
+{{ example.query }}{{ continuation_delimiter }}{{ example.choices[example.gold] }}
+
+{% endfor -%}
+{{ item.query }}{{ continuation_delimiter }}{{ choice }}""".strip()
+    template = Template(template_str)
+    fewshot_examples = fewshot_examples or []
+    context = {
+        'fewshot_examples': fewshot_examples,
+        'continuation_delimiter': continuation_delimiter,
+        'item': item
+    }
+    prompts = [template.render(choice=choice, **context) for choice in item['choices']]
+    return prompts
+
+
+def render_prompts_schema(item, continuation_delimiter, fewshot_examples=None):
+    """Render complete prompts for a schema question"""
+    template_str = """
+{%- for example in fewshot_examples -%}
+{{ example.context_options[example.gold] }}{{ continuation_delimiter }}{{ example.continuation }}
+
+{% endfor -%}
+{{ context }}{{ continuation_delimiter }}{{ item.continuation }}""".strip()
+    template = Template(template_str)
+    fewshot_examples = fewshot_examples or []
+    context = {
+        'fewshot_examples': fewshot_examples,
+        'continuation_delimiter': continuation_delimiter,
+        'item': item
+    }
+    prompts = [template.render(context=context_option, **context)
+               for context_option in item['context_options']]
+    return prompts
+
+
+def render_prompts_lm(item, continuation_delimiter, fewshot_examples=None):
+    """
+    Render complete prompt for a language modeling task.
+    Notice that we manually trim the context in the template,
+    which in some datasets seems to have trailing whitespace (which we don't want).
+    """
+    template_str = """
+{%- for example in fewshot_examples -%}
+{{ example.context | trim }}{{ continuation_delimiter }}{{ example.continuation }}
+
+{% endfor -%}
+{{ item.context | trim }}{{ continuation_delimiter }}{% if include_continuation %}{{ item.continuation }}{% endif %}""".strip()
+    template = Template(template_str)
+    fewshot_examples = fewshot_examples or []
+    context = {
+        'fewshot_examples': fewshot_examples,
+        'continuation_delimiter': continuation_delimiter,
+        'item': item
+    }
+    # Return two prompts: without and with the continuation
+    prompt_without = template.render(include_continuation=False, **context)
+    prompt_with = template.render(include_continuation=True, **context)
+    # Due to the way the data seems to be stored, I think I need to strip in the case of LM here.
+    # Otherwise we may get trailing whitespaces in prompt_without (which get absorbed into the next
+    # token in prompt_with), meaning we don't get a nice and clean prefix in the token space
+    # to detect the final continuation. Tokenizers...
+    prompt_without = prompt_without.strip()
+    return [prompt_without, prompt_with]
+
+
+def find_common_length(token_sequences, direction='left'):
+    """
+    Find the length of the common prefix or suffix across token sequences
+    - direction: 'left' for prefix, 'right' for suffix
+    """
+    min_len = min(len(seq) for seq in token_sequences)
+    indices = {
+        'left': range(min_len),
+        'right': range(-1, -min_len-1, -1)
+    }[direction]
+    # Find the first position where the token sequences differ
+    for i, idx in enumerate(indices):
+        token = token_sequences[0][idx]
+        if not all(seq[idx] == token for seq in token_sequences):
+            return i
+    return min_len
+
+
+def stack_sequences(tokens, pad_token_id):
+    """Stack up a list of token sequences, pad to longest on the right"""
+    bsz, seq_len = len(tokens), max(len(x) for x in tokens)
+    input_ids = torch.full((bsz, seq_len), pad_token_id, dtype=torch.long)
+    for i, x in enumerate(tokens):
+        input_ids[i, :len(x)] = torch.tensor(x, dtype=torch.long)
+    return input_ids
+
+
+def batch_sequences_mc(tokenizer, prompts):
+    # In multiple choice, contexts are the same but the continuation is different (common prefix)
+    tokens = tokenizer(prompts, prepend=tokenizer.get_bos_token_id())
+    # figure out the start and end of each continuation
+    answer_start_idx = find_common_length(tokens, direction='left')
+    start_indices = [answer_start_idx] * len(prompts)
+    end_indices = [len(x) for x in tokens]
+    return tokens, start_indices, end_indices
+
+
+def batch_sequences_schema(tokenizer, prompts):
+    # In schema tasks, contexts vary but continuation is the same (common suffix)
+    tokens = tokenizer(prompts, prepend=tokenizer.get_bos_token_id())
+    # figure out the start and end of each context
+    suffix_length = find_common_length(tokens, direction='right')
+    end_indices = [len(x) for x in tokens]
+    start_indices = [ei - suffix_length for ei in end_indices]
+    return tokens, start_indices, end_indices
+
+
+def batch_sequences_lm(tokenizer, prompts):
+    # In LM tasks, we have two prompts: without and with continuation
+    tokens = tokenizer(prompts, prepend=tokenizer.get_bos_token_id())
+    tokens_without, tokens_with = tokens
+    start_idx, end_idx = len(tokens_without), len(tokens_with)
+    assert start_idx < end_idx, "prompt without is supposed to be a prefix of prompt with"
+    assert tokens_without == tokens_with[:start_idx], "prompt without is supposed to be a prefix of prompt with"
+    # we only need the with continuation prompt in the LM task, i.e. batch size of 1
+    return [tokens_with], [start_idx], [end_idx]
+
+
+@torch.no_grad()
+def forward_model(model, input_ids):
+    """
+    Take BxT tensor of token ids, return BxT tensor of losses and argmax predictions.
+    The last column of losses is set to nan because we don't have autoregressive targets there.
+    """
+    batch_size, seq_len = input_ids.size()
+    outputs = model(input_ids)
+    # Roll the tensor to the left by one position to get the (autoregressive) target ids
+    target_ids = torch.roll(input_ids, shifts=-1, dims=1)
+    # Calculate cross entropy at all positions
+    losses = torch.nn.functional.cross_entropy(
+        outputs.view(batch_size * seq_len, -1),
+        target_ids.view(batch_size * seq_len),
+        reduction='none'
+    ).view(batch_size, seq_len)
+    # Set the last column to be nan because there is no autoregressive loss there
+    losses[:, -1] = float('nan')
+    # Get the argmax predictions at each position
+    predictions = outputs.argmax(dim=-1)
+    return losses, predictions
+
+
+@torch.no_grad()
+def evaluate_example(idx, model, tokenizer, data, device, task_meta):
+    """Evaluate a single example, return True if correct, False otherwise"""
+    item = data[idx]
+    task_type = task_meta['task_type']
+    num_fewshot = task_meta['num_fewshot']
+    continuation_delimiter = task_meta['continuation_delimiter']
+
+    # Sample few-shot examples (excluding current item)
+    fewshot_examples = []
+    if num_fewshot > 0:
+        rng = random.Random(1234 + idx)
+        available_indices = [i for i in range(len(data)) if i != idx]
+        fewshot_indices = rng.sample(available_indices, num_fewshot)
+        fewshot_examples = [data[i] for i in fewshot_indices]
+
+    # Render prompts and batch sequences based on task type
+    if task_type == 'multiple_choice':
+        prompts = render_prompts_mc(item, continuation_delimiter, fewshot_examples)
+        tokens, start_idxs, end_idxs = batch_sequences_mc(tokenizer, prompts)
+    elif task_type == 'schema':
+        prompts = render_prompts_schema(item, continuation_delimiter, fewshot_examples)
+        tokens, start_idxs, end_idxs = batch_sequences_schema(tokenizer, prompts)
+    elif task_type == 'language_modeling':
+        prompts = render_prompts_lm(item, continuation_delimiter, fewshot_examples)
+        tokens, start_idxs, end_idxs = batch_sequences_lm(tokenizer, prompts)
+    else:
+        raise ValueError(f"Unsupported task type: {task_type}")
+
+    # Some models can't forward sequences beyond a certain length (e.g. GPT-2)
+    # In these cases, we have to truncate sequences to max length and adjust the indices
+    if hasattr(model, 'max_seq_len') and model.max_seq_len is not None:
+        max_tokens = model.max_seq_len
+        new_tokens, new_start_idxs, new_end_idxs = [], [], []
+        for t, s, e in zip(tokens, start_idxs, end_idxs):
+            if len(t) > max_tokens:
+                num_to_crop = len(t) - max_tokens
+                new_tokens.append(t[-max_tokens:]) # take the last max_tokens tokens
+                new_start_idxs.append(s - num_to_crop) # shift the indices down
+                new_end_idxs.append(e - num_to_crop)
+                assert s - num_to_crop >= 0, "this should never happen right?"
+                assert e - num_to_crop >= 0, "this should never happen right?"
+            else:
+                new_tokens.append(t) # keep unchanged
+                new_start_idxs.append(s)
+                new_end_idxs.append(e)
+        tokens, start_idxs, end_idxs = new_tokens, new_start_idxs, new_end_idxs
+
+    # Stack up all the sequences into a batch
+    pad_token_id = tokenizer.get_bos_token_id() # use BOS as pad token is ok
+    input_ids = stack_sequences(tokens, pad_token_id)
+    input_ids = input_ids.to(device)
+
+    # Forward the model, get the autoregressive loss and argmax prediction at each token
+    losses, predictions = forward_model(model, input_ids)
+
+    # See if the losses/predictions come out correctly
+    if task_type == 'language_modeling':
+        # language modeling task is currently always batch size 1
+        si = start_idxs[0]
+        ei = end_idxs[0]
+        # predictions[i] predict input_ids[i+1] autoregressively
+        predicted_tokens = predictions[0, si-1:ei-1]
+        actual_tokens = input_ids[0, si:ei]
+        is_correct = torch.all(predicted_tokens == actual_tokens).item()
+    elif task_type in ['multiple_choice', 'schema']:
+        # For MC/schema: find the option with lowest average loss
+        mean_losses = [losses[i, si-1:ei-1].mean().item()
+                        for i, (si, ei) in enumerate(zip(start_idxs, end_idxs))]
+        pred_idx = mean_losses.index(min(mean_losses))
+        is_correct = pred_idx == item['gold']
+    else:
+        raise ValueError(f"Unsupported task type: {task_type}")
+
+    return is_correct
+
+
+def evaluate_task(model, tokenizer, data, device, task_meta):
+    """
+    This function is responsible for evaluating one task across many examples.
+    It also handles dispatch to all processes if the script is run with torchrun.
+    """
+    rank = dist.get_rank() if dist.is_initialized() else 0
+    world_size = dist.get_world_size() if dist.is_initialized() else 1
+    correct = torch.zeros(len(data), dtype=torch.float32, device=device)
+    # stride the examples to each rank
+    for idx in range(rank, len(data), world_size):
+        is_correct = evaluate_example(idx, model, tokenizer, data, device, task_meta)
+        correct[idx] = float(is_correct)
+    # sync results across all the processes if running distributed
+    if world_size > 1:
+        dist.barrier()
+        dist.all_reduce(correct, op=dist.ReduceOp.SUM)
+    # compute the mean
+    mean_correct = correct.mean().item()
+    return mean_correct

nanochat/dataloader.py

@@ -0,0 +1,166 @@
+"""
+Distributed dataloaders for pretraining.
+
+BOS-aligned bestfit:
+   - Every row starts with BOS token
+   - Documents packed using best-fit algorithm to minimize cropping
+   - When no document fits remaining space, crops a document to fill exactly
+   - 100% utilization (no padding), ~35% tokens cropped at T=2048
+
+Compared to the original tokenizing_distributed_data_loader:
+BOS-aligned loses ~35% of tokens to cropping, but ensures that
+there are fewer "confusing" tokens in the train/val batches as every token can
+now attend back to the BOS token and sees the full context of the document.
+
+Fallback to the original if you have very limited data AND long documents:
+https://github.com/karpathy/nanochat/blob/3c3a3d7/nanochat/dataloader.py#L78-L117
+"""
+
+import torch
+import pyarrow.parquet as pq
+
+from nanochat.common import get_dist_info
+from nanochat.dataset import list_parquet_files
+
+def _document_batches(split, resume_state_dict, tokenizer_batch_size):
+    """
+    Infinite iterator over document batches (list of text strings) from parquet files.
+
+    Handles DDP sharding and approximate resume. Each yield is (text_batch, (pq_idx, rg_idx, epoch))
+    where text_batch is a list of document strings, indices track position for resumption,
+    and epoch counts how many times we've cycled through the dataset (starts at 1).
+    """
+    ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info()
+
+    warn_on_legacy = ddp_rank == 0 and split == "train" # rank 0 on train split will warn on legacy
+    parquet_paths = list_parquet_files(warn_on_legacy=warn_on_legacy)
+    assert len(parquet_paths) != 0, "No dataset parquet files found, did you run dataset.py?"
+    parquet_paths = parquet_paths[:-1] if split == "train" else parquet_paths[-1:]
+
+    resume_pq_idx = resume_state_dict["pq_idx"] if resume_state_dict is not None else 0
+    resume_rg_idx = resume_state_dict["rg_idx"] if resume_state_dict is not None else None
+    resume_epoch = resume_state_dict.get("epoch", 1) if resume_state_dict is not None else 1
+    first_pass = True
+    pq_idx = resume_pq_idx
+    epoch = resume_epoch
+
+    while True:  # iterate infinitely (multi-epoch)
+        pq_idx = resume_pq_idx if first_pass else 0
+        while pq_idx < len(parquet_paths):
+            filepath = parquet_paths[pq_idx]
+            pf = pq.ParquetFile(filepath)
+            # Start from resume point if resuming on same file, otherwise from DDP rank
+            if first_pass and (resume_rg_idx is not None) and (pq_idx == resume_pq_idx):
+                base_idx = resume_rg_idx // ddp_world_size
+                base_idx += 1  # advance by 1 so we don't repeat data after resuming
+                rg_idx = base_idx * ddp_world_size + ddp_rank
+                if rg_idx >= pf.num_row_groups:
+                    pq_idx += 1
+                    continue
+                resume_rg_idx = None  # only do this once
+            else:
+                rg_idx = ddp_rank
+            while rg_idx < pf.num_row_groups:
+                rg = pf.read_row_group(rg_idx)
+                batch = rg.column('text').to_pylist()
+                for i in range(0, len(batch), tokenizer_batch_size):
+                    yield batch[i:i+tokenizer_batch_size], (pq_idx, rg_idx, epoch)
+                rg_idx += ddp_world_size
+            pq_idx += 1
+        first_pass = False
+        epoch += 1
+
+
+def tokenizing_distributed_data_loader_with_state_bos_bestfit(
+    tokenizer, B, T, split,
+    tokenizer_threads=4, tokenizer_batch_size=128,
+    device="cuda", resume_state_dict=None,
+    buffer_size=1000
+):
+    """
+    BOS-aligned dataloader with Best-Fit Cropping.
+
+    Reduces token waste compared to simple greedy cropping by searching a buffer
+    for documents that fit well, while maintaining 100% utilization (no padding).
+
+    Algorithm for each row:
+    1. From buffered docs, pick the LARGEST doc that fits entirely
+    2. Repeat until no doc fits
+    3. When nothing fits, crop a doc to fill remaining space exactly
+
+    Key properties:
+    - Every row starts with BOS
+    - 100% utilization (no padding, every token is trained on)
+    - Approximately 35% of all tokens are discarded due to cropping
+    """
+    assert split in ["train", "val"], "split must be 'train' or 'val'"
+
+    row_capacity = T + 1
+    batches = _document_batches(split, resume_state_dict, tokenizer_batch_size)
+    bos_token = tokenizer.get_bos_token_id()
+    doc_buffer = []
+    pq_idx, rg_idx, epoch = 0, 0, 1
+
+    def refill_buffer():
+        nonlocal pq_idx, rg_idx, epoch
+        doc_batch, (pq_idx, rg_idx, epoch) = next(batches)
+        token_lists = tokenizer.encode(doc_batch, prepend=bos_token, num_threads=tokenizer_threads)
+        for tokens in token_lists:
+            doc_buffer.append(tokens)
+
+    # Pre-allocate buffers once: layout is [inputs (B*T) | targets (B*T)]
+    # This gives us contiguous views and a single HtoD transfer
+    use_cuda = device == "cuda"
+    row_buffer = torch.empty((B, row_capacity), dtype=torch.long) # for building rows without creating Python lists
+    cpu_buffer = torch.empty(2 * B * T, dtype=torch.long, pin_memory=use_cuda) # staging area (CPU)
+    gpu_buffer = torch.empty(2 * B * T, dtype=torch.long, device=device) # on-device buffer
+    cpu_inputs = cpu_buffer[:B * T].view(B, T) # a few views into these buffers just for convenience
+    cpu_targets = cpu_buffer[B * T:].view(B, T)
+    inputs = gpu_buffer[:B * T].view(B, T)
+    targets = gpu_buffer[B * T:].view(B, T)
+
+    while True:
+        for row_idx in range(B):
+            pos = 0
+            while pos < row_capacity:
+                # Ensure buffer has documents
+                while len(doc_buffer) < buffer_size:
+                    refill_buffer()
+
+                remaining = row_capacity - pos
+
+                # Find largest doc that fits entirely
+                best_idx = -1
+                best_len = 0
+                for i, doc in enumerate(doc_buffer):
+                    doc_len = len(doc)
+                    if doc_len <= remaining and doc_len > best_len:
+                        best_idx = i
+                        best_len = doc_len
+
+                if best_idx >= 0:
+                    doc = doc_buffer.pop(best_idx)
+                    doc_len = len(doc)
+                    row_buffer[row_idx, pos:pos + doc_len] = torch.tensor(doc, dtype=torch.long)
+                    pos += doc_len
+                else:
+                    # No doc fits - crop shortest in buffer to fill remaining and minimize waste
+                    shortest_idx = min(range(len(doc_buffer)), key=lambda i: len(doc_buffer[i]))
+                    doc = doc_buffer.pop(shortest_idx)
+                    row_buffer[row_idx, pos:pos + remaining] = torch.tensor(doc[:remaining], dtype=torch.long)
+                    pos += remaining
+
+        # Copy to pinned CPU buffer, then single HtoD transfer
+        cpu_inputs.copy_(row_buffer[:, :-1])
+        cpu_targets.copy_(row_buffer[:, 1:])
+
+        state_dict = {"pq_idx": pq_idx, "rg_idx": rg_idx, "epoch": epoch}
+
+        # Single HtoD copy into persistent GPU buffer and yield
+        gpu_buffer.copy_(cpu_buffer, non_blocking=use_cuda)
+        yield inputs, targets, state_dict
+
+def tokenizing_distributed_data_loader_bos_bestfit(*args, **kwargs):
+    """Helper that omits state_dict from yields."""
+    for inputs, targets, state_dict in tokenizing_distributed_data_loader_with_state_bos_bestfit(*args, **kwargs):
+        yield inputs, targets

nanochat/dataset.py

@@ -0,0 +1,160 @@
+"""
+The base/pretraining dataset is a set of parquet files.
+This file contains utilities for:
+- iterating over the parquet files and yielding documents from it
+- download the files on demand if they are not on disk
+
+For details of how the dataset was prepared, see `repackage_data_reference.py`.
+"""
+
+import os
+import argparse
+import time
+import requests
+import pyarrow.parquet as pq
+from multiprocessing import Pool
+
+from nanochat.common import get_base_dir
+
+# -----------------------------------------------------------------------------
+# The specifics of the current pretraining dataset
+
+# The URL on the internet where the data is hosted and downloaded from on demand
+BASE_URL = "https://huggingface.co/datasets/karpathy/climbmix-400b-shuffle/resolve/main"
+MAX_SHARD = 6542 # the last datashard is shard_06542.parquet
+index_to_filename = lambda index: f"shard_{index:05d}.parquet" # format of the filenames
+base_dir = get_base_dir()
+DATA_DIR = os.path.join(base_dir, "base_data_climbmix")
+
+# -----------------------------------------------------------------------------
+# These functions are useful utilities to other modules, can/should be imported
+
+def list_parquet_files(data_dir=None, warn_on_legacy=False):
+    """ Looks into a data dir and returns full paths to all parquet files. """
+    data_dir = DATA_DIR if data_dir is None else data_dir
+
+    # Legacy-supporting code due to the upgrade from FinewebEdu-100B to ClimbMix-400B
+    # This code will eventually be deleted.
+    if not os.path.exists(data_dir):
+        if warn_on_legacy:
+            print()
+            print("=" * 80)
+            print("  WARNING: DATASET UPGRADE REQUIRED")
+            print("=" * 80)
+            print()
+            print(f"  Could not find: {data_dir}")
+            print()
+            print("  nanochat recently switched from FinewebEdu-100B to ClimbMix-400B.")
+            print("  Everyone who does `git pull` as of March 4, 2026 is expected to see this message.")
+            print("  To upgrade to the new ClimbMix-400B dataset, run these two commands:")
+            print()
+            print("    python -m nanochat.dataset -n 170     # download ~170 shards, enough for GPT-2, adjust as desired")
+            print("    python -m scripts.tok_train           # re-train tokenizer on new ClimbMix data")
+            print()
+            print("  For now, falling back to your old FinewebEdu-100B dataset...")
+            print("=" * 80)
+            print()
+        # attempt a fallback to the legacy data directory
+        data_dir = os.path.join(base_dir, "base_data")
+
+    parquet_files = sorted([
+        f for f in os.listdir(data_dir)
+        if f.endswith('.parquet') and not f.endswith('.tmp')
+    ])
+    parquet_paths = [os.path.join(data_dir, f) for f in parquet_files]
+    return parquet_paths
+
+def parquets_iter_batched(split, start=0, step=1):
+    """
+    Iterate through the dataset, in batches of underlying row_groups for efficiency.
+    - split can be "train" or "val". the last parquet file will be val.
+    - start/step are useful for skipping rows in DDP. e.g. start=rank, step=world_size
+    """
+    assert split in ["train", "val"], "split must be 'train' or 'val'"
+    parquet_paths = list_parquet_files()
+    parquet_paths = parquet_paths[:-1] if split == "train" else parquet_paths[-1:]
+    for filepath in parquet_paths:
+        pf = pq.ParquetFile(filepath)
+        for rg_idx in range(start, pf.num_row_groups, step):
+            rg = pf.read_row_group(rg_idx)
+            texts = rg.column('text').to_pylist()
+            yield texts
+
+# -----------------------------------------------------------------------------
+def download_single_file(index):
+    """ Downloads a single file index, with some backoff """
+
+    # Construct the local filepath for this file and skip if it already exists
+    filename = index_to_filename(index)
+    filepath = os.path.join(DATA_DIR, filename)
+    if os.path.exists(filepath):
+        print(f"Skipping {filepath} (already exists)")
+        return True
+
+    # Construct the remote URL for this file
+    url = f"{BASE_URL}/{filename}"
+    print(f"Downloading {filename}...")
+
+    # Download with retries
+    max_attempts = 5
+    for attempt in range(1, max_attempts + 1):
+        try:
+            response = requests.get(url, stream=True, timeout=30)
+            response.raise_for_status()
+            # Write to temporary file first
+            temp_path = filepath + f".tmp"
+            with open(temp_path, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=1024 * 1024):  # 1MB chunks
+                    if chunk:
+                        f.write(chunk)
+            # Move temp file to final location
+            os.rename(temp_path, filepath)
+            print(f"Successfully downloaded {filename}")
+            return True
+
+        except (requests.RequestException, IOError) as e:
+            print(f"Attempt {attempt}/{max_attempts} failed for {filename}: {e}")
+            # Clean up any partial files
+            for path in [filepath + f".tmp", filepath]:
+                if os.path.exists(path):
+                    try:
+                        os.remove(path)
+                    except:
+                        pass
+            # Try a few times with exponential backoff: 2^attempt seconds
+            if attempt < max_attempts:
+                wait_time = 2 ** attempt
+                print(f"Waiting {wait_time} seconds before retry...")
+                time.sleep(wait_time)
+            else:
+                print(f"Failed to download {filename} after {max_attempts} attempts")
+                return False
+
+    return False
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Download pretraining dataset shards")
+    parser.add_argument("-n", "--num-files", type=int, default=-1, help="Number of train shards to download (default: -1), -1 = disable")
+    parser.add_argument("-w", "--num-workers", type=int, default=4, help="Number of parallel download workers (default: 4)")
+    args = parser.parse_args()
+
+    # Prepare the output directory
+    os.makedirs(DATA_DIR, exist_ok=True)
+
+    # The way this works is that the user specifies the number of train shards to download via the -n flag.
+    # In addition to that, the validation shard is *always* downloaded and is pinned to be the last shard.
+    num_train_shards = MAX_SHARD if args.num_files == -1 else min(args.num_files, MAX_SHARD)
+    ids_to_download = list(range(num_train_shards))
+    ids_to_download.append(MAX_SHARD) # always download the validation shard
+
+    # Download the shards
+    print(f"Downloading {len(ids_to_download)} shards using {args.num_workers} workers...")
+    print(f"Target directory: {DATA_DIR}")
+    print()
+    with Pool(processes=args.num_workers) as pool:
+        results = pool.map(download_single_file, ids_to_download)
+
+    # Report results
+    successful = sum(1 for success in results if success)
+    print(f"Done! Downloaded: {successful}/{len(ids_to_download)} shards to {DATA_DIR}")

nanochat/distill/__init__.py

@@ -0,0 +1,6 @@
+"""
+Cross-Architecture Weight Translation (CAWT): distills a large teacher transformer
+(Model A, e.g. Gemma-12B) into the LFM2-Quantum student architecture (Model B) by
+training a converter network (Model C) that maps A's parameters into an
+initialization for B, then fine-tuning B directly. See scripts/distill_train.py.
+"""

nanochat/distill/converter.py

@@ -0,0 +1,228 @@
+"""
+Converter Model C (CAWT Sec 2.2-2.3): translates Model A's parameters into an
+initialization for Model B's weight tensors via cross-attention, so B "inherits"
+A's weights instead of being trained on A's outputs alone.
+
+Design, matching the algorithm description:
+  - Source weight tokens (Sec 2.1) are embedded with a trainable value projection
+    plus role/layer/chunk structural embeddings -- see weight_tokens.WeightCorpus for
+    the (frozen) chunking step this consumes.
+  - Each target tensor is produced by a bank of learned query tokens (one per output
+    chunk) that cross-attend over a *relevance window* of source tokens (Sec 2.2): the
+    top-k source layers under a learned soft layer-correspondence matrix M (Sec 2.3),
+    weighted into the attention scores as an additive log-bias so gradient descent can
+    still refine which of the top-k layers matters most. Role compatibility is not
+    hand-coded (Sec 2.4): query and source tokens both carry role embeddings, so the
+    cross-attention itself learns which teacher roles are relevant to which student
+    roles (e.g. a quantum FFN's theta angles have no analytic teacher counterpart, so
+    this is left entirely to learned attention + Phase 2 fine-tuning).
+  - This bounds cost to O(|theta_B| * k) source tokens attended per target tensor,
+    not O(|theta_A| * |theta_B|), which is what makes a 12B-parameter teacher feasible.
+"""
+
+import math
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from nanochat.gpt import norm  # parameter-free RMSNorm, reused for style consistency
+from .roles import NUM_ROLES, ROLE_TO_ID
+from .weight_tokens import WeightCorpus
+
+
+def sinusoidal_features(frac: torch.Tensor, dim: int) -> torch.Tensor:
+    """Continuous positional features for a scalar in [0, 1] (layer depth fraction or
+    intra-tensor chunk fraction). Used instead of a per-absolute-position embedding
+    table so the converter's size does not depend on the teacher's or student's depth
+    (needed for the Phase 3 amortization check: reusing C on a different teacher)."""
+    half = dim // 2
+    freqs = torch.exp(torch.linspace(0, math.log(1000.0), half, device=frac.device, dtype=torch.float32))
+    angles = frac.float().unsqueeze(-1) * freqs
+    feats = torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1)
+    if feats.shape[-1] < dim:
+        feats = F.pad(feats, (0, dim - feats.shape[-1]))
+    return feats
+
+
+def gaussian_bump_layer_map(student_layers: int, teacher_layers: int, tau: float) -> torch.Tensor:
+    """Pre-softmax init for the layer-correspondence matrix M (Sec 2.3): student layer j
+    starts out attending most to the teacher layer at the same relative depth."""
+    j = torch.arange(student_layers).float().unsqueeze(1) / max(student_layers - 1, 1)
+    i = torch.arange(teacher_layers).float().unsqueeze(0) / max(teacher_layers - 1, 1)
+    return -((i - j) ** 2) / max(tau, 1e-6)
+
+
+@dataclass
+class ConverterConfig:
+    chunk_dim: int = 1024
+    d_model: int = 512
+    n_heads: int = 8
+    n_decoder_blocks: int = 4
+    svd_rank: int = 64
+    top_k_source_layers: int = 4
+    layer_map_tau: float = 0.5
+    translate_quantum_theta: bool = True
+    reg_target_std: float = 0.02
+
+    def __post_init__(self):
+        assert self.d_model % self.n_heads == 0, "d_model must be divisible by n_heads"
+
+
+class _DecoderBlock(nn.Module):
+    """Self-attention among a tensor's own queries (keeps chunks of one output
+    coherent with each other) + cross-attention into the source weight-token window
+    + a small MLP. Pre-norm, no biases, no learnable norm scale -- mirrors the style of
+    nanochat/gpt.py's transformer blocks."""
+
+    def __init__(self, d_model: int, n_heads: int):
+        super().__init__()
+        self.n_heads = n_heads
+        self.self_q = nn.Linear(d_model, d_model, bias=False)
+        self.self_k = nn.Linear(d_model, d_model, bias=False)
+        self.self_v = nn.Linear(d_model, d_model, bias=False)
+        self.self_o = nn.Linear(d_model, d_model, bias=False)
+        self.cross_q = nn.Linear(d_model, d_model, bias=False)
+        self.cross_k = nn.Linear(d_model, d_model, bias=False)
+        self.cross_v = nn.Linear(d_model, d_model, bias=False)
+        self.cross_o = nn.Linear(d_model, d_model, bias=False)
+        self.mlp_fc = nn.Linear(d_model, 4 * d_model, bias=False)
+        self.mlp_proj = nn.Linear(4 * d_model, d_model, bias=False)
+
+    def _mha(self, q_in, kv_in, q_lin, k_lin, v_lin, o_lin, attn_bias=None):
+        Tq, Tk = q_in.shape[0], kv_in.shape[0]
+        H = self.n_heads
+        d = q_in.shape[-1]
+        hd = d // H
+        q = q_lin(q_in).view(Tq, H, hd).permute(1, 0, 2).unsqueeze(0)  # (1, H, Tq, hd)
+        k = k_lin(kv_in).view(Tk, H, hd).permute(1, 0, 2).unsqueeze(0)
+        v = v_lin(kv_in).view(Tk, H, hd).permute(1, 0, 2).unsqueeze(0)
+        mask = None
+        if attn_bias is not None:
+            mask = attn_bias.view(1, 1, Tq, Tk).expand(1, H, Tq, Tk)
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        out = out.squeeze(0).permute(1, 0, 2).reshape(Tq, d)
+        return o_lin(out)
+
+    def forward(self, queries, source=None, source_bias=None):
+        x = queries
+        x = x + self._mha(norm(x), norm(x), self.self_q, self.self_k, self.self_v, self.self_o)
+        if source is not None:
+            x = x + self._mha(norm(x), norm(source), self.cross_q, self.cross_k, self.cross_v, self.cross_o, attn_bias=source_bias)
+        x = x + self.mlp_proj(F.gelu(self.mlp_fc(norm(x))))
+        return x
+
+
+class ConverterC(nn.Module):
+    """
+    Phi: the trainable parameters of C. Deliberately small relative to either A or B
+    (Sec 5 budget: 50-300M) -- its size is independent of the student's depth/width
+    because queries and source tokens are built from role/position embeddings rather
+    than per-tensor learned banks.
+    """
+
+    def __init__(self, cfg: ConverterConfig, teacher_num_layers: int, student_num_layers: int):
+        super().__init__()
+        self.cfg = cfg
+        self.teacher_num_layers = teacher_num_layers
+        self.student_num_layers = student_num_layers
+        d = cfg.d_model
+
+        self.role_embed = nn.Embedding(NUM_ROLES, d)
+        self.value_proj = nn.Linear(cfg.chunk_dim, d)
+        self.src_layer_mlp = nn.Linear(d, d)
+        self.src_chunk_mlp = nn.Linear(d, d)
+        self.query_layer_mlp = nn.Linear(d, d)
+        self.query_chunk_mlp = nn.Linear(d, d)
+        self.blocks = nn.ModuleList(_DecoderBlock(d, cfg.n_heads) for _ in range(cfg.n_decoder_blocks))
+        self.out_head = nn.Linear(d, cfg.chunk_dim)
+
+        init_logits = gaussian_bump_layer_map(student_num_layers, teacher_num_layers, cfg.layer_map_tau)
+        self.layer_map_logits = nn.Parameter(init_logits)  # (L_B, L_A), pre-softmax
+
+    def layer_map(self) -> torch.Tensor:
+        """M in the algorithm description: row-softmax, (L_B, L_A)."""
+        return torch.softmax(self.layer_map_logits, dim=-1)
+
+    def encode_layer_tokens(self, corpus: WeightCorpus, layer_idx: int) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+        """Embeds one teacher layer's raw chunks into converter-space tokens. Done
+        lazily (per training step, only for layers actually selected by the top-k
+        relevance window) rather than once upfront, so value_proj / role_embed /
+        src_layer_mlp / src_chunk_mlp all receive gradients from the training loss."""
+        entries = corpus.layers.get(layer_idx, [])
+        if not entries:
+            return None, None
+        device = self.value_proj.weight.device
+        dtype = self.value_proj.weight.dtype
+        tokens, roles = [], []
+        for role_id, chunks in entries:
+            chunks = chunks.to(device=device, dtype=dtype)
+            n = chunks.shape[0]
+            vals = self.value_proj(chunks)
+            role_e = self.role_embed(torch.full((n,), role_id, device=device, dtype=torch.long))
+            layer_frac = torch.full((n,), layer_idx / max(self.teacher_num_layers - 1, 1), device=device)
+            chunk_frac = torch.arange(n, device=device).float() / max(n - 1, 1)
+            tok = (
+                vals
+                + role_e
+                + self.src_layer_mlp(sinusoidal_features(layer_frac, self.cfg.d_model).to(dtype))
+                + self.src_chunk_mlp(sinusoidal_features(chunk_frac, self.cfg.d_model).to(dtype))
+            )
+            tokens.append(tok)
+            roles.append(torch.full((n,), role_id, device=device, dtype=torch.long))
+        return torch.cat(tokens, dim=0), torch.cat(roles, dim=0)
+
+    def build_queries(self, role_id: int, layer_idx: int, num_chunks: int, device, dtype) -> torch.Tensor:
+        role_e = self.role_embed(torch.full((num_chunks,), role_id, device=device, dtype=torch.long))
+        layer_frac = torch.full((num_chunks,), layer_idx / max(self.student_num_layers - 1, 1), device=device)
+        chunk_frac = torch.arange(num_chunks, device=device).float() / max(num_chunks - 1, 1)
+        return (
+            role_e
+            + self.query_layer_mlp(sinusoidal_features(layer_frac, self.cfg.d_model).to(dtype))
+            + self.query_chunk_mlp(sinusoidal_features(chunk_frac, self.cfg.d_model).to(dtype))
+        )
+
+    def generate_tensor(self, corpus: WeightCorpus, target, source_cache: Dict[int, Tuple]) -> torch.Tensor:
+        """target: a student_targets.TargetSpec (key, shape, role, layer_idx)."""
+        numel = 1
+        for s in target.shape:
+            numel *= s
+        num_chunks = -(-numel // self.cfg.chunk_dim)
+        device = self.value_proj.weight.device
+        dtype = self.value_proj.weight.dtype
+
+        queries = self.build_queries(ROLE_TO_ID[target.role], target.layer_idx, num_chunks, device, dtype)
+
+        row = self.layer_map()[target.layer_idx]  # (L_A,)
+        k = min(self.cfg.top_k_source_layers, self.teacher_num_layers)
+        topk_w, topk_idx = torch.topk(row, k)
+
+        src_tokens, src_bias = [], []
+        for w, li in zip(topk_w, topk_idx.tolist()):
+            if li not in source_cache:
+                source_cache[li] = self.encode_layer_tokens(corpus, li)
+            toks, _ = source_cache[li]
+            if toks is None:
+                continue
+            src_tokens.append(toks)
+            src_bias.append(torch.log(w + 1e-8).expand(toks.shape[0]))
+
+        x = queries
+        if src_tokens:
+            source = torch.cat(src_tokens, dim=0)
+            bias = torch.cat(src_bias, dim=0).unsqueeze(0).expand(num_chunks, -1).to(dtype)
+        else:
+            source, bias = None, None
+        for block in self.blocks:
+            x = block(x, source, bias)
+
+        out = self.out_head(x)  # (num_chunks, chunk_dim)
+        flat = out.reshape(-1)[:numel]
+        return flat.view(target.shape).float()
+
+    def generate_state_dict(self, corpus: WeightCorpus, targets) -> Dict[str, torch.Tensor]:
+        """One call = one weight generation (theta_B <- C_phi(theta_A)), Sec 4 Phase 1/2."""
+        source_cache: Dict[int, Tuple] = {}
+        return {t.key: self.generate_tensor(corpus, t, source_cache) for t in targets}

nanochat/distill/kd_corpus.py

@@ -0,0 +1,67 @@
+"""
+Builds/caches the teacher-generated distillation corpus and serves aligned
+(student-tokenized, teacher-tokenized) examples of it. This corpus backs both:
+  - L_KL's tokenizer-mismatch fallback: sequence-level KD, i.e. plain LM loss on
+    Model A's own generations (Sec 3).
+  - L_hidden: pooled hidden-state alignment needs A's and B's hidden states for
+    "the same" text; running each model's own tokenizer over one shared string is the
+    simplest way to get that without a shared vocabulary.
+Examples are kept unbatched (list of single-sequence tensors) rather than padded into
+a batch, since the two tokenizers produce different lengths for the same text and
+padding would corrupt the mean-pooled hidden-state alignment in losses.py.
+"""
+
+import json
+import os
+
+import torch
+
+
+def build_or_load(teacher, prompts, cache_path: str = None, max_new_tokens: int = 256, temperature: float = 0.8):
+    if cache_path and os.path.exists(cache_path):
+        with open(cache_path) as f:
+            return json.load(f)
+    texts = teacher.generate_completions(prompts, max_new_tokens=max_new_tokens, temperature=temperature)
+    records = [{"prompt": p, "text": t} for p, t in zip(prompts, texts)]
+    if cache_path:
+        os.makedirs(os.path.dirname(cache_path), exist_ok=True)
+        with open(cache_path, "w") as f:
+            json.dump(records, f)
+    return records
+
+
+class KDBatcher:
+    """Cycles through the cached (prompt, completion) records, tokenizing each with
+    both tokenizers on demand."""
+
+    def __init__(self, records, student_tokenizer, teacher_hf_tokenizer, student_seq_len: int, teacher_seq_len: int, device):
+        assert records, "KD corpus is empty; build it first with kd_corpus.build_or_load"
+        self.records = records
+        self.student_tokenizer = student_tokenizer
+        self.teacher_hf_tokenizer = teacher_hf_tokenizer
+        self.student_seq_len = student_seq_len
+        self.teacher_seq_len = teacher_seq_len
+        self.device = device
+        self._idx = 0
+
+    def next_examples(self, n: int):
+        """Returns a list of up to n (x, y, teacher_input_ids) tuples, each shaped
+        (1, T) -- one unpadded sequence per example."""
+        examples = []
+        attempts = 0
+        while len(examples) < n and attempts < 4 * n:
+            attempts += 1
+            rec = self.records[self._idx % len(self.records)]
+            self._idx += 1
+            text = rec["prompt"] + rec["text"]
+            s_ids = self.student_tokenizer.encode(text, prepend="<|bos|>")[: self.student_seq_len + 1]
+            if len(s_ids) < 2:
+                continue
+            t_ids = self.teacher_hf_tokenizer(text, truncation=True, max_length=self.teacher_seq_len, return_tensors="pt").input_ids[0]
+            if t_ids.numel() < 1:
+                continue
+            x = torch.tensor(s_ids[:-1], dtype=torch.long, device=self.device).unsqueeze(0)
+            y = torch.tensor(s_ids[1:], dtype=torch.long, device=self.device).unsqueeze(0)
+            t = t_ids.to(self.device).unsqueeze(0)
+            examples.append((x, y, t))
+        return examples

nanochat/distill/losses.py

@@ -0,0 +1,76 @@
+"""
+Phase 1 training losses (CAWT Sec 3) beyond the plain LM cross-entropy that
+GPT.forward already computes (used directly for L_task and, on the teacher-generated
+corpus, for the L_KL tokenizer-mismatch fallback -- see scripts/distill_train.py).
+
+This module implements the two losses that don't already exist elsewhere:
+  - L_hidden: alignment between student and teacher hidden states, reusing the
+    converter's layer map M as the correspondence weighting (Sec 3: "hidden alignment
+    and weight alignment share correspondence, a useful inductive bias").
+  - L_reg: magnitude regularization on generated weights, so Phase 1's B doesn't
+    initialize with exploding attention logits (Sec 3).
+"""
+
+from typing import List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def pool_hidden_states(hidden: torch.Tensor, num_pools: int) -> torch.Tensor:
+    """Mean-pools a (T, D) sequence into (num_pools, D) chunks along the sequence
+    axis. This is the tokenizer-agnostic stand-in for exact per-token alignment: A and
+    B tokenize the same string into different-length sequences, so there is no shared
+    position index to match token-for-token. Pooling to a fixed number of
+    position-normalized chunks gives a coarse but tokenizer-independent alignment."""
+    T, D = hidden.shape
+    if T < num_pools:
+        hidden = F.pad(hidden, (0, 0, 0, num_pools - T))
+    pooled = F.adaptive_avg_pool1d(hidden.float().t().unsqueeze(0), num_pools)
+    return pooled.squeeze(0).t()  # (num_pools, D)
+
+
+def hidden_alignment_loss(
+    student_hidden_by_layer: List[torch.Tensor],
+    teacher_hidden_by_layer: List[torch.Tensor],
+    layer_map: torch.Tensor,
+    projections: nn.ModuleList,
+    num_pools: int = 8,
+) -> torch.Tensor:
+    """
+    student_hidden_by_layer: list[L_B] of (1, T_B, D_B) -- one entry per student block,
+        e.g. captured via forward hooks on model.transformer.h[j] (see distill_train.py).
+    teacher_hidden_by_layer: list[L_A] of (1, T_A, D_A).
+    layer_map: (L_B, L_A) row-softmax correspondence, i.e. converter.layer_map().
+    projections: nn.ModuleList of L_B Linear(D_B -> D_A) modules (P_j in the algorithm).
+    """
+    device = layer_map.device
+    teacher_pooled = torch.stack([pool_hidden_states(h[0], num_pools) for h in teacher_hidden_by_layer]).to(device)  # (L_A, P, D_A)
+
+    total = None
+    for j, h_b in enumerate(student_hidden_by_layer):
+        pooled_b = pool_hidden_states(h_b[0], num_pools).to(device)  # (P, D_B)
+        projected = projections[j](pooled_b)  # (P, D_A)
+        target = torch.einsum("i,ipd->pd", layer_map[j], teacher_pooled)  # soft-mixed teacher target
+        term = F.mse_loss(projected, target)
+        total = term if total is None else total + term
+    return total / max(len(student_hidden_by_layer), 1)
+
+
+def weight_magnitude_reg(generated_state_dict: dict, target_std: float = 0.02) -> torch.Tensor:
+    """Penalizes generated matmul weights whose std drifts *above* the std nanochat's
+    own from-scratch initializer would use for a matrix of that role (Sec 3 L_reg:
+    'avoid exploding attention logits at init'). Shrinking toward zero is not
+    penalized -- an under-confident init is safe and gets corrected by Phase 2
+    fine-tuning; an over-scaled one can blow up the forward pass before that ever
+    happens."""
+    terms = []
+    for tensor in generated_state_dict.values():
+        if tensor.ndim < 2:
+            continue
+        std = tensor.float().std()
+        terms.append(torch.clamp(std - target_std, min=0.0) ** 2)
+    if not terms:
+        return torch.zeros((), device=next(iter(generated_state_dict.values())).device)
+    return torch.stack(terms).mean()

nanochat/distill/roles.py

@@ -0,0 +1,24 @@
+"""
+Canonical weight-tensor role vocabulary shared between teacher (Model A) and student
+(Model B) parameters (CAWT Sec 2.1). Cross-attention in the converter learns role
+*compatibility* through these embeddings rather than a hand-written per-pair mapping
+table -- e.g. nothing here hard-codes "SwiGLU up_proj feeds quantum theta", the
+converter discovers whatever correspondence minimizes the training loss.
+"""
+
+ROLES = [
+    "attn_q",
+    "attn_k",
+    "attn_v",
+    "attn_o",
+    "ffn_gate",
+    "ffn_up",
+    "ffn_down",
+    "conv_in",
+    "conv_out",
+    "conv_depthwise",
+    "quantum_theta",
+]
+
+ROLE_TO_ID = {name: i for i, name in enumerate(ROLES)}
+NUM_ROLES = len(ROLES)

nanochat/distill/shuffle.py

@@ -0,0 +1,39 @@
+"""
+Diagnostic ablation (CAWT Sec 5 & 6): shuffle the *identity* of weight tokens fed to
+the converter while keeping their structural metadata (layer/role/position) intact.
+If C is actually reading theta_A, shuffling should crater downstream quality. If
+performance survives the shuffle, C has collapsed into a hypernetwork that memorizes
+theta_B purely from the training loss and ignores theta_A entirely -- i.e. the
+"translation" is fake. Run this (scripts/distill_shuffle_ablation.py) before investing
+in a full-scale run.
+"""
+
+import copy
+
+import torch
+
+from .weight_tokens import WeightCorpus
+
+
+def shuffle_weight_corpus(corpus: WeightCorpus, seed: int = 0) -> WeightCorpus:
+    """Returns a deep copy of `corpus` with every chunk's *values* permuted globally
+    across the whole teacher, independent of which (layer, role) they originally
+    belonged to. Shapes and per-entry chunk counts are preserved exactly, so this is a
+    drop-in replacement for the real corpus in ConverterC.generate_state_dict."""
+    rng = torch.Generator().manual_seed(seed)
+    shuffled = copy.deepcopy(corpus)
+
+    all_chunks = [chunks for entries in shuffled.layers.values() for _, chunks in entries]
+    if not all_chunks:
+        return shuffled
+    flat = torch.cat([c.reshape(-1) for c in all_chunks])
+    perm = torch.randperm(flat.numel(), generator=rng)
+    flat = flat[perm]
+
+    offset = 0
+    for entries in shuffled.layers.values():
+        for idx, (role_id, chunks) in enumerate(entries):
+            n = chunks.numel()
+            entries[idx] = (role_id, flat[offset : offset + n].view_as(chunks))
+            offset += n
+    return shuffled

nanochat/distill/student_targets.py

@@ -0,0 +1,89 @@
+"""
+Enumerates Model B's (the LFM2-Quantum student's) translatable weight tensors given a
+GPTConfig, and merges converter-generated tensors with a natively-initialized
+reference for everything the converter does not translate.
+
+Translated (matmul weights participating in the forward pass -- attention Q/K/V/O,
+short-conv in/depthwise/out, and FFN up/gate/down or the quantum circuit's theta
+angles): these are exactly the tensors that can plausibly carry transferred knowledge.
+
+Left at nanochat's own from-scratch init (GPT.init_weights): resid_lambdas,
+x0_lambdas, smear_gate/smear_lambda, backout_lambda, ve_gate, value_embeds, and the
+quantum FFN's input/output scale+bias. None of these have a teacher analog -- the
+quantum calibration params in particular are deliberately near-identity/zero at init
+(see QuantumMLP.reset_parameters) to preserve the model's zero-residual training
+stability property, so CAWT leaves them alone rather than learning to override an
+analytic choice (Sec 2.4: "do the linear algebra, don't make C learn it").
+Embeddings (wte, lm_head) are not enumerated here at all -- Sec 2.5 handles those
+through vocab_align.py instead of the chunked cross-attention pathway.
+"""
+
+from dataclasses import dataclass
+from typing import Iterator, List, Tuple
+
+import torch
+
+from nanochat.gpt import GPT
+
+
+@dataclass
+class TargetSpec:
+    key: str
+    shape: Tuple[int, ...]
+    role: str
+    layer_idx: int
+
+
+def iter_translatable_targets(config, translate_quantum_theta: bool = True) -> Iterator[TargetSpec]:
+    head_dim = config.n_embd // config.n_head
+    mixers = config.mixer_types()
+    ffns = config.ffn_types()
+    for i in range(config.n_layer):
+        if mixers[i] == "A":
+            yield TargetSpec(f"transformer.h.{i}.attn.c_q.weight", (config.n_head * head_dim, config.n_embd), "attn_q", i)
+            yield TargetSpec(f"transformer.h.{i}.attn.c_k.weight", (config.n_kv_head * head_dim, config.n_embd), "attn_k", i)
+            yield TargetSpec(f"transformer.h.{i}.attn.c_v.weight", (config.n_kv_head * head_dim, config.n_embd), "attn_v", i)
+            yield TargetSpec(f"transformer.h.{i}.attn.c_proj.weight", (config.n_embd, config.n_embd), "attn_o", i)
+        else:
+            yield TargetSpec(f"transformer.h.{i}.conv.in_proj.weight", (3 * config.n_embd, config.n_embd), "conv_in", i)
+            yield TargetSpec(f"transformer.h.{i}.conv.conv.weight", (config.n_embd, 1, config.conv_kernel), "conv_depthwise", i)
+            yield TargetSpec(f"transformer.h.{i}.conv.out_proj.weight", (config.n_embd, config.n_embd), "conv_out", i)
+
+        ffn_type = ffns[i]
+        if ffn_type == "Q":
+            if translate_quantum_theta:
+                num_qubits = config.quantum_num_qubits
+                num_registers = -(-config.n_embd // num_qubits)  # ceil div, mirrors QuantumMLP
+                yield TargetSpec(f"transformer.h.{i}.mlp.theta", (config.quantum_depth, num_registers, num_qubits), "quantum_theta", i)
+        elif ffn_type == "S":
+            hidden = int(config.n_embd * 8 / 3)
+            hidden = ((hidden + 127) // 128) * 128  # mirrors ClassicalSwiGLU
+            yield TargetSpec(f"transformer.h.{i}.mlp.w1.weight", (hidden, config.n_embd), "ffn_gate", i)
+            yield TargetSpec(f"transformer.h.{i}.mlp.w3.weight", (hidden, config.n_embd), "ffn_up", i)
+            yield TargetSpec(f"transformer.h.{i}.mlp.w2.weight", (config.n_embd, hidden), "ffn_down", i)
+        elif ffn_type == "C":
+            yield TargetSpec(f"transformer.h.{i}.mlp.c_fc.weight", (4 * config.n_embd, config.n_embd), "ffn_up", i)
+            yield TargetSpec(f"transformer.h.{i}.mlp.c_proj.weight", (config.n_embd, 4 * config.n_embd), "ffn_down", i)
+        else:
+            raise ValueError(f"Unknown ffn type {ffn_type!r}")
+
+
+def reference_state_dict(config, device="cpu") -> dict:
+    """A natively-initialized student (GPT.init_weights), used to source every
+    parameter the converter does not translate."""
+    with torch.device("meta"):
+        model = GPT(config)
+    model.to_empty(device=device)
+    model.init_weights()
+    return {k: v.detach().clone() for k, v in model.state_dict().items()}
+
+
+def merge_generated(config, generated: dict, reference: dict = None, translate_quantum_theta: bool = True) -> dict:
+    """Combine converter-generated tensors with the native-init reference for every
+    other key, producing a complete state_dict ready for GPT.load_state_dict."""
+    ref = reference if reference is not None else reference_state_dict(config)
+    merged = dict(ref)
+    for spec in iter_translatable_targets(config, translate_quantum_theta):
+        if spec.key in generated:
+            merged[spec.key] = generated[spec.key].to(ref[spec.key].dtype)
+    return merged

nanochat/distill/teacher.py

@@ -0,0 +1,270 @@
+"""
+Frozen teacher (Model A) wrapper around a HuggingFace causal LM -- Gemma-12B by
+default, but anything with a Llama-family decoder stack works (Gemma, Llama, Mistral,
+Qwen, ...), including natively multimodal checkpoints like google/gemma-3-12b-pt whose
+text config/decoder layers are nested under a `text_config`/`language_model` submodule
+rather than sitting at the top level -- see `_text_config` and `_find_decoder_layers`.
+
+`transformers` is only imported inside this module (lazily, inside __init__), so the
+rest of nanochat/distill and its tests do not require it to be installed -- it is an
+optional extra (`uv sync --extra distill`) since most of nanochat's core workflow
+never touches a HuggingFace model.
+"""
+
+import os
+import time
+
+import torch
+import torch.nn as nn
+
+from nanochat.common import print0
+from .roles import ROLE_TO_ID
+from .weight_tokens import WeightCorpus
+
+# Candidate attribute paths to a Llama-family decoder layer stack, tried in order.
+# Composite/multimodal checkpoints (e.g. Gemma3's *-pt/-it 4B/12B/27B sizes, which are
+# vision+text) nest the text decoder under a `language_model` submodule; plain
+# text-only checkpoints (Llama, Mistral, Gemma3 1B, ...) have it directly under `model`.
+_DECODER_LAYER_PATHS = [
+    "model.layers",
+    "model.language_model.layers",
+    "language_model.model.layers",
+    "language_model.layers",
+    "model.model.layers",
+    "model.text_model.layers",
+    "transformer.h",
+]
+
+
+def _text_config(config):
+    """Resolves the language-model sub-config for composite/multimodal HF configs
+    (e.g. Gemma3Config wraps a Gemma3TextConfig under `.text_config` and exposes it
+    via the standard `get_text_config()` accessor); falls back to `config` itself for
+    plain text-only models where there is nothing to unwrap."""
+    if hasattr(config, "get_text_config"):
+        return config.get_text_config()
+    return getattr(config, "text_config", config)
+
+
+def _find_decoder_layers(model, expected_num_layers: int) -> nn.ModuleList:
+    """Locates the `expected_num_layers`-long ModuleList of decoder layers, trying
+    known attribute paths first and falling back to a generic search over every
+    ModuleList in the model for one of the right length whose elements look like
+    decoder layers (have `self_attn` + `mlp` submodules) -- guards against silently
+    grabbing an unrelated stack (e.g. a vision tower) of a different length."""
+    for path in _DECODER_LAYER_PATHS:
+        obj = model
+        for part in path.split("."):
+            obj = getattr(obj, part, None)
+            if obj is None:
+                break
+        if obj is not None and hasattr(obj, "__len__") and len(obj) == expected_num_layers:
+            return obj
+    for _, module in model.named_modules():
+        if isinstance(module, nn.ModuleList) and len(module) == expected_num_layers:
+            first = module[0]
+            if hasattr(first, "self_attn") and hasattr(first, "mlp"):
+                return module
+    raise ValueError(
+        f"Could not locate the teacher's {expected_num_layers}-layer decoder stack "
+        f"(tried {_DECODER_LAYER_PATHS} and a generic module search). This teacher's "
+        f"architecture may nest its language-model layers differently than the "
+        f"Llama/Gemma family -- extend _DECODER_LAYER_PATHS in nanochat/distill/teacher.py."
+    )
+
+
+class Teacher:
+    """Model A. Always frozen (requires_grad_(False) on every parameter) -- CAWT never
+    updates the teacher, only reads its weights and activations."""
+
+    def __init__(
+        self,
+        model_name: str,
+        device_map="auto",
+        dtype=torch.bfloat16,
+        load_in_4bit: bool = False,
+        trust_remote_code: bool = False,
+    ):
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        from huggingface_hub.utils import enable_progress_bars
+
+        # HF/hub progress bars (download + "Loading checkpoint shards") are on by
+        # default; only force them on if the user hasn't explicitly opted out via
+        # HF_HUB_DISABLE_PROGRESS_BARS (calling enable_progress_bars() when that's set
+        # just raises a UserWarning and does nothing).
+        if os.environ.get("HF_HUB_DISABLE_PROGRESS_BARS") != "1":
+            enable_progress_bars()
+
+        kwargs = dict(device_map=device_map, trust_remote_code=trust_remote_code)
+        if load_in_4bit:
+            from transformers import BitsAndBytesConfig
+
+            kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=dtype)
+        else:
+            kwargs["dtype"] = dtype  # `torch_dtype` is the pre-5.x name; fall back to it below
+
+        print0(f"Loading teacher weights for {model_name} (this can take a while for a multi-GB checkpoint)...")
+        t0 = time.time()
+        try:
+            self.model = AutoModelForCausalLM.from_pretrained(model_name, **kwargs)
+        except TypeError:
+            if "dtype" in kwargs:
+                kwargs["torch_dtype"] = kwargs.pop("dtype")
+                self.model = AutoModelForCausalLM.from_pretrained(model_name, **kwargs)
+            else:
+                raise
+        self.model.eval()
+        for p in self.model.parameters():
+            p.requires_grad_(False)
+        print0(f"Teacher weights loaded in {time.time() - t0:.1f}s")
+
+        print0(f"Loading teacher tokenizer for {model_name}...")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=trust_remote_code)
+        self.model_name = model_name
+        self.config = _text_config(self.model.config)  # unwraps composite/multimodal configs (e.g. Gemma3)
+        self.num_layers = self.config.num_hidden_layers
+        self.hidden_size = self.config.hidden_size
+        self.layers = _find_decoder_layers(self.model, self.num_layers)
+
+    def get_device(self):
+        return next(self.model.parameters()).device
+
+    @staticmethod
+    def _dense_weight(module: nn.Module) -> torch.Tensor:
+        """Returns `module.weight` as a normal, correctly-shaped (out, in) float
+        tensor. Under --teacher-load-in-4bit, transformers replaces every nn.Linear
+        with bitsandbytes' Linear4bit, whose `.weight` is a `Params4bit` -- a *packed*
+        buffer whose physical shape does not match the logical (out, in) matrix.
+        Treating it as a plain tensor (e.g. multiplying it against a norm's scale
+        vector) doesn't raise a shape error; broadcasting silently expands it into an
+        enormous, nonsensical tensor instead. Detect and dequantize it first."""
+        weight = module.weight
+        quant_state = getattr(weight, "quant_state", None)
+        if quant_state is not None:
+            # Params4bit.dequantize() (the convenience method) returns the packed
+            # flat shape rather than the logical (out, in) matrix on this bitsandbytes
+            # version -- use the lower-level functional API, which reshapes correctly.
+            from bitsandbytes.functional import dequantize_4bit
+
+            return dequantize_4bit(weight.data, quant_state).float()
+        return weight.data.float()
+
+    @staticmethod
+    def _fold_input_norm(linear_weight: torch.Tensor, norm_weight: torch.Tensor, gemma_style: bool = True) -> torch.Tensor:
+        """Absorbs an RMSNorm gamma into the Linear that CONSUMES its output:
+        y = norm(x) * gamma @ W^T == x_normed @ (W * gamma)^T (broadcast over the
+        Linear's input dim, i.e. column-scale). Gemma's RMSNorm uses a (1 + weight)
+        scale (its weight is init'd to 0 = identity); set gemma_style=False for a
+        plain-gamma RMSNorm (e.g. Llama)."""
+        scale = (1.0 + norm_weight.data.float()) if gemma_style else norm_weight.data.float()
+        return linear_weight.data.float() * scale.unsqueeze(0)
+
+    @staticmethod
+    def _fold_output_norm(linear_weight: torch.Tensor, norm_weight: torch.Tensor, gemma_style: bool = True) -> torch.Tensor:
+        """Absorbs an RMSNorm gamma applied to a Linear's OUTPUT (Gemma2/3's
+        'sandwich norm': post_attention_layernorm/post_feedforward_layernorm scale the
+        attention/MLP output before the residual add, not the next layer's input) into
+        that Linear's weight, row-scaling instead of column-scaling."""
+        scale = (1.0 + norm_weight.data.float()) if gemma_style else norm_weight.data.float()
+        return linear_weight.data.float() * scale.unsqueeze(1)
+
+    @torch.no_grad()
+    def extract_weight_corpus(self, chunk_dim: int, svd_rank: int, cache_path: str = None, gemma_style_norm: bool = True) -> WeightCorpus:
+        """Phase 0.1 (norm folding) + Phase 0.2 (low-rank compression, via
+        WeightCorpus.add) applied to every attention/FFN matrix in the teacher.
+
+        Handles both the plain Llama-style single-norm-per-sublayer layout (one norm
+        feeding attention, one feeding the MLP) and Gemma2/3's 4-norm "sandwich" layout
+        (input_layernorm -> attn -> post_attention_layernorm -> +residual;
+        pre_feedforward_layernorm -> mlp -> post_feedforward_layernorm -> +residual):
+        the RMSNorm that *feeds* a Linear is folded into its input (column-scale); a
+        norm applied to a sublayer's *output* before the residual add (sandwich-only)
+        is folded into that sublayer's last Linear's output instead (row-scale).
+        """
+        if cache_path and os.path.exists(cache_path):
+            print0(f"Loading cached weight corpus from {cache_path}")
+            return torch.load(cache_path, weights_only=False)
+
+        from tqdm.auto import tqdm
+
+        corpus = WeightCorpus(chunk_dim, self.num_layers)
+        for i, layer in enumerate(tqdm(self.layers, desc="Extracting teacher weight corpus (norm-fold + low-rank SVD)")):
+            input_ln_w = layer.input_layernorm.weight
+            post_attn_ln = getattr(layer, "post_attention_layernorm", None)
+            pre_ffn_ln = getattr(layer, "pre_feedforward_layernorm", None)
+            post_ffn_ln = getattr(layer, "post_feedforward_layernorm", None)
+            sandwich = pre_ffn_ln is not None  # Gemma2/3-style 4-norm layout
+            if not sandwich and post_attn_ln is None:
+                raise ValueError(
+                    f"Layer {i} has neither pre_feedforward_layernorm nor "
+                    f"post_attention_layernorm; cannot determine the MLP's input norm "
+                    f"for this architecture."
+                )
+
+            q = self._fold_input_norm(self._dense_weight(layer.self_attn.q_proj), input_ln_w, gemma_style_norm)
+            k = self._fold_input_norm(self._dense_weight(layer.self_attn.k_proj), input_ln_w, gemma_style_norm)
+            v = self._fold_input_norm(self._dense_weight(layer.self_attn.v_proj), input_ln_w, gemma_style_norm)
+            o = self._dense_weight(layer.self_attn.o_proj)
+            if sandwich and post_attn_ln is not None:
+                o = self._fold_output_norm(o, post_attn_ln.weight, gemma_style_norm)
+            corpus.add(i, ROLE_TO_ID["attn_q"], q, rank=svd_rank)
+            corpus.add(i, ROLE_TO_ID["attn_k"], k, rank=svd_rank)
+            corpus.add(i, ROLE_TO_ID["attn_v"], v, rank=svd_rank)
+            corpus.add(i, ROLE_TO_ID["attn_o"], o, rank=svd_rank)
+
+            # Non-sandwich (Llama/Gemma1-style) models feed the MLP from the same
+            # single norm that also (post-)follows attention (post_attention_layernorm).
+            ffn_input_ln_w = pre_ffn_ln.weight if sandwich else post_attn_ln.weight
+            gate = self._fold_input_norm(self._dense_weight(layer.mlp.gate_proj), ffn_input_ln_w, gemma_style_norm)
+            up = self._fold_input_norm(self._dense_weight(layer.mlp.up_proj), ffn_input_ln_w, gemma_style_norm)
+            down = self._dense_weight(layer.mlp.down_proj)
+            if sandwich and post_ffn_ln is not None:
+                down = self._fold_output_norm(down, post_ffn_ln.weight, gemma_style_norm)
+            corpus.add(i, ROLE_TO_ID["ffn_gate"], gate, rank=svd_rank)
+            corpus.add(i, ROLE_TO_ID["ffn_up"], up, rank=svd_rank)
+            corpus.add(i, ROLE_TO_ID["ffn_down"], down, rank=svd_rank)
+
+        if cache_path:
+            os.makedirs(os.path.dirname(cache_path), exist_ok=True)
+            torch.save(corpus, cache_path)
+        return corpus
+
+    def get_embedding_matrix(self) -> torch.Tensor:
+        return self.model.get_input_embeddings().weight.detach()
+
+    def get_unembedding_matrix(self) -> torch.Tensor:
+        head = self.model.get_output_embeddings()
+        return head.weight.detach() if head is not None else self.get_embedding_matrix()
+
+    @torch.no_grad()
+    def forward_hidden(self, input_ids: torch.Tensor, attention_mask: torch.Tensor = None):
+        """Returns (hidden_states, logits). hidden_states is a tuple of L_A tensors
+        (B, T, d_A), one per transformer layer -- the embedding-layer output
+        (hidden_states[0] in HF's convention) is dropped since it has no student-side
+        counterpart in the per-layer correspondence."""
+        input_ids = input_ids.to(self.get_device())
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(self.get_device())
+        out = self.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True, use_cache=False)
+        return out.hidden_states[1:], out.logits
+
+    @torch.no_grad()
+    def generate_completions(self, prompts, max_new_tokens: int = 256, temperature: float = 0.8):
+        """Builds the sequence-level KD corpus (Sec 3 L_KL fallback: 'train on A's
+        generations', used because A and B use different tokenizers so per-token KL
+        has no shared position index to align on)."""
+        from tqdm.auto import tqdm
+
+        texts = []
+        device = self.get_device()
+        for prompt in tqdm(prompts, desc="Generating teacher completions (KD corpus)"):
+            enc = self.tokenizer(prompt, return_tensors="pt").to(device)
+            gen = self.model.generate(
+                **enc,
+                max_new_tokens=max_new_tokens,
+                do_sample=temperature > 0,
+                temperature=max(temperature, 1e-4),
+                pad_token_id=self.tokenizer.pad_token_id or self.tokenizer.eos_token_id,
+            )
+            texts.append(self.tokenizer.decode(gen[0], skip_special_tokens=True))
+        return texts

nanochat/distill/vocab_align.py

@@ -0,0 +1,101 @@
+"""
+Embedding/tokenizer alignment (CAWT Sec 2.5). Model A (Gemma) and Model B (the
+nanochat rustbpe/tiktoken tokenizer) have different vocabularies, so embeddings are
+handled explicitly rather than through the chunked cross-attention pathway used for
+matmul weights:
+
+  1. Shared-token overlap: for a student token id whose surface string is a *single*
+     token under the teacher's tokenizer too, its teacher correspondence is that one
+     token id.
+  2. FOCUS/WECHSEL-style fallback: for student-only tokens, the correspondence is
+     whatever (possibly multi-token) sequence the teacher's tokenizer produces for the
+     same string; the student embedding is initialized from the mean of those teacher
+     embeddings.
+  3. Token *correspondence* (which teacher token ids map to which student token) is a
+     one-time, tokenizer-only computation with no learnable parameters -- see
+     `build_token_correspondence`. Turning that correspondence into an actual student
+     embedding requires a d_A -> d_B projection (`EmbeddingAligner`), which stays
+     trainable through Phase 1 as part of phi; Phase 2 bakes the result into a real
+     nn.Embedding and fine-tunes it freely.
+
+Splitting it this way keeps the expensive part (looping over the whole vocab, calling
+the HF tokenizer per token) a Phase-0-only cost: every Phase 1 step just re-applies the
+small trainable projection to an already-gathered (vocab_size, teacher_dim) constant.
+"""
+
+from typing import List, Tuple
+
+import torch
+import torch.nn as nn
+
+
+class EmbeddingAligner(nn.Module):
+    """Learned d_A -> d_B projection applied to teacher embedding vectors."""
+
+    def __init__(self, teacher_dim: int, student_dim: int):
+        super().__init__()
+        self.proj = nn.Linear(teacher_dim, student_dim, bias=False)
+        nn.init.orthogonal_(self.proj.weight)
+
+    def forward(self, teacher_vectors: torch.Tensor) -> torch.Tensor:
+        return self.proj(teacher_vectors.float())
+
+
+def _is_special(piece: str) -> bool:
+    return not piece or piece.startswith("<|")
+
+
+def build_token_correspondence(student_tokenizer, teacher_hf_tokenizer, max_vocab: int = None) -> Tuple[List[List[int]], torch.Tensor]:
+    """One-time (Phase 0) pass over the student vocabulary. Returns:
+      - mapping: list of length vocab_size, mapping[sid] = list of teacher token ids
+        whose mean embedding should seed student token sid (empty if no correspondence
+        was found, e.g. for a special token).
+      - filled: bool tensor, filled[sid] = mapping[sid] is non-empty.
+    """
+    from tqdm.auto import tqdm
+
+    vocab_size = student_tokenizer.get_vocab_size()
+    n = vocab_size if max_vocab is None else min(vocab_size, max_vocab)
+    mapping: List[List[int]] = [[] for _ in range(vocab_size)]
+    filled = torch.zeros(vocab_size, dtype=torch.bool)
+    for sid in tqdm(range(n), desc="Building student<->teacher token correspondence"):
+        try:
+            piece = student_tokenizer.id_to_token(sid)
+        except Exception:
+            continue
+        if _is_special(piece):
+            continue
+        try:
+            teacher_ids = teacher_hf_tokenizer.encode(piece, add_special_tokens=False)
+        except Exception:
+            teacher_ids = []
+        if teacher_ids:
+            mapping[sid] = teacher_ids
+            filled[sid] = True
+    return mapping, filled
+
+
+@torch.no_grad()
+def gather_teacher_vectors(mapping: List[List[int]], teacher_embed_matrix: torch.Tensor, device="cpu") -> torch.Tensor:
+    """Applies a token correspondence to one teacher embedding matrix (input or
+    output embeddings -- call this twice, once per matrix, on the same `mapping`).
+    Rows with an empty correspondence are left as zero; the caller is expected to mask
+    them out with `filled` and fall back to the student's own native init there."""
+    d = teacher_embed_matrix.shape[-1]
+    out = torch.zeros(len(mapping), d)
+    teacher_embed_matrix = teacher_embed_matrix.to(device)
+    for sid, ids in enumerate(mapping):
+        if not ids:
+            continue
+        idx = torch.tensor(ids, device=device)
+        out[sid] = teacher_embed_matrix[idx].mean(dim=0).float().cpu()
+    return out
+
+
+def apply_alignment(aligner: EmbeddingAligner, teacher_vectors: torch.Tensor, filled: torch.Tensor, native_init: torch.Tensor) -> torch.Tensor:
+    """Per-training-step application (differentiable w.r.t. `aligner`): projects the
+    precomputed teacher vectors and splices them into the rows that have a
+    correspondence, keeping the student's own native init everywhere else."""
+    projected = aligner(teacher_vectors.to(native_init.device))
+    mask = filled.to(native_init.device).unsqueeze(-1)
+    return torch.where(mask, projected.to(native_init.dtype), native_init)

nanochat/distill/weight_tokens.py

@@ -0,0 +1,78 @@
+"""
+Chunked weight tokenization (CAWT Sec 2.1) and low-rank compression (Phase 0.2).
+
+A full dense mapping from a 12B teacher to a much smaller student is intractable, so
+every source tensor is flattened and sliced into fixed-size "weight token" chunks
+instead of being consumed whole. 2D matrices are additionally factored to a low rank
+via randomized SVD before chunking -- this is what keeps the token budget (and
+therefore the converter's attention cost) tractable for a teacher the size of
+Gemma-12B: a (4096, 4096) matrix at rank 64 is ~31x fewer elements than the dense
+tensor.
+"""
+
+import torch
+import torch.nn.functional as F
+
+
+def lowrank_factors(weight: torch.Tensor, rank: int):
+    """Randomized low-rank SVD: weight (out, in) ~= U_r @ V_r, U_r (out, r), V_r (r, in).
+    Singular values are folded into U_r so V_r's rows are directly comparable across
+    tensors (unit-ish scale), which keeps the weight-token embedder simpler."""
+    assert weight.ndim == 2, f"lowrank_factors expects a 2D matrix, got shape {tuple(weight.shape)}"
+    r = max(1, min(rank, min(weight.shape) - 1 if min(weight.shape) > 1 else 1))
+    w = weight.float()
+    q = min(r + 10, min(w.shape))
+    U, S, V = torch.svd_lowrank(w, q=q, niter=4)
+    U_r = U[:, :r] * S[:r].unsqueeze(0)  # (out, r)
+    V_r = V[:, :r].t()  # (r, in), since svd_lowrank gives w ~= U @ diag(S) @ V.T
+    return U_r.contiguous(), V_r.contiguous()
+
+
+def chunk_tensor(flat: torch.Tensor, chunk_dim: int) -> torch.Tensor:
+    """Pad a 1D tensor to a multiple of chunk_dim and reshape to (num_chunks, chunk_dim)."""
+    assert flat.ndim == 1
+    n = flat.numel()
+    num_chunks = -(-n // chunk_dim)  # ceil div
+    pad = num_chunks * chunk_dim - n
+    if pad > 0:
+        flat = F.pad(flat, (0, pad))
+    return flat.view(num_chunks, chunk_dim)
+
+
+class WeightCorpus:
+    """
+    Frozen source weight tokens extracted from a teacher, grouped by (0-indexed)
+    layer. Chunks are stored on CPU in fp16 -- this is a passive, non-trainable cache;
+    the converter's value/role/position embeddings (which ARE trainable) are applied
+    lazily when a layer's tokens are needed during a training step, so gradients still
+    reach the embedder even though the raw chunk values themselves are frozen constants
+    (Model A's parameters are never updated by CAWT).
+    """
+
+    def __init__(self, chunk_dim: int, num_layers: int):
+        self.chunk_dim = chunk_dim
+        self.num_layers = num_layers
+        # layer_idx -> list[(role_id, chunks (n_i, chunk_dim) fp16 CPU tensor)]
+        self.layers = {i: [] for i in range(num_layers)}
+
+    def add(self, layer_idx: int, role_id: int, tensor: torch.Tensor, rank: int = None):
+        """Add one teacher tensor under (layer_idx, role_id). 2D tensors are
+        low-rank-compressed first (Phase 0.2) when `rank` is given; everything else
+        (1D vectors, small 3D depthwise-conv kernels, etc.) is stored densely."""
+        tensor = tensor.detach()
+        if tensor.ndim == 2 and rank is not None and rank < min(tensor.shape):
+            U, V = lowrank_factors(tensor, rank)
+            self._add_flat(layer_idx, role_id, U.reshape(-1))
+            self._add_flat(layer_idx, role_id, V.reshape(-1))
+        else:
+            self._add_flat(layer_idx, role_id, tensor.float().reshape(-1))
+
+    def _add_flat(self, layer_idx, role_id, flat):
+        chunks = chunk_tensor(flat, self.chunk_dim).to(torch.float16).cpu()
+        self.layers[layer_idx].append((role_id, chunks))
+
+    def num_tokens(self) -> int:
+        return sum(c.shape[0] for entries in self.layers.values() for _, c in entries)
+
+    def num_source_layers_with_data(self) -> int:
+        return sum(1 for entries in self.layers.values() if entries)

nanochat/engine.py

@@ -0,0 +1,367 @@
+"""
+Engine for efficient inference of our models.
+
+Everything works around token sequences:
+- The user can send token sequences to the engine
+- The engine returns the next token
+
+Notes:
+- The engine knows nothing about tokenization, it's purely token id sequences.
+
+The whole thing is made as efficient as possible.
+"""
+
+import torch
+import torch.nn.functional as F
+import signal
+import warnings
+from contextlib import contextmanager
+from collections import deque
+from nanochat.common import compute_init, autodetect_device_type, COMPUTE_DTYPE
+from nanochat.checkpoint_manager import load_model
+
+# -----------------------------------------------------------------------------
+# Calculator tool helpers
+@contextmanager
+def timeout(duration, formula):
+    def timeout_handler(signum, frame):
+        raise Exception(f"'{formula}': timed out after {duration} seconds")
+
+    signal.signal(signal.SIGALRM, timeout_handler)
+    signal.alarm(duration)
+    yield
+    signal.alarm(0)
+
+def eval_with_timeout(formula, max_time=3):
+    try:
+        with timeout(max_time, formula):
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore", SyntaxWarning)
+                return eval(formula, {"__builtins__": {}}, {})
+    except Exception as e:
+        signal.alarm(0)
+        # print(f"Warning: Failed to eval {formula}, exception: {e}") # it's ok ignore wrong calculator usage
+        return None
+
+def use_calculator(expr):
+    """
+    Evaluate a Python expression safely.
+    Supports both math expressions and string operations like .count()
+    """
+    # Remove commas from numbers
+    expr = expr.replace(",", "")
+
+    # Check if it's a pure math expression (old behavior)
+    if all([x in "0123456789*+-/.() " for x in expr]):
+        if "**" in expr:  # disallow power operator
+            return None
+        return eval_with_timeout(expr)
+
+    # Check if it's a string operation we support
+    # Allow: strings (single/double quotes), .count(), letters, numbers, spaces, parens
+    allowed_chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789'\"()._ "
+    if not all([x in allowed_chars for x in expr]):
+        return None
+
+    # Disallow dangerous patterns
+    dangerous_patterns = ['__', 'import', 'exec', 'eval', 'compile', 'open', 'file',
+                         'input', 'raw_input', 'globals', 'locals', 'vars', 'dir',
+                         'getattr', 'setattr', 'delattr', 'hasattr']
+    expr_lower = expr.lower()
+    if any(pattern in expr_lower for pattern in dangerous_patterns):
+        return None
+
+    # Only allow .count() method for now (can expand later)
+    if '.count(' not in expr:
+        return None
+
+    # Evaluate with timeout
+    return eval_with_timeout(expr)
+
+# -----------------------------------------------------------------------------
+class KVCache:
+    """
+    KV Cache designed for Flash Attention 3's flash_attn_with_kvcache API.
+
+    Key differences from FA2-style cache:
+    - Tensors are (B, T, H, D) not (B, H, T, D)
+    - FA3 updates the cache in-place during flash_attn_with_kvcache
+    - Position tracked per batch element via cache_seqlens tensor
+    """
+
+    def __init__(self, batch_size, num_heads, seq_len, head_dim, num_layers, device, dtype):
+        self.batch_size = batch_size
+        self.max_seq_len = seq_len
+        self.n_layers = num_layers
+        self.n_heads = num_heads
+        self.head_dim = head_dim
+        # Pre-allocate cache tensors: (n_layers, B, T, H, D)
+        self.k_cache = torch.zeros(num_layers, batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
+        self.v_cache = torch.zeros(num_layers, batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
+        # Current sequence length per batch element (FA3 needs int32)
+        self.cache_seqlens = torch.zeros(batch_size, dtype=torch.int32, device=device)
+        # Previous token's normalized embedding for smear (set by model forward pass)
+        self.prev_embedding = None
+        # Rolling short-conv state per conv layer: {layer_idx: (B, C, kernel-1)}.
+        # Empty for all-attention models; populated by ShortConvBlock during decode.
+        self.conv_state = {}
+
+    def reset(self):
+        """Reset cache to empty state."""
+        self.cache_seqlens.zero_()
+        self.prev_embedding = None
+        self.conv_state = {}
+
+    def get_conv_state(self, layer_idx):
+        """Return the cached short-conv history for a layer, or None if unset."""
+        return self.conv_state.get(layer_idx)
+
+    def set_conv_state(self, layer_idx, state):
+        """Store the short-conv history (last kernel-1 gated inputs) for a layer."""
+        self.conv_state[layer_idx] = state
+
+    def get_pos(self):
+        """Get current position (assumes all batch elements at same position)."""
+        return self.cache_seqlens[0].item()
+
+    def get_layer_cache(self, layer_idx):
+        """Return (k_cache, v_cache) views for a specific layer."""
+        return self.k_cache[layer_idx], self.v_cache[layer_idx]
+
+    def advance(self, num_tokens):
+        """Advance the cache position by num_tokens."""
+        self.cache_seqlens += num_tokens
+
+    def prefill(self, other):
+        """
+        Copy cached KV from another cache into this one.
+        Used when we do batch=1 prefill and then want to generate multiple samples in parallel.
+        """
+        assert self.get_pos() == 0, "Cannot prefill a non-empty KV cache"
+        assert self.n_layers == other.n_layers and self.n_heads == other.n_heads and self.head_dim == other.head_dim
+        assert self.max_seq_len >= other.max_seq_len
+        other_pos = other.get_pos()
+        self.k_cache[:, :, :other_pos, :, :] = other.k_cache[:, :, :other_pos, :, :]
+        self.v_cache[:, :, :other_pos, :, :] = other.v_cache[:, :, :other_pos, :, :]
+        self.cache_seqlens.fill_(other_pos)
+        # Copy smear state: expand batch=1 prev_embedding to num_samples
+        if other.prev_embedding is not None:
+            self.prev_embedding = other.prev_embedding.expand(self.batch_size, -1, -1).clone()
+        # Copy short-conv state, expanding batch=1 prefill history to num_samples
+        for layer_idx, state in other.conv_state.items():
+            self.conv_state[layer_idx] = state.expand(self.batch_size, -1, -1).clone()
+
+# -----------------------------------------------------------------------------
+@torch.inference_mode()
+def sample_next_token(logits, rng, temperature=1.0, top_k=None):
+    """Sample a single next token from given logits of shape (B, vocab_size). Returns (B, 1)."""
+    assert temperature >= 0.0, "temperature must be non-negative"
+    if temperature == 0.0:
+        return torch.argmax(logits, dim=-1, keepdim=True)
+    if top_k is not None and top_k > 0:
+        k = min(top_k, logits.size(-1))
+        vals, idx = torch.topk(logits, k, dim=-1)
+        vals = vals / temperature
+        probs = F.softmax(vals, dim=-1)
+        choice = torch.multinomial(probs, num_samples=1, generator=rng)
+        return idx.gather(1, choice)
+    else:
+        logits = logits / temperature
+        probs = F.softmax(logits, dim=-1)
+        return torch.multinomial(probs, num_samples=1, generator=rng)
+
+# -----------------------------------------------------------------------------
+
+class RowState:
+    # Per-row state tracking during generation
+    def __init__(self, current_tokens=None):
+        self.current_tokens = current_tokens or [] # Current token sequence for this row
+        self.forced_tokens = deque() # Queue of tokens to force inject
+        self.in_python_block = False # Whether we are inside a python block
+        self.python_expr_tokens = [] # Tokens of the current python expression
+        self.completed = False # Whether this row has completed generation
+
+class Engine:
+
+    def __init__(self, model, tokenizer):
+        self.model = model
+        self.tokenizer = tokenizer # needed for tool use
+
+    @torch.inference_mode()
+    def generate(self, tokens, num_samples=1, max_tokens=None, temperature=1.0, top_k=None, seed=42):
+        """Same as generate, but does single prefill and then clones the KV cache."""
+        assert isinstance(tokens, list) and isinstance(tokens[0], int), "expecting list of ints"
+        device = self.model.get_device()
+        # Allocate the KV cache in the compute dtype so it matches what the forward pass emits
+        dtype = COMPUTE_DTYPE
+        rng = torch.Generator(device=device)
+        rng.manual_seed(seed)
+
+        # Get the special tokens we need to coordinate the tool use state machine
+        get_special = lambda s: self.tokenizer.encode_special(s)
+        python_start = get_special("<|python_start|>")
+        python_end = get_special("<|python_end|>")
+        output_start = get_special("<|output_start|>")
+        output_end = get_special("<|output_end|>")
+        assistant_end = get_special("<|assistant_end|>") # if sampled, ends row
+        bos = self.tokenizer.get_bos_token_id() # if sampled, ends row
+
+        # 1) Run a batch 1 prefill of the prompt tokens
+        m = self.model.config
+        kv_model_kwargs = {"num_heads": m.n_kv_head, "head_dim": m.n_embd // m.n_head, "num_layers": m.n_layer}
+        kv_cache_prefill = KVCache(
+            batch_size=1,
+            seq_len=len(tokens),
+            device=device,
+            dtype=dtype,
+            **kv_model_kwargs,
+        )
+        ids = torch.tensor([tokens], dtype=torch.long, device=device)
+        logits = self.model.forward(ids, kv_cache=kv_cache_prefill)
+        logits = logits[:, -1, :].expand(num_samples, -1)  # (num_samples, vocab_size)
+
+        # 2) Replicate the KV cache for each sample/row
+        kv_length_hint = (len(tokens) + max_tokens) if max_tokens is not None else self.model.config.sequence_len
+        kv_cache_decode = KVCache(
+            batch_size=num_samples,
+            seq_len=kv_length_hint,
+            device=device,
+            dtype=dtype,
+            **kv_model_kwargs,
+        )
+        kv_cache_decode.prefill(kv_cache_prefill)
+        del kv_cache_prefill # no need to keep this memory around
+
+        # 3) Initialize states for each sample
+        row_states = [RowState(tokens.copy()) for _ in range(num_samples)]
+
+        # 4) Main generation loop
+        num_generated = 0
+        while True:
+            # Stop condition: we've reached max tokens
+            if max_tokens is not None and num_generated >= max_tokens:
+                break
+            # Stop condition: all rows are completed
+            if all(state.completed for state in row_states):
+                break
+
+            # Sample the next token for each row
+            next_ids = sample_next_token(logits, rng, temperature, top_k)  # (B, 1)
+            sampled_tokens = next_ids[:, 0].tolist()
+
+            # Process each row: choose the next token, update state, optional tool use
+            token_column = [] # contains the next token id along each row
+            token_masks = [] # contains the mask (was it sampled (1) or forced (0)?) along each row
+            for i, state in enumerate(row_states):
+                # Select the next token in this row
+                is_forced = len(state.forced_tokens) > 0 # are there tokens waiting to be forced in deque?
+                token_masks.append(0 if is_forced else 1) # mask is 0 if forced, 1 if sampled
+                next_token = state.forced_tokens.popleft() if is_forced else sampled_tokens[i]
+                token_column.append(next_token)
+                # Update the state of this row to include the next token
+                state.current_tokens.append(next_token)
+                # On <|assistant_end|> or <|bos|>, mark the row as completed
+                if next_token == assistant_end or next_token == bos:
+                    state.completed = True
+                # Handle tool logic
+                if next_token == python_start:
+                    state.in_python_block = True
+                    state.python_expr_tokens = []
+                elif next_token == python_end and state.in_python_block:
+                    state.in_python_block = False
+                    if state.python_expr_tokens:
+                        expr = self.tokenizer.decode(state.python_expr_tokens)
+                        result = use_calculator(expr)
+                        if result is not None:
+                            result_tokens = self.tokenizer.encode(str(result))
+                            state.forced_tokens.append(output_start)
+                            state.forced_tokens.extend(result_tokens)
+                            state.forced_tokens.append(output_end)
+                    state.python_expr_tokens = []
+                elif state.in_python_block:
+                    state.python_expr_tokens.append(next_token)
+
+            # Yield the token column
+            yield token_column, token_masks
+            num_generated += 1
+
+            # Prepare logits for next iteration
+            ids = torch.tensor(token_column, dtype=torch.long, device=device).unsqueeze(1)
+            logits = self.model.forward(ids, kv_cache=kv_cache_decode)[:, -1, :]  # (B, vocab_size)
+
+    def generate_batch(self, tokens, num_samples=1, **kwargs):
+        """
+        Non-streaming batch generation that just returns the final token sequences.
+        Returns a list of token sequences (list of lists of ints).
+        Terminal tokens (assistant_end, bos) are not included in the results.
+        """
+        assistant_end = self.tokenizer.encode_special("<|assistant_end|>")
+        bos = self.tokenizer.get_bos_token_id()
+        results = [tokens.copy() for _ in range(num_samples)]
+        masks = [[0] * len(tokens) for _ in range(num_samples)]
+        completed = [False] * num_samples
+        for token_column, token_masks in self.generate(tokens, num_samples, **kwargs):
+            for i, (token, mask) in enumerate(zip(token_column, token_masks)):
+                if not completed[i]:
+                    if token == assistant_end or token == bos:
+                        completed[i] = True
+                    else:
+                        results[i].append(token)
+                        masks[i].append(mask)
+            # Stop if all rows are completed
+            if all(completed):
+                break
+        return results, masks
+
+
+if __name__ == "__main__":
+    """
+    Quick inline test to make sure that the naive/slow model.generate function
+    is equivalent to the faster Engine.generate function here.
+    """
+    import time
+    # init compute
+    device_type = autodetect_device_type()
+    ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type)
+    # load the model and tokenizer
+    model, tokenizer, meta = load_model("base", device, phase="eval")
+    bos_token_id = tokenizer.get_bos_token_id()
+    # common hyperparameters
+    kwargs = dict(max_tokens=64, temperature=0.0)
+    # set the starting prompt
+    prompt_tokens = tokenizer.encode("The chemical formula of water is", prepend=bos_token_id)
+    # generate the reference sequence using the model.generate() function
+    generated_tokens = []
+    torch.cuda.synchronize()
+    t0 = time.time()
+    stream = model.generate(prompt_tokens, **kwargs)
+    for token in stream:
+        generated_tokens.append(token)
+        chunk = tokenizer.decode([token])
+        print(chunk, end="", flush=True)
+    print()
+    torch.cuda.synchronize()
+    t1 = time.time()
+    print(f"Reference time: {t1 - t0:.2f}s")
+    reference_ids = generated_tokens
+    # generate tokens with Engine
+    generated_tokens = []
+    engine = Engine(model, tokenizer)
+    stream = engine.generate(prompt_tokens, num_samples=1, **kwargs) # note: runs in fp32
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for token_column, token_masks in stream:
+        token = token_column[0] # only print out the first row
+        generated_tokens.append(token)
+        chunk = tokenizer.decode([token])
+        print(chunk, end="", flush=True)
+    print()
+    torch.cuda.synchronize()
+    t1 = time.time()
+    print(f"Engine time: {t1 - t0:.2f}s")
+    # compare the two sequences
+    for i in range(len(reference_ids)):
+        if reference_ids[i] != generated_tokens[i]:
+            print(f"Mismatch at {i}: {reference_ids[i]} != {generated_tokens[i]}")
+            break
+    print(f"Match: {reference_ids == generated_tokens}")

nanochat/execution.py

@@ -0,0 +1,134 @@
+"""
+Sandboxed execution utilities for running Python code that comes out of an LLM.
+Inspired by the OpenAI HumanEval code:
+https://github.com/openai/human-eval/blob/master/human_eval/execution.py
+
+The code runs in a fresh Python subprocess. What is covered:
+- Each execution runs in its own process (killed hard by the parent on timeout)
+- A fresh interpreter: no access to the parent process memory, and a scrubbed environment
+- Memory limits are enforced via rlimits (256MB by default)
+- stdout and stderr are captured, stdin is disabled
+- Code runs in a temporary directory that is deleted afterwards
+- Destructive functions are disabled (examples: os.system, os.kill, shutil.rmtree, subprocess.Popen)
+
+What is not covered:
+- Not a true security sandbox
+- Network access is not blocked (e.g. sockets could be opened)
+- Python's dynamic features (e.g. ctypes) could bypass restrictions
+- No kernel-level isolation (no seccomp, no containers, no virtualization)
+
+Overall this sandbox is good for evaluation of generated code and protects against
+accidental destructive behavior, but it is not safe against malicious adversarial code.
+"""
+
+import subprocess
+import sys
+import tempfile
+from dataclasses import dataclass
+from typing import Optional
+
+# -----------------------------------------------------------------------------
+
+@dataclass
+class ExecutionResult:
+    """Result of executing Python code in a sandbox."""
+    success: bool
+    stdout: str
+    stderr: str
+    error: Optional[str] = None
+    timeout: bool = False
+    memory_exceeded: bool = False
+
+
+# The guard runs in the subprocess before the untrusted code. It applies the
+# resource limits and disables destructive functions to protect against
+# accidents (a fork bomb, deleting files, killing other processes, ...).
+# It is trivially bypassable by adversarial code, see docstring above.
+GUARD = r"""
+import faulthandler, builtins, os, shutil, subprocess, sys
+maximum_memory_bytes = {maximum_memory_bytes}
+if maximum_memory_bytes is not None and sys.platform != "darwin":
+    # (the resource limit calls seem to fail on macOS, skip them there)
+    import resource
+    resource.setrlimit(resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes))
+    resource.setrlimit(resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes))
+    resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes))
+faulthandler.disable()
+builtins.exit = None
+builtins.quit = None
+builtins.help = None
+os.environ["OMP_NUM_THREADS"] = "1"
+for name in ("kill", "system", "putenv", "remove", "removedirs", "rmdir", "fchdir",
+             "setuid", "fork", "forkpty", "killpg", "rename", "renames", "truncate",
+             "replace", "unlink", "fchmod", "fchown", "chmod", "chown", "chroot",
+             "lchflags", "lchmod", "lchown", "getcwd", "chdir"):
+    setattr(os, name, None)
+for name in ("rmtree", "move", "chown"):
+    setattr(shutil, name, None)
+subprocess.Popen = None
+for name in ("ipdb", "joblib", "resource", "psutil", "tkinter"):
+    sys.modules[name] = None
+"""
+
+
+def execute_code(
+    code: str,
+    timeout: float = 5.0, # 5 seconds default
+    maximum_memory_bytes: Optional[int] = 256 * 1024 * 1024, # 256MB default
+) -> ExecutionResult:
+    """
+    Execute Python code in a sandboxed environment.
+
+    Args:
+        code: Python code to execute as a string
+        timeout: Maximum execution time in seconds (default: 5.0)
+        maximum_memory_bytes: Memory limit in bytes (default: 256MB, None to disable)
+
+    Returns:
+        ExecutionResult with success status, stdout/stderr, and error information
+
+    Example:
+        >>> result = execute_code("print('hello world')")
+        >>> result.success
+        True
+        >>> result.stdout
+        'hello world\\n'
+    """
+    # the guard runs first, then the untrusted code (with fresh globals, as a repr'd literal)
+    guard = GUARD.format(maximum_memory_bytes=maximum_memory_bytes)
+    program = guard + f"\nexec(compile({code!r}, '<llm>', 'exec'), {{'__name__': '__main__'}})\n"
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        try:
+            process = subprocess.run(
+                [sys.executable, "-c", program],
+                cwd=tmpdir, # writes land in the tempdir, deleted afterwards
+                env={"PATH": "/usr/bin:/bin"}, # scrub the environment
+                stdin=subprocess.DEVNULL,
+                capture_output=True,
+                text=True,
+                timeout=timeout,
+            )
+        except subprocess.TimeoutExpired:
+            # subprocess.run kills the child process on timeout
+            return ExecutionResult(
+                success=False,
+                stdout="",
+                stderr="",
+                error="Execution timed out (process killed)",
+                timeout=True,
+            )
+
+    success = process.returncode == 0
+    stderr = process.stderr
+    # the last line of the traceback identifies the exception, e.g. "TypeError: ..."
+    error = None if success else (stderr.strip().splitlines() or ["Execution failed"])[-1]
+    memory_exceeded = "MemoryError" in stderr
+    result = ExecutionResult(
+        success=success,
+        stdout=process.stdout,
+        stderr=stderr,
+        error=error,
+        memory_exceeded=memory_exceeded,
+    )
+    return result

nanochat/flash_attention.py

@@ -0,0 +1,195 @@
+"""
+Unified Flash Attention interface with automatic FA3/SDPA switching.
+
+Exports `flash_attn` module that matches the FA3 API exactly, but falls back
+to PyTorch SDPA on incompatible CUDA GPUs, MPS, and CPU.
+
+Usage (drop-in replacement for FA3):
+    from nanochat.flash_attention import flash_attn
+
+    # Training (no KV cache)
+    y = flash_attn.flash_attn_func(q, k, v, causal=True, window_size=window_size)
+
+    # Inference (with KV cache)
+    y = flash_attn.flash_attn_with_kvcache(q, k_cache, v_cache, k=k, v=v, ...)
+"""
+import torch
+import torch.nn.functional as F
+
+
+# =============================================================================
+# Detection: Try to load FA3 on CUDA GPUs
+# =============================================================================
+def _load_flash_attention_3():
+    """Try to load Flash Attention 3."""
+    if not torch.cuda.is_available():
+        return None
+    try:
+        major, _ = torch.cuda.get_device_capability()
+        # FA3 kernels are currently compiled for Hopper (sm90), Ada (sm89) and Ampere (sm80/sm86)
+        # Blackwell (sm100) needs SDPA fallback until FA3 is recompiled or FA4 is released
+        import os
+        os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
+        from kernels import get_kernel, has_kernel
+        # The varunneal kernel obtains better results for H100/Hopper
+        if major == 9:
+            hf_kernel = "varunneal/flash-attention-3"
+            return get_kernel(hf_kernel).flash_attn_interface
+        else:
+            hf_kernel = "kernels-community/flash-attn3"
+            if has_kernel(hf_kernel):
+                return get_kernel(hf_kernel).flash_attn_interface
+            else:
+                return None
+
+    except Exception:
+        return None
+
+
+_fa3 = _load_flash_attention_3()
+HAS_FA3 = _fa3 is not None
+
+# Override for testing: set to 'fa3', 'sdpa', or None (auto)
+_override_impl = None
+
+
+def _resolve_use_fa3():
+    """Decide once whether to use FA3, based on availability, override, and dtype."""
+    if _override_impl == 'fa3':
+        assert HAS_FA3, "Cannot override to FA3: not available on this hardware"
+        return True
+    if _override_impl == 'sdpa':
+        return False
+    if HAS_FA3:
+        # FA3 Hopper kernels only support bf16 and fp8; fp16/fp32 must use SDPA fallback
+        from nanochat.common import COMPUTE_DTYPE
+        if COMPUTE_DTYPE == torch.bfloat16:
+            return True
+        return False
+    return False
+
+USE_FA3 = _resolve_use_fa3()
+
+
+# =============================================================================
+# SDPA helpers
+# =============================================================================
+def _sdpa_attention(q, k, v, window_size, enable_gqa):
+    """
+    SDPA attention with sliding window support.
+    q, k, v are (B, H, T, D) format.
+    """
+    Tq = q.size(2)
+    Tk = k.size(2)
+    window = window_size[0]
+
+    # Full context, same length
+    if (window < 0 or window >= Tq) and Tq == Tk:
+        return F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=enable_gqa)
+
+    # Single token generation
+    if Tq == 1:
+        if window >= 0 and window < Tk:
+            # window is "left" tokens we need to include (window + 1) keys total
+            start = max(0, Tk - (window + 1))
+            k = k[:, :, start:, :]
+            v = v[:, :, start:, :]
+        return F.scaled_dot_product_attention(q, k, v, is_causal=False, enable_gqa=enable_gqa)
+
+    # Need explicit mask for sliding window/chunk inference
+    device = q.device
+    # For chunk inference (Tq != Tk), is_causal is not aligned to cache position => build an explicit bool mask
+    row_idx = (Tk - Tq) + torch.arange(Tq, device=device).unsqueeze(1)
+    col_idx = torch.arange(Tk, device=device).unsqueeze(0)
+    mask = col_idx <= row_idx
+
+    # sliding window (left)
+    if window >= 0 and window < Tk:
+        mask = mask & ((row_idx - col_idx) <= window)
+
+    return F.scaled_dot_product_attention(q, k, v, attn_mask=mask, enable_gqa=enable_gqa)
+
+# =============================================================================
+# Public API: Same interface as FA3
+# =============================================================================
+def flash_attn_func(q, k, v, causal=False, window_size=(-1, -1)):
+    """
+    Flash Attention for training (no KV cache).
+
+    Args:
+        q, k, v: Tensors of shape (B, T, H, D)
+        causal: Whether to use causal masking
+        window_size: (left, right) sliding window. -1 means unlimited.
+
+    Returns:
+        Output tensor of shape (B, T, H, D)
+    """
+    if USE_FA3:
+        return _fa3.flash_attn_func(q, k, v, causal=causal, window_size=window_size)
+
+    # SDPA fallback: transpose (B, T, H, D) -> (B, H, T, D)
+    q = q.transpose(1, 2)
+    k = k.transpose(1, 2)
+    v = v.transpose(1, 2)
+    enable_gqa = q.size(1) != k.size(1)
+    y = _sdpa_attention(q, k, v, window_size, enable_gqa)
+    return y.transpose(1, 2)  # back to (B, T, H, D)
+
+
+def flash_attn_with_kvcache(q, k_cache, v_cache, k=None, v=None, cache_seqlens=None,
+                            causal=False, window_size=(-1, -1)):
+    """
+    Flash Attention with KV cache for inference.
+
+    FA3 updates k_cache/v_cache in-place. Our SDPA fallback does the same.
+
+    Args:
+        q: Queries, shape (B, T_new, H, D)
+        k_cache, v_cache: Pre-allocated cache tensors, shape (B, T_max, H_kv, D)
+        k, v: New keys/values to insert, shape (B, T_new, H_kv, D)
+        cache_seqlens: Current position in cache, shape (B,) int32
+        causal: Whether to use causal masking
+        window_size: (left, right) sliding window. -1 means unlimited.
+
+    Returns:
+        Output tensor of shape (B, T_new, H, D)
+    """
+    if USE_FA3:
+        return _fa3.flash_attn_with_kvcache(
+            q, k_cache, v_cache, k=k, v=v, cache_seqlens=cache_seqlens,
+            causal=causal, window_size=window_size
+        )
+
+    # SDPA fallback: manually manage KV cache
+    B, T_new, H, D = q.shape
+    pos = cache_seqlens[0].item()  # assume uniform position across batch
+
+    # Insert new k, v into cache (in-place, matching FA3 behavior)
+    if k is not None and v is not None:
+        k_cache[:, pos:pos+T_new, :, :] = k
+        v_cache[:, pos:pos+T_new, :, :] = v
+
+    # Get full cache up to current position + new tokens
+    end_pos = pos + T_new
+    k_full = k_cache[:, :end_pos, :, :]
+    v_full = v_cache[:, :end_pos, :, :]
+
+    # Transpose to SDPA layout: (B, T, H, D) -> (B, H, T, D)
+    q_sdpa = q.transpose(1, 2)
+    k_sdpa = k_full.transpose(1, 2)
+    v_sdpa = v_full.transpose(1, 2)
+
+    enable_gqa = q_sdpa.size(1) != k_sdpa.size(1)
+    y_sdpa = _sdpa_attention(q_sdpa, k_sdpa, v_sdpa, window_size, enable_gqa)
+
+    return y_sdpa.transpose(1, 2)  # back to (B, T, H, D)
+
+
+# =============================================================================
+# Export: flash_attn module interface (drop-in replacement for FA3)
+# =============================================================================
+from types import SimpleNamespace
+flash_attn = SimpleNamespace(
+    flash_attn_func=flash_attn_func,
+    flash_attn_with_kvcache=flash_attn_with_kvcache,
+)

nanochat/fp8.py

@@ -0,0 +1,266 @@
+"""Minimal FP8 training for nanochat — tensorwise dynamic scaling only.
+
+Drop-in replacement for torchao's Float8Linear (~2000 lines) with ~150 lines.
+We only need the "tensorwise" recipe (one scalar scale per tensor), not the full
+generality of torchao (rowwise scaling, FSDP float8 all-gather, DTensor, tensor
+subclass dispatch tables, etc.)
+
+How FP8 training works
+======================
+A standard Linear layer does one matmul in forward and two in backward:
+  forward:      output     = input      @ weight.T
+  backward:     grad_input = grad_output @ weight
+                grad_weight= grad_output.T @ input
+
+FP8 training wraps each of these three matmuls with:
+  1. Compute scale = FP8_MAX / max(|tensor|)  for each operand
+  2. Quantize: fp8_tensor = clamp(tensor * scale, -FP8_MAX, FP8_MAX).to(fp8)
+  3. Matmul via torch._scaled_mm (cuBLAS FP8 kernel, ~2x faster than bf16)
+  4. Dequantize: _scaled_mm handles this internally using the inverse scales
+
+The key insight: torch._scaled_mm and the float8 dtypes are PyTorch built-ins.
+torchao is just orchestration around these primitives. We can call them directly.
+
+FP8 dtype choice
+================
+There are two FP8 formats. We use both, following the standard convention:
+  - float8_e4m3fn: 4-bit exponent, 3-bit mantissa, range [-448, 448]
+    Higher precision (more mantissa bits), used for input and weight.
+  - float8_e5m2:   5-bit exponent, 2-bit mantissa, range [-57344, 57344]
+    Wider range (more exponent bits), used for gradients which can be large.
+
+torch._scaled_mm layout requirements
+=====================================
+The cuBLAS FP8 kernel requires specific memory layouts:
+  - First argument (A):  must be row-major (contiguous)
+  - Second argument (B): must be column-major (B.t().contiguous().t())
+If B is obtained by transposing a contiguous tensor (e.g. weight.t()), it is
+already column-major — no copy needed. Otherwise we use _to_col_major().
+
+How this differs from torchao's approach
+========================================
+torchao uses a "tensor subclass" architecture: Float8TrainingTensor is a subclass
+of torch.Tensor that bundles FP8 data + scale + metadata. It implements
+__torch_dispatch__ with a dispatch table that intercepts every aten op (mm, t,
+reshape, clone, ...) and handles it in FP8-aware fashion. When you call
+  output = input @ weight.T
+the @ operator dispatches to aten.mm, which gets intercepted and routed to
+torch._scaled_mm behind the scenes. This is ~2000 lines of code because you need
+a handler for every tensor operation that might touch an FP8 tensor.
+
+We take a simpler approach: a single autograd.Function (_Float8Matmul) that takes
+full-precision inputs, quantizes to FP8 internally, calls _scaled_mm, and returns
+full-precision outputs. Marked @allow_in_graph so torch.compile treats it as one
+opaque node rather than trying to trace inside.
+
+The trade-off is in how torch.compile sees the two approaches:
+  - torchao: compile decomposes the tensor subclass (via __tensor_flatten__) and
+    sees every individual op (amax, scale, cast, _scaled_mm) as separate graph
+    nodes. Inductor can fuse these with surrounding operations (e.g. fuse the
+    amax computation with the preceding layer's activation function).
+  - ours: compile sees a single opaque call. It can optimize everything around
+    the FP8 linear (attention, norms, etc.) but cannot fuse across the boundary.
+
+Both call the exact same cuBLAS _scaled_mm kernel — the GPU matmul is identical.
+The difference is only in the "glue" ops (amax, scale, cast) which are tiny
+compared to the matmul. In practice this means our version is slightly faster
+(less compilation overhead, no tensor subclass dispatch cost) but can produce
+subtly different floating-point rounding paths under torch.compile, since Inductor
+generates a different graph. Numerics are bitwise identical in eager mode.
+"""
+
+import torch
+import torch.nn as nn
+
+from nanochat.common import COMPUTE_DTYPE
+
+# Avoid division by zero when computing scale from an all-zeros tensor
+EPS = 1e-12
+
+
+@torch.no_grad()
+def _to_fp8(x, fp8_dtype):
+    """Dynamically quantize a tensor to FP8 using tensorwise scaling.
+
+    "Tensorwise" means one scalar scale for the entire tensor (as opposed to
+    "rowwise" which computes a separate scale per row). Tensorwise is faster
+    because cuBLAS handles the scaling; rowwise needs the CUTLASS kernel.
+
+    Returns (fp8_data, inverse_scale) for use with torch._scaled_mm.
+    """
+    fp8_max = torch.finfo(fp8_dtype).max
+    # Compute the max absolute value across the entire tensor
+    amax = x.float().abs().max()
+    # Scale maps [0, amax] -> [0, fp8_max]. Use float64 for the division to
+    # ensure consistent numerics between torch.compile and eager mode.
+    # (torchao does the same upcast — without it, compile/eager can diverge)
+    scale = fp8_max / amax.double().clamp(min=EPS)
+    scale = scale.float()
+    # Quantize: scale into FP8 range, saturate (clamp prevents overflow when
+    # casting — PyTorch's default is to wrap, not saturate), then cast to FP8
+    x_scaled = x.float() * scale
+    x_clamped = x_scaled.clamp(-fp8_max, fp8_max)
+    x_fp8 = x_clamped.to(fp8_dtype)
+    # _scaled_mm expects the *inverse* of our scale (it multiplies by this to
+    # convert FP8 values back to the original range during the matmul)
+    inv_scale = scale.reciprocal()
+    return x_fp8, inv_scale
+
+
+def _to_col_major(x):
+    """Rearrange a 2D tensor's memory to column-major layout.
+
+    torch._scaled_mm requires its second operand in column-major layout.
+    The trick: transpose -> contiguous (forces a copy in transposed order)
+    -> transpose back. The result has the same logical shape but column-major
+    strides, e.g. a [M, N] tensor gets strides (1, M) instead of (N, 1).
+    """
+    return x.t().contiguous().t()
+
+
+# allow_in_graph tells torch.compile to treat this as an opaque operation —
+# dynamo won't try to decompose it into smaller ops. See the module docstring
+# for how this differs from torchao's tensor subclass approach.
+@torch._dynamo.allow_in_graph
+class _Float8Matmul(torch.autograd.Function):
+    """Custom autograd for the three FP8 GEMMs of a Linear layer.
+
+    The forward quantizes input and weight to FP8 and saves
+    the quantized tensors + scales for backward.
+    """
+
+    @staticmethod
+    def forward(ctx, input_2d, weight):
+        # Quantize both operands to e4m3 (higher precision format)
+        input_fp8, input_inv = _to_fp8(input_2d, torch.float8_e4m3fn)
+        weight_fp8, weight_inv = _to_fp8(weight, torch.float8_e4m3fn)
+        ctx.save_for_backward(input_fp8, input_inv, weight_fp8, weight_inv)
+
+        # output = input @ weight.T
+        # input_fp8 is [B, K] contiguous = row-major (good for first arg)
+        # weight_fp8 is [N, K] contiguous, so weight_fp8.t() is [K, N] with
+        # strides (1, K) = column-major (good for second arg, no copy needed!)
+        output = torch._scaled_mm(
+            input_fp8,
+            weight_fp8.t(),
+            scale_a=input_inv,
+            scale_b=weight_inv,
+            out_dtype=input_2d.dtype,
+            # use_fast_accum=True accumulates the dot products in lower precision.
+            # Slightly less accurate but measurably faster. Standard practice for
+            # the forward pass; we use False in backward for more precise gradients.
+            use_fast_accum=True,
+        )
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        in_fp8, in_inv, w_fp8, w_inv = ctx.saved_tensors
+
+        # === GEMM 1: grad_input = grad_output @ weight ===
+        # Shapes: [B, N] @ [N, K] -> [B, K]
+        # Gradients use e5m2 (wider range), weights use e4m3 (higher precision)
+        go_fp8, go_inv = _to_fp8(grad_output, torch.float8_e5m2)
+        # go_fp8 is [B, N] contiguous = row-major, good for first arg
+        # w_fp8 is [N, K] contiguous = row-major, need column-major for second arg
+        w_col = _to_col_major(w_fp8)
+        grad_input = torch._scaled_mm(
+            go_fp8,
+            w_col,
+            scale_a=go_inv,
+            scale_b=w_inv,
+            out_dtype=grad_output.dtype,
+            use_fast_accum=False,
+        )
+
+        # === GEMM 2: grad_weight = grad_output.T @ input ===
+        # Shapes: [N, B] @ [B, K] -> [N, K]
+        # go_fp8 is [B, N] contiguous, we need go.T = [N, B] as first arg.
+        # Transposing gives column-major, but first arg needs row-major,
+        # so we must call .contiguous() to physically rearrange the memory.
+        go_T = go_fp8.t().contiguous()  # [N, B] row-major
+        in_col = _to_col_major(in_fp8)    # [B, K] column-major
+        grad_weight = torch._scaled_mm(
+            go_T,
+            in_col,
+            scale_a=go_inv,
+            scale_b=in_inv,
+            out_dtype=grad_output.dtype,
+            use_fast_accum=False,
+        )
+
+        return grad_input, grad_weight
+
+
+class Float8Linear(nn.Linear):
+    """Drop-in nn.Linear replacement that does FP8 compute.
+
+    Weights and biases remain in their original precision (e.g. fp32/bf16).
+    Only the matmul is performed in FP8 via the _Float8Matmul autograd function.
+    """
+
+    def forward(self, input):
+        # Cast input to COMPUTE_DTYPE (typically bf16) since _scaled_mm expects
+        # reduced precision input, and we no longer rely on autocast to do this.
+        input = input.to(COMPUTE_DTYPE)
+        # _scaled_mm only works on 2D tensors, so flatten batch dimensions
+        orig_shape = input.shape
+        input_2d = input.reshape(-1, orig_shape[-1])
+        output = _Float8Matmul.apply(input_2d, self.weight)
+        output = output.reshape(*orig_shape[:-1], output.shape[-1])
+        if self.bias is not None:
+            output = output + self.bias.to(output.dtype)
+        return output
+
+    @classmethod
+    def from_float(cls, mod):
+        """Create Float8Linear from nn.Linear, sharing the same weight and bias.
+
+        Uses meta device to avoid allocating a temporary weight tensor — we
+        create the module shell on meta (shapes/dtypes only, no memory), then
+        point .weight and .bias to the original module's parameters.
+        """
+        with torch.device("meta"):
+            new_mod = cls(mod.in_features, mod.out_features, bias=False)
+        new_mod.weight = mod.weight
+        new_mod.bias = mod.bias
+        return new_mod
+
+
+class Float8LinearConfig:
+    """Minimal config matching torchao's API. Only tensorwise recipe is supported."""
+
+    @staticmethod
+    def from_recipe_name(recipe_name):
+        if recipe_name != "tensorwise":
+            raise ValueError(
+                f"Only 'tensorwise' recipe is supported, got '{recipe_name}'. "
+                f"Rowwise/axiswise recipes require the full torchao library."
+            )
+        return Float8LinearConfig()
+
+
+def convert_to_float8_training(module, *, config=None, module_filter_fn=None):
+    """Replace nn.Linear layers with Float8Linear throughout a module.
+
+    Walks the module tree in post-order (children before parents) and swaps
+    each nn.Linear that passes the optional filter. The new Float8Linear shares
+    the original weight and bias tensors — no copies, no extra memory.
+
+    Args:
+        module: Root module to convert.
+        config: Float8LinearConfig (accepted for API compat, only tensorwise supported).
+        module_filter_fn: Optional filter(module, fqn) -> bool. Only matching Linears
+            are converted. Common use: skip layers with dims not divisible by 16
+            (hardware requirement for FP8 matmuls on H100).
+    """
+    def _convert(mod, prefix=""):
+        for name, child in mod.named_children():
+            fqn = f"{prefix}.{name}" if prefix else name
+            _convert(child, fqn)
+            if isinstance(child, nn.Linear) and not isinstance(child, Float8Linear):
+                if module_filter_fn is None or module_filter_fn(child, fqn):
+                    setattr(mod, name, Float8Linear.from_float(child))
+
+    _convert(module)
+    return module

nanochat/gpt.py

@@ -0,0 +1,916 @@
+"""
+GPT model (rewrite, a lot simpler)
+Notable features:
+- rotary embeddings (and no positional embeddings)
+- QK norm
+- untied weights for token embedding and lm_head
+- PyTorch-native variational quantum feed-forward network
+- norm after token embedding
+- no learnable params in rmsnorm
+- no bias in linear layers
+- Group-Query Attention (GQA) support for more efficient inference
+- Flash Attention 3 integration
+"""
+
+from functools import partial
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+from nanochat.common import get_dist_info, print0, COMPUTE_DTYPE
+from nanochat.optim import MuonAdamW
+
+# Our custom Flash Attention module that automatically uses FA3 when compatible and SDPA fallback otherwise
+from nanochat.flash_attention import flash_attn
+
+@dataclass
+class GPTConfig:
+    sequence_len: int = 2048
+    vocab_size: int = 32768
+    n_layer: int = 12
+    n_head: int = 6 # number of query heads
+    n_kv_head: int = 6 # number of key/value heads (GQA)
+    n_embd: int = 768
+    # Feed-forward network. "quantum" is a PyTorch-native exact simulator;
+    # "classical" retains the old ReLU² MLP for loading legacy checkpoints.
+    mlp_type: str = "quantum"
+    quantum_num_qubits: int = 4
+    quantum_depth: int = 2
+    # Sliding window attention pattern string, tiled across layers. Final layer always L.
+    # Characters: L=long (full context), S=short (quarter context)
+    # Examples: "L"=all full context, "SL"=alternating, "SSL"=two short then one long
+    window_pattern: str = "SSSL"
+    # LFM2-style hybrid backbone. mixer_pattern selects the per-layer token mixer,
+    # tiled across layers: 'A'=grouped-query attention, 'C'=gated short convolution.
+    # Default "A" reproduces the original all-attention model (back-compat).
+    mixer_pattern: str = "A"
+    conv_kernel: int = 3  # depthwise short-conv kernel size for 'C' layers
+    rope_theta: float = 100000.0  # RoPE base frequency (LFM2 uses 1e6)
+    # Per-layer feed-forward selection, tiled across layers: 'Q'=quantum circuit,
+    # 'C'=classical ReLU² MLP, 'S'=classical SwiGLU. Empty ⇒ derive from mlp_type.
+    ffn_pattern: str = ""
+    use_value_embeddings: bool = True  # ResFormer value embeddings (attention layers only)
+
+    def mixer_types(self):
+        return _tile_pattern(self.mixer_pattern, self.n_layer, "AC", "mixer_pattern")
+
+    def ffn_types(self):
+        pattern = self.ffn_pattern
+        if not pattern:
+            pattern = {"quantum": "Q", "classical": "C"}[self.mlp_type]
+        return _tile_pattern(pattern, self.n_layer, "QCS", "ffn_pattern")
+
+
+def _tile_pattern(pattern, n_layer, valid_chars, name):
+    """Tile/validate a pattern string to exactly n_layer characters."""
+    pattern = pattern.upper()
+    assert pattern, f"{name} must be non-empty"
+    assert all(c in valid_chars for c in pattern), \
+        f"Invalid {name}={pattern!r}; use only characters from {valid_chars!r}"
+    return [pattern[i % len(pattern)] for i in range(n_layer)]
+
+
+def _layer_has_ve(config, layer_idx):
+    """Value embeddings apply only to attention layers, when enabled, following
+    the alternating (last-layer-included) ResFormer schedule."""
+    if not getattr(config, "use_value_embeddings", True):
+        return False
+    if config.mixer_types()[layer_idx] != "A":
+        return False
+    return has_ve(layer_idx, config.n_layer)
+
+
+def norm(x):
+    return F.rms_norm(x, (x.size(-1),)) # note that this will run in bf16, seems ok
+
+class Linear(nn.Linear):
+    """nn.Linear that casts weights to match input dtype in forward.
+    Replaces autocast: master weights stay fp32 for optimizer precision,
+    but matmuls run in the activation dtype (typically bf16 from embeddings)."""
+    def forward(self, x):
+        return F.linear(x, self.weight.to(dtype=x.dtype))
+
+
+def has_ve(layer_idx, n_layer):
+    """Returns True if GPT layer should have Value Embedding (alternating, last layer always included)."""
+    return layer_idx % 2 == (n_layer - 1) % 2
+
+def apply_rotary_emb(x, cos, sin):
+    # note: this rotates by -theta, the transpose of the textbook convention. Functionally
+    # equivalent (only the relative q/k rotation matters), kept for checkpoint compatibility.
+    assert x.ndim == 4  # multihead attention
+    d = x.shape[3] // 2
+    x1, x2 = x[..., :d], x[..., d:] # split up last dim into two halves
+    y1 = x1 * cos + x2 * sin # rotate pairs of dims
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], 3)
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config, layer_idx, has_value_embed=None):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.n_embd = config.n_embd
+        self.head_dim = self.n_embd // self.n_head
+        assert self.n_embd % self.n_head == 0
+        assert self.n_kv_head <= self.n_head and self.n_head % self.n_kv_head == 0
+        self.c_q = Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
+        self.c_k = Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_v = Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = Linear(self.n_embd, self.n_embd, bias=False)
+        self.ve_gate_channels = 12
+        # has_value_embed defaults to the original alternating schedule when constructed
+        # directly (e.g. in tests); GPT passes the mixer-aware value from _layer_has_ve.
+        if has_value_embed is None:
+            has_value_embed = has_ve(layer_idx, config.n_layer)
+        self.ve_gate = Linear(self.ve_gate_channels, self.n_kv_head, bias=False) if has_value_embed else None
+
+    def forward(self, x, ve, cos_sin, window_size, kv_cache):
+        B, T, C = x.size()
+
+        # Project the input to get queries, keys, and values
+        # Shape: (B, T, H, D) - FA3's native layout, no transpose needed!
+        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
+        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
+        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
+
+        # Value residual (ResFormer): mix in value embedding with input-dependent gate per head
+        if ve is not None:
+            ve = ve.view(B, T, self.n_kv_head, self.head_dim)
+            gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))  # (B, T, n_kv_head), range (0, 3)
+            v = v + gate.unsqueeze(-1) * ve
+
+        # Apply Rotary Embeddings to queries and keys to get relative positional encoding
+        cos, sin = cos_sin
+        q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin)
+        q, k = norm(q), norm(k) # QK norm
+        q = q * 1.2  # sharper attention (split scale between Q and K), TODO think through better
+        k = k * 1.2
+
+        # Flash Attention (FA3 or SDPA fallback)
+        # window_size is (left, right) tuple: (N, 0) for causal, (-1, 0) for full context
+        if kv_cache is None:
+            # Training: causal attention with optional sliding window
+            y = flash_attn.flash_attn_func(q, k, v, causal=True, window_size=window_size)
+        else:
+            # Inference: use flash_attn_with_kvcache which handles cache management
+            k_cache, v_cache = kv_cache.get_layer_cache(self.layer_idx)
+            y = flash_attn.flash_attn_with_kvcache(
+                q, k_cache, v_cache,
+                k=k, v=v,
+                cache_seqlens=kv_cache.cache_seqlens,
+                causal=True,
+                window_size=window_size,
+            )
+            # Advance position after last layer processes
+            if self.layer_idx == kv_cache.n_layers - 1:
+                kv_cache.advance(T)
+
+        # Re-assemble the heads and project back to residual stream
+        y = y.contiguous().view(B, T, -1)
+        y = self.c_proj(y)
+        return y
+
+
+class ClassicalMLP(nn.Module):
+    """The original dense ReLU² MLP, retained for legacy checkpoints."""
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = Linear(config.n_embd, 4 * config.n_embd, bias=False)
+        self.c_proj = Linear(4 * config.n_embd, config.n_embd, bias=False)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = F.relu(x).square()
+        x = self.c_proj(x)
+        return x
+
+
+class QuantumMLP(nn.Module):
+    """
+    A differentiable variational quantum circuit implemented with PyTorch.
+
+    Embedding channels are split into independent small quantum registers. Each
+    register angle-encodes its inputs with RY gates, applies trainable RY gates
+    and a ring of CNOT entanglers, then returns Pauli-Z expectation values. The
+    simulator is exact and uses ordinary PyTorch tensor operations, so gradients,
+    device placement, state dicts, and torch.compile work without PennyLane/Qiskit.
+
+    This is a simulator: memory is O(2**num_qubits) per register. Keeping
+    num_qubits small (the default is 4) makes it practical inside a language model.
+    """
+    def __init__(self, config):
+        super().__init__()
+        self.n_embd = config.n_embd
+        self.num_qubits = config.quantum_num_qubits
+        self.depth = config.quantum_depth
+        if not 2 <= self.num_qubits <= 8:
+            raise ValueError("quantum_num_qubits must be between 2 and 8")
+        if self.depth < 1:
+            raise ValueError("quantum_depth must be at least 1")
+
+        self.num_registers = (self.n_embd + self.num_qubits - 1) // self.num_qubits
+        self.padded_embd = self.num_registers * self.num_qubits
+        parameter_shape = (self.depth, self.num_registers, self.num_qubits)
+        self.theta = nn.Parameter(torch.empty(parameter_shape))
+        self.input_scale = nn.Parameter(torch.empty(self.num_registers, self.num_qubits))
+        self.input_bias = nn.Parameter(torch.empty(self.num_registers, self.num_qubits))
+        self.output_scale = nn.Parameter(torch.empty(self.num_registers, self.num_qubits))
+        self.output_bias = nn.Parameter(torch.empty(self.num_registers, self.num_qubits))
+
+        cnot_permutations, z_signs = self._make_quantum_buffers(self.theta.device)
+        self.register_buffer("cnot_permutations", cnot_permutations, persistent=False)
+        self.register_buffer("z_signs", z_signs, persistent=False)
+        self.reset_parameters()
+
+    def _make_quantum_buffers(self, device):
+        """Build basis permutations for CNOTs and Pauli-Z measurement signs."""
+        num_states = 1 << self.num_qubits
+        basis = torch.arange(num_states, device=device, dtype=torch.long)
+        permutations = []
+        for control in range(self.num_qubits):
+            target = (control + 1) % self.num_qubits
+            control_is_one = ((basis >> control) & 1).bool()
+            permutations.append(torch.where(control_is_one, basis ^ (1 << target), basis))
+        cnot_permutations = torch.stack(permutations)
+
+        signs = []
+        for qubit in range(self.num_qubits):
+            signs.append(1.0 - 2.0 * ((basis >> qubit) & 1).float())
+        return cnot_permutations, torch.stack(signs)
+
+    @torch.no_grad()
+    def reset_parameters(self):
+        # Near-identity input encoding with small trainable variational angles.
+        nn.init.normal_(self.theta, mean=0.0, std=0.02)
+        nn.init.ones_(self.input_scale)
+        nn.init.zeros_(self.input_bias)
+        # Zero output preserves the transformer's residual-path initialization.
+        nn.init.zeros_(self.output_scale)
+        nn.init.zeros_(self.output_bias)
+        # __init__ may run under a meta-device context. GPT.init_weights calls
+        # this again after to_empty(), at which point real buffers are required.
+        if not self.theta.is_meta:
+            cnot_permutations, z_signs = self._make_quantum_buffers(self.theta.device)
+            self.cnot_permutations = cnot_permutations
+            self.z_signs = z_signs
+
+    @staticmethod
+    def _apply_ry(state, angle, qubit):
+        """Apply a batched RY(angle) gate to one qubit of a real statevector."""
+        low_dim = 1 << qubit
+        high_dim = state.size(-1) // (2 * low_dim)
+        paired = state.reshape(*state.shape[:-1], high_dim, 2, low_dim)
+        amplitude_zero, amplitude_one = paired.unbind(dim=-2)
+        half_angle = angle.unsqueeze(-1).unsqueeze(-1) * 0.5
+        cosine, sine = torch.cos(half_angle), torch.sin(half_angle)
+        rotated_zero = cosine * amplitude_zero - sine * amplitude_one
+        rotated_one = sine * amplitude_zero + cosine * amplitude_one
+        return torch.stack((rotated_zero, rotated_one), dim=-2).flatten(-3)
+
+    def forward(self, x):
+        input_dtype = x.dtype
+        # Trigonometric statevector evolution is kept in fp32 for stability.
+        angles = x.float()
+        if self.padded_embd != self.n_embd:
+            angles = F.pad(angles, (0, self.padded_embd - self.n_embd))
+        angles = angles.view(*x.shape[:-1], self.num_registers, self.num_qubits)
+        angles = angles * self.input_scale + self.input_bias
+
+        state = angles.new_zeros(*angles.shape[:-1], 1 << self.num_qubits)
+        state[..., 0] = 1.0  # |00...0>
+
+        # Data encoding.
+        for qubit in range(self.num_qubits):
+            state = self._apply_ry(state, angles[..., qubit], qubit)
+
+        # Hardware-efficient variational layers: rotations followed by a CNOT ring.
+        parameter_prefix = (1,) * (angles.ndim - 2)
+        for layer in range(self.depth):
+            for qubit in range(self.num_qubits):
+                angle = self.theta[layer, :, qubit].view(*parameter_prefix, self.num_registers)
+                state = self._apply_ry(state, angle, qubit)
+            for permutation in self.cnot_permutations:
+                state = state.index_select(-1, permutation)
+
+        probabilities = state.square()
+        expectations = torch.einsum("...s,qs->...q", probabilities, self.z_signs)
+        output = expectations * self.output_scale + self.output_bias
+        output = output.flatten(-2)[..., :self.n_embd]
+        return output.to(input_dtype)
+
+
+class ClassicalSwiGLU(nn.Module):
+    """SwiGLU feed-forward (as used by LFM2). Kept available as a classical
+    alternative to the quantum FFN; not used by the default quantum preset."""
+    def __init__(self, config):
+        super().__init__()
+        # ~8/3·d_model rounded to a multiple of 128 keeps the SwiGLU param budget
+        # comparable to a 4·d ReLU MLP (two input projections instead of one).
+        hidden = int(config.n_embd * 8 / 3)
+        hidden = ((hidden + 127) // 128) * 128
+        self.w1 = Linear(config.n_embd, hidden, bias=False)  # gate
+        self.w3 = Linear(config.n_embd, hidden, bias=False)  # up
+        self.w2 = Linear(hidden, config.n_embd, bias=False)  # down
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class ShortConvBlock(nn.Module):
+    """LFM2-style double-gated short-range convolution token mixer.
+
+    A single input projection produces three gates (B, C, x). The value x is
+    gated by B, passed through a causal depthwise short convolution, gated by C,
+    then projected out:  y = out_proj( C * depthwise_conv( B * x ) ).
+    Complexity is O(n) in sequence length (vs O(n²) attention). Causality is
+    enforced by left-padding the convolution by kernel-1 so position t never
+    reads position t+1. During KV-cache decoding the last kernel-1 gated inputs
+    are cached so incremental single-token steps match the full-sequence result.
+    """
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_embd = config.n_embd
+        self.kernel = config.conv_kernel
+        self.in_proj = Linear(config.n_embd, 3 * config.n_embd, bias=False)
+        self.conv = nn.Conv1d(config.n_embd, config.n_embd, kernel_size=self.kernel,
+                              groups=config.n_embd, bias=False)
+        self.out_proj = Linear(config.n_embd, config.n_embd, bias=False)
+
+    def forward(self, x, kv_cache=None):
+        B, T, C = x.size()
+        b_gate, c_gate, value = self.in_proj(x).chunk(3, dim=-1)
+        conv_in = (b_gate * value).transpose(1, 2)  # (B, C, T)
+        pad = self.kernel - 1
+        prev = None if kv_cache is None else kv_cache.get_conv_state(self.layer_idx)
+        if prev is None:
+            padded = F.pad(conv_in, (pad, 0))  # causal left pad
+        else:
+            padded = torch.cat([prev.to(conv_in.dtype), conv_in], dim=-1)
+        y = F.conv1d(padded, self.conv.weight.to(dtype=padded.dtype), groups=self.n_embd)  # valid conv over (pad + T) -> length T
+        if kv_cache is not None and pad > 0:
+            kv_cache.set_conv_state(self.layer_idx, padded[:, :, -pad:].detach())
+        y = y.transpose(1, 2)  # (B, T, C)
+        y = c_gate * y
+        return self.out_proj(y)
+
+
+def _make_ffn(config, layer_idx):
+    ffn_type = config.ffn_types()[layer_idx]
+    if ffn_type == "Q":
+        return QuantumMLP(config)
+    elif ffn_type == "C":
+        return ClassicalMLP(config)
+    elif ffn_type == "S":
+        return ClassicalSwiGLU(config)
+    raise ValueError(f"Unknown ffn type {ffn_type!r}; expected 'Q', 'C' or 'S'")
+
+
+class Block(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        # Token mixer: attention or short convolution. Submodule is named `attn`
+        # for attention (preserving legacy checkpoint keys) and `conv` for
+        # convolution; exactly one is non-None.
+        mixer_type = config.mixer_types()[layer_idx]
+        if mixer_type == "A":
+            self.attn = CausalSelfAttention(config, layer_idx, has_value_embed=_layer_has_ve(config, layer_idx))
+            self.conv = None
+        else:
+            self.attn = None
+            self.conv = ShortConvBlock(config, layer_idx)
+        self.mlp = _make_ffn(config, layer_idx)
+
+    def forward(self, x, ve, cos_sin, window_size, kv_cache):
+        if self.attn is not None:
+            x = x + self.attn(norm(x), ve, cos_sin, window_size, kv_cache)
+        else:
+            x = x + self.conv(norm(x), kv_cache)
+        x = x + self.mlp(norm(x))
+        return x
+
+
+class GPT(nn.Module):
+    def __init__(self, config, pad_vocab_size_to=64):
+        """
+        NOTE a major footgun: this __init__ function runs in meta device context (!!)
+        Therefore, any calculations inside here are shapes and dtypes only, no actual data.
+        => We actually initialize all data (parameters, buffers, etc.) in init_weights() instead.
+        """
+        super().__init__()
+        self.config = config
+        self.gradient_checkpointing = False
+        # Compute per-layer window sizes for sliding window attention
+        # window_size is (left, right) tuple: (-1, 0) for full context, (N, 0) for sliding window
+        self.window_sizes = self._compute_window_sizes(config)
+        # Pad vocab for efficiency (DDP, tensor cores). This is just an optimization - outputs are cropped in forward().
+        # https://huggingface.co/docs/transformers/main_classes/model#transformers.PreTrainedModel.resize_token_embeddings
+        padded_vocab_size = ((config.vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to) * pad_vocab_size_to
+        if padded_vocab_size != config.vocab_size:
+            print0(f"Padding vocab_size from {config.vocab_size} to {padded_vocab_size} for efficiency")
+        self.transformer = nn.ModuleDict({
+            "wte": nn.Embedding(padded_vocab_size, config.n_embd),
+            "h": nn.ModuleList([Block(config, layer_idx) for layer_idx in range(config.n_layer)]),
+        })
+        self.lm_head = Linear(config.n_embd, padded_vocab_size, bias=False)
+        # Per-layer learnable scalars (inspired by modded-nanogpt)
+        # resid_lambdas: scales the residual stream at each layer (init 1.0 = neutral)
+        # x0_lambdas: blends initial embedding back in at each layer (init 0.0 = disabled)
+        # Separate parameters so they can have different optimizer treatment
+        self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer))   # fake init, real init in init_weights()
+        self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer))     # fake init, real init in init_weights()
+        # Smear: mix previous token's embedding into current token (cheap bigram-like info)
+        self.smear_gate = Linear(24, 1, bias=False)
+        self.smear_lambda = nn.Parameter(torch.zeros(1))
+        # Backout: subtract cached mid-layer residual before final norm to remove low-level features
+        self.backout_lambda = nn.Parameter(0.2 * torch.ones(1))
+        # Value embeddings (ResFormer-style): alternating layers, last layer always included
+        head_dim = config.n_embd // config.n_head
+        kv_dim = config.n_kv_head * head_dim
+        self.value_embeds = nn.ModuleDict({str(i): nn.Embedding(padded_vocab_size, kv_dim) for i in range(config.n_layer) if _layer_has_ve(config, i)})
+        # To support meta device initialization, we init the rotary embeddings here, but it's just "fake" meta tensors only.
+        # As for rotary_seq_len, these rotary embeddings are pretty small/cheap in memory,
+        # so let's just over-compute them by 10X, but assert fail if we ever reach that amount.
+        # In the future we can dynamically grow the cache, for now it's fine.
+        self.rotary_seq_len = config.sequence_len * 10 # 10X over-compute should be enough, TODO make nicer?
+        head_dim = config.n_embd // config.n_head
+        cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim)
+        self.register_buffer("cos", cos, persistent=False) # persistent=False means it's not saved to the checkpoint
+        self.register_buffer("sin", sin, persistent=False)
+
+    @torch.no_grad()
+    def init_weights(self):
+        """
+        Initialize the full model in this one function for maximum clarity.
+
+        wte (embedding):     normal, std=1.0
+        lm_head:             normal, std=0.001
+        for each block:
+            attn.c_q:        uniform, std=1/sqrt(n_embd)
+            attn.c_k:        uniform, std=1/sqrt(n_embd)
+            attn.c_v:        uniform, std=1/sqrt(n_embd)
+            attn.c_proj:     zeros
+            mlp.c_fc:        uniform, std=1/sqrt(n_embd)
+            mlp.c_proj:      zeros
+        """
+
+        # Embedding and unembedding
+        torch.nn.init.normal_(self.transformer.wte.weight, mean=0.0, std=0.8)
+        torch.nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.001)
+
+        # Transformer blocks: uniform init with bound = sqrt(3) * std (same standard deviation as normal)
+        n_embd = self.config.n_embd
+        s = 3**0.5 * n_embd**-0.5 # sqrt(3) multiplier makes sure Uniform achieves the same std as Normal
+        for block in self.transformer.h:
+            # Token mixer
+            if block.attn is not None:
+                torch.nn.init.uniform_(block.attn.c_q.weight, -s, s) # weights use Uniform to avoid outliers
+                torch.nn.init.uniform_(block.attn.c_k.weight, -s, s)
+                torch.nn.init.uniform_(block.attn.c_v.weight, -s, s)
+                torch.nn.init.zeros_(block.attn.c_proj.weight) # projections are zero
+            else:  # ShortConvBlock
+                torch.nn.init.uniform_(block.conv.in_proj.weight, -s, s)
+                torch.nn.init.uniform_(block.conv.conv.weight, -s, s)  # depthwise (C,1,k)
+                torch.nn.init.zeros_(block.conv.out_proj.weight)  # zero-residual init like c_proj
+            # Feed-forward network
+            if isinstance(block.mlp, QuantumMLP):
+                block.mlp.reset_parameters()
+            elif isinstance(block.mlp, ClassicalSwiGLU):
+                torch.nn.init.uniform_(block.mlp.w1.weight, -s * 0.4, s * 0.4)
+                torch.nn.init.uniform_(block.mlp.w3.weight, -s * 0.4, s * 0.4)
+                torch.nn.init.zeros_(block.mlp.w2.weight)  # zero-residual init
+            else:  # ClassicalMLP (ReLU²)
+                torch.nn.init.uniform_(block.mlp.c_fc.weight, -s * 0.4, s * 0.4)  # 0.4x init scale for c_fc
+                torch.nn.init.zeros_(block.mlp.c_proj.weight)
+
+        # Per-layer scalars
+        # Per-layer resid init: stronger residual at early layers, weaker at deep layers
+        n_layer = self.config.n_layer
+        for i in range(n_layer):
+            self.resid_lambdas.data[i] = 1.15 - (0.10 * i / max(n_layer - 1, 1))
+        # Decaying x0 init: earlier layers get more input embedding blending
+        for i in range(n_layer):
+            self.x0_lambdas.data[i] = 0.20 - (0.15 * i / max(n_layer - 1, 1))
+
+        # Smear/backout scalars and smear gate must be explicitly initialized 
+        torch.nn.init.zeros_(self.smear_lambda)
+        torch.nn.init.constant_(self.backout_lambda, 0.2)
+        torch.nn.init.uniform_(self.smear_gate.weight, 0.0, 0.02)
+
+        # Value embeddings (init like c_v: uniform with same std)
+        for ve in self.value_embeds.values():
+            torch.nn.init.uniform_(ve.weight, -s, s)
+
+        # Gate weights init with small positive values so gates start slightly above neutral
+        for block in self.transformer.h:
+            if block.attn is not None and block.attn.ve_gate is not None:
+                torch.nn.init.uniform_(block.attn.ve_gate.weight, 0.0, 0.02)
+
+        # Rotary embeddings
+        head_dim = self.config.n_embd // self.config.n_head
+        cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim)
+        self.cos, self.sin = cos, sin
+
+        # Cast embeddings to COMPUTE_DTYPE: optimizer can tolerate reduced-precision
+        # embeddings and it saves memory. Exception: fp16 requires fp32 embeddings
+        # because GradScaler cannot unscale fp16 gradients.
+        if COMPUTE_DTYPE != torch.float16:
+            self.transformer.wte.to(dtype=COMPUTE_DTYPE)
+            for ve in self.value_embeds.values():
+                ve.to(dtype=COMPUTE_DTYPE)
+
+    def _precompute_rotary_embeddings(self, seq_len, head_dim, base=None, device=None):
+        # Base (theta) defaults to the model config (LFM2 uses 1e6; legacy 1e5).
+        if base is None:
+            base = self.config.rope_theta
+        # autodetect the device from model embeddings
+        if device is None:
+            device = self.transformer.wte.weight.device
+        # stride the channels
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        # stride the time steps
+        t = torch.arange(seq_len, dtype=torch.float32, device=device)
+        # calculate the rotation frequencies at each (time, channel) pair
+        freqs = torch.outer(t, inv_freq)
+        cos, sin = freqs.cos(), freqs.sin()
+        cos, sin = cos.to(COMPUTE_DTYPE), sin.to(COMPUTE_DTYPE)
+        cos, sin = cos[None, :, None, :], sin[None, :, None, :] # add batch and head dims for later broadcasting
+        return cos, sin
+
+    def _compute_window_sizes(self, config):
+        """
+        Compute per-layer window sizes for sliding window attention.
+
+        Returns list of (left, right) tuples for FA3's window_size parameter:
+        - left: how many tokens before current position to attend to (-1 = unlimited)
+        - right: how many tokens after current position to attend to (0 for causal)
+
+        Pattern string is tiled across layers. Final layer always gets L (full context).
+        Characters: L=long (full context), S=short (quarter context)
+        """
+        pattern = config.window_pattern.upper()
+        assert all(c in "SL" for c in pattern), f"Invalid window_pattern: {pattern}. Use only S and L."
+        # Map characters to window sizes
+        long_window = config.sequence_len
+        short_window = -(-long_window // 4 // 128) * 128  # ceil to FA3 tile size (2048 -> 768)
+        char_to_window = {
+            "L": (long_window, 0),
+            "S": (short_window, 0),
+        }
+        # Tile pattern across layers
+        window_sizes = []
+        for layer_idx in range(config.n_layer):
+            char = pattern[layer_idx % len(pattern)]
+            window_sizes.append(char_to_window[char])
+        # Final layer always gets full context
+        window_sizes[-1] = (long_window, 0)
+        return window_sizes
+
+    def get_device(self):
+        return self.transformer.wte.weight.device
+
+    def estimate_flops(self):
+        """
+        Return the estimated FLOPs per token for the model (forward + backward).
+        Each matmul weight parameter contributes 2 FLOPs (multiply *, accumulate +) in forward, and 2X that in backward => 2+4=6.
+        Cleanest explanation of this: https://medium.com/@dzmitrybahdanau/the-flops-calculus-of-language-model-training-3b19c1f025e4
+        On top of that, 12 * h * q * effective_seq_len accounts for key @ query matmul flops inside attention.
+        With sliding windows, effective_seq_len varies per layer (capped by window size).
+        Ref: https://arxiv.org/abs/2204.02311 (PaLM paper).
+        This is ~1% off from the exact formulas of Chinchilla paper, the difference is:
+        - Chinchilla counts the embedding layer as flops (? weird, it's just a lookup => we ignore)
+        - Chinchilla counts exp/sum/divide in attention softmax as flops (a little sus and very tiny => we ignore)
+        """
+        h, q, t = self.config.n_head, self.config.n_embd // self.config.n_head, self.config.sequence_len
+        # Sum attention FLOPs over attention layers only (conv layers have no attention),
+        # accounting for sliding window.
+        attn_flops = 0
+        for i, block in enumerate(self.transformer.h):
+            if block.attn is None:
+                continue
+            window = self.window_sizes[i][0]  # (left, right) tuple, we use left
+            effective_seq = t if window < 0 else min(window, t)
+            attn_flops += 12 * h * q * effective_seq
+        num_flops_per_token = (
+            6 * self.num_matmul_params()
+            + 3 * self.quantum_forward_flops_per_token()
+            + 3 * self.conv_forward_flops_per_token()
+            + attn_flops
+        )
+        return num_flops_per_token
+
+    def num_matmul_params(self):
+        """
+        The number of parameters that participate in matmuls with the token stream,
+        i.e. contribute 2 FLOPs/param to the forward pass. Counted structurally: every
+        matmul in this model goes through the Linear class, while non-matmul params
+        (embeddings = lookups, per-layer scalars) are nn.Embedding or raw Parameters.
+        """
+        matmul_params = sum(m.weight.numel() for m in self.modules() if isinstance(m, Linear))
+        return matmul_params
+
+    def quantum_forward_flops_per_token(self):
+        """
+        Approximate arithmetic cost of the simulated quantum circuits.
+
+        RY gates and Z measurements each scale with register state size. CNOT
+        gates are basis permutations, so they contribute memory traffic but no
+        floating-point arithmetic here.
+        """
+        total = 0
+        for module in self.modules():
+            if isinstance(module, QuantumMLP):
+                states = 1 << module.num_qubits
+                rotations_and_measurement = module.depth + 2
+                total += (
+                    3
+                    * module.num_registers
+                    * module.num_qubits
+                    * states
+                    * rotations_and_measurement
+                )
+        return total
+
+    def conv_forward_flops_per_token(self):
+        """Arithmetic cost of the depthwise short convolutions (per token). The
+        in/out projections are Linear and already counted in num_matmul_params."""
+        total = 0
+        for block in self.transformer.h:
+            if block.conv is not None:
+                total += 2 * self.config.n_embd * block.conv.kernel
+        return total
+
+    def estimate_decode_flops(self, context_len):
+        """
+        Forward FLOPs to decode one token at a given context length during inference:
+        2 FLOPs per matmul param, plus attention over min(context, window) per attention layer.
+        """
+        h = self.config.n_head
+        q = self.config.n_embd // self.config.n_head
+        attn_flops = 0
+        for i, block in enumerate(self.transformer.h):
+            if block.attn is None:
+                continue
+            window = self.window_sizes[i][0]
+            attn_flops += 4 * h * q * min(context_len, window)
+        decode_flops = (
+            2 * self.num_matmul_params()
+            + self.quantum_forward_flops_per_token()
+            + self.conv_forward_flops_per_token()
+            + attn_flops
+        )
+        return decode_flops
+
+    def estimate_prefill_flops(self, num_tokens):
+        """Forward FLOPs to prefill a prompt: causal, so token t attends to min(t, window)."""
+        h = self.config.n_head
+        q = self.config.n_embd // self.config.n_head
+        attn_flops = 0
+        for i, block in enumerate(self.transformer.h):
+            if block.attn is None:
+                continue
+            window = self.window_sizes[i][0]
+            w = min(window, num_tokens)
+            attended_tokens = w * (w + 1) // 2 + (num_tokens - w) * w # ramp up to w, then flat
+            attn_flops += 4 * h * q * attended_tokens
+        prefill_flops = (
+            (2 * self.num_matmul_params()
+             + self.quantum_forward_flops_per_token()
+             + self.conv_forward_flops_per_token())
+            * num_tokens
+            + attn_flops
+        )
+        return prefill_flops
+
+    def kv_bytes_per_token(self):
+        """Bytes to *store* one token of KV cache during inference, per row (all layers)."""
+        head_dim = self.config.n_embd // self.config.n_head
+        kv_dtype_bytes = COMPUTE_DTYPE.itemsize # the KV cache is kept in the compute dtype
+        return self.config.n_layer * 2 * self.config.n_kv_head * head_dim * kv_dtype_bytes
+
+    def kv_read_bytes(self, context_len):
+        """Bytes of KV cache *read* by one decode step at a given context length, per row.
+        Sliding window layers only attend to (and read) the last `window` tokens."""
+        head_dim = self.config.n_embd // self.config.n_head
+        kv_dtype_bytes = COMPUTE_DTYPE.itemsize
+        total = 0
+        for window, _ in self.window_sizes:
+            total += 2 * self.config.n_kv_head * head_dim * kv_dtype_bytes * min(context_len, window)
+        return total
+
+    def num_scaling_params(self):
+        """
+        Return detailed parameter counts for scaling law analysis.
+        Different papers use different conventions:
+        - Kaplan et al. excluded embedding parameters
+        - Chinchilla included all parameters
+        Ref: https://arxiv.org/abs/2203.15556 (Chinchilla paper)
+        Ref: https://arxiv.org/abs/2001.08361 (Kaplan et al. original scaling laws paper)
+
+        Returns a dict with counts for each parameter group, so downstream analysis
+        can experiment with which combination gives the cleanest scaling laws.
+        """
+        # Count each group separately (mirrors the grouping in setup_optimizers)
+        wte = sum(p.numel() for p in self.transformer.wte.parameters())
+        value_embeds = sum(p.numel() for p in self.value_embeds.parameters())
+        lm_head = sum(p.numel() for p in self.lm_head.parameters())
+        transformer_matrices = sum(p.numel() for p in self.transformer.h.parameters())
+        scalars = self.resid_lambdas.numel() + self.x0_lambdas.numel() + self.smear_gate.weight.numel() + self.smear_lambda.numel() + self.backout_lambda.numel()
+        total = wte + value_embeds + lm_head + transformer_matrices + scalars
+        assert total == sum(p.numel() for p in self.parameters()), "Parameter count mismatch"
+        return {
+            'wte': wte,
+            'value_embeds': value_embeds,
+            'lm_head': lm_head,
+            'transformer_matrices': transformer_matrices,
+            'scalars': scalars,
+            'total': total,
+        }
+
+    def setup_optimizer(self, unembedding_lr=0.004, embedding_lr=0.2, matrix_lr=0.02, weight_decay=0.0, scalar_lr=0.5, muon_group_size=-1, circuit_lr=None):
+        model_dim = self.config.n_embd
+
+        # Separate out all parameters into groups
+        transformer_params = list(self.transformer.h.parameters())
+        circuit_params = [
+            p
+            for module in self.transformer.h.modules()
+            if isinstance(module, QuantumMLP)
+            for p in module.parameters()
+        ]
+        circuit_param_ids = {id(p) for p in circuit_params}
+        non_circuit = [p for p in transformer_params if id(p) not in circuit_param_ids]
+        # Muon only handles 2D matrices. The depthwise short-conv weight is 3D
+        # (channels, 1, kernel), so it joins the quantum-circuit params in AdamW.
+        matrix_params = [p for p in non_circuit if p.ndim == 2]
+        conv_params = [p for p in non_circuit if p.ndim != 2]
+        assert all(p.ndim == 2 for p in matrix_params), "Muon only supports matrix parameters"
+        # The quantum circuit's parameters are rotation angles, not linear weights: they
+        # are periodic and much more LR-sensitive, so circuit_lr optionally puts them in
+        # their own AdamW group. Leaving it None keeps the single combined group, which
+        # matters because splitting changes the param_group layout and therefore breaks
+        # warm-starting the optimizer from a checkpoint saved without the split.
+        if circuit_lr is None:
+            adamw_extra_params = circuit_params + conv_params  # non-matrix transformer params
+        else:
+            adamw_extra_params = conv_params
+        value_embeds_params = list(self.value_embeds.parameters())
+        embedding_params = list(self.transformer.wte.parameters())
+        lm_head_params = list(self.lm_head.parameters())
+        resid_params = [self.resid_lambdas]
+        x0_params = [self.x0_lambdas]
+        smear_params = [self.smear_gate.weight, self.smear_lambda, self.backout_lambda]
+        split_circuit_params = [] if circuit_lr is None else circuit_params # counted separately below when split out
+        assert len(list(self.parameters())) == len(matrix_params) + len(adamw_extra_params) + len(split_circuit_params) + len(embedding_params) + len(lm_head_params) + len(value_embeds_params) + len(resid_params) + len(x0_params) + len(smear_params)
+
+        # Scale the LR for the AdamW parameters by ∝1/√dmodel (tuned for 768 dim model)
+        dmodel_lr_scale = (model_dim / 768) ** -0.5
+        print0(f"Scaling the LR for the AdamW parameters ∝1/√({model_dim}/768) = {dmodel_lr_scale:.6f}")
+
+        # Build param_groups with all required fields explicit
+        param_groups = [
+            # AdamW groups (embeddings, lm_head, scalars)
+            dict(kind='adamw', params=lm_head_params, lr=unembedding_lr * dmodel_lr_scale, betas=(0.8, 0.96), eps=1e-10, weight_decay=0.01),
+            dict(kind='adamw', params=embedding_params, lr=embedding_lr * dmodel_lr_scale, betas=(0.8, 0.995), eps=1e-10, weight_decay=0.001),
+            dict(kind='adamw', params=value_embeds_params, lr=embedding_lr * dmodel_lr_scale * 0.5, betas=(0.8, 0.995), eps=1e-10, weight_decay=0.01),
+            dict(kind='adamw', params=resid_params, lr=scalar_lr * 0.01, betas=(0.8, 0.95), eps=1e-10, weight_decay=0.05),
+            dict(kind='adamw', params=x0_params, lr=scalar_lr, betas=(0.96, 0.95), eps=1e-10, weight_decay=0.0),  # higher beta1 for x0
+            dict(kind='adamw', params=smear_params, lr=0.2, betas=(0.8, 0.95), eps=1e-10, weight_decay=0.0),
+        ]
+        if adamw_extra_params:
+            param_groups.append(dict(
+                kind='adamw', params=adamw_extra_params, lr=matrix_lr,
+                betas=(0.9, 0.99), eps=1e-8, weight_decay=weight_decay,
+            ))
+        if circuit_lr is not None and circuit_params:
+            print0(f"Quantum circuit parameters in their own AdamW group at lr={circuit_lr}")
+            param_groups.append(dict(
+                kind='adamw', params=circuit_params, lr=circuit_lr,
+                betas=(0.9, 0.99), eps=1e-8, weight_decay=0.0,
+            ))
+        # Muon groups (matrix params, grouped by shape for stacking)
+        for shape in sorted({p.shape for p in matrix_params}):
+            shape_params = [p for p in matrix_params if p.shape == shape]
+            group_size = len(shape_params) if muon_group_size <= 0 else muon_group_size
+            for start in range(0, len(shape_params), group_size):
+                group_params = shape_params[start:start + group_size]
+                param_groups.append(dict(
+                    kind='muon', params=group_params, lr=matrix_lr,
+                    momentum=0.95, ns_steps=5, beta2=0.9, weight_decay=weight_decay,
+                ))
+
+        optimizer = MuonAdamW(param_groups)
+        for group in optimizer.param_groups:
+            group["initial_lr"] = group["lr"]
+        return optimizer
+
+    def forward(self, idx, targets=None, kv_cache=None, loss_reduction='mean'):
+        B, T = idx.size()
+
+        # Grab the rotary embeddings for the current sequence length (they are of shape (1, seq_len, 1, head_dim/2))
+        assert T <= self.cos.size(1), f"Sequence length grew beyond the rotary embeddings cache: {T} > {self.cos.size(1)}"
+        assert idx.device == self.cos.device, f"Rotary embeddings and idx are on different devices: {idx.device} != {self.cos.device}"
+        assert self.cos.dtype == COMPUTE_DTYPE, f"Rotary embeddings must be in {COMPUTE_DTYPE}, got {self.cos.dtype}"
+        # if kv cache exists, we need to offset the rotary embeddings to the current position in the cache
+        T0 = 0 if kv_cache is None else kv_cache.get_pos()
+        cos_sin = self.cos[:, T0:T0+T], self.sin[:, T0:T0+T] # truncate cache to current sequence length
+
+        # Embed the tokens
+        x = self.transformer.wte(idx) # embed current token
+        x = x.to(COMPUTE_DTYPE) # ensure activations are in compute dtype (no-op usually, but active for fp16 code path)
+        x = norm(x)
+
+        # Smear: mix previous token's embedding into current position (cheap bigram info)
+        if kv_cache is None:
+            # Training / naive generate: full sequence available, use fast slice
+            assert T > 1, "Training forward pass should have T > 1"
+            gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
+            x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
+        else:
+            # KV cache inference: read prev embedding from cache, store current for next step
+            x_pre_smear = kv_cache.prev_embedding
+            kv_cache.prev_embedding = x[:, -1:, :]
+            if T > 1:
+                # Prefill: apply smear to positions 1+, same as training
+                gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24]))
+                x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1)
+            elif x_pre_smear is not None:
+                # Decode: single token, use cached prev embedding
+                gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, :, :24]))
+                x = x + gate * x_pre_smear
+
+        # Forward the trunk of the Transformer
+        x0 = x  # save initial normalized embedding for x0 residual
+        n_layer = self.config.n_layer
+        backout_layer = n_layer // 2  # cache at halfway point
+        x_backout = None
+        for i, block in enumerate(self.transformer.h):
+            x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
+            ve = self.value_embeds[str(i)](idx).to(x.dtype) if str(i) in self.value_embeds else None
+            if self.gradient_checkpointing and self.training and kv_cache is None:
+                # Recompute the block during backward instead of retaining its
+                # large attention and quantum-statevector activations.
+                def block_forward(block_input, value_embedding, block=block, window_size=self.window_sizes[i]):
+                    return block(block_input, value_embedding, cos_sin, window_size, None)
+                x = checkpoint(block_forward, x, ve, use_reentrant=False)
+            else:
+                x = block(x, ve, cos_sin, self.window_sizes[i], kv_cache)
+            if i == backout_layer:
+                x_backout = x
+        # Subtract mid-layer residual to remove low-level features before logit projection
+        if x_backout is not None:
+            x = x - self.backout_lambda.to(x.dtype) * x_backout
+        x = norm(x)
+
+        # Forward the lm_head (compute logits)
+        softcap = 15 # smoothly cap the logits to the range [-softcap, softcap]
+        logits = self.lm_head(x) # (B, T, padded_vocab_size) <- very big tensor, large amount of memory
+        logits = logits[..., :self.config.vocab_size] # slice to remove padding
+        logits = logits.float() # switch to fp32 for logit softcap and loss computation
+        logits = softcap * torch.tanh(logits / softcap) # squash the logits
+
+        if targets is not None:
+            # training: given the targets, compute and return the loss
+            # TODO experiment with chunked cross-entropy?
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1, reduction=loss_reduction)
+            return loss
+        else:
+            # inference: just return the logits directly
+            return logits
+
+    @torch.inference_mode()
+    def generate(self, tokens, max_tokens, temperature=1.0, top_k=None, seed=42):
+        """
+        Naive autoregressive streaming inference.
+        To make it super simple, let's assume:
+        - batch size is 1
+        - ids and the yielded tokens are simple Python lists and ints
+        """
+        assert isinstance(tokens, list)
+        device = self.get_device()
+        rng = None
+        if temperature > 0:
+            rng = torch.Generator(device=device)
+            rng.manual_seed(seed)
+        ids = torch.tensor([tokens], dtype=torch.long, device=device) # add batch dim
+        for _ in range(max_tokens):
+            logits = self.forward(ids) # (B, T, vocab_size)
+            logits = logits[:, -1, :] # (B, vocab_size)
+            if top_k is not None and top_k > 0:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            if temperature > 0:
+                logits = logits / temperature
+                probs = F.softmax(logits, dim=-1)
+                next_ids = torch.multinomial(probs, num_samples=1, generator=rng)
+            else:
+                next_ids = torch.argmax(logits, dim=-1, keepdim=True)
+            ids = torch.cat((ids, next_ids), dim=1)
+            token = next_ids.item()
+            yield token

nanochat/loss_eval.py

@@ -0,0 +1,65 @@
+"""
+A number of functions that help with evaluating a base model.
+"""
+import math
+import torch
+import torch.distributed as dist
+
+@torch.no_grad()
+def evaluate_bpb(model, batches, steps, token_bytes):
+    """
+    Instead of the naive 'mean loss', this function returns the bits per byte (bpb),
+    which is a tokenization vocab size-independent metric, meaning you are still comparing
+    apples:apples if you change the vocab size. The way this works is that instead of just
+    calculating the average loss as usual, you calculate the sum loss, and independently
+    also the sum bytes (of all the target tokens), and divide. This normalizes the loss by
+    the number of bytes that the target tokens represent.
+
+    The added complexity is so that:
+    1) All "normal" tokens are normalized by the length of the token in bytes
+    2) No special tokens (e.g. <|bos|>) are included in the metric - they are masked out.
+    3) No actively masked tokens (using ignore_index of e.g. -1) are included in the metric.
+
+    In addition to evaluate_loss, we need the token_bytes tensor:
+    It is a 1D tensor of shape (vocab_size,), indicating the number of bytes for
+    each token id, or 0 if the token is to not be counted (e.g. special tokens).
+    """
+    # record the losses
+    total_nats = torch.tensor(0.0, dtype=torch.float32, device=model.get_device())
+    total_bytes = torch.tensor(0, dtype=torch.int64, device=model.get_device())
+    batch_iter = iter(batches)
+    for _ in range(steps):
+        x, y = next(batch_iter)
+        loss2d = model(x, y, loss_reduction='none') # (B, T)
+        loss2d = loss2d.view(-1) # flatten
+        y = y.view(-1) # flatten
+        if (y.int() < 0).any(): # mps does not currently have kernel for < 0 for int64, only int32
+            # slightly more complex code path if some target tokens are ignore_index (e.g. -1)
+            # any target token < 0 is to be ignored: do NOT index token_bytes with negatives
+            valid = y >= 0
+            y_safe = torch.where(valid, y, torch.zeros_like(y))
+            # map valid targets to their byte length; ignored targets contribute 0 bytes
+            num_bytes2d = torch.where(
+                valid,
+                token_bytes[y_safe],
+                torch.zeros_like(y, dtype=token_bytes.dtype)
+            )
+            total_nats += (loss2d * (num_bytes2d > 0)).sum()
+            total_bytes += num_bytes2d.sum()
+        else:
+            # fast path: no ignored targets, safe to index directly
+            num_bytes2d = token_bytes[y]
+            total_nats += (loss2d * (num_bytes2d > 0)).sum()
+            total_bytes += num_bytes2d.sum()
+    # sum reduce across all ranks
+    world_size = dist.get_world_size() if dist.is_initialized() else 1
+    if world_size > 1:
+        dist.all_reduce(total_nats, op=dist.ReduceOp.SUM)
+        dist.all_reduce(total_bytes, op=dist.ReduceOp.SUM)
+    # move both to cpu, calculate bpb and return
+    total_nats = total_nats.item()
+    total_bytes = total_bytes.item()
+    if total_bytes == 0:
+        return float('inf')
+    bpb = total_nats / (math.log(2) * total_bytes)
+    return bpb

nanochat/optim.py

@@ -0,0 +1,471 @@
+"""
+A nice and efficient mixed AdamW/Muon Combined Optimizer.
+Usually the embeddings and scalars go into AdamW, and the matrix parameters go into Muon.
+The same class handles both single GPU and distributed training: when there is no
+multi-rank process group, the communication ops are simply skipped and every rank
+(i.e. the only rank) owns all of the parameters.
+
+Adapted from: https://github.com/KellerJordan/modded-nanogpt
+Further contributions from @karpathy and @chrisjmccormick.
+"""
+
+import torch
+import torch.distributed as dist
+from torch import Tensor
+from nanochat.common import COMPUTE_DTYPE
+
+# -----------------------------------------------------------------------------
+"""
+Good old AdamW optimizer, fused kernel.
+https://arxiv.org/abs/1711.05101
+"""
+
+@torch.compile(dynamic=False, fullgraph=True)
+def adamw_step_fused(
+    p: Tensor,              # (32768, 768) - parameter tensor
+    grad: Tensor,           # (32768, 768) - gradient, same shape as p
+    exp_avg: Tensor,        # (32768, 768) - first moment, same shape as p
+    exp_avg_sq: Tensor,     # (32768, 768) - second moment, same shape as p
+    step_t: Tensor,         # () - 0-D CPU tensor, step count
+    lr_t: Tensor,           # () - 0-D CPU tensor, learning rate
+    beta1_t: Tensor,        # () - 0-D CPU tensor, beta1
+    beta2_t: Tensor,        # () - 0-D CPU tensor, beta2
+    eps_t: Tensor,          # () - 0-D CPU tensor, epsilon
+    wd_t: Tensor,           # () - 0-D CPU tensor, weight decay
+) -> None:
+    """
+    Fused AdamW step: weight_decay -> momentum_update -> bias_correction -> param_update
+    All in one compiled graph to eliminate Python overhead between ops.
+    The 0-D CPU tensors avoid recompilation when hyperparameter values change.
+    """
+    # Some params (wte, value_embeds) are stored in bf16, so do the math in fp32 and
+    # cast back at the end. MPS errors on mixed-dtype ops (CUDA promotes them), and
+    # scalar arithmetic like 1 - beta2 loses all precision in bf16. compile fuses the casts.
+    p32 = p.float()
+    exp_avg32 = exp_avg.float()
+    exp_avg_sq32 = exp_avg_sq.float()
+    grad32 = grad.float()
+    # Weight decay (decoupled, applied before the update)
+    p32.mul_(1 - lr_t * wd_t)
+    # Update running averages (lerp_ is cleaner and fuses well)
+    exp_avg32.lerp_(grad32, 1 - beta1_t)
+    exp_avg_sq32.lerp_(grad32.square(), 1 - beta2_t)
+    # Bias corrections
+    bias1 = 1 - beta1_t ** step_t
+    bias2 = 1 - beta2_t ** step_t
+    # Compute update and apply
+    denom = (exp_avg_sq32 / bias2).sqrt() + eps_t
+    step_size = lr_t / bias1
+    p32.add_(exp_avg32 / denom, alpha=-step_size)
+    # Write back (no-ops in the common case where everything is already fp32)
+    p.copy_(p32)
+    exp_avg.copy_(exp_avg32)
+    exp_avg_sq.copy_(exp_avg_sq32)
+
+# -----------------------------------------------------------------------------
+"""
+Muon optimizer adapted and simplified from modded-nanogpt.
+https://github.com/KellerJordan/modded-nanogpt
+
+Background:
+Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
+quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
+of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
+zero even beyond the point where the iteration no longer converges all the way to one everywhere
+on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
+where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
+performance at all relative to UV^T, where USV^T = G is the SVD.
+
+Here, an alternative to Newton-Schulz iteration with potentially better convergence properties:
+Polar Express Sign Method for orthogonalization.
+https://arxiv.org/pdf/2505.16932
+by Noah Amsel, David Persson, Christopher Musco, Robert M. Gower.
+
+NorMuon variance reduction: per-neuron/column adaptive learning rate that normalizes
+update scales after orthogonalization (Muon's output has non-uniform scales across neurons).
+https://arxiv.org/pdf/2510.05491
+
+Two more (very) slight and optional improvements:
+1) MuonEq row equilibration: rescale each row to the mean row norm so the spectrum
+entering orthogonalization is better conditioned (https://arxiv.org/abs/2603.28254)
+2) Muon+ renormalization: snap the Frobenius norm to sqrt(min(m, n)), the norm of an exactly
+semi-orthogonal matrix, correcting for under-convergence of the polar iteration (https://arxiv.org/abs/2602.21545)
+
+Some of the changes in nanochat implementation:
+- Uses a simpler, more general approach to parameter grouping and stacking
+- Uses a single fused kernel for the momentum -> polar_express -> variance_reduction -> update step
+- Makes no assumptions about model architecture (e.g. that attention weights are fused into QKVO format)
+"""
+
+# Coefficients for Polar Express (computed for num_iters=5, safety_factor=2e-2, cushion=2)
+# From https://arxiv.org/pdf/2505.16932
+polar_express_coeffs = [
+    (8.156554524902461, -22.48329292557795, 15.878769915207462),
+    (4.042929935166739, -2.808917465908714, 0.5000178451051316),
+    (3.8916678022926607, -2.772484153217685, 0.5060648178503393),
+    (3.285753657755655, -2.3681294933425376, 0.46449024233003106),
+    (2.3465413258596377, -1.7097828382687081, 0.42323551169305323),
+]
+
+
+@torch.compile(dynamic=False, fullgraph=True)
+def muon_step_fused(
+    stacked_grads: Tensor,          # (12, 768, 3072) - stacked gradients
+    stacked_params: Tensor,         # (12, 768, 3072) - stacked parameters
+    momentum_buffer: Tensor,        # (12, 768, 3072) - first moment buffer
+    second_momentum_buffer: Tensor, # (12, 768, 1) or (12, 1, 3072) - factored second moment
+    momentum_t: Tensor,             # () - 0-D CPU tensor, momentum coefficient
+    lr_t: Tensor,                   # () - 0-D CPU tensor, learning rate
+    wd_t: Tensor,                   # () - 0-D CPU tensor, weight decay
+    beta2_t: Tensor,                # () - 0-D CPU tensor, beta2 for second moment
+    ns_steps: int,                  # 5 - number of Newton-Schulz/Polar Express iterations
+    red_dim: int,                   # -1 or -2 - reduction dimension for variance
+) -> None:
+    """
+    Fused Muon step: momentum -> polar_express -> variance_reduction -> cautious_update
+    All in one compiled graph to eliminate Python overhead between ops.
+    Some of the constants are 0-D CPU tensors to avoid recompilation when values change.
+    """
+
+    # Nesterov momentum
+    momentum = momentum_t.to(stacked_grads.dtype)
+    momentum_buffer.lerp_(stacked_grads, 1 - momentum)
+    g = stacked_grads.lerp_(momentum_buffer, momentum)
+
+    # Cast to bf16 for speed when available; skip cast otherwise (fp16 is unstable here due to limited exponent range)
+    X = g.bfloat16() if COMPUTE_DTYPE == torch.bfloat16 else g
+
+    # MuonEq row equilibration: rescale each row to the mean row norm so the spectrum entering orthogonalization is better conditioned
+    target = X.float().norm(dim=(-2, -1), keepdim=True) / (X.size(-2) ** 0.5)
+    row_norm = X.float().norm(dim=-1, keepdim=True).clamp_min(1e-6)
+    X = X * (target / row_norm).to(X.dtype)
+
+    # Polar Express orthogonalization: replace each update with the nearest orthogonal matrix
+    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.01 + 1e-6)
+    if g.size(-2) > g.size(-1): # Tall matrix
+        for a, b, c in polar_express_coeffs[:ns_steps]:
+            A = X.mT @ X
+            B = b * A + c * (A @ A)
+            X = a * X + X @ B
+    else: # Wide matrix (original math)
+        for a, b, c in polar_express_coeffs[:ns_steps]:
+            A = X @ X.mT
+            B = b * A + c * (A @ A)
+            X = a * X + B @ X
+    # Cast back to the param dtype (MPS errors on the mixed-dtype ops below when X is bf16)
+    g = X.to(stacked_params.dtype)
+
+    # Muon+ renormalization: snap Frobenius norm to sqrt(min(m, n))
+    target_norm = min(g.size(-2), g.size(-1)) ** 0.5
+    current_norm = g.float().norm(dim=(-2, -1), keepdim=True).clamp_min(1e-6)
+    g = g * (target_norm / current_norm).to(g.dtype)
+
+    # Variance reduction
+    beta2 = beta2_t.to(g.dtype)
+    v_mean = g.float().square().mean(dim=red_dim, keepdim=True)
+    red_dim_size = g.size(red_dim)
+    v_norm_sq = v_mean.sum(dim=(-2, -1), keepdim=True) * red_dim_size
+    v_norm = v_norm_sq.sqrt()
+    second_momentum_buffer.lerp_(v_mean.to(dtype=second_momentum_buffer.dtype), 1 - beta2)
+    step_size = second_momentum_buffer.clamp_min(1e-10).rsqrt()
+    scaled_sq_sum = (v_mean * red_dim_size) * step_size.float().square()
+    v_norm_new = scaled_sq_sum.sum(dim=(-2, -1), keepdim=True).sqrt()
+    final_scale = step_size * (v_norm / v_norm_new.clamp_min(1e-10))
+    g = g * final_scale.to(g.dtype)
+
+    # Cautious weight decay + parameter update
+    lr = lr_t.to(g.dtype)
+    wd = wd_t.to(g.dtype)
+    mask = (g * stacked_params) >= 0
+    if torch.compiler.is_compiling():
+        # Inductor fuses this expression into a single kernel.
+        stacked_params.sub_(lr * g + lr * wd * stacked_params * mask)
+    else:
+        # Without torch.compile the expression above materializes a full-size
+        # temporary update tensor. On large models that can cost multiple GiB. Apply
+        # decay before the gradient update so both operations are allocation-free
+        # while preserving p - lr*g - lr*wd*p*mask.
+        stacked_params.addcmul_(stacked_params, mask, value=-(lr * wd))
+        stacked_params.add_(g, alpha=-lr)
+
+# -----------------------------------------------------------------------------
+
+class MuonAdamW(torch.optim.Optimizer):
+    """
+    Combined optimizer: Muon for 2D matrix params, AdamW for others.
+
+    AdamW - Fused AdamW optimizer step.
+
+    Muon - MomentUm Orthogonalized by Newton-schulz
+    https://kellerjordan.github.io/posts/muon/
+
+    Muon internally runs standard SGD-momentum, and then performs an orthogonalization post-
+    processing step, in which each 2D parameter's update is replaced with the nearest orthogonal
+    matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has
+    the advantage that it can be stably run in bfloat16 on the GPU.
+
+    Some warnings:
+    - The Muon optimizer should not be used for the embedding layer, the final fully connected layer,
+    or any {0,1}-D parameters; those should all be optimized by a standard method (e.g., AdamW).
+    - To use it with 4D convolutional filters, it works well to just flatten their last 3 dimensions.
+
+    The same class covers single GPU and distributed training. In the distributed setting
+    (a multi-rank process group is initialized), gradients are synchronized here in the
+    optimizer (nanochat does not use DDP) and optimizer states are sharded across ranks
+    (ZeRO-2 style). On a single rank, all communication is skipped and the rank owns all
+    parameters, so the sharded code paths degenerate to plain full-tensor updates.
+
+    Design Goals:
+    - Overlap communication with computation (async ops)
+    - Minimize memory by sharding optimizer states across ranks (ZeRO-2 style)
+    - Batch small tensors into single comm ops where possible
+
+    Communication Pattern (3-phase async):
+    We use a 3-phase structure to maximize overlap between communication and compute:
+
+        Phase 1: Launch all async reduce ops
+            - Kick off all reduce_scatter/all_reduce operations
+            - Don't wait - let them run in background while we continue
+
+        Phase 2: Wait for reduces, compute updates, launch gathers
+            - For each group: wait for its reduce, compute the update, launch gather
+            - By processing groups in order, earlier gathers run while later computes happen
+
+        Phase 3: Wait for gathers, copy back
+            - Wait for all gathers to complete
+            - Copy updated params back to original tensors (Muon only)
+
+    AdamW Communication (ZeRO-2 style):
+    - Small params (<1024 elements): all_reduce gradients, update full param on each rank.
+      Optimizer state is replicated but these params are tiny (scalars, biases).
+    - Large params: reduce_scatter gradients so each rank gets 1/N of the grad, update
+      only that slice, then all_gather the updated slices. Optimizer state (exp_avg,
+      exp_avg_sq) is sharded - each rank only stores state for its slice.
+      Requires param.shape[0] divisible by world_size.
+
+    Muon Communication (stacked + chunked):
+    - All params in a Muon group must have the same shape (caller's responsibility).
+    - Stack all K params into a single (K, *shape) tensor for efficient comm.
+    - Divide K params across N ranks: each rank "owns" ceil(K/N) params.
+    - reduce_scatter the stacked grads so each rank gets its chunk.
+    - Each rank computes Muon update only for params it owns.
+    - all_gather the updated params back to all ranks.
+    - Optimizer state (momentum_buffer, second_momentum_buffer) is sharded by chunk.
+    - Padding: if K doesn't divide evenly, we zero-pad to (ceil(K/N) * N) for comm,
+      then ignore the padding when copying back.
+
+    Buffer Reuse:
+    - For Muon, we allocate stacked_grads for reduce_scatter input, then reuse the
+      same buffer as the output for all_gather (stacked_params). This saves memory
+      since we don't need both buffers simultaneously.
+
+    Arguments:
+        param_groups: List of dicts, each containing:
+            - 'params': List of parameters
+            - 'kind': 'adamw' or 'muon'
+            - For AdamW groups: 'lr', 'betas', 'eps', 'weight_decay'
+            - For Muon groups: 'lr', 'momentum', 'ns_steps', 'beta2', 'weight_decay'
+    """
+    def __init__(self, param_groups: list[dict]):
+        super().__init__(param_groups, defaults={})
+        # 0-D CPU tensors to avoid torch.compile recompilation when values change
+        self._adamw_step_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._adamw_lr_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._adamw_beta1_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._adamw_beta2_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._adamw_eps_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._adamw_wd_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._muon_momentum_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._muon_lr_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._muon_wd_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+        self._muon_beta2_t = torch.tensor(0.0, dtype=torch.float32, device="cpu")
+
+    def _reduce_adamw(self, group: dict, world_size: int) -> dict:
+        """Launch async reduce ops for AdamW group. Returns info dict with per-param infos."""
+        param_infos = {}
+        for p in group['params']:
+            grad = p.grad
+            if world_size == 1:
+                # Single rank: no communication, update the full param in place
+                param_infos[p] = dict(future=None, grad_slice=grad, is_small=True)
+            elif p.numel() < 1024:
+                # Small params: all_reduce (no scatter/gather needed)
+                future = dist.all_reduce(grad, op=dist.ReduceOp.AVG, async_op=True).get_future()
+                param_infos[p] = dict(future=future, grad_slice=grad, is_small=True)
+            else:
+                # Large params: reduce_scatter
+                assert grad.shape[0] % world_size == 0, f"AdamW reduce_scatter requires shape[0] ({grad.shape[0]}) divisible by world_size ({world_size})"
+                rank_size = grad.shape[0] // world_size
+                grad_slice = torch.empty_like(grad[:rank_size])
+                future = dist.reduce_scatter_tensor(grad_slice, grad, op=dist.ReduceOp.AVG, async_op=True).get_future()
+                param_infos[p] = dict(future=future, grad_slice=grad_slice, is_small=False)
+        return dict(param_infos=param_infos)
+
+    def _reduce_muon(self, group: dict, world_size: int) -> dict:
+        """Launch async reduce op for Muon group. Returns info dict."""
+        params = group['params']
+        if world_size == 1:
+            # Single rank: this rank owns all params, the stacked grads are the "chunk"
+            grad_chunk = torch.stack([p.grad for p in params])
+            return dict(future=None, grad_chunk=grad_chunk, stacked_grads=None, chunk_size=len(params))
+        chunk_size = (len(params) + world_size - 1) // world_size
+        padded_num_params = chunk_size * world_size
+        p = params[0]
+        shape, device, dtype = p.shape, p.device, p.dtype
+
+        # Stack grads and zero-pad to padded_num_params
+        grad_stack = torch.stack([p.grad for p in params])
+        stacked_grads = torch.empty(padded_num_params, *shape, dtype=dtype, device=device)
+        stacked_grads[:len(params)].copy_(grad_stack)
+        if len(params) < padded_num_params:
+            stacked_grads[len(params):].zero_()
+
+        # Reduce_scatter to get this rank's chunk
+        grad_chunk = torch.empty(chunk_size, *shape, dtype=dtype, device=device)
+        future = dist.reduce_scatter_tensor(grad_chunk, stacked_grads, op=dist.ReduceOp.AVG, async_op=True).get_future()
+
+        return dict(future=future, grad_chunk=grad_chunk, stacked_grads=stacked_grads, chunk_size=chunk_size)
+
+    def _compute_adamw(self, group: dict, info: dict, gather_list: list, rank: int, world_size: int) -> None:
+        """Wait for reduce, compute AdamW updates, launch gathers for large params."""
+        param_infos = info['param_infos']
+        for p in group['params']:
+            pinfo = param_infos[p]
+            if pinfo['future'] is not None:
+                pinfo['future'].wait()
+            grad_slice = pinfo['grad_slice']
+            state = self.state[p]
+
+            # For small params, operate on full param; for large, operate on slice
+            if pinfo['is_small']:
+                p_slice = p
+            else:
+                rank_size = p.shape[0] // world_size
+                p_slice = p[rank * rank_size:(rank + 1) * rank_size]
+
+            # State init
+            if not state:
+                state['step'] = 0
+                state['exp_avg'] = torch.zeros_like(p_slice)
+                state['exp_avg_sq'] = torch.zeros_like(p_slice)
+            state['step'] += 1
+
+            # Fill 0-D tensors and run fused kernel
+            self._adamw_step_t.fill_(state['step'])
+            self._adamw_lr_t.fill_(group['lr'])
+            self._adamw_beta1_t.fill_(group['betas'][0])
+            self._adamw_beta2_t.fill_(group['betas'][1])
+            self._adamw_eps_t.fill_(group['eps'])
+            self._adamw_wd_t.fill_(group['weight_decay'])
+            adamw_step_fused(
+                p_slice, grad_slice, state['exp_avg'], state['exp_avg_sq'],
+                self._adamw_step_t, self._adamw_lr_t, self._adamw_beta1_t,
+                self._adamw_beta2_t, self._adamw_eps_t, self._adamw_wd_t,
+            )
+
+            # Large params need all_gather
+            if not pinfo['is_small']:
+                future = dist.all_gather_into_tensor(p, p_slice, async_op=True).get_future()
+                gather_list.append(dict(future=future, params=None))
+
+    def _compute_muon(self, group: dict, info: dict, gather_list: list, rank: int) -> None:
+        """Wait for reduce, compute Muon updates, launch gather."""
+        if info['future'] is not None:
+            info['future'].wait()
+        params = group['params']
+        chunk_size = info['chunk_size']
+        grad_chunk = info['grad_chunk']
+        p = params[0]
+        shape, device, dtype = p.shape, p.device, p.dtype
+
+        # How many params does this rank own?
+        start_idx = rank * chunk_size
+        num_owned = min(chunk_size, max(0, len(params) - start_idx))
+
+        # Get or create group-level state
+        state = self.state[p]
+        if "momentum_buffer" not in state:
+            state["momentum_buffer"] = torch.zeros(chunk_size, *shape, dtype=dtype, device=device)
+        if "second_momentum_buffer" not in state:
+            state_shape = (chunk_size, shape[-2], 1) if shape[-2] >= shape[-1] else (chunk_size, 1, shape[-1])
+            state["second_momentum_buffer"] = torch.zeros(state_shape, dtype=dtype, device=device)
+        red_dim = -1 if shape[-2] >= shape[-1] else -2
+
+        stacked_owned = None
+        if num_owned > 0:
+            owned_params = [params[start_idx + i] for i in range(num_owned)]
+            stacked_owned = torch.stack(owned_params)
+
+            # Fill 0-D tensors and run fused kernel
+            self._muon_momentum_t.fill_(group["momentum"])
+            self._muon_beta2_t.fill_(group["beta2"])
+            self._muon_lr_t.fill_(group["lr"] * max(1.0, shape[-2] / shape[-1])**0.5)
+            self._muon_wd_t.fill_(group["weight_decay"])
+            muon_step_fused(
+                grad_chunk[:num_owned], stacked_owned,
+                state["momentum_buffer"][:num_owned], state["second_momentum_buffer"][:num_owned],
+                self._muon_momentum_t, self._muon_lr_t, self._muon_wd_t, self._muon_beta2_t,
+                group["ns_steps"], red_dim,
+            )
+
+        if info['stacked_grads'] is None:
+            # Single rank: copy back immediately so each temporary parameter
+            # stack can be released before the next Muon group is processed.
+            torch._foreach_copy_(params, list(stacked_owned.unbind(0)))
+            info["grad_chunk"] = None
+            return
+
+        # Build the input buffer for all_gather
+        updated_params = torch.empty(chunk_size, *shape, dtype=dtype, device=device)
+        if num_owned > 0:
+            updated_params[:num_owned].copy_(stacked_owned)
+        if num_owned < chunk_size:
+            updated_params[num_owned:].zero_()
+
+        # Reuse stacked_grads buffer for all_gather output
+        stacked_params = info["stacked_grads"]
+        future = dist.all_gather_into_tensor(stacked_params, updated_params, async_op=True).get_future()
+        gather_list.append(dict(future=future, stacked_params=stacked_params, params=params))
+        info["grad_chunk"] = None
+
+    def _finish_gathers(self, gather_list: list) -> None:
+        """Wait for all gathers and copy Muon params back."""
+        for info in gather_list:
+            if info["future"] is not None:
+                info["future"].wait()
+            if info["params"] is not None:
+                # Muon: copy from stacked buffer back to individual params
+                torch._foreach_copy_(info["params"], list(info["stacked_params"][:len(info["params"])].unbind(0)))
+
+    @torch.no_grad()
+    def step(self):
+        # On a single rank (no multi-rank process group), all communication is skipped
+        if dist.is_available() and dist.is_initialized():
+            rank = dist.get_rank()
+            world_size = dist.get_world_size()
+        else:
+            rank = 0
+            world_size = 1
+
+        # Phase 1: launch all async reduce ops
+        reduce_infos: list[dict] = []
+        for group in self.param_groups:
+            if group['kind'] == 'adamw':
+                reduce_infos.append(self._reduce_adamw(group, world_size))
+            elif group['kind'] == 'muon':
+                reduce_infos.append(self._reduce_muon(group, world_size))
+            else:
+                raise ValueError(f"Unknown optimizer kind: {group['kind']}")
+
+        # Phase 2: wait for reduces, compute updates, launch gathers
+        gather_list: list[dict] = []
+        for group, info in zip(self.param_groups, reduce_infos):
+            if group['kind'] == 'adamw':
+                self._compute_adamw(group, info, gather_list, rank, world_size)
+            elif group['kind'] == 'muon':
+                self._compute_muon(group, info, gather_list, rank)
+            else:
+                raise ValueError(f"Unknown optimizer kind: {group['kind']}")
+
+        # Phase 3: wait for gathers, copy back
+        self._finish_gathers(gather_list)

nanochat/tokenizer.py

@@ -0,0 +1,279 @@
+"""
+BPE Tokenizer in the style of GPT-4: train with rustbpe, inference with tiktoken.
+"""
+
+import os
+import copy
+from functools import lru_cache
+
+SPECIAL_TOKENS = [
+    # every document begins with the Beginning of Sequence (BOS) token that delimits documents
+    "<|bos|>",
+    # tokens below are only used during finetuning to render Conversations into token ids
+    "<|user_start|>", # user messages
+    "<|user_end|>",
+    "<|assistant_start|>", # assistant messages
+    "<|assistant_end|>",
+    "<|python_start|>", # assistant invokes python REPL tool
+    "<|python_end|>",
+    "<|output_start|>", # python REPL outputs back to assistant
+    "<|output_end|>",
+]
+
+# NOTE: this split pattern deviates from GPT-4 in that we use \p{N}{1,2} instead of \p{N}{1,3}
+# I did this because I didn't want to "waste" too many tokens on numbers for smaller vocab sizes.
+# I verified that 2 is the sweet spot for vocab size of 32K. 1 is a bit worse, 3 was worse still.
+SPLIT_PATTERN = r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,2}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"""
+
+# -----------------------------------------------------------------------------
+# Tokenizer based on rustbpe + tiktoken combo
+import pickle
+import rustbpe
+import tiktoken
+
+class RustBPETokenizer:
+    """Light wrapper around tiktoken (for efficient inference) but train with rustbpe"""
+
+    def __init__(self, enc, bos_token):
+        self.enc = enc
+        self.bos_token_id = self.encode_special(bos_token)
+
+    @classmethod
+    def train_from_iterator(cls, text_iterator, vocab_size):
+        # 1) train using rustbpe
+        tokenizer = rustbpe.Tokenizer()
+        # the special tokens are inserted later in __init__, we don't train them here
+        vocab_size_no_special = vocab_size - len(SPECIAL_TOKENS)
+        assert vocab_size_no_special >= 256, f"vocab_size_no_special must be at least 256, got {vocab_size_no_special}"
+        tokenizer.train_from_iterator(text_iterator, vocab_size_no_special, pattern=SPLIT_PATTERN)
+        # 2) construct the associated tiktoken encoding for inference
+        pattern = tokenizer.get_pattern()
+        mergeable_ranks_list = tokenizer.get_mergeable_ranks()
+        mergeable_ranks = {bytes(k): v for k, v in mergeable_ranks_list}
+        tokens_offset = len(mergeable_ranks)
+        special_tokens = {name: tokens_offset + i for i, name in enumerate(SPECIAL_TOKENS)}
+        enc = tiktoken.Encoding(
+            name="rustbpe",
+            pat_str=pattern,
+            mergeable_ranks=mergeable_ranks, # dict[bytes, int] (token bytes -> merge priority rank)
+            special_tokens=special_tokens, # dict[str, int] (special token name -> token id)
+        )
+        return cls(enc, "<|bos|>")
+
+    @classmethod
+    def from_directory(cls, tokenizer_dir):
+        pickle_path = os.path.join(tokenizer_dir, "tokenizer.pkl")
+        with open(pickle_path, "rb") as f:
+            enc = pickle.load(f)
+        return cls(enc, "<|bos|>")
+
+    @classmethod
+    def from_pretrained(cls, tiktoken_name):
+        # https://github.com/openai/tiktoken/blob/eedc8563/tiktoken_ext/openai_public.py
+        enc = tiktoken.get_encoding(tiktoken_name)
+        # tiktoken calls the special document delimiter token "<|endoftext|>"
+        # yes this is confusing because this token is almost always PREPENDED to the beginning of the document
+        # it most often is used to signal the start of a new sequence to the LLM during inference etc.
+        # so in nanoChat we always use "<|bos|>" short for "beginning of sequence", but historically it is often called "<|endoftext|>".
+        return cls(enc, "<|endoftext|>")
+
+    def get_vocab_size(self):
+        return self.enc.n_vocab
+
+    def get_special_tokens(self):
+        return self.enc.special_tokens_set
+
+    def id_to_token(self, id):
+        return self.enc.decode([id])
+
+    @lru_cache(maxsize=32)
+    def encode_special(self, text):
+        return self.enc.encode_single_token(text)
+
+    def get_bos_token_id(self):
+        return self.bos_token_id
+
+    def encode(self, text, prepend=None, append=None, num_threads=8):
+        # text can be either a string or a list of strings
+
+        if prepend is not None:
+            prepend_id = prepend if isinstance(prepend, int) else self.encode_special(prepend)
+        if append is not None:
+            append_id = append if isinstance(append, int) else self.encode_special(append)
+
+        if isinstance(text, str):
+            ids = self.enc.encode_ordinary(text)
+            if prepend is not None:
+                ids.insert(0, prepend_id) # TODO: slightly inefficient here? :( hmm
+            if append is not None:
+                ids.append(append_id)
+        elif isinstance(text, list):
+            ids = self.enc.encode_ordinary_batch(text, num_threads=num_threads)
+            if prepend is not None:
+                for ids_row in ids:
+                    ids_row.insert(0, prepend_id) # TODO: same
+            if append is not None:
+                for ids_row in ids:
+                    ids_row.append(append_id)
+        else:
+            raise ValueError(f"Invalid input type: {type(text)}")
+
+        return ids
+
+    def __call__(self, *args, **kwargs):
+        return self.encode(*args, **kwargs)
+
+    def decode(self, ids):
+        return self.enc.decode(ids)
+
+    def decode_single_token_bytes(self, token_id):
+        return self.enc.decode_single_token_bytes(token_id)
+
+    def save(self, tokenizer_dir):
+        # save the encoding object to disk
+        os.makedirs(tokenizer_dir, exist_ok=True)
+        pickle_path = os.path.join(tokenizer_dir, "tokenizer.pkl")
+        with open(pickle_path, "wb") as f:
+            pickle.dump(self.enc, f)
+        print(f"Saved tokenizer encoding to {pickle_path}")
+
+    def render_conversation(self, conversation, max_tokens=2048):
+        """
+        Tokenize a single Chat conversation (which we call a "doc" or "document" here).
+        Returns:
+        - ids: list[int] is a list of token ids of this rendered conversation
+        - mask: list[int] of same length, mask = 1 for tokens that the Assistant is expected to train on.
+        """
+        # ids, masks that we will return and a helper function to help build them up.
+        ids, mask = [], []
+        def add_tokens(token_ids, mask_val):
+            if isinstance(token_ids, int):
+                token_ids = [token_ids]
+            ids.extend(token_ids)
+            mask.extend([mask_val] * len(token_ids))
+
+        # sometimes the first message is a system message...
+        # => just merge it with the second (user) message
+        if conversation["messages"][0]["role"] == "system":
+            # some conversation surgery is necessary here for now...
+            conversation = copy.deepcopy(conversation) # avoid mutating the original
+            messages = conversation["messages"]
+            assert messages[1]["role"] == "user", "System message must be followed by a user message"
+            messages[1]["content"] = messages[0]["content"] + "\n\n" + messages[1]["content"]
+            messages = messages[1:]
+        else:
+            messages = conversation["messages"]
+        assert len(messages) >= 1, f"Conversation has less than 1 message: {messages}"
+
+        # fetch all the special tokens we need
+        bos = self.get_bos_token_id()
+        user_start, user_end = self.encode_special("<|user_start|>"), self.encode_special("<|user_end|>")
+        assistant_start, assistant_end = self.encode_special("<|assistant_start|>"), self.encode_special("<|assistant_end|>")
+        python_start, python_end = self.encode_special("<|python_start|>"), self.encode_special("<|python_end|>")
+        output_start, output_end = self.encode_special("<|output_start|>"), self.encode_special("<|output_end|>")
+
+        # now we can tokenize the conversation
+        add_tokens(bos, 0)
+        for i, message in enumerate(messages):
+
+            # some sanity checking here around assumptions, to prevent footguns
+            must_be_from = "user" if i % 2 == 0 else "assistant"
+            assert message["role"] == must_be_from, f"Message {i} is from {message['role']} but should be from {must_be_from}"
+
+            # content can be either a simple string or a list of parts (e.g. containing tool calls)
+            content = message["content"]
+
+            if message["role"] == "user":
+                assert isinstance(content, str), "User messages are simply expected to be strings"
+                value_ids = self.encode(content)
+                add_tokens(user_start, 0)
+                add_tokens(value_ids, 0)
+                add_tokens(user_end, 0)
+            elif message["role"] == "assistant":
+                add_tokens(assistant_start, 0)
+                if isinstance(content, str):
+                    # simple string => simply add the tokens
+                    value_ids = self.encode(content)
+                    add_tokens(value_ids, 1)
+                elif isinstance(content, list):
+                    for part in content:
+                        value_ids = self.encode(part["text"])
+                        if part["type"] == "text":
+                            # string part => simply add the tokens
+                            add_tokens(value_ids, 1)
+                        elif part["type"] == "python":
+                            # python tool call => add the tokens inside <|python_start|> and <|python_end|>
+                            add_tokens(python_start, 1)
+                            add_tokens(value_ids, 1)
+                            add_tokens(python_end, 1)
+                        elif part["type"] == "python_output":
+                            # python output => add the tokens inside <|output_start|> and <|output_end|>
+                            # none of these tokens are supervised because the tokens come from Python at test time
+                            add_tokens(output_start, 0)
+                            add_tokens(value_ids, 0)
+                            add_tokens(output_end, 0)
+                        else:
+                            raise ValueError(f"Unknown part type: {part['type']}")
+                else:
+                    raise ValueError(f"Unknown content type: {type(content)}")
+                add_tokens(assistant_end, 1)
+
+        # truncate to max_tokens tokens MAX (helps prevent OOMs)
+        ids = ids[:max_tokens]
+        mask = mask[:max_tokens]
+        return ids, mask
+
+    def visualize_tokenization(self, ids, mask, with_token_id=False):
+        """Small helper function useful in debugging: visualize the tokenization of render_conversation"""
+        RED = '\033[91m'
+        GREEN = '\033[92m'
+        RESET = '\033[0m'
+        GRAY = '\033[90m'
+        tokens = []
+        for i, (token_id, mask_val) in enumerate(zip(ids, mask)):
+            token_str = self.decode([token_id])
+            color = GREEN if mask_val == 1 else RED
+            tokens.append(f"{color}{token_str}{RESET}")
+            if with_token_id:
+                tokens.append(f"{GRAY}({token_id}){RESET}")
+        return '|'.join(tokens)
+
+    def render_for_completion(self, conversation):
+        """
+        Used during Reinforcement Learning. In that setting, we want to
+        render the conversation priming the Assistant for a completion.
+        Unlike the Chat SFT case, we don't need to return the mask.
+        """
+        # We have some surgery to do: we need to pop the last message (of the Assistant)
+        conversation = copy.deepcopy(conversation) # avoid mutating the original
+        messages = conversation["messages"]
+        assert messages[-1]["role"] == "assistant", "Last message must be from the Assistant"
+        messages.pop() # remove the last message (of the Assistant) inplace
+
+        # Now tokenize the conversation
+        ids, mask = self.render_conversation(conversation)
+
+        # Finally, to prime the Assistant for a completion, append the Assistant start token
+        assistant_start = self.encode_special("<|assistant_start|>")
+        ids.append(assistant_start)
+        return ids
+
+# -----------------------------------------------------------------------------
+# nanochat-specific convenience functions
+
+def get_tokenizer():
+    from nanochat.common import get_base_dir
+    base_dir = get_base_dir()
+    tokenizer_dir = os.path.join(base_dir, "tokenizer")
+    return RustBPETokenizer.from_directory(tokenizer_dir)
+
+def get_token_bytes(device="cpu"):
+    import torch
+    from nanochat.common import get_base_dir
+    base_dir = get_base_dir()
+    tokenizer_dir = os.path.join(base_dir, "tokenizer")
+    token_bytes_path = os.path.join(tokenizer_dir, "token_bytes.pt")
+    assert os.path.exists(token_bytes_path), f"Token bytes not found at {token_bytes_path}? It gets written by tok_train.py"
+    with open(token_bytes_path, "rb") as f:
+        token_bytes = torch.load(f, map_location=device)
+    return token_bytes