File size: 3,967 Bytes
b4194d0 | 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 | """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
|