Method test_one_per_process

axlearn/common/utils_test.py:2008–2039

Test a case where every process produces a slice. This is recommended to run on >1 process, e.g. v5e-16.

(self)

        if isinstance(expected, Exception):
            with self.assertRaisesRegex(type(expected), str(expected)):
                create_device_mesh(mesh_shape=logical_mesh, devices=devices)
        else:
            # Check that the constructed mesh has the expected shape.
            device_mesh = create_device_mesh(mesh_shape=logical_mesh, devices=devices)
            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": (2, 128, 2)},
        {"logical_mesh": (2, 16, 16)},
        # Use the first axis that divides number of granules for DCN mesh.
        {"logical_mesh": (1, 2, 16, 16)},
        # First non-singleton dim does not divide number of granules.
        {"logical_mesh": (3, 2, 16, 16), "expected": ValueError("First non-singleton")},
        # At least one ICI mesh should divide number of granules.
        {"logical_mesh": (1, 1), "expected": ValueError("At least one")},
        # Test a case where we infer -1 in ICI mesh.
        {"logical_mesh": (2, -1, 2), "expected": (2, 128, 2)},
        # Test a case where we infer -1 in DCN mesh.
        {"logical_mesh": (-1, 16, 16), "expected": (2, 16, 16)},
        # Test a basic hybrid mesh case.
        {
            "logical_mesh": HybridMeshShape(ici_mesh_shape=(1, 128, 2), dcn_mesh_shape=(2, 1, 1)),
            "expected": (2, 128, 2),
        },
        # Test that ICI mesh should respect the number of devices.