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    \"\"\"Warmup-Stable-Cosine-Decay (WSCD) schedule.\n    \n    Three-phase schedule designed for optimal LLM pretraining:\n    1. Smooth quadratic warmup (sin² shape) - reduces gradient shock\n    2. Extended stable phase (70% of training) - maximizes gradient accumulation\n    3. Smooth cosine decay (30% of training) - gentler than linear, preserves representations\n    \"\"\"\n    # Phase boundaries\n    stable_end = int(lr_decay_iters * 0.7)  # Stable phase ends at 70%\n    \n    # Phase 1: Smooth quadratic warmup (sin² shape)\n    if it < warmup_iters:\n        # sin² provides smoother transition than linear, reducing initial gradient shock\n        progress = (it + 1) / (warmup_iters + 1)\n        return learning_rate * math.sin(math.pi / 2 * progress) ** 2\n    \n    # Phase 2: Stable phase - maintain peak learning rate\n    if it < stable_end:\n        return learning_rate\n    \n    # Phase 3: Smooth cosine decay to min_lr\n    decay_progress = (it - stable_end) / (lr_decay_iters - stable_end)\n    # Cosine decay: coeff goes from 1.0 to 0.0 smoothly\n    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_progress))\n    return min_lr + coeff * (learning_rate - min_lr)"
7}

1{
2  "content": "def get_lr(it, warmup_iters, lr_decay_iters, learning_rate, min_lr):\n    \"\"\"Multi-Scale Smooth Warmup-Stable-Decay (MSSWSD) schedule.\n    \n    Key innovations:\n    1. Sin² warmup: smoother gradient transition than linear\n    2. Smooth sigmoid transition between stable and decay phases\n    3. Extended effective warmup with gradual ramp-up\n    4. Adaptive decay rate that's gentler early, steeper late\n    \"\"\"\n    import math\n    \n    # Phase 1: Smooth quadratic warmup (sin² shape)\n    if it < warmup_iters:\n        progress = (it + 1) / (warmup_iters + 1)\n        return learning_rate * math.sin(math.pi / 2 * progress) ** 2\n    \n    # Phase 2 & 3: Smooth transition from stable to decay\n    # Use a smooth sigmoid-like transition starting at 40% of training\n    # This avoids the abrupt change that can destabilize training\n    stable_fraction = 0.4  # Start decay after 40% of training\n    transition_width = 0.15  # Width of smooth transition region\n    \n    decay_start = int(lr_decay_iters * stable_fraction)\n    transition_start = decay_start - int(lr_decay_iters * transition_width / 2)\n    transition_end = int(lr_decay_iters * (stable_fraction + transition_width / 2))\n    \n    if it < transition_start:\n        # Full stable phase\n        return learning_rate\n    \n    if it < transition_end:\n        # Smooth sigmoid transition from stable to decay\n        t = (it - transition_start) / (transition_end - transition_start)\n        # Sigmoid: smooth step function\n        sigmoid_t = 1.0 / (1.0 + math.exp(-12 * (t - 0.5)))\n        # Blend between stable and cosine decay\n        coeff = 0.5 * (1.0 + math.cos(math.pi * sigmoid_t))\n        return min_lr + coeff * (learning_rate - min_lr)\n    \n    if it <= lr_decay_iters:\n        # Pure cosine decay phase\n        decay_progress = (it - transition_end) / (lr_decay_iters - transition_end)\n        coeff = 0.5 * (1.0 + math.cos(math.pi * decay_progress))\n        return min_lr + coeff * (learning_rate - min_lr)\n    \n    return min_lr",
3  "end_line": 216,
4  "filename": "nanoGPT/custom_pretrain.py",
5  "start_line": 191
6}

1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "content": "def get_lr(it, warmup_iters, lr_decay_iters, learning_rate, min_lr):\n    \"\"\"Multi-Scale Smooth Warmup-Stable-Decay (MSSWSD) schedule.\n    \n    Key innovations:\n    1. Sin² warmup: smoother gradient transition than linear\n    2. Smooth sigmoid transition between stable and decay phases\n    3. Extended effective warmup with gradual ramp-up\n    4. Adaptive decay rate that's gentler early, steeper late\n    \"\"\"\n    import math\n    \n    # Phase 1: Smooth quadratic warmup (sin² shape)\n    if it < warmup_iters:\n        progress = (it + 1) / (warmup_iters + 1)\n        return learning_rate * math.sin(math.pi / 2 * progress) ** 2\n    \n    # Phase 2 & 3: Smooth transition from stable to decay\n    # Use a smooth sigmoid-like transition starting at 40% of training\n    # This avoids the abrupt change that can destabilize training\n    stable_fraction = 0.4  # Start decay after 40% of training\n    transition_width = 0.15  # Width of smooth transition region\n    \n    decay_start = int(lr_decay_iters * stable_fraction)\n    transition_start = decay_start - int(lr_decay_iters * transition_width / 2)\n    transition_end = int(lr_decay_iters * (stable_fraction + transition_width / 2))\n    \n    if it < transition_start:\n        # Full stable phase\n        return learning_rate\n    \n    if it < transition_end:\n        # Smooth sigmoid transition from stable to decay\n        t = (it - transition_start) / (transition_end - transition_start)\n        # Sigmoid: smooth step function\n        sigmoid_t = 1.0 / (1.0 + math.exp(-12 * (t - 0.5)))\n        # Blend between stable and cosine decay\n        coeff = 0.5 * (1.0 + math.cos(math.pi * sigmoid_t))\n        return min_lr + coeff * (learning_rate - min_lr)\n    \n    if it <= lr_decay_iters:\n        # Pure cosine decay phase\n        decay_progress = (it - transition_end) / (lr_decay_iters - transition_end)\n        coeff = 0.5 * (1.0 + math.cos(math.pi * decay_progress))\n        return min_lr + coeff * (learning_rate - min_lr)\n    \n    return min_lr",
5  "start_line": 191,
6  "end_line": 216
7}

1{
2  "op": "replace",
3  "content": "def get_lr(it, warmup_iters, lr_decay_iters, learning_rate, min_lr):\n    \"\"\"Warmup-Stable-PowerDecay (WSPD) schedule.\n    \n    Algorithmic improvements over standard cosine:\n    1. Sin² warmup: smoother gradient transition than linear, reduces initial instability\n    2. Extended stable phase (60%): maximizes training efficiency at peak LR\n    3. Power-law decay: naturally tapers, shown effective for large-scale pretraining\n       Uses 1/(1 + kt)^α form with parameters tuned to reach min_lr exactly at end\n    \"\"\"\n    # Phase 1: Smooth quadratic warmup (sin² shape)\n    if it < warmup_iters:\n        progress = (it + 1) / (warmup_iters + 1)\n        return learning_rate * math.sin(math.pi / 2 * progress) ** 2\n    \n    # Phase 2: Stable phase at peak learning rate\n    stable_end = int(lr_decay_iters * 0.6)  # 60% of training at peak LR\n    if it < stable_end:\n        return learning_rate\n    \n    # Phase 3: Power-law decay from stable_end to lr_decay_iters\n    # Form: min_lr + (lr - min_lr) / (1 + k*t)^α\n    # Calibrated so that at t=1 (end of training), we reach exactly min_lr\n    t = (it - stable_end) / (lr_decay_iters - stable_end)  # t in [0, 1]\n    alpha = 0.75  # Decay exponent - gentler than linear initially\n    # With α=0.75, we need k such that (1+k)^α = (lr-min_lr)/(min_lr-min_lr) → inf\n    # Instead, use: coeff = (1 - t)^α which goes from 1 to 0\n    # But (1-t)^α has infinite derivative at t=1, so use: (1 - t^β)^α\n    beta = 0.5  # Shape parameter - makes decay steeper at start, gentler at end\n    coeff = (1.0 - math.pow(t, beta)) ** alpha\n    return min_lr + coeff * (learning_rate - min_lr)",
4  "end_line": 236,
5  "filename": "nanoGPT/custom_pretrain.py",
6  "start_line": 191
7}

1{
2  "content": "def get_lr(it, warmup_iters, lr_decay_iters, learning_rate, min_lr):\n    \"\"\"Smooth Warmup-Stable-SharpDecay (SWSSD) schedule.\n    \n    Key idea: Maximize training efficiency by maintaining peak LR as long as \n    possible, then decay sharply but smoothly at the end. This exploits the \n    observation that LLM training benefits from long periods at high LR, \n    with a short but smooth convergence phase.\n    \n    Schedule phases:\n    1. Smooth sin² warmup (4% of training)\n    2. Extended stable plateau (80% of training at peak LR)\n    3. Sharp but smooth cosine decay (16% of training)\n    \"\"\"\n    # Phase 1: Smooth quadratic warmup (sin² shape) - smoother than linear\n    if it < warmup_iters:\n        progress = (it + 1) / (warmup_iters + 1)\n        return learning_rate * math.sin(math.pi / 2 * progress) ** 2\n    \n    # Phase 2: Extended stable phase at peak learning rate\n    # Keep LR high for 80% of total training\n    decay_start = int(lr_decay_iters * 0.8)\n    if it < decay_start:\n        return learning_rate\n    \n    # Phase 3: Smooth cosine decay\n    t = (it - decay_start) / (lr_decay_iters - decay_start)  # t in [0, 1]\n    coeff = 0.5 * (1.0 + math.cos(math.pi * t))\n    return min_lr + coeff * (learning_rate - min_lr)",
3  "end_line": 220,
4  "filename": "nanoGPT/custom_pretrain.py",
5  "start_line": 191
6}

1{
2  "n": 1
3}