mlx.nn.layers.distributed.fully_shard

Contents

mlx.nn.layers.distributed.fully_shard#

fully_shard(module: Module, *, group: Group | None = None, compute_dtype: Dtype | None = None) → Module#

Wrap module in a FullyShardedModule.

Parameters:

  • module (Module) – The module to wrap.

  • group (Group, optional) – The group to shard across. If not set, the global group is used. Default: None.

  • compute_dtype (Dtype, optional) – If set, the gathered parameters are cast to this dtype for the forward pass. Default: None.

Returns:

The wrapped FullyShardedModule, or module unchanged.

Contents