Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion src/api/inference.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,10 @@
.env({"IMAGE_REV": "6"})
.add_local_dir(str(Path(__file__).resolve().parent.parent), "/root/interscript-ml", copy=True)
# deterministic mounts for the index-driven resolver + the index itself
.add_local_file(str(Path(__file__).resolve().parent / "model_resolution.py"), "/root/model_resolution.py")
.add_local_file(
str(Path(__file__).resolve().parent / "model_resolution.py"),
"/root/model_resolution.py",
)
.add_local_file(str(_REPO_ROOT / "models.yaml"), "/root/models.yaml")
)

Expand Down
3 changes: 2 additions & 1 deletion src/api/model_resolution.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,10 @@

from __future__ import annotations

import yaml
from pathlib import Path

import yaml


def load_index(path: str | Path = "models.yaml") -> dict:
with open(path, encoding="utf-8") as fh:
Expand Down
2 changes: 1 addition & 1 deletion src/gpu/modal_distill.py
Original file line number Diff line number Diff line change
Expand Up @@ -1617,7 +1617,7 @@ def probe_pkm_gates(spec_id: str = "ara-diac-small-pkm") -> dict:
raise RuntimeError(f"no step checkpoints under {out_root}")

_ensure_src_path()
from gpu.pkm import ProductKeyMemory, inject_pkm
from gpu.pkm import inject_pkm

student = AutoModelForSeq2SeqLM.from_pretrained(spec["student_init"])
inject_pkm(student, **spec["pkm"])
Expand Down
19 changes: 15 additions & 4 deletions src/gpu/muon.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,16 +41,27 @@ def __init__(self, params, lr: float = 0.01, momentum: float = 0.95,
weight_decay: float = 0.0) -> None:
super().__init__(
list(params),
dict(lr=lr, momentum=momentum, nesterov=nesterov,
ns_steps=ns_steps, weight_decay=weight_decay, adamw=False),
{
"lr": lr,
"momentum": momentum,
"nesterov": nesterov,
"ns_steps": ns_steps,
"weight_decay": weight_decay,
"adamw": False,
},
)

def add_adamw_group(self, params, lr: float = 1e-4, betas=(0.9, 0.999),
weight_decay: float = 0.0) -> None:
"""Embedding-like parameters: standard AdamW math, shared
scheduler (the cosine scales every group's lr)."""
self.add_param_group(dict(params=list(params), lr=lr, betas=tuple(betas),
weight_decay=weight_decay, adamw=True))
self.add_param_group({
"params": list(params),
"lr": lr,
"betas": tuple(betas),
"weight_decay": weight_decay,
"adamw": True,
})

@torch.no_grad()
def step(self, closure=None): # noqa: ARG002
Expand Down
3 changes: 1 addition & 2 deletions src/imf/parity.py
Original file line number Diff line number Diff line change
Expand Up @@ -277,9 +277,8 @@ def run_margin_analysis(model, zip_path, pairs, max_len: int = 256) -> MarginRep
def write_margin_report(report: MarginReport, out_path: Path | str) -> Path:
"""Emit the margin analysis as JSON next to a release zip (diagnostic
artifact; the release gate remains the CER parity block)."""
from dataclasses import asdict

import json
from dataclasses import asdict

out_path = Path(out_path)
out_path.parent.mkdir(parents=True, exist_ok=True)
Expand Down
14 changes: 7 additions & 7 deletions tests/test_decode_guard.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,24 +22,24 @@ def __init__(self):
self.step = 0

def get_outputs(self):
class O:
class Out:
def __init__(self, name):
self.name = name

return [O("logits"), O("present_k"), O("present_v")]
return [Out("logits"), Out("present_k"), Out("present_v")]

def get_inputs(self):
class I:
class In:
def __init__(self, name, shape, typ="tensor(float)"):
self.name = name
self.shape = shape
self.type = typ

return [
I("input_ids", [1, "seq"]),
I("encoder_hidden_states", [1, "seq", 4]),
I("past_k", [1, 4, "past", 8]),
I("past_v", [1, 4, "past", 8]),
In("input_ids", [1, "seq"]),
In("encoder_hidden_states", [1, "seq", 4]),
In("past_k", [1, 4, "past", 8]),
In("past_v", [1, 4, "past", 8]),
]

def run(self, _, feeds):
Expand Down
1 change: 0 additions & 1 deletion tests/test_imf_parity.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@
make_fixture_checkpoint,
)
from imf.parity import ( # noqa: E402
MarginReport,
ParityReport,
_margin_stats,
run_margin_analysis,
Expand Down
Loading