Training in progress - step 1000

config.json

@@ -179,7 +179,7 @@
   "projector_hidden_dim": null,
   "projector_init_std": 0.02,
   "projector_num_layers": 2,
-  "projector_pool_stride": 2,
+  "projector_pool_stride": 6,
   "projector_type": "mosa",
   "qformer_hidden_size": null,
   "qformer_intermediate_size": null,

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7c92dfb13651aa8309de3cde22e20237407f7593163b9c3d668c5be22a99be7e
-size 152361992
+oid sha256:2427b0eb3d43c1582d4e52125d39843fa62119c65b9d4b79f410f0de1c4386da
+size 320134160

projectors.py

@@ -104,6 +104,11 @@ class MOSAProjector(nn.Module):
         self.num_experts = getattr(config, "num_experts", None) or 4  # MOSA-Base uses 4
         adapter_hidden = getattr(config, "adapter_hidden_dim", None) or 4096
 
+        # Optional pre-norm before projection
+        self.use_pre_norm = getattr(config, "projector_pre_norm", False)
+        if self.use_pre_norm:
+            self.pre_norm = LlamaRMSNorm(self.encoder_dim, eps=1e-8)
+
         # Frame stacking: concat k adjacent frames then project
         in_dim = self.encoder_dim * self.k
 
@@ -126,6 +131,10 @@ class MOSAProjector(nn.Module):
         # x: (B, S, encoder_dim)
         batch_size, seq_len, dim = x.shape
 
+        # Apply pre-norm if enabled
+        if self.use_pre_norm:
+            x = self.pre_norm(x)
+
         # --- 1. Router Branch ---
         # Mean pool encoder outputs for routing decisions
         x_pooled = x.reshape(batch_size, -1, self.k, self.encoder_dim).mean(dim=2)  # (B, S//k, D)