diff --git a/src/transformers/distributed/configuration_utils.py b/src/transformers/distributed/configuration_utils.py index c84955bbe816..7eec6c4130a1 100644 --- a/src/transformers/distributed/configuration_utils.py +++ b/src/transformers/distributed/configuration_utils.py @@ -193,10 +193,10 @@ def _validate_mesh_config(self): if self.cp_size > 1: if not is_torch_greater_or_equal("2.10"): raise OSError("Context parallelism requires `torch>=2.10` (`DeviceMesh._unflatten`).") - if unsupported := [name for name in ("pp_size", "ep_size") if getattr(self, name) > 1]: + if self.pp_size > 1: raise ValueError( - f"Context parallelism (`cp_size={self.cp_size}`) is not supported yet with `{unsupported[0]}`. " - "Use `DistributedConfig(fsdp_size=N, tp_size=T, cp_size=M)`." + f"Context parallelism (`cp_size={self.cp_size}`) is not supported yet with pipeline parallelism " + f"(`pp_size={self.pp_size}`)." ) if self.dp_size > 1 and self.pp_size > 1: diff --git a/src/transformers/distributed/fsdp.py b/src/transformers/distributed/fsdp.py index 57065898e057..43b23c2391ff 100644 --- a/src/transformers/distributed/fsdp.py +++ b/src/transformers/distributed/fsdp.py @@ -202,10 +202,10 @@ def apply_fully_sharded_data_parallelism( for module_name, module in model.named_modules(): if ep_plan.get(module_name) == "ep_dispatch_experts": fully_shard(module, mesh=expert_mesh, reshard_after_forward=True, **fsdp_policy_kwargs) - # an EP group holds ep_size / tp_size distinct batches + # an EP group holds ep_size / tp_size distinct batches, or cp slices of them under context parallelism # the efsdp reduce the sums on efsdp_size copies of that expert. - # The expert gradient therefore covers ep_size / tp_size * efsdp_size = fsdp_size batches - # not efsdp_size batches (the relationship is efsdp_size = fsdp_size * tp_size / ep_size) + # The expert gradient therefore covers ep_size / tp_size * efsdp_size = fsdp_size * cp_size batches + # not efsdp_size batches (the relationship is efsdp_size = fsdp_size * cp_size * tp_size / ep_size) # With HSDP, the all-reduce across replicas adds dp_replicate_size times more: dp_size batches. module.set_gradient_divide_factor(float(distributed_config.dp_size)) # FSDP2 applies a divide factor other than the group size through PREMUL_SUM, which non-NCCL backends diff --git a/tests/context_parallel/test_context_parallel.py b/tests/context_parallel/test_context_parallel.py index c759705c4ba8..7c8c1e5efdc6 100644 --- a/tests/context_parallel/test_context_parallel.py +++ b/tests/context_parallel/test_context_parallel.py @@ -33,6 +33,7 @@ LlamaForSequenceClassification, LlamaModel, MistralConfig, + Qwen3MoeConfig, ) from transformers.distributed import DistributedConfig from transformers.distributed.context_parallel import ( @@ -576,6 +577,9 @@ def test_a_registered_attention_style_is_a_context_parallel_style(self): CP_RTOL = 1e-5 CP_ATOL = 1e-5 +# The all-reduce plan of `Qwen3Moe`, which overrides the token-dispatch rules of its own `ep_plan`: the router masks the +# experts of the other ranks and the partial expert outputs are all-reduced. +ALL_REDUCE_EP_PLAN = {"model.layers.*.mlp.gate": "ep_router", "model.layers.*.mlp.experts": "moe_tp_experts"} def _tiny_batch(config, batch_size, seq_len, packed=False, seed=0): @@ -675,6 +679,7 @@ def _compare_with_unsharded(init_dir, batch, distributed_config, keep=None): continue grad = _average_over_group(grad, groups, group_size) torch.testing.assert_close(grad, expected, rtol=CP_RTOL, atol=CP_ATOL, msg=f"Gradient mismatch for {name}") + return model def _test_matches_unsharded_with_tp(rank, tp_size, cp_size, fsdp_size, packed): @@ -754,6 +759,64 @@ def _test_padded(rank, cp_size, config=None, packed=False): dist.barrier() +def _tiny_moe_config(**kwargs): + """A `Qwen3MoeConfig` small enough for CPU tests, whose default `ep_plan` dispatches the tokens to the experts.""" + return Qwen3MoeConfig( + **{ + "hidden_size": 32, + "moe_intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "vocab_size": 99, + "num_experts": 8, + "num_experts_per_tok": 2, + "attn_implementation": "sdpa", + **kwargs, + } + ) + + +def _test_matches_unsharded_with_ep(rank, cp_size, fsdp_size, ep_size, tp_size=1, dp_replicate_size=1): + """ + Token dispatch with context parallelism: every rank sends the tokens of its own slice to the owners of the experts + they selected, and the experts' gradients reduce over the ranks that hold a copy of them. + """ + config = _tiny_moe_config() + batch_size, seq_len = 2, 8 + with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir: + rows = batch_size * fsdp_size * dp_replicate_size + input_ids, _ = _tiny_batch(config, rows, seq_len) + distributed_config = DistributedConfig( + cp_size=cp_size, + fsdp_size=fsdp_size, + ep_size=ep_size, + tp_size=tp_size, + dp_replicate_size=dp_replicate_size, + ) + batch = {"input_ids": input_ids, "labels": input_ids.clone(), "use_cache": False} + model = _compare_with_unsharded(init_dir, batch, distributed_config) + assert model.model.layers[0].mlp.experts.num_experts == config.num_experts // ep_size + dist.barrier() + + +def _test_matches_unsharded_with_all_reduce_ep(rank, cp_size, ep_size, tp_size): + """ + All-reduce experts need every rank of an ep group to see the same tokens: with `ep_size == tp_size` that group is + the tp group, whose ranks hold the same cp slice. + """ + config = _tiny_moe_config() + batch_size, seq_len = 2, 8 + with _deterministic_init_model_dir(rank, config, torch.float32) as init_dir: + input_ids, _ = _tiny_batch(config, batch_size, seq_len) + distributed_config = DistributedConfig( + cp_size=cp_size, ep_size=ep_size, tp_size=tp_size, ep_plan=ALL_REDUCE_EP_PLAN + ) + batch = {"input_ids": input_ids, "labels": input_ids.clone(), "use_cache": False} + _compare_with_unsharded(init_dir, batch, distributed_config) + dist.barrier() + + def _test_inputs_embeds(rank, cp_size): """`inputs_embeds` is sharded along its sequence dim, not along its last, hidden one.""" config = _tiny_config() @@ -847,7 +910,7 @@ def _test_save_load_distributed_checkpoint(rank, cp_size): @unittest.skipUnless(is_torch_greater_or_equal("2.10"), "Context parallelism requires torch>=2.10") class TestContextParallelConfig(unittest.TestCase): def test_refuses_pipeline_parallelism(self): - with self.assertRaisesRegex(ValueError, "not supported yet with `pp_size`"): + with self.assertRaisesRegex(ValueError, "not supported yet with pipeline parallelism"): DistributedConfig(cp_size=2, pp_size=2) def test_refuses_a_model_without_a_cp_plan(self): @@ -945,6 +1008,29 @@ def test_matches_unsharded_with_tensor_parallelism(self): def test_tensor_parallelism_and_context_parallelism_compose(self): run_distributed(world_size=4, backend="gloo")(_test_tp_and_cp_compose)(2, 2) + def test_matches_unsharded_with_expert_parallelism(self): + run_distributed(world_size=2, backend="gloo")(_test_matches_unsharded_with_ep)(2, 1, 2) + + def test_matches_unsharded_with_fsdp_and_expert_parallelism(self): + # ep_size = fsdp_size * cp_size, so efsdp_size is 1 and every expert lives whole on a single rank + run_distributed(world_size=4, backend="gloo")(_test_matches_unsharded_with_ep)(2, 2, 4) + + def test_matches_unsharded_with_expert_fsdp(self): + # efsdp_size = fsdp_size * cp_size * tp_size / ep_size = 2: pairs of ranks hold the same experts, sharded + # between them + run_distributed(world_size=4, backend="gloo")(_test_matches_unsharded_with_ep)(2, 2, 2) + + def test_matches_unsharded_with_tensor_parallelism_and_expert_parallelism(self): + # ep_size = 2 * tp_size puts both cp slices of the batch, two ranks each, in the same ep group + run_distributed(world_size=4, backend="gloo")(_test_matches_unsharded_with_ep)(2, 1, 4, 2) + + def test_matches_unsharded_with_hsdp_and_expert_parallelism(self): + # the expert gradient reduces over efsdp and then over the two replicas: dp_size = 2 * 1 * 2 ranks + run_distributed(world_size=4, backend="gloo")(_test_matches_unsharded_with_ep)(2, 1, 2, 1, 2) + + def test_matches_unsharded_with_all_reduce_expert_parallelism(self): + run_distributed(world_size=4, backend="gloo")(_test_matches_unsharded_with_all_reduce_ep)(2, 2, 2) + def test_inputs_embeds(self): run_distributed(world_size=2, backend="gloo")(_test_inputs_embeds)(2)