From 81166a87cc3f7fd30de42e370423d8d4358041a5 Mon Sep 17 00:00:00 2001 From: Robert Moseley Date: Tue, 23 Sep 2025 15:42:27 -0700 Subject: [PATCH] Ensure fusion gating is identity when audio disabled --- safe/models/audio_encoders.py | 99 +++++++++++-- safe/models/fusion_adapter.py | 173 +++++++++++++++++++---- safe/models/safe_model.py | 55 +++++++- safe/training/stage_a.py | 252 +++++++++++++++++++++++++++++++++- 4 files changed, 527 insertions(+), 52 deletions(-) diff --git a/safe/models/audio_encoders.py b/safe/models/audio_encoders.py index ce2c80a..bfde769 100644 --- a/safe/models/audio_encoders.py +++ b/safe/models/audio_encoders.py @@ -1,6 +1,6 @@ import torch -import torch.nn as nn import numpy as np +import torch.nn as nn import librosa from typing import Optional, Union, List, Tuple import laion_clap @@ -21,11 +21,14 @@ def __init__( max_length: float = 10.0 # seconds ): super().__init__() - + self.sample_rate = sample_rate self.max_length = max_length self.max_samples = int(sample_rate * max_length) - + self.debug_logging = False + self._waveform_log_limit = 5 + self._waveform_logs_emitted = 0 + # Load CLAP model (suppress verbose output) import logging import sys @@ -51,7 +54,40 @@ def __init__( # Get audio embedding dimension self.audio_embed_dim = 512 # CLAP audio embedding dimension - + + def set_debug_logging(self, enabled: bool, max_waveform_logs: int = 5) -> None: + """Enable or disable verbose waveform statistics logging.""" + + self.debug_logging = bool(enabled) + self._waveform_log_limit = int(max(0, max_waveform_logs)) + self._waveform_logs_emitted = 0 + + # ------------------------------------------------------------------ + def _log_waveform_stats(self, audio_data: np.ndarray, source: str) -> None: + if not self.debug_logging: + return + if self._waveform_logs_emitted >= self._waveform_log_limit: + return + + flattened = audio_data.astype(np.float32).ravel() + if flattened.size == 0: + print(f"[AudioDebug] {source}: empty waveform") + self._waveform_logs_emitted += 1 + return + + max_val = float(np.max(flattened)) + min_val = float(np.min(flattened)) + mean_val = float(np.mean(flattened)) + mean_abs = float(np.mean(np.abs(flattened))) + zero_fraction = float(np.mean(np.isclose(flattened, 0.0))) + + print( + f"[AudioDebug] {source}: max={max_val:.6f} min={min_val:.6f} " + f"mean={mean_val:.6f} mean|x|={mean_abs:.6f} zero_frac={zero_fraction:.3f}", + flush=True, + ) + self._waveform_logs_emitted += 1 + def preprocess_audio(self, audio: Union[torch.Tensor, np.ndarray, str]) -> torch.Tensor: """ Preprocess audio to the format expected by CLAP. @@ -68,20 +104,22 @@ def preprocess_audio(self, audio: Union[torch.Tensor, np.ndarray, str]) -> torch elif isinstance(audio, np.ndarray): audio_data = audio elif isinstance(audio, torch.Tensor): - audio_data = audio.cpu().numpy() + audio_data = audio.detach().cpu().numpy() else: raise ValueError(f"Unsupported audio type: {type(audio)}") - + # Ensure correct sample rate if len(audio_data.shape) > 1: audio_data = audio_data.mean(axis=0) # Convert to mono - + # Truncate or pad to max_length if len(audio_data) > self.max_samples: audio_data = audio_data[:self.max_samples] else: audio_data = np.pad(audio_data, (0, self.max_samples - len(audio_data))) - + + self._log_waveform_stats(audio_data, "CLAP.preprocess_audio") + return torch.from_numpy(audio_data).float().unsqueeze(0) # (1, max_samples) def forward(self, audio: Union[torch.Tensor, List[str], List[np.ndarray]]) -> torch.Tensor: @@ -97,7 +135,7 @@ def forward(self, audio: Union[torch.Tensor, List[str], List[np.ndarray]]) -> to if isinstance(audio, list): # Process batch of audio files/arrays processed_audio = [] - for a in audio: + for idx, a in enumerate(audio): processed = self.preprocess_audio(a) processed_audio.append(processed) audio_batch = torch.stack(processed_audio) # (batch_size, 1, max_samples) @@ -125,7 +163,7 @@ def forward(self, audio: Union[torch.Tensor, List[str], List[np.ndarray]]) -> to audio_numpy = audio_batch.detach().cpu().numpy() else: audio_numpy = audio_batch - + audio_embeddings = self.model.get_audio_embedding_from_data( x=audio_numpy, use_tensor=False @@ -157,11 +195,14 @@ def __init__( max_length: float = 30.0 # seconds ): super().__init__() - + self.sample_rate = sample_rate self.max_length = max_length self.extract_transcript = extract_transcript self.max_samples = int(sample_rate * max_length) + self.debug_logging = False + self._waveform_log_limit = 5 + self._waveform_logs_emitted = 0 # Load Whisper model self.model = whisper.load_model(model_name) @@ -174,6 +215,36 @@ def __init__( # Get embedding dimension from Whisper encoder self.audio_embed_dim = self.model.dims.n_audio_state # 512 for small, 768 for base, etc. + + def set_debug_logging(self, enabled: bool, max_waveform_logs: int = 5) -> None: + self.debug_logging = bool(enabled) + self._waveform_log_limit = int(max(0, max_waveform_logs)) + self._waveform_logs_emitted = 0 + + def _log_waveform_stats(self, audio_data: np.ndarray, source: str) -> None: + if not self.debug_logging: + return + if self._waveform_logs_emitted >= self._waveform_log_limit: + return + + flattened = audio_data.astype(np.float32).ravel() + if flattened.size == 0: + print(f"[AudioDebug] {source}: empty waveform", flush=True) + self._waveform_logs_emitted += 1 + return + + max_val = float(np.max(flattened)) + min_val = float(np.min(flattened)) + mean_val = float(np.mean(flattened)) + mean_abs = float(np.mean(np.abs(flattened))) + zero_fraction = float(np.mean(np.isclose(flattened, 0.0))) + + print( + f"[AudioDebug] {source}: max={max_val:.6f} min={min_val:.6f} " + f"mean={mean_val:.6f} mean|x|={mean_abs:.6f} zero_frac={zero_fraction:.3f}", + flush=True, + ) + self._waveform_logs_emitted += 1 def preprocess_audio(self, audio: Union[torch.Tensor, np.ndarray, str]) -> np.ndarray: """ @@ -188,7 +259,7 @@ def preprocess_audio(self, audio: Union[torch.Tensor, np.ndarray, str]) -> np.nd if isinstance(audio, str): audio_data, sr = librosa.load(audio, sr=self.sample_rate) elif isinstance(audio, torch.Tensor): - audio_data = audio.cpu().numpy() + audio_data = audio.detach().cpu().numpy() else: audio_data = audio @@ -201,7 +272,9 @@ def preprocess_audio(self, audio: Union[torch.Tensor, np.ndarray, str]) -> np.nd audio_data = audio_data[:self.max_samples] else: audio_data = np.pad(audio_data, (0, self.max_samples - len(audio_data))) - + + self._log_waveform_stats(audio_data, "Whisper.preprocess_audio") + return audio_data def forward( diff --git a/safe/models/fusion_adapter.py b/safe/models/fusion_adapter.py index 1bd7e44..d537ecb 100644 --- a/safe/models/fusion_adapter.py +++ b/safe/models/fusion_adapter.py @@ -1,9 +1,11 @@ +import torch +import math +from typing import Optional, Tuple + import torch import torch.nn as nn import torch.nn.functional as F from peft import LoraConfig, get_peft_model, LoraModel -from typing import Optional, Tuple -import math class CrossAttentionBlock(nn.Module): @@ -18,7 +20,8 @@ def __init__( num_attention_heads: int = 8, attention_dropout: float = 0.1, output_dropout: float = 0.1, - layer_norm_eps: float = 1e-5 + layer_norm_eps: float = 1e-5, + apply_post_layernorm: bool = False, ): super().__init__() @@ -26,6 +29,10 @@ def __init__( self.num_attention_heads = num_attention_heads self.attention_head_size = hidden_size // num_attention_heads self.all_head_size = self.num_attention_heads * self.attention_head_size + self.debug_logging = False + self._last_attention_summary: Optional[dict] = None + self._attention_log_limit = 5 + self._attention_logs_emitted = 0 # Query projection (from LLM hidden states) self.query = nn.Linear(hidden_size, self.all_head_size) @@ -39,7 +46,12 @@ def __init__( self.output_dropout = nn.Dropout(output_dropout) # Layer normalization - self.layer_norm = nn.LayerNorm(hidden_size, eps=layer_norm_eps) + self.apply_post_layernorm = bool(apply_post_layernorm) + self.layer_norm = ( + nn.LayerNorm(hidden_size, eps=layer_norm_eps) + if self.apply_post_layernorm + else None + ) # Attention dropout self.attention_dropout = nn.Dropout(attention_dropout) @@ -55,7 +67,8 @@ def forward( hidden_states: torch.Tensor, audio_tokens: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, - **kwargs # Accept and ignore extra kwargs to prevent PEFT errors + supervised_mask: Optional[torch.Tensor] = None, + **kwargs, # Accept and ignore extra kwargs to prevent PEFT errors ) -> torch.Tensor: """ Forward pass of cross-attention. @@ -110,15 +123,51 @@ def forward( context_layer = context_layer.view(*new_context_layer_shape) # Output projection - output = self.output_dense(context_layer) - output = torch.nan_to_num(output, nan=0.0, posinf=1e4, neginf=-1e4) - output = self.output_dropout(output) - output = output.to(input_dtype) - - # Residual connection and layer norm - hidden_states = hidden_states.to(output.dtype) - output = self.layer_norm(output + hidden_states) - + delta = self.output_dense(context_layer) + delta = torch.nan_to_num(delta, nan=0.0, posinf=1e4, neginf=-1e4) + delta = self.output_dropout(delta) + delta = delta.to(input_dtype) + + # Residual connection (optionally followed by layer norm) + hidden_states = hidden_states.to(delta.dtype) + output = hidden_states + delta + if self.layer_norm is not None: + output = self.layer_norm(output) + + if getattr(self, "debug_logging", False): + summary = { + "overall_mean": float(attention_probs.mean().item()), + "overall_max": float(attention_probs.max().item()), + } + + if supervised_mask is not None: + mask = supervised_mask + if mask.dim() == 2: + mask = mask.unsqueeze(1).unsqueeze(-1) + elif mask.dim() == 3: + mask = mask.unsqueeze(-1) + mask = mask.to(attention_probs.device, attention_probs.dtype) + denom = mask.sum(dim=(1, 2, 3)).clamp_min(1e-6) + weighted = (attention_probs * mask).sum(dim=(1, 2, 3)) / denom + summary["supervised_mean_per_sample"] = weighted.detach().cpu() + summary["supervised_mean"] = float(weighted.mean().item()) + else: + per_sample = attention_probs.mean(dim=(1, 2, 3)) + summary["per_sample_mean"] = per_sample.detach().cpu() + + self._last_attention_summary = summary + if self._attention_logs_emitted < self._attention_log_limit: + message = ( + f"[AttentionProbe] mean={summary['overall_mean']:.6f} " + f"max={summary['overall_max']:.6f}" + ) + if "supervised_mean" in summary: + message += f" supervised_mean={summary['supervised_mean']:.6f}" + print(message, flush=True) + self._attention_logs_emitted += 1 + else: + self._last_attention_summary = None + return output @@ -143,6 +192,10 @@ def __init__( self.hidden_size = hidden_size self.lora_rank = lora_rank self.lora_alpha = lora_alpha + self.debug_logging = False + self.last_attention_summary: Optional[dict] = None + self._attention_log_limit = 5 + self._attention_logs_emitted = 0 # Base cross-attention block self.cross_attention = CrossAttentionBlock( @@ -163,16 +216,31 @@ def __init__( bias="none", task_type="FEATURE_EXTRACTION" ) - + # Apply LoRA to cross-attention self.cross_attention = get_peft_model(self.cross_attention, self.lora_config) - + + def set_debug_logging(self, enabled: bool, log_limit: int = 5) -> None: + self.debug_logging = bool(enabled) + self._attention_log_limit = int(max(0, log_limit)) + self._attention_logs_emitted = 0 + + base_model = getattr(self.cross_attention, "base_model", None) + if base_model is not None: + base_model.debug_logging = self.debug_logging + base_model._attention_log_limit = self._attention_log_limit + base_model._attention_logs_emitted = 0 + + def configure_attention_probe(self, enabled: bool, log_limit: int = 5) -> None: + self.set_debug_logging(enabled, log_limit) + def forward( self, hidden_states: torch.Tensor, audio_tokens: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, - gate: float = 1.0 + gate: float = 1.0, + supervised_mask: Optional[torch.Tensor] = None ) -> torch.Tensor: """ Forward pass with LoRA fusion and gating. @@ -202,12 +270,35 @@ def forward( fused_states = self.cross_attention( hidden_states=hidden_states, audio_tokens=audio_tokens, - attention_mask=attention_mask + attention_mask=attention_mask, + supervised_mask=supervised_mask, ) - - # Apply gating: interpolate between original and fused states + + base_model = getattr(self.cross_attention, "base_model", None) + if base_model is not None: + self.last_attention_summary = getattr(base_model, "_last_attention_summary", None) + else: + self.last_attention_summary = None + if self.debug_logging and self.last_attention_summary is not None: + if self._attention_logs_emitted < self._attention_log_limit: + summary = self.last_attention_summary + msg = "[AttentionProbe]" + overall_mean = summary.get("overall_mean") + overall_max = summary.get("overall_max") + if overall_mean is not None: + msg += f" mean={overall_mean:.6f}" + if overall_max is not None: + msg += f" max={overall_max:.6f}" + supervised_mean = summary.get("supervised_mean") + if supervised_mean is not None: + msg += f" supervised_mean={supervised_mean:.6f}" + print(msg, flush=True) + self._attention_logs_emitted += 1 + + # Apply gating on the residual update so gate=0.0 is an exact no-op if gate != 1.0: - fused_states = gate * fused_states + (1.0 - gate) * hidden_states + delta = fused_states - hidden_states + fused_states = hidden_states + gate * delta return fused_states @@ -257,7 +348,8 @@ def forward( layer_idx: int, attention_mask: Optional[torch.Tensor] = None, gate: float = 1.0, - active_fusion_layer: Optional[int] = None + active_fusion_layer: Optional[int] = None, + supervised_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Apply fusion at specified layer. @@ -288,11 +380,19 @@ def forward( hidden_states=hidden_states, audio_tokens=audio_tokens, attention_mask=attention_mask, - gate=gate + gate=gate, + supervised_mask=supervised_mask, ) - + return fused_states + def set_debug_logging(self, enabled: bool, log_limit: int = 5) -> None: + for adapter in self.fusion_adapters.values(): + adapter.set_debug_logging(enabled, log_limit) + + def configure_attention_probe(self, enabled: bool, log_limit: int = 5) -> None: + self.set_debug_logging(enabled, log_limit) + class GatedFusionAdapter(nn.Module): """ @@ -334,7 +434,8 @@ def forward( hidden_states: torch.Tensor, audio_tokens: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, - force_gate: Optional[float] = None + force_gate: Optional[float] = None, + supervised_mask: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Forward with learnable gating. @@ -365,11 +466,18 @@ def forward( # Batch-wise gating fused_states = [] for i in range(hidden_states.shape[0]): + sample_attention = ( + attention_mask[i:i+1] + if attention_mask is not None + else None + ) + sample_mask = supervised_mask[i:i+1] if supervised_mask is not None else None sample_fused = self.fusion_adapter( hidden_states=hidden_states[i:i+1], audio_tokens=audio_tokens[i:i+1], - attention_mask=attention_mask[i:i+1] if attention_mask is not None else None, - gate=gate[i].item() + attention_mask=sample_attention, + gate=gate[i].item(), + supervised_mask=sample_mask, ) fused_states.append(sample_fused) fused_states = torch.cat(fused_states, dim=0) @@ -379,7 +487,14 @@ def forward( hidden_states=hidden_states, audio_tokens=audio_tokens, attention_mask=attention_mask, - gate=gate.item() if isinstance(gate, torch.Tensor) else gate + gate=gate.item() if isinstance(gate, torch.Tensor) else gate, + supervised_mask=supervised_mask, ) - + return fused_states, gate + + def set_debug_logging(self, enabled: bool, log_limit: int = 5) -> None: + self.fusion_adapter.set_debug_logging(enabled, log_limit) + + def configure_attention_probe(self, enabled: bool, log_limit: int = 5) -> None: + self.set_debug_logging(enabled, log_limit) diff --git a/safe/models/safe_model.py b/safe/models/safe_model.py index 120edf7..c142ba7 100644 --- a/safe/models/safe_model.py +++ b/safe/models/safe_model.py @@ -159,6 +159,40 @@ def __init__( # Update our stored hidden size to match the actual model self.llm_hidden_size = actual_hidden_size + + self.debug_logging = False + + def set_debug_logging(self, enabled: bool) -> None: + """Enable or disable verbose debugging across SAFE components.""" + + self.debug_logging = bool(enabled) + + if hasattr(self.audio_projector, "debug_logging"): + self.audio_projector.debug_logging = self.debug_logging + + if hasattr(self.audio_encoder, "set_debug_logging") and not self.debug_logging: + # Reset waveform logging counter when disabling global debug + self.audio_encoder.set_debug_logging(False) + + if hasattr(self.fusion_adapter, "set_debug_logging"): + self.fusion_adapter.set_debug_logging(self.debug_logging) + + def configure_audio_debug( + self, + waveform_stats: bool = False, + waveform_log_limit: int = 5, + ) -> None: + """Configure waveform statistics logging for the audio encoder.""" + + if hasattr(self.audio_encoder, "set_debug_logging"): + self.audio_encoder.set_debug_logging(waveform_stats, waveform_log_limit) + + def configure_attention_probe(self, enabled: bool, log_limit: int = 5) -> None: + if hasattr(self.fusion_adapter, "configure_attention_probe"): + self.fusion_adapter.configure_attention_probe(enabled, log_limit) + + def get_last_attention_summary(self) -> Optional[dict]: + return getattr(self.fusion_adapter, "last_attention_summary", None) def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: """ @@ -1065,11 +1099,16 @@ def forward( if audio_attention_mask is not None: audio_attention_mask = audio_attention_mask.to(attention_mask.device) + supervised_mask = None + if labels is not None: + supervised_mask = (labels != -100).to(inputs_embeds.device) + inputs_embeds = self.fusion_adapter( hidden_states=inputs_embeds, audio_tokens=audio_tokens, attention_mask=audio_attention_mask, gate=gate, + supervised_mask=supervised_mask, ) model_inputs = { @@ -1139,10 +1178,13 @@ def forward( # Forward through LLM with audio fusion (simplified implementation) if hasattr(self.base_vl.llm, 'transformer'): hidden_states = inputs_embeds - + supervised_mask = None + if labels is not None: + supervised_mask = (labels != -100).to(hidden_states.device) + for i, layer in enumerate(self.base_vl.llm.transformer.h): hidden_states = layer(hidden_states)[0] - + # Apply audio fusion at appropriate layers if self.fusion_type == "multilayer": hidden_states = self.fusion_adapter( @@ -1151,21 +1193,24 @@ def forward( layer_idx=i, attention_mask=None, gate=gate, - active_fusion_layer=fusion_layer + active_fusion_layer=fusion_layer, + supervised_mask=supervised_mask, ) elif i == len(self.base_vl.llm.transformer.h) // 2: # Mid-layer fusion if self.fusion_type == "gated": hidden_states, _ = self.fusion_adapter( hidden_states=hidden_states, audio_tokens=audio_tokens, - force_gate=gate if gate != 1.0 else None + force_gate=gate if gate != 1.0 else None, + supervised_mask=supervised_mask, ) else: hidden_states = self.fusion_adapter( hidden_states, audio_tokens, None, - gate + gate, + supervised_mask=supervised_mask, ) hidden_states = self.base_vl.llm.transformer.ln_f(hidden_states) diff --git a/safe/training/stage_a.py b/safe/training/stage_a.py index 77590d9..c0e3e31 100644 --- a/safe/training/stage_a.py +++ b/safe/training/stage_a.py @@ -1,3 +1,4 @@ +import math import torch import torch.nn as nn from torch.utils.data import DataLoader @@ -90,6 +91,14 @@ def __init__( "null_space_refresh_interval": 2000, "null_space_verbose": False, "audio_label_smoothing": 0.1, + "enable_audio_sanity_checks": True, + "log_waveform_stats": True, + "waveform_log_limit": 8, + "log_attention_probe": True, + "attention_probe_log_limit": 8, + "log_grad_norms": True, + "grad_log_interval": 1, + "grad_log_limit": 20, } if config: @@ -148,7 +157,25 @@ def __init__( self._teacher_broadcast_warned = False self._shape_debug_once = False self._eval_shape_debug_once = False - setattr(self.safe_model, "debug_logging", self.debug_logging) + self._sanity_checks_completed = False + self._attention_logs_emitted = 0 + self._grad_logs_emitted = 0 + + waveform_logs = int(self.config.get("waveform_log_limit", 8)) + waveform_debug = bool(self.config.get("log_waveform_stats", False)) + attention_probe = bool(self.config.get("log_attention_probe", False)) + attention_log_limit = int(self.config.get("attention_probe_log_limit", 5)) + + if hasattr(self.safe_model, "set_debug_logging"): + self.safe_model.set_debug_logging(self.debug_logging) + else: + setattr(self.safe_model, "debug_logging", self.debug_logging) + + if hasattr(self.safe_model, "configure_audio_debug"): + self.safe_model.configure_audio_debug(waveform_debug, waveform_logs) + + if hasattr(self.safe_model, "configure_attention_probe"): + self.safe_model.configure_attention_probe(attention_probe, attention_log_limit) # Move model to GPU if available if torch.cuda.is_available(): @@ -544,6 +571,202 @@ def _apply_warmup(self, step: int): print(f"Warning: base_lr is a sequence: {base_lr}, using first element", flush=True) base_lr = base_lr[0] g["lr"] = float(base_lr) * factor + + def _clone_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]: + cloned: Dict[str, Any] = {} + for key, value in inputs.items(): + if isinstance(value, torch.Tensor): + cloned[key] = value.detach().clone() + elif isinstance(value, list): + cloned[key] = list(value) + else: + cloned[key] = value + return cloned + + def _select_batch_indices( + self, + inputs: Dict[str, Any], + indices: torch.Tensor, + clone: bool = True, + ) -> Dict[str, Any]: + if indices.numel() == 0: + return {} + + if indices.dtype != torch.long: + indices = indices.to(torch.long) + + base_tensor = inputs.get("input_ids") + if not isinstance(base_tensor, torch.Tensor): + base_tensor = next( + (v for v in inputs.values() if isinstance(v, torch.Tensor) and v.dim() > 0), + None, + ) + batch_size = base_tensor.size(0) if isinstance(base_tensor, torch.Tensor) else None + + subset: Dict[str, Any] = {} + index_list = indices.tolist() + + for key, value in inputs.items(): + if isinstance(value, torch.Tensor) and batch_size is not None and value.size(0) == batch_size: + device_indices = indices.to(value.device) + sliced = value.index_select(0, device_indices) + subset[key] = sliced.detach().clone() if clone else sliced + elif isinstance(value, list) and batch_size is not None and len(value) == batch_size: + subset[key] = [value[int(i)] for i in index_list] + else: + if isinstance(value, torch.Tensor) and clone: + subset[key] = value.detach().clone() + else: + subset[key] = value + + return subset + + def _compute_audio_loss(self, inputs: Dict[str, Any], gate: float) -> Optional[float]: + labels = inputs.get("labels") + attention_mask = inputs.get("attention_mask") + if labels is None: + return None + + was_training = self.safe_model.training + self.safe_model.eval() + try: + with torch.no_grad(): + outputs = self.safe_model(**inputs, gate=gate) + logits = outputs.get("logits") if isinstance(outputs, dict) else getattr(outputs, "logits", None) + if logits is None: + return None + loss = self.audio_task_loss(logits, labels, attention_mask) + return float(loss.detach().item()) + finally: + if was_training: + self.safe_model.train() + + def _run_audio_sanity_checks(self, inputs: Dict[str, Any], has_audio: torch.Tensor) -> None: + if self._sanity_checks_completed: + return + if not self.config.get("enable_audio_sanity_checks", False): + self._sanity_checks_completed = True + return + + if has_audio.numel() == 0 or not torch.any(has_audio): + print("[SanityCheck] Skipping audio sanity checks: no audio samples in batch.", flush=True) + self._sanity_checks_completed = True + return + + audio_indices = torch.nonzero(has_audio, as_tuple=False).flatten() + audio_inputs = self._select_batch_indices(inputs, audio_indices) + if not audio_inputs or audio_inputs.get("audio_tokens") is None: + print("[SanityCheck] Audio tokens missing; cannot run sanity checks.", flush=True) + self._sanity_checks_completed = True + return + + print("\n[SanityCheck] Running audio fusion ablations on first audio batch...", flush=True) + + baseline_loss = self._compute_audio_loss(self._clone_inputs(audio_inputs), gate=1.0) + gate_off_loss = self._compute_audio_loss(self._clone_inputs(audio_inputs), gate=0.0) + + shuffled_inputs = self._clone_inputs(audio_inputs) + audio_tokens = shuffled_inputs.get("audio_tokens") + if isinstance(audio_tokens, torch.Tensor): + perm = torch.randperm(audio_tokens.size(0)).to(audio_tokens.device) + shuffled_inputs["audio_tokens"] = audio_tokens.index_select(0, perm) + shuffled_loss = self._compute_audio_loss(shuffled_inputs, gate=1.0) + + zero_inputs = self._clone_inputs(audio_inputs) + if isinstance(zero_inputs.get("audio_tokens"), torch.Tensor): + zero_inputs["audio_tokens"] = torch.zeros_like(zero_inputs["audio_tokens"]) + zero_loss = self._compute_audio_loss(zero_inputs, gate=1.0) + + def fmt(value: Optional[float]) -> str: + return "N/A" if value is None else f"{value:.6f}" + + print("[SanityCheck] Loss comparison (lower is better):", flush=True) + print(f" gate=1.0 (baseline): {fmt(baseline_loss)}", flush=True) + print(f" gate=0.0: {fmt(gate_off_loss)}", flush=True) + print(f" shuffled audio: {fmt(shuffled_loss)}", flush=True) + print(f" zeroed audio: {fmt(zero_loss)}", flush=True) + + self._sanity_checks_completed = True + + def _log_attention_probe(self) -> None: + if not self.config.get("log_attention_probe", False): + return + log_limit = int(self.config.get("attention_probe_log_limit", 0)) + if log_limit and self._attention_logs_emitted >= log_limit: + return + + summary = None + if hasattr(self.safe_model, "get_last_attention_summary"): + summary = self.safe_model.get_last_attention_summary() + elif hasattr(self.safe_model, "fusion_adapter"): + summary = getattr(self.safe_model.fusion_adapter, "last_attention_summary", None) + + if not summary: + return + + msg = f"[AttentionProbe][step {self.global_step}]" + if "overall_mean" in summary: + msg += f" overall_mean={summary['overall_mean']:.6f}" + if "overall_max" in summary: + msg += f" overall_max={summary['overall_max']:.6f}" + if "supervised_mean" in summary: + msg += f" supervised_mean={summary['supervised_mean']:.6f}" + print(msg, flush=True) + self._attention_logs_emitted += 1 + + def _log_grad_norms(self) -> None: + if not self.config.get("log_grad_norms", False): + return + log_limit = int(self.config.get("grad_log_limit", 0)) + if log_limit and self._grad_logs_emitted >= log_limit: + return + interval = max(1, int(self.config.get("grad_log_interval", 1))) + if self.global_step % interval != 0: + return + + projector_sq = 0.0 + fusion_sq = 0.0 + token_sq = 0.0 + overall_sq = 0.0 + projector_params = fusion_params = token_params = overall_params = 0 + + for name, param in self.safe_model.named_parameters(): + if not param.requires_grad or param.grad is None: + continue + grad = param.grad.detach() + norm = grad.norm(2).item() + overall_sq += norm ** 2 + overall_params += 1 + lower = name.lower() + if "audio_projector" in lower: + projector_sq += norm ** 2 + projector_params += 1 + elif "fusion_adapter" in lower or "lora" in lower: + fusion_sq += norm ** 2 + fusion_params += 1 + elif "audio_token_embeddings" in lower: + token_sq += norm ** 2 + token_params += 1 + + def safe_sqrt(value: float) -> float: + return float(math.sqrt(value)) if value > 0 else 0.0 + + print( + "[GradDebug] step {}: projector_norm={:.6f} (n={}) fusion_norm={:.6f} (n={}) " + "audio_token_norm={:.6f} (n={}) overall_norm={:.6f} (n={})".format( + self.global_step, + safe_sqrt(projector_sq), + projector_params, + safe_sqrt(fusion_sq), + fusion_params, + safe_sqrt(token_sq), + token_params, + safe_sqrt(overall_sq), + overall_params, + ), + flush=True, + ) + self._grad_logs_emitted += 1 def train_step(self, batch: Dict) -> Dict[str, float]: """ @@ -605,7 +828,14 @@ def train_step(self, batch: Dict) -> Dict[str, float]: print(f"[AUDIO_DEBUG] Answers data type: {type(answers_data)}", flush=True) # Prepare inputs for SAFE model (device-consistent masks) - has_audio = batch.get("has_audio", torch.zeros(len(batch["questions"]), dtype=torch.bool, device=device)) + has_audio = batch.get( + "has_audio", + torch.zeros(len(batch["questions"]), dtype=torch.bool, device=device), + ) + if isinstance(has_audio, torch.Tensor): + has_audio = has_audio.to(device=device, dtype=torch.bool) + else: + has_audio = torch.tensor(has_audio, dtype=torch.bool, device=device) print(f"[AUDIO_DEBUG] has_audio flag: {has_audio.sum().item()}/{len(has_audio)} samples marked as having audio", flush=True) # Create input tensors for training - apply answers for supervised learning @@ -638,13 +868,16 @@ def train_step(self, batch: Dict) -> Dict[str, float]: f"[TrainDebug] step {self.global_step}: prepared inputs input_ids={input_shape} has_audio={int(has_audio.sum().item())}", flush=True, ) - + + self._run_audio_sanity_checks(inputs, has_audio) + # Optional debug: Check multimodal input preparation (disabled for clean logs) # print(f"DEBUG: After prepare_multimodal_inputs, keys: {inputs.keys()}") # print(f"DEBUG: Has pixel_values: {'pixel_values' in inputs}") - + # Forward pass through SAFE model safe_outputs = self.safe_model(**inputs) + self._log_attention_probe() if self.debug_logging: if torch.cuda.is_available(): torch.cuda.synchronize() @@ -824,6 +1057,8 @@ def train_step(self, batch: Dict) -> Dict[str, float]: # Backward pass total_loss.backward() + self._log_grad_norms() + if self.null_space_projector is not None: self.null_space_projector.observe(step=self.global_step, has_audio=has_audio) self.null_space_projector.project() @@ -929,7 +1164,14 @@ def evaluate(self, max_batches: Optional[int] = None) -> Dict[str, float]: if isinstance(batch[key], torch.Tensor): batch[key] = batch[key].to(device) - has_audio = batch.get("has_audio", torch.zeros(len(batch["questions"]), dtype=torch.bool, device=device)) + has_audio = batch.get( + "has_audio", + torch.zeros(len(batch["questions"]), dtype=torch.bool, device=device), + ) + if isinstance(has_audio, torch.Tensor): + has_audio = has_audio.to(device=device, dtype=torch.bool) + else: + has_audio = torch.tensor(has_audio, dtype=torch.bool, device=device) # Prepare inputs - always pass images/audio if available, don't gate on flags if self.debug_logging: