From 54b3ad89249a23221dda036ed598dade06808011 Mon Sep 17 00:00:00 2001 From: littsk <1214689160@qq.com> Date: Wed, 27 Sep 2023 10:35:24 +0800 Subject: [PATCH] [hotfix] fix norm type error in zero optimizer (#4795) --- 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 ba1135940..0a15f8ddd 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)])