train_coder_30b.py

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.11,<3.14"
# dependencies = [
#   "datasets>=4.4,<5",
#   "huggingface-hub==1.27.0",
#   "numpy>=2,<3",
#   "safetensors>=0.5,<1",
#   "tokenizers>=0.21,<1",
#   "torch==2.9.0",
#   "transformers>=5.0,<6",
#   "zstandard>=0.23,<1",
# ]
# ///
"""Train BananaMind 2.1 NanoCoder or MiniCoder on 30B streamed tokens."""

from __future__ import annotations

import argparse
import gc
import importlib
import json
import math
import os
import queue
import shutil
import socket
import sys
import threading
import time
import traceback
from itertools import chain
from pathlib import Path
from typing import Any, Iterator

import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.multiprocessing as mp
from datasets import load_dataset
from datasets.distributed import split_dataset_by_node
from huggingface_hub import HfApi, hf_hub_download
from safetensors.torch import save_file
from tokenizers import Tokenizer
from torch.nn.parallel import DistributedDataParallel as DDP


TOTAL_TOKENS = 30_000_000_000
TOKENIZER_VOCAB_SIZE = 8192
ARCHITECTURE_FILES = (
    "configuration_bananamind21_coder.py",
    "modeling_bananamind21_coder.py",
    "curriculum_coder_30b.py",
)
EXPECTED = {
    "nanocoder": {
        "repo_id": "Banaxi-Tech/BananaMind-2.1-NanoCoder",
        "parameters": 9_895_690,
        "transformer": 7_975_688,
        "ngram": 1_920_002,
    },
    "minicoder": {
        "repo_id": "Banaxi-Tech/BananaMind-2.1-MiniCoder",
        "parameters": 24_949_999,
        "transformer": 19_950_029,
        "ngram": 4_999_970,
    },
}
DATASET_IDS = {
    "stack_v3": "HuggingFaceCode/stack-v3-train",
    "fineweb_edu": "HuggingFaceFW/fineweb-edu",
}


def retry(action, description: str, attempts: int = 6):
    for attempt in range(1, attempts + 1):
        try:
            return action()
        except Exception:
            if attempt == attempts:
                raise
            delay = min(60, 2**attempt)
            print(
                f"{description} failed ({attempt}/{attempts}); "
                f"retrying in {delay}s",
                flush=True,
            )
            time.sleep(delay)


def prepare_runtime_assets(args: argparse.Namespace) -> tuple[str, str, str, dict]:
    local_source = Path(__file__).resolve().parent
    token = os.environ.get("HF_TOKEN")
    api = HfApi(token=token)
    revision = retry(
        lambda: api.model_info(args.repo_id).sha,
        "resolve model repository revision",
    )
    destination = Path(args.output_dir) / "runtime"
    destination.mkdir(parents=True, exist_ok=True)

    for filename in (*ARCHITECTURE_FILES, "README.md"):
        local_file = local_source / filename
        if filename != "README.md" and local_file.is_file():
            shutil.copy2(local_file, destination / filename)
            continue
        try:
            retry(
                lambda filename=filename: hf_hub_download(
                    repo_id=args.repo_id,
                    filename=filename,
                    revision=revision,
                    token=token,
                    local_dir=destination,
                ),
                f"download {filename}",
            )
        except Exception:
            if filename != "README.md":
                raise

    tokenizer_dir = destination / "tokenizer"
    tokenizer_dir.mkdir(exist_ok=True)
    for filename in (
        "tokenizer.json",
        "tokenizer_config.json",
        "special_tokens_map.json",
    ):
        try:
            retry(
                lambda filename=filename: hf_hub_download(
                    repo_id=args.repo_id,
                    filename=filename,
                    revision=revision,
                    token=token,
                    local_dir=tokenizer_dir,
                ),
                f"download {filename}",
            )
        except Exception:
            if filename != "special_tokens_map.json":
                raise
    tokenizer = Tokenizer.from_file(str(tokenizer_dir / "tokenizer.json"))
    if tokenizer.get_vocab_size() != TOKENIZER_VOCAB_SIZE:
        raise RuntimeError(
            f"Expected {TOKENIZER_VOCAB_SIZE} tokenizer entries, "
            f"found {tokenizer.get_vocab_size()}"
        )
    if tokenizer.token_to_id("<|eos|>") != 2:
        raise RuntimeError("Nano tokenizer must use EOS token ID 2")

    revisions = {}
    for key, dataset_id in DATASET_IDS.items():
        revisions[key] = retry(
            lambda dataset_id=dataset_id: api.dataset_info(dataset_id).sha,
            f"resolve {key} revision",
        )
    return str(destination), revision, str(tokenizer_dir), revisions


def normalize_files(value: Any) -> list[dict[str, Any]]:
    if isinstance(value, list):
        return [item for item in value if isinstance(item, dict)]
    if isinstance(value, dict):
        lengths = [len(column) for column in value.values() if isinstance(column, list)]
        if not lengths:
            return []
        result = []
        for index in range(min(lengths)):
            result.append(
                {
                    key: column[index] if isinstance(column, list) else column
                    for key, column in value.items()
                }
            )
        return result
    return []


def repository_documents(row: dict[str, Any]) -> Iterator[str]:
    repo_path = str(row.get("repo_path") or "unknown/repository")
    for file in normalize_files(row.get("files")):
        if file.get("is_vendor"):
            continue
        content = file.get("content")
        if not isinstance(content, str) or not content.strip():
            continue
        path = str(file.get("file_path") or "unknown")
        language = str(file.get("language") or "Unknown")
        header = f"Repository: {repo_path}\nFile: {path}\nLanguage: {language}\n"
        # Avoid handing a tokenizer one unbounded generated or data file. Each
        # chunk remains adjacent and explicitly carries its repository path.
        for start in range(0, len(content), 200_000):
            chunk = content[start : start + 200_000]
            if chunk.strip():
                yield header + chunk


class StreamedSourceBatcher:
    def __init__(
        self,
        source_key: str,
        revision: str,
        tokenizer_path: Path,
        rank: int,
        world_size: int,
        local_batch: int,
        sequence_length: int,
        encode_batch_size: int,
        prefetch_batches: int,
        shuffle_buffer: int,
        seed: int,
    ):
        from curriculum_coder_30b import SOURCES

        self.source_key = source_key
        self.source = SOURCES[source_key]
        self.revision = revision
        self.tokenizer_path = tokenizer_path
        self.rank = rank
        self.world_size = world_size
        self.local_batch = local_batch
        self.sequence_length = sequence_length
        self.encode_batch_size = encode_batch_size
        self.prefetch_batches = prefetch_batches
        self.shuffle_buffer = shuffle_buffer
        self.seed = seed
        self.queue: queue.Queue[tuple[str, Any]] = queue.Queue(prefetch_batches)
        self.stop_event = threading.Event()
        self.thread: threading.Thread | None = None
        self.restart_count = 0

    def start(self) -> None:
        if self.thread is None:
            self.thread = threading.Thread(target=self._produce, daemon=True)
            self.thread.start()

    def _put(self, item: tuple[str, Any]) -> bool:
        while not self.stop_event.is_set():
            try:
                self.queue.put(item, timeout=1.0)
                return True
            except queue.Full:
                continue
        return False

    def _produce(self) -> None:
        try:
            tokenizer = Tokenizer.from_file(str(self.tokenizer_path))
            eos_id = tokenizer.token_to_id("<|eos|>")
            if eos_id != 2:
                raise ValueError(f"Expected EOS token 2, found {eos_id}")
            batch_tokens = self.local_batch * self.sequence_length
            required_tokens = batch_tokens + 1
            pending = np.empty(0, dtype=np.int64)
            texts: list[str] = []

            def encode_texts() -> None:
                nonlocal pending, texts
                if not texts:
                    return
                encodings = tokenizer.encode_batch(texts)
                values = chain.from_iterable(
                    chain(encoding.ids, (eos_id,))
                    for encoding in encodings
                    if encoding.ids
                )
                encoded = np.fromiter(values, dtype=np.int64)
                if encoded.size:
                    pending = (
                        encoded
                        if pending.size == 0
                        else np.concatenate((pending, encoded))
                    )
                texts = []

            def emit_ready_batches() -> bool:
                nonlocal pending
                while pending.size >= required_tokens:
                    packed = pending[:required_tokens].copy()
                    pending = pending[batch_tokens:]
                    inputs = torch.from_numpy(
                        packed[:-1].reshape(self.local_batch, self.sequence_length)
                    ).pin_memory()
                    labels = torch.from_numpy(
                        packed[1:].reshape(self.local_batch, self.sequence_length)
                    ).pin_memory()
                    if not self._put(("batch", (inputs, labels))):
                        return False
                return True

            epoch = self.restart_count * 10_000
            while not self.stop_event.is_set():
                kwargs: dict[str, Any] = {
                    "path": self.source["dataset_id"],
                    "split": "train",
                    "streaming": True,
                    "revision": self.revision,
                    "token": os.environ.get("HF_TOKEN"),
                }
                if self.source["config_name"]:
                    kwargs["name"] = self.source["config_name"]
                dataset = load_dataset(**kwargs)
                if self.source["kind"] == "repository":
                    dataset = dataset.select_columns(["repo_path", "files"])
                    buffer_size = min(self.shuffle_buffer, 512)
                else:
                    dataset = dataset.select_columns(["text"])
                    buffer_size = self.shuffle_buffer
                dataset = dataset.shuffle(
                    seed=self.seed + epoch * 1_000_003,
                    buffer_size=buffer_size,
                )
                dataset = split_dataset_by_node(
                    dataset,
                    rank=self.rank,
                    world_size=self.world_size,
                )
                rows_seen = 0
                for row in dataset:
                    if self.stop_event.is_set():
                        return
                    rows_seen += 1
                    documents = (
                        repository_documents(row)
                        if self.source["kind"] == "repository"
                        else iter((row.get("text"),))
                    )
                    for document in documents:
                        if not isinstance(document, str) or not document.strip():
                            continue
                        texts.append(document)
                        if len(texts) >= self.encode_batch_size:
                            encode_texts()
                            if not emit_ready_batches():
                                return
                encode_texts()
                if not emit_ready_batches():
                    return
                if rows_seen == 0:
                    raise RuntimeError(f"{self.source['label']} yielded no rows")
                epoch += 1
                if self.rank == 0:
                    print(
                        f"{self.source['label']} stream exhausted; restarting",
                        flush=True,
                    )
        except BaseException:
            self._put(("error", traceback.format_exc()))

    def next_batch(self) -> tuple[torch.Tensor, torch.Tensor]:
        for attempt in range(1, 7):
            self.start()
            kind, payload = self.queue.get()
            if kind == "batch":
                return payload
            if self.thread is not None:
                self.thread.join(timeout=1.0)
            self.thread = None
            self.restart_count += 1
            if attempt == 6:
                raise RuntimeError(
                    f"{self.source['label']} failed on rank {self.rank}:\n{payload}"
                )
            delay = min(30, 2**attempt)
            print(
                f"{self.source['label']} stream failed on rank {self.rank} "
                f"({attempt}/6); retrying in {delay}s",
                flush=True,
            )
            time.sleep(delay)
        raise AssertionError("unreachable")

    def close(self) -> None:
        self.stop_event.set()
        if self.thread is not None:
            self.thread.join(timeout=10.0)


def unwrap_model(model: nn.Module) -> nn.Module:
    current = model
    while True:
        candidate = getattr(current, "module", None)
        if candidate is None:
            candidate = getattr(current, "_orig_mod", None)
        if candidate is None or candidate is current:
            return current
        current = candidate


def canonical_state_dict(model: nn.Module) -> dict[str, torch.Tensor]:
    return {
        name: tensor.detach().float().cpu().contiguous().clone()
        for name, tensor in unwrap_model(model).state_dict().items()
    }


def tree_to_cpu(value: Any) -> Any:
    if isinstance(value, torch.Tensor):
        return value.detach().cpu()
    if isinstance(value, dict):
        return {key: tree_to_cpu(item) for key, item in value.items()}
    if isinstance(value, list):
        return [tree_to_cpu(item) for item in value]
    if isinstance(value, tuple):
        return tuple(tree_to_cpu(item) for item in value)
    return value


def split_optimizer_parameters(model: nn.Module):
    token_embedding_id = id(model.transformer["wte"].weight)
    ngram_ids = {id(parameter) for parameter in model.transformer["ngram"].parameters()}
    groups = {"muon": [], "embeddings": [], "ngram": [], "controls": []}
    names = {key: [] for key in groups}
    for name, parameter in model.named_parameters():
        if not parameter.requires_grad:
            continue
        if id(parameter) in ngram_ids:
            group = "ngram"
        elif id(parameter) == token_embedding_id:
            group = "embeddings"
        elif parameter.ndim == 2:
            group = "muon"
        elif parameter.ndim <= 1:
            group = "controls"
        else:
            raise ValueError(f"No optimizer group for {name}: {parameter.shape}")
        groups[group].append(parameter)
        names[group].append(name)
    assigned = [id(parameter) for values in groups.values() for parameter in values]
    expected = {
        id(parameter) for parameter in model.parameters() if parameter.requires_grad
    }
    if len(assigned) != len(set(assigned)) or set(assigned) != expected:
        raise AssertionError("Optimizer groups overlap or omit parameters")
    expected_ngram = {
        "transformer.ngram.injection_scales",
        "transformer.ngram.bigram_table.weight",
        "transformer.ngram.fourgram_table.weight",
        "transformer.ngram.out_proj.weight",
    }
    if set(names["ngram"]) != expected_ngram:
        raise AssertionError("The complete n-gram module needs its separate LR")
    return groups, names


def scheduled_lr(
    step: int,
    total_steps: int,
    peak: float,
    warmup_steps: int,
    decay_ratio: float,
) -> float:
    if step < warmup_steps:
        return peak * (step + 1) / max(1, warmup_steps)
    decay_steps = max(1, int(total_steps * decay_ratio))
    decay_start = max(warmup_steps, total_steps - decay_steps)
    if step < decay_start:
        return peak
    progress = (step - decay_start) / max(1, total_steps - decay_start - 1)
    return peak * 0.5 * (1.0 + math.cos(math.pi * min(1.0, progress)))


def export_checkpoint(
    model: nn.Module,
    config,
    source_dir: Path,
    tokenizer_dir: Path,
    args: argparse.Namespace,
    metadata: dict[str, Any],
    metrics_path: Path,
    muon_optimizer: torch.optim.Optimizer,
    adamw_optimizer: torch.optim.Optimizer,
    source_scheduler,
) -> None:
    export_dir = Path(args.output_dir) / "hf-export"
    if export_dir.exists():
        shutil.rmtree(export_dir)
    export_dir.mkdir(parents=True)
    for filename in (*ARCHITECTURE_FILES, "README.md"):
        source = source_dir / filename
        if source.is_file():
            shutil.copy2(source, export_dir / filename)
    for filename in (
        "tokenizer.json",
        "tokenizer_config.json",
        "special_tokens_map.json",
    ):
        source = tokenizer_dir / filename
        if source.is_file():
            shutil.copy2(source, export_dir / filename)
    tokenizer_config_path = export_dir / "tokenizer_config.json"
    tokenizer_config = json.loads(tokenizer_config_path.read_text())
    tokenizer_config["model_max_length"] = config.max_position_embeddings
    tokenizer_config_path.write_text(json.dumps(tokenizer_config, indent=2) + "\n")

    state = canonical_state_dict(model)
    if not torch.equal(state["transformer.wte.weight"], state["lm_head.weight"]):
        raise RuntimeError("Tied input/output embeddings diverged")
    save_file(state, export_dir / "model.safetensors", metadata={"format": "pt"})
    config_json = config.to_dict()
    config_json.update(
        {
            "architectures": ["BananaMind21CoderForCausalLM"],
            "auto_map": {
                "AutoConfig": (
                    "configuration_bananamind21_coder.BananaMind21CoderConfig"
                ),
                "AutoModelForCausalLM": (
                    "modeling_bananamind21_coder."
                    "BananaMind21CoderForCausalLM"
                ),
            },
            "torch_dtype": "float32",
            "_name_or_path": args.repo_id,
        }
    )
    (export_dir / "config.json").write_text(
        json.dumps(config_json, indent=2) + "\n"
    )
    (export_dir / "generation_config.json").write_text(
        json.dumps(
            {
                "_from_model_config": True,
                "bos_token_id": config.bos_token_id,
                "eos_token_id": config.eos_token_id,
                "pad_token_id": config.pad_token_id,
                "transformers_version": "5",
            },
            indent=2,
        )
        + "\n"
    )
    (export_dir / "checkpoint_metadata.json").write_text(
        json.dumps(metadata, indent=2) + "\n"
    )
    if metrics_path.is_file():
        shutil.copy2(metrics_path, export_dir / "training_metrics.jsonl")
    training_state = {
        "format_version": 1,
        "model_type": config.model_type,
        "variant": config.variant,
        "step": metadata["step"],
        "tokens_seen": metadata["tokens_seen"],
        "model": state,
        "muon_optimizer": tree_to_cpu(muon_optimizer.state_dict()),
        "adamw_optimizer": tree_to_cpu(adamw_optimizer.state_dict()),
        "source_scheduler": source_scheduler.state_dict(),
        "metadata": metadata,
    }
    torch.save(training_state, export_dir / "training_state.pt")
    del state, training_state
    gc.collect()

    api = HfApi(token=os.environ["HF_TOKEN"])
    result = retry(
        lambda: api.upload_folder(
            repo_id=args.repo_id,
            repo_type="model",
            folder_path=export_dir,
            commit_message=(
                f"Save {metadata['training_percent']}% checkpoint at "
                f"{metadata['tokens_seen']:,} tokens"
            ),
        ),
        "checkpoint upload",
    )
    tag = f"checkpoint-{metadata['training_percent']:03d}pct"
    try:
        retry(
            lambda: api.create_tag(
                repo_id=args.repo_id,
                repo_type="model",
                tag=tag,
                revision=result.oid,
                exist_ok=True,
            ),
            f"create {tag}",
        )
    except Exception as error:
        print(f"Could not create {tag}: {error}", flush=True)
    print(f"Uploaded {tag}: {result.commit_url}", flush=True)


def setup_distributed(rank: int, world_size: int) -> None:
    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
    os.environ.setdefault("MASTER_PORT", "29611")
    torch.cuda.set_device(rank)
    dist.init_process_group(
        backend="nccl",
        rank=rank,
        world_size=world_size,
        timeout=__import__("datetime").timedelta(minutes=30),
    )


def train_worker(
    rank: int,
    world_size: int,
    args: argparse.Namespace,
    source_dir_string: str,
    architecture_revision: str,
    tokenizer_dir_string: str,
    dataset_revisions: dict[str, str],
) -> None:
    setup_distributed(rank, world_size)
    device = torch.device("cuda", rank)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    torch.set_float32_matmul_precision("high")
    torch.backends.cuda.matmul.allow_tf32 = True

    source_dir = Path(source_dir_string)
    tokenizer_dir = Path(tokenizer_dir_string)
    sys.path.insert(0, source_dir_string)
    config_module = importlib.import_module("configuration_bananamind21_coder")
    model_module = importlib.import_module("modeling_bananamind21_coder")
    curriculum_module = importlib.import_module("curriculum_coder_30b")
    config = config_module.BananaMind21CoderConfig(variant=args.variant)
    model = model_module.BananaMind21CoderForCausalLM(config).to(device)
    breakdown = config.parameter_counts()
    parameter_count = sum(parameter.numel() for parameter in model.parameters())
    expected = EXPECTED[args.variant]
    if parameter_count != expected["parameters"] or breakdown["total"] != parameter_count:
        raise RuntimeError(
            f"Expected {expected['parameters']:,} parameters, found {parameter_count:,}"
        )
    if breakdown["transformer"] != expected["transformer"]:
        raise RuntimeError("Transformer parameter budget changed")
    if breakdown["ngram"] != expected["ngram"]:
        raise RuntimeError("N-gram parameter budget changed")

    groups, group_names = split_optimizer_parameters(model)
    muon_optimizer = torch.optim.Muon(
        groups["muon"],
        lr=args.muon_peak_lr,
        momentum=args.muon_momentum,
        nesterov=True,
        ns_steps=args.muon_ns_steps,
        weight_decay=args.weight_decay,
        adjust_lr_fn="original",
    )
    adamw_optimizer = torch.optim.AdamW(
        [
            {
                "name": "embeddings",
                "params": groups["embeddings"],
                "lr": args.adamw_peak_lr,
                "weight_decay": args.weight_decay,
            },
            {
                "name": "ngram",
                "params": groups["ngram"],
                "lr": args.ngram_peak_lr,
                "weight_decay": args.weight_decay,
            },
            {
                "name": "controls",
                "params": groups["controls"],
                "lr": args.adamw_peak_lr,
                "weight_decay": 0.0,
            },
        ],
        lr=args.adamw_peak_lr,
        betas=(0.9, 0.95),
        eps=1e-8,
        fused=True,
    )
    source_scheduler = curriculum_module.TokenCreditScheduler()
    start_step = 0
    tokens_seen = 0

    resume_path = None
    if args.resume and rank == 0:
        try:
            resume_path = hf_hub_download(
                repo_id=args.repo_id,
                filename="training_state.pt",
                token=os.environ.get("HF_TOKEN"),
                local_dir=Path(args.output_dir) / "resume",
            )
        except Exception as error:
            print(f"No resumable state found; starting fresh ({error})", flush=True)
    resume_box = [resume_path]
    dist.broadcast_object_list(resume_box, src=0)
    if resume_box[0]:
        state = torch.load(resume_box[0], map_location=device, weights_only=False)
        if state.get("variant") != args.variant:
            raise RuntimeError("Uploaded training state belongs to another variant")
        model.load_state_dict(state["model"], strict=True)
        muon_optimizer.load_state_dict(state["muon_optimizer"])
        adamw_optimizer.load_state_dict(state["adamw_optimizer"])
        source_scheduler.load_state_dict(state["source_scheduler"])
        start_step = int(state["step"])
        tokens_seen = int(state["tokens_seen"])
        del state
        gc.collect()

    if args.compile:
        model = torch.compile(model, dynamic=False)
    ddp = DDP(
        model,
        device_ids=[rank],
        output_device=rank,
        broadcast_buffers=False,
        gradient_as_bucket_view=True,
        static_graph=True,
    )
    local_batch = args.global_batch_sequences // world_size
    streams = {
        key: StreamedSourceBatcher(
            source_key=key,
            revision=dataset_revisions[key],
            tokenizer_path=tokenizer_dir / "tokenizer.json",
            rank=rank,
            world_size=world_size,
            local_batch=local_batch,
            sequence_length=args.seq_len,
            encode_batch_size=(
                args.stack_encode_batch_size
                if key == "stack_v3"
                else args.web_encode_batch_size
            ),
            prefetch_batches=args.prefetch_batches,
            shuffle_buffer=args.shuffle_buffer,
            seed=args.seed + start_step * 17 + index * 100_003,
        )
        for index, key in enumerate(curriculum_module.SOURCE_KEYS)
    }

    tokens_per_step = args.global_batch_sequences * args.seq_len
    total_steps = math.ceil(args.total_tokens / tokens_per_step)
    warmup_steps = max(1, math.ceil(args.warmup_tokens / tokens_per_step))
    checkpoint_steps = {
        max(1, math.ceil(total_steps * percent / 100)): percent
        for percent in range(5, 101, 5)
    }
    checkpoint_steps[total_steps] = 100
    metrics_path = Path(args.output_dir) / "training_metrics.jsonl"
    if rank == 0:
        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
        if start_step == 0:
            metrics_path.write_text("")
        print(f"BananaMind 2.1 {args.variant} code pretraining", flush=True)
        print(f"host:             {socket.gethostname()}", flush=True)
        print(f"hardware:         {world_size} x {torch.cuda.get_device_name(0)}", flush=True)
        print(f"parameters:       {parameter_count:,}", flush=True)
        print(f"transformer:      {breakdown['transformer']:,}", flush=True)
        print(f"n-gram:           {breakdown['ngram']:,}", flush=True)
        print(f"physical layers:  {breakdown['physical_layers']}", flush=True)
        print(f"effective passes: {breakdown['effective_layer_passes']}", flush=True)
        print(f"loop schedule:    {config.loop_schedule}", flush=True)
        print("data:             75% Stack v3 / 25% FineWeb-Edu", flush=True)
        print(f"context:          {args.seq_len:,}", flush=True)
        print(f"local batch:      {local_batch}", flush=True)
        print(f"global batch:     {args.global_batch_sequences}", flush=True)
        print(f"tokens/step:      {tokens_per_step:,}", flush=True)
        print(f"steps:            {total_steps:,}", flush=True)
        print(f"resume step:      {start_step:,}", flush=True)
        print(f"Muon tensors:     {len(group_names['muon'])}", flush=True)

    if start_step >= total_steps:
        if rank == 0:
            print("The uploaded checkpoint already completed training.", flush=True)
        for stream in streams.values():
            stream.close()
        dist.destroy_process_group()
        return

    dist.barrier()
    ddp.train()
    started = time.time()
    log_started = started
    log_tokens = 0
    try:
        for step_index in range(start_step, total_steps):
            step = step_index + 1
            step_tokens = min(tokens_per_step, args.total_tokens - tokens_seen)
            if step_tokens <= 0 or step_tokens % world_size:
                raise RuntimeError("Final supervised-token count must divide by GPUs")
            source_key = source_scheduler.choose(step_tokens)
            data_started = time.time()
            input_ids, shifted_labels = streams[source_key].next_batch()
            data_wait = time.time() - data_started
            local_supervised_tokens = step_tokens // world_size
            if local_supervised_tokens < shifted_labels.numel():
                shifted_labels.view(-1)[local_supervised_tokens:] = -100
            input_ids = input_ids.to(device, non_blocking=True)
            shifted_labels = shifted_labels.to(device, non_blocking=True)

            muon_lr = scheduled_lr(
                step_index,
                total_steps,
                args.muon_peak_lr,
                warmup_steps,
                args.decay_ratio,
            )
            adamw_lr = scheduled_lr(
                step_index,
                total_steps,
                args.adamw_peak_lr,
                warmup_steps,
                args.decay_ratio,
            )
            ngram_lr = scheduled_lr(
                step_index,
                total_steps,
                args.ngram_peak_lr,
                warmup_steps,
                args.decay_ratio,
            )
            weight_decay = (
                args.weight_decay
                if tokens_seen < args.weight_decay_switch_tokens
                else args.final_weight_decay
            )
            for group in muon_optimizer.param_groups:
                group["lr"] = muon_lr
                group["weight_decay"] = weight_decay
            for group in adamw_optimizer.param_groups:
                group["lr"] = ngram_lr if group["name"] == "ngram" else adamw_lr
                if group["name"] != "controls":
                    group["weight_decay"] = weight_decay
            muon_optimizer.zero_grad(set_to_none=True)
            adamw_optimizer.zero_grad(set_to_none=True)
            z_coefficient = (
                args.z_loss_coeff
                if tokens_seen < args.z_loss_until_tokens
                else 0.0
            )
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                loss, ce_loss, z_loss = ddp(
                    input_ids,
                    shifted_labels=shifted_labels,
                    return_training_losses=True,
                    z_loss_coefficient=z_coefficient,
                    loss_chunk_tokens=args.loss_chunk_tokens,
                    use_cache=False,
                )
            loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(ddp.parameters(), args.grad_clip)
            muon_optimizer.step()
            adamw_optimizer.step()
            tokens_seen += step_tokens
            log_tokens += step_tokens

            if step % args.log_interval == 0 or step == start_step + 1:
                stats = torch.tensor(
                    [loss.item(), ce_loss.item(), z_loss.item(), data_wait, float(grad_norm)],
                    dtype=torch.float64,
                    device=device,
                )
                dist.all_reduce(stats, op=dist.ReduceOp.SUM)
                stats /= world_size
                if rank == 0:
                    now = time.time()
                    throughput = log_tokens / max(now - log_started, 1e-9)
                    record = {
                        "step": step,
                        "total_steps": total_steps,
                        "tokens": tokens_seen,
                        "source": source_key,
                        "source_tokens": dict(source_scheduler.consumed),
                        "loss": stats[0].item(),
                        "ce_loss": stats[1].item(),
                        "perplexity": math.exp(min(20.0, stats[1].item())),
                        "z_loss": stats[2].item(),
                        "grad_norm": stats[4].item(),
                        "muon_lr": muon_lr,
                        "adamw_lr": adamw_lr,
                        "ngram_lr": ngram_lr,
                        "weight_decay": weight_decay,
                        "tokens_per_second": throughput,
                        "data_wait_seconds": stats[3].item(),
                        "eta_seconds": (
                            args.total_tokens - tokens_seen
                        ) / max(throughput, 1e-9),
                    }
                    with metrics_path.open("a") as file:
                        file.write(json.dumps(record) + "\n")
                    print(
                        f"step={step:06d}/{total_steps} tokens={tokens_seen:,} "
                        f"src={source_key} loss={record['loss']:.4f} "
                        f"ppl={record['perplexity']:.2f} "
                        f"grad={record['grad_norm']:.3f} "
                        f"tok/s={throughput:,.0f} "
                        f"data={record['data_wait_seconds']:.2f}s "
                        f"eta={record['eta_seconds'] / 3600:.2f}h",
                        flush=True,
                    )
                    log_started = now
                    log_tokens = 0
            del input_ids, shifted_labels, loss, ce_loss, z_loss

            if step in checkpoint_steps:
                percent = checkpoint_steps[step]
                dist.barrier()
                if rank == 0:
                    metadata = {
                        "variant": args.variant,
                        "parameters": parameter_count,
                        "transformer_parameters": breakdown["transformer"],
                        "ngram_parameters": breakdown["ngram"],
                        "architecture": config.to_dict(),
                        "training_percent": percent,
                        "step": step,
                        "total_steps": total_steps,
                        "tokens_seen": tokens_seen,
                        "target_tokens": args.total_tokens,
                        "tokens_per_full_step": tokens_per_step,
                        "final_step_supervised_tokens": (
                            args.total_tokens - (total_steps - 1) * tokens_per_step
                        ),
                        "world_size": world_size,
                        "local_batch": local_batch,
                        "global_batch_sequences": args.global_batch_sequences,
                        "gpu_name": torch.cuda.get_device_name(0),
                        "architecture_revision": architecture_revision,
                        "dataset_revisions": dataset_revisions,
                        "target_source_shares": curriculum_module.TARGET_SHARES,
                        "target_source_tokens": curriculum_module.TARGET_SOURCE_TOKENS,
                        "actual_source_tokens": dict(source_scheduler.consumed),
                        "muon_peak_lr": args.muon_peak_lr,
                        "adamw_peak_lr": args.adamw_peak_lr,
                        "ngram_peak_lr": args.ngram_peak_lr,
                        "elapsed_seconds_this_job": time.time() - started,
                    }
                    export_checkpoint(
                        ddp,
                        config,
                        source_dir,
                        tokenizer_dir,
                        args,
                        metadata,
                        metrics_path,
                        muon_optimizer,
                        adamw_optimizer,
                        source_scheduler,
                    )
                    log_started = time.time()
                    log_tokens = 0
                dist.barrier()
    finally:
        for stream in streams.values():
            stream.close()
        if dist.is_initialized():
            dist.destroy_process_group()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--variant", choices=tuple(EXPECTED), required=True)
    parser.add_argument("--repo-id", required=True)
    parser.add_argument("--output-dir", default="/tmp/bananamind21-coder-training")
    parser.add_argument("--total-tokens", type=int, default=TOTAL_TOKENS)
    parser.add_argument("--seq-len", type=int, default=4096)
    parser.add_argument("--global-batch-sequences", type=int, default=128)
    parser.add_argument("--expected-world-size", type=int, choices=(4, 8), default=4)
    parser.add_argument("--muon-peak-lr", type=float, default=0.02)
    parser.add_argument("--adamw-peak-lr", type=float, default=0.002)
    parser.add_argument("--ngram-peak-lr", type=float, default=0.001)
    parser.add_argument("--muon-momentum", type=float, default=0.95)
    parser.add_argument("--muon-ns-steps", type=int, default=5)
    parser.add_argument("--warmup-tokens", type=int, default=600_000_000)
    parser.add_argument("--decay-ratio", type=float, default=0.15)
    parser.add_argument("--weight-decay", type=float, default=0.1)
    parser.add_argument("--final-weight-decay", type=float, default=0.01)
    parser.add_argument(
        "--weight-decay-switch-tokens",
        type=int,
        default=12_000_000_000,
    )
    parser.add_argument("--grad-clip", type=float, default=1.0)
    parser.add_argument("--z-loss-coeff", type=float, default=1e-4)
    parser.add_argument("--z-loss-until-tokens", type=int, default=12_000_000_000)
    parser.add_argument("--loss-chunk-tokens", type=int, default=16_384)
    parser.add_argument("--stack-encode-batch-size", type=int, default=128)
    parser.add_argument("--web-encode-batch-size", type=int, default=1024)
    parser.add_argument("--prefetch-batches", type=int, default=2)
    parser.add_argument("--shuffle-buffer", type=int, default=10_000)
    parser.add_argument("--log-interval", type=int, default=10)
    parser.add_argument("--seed", type=int, default=1337)
    parser.add_argument(
        "--compile",
        action=argparse.BooleanOptionalAction,
        default=True,
    )
    parser.add_argument(
        "--resume",
        action=argparse.BooleanOptionalAction,
        default=True,
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.repo_id != EXPECTED[args.variant]["repo_id"]:
        print(f"Using custom target repository {args.repo_id}", flush=True)
    if args.total_tokens != TOTAL_TOKENS:
        raise ValueError("Coder runs require exactly 30B supervised tokens")
    if args.seq_len != 4096:
        raise ValueError("Coder runs require 4,096-token sequences")
    if args.global_batch_sequences != 128:
        raise ValueError("Keep the global batch at 128 sequences")
    if args.global_batch_sequences % args.expected_world_size:
        raise ValueError("Global batch must divide by the GPU count")
    if not os.environ.get("HF_TOKEN"):
        raise RuntimeError("HF_TOKEN must be configured as a Job secret")
    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
    assets = prepare_runtime_assets(args)
    world_size = torch.cuda.device_count()
    if world_size != args.expected_world_size:
        raise RuntimeError(f"Expected {args.expected_world_size} GPUs, found {world_size}")
    mp.spawn(
        train_worker,
        args=(world_size, args, *assets),
        nprocs=world_size,
        join=True,
    )
    sys.stdout.flush()
    sys.stderr.flush()
    os._exit(0)


if __name__ == "__main__":
    main()