"""Placeholder trainer. The step is fake. The commands are real. ``Schedule`` holds the current learning rate, then ramps in a straight line to a target. ``Trainer.duplicate_layer`` copies that layer's weights and inserts the copy after it. Nothing here trains a model or contacts a server unless ``main`` is invoked. """ from __future__ import annotations import os import launchhelm from launchhelm import CommandContext from launchhelm.values import JSONValue class Schedule: """Learning rate that can be replaced at a finished step.""" def __init__(self, lr: float) -> None: self.lr = lr self._applied_at = 0 self._hold_steps = 0 self._ramp_steps = 0 self._origin = lr self._target = lr def override(self, applied_at: int, hold_steps: int, ramp_steps: int, target_lr: float) -> None: """Keep ``lr`` for ``hold_steps`` steps after ``applied_at``, then ramp. ``applied_at`` is the step that just finished. The hold covers the following steps. A zero hold starts the ramp on the next step. A zero ramp jumps to ``target_lr`` on the next step. """ if hold_steps < 0 or ramp_steps < 0: raise ValueError("hold and ramp lengths must be zero or positive") self._applied_at = applied_at self._hold_steps = hold_steps self._ramp_steps = ramp_steps self._origin = self.lr self._target = target_lr def at(self, step: int) -> float: """Return the learning rate to use for ``step`` and store it.""" flat_until = self._applied_at + self._hold_steps if step <= flat_until: return self.lr if self._ramp_steps == 0 or step >= flat_until + self._ramp_steps: self.lr = self._target return self.lr done = step - flat_until self.lr = self._origin + (self._target - self._origin) * done / self._ramp_steps return self.lr class Layer: """One named vector. A copy is a new layer with the same weights.""" def __init__(self, name: str, weight: list[float]) -> None: self.name = name self.weight = weight def copy(self) -> Layer: """Return a new layer with a copy of these weights.""" return Layer(f"{self.name}-copy", list(self.weight)) class Trainer: """Tiny stand-in for a training step and the commands it can apply.""" def __init__(self, *, lr: float = 0.0003, depth: int = 3, width: int = 4) -> None: self.schedule = Schedule(lr) self.layers = [Layer(f"layer{index}", [0.1] * width) for index in range(depth)] self.step_index = 0 def set_schedule( self, hold_steps: int, ramp_steps: int, target_lr: float, context: CommandContext, ) -> None: """Replace the schedule from the step recorded on the command.""" applied_at = int(context.coordinates[0]["value"]) self.schedule.override(applied_at, hold_steps, ramp_steps, target_lr) def duplicate_layer(self, index: int) -> None: """Copy ``layers[index]`` and insert that copy immediately after it.""" if index < 0 or index >= len(self.layers): raise IndexError(f"layer index {index} is outside 0..{len(self.layers) - 1}") self.layers.insert(index + 1, self.layers[index].copy()) def step(self) -> float: """Apply one placeholder update and return a stand-in loss. Each weight moves by ``-lr * weight``. The loss is the sum of squares before that move. This is not a gradient of a real model. """ lr = self.schedule.at(self.step_index) loss = 0.0 for layer in self.layers: for position, value in enumerate(layer.weight): loss += value * value layer.weight[position] = value - lr * value self.step_index += 1 return loss def dump(self) -> dict[str, JSONValue]: """Return the schedule, step, and layer weights to keep.""" return { "lr": self.schedule.lr, "step_index": self.step_index, "applied_at": self.schedule._applied_at, "hold_steps": self.schedule._hold_steps, "ramp_steps": self.schedule._ramp_steps, "origin": self.schedule._origin, "target": self.schedule._target, "layers": [{"name": layer.name, "weight": list(layer.weight)} for layer in self.layers], } def load(self, state: object) -> None: """Restore a snapshot produced by ``dump``.""" if not isinstance(state, dict): raise TypeError("trainer snapshot must be an object") self.schedule.lr = float(state["lr"]) self.schedule._applied_at = int(state["applied_at"]) self.schedule._hold_steps = int(state["hold_steps"]) self.schedule._ramp_steps = int(state["ramp_steps"]) self.schedule._origin = float(state["origin"]) self.schedule._target = float(state["target"]) self.step_index = int(state["step_index"]) layers = state["layers"] if not isinstance(layers, list): raise TypeError("trainer snapshot layers must be a list") self.layers = [] for item in layers: if not isinstance(item, dict): raise TypeError("trainer snapshot layer must be an object") weight = item["weight"] if not isinstance(weight, list): raise TypeError("trainer snapshot weight must be a list") self.layers.append(Layer(str(item["name"]), [float(value) for value in weight])) def main() -> None: """Log a short fake run and accept schedule and duplication commands.""" trainer = Trainer() resume = os.environ.get("LAUNCHHELM_RESUME", "") with launchhelm.init("quickstart", name="placeholder", resume=resume) as run: if run.restored: trainer.load(run.restored) run.command("after_step")(trainer.set_schedule) run.command("after_step")(trainer.duplicate_layer) for step in range(trainer.step_index, trainer.step_index + 8): loss = trainer.step() run.sync("after_step", step=step) run.log( { "loss": loss, "lr": trainer.schedule.lr, "layers": len(trainer.layers), }, step=step, ) run.keep(trainer.dump()) if __name__ == "__main__": main()