Update model weights after training (epoch 4, loss 3.0820)

audio_decoder.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:56f890d89fb28ae5df4b1ab79c42b0c29edc81203f2a99ea077d47909b6d128a
-size 1458415836
+oid sha256:840aaf132e4830dfcfa0634d27acab02841f8eb9fffbfe4f78377c1d50aa050a
+size 1458410612

audio_encoder.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:30c6b7e43c61588099a04e970a49bd38fd73465ef129d8e39ed9a1e8c45aeecf
-size 466150140
+oid sha256:68b7ed71f13950d93d17d4152e38cdcdd5e1a157729f4615ee38072473e8c12a
+size 466119380

cross_attention.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6d6a01050517c1c2762257ca7f0c03259704ec620070948f9cccd9c26476fcae
+oid sha256:5343b1fe1af46ca860a51de6f3bd51d1843f70998850084f805c875aec2de030
 size 174191400

generator.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:02673bab0a9e2949c2f3bfe0725ff77cd631fd89a559d86e9c886c99455a5e72
+oid sha256:499cee360f74c21e9e08624abd739f8cd982b339117fbcae3ae09433cdbebc71
 size 629440508

llm.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:816725a4e4eaaf6f5a2bb5c3cb678c13f298ccc280937e88c3c947d9fc052fb3
+oid sha256:4c2d81eae0ff676724bf38cf020b2e6317e609eb90d43150ffe91610e67864e7
 size 1506831304

model.safetensors.index.json

@@ -1,6 +1,6 @@
 {
   "metadata": {
-    "total_size": 7309396886,
+    "total_size": 7309365038,
     "format": "components"
   },
   "weight_map": {
@@ -1951,23 +1951,14 @@
     "audio_encoder.speaker_encoder.frame_encoder.0.bias": "audio_encoder.safetensors",
     "audio_encoder.speaker_encoder.frame_encoder.2.weight": "audio_encoder.safetensors",
     "audio_encoder.speaker_encoder.frame_encoder.2.bias": "audio_encoder.safetensors",
-    "audio_encoder.speaker_encoder.frame_encoder.2.running_mean": "audio_encoder.safetensors",
-    "audio_encoder.speaker_encoder.frame_encoder.2.running_var": "audio_encoder.safetensors",
-    "audio_encoder.speaker_encoder.frame_encoder.2.num_batches_tracked": "audio_encoder.safetensors",
     "audio_encoder.speaker_encoder.frame_encoder.3.weight": "audio_encoder.safetensors",
     "audio_encoder.speaker_encoder.frame_encoder.3.bias": "audio_encoder.safetensors",
     "audio_encoder.speaker_encoder.frame_encoder.5.weight": "audio_encoder.safetensors",
     "audio_encoder.speaker_encoder.frame_encoder.5.bias": "audio_encoder.safetensors",
-    "audio_encoder.speaker_encoder.frame_encoder.5.running_mean": "audio_encoder.safetensors",
-    "audio_encoder.speaker_encoder.frame_encoder.5.running_var": "audio_encoder.safetensors",
-    "audio_encoder.speaker_encoder.frame_encoder.5.num_batches_tracked": "audio_encoder.safetensors",
     "audio_encoder.speaker_encoder.frame_encoder.6.weight": "audio_encoder.safetensors",
     "audio_encoder.speaker_encoder.frame_encoder.6.bias": "audio_encoder.safetensors",
     "audio_encoder.speaker_encoder.frame_encoder.8.weight": "audio_encoder.safetensors",
     "audio_encoder.speaker_encoder.frame_encoder.8.bias": "audio_encoder.safetensors",
-    "audio_encoder.speaker_encoder.frame_encoder.8.running_mean": "audio_encoder.safetensors",
-    "audio_encoder.speaker_encoder.frame_encoder.8.running_var": "audio_encoder.safetensors",
-    "audio_encoder.speaker_encoder.frame_encoder.8.num_batches_tracked": "audio_encoder.safetensors",
     "audio_encoder.speaker_encoder.lstm.weight_ih_l0": "audio_encoder.safetensors",
     "audio_encoder.speaker_encoder.lstm.weight_hh_l0": "audio_encoder.safetensors",
     "audio_encoder.speaker_encoder.lstm.bias_ih_l0": "audio_encoder.safetensors",
@@ -2033,9 +2024,6 @@
     "audio_encoder.conformer_blocks.0.conv.depthwise_conv.bias": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.0.conv.batch_norm.weight": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.0.conv.batch_norm.bias": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.0.conv.batch_norm.running_mean": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.0.conv.batch_norm.running_var": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.0.conv.batch_norm.num_batches_tracked": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.0.conv.pointwise_conv2.weight": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.0.conv.pointwise_conv2.bias": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.0.ff2_norm.weight": "audio_encoder.safetensors",
@@ -2069,9 +2057,6 @@
     "audio_encoder.conformer_blocks.1.conv.depthwise_conv.bias": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.1.conv.batch_norm.weight": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.1.conv.batch_norm.bias": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.1.conv.batch_norm.running_mean": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.1.conv.batch_norm.running_var": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.1.conv.batch_norm.num_batches_tracked": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.1.conv.pointwise_conv2.weight": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.1.conv.pointwise_conv2.bias": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.1.ff2_norm.weight": "audio_encoder.safetensors",
@@ -2105,9 +2090,6 @@
     "audio_encoder.conformer_blocks.2.conv.depthwise_conv.bias": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.2.conv.batch_norm.weight": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.2.conv.batch_norm.bias": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.2.conv.batch_norm.running_mean": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.2.conv.batch_norm.running_var": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.2.conv.batch_norm.num_batches_tracked": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.2.conv.pointwise_conv2.weight": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.2.conv.pointwise_conv2.bias": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.2.ff2_norm.weight": "audio_encoder.safetensors",
@@ -2141,9 +2123,6 @@
     "audio_encoder.conformer_blocks.3.conv.depthwise_conv.bias": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.3.conv.batch_norm.weight": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.3.conv.batch_norm.bias": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.3.conv.batch_norm.running_mean": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.3.conv.batch_norm.running_var": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.3.conv.batch_norm.num_batches_tracked": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.3.conv.pointwise_conv2.weight": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.3.conv.pointwise_conv2.bias": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.3.ff2_norm.weight": "audio_encoder.safetensors",
@@ -2177,9 +2156,6 @@
     "audio_encoder.conformer_blocks.4.conv.depthwise_conv.bias": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.4.conv.batch_norm.weight": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.4.conv.batch_norm.bias": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.4.conv.batch_norm.running_mean": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.4.conv.batch_norm.running_var": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.4.conv.batch_norm.num_batches_tracked": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.4.conv.pointwise_conv2.weight": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.4.conv.pointwise_conv2.bias": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.4.ff2_norm.weight": "audio_encoder.safetensors",
@@ -2213,9 +2189,6 @@
     "audio_encoder.conformer_blocks.5.conv.depthwise_conv.bias": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.5.conv.batch_norm.weight": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.5.conv.batch_norm.bias": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.5.conv.batch_norm.running_mean": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.5.conv.batch_norm.running_var": "audio_encoder.safetensors",
-    "audio_encoder.conformer_blocks.5.conv.batch_norm.num_batches_tracked": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.5.conv.pointwise_conv2.weight": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.5.conv.pointwise_conv2.bias": "audio_encoder.safetensors",
     "audio_encoder.conformer_blocks.5.ff2_norm.weight": "audio_encoder.safetensors",
@@ -2588,30 +2561,18 @@
     "audio_decoder.postnet.0.0.bias": "audio_decoder.safetensors",
     "audio_decoder.postnet.0.1.weight": "audio_decoder.safetensors",
     "audio_decoder.postnet.0.1.bias": "audio_decoder.safetensors",
-    "audio_decoder.postnet.0.1.running_mean": "audio_decoder.safetensors",
-    "audio_decoder.postnet.0.1.running_var": "audio_decoder.safetensors",
-    "audio_decoder.postnet.0.1.num_batches_tracked": "audio_decoder.safetensors",
     "audio_decoder.postnet.1.0.weight": "audio_decoder.safetensors",
     "audio_decoder.postnet.1.0.bias": "audio_decoder.safetensors",
     "audio_decoder.postnet.1.1.weight": "audio_decoder.safetensors",
     "audio_decoder.postnet.1.1.bias": "audio_decoder.safetensors",
-    "audio_decoder.postnet.1.1.running_mean": "audio_decoder.safetensors",
-    "audio_decoder.postnet.1.1.running_var": "audio_decoder.safetensors",
-    "audio_decoder.postnet.1.1.num_batches_tracked": "audio_decoder.safetensors",
     "audio_decoder.postnet.2.0.weight": "audio_decoder.safetensors",
     "audio_decoder.postnet.2.0.bias": "audio_decoder.safetensors",
     "audio_decoder.postnet.2.1.weight": "audio_decoder.safetensors",
     "audio_decoder.postnet.2.1.bias": "audio_decoder.safetensors",
-    "audio_decoder.postnet.2.1.running_mean": "audio_decoder.safetensors",
-    "audio_decoder.postnet.2.1.running_var": "audio_decoder.safetensors",
-    "audio_decoder.postnet.2.1.num_batches_tracked": "audio_decoder.safetensors",
     "audio_decoder.postnet.3.0.weight": "audio_decoder.safetensors",
     "audio_decoder.postnet.3.0.bias": "audio_decoder.safetensors",
     "audio_decoder.postnet.3.1.weight": "audio_decoder.safetensors",
     "audio_decoder.postnet.3.1.bias": "audio_decoder.safetensors",
-    "audio_decoder.postnet.3.1.running_mean": "audio_decoder.safetensors",
-    "audio_decoder.postnet.3.1.running_var": "audio_decoder.safetensors",
-    "audio_decoder.postnet.3.1.num_batches_tracked": "audio_decoder.safetensors",
     "audio_decoder.postnet.4.weight": "audio_decoder.safetensors",
     "audio_decoder.postnet.4.bias": "audio_decoder.safetensors",
     "audio_decoder.waveform_decoder.input_proj.bias": "audio_decoder.safetensors",

modeling_xoron.py

@@ -4371,16 +4371,18 @@ class SpeakerEncoder(nn.Module):
         self.output_size = output_size
 
         # Frame-level encoder
+        # Use GroupNorm instead of BatchNorm1d to handle sequence length of 1
+        # GroupNorm(1, C) is equivalent to LayerNorm and works with any batch/seq size
         self.frame_encoder = nn.Sequential(
             nn.Conv1d(80, hidden_size, 5, 1, 2),
             nn.ReLU(),
-            nn.BatchNorm1d(hidden_size),
+            nn.GroupNorm(1, hidden_size),
             nn.Conv1d(hidden_size, hidden_size, 5, 1, 2),
             nn.ReLU(),
-            nn.BatchNorm1d(hidden_size),
+            nn.GroupNorm(1, hidden_size),
             nn.Conv1d(hidden_size, hidden_size, 5, 1, 2),
             nn.ReLU(),
-            nn.BatchNorm1d(hidden_size),
+            nn.GroupNorm(1, hidden_size),
         )
 
         # LSTM for temporal modeling
@@ -4853,7 +4855,8 @@ class ConvolutionModule(nn.Module):
             channels, channels, kernel_size=kernel_size,
             padding=(kernel_size - 1) // 2, groups=channels
         )
-        self.batch_norm = nn.BatchNorm1d(channels)
+        # Use GroupNorm instead of BatchNorm1d to handle sequence length of 1
+        self.batch_norm = nn.GroupNorm(1, channels)
         self.pointwise_conv2 = nn.Conv1d(channels, channels, kernel_size=1)
         self.dropout = nn.Dropout(dropout)
 
@@ -5544,25 +5547,27 @@ class AudioDecoder(nn.Module):
         self.mel_linear = nn.Linear(hidden_size, n_mels)
 
         # Postnet
+        # Use GroupNorm instead of BatchNorm1d to handle sequence length of 1
+        # GroupNorm(1, C) is equivalent to LayerNorm and works with any batch/seq size
         self.postnet = nn.ModuleList([
             nn.Sequential(
                 nn.Conv1d(n_mels, 256, kernel_size=5, padding=2),
-                nn.BatchNorm1d(256),
+                nn.GroupNorm(1, 256),
                 nn.Tanh(),
             ),
             nn.Sequential(
                 nn.Conv1d(256, 256, kernel_size=5, padding=2),
-                nn.BatchNorm1d(256),
+                nn.GroupNorm(1, 256),
                 nn.Tanh(),
             ),
             nn.Sequential(
                 nn.Conv1d(256, 256, kernel_size=5, padding=2),
-                nn.BatchNorm1d(256),
+                nn.GroupNorm(1, 256),
                 nn.Tanh(),
             ),
             nn.Sequential(
                 nn.Conv1d(256, 256, kernel_size=5, padding=2),
-                nn.BatchNorm1d(256),
+                nn.GroupNorm(1, 256),
                 nn.Tanh(),
             ),
             nn.Conv1d(256, n_mels, kernel_size=5, padding=2),
@@ -5770,9 +5775,8 @@ class AudioDecoder(nn.Module):
         energy_pred = F.softplus(self.energy_predictor(x))
 
         # Determine output length
-        # IMPORTANT: BatchNorm1d requires sequence length > 1 during training
-        # Enforce minimum length of 2 to avoid "Expected more than 1 value per channel" error
-        MIN_MEL_LENGTH = 2
+        # Note: We use GroupNorm instead of BatchNorm1d so any sequence length works
+        MIN_MEL_LENGTH = 1
         if target_length is not None:
             mel_length = max(MIN_MEL_LENGTH, target_length)
         else:
@@ -8608,6 +8612,7 @@ class AuxLosslessMoELayer(nn.Module):
 
     def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         batch_size, seq_len, hidden_size = hidden_states.shape
+        original_dtype = hidden_states.dtype
         hidden_flat = hidden_states.view(-1, hidden_size)
 
         top_k_probs, top_k_indices, _ = self.router(hidden_states)
@@ -8622,10 +8627,11 @@ class AuxLosslessMoELayer(nn.Module):
                     expert_input = hidden_flat[mask]
                     expert_output = expert(expert_input)
                     weight = top_k_probs[mask, k:k+1]
-                    final_output[mask] = final_output[mask] + weight * expert_output
+                    weighted_output = (weight * expert_output).to(original_dtype)
+                    final_output[mask] = final_output[mask] + weighted_output
 
         shared_output = self.shared_expert(hidden_flat)
-        final_output = final_output + shared_output
+        final_output = final_output + shared_output.to(original_dtype)
 
         final_output = final_output.view(batch_size, seq_len, hidden_size)
         
@@ -9125,7 +9131,14 @@ class XoronMultimodalModel(nn.Module):
         super().__init__()
         self.config = config
         self.device_map = device_map
-        self._model_parallel = device_map is not None and len(set(device_map.values())) > 1
+        
+        # Check for model parallelism - only consider string device values
+        # (device_map may contain metadata like 'training_gpus' list, 'dual_gpu_mode' bool)
+        if device_map is not None:
+            device_values = [v for v in device_map.values() if isinstance(v, str)]
+            self._model_parallel = len(set(device_values)) > 1
+        else:
+            self._model_parallel = False
 
         print("\n" + "=" * 60)
         print("🚀 BUILDING XORON-DEV MULTIMODAL MODEL")
@@ -9273,7 +9286,11 @@ class XoronMultimodalModel(nn.Module):
     def apply_model_parallel(self, device_map: Dict[str, str]):
         """Apply Model Parallelism by placing components on different devices."""
         self.device_map = device_map
-        self._model_parallel = len(set(device_map.values())) > 1
+        
+        # Check for model parallelism - only consider string device values
+        # (device_map may contain metadata like 'training_gpus' list, 'dual_gpu_mode' bool)
+        device_values = [v for v in device_map.values() if isinstance(v, str)]
+        self._model_parallel = len(set(device_values)) > 1
 
         if not self._model_parallel:
             print("   ℹ️ Single device - no model parallelism needed")

streaming_state.json

@@ -1,15 +1,15 @@
 {
-  "epoch": 12,
-  "unique_samples": 3029,
-  "total_yields": 6058,
+  "epoch": 19,
+  "unique_samples": 300,
+  "total_yields": 600,
   "dataset_positions": {
-    "WebSight": 186,
-    "ScienceQA": 164,
-    "InstructPix2Pix": 186,
-    "Flickr8k": 186,
-    "NewYorker": 186,
+    "WebSight": 386,
+    "ScienceQA": 364,
+    "InstructPix2Pix": 386,
+    "Flickr8k": 386,
+    "NewYorker": 386,
     "Football": 6,
-    "MagicBrush": 186,
+    "MagicBrush": 386,
     "WildChat": 200,
     "Synth-ShellExecution": 200,
     "Midjourney-Prompts": 200,
@@ -113,20 +113,20 @@
       "OpenAssistant": 200
     },
     "image": {
-      "WebSight": 186,
-      "ScienceQA": 164,
-      "InstructPix2Pix": 186,
-      "Flickr8k": 186,
-      "NewYorker": 186,
+      "WebSight": 386,
+      "ScienceQA": 364,
+      "InstructPix2Pix": 386,
+      "Flickr8k": 386,
+      "NewYorker": 386,
       "Football": 6,
-      "MagicBrush": 186
+      "MagicBrush": 386
     },
     "video": {},
     "audio": {}
   },
   "modality_counts": {
-    "text": 3029,
-    "image": 0,
+    "text": 0,
+    "image": 300,
     "video": 0,
     "audio": 0
   },

trainer_state.json

@@ -1,14 +1,14 @@
 {
   "best_model_checkpoint": "/kaggle/working/xoron-final",
-  "best_metric": 3.646694382440487,
+  "best_metric": 3.0820325045382684,
   "epoch": 4,
   "epochs_completed": 4,
-  "global_step": 1597,
+  "global_step": 148,
   "is_local_process_zero": true,
   "is_world_process_zero": true,
   "log_history": [],
   "logging_steps": 50,
-  "max_steps": 1597,
+  "max_steps": 148,
   "num_train_epochs": 4,
   "total_flos": 0,
   "train_batch_size": 1,
@@ -16,16 +16,16 @@
   "learning_rate": 0.0001,
   "max_grad_norm": 1.0,
   "trainable_components": [
+    "vision",
     "llm",
     "cross_attention",
+    "image_generation",
     "modality_markers"
   ],
   "frozen_components": [
-    "vision",
     "video",
     "audio",
     "speech",
-    "image_generation",
     "video_generation"
   ],
   "trial_name": null,

training_state.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9be0dff94c6d235091cae224c3e034a33fa84932af351d74ac37a512956c5486
-size 781495681
+oid sha256:d17b71b1b8d8d73a29371b107d2020d349cf453a9089b49b44d1b5cb446fba74
+size 1419723549