| """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", |
| ] |
|
|
