diff --git a/deepspeed/runtime/utils.py b/deepspeed/runtime/utils.py index 61879cdad560..a2ccab9ee5ff 100644 --- a/deepspeed/runtime/utils.py +++ b/deepspeed/runtime/utils.py @@ -376,20 +376,81 @@ def clip_grad_norm_(parameters, max_norm, norm_type=2, mpu=None): """ if isinstance(parameters, torch.Tensor): parameters = [parameters] - parameters = list(filter(lambda p: p.grad is not None, parameters)) + parameters = list(parameters) norm_type = float(norm_type) + + # Expert parameters are not replicated across the data parallel group, so the + # rank-averaging below cannot reconstruct a global norm from them: each rank's + # norm covers a different set of experts. Combine those the way the bf16 and + # fp16 optimizers already do, over the expert parallel group. + # + # Read ownership from the parameters, before dropping the ones with no gradient. + # The two branches below issue different collectives, so this decision has to + # come out the same on every rank -- and a rank whose experts received no tokens + # this step has expert parameters but no expert gradients. + # The registry of expert groups is built identically on every rank, so an empty one + # means no MoE is configured anywhere and the choice is already made: take the + # original path, with no collective and no host sync. This is a per-step call from + # `_take_model_step`, so the common case must not pay for the MoE one. + registry = sorted(getattr(groups, "_EXPERT_PARALLEL_GROUP", None) or {}) + expert_group_names = [] + if registry: + # A rank can legitimately own no expert at all -- under pipeline parallelism, a + # stage holding only dense layers -- so reading ownership locally cannot guarantee + # the same choice everywhere. One scalar settles it, and it has to be reduced over + # the world rather than the data parallel group: under pipeline parallelism that + # group is one stage wide (`PipelineParallelGrid.get_data_parallel_group`), while + # the expert norm below reduces over the pipeline group, which spans stages. A + # per-stage vote let an expert stage and a dense stage issue different collectives + # on that group. The second slot carries "an expert here has no group name", which + # has no expert parallel group to reduce over, so such a call stays on the + # original path rather than guessing which group the parameter belongs to. + vote = torch.tensor([ + float(any(is_moe_param(p) and getattr(p, "group_name", None) is not None for p in parameters)), + float(any(is_moe_param(p) and getattr(p, "group_name", None) is None for p in parameters)), + ], + device=get_accelerator().current_device_name(), + dtype=torch.float) + dist.all_reduce(vote, group=groups._clone_world_group()) + if vote[0].item() > 0 and vote[1].item() == 0: + expert_group_names = registry + + parameters = list(filter(lambda p: p.grad is not None, parameters)) + + expert_tensors: Dict[str, List[torch.Tensor]] = {} + if expert_group_names: + # `get_norm_with_moe_layers` issues one collective per entry, so every rank has to + # walk the groups in the same order. Start from the registry's order rather than + # the order this rank's parameters happen to come in. + expert_tensors = {group_name: [] for group_name in expert_group_names} + for p in parameters: + if is_moe_param(p): + expert_tensors.setdefault(p.group_name, []).append(p.grad.data) + # Every rank has to reduce over every group, whether or not its own experts + # were used this step; a zero contributes nothing to the sum of p-th powers + # and nothing to a max over absolute values. + zero = torch.zeros(1, device=get_accelerator().current_device_name()) + for tensors in expert_tensors.values(): + if not tensors: + tensors.append(zero) + + norm_parameters = [p for p in parameters if not is_moe_param(p)] if expert_tensors else parameters + all_norms = [] if norm_type == inf: - for p in parameters: + for p in norm_parameters: all_norms.append(p.grad.data.abs().max().float()) - total_norm = torch.stack(all_norms).max() + # Everything may be an expert, leaving nothing replicated to take a max over. + # Zero is the identity here: these are absolute values. + total_norm = torch.stack(all_norms).max() if all_norms \ + else torch.zeros((), device=get_accelerator().current_device_name()) total_norm = total_norm.to(get_accelerator().current_device_name()) # Take max across all GPUs. if mpu is not None: dist.all_reduce(total_norm, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group()) else: total_norm = 0 - for p in parameters: + for p in norm_parameters: if mpu is not None: if (mpu.get_model_parallel_rank() == 0) or is_model_parallel_parameter(p): param_norm = p.grad.data.detach().float().norm(norm_type) @@ -410,14 +471,26 @@ def clip_grad_norm_(parameters, max_norm, norm_type=2, mpu=None): dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group()) total_norm = total_norm.pow(1. / norm_type) - # Need to average total_norm across different GPUs due to the presence of moe params - pg = groups._get_data_parallel_group() - scaled_norm = total_norm * 1.0 / float(dist.get_world_size(group=pg)) - scaled_norm_tensor = scaled_norm - - dist.all_reduce(scaled_norm_tensor, group=pg) - total_norm = scaled_norm_tensor - total_norm = total_norm.to(parameters[0].device) + if expert_tensors: + # `total_norm` now covers only the replicated parameters, so it is already the + # same on every rank. Fold the experts in over their own group, counting each + # expert once, rather than averaging norms that describe different parameters. + moe_norm = get_norm_with_moe_layers(total_norm, mpu=mpu, expert_tensors=expert_tensors, norm_type=norm_type) + # That helper reports a non-finite norm as -1. Left as-is it would make + # `clip_coef` negative and flip the sign of every gradient; inf keeps the + # behaviour the averaging path already had, which is to scale them to zero. + if moe_norm == -1: + moe_norm = float('inf') + total_norm = torch.tensor([float(moe_norm)], device=total_norm.device, dtype=torch.float) + else: + # Need to average total_norm across different GPUs due to the presence of moe params + pg = groups._get_data_parallel_group() + scaled_norm = total_norm * 1.0 / float(dist.get_world_size(group=pg)) + scaled_norm_tensor = scaled_norm + + dist.all_reduce(scaled_norm_tensor, group=pg) + total_norm = scaled_norm_tensor + total_norm = total_norm.to(parameters[0].device) max_norm = torch.tensor([float(max_norm)], device=total_norm.device) clip_coef = max_norm / (total_norm + 1e-6) diff --git a/tests/unit/runtime/test_runtime_utils.py b/tests/unit/runtime/test_runtime_utils.py index a0bdebdacac6..e9216f465840 100644 --- a/tests/unit/runtime/test_runtime_utils.py +++ b/tests/unit/runtime/test_runtime_utils.py @@ -9,9 +9,12 @@ import pytest from typing import Dict +import deepspeed import deepspeed.runtime.utils as ds_utils import deepspeed.utils.groups as groups from deepspeed.accelerator import get_accelerator +from deepspeed.moe.layer import MoE +from deepspeed.moe.utils import is_moe_param, split_params_into_different_moe_groups_for_optimizer from unit.common import DistributedTest @@ -51,6 +54,129 @@ def test_gather(self): assert gathered_norm[0] == gathered_norm[1], "norm at rank 0 does not match the norm at rank 1" + def test_sharded_experts_count_once(self): + """The global norm must not depend on how experts are spread over ranks. + + Rank-averaging the per-rank norms only reconstructs a global norm when every + rank holds the same parameters. Under expert parallelism they do not, so the + four experts below gave sqrt(20) and sqrt(100) averaged to 7.236 instead of + the sqrt(120) = 10.954 that counting each expert once produces (#8469). + """ + groups._create_expert_and_data_parallel(2) + rank = dist.get_rank() + device = get_accelerator().device_name(rank) + + experts = [torch.full((4, ), float(v)) for v in (1.0, 2.0, 3.0, 4.0)] + expected = torch.cat(experts).norm(2).item() + + owned = experts[:2] if rank == 0 else experts[2:] + params = [] + for grad in owned: + param = torch.nn.Parameter(torch.zeros(4, device=device)) + param.grad = grad.clone().to(device) + param.allreduce = False + param.group_name = "ep_size_2" + params.append(param) + + # A max_norm far above the norm leaves the gradients alone, so this reads the + # computed norm rather than the clipped result. + norm = ds_utils.clip_grad_norm_(params, max_norm=1e9) + + assert abs(float(norm) - + expected) < 1e-4, (f"global norm {float(norm)} should be {expected} regardless of expert placement") + + def _expert(self, value, device, with_grad=True): + param = torch.nn.Parameter(torch.zeros(4, device=device)) + if with_grad: + param.grad = torch.full((4, ), value, device=device) + param.allreduce = False + param.group_name = "ep_size_2" + return param + + def _plain(self, value, device): + param = torch.nn.Parameter(torch.zeros(4, device=device)) + param.grad = torch.full((4, ), value, device=device) + return param + + def test_experts_idle_on_one_rank_do_not_deadlock(self): + """A rank whose experts got no tokens must take the same collectives (#8469). + + Ownership has to be read from the parameters, not the gradients: `p.grad is + not None` filtering drops an unused expert, and branching on what is left made + one rank reduce over the expert group while its peer reduced over the data + parallel group. They then waited on each other. + """ + groups._create_expert_and_data_parallel(2) + rank = dist.get_rank() + device = get_accelerator().device_name(rank) + + expert = self._expert(2.0, device, with_grad=(rank == 0)) + norm = ds_utils.clip_grad_norm_([self._plain(1.0, device), expert], max_norm=1e9) + + # non-expert 2.0 replicated, one expert of norm 4.0 owned by rank 0: + # sqrt(2**2 + 4**2). + assert abs(float(norm) - 20**0.5) < 1e-4 + + def test_a_rank_owning_no_expert_does_not_deadlock(self): + """Under pipeline parallelism a rank can hold no expert at all. + + Local ownership then differs in kind rather than in degree, so the branch is + settled by one all-reduced flag instead of by what this rank happens to hold. + """ + groups._create_expert_and_data_parallel(2) + rank = dist.get_rank() + device = get_accelerator().device_name(rank) + + params = [self._plain(1.0, device)] + if rank == 0: + params.append(self._expert(2.0, device)) + + norm = ds_utils.clip_grad_norm_(params, max_norm=1e9) + assert abs(float(norm) - 20**0.5) < 1e-4 + + def test_expert_only_inf_norm(self): + """Nothing replicated leaves no tensor to take a max over; zero is the identity.""" + groups._create_expert_and_data_parallel(2) + rank = dist.get_rank() + device = get_accelerator().device_name(rank) + + norm = ds_utils.clip_grad_norm_([self._expert(float(rank + 1), device)], max_norm=1e9, norm_type=float('inf')) + + assert abs(float(norm) - 2.0) < 1e-4 + + def test_a_rank_with_no_gradient_at_all_still_gets_the_norm(self): + """Every local gradient can be absent, expert ones included, and the rank still has to answer.""" + groups._create_expert_and_data_parallel(2) + rank = dist.get_rank() + device = get_accelerator().device_name(rank) + + norm = ds_utils.clip_grad_norm_([self._expert(2.0, device, with_grad=(rank == 1))], max_norm=1e9) + + # Only rank 1's expert has a gradient, of norm 4.0, and both ranks report it. + assert abs(float(norm) - 4.0) < 1e-4 + + def test_expert_groups_are_visited_in_registry_order(self, monkeypatch): + """Ranks owning different expert groups must still walk the groups in one order.""" + groups._create_expert_and_data_parallel(2) + groups._create_expert_and_data_parallel(1) + rank = dist.get_rank() + device = get_accelerator().device_name(rank) + + visited = [] + real = ds_utils.get_norm_with_moe_layers + + def spy(non_expert_norm, mpu, expert_tensors, norm_type=2): + visited.append(list(expert_tensors)) + return real(non_expert_norm, mpu=mpu, expert_tensors=expert_tensors, norm_type=norm_type) + + monkeypatch.setattr(ds_utils, "get_norm_with_moe_layers", spy) + + expert = self._expert(2.0, device) + expert.group_name = "ep_size_2" if rank == 0 else "ep_size_1" + ds_utils.clip_grad_norm_([expert], max_norm=1e9) + + assert visited == [["ep_size_1", "ep_size_2"]] + def test_clipped_val(self): max_norm = 0.1 @@ -73,6 +199,84 @@ def test_params(): assert torch.equal(params_expected[1].grad, params_actual[1].grad) +class _OneMoELayer(torch.nn.Module): + """A dense layer, then one expert per rank behind a gate.""" + + def __init__(self, hidden_dim): + super().__init__() + self.dense = torch.nn.Linear(hidden_dim, hidden_dim) + expert = torch.nn.Linear(hidden_dim, hidden_dim) + # k=2 sends every token to both experts, so each rank's expert gets a gradient. + self.moe = MoE(hidden_size=hidden_dim, expert=expert, num_experts=2, ep_size=2, k=2, capacity_factor=2.0) + + def forward(self, x): + output, _, _ = self.moe(self.dense(x)) + return output.pow(2).mean() + + +class TestClipFp32GradientsThroughEngine(DistributedTest): + """The clip `DeepSpeedEngine.step()` applies in an fp32 run without ZeRO, with experts on two ranks. + + `clip_grad_norm_` returns its norm and the engine drops it, so the norm is read from its effect. With + plain SGD the update is `lr * clip_coef * grad`, and the coefficient has to come from the norm of every + gradient with each expert counted once, not from the ranks' own norms averaged. + """ + + world_size = 2 + + def test_the_update_is_clipped_by_the_global_norm(self): + hidden_dim, max_norm, lr = 8, 1e-3, 1.0 + torch.manual_seed(0) + model = _OneMoELayer(hidden_dim) + param_group = {'params': list(model.parameters()), 'name': 'all'} + optimizer = torch.optim.SGD(split_params_into_different_moe_groups_for_optimizer(param_group), lr=lr) + config = { + "train_micro_batch_size_per_gpu": 2, + "gradient_clipping": max_norm, + "zero_optimization": { + "stage": 0 + } + } + engine, _, _, _ = deepspeed.initialize(config=config, model=model, optimizer=optimizer) + + # A different loss scale on each rank gives the two experts gradients of different size. + rank = dist.get_rank() + generator = torch.Generator().manual_seed(100 + rank) + x = torch.randn(2, 4, hidden_dim, generator=generator).to(engine.device) + engine.backward(engine(x) * (1 + 3 * rank)) + + params = list(engine.module.parameters()) + dense = [p for p in params if not is_moe_param(p)] + experts = [p for p in params if is_moe_param(p)] + assert dense and experts + dense_sq = torch.stack([p.grad.float().pow(2).sum() for p in dense]).sum() + expert_sq = torch.stack([p.grad.float().pow(2).sum() for p in experts]).sum() + + # The oracle counts a dense gradient once because it is the same on both ranks, and an expert once + # because each rank owns a different one. + dense_on_every_rank = dense_sq.clone() + dist.all_reduce(dense_on_every_rank, op=dist.ReduceOp.MIN) + assert torch.equal(dense_on_every_rank, dense_sq), "dense gradients should match across ranks" + all_experts_sq = expert_sq.clone() + dist.all_reduce(all_experts_sq) + true_norm = (dense_sq + all_experts_sq).sqrt().item() + + # Averaging the ranks' own norms is the answer this test exists to rule out, so it has to differ. + own_norm = (dense_sq + expert_sq).sqrt() + dist.all_reduce(own_norm) + averaged_norm = own_norm.item() / dist.get_world_size() + assert true_norm > 5 * max_norm, "the clip must be active for this to say anything" + assert abs(averaged_norm - true_norm) > 0.05 * true_norm, "the two norms must be distinguishable" + + clip_coef = max_norm / (true_norm + 1e-6) + before = [p.detach().clone() for p in params] + grads = [p.grad.detach().clone() for p in params] + engine.step() + + for p, start, grad in zip(params, before, grads): + torch.testing.assert_close(p.detach(), start - lr * clip_coef * grad, rtol=1e-4, atol=1e-6) + + class TestClipGradNormPNorm(DistributedTest): # world_size 1 so this runs wherever the suite runs; the bug is in the per-rank # recombination of the norms, which is independent of the group size.