Use HybridMesh

jonb377 · Jul 24, 2023 · 7b90ea5 · 7b90ea5
1 parent 3799ac9
commit 7b90ea5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/pytorch/language-modeling/run_clm.py b/examples/pytorch/language-modeling/run_clm.py
@@ -488,7 +488,7 @@ def main():
             max_dim = np.argmax(param.shape)
             shape = [1] * len(param.shape)
             shape[max_dim] = num_devices
-            mesh = xs.Mesh(device_ids, tuple(shape))
+            mesh = xs.HybridMesh(ici_mesh_shape=tuple(shape))
             xs.mark_sharding(param, mesh, range(len(param.shape)))