Method test_every_other_process

axlearn/common/utils_test.py:2091–2137

Test a case where every other process produces a slice. We build the array directly with `global_to_host_array`.

(self, divisor: int)

        self,
        logical_mesh: Union[MeshShape, HybridMeshShape],
        expected: Optional[Union[MeshShape, Exception]] = None,
        is_custom: bool = False,
    ):
        slice_physical_mesh = (16, 16, 1)
        num_slices = 2
        coords = [
            (x, y, z)
            for x in range(slice_physical_mesh[0])
            for y in range(slice_physical_mesh[1])
            for z in range(slice_physical_mesh[2])
        ]
        devices = [
            DummyMultiSliceTpuDevice(
                platform="tpu",
                device_kind="TPU v5 lite",
                process_index=(len(coords) * slice_index + ix) // 4,
                coords=coord,
                slice_index=slice_index,
            )
            for ix, coord in enumerate(coords)
            for slice_index in range(num_slices)
        ]
        if isinstance(expected, Exception):
            with self.assertRaisesRegex(type(expected), str(expected)):
                create_device_mesh(mesh_shape=logical_mesh, devices=devices)
        else:
            # pylint: disable-next=protected-access
            custom_mesh_fn = mock.Mock(wraps=utils._reshape_mesh_to_rings)
            with mock.patch.object(utils, "_reshape_mesh_to_rings", custom_mesh_fn):
                device_mesh = create_device_mesh(mesh_shape=logical_mesh, devices=devices)
            if is_custom:
                self.assertEqual(custom_mesh_fn.call_count, num_slices)
            else:
                custom_mesh_fn.assert_not_called()
            # Check that the constructed mesh has the expected shape.
            self.assertEqual(expected or logical_mesh, device_mesh.shape)

            # Check that the sub_mesh along the first non-singleton mesh axis only contains devices
            # from one of the slices.
            mesh_shape = device_mesh.shape
            for dim in mesh_shape:
                if dim != 1:
                    break
                device_mesh = device_mesh[0]
            for ix, sub_mesh in enumerate(device_mesh):
                self.assertTrue(all(el.slice_index == ix for el in sub_mesh.flatten()))

    @parameterized.parameters(
        {"logical_mesh": (8, 2, 4)},
        {"logical_mesh": (16, 4)},