changed nn.linear to use 4-bit quant

modeling_molmo2.py

@@ -7,6 +7,8 @@ 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
@@ -86,9 +88,10 @@ 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 = nn.Linear(dim, hidden_dim, bias=True, device=device)
+        self.w1  = bnb.nn.Linear4bit(dim, hidden_dim, bias=True, quant_type="nf4", device=device)
+
         self.act = ACT2FN[hidden_act]
-        self.w2 = nn.Linear(hidden_dim, dim, bias=True, device=device)
+        self.w2  = bnb.nn.Linear4bit(dim, hidden_dim, bias=True, quant_type="nf4", device=device)
     
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.w2(self.act(self.w1(x)))
@@ -121,28 +124,36 @@ class ViTMultiHeadDotProductAttention(nn.Module):
 
         input_dim = input_dim or hidden_size
 
-        self.wq = nn.Linear(
+        self.wq  = bnb.nn.Linear4bit(
             input_dim,
             self.num_heads * self.head_dim,
             bias=use_bias,
+            quant_type="nf4",
             device=device,
         )
-        self.wk = nn.Linear(
+
+        self.wk = bnb.nn.Linear4bit(
             input_dim,
             self.num_key_value_heads * self.head_dim,
             bias=use_bias,
+            quant_type="nf4",
             device=device,
         )
-        self.wv = nn.Linear(
+
+        self.wv = bnb.nn.Linear4bit(
             input_dim,
             self.num_key_value_heads * self.head_dim,
             bias=use_bias,
+            quant_type="nf4",
             device=device,
         )
-        self.wo = nn.Linear(
+
+        self.wo = bnb.nn.Linear4bit(
             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)
@@ -247,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=config.float32_attention,
+            float32_attention=False,
             attention_dropout=config.attention_dropout,
             residual_dropout=config.residual_dropout,
             device=device,
@@ -258,7 +269,6 @@ class Molmo2VisionBlock(nn.Module):
         self.ffn_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, device=device)
     
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        print("dtype before norm:", x.dtype)
         x = x + self.attention(self.attention_norm(x))
         x = x + self.feed_forward(self.ffn_norm(x))
         return x
@@ -295,10 +305,12 @@ class Molmo2VisionTransformer(nn.Module):
         )
 
         image_patch_size = config.image_patch_size
-        self.patch_embedding = nn.Linear(
+
+        self.patch_embedding = bnb.nn.Linear4bit(
             image_patch_size * image_patch_size * 3,
             config.hidden_size,
             bias=True,
+            quant_type="nf4",
             device=device,
         )
 
@@ -355,9 +367,10 @@ class ImageProjectorMLP(nn.Module):
         device: Union[str, torch.device] = None,
     ):
         super().__init__()
-        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.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.act = ACT2FN[hidden_act]
     
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -636,11 +649,12 @@ class Molmo2Attention(nn.Module):
             config.head_dim * config.num_key_value_heads,
             config.head_dim * config.num_key_value_heads,
         )
-        self.att_proj = nn.Linear(
+        self.att_proj = bnb.nn.Linear4bit(
             config.hidden_size,
             sum(self.fused_dims),
             bias=config.qkv_bias,
-        )
+            quant_type="nf4",
+            )
 
         # Layer norms.
         self.k_norm: Optional[Molmo2RMSNorm] = None
@@ -662,11 +676,12 @@ class Molmo2Attention(nn.Module):
             self.qk_norm_type = config.qk_norm_type
 
         self.attention_dropout = config.attention_dropout
-        
-        self.attn_out = nn.Linear(
+
+        self.attn_out = bnb.nn.Linear4bit(
             config.head_dim * config.num_attention_heads,
             config.hidden_size,
             bias=False,
+            quant_type="nf4",
         )
 
     def forward(
@@ -737,8 +752,9 @@ class LanguageModelMLP(nn.Module):
         device: Union[str, torch.device] = None,
     ):
         super().__init__()
-        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.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.act = ACT2FN[hidden_act]
     
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -904,7 +920,7 @@ class Molmo2PreTrainedModel(PreTrainedModel):
 
     def _init_weights(self, module):
         std = self.config.initializer_range
-        if isinstance(module, (nn.Linear,)):
+        if isinstance(module, (bnb.nn.Linear4bit,)):
             module.weight.data.normal_(mean=0.0, std=std)
             if module.bias is not None:
                 module.bias.data.zero_()
@@ -1576,7 +1592,7 @@ class Molmo2ForConditionalGeneration(Molmo2PreTrainedModel, GenerationMixin):
         super().__init__(config)
 
         self.model = Molmo2Model(config)
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.lm_head = bnb.nn.Linear4bit(config.hidden_size, config.vocab_size, bias=False, quant_type="nf4")
         self.vocab_size = config.vocab_size
         
         # Initialize weights and apply final processing