`.typings/mlx/nn/layers/distributed.pyi`

228 lines
"""
This type stub file was generated by pyright.
"""

from functools import lru_cache
from typing import Callable, Optional, Union

import mlx.core as mx
from base import Module
from mlx.nn.layers.linear import Linear

@lru_cache
def sum_gradients(
    group: mx.distributed.Group,
) -> Callable[..., mx.array]:  # -> Callable[..., Any] | Callable[..., array]:
    ...
def shard_inplace(
    module: Module,
    sharding: str,
    *,
    segments: Union[int, list[int]] = ...,
    group: Optional[mx.distributed.Group] = ...,
) -> None:
    """Shard a module in-place by updating its parameter dictionary with the
    sharded parameter dictionary.

    The ``sharding`` argument can be any callable that given the path and the
    weight returns the sharding axis and optionally also the segments that
    comprise the unsharded weight. For instance if the weight is a fused QKV
    matrix the segments should be 3.

    .. note::
        The module doesn't change so in order for distributed communication to
        happen the module needs to natively support it and for it to be enabled.

    Args:
        module (Module): The parameters of this module will be sharded
            in-place.
        sharding (str or callable): One of "all-to-sharded" and
            "sharded-to-all" or a callable that returns the sharding axis and
            segments.
        segments (int or list): The segments to use if ``sharding`` is a
            string. Default: ``1``.
        group (mlx.core.distributed.Group): The distributed group to shard
            across. If not set, the global group will be used. Default: ``None``.
    """

def shard_linear(
    module: Module,
    sharding: str,
    *,
    segments: Union[int, list[int]] = ...,
    group: Optional[mx.distributed.Group] = ...,
) -> Linear:
    """Create a new linear layer that has its parameters sharded and also
    performs distributed communication either in the forward or backward
    pass.

    .. note::
        Contrary to ``shard_inplace``, the original layer is not changed but a
        new layer is returned.

    Args:
        module (Module): The linear layer to be sharded.
        sharding (str): One of "all-to-sharded" and
            "sharded-to-all" that defines the type of sharding to perform.
        segments (int or list): The segments to use. Default: ``1``.
        group (mlx.core.distributed.Group): The distributed group to shard
            across. If not set, the global group will be used. Default: ``None``.
    """

class AllToShardedLinear(Module):
    """Each member of the group applies part of the affine transformation such
    that the result is sharded across the group.

    The gradients are automatically aggregated from each member of the group.

    Args:
        input_dims (int): The dimensionality of the input features
        output_dims (int): The dimensionality of the output features
        bias (bool, optional): If set to ``False`` the the layer will not use a
            bias. Default is ``True``.
        group (mx.distributed.Group, optional): The sharding will happen across
            this group. If not set then the global group is used. Default is
            ``None``.
    """
    def __init__(
        self,
        input_dims: int,
        output_dims: int,
        bias: bool = ...,
        group: Optional[mx.distributed.Group] = ...,
    ) -> None: ...
    def __call__(self, x: mx.array) -> mx.array: ...
    @classmethod
    def from_linear(
        cls,
        linear_layer: Module,
        *,
        segments: Union[int, list[int]] = ...,
        group: Optional[mx.distributed.Group] = ...,
    ) -> AllToShardedLinear: ...

class ShardedToAllLinear(Module):
    """Each member of the group applies part of the affine transformation and
    then aggregates the results.

    All nodes will have the same exact result after this layer.

    :class:`ShardedToAllLinear` provides a classmethod :meth:`from_linear` to
    convert linear layers to sharded :obj:`ShardedToAllLinear` layers.

    Args:
        input_dims (int): The dimensionality of the input features
        output_dims (int): The dimensionality of the output features
        bias (bool, optional): If set to ``False`` the the layer will not use a
            bias. Default is ``True``.
        group (mx.distributed.Group, optional): The sharding will happen across
            this group. If not set then the global group is used. Default is
            ``None``.
    """
    def __init__(
        self,
        input_dims: int,
        output_dims: int,
        bias: bool = ...,
        group: Optional[mx.distributed.Group] = ...,
    ) -> None: ...
    def __call__(self, x: mx.array) -> mx.array: ...
    @classmethod
    def from_linear(
        cls,
        linear_layer: Module,
        *,
        segments: Union[int, list[int]] = ...,
        group: Optional[mx.distributed.Group] = ...,
    ) -> ShardedToAllLinear: ...

class QuantizedAllToShardedLinear(Module):
    """Each member of the group applies part of the affine transformation with
    a quantized matrix such that the result is sharded across the group.

    It is the quantized equivalent of :class:`AllToShardedLinear`.
    Similar to :class:`QuantizedLinear` its parameters are frozen and
    will not be included in any gradient computation.

    Args:
        input_dims (int): The dimensionality of the input features.
        output_dims (int): The dimensionality of the output features.
        bias (bool, optional): If set to ``False`` then the layer will not use
            a bias. Default: ``True``.
        group_size (int, optional): The group size to use for the quantized
            weight. See :func:`~mlx.core.quantize`. Default: ``64``.
        bits (int, optional): The bit width to use for the quantized weight.
            See :func:`~mlx.core.quantize`. Default: ``4``.
        group (mx.distributed.Group, optional): The sharding will happen across
            this group. If not set then the global group is used. Default is
            ``None``.
    """
    def __init__(
        self,
        input_dims: int,
        output_dims: int,
        bias: bool = ...,
        group_size: int = ...,
        bits: int = ...,
        group: Optional[mx.distributed.Group] = ...,
    ) -> None: ...
    def unfreeze(self, *args, **kwargs) -> None:
        """Wrap unfreeze so that we unfreeze any layers we might contain but
        our parameters will remain frozen."""

    def __call__(self, x: mx.array) -> mx.array: ...
    @classmethod
    def from_quantized_linear(
        cls,
        quantized_linear_layer: Module,
        *,
        segments: Union[int, list[int]] = ...,
        group: Optional[mx.distributed.Group] = ...,
    ) -> QuantizedAllToShardedLinear: ...

class QuantizedShardedToAllLinear(Module):
    """Each member of the group applies part of the affine transformation using
    the quantized matrix and then aggregates the results.

    All nodes will have the same exact result after this layer.

    It is the quantized equivalent of :class:`ShardedToAllLinear`.
    Similar to :class:`QuantizedLinear` its parameters are frozen and
    will not be included in any gradient computation.

    Args:
        input_dims (int): The dimensionality of the input features.
        output_dims (int): The dimensionality of the output features.
        bias (bool, optional): If set to ``False`` then the layer will not use
            a bias. Default: ``True``.
        group_size (int, optional): The group size to use for the quantized
            weight. See :func:`~mlx.core.quantize`. Default: ``64``.
        bits (int, optional): The bit width to use for the quantized weight.
            See :func:`~mlx.core.quantize`. Default: ``4``.
        group (mx.distributed.Group, optional): The sharding will happen across
            this group. If not set then the global group is used. Default is
            ``None``.
    """
    def __init__(
        self,
        input_dims: int,
        output_dims: int,
        bias: bool = ...,
        group_size: int = ...,
        bits: int = ...,
        group: Optional[mx.distributed.Group] = ...,
    ) -> None: ...
    def unfreeze(self, *args, **kwargs):  # -> None:
        """Wrap unfreeze so that we unfreeze any layers we might contain but
        our parameters will remain frozen."""

    def __call__(self, x: mx.array) -> mx.array: ...
    @classmethod
    def from_quantized_linear(
        cls,
        quantized_linear_layer: Module,
        *,
        segments: Union[int, list[int]] = ...,
        group: Optional[mx.distributed.Group] = ...,
    ) -> QuantizedShardedToAllLinear: ...