split activation_batch across data.

Author: jfacevedo-google

src/maxdiffusion/configs/base_wan_14b.yml

@@ -127,9 +127,7 @@ logical_axis_rules: [
                       ['batch', 'data'],
                       ['activation_heads', 'tensor'],
                       ['activation_length', 'fsdp'],
-                      #['activation_heads', 'fsdp'],
-                      #['activation_heads', 'fsdp'],
-                      #['activation_batch', ['data','fsdp']],
+                      ['activation_batch', 'data'],
                       ['activation_kv', 'tensor'],
                       ['mlp','tensor'],
                       ['embed','fsdp'],