| """Server-only runtime shared by the SDTT and Di4C sampling entrypoints.""" |
|
|
| from __future__ import annotations |
|
|
| from contextlib import contextmanager |
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
| import random |
| import re |
| import sys |
| import tempfile |
| from collections.abc import Mapping |
|
|
| import yaml |
|
|
| from dlb.adapters.conditional_runtime import load_conditioning_batch, token_project_fn |
|
|
|
|
| class TokenizerBinding: |
| __slots__ = ("tokenizer_id", "revision", "snapshot") |
|
|
| def __init__(self, tokenizer_id: str, revision: str, snapshot: Path) -> None: |
| self.tokenizer_id = tokenizer_id |
| self.revision = revision |
| self.snapshot = snapshot |
|
|
|
|
| def require_file(path: Path, label: str) -> Path: |
| path = path.absolute() |
| if path.is_symlink() or not path.is_file() or path.stat().st_size <= 0: |
| raise ValueError(f"{label} is missing or unsafe: {path}") |
| return path |
|
|
|
|
| def require_directory(path: Path, label: str) -> Path: |
| path = path.absolute() |
| if path.is_symlink() or not path.is_dir(): |
| raise ValueError(f"{label} is missing or unsafe: {path}") |
| return path |
|
|
|
|
| def sha256_file(path: Path) -> str: |
| digest = hashlib.sha256() |
| with require_file(path, "hashed file").open("rb") as handle: |
| for chunk in iter(lambda: handle.read(1024 * 1024), b""): |
| digest.update(chunk) |
| return digest.hexdigest() |
|
|
|
|
| def verify_sha256(path: Path, expected: str, label: str) -> Path: |
| path = require_file(path, label) |
| if ( |
| len(expected) != 64 |
| or any(character not in "0123456789abcdef" for character in expected) |
| or sha256_file(path) != expected |
| ): |
| raise ValueError(f"{label} SHA-256 differs from canonical provenance") |
| return path |
|
|
|
|
| def parse_bool(value: str) -> bool: |
| if value == "true": |
| return True |
| if value == "false": |
| return False |
| raise ValueError("boolean values must be 'true' or 'false'") |
|
|
|
|
| @contextmanager |
| def offline_huggingface(enabled: bool): |
| names = ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE") |
| previous = {name: os.environ.get(name) for name in names} |
| if enabled: |
| os.environ.update({name: "1" for name in names}) |
| try: |
| yield |
| finally: |
| for name, value in previous.items(): |
| if value is None: |
| os.environ.pop(name, None) |
| else: |
| os.environ[name] = value |
|
|
|
|
| def install_upstream(upstream_root: Path) -> None: |
| source = require_directory(upstream_root / "src", "upstream Python source") |
| source_text = str(source) |
| if source_text not in sys.path: |
| sys.path.insert(0, source_text) |
|
|
|
|
| def seed_everything(seed: int) -> None: |
| random.seed(seed) |
| try: |
| import numpy |
|
|
| numpy.random.seed(seed % (2**32)) |
| except ImportError: |
| pass |
| import torch |
|
|
| torch.manual_seed(seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(seed) |
|
|
|
|
| def load_tokenizer(snapshot: Path): |
| from transformers import AutoTokenizer |
|
|
| return AutoTokenizer.from_pretrained( |
| str(require_directory(snapshot, "locked tokenizer snapshot")), |
| local_files_only=True, |
| trust_remote_code=False, |
| ) |
|
|
|
|
| def load_config(path: Path, expected_sha256: str): |
| from omegaconf import OmegaConf |
|
|
| return OmegaConf.load(verify_sha256(path, expected_sha256, "sampling config")) |
|
|
|
|
| def checkpoint_state( |
| path: Path, expected_sha256: str |
| ) -> tuple[dict[str, object], object | None]: |
| path = verify_sha256(path, expected_sha256, "student checkpoint") |
| if path.suffix == ".safetensors": |
| from safetensors.torch import load_file |
|
|
| return dict(load_file(str(path))), None |
| import torch |
|
|
| |
| |
| |
| payload = torch.load(str(path), map_location="cpu", weights_only=False) |
| if not isinstance(payload, dict): |
| raise ValueError("student checkpoint is not a mapping") |
| state = payload.get("state_dict", payload) |
| if not isinstance(state, dict): |
| raise ValueError("student checkpoint state_dict is invalid") |
| config = None |
| hyperparameters = payload.get("hyper_parameters") |
| if isinstance(hyperparameters, dict): |
| config = hyperparameters.get("config") |
| return dict(state), config |
|
|
|
|
| _MISSING = object() |
| _STANDARD_REFERENCE = re.compile(r"^\$\{([A-Za-z_][A-Za-z0-9_.]*)\}$") |
|
|
|
|
| def _select_config_value( |
| config: object, path: str, references: frozenset[str] = frozenset() |
| ) -> object: |
| """Resolve one selected scalar without walking unrelated Hydra config nodes.""" |
|
|
| try: |
| from omegaconf import OmegaConf |
|
|
| if OmegaConf.is_config(config): |
| try: |
| return OmegaConf.select( |
| config, |
| path, |
| default=_MISSING, |
| throw_on_resolution_failure=True, |
| ) |
| except Exception as error: |
| raise ValueError( |
| f"sampling config field {path} cannot be resolved" |
| ) from error |
| except ImportError: |
| pass |
| value = config |
| for component in path.split("."): |
| try: |
| value = ( |
| value[component] |
| if isinstance(value, Mapping) |
| else getattr(value, component) |
| ) |
| except (AttributeError, KeyError, TypeError): |
| return _MISSING |
| if isinstance(value, str): |
| reference = _STANDARD_REFERENCE.fullmatch(value) |
| if reference is not None: |
| target = reference.group(1) |
| if target in references or target == path: |
| raise ValueError(f"sampling config has a cyclic reference at {path}") |
| return _select_config_value(config, target, references | {path}) |
| return value |
|
|
|
|
| def _nested(config: object, path: str) -> object: |
| value = _select_config_value(config, path) |
| if value is _MISSING: |
| raise ValueError(f"sampling config is missing required field {path}") |
| return value |
|
|
|
|
| def _optional_nested(config: object, path: str) -> object: |
| return _select_config_value(config, path) |
|
|
|
|
| _CHECKPOINT_CONFIG_FIELDS = ( |
| "T", |
| "time_conditioning", |
| "model.type", |
| "model.hidden_size", |
| "model.cond_dim", |
| "model.length", |
| "model.n_blocks", |
| "model.n_heads", |
| "model.scale_by_sigma", |
| "model.dropout", |
| "model.tie_word_embeddings", |
| "noise.type", |
| "noise.sigma_min", |
| "noise.sigma_max", |
| "training.ema", |
| "training.antithetic_sampling", |
| "training.importance_sampling", |
| "training.sampling_eps", |
| "training.change_of_variables", |
| "tokenizer.name", |
| "parameterization.name", |
| "parameterization.log_loss_buckets", |
| "parameterization.start_from_hf", |
| "parameterization.distill_mode", |
| "parameterization.num_distill_steps", |
| "parameterization.min_num_sampling_steps", |
| "parameterization.grow_dt_every", |
| "parameterization.orig_num_sampling_steps", |
| "parameterization.sampling_mode", |
| "parameterization.loss_precision", |
| "parameterization.reset_optimizer_on_growth", |
| "parameterization.use_ema_on_growth", |
| ) |
|
|
| _OPTIONAL_CHECKPOINT_CONFIG_FIELDS = ("model.causal", "is_di4c") |
|
|
|
|
| def validate_embedded_config( |
| authoritative: object, |
| embedded: object | None, |
| *, |
| allow_missing_embedded_config: bool = False, |
| ) -> None: |
| if embedded is None: |
| if allow_missing_embedded_config: |
| return |
| raise ValueError("Lightning checkpoint does not embed its training config") |
| for path in (*_CHECKPOINT_CONFIG_FIELDS, *_OPTIONAL_CHECKPOINT_CONFIG_FIELDS): |
| expected = _optional_nested(authoritative, path) |
| if expected is _MISSING: |
| continue |
| observed = _optional_nested(embedded, path) |
| if observed is _MISSING: |
| raise ValueError(f"embedded checkpoint config is missing required field {path}") |
| if observed != expected: |
| raise ValueError( |
| f"embedded checkpoint config differs at {path}: " |
| f"expected {expected!r}, found {observed!r}" |
| ) |
|
|
|
|
| def validate_sampling_config( |
| config: object, |
| *, |
| binding: TokenizerBinding, |
| sequence_length: int, |
| require_di4c: bool, |
| ) -> None: |
| for path in _CHECKPOINT_CONFIG_FIELDS: |
| _nested(config, path) |
| tokenizer_name = _nested(config, "tokenizer.name") |
| if tokenizer_name != binding.tokenizer_id and Path(str(tokenizer_name)).resolve() != binding.snapshot.resolve(): |
| raise ValueError("sampling config tokenizer differs from canonical data contract") |
| if _nested(config, "model.length") != sequence_length: |
| raise ValueError("sampling config sequence length differs from canonical dataset") |
| if _nested(config, "data_preprocess.seq_len") != sequence_length: |
| raise ValueError("sampling config preprocessing length differs from canonical dataset") |
| if _nested(config, "parameterization.sampling_mode") != "ancestral": |
| raise ValueError("sampling config does not use ancestral sampling") |
| di4c_identity = _optional_nested(config, "is_di4c") |
| if require_di4c: |
| if di4c_identity is not True: |
| raise ValueError("sampling config Di4C identity is inconsistent") |
| if _nested(config, "parameterization.start_from_hf") is not False: |
| raise ValueError("Di4C sampling config must instantiate the local student") |
| elif di4c_identity not in {_MISSING, False}: |
| raise ValueError("sampling config Di4C identity is inconsistent") |
|
|
|
|
| def load_tokenizer_binding( |
| data_config_path: Path, |
| downloads_manifest_path: Path, |
| dataset_id: str, |
| requested_snapshot: Path, |
| ) -> TokenizerBinding: |
| data_config_path = require_file(data_config_path, "data config") |
| downloads_manifest_path = require_file( |
| downloads_manifest_path, "download manifest" |
| ) |
| try: |
| configuration = yaml.safe_load(data_config_path.read_text(encoding="utf-8")) |
| downloads = json.loads(downloads_manifest_path.read_text(encoding="utf-8")) |
| tokenizer_id = configuration["datasets"][dataset_id]["tokenizer"] |
| revision = configuration["models"][tokenizer_id] |
| record = downloads["models"][tokenizer_id] |
| except (OSError, UnicodeDecodeError, json.JSONDecodeError, KeyError, TypeError) as error: |
| raise ValueError("locked tokenizer metadata is invalid") from error |
| if downloads.get("schema_version") != 1: |
| raise ValueError("download manifest schema is not supported") |
| if ( |
| not isinstance(tokenizer_id, str) |
| or not isinstance(revision, str) |
| or len(revision) != 40 |
| ): |
| raise ValueError("data config tokenizer binding is invalid") |
| if record.get("repo_id") != tokenizer_id or record.get("revision") != revision: |
| raise ValueError("download manifest tokenizer binding differs from data config") |
| recorded = record.get("snapshot_path") |
| if not isinstance(recorded, str) or not recorded: |
| raise ValueError("download manifest tokenizer snapshot is missing") |
| snapshot = Path(recorded) |
| if not snapshot.is_absolute(): |
| snapshot = data_config_path.parents[1] / snapshot |
| requested_snapshot = requested_snapshot.absolute() |
| if requested_snapshot != snapshot: |
| raise ValueError("requested tokenizer snapshot differs from download manifest") |
| parts = snapshot.parts |
| try: |
| recorded_revision = parts[parts.index("snapshots") + 1] |
| except (ValueError, IndexError) as error: |
| raise ValueError("download manifest tokenizer snapshot path is invalid") from error |
| if recorded_revision != revision: |
| raise ValueError("download manifest tokenizer snapshot revision differs from data config") |
| require_directory(snapshot, "locked tokenizer snapshot") |
| return TokenizerBinding(tokenizer_id, revision, snapshot) |
|
|
|
|
| def instantiate_without_training_teacher(model_type, config, tokenizer): |
| """Instantiate the released student architecture without loading a network teacher.""" |
|
|
| original = model_type.prepare_teacher_and_student |
|
|
| def sampling_only(self): |
| self.teacher = [] |
|
|
| model_type.prepare_teacher_and_student = sampling_only |
| try: |
| return model_type(config, tokenizer, verbose=False) |
| finally: |
| model_type.prepare_teacher_and_student = original |
|
|
|
|
| def normalize_state_dict(state: dict[str, object]) -> dict[str, object]: |
| return {key.replace("_orig_mod.", ""): value for key, value in state.items()} |
|
|
|
|
| def write_capture_atomic( |
| output: Path, |
| *, |
| model, |
| tokenizer, |
| sample_count: int, |
| batch_size: int, |
| num_steps: int, |
| seq_len: int, |
| sampler: str, |
| generation_mode: str = "unconditional", |
| conditioning_manifest: Path | None = None, |
| conditioning_manifest_sha256: str | None = None, |
| prompt_count: int | None = None, |
| diversity_prompt_count: int | None = None, |
| completions_per_diversity_prompt: int | None = None, |
| ) -> None: |
| """Sample exact batches, immediately move IDs to CPU, and atomically publish JSON.""" |
|
|
| if min(sample_count, batch_size, num_steps, seq_len) <= 0: |
| raise ValueError("sampling counts and lengths must be positive") |
| if sampler != "ancestral": |
| raise ValueError("distilled language baselines require ancestral sampling") |
| if generation_mode not in {"unconditional", "conditional_prefix"}: |
| raise ValueError(f"unsupported generation mode: {generation_mode!r}") |
| output = output.absolute() |
| if not output.is_absolute() or output.is_symlink(): |
| raise ValueError(f"output path is unsafe: {output}") |
| output.parent.mkdir(parents=True, exist_ok=True) |
| descriptor, temporary_name = tempfile.mkstemp( |
| prefix=f".{output.name}.", suffix=".partial", dir=output.parent |
| ) |
| temporary = Path(temporary_name) |
| written = 0 |
| try: |
| with os.fdopen(descriptor, "w", encoding="utf-8") as handle: |
| handle.write('{"schema":"dlb-upstream-token-capture-v1","samples":[') |
| if generation_mode == "conditional_prefix": |
| if ( |
| conditioning_manifest is None |
| or conditioning_manifest_sha256 is None |
| or prompt_count is None |
| or diversity_prompt_count is None |
| or completions_per_diversity_prompt is None |
| ): |
| raise ValueError("conditional distilled sampling requires a prompt contract") |
| batches: list[tuple[int, int, int]] = [] |
| for start in range(0, prompt_count, batch_size): |
| batches.append((0, start, min(batch_size, prompt_count - start))) |
| for completion_id in range(1, completions_per_diversity_prompt): |
| for start in range(0, diversity_prompt_count, batch_size): |
| batches.append( |
| ( |
| completion_id, |
| start, |
| min(batch_size, diversity_prompt_count - start), |
| ) |
| ) |
| else: |
| batches = [(-1, written, min(batch_size, sample_count - written)) for written in range(0, sample_count, batch_size)] |
| for completion_id, prompt_start, current in batches: |
| current = min(batch_size, sample_count - written) |
| batch = None |
| project_fn = lambda value: value |
| if generation_mode == "conditional_prefix": |
| batch = load_conditioning_batch( |
| conditioning_manifest, |
| conditioning_manifest_sha256, |
| completion_id=completion_id, |
| prompt_start=prompt_start, |
| batch_size=current, |
| device=str(model.device), |
| vocab_size=len(tokenizer), |
| ) |
| project_fn = token_project_fn(batch.prefix_token_ids) |
| result = project_fn( |
| model.sample( |
| n_samples=current, |
| num_steps=num_steps, |
| seq_len=seq_len, |
| sampler=sampler, |
| verbose=False, |
| project_fn=project_fn, |
| ) |
| ) |
| rows = result.detach().cpu().tolist() |
| if not isinstance(rows, list) or len(rows) != current: |
| raise ValueError("upstream sampler returned the wrong batch size") |
| texts = list(tokenizer.batch_decode(rows)) |
| if len(texts) != current: |
| raise ValueError("tokenizer returned the wrong decoded batch size") |
| for batch_index, (row, text) in enumerate(zip(rows, texts, strict=True)): |
| if not isinstance(row, list) or len(row) != seq_len: |
| raise ValueError("upstream sampler returned a noncanonical sequence") |
| if not isinstance(text, str) or not text.strip(): |
| raise ValueError("upstream tokenizer returned empty text") |
| payload = {"sample_id": written, "text": text, "token_ids": row} |
| if batch is not None: |
| prefix = batch.prefix_token_ids[batch_index].detach().cpu().tolist() |
| reference = batch.reference_token_ids[batch_index].detach().cpu().tolist() |
| if row[: len(prefix)] != prefix: |
| raise ValueError("conditional sampler returned a prefix mismatch") |
| payload.update( |
| { |
| "prompt_id": batch.prompt_ids[batch_index], |
| "completion_id": completion_id, |
| "source_index": batch.source_indices[batch_index], |
| "prefix_token_ids": prefix, |
| "reference_token_ids": reference, |
| "full_token_ids": row, |
| "prefix_text": tokenizer.decode(prefix), |
| "continuation_text": tokenizer.decode(row[len(prefix) :]), |
| "reference_text": tokenizer.decode(reference), |
| "full_text": text, |
| "prefix_exact_match": True, |
| } |
| ) |
| if written: |
| handle.write(",") |
| json.dump( |
| payload, |
| handle, |
| ensure_ascii=False, |
| separators=(",", ":"), |
| ) |
| written += 1 |
| del result, rows, texts |
| handle.write("]}") |
| handle.flush() |
| os.fsync(handle.fileno()) |
| if written != sample_count: |
| raise ValueError(f"expected {sample_count} samples, wrote {written}") |
| os.replace(temporary, output) |
| except BaseException: |
| temporary.unlink(missing_ok=True) |
| raise |
|
|
|
|
| def benchmark_model( |
| *, |
| model, |
| output: Path, |
| metadata_path: Path, |
| precision: str, |
| num_steps: int, |
| seq_len: int, |
| sampler: str, |
| generation_mode: str = "unconditional", |
| conditioning_manifest: Path | None = None, |
| conditioning_manifest_sha256: str | None = None, |
| prompt_count: int | None = None, |
| diversity_prompt_count: int | None = None, |
| completions_per_diversity_prompt: int | None = None, |
| ): |
| """Time the already-materialized student's real one-sample tensor callback.""" |
|
|
| if min(num_steps, seq_len) <= 0 or sampler != "ancestral": |
| raise ValueError("distilled benchmark sampling arguments are invalid") |
| from dlb.timing import benchmark_and_publish |
|
|
| project_fn = lambda value: value |
| if generation_mode == "conditional_prefix": |
| if conditioning_manifest is None or conditioning_manifest_sha256 is None: |
| raise ValueError("conditional benchmark requires a prompt manifest") |
| del prompt_count, diversity_prompt_count, completions_per_diversity_prompt |
| batch = load_conditioning_batch( |
| conditioning_manifest, |
| conditioning_manifest_sha256, |
| completion_id=0, |
| prompt_start=0, |
| batch_size=1, |
| device=str(model.device), |
| vocab_size=len(model.tokenizer), |
| ) |
| project_fn = token_project_fn(batch.prefix_token_ids) |
| return benchmark_and_publish( |
| lambda: project_fn(model.sample( |
| n_samples=1, |
| num_steps=num_steps, |
| seq_len=seq_len, |
| sampler=sampler, |
| verbose=False, |
| project_fn=project_fn, |
| )), |
| model=model, |
| output=output, |
| metadata_path=metadata_path, |
| precision=precision, |
| ) |
|
|
|
|
| def configure_for_sampling(config, *, tokenizer_snapshot: Path, seq_len: int) -> None: |
| config.tokenizer.name = str(tokenizer_snapshot) |
| config.model.length = seq_len |
| config.data_preprocess.seq_len = seq_len |
| config.mode = "sample" |
| config.compile = False |
|
|
|
|
| def materialize_model( |
| *, |
| model_type, |
| config, |
| tokenizer, |
| state: dict[str, object], |
| strict: bool, |
| ): |
| model = instantiate_without_training_teacher(model_type, config, tokenizer) |
| result = model.load_state_dict(normalize_state_dict(state), strict=strict) |
| if not strict: |
| missing = [key for key in result.missing_keys if not key.startswith("teacher.")] |
| unexpected = [key for key in result.unexpected_keys if not key.startswith("teacher.")] |
| if missing or unexpected: |
| raise ValueError( |
| "student checkpoint does not match the configured architecture: " |
| f"missing={missing[:8]}, unexpected={unexpected[:8]}" |
| ) |
| model.eval() |
| if not __import__("torch").cuda.is_available(): |
| raise RuntimeError("distilled sampling requires a CUDA server") |
| model.cuda() |
| return model |
|
|
