| """Exact tensor accounting; never equate payload with measured loaded memory.""" |
| import math |
| from collections.abc import Mapping |
|
|
| import torch |
|
|
|
|
| def packed_w4_bytes(rows: int, width: int) -> int: |
| if rows <= 0 or width <= 0: |
| raise ValueError("weight dimensions must be positive") |
| return rows * (math.ceil(width / 2) + 4 * math.ceil(width / 32)) |
|
|
|
|
| def kronecker_counts(left: int, right: int) -> dict: |
| if min(left, right) <= 0: |
| raise ValueError("factor dimensions must be positive") |
| return dict(effective_factors=left * left + right * right, |
| svd_trainable=2 * left * left + left + 2 * right * right + right) |
|
|
|
|
| def ouro_analytic_accounting(*, layers=24, hidden=2048, intermediate=5632, |
| vocab=49152, heads=16, kv_heads=16, |
| hidden_factors=(32, 64), intermediate_factors=(64, 88), |
| loops=4, selected_groups=4, rank=8) -> dict: |
| """Known Ouro layout, untied embeddings, bias-free projections, gate h+1. |
| |
| Bounds allow every selected group to be either the smallest or largest |
| candidate. They include all T copies, NOT an extra-copy delta. Shared and |
| loop-dependent counts remain separate from the BF16 backbone tensor count. |
| """ |
| if hidden % heads or hidden_factors[0] * hidden_factors[1] != hidden or intermediate_factors[0] * intermediate_factors[1] != intermediate: |
| raise ValueError("invalid head/factor configuration") |
| kv = hidden // heads * kv_heads |
| shapes = [(hidden, hidden), (kv, hidden), (kv, hidden), (hidden, hidden), |
| (intermediate, hidden), (intermediate, hidden), (hidden, intermediate)] |
| weight_elements = layers * sum(r * c for r, c in shapes) |
| scales = layers * sum(r * math.ceil(c / 32) for r, c in shapes) |
| |
| |
| non_projection = 2 * vocab * hidden + (4 * layers + 1) * hidden + hidden + 1 |
| base = weight_elements + non_projection |
| h, m = kronecker_counts(*hidden_factors), kronecker_counts(*intermediate_factors) |
| shared_uv = 2 * hidden * rank |
| dependent_cta = (loops - 1) * (2 * hidden + rank) |
| las = layers * 4 * loops |
| bounds = {name: [selected_groups * loops * min(h[name], m[name]) + las + dependent_cta, |
| selected_groups * loops * max(h[name], m[name]) + las + dependent_cta] |
| for name in h} |
| shared = {name: layers * (3 * h[name] + m[name]) for name in h} |
| packed_base = layers * sum(packed_w4_bytes(*shape) for shape in shapes) |
| return dict(projection_elements=weight_elements, group32_weight_scales=scales, |
| backbone_stored_elements=base, bf16_payload_bytes=base * 2, |
| bf16_payload_GiB=base * 2 / 2 ** 30, |
| packed_projection_payload_bytes=packed_base, |
| non_projection_bf16_payload_bytes=non_projection * 2, |
| shared_transform_counts=shared, shared_cta_uv=shared_uv, |
| loop_dependent_las=las, loop_dependent_cta=dependent_cta, |
| loop_dependent_bounds=bounds, |
| assumptions="fixed local Kronecker layout, all loop copies, no optional diagonals; not measured loaded memory") |
|
|
|
|
| def tensor_tree_inventory(tree) -> dict: |
| """Count tensor objects once; storage bytes deduplicate aliased views.""" |
| objects, storages = {}, {} |
|
|
| def visit(value): |
| if isinstance(value, torch.Tensor): |
| objects[id(value)] = value |
| storage = value.untyped_storage() |
| storages[(str(value.device), storage.data_ptr(), storage.nbytes())] = storage.nbytes() |
| elif isinstance(value, Mapping): |
| for child in value.values(): |
| visit(child) |
| elif isinstance(value, (list, tuple)): |
| for child in value: |
| visit(child) |
|
|
| visit(tree) |
| return dict(tensor_objects=len(objects), tensor_elements=sum(t.numel() for t in objects.values()), |
| logical_tensor_bytes=sum(t.numel() * t.element_size() for t in objects.values()), |
| unique_storage_bytes=sum(storages.values()), |
| excludes="allocator overhead, runtime workspaces, KV cache, Python containers") |
|
|
