changed nn.linear to use 4-bit quant

modeling_molmo2.py

@@ -91,7 +91,7 @@ class ViTMLP(nn.Module):
         self.w1  = bnb.nn.Linear4bit(dim, hidden_dim, bias=True, quant_type="nf4", device=device)
 
         self.act = ACT2FN[hidden_act]
-        self.w2  = bnb.nn.Linear4bit(dim, hidden_dim, bias=True, quant_type="nf4", device=device)
+        self.w2  = bnb.nn.Linear4bit(hidden_dim, dim, bias=True, quant_type="nf4", device=device)
     
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.w2(self.act(self.w1(x)))
@@ -258,7 +258,7 @@ class Molmo2VisionBlock(nn.Module):
             num_heads=config.num_attention_heads,
             num_key_value_heads=config.num_key_value_heads,
             head_dim=config.head_dim,
-            float32_attention=False,
+            float32_attention=config.float32_attention,
             attention_dropout=config.attention_dropout,
             residual_dropout=config.residual_dropout,
             device=device,