generalize across sharding parallelisms

Author: cpersson-amd

src/maxdiffusion/max_utils.py

@@ -654,4 +654,16 @@ 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()
+  
+def get_axis_names(axis_key: str, config=None) -> str:
+  """Returns the mesh axis names given the logical axis key from config.logical_axis_rules."""
+  axis_name = ''
+  if config:
+    axis_rules = config.logical_axis_rules
+  else:
+    axis_rules = nn.get_logical_axis_rules()
+  for rules in axis_rules:
+    if rules[0] == axis_key:
+      axis_name = rules[1]
+  return axis_name

src/maxdiffusion/models/attention_flax.py

@@ -27,6 +27,7 @@
 from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_kernel
 from einops import rearrange
 from .. import common_types, max_logging
+from .. import max_utils
 
 from . import quantizations
 
@@ -205,10 +206,7 @@ def _tpu_flash_attention(
         block_kv_dq=None if attention_kernel == "tokamax_flash" else min(kv_max_block_size, query.shape[2]),
         use_fused_bwd_kernel=True if attention_kernel == "tokamax_flash" else False,
     )
-  fsdp_key = "fsdp"
-  if "fsdp_tpu" in mesh.shape.keys():
-    fsdp_key = "fsdp_tpu"
-  
+  fsdp_key = max_utils.get_axis_names("activation_length")
   num_fsdp_shards = mesh.shape[fsdp_key]
   query = _reshape_data_for_flash(query, heads)
   key = _reshape_data_for_flash(key, heads)

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -20,6 +20,7 @@
 import jax.numpy as jnp
 from flax import nnx
 from ...configuration_utils import ConfigMixin
+from ... import max_utils
 from ..modeling_flax_utils import FlaxModelMixin, get_activation
 from ... import common_types
 from ..vae_flax import (FlaxAutoencoderKLOutput, FlaxDiagonalGaussianDistribution, FlaxDecoderOutput)
@@ -72,7 +73,10 @@ def __init__(
     self._depth_padding_before = self._causal_padding[1][0]  # 2 * padding_tuple[0]
 
     # Set sharding dynamically based on out_channels.
-    num_fsdp_axis_devices = mesh.device_ids.shape[2]
+    fspd_key = max_utils.get_axis_names("activation_length")
+    if not fspd_key:
+      fspd_key = "fsdp"
+    num_fsdp_axis_devices = mesh.shape[fspd_key]
     kernel_sharding = (None, None, None, None, None)
     if out_channels % num_fsdp_axis_devices == 0:
       kernel_sharding = (None, None, None, None, "conv_out")

src/maxdiffusion/trainers/wan_trainer.py

@@ -211,8 +211,8 @@ def prepare_sample_eval(features):
     return data_iterator
 
   def start_training(self):
-
-    pipeline, opt_state, step = self.checkpointer.load_checkpoint()
+    with nn_partitioning.axis_rules(self.config.logical_axis_rules):
+      pipeline, opt_state, step = self.checkpointer.load_checkpoint()
     restore_args = {}
     if opt_state and step:
       restore_args = {"opt_state": opt_state, "step": step}
@@ -362,13 +362,7 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data
     example_batch = load_next_batch(train_data_iterator, None, self.config)
 
     # Designate the context parallel axis for sharding
-    cp_resource = ''
-    for rules in self.config.logical_axis_rules:
-      if rules[0] == "activation_length":
-        if isinstance(rules[1], list):
-          cp_resource = rules[1][0]
-        else:
-          cp_resource = rules[1]
+    cp_resource = max_utils.get_axis_names("activation_length", config=self.config)
     mesh_resource = MeshResource(cp_resource=cp_resource)
 
     with ThreadPoolExecutor(max_workers=1) as executor: