Fix embed axis and add I2V configs to TP strategy

Author: Perseus14

src/maxdiffusion/configs/base_wan_14b.yml

@@ -177,7 +177,7 @@ logical_axis_rules: [
                       ['activation_length', None],
                       ['activation_heads', 'tensor'],
                       ['mlp','tensor'],
-                      ['embed', ['fsdp', 'tensor']],
+                      ['embed', 'fsdp'],
                       ['heads', 'tensor'],
                       ['norm', 'tensor'],
                       ['conv_batch', ['data', 'fsdp']],

src/maxdiffusion/configs/base_wan_1_3b.yml

@@ -150,7 +150,7 @@ logical_axis_rules: [
                       ['activation_length', None],
                       ['activation_heads', 'tensor'],
                       ['mlp','tensor'],
-                      ['embed', ['fsdp', 'tensor']],
+                      ['embed', 'fsdp'],
                       ['heads', 'tensor'],
                       ['norm', 'tensor'],
                       ['conv_batch', ['data', 'fsdp']],

src/maxdiffusion/configs/base_wan_27b.yml

@@ -161,7 +161,7 @@ logical_axis_rules: [
                       ['activation_length', None],
                       ['activation_heads', 'tensor'],
                       ['mlp','tensor'],
-                      ['embed', ['fsdp', 'tensor']],
+                      ['embed', 'fsdp'],
                       ['heads', 'tensor'],
                       ['norm', 'tensor'],
                       ['conv_batch', ['data', 'fsdp']],

src/maxdiffusion/configs/base_wan_i2v_14b.yml

@@ -151,17 +151,17 @@ mesh_axes: ['data', 'fsdp', 'context', 'tensor']
 logical_axis_rules: [
                       ['batch', ['data', 'fsdp']],
                       ['activation_batch', ['data', 'fsdp']],
-                      ['activation_self_attn_heads', ['context', 'tensor']],
-                      ['activation_cross_attn_q_length', ['context', 'tensor']],
-                      ['activation_length', 'context'],
+                      ['activation_self_attn_heads', 'tensor'],
+                      ['activation_cross_attn_q_length', 'tensor'],
+                      ['activation_length', None],
                       ['activation_heads', 'tensor'],
                       ['mlp','tensor'],
-                      ['embed', ['context', 'fsdp']],
+                      ['embed', 'fsdp'],
                       ['heads', 'tensor'],
                       ['norm', 'tensor'],
-                      ['conv_batch', ['data', 'context', 'fsdp']],
+                      ['conv_batch', ['data', 'fsdp']],
                       ['out_channels', 'tensor'],
-                      ['conv_out', 'context'],
+                      ['conv_out', 'tensor'],
                     ]
 data_sharding: [['data', 'fsdp', 'context', 'tensor']]
 
@@ -174,9 +174,9 @@ dcn_fsdp_parallelism: -1
 dcn_context_parallelism: 1
 dcn_tensor_parallelism: 1
 ici_data_parallelism: 1
-ici_fsdp_parallelism: -1  # recommended ICI axis to be auto-sharded
+ici_fsdp_parallelism: 1
 ici_context_parallelism: 1
-ici_tensor_parallelism: 1
+ici_tensor_parallelism: -1  # recommended ICI axis to be auto-sharded
 
 allow_split_physical_axes: False
 

src/maxdiffusion/configs/base_wan_i2v_27b.yml

@@ -152,17 +152,17 @@ mesh_axes: ['data', 'fsdp', 'context', 'tensor']
 logical_axis_rules: [
                       ['batch', ['data', 'fsdp']],
                       ['activation_batch', ['data', 'fsdp']],
-                      ['activation_self_attn_heads', ['context', 'tensor']],
-                      ['activation_cross_attn_q_length', ['context', 'tensor']],
-                      ['activation_length', 'context'],
+                      ['activation_self_attn_heads', 'tensor'],
+                      ['activation_cross_attn_q_length', 'tensor'],
+                      ['activation_length', None],
                       ['activation_heads', 'tensor'],
                       ['mlp','tensor'],
-                      ['embed', ['context', 'fsdp']],
+                      ['embed', 'fsdp'],
                       ['heads', 'tensor'],
                       ['norm', 'tensor'],
-                      ['conv_batch', ['data', 'context', 'fsdp']],
+                      ['conv_batch', ['data', 'fsdp']],
                       ['out_channels', 'tensor'],
-                      ['conv_out', 'context'],
+                      ['conv_out', 'tensor'],
                     ]
 data_sharding: [['data', 'fsdp', 'context', 'tensor']]
 
@@ -175,9 +175,9 @@ dcn_fsdp_parallelism: -1
 dcn_context_parallelism: 1
 dcn_tensor_parallelism: 1
 ici_data_parallelism: 1
-ici_fsdp_parallelism: -1  # recommended ICI axis to be auto-sharded
+ici_fsdp_parallelism: 1
 ici_context_parallelism: 1
-ici_tensor_parallelism: 1
+ici_tensor_parallelism: -1  # recommended ICI axis to be auto-sharded
 
 allow_split_physical_axes: False
 

src/maxdiffusion/configs/ltx2_video.yml

@@ -49,7 +49,7 @@ logical_axis_rules: [
                       ['activation_length', None],
                       ['activation_heads', 'tensor'],
                       ['mlp','tensor'],
-                      ['embed', ['fsdp', 'tensor']],
+                      ['embed', 'fsdp'],
                       ['heads', 'tensor'],
                       ['norm', 'tensor'],
                       ['conv_batch', ['data', 'fsdp']],