File size: 5,022 Bytes
e69b72a | 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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """Shared checkpoint construction for the STRATA-COMPOSE replacement frontier."""
from __future__ import annotations
import hashlib
from pathlib import Path
import torch
from strata.modeling.compose import ExportedAlgebraBlock, PrunableAlgebraLM
from strata.modeling.ph_core import DenseContinuationForCausalLM
from strata.modeling.ph_pat import PHPATConfig
def sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(16 << 20), b""):
digest.update(chunk)
return digest.hexdigest()
def _load_with_periodic_position_extension(
model: torch.nn.Module,
checkpoint_path: Path,
*,
position_name: str,
) -> dict[str, object]:
"""Load a checkpoint while periodically extending its sole position table."""
source = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
target = model.state_dict()
if set(source) != set(target):
missing = sorted(set(target) - set(source))
extra = sorted(set(source) - set(target))
raise ValueError(f"checkpoint keys differ: missing={missing[:3]}, extra={extra[:3]}")
source_positions = source[position_name]
target_positions = target[position_name]
if (
source_positions.ndim != 2
or target_positions.ndim != 2
or source_positions.shape[1] != target_positions.shape[1]
or source_positions.shape[0] > target_positions.shape[0]
):
raise ValueError("position embedding shapes cannot be extended")
state = {}
for name, value in source.items():
if name == position_name:
positions = target_positions.clone()
source_length = source_positions.shape[0]
for start in range(0, target_positions.shape[0], source_length):
count = min(source_length, target_positions.shape[0] - start)
positions[start:start + count] = source_positions[:count]
state[name] = positions
else:
if value.shape != target[name].shape:
raise ValueError(f"non-position tensor changed shape: {name}")
state[name] = value
model.load_state_dict(state, strict=True)
return {
"source_context": source_positions.shape[0],
"target_context": target_positions.shape[0],
"extension": "periodic_copy",
"non_position_tensors_exact": True,
}
def load_dense_base(
config_path: Path,
checkpoint_path: Path,
device: torch.device,
*,
allow_position_extension: bool = False,
) -> tuple[DenseContinuationForCausalLM, PHPATConfig]:
config = PHPATConfig.from_json_file(config_path).with_arm("dense_global")
model = DenseContinuationForCausalLM(config, gradient_checkpointing=False)
if allow_position_extension:
_load_with_periodic_position_extension(
model, checkpoint_path, position_name="position_embeddings.weight",
)
else:
model.load_state_dict(
torch.load(checkpoint_path, map_location="cpu", weights_only=True),
strict=True,
)
dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
return model.to(device=device, dtype=dtype), config
def build_exported_stage(
*,
config_path: Path,
base_checkpoint: Path,
stage_checkpoint: Path,
removed_layers: tuple[int, ...],
device: torch.device,
gamma_max: float = 0.02,
allow_position_extension: bool = False,
) -> tuple[PrunableAlgebraLM, PHPATConfig]:
base, config = load_dense_base(
config_path, base_checkpoint, device,
allow_position_extension=allow_position_extension,
)
model = PrunableAlgebraLM(
base,
config,
removed_layers,
gamma_max=gamma_max,
)
model.export_layers(removed_layers)
if allow_position_extension:
_load_with_periodic_position_extension(
model, stage_checkpoint,
position_name="base_model.position_embeddings.weight",
)
else:
model.load_state_dict(
torch.load(stage_checkpoint, map_location=device, weights_only=True),
strict=True,
)
if model.physically_removed_layers() != tuple(sorted(removed_layers)):
raise AssertionError("stage checkpoint does not match its physical-removal manifest")
if model.dense_modules_in_replacement_layers():
raise AssertionError("dense attention remains in a physically removed layer")
return model, config
def graph_adapter_state(model: PrunableAlgebraLM, layer: int) -> dict[str, torch.Tensor]:
block = model.base_model.blocks[layer]
if not isinstance(block, ExportedAlgebraBlock):
raise TypeError(f"graph-read layer {layer} is not physically exported")
return {name: value.detach().cpu() for name, value in block.graph_adapter.state_dict().items()}
__all__ = [
"build_exported_stage",
"graph_adapter_state",
"load_dense_base",
"sha256_file",
]
|