Skip to content

【训练营】分布式梯度范数与梯度裁剪实现 #222

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zjhmax777 wants to merge 7 commits into
base: master
Choose a base branch
from
Open

【训练营】分布式梯度范数与梯度裁剪实现 #222

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
e4b3fa0
feat: implement distributed gradient norm clipping
Sep 14, 2026
2dc16ad
feat: integrate gradient clipping into training examples
Sep 14, 2026
e2a845a
style: format gradient clipping changes
Sep 14, 2026
7ada2df
fix: preserve foreach option during formatting
Sep 14, 2026
d53d0a4
fix: report clipped gradient norm in examples
Sep 15, 2026
b350e85
update
Sep 15, 2026
3374208
更新报告
Sep 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
116 changes: 116 additions & 0 deletions docs/gradient_norm_clipping.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,116 @@
# 分布式梯度范数计算与梯度裁剪

## 公开接口

公开调用入口如下：

```cpp
auto total_norm = optimizer->ClipGradNorm_(
parameters,
max_norm,
norm_type,
error_if_nonfinite,
std::nullopt);
```

该接口返回一个位于 CPU 上的 FP32 标量，表示**裁剪前的梯度总范数**，并对选中的梯度进行原地缩放。

- 没有梯度的参数会被忽略。
- 对于重复出现的参数指针，仅计算一次。

## 梯度范数类型

`norm_type` 遵循 PyTorch 风格的向量范数定义：

- `0`：首先判断每个梯度张量是否至少包含一个非零元素，然后对这些梯度张量组成的列表计算 0-范数，即统计包含非零元素的梯度张量数量。
- 有限的正数或负数 `p`：按照 p-范数公式进行计算。
- `+inf`：取所有梯度元素绝对值的最大值。
- `-inf`：取所有梯度元素绝对值的最小值。

梯度缩放系数的计算公式为：

```text
min(max_norm / (total_norm + 1e-6), 1)
```

因此，当 `max_norm == 0` 时，所有梯度都会被置零，而不是禁用梯度裁剪。

## 多张量缩放

当设置 `foreach=true` 时，将使用多张量批量缩放路径。

CPU 和 CUDA 的 `ScaleInplaceMulti` 内核会直接写入现有的梯度视图，从而保留其与以下存储区域之间的别名关系：

- 扁平缓冲区（flat buffer）
- ZeRO 梯度分片

同时，每个设备和数据类型分组仅进行一次批量分派。

CUDA 梯度范数检查仅在获取标量统计信息时执行同步；梯度缩放操作仍在 CUDA 流上执行。

## 分布式梯度裁剪

`DistributedOptimizer::ClipGradNorm_` 的处理流程如下：

1. 等待所有尚未完成的梯度集合通信操作结束。
2. 针对底层优化器实际使用的分片参数计算梯度范数。
3. 在数据并行（DP）进程组中对范数统计量执行归约。
4. 使用相同的裁剪系数缩放每个进程上的本地梯度分片。

启用流水线并行（PP）后，范数统计量还会在 PP 进程组中进一步归约。

流水线调度器会在所有微批次执行完成后调用 `ClipGradNormConfigured`。

## GPT-2 和 LLaMA-3 配置参数

GPT-2 和 LLaMA-3 示例提供了以下命令行选项。

### 梯度裁剪阈值

```text
--clip_grad_norm=-1
```

负值表示禁用梯度裁剪。

### 梯度范数类型

```text
--grad_norm_type=2
```

指定梯度范数类型，此处表示使用 2-范数。

### 非有限值检查

```text
--clip_grad_error_if_nonfinite=true
```

当梯度总范数为 `NaN`、`+inf` 或 `-inf` 等非有限值时触发错误。

### 多张量缩放模式

```text
--clip_grad_foreach=auto|true|false
```

控制是否使用多张量批量缩放路径：

- `auto`：自动选择。
- `true`：强制启用。
- `false`：禁用。

## 日志记录

启用梯度裁剪后，日志中会输出一条 `total_grad_norm` 记录，用于显示裁剪前的梯度总范数。

## 并行训练支持

对于张量并行（TP）中的复制参数，`DistributedOptimizer` 会根据参数所有权进行过滤，以避免重复统计。

完整的 TP/SP/PP/vPP 混合并行训练测试由以下文件中的可选测试矩阵提供：

```text
scripts/test_config.json
```
48 changes: 43 additions & 5 deletions example/gpt2/main.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,10 @@ DEFINE_uint32(text_length, 64, "the length of the generated text");
// optimization
DEFINE_double(learning_rate, 1e-4, "Peak learning rate.");
DEFINE_int32(zero_stage, 0, "ZeRO stage (0/1/2/3); 0 disables DistributedOptimizer");
DEFINE_double(clip_grad_norm, -1.0, "Maximum gradient norm; negative disables clipping.");
DEFINE_double(grad_norm_type, 2.0, "Gradient norm type (finite p, inf, or -inf).");
DEFINE_bool(clip_grad_error_if_nonfinite, true, "Fail if the pre-clipping gradient norm is NaN or Inf.");
DEFINE_string(clip_grad_foreach, "auto", "Gradient clipping path: auto|true|false.");
// lr scheduler
DEFINE_double(min_lr, 0.0, "Minimum learning rate.");
DEFINE_string(lr_decay_style, "constant", "LR decay style: none|constant|linear|cosine|inverse-square-root");
Expand Down Expand Up @@ -352,6 +356,20 @@ void Train(const nn::parallel::Rank &rank) {
optimizer = optimizer_creator(named_parameters);
}

if (FLAGS_clip_grad_norm >= 0.0) {
std::optional<bool> foreach_option = std::nullopt;
if (FLAGS_clip_grad_foreach == "true") {
foreach_option = true;
} else if (FLAGS_clip_grad_foreach == "false") {
foreach_option = false;
} else {
CHECK_EQ(FLAGS_clip_grad_foreach, "auto");
}
optimizer->SetClipGradNormConfig(static_cast<float>(FLAGS_clip_grad_norm),
static_cast<float>(FLAGS_grad_norm_type), FLAGS_clip_grad_error_if_nonfinite,
foreach_option);
}

const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration;
TrainingLRSchedulerConfig sched_config;
sched_config.lr = static_cast<float>(FLAGS_learning_rate);
Expand Down Expand Up @@ -465,6 +483,7 @@ void Train(const nn::parallel::Rank &rank) {

const float current_lr = scheduler ? scheduler->learning_rate() : static_cast<float>(FLAGS_learning_rate);
float lossf = 0.0f;
std::optional<float> total_grad_norm;
// model->Train();
if (pp_world_size == 1) {
optimizer->ZeroGrad();
Expand Down Expand Up @@ -510,6 +529,13 @@ void Train(const nn::parallel::Rank &rank) {
LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward";
}

if (optimizer->HasClipGradNormConfig()) {
auto norm_tensor = optimizer->ClipGradNormConfigured();
if (norm_tensor) {
auto total_grad_norm_cpu = norm_tensor->To(Device());
total_grad_norm = *static_cast<const float *>(total_grad_norm_cpu.DataPtr());
}
}
optimizer->Step();
if (scheduler) {
scheduler->Step();
Expand All @@ -520,6 +546,14 @@ void Train(const nn::parallel::Rank &rank) {
y = std::make_shared<Tensor>(y->To(device));

lossf = model->TrainStep({x}, {y}, optimizer, loss_fn, dtype);
if (optimizer->HasClipGradNormConfig()) {
auto *pp_model = dynamic_cast<nn::parallel::PipelineParallel *>(model.get());
auto norm_tensor = pp_model ? pp_model->last_grad_norm() : nullptr;
if (norm_tensor) {
auto total_grad_norm_cpu = norm_tensor->To(Device());
total_grad_norm = *static_cast<const float *>(total_grad_norm_cpu.DataPtr());
}
}
if (scheduler) {
scheduler->Step();
}
Expand All @@ -538,11 +572,15 @@ void Train(const nn::parallel::Rank &rank) {
if (rank.IsLastRank()) {
size_t used_mb = 0, reserved_mb = 0;
std::tie(used_mb, reserved_mb) = impl->GetMemPoolPeakMB(device);
LOG(ERROR) << std::format("step {:4d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} tok/s | "
"peak used: {:5d} MB | peak reserved: {:5d} MB, DP={}, TP={}, SP={}, PP={})",
step + 1, FLAGS_num_iteration, lossf, current_lr, duration_us / 1e3f, tps,
used_mb, reserved_mb, ddp_world_size, tp_world_size, sp_world_size,
pp_world_size);
auto message = std::format(
"step {:4d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} tok/s | "
"peak used: {:5d} MB | peak reserved: {:5d} MB, DP={}, TP={}, SP={}, PP={})",
step + 1, FLAGS_num_iteration, lossf, current_lr, duration_us / 1e3f, tps, used_mb, reserved_mb,
ddp_world_size, tp_world_size, sp_world_size, pp_world_size);
if (total_grad_norm) {
message += std::format(" | total_grad_norm {:.6f}", *total_grad_norm);
}
LOG(ERROR) << message;

if ((step + 1) % FLAGS_freq_generate_txt == 0) {
if (tokenizer) {
Expand Down
48 changes: 43 additions & 5 deletions example/llama3/main.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,10 @@ DEFINE_uint32(text_length, 64, "the length of the generated text");
// optimization
DEFINE_double(learning_rate, 1e-5, "Peak learning rate.");
DEFINE_int32(zero_stage, 0, "ZeRO stage (0/1/2/3); 0 disables DistributedOptimizer");
DEFINE_double(clip_grad_norm, -1.0, "Maximum gradient norm; negative disables clipping.");
DEFINE_double(grad_norm_type, 2.0, "Gradient norm type (finite p, inf, or -inf).");
DEFINE_bool(clip_grad_error_if_nonfinite, true, "Fail if the pre-clipping gradient norm is NaN or Inf.");
DEFINE_string(clip_grad_foreach, "auto", "Gradient clipping path: auto|true|false.");
// lr scheduler
DEFINE_double(min_lr, 0.0, "Minimum learning rate.");
DEFINE_string(lr_decay_style, "constant", "LR decay style: none|constant|linear|cosine|inverse-square-root");
Expand Down Expand Up @@ -334,6 +338,20 @@ void Train(const nn::parallel::Rank &rank) {
optimizer = optimizer_creator(named_parameters);
}

if (FLAGS_clip_grad_norm >= 0.0) {
std::optional<bool> foreach_option = std::nullopt;
if (FLAGS_clip_grad_foreach == "true") {
foreach_option = true;
} else if (FLAGS_clip_grad_foreach == "false") {
foreach_option = false;
} else {
CHECK_EQ(FLAGS_clip_grad_foreach, "auto");
}
optimizer->SetClipGradNormConfig(static_cast<float>(FLAGS_clip_grad_norm),
static_cast<float>(FLAGS_grad_norm_type), FLAGS_clip_grad_error_if_nonfinite,
foreach_option);
}

const int64_t lr_decay_iters = FLAGS_lr_decay_iters > 0 ? FLAGS_lr_decay_iters : FLAGS_num_iteration;
TrainingLRSchedulerConfig sched_config;
sched_config.lr = static_cast<float>(FLAGS_learning_rate);
Expand Down Expand Up @@ -445,6 +463,7 @@ void Train(const nn::parallel::Rank &rank) {

const float current_lr = scheduler ? scheduler->learning_rate() : static_cast<float>(FLAGS_learning_rate);
float lossf = 0.0f;
std::optional<float> total_grad_norm;
if (pp_world_size == 1) {
// model->Train();
optimizer->ZeroGrad();
Expand Down Expand Up @@ -489,6 +508,13 @@ void Train(const nn::parallel::Rank &rank) {
LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish backward";
}

if (optimizer->HasClipGradNormConfig()) {
auto norm_tensor = optimizer->ClipGradNormConfigured();
if (norm_tensor) {
auto total_grad_norm_cpu = norm_tensor->To(Device());
total_grad_norm = *static_cast<const float *>(total_grad_norm_cpu.DataPtr());
}
}
optimizer->Step();
if (scheduler) {
scheduler->Step();
Expand All @@ -499,6 +525,14 @@ void Train(const nn::parallel::Rank &rank) {
y = std::make_shared<Tensor>(y->To(device));

lossf = model->TrainStep({x}, {y}, optimizer, loss_fn, dtype);
if (optimizer->HasClipGradNormConfig()) {
auto *pp_model = dynamic_cast<nn::parallel::PipelineParallel *>(model.get());
auto norm_tensor = pp_model ? pp_model->last_grad_norm() : nullptr;
if (norm_tensor) {
auto total_grad_norm_cpu = norm_tensor->To(Device());
total_grad_norm = *static_cast<const float *>(total_grad_norm_cpu.DataPtr());
}
}
if (scheduler) {
scheduler->Step();
}
Expand All @@ -517,11 +551,15 @@ void Train(const nn::parallel::Rank &rank) {
if (rank.IsLastRank()) {
size_t used_mb = 0, reserved_mb = 0;
std::tie(used_mb, reserved_mb) = impl->GetMemPoolPeakMB(device);
LOG(ERROR) << std::format("step {:4d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} tok/s | "
"peak used: {:5d} MB | peak reserved: {:5d} MB, DP={}, TP={}, SP={}, PP={})",
step + 1, FLAGS_num_iteration, lossf, current_lr, duration_us / 1e3f, tps,
used_mb, reserved_mb, ddp_world_size, tp_world_size, sp_world_size,
pp_world_size);
auto message = std::format(
"step {:4d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} tok/s | "
"peak used: {:5d} MB | peak reserved: {:5d} MB, DP={}, TP={}, SP={}, PP={})",
step + 1, FLAGS_num_iteration, lossf, current_lr, duration_us / 1e3f, tps, used_mb, reserved_mb,
ddp_world_size, tp_world_size, sp_world_size, pp_world_size);
if (total_grad_norm) {
message += std::format(" | total_grad_norm {:.6f}", *total_grad_norm);
}
LOG(ERROR) << message;

if ((step + 1) % FLAGS_freq_generate_txt == 0) {
// FIXME(jym): to support PP
Expand Down
Loading