Function host_to_global_array

axlearn/common/utils.py:842–900

Converts the given host device arrays to global device arrays. Must be called within the context of a Mesh. Args: host_arrays: A nested tree of device arrays in host memory. Usually these present the per-host portion of the global input batch. We currently assume that p

(
    host_arrays: Nested[Union[np.ndarray, Tensor]],
    *,
    partition: Union[Nested[PartitionSpec], DataPartitionType] = DataPartitionType.FULL,
)

            f"corresponding value, got specification {partition_spec_tree} "
            f"for value tree {treedef}. Original ValueError: {err}"
        ) from None
    axes = [None if a is proxy else a for a in axes]
    assert len(axes) == treedef.num_leaves, (
        f"({len(axes)} vs. {treedef.num_leaves}) {axes} {treedef}"
    )
    return jax.tree_util.tree_unflatten(treedef, axes)


def input_partition_spec() -> PartitionSpec:
    """Returns partition spec for the input batch.

    We partition the inputs along all axes. For example, if the mesh has shape (64, 4) and axis
    names of ("data", "model"), the partition spec will be (("data", "model"), None...) so that the
    batch axis of every global tensor will be partitioned 256 (= 64 * 4) ways.

    Must be called within the context of a Mesh.
    """
    mesh = thread_resources.env.physical_mesh
    return PartitionSpec(
        mesh.axis_names,
    )


# Key associated with per-example dataset dispatch index tensor, indicating which logical
# batch index the example maps to.
PHYSICAL_TO_LOGICAL_DISPATCH_KEY = "__physical_to_logical_batch_dispatch"


def dispatch_input_batch(
    input_batch: NestedTensor, *, batch_axis_names: Union[str, Sequence[str]] = "data"
) -> NestedTensor:
    """Constrains all leaf values in the input batch, then (optionally) dispatches examples
    to a subset along the batch axis.

    The dispatchings are applied to all nested dicts which contain a special dispatching key in
    their root.

    This is deprecated in favor of `axlearn.common.input_dispatch`.

    Args:
        input_batch: The input batch, where the first dimension of each leaf is the batch dim.
        batch_axis_names: The name(s) of the batch axes.

    Returns:
        A nested tensor like the input batch, where each leaf contains
            a subset of the input batch, and has been wrapped with sharding annotations.
            N.B. some internal key-value pairs (like PHYSICAL_TO_LOGICAL_DISPATCH_KEY)
            may be dropped after use if present.
    """
    logging.log_first_n(
        logging.WARNING,
        "dispatch_input_batch is deprecated. Please use `axlearn.common.input_dispatch` instead.",
        n=1,
    )

    # Constrain the input batch.
    input_batch = jax.tree.map(
        lambda x: maybe_shard(x, PartitionSpec(batch_axis_names)), input_batch