File size: 3,427 Bytes
ec24fb0 | 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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | """Training speed/memory options shared by the benchmark and the training scripts.
loss="hf" transformers' built-in loss: materialises the full [tokens x vocab] logits (Gemma: 262k vocab)
loss="chunked" plain-PyTorch chunked cross-entropy; each chunk's logits are recomputed in backward
loss="liger" Liger fused linear + cross-entropy Triton kernel (never materialises the logits)
ckpt_every=N recompute activations of every Nth decoder layer only (1 = all layers, 0 = none)
compile=True torch.compile each decoder layer (regional compilation)
"""
from __future__ import annotations
import torch
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
def lm_parts(model):
"""(decoder body, lm_head, final logit softcap) for plain and PEFT-wrapped causal LMs."""
lm = model.get_base_model() if hasattr(model, "get_base_model") else model
return lm.model, lm.lm_head, getattr(lm.config, "final_logit_softcapping", None)
def _ce_chunk(h, w, y, softcap):
logits = (h @ w.t()).float()
if softcap:
logits = torch.tanh(logits / softcap) * softcap
return F.cross_entropy(logits, y, ignore_index=-100, reduction="sum")
def chunked_ce(h, w, y, softcap, chunk: int = 4096):
total = h.new_zeros((), dtype=torch.float32)
for i in range(0, h.shape[0], chunk):
total = total + checkpoint(_ce_chunk, h[i:i + chunk], w, y[i:i + chunk], softcap, use_reentrant=False)
return total / (y != -100).sum().clamp(min=1)
_liger = {}
def lm_loss(model, input_ids, labels, loss: str = "hf", attention_mask=None):
"""Next-token cross-entropy, mean over non-ignored target tokens. All three variants give the same value."""
if loss == "hf":
return model(input_ids=input_ids, attention_mask=attention_mask, labels=labels).loss
body, head, softcap = lm_parts(model)
h = body(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
h = h[:, :-1].reshape(-1, h.shape[-1])
y = labels[:, 1:].reshape(-1)
if loss == "chunked":
return chunked_ce(h, head.weight, y, softcap)
if loss == "liger":
from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss
fn = _liger.setdefault(softcap, LigerFusedLinearCrossEntropyLoss(softcap=softcap))
return fn(head.weight, h, y)
raise ValueError(loss)
def configure(model, ckpt_every: int = 1, compile_layers: bool = False) -> None:
body = lm_parts(model)[0]
if ckpt_every > 0:
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
for i, layer in enumerate(body.layers):
layer.gradient_checkpointing = i % ckpt_every == 0
else:
model.gradient_checkpointing_disable()
if compile_layers:
# Layers differ (sliding vs global attention, linear-attention blocks), so dynamo compiles one graph per
# layer variant; past the default limit of 8 it silently falls back to eager. The limit is passed per call:
# the global config value was not honoured in the 11 Oct rerun. dynamic=False keeps one static graph per
# shape, so activation recomputation cannot pick a different (dynamic) graph than the forward pass used
# (pytorch/pytorch#166926).
limit = max(64, 4 * len(body.layers))
for layer in body.layers:
layer.compile(dynamic=False, recompile_limit=limit)
|