Create data parallelism meshes for each Inference worker.
(devices, config)
| 940 | |
| 941 | @staticmethod |
| 942 | def create_mesh(devices, config): |
| 943 | """Create data parallelism meshes for each Inference worker.""" |
| 944 | ici_parallelism = max_utils.fill_unspecified_mesh_axes(config.ici_parallelism.copy(), len(devices), "ICI") |
| 945 | devices_array = mesh_utils.create_device_mesh( |
| 946 | ici_parallelism, |
| 947 | devices, |
| 948 | contiguous_submeshes=False, |
| 949 | allow_split_physical_axes=config.allow_split_physical_axes or False, |
| 950 | ) |
| 951 | mesh = Mesh(devices_array.reshape(ici_parallelism), config.mesh_axes) |
| 952 | return mesh |
| 953 | |
| 954 | def _validate_config(self): |
| 955 | """Validate configuration parameters and check for incompatible settings.""" |