#!/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()