From a3c6c9eb4c254acdc2eba782456e443406e5ee8d Mon Sep 17 00:00:00 2001 From: littsk <1214689160@qq.com> Date: Mon, 25 Sep 2023 10:45:50 +0800 Subject: [PATCH] fix norm type error in zero optimizer --- colossalai/zero/low_level/_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/colossalai/zero/low_level/_utils.py b/colossalai/zero/low_level/_utils.py index ba1135940df0..0a15f8ddd718 100644 --- a/colossalai/zero/low_level/_utils.py +++ b/colossalai/zero/low_level/_utils.py @@ -221,8 +221,8 @@ def compute_norm(gradients: Tensor, dp_group: ProcessGroup, tp_group: ProcessGro else: total_norm = 0.0 for g in gradients: - param_norm = g.data.double().norm(2) - total_norm += param_norm.item() ** 2 + param_norm = g.data.double().norm(norm_type) + total_norm += param_norm.item() ** norm_type # Sum across all model parallel GPUs. total_norm_cuda = torch.cuda.FloatTensor([float(total_norm)])