diff --git a/swift/sequence_parallel/sequence_parallel.py b/swift/sequence_parallel/sequence_parallel.py index 4b27488721..8655a32067 100644 --- a/swift/sequence_parallel/sequence_parallel.py +++ b/swift/sequence_parallel/sequence_parallel.py @@ -332,7 +332,12 @@ def moe_aux_loss_hook(module, args, kwargs, output): _gathered_logits.append(_bs_logits) router_logits = torch.stack(_gathered_logits, dim=0) if self.real_position_ids is not None: - router_logits = router_logits[:, :, :self.real_position_ids.shape[1], :] + if self.rp_world_size > 1: + position_ids = self.pad( + self.real_position_ids, padding_value=-1, position_ids=self.real_position_ids) + router_logits = router_logits[:, :, position_ids[0] >= 0, :] + else: + router_logits = router_logits[:, :, :self.real_position_ids.shape[1], :] output['router_logits'] = tuple( [logit.reshape(-1, logit.shape[-1]) for logit in router_logits.split(1, dim=1)]) return output diff --git a/tests/sequence_parallel/test_moe_aux_loss.py b/tests/sequence_parallel/test_moe_aux_loss.py new file mode 100644 index 0000000000..a1c86c5533 --- /dev/null +++ b/tests/sequence_parallel/test_moe_aux_loss.py @@ -0,0 +1,62 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from datetime import timedelta +from torch.distributed import init_device_mesh +from transformers.modeling_outputs import MoeModelOutputWithPast +from transformers.models.qwen3_moe.modeling_qwen3_moe import load_balancing_loss_func + +from swift.sequence_parallel import sequence_parallel + + +class RouterOutputModel(torch.nn.Module): + + def forward(self, router_logits, attention_mask=None): + return MoeModelOutputWithPast(router_logits=router_logits) + + +def _check_moe_aux_loss(rank, rendezvous, ring_size, sequence_size): + world_size = ring_size * sequence_size + dist.init_process_group( + 'gloo', init_method=rendezvous, rank=rank, world_size=world_size, timeout=timedelta(seconds=60)) + try: + sp = sequence_parallel + sp.world_size = world_size + sp.rp_world_size = ring_size + sp.sp_world_size = sequence_size + sp.device_mesh = init_device_mesh( + 'cpu', (1, ring_size, sequence_size), mesh_dim_names=('data', 'ring', 'sequence')) + model = RouterOutputModel() + sp._prepare_moe_aux_loss(model) + # Internal padding, tail-only padding, aligned sequences, and a one-token sample. + for lengths in ([3, 5], [8, 3], [8, 8], [1, 3, 5], [5]): + positions = torch.cat([torch.arange(length) for length in lengths]).unsqueeze(0) + sp.extra_kwargs['text_position_ids'] = positions + padded_positions = sp.pad(positions, padding_value=-1, position_ids=positions) + reference = torch.randn(2, sum(lengths), 4, generator=torch.Generator().manual_seed(42)) + reference.requires_grad_() + padded = sp.pad(reference.detach(), padding_value=0, position_ids=positions) + local = sp.split(padded, dim=1, position_ids=padded_positions).detach().requires_grad_() + output = model(tuple(local.unbind(0)), attention_mask=None) + actual = torch.stack(output.router_logits) + torch.testing.assert_close(actual, reference) + actual_loss = load_balancing_loss_func(output.router_logits, num_experts=4, top_k=2) + expected_loss = load_balancing_loss_func(tuple(reference.unbind(0)), num_experts=4, top_k=2) + torch.testing.assert_close(actual_loss, expected_loss) + actual_loss.backward() + expected_loss.backward() + padded_grad = sp.pad(reference.grad, padding_value=0, position_ids=positions) + expected_grad = sp.split(padded_grad, dim=1, position_ids=padded_positions) * world_size + # GatherLoss compensates for the subsequent distributed gradient average. + torch.testing.assert_close(local.grad, expected_grad) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not dist.is_available() or not dist.is_gloo_available(), reason='Gloo is not available') +@pytest.mark.parametrize(('ring_size', 'sequence_size'), [(2, 1), (2, 2), (1, 2), (1, 1)]) +def test_moe_aux_loss_excludes_sequence_parallel_padding(tmp_path, ring_size, sequence_size): + rendezvous = (tmp_path / 'rendezvous').as_uri() + mp.spawn(_check_moe_aux_loss, args=(rendezvous, ring_size, sequence_size), nprocs=ring_size * sequence_size)