| """Official forward values; STE QDQ and recomputed frozen-weight input gradients. |
| |
| No decoded routed matrix survives a call. Packed expert bytes alone are saved |
| for backward. BF16 decoding is exact for the pinned 1017 tensors (N1 audit). |
| FP64 accumulation is an oracle option, FP32 is the training path; TF32 is off. |
| """ |
| from __future__ import annotations |
|
|
| from typing import Any, Callable, cast |
| from contextlib import contextmanager |
| from collections.abc import Iterator |
| import threading |
|
|
| import torch |
| from torch import Tensor, nn |
| from torch.nn import functional as F |
|
|
| from ...native.formats import act_qdq, decode_fp4 |
|
|
| _STOP: threading.Event | None = None |
|
|
|
|
| class TrainingCancelled(Exception): |
| """Discard the current, uncommitted optimizer step and checkpoint its boundary.""" |
|
|
|
|
| @contextmanager |
| def cancellation(stop: threading.Event) -> Iterator[None]: |
| global _STOP |
| old, _STOP = _STOP, stop |
| try: |
| yield |
| finally: |
| _STOP = old |
|
|
|
|
| def check_cancelled() -> None: |
| if _STOP is not None and _STOP.is_set(): |
| raise TrainingCancelled() |
|
|
|
|
| class _QDQ(torch.autograd.Function): |
| @staticmethod |
| def forward(ctx: Any, x: Tensor) -> Tensor: |
| return act_qdq(x) |
|
|
| @staticmethod |
| def backward(ctx: Any, grad: Tensor) -> Tensor: |
| return grad |
|
|
|
|
| def qdq(x: Tensor, fast: bool = False) -> Tensor: |
| if fast: |
| from .fast_gpu import fast_qdq |
| return fast_qdq(x) |
| return cast(Callable[[Tensor], Tensor], _QDQ.apply)(x) |
|
|
|
|
| class _FrozenLinear(torch.autograd.Function): |
| @staticmethod |
| def forward(ctx: Any, x: Tensor, w: Tensor, scale: Tensor | None, acc: torch.dtype) -> Tensor: |
| check_cancelled() |
| ctx.save_for_backward(w, *(() if scale is None else (scale,))) |
| ctx.acc, ctx.input_dtype, ctx.input_device = acc, x.dtype, x.device |
| if scale is not None and w.is_cuda: |
| from .packed_forward import packed_matmul |
| result = packed_matmul(x.to(w.device), w, scale) |
| return result.to(device=x.device, dtype=x.dtype) |
| decoded = w if scale is None else decode_fp4(w, scale).to(torch.bfloat16) |
| return F.linear(x.to(device=w.device, dtype=acc), decoded.to(acc)).to(device=x.device, dtype=x.dtype) |
|
|
| @staticmethod |
| def backward(ctx: Any, grad: Tensor) -> tuple[Tensor, None, None, None]: |
| check_cancelled() |
| saved = ctx.saved_tensors |
| decoded = saved[0] if len(saved) == 1 else decode_fp4(*saved).to(torch.bfloat16) |
| |
| gradient_acc = torch.float32 if decoded.is_cuda else ctx.acc |
| dx = (grad.to(device=decoded.device, dtype=gradient_acc) @ decoded.to(gradient_acc)).to(device=ctx.input_device, dtype=ctx.input_dtype) |
| return dx, None, None, None |
|
|
|
|
| def frozen_linear(x: Tensor, w: Tensor, scale: Tensor | None = None, |
| acc: torch.dtype = torch.float32) -> Tensor: |
| return cast(Callable[[Tensor, Tensor, Tensor | None, torch.dtype], Tensor], _FrozenLinear.apply)(x, w, scale, acc) |
|
|
|
|
| def head_product(a: Tensor, b: Tensor) -> Tensor: |
| """FP32 head output from a two-part BF16 activation and exact BF16 weights. |
| |
| No QDQ follows the vocabulary head, so its small input decomposition error |
| cannot cross any internal quantizer. Whole-model KL/loss/gradient gates |
| explicitly include this arithmetic. CPU is the FP32 oracle for tests. |
| """ |
| if not a.is_cuda or b.dtype != torch.bfloat16: |
| return a.float() @ b.float() |
| high = a.bfloat16() |
| low = (a.float() - high.float()).bfloat16() |
| return torch.mm(high, b, out_dtype=torch.float32) + torch.mm(low, b, out_dtype=torch.float32) |
|
|
|
|
| class _HeadLinear(torch.autograd.Function): |
| @staticmethod |
| def forward(ctx: Any, x: Tensor, w: Tensor) -> Tensor: |
| ctx.save_for_backward(w) |
| return head_product(x.reshape(-1, x.shape[-1]), w.T).reshape(*x.shape[:-1], w.shape[0]) |
|
|
| @staticmethod |
| def backward(ctx: Any, grad: Tensor) -> tuple[Tensor, None]: |
| (w,) = ctx.saved_tensors |
| return head_product(grad.reshape(-1, grad.shape[-1]), w).reshape(*grad.shape[:-1], w.shape[1]), None |
|
|
|
|
| def fast_head(x: Tensor, w: Tensor) -> Tensor: |
| return cast(Callable[[Tensor, Tensor], Tensor], _HeadLinear.apply)(x, w) |
|
|
|
|
| class _GroupedLinear(torch.autograd.Function): |
| @staticmethod |
| def forward(ctx: Any, x: Tensor, w: Tensor) -> Tensor: |
| ctx.save_for_backward(w) |
| expanded = w.float() |
| if x.device.type == "cpu" and x.shape[0] > 1: |
| return torch.cat([torch.einsum("bsgd,grd->bsgr", x[b:b+1], expanded) for b in range(x.shape[0])], 0) |
| return torch.einsum("bsgd,grd->bsgr", x, expanded) |
|
|
| @staticmethod |
| def backward(ctx: Any, grad: Tensor) -> tuple[Tensor, None]: |
| (w,) = ctx.saved_tensors |
| return torch.einsum("bsgr,grd->bsgd", grad, w.float()), None |
|
|
|
|
| def grouped_linear(x: Tensor, w: Tensor) -> Tensor: |
| |
| |
| return cast(Callable[[Tensor, Tensor], Tensor], _GroupedLinear.apply)(x, w) |
|
|
|
|
| class Dense(nn.Module): |
| """Frozen decoded BF16 master, optional r32/64 LoRA, official input QDQ. |
| |
| LoRA is added before the linear output rounding. This is a floating adapter; |
| the merged FP8 path is deliberately a separate packaging/evaluation gate. |
| """ |
| weight: Tensor |
|
|
| def __init__(self, weight: Tensor, rank: int = 0, quant_input: bool = True, |
| acc: torch.dtype = torch.float32): |
| super().__init__() |
| self.register_buffer("weight", weight.to(torch.bfloat16)) |
| self.quant_input, self.acc = quant_input, acc |
| self.fast_training = False |
| self.rank = rank |
| if rank: |
| if rank not in (32, 64): |
| raise ValueError("Phase 2a supports only LoRA r32 or r64") |
| self.a = nn.Parameter(torch.empty(rank, weight.shape[1], device=weight.device, dtype=torch.float32)) |
| self.b = nn.Parameter(torch.zeros(weight.shape[0], rank, device=weight.device, dtype=torch.float32)) |
| nn.init.kaiming_uniform_(self.a, a=5 ** .5) |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| x = qdq(x, self.fast_training) if self.quant_input else x |
| if self.fast_training: |
| from .fast_gpu import fast_linear |
| base = fast_linear(x, self.weight, self.quant_input) |
| else: |
| base = frozen_linear(x, self.weight, acc=self.acc) |
| if self.rank: |
| base = base + F.linear(F.linear(x.to(device=self.a.device, dtype=torch.float32), self.a), self.b).to(x.device) |
| return base.to(x.dtype) |
|
|
| def merged(self) -> Tensor: |
| w = self.weight.float() |
| return w + self.b @ self.a if self.rank else w |
|
|
|
|
| class PackedLinear(nn.Module): |
| weight: Tensor |
| scale: Tensor |
|
|
| def __init__(self, weight: Tensor, scale: Tensor, acc: torch.dtype): |
| super().__init__() |
| self.register_buffer("weight", weight) |
| self.register_buffer("scale", scale) |
| self.acc = acc |
|
|
| def forward(self, x: Tensor) -> Tensor: |
| return frozen_linear(qdq(x), self.weight, self.scale, self.acc) |
|
|
|
|
| def rms(x: Tensor, weight: Tensor, eps: float) -> Tensor: |
| z = x.float() |
| z = z * torch.rsqrt(z.square().mean(-1, keepdim=True) + eps) |
| return (weight * z).to(x.dtype) |
|
|
|
|
| def rope(x: Tensor, positions: Tensor, rd: int, theta: float, inverse: bool = False) -> Tensor: |
| |
| freq = 1 / theta ** (torch.arange(0, rd, 2, device=x.device, dtype=torch.float32) / rd) |
| angles = positions.float().unsqueeze(-1) * freq |
| phase = torch.polar(torch.ones_like(angles), -angles if inverse else angles) |
| if x.ndim == 4: |
| phase = phase.unsqueeze(-2) |
| tail = torch.view_as_complex(x[..., -rd:].float().reshape(*x.shape[:-1], rd//2, 2).contiguous()) |
| rotated = torch.view_as_real(tail * phase).flatten(-2).to(x.dtype) |
| return torch.cat((x[..., :-rd], rotated), -1) |
|
|
|
|
| def mixes(x: Tensor, fn: Tensor, scale: Tensor, base: Tensor, hc: int, iters: int, |
| norm_eps: float, eps: float, acc: torch.dtype) -> tuple[Tensor, Tensor, Tensor]: |
| z = x.flatten(-2).float() |
| z = F.linear(z, fn) * torch.rsqrt(z.square().mean(-1, keepdim=True) + norm_eps) |
| z, scale, base = z.to(acc), scale.to(acc), base.to(acc) |
| pre = torch.sigmoid(z[..., :hc] * scale[0] + base[:hc]) + eps |
| post = 2 * torch.sigmoid(z[..., hc:2*hc] * scale[1] + base[hc:2*hc]) |
| comb = (z[..., 2*hc:] * scale[2] + base[2*hc:]).reshape(*z.shape[:-1], hc, hc) |
| comb = comb.softmax(-1) + eps |
| comb = comb / (comb.sum(-2, keepdim=True) + eps) |
| for _ in range(iters - 1): |
| comb = comb / (comb.sum(-1, keepdim=True) + eps) |
| comb = comb / (comb.sum(-2, keepdim=True) + eps) |
| return pre.float(), post.float(), comb.float() |
|
|
|
|
| def collapse(x: Tensor, pre: Tensor) -> Tensor: |
| return (pre.unsqueeze(-1) * x.float()).sum(-2).to(x.dtype) |
|
|
|
|
| def expand(x: Tensor, residual: Tensor, post: Tensor, comb: Tensor) -> Tensor: |
| return (post.unsqueeze(-1) * x.unsqueeze(-2) |
| + (comb.unsqueeze(-1) * residual.unsqueeze(-2)).sum(-3)).to(x.dtype) |
|
|
