Extracts host addressable data from each Tensor in `global_arrays`. Args: global_arrays: A nested Tensor. Each leaf Tensor must be uniformly partitioned across each dim. partition: Deprecated. Returns: A nested Tensor with the same structure as `global_a
(
global_arrays: Nested[Tensor],
*,
partition: Optional[DataPartitionType] = DataPartitionType.FULL,
)
| 939 | seq_axis_names = tuple(name for name in mesh.axis_names if name in _SEQ_AXES) |
| 940 | return PartitionSpec(batch_axis_names, seq_axis_names) |
| 941 | elif isinstance(partition, PartitionSpec): |
| 942 | return partition |
| 943 | elif isinstance(partition, dict): |
| 944 | return {k: data_partition_type_to_spec(v) for k, v in partition.items()} |
| 945 | else: |
| 946 | raise NotImplementedError(f"Unsupported partition: {partition}") |
| 947 | |
| 948 | |
| 949 | def host_to_global_array( |
| 950 | host_arrays: Nested[Union[np.ndarray, Tensor]], |
| 951 | *, |
| 952 | partition: Union[Nested[PartitionSpec], DataPartitionType] = DataPartitionType.FULL, |
| 953 | ) -> Nested[Tensor]: |
| 954 | """Converts the given host device arrays to global device arrays. |
| 955 | |
| 956 | Must be called within the context of a Mesh. |
| 957 | |
| 958 | Args: |
| 959 | host_arrays: A nested tree of device arrays in host memory. Usually these present the |
| 960 | per-host portion of the global input batch. We currently assume that per-host portions |
| 961 | form a uniform sharding across the batch. |
| 962 | partition: How the global array should be partitioned. |
| 963 | |
| 964 | Returns: |
| 965 | A nested tree with the same structure as `host_arrays`, but global device arrays at the |
| 966 | leaves. Each global device array is partitioned according to `partition`. |
| 967 | |
| 968 | Raises: |
| 969 | NotImplementedError: If the given `partition` type is not supported. |
| 970 | """ |
| 971 | if isinstance(partition, DataPartitionType): |
| 972 | logging.log_first_n( |
| 973 | logging.WARNING, |
| 974 | "Passing DataPartitionType is deprecated. Please specify a PartitionSpec directly.", |
| 975 | n=1, |
| 976 | ) |
| 977 | |
| 978 | mesh = thread_resources.env.physical_mesh |
| 979 | partition_specs = complete_partition_spec_tree( |
| 980 | jax.tree_util.tree_structure(host_arrays), |
| 981 | data_partition_type_to_spec(partition), |
| 982 | ) |
| 983 | process_count = jax.process_count() |
| 984 | |
| 985 | def make_array(x: np.ndarray, partition_spec: PartitionSpec): |
| 986 | if partition == DataPartitionType.FULL: |
| 987 | global_shape = (x.shape[0] * process_count, *x.shape[1:]) |
| 988 | elif partition == DataPartitionType.REPLICATED: |
| 989 | global_shape = (x.shape[0], *x.shape[1:]) |
| 990 | elif partition == DataPartitionType.BATCH: |
| 991 | if not any(axis in mesh.shape for axis in _BATCH_AXES): |
| 992 | raise ValueError( |
| 993 | f"Mesh {mesh} does not have 'data' or 'fsdp' axis for batch partitioning." |
| 994 | ) |
| 995 | batch_count = math.prod(mesh.shape.get(axis, 1) for axis in _BATCH_AXES) |
| 996 | global_shape = (x.shape[0] * batch_count, *x.shape[1:]) |
| 997 | elif isinstance(partition, (PartitionSpec, dict)): |
| 998 | global_shape = None # Allow jax to infer. |