Tests roundtrip host-to-global and global-to-host with fixed process shape.
(
self,
platform: str,
mesh_shape: tuple[int, int],
process_shape: Sequence[int],
partition: Union[DataPartitionType, PartitionSpec],
)
| 113 | # NOTE: while annotated with `for_8_devices`, this runs on other configurations. |
| 114 | @pytest.mark.for_8_devices |
| 115 | def test_fixed_process_shape( |
| 116 | self, |
| 117 | platform: str, |
| 118 | mesh_shape: tuple[int, int], |
| 119 | process_shape: Sequence[int], |
| 120 | partition: Union[DataPartitionType, PartitionSpec], |
| 121 | ): |
| 122 | """Tests roundtrip host-to-global and global-to-host with fixed process shape.""" |
| 123 | |
| 124 | mesh_shape = infer_mesh_shape(mesh_shape) |
| 125 | if not _is_supported(platform=platform, mesh_shape=mesh_shape): |
| 126 | self.skipTest("Unsupported platform/mesh.") |
| 127 | |
| 128 | devices = mesh_utils.create_device_mesh(mesh_shape, allow_split_physical_axes=True) |
| 129 | mesh = jax.sharding.Mesh(devices, ("data", "model")) |
| 130 | |
| 131 | partition = data_partition_type_to_spec(partition) |
| 132 | sharding = jax.NamedSharding(mesh, partition) |
| 133 | |
| 134 | # Number of dims should match number of partitioned axes. |
| 135 | if len(process_shape) < len(partition): |
| 136 | self.skipTest("Incompatible process_shape/partition.") |
| 137 | |
| 138 | # Infer global shape from local_shape and number of processes. |
| 139 | global_shape = local_to_global_shape(sharding, process_shape) |
| 140 | # Partition should divide global_shape evenly. |
| 141 | partitions = _infer_num_partitions(global_shape, mesh=mesh, partition=partition) |
| 142 | if any(dim % num_parts != 0 for dim, num_parts in zip(global_shape, partitions)): |
| 143 | self.skipTest("Incompatible global_shape/partitioning.") |
| 144 | |
| 145 | with mesh: |
| 146 | host_arrays = dict( |
| 147 | x=jax.random.uniform(jax.random.PRNGKey(jax.process_count()), shape=process_shape) |
| 148 | ) |
| 149 | |
| 150 | global_arrays = host_to_global_device_array(host_arrays, partition=partition) |
| 151 | for path, value in flatten_items(global_arrays): |
| 152 | self.assertEqual(tuple(global_shape), value.shape, msg=path) |
| 153 | global_arrays["y"] = 2 * global_arrays["x"] |
| 154 | restored_host_arrays = global_to_host_array(global_arrays) |
| 155 | for path, restored_value in flatten_items(restored_host_arrays): |
| 156 | self.assertEqual(tuple(process_shape), restored_value.shape, msg=path) |
| 157 | |
| 158 | # "x" and "y" are partitioned consistently. |
| 159 | np.testing.assert_array_equal(restored_host_arrays["y"], 2 * restored_host_arrays["x"]) |
| 160 | |
| 161 | # Check round-trip equality of host_to_global_device_array and global_to_host_array. |
| 162 | np.testing.assert_array_equal(host_arrays["x"], restored_host_arrays["x"]) |
| 163 | |
| 164 | @parameterized.product( |
| 165 | platform=["cpu", "tpu"], |
nothing calls this directly
no test coverage detected