diff --git a/README.md b/README.md index ab52e606..f327c81e 100644 --- a/README.md +++ b/README.md @@ -58,6 +58,10 @@ pip install . Refer to tests for usage of different optimizers, e.g. [`tests/test_orthogonalized_optimizer.py::MuonTest`](tests/test_orthogonalized_optimizer.py). +Train the upstream nanoGPT model with absl flags and a choice of optimizers using the +[minimal-dependency nanoGPT example](examples/nanogpt/README.md). It supports checkpoint resume, +mixed precision, generation, and PyTorch DDP. + ### Integration with Megatron Core Emerging optimizer has been integrated into Megatron Core with tensor parallel support and many other features (split QKV for example), check out [emerging_optimizers.py](https://github.com/NVIDIA/Megatron-LM/blob/core_v0.18.0/megatron/core/optimizer/emerging_optimizers.py). diff --git a/examples/nanogpt/README.md b/examples/nanogpt/README.md new file mode 100644 index 00000000..6ff92ba7 --- /dev/null +++ b/examples/nanogpt/README.md @@ -0,0 +1,29 @@ +# nanoGPT with Emerging Optimizers + +Run from the repository root in a PyTorch container with `emerging-optimizers` +installed, a local [nanoGPT](https://github.com/karpathy/nanoGPT) checkout, and +prepared `train.bin`/`val.bin` data (`prepare.py` prepares character data). + +```bash +python examples/nanogpt/train.py \ + --nanogpt_dir=/opt/nanoGPT --data_dir=/tmp/shakespeare_char \ + --out_dir=/tmp/nanogpt-muon --optimizer=muon \ + --n_layer=4 --n_head=4 --n_embd=128 --block_size=128 \ + --batch_size=8 --gradient_accumulation_steps=4 \ + --learning_rate=0.003 --adamw_learning_rate=0.0003 \ + --optimizer_kwargs='{"momentum":0.95,"extra_scale_factor":0.2}' \ + --max_iters=2000 --warmup_iters=100 --lr_decay_iters=2000 \ + --eval_interval=100 --dtype=bfloat16 +``` + +Features: + +- Imports the upstream GPT model directly; standalone scripts with absl flags and flagfiles. +- 27 optimizer choices, including Muown and AdamW; excludes TP REKLS and StackedSoap. See `--list_optimizers`. +- Matrix optimizers paired with AdamW for embeddings, output head, biases, and normalization parameters. +- Checkpoint resume, validation, and text generation through `sample.py`. +- Gradient accumulation, clipping, and cosine learning-rate decay with warmup. +- CPU/GPU training, FP32/BF16/FP16, optional `torch.compile`, and PyTorch DDP. +- Upstream binary datasets and character-data preparation without NumPy. +- Console and JSONL metrics; optional W&B logging. +- No additional required dependencies beyond PyTorch and absl-py. Optional GPT-2 initialization, GPT-2 sampling, and W&B use `transformers`, `tiktoken`, and `wandb`, respectively. diff --git a/examples/nanogpt/common.py b/examples/nanogpt/common.py new file mode 100644 index 00000000..02001f9a --- /dev/null +++ b/examples/nanogpt/common.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Shared upstream imports, data loading, and precision configuration.""" + +import importlib.util +import pickle +import sys +from collections.abc import Iterator +from contextlib import contextmanager, nullcontext +from pathlib import Path +from types import ModuleType +from typing import Any + +import torch + + +def load_nanogpt(directory: str) -> ModuleType: + """Import model.py from an upstream checkout without executing train.py.""" + path = Path(directory).expanduser().resolve() / "model.py" + if not path.is_file(): + raise FileNotFoundError(f"No nanoGPT model at {path}. Set --nanogpt_dir to a karpathy/nanoGPT checkout.") + name = "_emerging_upstream_nanogpt_model" + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ImportError(f"Cannot import {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module # dataclasses resolves annotations through sys.modules. + spec.loader.exec_module(module) + return module + + +def autocast_context(device: torch.device, dtype: str) -> Any: + """Keep weights in FP32 and optionally autocast forward/backward operations.""" + if dtype == "float32": + return nullcontext() + if device.type == "cpu" and dtype == "float16": + raise ValueError("CPU float16 is unsupported; use float32 or bfloat16.") + return torch.autocast(device_type=device.type, dtype=getattr(torch, dtype)) + + +def load_checkpoint(path: Path) -> dict[str, Any]: + """Read a nanoGPT checkpoint on CPU, including upstream configuration metadata.""" + return torch.load(path, map_location="cpu", weights_only=False) + + +def load_model_state(model: torch.nn.Module, state: dict[str, torch.Tensor]) -> None: + """Accept both compiled upstream checkpoints and ordinary state dictionaries.""" + model.load_state_dict({key.removeprefix("_orig_mod."): value for key, value in state.items()}) + + +class TokenData: + """Memory-map upstream uint16 train.bin/val.bin without requiring NumPy.""" + + def __init__(self, directory: str, block_size: int) -> None: + self.directory = Path(directory) + self.block_size = block_size + self.meta: dict[str, Any] = {} + meta_path = self.directory / "meta.pkl" + if meta_path.exists(): + with meta_path.open("rb") as stream: + self.meta = pickle.load(stream) + self.tokens = {} + if sys.byteorder != "little": + raise ValueError("The memory-mapped uint16 dataset requires a little-endian host.") + for split in ("train", "val"): + path = self.directory / f"{split}.bin" + size = path.stat().st_size + if size % 2 or size // 2 <= block_size: + raise ValueError(f"{path} must contain more than {block_size} uint16 tokens.") + self.tokens[split] = torch.from_file(str(path), shared=False, size=size // 2, dtype=torch.int16) + + def batch( + self, split: str, batch_size: int, device: torch.device, generator: torch.Generator + ) -> tuple[torch.Tensor, torch.Tensor]: + """Draw contiguous next-token batches using an independent CPU RNG.""" + data = self.tokens[split] + starts = torch.randint(len(data) - self.block_size, (batch_size,), generator=generator) + windows = torch.stack([data[i : i + self.block_size + 1] for i in starts.tolist()]) + windows = windows.long().bitwise_and_(0xFFFF) + x, y = windows[:, :-1].contiguous(), windows[:, 1:].contiguous() + if device.type == "cuda": + x, y = x.pin_memory(), y.pin_memory() + return x.to(device, non_blocking=True), y.to(device, non_blocking=True) + + +@contextmanager +def optimizer_rng(device: torch.device, seed: int) -> Iterator[None]: + """Give stochastic updates the same RNG on each DDP rank without changing dropout RNG.""" + devices = ( + [device.index if device.index is not None else torch.cuda.current_device()] if device.type == "cuda" else [] + ) + with torch.random.fork_rng(devices=devices): + torch.random.default_generator.manual_seed(seed) + if device.type == "cuda": + with torch.cuda.device(device): + torch.cuda.manual_seed(seed) + yield diff --git a/examples/nanogpt/optimizers.py b/examples/nanogpt/optimizers.py new file mode 100644 index 00000000..5ea6b1c8 --- /dev/null +++ b/examples/nanogpt/optimizers.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Optimizer selection and disjoint parameter ownership for upstream GPT.""" + +import importlib +import inspect +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, override + +import torch +from absl import logging + +from emerging_optimizers import registry + + +if TYPE_CHECKING: + from typing import overload + + +SCALAR_OPTIMIZERS = {"adamw", "lion", "signum", "laprop", "sim_ademamix"} +EXCLUDED_OPTIMIZERS = {"tp_rekls", "stacked_soap"} + + +def optimizer_classes() -> dict[str, type[torch.optim.Optimizer]]: + """Register supported families, including the unregistered contributed Muown.""" + for module in ( + "orthogonalized_optimizers", + "scalar_optimizers", + "legacy_soap", + "shampoo", + "psgd", + "riemannian_optimizers", + "riemannian_optimizers.isospectral", + ): + importlib.import_module(f"emerging_optimizers.{module}") + from emerging_optimizers.contrib.muown import Muown + + classes = {name: registry.get_optimizer_cls(name) for name in registry.get_optimizer_name_list()} + return {name: cls for name, cls in classes.items() if name not in EXCLUDED_OPTIMIZERS} | { + "adamw": torch.optim.AdamW, + "muown": Muown, + } + + +class OptimizerGroup(torch.optim.Optimizer): + """Expose child optimizers as one optimizer, including one AMP overflow decision.""" + + def __init__(self, optimizers: list[torch.optim.Optimizer]) -> None: + self.optimizers = optimizers + super().__init__([group for opt in optimizers for group in opt.param_groups], {}) + + if TYPE_CHECKING: + + @overload + def step(self, closure: None = ...) -> None: ... + + @overload + def step(self, closure: Callable[[], float]) -> float: ... + + @torch.no_grad() + @override + def step(self, closure: Callable[[], float] | None = None) -> float | None: + """Apply all child updates; GradScaler checks all parameter groups first.""" + if closure is not None: + raise ValueError("closure is not supported") + for optimizer in self.optimizers: + optimizer.step() + return None + + @override + def state_dict(self) -> dict[str, Any]: + """Preserve each optimizer's own serialization format.""" + return {"optimizers": [optimizer.state_dict() for optimizer in self.optimizers]} + + @override + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + """Restore children and refresh group references replaced during loading.""" + states = state_dict["optimizers"] + if len(states) != len(self.optimizers): + raise ValueError("Checkpoint optimizer layout does not match the configured optimizer.") + for optimizer, state in zip(self.optimizers, states, strict=True): + optimizer.load_state_dict(state) + self.param_groups = [group for opt in self.optimizers for group in opt.param_groups] + + +def decay_groups(params: list[torch.nn.Parameter], weight_decay: float) -> list[dict[str, Any]]: + """Apply weight decay to matrices, leaving biases and normalization scales alone.""" + return [ + {"params": selected, "weight_decay": decay} + for selected, decay in ( + ([p for p in params if p.ndim >= 2], weight_decay), + ([p for p in params if p.ndim < 2], 0.0), + ) + if selected + ] + + +def build_optimizer( + model: torch.nn.Module, + name: str, + learning_rate: float, + adamw_learning_rate: float, + weight_decay: float, + optimizer_kwargs: dict[str, Any] | None = None, + adamw_betas: tuple[float, float] = (0.9, 0.95), + *, + initialize: bool = True, +) -> OptimizerGroup: + """Use scalar optimizers globally or pair a hidden-matrix optimizer with AdamW.""" + classes = optimizer_classes() + if name not in classes: + raise ValueError(f"Unknown optimizer {name!r}; choose from {sorted(classes)}") + cls = classes[name] + kwargs = dict(optimizer_kwargs or {}) + if {"params", "lr", "weight_decay"} & kwargs.keys(): + raise ValueError("Set params, lr, and weight_decay through the model and dedicated flags.") + signature = inspect.signature(cls) + allowed = set(signature.parameters) + if name == "muon_hyperball": + allowed |= set(inspect.signature(classes["muon"]).parameters) + allowed.discard("weight_update_hook") + unknown = kwargs.keys() - allowed + if unknown: + raise ValueError(f"Unsupported options for {name}: {sorted(unknown)}; valid options: {sorted(allowed)}") + kwargs["lr"] = learning_rate + if "weight_decay" in allowed: + kwargs["weight_decay"] = weight_decay + if name == "adaptive_muon": + kwargs.setdefault("momentum", 0.95) + if name == "adamw": + kwargs.setdefault("betas", adamw_betas) + + params = [p for p in model.parameters() if p.requires_grad] + if name in SCALAR_OPTIMIZERS: + children = [cls(decay_groups(params, weight_decay), **kwargs)] + else: + # Collect by identity: the LM head and token embedding normally share one Parameter. + auxiliary_ids = { + id(p) for module in model.modules() if isinstance(module, torch.nn.Embedding) for p in module.parameters() + } + auxiliary_ids.update(id(p) for p in model.get_submodule("lm_head").parameters()) + matrices = [p for p in params if p.ndim == 2 and id(p) not in auxiliary_ids] + matrix_ids = {id(p) for p in matrices} + auxiliary = [p for p in params if id(p) not in matrix_ids] + if not matrices: + raise ValueError("The model has no hidden matrices to optimize.") + if name == "muon_hyperball": + radius = kwargs.setdefault("hyperball_radius", 1.0) + if radius <= 0: + raise ValueError("hyperball_radius must be positive") + if initialize: + logging.info("Rescaling hidden matrices to MuonHyperball radius %s", radius) + with torch.no_grad(): + for p in matrices: + norm = p.norm() + if norm == 0: + raise ValueError("MuonHyperball cannot initialize a zero matrix.") + p.mul_(radius / norm) + children = [cls(matrices, **kwargs)] + if auxiliary: + children.append( + torch.optim.AdamW(decay_groups(auxiliary, weight_decay), lr=adamw_learning_rate, betas=adamw_betas) + ) + optimizer = OptimizerGroup(children) + owned = [id(p) for group in optimizer.param_groups for p in group["params"]] + if len(owned) != len(set(owned)) or set(owned) != {id(p) for p in params}: + raise ValueError("Every trainable parameter must belong to exactly one optimizer group.") + for group in optimizer.param_groups: + group["initial_lr"] = group["lr"] + for child in children: + count = sum(p.numel() for group in child.param_groups for p in group["params"]) + logging.info("%s owns %d parameters", type(child).__name__, count) + return optimizer diff --git a/examples/nanogpt/prepare.py b/examples/nanogpt/prepare.py new file mode 100644 index 00000000..99e89b2a --- /dev/null +++ b/examples/nanogpt/prepare.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Prepare upstream-compatible character data using only the standard library.""" + +import array +import pickle +import sys +import urllib.request +from pathlib import Path + +from absl import app, flags, logging + + +FLAGS = flags.FLAGS +flags.DEFINE_string("input_file", None, "Local UTF-8 input text; otherwise download Tiny Shakespeare.") +flags.DEFINE_string("data_dir", "data/shakespeare_char", "Output dataset directory.") +flags.DEFINE_float("train_fraction", 0.9, "Fraction of text used for training.", lower_bound=0.0, upper_bound=1.0) + + +def prepare(text: str, directory: Path, train_fraction: float = 0.9) -> None: + """Encode text to little-endian uint16 bins and upstream's character metadata.""" + chars = sorted(set(text)) + if not chars or len(chars) > 65536: + raise ValueError("Character vocabulary must contain between 1 and 65536 entries.") + boundary = int(len(text) * train_fraction) + if not 0 < boundary < len(text): + raise ValueError("Training and validation splits must both be nonempty.") + stoi = {char: index for index, char in enumerate(chars)} + directory.mkdir(parents=True, exist_ok=True) + for split, content in (("train", text[:boundary]), ("val", text[boundary:])): + tokens = array.array("H", (stoi[char] for char in content)) + if sys.byteorder != "little": + tokens.byteswap() + (directory / f"{split}.bin").write_bytes(tokens.tobytes()) + with (directory / "meta.pkl").open("wb") as stream: + pickle.dump({"vocab_size": len(chars), "stoi": stoi, "itos": dict(enumerate(chars))}, stream) + logging.info( + "Wrote %d training tokens, %d validation tokens, vocabulary %d", boundary, len(text) - boundary, len(chars) + ) + + +def main(argv: list[str]) -> None: + """Prepare local text or the same Tiny Shakespeare corpus used by nanoGPT.""" + if len(argv) > 1: + raise app.UsageError("Use named flags; positional arguments are unsupported.") + if FLAGS.input_file: + text = Path(FLAGS.input_file).read_text(encoding="utf-8") + else: + url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt" + with urllib.request.urlopen(url, timeout=60) as response: + text = response.read().decode("utf-8") + prepare(text, Path(FLAGS.data_dir), FLAGS.train_fraction) + + +if __name__ == "__main__": + app.run(main) diff --git a/examples/nanogpt/sample.py b/examples/nanogpt/sample.py new file mode 100644 index 00000000..c0e726c2 --- /dev/null +++ b/examples/nanogpt/sample.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Generate text with upstream GPT.generate and absl flags.""" + +import pickle +from pathlib import Path +from typing import Any + +import torch +from absl import app, flags +from common import autocast_context, load_checkpoint, load_model_state, load_nanogpt + + +FLAGS = flags.FLAGS +flags.DEFINE_string("nanogpt_dir", "../nanoGPT", "Upstream nanoGPT checkout.") +flags.DEFINE_string("checkpoint", "out/nanogpt/ckpt.pt", "Checkpoint for init_from=resume.") +flags.DEFINE_enum("init_from", "resume", ["resume", "gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl"], "Initialization.") +flags.DEFINE_string("data_dir", None, "Optional upstream dataset directory containing character meta.pkl.") +flags.DEFINE_string("start", "\n", "Prompt, or FILE:path to read a prompt file.") +flags.DEFINE_integer("num_samples", 1, "Number of samples.", lower_bound=1) +flags.DEFINE_integer("max_new_tokens", 200, "Maximum new tokens per sample.", lower_bound=0) +flags.DEFINE_float("temperature", 0.8, "Sampling temperature; must be positive.") +flags.DEFINE_integer("top_k", 200, "Top-k filtering; zero disables it.", lower_bound=0) +flags.DEFINE_integer("seed", 1337, "Sampling seed.") +flags.DEFINE_string("device", "cuda", "cpu, cuda, or cuda:N.") +flags.DEFINE_enum("dtype", "bfloat16", ["float32", "bfloat16", "float16"], "Inference precision.") +flags.DEFINE_bool("compile", False, "Compile the model.") + + +def main(argv: list[str]) -> None: + """Load a trained or pretrained model and reuse upstream generation.""" + if len(argv) > 1: + raise app.UsageError("Use named flags; positional arguments are unsupported.") + if FLAGS.temperature <= 0: + raise app.UsageError("temperature must be positive") + torch.manual_seed(FLAGS.seed) + device = torch.device(FLAGS.device) + meta = {} + if FLAGS.data_dir: + with (Path(FLAGS.data_dir) / "meta.pkl").open("rb") as stream: + meta = pickle.load(stream) + if FLAGS.init_from != "resume" and ("stoi" in meta or "itos" in meta): + raise app.UsageError("Character metadata cannot be used with pretrained GPT-2; omit --data_dir.") + upstream = load_nanogpt(FLAGS.nanogpt_dir) + model: Any + if FLAGS.init_from == "resume": + checkpoint = load_checkpoint(Path(FLAGS.checkpoint)) + model = upstream.GPT(upstream.GPTConfig(**checkpoint["model_args"])) + load_model_state(model, checkpoint["model"]) + saved_meta = checkpoint.get("meta", {}) + if not FLAGS.data_dir: + meta = saved_meta + elif "stoi" in saved_meta and meta.get("stoi") != saved_meta["stoi"]: + raise app.UsageError("Character metadata does not match the checkpoint's saved mapping.") + else: + model = upstream.GPT.from_pretrained(FLAGS.init_from, {"dropout": 0.0}) + if "stoi" in meta or "itos" in meta: + stoi, itos = meta.get("stoi"), meta.get("itos") + if not isinstance(stoi, dict) or not isinstance(itos, dict): + raise app.UsageError("Character metadata must include both stoi and itos dictionaries.") + vocab_size = model.config.vocab_size + if meta.get("vocab_size") != vocab_size or len(stoi) != vocab_size or set(itos) != set(range(vocab_size)): + raise app.UsageError("Character metadata must cover exactly the loaded model's vocabulary.") + if any( + not isinstance(char, str) or len(char) != 1 or not isinstance(token, int) or itos.get(token) != char + for char, token in stoi.items() + ): + raise app.UsageError("Character metadata stoi and itos must be inverse character mappings.") + encode = lambda text: [stoi[char] for char in text] + decode = lambda tokens: "".join(itos[token] for token in tokens) + else: + import tiktoken + + encoding = tiktoken.get_encoding("gpt2") + encode = lambda text: encoding.encode(text, allowed_special={"<|endoftext|>"}) + decode = encoding.decode + prompt = Path(FLAGS.start[5:]).read_text() if FLAGS.start.startswith("FILE:") else FLAGS.start + tokens = encode(prompt) + if not tokens: + raise app.UsageError("The prompt must contain at least one token.") + x = torch.tensor([tokens], device=device, dtype=torch.long) + model.eval().to(device) + if FLAGS.compile: + model = torch.compile(model) + with torch.no_grad(), autocast_context(device, FLAGS.dtype): + for _ in range(FLAGS.num_samples): + result = model.generate(x, FLAGS.max_new_tokens, temperature=FLAGS.temperature, top_k=FLAGS.top_k or None) + print(decode(result[0].tolist())) + + +if __name__ == "__main__": + app.run(main) diff --git a/examples/nanogpt/train.py b/examples/nanogpt/train.py new file mode 100644 index 00000000..7afb69ca --- /dev/null +++ b/examples/nanogpt/train.py @@ -0,0 +1,379 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Train an imported nanoGPT model with absl flags and Emerging Optimizers.""" + +import json +import math +import os +import time +from contextlib import nullcontext +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist +from absl import app, flags, logging +from common import ( + TokenData, + autocast_context, + load_checkpoint, + load_model_state, + load_nanogpt, + optimizer_rng, +) +from optimizers import build_optimizer, optimizer_classes +from torch.nn.parallel import DistributedDataParallel + + +FLAGS = flags.FLAGS +_EXISTING_FLAGS = set(FLAGS) +flags.DEFINE_string("nanogpt_dir", "../nanoGPT", "Path to an upstream karpathy/nanoGPT checkout.") +flags.DEFINE_string( + "data_dir", "data/shakespeare_char", "Directory containing train.bin, val.bin, and optional meta.pkl." +) +flags.DEFINE_string("out_dir", "out/nanogpt", "Checkpoint and metrics directory.") +flags.DEFINE_enum( + "init_from", "scratch", ["scratch", "resume", "gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl"], "Initialization." +) +flags.DEFINE_string("checkpoint", None, "Checkpoint to resume; defaults to out_dir/ckpt.pt.") +flags.DEFINE_string("optimizer", "muon", "Optimizer name; see --list_optimizers.") +flags.DEFINE_bool("list_optimizers", False, "List supported optimizer names and exit.") +flags.DEFINE_string("optimizer_kwargs", "{}", 'JSON constructor options, e.g. {"momentum": 0.95, "num_ns_steps": 5}.') +flags.DEFINE_float("learning_rate", 3e-4, "Main optimizer peak learning rate.", lower_bound=0.0) +flags.DEFINE_float("adamw_learning_rate", 3e-4, "Auxiliary AdamW peak learning rate.", lower_bound=0.0) +flags.DEFINE_float( + "weight_decay", 0.1, "Matrix weight decay; biases and normalization scales use zero.", lower_bound=0.0 +) +flags.DEFINE_float("beta1", 0.9, "AdamW first moment beta.") +flags.DEFINE_float("beta2", 0.95, "AdamW second moment beta.") +flags.DEFINE_integer("n_layer", 6, "Transformer layers.", lower_bound=1) +flags.DEFINE_integer("n_head", 6, "Attention heads.", lower_bound=1) +flags.DEFINE_integer("n_embd", 384, "Embedding dimension.", lower_bound=1) +flags.DEFINE_integer("block_size", 256, "Context length; may crop a pretrained model.", lower_bound=1) +flags.DEFINE_integer("vocab_size", None, "Vocabulary size; otherwise metadata or GPT-2's padded 50304.", lower_bound=1) +flags.DEFINE_float("dropout", 0.0, "Model dropout.", lower_bound=0.0, upper_bound=1.0) +flags.DEFINE_bool("bias", False, "Use linear and LayerNorm biases.") +flags.DEFINE_integer("batch_size", 12, "Microbatch size per rank.", lower_bound=1) +flags.DEFINE_integer( + "gradient_accumulation_steps", 1, "Global microbatches per update; divisible by world size.", lower_bound=1 +) +flags.DEFINE_integer( + "max_iters", 5000, "Total optimizer iterations, including previously completed iterations.", lower_bound=0 +) +flags.DEFINE_float("grad_clip", 1.0, "Global gradient norm clip; zero disables clipping.", lower_bound=0.0) +flags.DEFINE_bool("decay_lr", True, "Use linear warmup and cosine learning-rate decay.") +flags.DEFINE_integer("warmup_iters", 100, "Learning-rate warmup iterations.", lower_bound=0) +flags.DEFINE_integer("lr_decay_iters", 5000, "Iteration at which cosine decay reaches its minimum.", lower_bound=1) +flags.DEFINE_float("min_lr", 3e-5, "Main optimizer minimum LR; auxiliary groups keep their LR ratio.", lower_bound=0.0) +flags.DEFINE_integer("eval_interval", 250, "Evaluate and consider checkpointing every N updates.", lower_bound=1) +flags.DEFINE_integer("eval_iters", 20, "Batches per split and rank during evaluation.", lower_bound=1) +flags.DEFINE_bool("eval_only", False, "Evaluate once and exit, including when resuming.") +flags.DEFINE_bool( + "always_save_checkpoint", True, "Save at each evaluation; otherwise save improvements and the final state." +) +flags.DEFINE_integer("log_interval", 10, "Log training metrics every N updates.", lower_bound=1) +flags.DEFINE_string("device", "cuda", "cpu, cuda, or a CUDA device such as cuda:0.") +flags.DEFINE_enum( + "dtype", "bfloat16", ["float32", "bfloat16", "float16"], "Forward/backward precision; parameters remain FP32." +) +flags.DEFINE_bool("compile", False, "Compile the model with torch.compile.") +flags.DEFINE_enum("backend", "nccl", ["nccl", "gloo"], "DDP process-group backend; use gloo for CPU.") +flags.DEFINE_integer("seed", 1337, "Random seed.") +flags.DEFINE_bool("wandb_log", False, "Enable optional Weights & Biases logging.") +flags.DEFINE_string("wandb_project", "nanogpt", "W&B project.") +flags.DEFINE_string("wandb_run_name", None, "W&B run name.") +_CONFIG_FLAGS = set(FLAGS) - _EXISTING_FLAGS + + +def lr_multiplier(step: int, config: dict[str, Any]) -> float: + """Scale all groups together while preserving their configured LR ratios.""" + if not config["decay_lr"]: + return 1.0 + if step < config["warmup_iters"]: + return (step + 1) / (config["warmup_iters"] + 1) + minimum = config["min_lr"] / config["learning_rate"] if config["learning_rate"] else 0.0 + if step >= config["lr_decay_iters"]: + return minimum + progress = (step - config["warmup_iters"]) / (config["lr_decay_iters"] - config["warmup_iters"]) + return minimum + (1 - minimum) * 0.5 * (1 + math.cos(math.pi * progress)) + + +def train(config: dict[str, Any]) -> None: + """Run single-device or torchrun DDP training, evaluation, and checkpointing.""" + ddp = "RANK" in os.environ + rank = int(os.environ.get("RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + device = torch.device(config["device"]) + if ddp and device.type == "cuda": + device = torch.device("cuda", int(os.environ["LOCAL_RANK"])) + if device.type == "cuda": + device = torch.device("cuda", device.index if device.index is not None else 0) + torch.cuda.set_device(device) + if ddp: + dist.init_process_group(backend=config["backend"]) + try: + _train(config, device, rank, world_size, ddp) + finally: + if ddp: + dist.destroy_process_group() + + +def _train(config: dict[str, Any], device: torch.device, rank: int, world_size: int, ddp: bool) -> None: + output = Path(config["out_dir"]) + if rank == 0: + output.mkdir(parents=True, exist_ok=True) + checkpoint = None + if config["init_from"] == "resume": + checkpoint = load_checkpoint(Path(config["checkpoint"] or output / "ckpt.pt")) + if "emerging_optimizer" not in checkpoint: + raise ValueError( + "Resume needs an Emerging Optimizers checkpoint; upstream checkpoints can be used by sample.py." + ) + # Build the data loader with the context length stored in the model checkpoint. + config["block_size"] = checkpoint["model_args"]["block_size"] + # Constructor choices and the schedule are part of the optimizer state, not restart overrides. + for key in ( + "seed", + "optimizer", + "optimizer_kwargs", + "learning_rate", + "adamw_learning_rate", + "weight_decay", + "beta1", + "beta2", + "decay_lr", + "warmup_iters", + "lr_decay_iters", + "min_lr", + ): + config[key] = checkpoint["config"][key] + if config["gradient_accumulation_steps"] % world_size: + raise ValueError("gradient_accumulation_steps must be divisible by the DDP world size.") + accumulation = config["gradient_accumulation_steps"] // world_size + if config["decay_lr"] and ( + config["lr_decay_iters"] <= config["warmup_iters"] or config["min_lr"] > config["learning_rate"] + ): + raise ValueError("Require lr_decay_iters > warmup_iters and min_lr <= learning_rate.") + kwargs = json.loads(config["optimizer_kwargs"]) + if not isinstance(kwargs, dict): + raise ValueError("optimizer_kwargs must be a JSON object.") + + torch.manual_seed(config["seed"]) + upstream = load_nanogpt(config["nanogpt_dir"]) + data = TokenData(config["data_dir"], config["block_size"]) + model_args = {key: config[key] for key in ("n_layer", "n_head", "n_embd", "block_size", "dropout", "bias")} + model_args["vocab_size"] = config["vocab_size"] or data.meta.get("vocab_size", 50304) + model: Any + if checkpoint is not None: + model_args = dict(checkpoint["model_args"]) + model = upstream.GPT(upstream.GPTConfig(**model_args)) + load_model_state(model, checkpoint["model"]) + elif config["init_from"] == "scratch": + model = upstream.GPT(upstream.GPTConfig(**model_args)) + else: + model = upstream.GPT.from_pretrained(config["init_from"], {"dropout": config["dropout"]}) + model_args = vars(model.config).copy() + if config["block_size"] > model.config.block_size: + raise ValueError("block_size cannot exceed the resumed/pretrained model's context length.") + if config["block_size"] < model.config.block_size: + if checkpoint is not None: + raise ValueError( + "Changing block_size on resume would invalidate optimizer state; keep the saved block_size." + ) + model.crop_block_size(config["block_size"]) + model_args["block_size"] = config["block_size"] + if data.meta.get("vocab_size", 0) > model.config.vocab_size: + raise ValueError("Dataset vocabulary exceeds the model vocabulary.") + model.to(device) + optimizer = build_optimizer( + model, + config["optimizer"], + config["learning_rate"], + config["adamw_learning_rate"], + config["weight_decay"], + kwargs, + (config["beta1"], config["beta2"]), + initialize=checkpoint is None, + ) + scaler = torch.amp.GradScaler(device.type, enabled=config["dtype"] == "float16") + step = 0 + best_val_loss = float("inf") + if checkpoint is not None: + optimizer.load_state_dict(checkpoint["optimizer"]) + if checkpoint["scaler"]: + scaler.load_state_dict(checkpoint["scaler"]) + step = checkpoint["iter_num"] + best_val_loss = checkpoint["best_val_loss"] + + raw_model = model + if config["compile"]: + model = torch.compile(model) + if ddp: + model = DistributedDataParallel(model, device_ids=[device.index] if device.type == "cuda" else None) + torch.manual_seed(config["seed"] + rank) + train_rng = torch.Generator().manual_seed(config["seed"] + rank) + eval_rng = torch.Generator().manual_seed(config["seed"] + 10000 + rank) + if checkpoint is not None and len(checkpoint["rng_states"]) == world_size: + state = checkpoint["rng_states"][rank] + torch.set_rng_state(state["torch"]) + train_rng.set_state(state["train"]) + eval_rng.set_state(state["eval"]) + if device.type == "cuda" and state["cuda"] is not None: + torch.cuda.set_rng_state(state["cuda"], device) + elif checkpoint is not None: + logging.warning("World size changed; optimizer state resumes but rank RNG streams restart.") + checkpoint = None + wandb_run = None + if config["wandb_log"] and rank == 0: + import wandb + + wandb_run = wandb.init(project=config["wandb_project"], name=config["wandb_run_name"], config=config) + + def log_metrics(metrics: dict[str, Any]) -> None: + if rank == 0: + logging.info("%s", metrics) + with (output / "metrics.jsonl").open("a") as stream: + stream.write(json.dumps(metrics) + "\n") + if wandb_run is not None: + wandb_run.log(metrics, step=step) + + @torch.no_grad() + def evaluate() -> dict[str, float]: + model.eval() + losses = {} + for split in ("train", "val"): + total = torch.zeros((), device=device) + for _ in range(config["eval_iters"]): + x, y = data.batch(split, config["batch_size"], device, eval_rng) + with autocast_context(device, config["dtype"]): + _, loss = model(x, y) + total += loss.detach() + if ddp: + dist.all_reduce(total) + losses[f"{split}/loss"] = total.item() / (config["eval_iters"] * world_size) + model.train() + return losses + + def save() -> None: + state = { + "torch": torch.get_rng_state(), + "train": train_rng.get_state(), + "eval": eval_rng.get_state(), + "cuda": torch.cuda.get_rng_state(device) if device.type == "cuda" else None, + } + states: list[Any] = [None] * world_size + if ddp: + dist.all_gather_object(states, state) + else: + states[0] = state + if rank == 0: + payload = { + "model": raw_model.state_dict(), + "model_args": model_args, + "optimizer": optimizer.state_dict(), + "emerging_optimizer": config["optimizer"], + "scaler": scaler.state_dict(), + "iter_num": step, + "best_val_loss": best_val_loss, + "config": config, + "rng_states": states, + "meta": data.meta, + } + temporary = output / "ckpt.pt.tmp" + torch.save(payload, temporary) + temporary.replace(output / "ckpt.pt") + + optimizer.zero_grad(set_to_none=True) + if rank == 0: + logging.info("Tokens per update: %d", config["batch_size"] * config["block_size"] * accumulation * world_size) + # A resumed run continues directly; reevaluation would advance its saved evaluation RNG. + if step == 0 or config["eval_only"]: + losses = evaluate() + best_val_loss = min(best_val_loss, losses["val/loss"]) + log_metrics({"iter": step, **losses}) + if not config["eval_only"]: + while step < config["max_iters"]: + if device.type == "cuda" and (step + 1) % config["log_interval"] == 0: + torch.cuda.synchronize(device) + start = time.perf_counter() + multiplier = lr_multiplier(step, config) + for group in optimizer.param_groups: + group["lr"] = group["initial_lr"] * multiplier + accumulated_loss = torch.zeros((), device=device) + for micro_step in range(accumulation): + sync = model.no_sync() if ddp and micro_step < accumulation - 1 else nullcontext() + with sync: + x, y = data.batch("train", config["batch_size"], device, train_rng) + with autocast_context(device, config["dtype"]): + _, loss = model(x, y) + accumulated_loss += loss.detach() / accumulation + scaler.scale(loss / accumulation).backward() + if config["grad_clip"]: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(model.parameters(), config["grad_clip"]) + # PSGD samples damping noise; DDP replicas must use identical optimizer randomness. + rng_context = ( + optimizer_rng(device, config["seed"] + step + 1000000) + if config["optimizer"] == "psgd_pro" + else nullcontext() + ) + with rng_context: + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad(set_to_none=True) + step += 1 + if step % config["log_interval"] == 0: + if ddp: + dist.all_reduce(accumulated_loss) + loss_value = accumulated_loss.item() / world_size + if device.type == "cuda": + torch.cuda.synchronize(device) + elapsed = time.perf_counter() - start + tokens = config["batch_size"] * config["block_size"] * accumulation * world_size + log_metrics( + { + "iter": step, + "loss": loss_value, + "lr": config["learning_rate"] * multiplier, + "tokens_per_second": tokens / elapsed, + "step_ms": elapsed * 1000, + } + ) + final = step == config["max_iters"] + if step % config["eval_interval"] == 0 or final: + losses = evaluate() + improved = losses["val/loss"] < best_val_loss + best_val_loss = min(best_val_loss, losses["val/loss"]) + log_metrics({"iter": step, **losses}) + if improved or config["always_save_checkpoint"] or final: + save() + if config["max_iters"] == 0: + save() + if wandb_run is not None: + wandb_run.finish() + + +def main(argv: list[str]) -> None: + """Parse absl flags and launch training; --flagfile supports reusable recipes.""" + if len(argv) > 1: + raise app.UsageError("Use named flags or --flagfile; Python configuration files are unsupported.") + if FLAGS.list_optimizers: + print("\n".join(sorted(optimizer_classes()))) + return + train({key: FLAGS[key].value for key in _CONFIG_FLAGS}) + + +if __name__ == "__main__": + app.run(main)