layers.py

# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import torch
import re
from deepspeed import comm as dist
from torch import nn
from torch.nn import functional as F
from torch.nn.parameter import Parameter
from deepspeed.accelerator import get_accelerator
from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list
from deepspeed.runtime.zero.utils import is_zero_param
from abc import ABC, abstractmethod
from typing import Iterable, Any, Optional, List, Tuple, Dict
from .fusedqkv_utils import shard_value_with_share_qk, shard_chunk_mlp, prepare_tp_fused_qkvw
from deepspeed.runtime.tensor_parallel import AUTOTP_MODE
from deepspeed.checkpoint.constants import DS_AUTOTP_UC_META
from copy import deepcopy
from typing import Union

__all__ = [
    "TensorParallel_Layer", "LinearAllreduce", "LinearLayer", "LmHeadLinearAllreduce", "Yuan_LinearAllreduce",
    "Yuan_LinearLayer", "GateUpPack_LinearLayer", "Conv_LinearALlreduce", "fused_LinearLayer", "conv_LinearLayer",
    "SubParamLinearLayer", "SubParamLinearAllreduce"
]

DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.INFERENCE
DS_IS_REPLACED_MODULE = 'ds_is_replaced_module'
DS_TENSOR_MODEL_PARALLEL = 'tensor_model_parallel'


def _normalize_uc_shape(value):
    return tuple(value) if value is not None else None


def _build_param_uc_conversion_meta(*,
                                    partition_type,
                                    partition_dim=None,
                                    sub_param_shape=None,
                                    original_shape=None,
                                    is_bias=False,
                                    replicated=False):
    """Build the conversion-facing subset of parameter UC metadata.

    This is the only schema that should flow into model-level
    `UNIVERSAL_CHECKPOINT_INFO` via `collect_autotp_universal_checkpoint_info()`.
    """
    return {
        'partition_type': partition_type,
        'partition_dim': partition_dim,
        'sub_param_shape': _normalize_uc_shape(sub_param_shape),
        'original_shape': _normalize_uc_shape(original_shape),
        'is_bias': is_bias,
        'replicated': replicated,
    }


def _build_param_uc_restore_meta(*,
                                 partition_type,
                                 partition_dim=None,
                                 logical_shape=None,
                                 output_shape=None,
                                 sub_param_shape=None,
                                 sub_param_sizes=None,
                                 target_partition_shape=None,
                                 original_shape=None,
                                 is_bias=False,
                                 replicated=False):
    """Build the restore-facing parameter UC metadata.

    Restore metadata stays on the parameter object and may include details that
    are intentionally omitted from model-level conversion schema.
    """
    return {
        'partition_type':
        partition_type,
        'partition_dim':
        partition_dim,
        'logical_shape':
        _normalize_uc_shape(logical_shape),
        'output_shape':
        _normalize_uc_shape(output_shape),
        'sub_param_shape':
        _normalize_uc_shape(sub_param_shape),
        'sub_param_sizes':
        _normalize_uc_shape(sub_param_sizes),
        'target_partition_shape':
        _normalize_uc_shape(target_partition_shape),
        'original_shape':
        _normalize_uc_shape(original_shape),
        'is_bias':
        is_bias,
        'replicated':
        replicated,
        'conversion':
        _build_param_uc_conversion_meta(partition_type=partition_type,
                                        partition_dim=partition_dim,
                                        sub_param_shape=sub_param_shape,
                                        original_shape=original_shape,
                                        is_bias=is_bias,
                                        replicated=replicated),
    }


def get_auto_tp_mode():
    global DEEPSPEED_AUTOTP_MODE
    return DEEPSPEED_AUTOTP_MODE


def is_autotp_training_mode():
    global DEEPSPEED_AUTOTP_MODE
    return DEEPSPEED_AUTOTP_MODE == AUTOTP_MODE.TRAINING


def set_autotp_mode(training=False):
    """
    Set the DEEPSPEED_AUTOTP_MODE based on the training flag
    """
    global DEEPSPEED_AUTOTP_MODE
    if training:
        DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.TRAINING
    else:
        DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.INFERENCE


def add_bias(input, bias):
    if bias is None:
        return input
    if is_autotp_training_mode():
        # Training mode - avoid inplace to ensure correct autograd
        input = input + bias
        return input
    else:
        input += bias
        return input


class RowParallel(torch.autograd.Function):
    """
    A custom autograd function for performing row-wise parallelism.
    """

    @staticmethod
    def symbolic(graph, input):
        """Symbolic function for tracing."""
        return input

    @staticmethod
    def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor, is_inference_mode: bool) -> torch.Tensor:
        """
        Forward pass.
        """
        ctx.group = group
        if group == None:
            return input
        if is_inference_mode:
            dist.inference_all_reduce(input, group=group)
        else:
            dist.all_reduce(input.contiguous(), group=group)
        return input

    @staticmethod
    def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor, None]:
        """
        Backward pass.
        """
        return None, grad_output, None


class AsyncColumnParallel(torch.autograd.Function):

    @staticmethod
    def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor, weight, bias) -> torch.Tensor:
        """
        Forward pass.
        """
        ctx.use_bias = bias is not None
        ctx.group = group
        output = torch.matmul(input, weight.transpose(-1, -2))
        if bias is not None:
            output = add_bias(output, bias)

        ctx.save_for_backward(input, weight)

        return output

    @staticmethod
    def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor]:

        input, weight = ctx.saved_tensors
        grad_input = grad_output.matmul(weight)
        handle = dist.all_reduce(grad_input.contiguous(), group=ctx.group, async_op=True)
        grad_weight = grad_output.view(-1, grad_output.shape[-1]).t().matmul(input.view(-1, input.shape[-1]))
        grad_bias = grad_output.sum(0) if ctx.use_bias else None
        handle.wait()
        return None, grad_input, grad_weight, grad_bias


class ColumnParallel(torch.autograd.Function):
    """
    Custom autograd function for column-wise parallelism.
    """

    @staticmethod
    def symbolic(graph, input):
        """Symbolic function for tracing."""
        return dist.all_reduce(input.contiguous(), dist.get_tensor_model_parallel_group())

    @staticmethod
    def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor) -> torch.Tensor:
        """
        Forward pass.
        """
        ctx.group = group
        return input

    @staticmethod
    def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor]:
        """
        Backward pass.
        """
        if ctx.group == None:
            return None, grad_output

        dist.all_reduce(grad_output.contiguous(), group=ctx.group)
        return None, grad_output


class GatherFromTensorParallelRegion(torch.autograd.Function):
    """Gather last-dimension shards while keeping the output replicated."""

    @staticmethod
    def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor) -> torch.Tensor:
        ctx.group = group
        if group is None:
            ctx.partition_sizes = (input.shape[-1], )
            ctx.tp_index = 0
            return input

        tp_world_size = dist.get_world_size(group=group)
        ctx.tp_index = dist.get_rank(group=group)
        if tp_world_size == 1:
            ctx.partition_sizes = (input.shape[-1], )
            return input

        local_size = torch.tensor([input.shape[-1]], dtype=torch.long, device=input.device)
        gathered_sizes = [torch.empty_like(local_size) for _ in range(tp_world_size)]
        dist.all_gather(gathered_sizes, local_size, group=group)
        ctx.partition_sizes = tuple(int(size.item()) for size in gathered_sizes)

        max_partition_size = max(ctx.partition_sizes)
        if input.shape[-1] == max_partition_size:
            input_padded = input.contiguous()
        else:
            padded_shape = (*input.shape[:-1], max_partition_size)
            input_padded = input.new_zeros(padded_shape)
            input_padded[..., :input.shape[-1]].copy_(input)

        gathered = [torch.empty_like(input_padded) for _ in range(tp_world_size)]
        dist.all_gather(gathered, input_padded, group=group)
        return torch.cat([shard[..., :size] for shard, size in zip(gathered, ctx.partition_sizes)], dim=-1)

    @staticmethod
    def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor]:
        shard_offset = sum(ctx.partition_sizes[:ctx.tp_index])
        shard_size = ctx.partition_sizes[ctx.tp_index]
        grad_input = grad_output.narrow(-1, shard_offset, shard_size).contiguous()
        return None, grad_input


class TensorParallel_Layer(nn.Module, ABC):
    """
    A base class for model layers with  tensor parallelism support.
    This class is designed to be extended by specific layers that require distributed
    operations and parameter gather/partitioning during inference or training.

    Attributes:
        mode (str): The mode of operation[INFERENCE or TRAINING], default is "INFERENCE".
        mp_group (Optional[dist.ProcessGroup]): The process group used for model parallelism.
        tp_world_size (int): The world size of tensor parallelism, i.e., the number of parallel workers.
        tp_index (int): The rank (ID) of the current worker in tensor parallelism.
        support_training (bool): Flag indicating whether the layer supports training (default: False).
        name (Optional[str]): The name of the layer, if provided.
    """
    ##### Initialize Parameter List #####

    # keep_module_on_host determines whether to keep the module on the host.
    # Checkpoints are first loaded to the host (sometimes directly from disk to avoid filling host memory),
    # so an additional copy is unnecessary.
    keep_module_on_host: bool = False

    ##### Runtime Parameter List #####
    tp_overlap_comm: bool = False
    """ Whether to overlap communication with computation. Currently, only allreduce supports overlap. """

    def __init__(self, mp_group: Optional[dist.ProcessGroup], **kwargs: Any):
        """
        Initializes the TensorParallel_Layer with optional model parallelism group and layer name.

        Args:
            mp_group (Optional[dist.ProcessGroup]): The process group for model parallelism.
                                                    If None, no model parallelism is set.
        """
        super().__init__()
        self.support_training: bool = False
        self.mp_group = mp_group
        if mp_group is not None:
            self.tp_world_size: int = dist.get_world_size(self.mp_group)
            self.tp_index: int = dist.get_rank(self.mp_group)
        else:
            self.tp_world_size: int = 1
            self.tp_index: int = 0

        # backward compatibility
        self.world_size = self.tp_world_size
        self.rank = self.tp_index

        self.name = getattr(self, 'name', None)
        if kwargs.get('name') is not None:
            self.name = kwargs.get('name')  # Set the layer name if provided.

    @classmethod
    def set_keep_module_on_host(cls, value: bool):
        """
        Set the static variable keep_module_on_host.

        Args:
            value (bool): The new value for keep_module_on_host.
        """
        cls.keep_module_on_host = value

    @abstractmethod
    def forward(self, input):
        """
        Forward pass method. Must be implemented by subclasses to define layer-specific operations.
        """
        pass

    @abstractmethod
    def gather_params(self, params_list):
        """
        Gathers parameters across devices for distributed training. Must be implemented by subclasses in "TRAINING" mode.
        """
        pass

    @abstractmethod
    def _tp_partition(self, params_list: List[torch.Tensor]):
        """
        Partitions the parameters for tensor parallelism.
        It is necessary to ensure that this function only involves the logic of params partitioning.
        """
        pass

    def config_requires_grad(self, weight):
        if weight is not None:
            if self.is_training_mode():
                if weight.requires_grad is None:
                    weight.requires_grad = True
            else:
                weight.requires_grad = False

    def config_tp_params(self, weight):
        """
        Configures the weight tensor for training with tensor parallelism. This includes enabling gradients
        and associating necessary methods for parameter gathering and partitioning.

        Args:
            weight (Optional[torch.Tensor]): The weight tensor to configure for tensor parallelism.
                                              If None, no action is taken.
        """
        # # The RNG states have already been synchronized in init_inference.
        if self.is_training_mode():
            assert self.support_training, "No implementation of backward."
        if weight is not None:
            self.config_requires_grad(weight)
            weight.gather_params = self.gather_params
            weight._tp_partition = self._tp_partition
            setattr(weight, DS_TENSOR_MODEL_PARALLEL, True)
            setattr(weight, DS_IS_REPLACED_MODULE, True)

    def _set_param_uc_meta(self,
                           param,
                           *,
                           partition_type,
                           partition_dim=None,
                           logical_shape=None,
                           output_shape=None,
                           sub_param_shape=None,
                           sub_param_sizes=None,
                           target_partition_shape=None,
                           original_shape=None,
                           is_bias=False,
                           replicated=False):
        if param is None:
            return
        setattr(
            param, DS_AUTOTP_UC_META,
            _build_param_uc_restore_meta(partition_type=partition_type,
                                         partition_dim=partition_dim,
                                         logical_shape=logical_shape,
                                         output_shape=output_shape,
                                         sub_param_shape=sub_param_shape,
                                         sub_param_sizes=sub_param_sizes,
                                         target_partition_shape=target_partition_shape,
                                         original_shape=original_shape,
                                         is_bias=is_bias,
                                         replicated=replicated))

    def _mark_uc_metadata(self):
        return

    def _should_materialize_tp_partition(self):
        # AutoTP partitioning should only materialize parameters when an actual
        # TP process group is present. Metadata-only construction with
        # mp_group=None should not touch device placement.
        return self.mp_group is not None

    def is_training_mode(self):
        global DEEPSPEED_AUTOTP_MODE
        return DEEPSPEED_AUTOTP_MODE == AUTOTP_MODE.TRAINING

    def __deepcopy__(self, memo):
        # This function is designed for
        # 'mp_group' (a 'ProcessGroup') cannot be pickled during deepcopy in some usage.
        cls = self.__class__
        new_obj = cls.__new__(cls)

        for key, value in vars(self).items():
            if key == 'mp_group':
                new_obj.mp_group = self.mp_group
            else:
                setattr(new_obj, key, deepcopy(value, memo))

        memo[id(self)] = new_obj
        return new_obj

    def extra_repr(self):
        out_features, in_features = None, None
        if self.weight is not None:
            out_features, in_features = self.weight.ds_shape[-2:] if is_zero_param(
                self.weight) else self.weight.shape[-2:]
        dtype = self.weight.dtype if self.weight is not None else None
        return "in_features={}, out_features={}, bias={}, dtype={}".format(in_features, out_features, self.bias
                                                                           is not None, dtype)

    def move(self, tensor):
        # TODO: consider the timing of deletion
        # to save host resources when DP > 1。

        # keep_module_on_host is used to keep the module on the host. Checkpoints are loaded to the host first (in some
        # cases it can be done from the disk even to prevent filling host's memory), thus no need to create a new copy.
        if tensor.is_meta:
            # Keep tensor in meta device if tensor is meta.
            return tensor
        else:
            device = 'cpu' if self.__class__.keep_module_on_host else get_accelerator().current_device_name()
            return_new_copy = not self.__class__.keep_module_on_host

            # Using new tensors help in freeing memory (after split for example) was done before by calling clone().
            # Using copy=True instead of clone() will help in case of cpu --> cpu.
            # Otherwise to() will not create a new copy for the view of the full tensor, and it will not be de-referenced.
            cloned_tensor = tensor.to(device, copy=return_new_copy)

            if return_new_copy:
                # free the memory of the original tensor to reduce memory peak
                # Equivalent to directly deleting the tensor reference outside the function.
                # see https://github.com/microsoft/DeepSpeed/pull/4353
                tensor.data = torch.empty(0, device=tensor.device)
            return cloned_tensor


def configure_tensor_parallel_runtime(config):
    runtime_keys = ['tp_overlap_comm']
    for key in runtime_keys:
        if hasattr(config, key):
            setattr(TensorParallel_Layer, key, getattr(config, key))


def _get_param_uc_conversion_meta(param: torch.Tensor) -> Optional[Dict[str, Any]]:
    """Return the conversion-facing view of AutoTP UC metadata for a parameter.

    AutoTP keeps a single parameter-level metadata object with two roles:
    - top-level fields: restore-time details consumed by `universal_checkpoint.py`
    - `conversion`: conversion-time details consumed by
      `collect_autotp_universal_checkpoint_info()` and then aggregated into
      model-level `UNIVERSAL_CHECKPOINT_INFO` for `ds_to_universal.py`
    """
    meta = getattr(param, DS_AUTOTP_UC_META, None)
    if not meta:
        return None
    return meta.get('conversion', None)


def collect_autotp_universal_checkpoint_info(model: nn.Module) -> Dict[str, Any]:
    """Collect the model-level conversion schema for AutoTP universal checkpoints.

    The returned `UNIVERSAL_CHECKPOINT_INFO` is intentionally limited to the
    pattern/schema data needed during checkpoint conversion. It does not include
    restore-time per-parameter details such as `sub_param_sizes` or
    `target_partition_shape`, which stay on the parameter metadata object.
    """
    from deepspeed.checkpoint.constants import (ORIGINAL_VOCAB_SIZE, PARAMETER_WITH_ROW_PARALLELISM_PATTERNS,
                                                PARAMETER_WITH_SUB_PARAMS, TP_REPLICATED_PARAMETER_PATTERNS,
                                                UNIVERSAL_CHECKPOINT_VERSION_KEY, UNIVERSAL_CHECKPOINT_VERSION_VALUE,
                                                VOCABULARY_PARAMETER_PATTERNS)

    row_parallel_patterns = []
    replicated_patterns = []
    vocabulary_patterns = []
    parameter_with_sub_params = []
    original_vocab_size = None

    for module_name, module in model.named_modules():
        marker = getattr(module, "_mark_uc_metadata", None)
        if marker is not None:
            marker()

        for param_name, param in module.named_parameters(recurse=False):
            conversion_meta = _get_param_uc_conversion_meta(param)
            if not conversion_meta:
                continue

            full_name = f"{module_name}.{param_name}" if module_name else param_name
            pattern = rf"^{re.escape(full_name)}$"

            if conversion_meta.get('replicated'):
                replicated_patterns.append(pattern)

            if conversion_meta.get('partition_type') == 'row' and not conversion_meta.get('is_bias', False):
                row_parallel_patterns.append(pattern)

            original_shape = conversion_meta.get('original_shape')
            if original_shape and len(original_shape) == 2 and ('embed' in full_name or 'lm_head' in full_name):
                vocabulary_patterns.append(pattern)
                if original_vocab_size is None:
                    original_vocab_size = original_shape[0]

            sub_param_shape = conversion_meta.get('sub_param_shape')
            partition_dim = conversion_meta.get('partition_dim')
            if sub_param_shape is not None and partition_dim is not None and not conversion_meta.get('is_bias', False):
                parameter_with_sub_params.append({
                    'patterns': [pattern],
                    'shape': list(sub_param_shape),
                    'partition_dim': partition_dim,
                })

    uc_info = {
        UNIVERSAL_CHECKPOINT_VERSION_KEY: UNIVERSAL_CHECKPOINT_VERSION_VALUE,
        PARAMETER_WITH_ROW_PARALLELISM_PATTERNS: sorted(set(row_parallel_patterns)),
        TP_REPLICATED_PARAMETER_PATTERNS: sorted(set(replicated_patterns)),
        VOCABULARY_PARAMETER_PATTERNS: sorted(set(vocabulary_patterns)),
        PARAMETER_WITH_SUB_PARAMS: parameter_with_sub_params,
    }
    if original_vocab_size is not None:
        uc_info[ORIGINAL_VOCAB_SIZE] = original_vocab_size
    return uc_info


class GatherReplacedLayerParams:
    """
    A context manager for gathering parameters of a replaced layer, enabling partitioning and gathering functionality
    based on the configuration of the model.
    """

    def __init__(self,
                 params: Union[Iterable[torch.Tensor], torch.Tensor],
                 module: torch.nn.Module,
                 enabled: bool = True):
        """
        Initialize the context manager to handle parameter gathering and partitioning for a replaced layer.

        Args:
            params (Iterable or torch.Tensor): A collection or single parameter to manage.
            module (torch.nn.Module): The module that these parameters belong to.
            enabled (bool): Flag indicating whether the parameter management is enabled (default: True).
        """
        self.enabled = enabled
        self.module = module
        if not enabled:
            return

        # Ensure params is a list, whether it's a single param or iterable (e.g., model.parameters())
        if isinstance(params, Iterable) and not isinstance(params, torch.Tensor):
            self.params: List[torch.Tensor] = list(params)  # Convert generators to a list for multiple iterations
        else:
            self.params: List[torch.Tensor] = [params]  # Wrap single parameter in a list for uniform processing

        # Check if the parameters belong to a replaced layer (indicated by a specific attribute)
        if not any(self._is_replaced_module_weight(p) for p in params):
            self.enabled = False
            return

    def _is_replaced_module_weight(self, param: torch.Tensor) -> bool:
        """
        Helper function to determine if a parameter belongs to a replaced module.

        Args:
            param (torch.Tensor): The parameter to check.

        Returns:
            bool: True if the parameter belongs to a replaced module, False otherwise.
        """
        return getattr(param, DS_IS_REPLACED_MODULE, False)

    def __enter__(self) -> None:
        """
        Enter the context manager. If enabled, gather parameters for the replaced module.
        """
        if self.enabled:
            self.params[0].gather_params(self.params)

    def __exit__(self, exc_type, exc_value, traceback) -> None:
        """
        Exit the context manager. If enabled, partition the parameters for the replaced module.
        """
        #TODO : Check whether there are any missing attributes.
        if self.enabled:
            self.params[0]._tp_partition(self.params)


class LinearAllreduce(TensorParallel_Layer):

    def __init__(self, module, mp_group, **kwargs):
        super(LinearAllreduce, self).__init__(mp_group, **kwargs)
        self.weight = module.weight
        self.bias = module.bias

        if self._should_materialize_tp_partition():
            self._tp_partition([self.weight, self.bias])
        self.support_training = True
        self.config_tp_params(self.weight)
        if self.bias is not None:
            # bias here is not tp params
            self.config_requires_grad(self.bias)
        self._mark_uc_metadata()

    def forward(self, input):
        output = torch.matmul(input, self.weight.transpose(-1, -2))
        output = RowParallel.apply(self.mp_group, output, not self.is_training_mode())
        if self.bias is not None:
            output = add_bias(output, self.bias)
        return output

    @torch.no_grad()
    def gather_params(self, params_list):

        for idx, param in enumerate(params_list):
            if param is None or idx > 0:
                # don't gather bias
                return
            params_list[idx].data_partition = param.data
            param = param.transpose(0, 1).contiguous()

            output_param = torch.empty(self.tp_world_size * param.shape[0],
                                       param.shape[1],
                                       dtype=param.dtype,
                                       device=param.device)
            dist.all_gather_into_tensor(output_param, param, group=self.mp_group)
            params_list[idx].data = output_param.transpose(0, 1).contiguous()
        return

    @torch.no_grad()
    def _tp_partition(self, params_list):

        if not self.is_training_mode():
            self.uneven_partition(params_list)
            return

        else:
            for idx, param in enumerate(params_list):
                if param is None:
                    # don't slipt bias
                    return
                if idx > 0:  # move bias to device at initialization
                    _partition = self.move(param).detach()
                    params_list[idx].data = _partition
                    return

                _partition = torch.chunk(param, self.tp_world_size, dim=-1)[self.tp_index]

                _partition = self.move(_partition).detach()

                params_list[idx].data = _partition

    def uneven_partition(self, params_list):
        for idx, param in enumerate(params_list):
            if param is None or idx > 0:
                # don't slipt bias
                return
            assert self.name is not None, "The module name must be provided in the initialization."
            _partition = params_list[idx].split(get_shard_size_list(params_list[idx].shape[1], self.tp_world_size,
                                                                    self.name),
                                                dim=1)[self.tp_index]

            _partition = self.move(_partition).detach()
            params_list[idx].data = _partition

    def _mark_uc_metadata(self):
        original_weight_shape = (self.weight.shape[0], self.weight.shape[1] * self.tp_world_size)
        self._set_param_uc_meta(self.weight,
                                partition_type='row',
                                partition_dim=1,
                                logical_shape=original_weight_shape,
                                output_shape=(original_weight_shape[0], ),
                                original_shape=original_weight_shape)
        if self.bias is not None:
            self._set_param_uc_meta(self.bias,
                                    partition_type='row',
                                    partition_dim=None,
                                    logical_shape=tuple(self.bias.shape),
                                    output_shape=tuple(self.bias.shape),
                                    original_shape=tuple(self.bias.shape),
                                    is_bias=True,
                                    replicated=True)


#remove kwargs from partition.
class LinearLayer(TensorParallel_Layer):

    def __init__(self, module, mp_group=None, skip_partition=False, gather_output=False, **kwargs):
        super(LinearLayer, self).__init__(mp_group, **kwargs)
        self.weight = module.weight
        self.bias = module.bias
        self.gather_output = gather_output
        if not skip_partition and self._should_materialize_tp_partition():
            self._tp_partition([self.weight, self.bias])
        self.support_training = True
        self.config_tp_params(self.weight)
        if self.bias is not None:
            self.config_tp_params(self.bias)
        self._mark_uc_metadata()

    def forward(self, input):
        if not self.__class__.tp_overlap_comm:
            if getattr(self, 'mp_group', None) is not None:
                input = ColumnParallel.apply(self.mp_group, input)
            output = torch.matmul(input, self.weight.transpose(-1, -2))
            if self.bias is not None:
                output = add_bias(output, self.bias)
        else:
            output = AsyncColumnParallel.apply(self.mp_group, input, self.weight, self.bias)

        if self.gather_output:
            output = GatherFromTensorParallelRegion.apply(self.mp_group, output)

        return output

    @torch.no_grad()
    def gather_params(self, params_list):
        #  Does not support uneven shard.
        for idx, param in enumerate(params_list):

            params_list[idx].data_partition = param.data
            output_param = torch.empty((self.tp_world_size * param.shape[0], *param.shape[1:]),
                                       dtype=param.dtype,
                                       device=param.device)
            dist.all_gather_into_tensor(output_param, param, group=self.mp_group)
            params_list[idx].data = output_param.contiguous()

    @torch.no_grad()
    def _tp_partition(self, params_list):

        if not self.is_training_mode():
            self.uneven_partition(params_list)
            return
        for idx, param in enumerate(params_list):
            if param is None:
                return
            #split bias if provide
            _partition = torch.chunk(param, self.tp_world_size, dim=0)[self.tp_index]

            _partition = self.move(_partition).detach()

            params_list[idx].data = _partition

    def uneven_partition(self, params_list):

        for idx, param in enumerate(params_list):
            if param is None:
                #split bias if provide
                return
            assert self.name is not None, "The module name must be provided in the initialization."
            _partition = params_list[idx].split(get_shard_size_list(params_list[idx].shape[0], self.tp_world_size,
                                                                    self.name),
                                                dim=0)[self.tp_index]

            _partition = self.move(_partition).detach()

            params_list[idx].data = _partition

    def _mark_uc_metadata(self):
        original_out_dim = self.weight.shape[0] * self.tp_world_size
        original_weight_shape = (original_out_dim, self.weight.shape[1])
        self._set_param_uc_meta(self.weight,
                                partition_type='column',
                                partition_dim=0,
                                logical_shape=original_weight_shape,
                                output_shape=(original_out_dim, ),
                                original_shape=original_weight_shape)
        if self.bias is not None:
            original_bias_shape = (self.bias.shape[0] * self.tp_world_size, )
            self._set_param_uc_meta(self.bias,
                                    partition_type='column',
                                    partition_dim=0,
                                    logical_shape=original_bias_shape,
                                    output_shape=original_bias_shape,
                                    original_shape=original_bias_shape,
                                    is_bias=True)

    # for bwc
    @classmethod
    def from_weights(cls, weight_shape=None, dtype=torch.half, weight=None, bias=None, gather_output=False):
        if weight is not None:
            in_features = weight.shape[1]
            out_features = weight.shape[0]
            linear = nn.Linear(in_features, out_features, bias=(bias is not None))
            linear.weight.data = weight
            if bias is not None:
                linear.bias.data = bias
        else:
            in_features = weight_shape[1]
            out_features = weight_shape[0]
            linear = nn.Linear(in_features, out_features, bias=(bias is not None))
        return cls(linear, skip_partition=True, gather_output=gather_output)


class FusedModuleWrapper:

    def __init__(self, fused_module: nn.Module):
        self.fused_module = fused_module

    def __getattr__(self, module):
        return self.fused_module


class fused_LinearLayer(LinearLayer):

    def __init__(self, module, mp_group, skip_partition=False, **kwargs):
        assert kwargs.get('fused_module') is not None, "'fused_module' is required but not provided"
        # Use the warp class to avoid module circular references.
        self.fused_module = FusedModuleWrapper(kwargs.get('fused_module'))
        super().__init__(module, mp_group, skip_partition, **kwargs)

    @torch.no_grad()
    def _tp_partition(self, params_list):
        for idx, param in enumerate(params_list):
            if param is None:
                return

            _partition = prepare_tp_fused_qkvw(self.fused_module.module, param, self.tp_world_size, self.tp_index)

            _partition = self.move(_partition).detach()

            params_list[idx].data = _partition


class conv_LinearLayer(LinearLayer):

    @torch.no_grad()
    def _tp_partition(self, params_list):
        weight = None
        bias = None
        if len(params_list) == 1:
            weight = params_list[0]
        elif len(params_list) == 2:
            weight, bias = params_list[0], params_list[1]
        _partition = weight.data.split(get_shard_size_list(weight.shape[0], self.tp_world_size, self.name),
                                       dim=1)[self.tp_index]
        _partition = self.move(_partition).detach()
        weight.data = _partition

        if bias is not None:
            _partition = bias.data.split(get_shard_size_list(weight.shape[1], self.tp_world_size, self.name),
                                         dim=0)[self.tp_index]
            _partition = self.move(_partition).detach()

            bias.data = _partition


#override the subclasses related to weight splitting.
class Yuan_LinearAllreduce(LinearAllreduce):

    #Yuan2
    @torch.no_grad()
    def _tp_partition(self, params_list):
        weight, bias = shard_value_with_share_qk(params_list[0].data, params_list[1], self.tp_index,
                                                 self.tp_world_size, False)
        params_list[0].data = weight
        if bias is not None:
            params_list[1].data = bias


class Yuan_LinearLayer(LinearLayer):
    #Yuan2
    @torch.no_grad()
    def _tp_partition(self, params_list):
        weight, bias = shard_value_with_share_qk(params_list[0].data, params_list[1], self.tp_index,
                                                 self.tp_world_size, True)
        params_list[0].data = self.move(weight).detach()
        if bias is not None:
            params_list[1].data = self.move(bias).detach()


class GateUpPack_LinearLayer(LinearLayer):
    # chatGLM2, chatGLM2
    @torch.no_grad()
    def _tp_partition(self, params_list):
        weight, bias = shard_chunk_mlp(params_list[0].data, params_list[1], self.tp_index, self.tp_world_size)
        params_list[0].data = self.move(weight).detach()
        if bias is not None:
            params_list[1].data = self.move(bias).detach()


class Conv_LinearALlreduce(LinearAllreduce):

    @torch.no_grad()
    def _tp_partition(self, params_list):
        for idx, param in enumerate(params_list):
            if param is None:
                return
            param.data = param.data.transpose(-1, -2).contiguous()

            _partition = param.split(get_shard_size_list(param.shape[0], self.tp_world_size, self.name),
                                     dim=1)[self.tp_index]

            _partition = self.move(_partition).detach()

            params_list[idx].data = _partition


#override the subclasses related to fwd/bwd.
class LmHeadLinearAllreduce(LinearAllreduce):

    def __init__(self, module, mp_group, **kwargs):
        # set the fixed name before partition
        self.name = "lm_head"

        # In some tied_embedding cases, only the lm head is sharded, while the word embedding is not.
        # Reinitialization is used to decouple them and prevent the word embedding from being sharded.
        # This should also be effective for cases where both are sharded in tied_embedding scenarios.

        # TODO: Training scenario-related tests, is it necessary to re-implement the vocab parallel module?
        module.weight = nn.Parameter(module.weight.clone().detach())
        if hasattr(module, 'bias') and module.bias is not None:
            module.bias = nn.Parameter(module.bias.clone().detach())
        super().__init__(module, mp_group, **kwargs)

    def forward(self, input):
        input_shard_size = get_shard_size(input.shape[-1], self.tp_world_size, "lm_head")
        input_shard_offset = sum(get_shard_size_list(input.shape[-1], self.tp_world_size, "lm_head")[0:self.tp_index])
        output = torch.matmul(input[:, :, input_shard_offset:input_shard_offset + input_shard_size],
                              self.weight.transpose(-1, -2))
        if self.mp_group is not None:
            dist.inference_all_reduce(output, group=self.mp_group)
        if self.bias is not None:
            output = add_bias(output, self.bias)
        return output


class TensorParallelConv2d(nn.Module):

    def __init__(self, conv, rank, world_size, shard_by_oc):
        super().__init__()
        self.rank = rank
        self.world_size = world_size
        self.shard_by_oc = shard_by_oc
        self.shard_weights(conv)

    # Split along the input/output channel depending on whether it is the last conv layer.
    def shard_weights(self, conv):
        if self.shard_by_oc:
            total_size = conv.weight.shape[0]
        else:
            total_size = conv.weight.shape[1]
        bias_data = None
        cols_per_rank = [0]
        for i in range(self.world_size - 1, -1, -1):
            cols = total_size // self.world_size
            if i < total_size % self.world_size:
                cols += 1
            cols_per_rank.append(cols_per_rank[-1] + cols)
        weight_data = conv.weight.data
        if self.shard_by_oc:
            # not last conv layer, split output channel
            weight_data = weight_data[cols_per_rank[self.rank]:cols_per_rank[self.rank + 1]]
            if conv.bias is not None:
                bias_data = conv.bias.data[cols_per_rank[self.rank]:cols_per_rank[self.rank + 1]]
        else:
            # last conv layer, split input channel
            weight_data = weight_data[:, cols_per_rank[self.rank]:cols_per_rank[self.rank + 1]]
            if conv.bias is not None:
                bias_data = conv.bias.data / float(self.world_size)
        self.conv = nn.Conv2d(weight_data.shape[1], weight_data.shape[0], conv.kernel_size, conv.stride, conv.padding,
                              conv.dilation, conv.groups, conv.bias is not None, conv.padding_mode)
        self.conv.weight = torch.nn.Parameter(weight_data)
        if conv.bias is not None:
            self.conv.bias = torch.nn.Parameter(bias_data)
        del conv

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return self.conv(input)


class TensorParallelOcShardConv2d(TensorParallelConv2d):

    def __init__(self, conv, rank, world_size):
        super().__init__(conv, rank, world_size, True)


class TensorParallelIcShardConv2d(TensorParallelConv2d):

    def __init__(self, conv, rank, world_size):
        super().__init__(conv, rank, world_size, False)

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        out = self.conv(input)
        if self.world_size > 1:
            dist.inference_all_reduce(out)
        return out


class Normalize(nn.Module):

    def __init__(self, dim=None, dtype=torch.float, eps=1e-5, weight=None, bias=None):
        super(Normalize, self).__init__()
        if weight is not None:
            self.weight = weight
            self.bias = bias
        else:
            self.norm = nn.LayerNorm(dim, eps=eps).to(dtype).to(get_accelerator().current_device_name())
            self.weight = self.norm.weight
            self.bias = self.norm.bias

        self.eps = eps

    def forward(self, input):
        return nn.functional.layer_norm(input, input.shape[-1:], self.weight, self.bias, eps=self.eps)


class EmbeddingLayer(nn.Module):

    def __init__(self, weight_shape=None, dtype=torch.half, weight=None, bias=None):
        super(EmbeddingLayer, self).__init__()
        if weight is None:
            self.weight = Parameter(
                torch.empty(weight_shape[0],
                            weight_shape[1],
                            dtype=dtype,
                            device=get_accelerator().current_device_name()))
        else:
            self.weight = weight

    def forward(self, input):
        return F.embedding(input, self.weight)


class OPTEmbedding(EmbeddingLayer):
    """
    This module learns positional embeddings up to a fixed maximum size.
    """

    def __init__(self, weight_shape=None, weight=None, bias=None):
        # OPT is set up so that if padding_idx is specified then offset the embedding ids by 2
        # and adjust num_embeddings appropriately. Other models don't have this hack
        self.offset = 2
        super().__init__(weight_shape, weight=weight)

    def forward(self, attention_mask: torch.LongTensor, past_key_values_length: int = 0, position_ids: int = 0):
        """`input_ids_shape` is expected to be [bsz x seqlen]."""
        attention_mask = attention_mask.long()

        # create positions depending on attention_mask
        positions = (torch.cumsum(attention_mask, dim=1).type_as(attention_mask) * attention_mask).long() - 1

        # cut positions if `past_key_values_length` is > 0
        positions = positions[:, past_key_values_length:]

        return super().forward(positions + self.offset)


def _shape_prod(values):
    result = 1
    for val in values:
        result *= val
    return result


def _normalize_shape_spec(shape):
    if isinstance(shape, list):
        return tuple(_normalize_shape_spec(item) for item in shape)
    if isinstance(shape, tuple):
        return tuple(_normalize_shape_spec(item) if isinstance(item, list) else item for item in shape)
    return shape


def _infer_subparam_logical_shapes(weight_shape, shape, partition_dim, name=None):
    shape = _normalize_shape_spec(shape)
    if not isinstance(shape, tuple):
        raise ValueError("AutoTP shape must be a tuple for sub-parameter partitioning.")
    if partition_dim < 0 or partition_dim >= len(shape):
        raise ValueError(f"AutoTP partition_dim {partition_dim} is out of range for shape length {len(shape)}.")

    layer_label = f"AutoTP layer '{name}'" if name else "AutoTP layer"
    partition_elem = shape[partition_dim]
    subparam_sizes = None
    num_subparams = None

    if isinstance(partition_elem, tuple):
        if len(partition_elem) == 0:
            raise ValueError(f"{layer_label} sub-parameter size tuple cannot be empty.")
        if any(isinstance(val, tuple) for val in partition_elem):
            raise ValueError(f"{layer_label} supports only 1-level nesting at partition_dim.")
        if any((not isinstance(val, int)) or val <= 0 for val in partition_elem):
            raise ValueError(f"{layer_label} sub-parameter sizes must be positive integers.")
        subparam_sizes = tuple(int(val) for val in partition_elem)
        partition_dim_size = sum(subparam_sizes)
    elif isinstance(partition_elem, int):
        if partition_elem == -1:
            partition_dim_size = None
        elif partition_elem > 0:
            num_subparams = partition_elem
            partition_dim_size = None
        else:
            raise ValueError(f"{layer_label} partition_dim spec must be positive integer or -1.")
    else:
        raise ValueError(f"{layer_label} partition_dim spec must be int or tuple.")

    logical_dims = []
    for idx, dim in enumerate(shape):
        if idx == partition_dim:
            logical_dims.append(partition_dim_size)
            continue
        if isinstance(dim, tuple):
            raise ValueError(f"{layer_label} nested tuple only allowed at partition_dim={partition_dim}.")
        if isinstance(dim, int):
            if dim == -1:
                logical_dims.append(None)
            elif dim > 0:
                logical_dims.append(dim)
            else:
                raise ValueError(f"{layer_label} shape dimensions must be positive integers or -1.")
        else:
            raise ValueError(f"{layer_label} shape dimensions must be integers.")

    total_numel = _shape_prod(weight_shape)
    known_product = _shape_prod([dim for dim in logical_dims if dim is not None])
    unknown_indices = [idx for idx, dim in enumerate(logical_dims) if dim is None]

    if len(unknown_indices) == 0:
        if known_product != total_numel:
            raise ValueError(f"{layer_label} shape product {known_product} != weight numel {total_numel}.")
    elif len(unknown_indices) == 1:
        inferred = total_numel // known_product
        if inferred * known_product != total_numel:
            raise ValueError(f"{layer_label} cannot infer shape for weight with numel {total_numel}.")
        logical_dims[unknown_indices[0]] = inferred
    else:
        if len(shape) == len(weight_shape):
            for idx in unknown_indices:
                logical_dims[idx] = weight_shape[idx]
            if _shape_prod(logical_dims) != total_numel:
                raise ValueError(
                    f"{layer_label} shape product {_shape_prod(logical_dims)} != weight numel {total_numel}.")
        else:
            raise ValueError(f"{layer_label} shape has multiple inferred dims and is ambiguous for weight.")

    logical_shape = tuple(logical_dims)
    if logical_shape[-1] != weight_shape[-1]:
        raise ValueError(
            f"{layer_label} shape last dim {logical_shape[-1]} must match weight input dim {weight_shape[-1]}.")

    output_shape = logical_shape[:-1]
    if len(output_shape) == 0:
        raise ValueError(f"{layer_label} shape must include at least one output dimension.")
    if _shape_prod(output_shape) != weight_shape[0]:
        raise ValueError(
            f"{layer_label} output shape product {_shape_prod(output_shape)} != weight output dim {weight_shape[0]}.")

    partition_dim_size = logical_shape[partition_dim]
    if partition_dim_size is None or partition_dim_size <= 0:
        raise ValueError(f"{layer_label} partition_dim size must be a positive integer.")

    if num_subparams is not None:
        if partition_dim_size % num_subparams != 0:
            raise ValueError(
                f"{layer_label} partition_dim size {partition_dim_size} not divisible by sub-param count {num_subparams}."
            )
        subparam_sizes = tuple([partition_dim_size // num_subparams] * num_subparams)

    if subparam_sizes is not None and sum(subparam_sizes) != partition_dim_size:
        raise ValueError(
            f"{layer_label} sub-parameter sizes sum {sum(subparam_sizes)} != partition_dim size {partition_dim_size}.")

    bias_partition_dim = partition_dim if partition_dim < len(output_shape) else None
    return logical_shape, output_shape, subparam_sizes, bias_partition_dim


def _partition_logical_tensor(tensor, partition_dim, tp_world_size, tp_index, name=None, subparam_sizes=None):
    if tp_world_size == 1:
        return tensor
    layer_label = f"AutoTP layer '{name}'" if name else "AutoTP layer"
    if subparam_sizes:
        for size in subparam_sizes:
            if size % tp_world_size != 0:
                raise ValueError(f"{layer_label} sub-parameter size {size} not divisible by tp_size {tp_world_size}.")
        sub_params = torch.split(tensor, subparam_sizes, dim=partition_dim)
        partitioned_sub_params = [torch.chunk(sp, tp_world_size, dim=partition_dim)[tp_index] for sp in sub_params]
        return torch.cat(partitioned_sub_params, dim=partition_dim)
    if tensor.shape[partition_dim] % tp_world_size != 0:
        raise ValueError(
            f"{layer_label} partition_dim size {tensor.shape[partition_dim]} not divisible by tp_size {tp_world_size}."
        )
    return torch.chunk(tensor, tp_world_size, dim=partition_dim)[tp_index]


def _all_gather_along_dim(tensor, partition_dim, mp_group, tp_world_size):
    if mp_group is None or tp_world_size == 1:
        return tensor
    perm = [partition_dim] + [idx for idx in range(tensor.dim()) if idx != partition_dim]
    inv_perm = [0] * len(perm)
    for idx, dim in enumerate(perm):
        inv_perm[dim] = idx
    tensor_perm = tensor.permute(perm).contiguous()
    output = torch.empty((tp_world_size * tensor_perm.shape[0], *tensor_perm.shape[1:]),
                         dtype=tensor.dtype,
                         device=tensor.device)
    dist.all_gather_into_tensor(output, tensor_perm, group=mp_group)
    return output.permute(inv_perm).contiguous()


def _gather_logical_tensor(tensor,
                           logical_shape,
                           partition_dim,
                           mp_group,
                           tp_world_size,
                           name=None,
                           subparam_sizes=None):
    if mp_group is None or tp_world_size == 1:
        return tensor.reshape(logical_shape)
    layer_label = f"AutoTP layer '{name}'" if name else "AutoTP layer"
    if logical_shape[partition_dim] % tp_world_size != 0:
        raise ValueError(
            f"{layer_label} partition_dim size {logical_shape[partition_dim]} not divisible by tp_size {tp_world_size}."
        )
    partitioned_shape = list(logical_shape)
    partitioned_shape[partition_dim] = logical_shape[partition_dim] // tp_world_size
    tensor_view = tensor.reshape(partitioned_shape)

    if subparam_sizes:
        for size in subparam_sizes:
            if size % tp_world_size != 0:
                raise ValueError(f"{layer_label} sub-parameter size {size} not divisible by tp_size {tp_world_size}.")
        partitioned_sizes = [size // tp_world_size for size in subparam_sizes]
        sub_params = torch.split(tensor_view, partitioned_sizes, dim=partition_dim)
        gathered_sub_params = [_all_gather_along_dim(sp, partition_dim, mp_group, tp_world_size) for sp in sub_params]
        return torch.cat(gathered_sub_params, dim=partition_dim)
    return _all_gather_along_dim(tensor_view, partition_dim, mp_group, tp_world_size)


class SubParamLinearLayer(TensorParallel_Layer):
    """
    Column-parallel linear layer with sub-parameter support.

    Handles cases where weights contain multiple logical sub-parameters
    that need to be partitioned separately (e.g., fused QKV, chunked MLP, GQA).

    The `shape` parameter controls how the weight is viewed and partitioned:
    - (3, -1) with partition_dim=0: 3 equal sub-params, partition each at dim 0
    - ((q, k, v), -1) with partition_dim=0: 3 unequal sub-params (1-level nesting)
    """

    def __init__(self, module, mp_group, shape, partition_dim=0, **kwargs):
        super(SubParamLinearLayer, self).__init__(mp_group, **kwargs)
        self.weight = module.weight
        self.bias = module.bias
        self.shape = shape
        self.partition_dim = partition_dim

        self._orig_weight_shape = tuple(module.weight.shape)
        self._orig_bias_shape = tuple(module.bias.shape) if self.bias is not None else None
        (self._logical_shape, self._output_shape, self._subparam_sizes,
         self._bias_partition_dim) = _infer_subparam_logical_shapes(self._orig_weight_shape, self.shape,
                                                                    self.partition_dim, self.name)
        if self.bias is not None and self.bias.numel() != _shape_prod(self._output_shape):
            raise ValueError(f"AutoTP layer '{self.name}' bias size {self.bias.numel()} does not match output shape "
                             f"{self._output_shape}.")

        if self._should_materialize_tp_partition():
            self._tp_partition([self.weight, self.bias])
        self.support_training = True
        self.config_tp_params(self.weight)
        if self.bias is not None:
            self.config_tp_params(self.bias)
        self._mark_uc_metadata()

    def forward(self, input):
        if getattr(self, 'mp_group', None) is not None:
            input = ColumnParallel.apply(self.mp_group, input)
        output = torch.matmul(input, self.weight.transpose(-1, -2))
        if self.bias is not None:
            output = add_bias(output, self.bias)
        return output

    @torch.no_grad()
    def gather_params(self, params_list):
        """Gather partitioned parameters back to full size."""
        for idx, param in enumerate(params_list):
            if param is None:
                continue
            params_list[idx].data_partition = param.data
            if idx == 0:
                full_view = _gather_logical_tensor(param,
                                                   self._logical_shape,
                                                   self.partition_dim,
                                                   self.mp_group,
                                                   self.tp_world_size,
                                                   name=self.name,
                                                   subparam_sizes=self._subparam_sizes)
                params_list[idx].data = full_view.reshape(self._orig_weight_shape)
            else:
                if self._bias_partition_dim is None:
                    params_list[idx].data = param.data
                else:
                    full_bias_view = _gather_logical_tensor(param,
                                                            self._output_shape,
                                                            self._bias_partition_dim,
                                                            self.mp_group,
                                                            self.tp_world_size,
                                                            name=self.name,
                                                            subparam_sizes=self._subparam_sizes)
                    params_list[idx].data = full_bias_view.reshape(self._orig_bias_shape)

    @torch.no_grad()
    def _tp_partition(self, params_list):
        weight = params_list[0]
        if weight is None:
            return

        weight_view = weight.reshape(self._logical_shape)
        partitioned_view = _partition_logical_tensor(weight_view,
                                                     self.partition_dim,
                                                     self.tp_world_size,
                                                     self.tp_index,
                                                     name=self.name,
                                                     subparam_sizes=self._subparam_sizes)
        params_list[0].data = self.move(partitioned_view.reshape(-1, partitioned_view.shape[-1])).detach()

        if params_list[1] is not None:
            if self._bias_partition_dim is None:
                params_list[1].data = self.move(params_list[1]).detach()
            else:
                bias_view = params_list[1].reshape(self._output_shape)
                bias_partitioned = _partition_logical_tensor(bias_view,
                                                             self._bias_partition_dim,
                                                             self.tp_world_size,
                                                             self.tp_index,
                                                             name=self.name,
                                                             subparam_sizes=self._subparam_sizes)
                params_list[1].data = self.move(bias_partitioned.reshape(-1)).detach()

    def _mark_uc_metadata(self):
        self._set_param_uc_meta(self.weight,
                                partition_type='column',
                                partition_dim=self.partition_dim,
                                logical_shape=self._logical_shape,
                                output_shape=self._output_shape,
                                sub_param_shape=self.shape,
                                sub_param_sizes=self._subparam_sizes,
                                target_partition_shape=self.weight.shape,
                                original_shape=self._orig_weight_shape)
        if self.bias is not None:
            self._set_param_uc_meta(
                self.bias,
                partition_type='column',
                partition_dim=self._bias_partition_dim,
                logical_shape=self._output_shape,
                output_shape=self._output_shape,
                sub_param_shape=self.shape if self._bias_partition_dim is not None else None,
                sub_param_sizes=self._subparam_sizes if self._bias_partition_dim is not None else None,
                target_partition_shape=self.bias.shape,
                original_shape=self._orig_bias_shape,
                is_bias=True,
                replicated=self._bias_partition_dim is None)


class SubParamLinearAllreduce(TensorParallel_Layer):
    """
    Row-parallel linear layer with sub-parameter support (AllReduce after forward).

    Handles cases where weights contain multiple logical sub-parameters
    that need to be partitioned separately.
    """

    def __init__(self, module, mp_group, shape, partition_dim=1, **kwargs):
        super(SubParamLinearAllreduce, self).__init__(mp_group, **kwargs)
        self.weight = module.weight
        self.bias = module.bias
        self.shape = shape
        self.partition_dim = partition_dim

        self._orig_weight_shape = tuple(module.weight.shape)
        self._orig_bias_shape = tuple(module.bias.shape) if self.bias is not None else None
        (self._logical_shape, self._output_shape, self._subparam_sizes,
         self._bias_partition_dim) = _infer_subparam_logical_shapes(self._orig_weight_shape, self.shape,
                                                                    self.partition_dim, self.name)

        if self._should_materialize_tp_partition():
            self._tp_partition([self.weight, self.bias])
        self.support_training = True
        self.config_tp_params(self.weight)
        if self.bias is not None:
            self.config_requires_grad(self.bias)
        self._mark_uc_metadata()

    def forward(self, input):
        output = torch.matmul(input, self.weight.transpose(-1, -2))
        output = RowParallel.apply(self.mp_group, output, not self.is_training_mode())
        if self.bias is not None:
            output = add_bias(output, self.bias)
        return output

    @torch.no_grad()
    def gather_params(self, params_list):
        """Gather partitioned parameters back to full size."""
        for idx, param in enumerate(params_list):
            if param is None or idx > 0:
                # don't gather bias for row parallel
                return
            params_list[idx].data_partition = param.data
            full_view = _gather_logical_tensor(param,
                                               self._logical_shape,
                                               self.partition_dim,
                                               self.mp_group,
                                               self.tp_world_size,
                                               name=self.name,
                                               subparam_sizes=self._subparam_sizes)
            params_list[idx].data = full_view.reshape(self._orig_weight_shape)

    @torch.no_grad()
    def _tp_partition(self, params_list):
        weight = params_list[0]
        if weight is None:
            return

        weight_view = weight.reshape(self._logical_shape)
        partitioned_view = _partition_logical_tensor(weight_view,
                                                     self.partition_dim,
                                                     self.tp_world_size,
                                                     self.tp_index,
                                                     name=self.name,
                                                     subparam_sizes=self._subparam_sizes)
        params_list[0].data = self.move(partitioned_view.reshape(-1, partitioned_view.shape[-1])).detach()

        # Bias is not partitioned for row parallel (it's applied after all-reduce)
        if params_list[1] is not None:
            params_list[1].data = self.move(params_list[1]).detach()

    def _mark_uc_metadata(self):
        self._set_param_uc_meta(self.weight,
                                partition_type='row',
                                partition_dim=self.partition_dim,
                                logical_shape=self._logical_shape,
                                output_shape=self._output_shape,
                                sub_param_shape=self.shape,
                                sub_param_sizes=self._subparam_sizes,
                                target_partition_shape=self.weight.shape,
                                original_shape=self._orig_weight_shape)
        if self.bias is not None:
            self._set_param_uc_meta(self.bias,
                                    partition_type='row',
                                    partition_dim=None,
                                    logical_shape=self._orig_bias_shape,
                                    output_shape=self._orig_bias_shape,
                                    original_shape=self._orig_bias_shape,
                                    target_partition_shape=self.bias.shape,
                                    is_bias=True,
                                    replicated=True)


class RMSNormalize(nn.Module):

    def __init__(self, dim=None, dtype=torch.float, eps=1e-5, weight=None):
        super(RMSNormalize, self).__init__()
        if weight is not None:
            self.weight = weight
        else:
            self.weight = nn.Parameter(torch.ones(dim, dtype=dtype, device=get_accelerator().current_device_name()))

        self.eps = eps

    def forward(self, hidden_states):
        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
        if self.weight.dtype in [torch.float16, torch.bfloat16]:
            hidden_states = hidden_states.to(self.weight.dtype)

        return hidden_states * self.weight