Tests roundtrip host-to-global and global-to-host with fixed global shape.
(
self,
platform: str,
mesh_shape: tuple[int, int],
global_shape: Sequence[int],
partition: Union[PartitionSpec, DataPartitionType],
)
| 175 | # NOTE: while annotated with `for_8_devices`, this runs on other configurations. |
| 176 | @pytest.mark.for_8_devices |
| 177 | def test_fixed_global_shape( |
| 178 | self, |
| 179 | platform: str, |
| 180 | mesh_shape: tuple[int, int], |
| 181 | global_shape: Sequence[int], |
| 182 | partition: Union[PartitionSpec, DataPartitionType], |
| 183 | ): |
| 184 | """Tests roundtrip host-to-global and global-to-host with fixed global shape.""" |
| 185 | |
| 186 | mesh_shape = infer_mesh_shape(mesh_shape) |
| 187 | if not _is_supported(platform=platform, mesh_shape=mesh_shape): |
| 188 | self.skipTest("Unsupported platform/mesh.") |
| 189 | logging.info( |
| 190 | "platform=%s mesh_shape=%s global_shape=%s data_partition=%s", |
| 191 | platform, |
| 192 | mesh_shape, |
| 193 | global_shape, |
| 194 | partition, |
| 195 | ) |
| 196 | devices = mesh_utils.create_device_mesh(mesh_shape, allow_split_physical_axes=True) |
| 197 | mesh = jax.sharding.Mesh(devices, ("data", "model")) |
| 198 | logging.info("Global mesh: %s", mesh) |
| 199 | |
| 200 | partition = data_partition_type_to_spec(partition) |
| 201 | # Number of dims should match number of partitioned axes. |
| 202 | if len(global_shape) < len(partition): |
| 203 | self.skipTest("Incompatible process_shape/partition.") |
| 204 | |
| 205 | partitions = _infer_num_partitions(global_shape, mesh=mesh, partition=partition) |
| 206 | if any(dim % num_parts != 0 for dim, num_parts in zip(global_shape, partitions)): |
| 207 | self.skipTest("Incompatible global_shape/partitioning.") |
| 208 | |
| 209 | with mesh: |
| 210 | sharding = jax.sharding.NamedSharding(mesh, partition) |
| 211 | |
| 212 | ndim = len(global_shape) |
| 213 | process_shape = [] |
| 214 | for dim in range(ndim): |
| 215 | _, num_shards = get_process_index_and_count(sharding, dim=dim, ndims=ndim) |
| 216 | process_shape.append(global_shape[dim] // num_shards) |
| 217 | |
| 218 | host_arrays = dict( |
| 219 | x=jax.random.uniform(jax.random.PRNGKey(jax.process_index()), shape=process_shape) |
| 220 | ) |
| 221 | global_arrays = host_to_global_device_array(host_arrays, partition=partition) |
| 222 | for path, value in flatten_items(global_arrays): |
| 223 | self.assertEqual(tuple(global_shape), value.shape, msg=path) |
| 224 | global_arrays["y"] = 2 * global_arrays["x"] |
| 225 | restored_host_arrays = global_to_host_array(global_arrays, partition=partition) |
| 226 | for path, restored_value in flatten_items(restored_host_arrays): |
| 227 | restored_shape = restored_value.shape |
| 228 | self.assertEqual(tuple(process_shape), restored_shape, msg=path) |
| 229 | |
| 230 | # "x" and "y" are partitioned consistently. |
| 231 | np.testing.assert_array_equal(restored_host_arrays["y"], 2 * restored_host_arrays["x"]) |
| 232 | |
| 233 | # Check round-trip equality of host_to_global_device_array and global_to_host_array. |
| 234 | np.testing.assert_array_equal(host_arrays["x"], restored_host_arrays["x"]) |
nothing calls this directly
no test coverage detected