Upload modeling_molmo.py with huggingface_hub

modeling_molmo.py

@@ -762,7 +762,6 @@ class ViTMLP(nn.Module):
         return x
 
 
-
 class ResidualAttentionBlock(nn.Module):
 
     def __init__(self, config: FullMolmoConfig):
@@ -819,6 +818,14 @@ class BlockCollection(nn.Module):
         return hidden_states
 
 
+class LayerNormFp32(nn.LayerNorm):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        orig_type = x.dtype
+        x = F.layer_norm(x.to(torch.float32), self.normalized_shape, self.weight.to(torch.float32),
+                         self.bias.to(torch.float32), self.eps)
+        return x.to(orig_type)
+
+
 class VisionTransformer(nn.Module):
 
     def __init__(self, config: FullMolmoConfig):
@@ -844,7 +851,7 @@ class VisionTransformer(nn.Module):
             device=config.init_device,
             )
 
-        self.pre_ln = nn.LayerNorm(
+        self.pre_ln = LayerNormFp32(
             v_cfg.image_emb_dim,
             eps=v_cfg.image_norm_eps,
         )