Test a case where each process produces a slice.
(self)
| 1981 | ) |
| 1982 | def test_create_device_mesh_multi_slice_tpuv4( |
| 1983 | self, |
| 1984 | logical_mesh: Union[MeshShape, HybridMeshShape], |
| 1985 | expected: Optional[Union[MeshShape, Exception]] = None, |
| 1986 | ): |
| 1987 | slice_physical_mesh = (4, 4, 1) |
| 1988 | num_slices = 2 |
| 1989 | coords = [ |
| 1990 | (x, y, z) |
| 1991 | for x in range(slice_physical_mesh[0]) |
| 1992 | for y in range(slice_physical_mesh[1]) |
| 1993 | for z in range(slice_physical_mesh[2]) |
| 1994 | ] |
| 1995 | devices = [ |
| 1996 | DummyMultiSliceTpuDevice( |
| 1997 | platform="tpu", |
| 1998 | device_kind="TPU v4", |
| 1999 | process_index=(len(coords) * slice_index + ix) // 4, |
| 2000 | coords=coord, |
| 2001 | slice_index=slice_index, |
| 2002 | ) |
| 2003 | for ix, coord in enumerate(coords) |
| 2004 | for slice_index in range(num_slices) |
| 2005 | ] |
| 2006 | if isinstance(expected, Exception): |
| 2007 | with self.assertRaisesRegex(type(expected), str(expected)): |
| 2008 | create_device_mesh(mesh_shape=logical_mesh, devices=devices) |
nothing calls this directly
no test coverage detected