Repository: huggingface/transformers
File: src/transformers/trainer.py (Python)

Code before:

        Returns:
            The loss of the model along with its output if return_outputs was set to True
        """

        # DeepSpeed SP automatically injects shift_labels into inputs (pre-shifted labels for SP).
        # The model's forward pass receives shift_labels via **kwargs and passes it to the loss function.
        # Both standard transformer models and Liger-patched models handle shift_labels correctly,
        # so we can directly use the computed loss from the model output.
        # See: https://huggingface.co/docs/accelerate/en/concept_guides/sequence_parallelism
        if "labels" not in inputs and "shift_labels" in inputs:
            # DeepSpeed SP Dataloader removes "labels" but we need it, otherwise, we won't compute the loss.
            inputs["labels"] = inputs["shift_labels"]
        outputs = model(**inputs)
        loss = outputs.loss

        # Use torch device mesh when available; fallback to DeepSpeed SP groups only for Ulysses SP
        if self.accelerator.torch_device_mesh is not None:
            sp_group = self.accelerator.torch_device_mesh["sp"].get_group()
            sp_world_size = pc.sp_size
        elif pc.sp_backend == "deepspeed" and pc.sp_size > 1:
            from deepspeed.utils import groups

            sp_group = groups._get_sequence_parallel_group()
            sp_world_size = groups._get_sequence_parallel_world_size()
        else:
            raise ValueError(
                "Sequence parallelism is enabled but no SP process group is available. "
                "Ensure torch_device_mesh is initialized or sp_backend='deepspeed' with sp_size > 1."
            )
        # differentiable weighted per-shard-loss aggregation across ranks
        losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)
        # special dealing with SFT that has prompt tokens that aren't used in loss computation
        good_tokens = (inputs["shift_labels"] != -100).view(-1).sum()
        good_tokens_per_rank = torch.distributed.nn.functional.all_gather(good_tokens, group=sp_group)
        # Skip ranks with zero valid tokens
        total_loss = sum(
            losses_per_rank[rank] * good_tokens_per_rank[rank]
            for rank in range(sp_world_size)
            if good_tokens_per_rank[rank] > 0
        )
        total_good_tokens = sum(good_tokens_per_rank)
        loss = total_loss / max(total_good_tokens, 1)

        return (loss, outputs) if return_outputs else loss

    def is_local_process_zero(self) -> bool:
        """
        Whether or not this process is the local (e.g., on one machine if training in a distributed fashion on several
        machines) main process.
        """

Review comment (bug):
Maybe switch the order so that we first check if we are using deepspeed ? I'm afraid of silent bug where `torch_device_mesh` is actually defined. 

Code after:
        losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)
        # special dealing with SFT that has prompt tokens that aren't used in loss computation
        good_tokens = (inputs["shift_labels"] != -100).view(-1).sum()
        good_tokens_per_rank = torch.distributed.nn.functional.all_gather(good_tokens, group=sp_group)
        # Skip ranks with zero valid tokens
        total_loss = sum(
            losses_per_rank[rank] * good_tokens_per_rank[rank]
            for rank in range(sp_world_size)
            if good_tokens_per_rank[rank] > 0
        )
        total_good_tokens = sum(good_tokens_per_rank)
        loss = total_loss / max(total_good_tokens, 1)

        return (loss, outputs) if return_outputs else loss

    def is_local_process_zero(self) -> bool:
        """
        Whether or not this process is the local (e.g., on one machine if training in a distributed fashion on several
        machines) main process.
        """
        return self.args.local_process_index == 0

    def is_world_process_zero(self) -> bool:
        """
        Whether or not this process is the global main process (when training in a distributed fashion on several
        machines, this is only going to be `True` for one process).
        """
        # Special case for SageMaker ModelParallel since there process_index is dp_process_index, not the global
        # process index.
        if is_sagemaker_mp_enabled():
            return smp.rank() == 0
        else:
            return self.args.process_index == 0

    def save_model(self, output_dir: str | None = None, _internal_call: bool = False):
        """
        Will save the model, so you can reload it using `from_pretrained()`.

        Will only save from the main process.
        """

        if output_dir is None:
            output_dir = self.args.output_dir

        if is_torch_xla_available():
            self._save_tpu(output_dir)
        elif is_sagemaker_mp_enabled():
            # Calling the state_dict needs to be done on the wrapped model and on all processes.
            os.makedirs(output_dir, exist_ok=True)
            state_dict = self.model_wrapped.state_dict()
            if self.args.should_save: