feat(training): polynomial, REX, RAWR schedulers and Simplified AdEMAMix

runtime/training/optimizers/__init__.py

@@ -24,6 +24,7 @@
     lion,
     prodigy,
     prodigy_plus_schedulefree,
+    simplified_ademamix,
     soap,
     soap_sf,
 )
@@ -40,6 +41,7 @@
     "lion": lion.build,
     "prodigy": prodigy.build,
     "prodigy_plus_schedulefree": prodigy_plus_schedulefree.build,
+    "simplified_ademamix": simplified_ademamix.build,
     "soap": soap.build,
     "soap_sf": soap_sf.build,
 }

runtime/training/optimizers/simplified_ademamix.py

@@ -0,0 +1,28 @@
+"""Simplified AdEMAMix optimizer build wrapper.
+
+Simplified AdEMAMix (Morwani et al., 2025, arxiv 2502.02431): one momentum
+buffer that sums raw gradients plus alpha * current gradient. A regular
+optimizer with a configurable lr_scheduler, but its lr is ~(1 - beta1) times
+the AdamW lr (1e-6 at beta1=0.99 for AdamW 1e-4). Implemented in
+utils/optimizer_utils.py `class SimplifiedAdEMAMix`.
+"""
+
+from __future__ import annotations
+
+
+def build(args, params, lr: float, weight_decay: float):
+    from utils.optimizer_utils import create_optimizer
+
+    return create_optimizer(
+        optimizer_type="simplified_ademamix",
+        params=params,
+        learning_rate=lr,
+        weight_decay=weight_decay,
+        betas=(
+            float(getattr(args, "ademamix_beta1", 0.99)),
+            float(getattr(args, "ademamix_beta2", 0.999)),
+        ),
+        alpha=float(getattr(args, "ademamix_alpha", 0.0)),
+        beta1_warmup=int(getattr(args, "ademamix_beta1_warmup", 0) or 0),
+        min_beta1=float(getattr(args, "ademamix_min_beta1", 0.9)),
+    )

runtime/training/schedulers/__init__.py

@@ -21,6 +21,9 @@
     cosine_cycles,
     cosine_with_restart,
     cosine_with_warmup,
+    polynomial,
+    rex,
+    rex_annealing_warm_restarts,
 )
 
 __all__ = ["BUILDERS", "build_scheduler", "validate_schema_consistency"]
@@ -33,6 +36,9 @@
     "cosine_cycles": cosine_cycles.build,
     "cosine_with_restart": cosine_with_restart.build,
     "cosine_with_warmup": cosine_with_warmup.build,
+    "polynomial": polynomial.build,
+    "rex": rex.build,
+    "rex_annealing_warm_restarts": rex_annealing_warm_restarts.build,
 }
 
 # schema allows "none" but BUILDERS doesn't list it; validate_schema_consistency exempts it

runtime/training/schedulers/polynomial.py

@@ -0,0 +1,52 @@
+"""Polynomial decay scheduler with linear warmup.
+
+Matches transformers' get_polynomial_decay_schedule_with_warmup: after warmup,
+lr = min_lr + (base_lr - min_lr) * (1 - progress) ** power. power=1 is linear
+decay; power>1 drops faster early and flattens out near the end.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+
+logger = logging.getLogger(__name__)
+
+
+def build(args, optimizer, total_steps: Optional[int]):
+    if total_steps is None or int(total_steps) <= 0:
+        logger.warning("polynomial scheduler needs a known total_steps, falling back to none")
+        return None
+
+    import torch
+
+    total_steps = int(total_steps)
+    warmup_steps = int(getattr(args, "lr_scheduler_warmup_steps", 0) or 0)
+    warmup_steps = max(0, min(warmup_steps, total_steps))
+    power = float(getattr(args, "lr_scheduler_power", 1.0) or 1.0)
+    eta_min = max(0.0, float(getattr(args, "lr_scheduler_eta_min", 0.0) or 0.0))
+    if power <= 0.0:
+        raise ValueError("polynomial: lr_scheduler_power must be > 0")
+
+    def make_lambda(base_lr: float):
+        min_factor = min(1.0, eta_min / base_lr) if base_lr > 0 else 0.0
+
+        def lr_lambda(step: int) -> float:
+            if warmup_steps > 0 and step < warmup_steps:
+                return float(step) / float(warmup_steps)
+            decay_steps = max(1, total_steps - warmup_steps)
+            progress = min(1.0, max(0.0, float(step - warmup_steps) / float(decay_steps)))
+            return min_factor + (1.0 - min_factor) * (1.0 - progress) ** power
+
+        return lr_lambda
+
+    scheduler = torch.optim.lr_scheduler.LambdaLR(
+        optimizer,
+        lr_lambda=[make_lambda(group["lr"]) for group in optimizer.param_groups],
+    )
+    logger.info(
+        "LR schedule: polynomial (total_steps=%s, warmup_steps=%s, power=%s, eta_min=%s)",
+        total_steps, warmup_steps, power, eta_min,
+    )
+    return scheduler

runtime/training/schedulers/rex.py

@@ -0,0 +1,56 @@
+"""REX scheduler (Chen et al., 2021, "REX: Revisiting Budgeted Training with an
+Improved Schedule", arxiv 2107.04197) with optional linear warmup.
+
+After warmup, with z = 1 - progress:
+    factor = z / ((1 - d) + d * z)
+d = 0.5 is the paper's curve; larger d (LoRA Easy Training Scripts uses 0.9) holds
+the peak lr longer and drops more sharply at the end.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+from training.schedulers.rex_annealing_warm_restarts import rex_factor
+
+
+logger = logging.getLogger(__name__)
+
+
+def build(args, optimizer, total_steps: Optional[int]):
+    if total_steps is None or int(total_steps) <= 0:
+        logger.warning("rex scheduler needs a known total_steps, falling back to none")
+        return None
+
+    import torch
+
+    total_steps = int(total_steps)
+    warmup_steps = int(getattr(args, "lr_scheduler_warmup_steps", 0) or 0)
+    warmup_steps = max(0, min(warmup_steps, total_steps))
+    eta_min = max(0.0, float(getattr(args, "lr_scheduler_eta_min", 0.0) or 0.0))
+    d = float(getattr(args, "lr_scheduler_rex_d", 0.9))
+    if not 0.0 < d <= 1.0:
+        raise ValueError("rex: lr_scheduler_rex_d must be in (0, 1]")
+
+    def make_lambda(base_lr: float):
+        min_factor = min(1.0, eta_min / base_lr) if base_lr > 0 else 0.0
+
+        def lr_lambda(step: int) -> float:
+            if warmup_steps > 0 and step < warmup_steps:
+                return float(step) / float(warmup_steps)
+            decay_steps = max(1, total_steps - warmup_steps)
+            progress = min(1.0, max(0.0, float(step - warmup_steps) / float(decay_steps)))
+            return min_factor + (1.0 - min_factor) * rex_factor(progress, d)
+
+        return lr_lambda
+
+    scheduler = torch.optim.lr_scheduler.LambdaLR(
+        optimizer,
+        lr_lambda=[make_lambda(group["lr"]) for group in optimizer.param_groups],
+    )
+    logger.info(
+        "LR schedule: rex (total_steps=%s, warmup_steps=%s, d=%s, eta_min=%s)",
+        total_steps, warmup_steps, d, eta_min,
+    )
+    return scheduler

runtime/training/schedulers/rex_annealing_warm_restarts.py

@@ -0,0 +1,104 @@
+"""REX annealing with warm restarts (RAWR).
+
+Port of RexAnnealingWarmRestarts from derrian-distro/LoRA_Easy_Training_scripts_Backend,
+rewritten as a stateless LambdaLR (cycle position derived from the step index, so
+resume needs no extra scheduler state). Each cycle:
+  - warmup: linear ramp from min_lr to the cycle peak over lr_scheduler_warmup_steps;
+  - decay: REX curve min_lr + (peak - min_lr) * z / ((1 - d) + d * z), z = 1 - progress.
+After each restart the peak is multiplied by gamma and the decay part of the cycle
+by cycle_multiplier. The first cycle length is chosen so lr_scheduler_cycle_count
+cycles fill total_steps.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+
+logger = logging.getLogger(__name__)
+
+
+def rex_factor(progress: float, d: float) -> float:
+    z = 1.0 - min(1.0, max(0.0, progress))
+    return z / ((1.0 - d) + d * z)
+
+
+def cycle_lengths(total_steps: int, cycles: int, multiplier: float, warmup: int) -> list[int]:
+    """Lengths (warmup included) of `cycles` cycles covering ~total_steps."""
+    geometric = sum(multiplier ** k for k in range(cycles))
+    decay0 = max(1, round((total_steps - cycles * warmup) / geometric))
+    lengths = [decay0 + warmup]
+    while len(lengths) < cycles:
+        lengths.append(max(1, round((lengths[-1] - warmup) * multiplier)) + warmup)
+    return lengths
+
+
+def build(args, optimizer, total_steps: Optional[int]):
+    if total_steps is None or int(total_steps) <= 0:
+        logger.warning("rex_annealing_warm_restarts needs a known total_steps, falling back to none")
+        return None
+
+    import torch
+
+    total_steps = int(total_steps)
+    cycles = max(1, int(getattr(args, "lr_scheduler_cycle_count", 3) or 1))
+    multiplier = float(getattr(args, "lr_scheduler_cycle_multiplier", 1.0) or 1.0)
+    gamma = float(getattr(args, "lr_scheduler_gamma", 0.9))
+    d = float(getattr(args, "lr_scheduler_rex_d", 0.9))
+    eta_min = max(0.0, float(getattr(args, "lr_scheduler_eta_min", 0.0) or 0.0))
+    warmup = max(0, int(getattr(args, "lr_scheduler_warmup_steps", 0) or 0))
+    if multiplier <= 0.0:
+        raise ValueError("rex_annealing_warm_restarts: lr_scheduler_cycle_multiplier must be > 0")
+    if not 0.0 < d <= 1.0:
+        raise ValueError("rex_annealing_warm_restarts: lr_scheduler_rex_d must be in (0, 1]")
+    if gamma <= 0.0:
+        raise ValueError("rex_annealing_warm_restarts: lr_scheduler_gamma must be > 0")
+    # Warmup has to leave room for decay in every cycle.
+    max_warmup = max(0, total_steps // cycles - 1)
+    if warmup > max_warmup:
+        logger.warning(
+            "rex_annealing_warm_restarts: warmup_steps=%s doesn't fit %s cycles in %s steps, "
+            "clamping to %s", warmup, cycles, total_steps, max_warmup,
+        )
+        warmup = max_warmup
+    lengths = cycle_lengths(total_steps, cycles, multiplier, warmup)
+
+    def locate(step: int) -> tuple[int, int, int]:
+        """(cycle index, step within cycle, cycle length); continues past the planned cycles."""
+        cycle, start, length = 0, 0, lengths[0]
+        while step >= start + length:
+            start += length
+            cycle += 1
+            if cycle < len(lengths):
+                length = lengths[cycle]
+            else:
+                length = max(1, round((length - warmup) * multiplier)) + warmup
+        return cycle, step - start, length
+
+    def make_lambda(base_lr: float):
+        min_factor = min(1.0, eta_min / base_lr) if base_lr > 0 else 0.0
+
+        def lr_lambda(step: int) -> float:
+            cycle, pos, length = locate(max(0, step))
+            peak = gamma ** cycle
+            if peak <= min_factor:
+                return min_factor
+            span = peak - min_factor
+            if warmup > 0 and pos < warmup:
+                return min_factor + span * pos / warmup
+            progress = (pos - warmup) / max(1, length - warmup)
+            return min_factor + span * rex_factor(progress, d)
+
+        return lr_lambda
+
+    scheduler = torch.optim.lr_scheduler.LambdaLR(
+        optimizer,
+        lr_lambda=[make_lambda(group["lr"]) for group in optimizer.param_groups],
+    )
+    logger.info(
+        "LR schedule: rex_annealing_warm_restarts (total_steps=%s, cycles=%s, lengths=%s, "
+        "multiplier=%s, gamma=%s, d=%s, warmup_steps=%s, eta_min=%s)",
+        total_steps, cycles, lengths, multiplier, gamma, d, warmup, eta_min,
+    )
+    return scheduler

studio/domain/training.py

@@ -445,6 +445,9 @@ class TrainingConfig(BaseModel):
         "cosine_with_warmup",
         "cosine_cycles",
         "constant_then_cosine",
+        "polynomial",
+        "rex",
+        "rex_annealing_warm_restarts",
     ] = Field(
         "none",
         description="Learning rate schedule (none = constant; Prodigy / PPSF / Automagic / SOAP-SF are fixed to none)",
@@ -467,22 +470,46 @@ class TrainingConfig(BaseModel):
     )
     lr_scheduler_eta_min: float = Field(
         1e-6, ge=0.0,
-        description="Learning rate floor: cosine schedules stop decreasing once they reach this value; usually much smaller than the initial lr (e.g. 1e-4 initial with 1e-6 floor)",
+        description="Learning rate floor: cosine / polynomial / rex / RAWR schedules stop decreasing once they reach this value; usually much smaller than the initial lr (e.g. 1e-4 initial with 1e-6 floor)",
         json_schema_extra=_meta(
             "training",
-            show_when="lr_scheduler==cosine||lr_scheduler==cosine_with_restart||lr_scheduler==cosine_with_warmup",
+            show_when="lr_scheduler==cosine||lr_scheduler==cosine_with_restart||lr_scheduler==cosine_with_warmup||lr_scheduler==polynomial||lr_scheduler==rex||lr_scheduler==rex_annealing_warm_restarts",
             advanced=True,
         ),
     )
     lr_scheduler_warmup_steps: int = Field(
         100, ge=0,
-        description="cosine_with_warmup warmup steps",
-        json_schema_extra=_meta("training", show_when="lr_scheduler==cosine_with_warmup", advanced=True),
+        description="Warmup steps for cosine_with_warmup / polynomial / rex (linear ramp from 0 to the peak lr); for RAWR the warmup repeats at the start of every cycle. 0 = no warmup",
+        json_schema_extra=_meta(
+            "training",
+            show_when="lr_scheduler==cosine_with_warmup||lr_scheduler==polynomial||lr_scheduler==rex||lr_scheduler==rex_annealing_warm_restarts",
+            advanced=True,
+        ),
+    )
+    lr_scheduler_power: float = Field(
+        1.0, gt=0.0,
+        description="polynomial: decay power. 1 = linear decay to the floor; >1 drops faster early and flattens near the end; <1 holds the lr longer and drops at the end",
+        json_schema_extra=_meta("training", show_when="lr_scheduler==polynomial"),
+    )
+    lr_scheduler_rex_d: float = Field(
+        0.9, gt=0.0, le=1.0,
+        description="rex / RAWR curve shape d: factor = z / ((1 - d) + d·z), z = remaining fraction. 0.5 = the REX paper; 0.9 (LoRA Easy Training Scripts default) holds the peak longer and drops more sharply at the end",
+        json_schema_extra=_meta("training", show_when="lr_scheduler==rex||lr_scheduler==rex_annealing_warm_restarts", advanced=True),
+    )
+    lr_scheduler_cycle_multiplier: float = Field(
+        1.0, gt=0.0,
+        description="RAWR: each next cycle is this many times longer than the previous one (1 = equal cycles, 2 = doubling). Cycle lengths are fitted so all cycles fill the run",
+        json_schema_extra=_meta("training", show_when="lr_scheduler==rex_annealing_warm_restarts"),
+    )
+    lr_scheduler_gamma: float = Field(
+        0.9, gt=0.0, le=1.0,
+        description="RAWR: peak lr multiplier after each restart (1 = every cycle restarts at the full lr; 0.9 = each peak 10% lower)",
+        json_schema_extra=_meta("training", show_when="lr_scheduler==rex_annealing_warm_restarts"),
     )
     lr_scheduler_cycle_count: int = Field(
         3, ge=1, le=100,
-        description="Number of peaks for cosine_cycles; total training steps are divided evenly into this many periods",
-        json_schema_extra=_meta("training", show_when="lr_scheduler==cosine_cycles"),
+        description="Number of peaks for cosine_cycles / RAWR; the cycles together fill the whole run (cosine_cycles splits it evenly)",
+        json_schema_extra=_meta("training", show_when="lr_scheduler==cosine_cycles||lr_scheduler==rex_annealing_warm_restarts"),
     )
     lr_scheduler_cycle_max_lr: Optional[float] = Field(
         None, gt=0.0,
@@ -514,9 +541,9 @@ class TrainingConfig(BaseModel):
         description="constant_then_cosine: LR at the very end of training",
         json_schema_extra=_meta("training", show_when="lr_scheduler==constant_then_cosine"),
     )
-    optimizer_type: Literal["adamw", "adamw8bit", "automagic", "came", "lion", "prodigy", "prodigy_plus_schedulefree", "soap", "soap_sf"] = Field(
+    optimizer_type: Literal["adamw", "adamw8bit", "automagic", "came", "lion", "prodigy", "prodigy_plus_schedulefree", "simplified_ademamix", "soap", "soap_sf"] = Field(
         "adamw",
-        description="Optimizer. adamw standard baseline; adamw8bit same as AdamW but with state quantized to int8 (state VRAM ~1/4 of AdamW, hyperparameters carry over directly, requires bitsandbytes); automagic adaptive per-parameter lr (recommended lr=1e-6); came confidence-guided + factored second moment (lower state VRAM than AdamW, lr on the same scale as AdamW); lion roughly half AdamW's VRAM (recommended lr=AdamW lr / 3); prodigy / prodigy_plus_schedulefree adaptively estimates lr (set lr to 1.0); soap Adam-in-Shampoo-eigenbasis second-order preconditioning (fits faster, lr on the same scale as AdamW); soap_sf SOAP + Schedule-Free (lr_scheduler fixed to none)",
+        description="Optimizer. adamw standard baseline; adamw8bit same as AdamW but with state quantized to int8 (state VRAM ~1/4 of AdamW, hyperparameters carry over directly, requires bitsandbytes); automagic adaptive per-parameter lr (recommended lr=1e-6); came confidence-guided + factored second moment (lower state VRAM than AdamW, lr on the same scale as AdamW); lion roughly half AdamW's VRAM (recommended lr=AdamW lr / 3); prodigy / prodigy_plus_schedulefree adaptively estimates lr (set lr to 1.0); simplified_ademamix one raw-gradient momentum + current-gradient mix (lr ≈ AdamW lr × (1 − β1), e.g. 1e-6 at β1=0.99); soap Adam-in-Shampoo-eigenbasis second-order preconditioning (fits faster, lr on the same scale as AdamW); soap_sf SOAP + Schedule-Free (lr_scheduler fixed to none)",
         json_schema_extra=_meta("training"),
     )
     prodigy_d_coef: float = Field(
@@ -573,6 +600,34 @@ class TrainingConfig(BaseModel):
         description="Lion β2 (momentum accumulation coefficient)",
         json_schema_extra=_meta("training", show_when="optimizer_type==lion", advanced=True),
     )
+    # ------------------------- Simplified AdEMAMix-specific fields -------------------------
+    # Morwani et al. 2025 (arxiv 2502.02431). The momentum sums raw gradients
+    # (no 1-β1 factor), so lr is ~(1-β1) times the AdamW lr.
+    ademamix_beta1: float = Field(
+        0.99, ge=0.0, lt=1.0,
+        description="Simplified AdEMAMix β1 (momentum decay; higher = longer gradient memory, and lr must shrink by (1 − β1))",
+        json_schema_extra=_meta("training", show_when="optimizer_type==simplified_ademamix", advanced=True),
+    )
+    ademamix_beta2: float = Field(
+        0.999, ge=0.0, lt=1.0,
+        description="Simplified AdEMAMix β2 (second-moment EMA decay)",
+        json_schema_extra=_meta("training", show_when="optimizer_type==simplified_ademamix", advanced=True),
+    )
+    ademamix_alpha: float = Field(
+        0.0, ge=0.0,
+        description="Simplified AdEMAMix α: weight of the current gradient added on top of the momentum (0 = momentum only)",
+        json_schema_extra=_meta("training", show_when="optimizer_type==simplified_ademamix", advanced=True),
+    )
+    ademamix_beta1_warmup: int = Field(
+        0, ge=0,
+        description="Simplified AdEMAMix: steps over which β1 ramps up from the starting β1 to β1 (0 = off). Stabilizes the start with a high β1",
+        json_schema_extra=_meta("training", show_when="optimizer_type==simplified_ademamix", advanced=True),
+    )
+    ademamix_min_beta1: float = Field(
+        0.9, ge=0.0, lt=1.0,
+        description="Simplified AdEMAMix: starting β1 for the β1 warmup",
+        json_schema_extra=_meta("training", show_when="optimizer_type==simplified_ademamix", advanced=True),
+    )
     automagic_variant: Literal["v1", "v2"] = Field(
         "v1",
         description="v1: per-element lr mask (classic, recommended); v2: per-param scalar lr + fused backward (experimental, saves VRAM; incompatible with grad_accum / grad_clip / fp16)",