Function pipeline_auto_parallel

backend/python/mlx-distributed/sharding.py:97–136

Apply pipeline parallelism to a model. Each rank only keeps its slice of layers. The first layer receives from the previous rank, and the last layer sends to the next rank. Args: model: The MLX model (must have model.layers or similar) group: The distributed group

(model, group, start_layer=None, end_layer=None)

def pipeline_auto_parallel(model, group, start_layer=None, end_layer=None):
    """Apply pipeline parallelism to a model.

    Each rank only keeps its slice of layers.  The first layer receives from
    the previous rank, and the last layer sends to the next rank.

    Args:
        model: The MLX model (must have model.layers or similar)
        group: The distributed group
        start_layer: First layer index for this rank (auto-computed if None)
        end_layer: Last layer index (exclusive) for this rank (auto-computed if None)
    """
    rank = group.rank()
    world_size = group.size()

    inner = get_inner_model(model)
    layers = list(get_layers(inner))
    total_layers = len(layers)

    if start_layer is None or end_layer is None:
        layers_per_rank = total_layers // world_size
        remainder = total_layers % world_size
        start_layer = rank * layers_per_rank + min(rank, remainder)
        end_layer = start_layer + layers_per_rank + (1 if rank < remainder else 0)

    layers = layers[start_layer:end_layer]
    for layer in layers:
        mx.eval(layer)

    # Wrap first and last layers
    layers[0] = PipelineFirstLayer(layers[0], rank, group=group)
    layers[-1] = PipelineLastLayer(layers[-1], rank, world_size, group=group)

    # Replace layers on the inner model
    if hasattr(inner, "layers"):
        inner.layers = layers
    elif hasattr(inner, "h"):
        inner.h = layers

    return model