fix sharding

Author: mailvijayasingh

src/maxdiffusion/configs/ltx_video.yml

@@ -12,7 +12,7 @@ output_dir: 'ltx-video-output'
 save_config_to_gcs: False
 
 #parallelism
-mesh_axes: ['data', 'fsdp', 'tensor']
+mesh_axes: ['data', 'fsdp', 'tensor', 'fsdp_transpose', 'expert', 'tensor_transpose', 'tensor_sequence', 'sequence']
 logical_axis_rules: [
                       ['batch', 'data'],
                       ['activation_batch', ['data','fsdp']],
@@ -25,13 +25,19 @@ logical_axis_rules: [
                       ['out_channels', 'tensor'],
                       ['conv_out', 'fsdp'],
                     ]
-data_sharding: [['data', 'fsdp', 'tensor']]
+data_sharding: [['data', 'fsdp', 'tensor', 'fsdp_transpose', 'expert', 'tensor_transpose', 'tensor_sequence', 'sequence']]
 dcn_data_parallelism: 1  # recommended DCN axis to be auto-sharded
 dcn_fsdp_parallelism: -1
 dcn_tensor_parallelism: 1
+
 ici_data_parallelism: -1
 ici_fsdp_parallelism: 1  # recommended ICI axis to be auto-sharded
 ici_tensor_parallelism: 1
+ici_fsdp_transpose_parallelism: 1
+ici_sequence_parallelism: 1
+ici_tensor_transpose_parallelism: 1
+ici_expert_parallelism: 1
+ici_sequence_parallelism: 1
 
 
 

src/maxdiffusion/max_utils.py

@@ -251,46 +251,88 @@ def fill_unspecified_mesh_axes(parallelism_vals, target_product, parallelism_typ
 
   return parallelism_vals
 
-
-def create_device_mesh(config, devices=None, logging=True):
+def create_device_mesh(config, devices=None):
   """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_slices = 1
+  # if config.inference_benchmark_test else config.num_slices
   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,
-  ]
+  # multi_slice_env = num_slices > 1
 
   # 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}")
+  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}")
 
   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):
   """Unboxes the flax.LogicallyPartitioned pieces in a train state.
 

src/maxdiffusion/models/ltx_video/transformers/attention.py

@@ -631,27 +631,27 @@ def partial_flash_attention(q, k, v, q_segment_ids, kv_segment_ids):
                 raise ValueError(f"Expected mask with 2 dims, got {q_segment_ids.ndim}.")
             # Based on: ("activation_kv_batch", "activation_kv_heads", "activation_length", "activation_kv_head_dim")
             # Computation of the spec based on the logical constraints can be found in logical_axes_to_spec.py.
+            qkvo_sharding_spec = jax.sharding.PartitionSpec(
+                ("data", "fsdp", "fsdp_transpose", "expert"),
+                ("tensor", "tensor_transpose", "sequence", "tensor_sequence"),
+                None,
+                None,
+            )
             # qkvo_sharding_spec = jax.sharding.PartitionSpec(
             #     ("data", "fsdp", "fsdp_transpose", "expert"),
             #     ("tensor", "tensor_transpose", "sequence", "tensor_sequence"),
             #     None,
             #     None,
             # )
             # qkvo_sharding_spec = jax.sharding.PartitionSpec(
-            #     ("data", "fsdp", "fsdp_transpose", "expert"),
-            #     ("tensor", "tensor_transpose", "sequence", "tensor_sequence"),
+            #     None,
+            #     None,
             #     None,
             #     None,
             # )
-            qkvo_sharding_spec = jax.sharding.PartitionSpec(
-                None,
-                None,
-                None,
-                None,
-            )
             #Based on: ("activation_kv_batch", "activation_length")
-            # qkv_segment_ids_spec = jax.sharding.PartitionSpec(("data", "fsdp", "fsdp_transpose", "expert"), "sequence")
-            qkv_segment_ids_spec = jax.sharding.PartitionSpec(None, None)
+            qkv_segment_ids_spec = jax.sharding.PartitionSpec(("data", "fsdp", "fsdp_transpose", "expert"), "sequence")
+            # qkv_segment_ids_spec = jax.sharding.PartitionSpec(None, None)
             wrapped_flash_attention = shard_map(
                 partial_flash_attention,
                 mesh=sharding_mesh,

src/maxdiffusion/models/ltx_video/xora_v1.2-13B-balanced-128.json

@@ -1,5 +1,5 @@
 {
-    "ckpt_path": "/mnt/disks/diffusionproj/jax_weights",
+    "ckpt_path": "/dev/shm/ltx_converted",
     "activation_fn": "gelu-approximate",
     "attention_bias": true,
     "attention_head_dim": 128,

src/maxdiffusion/pyconfig.py

@@ -41,6 +41,21 @@ def string_to_bool(s: str) -> bool:
 config = None
 
 
+def create_parallelisms_list(raw_keys):
+  ici_parallelism = [
+      raw_keys["ici_data_parallelism"],
+      raw_keys["ici_fsdp_parallelism"],
+      raw_keys["ici_fsdp_transpose_parallelism"],
+      raw_keys["ici_sequence_parallelism"],
+      raw_keys["ici_tensor_parallelism"],
+      raw_keys["ici_tensor_transpose_parallelism"],
+      raw_keys["ici_expert_parallelism"],
+      raw_keys["ici_sequence_parallelism"],
+  ]
+  raw_keys["ici_parallelism"] = ici_parallelism
+  return raw_keys
+
+
 def print_system_information():
   max_logging.log(f"System Information: Jax Version: {jax.__version__}")
   max_logging.log(f"System Information: Jaxlib Version: {jax.lib.__version__}")
@@ -154,6 +169,7 @@ 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)
 
 
 def get_num_slices(raw_keys):