| """Memory-bounded tied language-head cross entropy for Multiscreen training.""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
|
|
|
|
| class _ChunkedLinearCrossEntropy(torch.autograd.Function): |
| @staticmethod |
| def forward( |
| ctx, |
| hidden: torch.Tensor, |
| weight: torch.Tensor, |
| targets: torch.Tensor, |
| chunk_tokens: int, |
| ) -> torch.Tensor: |
| hidden_flat = hidden.reshape(-1, hidden.size(-1)) |
| targets_flat = targets.reshape(-1) |
| valid_count = (targets_flat != -100).sum().clamp_min(1) |
| loss_sum = torch.zeros((), device=hidden.device, dtype=torch.float32) |
| for start in range(0, hidden_flat.size(0), chunk_tokens): |
| end = min(start + chunk_tokens, hidden_flat.size(0)) |
| chunk_targets = targets_flat[start:end] |
| logits = hidden_flat[start:end].float() @ weight.float().t() |
| loss_sum += torch.nn.functional.cross_entropy( |
| logits, |
| chunk_targets, |
| ignore_index=-100, |
| reduction="sum", |
| ) |
| ctx.save_for_backward(hidden, weight, targets, valid_count) |
| ctx.chunk_tokens = chunk_tokens |
| return loss_sum / valid_count |
|
|
| @staticmethod |
| def backward(ctx, grad_output: torch.Tensor): |
| hidden, weight, targets, valid_count = ctx.saved_tensors |
| hidden_flat = hidden.reshape(-1, hidden.size(-1)) |
| targets_flat = targets.reshape(-1) |
| grad_hidden = torch.empty_like(hidden_flat) |
| grad_weight = torch.zeros_like(weight, dtype=torch.float32) |
| scale = grad_output.float() / valid_count.float() |
|
|
| for start in range(0, hidden_flat.size(0), ctx.chunk_tokens): |
| end = min(start + ctx.chunk_tokens, hidden_flat.size(0)) |
| hidden_chunk = hidden_flat[start:end].float() |
| chunk_targets = targets_flat[start:end] |
| logits = hidden_chunk @ weight.float().t() |
| probabilities = logits.softmax(dim=-1) |
| valid = chunk_targets != -100 |
| safe_targets = chunk_targets.masked_fill(~valid, 0) |
| probabilities[ |
| torch.arange(end - start, device=hidden.device), safe_targets |
| ] -= valid.to(probabilities.dtype) |
| probabilities *= valid[:, None].to(probabilities.dtype) |
| probabilities *= scale |
| grad_hidden[start:end] = ( |
| probabilities @ weight.float() |
| ).to(hidden.dtype) |
| grad_weight.add_(probabilities.t() @ hidden_chunk) |
|
|
| return ( |
| grad_hidden.view_as(hidden), |
| grad_weight.to(weight.dtype), |
| None, |
| None, |
| ) |
|
|
|
|
| def chunked_linear_cross_entropy( |
| hidden: torch.Tensor, |
| weight: torch.Tensor, |
| targets: torch.Tensor, |
| chunk_tokens: int = 2048, |
| ) -> torch.Tensor: |
| if chunk_tokens <= 0: |
| raise ValueError("chunk_tokens must be positive") |
| return _ChunkedLinearCrossEntropy.apply(hidden, weight, targets, chunk_tokens) |
|
|
