File size: 4,292 Bytes
a311e94 | 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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | """Stability algorithms: clipping, accumulation, z-loss, noise, checkpointing."""
from __future__ import annotations
import math
from typing import Optional
import torch
from .base import AlgorithmContext, BaseAlgorithm, compose_loss
class GradientClip(BaseAlgorithm):
"""Clip gradients by global norm (and optionally value)."""
name = "gradient_clip"
category = "stability"
description = "Clips the global gradient norm before the optimiser step."
def __init__(self, max_norm: float = 1.0, norm_type: float = 2.0, clip_values: float = 0.0):
super().__init__(max_norm=max_norm, norm_type=norm_type, clip_values=clip_values)
self.max_norm = max_norm
self.norm_type = norm_type
self.clip_values = clip_values
def on_before_optimizer_step(self, ctx):
params = [p for p in ctx.model.parameters() if p.grad is not None]
if not params:
return
torch.nn.utils.clip_grad_norm_(params, self.max_norm, norm_type=self.norm_type)
if self.clip_values > 0:
torch.nn.utils.clip_grad_value_(params, self.clip_values)
class GradientAccumulation(BaseAlgorithm):
"""Accumulate gradients over ``steps`` micro-batches before stepping.
``train_step`` performs the scaled backward pass; the engine consults
:meth:`should_step` to decide whether to run ``optimizer.step()``.
"""
name = "gradient_accumulation"
category = "stability"
description = "Simulate larger batch sizes via gradient accumulation."
def __init__(self, steps: int = 4):
super().__init__(steps=steps)
self.steps = max(1, steps)
def train_step(self, ctx, batch):
batch = self.on_step_begin(ctx, batch)
outputs = ctx.model(**batch)
loss = outputs["loss"] if isinstance(outputs, dict) else outputs.loss
loss = self.compute_loss(ctx, batch, outputs, loss)
loss = compose_loss(ctx, batch, outputs, loss, skip=self)
if loss is not None:
(loss / self.steps).backward()
ctx.set("accum_step", (ctx.get("accum_step", 0) + 1) % self.steps)
return loss
def should_step(self, ctx) -> bool:
return ctx.get("accum_step", 0) == 0
class ZLoss(BaseAlgorithm):
"""Auxiliary z-loss to keep the logits well-behaved (PaLM style)."""
name = "z_loss"
category = "stability"
description = "Penalises large log-sum-exp values to stabilise logits."
def __init__(self, weight: float = 1e-4):
super().__init__(weight=weight)
self.weight = weight
def compute_loss(self, ctx, batch, outputs, loss):
if loss is None:
return None
logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits
z = logits.logsumexp(dim=-1).pow(2).mean()
return loss + self.weight * z
class GradientNoise(BaseAlgorithm):
"""Add calibrated Gaussian noise to gradients.
Useful both as a differential-privacy mechanism and as an exploration
regulariser for small models that overfit quickly.
"""
name = "gradient_noise"
category = "stability"
description = "Adds Gaussian noise to gradients (DP-style regularisation)."
def __init__(self, sigma: float = 1e-3):
super().__init__(sigma=sigma)
self.sigma = sigma
def on_before_optimizer_step(self, ctx):
if self.sigma <= 0:
return
for p in ctx.model.parameters():
if p.grad is not None:
p.grad.add_(torch.randn_like(p.grad) * self.sigma)
class ActivationCheckpointing(BaseAlgorithm):
"""Enable gradient checkpointing on compatible models."""
name = "activation_checkpointing"
category = "memory"
description = "Trades compute for memory by recomputing activations."
def on_train_begin(self, ctx):
model = ctx.model
if hasattr(model, "gradient_checkpointing_enable"):
try:
model.gradient_checkpointing_enable()
except Exception:
pass
def on_train_end(self, ctx):
model = ctx.model
if hasattr(model, "gradient_checkpointing_disable"):
try:
model.gradient_checkpointing_disable()
except Exception:
pass
|