wip - context paralleism - not working yet.

Author: jfacevedo-google

src/maxdiffusion/common_types.py

@@ -36,6 +36,8 @@
 
 BATCH = "activation_batch"
 LENGTH = "activation_length"
+Q_LENGTH = "activation_q_length"
+KV_LENGTH = "activation_kv_length"
 EMBED = "activation_embed"
 HEAD = "activation_heads"
 D_KV = "activation_kv"

src/maxdiffusion/configs/base_wan_14b.yml

@@ -96,7 +96,7 @@ hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu'
 skip_jax_distributed_system: False
 
 # Parallelism
-mesh_axes: ['data', 'fsdp', 'tensor']
+mesh_axes: ['data', 'fsdp', 'context', 'tensor']
 
 # batch : batch dimension of data and activations
 # hidden :
@@ -112,29 +112,33 @@ mesh_axes: ['data', 'fsdp', 'tensor']
 # conv_out : conv.shape[-1] weight
 logical_axis_rules: [
                       ['batch', 'data'],
-                      ['activation_batch', ['data','fsdp']],
+                      #['activation_batch', 'fsdp'],
                       ['activation_heads', 'tensor'],
+                      ['activation_q_length', ['context']],
+                      ['activatation_kv_length', []],
                       ['activation_kv', 'tensor'],
                       ['mlp','tensor'],
-                      ['embed','fsdp'],
+                      ['embed',['fsdp','context']],
                       ['heads', 'tensor'],
                       ['norm', 'fsdp'],
                       ['conv_batch', ['data','fsdp']],
                       ['out_channels', 'tensor'],
                       ['conv_out', 'fsdp'],
                       ['conv_in', 'fsdp']
                     ]
-data_sharding: [['data', 'fsdp', 'tensor']]
+data_sharding: [['data', 'fsdp', 'context', 'tensor']]
 
 # One axis for each parallelism type may hold a placeholder (-1)
 # value to auto-shard based on available slices and devices.
 # By default, product of the DCN axes should equal number of slices
 # and product of the ICI axes should equal number of devices per slice.
 dcn_data_parallelism: 1  # recommended DCN axis to be auto-sharded
-dcn_fsdp_parallelism: -1
+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  # recommended ICI axis to be auto-sharded
+ici_context_parallelism: -1
 ici_tensor_parallelism: 1
 
 # Dataset

src/maxdiffusion/max_utils.py

@@ -266,21 +266,13 @@ def create_device_mesh(config, devices=None, logging=True):
 
   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,
-  ]
+  ici_parallelism = fill_unspecified_mesh_axes(config.ici_parallelism.copy(), num_devices_per_slice, "ICI")
 
   # 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")
+    dcn_parallelism = fill_unspecified_mesh_axes(config.dcp_parallelism.copy(), 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)

src/maxdiffusion/models/attention_flax.py

@@ -39,6 +39,8 @@
 AxisNames = common_types.AxisNames
 BATCH = common_types.BATCH
 LENGTH = common_types.LENGTH
+Q_LENGTH = common_types.Q_LENGTH
+KV_LENGTH = common_types.KV_LENGTH
 HEAD = common_types.HEAD
 D_KV = common_types.D_KV
 EMBED = common_types.EMBED
@@ -139,50 +141,87 @@ def _tpu_flash_attention(
   value: jax.Array,
   heads: int,
   mesh: Mesh,
-  flash_axis_names: AxisNames,
-  flash_block_sizes: BlockSizes,
+  flash_block_sizes: BlockSizes = None,
+  flash_axis_names_kv: AxisNames = (BATCH, HEAD, KV_LENGTH, D_KV),
+  flash_axis_names_q: AxisNames = (BATCH, HEAD, LENGTH, D_KV),
+  flash_axis_names_splash_kernel: AxisNames = (HEAD, LENGTH),
   dtype: jnp.dtype = jnp.float32) -> jax.Array:
   """TPU Flash Attention"""
 
-  max_block_size = 1024 if dtype == jnp.bfloat16 else 512
+  cp_size = mesh.shape["context"]
+  #breakpoint()
+  axis_names_splash_kernel = nn.logical_to_mesh_axes(flash_axis_names_splash_kernel)
+  axis_names_q = nn.logical_to_mesh_axes(flash_axis_names_q)
+  axis_names_kv = nn.logical_to_mesh_axes(flash_axis_names_kv)
+  max_logging.log(f"axis_names_q: {axis_names_q}")
+  max_logging.log(f"axis_names_kv: {axis_names_kv}")
+  max_logging.log(f"axis_names_splash_kernel: {axis_names_splash_kernel}")
+
+  max_block_size = 256 if dtype == jnp.bfloat16 else 128
   if flash_block_sizes:
     block_sizes = flash_block_sizes
   else:
     block_sizes = splash_attention_kernel.BlockSizes(
         block_q=min(max_block_size, query.shape[2]),
-        block_kv_compute=min(max_block_size, key.shape[2]),
         block_kv=min(max_block_size, key.shape[2]),
+        block_kv_compute=min(max_block_size, key.shape[2]),
         block_q_dkv=min(max_block_size, query.shape[2]),
         block_kv_dkv=min(max_block_size, key.shape[2]),
         block_kv_dkv_compute=min(max_block_size, query.shape[2]),
         block_q_dq=min(max_block_size, query.shape[2]),
         block_kv_dq=min(max_block_size, query.shape[2]),
+        q_layout=splash_attention_kernel.QKVLayout["HEAD_DIM_MINOR"],
+        k_layout=splash_attention_kernel.QKVLayout["HEAD_DIM_MINOR"],
+        v_layout=splash_attention_kernel.QKVLayout["HEAD_DIM_MINOR"],
     )
 
   query, kv_size, query_seq_len = _reshape_data_for_flash(query, heads, block_sizes.block_q)
   key, _, _ = _reshape_data_for_flash(key, heads, block_sizes.block_kv_compute)
   value, _, _ = _reshape_data_for_flash(value, heads, block_sizes.block_kv_compute)
 
-  axis_names = nn.logical_to_mesh_axes(flash_axis_names)
-
   @functools.partial(
-      shard_map.shard_map,
-      mesh=mesh,
-      in_specs=(
-          axis_names,
-          axis_names,
-          axis_names,
-      ),
-      out_specs=axis_names,
-      check_rep=False,
+      jax.jit,
+      static_argnames=[
+        "multi_head_mask",
+        "shard_head_size"
+      ],
   )
-  def wrap_flash_attention(query, key, value):
-    masks = [splash_attention_mask.FullMask(_shape=(query.shape[2], query.shape[2])) for _ in range(query.shape[1])]
-    multi_head_mask = splash_attention_mask.MultiHeadMask(masks=masks)
+  def wrap_splash_kernel(multi_head_mask, shard_head_size=1):
+    # breakpoint()
     splash_kernel = splash_attention_kernel.make_splash_mha(
-        mask=multi_head_mask, head_shards=1, q_seq_shards=1, block_sizes=block_sizes
+      mask=multi_head_mask,
+      head_shards=shard_head_size, # the sizes of the axis is sharding over heads
+      q_seq_shards=cp_size,
+      block_sizes=block_sizes,
     )
-    return jax.vmap(splash_kernel)(query, key, value)
+    return splash_kernel
+
+  # logical_axis_rules_head = np.array(
+  #   [mesh.shape[physical_axes] for physical_axes in dict(config.logical_axis_rules)[HEAD]]
+  # )
+  shard_head_size = 1
+
+  masks = [splash_attention_mask.FullMask(_shape=(query.shape[2], query.shape[2])) for _ in range(query.shape[1])]
+  multi_head_mask = splash_attention_mask.MultiHeadMask(masks=masks)
+  splash_kernel = wrap_splash_kernel(multi_head_mask, int(shard_head_size))
+  named_sharding = jax.sharding.NamedSharding(mesh, axis_names_splash_kernel)
+  segment_axis_names_splash_kernel = splash_kernel.manual_sharding_spec(named_sharding)
+  @functools.partial(
+    shard_map.shard_map,
+    mesh=mesh,
+    in_specs=(
+      axis_names_q,
+      axis_names_kv,
+      axis_names_kv,
+      segment_axis_names_splash_kernel,
+      None
+    ),
+    out_specs=axis_names_q,
+    check_rep=False
+  )
+  def wrap_flash_attention(query, key, value, splash_kernel, cp_size):
+    attention_output = jax.vmap(splash_kernel)(query, key, value)
+    return attention_output
 
   devices_in_data_fsdp = mesh.shape["data"] * mesh.shape["fsdp"]
   # This warning might show up when doing model eval for example, when calculating model flops
@@ -192,7 +231,7 @@ def wrap_flash_attention(query, key, value):
         "Warning, batch dimension should be shardable among the devices in data and fsdp"
         f" axis, batch dimension: {query.shape[0]}, devices_in_data_fsdp: {devices_in_data_fsdp}"
     )
-  x = wrap_flash_attention(query, key, value)
+  x = wrap_flash_attention(query, key, value, splash_kernel, cp_size)
   x = x[:, :, :query_seq_len, :kv_size]
   x = _reshape_heads_to_head_dim(x)
 
@@ -343,7 +382,15 @@ def _apply_attention(
   if attention_kernel == "dot_product" or use_memory_efficient_attention or not can_use_flash_attention:
     return _apply_attention_dot(query, key, value, dtype, heads, dim_head, scale, split_head_dim, float32_qk_product, use_memory_efficient_attention)
   elif attention_kernel == "flash":
-    return _tpu_flash_attention(query, key * scale, value, heads, mesh, flash_axis_names, flash_block_sizes, dtype)
+    return _tpu_flash_attention(
+      query=query,
+      key=key * scale,
+      value=value,
+      heads=heads,
+      mesh=mesh,
+      flash_block_sizes=flash_block_sizes,
+      dtype=dtype
+    )
   elif attention_kernel == "cudnn_flash_te":
     return _cudnn_flash_attention(query, key, value, heads, mesh, dpa_layer)
   else:

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -21,6 +21,7 @@
 import flax
 import flax.linen as nn
 from flax import nnx
+from flax.linen import partitioning as nn_partitioning
 from ...pyconfig import HyperParameters
 from ... import max_logging
 from ... import max_utils
@@ -434,7 +435,7 @@ def __call__(
         num_transformer_layers=self.transformer.config.num_layers
       )
 
-      with self.mesh:
+      with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
         latents = p_run_inference(
           graphdef=graphdef,
           sharded_state=state,

src/maxdiffusion/pyconfig.py

@@ -155,6 +155,22 @@ def user_init(raw_keys):
     raw_keys["num_slices"] = get_num_slices(raw_keys)
     raw_keys["quantization_local_shard_count"] = get_quantization_local_shard_count(raw_keys)
 
+    ici_parallelism = [
+      raw_keys["ici_data_parallelism"],
+      raw_keys["ici_fsdp_parallelism"],
+      raw_keys["ici_context_parallelism"],
+      raw_keys["ici_tensor_parallelism"],
+      
+    ]
+    dcn_parallelism = [
+      raw_keys["dcn_data_parallelism"],
+      raw_keys["dcn_fsdp_parallelism"],
+      raw_keys["dcn_context_parallelism"],
+      raw_keys["dcn_tensor_parallelism"],
+    ]
+    raw_keys["ici_parallelism"] = ici_parallelism
+    raw_keys['dcn_parallelism'] = dcn_parallelism
+
 
 def get_num_slices(raw_keys):
   if int(raw_keys["compile_topology_num_slices"]) > 0:

src/maxdiffusion/trainers/wan_trainer.py

@@ -77,26 +77,24 @@ def start_training(self):
 
     pipeline = self.load_checkpoint()
     mesh = pipeline.mesh
+    #breakpoint()
+    # logical_axis_rules_head = np.array([mesh.shape[physical_axes] for physical_axes in dict(self.config.logical_axis_rules)["activation_heads"]])
+    # breakpoint()
 
     optimizer, learning_rate_scheduler = self._create_optimizer(pipeline.transformer, self.config, self.config.learning_rate)
 
-    # @nnx.jit
-    # def create_transformer_state(transformer):
-    #   optimizer = self._create_optimizer(transformer, self.config, self.config.learning_rate)
-    #   breakpoint()
-    #   _, state = nnx.split((transformer, optimizer))
-    
-    # with mesh:
-    #   create_transformer_state(pipeline.transformer)
-
-    #graphdef, state = nnx.plit((pipeline.transformer, optimizer))
     dummy_inputs = self.load_dataset(pipeline)
     dummy_inputs = tuple([jtu.tree_map_with_path(functools.partial(_form_global_array, global_mesh=mesh), input) for input in dummy_inputs])
 
     self.training_loop(pipeline, optimizer, learning_rate_scheduler, dummy_inputs)
 
   def training_loop(self, pipeline, optimizer, learning_rate_scheduler, data):
-    
+    # From Wan 2.1 paper https://arxiv.org/pdf/2503.20314
+    # Input shape of DiT block is (b, s, h)
+    # b corresponds to data parallelism.
+    # s represents the sequence length and sharding is achieved through context parallelism.
+    # Sharding long the h dimension primarily involves Megatron's tensor parallelism TP combined
+    # with sequence parallelism  which shards the hidden dimension of the activations by splitting the weights.
     graphdef, state = nnx.split((pipeline.transformer, optimizer))
     state = state.to_pure_dict()
     p_train_step = jax.jit(