| """Predictor architecture metadata carried inside safetensors checkpoints. |
| |
| Stage-1 trainers write ``predictor_config`` into the safetensors metadata so |
| Stage-2 DMD training and the evaluators can rebuild the exact module (source |
| layer, input variant, ATC hyper-parameters) without extra side files. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| from safetensors import safe_open |
|
|
| from predictor_training.single_block import ( |
| SingleBlockPredictor, |
| initialize_predictor_block, |
| ) |
|
|
|
|
| SUPPORTED_INPUT_VARIANTS = ("self_forcing", "disca", "atc") |
| LEGACY_PREDICTOR_CONFIG = {"source_layer": 17, "input_variant": "self_forcing"} |
|
|
|
|
| def predictor_config_metadata(config: dict[str, Any]) -> dict[str, str]: |
| """Serialize a Predictor config as safetensors metadata.""" |
| return { |
| "predictor_config": json.dumps( |
| config, sort_keys=True, separators=(",", ":") |
| ) |
| } |
|
|
|
|
| def read_predictor_config( |
| weights: str | Path, expected_input_variant: str = "auto" |
| ) -> dict[str, Any]: |
| """Recover the Predictor architecture from a checkpoint's metadata. |
| |
| Checkpoints written before metadata existed are the original concat |
| Predictor on Teacher layer 17; they are accepted only when the caller does |
| not demand a different variant. |
| """ |
| with safe_open(str(weights), framework="pt", device="cpu") as handle: |
| metadata = handle.metadata() or {} |
| raw = metadata.get("predictor_config") |
| if raw is None: |
| if expected_input_variant not in ("auto", "self_forcing"): |
| raise ValueError( |
| f"{weights} has no predictor_config metadata; cannot assume " |
| f"input_variant={expected_input_variant!r}" |
| ) |
| config = dict(LEGACY_PREDICTOR_CONFIG) |
| else: |
| config = json.loads(raw) |
| actual = str(config.get("input_variant", "self_forcing")) |
| if actual not in SUPPORTED_INPUT_VARIANTS: |
| raise ValueError(f"Unsupported input_variant={actual!r} in {weights}") |
| if expected_input_variant != "auto" and actual != expected_input_variant: |
| raise ValueError( |
| "Predictor input variant mismatch: " |
| f"expected={expected_input_variant} actual={actual}" |
| ) |
| config["input_variant"] = actual |
| config["source_layer"] = int(config.get("source_layer", 17)) |
| return config |
|
|
|
|
| def build_single_block_predictor( |
| teacher: torch.nn.Module, |
| config: dict[str, Any], |
| *, |
| gradient_checkpointing: bool, |
| collect_diagnostics: bool = False, |
| ) -> SingleBlockPredictor: |
| """Instantiate a SingleBlockPredictor matching ``config`` (weights not loaded).""" |
| source_layer = int(config.get("source_layer", 17)) |
| predictor = SingleBlockPredictor( |
| block=initialize_predictor_block(teacher.blocks[source_layer], "teacher_full"), |
| dim=teacher.dim, |
| gradient_checkpointing=gradient_checkpointing, |
| input_variant=str(config.get("input_variant", "self_forcing")), |
| gate_mode=str(config.get("gate_mode", "baseline")), |
| gate_hidden_dim=int(config.get("gate_hidden_dim", 128)), |
| gate_initial_bias=float(config.get("gate_initial_bias", 4.6)), |
| gate_floor=float(config.get("gate_floor", 0.0)), |
| constant_gate=float(config.get("constant_gate", 1.0)), |
| atc_previous_scope=str(config.get("atc_previous_scope", "chunk")), |
| atc_freq_dim=int(config.get("atc_freq_dim", 256)), |
| atc_mlp_hidden_dim=int(config.get("atc_mlp_hidden_dim", 3072)), |
| atc_gate_hidden_dim=int(config.get("atc_gate_hidden_dim", 512)), |
| atc_transport_residual_scale=float( |
| config.get("atc_transport_residual_scale", 0.1) |
| ), |
| atc_gate_initial_probability=float( |
| config.get("atc_gate_initial_probability", 0.3) |
| ), |
| atc_collect_diagnostics=collect_diagnostics, |
| ) |
| predictor.source_layer = source_layer |
| return predictor |
|
|
