From 4abd20b4d02e7717818b076c8cadcdd030120e90 Mon Sep 17 00:00:00 2001 From: JazyJiang Date: Thu, 15 Oct 2026 20:13:38 +0800 Subject: [PATCH] [zero] compute the grad norm over all gradient stores at once With expert parallelism, LowLevelZero keeps dense and expert params in two gradient stores. step() computed a norm per store and added them after the square root (D + E instead of sqrt(D^2 + E^2)), and the expert norm was only reduced over the moe_dp group, so experts held by other EP ranks were never counted. Each EP rank got a different, usually much too small norm, and clipping used it. Follow Megatron-Core: _compute_grad_norm now takes the grads of all stores, sums their p-th powers locally, all-reduces the sum once over _grad_norm_process_groups() and takes the root once. The groups default to a new grad_norm_process_group argument, which is the store's own group when there is only one store. MoeHybridParallelZeroOptimizer passes its dp group, which spans the moe_dp and ep axes. HybridParallelZeroOptimizer adds the tp and pp groups, and its existing duplicate handling (1/tp_size for params not split by TP, 1/num_stages for shared params) moves unchanged into a _grad_norm_weight hook. This also drops its inf branch, which called the parent with shifted arguments, and every rank now joins the collective even without grads. Add test_moe_ep_zero_grad_norm.py (ep=2/4, MoE-only and full training): the norm must be equal on all ranks and match an fp32 single-GPU reference. Update test_zero_optimizer.py, which called the private method with the old signature. Co-Authored-By: Claude Opus 5.5 --- .../booster/plugin/hybrid_parallel_plugin.py | 109 ++++-------------- .../plugin/moe_hybrid_parallel_plugin.py | 2 + colossalai/zero/low_level/low_level_optim.py | 97 +++++++--------- tests/test_moe/test_moe_ep_zero_grad_norm.py | 103 +++++++++++++++++ .../test_zero_optimizer.py | 3 +- 5 files changed, 171 insertions(+), 143 deletions(-) create mode 100644 tests/test_moe/test_moe_ep_zero_grad_norm.py diff --git a/colossalai/booster/plugin/hybrid_parallel_plugin.py b/colossalai/booster/plugin/hybrid_parallel_plugin.py index 083332f278c1..ea0a650d5051 100644 --- a/colossalai/booster/plugin/hybrid_parallel_plugin.py +++ b/colossalai/booster/plugin/hybrid_parallel_plugin.py @@ -692,6 +692,7 @@ def __init__( forced_dtype: Optional[torch.dtype] = None, overlap_allgather: bool = False, fp8_communication: bool = False, + grad_norm_process_group: Optional[ProcessGroup] = None, ): self.model = model self.param_info = param_info @@ -723,6 +724,7 @@ def __init__( overlap_allgather=overlap_allgather, fp8_communication=fp8_communication, backward_context=model._hook_context, + grad_norm_process_group=grad_norm_process_group, ) def sync_dp_grads(self): @@ -840,90 +842,29 @@ def backward_by_grad(self, tensor, grad, inputs: Tensor = None, retain_graph: bo # If gradient synchronization is is not required, return. return - def _compute_grad_norm(self, dp_pg, gradients: List[Tensor], norm_type: int = 2) -> float: - r""" - Compute and return the gradient norm for gradient clipping. - - Args: - gradients (List[Tensor]): A list of tensors containing gradients. - norm_type (int, optional): Type of the p-norm to be computed. Defaults to 2. - - Returns: - float: The computed gradient norm. - """ - - # Check if the list of gradients is empty - if len(gradients) == 0: - return 0.0 - - dp_size = get_world_size(dp_pg) if dp_pg is not None else 1 - tp_size = get_world_size(self.tp_pg) if self.tp_pg is not None else 1 - pp_size = get_world_size(self.pp_pg) if self.pp_pg is not None else 1 - norm_type = float(norm_type) - - if norm_type == inf: - # The parent class calculates the norm of 'dp' gradients, - # so we only need to calculate the norm 'tp' of 'pp' gradients. - total_norm = super()._compute_grad_norm(gradients, norm_type) - - total_norm_cuda = torch.tensor( - [float(total_norm)], device=get_accelerator().get_current_device(), dtype=torch.float32 - ) - - if tp_size > 1: - dist.all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.MAX, group=self.tp_pg) - if pp_size > 1: - dist.all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.MAX, group=self.pp_pg) - - total_norm = total_norm_cuda.item() - else: - total_norm_exponentiated = 0.0 - for grad in gradients: - grad_norm_exponentiated = grad.data.double().norm(norm_type) ** norm_type - - # If 'tp_size' is greater than 1 and the parameter for the gradient is not a distributed tensor, - # it indicates that the parameter is not distributed across devices of the 'tp_group'. - # Consequently, there is no need to perform an 'all_reduce' operation for 'grad_norm'. - # However, we still perform the 'all_reduce' operation for the sake of good coding practices. - # To ensure mathematical equivalence, we divide the 'grad_norm' by 'tp_size.' - if tp_size > 1: - param_id_for_grad = self.get_param_id_for_grad(grad) - param_for_grad = ctypes.cast(param_id_for_grad, ctypes.py_object).value - - if not is_distributed_tensor(param_for_grad): - grad_norm_exponentiated /= tp_size - - # If 'pp_size' is greater than 1 and the gradient belongs to shared parameters, - # it means that this parameter is used in two different pipeline stages. - # To avoid redundant norm calculations, we divide the exponent of this norm by - # the number of shared stages. - if pp_size > 1: - for shared_param in self.shared_params: - if self.stage_manager.stage in shared_param: - stage_shared_param = shared_param[self.stage_manager.stage] - working_grad = self.get_working_grad_by_param_id(id(stage_shared_param)) - if grad is working_grad: - grad_norm_exponentiated /= len(shared_param) - - total_norm_exponentiated += grad_norm_exponentiated - - total_norm_exponentiated_cuda = torch.tensor( - [float(total_norm_exponentiated)], device=get_accelerator().get_current_device(), dtype=torch.float32 - ) - if dp_size > 1: - # compute norm in dp process group - dist.all_reduce(tensor=total_norm_exponentiated_cuda, op=dist.ReduceOp.SUM, group=dp_pg) - if tp_size > 1: - # compute norm in tp process group - dist.all_reduce(tensor=total_norm_exponentiated_cuda, op=dist.ReduceOp.SUM, group=self.tp_pg) - if pp_size > 1: - # compute norm in pp process group - dist.all_reduce(tensor=total_norm_exponentiated_cuda, op=dist.ReduceOp.SUM, group=self.pp_pg) - - # Compute the 'total_norm' from 'total_norm_exponentiated' - total_norm = total_norm_exponentiated_cuda.item() ** (1.0 / norm_type) - - return total_norm + def _grad_norm_process_groups(self) -> List[ProcessGroup]: + # gradients are also split across tp and pp ranks, so the norm is summed over them too + groups = super()._grad_norm_process_groups() + for group in (self.tp_pg, self.pp_pg): + if group is not None and get_world_size(group) > 1: + groups.append(group) + return groups + + def _grad_norm_weight(self, grad: Tensor) -> float: + weight = 1.0 + # a parameter not split by tp is replicated on every tp rank: each copy counts 1 / tp_size + if self.tp_pg is not None and get_world_size(self.tp_pg) > 1: + param_for_grad = ctypes.cast(self.get_param_id_for_grad(grad), ctypes.py_object).value + if not is_distributed_tensor(param_for_grad): + weight /= get_world_size(self.tp_pg) + # a parameter shared by several pipeline stages (e.g. tied embeddings) counts 1 / num_stages per stage + if self.pp_pg is not None and get_world_size(self.pp_pg) > 1: + for shared_param in self.shared_params: + if self.stage_manager.stage in shared_param: + stage_shared_param = shared_param[self.stage_manager.stage] + if grad is self.get_working_grad_by_param_id(id(stage_shared_param)): + weight /= len(shared_param) + return weight class HybridParallelPlugin(PipelinePluginBase): diff --git a/colossalai/booster/plugin/moe_hybrid_parallel_plugin.py b/colossalai/booster/plugin/moe_hybrid_parallel_plugin.py index a733fc5f5261..f4337ef14211 100644 --- a/colossalai/booster/plugin/moe_hybrid_parallel_plugin.py +++ b/colossalai/booster/plugin/moe_hybrid_parallel_plugin.py @@ -101,6 +101,8 @@ def __init__( forced_dtype=forced_dtype, pg_to_param_list=pg_param_list, overlap_allgather=overlap_allgather, + # the dp group spans the moe_dp and ep axes, so it covers every shard of dense and expert grads + grad_norm_process_group=dp_process_group, ) diff --git a/colossalai/zero/low_level/low_level_optim.py b/colossalai/zero/low_level/low_level_optim.py index c530ff009fb1..28161046af98 100644 --- a/colossalai/zero/low_level/low_level_optim.py +++ b/colossalai/zero/low_level/low_level_optim.py @@ -99,6 +99,7 @@ def __init__( overlap_allgather: bool = False, fp8_communication: bool = False, backward_context=None, + grad_norm_process_group: Optional[Union[ProcessGroup, Tuple[ProcessGroup, ...]]] = None, ): super(LowLevelZeroOptimizer, self).__init__(optim=optimizer) @@ -124,6 +125,15 @@ def __init__( pg_to_param_list[unique_dp_group].extend(group["params"]) self.pg_to_param_list = pg_to_param_list + # The grad norm is reduced once over this group (see _compute_grad_norm), so it must contain every rank that + # owns a distinct gradient shard of any store. With a single store this is the store's own group. + if grad_norm_process_group is None: + if len(pg_to_param_list) > 1: + raise ValueError( + "grad_norm_process_group should be provided when parameters are split across several process groups." + ) + grad_norm_process_group = next(iter(pg_to_param_list)) + self._grad_norm_process_group = grad_norm_process_group param_to_pg = {} for grp, param_list in pg_to_param_list.items(): for p in param_list: @@ -550,7 +560,6 @@ def step(self, closure=None): working_params = self._working_param_groups[group_id] real_working_params[group_id] = [] real_master_params[group_id] = [] - working_grads = [] for working_param, master_param in zip(working_params, master_params): # if a working param requires grad and has no grad # it is not 'really' working, e.g. the droped layer @@ -568,13 +577,8 @@ def step(self, closure=None): grad_partition_groups.append(grad) real_master_params[group_id].append(master_param) - # compute norm - norm_group = 0 - for grad_store in self.pg_to_grad_store.values(): - working_grads = grad_store.get_working_grads_by_group_id(group_id) - norm_group += self._compute_grad_norm(dp_pg=grad_store.torch_pg, gradients=working_grads) - - norm_groups.append(norm_group) + # one norm over the grads of all stores, so that experts and dense params are combined correctly + norm_groups.append(self._compute_grad_norm(self.get_working_grads_by_group_id(group_id))) # update the params in the optimizer self.optim.param_groups[group_id]["params"] = real_master_params[group_id] @@ -634,66 +638,45 @@ def step(self, closure=None): if not tensor_bucket.is_empty(): tensor_bucket.all_gather(pg, fp8_communication=self._fp8_communication) - def _compute_grad_norm( - self, dp_pg: Union[ProcessGroup, Tuple[ProcessGroup, ...]], gradients: List[Tensor], norm_type: int = 2 - ) -> float: + def _compute_grad_norm(self, gradients: List[Tensor], norm_type: int = 2) -> float: r""" - Compute and return the gradient norm for gradient clipping. + Compute the norm of the whole gradient, viewed as a single vector, for gradient clipping. + + Each rank passes the gradient shards it owns, from all gradient stores. Their p-th power sum (or max for the + infinity norm) is reduced once over the groups of ``_grad_norm_process_groups`` and the root is taken once, + so every element is counted once and all ranks get the same value. Copies of the same element held by several + ranks are handled by ``_grad_norm_weight``. Every rank must call this, even without gradients. Args: - gradients (List[Tensor]): The gradients to compute norm + gradients (List[Tensor]): The gradient shards owned by this rank norm_type (int, optional): type of the used p-norm, Can be ``'inf'`` for infinity norm. Defaults to 2. Returns: - float: The total norm of given gradients + float: The total norm of the gradients """ - - if len(gradients) == 0: - return 0.0 - norm_type = float(norm_type) if norm_type == inf: - total_norm = max(grad.data.abs().max() for grad in gradients) - total_norm_cuda = torch.tensor( - [float(total_norm)], - device=get_accelerator().get_current_device(), - dtype=torch.float, - ) - if isinstance(dp_pg, tuple): - for grp in dp_pg: - dist.all_reduce(total_norm_cuda, op=torch.distributed.ReduceOp.MAX, group=grp) - else: - dist.all_reduce(total_norm_cuda, op=torch.distributed.ReduceOp.MAX, group=dp_pg) - total_norm = total_norm_cuda.item() - + local = max((float(grad.data.abs().max()) for grad in gradients), default=0.0) + op = dist.ReduceOp.MAX else: - total_norm_exponentiated = 0.0 + local = 0.0 for grad in gradients: - grad_norm_exponentiated = grad.data.double().norm(norm_type) ** norm_type - total_norm_exponentiated += grad_norm_exponentiated - - # Sum across all model parallel GPUs. - total_norm_exponentiated_cuda = torch.tensor( - [float(total_norm_exponentiated)], - device=get_accelerator().get_current_device(), - dtype=torch.float, - ) - if isinstance(dp_pg, tuple): - for grp in dp_pg: - dist.all_reduce( - total_norm_exponentiated_cuda, - op=torch.distributed.ReduceOp.SUM, - group=grp, - ) - else: - torch.distributed.all_reduce( - total_norm_exponentiated_cuda, - op=torch.distributed.ReduceOp.SUM, - group=dp_pg, - ) - total_norm = total_norm_exponentiated_cuda.item() ** (1.0 / norm_type) - - return total_norm + local += self._grad_norm_weight(grad) * grad.data.double().norm(norm_type) ** norm_type + op = dist.ReduceOp.SUM + + total = torch.tensor([float(local)], device=get_accelerator().get_current_device(), dtype=torch.float) + for group in self._grad_norm_process_groups(): + dist.all_reduce(total, op=op, group=group) + return total.item() if norm_type == inf else total.item() ** (1.0 / norm_type) + + def _grad_norm_process_groups(self) -> List[ProcessGroup]: + """The groups the local grad norm is reduced over, one after another.""" + group = self._grad_norm_process_group + return list(group) if isinstance(group, tuple) else [group] + + def _grad_norm_weight(self, grad: Tensor) -> float: + """How much this rank's copy of a gradient counts towards the norm; copies of one gradient sum to 1.""" + return 1.0 ############################# # Mixed Precision Utilities # diff --git a/tests/test_moe/test_moe_ep_zero_grad_norm.py b/tests/test_moe/test_moe_ep_zero_grad_norm.py new file mode 100644 index 000000000000..b66378833902 --- /dev/null +++ b/tests/test_moe/test_moe_ep_zero_grad_norm.py @@ -0,0 +1,103 @@ +from copy import deepcopy + +import pytest +import torch +import torch.distributed as dist +from transformers.models.mixtral.configuration_mixtral import MixtralConfig +from transformers.models.mixtral.modeling_mixtral import MixtralModel + +import colossalai +from colossalai.booster.booster import Booster +from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.testing.random import seed_all + +# Relative error allowed between the ZeRO grad norm (bf16 grads) and the norm of an fp32 single-GPU reference. With +# the norm combined correctly it is below 1.2e-3 here; the old per-store norms were off by 0.6% to 67%. +NORM_RTOL = 5e-3 + +HIDDEN_SIZE = 8 +NUM_EXPERTS = 8 + + +@parameterize("ep_size", [2, 4]) +@parameterize("moe_only", [True, False]) +def run_grad_norm(ep_size: int, moe_only: bool): + """The ZeRO grad norm must be the norm of the whole gradient and equal on every rank. + + Dense and expert grads live in two ZeRO stores, and every EP rank holds different experts, so the norm has to + combine both stores and all EP ranks. Training only the MoE blocks makes experts carry most of the norm, as in a + real MoE model; training everything mixes dense and expert grads. Routing is random, so every expert is used on + every moe_dp replica and the moe_dp scaling of #6464 cannot interfere. + """ + if dist.get_world_size() % ep_size != 0: + return + dtype = torch.bfloat16 + rank = dist.get_rank() + + plugin = MoeHybridParallelPlugin( + pp_size=1, + tp_size=1, + ep_size=ep_size, + zero_stage=1, + precision="bf16", + overlap_communication=False, + ) + booster = Booster(plugin=plugin) + + seed_all(10086) + config = MixtralConfig( + hidden_size=HIDDEN_SIZE, + intermediate_size=HIDDEN_SIZE * 64, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + num_local_experts=NUM_EXPERTS, + num_experts_per_tok=2, + ) + torch_model = MixtralModel(config).to(dtype).cuda() + for name, param in torch_model.named_parameters(): + param.requires_grad_(not moe_only or "block_sparse_moe" in name) + + # single-GPU reference: same weights in fp32 on the concatenation of every rank's batch (= global mean loss) + single_model = deepcopy(torch_model).float() + zero_model = deepcopy(torch_model) + zero_optimizer = torch.optim.SGD(zero_model.parameters(), lr=0.0) + zero_model, zero_optimizer, _, _, _ = booster.boost(zero_model, zero_optimizer) + + seed_all(1453 + rank) + input_data = torch.rand(8, 16, HIDDEN_SIZE).cuda().to(dtype) + zero_optimizer.backward(zero_model(inputs_embeds=input_data).last_hidden_state.mean()) + zero_optimizer.step() # lr = 0: nothing moves, but this computes the norm used for clipping + zero_norm = zero_optimizer.get_grad_norm() + + all_inputs = [torch.empty_like(input_data) for _ in range(dist.get_world_size())] + dist.all_gather(all_inputs, input_data) + single_model(inputs_embeds=torch.cat(all_inputs).float()).last_hidden_state.mean().backward() + grads = [p.grad.double() for p in single_model.parameters() if p.grad is not None] + ref_norm = sum(g.pow(2).sum() for g in grads).sqrt().item() + + norms = [None] * dist.get_world_size() + dist.all_gather_object(norms, zero_norm) + assert len(set(norms)) == 1, f"ep_size={ep_size} moe_only={moe_only}: grad norm differs across ranks: {norms}" + assert abs(zero_norm - ref_norm) <= NORM_RTOL * ref_norm, ( + f"ep_size={ep_size} moe_only={moe_only}: ZeRO grad norm {zero_norm:.6g}, reference {ref_norm:.6g}, " + f"ratio {zero_norm / ref_norm:.4f}" + ) + + +def run_dist(rank, world_size, port): + colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") + torch.cuda.set_device(rank) + run_grad_norm() + + +@pytest.mark.dist +@pytest.mark.parametrize("world_size", [4]) +@rerun_if_address_is_in_use() +def test_moe_ep_zero_grad_norm(world_size): + spawn(run_dist, world_size) + + +if __name__ == "__main__": + test_moe_ep_zero_grad_norm(world_size=4) diff --git a/tests/test_shardformer/test_hybrid_parallel_grad_clip_norm/test_zero_optimizer.py b/tests/test_shardformer/test_hybrid_parallel_grad_clip_norm/test_zero_optimizer.py index e37a050e3dbe..c96a4a68ed06 100644 --- a/tests/test_shardformer/test_hybrid_parallel_grad_clip_norm/test_zero_optimizer.py +++ b/tests/test_shardformer/test_hybrid_parallel_grad_clip_norm/test_zero_optimizer.py @@ -32,7 +32,6 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, stage_manager = booster.plugin.stage_manager tp_group = booster.plugin.tp_group - dp_group = booster.plugin.dp_group bert = unwrap_model(org_model, "BertModel", "bert") sharded_bert = unwrap_model(sharded_model, "BertModel", "bert") @@ -55,7 +54,7 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, norm_groups = [] for group_id in range(sharded_optimizer.num_param_groups): working_grads = sharded_optimizer.get_working_grads_by_group_id(group_id) - norm_group = sharded_optimizer._compute_grad_norm(dp_group, gradients=working_grads) + norm_group = sharded_optimizer._compute_grad_norm(gradients=working_grads) norm_groups.append(norm_group) total_norm = 0.0 for norm in norm_groups: