Duplicate from allenai/Molmo2-8B

Co-authored-by: Sangho Lee <sanghol@users.noreply.huggingface.co>

(cherry picked from commit 3d666677251284f0f13befe3bbd34062af653bbc)

modeling_molmo2.py

@@ -7,8 +7,6 @@ import torch
 from torch import nn
 from torch.nn import functional as F
 
-import bitsandbytes as bnb
-
 from transformers.models.auto import AutoModelForImageTextToText
 from transformers.activations import ACT2FN
 from transformers.configuration_utils import PretrainedConfig
@@ -88,10 +86,9 @@ class Molmo2ModelOutputWithPast(BaseModelOutputWithPast):
 class ViTMLP(nn.Module):
     def __init__(self, dim: int, hidden_dim: int, hidden_act: str, device: Union[str, torch.device] = None):
         super().__init__()
-        self.w1  = bnb.nn.Linear4bit(dim, hidden_dim, bias=True, quant_type="nf4", device=device)
-
+        self.w1 = nn.Linear(dim, hidden_dim, bias=True, device=device)
         self.act = ACT2FN[hidden_act]
-        self.w2  = bnb.nn.Linear4bit(hidden_dim, dim, bias=True, quant_type="nf4", device=device)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=True, device=device)
     
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.w2(self.act(self.w1(x)))
@@ -124,36 +121,28 @@ class ViTMultiHeadDotProductAttention(nn.Module):
 
         input_dim = input_dim or hidden_size
 
-        self.wq  = bnb.nn.Linear4bit(
+        self.wq = nn.Linear(
             input_dim,
             self.num_heads * self.head_dim,
             bias=use_bias,
-            quant_type="nf4",
             device=device,
         )
-
-        self.wk = bnb.nn.Linear4bit(
+        self.wk = nn.Linear(
             input_dim,
             self.num_key_value_heads * self.head_dim,
             bias=use_bias,
-            quant_type="nf4",
             device=device,
         )
-
-        self.wv = bnb.nn.Linear4bit(
+        self.wv = nn.Linear(
             input_dim,
             self.num_key_value_heads * self.head_dim,
             bias=use_bias,
-            quant_type="nf4",
             device=device,
         )
-
-        self.wo = bnb.nn.Linear4bit(
+        self.wo = nn.Linear(
             self.num_heads * self.head_dim,
             self.hidden_size,
-            quant_type="nf4",
         )
-
         self.float32_attention = float32_attention
         self.attention_dropout = attention_dropout
         self.residual_dropout = nn.Dropout(residual_dropout)
@@ -305,12 +294,10 @@ class Molmo2VisionTransformer(nn.Module):
         )
 
         image_patch_size = config.image_patch_size
-
-        self.patch_embedding = bnb.nn.Linear4bit(
+        self.patch_embedding = nn.Linear(
             image_patch_size * image_patch_size * 3,
             config.hidden_size,
             bias=True,
-            quant_type="nf4",
             device=device,
         )
 
@@ -367,10 +354,9 @@ class ImageProjectorMLP(nn.Module):
         device: Union[str, torch.device] = None,
     ):
         super().__init__()
-
-        self.w1 = bnb.nn.Linear4bit(input_dim, hidden_dim, bias=False, quant_type="nf4", device=device)
-        self.w2 = bnb.nn.Linear4bit(hidden_dim, output_dim, bias=False, quant_type="nf4", device=device)
-        self.w3 = bnb.nn.Linear4bit(input_dim, hidden_dim, bias=False, quant_type="nf4", device=device)
+        self.w1 = nn.Linear(input_dim, hidden_dim, bias=False, device=device)
+        self.w2 = nn.Linear(hidden_dim, output_dim, bias=False, device=device)
+        self.w3 = nn.Linear(input_dim, hidden_dim, bias=False, device=device)
         self.act = ACT2FN[hidden_act]
     
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -649,12 +635,11 @@ class Molmo2Attention(nn.Module):
             config.head_dim * config.num_key_value_heads,
             config.head_dim * config.num_key_value_heads,
         )
-        self.att_proj = bnb.nn.Linear4bit(
+        self.att_proj = nn.Linear(
             config.hidden_size,
             sum(self.fused_dims),
             bias=config.qkv_bias,
-            quant_type="nf4",
-            )
+        )
 
         # Layer norms.
         self.k_norm: Optional[Molmo2RMSNorm] = None
@@ -676,12 +661,11 @@ class Molmo2Attention(nn.Module):
             self.qk_norm_type = config.qk_norm_type
 
         self.attention_dropout = config.attention_dropout
-
-        self.attn_out = bnb.nn.Linear4bit(
+        
+        self.attn_out = nn.Linear(
             config.head_dim * config.num_attention_heads,
             config.hidden_size,
             bias=False,
-            quant_type="nf4",
         )
 
     def forward(
@@ -752,9 +736,8 @@ class LanguageModelMLP(nn.Module):
         device: Union[str, torch.device] = None,
     ):
         super().__init__()
-        self.ff_proj = bnb.nn.Linear4bit(input_dim, intermediate_size * 2, bias=False, quant_type="nf4", device=device)
-        self.ff_out = bnb.nn.Linear4bit(intermediate_size, input_dim, bias=False, quant_type="nf4", device=device)
-
+        self.ff_proj = nn.Linear(input_dim, intermediate_size * 2, bias=False, device=device)
+        self.ff_out = nn.Linear(intermediate_size, input_dim, bias=False, device=device)
         self.act = ACT2FN[hidden_act]
     
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -920,7 +903,7 @@ class Molmo2PreTrainedModel(PreTrainedModel):
 
     def _init_weights(self, module):
         std = self.config.initializer_range
-        if isinstance(module, (bnb.nn.Linear4bit,)):
+        if isinstance(module, (nn.Linear,)):
             module.weight.data.normal_(mean=0.0, std=std)
             if module.bias is not None:
                 module.bias.data.zero_()
@@ -1592,7 +1575,7 @@ class Molmo2ForConditionalGeneration(Molmo2PreTrainedModel, GenerationMixin):
         super().__init__(config)
 
         self.model = Molmo2Model(config)
-        self.lm_head = bnb.nn.Linear4bit(config.hidden_size, config.vocab_size, bias=False, quant_type="nf4")
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.vocab_size = config.vocab_size
         
         # Initialize weights and apply final processing