project complete

modeling_molmo2.py

@@ -758,7 +758,7 @@ class Molmo2DecoderLayer(GradientCheckpointingLayer):
     ):
         super().__init__()
         self.config = config
-
+        self.device = device
         self.self_attn = Molmo2Attention(config, layer_idx)
         self.attn_norm = Molmo2RMSNorm(
             config.hidden_size, eps=config.layer_norm_eps, device=device)
@@ -782,6 +782,7 @@ class Molmo2DecoderLayer(GradientCheckpointingLayer):
     ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
 
         residual = hidden_states
+        block_device = self.device
         hidden_states = self.attn_norm(hidden_states)
 
         # Self Attention
@@ -797,12 +798,23 @@ class Molmo2DecoderLayer(GradientCheckpointingLayer):
             **kwargs,
         )
 
+        def move_if_needed(x, device):
+            if x is None:
+                return x
+            if isinstance(x, tuple):
+                return tuple(move_if_needed(xx, device) for xx in x)
+            if hasattr(x, "device") and x.device != device:
+                return x.to(device)
+            return x
+
+        residual = move_if_needed(residual, block_device)
         hidden_states = residual + self.dropout(hidden_states)
 
         # Fully Connected
         residual = hidden_states
         hidden_states = self.ff_norm(hidden_states)
         hidden_states = self.mlp(hidden_states)
+        residual = move_if_needed(residual, block_device)
 
         hidden_states = residual + self.dropout(hidden_states)