| """Fused CUDA kernels for memory-bound blocks: SiTU-GLU, the LM loss, token log-probs.""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn.functional as F |
| import triton |
| import triton.language as tl |
| from triton.language.extra import libdevice |
|
|
| from tiny_gdn.nn_common import ( |
| linear_cross_entropy_reference, |
| linear_token_logprobs_reference, |
| situ_glu_reference, |
| ) |
|
|
| BLOCK_SIZE = 1024 |
| VOCAB_BLOCK_SIZE = 4096 |
| LOGITS_CHUNK_BYTES = 1 << 30 |
|
|
|
|
| @triton.jit |
| def _situ_glu_forward_kernel( |
| gate_up, |
| output, |
| total, |
| width, |
| gate_cap, |
| up_cap, |
| BLOCK: tl.constexpr, |
| ): |
| offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) |
| mask = offsets < total |
| row = offsets // width |
| gate_offsets = offsets + row * width |
| gate = tl.load(gate_up + gate_offsets, mask=mask, other=0.0).to(tl.float32) |
| up = tl.load(gate_up + gate_offsets + width, mask=mask, other=0.0).to(tl.float32) |
| gated = gate_cap * libdevice.tanh(gate / gate_cap) * tl.sigmoid(gate) |
| result = gated * up_cap * libdevice.tanh(up / up_cap) |
| tl.store(output + offsets, result.to(output.dtype.element_ty), mask=mask) |
|
|
|
|
| @triton.jit |
| def _situ_glu_backward_kernel( |
| gate_up, |
| grad_output, |
| grad_gate_up, |
| total, |
| width, |
| gate_cap, |
| up_cap, |
| BLOCK: tl.constexpr, |
| ): |
| offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) |
| mask = offsets < total |
| row = offsets // width |
| gate_offsets = offsets + row * width |
| gate = tl.load(gate_up + gate_offsets, mask=mask, other=0.0).to(tl.float32) |
| up = tl.load(gate_up + gate_offsets + width, mask=mask, other=0.0).to(tl.float32) |
| grad = tl.load(grad_output + offsets, mask=mask, other=0.0).to(tl.float32) |
| gate_tanh = libdevice.tanh(gate / gate_cap) |
| up_tanh = libdevice.tanh(up / up_cap) |
| sigmoid = tl.sigmoid(gate) |
| gated = gate_cap * gate_tanh * sigmoid |
| capped_up = up_cap * up_tanh |
| gate_slope = sigmoid * ((1.0 - gate_tanh * gate_tanh) + gate_cap * gate_tanh * (1.0 - sigmoid)) |
| grad_gate = grad * capped_up * gate_slope |
| grad_up = grad * gated * (1.0 - up_tanh * up_tanh) |
| dtype = grad_gate_up.dtype.element_ty |
| tl.store(grad_gate_up + gate_offsets, grad_gate.to(dtype), mask=mask) |
| tl.store(grad_gate_up + gate_offsets + width, grad_up.to(dtype), mask=mask) |
|
|
|
|
| def _launch_grid(total: int) -> tuple[int]: |
| return (triton.cdiv(total, BLOCK_SIZE),) |
|
|
|
|
| class _SiTUGLUFunction(torch.autograd.Function): |
| @staticmethod |
| def forward(ctx, gate_up: torch.Tensor, gate_cap: float, up_cap: float) -> torch.Tensor: |
| gate_up = gate_up.contiguous() |
| width = gate_up.shape[-1] // 2 |
| output = gate_up.new_empty((*gate_up.shape[:-1], width)) |
| total = output.numel() |
| if total: |
| _situ_glu_forward_kernel[_launch_grid(total)]( |
| gate_up, output, total, width, gate_cap, up_cap, BLOCK=BLOCK_SIZE |
| ) |
| ctx.save_for_backward(gate_up) |
| ctx.caps = (gate_cap, up_cap) |
| return output |
|
|
| @staticmethod |
| def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]: |
| (gate_up,) = ctx.saved_tensors |
| gate_cap, up_cap = ctx.caps |
| grad_output = grad_output.contiguous() |
| grad_gate_up = torch.empty_like(gate_up) |
| width = gate_up.shape[-1] // 2 |
| total = grad_output.numel() |
| if total: |
| _situ_glu_backward_kernel[_launch_grid(total)]( |
| gate_up, |
| grad_output, |
| grad_gate_up, |
| total, |
| width, |
| gate_cap, |
| up_cap, |
| BLOCK=BLOCK_SIZE, |
| ) |
| return grad_gate_up, None, None |
|
|
|
|
| def situ_glu(gate_up: torch.Tensor, gate_cap: float, up_cap: float) -> torch.Tensor: |
| """``softcap(gate) * sigmoid(gate) * softcap(up)`` over ``[..., gate | up]``. |
| |
| The CUDA path evaluates in FP32 and rounds once; backward recomputes the |
| activations from ``gate_up`` instead of saving intermediates. |
| """ |
|
|
| if gate_up.shape[-1] % 2: |
| raise ValueError("gate_up must have an even last dimension") |
| if not gate_up.is_cuda: |
| return situ_glu_reference(gate_up, gate_cap, up_cap) |
| return _SiTUGLUFunction.apply(gate_up, gate_cap, up_cap) |
|
|
|
|
| @triton.jit |
| def _cross_entropy_rows_kernel( |
| logits, |
| targets, |
| losses, |
| log_partitions, |
| valid_count, |
| row_stride, |
| vocab_size, |
| z_loss_coefficient, |
| ignore_index, |
| COMPUTE_GRAD: tl.constexpr, |
| BLOCK: tl.constexpr, |
| ): |
| """Per row: logsumexp and NLL; optionally overwrite the row with d(mean loss)/d(logits).""" |
|
|
| row = tl.program_id(0).to(tl.int64) |
| row_logits = logits + row * row_stride |
| target = tl.load(targets + row) |
| columns = tl.arange(0, BLOCK) |
| if target == ignore_index: |
| tl.store(losses + row, 0.0) |
| tl.store(log_partitions + row, 0.0) |
| if COMPUTE_GRAD: |
| for start in range(0, vocab_size, BLOCK): |
| offsets = start + columns |
| zeros = tl.zeros((BLOCK,), dtype=logits.dtype.element_ty) |
| tl.store(row_logits + offsets, zeros, mask=offsets < vocab_size) |
| else: |
| running_max = tl.full((), float("-inf"), tl.float32) |
| running_sum = tl.zeros((), dtype=tl.float32) |
| for start in range(0, vocab_size, BLOCK): |
| offsets = start + columns |
| values = tl.load( |
| row_logits + offsets, mask=offsets < vocab_size, other=float("-inf") |
| ).to(tl.float32) |
| new_max = tl.maximum(running_max, tl.max(values, axis=0)) |
| running_sum = running_sum * tl.exp(running_max - new_max) + tl.sum( |
| tl.exp(values - new_max), axis=0 |
| ) |
| running_max = new_max |
| log_partition = running_max + tl.log(running_sum) |
| target_logit = tl.load(row_logits + target).to(tl.float32) |
| tl.store(losses + row, log_partition - target_logit) |
| tl.store(log_partitions + row, log_partition) |
| if COMPUTE_GRAD: |
| grad_scale = 1.0 / tl.load(valid_count).to(tl.float32) |
| softmax_scale = grad_scale * (1.0 + 2.0 * z_loss_coefficient * log_partition) |
| for start in range(0, vocab_size, BLOCK): |
| offsets = start + columns |
| in_vocab = offsets < vocab_size |
| values = tl.load(row_logits + offsets, mask=in_vocab, other=0.0).to(tl.float32) |
| grad = tl.exp(values - log_partition) * softmax_scale |
| grad = tl.where(offsets == target, grad - grad_scale, grad) |
| tl.store(row_logits + offsets, grad.to(logits.dtype.element_ty), mask=in_vocab) |
|
|
|
|
| class _LinearCrossEntropyFunction(torch.autograd.Function): |
| """Chunked ``states @ weight.T`` + cross-entropy + z-loss with gradients precomputed. |
| |
| Each logits chunk is consumed in place, so peak extra memory is one chunk |
| instead of ``[tokens, vocab]`` logits plus their softmax for backward. |
| """ |
|
|
| @staticmethod |
| def forward( |
| ctx, |
| states: torch.Tensor, |
| weight: torch.Tensor, |
| targets: torch.Tensor, |
| z_loss_coefficient: float, |
| ignore_index: int, |
| compute_grad: bool, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| states = states.contiguous() |
| targets = targets.contiguous() |
| rows, vocab_size = states.shape[0], weight.shape[0] |
| valid = targets.ne(ignore_index) |
| valid_count = valid.sum(dtype=torch.float32) |
| losses = states.new_empty(rows, dtype=torch.float32) |
| log_partitions = states.new_empty(rows, dtype=torch.float32) |
| grad_states = torch.empty_like(states) if compute_grad else None |
| grad_weight = ( |
| torch.zeros(weight.shape, dtype=torch.float32, device=weight.device) |
| if compute_grad |
| else None |
| ) |
| chunk_rows = max(1, LOGITS_CHUNK_BYTES // (vocab_size * states.element_size())) |
| for start in range(0, rows, chunk_rows): |
| end = min(start + chunk_rows, rows) |
| chunk = states[start:end] |
| logits = chunk @ weight.t() |
| _cross_entropy_rows_kernel[(end - start,)]( |
| logits, |
| targets[start:end], |
| losses[start:end], |
| log_partitions[start:end], |
| valid_count, |
| logits.stride(0), |
| vocab_size, |
| z_loss_coefficient, |
| ignore_index, |
| COMPUTE_GRAD=compute_grad, |
| BLOCK=VOCAB_BLOCK_SIZE, |
| num_warps=8, |
| ) |
| if compute_grad: |
| torch.mm(logits, weight, out=grad_states[start:end]) |
| if chunk.dtype == torch.float32: |
| grad_weight.addmm_(logits.t(), chunk) |
| else: |
| torch.addmm( |
| grad_weight, logits.t(), chunk, out_dtype=torch.float32, out=grad_weight |
| ) |
| cross_entropy = losses.sum() / valid_count |
| z_loss = log_partitions.square().sum() / valid_count |
| objective = cross_entropy + z_loss_coefficient * z_loss |
| if compute_grad: |
| ctx.save_for_backward(grad_states, grad_weight) |
| ctx.weight_dtype = weight.dtype |
| ctx.mark_non_differentiable(cross_entropy, z_loss) |
| return objective, cross_entropy, z_loss |
|
|
| @staticmethod |
| def backward( |
| ctx, grad_objective: torch.Tensor, *_: torch.Tensor |
| ) -> tuple[torch.Tensor, torch.Tensor, None, None, None, None]: |
| grad_states, grad_weight = ctx.saved_tensors |
| scale = grad_objective.to(torch.float32) |
| grad_states = (grad_states * scale).to(grad_states.dtype) |
| grad_weight = (grad_weight * scale).to(ctx.weight_dtype) |
| return grad_states, grad_weight, None, None, None, None |
|
|
|
|
| class _WeightedLinearCrossEntropyFunction(torch.autograd.Function): |
| """``Σ_i w_i (CE_i + z log Z_i²)`` over chunked ``states @ weight.T``, gradients precomputed. |
| |
| The row kernel runs with unit scale and each row's gradient is then |
| multiplied by its weight, so per-token weights cost one in-place multiply. |
| """ |
|
|
| @staticmethod |
| def forward( |
| ctx, |
| states: torch.Tensor, |
| weight: torch.Tensor, |
| targets: torch.Tensor, |
| token_weights: torch.Tensor, |
| z_loss_coefficient: float, |
| compute_grad: bool, |
| excluded_token_id: int | None, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| states = states.contiguous() |
| targets = targets.contiguous() |
| token_weights = token_weights.float().contiguous() |
| rows, vocab_size = states.shape[0], weight.shape[0] |
| unit = torch.ones((), dtype=torch.float32, device=states.device) |
| losses = states.new_empty(rows, dtype=torch.float32) |
| log_partitions = states.new_empty(rows, dtype=torch.float32) |
| grad_states = torch.empty_like(states) if compute_grad else None |
| grad_weight = ( |
| torch.zeros(weight.shape, dtype=torch.float32, device=weight.device) |
| if compute_grad |
| else None |
| ) |
| chunk_rows = max(1, LOGITS_CHUNK_BYTES // (vocab_size * states.element_size())) |
| for start in range(0, rows, chunk_rows): |
| end = min(start + chunk_rows, rows) |
| chunk = states[start:end] |
| logits = chunk @ weight.t() |
| if excluded_token_id is not None: |
| logits[:, excluded_token_id] = float("-inf") |
| _cross_entropy_rows_kernel[(end - start,)]( |
| logits, |
| targets[start:end], |
| losses[start:end], |
| log_partitions[start:end], |
| unit, |
| logits.stride(0), |
| vocab_size, |
| z_loss_coefficient, |
| -100, |
| COMPUTE_GRAD=compute_grad, |
| BLOCK=VOCAB_BLOCK_SIZE, |
| num_warps=8, |
| ) |
| if compute_grad: |
| logits.mul_(token_weights[start:end, None].to(logits.dtype)) |
| torch.mm(logits, weight, out=grad_states[start:end]) |
| if chunk.dtype == torch.float32: |
| grad_weight.addmm_(logits.t(), chunk) |
| else: |
| torch.addmm( |
| grad_weight, logits.t(), chunk, out_dtype=torch.float32, out=grad_weight |
| ) |
| objective = (token_weights * (losses + z_loss_coefficient * log_partitions.square())).sum() |
| if compute_grad: |
| ctx.save_for_backward(grad_states, grad_weight) |
| ctx.weight_dtype = weight.dtype |
| ctx.mark_non_differentiable(losses, log_partitions) |
| return objective, losses, log_partitions |
|
|
| @staticmethod |
| def backward( |
| ctx, grad_objective: torch.Tensor, *_: torch.Tensor |
| ) -> tuple[torch.Tensor, torch.Tensor, None, None, None, None, None]: |
| grad_states, grad_weight = ctx.saved_tensors |
| scale = grad_objective.to(torch.float32) |
| grad_states = (grad_states * scale).to(grad_states.dtype) |
| grad_weight = (grad_weight * scale).to(ctx.weight_dtype) |
| return grad_states, grad_weight, None, None, None, None, None |
|
|
|
|
| def weighted_linear_cross_entropy( |
| states: torch.Tensor, |
| weight: torch.Tensor, |
| targets: torch.Tensor, |
| token_weights: torch.Tensor, |
| *, |
| z_loss_coefficient: float = 0.0, |
| excluded_token_id: int | None = None, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| """``(Σ_i w_i (CE_i + z log Z_i²), CE [rows], log Z [rows])`` without full logits. |
| |
| Every row is scored; only the weighted sum carries gradients. |
| ``excluded_token_id`` gets probability zero (its logit is -inf). |
| """ |
|
|
| if states.ndim != 2 or targets.shape != states.shape[:1] or token_weights.shape != targets.shape: |
| raise ValueError("states must be [tokens, hidden], targets and token_weights [tokens]") |
| if not states.is_cuda: |
| logits = F.linear(states, weight).float() |
| if excluded_token_id is not None: |
| logits[:, excluded_token_id] = float("-inf") |
| log_partitions = torch.logsumexp(logits, dim=-1) |
| losses = log_partitions - logits.gather(1, targets[:, None]).squeeze(1) |
| weights = token_weights.float() |
| objective = (weights * (losses + z_loss_coefficient * log_partitions.square())).sum() |
| return objective, losses.detach(), log_partitions.detach() |
| compute_grad = torch.is_grad_enabled() and (states.requires_grad or weight.requires_grad) |
| return _WeightedLinearCrossEntropyFunction.apply( |
| states, weight, targets, token_weights, z_loss_coefficient, compute_grad, excluded_token_id |
| ) |
|
|
|
|
| def linear_cross_entropy( |
| states: torch.Tensor, |
| weight: torch.Tensor, |
| targets: torch.Tensor, |
| *, |
| z_loss_coefficient: float = 0.0, |
| ignore_index: int = -100, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| """Mean ``cross_entropy(states @ weight.T, targets)`` without full logits. |
| |
| Returns ``(objective, cross_entropy, z_loss)``. Only ``objective = |
| cross_entropy + z_loss_coefficient * z_loss`` carries gradients; ``z_loss`` |
| is the mean squared log-partition over non-ignored targets. |
| """ |
|
|
| if states.ndim != 2 or targets.shape != states.shape[:1]: |
| raise ValueError("states must be [tokens, hidden] and targets [tokens]") |
| if not states.is_cuda: |
| return linear_cross_entropy_reference( |
| states, |
| weight, |
| targets, |
| z_loss_coefficient=z_loss_coefficient, |
| ignore_index=ignore_index, |
| ) |
| compute_grad = torch.is_grad_enabled() and (states.requires_grad or weight.requires_grad) |
| return _LinearCrossEntropyFunction.apply( |
| states, weight, targets, z_loss_coefficient, ignore_index, compute_grad |
| ) |
|
|
|
|
| @triton.jit |
| def _token_logprob_grad_rows_kernel( |
| logits, |
| targets, |
| log_partitions, |
| row_grads, |
| row_stride, |
| vocab_size, |
| ignore_index, |
| BLOCK: tl.constexpr, |
| ): |
| """Overwrite each row with ``g * (onehot(target) - softmax)``, the gradient of ``g * log p``.""" |
|
|
| row = tl.program_id(0).to(tl.int64) |
| row_logits = logits + row * row_stride |
| target = tl.load(targets + row) |
| columns = tl.arange(0, BLOCK) |
| if target == ignore_index: |
| for start in range(0, vocab_size, BLOCK): |
| offsets = start + columns |
| zeros = tl.zeros((BLOCK,), dtype=logits.dtype.element_ty) |
| tl.store(row_logits + offsets, zeros, mask=offsets < vocab_size) |
| else: |
| scale = tl.load(row_grads + row).to(tl.float32) |
| log_partition = tl.load(log_partitions + row) |
| for start in range(0, vocab_size, BLOCK): |
| offsets = start + columns |
| in_vocab = offsets < vocab_size |
| values = tl.load(row_logits + offsets, mask=in_vocab, other=0.0).to(tl.float32) |
| grad = -scale * tl.exp(values - log_partition) |
| grad = tl.where(offsets == target, grad + scale, grad) |
| tl.store(row_logits + offsets, grad.to(logits.dtype.element_ty), mask=in_vocab) |
|
|
|
|
| class _LinearTokenLogprobFunction(torch.autograd.Function): |
| """Per-row ``log_softmax(states @ weight.T)[target]`` without keeping logits. |
| |
| Preference losses weight every sequence differently, so the gradient |
| cannot be precomputed in forward like the mean LM loss. Forward keeps one |
| log-partition per row; backward recomputes each logits chunk and turns it |
| into the scaled softmax gradient in place. |
| """ |
|
|
| @staticmethod |
| def forward( |
| ctx, |
| states: torch.Tensor, |
| weight: torch.Tensor, |
| targets: torch.Tensor, |
| ignore_index: int, |
| ) -> torch.Tensor: |
| states = states.contiguous() |
| targets = targets.contiguous() |
| rows, vocab_size = states.shape[0], weight.shape[0] |
| negative_logprobs = states.new_empty(rows, dtype=torch.float32) |
| log_partitions = states.new_empty(rows, dtype=torch.float32) |
| unused_count = states.new_ones((), dtype=torch.float32) |
| chunk_rows = max(1, LOGITS_CHUNK_BYTES // (vocab_size * states.element_size())) |
| for start in range(0, rows, chunk_rows): |
| end = min(start + chunk_rows, rows) |
| logits = states[start:end] @ weight.t() |
| _cross_entropy_rows_kernel[(end - start,)]( |
| logits, |
| targets[start:end], |
| negative_logprobs[start:end], |
| log_partitions[start:end], |
| unused_count, |
| logits.stride(0), |
| vocab_size, |
| 0.0, |
| ignore_index, |
| COMPUTE_GRAD=False, |
| BLOCK=VOCAB_BLOCK_SIZE, |
| num_warps=8, |
| ) |
| ctx.save_for_backward(states, weight, targets, log_partitions) |
| ctx.ignore_index = ignore_index |
| ctx.chunk_rows = chunk_rows |
| return negative_logprobs.neg_() |
|
|
| @staticmethod |
| def backward( |
| ctx, grad_logprobs: torch.Tensor |
| ) -> tuple[torch.Tensor | None, torch.Tensor | None, None, None]: |
| states, weight, targets, log_partitions = ctx.saved_tensors |
| row_grads = grad_logprobs.to(torch.float32).contiguous() |
| rows, vocab_size = states.shape[0], weight.shape[0] |
| grad_states = torch.empty_like(states) if ctx.needs_input_grad[0] else None |
| grad_weight = ( |
| torch.zeros(weight.shape, dtype=torch.float32, device=weight.device) |
| if ctx.needs_input_grad[1] |
| else None |
| ) |
| for start in range(0, rows, ctx.chunk_rows): |
| end = min(start + ctx.chunk_rows, rows) |
| chunk = states[start:end] |
| logits = chunk @ weight.t() |
| _token_logprob_grad_rows_kernel[(end - start,)]( |
| logits, |
| targets[start:end], |
| log_partitions[start:end], |
| row_grads[start:end], |
| logits.stride(0), |
| vocab_size, |
| ctx.ignore_index, |
| BLOCK=VOCAB_BLOCK_SIZE, |
| num_warps=8, |
| ) |
| if grad_states is not None: |
| torch.mm(logits, weight, out=grad_states[start:end]) |
| if grad_weight is not None: |
| if chunk.dtype == torch.float32: |
| grad_weight.addmm_(logits.t(), chunk) |
| else: |
| torch.addmm( |
| grad_weight, logits.t(), chunk, out_dtype=torch.float32, out=grad_weight |
| ) |
| return ( |
| grad_states, |
| None if grad_weight is None else grad_weight.to(weight.dtype), |
| None, |
| None, |
| ) |
|
|
|
|
| def linear_token_logprobs( |
| states: torch.Tensor, |
| weight: torch.Tensor, |
| targets: torch.Tensor, |
| *, |
| ignore_index: int = -100, |
| ) -> torch.Tensor: |
| """fp32 ``log_softmax(states @ weight.T)[row, target]`` per row without full logits. |
| |
| Rows whose target is ``ignore_index`` return 0 and receive no gradient. |
| """ |
|
|
| if states.ndim != 2 or targets.shape != states.shape[:1]: |
| raise ValueError("states must be [tokens, hidden] and targets [tokens]") |
| if not states.is_cuda: |
| return linear_token_logprobs_reference( |
| states, weight, targets, ignore_index=ignore_index |
| ) |
| return _LinearTokenLogprobFunction.apply(states, weight, targets, ignore_index) |
|
|
