mesh edit

Author: Serenagu525

src/maxdiffusion/max_utils.py

@@ -252,86 +252,115 @@ def fill_unspecified_mesh_axes(parallelism_vals, target_product, parallelism_typ
   return parallelism_vals
 
 
-def create_device_mesh(config, devices=None):
+def create_device_mesh(config, devices=None, logging=True):
   """Creates a device mesh with each slice in its own data parallel group. If there is only one slice, uses two replicas"""
   if devices is None:
     devices = jax.devices()
   num_devices = len(devices)
-  num_slices = 1
-  # if config.inference_benchmark_test else config.num_slices
+  ##special case for ltx-video
+  if config.ici_fsdp_transpose_parallelism:
+    num_slices = 1
+    # if config.inference_benchmark_test else config.num_slices
+    num_devices_per_slice = num_devices // num_slices
+    # Find possible unspecified parallelisms
+    ici_parallelism = fill_unspecified_mesh_axes(config.ici_parallelism.copy(), num_devices_per_slice, "ICI")
+    mesh = mesh_utils.create_device_mesh(
+        ici_parallelism,
+        devices,
+    )
+    max_logging.log(f"Num_devices: {num_devices}, shape {mesh.shape}")
+
+    return mesh
+    
+  try:
+    num_slices = 1 + max([d.slice_index for d in devices])
+  except:
+    num_slices = 1
   num_devices_per_slice = num_devices // num_slices
+  max_logging.log(f"Devices: {devices} (num_devices: {num_devices})")
+
+  multi_slice_env = num_slices > 1
 
-  # multi_slice_env = num_slices > 1
+  dcn_parallelism = [
+      config.dcn_data_parallelism,
+      config.dcn_fsdp_parallelism,
+      config.dcn_tensor_parallelism,
+  ]
+  ici_parallelism = [
+      config.ici_data_parallelism,
+      config.ici_fsdp_parallelism,
+      config.ici_tensor_parallelism,
+  ]
 
   # Find possible unspecified parallelisms
-  ici_parallelism = fill_unspecified_mesh_axes(config.ici_parallelism.copy(), num_devices_per_slice, "ICI")
-
-  # allow_split_physical_axes = config.allow_split_physical_axes if config.allow_split_physical_axes else False
-
-  # if allow_split_physical_axes:
-  #   if max_utils.is_valid_custom_mesh(ici_parallelism, config.custom_mesh):
-  #     mesh = mesh_utils.create_device_mesh(
-  #         [16, 16],
-  #         devices,
-  #         contiguous_submeshes=False,
-  #         allow_split_physical_axes=False,
-  #     )
-  #     mesh = max_utils.reshape_mesh_to_rings(mesh, config.custom_mesh)
-  #     mesh = np.reshape(mesh, ici_parallelism)
-  #   else:
-  #     mesh = mesh_utils.create_device_mesh(
-  #         ici_parallelism,
-  #         devices,
-  #         contiguous_submeshes=False,
-  #         allow_split_physical_axes=allow_split_physical_axes,
-  #     )
-  # else:
-  mesh = mesh_utils.create_device_mesh(
-      ici_parallelism,
-      devices,
-  )
-  max_logging.log(f"Num_devices: {num_devices}, shape {mesh.shape}")
+  ici_parallelism = fill_unspecified_mesh_axes(ici_parallelism, num_devices_per_slice, "ICI")
+  if multi_slice_env:
+    dcn_parallelism = fill_unspecified_mesh_axes(dcn_parallelism, num_slices, "DCN")
+    mesh = mesh_utils.create_hybrid_device_mesh(ici_parallelism, dcn_parallelism, devices)
+  else:
+    mesh = mesh_utils.create_device_mesh(ici_parallelism, devices)
+
+  if logging:
+    max_logging.log(f"Decided on mesh: {mesh}")
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
 
   return mesh
 
 
-# def create_device_mesh(config, devices=None, logging=True):
-#   """Creates a device mesh with each slice in its own data parallel group. If there is only one slice, uses two replicas"""
-#   if devices is None:
-#     devices = jax.devices()
-#   num_devices = len(devices)
-#   try:
-#     num_slices = 1 + max([d.slice_index for d in devices])
-#   except:
-#     num_slices = 1
-#   num_devices_per_slice = num_devices // num_slices
-#   max_logging.log(f"Devices: {devices} (num_devices: {num_devices})")
-
-#   multi_slice_env = num_slices > 1
-
-#   dcn_parallelism = [
-#       config.dcn_data_parallelism,
-#       config.dcn_fsdp_parallelism,
-#       config.dcn_tensor_parallelism,
-#   ]
-#   ici_parallelism = [
-#       config.ici_data_parallelism,
-#       config.ici_fsdp_parallelism,
-#       config.ici_tensor_parallelism,
-#   ]
-
-#   # Find possible unspecified parallelisms
-#   ici_parallelism = fill_unspecified_mesh_axes(ici_parallelism, num_devices_per_slice, "ICI")
-#   if multi_slice_env:
-#     dcn_parallelism = fill_unspecified_mesh_axes(dcn_parallelism, num_slices, "DCN")
-#     mesh = mesh_utils.create_hybrid_device_mesh(ici_parallelism, dcn_parallelism, devices)
-#   else:
-#     mesh = mesh_utils.create_device_mesh(ici_parallelism, devices)
-
-#   if logging:
-#     max_logging.log(f"Decided on mesh: {mesh}")
-
-#   return mesh
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
 
 
 def unbox_logicallypartioned_trainstate(boxed_train_state: train_state.TrainState):
@@ -445,6 +474,7 @@ def setup_initial_state(
           config.enable_single_replica_ckpt_restoring,
       )
       if state:
+        ###!Edited
         if checkpoint_item == " ":
           state = state
         else:
@@ -655,4 +685,4 @@ def maybe_initialize_jax_distributed_system(raw_keys):
     initialize_jax_for_gpu()
     max_logging.log("Jax distributed system initialized on GPU!")
   else:
-    jax.distributed.initialize()
+    jax.distributed.initialize()

src/maxdiffusion/pyconfig.py

@@ -169,7 +169,8 @@ def user_init(raw_keys):
     raw_keys["total_train_batch_size"] = max_utils.get_global_batch_size(raw_keys["per_device_batch_size"])
     raw_keys["num_slices"] = get_num_slices(raw_keys)
     raw_keys["quantization_local_shard_count"] = get_quantization_local_shard_count(raw_keys)
-    raw_keys = create_parallelisms_list(raw_keys)
+    if "ici_fsdp_transpose_parallelism" in raw_keys:
+      raw_keys = create_parallelisms_list(raw_keys)
 
 
 def get_num_slices(raw_keys):