Repository: Physical-Intelligence/openpi
File: src/openpi/training/sharding.py (Python)

Code before:
import logging

from flax import linen as nn
import jax
import numpy as np

import openpi.shared.array_typing as at

# logical sharding annotations
BATCH_AXIS = "batch_axis"


def ensure_data_parallel_first_dim(x: at.ArrayLike) -> at.ArrayLike:
    return nn.with_logical_constraint(x, (BATCH_AXIS,))


def fsdp_sharding(
    pytree,
    mesh: jax.sharding.Mesh,
    *,
    min_size_mbytes: int = 4,  # 4 MiB
    log: bool = False,
    fsdp_dim_name: str = "model",
):
    """Apply FSDP sharding to a pytree of arrays based on the mesh shape.

    Args:
        pytree: A pytree to be apply sharding specified by the mesh, note that only array types (eg. contains .shape attr)
          will be considered for sharding.
        mesh: The mesh being used for applying sharding on to pytree.
        min_size_mbytes: The minimum size of the array in MiB to be considered for sharding, any array smaller than this
          will be replicated.
        log: If true, will log the sharding decisions for arrays that are being considered for sharding.
        fsdp_dim_name: The name of the dimension to shard on, this is the name of the dimension in the mesh shape. default
          is "model" which is consistent with the mesh created in the main function.

    Returns:
        The sharded pytree.

Review comment (performance):
The name is breaking the abstraction. The code is only responsible for saying that the first dimension needs to be sharded across the logical `BATCH_AXIS`. The logic that actually enforces the "data parallelism" lives somewhere else and depending on what the logic does, the actual operation may change. 

I don't have a great suggestion but I am also worried that the code as-is may be misleading. Happy to discuss offline.



Code after:
import logging

from flax import linen as nn
import jax
import numpy as np

import openpi.shared.array_typing as at

# logical sharding annotations
BATCH_AXIS = "batch_axis"


def annotate_batch_axis_sharding_on_first_dim(x: at.ArrayLike) -> at.ArrayLike:
    """Annotate and apply compile constraint that the first dimension of the array is annotated as BATCH_AXIS.
    Training step compilation will then be able to map the logical BATCH_AXIS to the defined mesh axis.
    """
    return nn.with_logical_constraint(x, (BATCH_AXIS,))


def fsdp_sharding(
    pytree,
    mesh: jax.sharding.Mesh,
    *,
    min_size_mbytes: int = 4,  # 4 MiB
    log: bool = False,
    fsdp_dim_name: str = "model",
):
    """Apply FSDP sharding to a pytree of arrays based on the mesh shape.

    Args:
        pytree: A pytree to be apply sharding specified by the mesh, note that only array types (eg. contains .shape attr)
          will be considered for sharding.
        mesh: The mesh being used for applying sharding on to pytree.
        min_size_mbytes: The minimum size of the array in MiB to be considered for sharding, any array smaller than this
          will be replicated.
        log: If true, will log the sharding decisions for arrays that are being considered for sharding.
        fsdp_dim_name: The name of the dimension to shard on, this is the name of the dimension in the mesh shape. default
          is "model" which is consistent with the mesh created in the main function.