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

Code before:
        return outputs

    def shard_tensor(
        self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None
    ) -> torch.Tensor:
        # If only 1 dim, shard this one (usually it's a `bias`)
        dim = param.dim() if isinstance(param, torch.Tensor) else len(param.get_shape())
        if dim == 1:
            parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -1, tensor_idx)
        else:
            parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -2, tensor_idx)
        return parameter.to(device=device, dtype=dtype)

    def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]:
        world_size = self.device_mesh.size()
        shape = list(full_shape)
        # Colwise shards dim -2, but 1D tensors (bias) shard on dim -1
        dim = -1 if len(shape) == 1 else -2
        dim = len(shape) + dim if dim < 0 else dim
        shard_size = math.ceil(shape[dim] / world_size)
        start = self.rank * shard_size
        end = min(start + shard_size, shape[dim])
        shape[dim] = end - start
        return tuple(shape)

class AllReduce(TensorParallelLayer):
    """
    Column-wise parallel: weight is sharded on dim -2 (output features).
    Forward: input replicated -> output sharded on last dim.
    If gather_output=True, output is all-gathered to produce full tensor.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def _prepare_output_fn(self, mod, outputs, device_mesh):
        return all_reduce_forward(outputs, device_mesh)



class RowwiseParallel(TensorParallelLayer):
    """
    Row-wise parallel: weight is sharded on dim -1 (input features).
    Forward: input (optionally split) -> output partial -> all-reduce to replicate.

    Args:
        split_input: If True, splits replicated input before matmul. Use when input
                     comes from a non-parallelizable operation (chunk/slice).
                     Default False (expects pre-sharded input from colwise layer).
    """


Review comment (suggestion):
we should rename it else it becomes hard to understand given we have

```
def all_reduce_forward(x, device_mesh):
    """All-reduce forward, identity backward. Use after rowwise layers."""
    return _AllReduceForward.apply(x, device_mesh)
```

Code after:

    def shard_tensor(
        self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None
    ) -> torch.Tensor:
        # If only 1 dim, shard this one (usually it's a `bias`)
        dim = param.dim() if isinstance(param, torch.Tensor) else len(param.get_shape())
        if dim == 1:
            parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -1)
        else:
            parameter = get_tensor_shard(param, self.empty_param, self.device_mesh, self.rank, -2)
        return parameter.to(device=device, dtype=dtype)

    def get_expected_sharded_shape(self, full_shape: tuple[int, ...] | torch.Size) -> tuple[int, ...]:
        world_size = self.device_mesh.size()
        shape = list(full_shape)
        # Colwise shards dim -2, but 1D tensors (bias) shard on dim -1
        dim = -1 if len(shape) == 1 else -2
        dim = len(shape) + dim if dim < 0 else dim
        shard_size = math.ceil(shape[dim] / world_size)
        start = self.rank * shard_size
        end = min(start + shard_size, shape[dim])
        shape[dim] = end - start
        return tuple(shape)


class RowwiseParallel(TensorParallelLayer):
    """
    Row-wise parallel: weight is sharded on dim -1 (input features).
    Forward: input (optionally split) -> output partial -> all-reduce to replicate.

    Args:
        split_input: If True, splits replicated input before matmul. Use when input
                     comes from a non-parallelizable operation (chunk/slice).
                     Default False (expects pre-sharded input from colwise layer).
    """

    def __init__(self, split_input: bool = False, **kwargs):
        super().__init__(**kwargs)
        self.split_input = split_input

    def _prepare_input_fn(self, mod, inputs, device_mesh):
        if hasattr(mod, "bias") and mod.bias is not None:
            mod._bias = mod.bias
            mod.bias = None

        input_tensor = inputs[0] if inputs else inputs

        if self.split_input:
            # Input is replicated, split it to match sharded weight
            return split(input_tensor, device_mesh)
        return input_tensor