# Copyright (c) ModelScope Contributors. All rights reserved.
import inspect
import math
import torch
import torch.distributed as dist
from functools import lru_cache, partial
from torch.distributed import init_device_mesh
from transformers import PreTrainedTokenizer
from types import SimpleNamespace
from typing import Optional

from swift.utils import HfConfigFactory, get_cu_seqlens_from_position_ids, get_device, get_dist_setting
from .ulysses import DistributedAttention
from .zigzag_ring_attn import zigzag_ring_flash_attn_varlen_func


@lru_cache(maxsize=None)
def get_signature_parameters(fn):
    try:
        return inspect.signature(fn).parameters
    except (TypeError, ValueError):
        return None


def has_signature_parameter(fn, parameter: str) -> bool:
    parameters = get_signature_parameters(fn)
    return parameters is not None and parameter in parameters


def call_with_supported_kwargs(fn, *args, **kwargs):
    parameters = get_signature_parameters(fn)
    if parameters is None:
        return fn(*args, **kwargs)
    if not any(param.kind == inspect.Parameter.VAR_KEYWORD for param in parameters.values()):
        kwargs = {key: value for key, value in kwargs.items() if key in parameters}
    return fn(*args, **kwargs)


def _call_create_causal_mask(fn, config, input_embeds, attention_mask, cache_position_or_past_key_values, *args,
                             **kwargs):
    if has_signature_parameter(fn, 'cache_position'):
        return call_with_supported_kwargs(
            fn,
            config,
            input_embeds,
            attention_mask,
            cache_position_or_past_key_values,
            *args,
            **kwargs,
        )
    if cache_position_or_past_key_values is None and 'past_key_values' in kwargs:
        return call_with_supported_kwargs(fn, config, input_embeds, attention_mask, *args, **kwargs)
    return call_with_supported_kwargs(
        fn,
        config,
        input_embeds,
        attention_mask,
        cache_position_or_past_key_values,
        *args,
        **kwargs,
    )


class SequenceParallel:

    _global_inited: bool = False

    def __init__(self):
        self.sp_world_size = None
        self.dp_world_size = None
        self.rp_world_size = None
        self.world_size = None
        self.model_dtype = None
        self.tokenizer = None
        self.device_mesh = None
        self.num_heads = None
        self.causal_mask_func = None
        self.extra_kwargs = {}

    @property
    def real_position_ids(self) -> torch.Tensor:
        """The real position ids, this is different from the position_ids in mrope"""
        return self.extra_kwargs.get('text_position_ids')

    def _prepare_flash_attn(self, base_model: torch.nn.Module):
        try:
            from transformers import masking_utils

            _origin_flash_attention_mask = masking_utils.flash_attention_mask

            def flash_attention_mask(*args, **kwargs):
                if self.world_size == 1:
                    return _origin_flash_attention_mask(*args, **kwargs)
                attention_mask = kwargs.get('attention_mask')
                if attention_mask is not None:
                    if attention_mask.all():
                        attention_mask = None

                return attention_mask

            masking_utils.flash_attention_mask = flash_attention_mask
            masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['flash_attention_2'] = flash_attention_mask

            def sdpa_mask(*args, **kwargs):
                if self.world_size == 1:
                    return masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa_origin'](*args, **kwargs)
                if 'cache_position' in kwargs:
                    device = kwargs['cache_position'].device
                else:
                    # transformers>=5.4.0
                    device = kwargs['device']
                cache_position = self.real_position_ids[0]
                cache_position = self.pad(cache_position, padding_value=-1, position_ids=self.real_position_ids, dim=0)
                cache_position = torch.arange(0, cache_position.shape[0], device=device)
                kwargs['kv_length'] = cache_position.shape[0]
                if 'cache_position' in kwargs:
                    kwargs['cache_position'] = cache_position
                else:
                    kwargs['q_length'] = kwargs['kv_length']
                return masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa_origin'](*args, **kwargs)

            masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping[
                'sdpa_origin'] = masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa']
            masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa'] = sdpa_mask

            def create_causal_mask(config,
                                   input_embeds,
                                   attention_mask,
                                   cache_position_or_past_key_values=None,
                                   *args,
                                   **kwargs):
                if self.world_size == 1:
                    return _call_create_causal_mask(masking_utils.origin_create_causal_mask, config, input_embeds,
                                                    attention_mask, cache_position_or_past_key_values, *args, **kwargs)
                input_embeds = torch.ones(
                    (input_embeds.shape[0], input_embeds.shape[1] * self.sp_world_size, input_embeds.shape[2]),
                    dtype=input_embeds.dtype,
                    device=input_embeds.device)
                if has_signature_parameter(masking_utils.origin_create_causal_mask, 'cache_position'):
                    cache_position_or_past_key_values = torch.arange(
                        0,
                        input_embeds.shape[1],
                        device=input_embeds.device,
                    )
                return _call_create_causal_mask(masking_utils.origin_create_causal_mask, config, input_embeds,
                                                attention_mask, cache_position_or_past_key_values, *args, **kwargs)

            masking_utils.origin_create_causal_mask = masking_utils.create_causal_mask
            masking_utils.create_causal_mask = create_causal_mask
        except ImportError:
            pass

        if hasattr(base_model, 'language_model'):
            text_model = base_model.language_model
        else:
            text_model = base_model

        from transformers.modeling_flash_attention_utils import is_flash_attn_available
        if is_flash_attn_available():
            # TODO this works for multi-modal models like qwen2.5-vl
            # SDPA is not supported, because we need to copy the code to our project, which will bring
            # more works for maintaining.
            from transformers import modeling_flash_attention_utils
            from transformers.modeling_flash_attention_utils import _flash_attention_forward
            _distributed_flash_attention = DistributedAttention(_flash_attention_forward, self)

            modeling_flash_attention_utils._flash_attention_forward_origin = _flash_attention_forward

            def flash_attention_forward(query_states: torch.Tensor, key_states: torch.Tensor,
                                        value_states: torch.Tensor, attention_mask: Optional[torch.Tensor], q_len,
                                        *args, **kwargs):
                if self.world_size == 1:
                    return _flash_attention_forward(query_states, key_states, value_states, attention_mask, q_len,
                                                    *args, **kwargs)
                return _distributed_flash_attention(query_states, key_states, value_states, attention_mask,
                                                    q_len * self.sp_world_size, *args, **kwargs)

            modeling_flash_attention_utils._flash_attention_forward = flash_attention_forward

        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS

        def local_flash_attn(module: torch.nn.Module, query_states, key_states, value_states, attention_mask, *args,
                             dist_attn, **kwargs):
            if self.world_size == 1 or module.__class__ not in [m.__class__ for m in text_model.modules()]:
                return ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'](module, query_states, key_states,
                                                                           value_states, attention_mask, *args,
                                                                           **kwargs)
            if dist_attn.local_attn is None:

                def _attention(query, key, value, *args, **kwargs):
                    query = query.transpose(1, 2)
                    key = key.transpose(1, 2)
                    value = value.transpose(1, 2)
                    if self.rp_world_size is not None and self.rp_world_size > 1:
                        position_ids = kwargs['position_ids']
                        cu_seqlens = get_cu_seqlens_from_position_ids(position_ids)
                        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
                        position_ids = self._split_packed(position_ids, cu_seqlens)
                        mask = position_ids != -1
                        query = query.transpose(1, 2)
                        key = key.transpose(1, 2)
                        value = value.transpose(1, 2)
                        # this is important
                        # the length is correct, but the value is wrong, by mask qkv, we do:
                        # mask the padded values of q and v to zero
                        # mask the padded values of k to -1e5
                        # to keep the attention correct
                        query, key, value = self._mask_qkv(query, key, value, mask)
                        output = zigzag_ring_flash_attn_varlen_func(
                            query,
                            key,
                            value,
                            cu_seqlens=cu_seqlens,
                            max_seqlen=max_seqlen,
                            causal=module.is_causal,
                            dropout_p=kwargs.get('dropout', 0.0),
                            softmax_scale=kwargs.get('scaling', 0.0),
                            window_size=kwargs.get('sliding_window') or (-1, -1),
                            group=self.rp_group)
                        return output
                    else:
                        if 'cu_seq_lens_q' in kwargs:
                            position_ids = kwargs.get('position_ids')
                            if self.real_position_ids is not None:
                                position_ids = self.real_position_ids
                            position_ids = self.pad(position_ids, padding_value=-1, position_ids=position_ids)
                            cu_seqlens = get_cu_seqlens_from_position_ids(position_ids)
                            max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
                            assert query.shape[2] == cu_seqlens[-1]
                            kwargs['cu_seq_lens_q'] = cu_seqlens
                            kwargs['cu_seq_lens_k'] = cu_seqlens
                            kwargs['max_length_q'] = max_seqlen
                            kwargs['max_length_k'] = max_seqlen
                        return ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'](module, query, key, value, *args,
                                                                                   **kwargs)[0]

                dist_attn.local_attn = _attention

            return dist_attn(
                query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask,
                *args, **kwargs), None

        def local_sdpa_attn(module: torch.nn.Module, query_states, key_states, value_states, attention_mask, *args,
                            dist_attn, **kwargs):
            # Bypass SP logic when world_size == 1 (SP disabled) or module not in text_model
            if self.world_size == 1 or module.__class__ not in [m.__class__ for m in text_model.modules()]:
                return ALL_ATTENTION_FUNCTIONS['sdpa_origin'](module, query_states, key_states, value_states,
                                                              attention_mask, *args, **kwargs)
            if dist_attn.local_attn is None:

                def _attention(query, key, value, *args, **kwargs):
                    query = query.transpose(1, 2)
                    key = key.transpose(1, 2)
                    value = value.transpose(1, 2)
                    if self.rp_world_size > 1:
                        raise NotImplementedError('SDPA does not support Ring attention.')
                    return ALL_ATTENTION_FUNCTIONS['sdpa_origin'](module, query, key, value, *args, **kwargs)[0]

                dist_attn.local_attn = _attention
            return dist_attn(
                query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask,
                *args, **kwargs), None

        ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'] = ALL_ATTENTION_FUNCTIONS['flash_attention_2']
        ALL_ATTENTION_FUNCTIONS['sdpa_origin'] = ALL_ATTENTION_FUNCTIONS['sdpa']
        ALL_ATTENTION_FUNCTIONS['flash_attention_2'] = partial(
            local_flash_attn, dist_attn=DistributedAttention(None, self))
        ALL_ATTENTION_FUNCTIONS['sdpa'] = partial(local_sdpa_attn, dist_attn=DistributedAttention(None, self))

    def _prepare_forward_hook(self, base_model: torch.nn.Module):

        def pre_forward_split_hook(_self, args, kwargs):
            if self.world_size == 1:
                return args, kwargs
            input_ids = kwargs.get('input_ids', None)
            inputs_embeds = kwargs.get('inputs_embeds', None)
            position_ids = kwargs['position_ids']
            attention_mask = kwargs.get('attention_mask', None)
            if hasattr(_self, 'language_model'):
                embed_tokens = getattr(_self.language_model, 'embed_tokens', None)
            else:
                embed_tokens = getattr(_self, 'embed_tokens', None)
            input_ids, inputs_embeds, _, position_ids, attention_mask, _, _ = self.pad_and_split_inputs(
                input_ids,
                inputs_embeds,
                None,
                position_ids,
                attention_mask,
                None,
                embed_tokens=embed_tokens,
                real_position_ids=self.real_position_ids)
            kwargs['input_ids'] = input_ids
            kwargs['inputs_embeds'] = inputs_embeds
            kwargs['position_ids'] = position_ids
            kwargs['attention_mask'] = attention_mask
            return args, kwargs

        base_model.register_forward_pre_hook(pre_forward_split_hook, with_kwargs=True)

    def _prepare_moe_aux_loss(self, base_model: torch.nn.Module):
        from .utils import GatherLoss

        def moe_aux_loss_hook(module, args, kwargs, output):
            router_logits = getattr(output, 'router_logits', None)
            if router_logits is None:
                return output

            attention_mask = kwargs['attention_mask']
            if attention_mask is None:
                batch_size = 1
            else:
                batch_size = attention_mask.shape[0]

            assert router_logits[0].shape[0] % batch_size == 0
            seq_len = router_logits[0].shape[0] // batch_size

            _gathered_logits = []
            for i in range(batch_size):
                _slice = slice(i * seq_len, (i + 1) * seq_len)
                _bs_logits = [logit[_slice] for logit in router_logits]
                compute_device = _bs_logits[0].device
                _bs_logits = torch.stack([layer_gate.to(compute_device) for layer_gate in _bs_logits], dim=0)
                _bs_logits, _ = GatherLoss.apply(_bs_logits, None, 1, self.real_position_ids)
                _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], :]
            output['router_logits'] = tuple(
                [logit.reshape(-1, logit.shape[-1]) for logit in router_logits.split(1, dim=1)])
            return output

        base_model.register_forward_hook(moe_aux_loss_hook, with_kwargs=True)

    def prepare(self, sp_size: int, model: torch.nn.Module, tokenizer: PreTrainedTokenizer, padding_free: bool):
        from swift.model import get_llm_model
        self.num_heads = HfConfigFactory.get_config_attr(model.config, 'num_key_value_heads')
        if self.num_heads is None:
            self.num_heads = HfConfigFactory.get_config_attr(model.config, 'num_attention_heads')
        assert self.num_heads is not None, 'Cannot find num_heads config in config.json'
        self.padding_free = padding_free
        self.world_size = sp_size

        llm_model = get_llm_model(model)

        if hasattr(llm_model, 'language_model'):
            if hasattr(llm_model.language_model, '_update_causal_mask'):
                self.causal_mask_func = llm_model.language_model._update_causal_mask
        else:
            if hasattr(llm_model, '_update_causal_mask'):
                self.causal_mask_func = llm_model._update_causal_mask

        if not SequenceParallel._global_inited:
            # these operations are global initializations and patches
            self._init_device_mesh()
            self._prepare_flash_attn(llm_model)
            SequenceParallel._global_inited = True

        self._prepare_forward_hook(llm_model)

        if model.model_info.is_moe_model:
            self._prepare_moe_aux_loss(llm_model)

        self.model_dtype = next(model.parameters()).dtype
        self.tokenizer = tokenizer
        if self.rp_world_size > 1 and not self.padding_free:
            raise NotImplementedError(
                f'The world_size {self.world_size} needs ulysses/ring-attention, which needs --padding_free true')

    def _mask_qkv(self, query, key, value, mask):
        mask = mask.unsqueeze(2).unsqueeze(3)
        query = query * mask
        value = value * mask
        mask = (~mask) * -1e5  # for bf16
        key = key + mask.to(key.dtype)
        return query, key, value

    def pad(self, tensor, padding_value, position_ids=None, dim=1):
        """Pad tensor for sequence parallel"""
        if self.rp_world_size > 1:
            world_size = self.world_size * 2
        else:
            world_size = self.world_size

        def _do_pad(tensor):
            length = tensor.shape[dim]
            pad_num = world_size - (length % world_size)
            if pad_num == 0 or pad_num == world_size:
                return tensor
            if not isinstance(padding_value, torch.Tensor):
                # ids
                pad_shape = ((*tensor.shape[:dim], pad_num, *tensor.shape[dim + 1:]) if dim != -1 else
                             (*tensor.shape[:dim], pad_num))
                pad = torch.full(pad_shape, padding_value, dtype=tensor.dtype, device=tensor.device)
                tensor = torch.cat([tensor, pad], dim=dim)
            else:
                # For embeddings
                tensor = torch.cat([tensor, padding_value.unsqueeze(0).repeat(tensor.shape[0], pad_num, 1)], dim=dim)
            return tensor

        if position_ids is not None and self.rp_world_size > 1:
            cu_seqlens = get_cu_seqlens_from_position_ids(position_ids)
            all_tensors = []
            for i in range(len(cu_seqlens) - 1):
                if dim == 1:
                    sub_tensor = tensor[:, cu_seqlens[i]:cu_seqlens[i + 1]]
                elif dim == -1:
                    sub_tensor = tensor[..., cu_seqlens[i]:cu_seqlens[i + 1]]
                else:
                    raise NotImplementedError()
                all_tensors.append(_do_pad(sub_tensor))
            tensor = torch.cat(all_tensors, dim=dim)

        return _do_pad(tensor)

    def gather(self, local_output, dim: int, position_ids=None):
        """Gather tensor for sequence parallel - reverse of split"""
        if self.world_size == 1:
            return local_output

        if self.rp_world_size > 1:
            input_dim = local_output.dim()
            assert input_dim >= 2

            if position_ids is not None:
                position_ids = self.pad(position_ids, padding_value=-1, position_ids=position_ids)

            # Step 1: Gather from all sequence parallel ranks
            # Each sp_rank has its own piece, we need to gather them first
            gathered_sp = [torch.zeros_like(local_output) for _ in range(self.sp_world_size)]
            torch.distributed.all_gather(gathered_sp, local_output.contiguous(), group=self.sp_group)

            # Concatenate the sp pieces to form the complete chunk for this rp_rank
            rp_chunk = torch.cat(gathered_sp, dim=dim)

            # Step 2: Gather all rp chunks
            gathered_rp = [torch.zeros_like(rp_chunk) for _ in range(self.rp_world_size)]
            torch.distributed.all_gather(gathered_rp, rp_chunk, group=self.rp_group)

            cu_seqlens = get_cu_seqlens_from_position_ids(position_ids)
            all_tensor_length = []
            for i in range(len(cu_seqlens) - 1):
                length = cu_seqlens[i + 1] - cu_seqlens[i]
                padding_length = math.ceil(length / (self.world_size * 2)) * (self.world_size * 2)
                all_tensor_length.append(padding_length)

            full_output = torch.zeros(
                [local_output.shape[0], sum(all_tensor_length), *local_output.shape[2:]], device=local_output.device)
            for idx_rp, rp_tensor in enumerate(gathered_rp):  # rp world size
                # re-group the zigzag to the correct order
                accumulated_length = 0
                for idx_seq, length in enumerate(all_tensor_length):  # sequence number
                    local_length = length // self.rp_world_size
                    local_tensor = rp_tensor[:, accumulated_length:local_length + accumulated_length]
                    chunk_size = local_length // 2
                    left_idx = accumulated_length * self.rp_world_size + idx_rp * chunk_size
                    right_idx = accumulated_length * self.rp_world_size + (idx_rp + 1) * chunk_size
                    full_output[:, left_idx:right_idx] = local_tensor[:, :chunk_size]
                    left_idx = accumulated_length * self.rp_world_size + (2 * self.rp_world_size - idx_rp
                                                                          - 1) * chunk_size
                    right_idx = accumulated_length * self.rp_world_size + (2 * self.rp_world_size - idx_rp) * chunk_size
                    full_output[:, left_idx:right_idx] = local_tensor[:, chunk_size:]
                    accumulated_length += local_length

            return full_output.contiguous()
        else:
            gathered_sp = torch.empty(
                [local_output.shape[0] * self.sp_world_size] + list(local_output.shape[1:]),
                dtype=local_output.dtype,
                device=local_output.device)
            dist.all_gather_into_tensor(gathered_sp, local_output, group=self.sp_group)
            gathered_sp = torch.cat(gathered_sp.split(local_output.shape[0], dim=0), dim=dim)
            return gathered_sp.contiguous()

    def _split_packed(self, value, cu_seqlens, dim=1):
        """Split and re-group in zigzag"""
        local_values = []
        for i in range(len(cu_seqlens) - 1):
            start, end = cu_seqlens[i], cu_seqlens[i + 1]
            if dim == 1:
                sub_value = value[:, start:end]
            elif dim == -1:
                sub_value = value[..., start:end]
            else:
                raise NotImplementedError()
            local_value = sub_value.chunk(2 * self.rp_world_size, dim=dim)
            local_values.extend([
                local_value[self.rp_rank],
                local_value[2 * self.rp_world_size - 1 - self.rp_rank],
            ])
        return torch.cat(local_values, dim=dim).contiguous()

    def split(self, input, dim: int, position_ids=None):
        """Split tensor for sequence parallel"""
        if self.world_size == 1:
            return input

        if self.rp_world_size > 1:
            input_dim = input.dim()
            assert input_dim >= 2
            cu_seqlens = get_cu_seqlens_from_position_ids(position_ids)
            assert torch.all(cu_seqlens % (2 * self.rp_world_size) == 0)
            value_chunks = self._split_packed(input, cu_seqlens, dim=dim)
            local_value = value_chunks.chunk(self.sp_world_size, dim=dim)[self.sp_rank].contiguous()
            return local_value
        else:
            rank = self.sp_rank
            dim_size = input.size(dim)
            assert dim_size % self.sp_world_size == 0, (
                f'The dimension to split ({dim_size}) is not a multiple of '
                f'world size ({self.sp_world_size}), cannot split tensor evenly')

            tensor_list = torch.split(input, dim_size // self.sp_world_size, dim=dim)
            output = tensor_list[rank].contiguous()
            return output

    def pad_and_split_mm_tokens(self, visual_mask, mm_embeds):
        input_ids = self.extra_kwargs['input_ids']
        empty_embeds = torch.empty(
            (input_ids.shape[0], input_ids.shape[1], mm_embeds.shape[-1])).to(mm_embeds.device).to(mm_embeds.dtype)
        empty_embeds[visual_mask] = mm_embeds

        embeds = SimpleNamespace(weight=mm_embeds)

        _, split_input_embeds, _, _, _, _, extra_values = self.pad_and_split_inputs(
            None,
            empty_embeds,
            None,
            None,
            None,
            None,
            embeds,
            self.real_position_ids,
            extra_split_values=[(visual_mask, 0, -1)])
        visual_mask = extra_values[0]
        return visual_mask, split_input_embeds[visual_mask]

    def pad_and_split_inputs(self,
                             input_ids,
                             input_embeds,
                             labels,
                             position_ids,
                             attention_mask,
                             loss_scale,
                             embed_tokens=None,
                             real_position_ids=None,
                             extra_split_values=None):
        """Common implementation for padding and splitting inputs

        When a sequence comes, it will be split into rp_world_size * 2 sub tensors, and group them as the
        zigzag order. So we get rp_world_size tensors, then we split each tensor to sp_world_size ones.
        So, we should first pad the original sequence to the length can be divided by 2 * world_size, then re-group it.

        Only support padding_free for ring-attention, because non-padding_free mode needs another pad/split workflow
        man, that's a lot of work...

        Args:
            input_ids: input_ids
            input_embeds: input_embeds
            labels: labels
            position_ids: position_ids or, position_ids for mrope
            attention_mask: attention_mask
            loss_scale: loss_scale
            embed_tokens: embed_tokens
            real_position_ids: the real position_ids to represent the seq length information
            extra_split_values: List of Tuples for extra split values, e.g.: (tensor, pad_value, split_dim)
        """
        tokenizer = self.tokenizer
        real_position_ids = real_position_ids if real_position_ids is not None else position_ids
        extra_values = []
        batch_size = input_ids.shape[
            0] if input_ids is not None else input_embeds.shape[0] if input_embeds is not None else None
        if real_position_ids is not None and batch_size is not None and real_position_ids.shape[0] == batch_size:
            # TODO clone everytime, but the position_ids is a small tensor
            self.extra_kwargs['text_position_ids'] = real_position_ids.clone()
        if input_ids is not None:
            input_ids = self.pad(input_ids, padding_value=tokenizer.pad_token_id, position_ids=real_position_ids)
            self.extra_kwargs['input_ids'] = input_ids.clone()
        if input_embeds is not None:
            pad_emb = torch.zeros(
                (1, embed_tokens.weight.shape[-1])).to(embed_tokens.weight.device).to(embed_tokens.weight.dtype)
            input_embeds = self.pad(input_embeds, padding_value=pad_emb, position_ids=real_position_ids)
        batch_size = input_ids.shape[
            0] if input_ids is not None else input_embeds.shape[0] if input_embeds is not None else 1
        if position_ids is not None:
            position_ids = self.pad(position_ids, padding_value=-1, position_ids=real_position_ids, dim=-1)
        if labels is not None:
            labels = self.pad(labels, padding_value=-100, position_ids=real_position_ids)
        if loss_scale is not None:
            loss_scale = self.pad(loss_scale, padding_value=0., position_ids=real_position_ids)
        if real_position_ids is not None:
            real_position_ids = self.pad(real_position_ids, padding_value=-1, position_ids=real_position_ids)
        if (input_ids is not None or input_embeds is not None) and batch_size > 1:
            # not padding_free, so not ring-attention
            inputs = input_ids if input_ids is not None else input_embeds
            attn_shape = inputs.shape[1]  # The sequence length
            if attention_mask is None:
                attention_mask = torch.ones_like(real_position_ids)
            # no need position_ids here, because padding_free does not need attention_mask,
            # so this is not ring-attention
            attention_mask = self.pad(attention_mask, padding_value=0)
            cache_position = torch.arange(0, attn_shape, device=inputs.device)
            # pad attention mask to 4d to avoid calculation errors
            if hasattr(self, 'causal_mask_func') and self.causal_mask_func is not None:
                attention_mask = self.causal_mask_func(attention_mask, inputs.to(self.model_dtype), cache_position,
                                                       None, None)
        if extra_split_values is not None:
            for (tensor, pad_value, split_dim) in extra_split_values:
                extra_values.append(
                    self.pad(tensor, padding_value=pad_value, position_ids=real_position_ids, dim=split_dim))
        if input_ids is not None:
            input_ids = self.split(input_ids, dim=1, position_ids=real_position_ids)
        if input_embeds is not None:
            input_embeds = self.split(input_embeds, dim=1, position_ids=real_position_ids)
        if labels is not None:
            labels = torch.roll(labels, shifts=-1, dims=-1)
            labels = self.split(labels, dim=-1, position_ids=real_position_ids)
        if loss_scale is not None:
            loss_scale = torch.roll(loss_scale, shifts=-1, dims=-1)
            loss_scale = self.split(loss_scale, dim=-1, position_ids=real_position_ids)

        if position_ids is not None:
            position_ids = self.split(position_ids, dim=-1, position_ids=real_position_ids)
        if extra_split_values is not None:
            for i in range(len(extra_values)):
                extra_values[i] = self.split(
                    extra_values[i], dim=extra_split_values[i][2], position_ids=real_position_ids)
        return input_ids, input_embeds, labels, position_ids, attention_mask, loss_scale, extra_values

    def _gather_object_dp(self, input_data):
        """Gather object for data parallel"""
        input_data_list = [None] * self.dp_world_size
        dist.all_gather_object(input_data_list, input_data, group=self.dp_group)
        return [x for y in input_data_list for x in y]

    def _init_device_mesh(self):
        """Initialize device mesh for sequence and ring parallel.

        The logic is unified:
        1. Determine the Sequence Parallel (SP) size first based on GCD to satisfy constraints.
        2. Allocate all remaining model parallelism to Ring Parallel (RP).
        """
        _, _, world_size, _ = get_dist_setting()
        self.dp_world_size = world_size // self.world_size
        # SP size is the GCD of num_heads and world_size, guaranteeing it divides both.
        self.sp_world_size = math.gcd(self.num_heads, self.world_size)
        # RP takes the remaining factor so all model-parallel GPUs are used.
        self.rp_world_size = self.world_size // self.sp_world_size

        if self.rp_world_size > 1:
            mesh_shape = (self.dp_world_size, self.rp_world_size, self.sp_world_size)
            mesh_dim_names = ('data', 'ring', 'sequence')
        else:
            mesh_shape = (self.dp_world_size, self.sp_world_size)
            mesh_dim_names = ('data', 'sequence')
        self.device_mesh = init_device_mesh(
            get_device().split(':')[0], mesh_shape=mesh_shape, mesh_dim_names=mesh_dim_names)

    def _dim_group(self, dim_name: str):
        """Return the process group of the given mesh dim, or None if absent."""
        if not self.device_mesh or dim_name not in self.device_mesh.mesh_dim_names:
            return None
        return self.device_mesh[dim_name].get_group()

    def _dim_rank(self, dim_name: str, default: int) -> int:
        """Return the rank within the given mesh dim, or `default` if absent."""
        group = self._dim_group(dim_name)
        return dist.get_rank(group) if group is not None else default

    @property
    def sp_group(self):
        """Return the sequence parallel group"""
        return self._dim_group('sequence')

    def enabled(self):
        return self.world_size is not None and self.world_size > 1

    @property
    def sp_rank(self):
        """Return the sequence parallel rank"""
        return self._dim_rank('sequence', default=0)

    @property
    def dp_group(self):
        """Return the data parallel group"""
        return self._dim_group('data')

    @property
    def dp_rank(self):
        """Return the data parallel rank"""
        return self._dim_rank('data', default=0)

    @property
    def rp_group(self):
        """Return the ring parallel group"""
        return self._dim_group('ring')

    @property
    def rp_rank(self):
        """Return the ring parallel rank"""
        return self._dim_rank('ring', default=-1)

    def prepare_inputs(self, inputs):
        """Prepare inputs

        1. set extra_kwargs['text_position_ids']
        2. split labels
        """
        position_ids = inputs.get('text_position_ids')
        input_ids = inputs.get('input_ids')
        if position_ids is None:
            position_ids = inputs.get('position_ids')
        if position_ids is not None and input_ids is not None and position_ids.shape[0] == input_ids.shape[0]:
            self.extra_kwargs['text_position_ids'] = position_ids.clone()
        if input_ids is not None:
            self.extra_kwargs['input_ids'] = input_ids.clone()
        if 'labels' in inputs:
            labels = inputs['labels']
            _, _, labels, _, _, _, _ = self.pad_and_split_inputs(
                None, None, labels, None, None, None, real_position_ids=position_ids)
            inputs['labels'] = labels


sequence_parallel = SequenceParallel()
