changed sharding back

Author: Serenagu525

src/maxdiffusion/configs/ltx_video.yml

@@ -40,20 +40,22 @@ output_dir: 'ltx-video-output'
 save_config_to_gcs: False
 
 #parallelism
-mesh_axes: ['data', 'fsdp', 'tensor', 'fsdp_transpose', 'expert', 'tensor_transpose', 'tensor_sequence', 'sequence']
+mesh_axes: ['data', 'fsdp', 'tensor']
 logical_axis_rules: [
                       ['batch', 'data'],
+                      ['activation_heads', 'fsdp'],
                       ['activation_batch', ['data','fsdp']],
-                      ['activation_heads', 'tensor'],
                       ['activation_kv', 'tensor'],
                       ['mlp','tensor'],
                       ['embed','fsdp'],
                       ['heads', 'tensor'],
+                      ['norm', 'fsdp'],
                       ['conv_batch', ['data','fsdp']],
                       ['out_channels', 'tensor'],
                       ['conv_out', 'fsdp'],
+                      ['conv_in', 'fsdp']
                     ]
-data_sharding: [['data', 'fsdp', 'tensor', 'fsdp_transpose', 'expert', 'tensor_transpose', 'tensor_sequence', 'sequence']]
+data_sharding: [['data', 'fsdp', 'tensor']]
 dcn_data_parallelism: 1  # recommended DCN axis to be auto-sharded
 dcn_fsdp_parallelism: -1
 dcn_tensor_parallelism: 1

src/maxdiffusion/max_utils.py

@@ -609,4 +609,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/models/ltx_video/transformers/attention.py

@@ -622,14 +622,21 @@ 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,
+          ("data", "fsdp", "tensor"),
           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("fsdp", None)
+      # qkv_segment_ids_spec = jax.sharding.PartitionSpec(None, None)
       wrapped_flash_attention = shard_map(
           partial_flash_attention,
           mesh=sharding_mesh,

src/maxdiffusion/pyconfig.py

@@ -226,4 +226,4 @@ def initialize(argv, **kwargs):
 if __name__ == "__main__":
   initialize(sys.argv)
   print(config.steps)
-  r = range(config.steps)
+  r = range(config.steps)

src/maxdiffusion/tests/ltx_transformer_step_test.py

@@ -104,7 +104,7 @@ def test_one_step_transformer(self):
     devices_array = create_device_mesh(config)
     mesh = Mesh(devices_array, config.mesh_axes)
     base_dir = os.path.dirname(__file__)
-    config_path = os.path.join(base_dir,  "../models/ltx_video/xora_v1.2-13B-balanced-128.json")
+    config_path = os.path.join(base_dir, "../models/ltx_video/xora_v1.2-13B-balanced-128.json")
 
     with open(config_path, "r") as f:
       model_config = json.load(f)