SolitaryThinker
270fae959d
A scoped, CPU-testable realization of design_v3.md — the architecture where the atomic unit is a typed (recipe, runtime) ModelCard, the model owns loop semantics while the runtime owns loop lifecycle, one resident instance runs many loops, and training records behavior on the same loops it serves. Implements: - card/ loop/ runtime/ cache/ memory/ parallel/ parity/ extend/ program/ request/ training/ spanning design_v3 §4-§13: ModelCard + validate(); driven loops (init/next/advance/finalize); step-interleaving Engine with reservation-before-admission + per-class caches keyed by CacheKey; the C0-C4 consistency ladder + the non-negotiable batch-of-N interleave parity gate. - Inference: Wan2.1-1.3B (T2V), LTX2.3 (two-stage distilled, shared transformer), Wan-causal (chunk rollout + slab-KV streaming). - Training (Wan2.1-1.3B), each driving the SAME loops the engine serves: finetune (flow-match), DMD2 (teacher/critic distribution matching), DiffusionNFT (likelihood-free C2, samples from the decay-blended old policy, group-relative advantages, shared-prompt cache reuse), self-forcing (causal chunk loop). The engine never imports training (the §10 dependency rule, grep-verified). numpy-only core (no torch/GPU here); heavy Wan/LTX forwards are deterministic toy stand-ins with lazy torch-adapter seams (ComponentSpec.load_id/factory) for a GPU box. 49 tests pass via pytest and a zero-dependency runner; interleave parity verified (and a buggy module-global interceptor provably breaks it). Omni-ready spine (ar_decode/chunk_step loop kinds, multi-loop instances, LoopState.extension) for the phase-2 Cosmos3 + vllm-omni omni ports. Run: python3 -m pytest mini_fastvideo/tests/ -q ; python3 -m mini_fastvideo.examples
82 lines
2.6 KiB
Python
82 lines
2.6 KiB
Python
"""Tagged memory pools with reservation-before-admission (design_v3 §6.2, §7.3).
|
|
|
|
> Reservation must be pre-flight. A scheduler that admits a diffusion step and then
|
|
> discovers it cannot allocate the VAE tile is wrong.
|
|
|
|
Sleep/wake is component-granular (tags are component names) for RL: drop DiT + caches,
|
|
keep VAE/text-encoder resident (design_v3 §7.3; CuMem-style).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import itertools
|
|
from dataclasses import dataclass, field
|
|
|
|
|
|
class OutOfMemory(Exception):
|
|
pass
|
|
|
|
|
|
_res_ctr = itertools.count(1)
|
|
|
|
|
|
@dataclass
|
|
class Reservation:
|
|
res_id: int
|
|
tag: str
|
|
nbytes: int
|
|
active: bool = True
|
|
|
|
|
|
class MemoryManager:
|
|
def __init__(self, total_bytes: int = 1 << 40, per_tag_budget: dict[str, int] | None = None):
|
|
self.total_bytes = total_bytes
|
|
self.per_tag_budget = dict(per_tag_budget or {})
|
|
self.reserved = 0
|
|
self._by_tag: dict[str, int] = {}
|
|
self._reservations: dict[int, Reservation] = {}
|
|
self._asleep: set[str] = set()
|
|
|
|
@property
|
|
def available(self) -> int:
|
|
return self.total_bytes - self.reserved
|
|
|
|
def can_reserve(self, tag: str, nbytes: int) -> bool:
|
|
if nbytes > self.available:
|
|
return False
|
|
budget = self.per_tag_budget.get(tag)
|
|
if budget is not None and self._by_tag.get(tag, 0) + nbytes > budget:
|
|
return False
|
|
return True
|
|
|
|
def reserve(self, tag: str, nbytes: int) -> Reservation:
|
|
if not self.can_reserve(tag, nbytes):
|
|
raise OutOfMemory(f"cannot reserve {nbytes} bytes for tag {tag!r} "
|
|
f"(available={self.available}, used={self._by_tag.get(tag, 0)})")
|
|
res = Reservation(next(_res_ctr), tag, nbytes)
|
|
self._reservations[res.res_id] = res
|
|
self.reserved += nbytes
|
|
self._by_tag[tag] = self._by_tag.get(tag, 0) + nbytes
|
|
return res
|
|
|
|
def release(self, res: Reservation) -> None:
|
|
if not res.active:
|
|
return
|
|
res.active = False
|
|
self._reservations.pop(res.res_id, None)
|
|
self.reserved -= res.nbytes
|
|
self._by_tag[res.tag] = max(0, self._by_tag.get(res.tag, 0) - res.nbytes)
|
|
|
|
# component-granular sleep/wake (tags = component names) ------------------- #
|
|
def sleep(self, tags: list[str]) -> int:
|
|
freed = 0
|
|
for tag in tags:
|
|
self._asleep.add(tag)
|
|
for res in [r for r in self._reservations.values() if r.tag == tag]:
|
|
freed += res.nbytes
|
|
self.release(res)
|
|
return freed
|
|
|
|
def wake(self, tags: list[str]) -> None:
|
|
for tag in tags:
|
|
self._asleep.discard(tag)
|