Add nanogpt example

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).

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.

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

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