diff --git a/swift/loss/causal_lm.py b/swift/loss/causal_lm.py index 78ac52dfd8..6bfe05a6a8 100644 --- a/swift/loss/causal_lm.py +++ b/swift/loss/causal_lm.py @@ -1,10 +1,26 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +import torch.distributed as dist + from .base import BaseLoss class CustomCrossEntropyLoss(BaseLoss): def __call__(self, outputs, labels, *, num_items_in_batch=None, loss_scale=None, **kwargs): + if self.trainer is not None and self.trainer.template.sequence_parallel_size > 1: + # The trainer already shifted, gathered and weighted the per-token loss. + loss = outputs.loss.sum() + if num_items_in_batch is None: + num_items_in_batch = (labels != -100).sum() + dist.all_reduce(num_items_in_batch, op=dist.ReduceOp.SUM) + # Counts inferred inside this callback are not visible to the trainer's rescaling step. + if (getattr(self.trainer.args, 'average_tokens_across_devices', False) + and self.trainer.model_accepts_loss_kwargs): + loss = loss * self.trainer.accelerator.num_processes + if not self.trainer.model.training: + loss = loss / self.trainer.template.sequence_parallel_size + return loss / num_items_in_batch + from swift.trainers import per_token_loss_func token_loss = per_token_loss_func(outputs, labels) if loss_scale is not None: diff --git a/tests/sequence_parallel/test_custom_cross_entropy.py b/tests/sequence_parallel/test_custom_cross_entropy.py new file mode 100644 index 0000000000..20a06f941c --- /dev/null +++ b/tests/sequence_parallel/test_custom_cross_entropy.py @@ -0,0 +1,119 @@ +# 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 itertools import product +from torch.distributed import init_device_mesh +from transformers.modeling_outputs import CausalLMOutputWithPast +from types import SimpleNamespace + +from swift.loss.causal_lm import CustomCrossEntropyLoss +from swift.sequence_parallel import sequence_parallel +from swift.trainers.seq2seq_trainer import Seq2SeqTrainer + + +class LocalLogitsModel(torch.nn.Module): + + def __init__(self, logits): + super().__init__() + self.logits = torch.nn.Parameter(logits) + self.model_info = SimpleNamespace(is_moe_model=False) + + def forward(self, **kwargs): + return CausalLMOutputWithPast(logits=self.logits) + + +def _check_loss(rank, rendezvous, ring_size, sequence_size, data_size): + parallel_size = ring_size * sequence_size + world_size = data_size * parallel_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 = parallel_size + sp.rp_world_size = ring_size + sp.sp_world_size = sequence_size + sp.device_mesh = init_device_mesh( + 'cpu', (data_size, ring_size, sequence_size), mesh_dim_names=('data', 'ring', 'sequence')) + for lengths in ([8], [5], [3, 5]): + positions = torch.cat([torch.arange(length) for length in lengths]).unsqueeze(0) + labels = (torch.arange(sum(lengths)).unsqueeze(0) % 6) + 1 + labels[positions < 2] = -100 + if rank // parallel_size > 0: + labels[:, -1] = -100 + padded_positions = sp.pad(positions, padding_value=-1, position_ids=positions) + logits = torch.randn( + 1, sum(lengths), 8, generator=torch.Generator().manual_seed(42 + rank // parallel_size)) + local_logits = sp.split(sp.pad(logits, 0, positions), 1, padded_positions) + for scale_mode in ('none', 'weighted', 'zero'): + for denominator, average_tokens, training in product((None, 20), (False, True), (False, True)): + results = [] + for custom in (False, True): + inputs = {'input_ids': labels.clamp_min(0), 'labels': labels.clone(), 'position_ids': positions} + if scale_mode != 'none': + weights = torch.arange(sum(lengths)).unsqueeze(0).float() / 3 + inputs['loss_scale'] = weights if scale_mode == 'weighted' else torch.zeros_like(weights) + sp.prepare_inputs(inputs) + model = LocalLogitsModel(local_logits.clone()) + model.train(training) + template = SimpleNamespace( + sequence_parallel_size=parallel_size, + padding_free=True, + compute_sft_loss=lambda model, inputs, **kwargs: model(**inputs)) + trainer = SimpleNamespace( + template=template, + model=model, + label_smoother=None, + model_accepts_loss_kwargs=True, + accelerator=SimpleNamespace(unwrap_model=lambda model: model, num_processes=world_size), + _compute_acc=lambda *args, **kwargs: None, + args=SimpleNamespace( + use_liger_kernel=False, + past_index=-1, + enable_dft_loss=False, + enable_channel_loss=False, + average_tokens_across_devices=average_tokens, + tuner_backend='peft', + acc_strategy='token')) + if custom: + inputs['compute_loss_func'] = CustomCrossEntropyLoss(None, trainer) + loss = Seq2SeqTrainer.compute_loss(trainer, model, inputs, num_items_in_batch=denominator) + loss.backward() + results.append((loss.detach(), model.logits.grad)) + for actual, expected in zip(results[1], results[0]): + torch.testing.assert_close(actual, expected) + reference_logits = logits.clone().requires_grad_() + token_loss = torch.nn.functional.cross_entropy( + reference_logits.reshape(-1, 8), labels.roll(-1, dims=1).reshape(-1), reduction='none') + if scale_mode != 'none': + weights = torch.arange(sum(lengths)).float() / 3 + if scale_mode == 'zero': + weights.zero_() + token_loss = token_loss * weights.roll(-1) + count = denominator + if count is None: + count = (labels != -100).sum() + dist.all_reduce(count) + count = count / parallel_size + reference_loss = token_loss.sum() / count + if average_tokens: + reference_loss = reference_loss * world_size + if not training: + reference_loss = reference_loss / parallel_size + reference_loss.backward() + torch.testing.assert_close(results[1][0], reference_loss) + expected_grad = sp.split(sp.pad(reference_logits.grad, 0, positions), 1, padded_positions) + torch.testing.assert_close(results[1][1], expected_grad * parallel_size) + 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', 'data_size'), [(1, 2, 1), (2, 1, 1), (2, 2, 1), (1, 2, 2)]) +def test_explicit_cross_entropy_matches_default_sequence_parallel_loss(tmp_path, ring_size, sequence_size, data_size): + mp.spawn( + _check_loss, + args=((tmp_path / 'rendezvous').as_uri(), ring_size, sequence_size, data_size), + nprocs=ring_size * sequence_size * data_size)