Uniform float32 weights, clean dtype handling: modeling_slip.py

modeling_slip.py

@@ -170,9 +170,9 @@ class MultiSizePatchEmbed(nn.Module):
         res_outputs = torch.empty(len(x_list), N, self.hidden_size, device=device, dtype=dtype)
         for psize in unique_sizes.tolist():
             idxs = (sizes == psize).nonzero(as_tuple=True)[0]
-            xs = torch.stack([x_list[i] for i in idxs]).to(device=device, dtype=dtype, non_blocking=True)
-            mask = torch.stack([attention_mask[i] for i in idxs]).to(device=device, dtype=dtype, non_blocking=True)
-            ti = torch.stack([time_idx[i] for i in idxs]).to(device=device, dtype=dtype, non_blocking=True)
+            xs = torch.stack([x_list[i] for i in idxs]).to(device=device, non_blocking=True)
+            mask = torch.stack([attention_mask[i] for i in idxs]).to(device=device, non_blocking=True)
+            ti = torch.stack([time_idx[i] for i in idxs]).to(device=device, non_blocking=True)
             xs = torch.cat([xs, mask, ti], dim=-1)
             w, b, r_w, r_b = self.resize_weight(psize * 3)
             res_outputs[idxs] = F.linear(xs, r_w, r_b)
@@ -193,10 +193,9 @@ class PatchEmbedding(nn.Module):
         self.residual_layer = nn.Linear(patch_size * 3, hidden_size)
 
     def forward(self, x, mask, time_idx):
-        dtype = self.hidden_layer.weight.dtype
-        x = rearrange(x, 'bs nvar (nump ps) -> (bs nvar) nump ps', ps=self.patch_size).to(dtype=dtype)
-        mask = rearrange(mask, 'bs nvar (nump ps) -> (bs nvar) nump ps', ps=self.patch_size).to(dtype=dtype)
-        time_idx = rearrange(time_idx, 'bs nvar (nump ps) -> (bs nvar) nump ps', ps=self.patch_size).to(dtype=dtype)
+        x = rearrange(x, 'bs nvar (nump ps) -> (bs nvar) nump ps', ps=self.patch_size)
+        mask = rearrange(mask, 'bs nvar (nump ps) -> (bs nvar) nump ps', ps=self.patch_size)
+        time_idx = rearrange(time_idx, 'bs nvar (nump ps) -> (bs nvar) nump ps', ps=self.patch_size)
         x = torch.cat([x, mask, time_idx], dim=-1)
         return self.dropout(self.output_layer(self.act(self.hidden_layer(x)))) + self.residual_layer(x)
 
@@ -244,10 +243,6 @@ class CrossAttention(nn.Module):
         self.o_proj = nn.Linear(dim, dim, bias=False)
 
     def forward(self, query, context, attention_mask=None, **kwargs):
-        # Cast inputs to match weight dtype (handles mixed-precision boundaries)
-        weight_dtype = self.q_proj.weight.dtype
-        query = query.to(dtype=weight_dtype)
-        context = context.to(dtype=weight_dtype)
         bsz, q_len, _ = query.size()
         assert context.size(0) == bsz, (
             f"Context batch size ({context.size(0)}) must match query batch size ({bsz}). "
@@ -424,8 +419,6 @@ class Gemma3MultimodalLayer(nn.Module):
         assert self.vis_x is not None, "vis_x must be set before forward pass."
         outputs = self.original_layer(hidden_states, **kwargs)
         hidden_states = self.cross_attn_block(outputs[0], context=self.vis_x)
-        # Match dtype of original layer output (cross-attn may use different dtype)
-        hidden_states = hidden_states.to(dtype=outputs[0].dtype)
         return (hidden_states,) + outputs[1:]
 
 
@@ -440,6 +433,9 @@ class Gemma3MultimodalModel(nn.Module):
                 model_id, trust_remote_code=True)
         else:
             config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
+            # Override Gemma's default bfloat16 so all params init in float32;
+            # users control runtime dtype via AutoModel.from_pretrained(torch_dtype=...)
+            config.torch_dtype = torch.float32
             self.model = AutoModelForCausalLM.from_config(
                 config, trust_remote_code=True)
 
@@ -455,8 +451,7 @@ class Gemma3MultimodalModel(nn.Module):
                 self.model.model.layers[i], Residual(cross_attn))
 
     def condition_image(self, image_embeds):
-        param = next(self.parameters())
-        self.image_embeds = image_embeds.to(device=param.device, dtype=param.dtype)
+        self.image_embeds = image_embeds
         for layer in self.model.model.layers:
             if isinstance(layer, Gemma3MultimodalLayer):
                 layer.condition_vis_x(self.image_embeds)
@@ -487,7 +482,7 @@ class EmbedToLatents(nn.Module):
         self.to_latents = nn.Linear(dim, dim_latents, bias=False)
 
     def forward(self, x):
-        return F.normalize(self.to_latents(x.to(dtype=self.to_latents.weight.dtype)), dim=-1)
+        return F.normalize(self.to_latents(x), dim=-1)
 
 
 class SLIPPreTrainedModel(PreTrainedModel):