mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-09-21 17:00:50 +02:00
473 lines
16 KiB
Python
473 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate the deterministic, synthetic Wave 2 checkpoint fixture corpus."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from hashlib import sha256
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
from safetensors.torch import save_file
|
|
|
|
from obliteratus.checkpoint_fragments import (
|
|
Padding,
|
|
Replica,
|
|
TensorFragment,
|
|
reconstruct_logical_tensor,
|
|
validate_fragments,
|
|
)
|
|
|
|
|
|
GENERATOR_VERSION = "1.0.0"
|
|
CORPUS_LIMITS = {
|
|
"max_case_bytes": 65536,
|
|
"max_cases": 16,
|
|
"max_files_per_case": 16,
|
|
"max_tensors_per_case": 16,
|
|
}
|
|
|
|
|
|
def _json_bytes(value: object) -> bytes:
|
|
return (
|
|
json.dumps(value, indent=2, sort_keys=True, ensure_ascii=True, allow_nan=False) + "\n"
|
|
).encode("utf-8")
|
|
|
|
|
|
def _write_json(path: Path, value: object) -> None:
|
|
path.write_bytes(_json_bytes(value))
|
|
|
|
|
|
def _tensor_digest(tensor: torch.Tensor) -> str:
|
|
raw = tensor.detach().contiguous().reshape(-1).view(torch.uint8).numpy().tobytes()
|
|
return f"sha256:{sha256(raw).hexdigest()}"
|
|
|
|
|
|
def _fragment(
|
|
fragment_id: str,
|
|
payload: torch.Tensor,
|
|
*,
|
|
payload_file: str,
|
|
global_shape: tuple[int, ...],
|
|
offset: tuple[int, ...],
|
|
extent: tuple[int, ...] | None = None,
|
|
logical_tensor_id: str,
|
|
component_id: str = "model",
|
|
role: str = "parameter",
|
|
padding: Padding | None = None,
|
|
replica: Replica | None = None,
|
|
partition_axes: tuple[int, ...] = (0,),
|
|
tie_group_id: str | None = None,
|
|
topology_coordinates: tuple[tuple[str, int], ...] = (("rank", 0),),
|
|
) -> tuple[TensorFragment, str]:
|
|
extent = extent if extent is not None else tuple(payload.shape)
|
|
padding = padding or Padding.zeros(len(global_shape))
|
|
replica = replica or Replica.unique()
|
|
logical_payload = payload
|
|
if global_shape:
|
|
logical_payload = payload[
|
|
tuple(
|
|
slice(before, before + size)
|
|
for before, size in zip(padding.before, extent, strict=True)
|
|
)
|
|
]
|
|
fragment = TensorFragment(
|
|
fragment_id=fragment_id,
|
|
component_id=component_id,
|
|
fqn=logical_tensor_id,
|
|
role=role,
|
|
dtype=str(payload.dtype).removeprefix("torch."),
|
|
global_shape=global_shape,
|
|
local_shape=tuple(payload.shape),
|
|
element_offset=offset,
|
|
element_extent=extent,
|
|
padding=padding,
|
|
shard_file_id=payload_file,
|
|
shard_digest_ref=f"source-digest:{payload_file}",
|
|
fragment_digest=_tensor_digest(logical_payload),
|
|
replica=replica,
|
|
partition_axes=partition_axes if global_shape else (),
|
|
logical_tensor_id=logical_tensor_id,
|
|
tie_group_id=tie_group_id,
|
|
shared_storage_id=tie_group_id,
|
|
topology_coordinates=topology_coordinates,
|
|
evidence_refs=(f"synthetic:{fragment_id}",),
|
|
payload=payload,
|
|
)
|
|
return fragment, payload_file
|
|
|
|
|
|
def _case_definitions() -> list[dict[str, Any]]:
|
|
world1: list[tuple[TensorFragment, str]] = []
|
|
world1.append(
|
|
_fragment(
|
|
"w1-weight",
|
|
torch.arange(6, dtype=torch.float32).reshape(2, 3),
|
|
payload_file="rank-00000.safetensors",
|
|
global_shape=(2, 3),
|
|
offset=(0, 0),
|
|
logical_tensor_id="model.weight",
|
|
partition_axes=(),
|
|
)
|
|
)
|
|
world1.append(
|
|
_fragment(
|
|
"w1-scalar",
|
|
torch.tensor(3, dtype=torch.int64),
|
|
payload_file="rank-00000.safetensors",
|
|
global_shape=(),
|
|
offset=(),
|
|
logical_tensor_id="model.step",
|
|
role="persistent_buffer",
|
|
partition_axes=(),
|
|
)
|
|
)
|
|
world1.append(
|
|
_fragment(
|
|
"w1-buffer",
|
|
torch.tensor([0.25, 0.5], dtype=torch.float32),
|
|
payload_file="rank-00000.safetensors",
|
|
global_shape=(2,),
|
|
offset=(0,),
|
|
logical_tensor_id="model.running_mean",
|
|
role="persistent_buffer",
|
|
partition_axes=(),
|
|
)
|
|
)
|
|
tied = torch.tensor([1.0, 2.0, 3.0, 4.0])
|
|
for fragment_id, tensor_id in (("w1-embed", "model.embed.weight"), ("w1-head", "lm_head.weight")):
|
|
world1.append(
|
|
_fragment(
|
|
fragment_id,
|
|
tied.clone(),
|
|
payload_file="rank-00000.safetensors",
|
|
global_shape=(4,),
|
|
offset=(0,),
|
|
logical_tensor_id=tensor_id,
|
|
partition_axes=(),
|
|
tie_group_id="tie-embedding-head",
|
|
)
|
|
)
|
|
world1.append(
|
|
_fragment(
|
|
"w1-expert",
|
|
torch.arange(4, dtype=torch.float32).reshape(2, 2),
|
|
payload_file="rank-00000.safetensors",
|
|
global_shape=(2, 2),
|
|
offset=(0, 0),
|
|
logical_tensor_id="model.experts.0.weight",
|
|
partition_axes=(),
|
|
topology_coordinates=(("ep", 0), ("rank", 0)),
|
|
)
|
|
)
|
|
|
|
uneven = torch.arange(7, dtype=torch.int64)
|
|
world2 = [
|
|
_fragment(
|
|
"w2-r0",
|
|
torch.cat((uneven[:3], torch.tensor([-1], dtype=torch.int64))),
|
|
payload_file="rank-00000.safetensors",
|
|
global_shape=(7,),
|
|
offset=(0,),
|
|
extent=(3,),
|
|
logical_tensor_id="model.weight",
|
|
padding=Padding((0,), (1,), "producer_declared"),
|
|
topology_coordinates=(("rank", 0), ("tp", 0)),
|
|
),
|
|
_fragment(
|
|
"w2-r1",
|
|
uneven[3:].clone(),
|
|
payload_file="rank-00001.safetensors",
|
|
global_shape=(7,),
|
|
offset=(3,),
|
|
logical_tensor_id="model.weight",
|
|
topology_coordinates=(("rank", 1), ("tp", 1)),
|
|
),
|
|
]
|
|
|
|
matrix = torch.arange(35, dtype=torch.float32).reshape(5, 7)
|
|
world4: list[tuple[TensorFragment, str]] = []
|
|
rank = 0
|
|
for row, (top, bottom) in enumerate(((0, 2), (2, 5))):
|
|
for column, (left, right) in enumerate(((0, 3), (3, 7))):
|
|
world4.append(
|
|
_fragment(
|
|
f"w4-r{rank}",
|
|
matrix[top:bottom, left:right].clone(),
|
|
payload_file=f"rank-{rank:05d}.safetensors",
|
|
global_shape=(5, 7),
|
|
offset=(top, left),
|
|
logical_tensor_id="model.weight",
|
|
partition_axes=(0, 1),
|
|
topology_coordinates=(("rank", rank), ("tp_row", row), ("tp_col", column)),
|
|
)
|
|
)
|
|
rank += 1
|
|
|
|
replica_value = torch.tensor([5.0, 6.0, 7.0])
|
|
replicas = [
|
|
_fragment(
|
|
f"dp-r{rank}",
|
|
replica_value.clone(),
|
|
payload_file=f"rank-{rank:05d}.safetensors",
|
|
global_shape=(3,),
|
|
offset=(0,),
|
|
logical_tensor_id="model.weight",
|
|
partition_axes=(),
|
|
replica=Replica("dp-full", rank, 4),
|
|
topology_coordinates=(("dp", rank), ("rank", rank)),
|
|
)
|
|
for rank in range(4)
|
|
]
|
|
|
|
topology_value = torch.arange(16, dtype=torch.float32).reshape(4, 4)
|
|
tp_pp: list[tuple[TensorFragment, str]] = []
|
|
rank = 0
|
|
for pp, (top, bottom) in enumerate(((0, 2), (2, 4))):
|
|
for tp, (left, right) in enumerate(((0, 2), (2, 4))):
|
|
tp_pp.append(
|
|
_fragment(
|
|
f"tp-pp-r{rank}",
|
|
topology_value[top:bottom, left:right].clone(),
|
|
payload_file=f"rank-{rank:05d}.safetensors",
|
|
global_shape=(4, 4),
|
|
offset=(top, left),
|
|
logical_tensor_id="model.weight",
|
|
partition_axes=(0, 1),
|
|
topology_coordinates=(("pp", pp), ("rank", rank), ("tp", tp)),
|
|
)
|
|
)
|
|
rank += 1
|
|
|
|
mixed = [
|
|
_fragment(
|
|
"mixed-model",
|
|
torch.arange(4, dtype=torch.float32).reshape(2, 2),
|
|
payload_file="rank-00000.safetensors",
|
|
global_shape=(2, 2),
|
|
offset=(0, 0),
|
|
logical_tensor_id="model.weight",
|
|
component_id="full-model",
|
|
partition_axes=(),
|
|
topology_coordinates=(("rank", 0),),
|
|
),
|
|
_fragment(
|
|
"mixed-adapter",
|
|
torch.tensor([[0.5, -0.5]], dtype=torch.float32),
|
|
payload_file="rank-00001.safetensors",
|
|
global_shape=(1, 2),
|
|
offset=(0, 0),
|
|
logical_tensor_id="adapter.lora_A.weight",
|
|
component_id="peft-adapter",
|
|
partition_axes=(),
|
|
topology_coordinates=(("rank", 1),),
|
|
),
|
|
]
|
|
|
|
return [
|
|
{
|
|
"case_id": "world1-complete",
|
|
"world_size": 1,
|
|
"features": ["buffer", "expert", "scalar", "tied_weight"],
|
|
"components": ["full_model"],
|
|
"topology": {"source": {"world_size": 1}, "target": {"world_size": 1}},
|
|
"fragments": world1,
|
|
},
|
|
{
|
|
"case_id": "world2-uneven-1d",
|
|
"world_size": 2,
|
|
"features": ["padding", "uneven_1d"],
|
|
"components": ["full_model"],
|
|
"topology": {"source": {"tp": 2}, "target": {"world_size": 1}},
|
|
"fragments": world2,
|
|
},
|
|
{
|
|
"case_id": "world4-uneven-2d",
|
|
"world_size": 4,
|
|
"features": ["uneven_2d"],
|
|
"components": ["full_model"],
|
|
"topology": {"source": {"tp_rows": 2, "tp_columns": 2}, "target": {"world_size": 1}},
|
|
"fragments": world4,
|
|
},
|
|
{
|
|
"case_id": "world4-dp-replicas",
|
|
"world_size": 4,
|
|
"features": ["dp_replica"],
|
|
"components": ["full_model"],
|
|
"topology": {"source": {"dp": 4}, "target": {"world_size": 1}},
|
|
"fragments": replicas,
|
|
},
|
|
{
|
|
"case_id": "tp2-pp2-to-single",
|
|
"world_size": 4,
|
|
"features": ["pipeline_parallel", "topology_a_to_b"],
|
|
"components": ["full_model"],
|
|
"topology": {"source": {"pp": 2, "tp": 2}, "target": {"world_size": 1}},
|
|
"fragments": tp_pp,
|
|
},
|
|
{
|
|
"case_id": "mixed-model-peft",
|
|
"world_size": 2,
|
|
"features": ["mixed_full_model_peft"],
|
|
"components": ["full_model", "peft_adapter"],
|
|
"topology": {"source": {"world_size": 2}, "target": {"world_size": 1}},
|
|
"fragments": mixed,
|
|
},
|
|
]
|
|
|
|
|
|
def _write_case(root: Path, definition: dict[str, Any]) -> dict[str, Any]:
|
|
case_id = definition["case_id"]
|
|
case_root = root / "cases" / case_id
|
|
case_root.mkdir(parents=True)
|
|
fragment_pairs: list[tuple[TensorFragment, str]] = definition["fragments"]
|
|
by_file: dict[str, dict[str, torch.Tensor]] = {}
|
|
fragments: list[TensorFragment] = []
|
|
fragment_records: list[dict[str, object]] = []
|
|
for fragment, payload_file in fragment_pairs:
|
|
fragments.append(fragment)
|
|
by_file.setdefault(payload_file, {})[fragment.fragment_id] = fragment.payload
|
|
record = fragment.manifest_record()
|
|
record["payload_file"] = payload_file
|
|
record["payload_key"] = fragment.fragment_id
|
|
fragment_records.append(record)
|
|
for filename, tensors in sorted(by_file.items()):
|
|
save_file(dict(sorted(tensors.items())), case_root / filename)
|
|
validation = validate_fragments(fragments)
|
|
oracle_tensors: dict[str, torch.Tensor] = {}
|
|
oracle_records: list[dict[str, object]] = []
|
|
for index, logical in enumerate(validation.logical_tensors):
|
|
value = reconstruct_logical_tensor(validation, logical.logical_tensor_id)
|
|
payload_key = f"tensor_{index:03d}"
|
|
oracle_tensors[payload_key] = value
|
|
oracle_records.append(
|
|
{
|
|
"logical_tensor_id": logical.logical_tensor_id,
|
|
"payload_key": payload_key,
|
|
"shape": list(value.shape),
|
|
"dtype": str(value.dtype).removeprefix("torch."),
|
|
"sha256": _tensor_digest(value),
|
|
}
|
|
)
|
|
save_file(dict(sorted(oracle_tensors.items())), case_root / "oracles.safetensors")
|
|
_write_json(
|
|
case_root / "case.json",
|
|
{
|
|
"schema_id": "obliteratus.checkpoint-fixture-case",
|
|
"schema_version": "1.0.0",
|
|
"case_id": case_id,
|
|
"world_size": definition["world_size"],
|
|
"features": sorted(definition["features"]),
|
|
"components": definition["components"],
|
|
"topology": definition["topology"],
|
|
"fragments": sorted(fragment_records, key=lambda item: item["fragment_id"]),
|
|
"oracle_file": "oracles.safetensors",
|
|
"oracles": oracle_records,
|
|
"expected_manifest_digest": validation.manifest_digest,
|
|
},
|
|
)
|
|
files = []
|
|
for path in sorted(item for item in case_root.iterdir() if item.is_file()):
|
|
payload = path.read_bytes()
|
|
files.append(
|
|
{
|
|
"relative_path": path.name,
|
|
"size_bytes": len(payload),
|
|
"sha256": f"sha256:{sha256(payload).hexdigest()}",
|
|
}
|
|
)
|
|
if len(files) > CORPUS_LIMITS["max_files_per_case"]:
|
|
raise ValueError(f"fixture case exceeds file limit: {case_id}")
|
|
if sum(item["size_bytes"] for item in files) > CORPUS_LIMITS["max_case_bytes"]:
|
|
raise ValueError(f"fixture case exceeds byte limit: {case_id}")
|
|
return {
|
|
"case_id": case_id,
|
|
"relative_path": f"cases/{case_id}",
|
|
"world_size": definition["world_size"],
|
|
"features": sorted(definition["features"]),
|
|
"files": files,
|
|
}
|
|
|
|
|
|
def _negative_catalog() -> dict[str, object]:
|
|
failures = {
|
|
"dimension_mismatch": "DCI_VALIDATION_FAILED",
|
|
"extra_shard": "DCI_SOURCE_BOUNDARY_VIOLATION",
|
|
"fragment_out_of_bounds": "DCI_VALIDATION_FAILED",
|
|
"integer_overflow": "DCI_VALIDATION_FAILED",
|
|
"missing_shard": "DCI_SOURCE_BOUNDARY_VIOLATION",
|
|
"negative_integer": "DCI_VALIDATION_FAILED",
|
|
"padding_shape_mismatch": "DCI_VALIDATION_FAILED",
|
|
"path_traversal": "DCI_SOURCE_BOUNDARY_VIOLATION",
|
|
"payload_dtype_mismatch": "DCI_VALIDATION_FAILED",
|
|
"payload_shape_mismatch": "DCI_VALIDATION_FAILED",
|
|
"replica_digest_mismatch": "DCI_VALIDATION_FAILED",
|
|
"resource_manifest_bomb": "DCI_RESOURCE_LIMIT",
|
|
"source_special_file": "DCI_SOURCE_BOUNDARY_VIOLATION",
|
|
"source_symlink": "DCI_SOURCE_BOUNDARY_VIOLATION",
|
|
"truncated_shard": "DCI_SOURCE_BOUNDARY_VIOLATION",
|
|
"coverage_gap": "DCI_VALIDATION_FAILED",
|
|
"coverage_overlap": "DCI_VALIDATION_FAILED",
|
|
}
|
|
return {
|
|
"schema_id": "obliteratus.checkpoint-negative-fixtures",
|
|
"schema_version": "1.0.0",
|
|
"cases": [
|
|
{
|
|
"case_id": f"negative-{index:02d}",
|
|
"failure": failure,
|
|
"expected_code": code,
|
|
"mutation": f"deterministic:{failure}",
|
|
}
|
|
for index, (failure, code) in enumerate(sorted(failures.items()), start=1)
|
|
],
|
|
}
|
|
|
|
|
|
def generate_corpus(destination: Path | str) -> Path:
|
|
"""Create a new bounded corpus; existing paths are never overwritten."""
|
|
root = Path(destination)
|
|
if root.exists() or root.is_symlink():
|
|
raise FileExistsError(f"fixture destination already exists: {root}")
|
|
root.mkdir(parents=True)
|
|
definitions = _case_definitions()
|
|
if len(definitions) > CORPUS_LIMITS["max_cases"]:
|
|
raise ValueError("fixture corpus exceeds case limit")
|
|
cases = [_write_case(root, definition) for definition in definitions]
|
|
_write_json(root / "negative-cases.json", _negative_catalog())
|
|
_write_json(
|
|
root / "fixture-corpus.json",
|
|
{
|
|
"schema_id": "obliteratus.checkpoint-fixture-corpus",
|
|
"schema_version": "1.0.0",
|
|
"generator": {
|
|
"path": "scripts/generate_checkpoint_fixtures.py",
|
|
"version": GENERATOR_VERSION,
|
|
},
|
|
"license": "AGPL-3.0-or-later",
|
|
"provenance": {
|
|
"kind": "deterministic_synthetic",
|
|
"seed": 0,
|
|
"third_party_data": False,
|
|
"third_party_weights": False,
|
|
},
|
|
"limits": CORPUS_LIMITS,
|
|
"cases": sorted(cases, key=lambda item: item["case_id"]),
|
|
},
|
|
)
|
|
return root
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("destination", type=Path)
|
|
arguments = parser.parse_args()
|
|
generate_corpus(arguments.destination)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__": # pragma: no cover - CLI wrapper
|
|
raise SystemExit(main())
|