File size: 3,056 Bytes
e1acc15 | 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 | """Independent CPU E2M1/E8M0 decoder and admission; no CUDA import or access."""
from __future__ import annotations
import torch
E2M1: tuple[float, ...] = (0., .5, 1., 1.5, 2., 3., 4., 6.,
-0., -.5, -1., -1.5, -2., -3., -4., -6.)
def decode_codes(codes: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
"""Float64 oracle, including signed zero, scale-0 and NaN scale-255."""
if codes.device.type != 'cpu' or scales.device.type != 'cpu':
raise ValueError('reference decoder requires CPU tensors')
values = torch.tensor(E2M1, dtype=torch.float64)[codes.long()]
values = torch.ldexp(values, scales.to(torch.int32) - 127)
return torch.where(scales == 255, float('nan'), values)
def unpack(packed: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
if packed.dtype != torch.uint8 or scales.dtype != torch.uint8:
raise ValueError('checkpoint codes and E8M0 scales must be uint8')
if packed.ndim != 3 or scales.shape != (*packed.shape[:2], packed.shape[2] // 16):
raise ValueError('expected [E,N,K/2] codes and [E,N,K/32] scales')
if packed.shape[-1] % 16:
raise ValueError('K must be a multiple of group size 32')
codes = torch.stack((packed & 15, packed >> 4), dim=-1).flatten(-2)
return decode_codes(codes, scales.repeat_interleave(32, dim=-1))
def admit(packed: torch.Tensor, scales: torch.Tensor) -> dict[str, int]:
"""Bounded load-time validation. Never narrow/repair unsupported native values.
Inspect one expert at a time so preparing a full bank never expands the
checkpoint in RAM. BF16 subnormals decode exactly but MMA may flush them:
reject them for computation rather than asserting unsupported fidelity.
"""
if packed.device.type != 'cpu' or scales.device.type != 'cpu':
raise ValueError('admission must run on the CPU before GPU allocation')
if packed.ndim != 3 or packed.shape[0] == 0 or packed.shape[1] == 0:
raise ValueError('expected nonempty expert bank')
result = {'values': 0, 'negative_zeros': 0, 'scale_min': 255, 'scale_max': 0}
for e in range(packed.shape[0]):
decoded = unpack(packed[e:e+1], scales[e:e+1])
if not bool(torch.isfinite(decoded).all()):
raise ValueError(f'expert {e}: NaN/nonfinite native scale/value')
widened = decoded.to(torch.bfloat16).to(torch.float64)
if not torch.equal(decoded, widened):
raise ValueError(f'expert {e}: native values not exactly representable in BF16')
subnormal = (decoded != 0) & (decoded.abs() < torch.finfo(torch.bfloat16).tiny)
if bool(subnormal.any()):
raise ValueError(f'expert {e}: BF16 subnormal operands require a wider arithmetic path')
result['values'] += decoded.numel()
result['negative_zeros'] += int(((decoded == 0) & decoded.signbit()).sum())
result['scale_min'] = min(result['scale_min'], int(scales[e].min()))
result['scale_max'] = max(result['scale_max'], int(scales[e].max()))
return result
|