From 408792a1fdbe5c91eb31af7932810c71d0ed7b32 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 20:40:10 -0700 Subject: [PATCH 1/9] Add the MuonClip optimizer extension and launcher --- .../src/rg_nanogpt_one_head/muonclip.py | 645 ++++++++++++++++++ 1 file changed, 645 insertions(+) create mode 100644 baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip.py diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip.py new file mode 100644 index 0000000..b074f72 --- /dev/null +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/muonclip.py @@ -0,0 +1,645 @@ +from __future__ import annotations + +"""MuonClip extension for the one-head nanoGPT baseline. + +This module intentionally leaves the historical SGD, AdamW, and Muon entry +point unchanged. The dedicated ``rg-onehead-muonclip`` command installs the +extension in-process and then delegates to the existing training launcher. +""" + +import csv +from copy import deepcopy +import math +from pathlib import Path +from typing import Any, Iterable + +import torch +import torch.nn.functional as F + +_INSTALLED = False +_CURRENT_RUN_DIR: Path | None = None + + +def _zeropower( + update: torch.Tensor, + *, + steps: int, + eps: float, +) -> torch.Tensor: + """Use the repository's existing Newton--Schulz implementation.""" + + from .optimizers import zeropower_via_newton_schulz_5 + + return zeropower_via_newton_schulz_5( + update, + steps=int(steps), + eps=float(eps), + ) + + +class MuonClip(torch.optim.Optimizer): + """Muon with decoupled weight decay, RMS matching, and QK-Clip. + + The hidden-matrix update is + + ``M_t = momentum * M_(t-1) + G_t`` + ``O_t = NS(M_t) * update_rms_scale * sqrt(max(n, m))`` + ``W_t = W_(t-1) - lr * (O_t + weight_decay * W_(t-1))`` + + QK-Clip is applied after the matrix update. For regular multi-head + attention, query and key rows belonging to head ``h`` are multiplied by + ``gamma_h ** balance`` and ``gamma_h ** (1-balance)`` respectively, where + ``gamma_h = min(1, threshold / max_logit_h)``. + """ + + def __init__( + self, + params: Iterable[torch.nn.Parameter], + *, + model, + lr: float, + momentum: float = 0.95, + nesterov: bool = True, + weight_decay: float = 0.1, + newton_schulz_steps: int = 5, + eps: float = 1e-7, + update_rms_scale: float = 0.2, + qk_clip_threshold: float = 100.0, + qk_clip_balance: float = 0.5, + diagnostics_interval: int = 250, + diagnostics_path: Path | None = None, + ) -> None: + params = list(params) + if not params: + raise ValueError("MuonClip requires at least one parameter") + if any(parameter.ndim != 2 for parameter in params): + raise ValueError("MuonClip accepts only 2-D parameters") + if not 0.0 <= float(momentum) < 1.0: + raise ValueError("momentum must be in [0, 1)") + if float(weight_decay) < 0.0: + raise ValueError("weight_decay must be nonnegative") + if int(newton_schulz_steps) < 1 or float(eps) <= 0.0: + raise ValueError("Newton--Schulz steps and epsilon must be positive") + if float(update_rms_scale) <= 0.0: + raise ValueError("update_rms_scale must be positive") + if float(qk_clip_threshold) <= 0.0: + raise ValueError("qk_clip_threshold must be positive") + if not 0.0 <= float(qk_clip_balance) <= 1.0: + raise ValueError("qk_clip_balance must be in [0, 1]") + if int(diagnostics_interval) < 1: + raise ValueError("diagnostics_interval must be positive") + + defaults = { + "lr": float(lr), + "momentum": float(momentum), + "nesterov": bool(nesterov), + "weight_decay": float(weight_decay), + "newton_schulz_steps": int(newton_schulz_steps), + "eps": float(eps), + "update_rms_scale": float(update_rms_scale), + "qk_clip_threshold": float(qk_clip_threshold), + "qk_clip_balance": float(qk_clip_balance), + } + super().__init__(params, defaults) + self.model = model + self.diagnostics_interval = int(diagnostics_interval) + self.diagnostics_path = diagnostics_path + self.step_index = 0 + self.last_diagnostics: dict[str, float] = {} + self._diagnostic_interval_state: dict[str, torch.Tensor] | None = None + + @property + def qk_clip_threshold(self) -> float: + return float(self.param_groups[0]["qk_clip_threshold"]) + + @property + def qk_clip_balance(self) -> float: + return float(self.param_groups[0]["qk_clip_balance"]) + + def reset_qk_tracking(self) -> None: + for block in self.model.blocks: + setattr(block.attn, "_muonclip_max_logits", None) + + def _consume_qk_logits(self) -> tuple[torch.Tensor, ...]: + values: list[torch.Tensor] = [] + for block_index, block in enumerate(self.model.blocks): + value = getattr(block.attn, "_muonclip_max_logits", None) + setattr(block.attn, "_muonclip_max_logits", None) + if value is None: + raise RuntimeError( + "MuonClip did not observe pre-softmax QK logits for " + f"block {block_index}; the dedicated MuonClip launcher " + "must be used for this optimizer" + ) + value = value.detach().float().reshape(-1) + if value.numel() != block.attn.n_head: + raise RuntimeError("QK-logit head inventory changed") + values.append(value) + return tuple(values) + + def _empty_diagnostic_interval( + self, + device: torch.device, + ) -> dict[str, torch.Tensor]: + zero = torch.zeros((), device=device, dtype=torch.float32) + return { + "steps": zero.clone(), + "head_observations": zero.clone(), + "active_heads": zero.clone(), + "sum_max_logit": zero.clone(), + "max_logit": torch.full( + (), + -float("inf"), + device=device, + dtype=torch.float32, + ), + "sum_gamma": zero.clone(), + "min_gamma": torch.ones( + (), + device=device, + dtype=torch.float32, + ), + } + + def _ensure_diagnostic_interval( + self, + device: torch.device, + ) -> dict[str, torch.Tensor]: + if self._diagnostic_interval_state is None: + self._diagnostic_interval_state = self._empty_diagnostic_interval( + device + ) + return self._diagnostic_interval_state + + @torch.no_grad() + def _apply_qk_clip(self) -> None: + observations = self._consume_qk_logits() + threshold = self.qk_clip_threshold + balance = self.qk_clip_balance + all_logits: list[torch.Tensor] = [] + all_gamma: list[torch.Tensor] = [] + + for block, max_logits in zip( + self.model.blocks, + observations, + strict=True, + ): + ones = torch.ones_like(max_logits) + gamma = torch.where( + max_logits > threshold, + threshold + / max_logits.clamp_min( + torch.finfo(max_logits.dtype).tiny + ), + ones, + ) + q_scale = gamma.pow(balance) + k_scale = gamma.pow(1.0 - balance) + head_width = block.attn.n_embd // block.attn.n_head + + q_weight = block.attn.q_proj.weight.view( + block.attn.n_head, + head_width, + -1, + ) + k_weight = block.attn.k_proj.weight.view( + block.attn.n_head, + head_width, + -1, + ) + q_weight.mul_(q_scale.to(q_weight)[:, None, None]) + k_weight.mul_(k_scale.to(k_weight)[:, None, None]) + if block.attn.q_proj.bias is not None: + block.attn.q_proj.bias.view( + block.attn.n_head, + head_width, + ).mul_(q_scale.to(block.attn.q_proj.bias)[:, None]) + if block.attn.k_proj.bias is not None: + block.attn.k_proj.bias.view( + block.attn.n_head, + head_width, + ).mul_(k_scale.to(block.attn.k_proj.bias)[:, None]) + + all_logits.append(max_logits) + all_gamma.append(gamma) + + logits = torch.cat(all_logits) + gamma = torch.cat(all_gamma) + interval = self._ensure_diagnostic_interval(logits.device) + interval["steps"].add_(1.0) + interval["head_observations"].add_(float(logits.numel())) + interval["active_heads"].add_((gamma < 1.0).float().sum()) + interval["sum_max_logit"].add_(logits.sum()) + interval["max_logit"].copy_( + torch.maximum(interval["max_logit"], logits.max()) + ) + interval["sum_gamma"].add_(gamma.sum()) + interval["min_gamma"].copy_( + torch.minimum(interval["min_gamma"], gamma.min()) + ) + + def _flush_diagnostics(self) -> None: + interval = self._diagnostic_interval_state + if interval is None: + return + values = { + key: float(value.detach().cpu()) + for key, value in interval.items() + } + observations = max(values["head_observations"], 1.0) + diagnostics = { + "step": float(self.step_index), + "threshold": float(self.qk_clip_threshold), + "steps_in_interval": values["steps"], + "head_observations": values["head_observations"], + "active_heads": values["active_heads"], + "active_fraction": values["active_heads"] / observations, + "mean_max_logit": values["sum_max_logit"] / observations, + "max_logit": values["max_logit"], + "mean_gamma": values["sum_gamma"] / observations, + "min_gamma": values["min_gamma"], + } + self.last_diagnostics = diagnostics + self._diagnostic_interval_state = None + self._write_diagnostics(diagnostics) + + def _write_diagnostics(self, values: dict[str, float]) -> None: + if self.diagnostics_path is None: + return + path = Path(self.diagnostics_path) + path.parent.mkdir(parents=True, exist_ok=True) + fields = [ + "step", + "threshold", + "steps_in_interval", + "head_observations", + "active_heads", + "active_fraction", + "mean_max_logit", + "max_logit", + "mean_gamma", + "min_gamma", + ] + write_header = not path.is_file() or path.stat().st_size == 0 + with path.open("a", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fields) + if write_header: + writer.writeheader() + writer.writerow(values) + + @torch.no_grad() + def step(self, closure=None): + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + lr = float(group["lr"]) + momentum = float(group["momentum"]) + rms_scale = float(group["update_rms_scale"]) + for parameter in group["params"]: + if parameter.grad is None: + continue + gradient = parameter.grad.detach() + if gradient.is_sparse: + raise RuntimeError("MuonClip does not support sparse gradients") + state = self.state[parameter] + buffer = state.get("momentum_buffer") + if buffer is None: + buffer = torch.zeros_like(gradient) + state["momentum_buffer"] = buffer + # Public Moonlight form: unnormalized momentum accumulation. + buffer.mul_(momentum).add_(gradient) + update_source = ( + gradient.add(buffer, alpha=momentum) + if bool(group["nesterov"]) + else buffer + ) + update = _zeropower( + update_source, + steps=int(group["newton_schulz_steps"]), + eps=float(group["eps"]), + ) + update.mul_( + rms_scale + * math.sqrt( + max(parameter.shape[0], parameter.shape[1]) + ) + ) + decay = float(group["weight_decay"]) + if decay: + parameter.mul_(max(0.0, 1.0 - lr * decay)) + parameter.add_(update, alpha=-lr) + + self.step_index += 1 + self._apply_qk_clip() + if self.step_index % self.diagnostics_interval == 0: + self._flush_diagnostics() + return loss + + def state_dict(self) -> dict[str, Any]: + payload = super().state_dict() + payload["muonclip_global_state"] = { + "step_index": int(self.step_index), + "last_diagnostics": dict(self.last_diagnostics), + "diagnostic_interval_state": self._diagnostic_interval_state, + } + return payload + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + payload = deepcopy(state_dict) + global_state = payload.pop("muonclip_global_state", {}) + super().load_state_dict(payload) + self.step_index = int(global_state.get("step_index", 0)) + self.last_diagnostics = dict( + global_state.get("last_diagnostics", {}) + ) + interval = global_state.get("diagnostic_interval_state") + if interval is None: + self._diagnostic_interval_state = None + else: + device = self.param_groups[0]["params"][0].device + self._diagnostic_interval_state = { + key: value.to(device) + if torch.is_tensor(value) + else torch.tensor(value, device=device, dtype=torch.float32) + for key, value in interval.items() + } + + +def _record_qk_logits(attention, scores: torch.Tensor) -> None: + if not attention.training: + return + value = ( + scores.detach() + .amax(dim=(-2, -1)) + .amax(dim=0) + .float() + ) + previous = getattr(attention, "_muonclip_max_logits", None) + setattr( + attention, + "_muonclip_max_logits", + value if previous is None else torch.maximum(previous, value), + ) + + +def _muonclip_attention_forward(attention, x: torch.Tensor) -> torch.Tensor: + """Track exact causal pre-softmax maxima while preserving native SDPA.""" + + batch, sequence, channels = x.shape + head_width = channels // attention.n_head + q = attention.q_proj(x).view( + batch, + sequence, + attention.n_head, + head_width, + ).transpose(1, 2) + k = attention.k_proj(x).view( + batch, + sequence, + attention.n_head, + head_width, + ).transpose(1, 2) + v = attention.v_proj(x).view( + batch, + sequence, + attention.n_head, + head_width, + ).transpose(1, 2) + dropout_p = attention.dropout if attention.training else 0.0 + + if q.device.type == "xla": + scores = (q @ k.transpose(-2, -1)) / math.sqrt(q.shape[-1]) + mask = attention.causal_mask[:, :, :sequence, :sequence] + scores = scores.masked_fill( + ~mask, + torch.finfo(scores.dtype).min, + ) + _record_qk_logits(attention, scores) + probabilities = F.softmax(scores, dim=-1) + if dropout_p: + probabilities = F.dropout( + probabilities, + p=dropout_p, + training=True, + ) + y = probabilities @ v + else: + if attention.training: + with torch.no_grad(): + scores = ( + q.detach() @ k.detach().transpose(-2, -1) + ) / math.sqrt(q.shape[-1]) + mask = attention.causal_mask[:, :, :sequence, :sequence] + scores = scores.masked_fill( + ~mask, + torch.finfo(scores.dtype).min, + ) + _record_qk_logits(attention, scores) + y = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=None, + dropout_p=dropout_p, + is_causal=True, + ) + + y = y.transpose(1, 2).contiguous().view( + batch, + sequence, + channels, + ) + return attention.resid_dropout(attention.out_proj(y)) + + +def _validate_muonclip_profile(profile: dict[str, Any]) -> None: + warmup = float(profile.get("warmup_fraction", -1.0)) + if not 0.0 <= warmup < 1.0: + raise ValueError("warmup_fraction must be in [0, 1)") + if str(profile.get("schedule")) != "warmup_cosine": + raise ValueError("MuonClip currently requires warmup_cosine") + peak = float(profile["learning_rate"]) + floor = float(profile["min_learning_rate"]) + if peak <= 0.0 or floor < 0.0 or floor > peak: + raise ValueError("MuonClip learning-rate peak/floor are inconsistent") + if not 0.0 <= float(profile["momentum"]) < 1.0: + raise ValueError("MuonClip momentum must be in [0, 1)") + if int(profile["newton_schulz_steps"]) < 1: + raise ValueError("MuonClip newton_schulz_steps must be positive") + if float(profile.get("muon_epsilon", 0.0)) <= 0.0: + raise ValueError("MuonClip muon_epsilon must be positive") + if float(profile["weight_decay"]) < 0.0: + raise ValueError("MuonClip weight_decay must be nonnegative") + if float(profile["update_rms_scale"]) <= 0.0: + raise ValueError("MuonClip update_rms_scale must be positive") + if float(profile["qk_clip_threshold"]) <= 0.0: + raise ValueError("MuonClip qk_clip_threshold must be positive") + balance = float(profile.get("qk_clip_balance", 0.5)) + if not 0.0 <= balance <= 1.0: + raise ValueError("MuonClip qk_clip_balance must be in [0, 1]") + + +def _partition(model): + named = [ + (name, parameter) + for name, parameter in model.named_parameters() + if parameter.requires_grad + ] + hidden = [ + parameter + for name, parameter in named + if name.startswith("blocks.") and parameter.ndim == 2 + ] + hidden_ids = {id(parameter) for parameter in hidden} + auxiliary = [ + (name, parameter) + for name, parameter in named + if id(parameter) not in hidden_ids + ] + if not hidden or not auxiliary: + raise ValueError( + "MuonClip partition must contain hidden matrices and auxiliary parameters" + ) + return hidden, auxiliary + + +def _decay_groups(named, weight_decay: float) -> list[dict[str, Any]]: + return [ + { + "params": [parameter for _, parameter in named if parameter.ndim >= 2], + "weight_decay": float(weight_decay), + }, + { + "params": [parameter for _, parameter in named if parameter.ndim < 2], + "weight_decay": 0.0, + }, + ] + + +def install_muonclip_extension() -> None: + global _INSTALLED + if _INSTALLED: + return + + from . import analysis as analysis_module + from . import config as config_module + from . import engine as engine_module + from . import model as model_module + from . import optimizers as optimizers_module + from . import training as training_module + from . import train_loop as train_loop_module + + all_optimizers = tuple( + dict.fromkeys((*config_module.SUPPORTED_OPTIMIZERS, "muon_clip")) + ) + config_module.SUPPORTED_OPTIMIZERS = all_optimizers + engine_module.SUPPORTED_OPTIMIZERS = all_optimizers + training_module.SUPPORTED_OPTIMIZERS = all_optimizers + analysis_module.OPTIMIZER_LABELS["muon_clip"] = ( + "MuonClip + auxiliary AdamW" + ) + analysis_module.OPTIMIZER_COLORS["muon_clip"] = "#CC79A7" + + original_validate = config_module.validate_optimizer_profile + original_make_handles = optimizers_module.make_optimizer_handles + original_zero_grad = optimizers_module.zero_grad + original_training_run_one = training_module.run_one + + def validate_optimizer_profile(profile: dict[str, Any]) -> None: + if str(profile.get("family", "")) == "muon_clip": + _validate_muonclip_profile(profile) + return + original_validate(profile) + + def make_optimizer_handles(model, profile: dict): + if str(profile.get("family", "")) != "muon_clip": + return original_make_handles(model, profile) + + hidden, auxiliary_named = _partition(model) + learning_rate = float(profile["learning_rate"]) + minimum = float(profile["min_learning_rate"]) + diagnostics_path = ( + None + if _CURRENT_RUN_DIR is None + else _CURRENT_RUN_DIR / "muonclip_qk.csv" + ) + primary = MuonClip( + hidden, + model=model, + lr=learning_rate, + momentum=float(profile["momentum"]), + nesterov=bool(profile.get("nesterov", True)), + weight_decay=float(profile["weight_decay"]), + newton_schulz_steps=int(profile["newton_schulz_steps"]), + eps=float(profile.get("muon_epsilon", 1e-7)), + update_rms_scale=float(profile["update_rms_scale"]), + qk_clip_threshold=float(profile["qk_clip_threshold"]), + qk_clip_balance=float(profile.get("qk_clip_balance", 0.5)), + diagnostics_interval=int( + profile.get("qk_diagnostics_interval", 250) + ), + diagnostics_path=diagnostics_path, + ) + auxiliary = torch.optim.AdamW( + _decay_groups(auxiliary_named, float(profile["weight_decay"])), + lr=learning_rate, + betas=(float(profile["beta1"]), float(profile["beta2"])), + eps=float(profile["epsilon"]), + ) + return [ + optimizers_module.OptimizerHandle( + role="primary", + optimizer=primary, + peak_lr=learning_rate, + min_lr=minimum, + ), + optimizers_module.OptimizerHandle( + role="auxiliary", + optimizer=auxiliary, + peak_lr=learning_rate, + min_lr=minimum, + ), + ] + + def zero_grad(handles) -> None: + for handle in handles: + if isinstance(handle.optimizer, MuonClip): + handle.optimizer.reset_qk_tracking() + original_zero_grad(handles) + + def run_one_with_context(*args, **kwargs): + global _CURRENT_RUN_DIR + optimizer_name = str(kwargs.get("optimizer_name", "")) + results_root = Path(kwargs["results_root"]) + seed = int(kwargs["seed"]) + _CURRENT_RUN_DIR = ( + results_root / optimizer_name / f"seed_{seed}" + ) + try: + return original_training_run_one(*args, **kwargs) + finally: + _CURRENT_RUN_DIR = None + + config_module.validate_optimizer_profile = validate_optimizer_profile + optimizers_module.make_optimizer_handles = make_optimizer_handles + engine_module.make_optimizer_handles = make_optimizer_handles + optimizers_module.zero_grad = zero_grad + train_loop_module.zero_grad = zero_grad + training_module.run_one = run_one_with_context + model_module.CausalSelfAttention.forward = _muonclip_attention_forward + + _INSTALLED = True + + +def main() -> None: + install_muonclip_extension() + from .training import main as training_main + + training_main() + + +if __name__ == "__main__": + main() From 868eb8d68060ada55e70a66976b4376259b5a68e Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 20:40:25 -0700 Subject: [PATCH 2/9] Add the one-epoch MuonClip reference protocol --- .../configs/muonclip_reference.yaml | 122 ++++++++++++++++++ 1 file changed, 122 insertions(+) create mode 100644 baseline/nanogpt_one_head/configs/muonclip_reference.yaml diff --git a/baseline/nanogpt_one_head/configs/muonclip_reference.yaml b/baseline/nanogpt_one_head/configs/muonclip_reference.yaml new file mode 100644 index 0000000..371bdc9 --- /dev/null +++ b/baseline/nanogpt_one_head/configs/muonclip_reference.yaml @@ -0,0 +1,122 @@ +protocol: + name: rg_nanogpt_one_head_muonclip_reference + version: 6 + description: Matched one-block, one-head nanoGPT MuonClip baseline with Kimi-K2 RMS matching, decoupled weight decay, and regular-MHA QK-Clip. + +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + train_tokens: 80000000 + val_tokens: 1000000 + test_tokens: 1000000 + +model: + vocab_size: 50257 + block_size: 256 + n_layer: 1 + n_head: 1 + n_embd: 128 + dropout: 0.0 + bias: false + tie_weights: true + +training: + seeds: [1337, 2027, 4099] + batch_size: 4 + grad_accum_steps: 8 + target_epochs: 1.0 + epoch_interval: 0.125 + eval_interval_steps: 250 + eval_batches: 64 + checkpoint_interval_steps: 250 + grad_clip: 1.0 + +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum + family: sgd + learning_rate: 0.05 + min_learning_rate: 0.005 + warmup_fraction: 0.10 + schedule: warmup_cosine + momentum: 0.90 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + + adamw: + display_name: AdamW + family: adamw + learning_rate: 0.0006 + min_learning_rate: 0.00006 + warmup_fraction: 0.01 + schedule: warmup_cosine + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + weight_decay: 0.10 + + muon: + display_name: Muon + auxiliary AdamW + family: muon + matrix_learning_rate: 0.02 + matrix_min_learning_rate: 0.002 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 0.00003 + warmup_fraction: 0.05 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-7 + matrix_weight_decay: 0.01 + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + aux_weight_decay: 0.01 + + muon_clip: + display_name: MuonClip + auxiliary AdamW + family: muon_clip + learning_rate: 0.0002 + min_learning_rate: 0.00002 + warmup_fraction: 0.0512 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-7 + weight_decay: 0.10 + update_rms_scale: 0.20 + qk_clip_threshold: 100.0 + qk_clip_balance: 0.50 + qk_diagnostics_interval: 250 + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 64 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 + +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + +runtime: + matmul_precision: high + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true From 81d9a63409369e172b55e747f91be15e40364d1f Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 20:40:47 -0700 Subject: [PATCH 3/9] Add the ten-epoch MuonClip spectral-relaxation protocol --- .../configs/muonclip_10epochs.yaml | 126 ++++++++++++++++++ 1 file changed, 126 insertions(+) create mode 100644 baseline/nanogpt_one_head/configs/muonclip_10epochs.yaml diff --git a/baseline/nanogpt_one_head/configs/muonclip_10epochs.yaml b/baseline/nanogpt_one_head/configs/muonclip_10epochs.yaml new file mode 100644 index 0000000..e11b668 --- /dev/null +++ b/baseline/nanogpt_one_head/configs/muonclip_10epochs.yaml @@ -0,0 +1,126 @@ +protocol: + name: rg_nanogpt_one_head_muonclip_10epochs + version: 6 + description: Ten-corpus-equivalent MuonClip spectral-relaxation run using the matched one-epoch warmup/cosine LR horizon and holding the LR floor thereafter. + +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + train_tokens: 80000000 + val_tokens: 1000000 + test_tokens: 1000000 + +model: + vocab_size: 50257 + block_size: 256 + n_layer: 1 + n_head: 1 + n_embd: 128 + dropout: 0.0 + bias: false + tie_weights: true + +training: + seeds: [1337] + batch_size: 4 + grad_accum_steps: 8 + target_epochs: 10.0 + epoch_interval: 0.25 + eval_interval_steps: 250 + eval_batches: 64 + checkpoint_interval_steps: 250 + grad_clip: 1.0 + +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum + family: sgd + learning_rate: 0.05 + min_learning_rate: 0.005 + warmup_fraction: 0.10 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + momentum: 0.90 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + + adamw: + display_name: AdamW + family: adamw + learning_rate: 0.0006 + min_learning_rate: 0.00006 + warmup_fraction: 0.01 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + weight_decay: 0.10 + + muon: + display_name: Muon + auxiliary AdamW + family: muon + matrix_learning_rate: 0.02 + matrix_min_learning_rate: 0.002 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 0.00003 + warmup_fraction: 0.05 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-7 + matrix_weight_decay: 0.01 + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + aux_weight_decay: 0.01 + + muon_clip: + display_name: MuonClip + auxiliary AdamW + family: muon_clip + learning_rate: 0.0002 + min_learning_rate: 0.00002 + warmup_fraction: 0.0512 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-7 + weight_decay: 0.10 + update_rms_scale: 0.20 + qk_clip_threshold: 100.0 + qk_clip_balance: 0.50 + qk_diagnostics_interval: 250 + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 64 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 + +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + +runtime: + matmul_precision: high + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true From 84eb9974578e28ac2d2f5e1b0ba6af9548191a54 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 20:41:09 -0700 Subject: [PATCH 4/9] Document the MuonClip nanoGPT baseline --- baseline/nanogpt_one_head/MUONCLIP.md | 172 ++++++++++++++++++++++++++ 1 file changed, 172 insertions(+) create mode 100644 baseline/nanogpt_one_head/MUONCLIP.md diff --git a/baseline/nanogpt_one_head/MUONCLIP.md b/baseline/nanogpt_one_head/MUONCLIP.md new file mode 100644 index 0000000..0d63195 --- /dev/null +++ b/baseline/nanogpt_one_head/MUONCLIP.md @@ -0,0 +1,172 @@ +# MuonClip nanoGPT baseline + +This baseline adds Moonshot AI's **MuonClip** optimizer to the existing +one-block, one-head FineWeb-Edu experiment without changing the historical +`rg-onehead-train` command or the committed SGD, AdamW, and ordinary-Muon +protocols. + +MuonClip combines: + +1. Muon momentum followed by Newton--Schulz orthogonalization; +2. decoupled weight decay; +3. RMS-matched matrix updates; +4. per-head QK-Clip. + +For a hidden matrix `W` with shape `n x m`, this implementation applies + +```text +M_t = momentum * M_(t-1) + G_t +O_t = NewtonSchulz(M_t) * 0.2 * sqrt(max(n, m)) +W_t = W_(t-1) - lr * (O_t + weight_decay * W_(t-1)) +``` + +For each attention head, the maximum finite causal pre-softmax logit is +measured over every gradient-accumulation micro-batch. With threshold `tau`, + +```text +gamma_h = min(1, tau / max_logit_h) +``` + +and regular multi-head attention uses the balanced scaling + +```text +W_Q[h] *= gamma_h ** 0.5 +W_K[h] *= gamma_h ** 0.5 +``` + +The committed reference value is `tau = 100`, matching the Kimi K2 report. +The compact nanoGPT baseline has only one head, so it is possible that QK-Clip +never activates. That is a valid experimental result and is recorded explicitly +rather than inferred from the loss curve. + +## Why a dedicated launcher? + +MuonClip is opt-in. The historical package and result tables retain exactly the +three original optimizer arms. The dedicated launcher installs the MuonClip +extension in the current process and then delegates to the same data, +checkpoint, evaluation, WeightWatcher, MPS/CUDA/TPU, and long-horizon code. + +```bash +rg-onehead-muonclip --help +``` + +## One-epoch reference run + +```bash +cd baseline/nanogpt_one_head +python -m pip install -e . + +rg-onehead-muonclip \ + --config configs/muonclip_reference.yaml \ + --optimizer muon_clip \ + --seeds 1337 \ + --data-root /tmp/rg-nanogpt-one-head/data \ + --results-root /tmp/rg-nanogpt-muonclip/results \ + --device auto \ + --no-resume +``` + +The reference profile uses: + +```text +base LR: 2e-4 +minimum LR: 2e-5 +warmup: 500 of 9,766 steps +weight decay: 0.1 +momentum: 0.95 +Newton--Schulz steps: 5 +RMS scale: 0.2 * sqrt(max(n, m)) +QK-Clip threshold: 100 +Q/K balance: 0.5 / 0.5 +auxiliary optimizer: AdamW with the same base LR and weight decay +``` + +The Kimi K2 production recipe used a much larger model, a 15.5-trillion-token +WSD schedule, and distributed bfloat16 training. This compact baseline preserves +the repository's matched warmup-plus-cosine protocol so optimizer comparisons +remain interpretable; it is not presented as a literal reproduction of K2 +pretraining. + +## Ten-epoch spectral-relaxation run + +```bash +rg-onehead-muonclip \ + --config configs/muonclip_10epochs.yaml \ + --optimizer muon_clip \ + --seeds 1337 \ + --data-root /tmp/rg-nanogpt-one-head/data \ + --results-root /tmp/rg-nanogpt-muonclip-long/results \ + --device auto \ + --no-resume +``` + +This follows the same long-horizon convention as ordinary Muon: + +```text +training steps: 97,657 +LR schedule steps: 9,766 +warmup steps: 500 +post-epoch-1 LR: 2e-5 +``` + +It is a long-time spectral-relaxation experiment, not the K2 WSD schedule. + +## Diagnostics + +The usual files remain unchanged: + +```text +metrics.csv +spectral/layers.csv +spectral/summary.csv +checkpoint_latest.pt +checkpoint_best.pt +checkpoint_final.pt +``` + +MuonClip additionally writes: + +```text +muonclip_qk.csv +``` + +at the configured diagnostic interval. Its columns are: + +```text +step +threshold +steps_in_interval +head_observations +active_heads +active_fraction +mean_max_logit +max_logit +mean_gamma +min_gamma +``` + +This file is monitoring-only. It does not select checkpoints or tune the +threshold. The existing WeightWatcher monitor can be used for alpha, fit `D`, +`rand_distance`, ERG gap, and traps: + +```bash +rg-onehead-monitor \ + --results-root /tmp/rg-nanogpt-muonclip/results \ + --optimizer muon_clip \ + --seed 1337 +``` + +## TPU + +The same launcher uses the automatic accelerator and persistent-storage logic: + +```bash +rg-onehead-muonclip \ + --config configs/muonclip_reference.yaml \ + --optimizer muon_clip \ + --device auto \ + --no-resume +``` + +Run the existing TPU smoke test first. The MuonClip-specific path should then be +qualified with a short one-epoch seed before starting the ten-epoch run. From ad71a141861b78f36f93e0701446adf59abbdebe Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 20:41:36 -0700 Subject: [PATCH 5/9] Test MuonClip RMS matching, QK-Clip, and config installation --- .../nanogpt_one_head/tests/test_muonclip.py | 207 ++++++++++++++++++ 1 file changed, 207 insertions(+) create mode 100644 baseline/nanogpt_one_head/tests/test_muonclip.py diff --git a/baseline/nanogpt_one_head/tests/test_muonclip.py b/baseline/nanogpt_one_head/tests/test_muonclip.py new file mode 100644 index 0000000..f892da2 --- /dev/null +++ b/baseline/nanogpt_one_head/tests/test_muonclip.py @@ -0,0 +1,207 @@ +from __future__ import annotations + +import math +from pathlib import Path +import subprocess +import sys + +import pytest +import torch +import yaml + +EXPERIMENT_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(EXPERIMENT_ROOT / "src")) + +from rg_nanogpt_one_head.config import SUPPORTED_OPTIMIZERS +from rg_nanogpt_one_head.model import GPT, GPTConfig +from rg_nanogpt_one_head.muonclip import ( + MuonClip, + _muonclip_attention_forward, +) +from rg_nanogpt_one_head.optimizers import zeropower_via_newton_schulz_5 + + +def small_model() -> GPT: + return GPT( + GPTConfig( + vocab_size=64, + block_size=8, + n_layer=1, + n_head=1, + n_embd=16, + dropout=0.0, + bias=False, + ) + ) + + +def hidden_matrices(model: GPT) -> list[torch.nn.Parameter]: + return [ + parameter + for name, parameter in model.named_parameters() + if name.startswith("blocks.") and parameter.ndim == 2 + ] + + +def test_historical_launcher_remains_three_optimizer_reference() -> None: + # Importing the class does not mutate the historical launcher. The extension + # is installed only by rg-onehead-muonclip. + assert SUPPORTED_OPTIMIZERS == ( + "sgd_momentum", + "adamw", + "muon", + ) + + +def test_muonclip_configs_use_reported_scaling_and_threshold() -> None: + reference = yaml.safe_load( + (EXPERIMENT_ROOT / "configs" / "muonclip_reference.yaml").read_text() + ) + profile = reference["optimizer_profiles"]["muon_clip"] + + assert profile["learning_rate"] == pytest.approx(2e-4) + assert profile["min_learning_rate"] == pytest.approx(2e-5) + assert profile["weight_decay"] == pytest.approx(0.1) + assert profile["update_rms_scale"] == pytest.approx(0.2) + assert profile["qk_clip_threshold"] == pytest.approx(100.0) + assert profile["qk_clip_balance"] == pytest.approx(0.5) + assert round(9766 * profile["warmup_fraction"]) == 500 + + long_cfg = yaml.safe_load( + (EXPERIMENT_ROOT / "configs" / "muonclip_10epochs.yaml").read_text() + ) + long_profile = long_cfg["optimizer_profiles"]["muon_clip"] + assert long_cfg["training"]["target_epochs"] == 10.0 + assert long_profile["lr_schedule_epochs"] == 1.0 + + +def test_rms_matched_update_uses_point_two_sqrt_max_dimension() -> None: + parameter = torch.nn.Parameter(torch.zeros(4, 2)) + gradient = torch.tensor( + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]] + ) + parameter.grad = gradient.clone() + + class DummyModel: + blocks = [] + + optimizer = MuonClip( + [parameter], + model=DummyModel(), + lr=1.0, + momentum=0.0, + nesterov=False, + weight_decay=0.0, + update_rms_scale=0.2, + qk_clip_threshold=100.0, + qk_clip_balance=0.5, + ) + optimizer._apply_qk_clip = lambda: None + expected = -zeropower_via_newton_schulz_5( + gradient, + steps=5, + eps=1e-7, + ) * (0.2 * math.sqrt(4)) + + optimizer.step() + + assert torch.allclose(parameter, expected) + + +def test_qk_clip_balances_query_and_key_scaling() -> None: + model = small_model() + optimizer = MuonClip( + hidden_matrices(model), + model=model, + lr=0.0, + momentum=0.0, + nesterov=False, + weight_decay=0.0, + update_rms_scale=0.2, + qk_clip_threshold=100.0, + qk_clip_balance=0.5, + diagnostics_interval=1, + ) + q_before = model.blocks[0].attn.q_proj.weight.detach().clone() + k_before = model.blocks[0].attn.k_proj.weight.detach().clone() + model.blocks[0].attn._muonclip_max_logits = torch.tensor([400.0]) + for parameter in hidden_matrices(model): + parameter.grad = torch.zeros_like(parameter) + + optimizer.step() + + assert torch.allclose( + model.blocks[0].attn.q_proj.weight, + 0.5 * q_before, + ) + assert torch.allclose( + model.blocks[0].attn.k_proj.weight, + 0.5 * k_before, + ) + assert optimizer.last_diagnostics["active_fraction"] == pytest.approx(1.0) + assert optimizer.last_diagnostics["min_gamma"] == pytest.approx(0.25) + + +def test_qk_clip_is_noop_below_threshold() -> None: + model = small_model() + optimizer = MuonClip( + hidden_matrices(model), + model=model, + lr=0.0, + momentum=0.0, + nesterov=False, + weight_decay=0.0, + update_rms_scale=0.2, + qk_clip_threshold=100.0, + qk_clip_balance=0.5, + diagnostics_interval=1, + ) + q_before = model.blocks[0].attn.q_proj.weight.detach().clone() + k_before = model.blocks[0].attn.k_proj.weight.detach().clone() + model.blocks[0].attn._muonclip_max_logits = torch.tensor([20.0]) + for parameter in hidden_matrices(model): + parameter.grad = torch.zeros_like(parameter) + + optimizer.step() + + assert torch.equal(model.blocks[0].attn.q_proj.weight, q_before) + assert torch.equal(model.blocks[0].attn.k_proj.weight, k_before) + assert optimizer.last_diagnostics["active_fraction"] == pytest.approx(0.0) + assert optimizer.last_diagnostics["min_gamma"] == pytest.approx(1.0) + + +def test_attention_observation_matches_native_sdpa_output() -> None: + model = small_model() + attention = model.blocks[0].attn + attention.train() + x = torch.randn(2, 8, 16) + + observed = _muonclip_attention_forward(attention, x) + tracked = attention._muonclip_max_logits + attention._muonclip_max_logits = None + expected = attention.__class__.__dict__["forward"](attention, x) + + assert tracked is not None + assert tracked.shape == (1,) + assert torch.allclose(observed, expected, atol=1e-6, rtol=1e-5) + + +def test_extension_validates_muonclip_config_in_isolated_process() -> None: + code = """ +from pathlib import Path +from rg_nanogpt_one_head.muonclip import install_muonclip_extension +install_muonclip_extension() +from rg_nanogpt_one_head.config import load_config, optimizer_profile, max_steps, lr_schedule_steps, warmup_steps +root = Path.cwd() +cfg = load_config(root / 'configs' / 'muonclip_reference.yaml') +p = optimizer_profile(cfg, 'muon_clip') +print(max_steps(cfg), lr_schedule_steps(cfg, p), warmup_steps(p, lr_schedule_steps(cfg, p))) +""" + completed = subprocess.run( + [sys.executable, "-c", code], + cwd=EXPERIMENT_ROOT, + check=True, + capture_output=True, + text=True, + ) + assert completed.stdout.strip().endswith("9766 9766 500") From b1592cd509400e9535693d14e2510d81cb422829 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 20:42:06 -0700 Subject: [PATCH 6/9] Expose the MuonClip nanoGPT launcher --- baseline/nanogpt_one_head/pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/baseline/nanogpt_one_head/pyproject.toml b/baseline/nanogpt_one_head/pyproject.toml index 4723343..ff7d56e 100644 --- a/baseline/nanogpt_one_head/pyproject.toml +++ b/baseline/nanogpt_one_head/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "rg-nanogpt-one-head" -version = "0.2.1" +version = "0.3.0" description = "Matched one-head nanoGPT optimizer baselines on pinned FineWeb-Edu" requires-python = ">=3.10" dependencies = [ @@ -35,6 +35,7 @@ rg-onehead-prepare = "rg_nanogpt_one_head.data:main" rg-onehead-train = "rg_nanogpt_one_head.training:main" rg-onehead-env = "rg_nanogpt_one_head.runtime:main" rg-onehead-monitor = "rg_nanogpt_one_head.monitor:main" +rg-onehead-muonclip = "rg_nanogpt_one_head.muonclip:main" [tool.setuptools.packages.find] where = ["src"] From cbb9a414df8342f6b3495e327ea2a81bea421879 Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 20:44:23 -0700 Subject: [PATCH 7/9] Add a tiny end-to-end MuonClip training integration test --- .../tests/test_muonclip_integration.py | 131 ++++++++++++++++++ 1 file changed, 131 insertions(+) create mode 100644 baseline/nanogpt_one_head/tests/test_muonclip_integration.py diff --git a/baseline/nanogpt_one_head/tests/test_muonclip_integration.py b/baseline/nanogpt_one_head/tests/test_muonclip_integration.py new file mode 100644 index 0000000..a9539e1 --- /dev/null +++ b/baseline/nanogpt_one_head/tests/test_muonclip_integration.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +from pathlib import Path +import subprocess +import sys + + +EXPERIMENT_ROOT = Path(__file__).resolve().parents[1] + + +def test_tiny_muonclip_training_writes_qk_diagnostics(tmp_path) -> None: + code = r''' +from copy import deepcopy +import hashlib +import json +from pathlib import Path +import sys + +import numpy as np +import pandas as pd + +from rg_nanogpt_one_head.muonclip import install_muonclip_extension +install_muonclip_extension() + +from rg_nanogpt_one_head.config import load_config +from rg_nanogpt_one_head.data import TOKEN_DTYPE +import rg_nanogpt_one_head.train_loop as train_loop +import rg_nanogpt_one_head.run_utils as run_utils +from rg_nanogpt_one_head.training import run_one + +root = Path(sys.argv[1]) +cfg = deepcopy(load_config(Path.cwd() / 'configs' / 'muonclip_reference.yaml')) +cfg['dataset'].update({ + 'name': 'unit/fineweb', + 'config': 'unit', + 'revision': 'unit-revision', + 'train_tokens': 2048, + 'val_tokens': 512, + 'test_tokens': 512, +}) +cfg['model'].update({ + 'vocab_size': 64, + 'block_size': 8, + 'n_layer': 1, + 'n_head': 1, + 'n_embd': 16, +}) +cfg['training'].update({ + 'seeds': [13], + 'batch_size': 2, + 'grad_accum_steps': 2, + 'target_epochs': 0.04, + 'epoch_interval': 1.0, + 'eval_interval_steps': 1, + 'eval_batches': 1, + 'checkpoint_interval_steps': 1, +}) +cfg['evaluation'].update({ + 'bleu_examples': 2, + 'bleu_prompt_tokens': 3, + 'bleu_continuation_tokens': 2, + 'bleu_batch_size': 2, +}) +cfg['optimizer_profiles']['muon_clip']['qk_diagnostics_interval'] = 1 + +data_root = root / 'data' +results_root = root / 'results' +data_root.mkdir(parents=True) +rng = np.random.default_rng(7) +splits = { + 'train': int(cfg['dataset']['train_tokens']), + 'val': int(cfg['dataset']['val_tokens']), + 'test': int(cfg['dataset']['test_tokens']), +} +files = {} +for split, size in splits.items(): + path = data_root / f'{split}.bin' + rng.integers(0, cfg['model']['vocab_size'], size=size, dtype=np.uint16).tofile(path) + files[split] = { + 'path': path.name, + 'sha256': hashlib.sha256(path.read_bytes()).hexdigest(), + 'bytes': path.stat().st_size, + } +(data_root / 'meta.json').write_text(json.dumps({ + 'schema_version': 2, + 'tokenizer': 'gpt2', + 'vocab_size': cfg['model']['vocab_size'], + 'dtype': TOKEN_DTYPE.name, + 'splits': splits, + 'document_disjoint_splits': True, + 'dataset_name': cfg['dataset']['name'], + 'dataset_config': cfg['dataset']['config'], + 'dataset_revision': cfg['dataset']['revision'], + 'files': files, +}), encoding='utf-8') + +train_loop.evaluate_bleu = lambda *args, **kwargs: {'bleu': 0.0} +run_utils.evaluate_bleu = lambda *args, **kwargs: {'bleu': 0.0} +train_loop.run_weightwatcher = lambda *args, **kwargs: { + 'alpha_median': 2.0, + 'rand_distance_median': 0.1, + 'ERG_gap_median': 0.0, + 'num_traps_mean': 0.0, +} + +run_dir = run_one( + cfg=cfg, + data_root=data_root, + results_root=results_root, + optimizer_name='muon_clip', + seed=13, + device='cpu', + resume=True, + progress=False, +) +assert (run_dir / 'run_complete.json').is_file() +assert (run_dir / 'checkpoint_final.pt').is_file() +qk = pd.read_csv(run_dir / 'muonclip_qk.csv') +assert len(qk) >= 1 +assert qk['head_observations'].gt(0).all() +assert qk['max_logit'].notna().all() +print(run_dir) +''' + completed = subprocess.run( + [sys.executable, "-c", code, str(tmp_path)], + cwd=EXPERIMENT_ROOT, + check=True, + capture_output=True, + text=True, + ) + assert "muon_clip/seed_13" in completed.stdout From 8288dd0247799a08aeea65a17578daa22526c67e Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 20:45:47 -0700 Subject: [PATCH 8/9] Match the Kimi K2 MuonClip momentum update without Nesterov --- baseline/nanogpt_one_head/configs/muonclip_reference.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/baseline/nanogpt_one_head/configs/muonclip_reference.yaml b/baseline/nanogpt_one_head/configs/muonclip_reference.yaml index 371bdc9..92b4359 100644 --- a/baseline/nanogpt_one_head/configs/muonclip_reference.yaml +++ b/baseline/nanogpt_one_head/configs/muonclip_reference.yaml @@ -86,7 +86,7 @@ optimizer_profiles: warmup_fraction: 0.0512 schedule: warmup_cosine momentum: 0.95 - nesterov: true + nesterov: false newton_schulz_steps: 5 muon_epsilon: 1.0e-7 weight_decay: 0.10 From 1d0fa7649594d54fcacfa1abd743a304e6b4f45a Mon Sep 17 00:00:00 2001 From: Charles Martin Date: Tue, 11 Aug 2026 20:46:09 -0700 Subject: [PATCH 9/9] Match the long MuonClip run to the Kimi K2 momentum update --- baseline/nanogpt_one_head/configs/muonclip_10epochs.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/baseline/nanogpt_one_head/configs/muonclip_10epochs.yaml b/baseline/nanogpt_one_head/configs/muonclip_10epochs.yaml index e11b668..f44077d 100644 --- a/baseline/nanogpt_one_head/configs/muonclip_10epochs.yaml +++ b/baseline/nanogpt_one_head/configs/muonclip_10epochs.yaml @@ -90,7 +90,7 @@ optimizer_profiles: lr_schedule_epochs: 1.0 schedule: warmup_cosine momentum: 0.95 - nesterov: true + nesterov: false newton_schulz_steps: 5 muon_epsilon: 1.0e-7 weight_decay: 0.10