Fix dtype mismatches: cast entire model to bfloat16 in init

modeling_slip.py

@@ -161,18 +161,18 @@ class MultiSizePatchEmbed(nn.Module):
         return new_w, base_b, new_res_w, res_b
 
     def forward(self, x_list, attention_mask, time_idx):
-        amp_dtype = torch.get_autocast_gpu_dtype() if torch.is_autocast_enabled() else torch.float32
-        device = torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else torch.device("cpu")
+        device = self.shared_linear.weight.device
+        dtype = self.shared_linear.weight.dtype
         sizes = torch.tensor([x.shape[-1] for x in x_list])
         unique_sizes = sizes.unique(sorted=True)
         N = x_list[0].shape[0]
-        outputs = torch.empty(len(x_list), N, self.intermediate_size, device=device, dtype=amp_dtype)
-        res_outputs = torch.empty(len(x_list), N, self.hidden_size, device=device, dtype=amp_dtype)
+        outputs = torch.empty(len(x_list), N, self.intermediate_size, device=device, dtype=dtype)
+        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, 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.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.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,9 +193,10 @@ class PatchEmbedding(nn.Module):
         self.residual_layer = nn.Linear(patch_size * 3, hidden_size)
 
     def forward(self, x, mask, time_idx):
-        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)
+        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 = torch.cat([x, mask, time_idx], dim=-1)
         return self.dropout(self.output_layer(self.act(self.hidden_layer(x)))) + self.residual_layer(x)
 
@@ -243,6 +244,9 @@ class CrossAttention(nn.Module):
         self.o_proj = nn.Linear(dim, dim, bias=False)
 
     def forward(self, query, context, attention_mask=None, **kwargs):
+        dtype = self.q_proj.weight.dtype
+        query = query.to(dtype=dtype)
+        context = context.to(dtype=dtype)
         bsz, q_len, _ = query.size()
         k_len = context.size(1)
         query = self.norm(query)
@@ -426,11 +430,11 @@ class Gemma3MultimodalModel(nn.Module):
         # Only init from pretrained when training from scratch.
         if init_from_pretrained:
             self.model = AutoModelForCausalLM.from_pretrained(
-                model_id, dtype=torch.bfloat16, trust_remote_code=True)
+                model_id, trust_remote_code=True)
         else:
             config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
             self.model = AutoModelForCausalLM.from_config(
-                config, torch_dtype=torch.bfloat16, trust_remote_code=True)
+                config, trust_remote_code=True)
 
         self.split_layer = split_layer
         hidden_size = self.model.config.hidden_size
@@ -442,10 +446,13 @@ class Gemma3MultimodalModel(nn.Module):
                 dim=hidden_size, context_dim=hidden_size, num_heads=num_heads, dropout_rate=0.1)
             self.model.model.layers[i] = Gemma3MultimodalLayer(
                 self.model.model.layers[i], Residual(cross_attn))
+
+        # Ensure all parameters (including new cross-attention modules) are bfloat16
         self.to(torch.bfloat16)
 
     def condition_image(self, image_embeds):
-        self.image_embeds = image_embeds.to(next(self.parameters()).device, dtype=torch.bfloat16)
+        param = next(self.parameters())
+        self.image_embeds = image_embeds.to(device=param.device, dtype=param.dtype)
         for layer in self.model.model.layers:
             if isinstance(layer, Gemma3MultimodalLayer):
                 layer.condition_vis_x(self.image_embeds)
@@ -536,6 +543,9 @@ class SLIPModel(SLIPPreTrainedModel):
         self.temperature = nn.Parameter(torch.tensor(math.log(1 / 0.07)))
         self.temperature_max = math.log(1 / 0.07)
 
+        # Cast entire model to bfloat16 (matches training checkpoint dtype)
+        self.to(torch.bfloat16)
+
     def embed_sensor(self, sensors, sensor_attn_mask=None, time_index=None):
         sensor_tokens, attn_mask = self.sensor_encoder(sensors, sensor_attn_mask, time_index=time_index)
         if hasattr(self, "img_attn_pool"):