File size: 4,028 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 | """Adapter for the pinned official RDLM LM1B SDE sampler."""
from pathlib import Path
import sys
from typing import Iterable
from dlb.adapters.base import AdapterError, BaseTeacherAdapter, CheckpointSelection
from dlb.runner import RunRequest
from dlb.schema import SampleRecord
class RDLMAdapter(BaseTeacherAdapter):
identity = "dlb.adapters.rdlm:v1"
upstream = "rdlm"
supported_models = frozenset({"rdlm"})
teacher_families = {"rdlm": "continuous_rdlm"}
batch_sizes = {("rdlm", "lm1b"): 8}
_LM1B_FILES = (
"LM1B/checkpoint.pth",
"LM1B/config.yaml",
"LM1B/sde.pkl",
)
def _asset_paths(
self, checkpoint: CheckpointSelection
) -> tuple[Path, Path, Path]:
missing = [name for name in self._LM1B_FILES if name not in checkpoint.required_files]
if missing:
raise AdapterError(
"canonical RDLM resource does not bind the saved LM1B trio: "
+ ", ".join(missing)
)
asset_root = checkpoint.path / "LM1B"
return (
asset_root / "checkpoint.pth",
asset_root / "config.yaml",
asset_root / "sde.pkl",
)
def render_command(
self, request: RunRequest, run_dir: Path, *, dry_run: bool
) -> list[str]:
root, length, _ = self._validate_request(request, run_dir)
if length != 128:
raise AdapterError(f"RDLM LM1B requires sequence length 128, found {length}")
checkpoint = self._resolve_checkpoint(root, request, dry_run=dry_run)
checkpoint_path, saved_config, saved_sde = self._asset_paths(checkpoint)
entrypoint = root / "upstreams" / "rdlm" / "main.py"
if not entrypoint.is_file() or entrypoint.is_symlink():
raise AdapterError(f"pinned upstream entrypoint is missing or unsafe: {entrypoint}")
capture_path = run_dir.resolve() / "upstream_token_ids.json"
hydra_dir = run_dir.resolve() / "upstream_hydra"
tokenizer_name = self._load_data_contract(root, request.dataset_id)["tokenizer"]
tokenizer_revision = self._tokenizer_revision(root, tokenizer_name)
tokenizer_snapshot = (
root
/ "data"
/ "raw"
/ "huggingface"
/ "hub"
/ f"models--{tokenizer_name}"
/ "snapshots"
/ tokenizer_revision
)
arguments = [
sys.executable,
"-B",
"-u",
"-m",
"dlb.adapters.capture",
f"--upstream-entrypoint={entrypoint}",
f"--capture-path={capture_path}",
"--capture-kind=rdlm",
f"--saved-config-path={saved_config}",
f"--saved-sde-path={saved_sde}",
f"--expected-samples={request.sample_count}",
f"--data-config-path={root / 'artifacts' / 'data.yaml'}",
f"--downloads-manifest-path={root / 'data' / 'manifests' / 'downloads.json'}",
f"--dataset-id={request.dataset_id}",
f"--tokenizer-snapshot={tokenizer_snapshot}",
*self._conditional_capture_flags(request),
"--",
"run_mode=sample",
"server=sample",
"exp=sample_lm1b",
"ngpus=1",
"use_wandb=False",
f"model_path={checkpoint_path}",
f"seed={request.seed}",
"sampling.predictor=grw",
f"sampling.steps={request.step_count}",
f"sampling.batch_per_gpu={min(8, request.sample_count)}",
"eval.entropy=False",
"eval.nll=False",
"eval.gen_ppl=False",
f"hydra.run.dir={hydra_dir}",
]
self._validate_argv(arguments)
return arguments
def convert_outputs(
self, request: RunRequest, run_dir: Path
) -> Iterable[SampleRecord]:
return self._convert_capture_outputs(
request, run_dir, retokenize_inexact_rows=True
)
adapter = RDLMAdapter()
|