| """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}" |
| ) |
|
|
