File size: 2,063 Bytes
5f0e4a2 | 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 | """Schema for exact all-layer BF16 Teacher Text K/V caches."""
from __future__ import annotations
from typing import Mapping
import torch
SCHEMA_VERSION_TEXT_KV_ALL54 = "hyworldplay_teacher_text_kv_exact_all54_bf16_v1"
TEXT_KV_LAYER_COUNT = 54
TEXT_KV_HEADS = 16
TEXT_KV_TOKENS = 903
TEXT_KV_HEAD_DIM = 128
def text_kv_key(block_id: int, kind: str) -> str:
if not 0 <= block_id < TEXT_KV_LAYER_COUNT:
raise ValueError(f"Invalid Text K/V block: {block_id}")
if kind not in {"k_txt", "v_txt"}:
raise ValueError(f"Invalid Text K/V kind: {kind}")
return f"block_{block_id:02d}_{kind}"
def expected_text_kv_keys() -> set[str]:
return {
text_kv_key(block_id, kind)
for block_id in range(TEXT_KV_LAYER_COUNT)
for kind in ("k_txt", "v_txt")
} | {"text_valid_mask"}
def validate_text_kv_all54(tensors: Mapping[str, torch.Tensor]) -> None:
missing = expected_text_kv_keys().difference(tensors)
unexpected = set(tensors).difference(expected_text_kv_keys())
if missing or unexpected:
raise ValueError(
f"Text K/V key mismatch: missing={sorted(missing)}, "
f"unexpected={sorted(unexpected)}"
)
expected_shape = (1, TEXT_KV_HEADS, TEXT_KV_TOKENS, TEXT_KV_HEAD_DIM)
for block_id in range(TEXT_KV_LAYER_COUNT):
for kind in ("k_txt", "v_txt"):
name = text_kv_key(block_id, kind)
value = tensors[name]
if tuple(value.shape) != expected_shape:
raise ValueError(f"Unexpected {name} shape: {tuple(value.shape)}")
if value.dtype != torch.bfloat16:
raise TypeError(f"{name} must be BF16, got {value.dtype}")
if not torch.isfinite(value).all():
raise ValueError(f"{name} contains non-finite values")
mask = tensors["text_valid_mask"]
if tuple(mask.shape) != (1, TEXT_KV_TOKENS) or mask.dtype != torch.bool:
raise ValueError(
f"Unexpected text_valid_mask: shape={tuple(mask.shape)}, dtype={mask.dtype}"
)
|