Add MP-01 (9500) 100M base: weights, config, tokenizer, model code, model card and eval results
11e2c7c verified | """Fused cross-entropy: same maths as FusedLinearCE, but the per-chunk softmax/grad chain is one compiled kernel set. | |
| On a 288 GB/s card the eager version streams ~5 GB of fp32 logits-sized tensors per chunk; here the bf16 logits are | |
| read twice and the bf16 gradient written once. | |
| """ | |
| import torch | |
| def _chunk(logits, tgt, z_loss: float, inv_n: float): | |
| lg = logits.float() | |
| lse = torch.logsumexp(lg, dim=-1) | |
| tl = lg.gather(1, tgt[:, None]).squeeze(1) | |
| loss = (lse - tl).sum() | |
| p = torch.exp(lg - lse[:, None]) | |
| if z_loss: | |
| loss = loss + z_loss * (lse * lse).sum() | |
| p = p * (1.0 + 2.0 * z_loss * lse[:, None]) | |
| cols = torch.arange(lg.shape[1], device=lg.device)[None, :] | |
| p = torch.where(cols == tgt[:, None], p - 1.0, p) | |
| return loss, (p * inv_n).to(torch.bfloat16) | |
| chunk_compiled = torch.compile(_chunk, dynamic=False) | |
| class FastLinearCE(torch.autograd.Function): | |
| def forward(ctx, h, w, targets, chunk: int, z_loss: float): | |
| N = h.shape[0] | |
| wb = w.to(h.dtype) | |
| grad_h = torch.empty_like(h) | |
| grad_w = torch.zeros_like(w, dtype=torch.float32) | |
| total = torch.zeros((), device=h.device, dtype=torch.float32) | |
| for i in range(0, N, chunk): | |
| hc, tc = h[i : i + chunk], targets[i : i + chunk] | |
| loss_i, pb = chunk_compiled(hc @ wb.t(), tc, z_loss, 1.0 / N) | |
| total += loss_i | |
| grad_h[i : i + chunk] = pb @ wb | |
| grad_w.add_((pb.t() @ hc).float()) | |
| ctx.save_for_backward(grad_h, grad_w) | |
| return total / N | |
| def backward(ctx, g): | |
| grad_h, grad_w = ctx.saved_tensors | |
| return grad_h * g.to(grad_h.dtype), grad_w * g, None, None, None | |
| def _chunk_multi(logits, tgt, z_loss: float, inv_n: float): | |
| """Patch-level loss: one distribution per position scored against all k tokens of the next patch. | |
| tgt [n, k]; loss = sum_k CE(logits, tgt[:, k]); grad = k*softmax - sum_k onehot.""" | |
| lg = logits.float() | |
| k = tgt.shape[1] | |
| lse = torch.logsumexp(lg, dim=-1) | |
| loss = (k * lse - lg.gather(1, tgt).sum(1)).sum() | |
| p = k * torch.exp(lg - lse[:, None]) | |
| if z_loss: | |
| loss = loss + k * z_loss * (lse * lse).sum() | |
| p = p * (1.0 + 2.0 * z_loss * lse[:, None]) | |
| cols = torch.arange(lg.shape[1], device=lg.device)[None, :] | |
| for j in range(k): # unrolled at trace time and fused: no logits-sized one-hot buffer | |
| p = p - (cols == tgt[:, j : j + 1]).to(p.dtype) | |
| return loss, (p * inv_n).to(torch.bfloat16) | |
| chunk_multi = torch.compile(_chunk_multi, dynamic=False) | |
| class PatchLinearCE(torch.autograd.Function): | |
| """FastLinearCE for patch-level training: h [N, d], targets [N, k]; returns the mean per-token CE.""" | |
| def forward(ctx, h, w, targets, chunk: int, z_loss: float): | |
| N, k = targets.shape | |
| wb = w.to(h.dtype) | |
| grad_h = torch.empty_like(h) | |
| grad_w = torch.zeros_like(w, dtype=torch.float32) | |
| total = torch.zeros((), device=h.device, dtype=torch.float32) | |
| for i in range(0, N, chunk): | |
| hc, tc = h[i : i + chunk], targets[i : i + chunk] | |
| loss_i, pb = chunk_multi(hc @ wb.t(), tc, z_loss, 1.0 / (N * k)) | |
| total += loss_i | |
| grad_h[i : i + chunk] = pb @ wb | |
| grad_w.add_((pb.t() @ hc).float()) | |
| ctx.save_for_backward(grad_h, grad_w) | |
| return total / (N * k) | |
| def backward(ctx, g): | |
| grad_h, grad_w = ctx.saved_tensors | |
| return grad_h * g.to(grad_h.dtype), grad_w * g, None, None, None | |
| # Variable assistant counts and CE tails must not specialize each compiled call. | |
| chunk_sft = torch.compile(_chunk, dynamic=True) | |
| class SFTLinearCE(torch.autograd.Function): | |
| def forward(ctx, h, w, targets, chunk: int, z_loss: float): | |
| N = h.shape[0] | |
| wb = w.to(h.dtype) | |
| grad_h = torch.empty_like(h) | |
| grad_w = torch.zeros_like(w, dtype=torch.float32) | |
| total = torch.zeros((), device=h.device, dtype=torch.float32) | |
| for i in range(0, N, chunk): | |
| hc, tc = h[i : i + chunk], targets[i : i + chunk] | |
| loss_i, pb = chunk_sft(hc @ wb.t(), tc, z_loss, torch.tensor(1.0 / N, device=h.device)) | |
| total += loss_i | |
| grad_h[i : i + chunk] = pb @ wb | |
| grad_w.add_((pb.t() @ hc).float()) | |
| ctx.save_for_backward(grad_h, grad_w) | |
| return total / N | |
| def backward(ctx, g): | |
| grad_h, grad_w = ctx.saved_tensors | |
| return grad_h * g.to(grad_h.dtype), grad_w * g, None, None, None | |