File size: 4,831 Bytes
784f572 | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """Validation and loading for the experiment coverage registry."""
from __future__ import annotations
from pathlib import Path
from typing import Literal
import yaml
from pydantic import BaseModel, ConfigDict, Field, model_validator
MANY_STEPS = [1, 2, 4, 8, 16, 32, 1024]
FEW_STEPS = [1, 2, 4, 8, 16, 32]
FIXED_1024_STEPS = [1024]
RDLM_OFFICIAL_STEPS = [1000, 1024]
MODEL_IDENTIFIERS = {
"flm": ("many", "dlb-flm", "flm", "flm"),
"fmlm": ("few", "dlb-flm", "flm", "flm"),
"langflow": ("many", "dlb-langflow", "langflow", "langflow"),
"duo": ("many", "dlb-duo", "duo", "duo"),
"duo_dcd": ("few", "dlb-duo", "duo", "duo"),
"mdlm": ("many", "dlb-mdlm", "mdlm", "mdlm"),
"candi": ("many", "dlb-candi", "candi", "candi"),
"rdlm": ("many", "dlb-rdlm", "rdlm", "rdlm"),
"mdlm_sdtt": ("few", "dlb-sdtt", "sdtt", "sdtt"),
"duo_di4c": ("few", "dlb-di4c", "di4c", "di4c"),
"mdlm_di4c": ("few", "dlb-di4c", "di4c", "di4c"),
}
Category = Literal["many", "few", "fixed_1024"]
SupportStatus = Literal["supported", "unsupported"]
Provenance = Literal["official", "reference_reproduction", "self_trained"]
class RegistryModel(BaseModel):
model_config = ConfigDict(extra="forbid")
class DatasetSupport(RegistryModel):
status: SupportStatus
provenance: Provenance | None = None
reason: str | None = None
train_recipe: str | None = Field(default=None, pattern=r"^[a-z][a-z0-9_]*$")
@model_validator(mode="after")
def validate_support_details(self) -> "DatasetSupport":
if self.status == "supported":
if self.provenance is None or self.reason is not None:
raise ValueError("supported datasets require provenance and no reason")
elif (
self.reason is None
or self.reason.strip() == ""
or self.provenance is not None
or self.train_recipe is not None
):
raise ValueError(
"unsupported datasets require a reason and no provenance or train recipe"
)
return self
class ModelRegistryEntry(RegistryModel):
category: Category
step_override: list[int] | None = None
environment: str = Field(pattern=r"^[a-z][a-z0-9-]*$")
adapter: str = Field(pattern=r"^[a-z][a-z0-9-]*$")
source: str = Field(pattern=r"^[a-z][a-z0-9-]*$")
datasets: dict[Literal["lm1b", "owt"], DatasetSupport]
@model_validator(mode="after")
def require_all_datasets(self) -> "ModelRegistryEntry":
if set(self.datasets) != {"lm1b", "owt"}:
raise ValueError("each model must declare lm1b and owt support")
if self.step_override is not None:
if self.step_override != sorted(set(self.step_override)):
raise ValueError("step overrides must be sorted and unique")
if any(step <= 0 for step in self.step_override):
raise ValueError("step overrides must be positive")
return self
class ExperimentRegistry(RegistryModel):
step_grids: dict[Category, list[int]]
models: dict[str, ModelRegistryEntry]
@model_validator(mode="after")
def validate_coverage(self) -> "ExperimentRegistry":
if self.step_grids != {
"many": MANY_STEPS,
"few": FEW_STEPS,
"fixed_1024": FIXED_1024_STEPS,
}:
raise ValueError("step grids must match the prescribed schedules")
if set(self.models) != set(MODEL_IDENTIFIERS):
raise ValueError("registry must contain the complete baseline model scope")
for model_id, model in self.models.items():
if (model.category, model.environment, model.adapter, model.source) != (
MODEL_IDENTIFIERS[model_id]
):
raise ValueError(f"registry identifiers do not match {model_id}")
if model_id == "rdlm":
if model.step_override != RDLM_OFFICIAL_STEPS:
raise ValueError("RDLM must use the official/default step override")
elif model.step_override is not None:
raise ValueError(f"{model_id} must use its category step grid")
return self
def load_registry(path: Path) -> ExperimentRegistry:
"""Load and validate a registry YAML document from *path*."""
with path.open(encoding="utf-8") as registry_file:
document = yaml.safe_load(registry_file)
if not isinstance(document, dict):
raise ValueError("registry document must be a mapping")
return ExperimentRegistry.model_validate(document)
def step_grid_for_model(registry: ExperimentRegistry, model_id: str) -> list[int]:
"""Return the canonical step grid for one logical baseline model."""
model = registry.models[model_id]
return list(model.step_override or registry.step_grids[model.category])
|