updated vae config logic to be the consistent, update xprof logic

Author: eltsai

src/maxdiffusion/generate_wan.py

@@ -313,13 +313,19 @@ def block_if_jax(x):
       # Block until warmup completes
       jax.tree_util.tree_map(block_if_jax, warmup_videos)
 
+    # Warm up GCS connection by flushing writer before starting profiler
+    if writer and jax.process_index() == 0:
+      max_logging.log("Flushing writer to warm up GCS connection before profiler...")
+      writer.flush()
+
     s0 = time.perf_counter()
     max_utils.activate_profiler(config)
     max_logging.log(f"Profiler: starting profiled run with {steps_for_profile} steps")
     profiled_videos = call_pipeline(config, pipeline, prompt, negative_prompt, num_inference_steps=steps_for_profile)
     # Wait for all computation to finish before stopping profiler
     jax.tree_util.tree_map(block_if_jax, profiled_videos)
     max_utils.deactivate_profiler(config)
+    max_utils.upload_profiler_traces(config)
     generation_time_with_profiler = time.perf_counter() - s0
     max_logging.log(f"generation_time_with_profiler: {generation_time_with_profiler}")
     if writer and jax.process_index() == 0:

src/maxdiffusion/max_utils.py

@@ -78,14 +78,30 @@ def l2norm_pytree(x):
 
 def activate_profiler(config):
   if jax.process_index() == 0 and config.enable_profiler:
-    jax.profiler.start_trace(config.tensorboard_dir)
+    # If tensorboard_dir is GCS, write profiler traces locally instead
+    profiler_path = config.tensorboard_dir
+    if config.tensorboard_dir.startswith("gs://"):
+      profiler_path = "/tmp/profiler_traces"
+      os.makedirs(profiler_path, exist_ok=True)
+      max_logging.log(f"Profiler: saving traces locally to {profiler_path} (GCS paths not supported)")
+    jax.profiler.start_trace(profiler_path)
 
 
 def deactivate_profiler(config):
   if jax.process_index() == 0 and config.enable_profiler:
     jax.profiler.stop_trace()
 
 
+def upload_profiler_traces(config):
+  """No-op for now - profiler traces are saved locally"""
+  if jax.process_index() == 0 and config.enable_profiler:
+    if config.tensorboard_dir.startswith("gs://"):
+      max_logging.log("Profiler traces saved to: /tmp/profiler_traces")
+      max_logging.log("You can download them manually or use: gsutil -m rsync -r /tmp/profiler_traces/ " + config.tensorboard_dir.rstrip("/") + "/")
+    else:
+      max_logging.log(f"Profiler traces saved to: {config.tensorboard_dir}")
+
+
 def initialize_summary_writer(config):
   return writer.SummaryWriter(config.tensorboard_dir) if jax.process_index() == 0 else None
 
@@ -94,7 +110,6 @@ def close_summary_writer(summary_writer):
   if jax.process_index() == 0:
     summary_writer.close()
 
-
 def _prepare_metrics_for_json(metrics, step, run_name):
   """Converts metric dictionary into json supported types (e.g. float)"""
   metrics_dict = {}

src/maxdiffusion/models/attention_flax.py

@@ -112,6 +112,7 @@ def _reshape_batch_dim_to_heads(tensor, heads):
   head_size = heads
   tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)
   tensor = jnp.transpose(tensor, (0, 2, 1, 3))
+  reshaped_tensor = tensor.reshape(batch_size // head_size, seq_len, dim * head_size)
   axis_names = nn.logical_to_mesh_axes((BATCH, LENGTH, HEAD))
   return jax.lax.with_sharding_constraint(reshaped_tensor, axis_names)
 

src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py

@@ -57,6 +57,7 @@ def _load_and_init(cls, config, restored_checkpoint=None, vae_only=False, load_t
           devices_array=common_components["devices_array"],
           mesh=common_components["mesh"],
           vae_mesh=common_components["vae_mesh"],
+          vae_logical_axis_rules=common_components["vae_logical_axis_rules"],
           config=config,
         )
 

src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py

@@ -74,6 +74,7 @@ def _load_and_init(cls, config, restored_checkpoint=None, vae_only=False, load_t
           devices_array=common_components["devices_array"],
           mesh=common_components["mesh"],
           vae_mesh=common_components["vae_mesh"],
+          vae_logical_axis_rules=common_components["vae_logical_axis_rules"],
           config=config,
       )
     return pipeline, low_noise_transformer, high_noise_transformer

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p1.py

@@ -61,6 +61,8 @@ def _load_and_init(cls, config, restored_checkpoint=None, vae_only=False, load_t
         scheduler_state=common_components["scheduler_state"],
         devices_array=common_components["devices_array"],
         mesh=common_components["mesh"],
+      vae_mesh=common_components["vae_mesh"],
+      vae_logical_axis_rules=common_components["vae_logical_axis_rules"],
         config=config,
     )
     return pipeline, transformer

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py

@@ -79,6 +79,8 @@ def _load_and_init(cls, config, restored_checkpoint=None, vae_only=False, load_t
         scheduler_state=common_components["scheduler_state"],
         devices_array=common_components["devices_array"],
         mesh=common_components["mesh"],
+      vae_mesh=common_components["vae_mesh"],
+      vae_logical_axis_rules=common_components["vae_logical_axis_rules"],
         config=config,
     )
     return pipeline, low_noise_transformer, high_noise_transformer

src/maxdiffusion/trainers/flux_trainer.py

@@ -386,6 +386,7 @@ def training_loop(self, p_train_step, pipeline, params, train_states, data_itera
 
       if self.config.enable_profiler and step == last_profiling_step:
         max_utils.deactivate_profiler(self.config)
+        max_utils.upload_profiler_traces(self.config)
 
     train_states[FLUX_STATE_KEY] = flux_state
     if len(times) > 0: