SolitaryThinker
db866dfb04
Sweep all 296 v2 modules: simplify verbose docstrings/comments and remove 376 dangling "(design_vN §X)" citations to the now-deleted design docs (v2/README.md is the source of truth). Comment/docstring-only — AST-verified code-identical; the CPU suite holds at 237 passed / 5 pre-existing. Also applies yapf + ruff --fix auto-fixes (import ordering, forward-ref annotation de-quoting under `from __future__ import annotations`; behavior- neutral, suite-confirmed) and adds the legitimate domain terms mot/clen/te to the codespell ignore-list. Remaining ruff (24) + mypy (68 no-untyped-call) findings are pre-existing v2 debt, untouched here.
72 lines
2.7 KiB
Python
72 lines
2.7 KiB
Python
"""TrainingMethod base (mirrors the repo's train/methods contract).
|
||
|
||
Every method's rollout drives the same loop the engine serves (via ``_rollout`` → ``rollout_loop``),
|
||
so there is one numerics surface. Methods differ only in loss math, roles, and capture — not in a
|
||
second sampler. ``manages_optimization`` lets RL own its sample→score→inner-train cadence.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from abc import ABC, abstractmethod
|
||
from typing import Any
|
||
|
||
import numpy as np
|
||
|
||
from v2._enums import ConsistencyLevel, ExecutionProfile
|
||
from v2.cache import CacheManager
|
||
from v2.card import load_card
|
||
from v2.recipes.common import text_encode_node_fn
|
||
from v2.training.rollout import rollout_loop
|
||
|
||
|
||
def new_instance(card) -> Any:
|
||
"""A fresh resident instance of a card (a training role: student/teacher/critic/old/reference)."""
|
||
return load_card(card, cache_manager=CacheManager.from_card(card), validate=False)
|
||
|
||
|
||
def predict_x0(velocity: np.ndarray, noised: np.ndarray, sigma: float) -> np.ndarray:
|
||
"""Flow prediction → clean sample: x0 = x_σ − σ·v (Wan ``x0 = x_t − σ·model_output``)."""
|
||
return noised - sigma * np.asarray(velocity, dtype="float32")
|
||
|
||
|
||
class TrainingMethod(ABC):
|
||
name: str = "base"
|
||
consistency: ConsistencyLevel = ConsistencyLevel.C1
|
||
student_loop_id: str = "diffusion_denoise"
|
||
|
||
def __init__(self, student_instance: Any, *, lr: float = 0.05):
|
||
self.student = student_instance
|
||
self.lr = lr
|
||
self.iteration = 0
|
||
|
||
@property
|
||
def student_dit(self) -> Any:
|
||
return self.student.component("transformer")
|
||
|
||
def manages_optimization(self) -> bool:
|
||
return False
|
||
|
||
def get_grad_clip_targets(self, iteration: int = 0) -> dict[str, Any]:
|
||
return {"student": self.student_dit}
|
||
|
||
def consistency_level(self) -> ConsistencyLevel:
|
||
return self.consistency
|
||
|
||
# --- rollout helpers: the SAME loop the engine serves ---------------------- #
|
||
def _encode_slots(self, request, instance: Any = None) -> dict:
|
||
slots: dict = {}
|
||
text_encode_node_fn(instance or self.student, slots, request, None)
|
||
return slots
|
||
|
||
def _rollout(self, request, *, instance: Any = None, loop_id: str | None = None):
|
||
inst = instance or self.student
|
||
slots = self._encode_slots(request, inst)
|
||
return rollout_loop(inst,
|
||
loop_id or self.student_loop_id,
|
||
request,
|
||
slots=slots,
|
||
profile=ExecutionProfile.ROLLOUT)
|
||
|
||
@abstractmethod
|
||
def train_step(self, batch: dict, iteration: int) -> tuple[dict, dict]:
|
||
"""Return (loss_map, metrics). loss_map values are floats; metrics includes grad_norm/*."""
|