debug

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -389,9 +389,15 @@ def __call__(
       a2v_cross_attention_mask: Optional[jax.Array] = None,
       v2a_cross_attention_mask: Optional[jax.Array] = None,
       perturbation_mask: Optional[jax.Array] = None,
+      layer_id: int = 0,
   ) -> Tuple[jax.Array, jax.Array]:
     batch_size = hidden_states.shape[0]
 
+    is_layer_0 = (layer_id == 0)
+    def _print_stats_layer(name, tensor):
+      jax.debug.print("DEBUG [BLOCK 0] {name} shape: {shape}, mean: {mean}, min: {min}, max: {max}, std: {std}",
+                      name=name, shape=tensor.shape, mean=jnp.round(jnp.mean(tensor), 6), min=jnp.round(jnp.min(tensor), 4), max=jnp.round(jnp.max(tensor), 4), std=jnp.round(jnp.std(tensor), 4))
+
     axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed"))
     hidden_states = jax.lax.with_sharding_constraint(hidden_states, axis_names)
     axis_names_audio = nn.logical_to_mesh_axes(("activation_batch", None, "activation_embed"))
@@ -438,6 +444,7 @@ def __call__(
     hidden_states = hidden_states + attn_hidden_states * gate_msa
 
     if self.use_audio and audio_hidden_states is not None:
+      jax.lax.cond(is_layer_0, lambda: _print_stats_layer("audio_in", audio_hidden_states), lambda: None)
       # Calculate Audio AdaLN values
       norm_audio_hidden_states = self.audio_norm1(audio_hidden_states)
 
@@ -459,6 +466,7 @@ def __call__(
         audio_gate_q = audio_ada_values[:, :, 8, :]
 
       norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_msa) + audio_shift_msa
+      jax.lax.cond(is_layer_0, lambda: _print_stats_layer("audio_norm1_out", norm_audio_hidden_states), lambda: None)
 
       with jax.named_scope("Audio Self-Attention"):
         attn_audio_hidden_states = self.audio_attn1(
@@ -467,7 +475,9 @@ def __call__(
             rotary_emb=audio_rotary_emb,
             perturbation_mask=perturbation_mask,
         )
+      jax.lax.cond(is_layer_0, lambda: _print_stats_layer("audio_attn1_out", attn_audio_hidden_states), lambda: None)
       audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * audio_gate_msa
+      jax.lax.cond(is_layer_0, lambda: _print_stats_layer("audio_attn1_residual", audio_hidden_states), lambda: None)
 
     # 2. Video and Audio Cross-Attention with the text embeddings
     norm_hidden_states = self.norm2(hidden_states)
@@ -496,6 +506,7 @@ def __call__(
       norm_audio_hidden_states = self.audio_norm2(audio_hidden_states)
       if getattr(self, "cross_attn_mod", False):
         norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_q) + audio_shift_q
+      jax.lax.cond(is_layer_0, lambda: _print_stats_layer("audio_norm2_out", norm_audio_hidden_states), lambda: None)
 
       if getattr(self, "cross_attn_mod", False) and temb_prompt_audio is not None:
         audio_prompt_table_reshaped = jnp.expand_dims(self.audio_prompt_scale_shift_table, axis=(0, 1))
@@ -513,11 +524,14 @@ def __call__(
       )
       if getattr(self, "cross_attn_mod", False):
         attn_audio_hidden_states = attn_audio_hidden_states * audio_gate_q
+      jax.lax.cond(is_layer_0, lambda: _print_stats_layer("audio_attn2_out", attn_audio_hidden_states), lambda: None)
       audio_hidden_states = audio_hidden_states + attn_audio_hidden_states
+      jax.lax.cond(is_layer_0, lambda: _print_stats_layer("audio_attn2_residual", audio_hidden_states), lambda: None)
 
       # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
       norm_hidden_states = self.audio_to_video_norm(hidden_states)
       norm_audio_hidden_states = self.video_to_audio_norm(audio_hidden_states)
+      jax.lax.cond(is_layer_0, lambda: _print_stats_layer("audio_v2a_norm_out", norm_audio_hidden_states), lambda: None)
 
       # Calculate Cross-Attention Modulation values
       # Video
@@ -582,9 +596,11 @@ def __call__(
             k_rotary_emb=ca_video_rotary_emb,
             attention_mask=v2a_cross_attention_mask,
         )
+      jax.lax.cond(is_layer_0, lambda: _print_stats_layer("audio_v2a_attn_out", v2a_attn_hidden_states), lambda: None)
       if modality_mask is not None:
         v2a_attn_hidden_states = v2a_attn_hidden_states * modality_mask
       audio_hidden_states = audio_hidden_states + v2a_gate * v2a_attn_hidden_states
+      jax.lax.cond(is_layer_0, lambda: _print_stats_layer("audio_v2a_residual", audio_hidden_states), lambda: None)
 
     # 4. Feedforward
     norm_hidden_states = self.norm3(hidden_states)
@@ -595,8 +611,11 @@ def __call__(
     if self.use_audio and audio_hidden_states is not None:
       norm_audio_hidden_states = self.audio_norm3(audio_hidden_states)
       norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_mlp) + audio_shift_mlp
+      jax.lax.cond(is_layer_0, lambda: _print_stats_layer("audio_norm3_out", norm_audio_hidden_states), lambda: None)
       audio_ff_output = self.audio_ff(norm_audio_hidden_states)
+      jax.lax.cond(is_layer_0, lambda: _print_stats_layer("audio_ff_out", audio_ff_output), lambda: None)
       audio_hidden_states = audio_hidden_states + audio_ff_output * audio_gate_mlp
+      jax.lax.cond(is_layer_0, lambda: _print_stats_layer("audio_block_out", audio_hidden_states), lambda: None)
 
     return hidden_states, audio_hidden_states
 
@@ -1193,6 +1212,7 @@ def scan_fn(carry, block_mask_and_id):
             v2a_cross_attention_mask=None,
             perturbation_mask=mask,
             modality_mask=modality_mask,
+            layer_id=layer_id,
         )
       return (
           hidden_states_out.astype(hidden_states.dtype),
@@ -1238,6 +1258,7 @@ def scan_fn(carry, block_mask_and_id):
               a2v_cross_attention_mask=None,
               v2a_cross_attention_mask=None,
               perturbation_mask=mask,
+              layer_id=i,
           )
 
     # 6. Output layers

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1398,26 +1398,7 @@ def __call__(
       prompt_attention_mask_jax = jnp.concatenate([negative_prompt_attention_mask_jax, prompt_attention_mask_jax], axis=0)
       latents_jax = jnp.concatenate([latents_jax] * 2, axis=0)
       audio_latents_jax = jnp.concatenate([audio_latents_jax] * 2, axis=0)
-    def _print_stats_gemma(name, tensor):
-      if isinstance(tensor, list):
-        for idx, t in enumerate(tensor):
-          print(
-              f"DEBUG {name}_{idx} shape: {t.shape}, mean: {jnp.round(jnp.mean(t), 6)}, min: {jnp.round(jnp.min(t), 4)}, max: {jnp.round(jnp.max(t), 4)}, std: {jnp.round(jnp.std(t), 4)}"
-          )
-      else:
-        print(
-            f"DEBUG {name} shape: {tensor.shape}, mean: {jnp.round(jnp.mean(tensor), 6)}, min: {jnp.round(jnp.min(tensor), 4)}, max: {jnp.round(jnp.max(tensor), 4)}, std: {jnp.round(jnp.std(tensor), 4)}"
-        )
 
-    if do_cfg and do_stg:
-      if isinstance(prompt_embeds_jax, list):
-        prompt_embeds_to_print = [x[:2] for x in prompt_embeds_jax]
-      else:
-        prompt_embeds_to_print = prompt_embeds_jax[:2]
-    else:
-      prompt_embeds_to_print = prompt_embeds_jax
-
-    _print_stats_gemma("text_encoder_output_flattened", prompt_embeds_to_print)
 
     if hasattr(self, "mesh") and self.mesh is not None:
       data_sharding_3d = NamedSharding(self.mesh, P())
@@ -1454,6 +1435,11 @@ def _print_stats_gemma(name, tensor):
       jax.block_until_ready(video_embeds)
       max_logging.log(f"⏱️ Connectors Time: {time.time() - connectors_start:.4f} seconds")
 
+      def _print_stats(name, tensor):
+        print(
+            f"DEBUG {name} shape: {tensor.shape}, mean: {jnp.round(jnp.mean(tensor), 6)}, min: {jnp.round(jnp.min(tensor), 4)}, max: {jnp.round(jnp.max(tensor), 4)}, std: {jnp.round(jnp.std(tensor), 4)}"
+        )
+
       video_embeds_sharded = video_embeds
       audio_embeds_sharded = audio_embeds
 
@@ -1462,32 +1448,14 @@ def _print_stats_gemma(name, tensor):
         spec = NamedSharding(self.mesh, P(*activation_axes))
         video_embeds_sharded = jax.device_put(video_embeds, spec)
         audio_embeds_sharded = jax.device_put(audio_embeds, spec)
-      def _print_stats(name, tensor):
-        print(
-            f"DEBUG {name} shape: {tensor.shape}, mean: {jnp.round(jnp.mean(tensor), 6)}, min: {jnp.round(jnp.min(tensor), 4)}, max: {jnp.round(jnp.max(tensor), 4)}, std: {jnp.round(jnp.std(tensor), 4)}"
-        )
       print(f"WEIGHT DEBUG: block 0 to_q kernel mean: {float(self.transformer.transformer_blocks.attn1.to_q.kernel.value[0].mean()):.6f}")
-      if do_cfg and do_stg:
-        _print_stats("video_text_embedding", video_embeds[:2])
-        _print_stats("audio_text_embedding", audio_embeds[:2])
-      else:
-        _print_stats("video_text_embedding", video_embeds)
-        _print_stats("audio_text_embedding", audio_embeds)
 
       timesteps_jax = jnp.array(timesteps, dtype=jnp.float32)
 
       diffusion_loop_start = time.time()
       scan_diffusion_loop = getattr(self.config, "scan_diffusion_loop", True)
 
-      if do_cfg and do_stg:
-        _print_stats("latents_jax_before_loop", latents_jax[:batch_size])
-        _print_stats("audio_latents_jax_before_loop", audio_latents_jax[:batch_size])
-      elif do_cfg:
-        _print_stats("latents_jax_before_loop", latents_jax[:batch_size])
-        _print_stats("audio_latents_jax_before_loop", audio_latents_jax[:batch_size])
-      else:
-        _print_stats("latents_jax_before_loop", latents_jax)
-        _print_stats("audio_latents_jax_before_loop", audio_latents_jax)
+
 
       if scan_diffusion_loop:
         latents_jax, audio_latents_jax = run_diffusion_loop(
@@ -1980,20 +1948,7 @@ def scan_body(carry, inputs):
 
       is_first_step = (t == timesteps_jax[0])
 
-      def print_raw_stats():
-        print_stats_jit("noise_pred_video_raw", noise_pred)
-        print_stats_jit("noise_pred_audio_raw", noise_pred_audio)
-        if do_cfg:
-          uncond_v = noise_pred[:batch_size]
-          cond_v = noise_pred[batch_size : 2 * batch_size]
-          uncond_a = noise_pred_audio[:batch_size]
-          cond_a = noise_pred_audio[batch_size : 2 * batch_size]
-          print_stats_jit("noise_pred_video_raw_uncond", uncond_v)
-          print_stats_jit("noise_pred_video_raw_cond", cond_v)
-          print_stats_jit("noise_pred_audio_raw_uncond", uncond_a)
-          print_stats_jit("noise_pred_audio_raw_cond", cond_a)
-
-      jax.lax.cond(is_first_step, print_raw_stats, lambda: None)
+
 
 
       # Extract latents_step based on stacking strategy