3.5 Gradient Clipping

Gradient clipping 是训练稳定性的保险丝。它不解决 objective 本身的问题，但能阻止异常 batch、长序列反传或 recurrent/attention dynamics 产生的极端梯度把参数一步推坏。

Exploding Gradients

考虑深层复合函数

\[ h_L=f_L\circ f_{L-1}\circ\cdots\circ f_1(x). \]

反向传播包含 Jacobian 连乘：

\[ \frac{\partial L}{\partial h_0} = \frac{\partial L}{\partial h_L} \prod_{\ell=1}^{L} \frac{\partial h_\ell}{\partial h_{\ell-1}}. \]

如果 Jacobian 的谱范数长期大于 \(1\)，梯度会指数增长；如果长期小于 \(1\)，梯度会消失。

Global Norm Clipping

Definition: Global Norm Clipping

Given all gradients concatenated as \(g\), global norm clipping with threshold \(\tau\) uses \[ \tilde{g} = g\cdot\min\left(1,\frac{\tau}{\|g\|_2}\right). \]

如果 \(\|g\|_2\leq\tau\)，不做任何事；如果超过阈值，只缩放方向长度，不改变方向。

这等价于把梯度投影到 \(\ell_2\) ball：

\[ \tilde{g} = \operatorname{Proj}_{\{u:\|u\|_2\leq\tau\}}(g). \]

Proof

要解投影问题

\[ \min_u \frac12\|u-g\|_2^2 \quad \text{s.t.} \quad \|u\|_2\le\tau. \]

如果 \(\|g\|_2\le\tau\)，显然 \(u=g\) 可行且最优。

若 \(\|g\|_2>\tau\)，最优点在边界 \(\|u\|_2=\tau\)。Lagrangian:

\[ \mathcal{L}(u,\lambda) = \frac12\|u-g\|_2^2 + \frac{\lambda}{2}(\|u\|_2^2-\tau^2). \]

一阶条件：

\[ u-g+\lambda u=0 \quad\Rightarrow\quad u=\frac{1}{1+\lambda}g. \]

再用 \(\|u\|_2=\tau\) 得

\[ u=\tau\frac{g}{\|g\|_2}. \]

合并两种情况：

\[ \tilde{g} = g\min\left(1,\frac{\tau}{\|g\|_2}\right). \]

Global norm clipping 保留方向，是因为所有参数梯度共享同一个缩放系数。它改变的是 step length，而不是每个坐标的相对比例。

Clipping as a Trust-Region Step

如果 optimizer 是 plain SGD：

\[ \theta_{t+1} = \theta_t-\eta\tilde{g}. \]

global norm clipping 等价于限制 SGD update 的长度：

\[ \|\theta_{t+1}-\theta_t\|_2 = \eta\|\tilde{g}\|_2 \leq \eta\tau. \]

所以它像一个非常简单的 trust region：这一步无论 raw gradient 多大，参数最多移动 \(\eta\tau\)。

Definition: Gradient Trust Region

A gradient trust region limits the norm of the update implied by the gradient, preventing a single batch from moving parameters too far.

但对 AdamW 来说，update 不是简单的 \(-\eta g\)：

\[ \Delta\theta_i = -\eta\frac{\hat{m}_{t,i}}{\sqrt{\hat{v}_{t,i}}+\epsilon}. \]

因此 clip gradient norm 不等价于 clip Adam update norm。它限制的是进入 moment estimates 的 raw gradient，而不是最终 preconditioned update。

Pitfall: Gradient Clipping Is Not Update Clipping for Adam

With adaptive optimizers, clipping gradients protects moment estimates but does not directly bound the final parameter update norm.

Preconditioned Update Clipping

上面的 warning 很重要，因为 AdamW、Adafactor、Shampoo、K-FAC 这类 optimizer 都不是直接沿着 \(g_t\) 走。更抽象地写，很多 optimizer 的 data update 可以看成

\[ \Delta\theta_t = -\eta_t P_t g_t, \]

其中 \(P_t\) 是某种 preconditioner。SGD 的 \(P_t=I\)；AdamW 的 \(P_t\) 近似为逐坐标矩阵

\[ P_t = \operatorname{diag} \left( \frac{1}{\sqrt{\hat v_t}+\epsilon} \right), \]

但实际 update 用的是一阶动量 \(\hat m_t\)：

\[ u_t = \frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon}, \qquad \theta_{t+1} = \theta_t-\eta_t u_t-\eta_t\lambda\theta_t. \]

因此可以区分三种不同的裁剪对象：

Object	Formula	Protects	Does not directly bound
raw gradient clipping	\(g_t\leftarrow C(g_t)\)	backward signal, moment states	final Adam update
preconditioned update clipping	\(u_t\leftarrow C(u_t)\)	actual adaptive update length	raw moment contamination
parameter delta clipping	\(\Delta\theta_t\leftarrow C(\Delta\theta_t)\)	total parameter movement	optimizer statistics

raw gradient clipping 发生在 optimizer.step() 之前，optimizer 看到的是 clipped .grad；update clipping 则要在 optimizer 已经计算出 \(u_t\) 之后、参数真正写回之前做。两者的语义不同：

\[ C(P_t g_t) \ne P_t C(g_t) \]

一般成立，因为 \(P_t\) 改变了坐标尺度。一个坐标如果历史方差很小，Adam 会给它更大的有效学习率；即使 raw gradient norm 不大，preconditioned update norm 也可能很大。

Definition: Update Norm

The update norm is the norm of the parameter delta actually applied by the optimizer, usually measured before decoupled weight decay is added unless stated otherwise.

Metric View

clipping 也可以理解成“在哪个几何里限制 step”。SGD 中限制 \(\|g\|_2\) 等价于限制 Euclidean update：

\[ \|\Delta\theta\|_2 = \eta\|g\|_2. \]

若 optimizer 使用正定 preconditioner \(P\)，真正的一阶下降方向是 \(Pg\)。如果我们希望直接限制参数空间里的移动，应当约束

\[ \|Pg\|_2\leq \rho. \]

如果我们从近似二阶优化看，常见的 trust-region 子问题是

\[ \min_\Delta \quad g^\top\Delta+\frac{1}{2\eta}\Delta^\top P^{-1}\Delta \quad \text{s.t.} \quad \|\Delta\|_{P^{-1}}^2 = \Delta^\top P^{-1}\Delta \leq \rho^2. \]

它的无约束解是 \(\Delta^\star=-\eta Pg\)。这说明 adaptive optimizer 隐含选择了一个 metric；在 raw gradient norm 上裁剪，是在 Euclidean gradient space 里做保护；在 update norm 上裁剪，是在 optimizer 真正使用的参数移动上做保护。

Practical Rule

For most LLM training loops, start with raw global norm clipping because it protects optimizer states. Add update-norm logging first; only introduce update clipping when logs show the adaptive update itself is the unstable quantity.

A Concrete Adam Example

设二维梯度为

\[ g=(1,1), \qquad \sqrt{\hat v}+\epsilon=(1,0.01). \]

raw gradient norm 是

\[ \|g\|_2=\sqrt{2}. \]

但 Adam 的 preconditioned direction 是

\[ u= \left( \frac{1}{1}, \frac{1}{0.01} \right) =(1,100), \qquad \|u\|_2\approx 100. \]

如果只看 \(\|g\|_2\)，这一步并不夸张；如果看 \(\|u\|_2\)，第二个坐标会产生极大的 adaptive update。反过来，一个 raw gradient spike 也可能因为 \(\hat v\) 同时变大而在本步 update 里被抑制，但它仍然污染了之后许多 step 的 \(v_t\)。这就是为什么训练日志里最好同时看 grad_norm 与 update_norm。

Value Clipping

Value clipping 对每个坐标截断：

\[ \tilde{g}_i = \operatorname{clip}(g_i,-c,c). \]

它实现简单，但会改变梯度方向。深度学习中更常用 global norm clipping，尤其是 Transformer、RNN 和 RL。

Pitfall: Clipping Is Not a Substitute for Bad LR

If clipping triggers on almost every step, the learning rate, loss scale, initialization, or data pipeline may be wrong. Clipping should catch spikes, not define the normal update magnitude.

Clipping Biases the Gradient

Clipping 会引入 bias。若随机梯度 \(g\) 是无偏估计：

\[ \mathbb{E}[g]=\nabla L(\theta), \]

clipped gradient

\[ \tilde{g}=C(g) \]

通常不满足

\[ \mathbb{E}[\tilde{g}]=\nabla L(\theta). \]

因为 \(C(\cdot)\) 是非线性操作：

\[ \mathbb{E}[C(g)]\ne C(\mathbb{E}[g]). \]

这不是说 clipping 不该用，而是说 clipping 是一个稳定性与优化偏差之间的 trade-off。阈值太低时，训练方向会长期被扭曲；阈值太高时，又无法拦住异常 spike。

Pitfall: Clipping Threshold Is an Optimization Hyperparameter

Changing max_grad_norm changes the effective optimizer. It should be logged and tuned like learning rate, not treated as harmless boilerplate.

Expected Bias Example

一维例子最直观。设随机梯度：

\[ g= \begin{cases} 10, & p=0.1,\\ 0, & p=0.9. \end{cases} \]

则

\[ \mathbb{E}[g]=1. \]

如果用阈值 \(\tau=1\) 裁剪：

\[ C(g)=\min(g,1), \]

则

\[ \mathbb{E}[C(g)] = 0.1. \]

这说明 clipping 会显著改变梯度估计，尤其是 heavy-tailed gradient noise 下。它降低 variance，但也引入 bias。

Definition: Clipping Ratio

The clipping ratio is the fraction of optimizer steps whose pre-clip gradient norm exceeds the clipping threshold.

如果 clipping ratio 长期接近 100%，说明 clipping 已经成为常规 optimizer 行为，而不是异常保护。

Choosing the Threshold

max_grad_norm=1.0 很常见，但它不是一个理论常数。阈值 \(\tau\) 的含义依赖于：

loss reduction 是按 sample mean、token mean，还是 sum；
global batch size 与 gradient accumulation；
模型宽度、层数、初始化尺度；
optimizer、LR、warmup、loss scale；
是否包含 embedding、adapter、LoRA、norm 参数；
是否有 packed sequence、padding mask、不同数据源混合。

所以更稳的做法是把 \(\tau\) 当成一个 calibration target。一个简单流程是：

先用较保守 LR 跑 warmup 的前几百步；
记录每一步 pre-clip global norm \(n_t=\|g_t\|_2\)；
计算分位数，例如 \(p_{90},p_{95},p_{99}\)；
选择 \(\tau\) 让 clipping ratio 落在一个小范围内；
随 LR、数据混合或序列长度变化重新检查。

例如，如果目标是“只裁掉明显 spike”，可以设

\[ \tau = \operatorname{Quantile}_{0.95} \left( \{n_t\}_{t=1}^{T} \right). \]

如果目标是“非常谨慎地保护 pretrained model”，可以选择更低的分位数，但要接受更大的 bias。

Training situation	Healthy clipping ratio	Warning sign
large-scale pretraining	low, often below 5%	sustained clipping after warmup
SFT on curated data	low to moderate	spikes tied to prompt/label masks
noisy instruction mix	moderate	one source dominates top norm batches
preference/RL training	moderate to high	advantage or reward outliers dominate
debugging new loss	should be inspected manually	clipping hides a broken objective

Definition: Norm Calibration Window

A norm calibration window is a short early training interval used to estimate typical gradient-norm percentiles before fixing or revising a clipping threshold.

Why Token Normalization Matters

语言模型的 loss 常写成

\[ L = \frac{1}{N_{\text{tok}}} \sum_{b,t} m_{b,t}\ell_{b,t}. \]

如果实现里误用了 batch mean，例如每个 sequence 先平均再 batch 平均：

\[ L_{\text{seq-mean}} = \frac{1}{B} \sum_{b=1}^{B} \frac{1}{N_b} \sum_t m_{b,t}\ell_{b,t}, \]

那么短序列和长序列的权重会改变，gradient norm 的尺度也会改变。更糟的是，当一个 batch 的有效 token 数很少时，loss 和 gradient 都可能异常放大。此时调低 \(\tau\) 只是掩盖 reduction bug。

一个更可诊断的日志是同时记录

\[ (\|g_t\|_2,\; N_{\text{tok}},\; \|g_t\|_2\sqrt{N_{\text{tok}}}). \]

在独立同分布、mean reduction 的粗略近似下，batch 平均梯度的噪声尺度约随 \(1/\sqrt{N_{\text{tok}}}\) 下降。因此 \(\|g_t\|_2\sqrt{N_{\text{tok}}}\) 可以帮助区分“token 少导致 norm 大”和“样本本身真的异常”。

Pitfall: Do Not Copy Thresholds Across Loss Semantics

The same numeric threshold can mean different things under sample-mean, token-mean, sequence-mean, or summed losses.

Dynamic Thresholds

有些训练系统会用动态阈值：

\[ \tau_t = \alpha\cdot \operatorname{EMAQuantile}_{q} \left( \|g_1\|_2,\ldots,\|g_t\|_2 \right). \]

它的好处是适应不同阶段；坏处是异常阶段本身会抬高阈值，让保护失效。工程上更常见的折中是 stage-wise threshold：pretraining、SFT、DPO/RL 各自固定一组阈值，并在切换数据和 objective 时重新校准。

Practical Rule

If changing max_grad_norm improves training, check whether the same improvement comes from lowering LR or fixing loss normalization. Clipping and LR are entangled through the effective step length.

Mixed Precision and Loss Scaling

在 FP16/BF16 训练里，梯度可能 underflow 或 overflow。AMP 常用 loss scaling：

\[ \tilde{L}=sL. \]

反传后梯度为 \(s\nabla L\)，optimizer step 前再 unscale：

\[ g=\frac{\tilde{g}}{s}. \]

正确顺序是：

scale loss；
backward；
unscale gradients；
clip gradients；
optimizer step。

如果先 clip 再 unscale，阈值含义就错了。

AMP 中还要处理 overflow。若 unscale 后发现 NaN/Inf gradient，正确行为通常是 skip optimizer step 并降低 loss scale，而不是先把 NaN clip 掉。NaN 不是“大梯度”，而是无效数值。

Pitfall: Clipping Does Not Fix NaNs

Clip finite large gradients. If gradients contain NaN or Inf, skip the step and debug overflow or invalid operations.

Clipping With AdamW

AdamW 使用 gradient 更新 moment estimates：

\[ m_t=\beta_1m_{t-1}+(1-\beta_1)g_t, \qquad v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2. \]

若出现异常大梯度并进入 \(v_t\)，二阶动量会在之后很多步保持偏大，导致有效学习率变小：

\[ \eta_{\text{eff},i} \approx \frac{\eta}{\sqrt{\hat{v}_{t,i}}+\epsilon}. \]

因此 clipping 通常在 optimizer.step() 前执行，让 optimizer 看到的是 clipped gradients。对 AdamW 来说，clipping 不只是限制这一步参数更新，也是在保护 moment states 不被异常 batch 污染。

Weight Decay Order

AdamW 的 decoupled weight decay 是：

\[ \theta \leftarrow \theta-\eta\lambda\theta \]

再加 adaptive gradient update。Gradient clipping 通常只作用于 data gradient，不应该把 decoupled weight decay 当成 gradient 一起 clip。若用的是 coupled L2 regularization，正则项已经加进 gradient：

\[ g_{\text{coupled}} = \nabla L(\theta)+\lambda\theta, \]

这时 clipping 会同时裁剪 loss gradient 和 L2 gradient，含义不同。

Regularization	Enters `.grad`?	Clipped?
coupled L2	yes	yes
AdamW decoupled weight decay	no	no

这也是 AdamW 推荐 decoupled decay 的一个工程好处：clipping threshold 更专注于 batch gradient spike。

Gradient Accumulation Order

有 gradient accumulation 时，应该先累积所有 micro-batch 的梯度，再在 optimizer step 前 clip：

optimizer.zero_grad(set_to_none=True)

for i, batch in enumerate(loader):
    loss = model(batch).loss / grad_accum
    loss.backward()

    if (i + 1) % grad_accum == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

如果每个 micro-batch 后都 clip，就不再等价于大 batch 的平均梯度：

\[ C\left(\frac1K\sum_k g_k\right) \ne \frac1K\sum_k C(g_k). \]

Pitfall: Clip After Accumulation

For gradient accumulation, clipping each micro-batch changes the training objective. Clip the accumulated gradient unless there is a deliberate reason not to.

如果 loss reduction 是 mean，还要确认每个 micro-batch 的 loss 是否按 accumulation 正确归一化。对 token-level LM，最好按有效 token 数归一：

\[ g = \frac{\sum_{k=1}^{K}\sum_{t}m_{k,t}\nabla \ell_{k,t}} {\sum_{k=1}^{K}\sum_t m_{k,t}}. \]

否则最后一个短 batch 或 padding 多的 batch 会改变有效梯度尺度，从而改变 clipping 触发频率。

Adaptive Gradient Clipping

global norm clipping 用同一个阈值 \(\tau\) 对所有参数组裁剪。Adaptive Gradient Clipping (AGC) 改用参数相对尺度：

\[ \|g_\ell\|_2 \leq \lambda\|\theta_\ell\|_2 \]

对每个 layer 或 parameter block \(\ell\)，若超过阈值则：

\[ \tilde{g}_\ell = g_\ell \cdot \min \left( 1, \frac{\lambda\|\theta_\ell\|_2}{\|g_\ell\|_2+\epsilon} \right). \]

直觉是：同样的 gradient norm，对大权重层可能正常，对小权重层可能过猛。AGC 在视觉模型和一些大规模训练中常用于更细粒度稳定更新。

Definition: Adaptive Gradient Clipping

Adaptive gradient clipping scales each parameter block’s gradient according to the ratio between gradient norm and parameter norm.

AGC 的坑也很直接：bias、LayerNorm/RMSNorm、embedding 等小参数或 scale 参数可能不适合按 \(\|\theta\|\) 裁剪，实践中常排除这些参数组。

Parameter Groups and Sparse Gradients

真实训练脚本里，model.parameters() 往往不是一个同质集合。LLM 里常见参数组包括：

Group	Typical examples	Clipping caveat
dense backbone weights	attention/MLP linear layers	global norm works well
embeddings	token embedding, position embedding	sparse or highly skewed updates
norm parameters	LayerNorm/RMSNorm scale	small parameter norm makes AGC fragile
LM head	tied or untied output projection	may dominate vocab-heavy losses
adapters	LoRA A/B matrices, prompt vectors	scale differs from frozen backbone
biases	bias terms if present	often excluded from weight decay, but not necessarily clipping

全局裁剪把所有参数拼成一个向量，因此缩放系数只有一个：

\[ s = \min\left(1,\frac{\tau}{\sqrt{\sum_j\|g_j\|_2^2}}\right), \qquad \tilde g_j=sg_j. \]

这保留了整体方向。但如果某一个参数组长期贡献了大部分 norm，其他组也会被一起缩小。例如词表 embedding 在少数 token 上出现大梯度时，attention/MLP 的 dense gradient 也会被同一个 \(s\) 压小。

参数组裁剪则是

\[ s_j = \min\left(1,\frac{\tau_j}{\|g_j\|_2}\right), \qquad \tilde g_j=s_jg_j. \]

它更可控，但会改变整体方向：

\[ \tilde g \not\parallel g \quad \text{when some }s_j\ne s_k. \]

因此 group clipping 更像“按模块限速”，不是单纯的 global projection。它适合有明确工程理由的场景，例如只训练 LoRA、embedding 扩词表、或者某些 auxiliary head 梯度尺度与主干差异很大。

Pitfall: Per-Group Clipping Changes the Direction

Per-group clipping can prevent one module from dominating the norm, but it no longer preserves the full-model gradient direction.

LoRA and Adapter-Only Training

LoRA 微调时，大部分 pretrained weights 冻结，只有低秩矩阵 \(A,B\) 可训练：

\[ W_{\text{eff}} = W_0+\frac{\alpha}{r}BA. \]

如果只把 trainable parameters 传给 optimizer，那么 global norm 只覆盖 LoRA 参数，这是合理的。若日志里误把 frozen parameters 也遍历进去，虽然 .grad is None 通常不会影响 norm，但会让代码审计变得混乱。更重要的是，LoRA 参数的 norm 和 full fine-tuning 的 backbone norm 不在同一尺度，full fine-tuning 的阈值不能直接照搬。

对 adapter 训练，更应该记录：

adapter grad norm；
adapter update norm；
backbone 是否有非空 gradient；
LoRA scaling \(\alpha/r\)；
每个 target module 的 norm top-k。

这样才能判断是某个 attention projection 的 adapter 在 spike，还是整个 objective 都不稳定。

Sparse Embedding Gradients

nn.Embedding(..., sparse=True) 的 .grad 可能是 sparse COO tensor，只包含本 batch 出现过的 token。norm 计算不能假设所有 gradient 都是 dense：

total_sq = torch.zeros([], device=device)
for p in model.parameters():
    if p.grad is None:
        continue
    grad = p.grad
    if grad.is_sparse:
        grad = grad.coalesce()
        values = grad.values()
        total_sq += values.float().pow(2).sum()
    else:
        total_sq += grad.float().pow(2).sum()
total_norm = total_sq.sqrt()

如果 sparse gradient 有重复 index，必须先 coalesce()，否则同一个 token 的多次贡献会分散在不同 entry 里。对 embedding 来说，还要区分两个 norm：

\[ \|g_{\text{sparse-values}}\|_2 \quad\text{versus}\quad \|g_{\text{dense-table}}\|_2. \]

数学上它们表示同一个稀疏向量的 norm，但实现上前者只遍历非零 entry，后者需要 materialize 整个表。大词表下不能为了日志把 sparse gradient densify。

Pitfall: Do Not Densify Huge Sparse Gradients for Logging

For sparse embeddings, compute norms from coalesced sparse values. Densifying a vocabulary-sized gradient can turn a logging line into an OOM.

Nonfinite Gradients

PyTorch 的 clip_grad_norm_ 有 error_if_nonfinite 参数。训练系统里建议让 nonfinite 尽早暴露：

total_norm = torch.nn.utils.clip_grad_norm_(
    model.parameters(),
    max_norm=1.0,
    error_if_nonfinite=True,
)

在 AMP 训练中，如果 unscale 后出现 Inf/NaN，GradScaler 通常会 skip step；在 BF16 或自定义训练循环中，则要显式检查。不要把 nonfinite 当成“大于阈值的数”来处理，因为

\[ \min\left(1,\frac{\tau}{\text{NaN}}\right) = \text{NaN}. \]

一旦 NaN 进入参数或 optimizer state，后续 checkpoint 也可能被污染。

Gradient Clipping in Sequence Models

RNN 中 exploding gradients 来自 recurrent Jacobian 的长时间连乘；Transformer 中则常来自长序列 attention、activation outliers、loss spikes 或 post-training 中偏好/奖励信号的高方差。

LLM SFT/RL 阶段常见做法：

Stage	Clipping threshold	Reason
pretraining	1.0 or tuned	rare but costly spikes
SFT	1.0	protect pretrained weights
preference training	0.5-1.0	noisy pairwise gradients
RL	often strict	reward variance and KL control

RL and Advantage Outliers

在 policy gradient / RLHF 中，梯度常含有 advantage 或 reward weight：

\[ g \approx -A_t\nabla_\theta\log\pi_\theta(a_t\mid s_t). \]

如果 \(A_t\) 有 heavy tail，一个 batch 就可能产生很大的 gradient norm。这里 clipping 不只是数值稳定，也在限制单个高 reward / 低 reward 样本对 policy 的影响。KL penalty、reward normalization、advantage whitening 和 gradient clipping 通常要一起看。

PyTorch Pattern

import torch

loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
optimizer.zero_grad(set_to_none=True)

对于 AMP：

scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)

如果要记录 clipping ratio：

total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
was_clipped = float(total_norm > 1.0)
logger.log({"grad_norm": float(total_norm), "was_clipped": was_clipped})

clip_grad_norm_ 返回的是 pre-clip norm，所以适合直接用于日志。

What to Log

只打开 clipping 而不记录 norm 是不够的。至少记录：

pre-clip global grad norm；
clipping ratio；
loss spike 发生的 step 和 batch；
optimizer LR；
AMP loss scale。

这些指标能区分三种情况：偶发异常、LR 过大、或者模型/数据已经系统性不稳定。

还可以记录：

Metric	Meaning
clipping ratio	clipping 是否成为常态
norm percentile	spike 是否来自长尾
per-layer norm	哪些层触发异常
update norm	Adam precondition 后实际移动多大
token count	norm spike 是否来自长序列或 padding mask

一个更工程化的日志 schema 可以是：

{
  "step": 18420,
  "stage": "sft",
  "lr": 0.00002,
  "loss": 1.87,
  "token_count": 98304,
  "grad_norm_pre": 3.42,
  "grad_norm_post": 1.0,
  "max_grad_norm": 1.0,
  "was_clipped": true,
  "update_norm": 0.018,
  "loss_scale": 65536,
  "data_source": "math_mix",
  "top_grad_layers": [
    ["model.embed_tokens.weight", 1.21],
    ["model.layers.17.mlp.down_proj.weight", 0.74]
  ]
}

这里 grad_norm_post 不一定要重新遍历所有参数计算；若使用 global norm clipping 且 pre-clip norm 为 \(n\)，post norm 可以写成

\[ n_{\text{post}} = n\cdot \min\left(1,\frac{\tau}{n+\epsilon}\right). \]

但在分布式、sparse gradient 或 per-group clipping 下，最好用框架返回的 global norm 和 scale 明确记录，避免日志语义漂移。

Distributed Global Norm

Data parallel 训练中，每张卡有一份梯度。若 DDP 已经 all-reduce 完成，那么每张卡上的 .grad 是同步后的全局平均梯度，可以直接 clip。若使用 ZeRO/FSDP，梯度可能是 shard，需要框架提供的 global norm：

\[ \|g\|_2 = \sqrt{\sum_{r=1}^{R}\|g^{(r)}\|_2^2}. \]

只在每个 rank 的 shard 上各自 clip，会得到错误的全局阈值。正确做法是先 all-reduce norm，再统一缩放。

Definition: Global Gradient Norm

The global gradient norm is the \(\ell_2\) norm of all trainable parameter gradients as if they were concatenated into one vector, even when those gradients are sharded across devices.

概念伪代码：

local_sq = torch.zeros([], device=device)
for p in sharded_parameters:
    if p.grad is not None:
        local_sq += p.grad.float().pow(2).sum()

global_sq = all_reduce_sum(local_sq)
global_norm = torch.sqrt(global_sq)
scale = min(1.0, max_norm / (global_norm + 1e-6))

for p in sharded_parameters:
    if p.grad is not None:
        p.grad.mul_(scale)

关键点是：先聚合 norm，再统一缩放。不能每张卡根据自己的 shard norm 各自决定 scale。

Per-Sample Clipping and DP-SGD

Differentially private SGD 使用 per-sample gradient clipping：

\[ \tilde{g}_i = g_i\min\left(1,\frac{C}{\|g_i\|_2}\right), \]

然后平均并加噪声：

\[ \bar{g} = \frac{1}{B} \left( \sum_{i=1}^{B}\tilde{g}_i +\mathcal{N}(0,\sigma^2C^2I) \right). \]

它和普通 global norm clipping 不同：

Method	Clips	Purpose
global norm clipping	batch gradient	training stability
per-sample clipping	each example gradient	bound individual influence
DP-SGD	per-sample gradient + noise	privacy guarantee

LLM 训练中常用的是 batch/global clipping；DP-SGD 的 per-sample gradient 成本更高，因为需要每个样本的梯度 norm。

Clipping Smoke Tests

gradient clipping 很容易“看起来有写，实际语义错”。下面这些测试不需要完整训练，只用小模型和固定随机种子就能抓出大部分问题。

Test 1: Return Value Is Pre-Clip Norm

PyTorch 的 clip_grad_norm_ 返回裁剪前的 norm。一个最小测试是手动设置梯度：

import torch

p = torch.nn.Parameter(torch.zeros(2))
p.grad = torch.tensor([3.0, 4.0])

norm = torch.nn.utils.clip_grad_norm_([p], max_norm=1.0)

assert torch.allclose(norm, torch.tensor(5.0))
assert torch.allclose(p.grad.norm(), torch.tensor(1.0), atol=1e-6)
assert torch.allclose(p.grad, torch.tensor([0.6, 0.8]), atol=1e-6)

这个测试同时确认三件事：返回值是 pre-clip norm，post norm 被限制到阈值，方向保持不变。

Test 2: Clip After Accumulation

构造两个 micro-batch 梯度：

\[ g_1=(10,0), \qquad g_2=(-10,1). \]

如果先平均再裁剪：

\[ \bar g = \frac{g_1+g_2}{2} = (0,0.5), \]

阈值 \(\tau=1\) 时不裁剪。若每个 micro-batch 先裁剪：

\[ \frac{C(g_1)+C(g_2)}{2} \approx \frac{(1,0)+(-0.995,0.0995)}{2} =(0.0025,0.04975). \]

方向和尺度都变了。测试里可以用两个 synthetic loss 分别 backward，比较“每步 clip”和“累积后 clip”的参数更新，确保训练循环采用你想要的语义。

Test 3: Sharded Norm Equals Full Norm

对 FSDP/ZeRO 类实现，拿一个已知 dense gradient 向量分成两个 shard：

\[ g^{(1)}=(3,4), \qquad g^{(2)}=(0,12). \]

局部 norm 分别是 \(5\) 和 \(12\)，全局 norm 应为

\[ \sqrt{5^2+12^2}=13. \]

如果每个 rank 各自用 \(\tau=1\) 裁剪，两个 shard 的 post norm 都变成 \(1\)，拼回去的全局 norm 是 \(\sqrt{2}\)，不是 \(1\)。正确做法是用同一个 scale：

\[ s=\frac{1}{13}. \]

测试应断言所有 rank 使用相同 scale，并且拼接后的全局 post norm 不超过 \(\tau\)。

Test 4: Nonfinite Fails Closed

给任意一个参数写入 NaN gradient：

p = torch.nn.Parameter(torch.zeros(1))
p.grad = torch.tensor([float("nan")])

try:
    torch.nn.utils.clip_grad_norm_([p], 1.0, error_if_nonfinite=True)
except RuntimeError:
    pass
else:
    raise AssertionError("nonfinite gradient should fail before optimizer.step")

这类测试比训练过程中才发现 checkpoint 全 NaN 要便宜得多。

Debugging With Norms

如果训练发散，grad norm 日志往往比 loss 更早报警：

Pattern	Likely cause
norm spikes on rare batches	bad samples, long sequences, label bugs
norm grows slowly over many steps	LR too high, instability accumulating
norm always clipped	threshold too low or LR/loss scale wrong
norm becomes NaN	overflow, invalid loss, bad input
norm exactly zero	detached graph, frozen params, empty loss mask

一个实用习惯是同时记录 pre-clip norm 和 post-clip norm。PyTorch 的 clip_grad_norm_ 返回的是 clipping 前的 total norm，这正适合日志。

Minimal Debug Protocol

遇到 gradient spike 时，按这个顺序排查：

打印 pre-clip norm、loss、有效 token 数；
检查 labels 是否有未 mask 的 padding 或 prompt；
检查 AMP 是否先 unscale 再 clip；
检查 spike 是否集中在某些样本长度、数据源或任务类型；
打印 per-layer grad norm，定位异常层；
暂时降低 LR，观察 clipping ratio 是否下降；
单 batch overfit，确认不是计算图或 loss reduction bug；
分布式训练中确认 global norm 是跨 shard/rank 计算。

Gradient clipping 最有用的地方，是把“偶发坏 batch”从“训练直接炸掉”降级成“日志里一个可调查的异常点”。如果它每天都在救火，那真正的问题还在 learning rate、数据、loss mask 或数值精度里。