| |
| |
| |
| |
| |
| |
|
|
| import torch |
|
|
|
|
| class L2Wrap(torch.autograd.Function): |
| r""" |
| This class of penalty prevents the model from becoming overconfident, |
| thereby mitigating precision loss in BF16. |
| |
| This version is memory-optimized by not storing the full logits tensor. |
| """ |
| @staticmethod |
| def forward( |
| ctx, |
| loss: torch.Tensor, |
| logits: torch.Tensor, |
| l2_penalty_factor: float = 1e-4, |
| ) -> torch.Tensor: |
| """ |
| Args: |
| loss (torch.Tensor): |
| The already-reduced (scalar) loss to wrap. |
| logits (torch.Tensor): |
| The logits of shape `[B, T, V]`. |
| l2_penalty_factor (float, Optional): |
| The strength of the L2 penalty on the max logit. Default: 1e-4. |
| """ |
| maxx, ids = torch.max(logits, dim=-1, keepdim=True) |
| ctx.logits_shape = logits.shape |
| factor = l2_penalty_factor / (logits.shape[0] * logits.shape[1]) |
| maxx = maxx * factor |
| ctx.save_for_backward(maxx, ids) |
| return loss |
|
|
| @staticmethod |
| def backward(ctx, grad_output: torch.Tensor): |
| maxx, ids = ctx.saved_tensors |
| glogits = torch.zeros(ctx.logits_shape, device=grad_output.device, dtype=grad_output.dtype) |
| |
| |
| glogits.scatter_(-1, ids, maxx * grad_output) |
| return grad_output, glogits, None |
|
|
|
|
| l2_warp = L2Wrap.apply |
|
|
