1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 191,
5  "end_line": 201,
6  "content": "def get_lr(it, warmup_iters, lr_decay_iters, learning_rate, min_lr):\n    \"\"\"Cosine learning rate schedule with linear warmup.\"\"\"\n    if it < warmup_iters:\n        return learning_rate * (it + 1) / (warmup_iters + 1)\n    if it > lr_decay_iters:\n        return min_lr\n    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)\n    assert 0 <= decay_ratio <= 1\n    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))\n    return min_lr + coeff * (learning_rate - min_lr)"
7}

1{
2  "content": "    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):\n        param_dict = {pn: p for pn, p in self.named_parameters() if p.requires_grad}\n        muon_p = [p for n, p in param_dict.items() if p.dim() >= 2 and 'wte' not in n and 'wpe' not in n and 'lm_head' not in n]\n        adam_decay = [p for n, p in param_dict.items() if p.dim() >= 2 and ('wte' in n or 'wpe' in n or 'lm_head' in n)]\n        adam_nodecay = [p for n, p in param_dict.items() if p.dim() < 2]\n        class Muon(torch.optim.Optimizer):\n            def __init__(self, params, lr=0.02, momentum=0.95, ns_steps=5, wd=0.0):\n                super().__init__(params, dict(lr=lr, momentum=momentum, ns_steps=ns_steps, wd=wd))\n            @staticmethod\n            def _ns(G, steps=5):\n                X = G.bfloat16() / (G.norm() + 1e-7)\n                if G.size(0) > G.size(1): X = X.T\n                a, b, c = 3.4445, -4.7750, 2.0315\n                for _ in range(steps):\n                    A = X @ X.T; X = a*X + b*(A@X) + c*(A@(A@X))\n                if G.size(0) > G.size(1): X = X.T\n                return X\n            @torch.no_grad()\n            def step(self):\n                for g in self.param_groups:\n                    for p in g['params']:\n                        if p.grad is None: continue\n                        if g['wd'] > 0: p.mul_(1 - g['lr'] * g['wd'])\n                        s = self.state.setdefault(p, {'b': torch.zeros_like(p.grad)})['b']\n                        s.lerp_(p.grad, 1 - g['momentum'])\n                        n = p.grad.lerp(s, g['momentum'])\n                        if n.dim() == 2:\n                            orig = n.shape\n                            if orig[0] == 3 * orig[1]:\n                                parts = n.split(orig[1])\n                                U = torch.cat([self._ns(part, g['ns_steps']) for part in parts])\n                                scale = max(1, orig[0] // orig[1]) ** 0.5\n                            else:\n                                U = self._ns(n, g['ns_steps'])\n                                scale = max(1, orig[0] / orig[1]) ** 0.5\n                            p.data.add_(U.to(p.dtype), alpha=-g['lr'] * scale)\n                        else:\n                            p.data.add_(n, alpha=-g['lr'])\n        fused = 'fused' in inspect.signature(torch.optim.AdamW).parameters\n        muon_opt = Muon([{'params': muon_p, 'lr_scale': 0.02 / learning_rate}], lr=0.02, wd=weight_decay)\n        adam_opt = torch.optim.AdamW(\n            [{'params': adam_decay, 'weight_decay': weight_decay},\n             {'params': adam_nodecay, 'weight_decay': 0.0}],\n            lr=learning_rate, betas=betas, fused=fused)\n        class CombinedOpt:\n            def __init__(self, opts):\n                self.optimizers = opts\n                self.param_groups = sum([o.param_groups for o in opts], [])\n            def zero_grad(self, set_to_none=True):\n                for o in self.optimizers: o.zero_grad(set_to_none=set_to_none)\n            def step(self):\n                for o in self.optimizers: o.step()\n        print(f\"using Muon (lr=0.02, wd={weight_decay}) + AdamW\")\n        return CombinedOpt([muon_opt, adam_opt])",
3  "end_line": 189,
4  "start_line": 171,
5  "filename": "nanoGPT/custom_pretrain.py"
6}

1{
2  "op": "replace",
3  "content": "    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):\n        param_dict = {pn: p for pn, p in self.named_parameters() if p.requires_grad}\n        muon_p = [p for n, p in param_dict.items() if p.dim() >= 2 and 'wte' not in n and 'wpe' not in n and 'lm_head' not in n]\n        adam_decay = [p for n, p in param_dict.items() if p.dim() >= 2 and ('wte' in n or 'wpe' in n or 'lm_head' in n)]\n        adam_nodecay = [p for n, p in param_dict.items() if p.dim() < 2]\n        class Muon(torch.optim.Optimizer):\n            def __init__(self, params, lr=0.02, momentum=0.95, ns_steps=5, wd=0.0):\n                super().__init__(params, dict(lr=lr, momentum=momentum, ns_steps=ns_steps, wd=wd))\n            @staticmethod\n            def _ns(G, steps=5):\n                X = G.bfloat16() / (G.norm() + 1e-7)\n                if G.size(0) > G.size(1): X = X.T\n                a, b, c = 3.4445, -4.7750, 2.0315\n                for _ in range(steps):\n                    A = X @ X.T; X = a*X + b*(A@X) + c*(A@(A@X))\n                if G.size(0) > G.size(1): X = X.T\n                return X\n            @torch.no_grad()\n            def step(self):\n                for g in self.param_groups:\n                    for p in g['params']:\n                        if p.grad is None: continue\n                        if g['wd'] > 0: p.mul_(1 - g['lr'] * g['wd'])\n                        s = self.state.setdefault(p, {'b': torch.zeros_like(p.grad)})['b']\n                        s.lerp_(p.grad, 1 - g['momentum'])\n                        n = p.grad.lerp(s, g['momentum'])\n                        if n.dim() == 2:\n                            orig = n.shape\n                            if orig[0] == 3 * orig[1]:\n                                parts = n.split(orig[1])\n                                U = torch.cat([self._ns(part, g['ns_steps']) for part in parts])\n                                scale = max(1, orig[0] // orig[1]) ** 0.5\n                            else:\n                                U = self._ns(n, g['ns_steps'])\n                                scale = max(1, orig[0] / orig[1]) ** 0.5\n                            p.data.add_(U.to(p.dtype), alpha=-g['lr'] * scale)\n                        else:\n                            p.data.add_(n, alpha=-g['lr'])\n        fused = 'fused' in inspect.signature(torch.optim.AdamW).parameters\n        muon_opt = Muon([{'params': muon_p, 'lr_scale': 0.02 / learning_rate}], lr=0.02, wd=weight_decay)\n        adam_opt = torch.optim.AdamW(\n            [{'params': adam_decay, 'weight_decay': weight_decay},\n             {'params': adam_nodecay, 'weight_decay': 0.0}],\n            lr=learning_rate, betas=betas, fused=fused)\n        class CombinedOpt:\n            def __init__(self, opts):\n                self.optimizers = opts\n                self.param_groups = sum([o.param_groups for o in opts], [])\n            def zero_grad(self, set_to_none=True):\n                for o in self.optimizers: o.zero_grad(set_to_none=set_to_none)\n            def step(self):\n                for o in self.optimizers: o.step()\n        print(f\"using Muon (lr=0.02, wd={weight_decay}) + AdamW\")\n        return CombinedOpt([muon_opt, adam_opt])\n",
4  "filename": "nanoGPT/custom_pretrain.py",
5  "start_line": 171,
6  "end_line": 189
7}

1{
2  "content": "    # CONFIG_OVERRIDES: override training hyperparameters for your method.\n    # Allowed keys: learning_rate, weight_decay, warmup_iters, min_lr, grad_clip.\n    CONFIG_OVERRIDES = {'learning_rate': 3e-4, 'warmup_iters': int(max_iters * 0.10)}",
3  "end_line": 282,
4  "filename": "nanoGPT/custom_pretrain.py",
5  "op": "replace",
6  "start_line": 280
7}

1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 171,
5  "end_line": 224,
6  "content": "    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):\n        param_dict = {pn: p for pn, p in self.named_parameters() if p.requires_grad}\n        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]\n        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]\n        num_decay_params = sum(p.numel() for p in decay_params)\n        num_nodecay_params = sum(p.numel() for p in nodecay_params)\n        print(f\"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters\")\n        print(f\"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters\")\n        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters\n        use_fused = fused_available and device_type == 'cuda'\n\n        class AdEMAMix(torch.optim.Optimizer):\n            \"\"\"AdEMAMix optimizer (arXiv:2409.03137).\n            Maintains two momentum buffers with different decay rates to capture\n            both short-term and long-term gradient structure. The long-term buffer\n            uses a very slow decay (e.g., 0.9999) to track the average gradient\n            direction, while the short-term buffer (e.g., 0.9) captures recent changes.\n            The update blends both, giving the best of both worlds.\n            \"\"\"\n            def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), beta3=0.9999,\n                         alpha=5.0, weight_decay=0.0, fused=True, eps=1e-8):\n                defaults = dict(lr=lr, betas=betas, beta3=beta3, alpha=alpha,\n                               weight_decay=weight_decay, eps=eps, fused=fused)\n                super().__init__(params, defaults)\n            @torch.no_grad()\n            def step(self):\n                for group in self.param_groups:\n                    beta1, beta2 = group['betas']\n                    beta3 = group['beta3']\n                    alpha = group['alpha']\n                    wd = group['weight_decay']\n                    eps = group['eps']\n                    lr = group['lr']\n                    for p in group['params']:\n                        if p.grad is None: continue\n                        grad = p.grad\n                        state = self.state[p]\n                        if len(state) == 0:\n                            state['step'] = 0\n                            state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)\n                            state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)\n                            state['exp_avg_slow'] = torch.zeros_like(p, memory_format=torch.preserve_format)\n                        step = state['step'] + 1\n                        state['step'] = step\n                        # Decoupled weight decay\n                        if wd > 0:\n                            p.mul_(1 - lr * wd)\n                        exp_avg = state['exp_avg']\n                        exp_avg_sq = state['exp_avg_sq']\n                        exp_avg_slow = state['exp_avg_slow']\n                        # Update momentums\n                        exp_avg.lerp_(grad, 1 - beta1)\n                        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)\n                        exp_avg_slow.lerp_(grad, 1 - beta3)\n                        # Bias correction\n                        bias1 = 1 - beta1 ** step\n                        bias2 = 1 - beta2 ** step\n                        bias3 = 1 - beta3 ** step\n                        # Combined momentum: short-term + alpha * long-term\n                        combined = exp_avg / bias1 + alpha * exp_avg_slow / bias3\n                        # RMS\n                        denom = exp_avg_sq.sqrt() / math.sqrt(bias2) + eps\n                        # Update\n                        p.addcdiv_(combined, denom, value=-lr)\n\n        optim_groups = [\n            {'params': decay_params, 'weight_decay': weight_decay},\n            {'params': nodecay_params, 'weight_decay': 0.0},\n        ]\n        optimizer = AdEMAMix(optim_groups, lr=learning_rate, betas=(beta1, beta2),\n                            beta3=0.9999, alpha=5.0, weight_decay=weight_decay,\n                            fused=use_fused, eps=1e-8)\n        print(f\"using AdEMAMix (beta3=0.9999, alpha=5.0)\")\n        return optimizer"
7}

1{
2  "content": "    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):\n        param_dict = {pn: p for pn, p in self.named_parameters() if p.requires_grad}\n        muon_p = [(n, p) for n, p in param_dict.items() if p.dim() >= 2 and 'wte' not in n and 'wpe' not in n and 'lm_head' not in n]\n        adam_decay = [p for n, p in param_dict.items() if p.dim() >= 2 and ('wte' in n or 'wpe' in n or 'lm_head' in n)]\n        adam_nodecay = [p for n, p in param_dict.items() if p.dim() < 2]\n        num_decay = sum(p.numel() for _, p in muon_p) + sum(p.numel() for p in adam_decay)\n        num_nodecay = sum(p.numel() for p in adam_nodecay)\n        print(f\"num decayed parameter tensors: {len(muon_p)+len(adam_decay)}, with {num_decay:,} parameters\")\n        print(f\"num non-decayed parameter tensors: {len(adam_nodecay)}, with {num_nodecay:,} parameters\")\n        fused = 'fused' in inspect.signature(torch.optim.AdamW).parameters\n        muon_base_lr = 0.02\n        muon_lr_scale = muon_base_lr / learning_rate\n\n        class Muon(torch.optim.Optimizer):\n            \"\"\"Muon with Nesterov momentum and Newton-Schulz orthogonalization.\"\"\"\n            def __init__(self, params, lr=0.02, momentum=0.95, ns_steps=5, weight_decay=0.0):\n                super().__init__(params, dict(lr=lr, momentum=momentum, ns_steps=ns_steps, weight_decay=weight_decay))\n            @staticmethod\n            def _ns(G, steps=5):\n                X = G.bfloat16() / (G.norm() + 1e-7)\n                if G.size(0) > G.size(1): X = X.T\n                a, b, c = 3.4445, -4.7750, 2.0315\n                for _ in range(steps):\n                    A = X @ X.T; X = a*X + b*(A@X) + c*(A@(A@X))\n                if G.size(0) > G.size(1): X = X.T\n                return X.to(G.dtype)\n            @torch.no_grad()\n            def step(self):\n                for g in self.param_groups:\n                    lr, mom, wd = g['lr'], g['momentum'], g['weight_decay']\n                    for p in g['params']:\n                        if p.grad is None: continue\n                        if wd > 0: p.mul_(1 - lr * wd)\n                        grad = p.grad\n                        state = self.state[p]\n                        if len(state) == 0:\n                            state['buf'] = torch.zeros_like(grad)\n                        buf = state['buf']\n                        buf.mul_(mom).add_(grad, alpha=1 - mom)\n                        nesterov = grad.lerp(buf, mom)\n                        orig_shape = nesterov.shape\n                        if nesterov.dim() == 2:\n                            if orig_shape[0] == 3 * orig_shape[1]:\n                                parts = nesterov.split(orig_shape[1])\n                                U = torch.cat([self._ns(part, g['ns_steps']) for part in parts])\n                                scale = max(1, orig_shape[0] // orig_shape[1]) ** 0.5\n                            else:\n                                U = self._ns(nesterov, g['ns_steps'])\n                                scale = max(1, orig_shape[0] / orig_shape[1]) ** 0.5\n                            p.add_(U, alpha=-lr * scale)\n                        else:\n                            p.add_(nesterov, alpha=-lr)\n\n        muon_opt = Muon([{'params': [p for _, p in muon_p], 'lr_scale': muon_lr_scale}],\n                        lr=muon_base_lr, momentum=0.95, weight_decay=weight_decay)\n        adam_opt = torch.optim.AdamW(\n            [{'params': adam_decay, 'weight_decay': weight_decay},\n             {'params': adam_nodecay, 'weight_decay': 0.0}],\n            lr=learning_rate, betas=betas, fused=fused)\n        class CombinedOpt:\n            def __init__(self, opts):\n                self.optimizers = opts\n                self.param_groups = sum([o.param_groups for o in opts], [])\n            def zero_grad(self, set_to_none=True):\n                for o in self.optimizers: o.zero_grad(set_to_none=set_to_none)\n            def step(self):\n                for o in self.optimizers: o.step()\n        print(f\"using Muon (lr={muon_base_lr}, scale={muon_lr_scale:.1f}, wd={weight_decay}) + AdamW\")\n        return CombinedOpt([muon_opt, adam_opt])",
3  "end_line": 224,
4  "filename": "nanoGPT/custom_pretrain.py",
5  "op": "replace",
6  "start_line": 171
7}

1{
2  "op": "replace",
3  "content": "    # CONFIG_OVERRIDES: override training hyperparameters for your method.\n    # Allowed keys: learning_rate, weight_decay, warmup_iters, min_lr, grad_clip.\n    CONFIG_OVERRIDES = {'learning_rate': 1e-3}",
4  "filename": "nanoGPT/custom_pretrain.py",
5  "start_line": 300,
6  "end_line": 302
7}