import dataclasses from functools import cached_property import torch __all__ = [] def _make_empty(shape, dtype, device, when=True): # when=False returns a rank-matching empty (shape (0,)*len(shape)). return torch.empty( shape if when else (0,) * len(shape), device=device, dtype=dtype, requires_grad=False, ) def _make_zeros(shape, dtype, device, when=True): return torch.zeros( shape if when else (0,) * len(shape), device=device, dtype=dtype, requires_grad=False, ) def _linear_cross_entropy_batch_chunked_setup_context(ctx, inputs, output): ( *_, allow_retain_graph, compute_input_grad, compute_linear_weight_grad, compute_linear_bias_grad, ) = inputs ctx.allow_retain_graph = allow_retain_graph ctx.compute_input_grad = compute_input_grad ctx.compute_linear_weight_grad = compute_linear_weight_grad ctx.compute_linear_bias_grad = compute_linear_bias_grad _, grad_input, grad_linear_weight, grad_linear_bias = output # ``.detach()`` the stashed grads: they are forward OUTPUTS, so they # carry ``grad_fn`` back to this custom-op node. Holding them as raw # ctx attributes would form a reference cycle (node -> ctx -> grad -> # grad_fn -> node) that refcounting can't break, leaking the tensors # until a GC pass -- which the CUDA mem-leak check flags. Detaching # drops the grad_fn link so the tensors free promptly when the loss's # graph is. Backward only multiplies them by the upstream grad, so the # lost autograd history is irrelevant (double-backward is unsupported). ctx._gi = grad_input.detach() if compute_input_grad else None ctx._gw = grad_linear_weight.detach() if compute_linear_weight_grad else None ctx._gb = grad_linear_bias.detach() if compute_linear_bias_grad else None @dataclasses.dataclass class _ChunkViews: """Per-iteration tensor views; each property picks the right operand or scratch from ctx dispatch flags so call sites read raw. """ ctx: "_ChunkContext" bchunk_start: int bchunk_size: int input_chunk: torch.Tensor target_chunk: torch.Tensor weight_chunk: torch.Tensor logits: torch.Tensor input_chunk_acc: torch.Tensor @property def input(self) -> torch.Tensor: return ( self.input_chunk_acc if self.ctx.forward_uses_acc_input else self.input_chunk ) @property def linear_weight(self) -> torch.Tensor: ctx = self.ctx return ( ctx.linear_weight if ctx.forward_uses_cuda_out_dtype else ctx.linear_weight_cast ) @property def input_grad_logits(self) -> torch.Tensor: if self.ctx.input_grad_uses_logits_lw: return self.logits_downcast return self.logits @property def input_grad_linear_weight(self) -> torch.Tensor: ctx = self.ctx return ( ctx.linear_weight if ctx.input_grad_uses_logits_lw else ctx.linear_weight_cast ) @property def weight_grad_input(self) -> torch.Tensor: ctx = self.ctx if ctx.is_cuda and not ctx.weight_grad_mm_same_dtype: return self.input_chunk return self.input_chunk_acc @property def logits_upcast(self) -> torch.Tensor: # Non-CUDA mixed-dtype: copy into logits_acc_buf for matching-dtype mm. ctx = self.ctx if ctx.is_cuda or ctx.weight_grad_mm_same_dtype: return self.logits return ctx.logits_acc_buf.narrow(0, 0, self.bchunk_size).copy_(self.logits) @cached_property def grad_input_chunk(self) -> torch.Tensor: # Buf-and-copy commit happens in ctx.chunks() post-yield. ctx = self.ctx if ctx.alloc_input_grad_acc_buf: return ctx.input_grad_acc_buf.narrow(0, 0, self.bchunk_size) return ctx.grad_input.narrow(0, self.bchunk_start, self.bchunk_size) @cached_property def bias_grad_acc(self) -> torch.Tensor: # ``ctx.linear_bias_grad_chunk`` is an acc_dtype scratch that # accumulates the per-chunk sum + index_add at acc precision, # then commits to ``grad_linear_bias`` (input dtype) via the # ``ctx.chunks()`` post-yield hook -- mirrors # ``weight_grad_chunk``'s bulk-+-correction pattern for # ``grad_linear_weight``. On the non-acc path, accumulation # goes directly into ``grad_linear_bias`` (no commit, no # post-yield action). ctx = self.ctx if ctx.alloc_linear_bias_grad_chunk: return ctx.linear_bias_grad_chunk return ctx.grad_linear_bias @cached_property def logits_downcast(self) -> torch.Tensor: # Lazy: first access must come AFTER in-place mods to self.logits. ctx = self.ctx if ctx.loop_caches_logits_downcast: return self.logits.to(ctx.linear_weight.dtype) return self.logits @dataclasses.dataclass class _ChunkContext: """Per-call state for the chunked loop, built once via ``build``. Methods hide dtype/device/acc_policy dispatch behind single math ops; dispatch-free per-iter math is inlined into the loop body. Buffers dispatch decided are not needed are rank-matching empty tensors (``when=False`` in ``_make_*``) so the dataclass surface stays uniform. """ dtype: torch.dtype num_batches: int in_features: int num_classes: int is_cuda: bool use_acc_dtype: bool acc_dtype: torch.dtype weight_chunk_dtype: torch.dtype grad_input_dtype: torch.dtype linear_weight_cast_dtype: torch.dtype input: torch.Tensor target: torch.Tensor # uncorrected; per-iter slice is from corrected_target weight: torch.Tensor | None ignore_index: int reduction: str linear_weight: torch.Tensor linear_bias: torch.Tensor | None # reduction='none' backward only: per-sample upstream grad (N,). When # set, neg_weight_target folds it in so the shared grad loop produces # the grad_output-weighted VJP. None on every other path. loss_grad_output: torch.Tensor | None # The optional "when=" buffers (weight_grad_chunk, logits_acc_buf, # input_grad_acc_buf, input_chunk_acc_buf, grad_input, # grad_linear_weight, grad_linear_bias) are cached_properties below. logits_buf: torch.Tensor tmp: torch.Tensor output: torch.Tensor compute_input_grad: bool compute_linear_weight_grad: bool compute_linear_bias_grad: bool alloc_weight_grad_chunk: bool alloc_linear_bias_grad_chunk: bool alloc_input_grad_acc_buf: bool alloc_input_chunk_acc_buf: bool alloc_logits_acc_buf: bool forward_uses_acc_input: bool forward_uses_cuda_out_dtype: bool weight_grad_mm_same_dtype: bool input_grad_uses_logits_lw: bool loop_caches_logits_downcast: bool weight_grad_uses_logits_buf_temp: bool @cached_property def _mask(self) -> torch.Tensor: return self.target == self.ignore_index @cached_property def corrected_target(self) -> torch.Tensor: # Replace out-of-range ignore_index values with 0 lazily so # downstream index_select / index_add_ stay in bounds. if self.ignore_index < 0 or self.ignore_index >= self.num_classes: return torch.where(self._mask, 0, self.target) return self.target @cached_property def neg_weight_target(self) -> torch.Tensor: # Per-row weighting: -(weight[target] if weight else 1) / d on # unmasked positions, 0 on masked positions. d = sum of unmasked # weights for "mean", 1 for "sum". mask = self._mask target = self.corrected_target weight = self.weight weight_chunk_dtype = self.weight_chunk_dtype if weight is None: neg_weight_target = (~mask).to(weight_chunk_dtype) elif target.numel() > weight.numel(): neg_weight_target = torch.where( mask, 0, weight.to(weight_chunk_dtype).index_select(0, target) ) else: neg_weight_target = torch.where( mask, 0, weight.index_select(0, target).to(weight_chunk_dtype) ) if self.reduction == "mean": d = neg_weight_target.sum() neg_weight_target.div_(torch.where(d == 0, torch.nan, -d)) elif self.reduction == "sum": neg_weight_target.neg_() else: # "none" # Forward (loss_grad_output is None) intentionally leaves the # bare +W[T] unsigned: the loss-only branch needs the positive # per-row loss weight. Only the backward (upstream grad set) # negates and folds in grad_output[n] for the VJP -- do NOT add # a .neg_() here to mirror the "sum" case. if self.loss_grad_output is not None: neg_weight_target.neg_().mul_( self.loss_grad_output.to(neg_weight_target.dtype) ) return neg_weight_target @cached_property def linear_weight_cast(self) -> torch.Tensor: if self.linear_weight_cast_dtype != self.dtype: return self.linear_weight.to(self.linear_weight_cast_dtype) return self.linear_weight @cached_property def linear_bias_cast(self) -> torch.Tensor | None: # ``linear_bias`` is added to per-chunk logits, so its dtype # must match the logits buffer (which may be acc_dtype on the # upcast path). ``None`` when the user passed no bias -- the # loop skips the add. if self.linear_bias is None: return None if self.linear_bias.dtype != self.logits_buf.dtype: return self.linear_bias.to(self.logits_buf.dtype) return self.linear_bias @cached_property def weight_grad_chunk(self) -> torch.Tensor: # (V, F) acc_dtype accumulator for keep-path weight-grad mm; # rank-empty on the direct (compact) path. return _make_empty( (self.num_classes, self.in_features), self.acc_dtype, self.input.device, when=self.alloc_weight_grad_chunk, ) @cached_property def logits_acc_buf(self) -> torch.Tensor: # (B, V) scratch for non-CUDA mixed-dtype logits_upcast. return _make_empty( (self.logits_buf.shape[0], self.num_classes), self.acc_dtype, self.input.device, when=self.alloc_logits_acc_buf, ) @cached_property def input_grad_acc_buf(self) -> torch.Tensor: # (B, F) scratch for input-grad addmm on the buf-and-copy path # (CPU mixed-dtype, MPS). return _make_empty( (self.logits_buf.shape[0], self.in_features), self.linear_weight_cast_dtype, self.input.device, when=self.alloc_input_grad_acc_buf, ) @cached_property def input_chunk_acc_buf(self) -> torch.Tensor: return _make_empty( (self.logits_buf.shape[0], self.in_features), self.acc_dtype, self.input.device, when=self.alloc_input_chunk_acc_buf, ) @cached_property def grad_input(self) -> torch.Tensor: return _make_empty( self.input.shape, self.grad_input_dtype, self.input.device, when=self.compute_input_grad, ) @cached_property def grad_linear_weight(self) -> torch.Tensor: return _make_zeros( self.linear_weight.shape, self.dtype, self.input.device, when=self.compute_linear_weight_grad, ) @cached_property def grad_linear_bias(self) -> torch.Tensor: # Same shape invariant as ``linear_bias``: # ``linear_bias.shape == linear_weight.shape[:-1]``. Final # accumulator in input dtype to match ``grad_linear_weight``. return _make_zeros( self.linear_weight.shape[:-1], self.dtype, self.input.device, when=self.compute_linear_bias_grad, ) @cached_property def linear_bias_grad_chunk(self) -> torch.Tensor: # ``linear_weight.shape[:-1]`` acc_dtype per-chunk scratch, # mirroring ``weight_grad_chunk``'s role for # ``grad_linear_weight``: the sub_ + index_add_ pair # accumulates per-chunk in acc_dtype, then a single downcast # commit (in ``chunks()`` post-yield) lands in # ``grad_linear_bias``. ``_make_zeros`` so the first chunk's # accumulators start at zero; the post-yield hook zeroes it # again for subsequent chunks. return _make_zeros( self.linear_weight.shape[:-1], self.acc_dtype, self.input.device, when=self.alloc_linear_bias_grad_chunk, ) @classmethod def build( cls, input: torch.Tensor, linear_weight: torch.Tensor, target: torch.Tensor, linear_bias: torch.Tensor | None, weight: torch.Tensor | None, reduction: str, ignore_index: int, label_smoothing: float, batch_chunk_size: int, acc_policy: str, acc_dtype: torch.dtype, compute_input_grad: bool, compute_linear_weight_grad: bool, compute_linear_bias_grad: bool, loss_grad_output: torch.Tensor | None = None, ) -> "_ChunkContext": # ===== Validation ===== if target.dtype != torch.int64: raise TypeError( f"linear_cross_entropy: target dtype must be torch.int64, got {target.dtype}." ) if label_smoothing > 0.0: raise NotImplementedError( "linear_cross_entropy does not support label smoothing" ) if reduction not in {"mean", "sum", "none"}: raise NotImplementedError( f"linear_cross_entropy does not support {reduction=}" ) if linear_bias is not None and linear_bias.shape != linear_weight.shape[:-1]: raise RuntimeError( "linear_cross_entropy: expected linear_bias shape " f"{tuple(linear_weight.shape[:-1])}, got {tuple(linear_bias.shape)}." ) device = input.device dtype = input.dtype num_batches, in_features = input.shape num_classes, _ = linear_weight.shape # CUDA gates the out_dtype= mm fast path; the non-CUDA path # (CPU, MPS, XPU, ...) routes mixed-dtype mm through explicit # casts until out_dtype= is validated on those backends. is_cuda = device.type == "cuda" is_mps = device.type == "mps" if dtype != acc_dtype and not ( dtype in {torch.float16, torch.bfloat16} and acc_dtype == torch.float32 ): raise RuntimeError( "linear_cross_entropy supports float32 acc_dtype with" f" float16/bfloat16 inputs, but got {acc_dtype} acc_dtype and {dtype} inputs." ) use_acc_dtype = dtype != acc_dtype # Internal dtype layout. ``compact`` reuses ``balanced``'s layout; # its savings come from skipping the weight_grad_chunk scratch, not dtype. is_memory_like = acc_policy in {"balanced", "compact"} if use_acc_dtype: output_dtype = acc_dtype if dtype == torch.float16 else dtype grad_input_dtype = dtype if is_memory_like else acc_dtype # fp16 + memory-like keeps logits at fp16 (2BC vs 4BC); # everything else upcasts for softmax stability. logits_buf_dtype = ( dtype if dtype == torch.float16 and is_memory_like else acc_dtype ) # acc_dtype mandatory: 1/N for reduction="mean" is subnormal # in fp16 at N >= 65536. weight_chunk_dtype = acc_dtype else: output_dtype = grad_input_dtype = logits_buf_dtype = weight_chunk_dtype = ( dtype ) # ===== Dispatch flags ===== # CUDA + balanced uses cuBLAS out_dtype= directly; no cast. needs_linear_weight_cast = use_acc_dtype and ( not is_cuda or (compute_input_grad and grad_input_dtype == logits_buf_dtype) ) linear_weight_cast_dtype = ( logits_buf_dtype if needs_linear_weight_cast else dtype ) alloc_weight_grad_chunk = compute_linear_weight_grad and not ( acc_policy == "compact" and (is_cuda or logits_buf_dtype == dtype) ) # Per-chunk acc_dtype scratch for grad_linear_bias; same # bulk-+-correction precision rationale as # ``alloc_weight_grad_chunk``. Fires only when there is # precision to gain: ``use_acc_dtype`` and the user asked for # the bias gradient. alloc_linear_bias_grad_chunk = compute_linear_bias_grad and use_acc_dtype alloc_input_grad_acc_buf = ( compute_input_grad and not is_cuda and (grad_input_dtype != linear_weight_cast_dtype or is_mps) ) alloc_input_chunk_acc_buf = use_acc_dtype and ( compute_linear_weight_grad or (not is_cuda and dtype != logits_buf_dtype) ) forward_uses_acc_input = ( use_acc_dtype and not is_cuda and dtype != logits_buf_dtype ) forward_uses_cuda_out_dtype = is_cuda and use_acc_dtype weight_grad_mm_same_dtype = logits_buf_dtype == acc_dtype # Storage-trick fires when logits' dtype differs from the input- # grad accumulator dtype. input_grad_uses_logits_lw = ( is_cuda and compute_input_grad and logits_buf_dtype != grad_input_dtype ) # Per-iter ``logits.to(linear_weight.dtype)`` shared between the # inlined input-grad addmm and the inlined direct weight-grad addmm_. loop_caches_logits_downcast = input_grad_uses_logits_lw or ( compute_linear_weight_grad and not alloc_weight_grad_chunk and logits_buf_dtype != dtype ) alloc_logits_acc_buf = ( alloc_weight_grad_chunk and not is_cuda and logits_buf_dtype != acc_dtype ) # Direct weight-grad path: reuse logits_buf storage for the # acc_dtype->dtype cast when its bytes/row fits the (B, F) cast. weight_grad_uses_logits_buf_temp = ( not alloc_weight_grad_chunk and use_acc_dtype and num_classes * logits_buf_dtype.itemsize >= in_features * dtype.itemsize ) return cls( dtype=dtype, num_batches=num_batches, in_features=in_features, num_classes=num_classes, is_cuda=is_cuda, use_acc_dtype=use_acc_dtype, acc_dtype=acc_dtype, weight_chunk_dtype=weight_chunk_dtype, grad_input_dtype=grad_input_dtype, linear_weight_cast_dtype=linear_weight_cast_dtype, input=input, target=target, weight=weight, ignore_index=ignore_index, reduction=reduction, linear_weight=linear_weight, linear_bias=linear_bias, loss_grad_output=loss_grad_output, logits_buf=torch.empty( (batch_chunk_size, num_classes), dtype=logits_buf_dtype, device=device, requires_grad=False, ), tmp=torch.empty( (batch_chunk_size,), dtype=logits_buf_dtype, device=device, requires_grad=False, ), output=_make_zeros((), output_dtype, device), compute_input_grad=compute_input_grad, compute_linear_weight_grad=compute_linear_weight_grad, compute_linear_bias_grad=compute_linear_bias_grad, alloc_weight_grad_chunk=alloc_weight_grad_chunk, alloc_linear_bias_grad_chunk=alloc_linear_bias_grad_chunk, alloc_input_grad_acc_buf=alloc_input_grad_acc_buf, alloc_input_chunk_acc_buf=alloc_input_chunk_acc_buf, alloc_logits_acc_buf=alloc_logits_acc_buf, forward_uses_acc_input=forward_uses_acc_input, forward_uses_cuda_out_dtype=forward_uses_cuda_out_dtype, weight_grad_mm_same_dtype=weight_grad_mm_same_dtype, input_grad_uses_logits_lw=input_grad_uses_logits_lw, loop_caches_logits_downcast=loop_caches_logits_downcast, weight_grad_uses_logits_buf_temp=weight_grad_uses_logits_buf_temp, ) def chunks(self): """Yield a ``_ChunkViews`` per iter. Post-yield, run the scratch-to-final commits that the loop body deferred: the buf-and-copy input-grad slice into ``grad_input`` (skipped on the fast path where ``grad_input_chunk`` aliases ``grad_input``), and the acc_dtype linear_bias scratch into ``grad_linear_bias`` (skipped when accumulation went straight into ``grad_linear_bias`` in input dtype). """ batch_chunk_size = self.logits_buf.shape[0] for bchunk_start in range(0, self.num_batches, batch_chunk_size): bchunk_size = min(batch_chunk_size, self.num_batches - bchunk_start) chunk = self.bind_chunk(bchunk_start, bchunk_size) yield chunk if self.alloc_input_grad_acc_buf: self.grad_input.narrow(0, bchunk_start, bchunk_size).copy_( chunk.grad_input_chunk ) if self.alloc_linear_bias_grad_chunk: # ``chunk.bias_grad_acc`` aliases ``linear_bias_grad_chunk`` # here; downcast-commit, then zero for the next chunk. self.grad_linear_bias.add_( chunk.bias_grad_acc.to(self.grad_linear_bias.dtype) ) chunk.bias_grad_acc.zero_() def bind_chunk(self, bchunk_start: int, bchunk_size: int) -> _ChunkViews: input_chunk = self.input.narrow(0, bchunk_start, bchunk_size) target_chunk = self.corrected_target.narrow(0, bchunk_start, bchunk_size) weight_chunk = self.neg_weight_target.narrow(0, bchunk_start, bchunk_size) logits = self.logits_buf.narrow(0, 0, bchunk_size) input_chunk_acc = ( self.input_chunk_acc_buf.narrow(0, 0, bchunk_size).copy_(input_chunk) if self.alloc_input_chunk_acc_buf else input_chunk ) return _ChunkViews( ctx=self, bchunk_start=bchunk_start, bchunk_size=bchunk_size, input_chunk=input_chunk, target_chunk=target_chunk, weight_chunk=weight_chunk, logits=logits, input_chunk_acc=input_chunk_acc, ) def mm(self, mat1: torch.Tensor, mat2: torch.Tensor, *, out: torch.Tensor) -> None: # Mismatched dtype: cuBLAS ``out_dtype=`` upcasts inside the matmul. if mat1.dtype == out.dtype: torch.mm(mat1, mat2, out=out) else: torch.mm(mat1, mat2, out_dtype=out.dtype, out=out) def amax(self, x: torch.Tensor) -> torch.Tensor: out = self.tmp.narrow(0, 0, x.shape[0]).unsqueeze(1) torch.amax(x, dim=1, keepdim=True, out=out) return out def dotgather( self, weight: torch.Tensor, x: torch.Tensor, indices: torch.Tensor ) -> torch.Tensor: return weight.dot(x.gather(1, indices.unsqueeze(1)).squeeze(1).to(weight.dtype)) def sumexp_(self, x: torch.Tensor, dim: int) -> torch.Tensor: # Result always in acc_dtype: fp16 softmax_denom can underflow # on rows with widely-spread logits, after which log(0)=-inf # poisons the loss and 1/0 the gradient. return x.exp_().sum(dim, dtype=self.acc_dtype) def sum(self, x: torch.Tensor, dim: int) -> torch.Tensor: # Reduce ``x`` along ``dim`` at acc_dtype on mixed-precision # paths and at input dtype otherwise -- the precision # contract callers see is "as accurate as the active policy # allows". Currently used by the bias-grad path so the per- # chunk sum carries full precision before the post-yield # downcast commit to grad_linear_bias. target_dtype = self.acc_dtype if self.use_acc_dtype else self.dtype return x.sum(dim, dtype=target_dtype) def div(self, num: torch.Tensor, den: torch.Tensor) -> torch.Tensor: factor = self.tmp.narrow(0, 0, num.shape[0]) # MPS lacks fused cross-dtype torch.div(out=); copy_ fallback. if self.input.device.type == "mps" and ( num.dtype != factor.dtype or den.dtype != factor.dtype ): factor.copy_(num / den) else: torch.div(num, den, out=factor) return factor def mul(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: # Pick a destination from in-flight scratch when possible. if self.use_acc_dtype: return x.mul_(w.unsqueeze(1)) if self.num_classes >= self.in_features: return torch.mul( x, w.unsqueeze(1), out=self.logits_buf.narrow(0, 0, x.shape[0]).narrow( 1, 0, self.in_features ), ) return x * w.unsqueeze(1) def to(self, x: torch.Tensor, *, dtype: torch.dtype) -> torch.Tensor: # Reuse logits_buf storage as cast destination when its bytes/row fits. if x.dtype == dtype: return x if self.weight_grad_uses_logits_buf_temp: x = ( self.logits_buf.view(dtype) .narrow(0, 0, x.shape[0]) .narrow(1, 0, x.shape[1]) .copy_(x) ) return x return x.to(dtype) def _linear_cross_entropy_batch_chunked_accumulator( input: torch.Tensor, linear_weight: torch.Tensor, target: torch.Tensor, linear_bias: torch.Tensor | None, weight: torch.Tensor | None, reduction: str, ignore_index: int, label_smoothing: float, batch_chunk_size: int, acc_policy: str, acc_dtype: torch.dtype, compute_input_grad: bool, compute_linear_weight_grad: bool, compute_linear_bias_grad: bool, loss_grad_output: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Chunked loop shared by all three custom-op entry points; returns ``(loss, grad_input, grad_linear_weight, grad_linear_bias)``. Caller -> branch: scalar-op forward (reduction in {mean,sum}) -> main grad loop no_reduction forward (none, loss_grad_output None) -> loss-only branch no_reduction backward (none, loss_grad_output set) -> main grad loop - ``_linear_cross_entropy_batch_chunked`` (scalar reduction): the forward precomputes the scalar loss and all gradients in one pass; backward just scales the stashed grads by the upstream gradient. - ``_linear_cross_entropy_batch_chunked_no_reduction`` (forward): ``reduction='none'`` with ``loss_grad_output is None`` takes the loss-only branch, returning the per-sample loss ``(N,)`` and no gradients (they are recomputed in its backward instead). - ``_linear_cross_entropy_batch_chunked_no_reduction`` (backward): ``loss_grad_output`` set folds the per-sample upstream grad into ``neg_weight_target``, so this same loop recomputes the chunked grads as the grad_output-weighted VJP (the scalar loss it also accumulates is then meaningless and ignored). Dispatch contract ----------------- The chunked loop body is the mathematical description of the algorithm. All policy / device / dtype dispatch lives in ``_ChunkContext`` (and its ``_ChunkViews`` peers): - **Buffer choices** -- whether a per-chunk acc_dtype scratch is used or accumulation runs directly into the final tensor -- live in ``alloc_*`` flags computed in ``_ChunkContext.build`` and surfaced as ``_ChunkViews`` properties (e.g. ``grad_input_chunk``, ``bias_grad_acc``). - **Scratch-to-final commits** for those buffers live in ``_ChunkContext.chunks()`` post-yield hooks (e.g. the ``grad_input`` buf-and-copy commit, the ``grad_linear_bias`` downcast-and-zero commit). - **Dtype dispatch on individual math operations** lives in ``_ChunkContext`` methods that mirror PyTorch ops -- ``mm``, ``amax``, ``dotgather``, ``sumexp_``, ``sum``, ``div``, ``mul``, ``to`` -- each one hiding ``out_dtype=`` / acc-dtype / buffer-reuse decisions behind a single math call. The function body therefore should not introduce inline ``if ctx.use_acc_dtype`` / ``if ctx.is_cuda`` / etc. branches; new policy-aware behaviour belongs in one of the three locations above. The two existing inline branches (``if compute_input_grad``, ``if compute_linear_weight_grad``, etc.) gate optional outputs, not dispatch. """ # Direct callers must resolve "auto" / None via _adjust first. if acc_policy == "auto" or acc_dtype is None: raise RuntimeError( f"unresolved acc_policy={acc_policy!r} or acc_dtype={acc_dtype!r};" " use F.linear_cross_entropy." ) ctx = _ChunkContext.build( input, linear_weight, target, linear_bias, weight, reduction, ignore_index, label_smoothing, batch_chunk_size, acc_policy, acc_dtype, compute_input_grad, compute_linear_weight_grad, compute_linear_bias_grad, loss_grad_output=loss_grad_output, ) dtype = ctx.dtype # reduction='none' forward (no upstream grad yet): per-sample loss # into an (N,) output, no gradient precompute. ``chunk.weight_chunk`` # is the unsigned masked class weight here (neg_weight_target's "none" # forward form), so each row's loss is W[T[n]] * (log denom - # shifted_logit[T[n]]) = W[T[n]] * (-log_softmax). The "none" backward # (loss_grad_output set) routes through the grad loop below instead. if reduction == "none" and loss_grad_output is None: linear_bias_cast = ctx.linear_bias_cast out = torch.empty(ctx.num_batches, dtype=dtype, device=ctx.input.device) for chunk in ctx.chunks(): logits = chunk.logits ctx.mm(chunk.input, chunk.linear_weight.T, out=logits) if linear_bias_cast is not None: logits.add_(linear_bias_cast) logits.sub_(ctx.amax(logits)) # Read the target logit BEFORE ``sumexp_`` -- it does ``exp_()`` # in place, overwriting logits with exp(shifted). ls_target = logits.gather(1, chunk.target_chunk.unsqueeze(1)).squeeze(1) softmax_denom = ctx.sumexp_(logits, dim=1) loss_chunk = softmax_denom.log_().sub_(ls_target.to(softmax_denom.dtype)) loss_chunk.mul_(chunk.weight_chunk.to(softmax_denom.dtype)) out.narrow(0, chunk.bchunk_start, chunk.bchunk_size).copy_(loss_chunk) return ( out, ctx.grad_input.to(dtype), ctx.grad_linear_weight, ctx.grad_linear_bias, ) output = ctx.output grad_input = ctx.grad_input grad_linear_weight = ctx.grad_linear_weight grad_linear_bias = ctx.grad_linear_bias weight_grad_chunk = ctx.weight_grad_chunk linear_weight_cast = ctx.linear_weight_cast linear_bias_cast = ctx.linear_bias_cast if reduction == "mean" and ctx.num_batches == 0: output.fill_(torch.nan) compute_grads = ( compute_input_grad or compute_linear_weight_grad or compute_linear_bias_grad ) # Do not ``break`` from this loop -- ``ctx.chunks()`` runs a # post-yield grad_input commit that the in-flight chunk needs. for chunk in ctx.chunks(): logits = chunk.logits weight_chunk = chunk.weight_chunk target_chunk = chunk.target_chunk ctx.mm(chunk.input, chunk.linear_weight.T, out=logits) if linear_bias_cast is not None: logits.add_(linear_bias_cast) # broadcasts (V,) -> (B_chunk, V) logits.sub_(ctx.amax(logits)) # softmax stability # output += output.add_(ctx.dotgather(weight_chunk, logits, target_chunk)) softmax_denom = ctx.sumexp_(logits, dim=1) if compute_grads: # logits *= weight_chunk / softmax_denom (= softmax(logits) * w) # MPS: in-place mul_ on a narrow view of logits_buf doesn't # propagate to parent storage; the out= form does. torch.mul( logits, ctx.div(weight_chunk, softmax_denom).unsqueeze(1), out=logits, ) # output -= # softmax_denom is in acc_dtype (see ``sumexp_``); promote weight_chunk # for the dot rather than rounding the wider denominator down. output.sub_(weight_chunk.to(softmax_denom.dtype).dot(softmax_denom.log_())) if compute_grads: if compute_linear_bias_grad: # Must run BEFORE the linear-weight-grad branch below. # The ordering invariant is defensive across all # ``ctx.mul`` paths: only the (not use_acc_dtype && # num_classes >= in_features) branch overwrites the # first ``in_features`` columns of ``logits_buf`` via a # narrow ``out=`` view today, but ``ctx.mul``'s contract # does not promise to leave ``logits_buf`` untouched on # other paths -- placing the bias-grad read here keeps # the loop correct regardless of which ``ctx.mul`` # variant fires. # # The inline form expresses # dL/db_j = sum_n W[T[n]] * (softmax(X)[n, j] - delta(T[n], j)) # as ``-logits.sum(dim=0)`` (the softmax-weighted term) # plus ``index_add_(0, target_chunk, weight_chunk)`` # (the one-hot-weighted correction). ``logits`` is # ``softmax(X) * weight_chunk[n]`` here; ``weight_chunk`` # already carries the sign and the 1/d mean # normalization (``W`` here is the CE class weight, not # ``linear_weight``), so no extra scaling is needed. bias_grad_acc = chunk.bias_grad_acc bias_grad_acc.sub_(ctx.sum(logits, dim=0)) bias_grad_acc.index_add_(0, target_chunk, weight_chunk) if compute_input_grad: grad_input_chunk = chunk.grad_input_chunk input_grad_logits = chunk.input_grad_logits input_grad_linear_weight = chunk.input_grad_linear_weight # grad_input_chunk = linear_weight_cast[target_chunk] * weight_chunk torch.mul( torch.index_select(linear_weight_cast, 0, target_chunk), weight_chunk.to(grad_input_chunk.dtype).unsqueeze(1), out=grad_input_chunk, ) # grad_input_chunk -= input_grad_logits @ input_grad_linear_weight torch.addmm( grad_input_chunk, input_grad_logits, input_grad_linear_weight, alpha=-1, out=grad_input_chunk, ) if compute_linear_weight_grad: input_chunk_acc = chunk.input_chunk_acc # grad_linear_weight += per-chunk weight grad # (= -logits.T @ input + input*weight scattered at target rows) if ctx.alloc_weight_grad_chunk: # Stage per-chunk weight grad in acc_dtype scratch; # one sub_ commit keeps bulk-minus-correction in fp32. logits_upcast = chunk.logits_upcast weight_grad_input = chunk.weight_grad_input ctx.mm( logits_upcast.T, weight_grad_input, out=weight_grad_chunk, ) temp = ctx.mul(input_chunk_acc, weight_chunk) weight_grad_chunk.index_add_(0, target_chunk, temp, alpha=-1) grad_linear_weight.sub_(weight_grad_chunk) else: # Stream bulk + correction directly into grad_lw without # a (V, F) scratch (the compact path). grad_linear_weight.addmm_( chunk.logits_downcast.T, chunk.input_chunk, alpha=-1 ) temp = ctx.to( ctx.mul(input_chunk_acc, weight_chunk), dtype=grad_linear_weight.dtype, ) grad_linear_weight.index_add_(0, target_chunk, temp, alpha=1) return ( output.to(dtype), grad_input.to(dtype), grad_linear_weight, grad_linear_bias, ) # Scalar-reduction (mean/sum) op. Returns # ``(loss, grad_input, grad_linear_weight, grad_linear_bias)`` with grads # precomputed in forward (the accumulator runs the full grad loop) and # stashed on ctx by ``setup_context``; backward mutates them in place via # ``.mul_()`` by the upstream gradient. ``torch.library.register_autograd`` # has no ``save_for_backward``-style state, hence the return-tuple + # ctx-stash dance. ``mutates_args=()`` reflects that inputs are not mutated. # Shape invariants: grad_input.shape == input.shape, # grad_linear_weight.shape == linear_weight.shape, grad_linear_bias.shape == # linear_weight.shape[:-1] (here (C,), since the upstream gate requires # out_features == ()). @torch.library.custom_op( "torch_nn::_linear_cross_entropy_batch_chunked", mutates_args=() ) def _linear_cross_entropy_batch_chunked( input: torch.Tensor, linear_weight: torch.Tensor, target: torch.Tensor, linear_bias: torch.Tensor | None, weight: torch.Tensor | None, reduction: str, ignore_index: int, label_smoothing: float, batch_chunk_size: int, acc_policy: str, acc_dtype: torch.dtype, allow_retain_graph: bool, compute_input_grad: bool, compute_linear_weight_grad: bool, compute_linear_bias_grad: bool, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Scalar-reduction (mean/sum) chunked op: precomputes the loss and gradients in forward, backward scales by the upstream gradient. The chunked math lives in ``_linear_cross_entropy_batch_chunked_accumulator``. """ # AOTAutograd/AOTInductor bake compute_*_grad at trace time; catch # the silent-corruption case (False at trace, but grad-enabled at # runtime with a requires_grad leaf). These guards are specific to # the precompute path -- the reduction='none' op recomputes grads in # its own backward, so it does not route through here. grad_enabled = torch.is_grad_enabled() if not compute_input_grad and input.requires_grad and grad_enabled: raise RuntimeError( "linear_cross_entropy chunked op: compute_input_grad was False at " "trace time but input.requires_grad is True at runtime; recompile " "the graph with the desired requires_grad." ) if not compute_linear_weight_grad and linear_weight.requires_grad and grad_enabled: raise RuntimeError( "linear_cross_entropy chunked op: compute_linear_weight_grad was " "False at trace time but linear_weight.requires_grad is True at " "runtime; recompile the graph with the desired requires_grad." ) if ( not compute_linear_bias_grad and linear_bias is not None and linear_bias.requires_grad and grad_enabled ): raise RuntimeError( "linear_cross_entropy chunked op: compute_linear_bias_grad was False at " "trace time but linear_bias.requires_grad is True at runtime; recompile " "the graph with the desired requires_grad." ) return _linear_cross_entropy_batch_chunked_accumulator( input, linear_weight, target, linear_bias, weight, reduction, ignore_index, label_smoothing, batch_chunk_size, acc_policy, acc_dtype, compute_input_grad, compute_linear_weight_grad, compute_linear_bias_grad, ) # Op arg count; backward returns this many slots. Hand-maintained # because CustomOpDef has no public arg-introspection. Update on # signature change. _NUM_OP_INPUTS = 15 @_linear_cross_entropy_batch_chunked.register_fake def _( input, linear_weight, target, linear_bias, weight, reduction, ignore_index, label_smoothing, batch_chunk_size, acc_policy: str, acc_dtype, allow_retain_graph, compute_input_grad, compute_linear_weight_grad, compute_linear_bias_grad, ): if reduction in {"mean", "sum"}: result = torch.empty((), dtype=input.dtype, device=input.device) else: raise NotImplementedError(f"linear_cross_entropy does not support {reduction=}") grad_input_shape = input.shape if compute_input_grad else (0, 0) grad_linear_weight_shape = ( linear_weight.shape if compute_linear_weight_grad else (0, 0) ) # ``linear_bias.shape == linear_weight.shape[:-1]`` (the (C, *out_features) # invariant); on this code path ``out_features == ()`` so it reduces # to (C,), but expressing it this way keeps the relationship visible. grad_linear_bias_shape = ( linear_weight.shape[:-1] if compute_linear_bias_grad else (0,) ) grad_input = torch.empty( grad_input_shape, dtype=input.dtype, device=input.device, requires_grad=False, ) # The real implementation allocates ``grad_linear_weight`` and # ``grad_linear_bias`` in ``self.dtype = input.dtype`` (chunked op # supports mixed-precision: ``input.dtype`` may differ from # ``linear_weight.dtype``). Mirror that here so the fake-tensor # shape/dtype matches what eager produces. grad_linear_weight = torch.empty( grad_linear_weight_shape, dtype=input.dtype, device=linear_weight.device, requires_grad=False, ) grad_linear_bias = torch.empty( grad_linear_bias_shape, dtype=input.dtype, device=input.device, requires_grad=False, ) return result, grad_input, grad_linear_weight, grad_linear_bias @_linear_cross_entropy_batch_chunked.register_vmap def _vmap(info, in_dims, *args): """vmap rule (slow path: per-sample Python loop). A fold-into-num_batches fast path is now possible for both ops via the reduction='none' chunked op (``_linear_cross_entropy_batch_chunked_no_reduction``); left as a follow-up. """ batch_size = info.batch_size moved_args = [ arg.movedim(in_dim, 0) if in_dim is not None and isinstance(arg, torch.Tensor) else arg for arg, in_dim in zip(args, in_dims) ] outputs = [ _linear_cross_entropy_batch_chunked( *[ arg[i] if in_dim is not None and isinstance(arg, torch.Tensor) else arg for arg, in_dim in zip(moved_args, in_dims) ] ) for i in range(batch_size) ] losses = torch.stack([o[0] for o in outputs]) grad_inputs = torch.stack([o[1] for o in outputs]) grad_linear_weights = torch.stack([o[2] for o in outputs]) grad_linear_biases = torch.stack([o[3] for o in outputs]) return (losses, grad_inputs, grad_linear_weights, grad_linear_biases), (0, 0, 0, 0) def _linear_cross_entropy_batch_chunked_backward( ctx, grad_output, _gi_grad, _gw_grad, _gb_grad ): result = [None] * _NUM_OP_INPUTS if ctx.compute_input_grad: if ctx.allow_retain_graph: result[0] = ctx._gi * grad_output else: gi, ctx._gi = ctx._gi, None if gi is None: raise RuntimeError( "linear_cross_entropy chunked backward called twice; " "retain_graph=True / double backward are not supported." ) result[0] = gi.mul_(grad_output) if ctx.compute_linear_weight_grad: if ctx.allow_retain_graph: result[1] = ctx._gw * grad_output else: gw, ctx._gw = ctx._gw, None if gw is None: raise RuntimeError( "linear_cross_entropy chunked backward called twice; " "retain_graph=True / double backward are not supported." ) result[1] = gw.mul_(grad_output) # ``linear_bias`` is forward arg position 3 (after input, linear_weight, target). if ctx.compute_linear_bias_grad: if ctx.allow_retain_graph: result[3] = ctx._gb * grad_output else: gb, ctx._gb = ctx._gb, None if gb is None: raise RuntimeError( "linear_cross_entropy chunked backward called twice; " "retain_graph=True / double backward are not supported." ) result[3] = gb.mul_(grad_output) return tuple(result) _linear_cross_entropy_batch_chunked.register_autograd( _linear_cross_entropy_batch_chunked_backward, setup_context=_linear_cross_entropy_batch_chunked_setup_context, ) # reduction='none' op. Separate from the scalar-reduction op because the # autograd model differs: per-sample loss makes ``grad_output`` a vector, # so grad_linear_weight / grad_linear_bias cannot be precomputed in forward # (they need the per-sample upstream grad before summing over the batch, and # storing the per-sample (N, V) softmax would defeat chunking). Instead the # forward returns only the loss and saves its inputs; backward recomputes the # chunked grads with ``grad_output`` folded into ``weight_chunk`` (both share # ``_linear_cross_entropy_batch_chunked_accumulator``). Peak memory stays bounded to # (batch_chunk_size, num_classes); the cost is one recomputed logits matmul # per chunk in backward. Recompute-from-saved-inputs supports retain_graph # naturally (no consumed buffers), so no allow_retain_graph flag is needed. @torch.library.custom_op( "torch_nn::_linear_cross_entropy_batch_chunked_no_reduction", mutates_args=() ) def _linear_cross_entropy_batch_chunked_no_reduction( input: torch.Tensor, linear_weight: torch.Tensor, target: torch.Tensor, linear_bias: torch.Tensor | None, weight: torch.Tensor | None, ignore_index: int, batch_chunk_size: int, acc_policy: str, acc_dtype: torch.dtype, ) -> torch.Tensor: """reduction='none' chunked forward: returns per-sample loss ``(N,)``. Gradients are not precomputed (the per-sample upstream grad is unknown at forward); the accumulator's loss-only branch runs here, and the backward recomputes the chunked grads. See ``_linear_cross_entropy_batch_chunked_accumulator``. """ # loss_grad_output=None -> the accumulator's reduction='none' forward # (loss-only) branch; compute_*=False since nothing is precomputed. return _linear_cross_entropy_batch_chunked_accumulator( input, linear_weight, target, linear_bias, weight, "none", ignore_index, 0.0, batch_chunk_size, acc_policy, acc_dtype, compute_input_grad=False, compute_linear_weight_grad=False, compute_linear_bias_grad=False, )[0] def _linear_cross_entropy_batch_chunked_no_reduction_setup_context(ctx, inputs, output): ( input, linear_weight, target, linear_bias, weight, ignore_index, batch_chunk_size, acc_policy, acc_dtype, ) = inputs # Detach before saving: the backward recomputes the grads analytically # from these values (it does not differentiate through them), and saving # ``linear_weight`` -- a differentiable ``reshape`` view of the # ``nn.Linear`` weight -- would make its SavedVariable retain the base # parameter, inflating its use_count past backward and breaking # swap-based ``module.to()`` (the custom-op node keeps saved tensors # alive while the loss output is still referenced). ctx.save_for_backward( input.detach(), linear_weight.detach(), target, linear_bias.detach() if linear_bias is not None else None, weight.detach() if weight is not None else None, ) ctx._ignore_index = ignore_index ctx._batch_chunk_size = batch_chunk_size ctx._acc_policy = acc_policy ctx._acc_dtype = acc_dtype @_linear_cross_entropy_batch_chunked_no_reduction.register_fake def _( input, linear_weight, target, linear_bias, weight, ignore_index, batch_chunk_size, acc_policy: str, acc_dtype, ): return torch.empty( input.shape[:1], dtype=input.dtype, device=input.device, requires_grad=False ) @_linear_cross_entropy_batch_chunked_no_reduction.register_vmap def _vmap_no_reduction(info, in_dims, *args): """vmap rule (slow path: per-sample Python loop, mirrors the scalar-reduction op). Stacks the per-sample ``(N,)`` losses. """ batch_size = info.batch_size moved_args = [ arg.movedim(in_dim, 0) if in_dim is not None and isinstance(arg, torch.Tensor) else arg for arg, in_dim in zip(args, in_dims) ] outputs = [ _linear_cross_entropy_batch_chunked_no_reduction( *[ arg[i] if in_dim is not None and isinstance(arg, torch.Tensor) else arg for arg, in_dim in zip(moved_args, in_dims) ] ) for i in range(batch_size) ] return torch.stack(outputs), 0 def _linear_cross_entropy_batch_chunked_no_reduction_backward(ctx, grad_output): # grad_output is the per-sample upstream grad (N,). Forward arg order: # (input, linear_weight, target, linear_bias, weight, ignore_index, # batch_chunk_size, acc_policy, acc_dtype). input, linear_weight, target, linear_bias, weight = ctx.saved_tensors needs = ctx.needs_input_grad compute_input_grad = needs[0] compute_linear_weight_grad = needs[1] compute_linear_bias_grad = linear_bias is not None and needs[3] result: list[torch.Tensor | None] = [None] * 9 if not ( compute_input_grad or compute_linear_weight_grad or compute_linear_bias_grad ): return tuple(result) # loss_grad_output=grad_output -> the accumulator's grad loop with the # per-sample upstream grad folded into neg_weight_target (the "none" # backward path), recomputing the chunked grads as the VJP. The loss # (result[0]) it also returns is the meaningless scalar and dropped. _, grad_input, grad_linear_weight, grad_linear_bias = ( _linear_cross_entropy_batch_chunked_accumulator( input, linear_weight, target, linear_bias, weight, "none", ctx._ignore_index, 0.0, # label_smoothing ctx._batch_chunk_size, ctx._acc_policy, ctx._acc_dtype, compute_input_grad, compute_linear_weight_grad, compute_linear_bias_grad, loss_grad_output=grad_output, ) ) if compute_input_grad: result[0] = grad_input if compute_linear_weight_grad: result[1] = grad_linear_weight if compute_linear_bias_grad: result[3] = grad_linear_bias return tuple(result) _linear_cross_entropy_batch_chunked_no_reduction.register_autograd( _linear_cross_entropy_batch_chunked_no_reduction_backward, setup_context=_linear_cross_entropy_batch_chunked_no_reduction_setup_context, )