File size: 3,743 Bytes
e1acc15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | """Matched pilot work units, including masked tails and end-to-end estimates."""
from __future__ import annotations
import math
from typing import Protocol, TypedDict
from .loop import BatchPlan, Run, lr_multiplier
QDQ_SEMANTICS = "official act_qdq before dense/expert, output dtype rounding, identity STE backward; additive FP32 LoRA before output rounding; raw uncached taps"
class TokenCounts(Protocol):
def __len__(self) -> int: ...
def supervised_tokens(self, index: int) -> int: ...
class Units(TypedDict):
blocks: int
valid_supervised_tokens: int
optimizer_updates: int
accumulation: int
batch_ceiling: int
lr_schedule: str
learning_rates: list[float]
qdq_semantics: str
def work_units(data: TokenCounts, plan: BatchPlan, run: Run, step: int | None = None) -> Units:
stop = len(plan.batches) if step is None else step
if not 0 < stop <= len(plan.batches):
raise ValueError("invalid matched-unit checkpoint")
counts = [data.supervised_tokens(i) for batch in plan.batches[:stop] for i in batch]
if any(type(n) is not int or not 1 <= n <= 5 for n in counts):
raise ValueError("invalid supervised position count")
return {"blocks": len(counts), "valid_supervised_tokens": sum(counts), "optimizer_updates": stop,
"accumulation": 1, "batch_ceiling": run.batch_size,
"lr_schedule": ("linear-warmup-10pct-then-cosine; " if run.warmup_fraction else "")+"cosine-per-optimizer-update; initial lr then scheduler.step after update; total="+str(len(plan.batches)),
"learning_rates": [run.lr*lr_multiplier(run, s, len(plan.batches)) for s in range(stop)],
"qdq_semantics": QDQ_SEMANTICS}
def require_matched(arms: dict[str, dict[str, Units]]) -> None:
if set(arms) != {"REF", "FAST", "FAST-B"} or not arms["REF"]:
raise ValueError("matched units require all arms/checkpoints")
if arms["REF"] != arms["FAST"] or arms["REF"] != arms["FAST-B"]:
raise ValueError("pilot blocks/tokens/updates/accumulation/LR/QDQ units differ")
class Estimate(TypedDict):
label: str
training_gpu_hours: float
evaluation_gpu_hours: float
replay_gpu_hours: float
serving_gpu_hours: float
total_gpu_hours: float
ideal_two_solo_wall_hours: float
training_blocks: int
optimizer_updates: int
valid_supervised_tokens: int
def estimate(units: Units, blocks_per_second: dict[str, float], evaluation_rounds: int,
eval_seconds_per_round: float, replay_gpu_seconds: float, serving_gpu_seconds: float) -> Estimate:
costs = [*blocks_per_second.values(), eval_seconds_per_round, replay_gpu_seconds, serving_gpu_seconds]
if (set(blocks_per_second) != {"REF", "FAST", "FAST-B"} or any(not math.isfinite(x) or x <= 0 for x in costs)
or evaluation_rounds < 1):
raise ValueError("positive measured/declared costs required including replay and serving qualification")
training = sum(units["blocks"]/rate for rate in blocks_per_second.values())/3600
evaluation = 13*evaluation_rounds*eval_seconds_per_round/3600
replay, serving = replay_gpu_seconds/3600, serving_gpu_seconds/3600
total = training+evaluation+replay+serving
return {"label": "ESTIMATED; additive costs, perfect two-solo balance lower bound; excludes load/I/O/queue overhead",
"training_gpu_hours": training, "evaluation_gpu_hours": evaluation, "replay_gpu_hours": replay,
"serving_gpu_hours": serving, "total_gpu_hours": total, "ideal_two_solo_wall_hours": total/2,
"training_blocks": 3*units["blocks"], "optimizer_updates": 3*units["optimizer_updates"],
"valid_supervised_tokens": 3*units["valid_supervised_tokens"]}
|