Update modeling_hf_nomic_bert.py

modeling_hf_nomic_bert.py

@@ -6,6 +6,9 @@
 import logging
 
 # Inspired by https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py
+import math
+import numpy as np
+import collections
 import os
 import re
 from collections import OrderedDict
@@ -17,7 +20,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange, repeat
 from safetensors.torch import load_file as safe_load_file
-from transformers import GPT2Config, PreTrainedModel
+from transformers import GPT2Config, PreTrainedModel, ViTModel, ViTConfig
 from transformers.models.bert.modeling_bert import (
     BaseModelOutputWithPoolingAndCrossAttentions,
     MaskedLMOutput,
@@ -25,6 +28,8 @@ from transformers.models.bert.modeling_bert import (
 )
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME
 from transformers.utils.hub import cached_file, get_checkpoint_shard_files
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from torch.nn.modules.utils import _pair
 
 from .configuration_hf_nomic_bert import NomicBertConfig
 
@@ -268,6 +273,68 @@ def remap_bert_state_dict(
 
     return state_dict
 
+    
+def _trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        print("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    # Values are generated by using a truncated uniform distribution and
+    # then using the inverse CDF for the normal distribution.
+    # Get upper and lower cdf values
+    l = norm_cdf((a - mean) / std)
+    u = norm_cdf((b - mean) / std)
+
+    # Uniformly fill tensor with values from [l, u], then translate to
+    # [2l-1, 2u-1].
+    tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+    # Use inverse cdf transform for normal distribution to get truncated
+    # standard normal
+    tensor.erfinv_()
+
+    # Transform to proper mean, std
+    tensor.mul_(std * math.sqrt(2.))
+    tensor.add_(mean)
+
+    # Clamp to ensure it's in the proper range
+    tensor.clamp_(min=a, max=b)
+    return tensor
+
+def trunc_normal_tf_(tensor, mean=0., std=1., a=-2., b=2.):
+    r"""Fills the input Tensor with values drawn from a truncated
+    normal distribution. The values are effectively drawn from the
+    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+
+    NOTE: this 'tf' variant behaves closer to Tensorflow / JAX impl where the
+    bounds [a, b] are applied when sampling the normal distribution with mean=0, std=1.0
+    and the result is subsquently scaled and shifted by the mean and std args.
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.trunc_normal_(w)
+    """
+    with torch.no_grad():
+        _trunc_normal_(tensor, 0, 1.0, a, b)
+        tensor.mul_(std).add_(mean)
+    return tensor
+
 
 class NomicBertPreTrainedModel(PreTrainedModel):
     """An abstract class to handle weights initialization and
@@ -382,6 +449,487 @@ def _init_weights(module, initializer_range=0.02):
         if module.padding_idx is not None:
             nn.init.zeros_(module.weight[module.padding_idx])
 
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+            return tuple(x)
+        return tuple(repeat(x, n))
+    return parse
+
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+to_ntuple = _ntuple
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, add_cls_token=False):
+    """
+    Create 2D sin/cos positional embeddings.
+
+    Args:
+        embed_dim (`int`):
+            Embedding dimension.
+        grid_size (`int`):
+            The grid height and width.
+        add_cls_token (`bool`, *optional*, defaults to `False`):
+            Whether or not to add a classification (CLS) token.
+
+    Returns:
+        (`torch.FloatTensor` of shape (grid_size*grid_size, embed_dim) or (1+grid_size*grid_size, embed_dim): the
+        position embeddings (with or without classification token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if add_cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    if embed_dim % 2 != 0:
+        raise ValueError("embed_dim must be even")
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)
+    """
+    if embed_dim % 2 != 0:
+        raise ValueError("embed_dim must be even")
+
+    omega = np.arange(embed_dim // 2, dtype=float)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+def ndgrid(*tensors) -> Tuple[torch.Tensor, ...]:
+    """generate N-D grid in dimension order.
+
+    The ndgrid function is like meshgrid except that the order of the first two input arguments are switched.
+
+    That is, the statement
+    [X1,X2,X3] = ndgrid(x1,x2,x3)
+
+    produces the same result as
+
+    [X2,X1,X3] = meshgrid(x2,x1,x3)
+
+    This naming is based on MATLAB, the purpose is to avoid confusion due to torch's change to make
+    torch.meshgrid behaviour move from matching ndgrid ('ij') indexing to numpy meshgrid defaults of ('xy').
+
+    """
+    try:
+        return torch.meshgrid(*tensors, indexing='ij')
+    except TypeError:
+        # old PyTorch < 1.10 will follow this path as it does not have indexing arg,
+        # the old behaviour of meshgrid was 'ij'
+        return torch.meshgrid(*tensors)
+
+def build_fourier_pos_embed(
+        feat_shape: List[int],
+        bands: Optional[torch.Tensor] = None,
+        num_bands: int = 64,
+        max_res: int = 224,
+        temperature: float = 10000.,
+        linear_bands: bool = False,
+        include_grid: bool = False,
+        in_pixels: bool = True,
+        ref_feat_shape: Optional[List[int]] = None,
+        dtype: torch.dtype = torch.float32,
+        device: Optional[torch.device] = None,
+) -> List[torch.Tensor]:
+    """
+
+    Args:
+        feat_shape: Feature shape for embedding.
+        bands: Pre-calculated frequency bands.
+        num_bands: Number of frequency bands (determines output dim).
+        max_res: Maximum resolution for pixel based freq.
+        temperature: Temperature for non-pixel freq.
+        linear_bands: Linear band spacing for pixel based freq.
+        include_grid: Include the spatial grid in output.
+        in_pixels: Output in pixel freq.
+        ref_feat_shape: Reference feature shape for resize / fine-tune.
+        dtype: Output dtype.
+        device: Output device.
+
+    Returns:
+
+    """
+    if bands is None:
+        if in_pixels:
+            bands = pixel_freq_bands(
+                num_bands,
+                float(max_res),
+                linear_bands=linear_bands,
+                device=device,
+            )
+        else:
+            bands = freq_bands(
+                num_bands,
+                temperature=temperature,
+                step=1,
+                device=device,
+            )
+    else:
+        if device is None:
+            device = bands.device
+        if dtype is None:
+            dtype = bands.dtype
+
+    if in_pixels:
+        t = [torch.linspace(-1., 1., steps=s, device=device, dtype=torch.float32) for s in feat_shape]
+    else:
+        t = [torch.arange(s, device=device, dtype=torch.int64).to(torch.float32) for s in feat_shape]
+
+    if ref_feat_shape is not None:
+        # eva's scheme for resizing rope embeddings (ref shape = pretrain)
+        t = [x / f * r for x, f, r in zip(t, feat_shape, ref_feat_shape)]
+
+    grid = torch.stack(ndgrid(t), dim=-1)
+    grid = grid.unsqueeze(-1)
+    pos = grid * bands
+
+    pos_sin, pos_cos = pos.sin().to(dtype=dtype), pos.cos().to(dtype)
+    out = [grid, pos_sin, pos_cos] if include_grid else [pos_sin, pos_cos]
+    return out
+
+
+def build_rotary_pos_embed(
+        feat_shape: List[int],
+        bands: Optional[torch.Tensor] = None,
+        dim: int = 64,
+        max_res: int = 224,
+        temperature: float = 10000.,
+        linear_bands: bool = False,
+        in_pixels: bool = True,
+        ref_feat_shape: Optional[List[int]] = None,
+        dtype: torch.dtype = torch.float32,
+        device: Optional[torch.device] = None,
+):
+    """
+
+    Args:
+        feat_shape: Spatial shape of the target tensor for embedding.
+        bands: Optional pre-generated frequency bands
+        dim: Output dimension of embedding tensor.
+        max_res: Maximum resolution for pixel mode.
+        temperature: Temperature (inv freq) for non-pixel mode
+        linear_bands: Linearly (instead of log) spaced bands for pixel mode
+        in_pixels: Pixel vs language (inv freq) mode.
+        dtype: Output dtype.
+        device: Output device.
+
+    Returns:
+
+    """
+    sin_emb, cos_emb = build_fourier_pos_embed(
+        feat_shape,
+        bands=bands,
+        num_bands=dim // 4,
+        max_res=max_res,
+        temperature=temperature,
+        linear_bands=linear_bands,
+        in_pixels=in_pixels,
+        ref_feat_shape=ref_feat_shape,
+        device=device,
+        dtype=dtype,
+    )
+    num_spatial_dim = 1
+    # this would be much nicer as a .numel() call to torch.Size(), but torchscript sucks
+    for x in feat_shape:
+        num_spatial_dim *= x
+    sin_emb = sin_emb.reshape(num_spatial_dim, -1).repeat_interleave(2, -1)
+    cos_emb = cos_emb.reshape(num_spatial_dim, -1).repeat_interleave(2, -1)
+    return sin_emb, cos_emb
+
+def freq_bands(
+        num_bands: int,
+        temperature: float = 10000.,
+        step: int = 2,
+        device: Optional[torch.device] = None,
+) -> torch.Tensor:
+    exp = torch.arange(0, num_bands, step, dtype=torch.int64, device=device).to(torch.float32) / num_bands
+    bands = 1. / (temperature ** exp)
+    return bands
+
+    
+def pixel_freq_bands(
+        num_bands: int,
+        max_freq: float = 224.,
+        linear_bands: bool = True,
+        device: Optional[torch.device] = None,
+):
+    if linear_bands:
+        bands = torch.linspace(1.0, max_freq / 2, num_bands, dtype=torch.float32, device=device)
+    else:
+        bands = 2 ** torch.linspace(0, math.log(max_freq, 2) - 1, num_bands, dtype=torch.float32, device=device)
+    return bands * torch.pi
+
+def rot(x):
+    return torch.stack([-x[..., 1::2], x[..., ::2]], -1).reshape(x.shape)
+
+def apply_rot_embed_cat(x: torch.Tensor, emb):
+    sin_emb, cos_emb = emb.tensor_split(2, -1)
+    if sin_emb.ndim == 3:
+        return x * cos_emb.unsqueeze(1).expand_as(x) + rot(x) * sin_emb.unsqueeze(1).expand_as(x)
+    return x * cos_emb + rot(x) * sin_emb
+
+# taken from https://github.com/huggingface/pytorch-image-models/blob/cb0e4391beedcc5ac3ae4bce16561b95c326f32c/timm/layers/pos_embed_sincos.py#L363
+class NomicVisionRotaryEmbeddingCat(nn.Module):
+    """ Rotary position embedding w/ concatenatd sin & cos
+
+    The following impl/resources were referenced for this impl:
+    * https://github.com/lucidrains/vit-pytorch/blob/6f3a5fcf0bca1c5ec33a35ef48d97213709df4ba/vit_pytorch/rvt.py
+    * https://blog.eleuther.ai/rotary-embeddings/
+    """
+
+    def __init__(
+            self,
+            dim,
+            max_res=224,
+            temperature=10000,
+            in_pixels=True,
+            linear_bands: bool = False,
+            feat_shape: Optional[List[int]] = None,
+            ref_feat_shape: Optional[List[int]] = None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.max_res = max_res
+        self.temperature = temperature
+        self.in_pixels = in_pixels
+        self.feat_shape = feat_shape
+        self.ref_feat_shape = ref_feat_shape
+
+        if feat_shape is None:
+            # only cache bands
+            if in_pixels:
+                bands = pixel_freq_bands(
+                    dim // 4,
+                    float(max_res),
+                    linear_bands=linear_bands,
+                )
+            else:
+                bands = freq_bands(
+                    dim // 4,
+                    temperature=temperature,
+                    step=1,
+                )
+            self.register_buffer(
+                'bands',
+                bands,
+                persistent=False,
+            )
+            self.pos_embed = None
+        else:
+            # cache full sin/cos embeddings if shape provided up front
+            embeds = build_rotary_pos_embed(
+                feat_shape=feat_shape,
+                dim=dim,
+                max_res=max_res,
+                linear_bands=linear_bands,
+                in_pixels=in_pixels,
+                ref_feat_shape=self.ref_feat_shape,
+            )
+            self.bands = None
+            self.register_buffer(
+                'pos_embed',
+                torch.cat(embeds, -1),
+                persistent=False,
+            )
+
+    def get_embed(self, shape: Optional[List[int]] = None):
+        if self.bands is not None and shape is not None:
+            # rebuild embeddings every call, use if target shape changes
+            embeds = build_rotary_pos_embed(
+                shape,
+                self.bands,
+                in_pixels=self.in_pixels,
+                ref_feat_shape=self.ref_feat_shape,
+            )
+            return torch.cat(embeds, -1)
+        elif self.pos_embed is not None:
+            return self.pos_embed
+        else:
+            assert False, "get_embed() requires pre-computed pos_embed or valid shape w/ pre-computed bands"
+
+    def forward(self, x):
+        # assuming channel-first tensor where spatial dim are >= 2
+        pos_embed = self.get_embed(x.shape[2:])
+        return apply_rot_embed_cat(x, pos_embed)
+
+class NomicVisionPatchEmbeddings(nn.Module):
+    def __init__(
+        self,
+        config,
+    ):
+        super().__init__()
+        img_size = _pair(config.img_size)
+        patch_size = _pair(config.patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+
+        self.proj = nn.Linear(
+            config.num_channels * patch_size[0] * patch_size[1], config.n_embd, bias=config.patch_embed_bias
+        )
+
+        self.learned_pos_embedding = False
+        self.sinusoidal_pos_embedding = False
+        self.no_embed_class = getattr(config, "no_embed_class", False)
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.n_embd)) if not getattr(config, "no_cls_token", False) else None
+        if config.learned_pos_embedding:
+            # this is the default in DINO
+            self.learned_pos_embedding = True
+            # hack for timm dinov2 with registers
+            num_patches = self.num_patches if getattr(config, "register_tokens", 0) > 0 else self.num_patches + 1
+            self.pos_embed = nn.Parameter(torch.randn(1, num_patches, config.n_embd) * 0.02) if getattr(config, "use_pos_embed", True) else None
+        elif getattr(config, "sinusoidal_pos_embedding", False):
+            self.sinusoidal_pos_embedding = True
+            if getattr(config, "use_pos_embed", True):
+                self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, config.n_embd), requires_grad=False)
+                pos_embed = get_2d_sincos_pos_embed(config.n_embd, self.grid_size[0], add_cls_token=True)
+                self.pos_embed.data.copy_(torch.from_numpy(pos_embed).to(self.pos_embed))
+            else:
+                self.pos_embed = None
+        else:
+            self.pos_embed = nn.Parameter(torch.randn(1, self.num_patches + 1, config.n_embd) * 0.02) if getattr(config, "use_pos_embed", True) else None
+
+        if getattr(config, "register_tokens", 0) > 0:
+            self.reg_token = nn.Parameter(torch.randn(1, config.register_tokens, config.n_embd) * 0.02)
+        else:
+            self.reg_token = None 
+
+        if config.mask_token:
+            self.mask_token = nn.Parameter(torch.zeros(1, config.n_embd))
+
+        self.patch_dropout = nn.Identity()
+
+        if getattr(config, "use_rotary_pos_emb", False):
+            ref_feat_shape = getattr(config, "ref_feat_shape", None)
+            ref_feat_shape = to_2tuple(ref_feat_shape) if ref_feat_shape is not None else None
+            self.rope = NomicVisionRotaryEmbeddingCat(
+                config.n_embd // config.n_head,
+                in_pixels=False,
+                feat_shape=self.grid_size,
+                ref_feat_shape=ref_feat_shape,
+            )
+        else:
+            self.rope = None
+
+
+    def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
+        """
+        This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
+        resolution images.
+
+        Source:
+        https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174
+        """
+        num_patches = embeddings.shape[1] - 1
+        num_positions = self.pos_embed.shape[1] - 1
+        if num_patches == num_positions and height == width:
+            return self.pos_embed
+        class_pos_embed = self.pos_embed[:, 0]
+        patch_pos_embed = self.pos_embed[:, 1:]
+        dim = embeddings.shape[-1]
+        height = height // self.patch_size[0]
+        width = width // self.patch_size[1]
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        height, width = height + 0.1, width + 0.1
+        patch_pos_embed = patch_pos_embed.reshape(1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim)
+        patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed,
+            scale_factor=(height / math.sqrt(num_positions), width / math.sqrt(num_positions)),
+            mode="bicubic",
+            align_corners=False,
+        )
+        if int(height) != patch_pos_embed.shape[-2] or int(width) != patch_pos_embed.shape[-1]:
+            raise ValueError("Width or height does not match with the interpolated position embeddings")
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+    def forward(self, x):
+        # deepspeed case where the input is in fp32
+        if x.dtype != self.proj.weight.dtype:
+            x = x.to(dtype=self.proj.weight.dtype)
+
+        _, _, height, width = x.shape
+        x = self.proj(
+            rearrange(
+                x,
+                "b c (h p1) (w p2) -> b h w (c p1 p2)",
+                p1=self.patch_size[0],
+                p2=self.patch_size[1],
+            )
+        )
+        embeddings = rearrange(x, "b h w c -> b (h w) c")
+
+        to_cat = []
+        if self.cls_token is not None:
+            if self.sinusoidal_pos_embedding:
+                cls_token = self.cls_token + self.pos_embed[:, 0]
+                cls_token = cls_token.expand(embeddings.shape[0], -1, -1)
+                to_cat += [cls_token]
+            else:
+                cls_token = self.cls_token.expand(embeddings.shape[0], 1, -1)
+                to_cat += [cls_token]
+
+        if self.reg_token is not None:
+            to_cat += [self.reg_token.expand(embeddings.shape[0], -1, -1)]
+
+        rot_pos_embed = self.rope.get_embed() if self.rope is not None else None
+
+        if self.no_embed_class:
+            if self.learned_pos_embedding:
+                embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
+            else:
+                if self.pos_embed is not None:
+                    embeddings = embeddings + self.pos_embed 
+            if to_cat:
+                embeddings = torch.cat(to_cat + [embeddings], dim=1)
+        else:
+            if to_cat:
+                embeddings = torch.cat(to_cat + [embeddings], dim=1)
+            if self.learned_pos_embedding:
+                if self.pos_embed is not None:
+                    embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
+            else:
+                if self.pos_embed is not None:
+                    embeddings = embeddings + self.pos_embed 
+
+        embeddings = self.patch_dropout(embeddings)
+
+        return embeddings, rot_pos_embed
+
 
 class NomicBertEmbeddings(nn.Module):
     def __init__(self, config):
@@ -466,17 +1014,19 @@ class NomciBertGatedMLP(nn.Module):
         fused_bias_fc=True,
         device=None,
         dtype=None,
+        norm_layer=False,
     ):
         super().__init__()
         out_features = out_features if out_features is not None else in_features
         hidden_features = hidden_features if hidden_features is not None else int(8 * in_features / 3)
-        hidden_features = (hidden_features + multiple_of - 1) // multiple_of * multiple_of
+        hidden_features = int((hidden_features + multiple_of - 1) // multiple_of * multiple_of)
         self.return_residual = return_residual
 
         self.fc11 = nn.Linear(in_features, hidden_features, bias=bias1)
         self.fc12 = nn.Linear(in_features, hidden_features, bias=bias1)
         self.activation = activation
         self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2)
+        self.norm = nn.LayerNorm(hidden_features) if norm_layer else nn.Identity()
 
     def forward(self, x):
         y = self.fc11(x)
@@ -485,6 +1035,10 @@ class NomciBertGatedMLP(nn.Module):
             y = F.glu(torch.cat([y, gate], dim=-1), dim=-1)
         else:
             y = y * self.activation(gate)
+
+        # eva uses layer norm after the activation
+        y = self.norm(y)
+
         y = self.fc2(y)
         return y if not self.return_residual else (y, x)
 
@@ -758,6 +1312,7 @@ class NomicBertAttention(nn.Module):
         self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.qkv_proj_bias)
         self.causal = config.causal
         self.drop = nn.Dropout(config.attn_pdrop)
+        self.num_prefix_tokens = max(getattr(config, "register_tokens", 1), 1)
 
     def forward(
         self,
@@ -770,6 +1325,7 @@ class NomicBertAttention(nn.Module):
         is_padded_inputs: Optional[bool] = True,
         cu_seqlens: Optional[torch.Tensor] = None,
         max_seq_len: Optional[int] = None,
+        rope: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
 
         has_layer_past = past_key_value is not None
@@ -792,6 +1348,13 @@ class NomicBertAttention(nn.Module):
 
             if self.rotary_head_dim:
                 qkv = rearrange(qkv, "b h three s d -> b s three h d")
+        elif rope is not None:
+            q, k, v = qkv.permute(0, 3, 1, 2, 4).unbind(dim=-2)
+            q = torch.cat([q[:, :, :self.num_prefix_tokens], apply_rot_embed_cat(q[:, :, self.num_prefix_tokens:], rope)], dim=2).type_as(q)
+            k = torch.cat([k[:, :, :self.num_prefix_tokens], apply_rot_embed_cat(k[:, :, self.num_prefix_tokens:], rope)], dim=2).type_as(q)
+
+            qkv = torch.stack([q, k, v], dim=-2)
+            qkv = rearrange(qkv, "b h s three d -> b s three h d")
 
         query, key, value = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
 
@@ -837,6 +1400,7 @@ class NomicBertBlock(NomicBertPreTrainedModel):
                 bias2=config.mlp_fc2_bias,
                 activation=activation,
                 fused_bias_fc=config.fused_bias_fc,
+                norm_layer=getattr(config, "norm_mlp", False),
             )
         else:
             self.mlp = NomicBertMLP(
@@ -866,6 +1430,7 @@ class NomicBertBlock(NomicBertPreTrainedModel):
         use_cache: Optional[bool] = False,
         cu_seqlens: Optional[torch.Tensor] = None,
         max_seq_len: Optional[int] = None,
+        rope: Optional[torch.Tensor] = None,
     ):
         r"""Pass the input through the encoder layer.
 
@@ -886,6 +1451,7 @@ class NomicBertBlock(NomicBertPreTrainedModel):
                 is_padded_inputs=is_padded_inputs,
                 cu_seqlens=cu_seqlens,
                 max_seq_len=max_seq_len,
+                rope=rope,
             )
 
             dropped = self.dropout2(hidden_states)
@@ -902,6 +1468,7 @@ class NomicBertBlock(NomicBertPreTrainedModel):
                 is_padded_inputs=is_padded_inputs,
                 cu_seqlens=cu_seqlens,
                 max_seq_len=max_seq_len,
+                rope=rope,
             )
             hidden_states = self.norm1((self.dropout1(attn_outputs) + hidden_states).to(dtype=self.norm1.weight.dtype))
             mlp_out = self.mlp(hidden_states)
@@ -929,6 +1496,7 @@ class NomicBertEncoder(nn.Module):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         is_padded_inputs: Optional[bool] = True,
+        rope: Optional[torch.Tensor] = None,
     ):
         """If subset_mask is not None, we only want output for the subset of the sequence.
         This means that we only compute the last layer output for these tokens.
@@ -953,9 +1521,14 @@ class NomicBertEncoder(nn.Module):
                     hidden_states2,
                     residual,
                     attention_mask,
+                    position_ids,
+                    past_key_values,
+                    is_padded_inputs,
+                    output_attentions,
+                    use_cache,
                     None,
                     None,
-                    is_padded_inputs,
+                    rope,
                     # if you freeze ANY layers, you need `use_reentrant=False`
                     # https://github.com/huggingface/transformers/issues/21381
                     # https://discuss.pytorch.org/t/checkpoint-with-no-grad-requiring-inputs-problem/19117/7
@@ -973,6 +1546,7 @@ class NomicBertEncoder(nn.Module):
                     is_padded_inputs,
                     output_attentions,
                     use_cache,
+                    rope=rope,
                 )
         return hidden_states
 
@@ -1232,3 +1806,266 @@ class NomicBertForSequenceClassification(NomicBertPreTrainedModel):
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
         )
+
+def hf_vit_config_to_vit_config(vit_config: ViTConfig) -> GPT2Config:
+    return GPT2Config(
+        n_embd=vit_config.hidden_size,
+        n_layer=vit_config.num_hidden_layers,
+        n_head=vit_config.num_attention_heads,
+        n_inner=vit_config.intermediate_size,
+        activation_function=vit_config.hidden_act,
+        vocab_size=0, # no vocab since using patches
+        n_positions=0,  # No absolute position embedding
+        resid_pdrop=0.0,  # No dropout
+        embd_pdrop=getattr(vit_config, "dropout", 0.0),
+        attn_pdrop=vit_config.attention_probs_dropout_prob,
+        layer_norm_epsilon=vit_config.layer_norm_eps,
+        initializer_range=vit_config.initializer_range,
+        bos_token_id=None,
+        eos_token_id=None,
+        # These are new arguments not in the original GPT2Config
+        drop_path_rate=0.0,
+        # Why is there double layer norm??
+        prepre_layernom=False,
+        layer_scale=False,
+        layer_scale_init=None,
+        img_size=vit_config.image_size,
+        patch_size=vit_config.patch_size,
+        num_channels=vit_config.num_channels,
+        prenorm=True,
+        parallel_block=False,
+        parallel_block_tied_norm=False,
+        rotary_emb_fraction=0,
+        tie_word_embeddings=False,
+        fused_dropout_add_ln=True,
+        fused_bias_fc=True,
+        patch_embed_bias=True,
+        use_flash_attn=True,
+        qkv_proj_bias=True,
+        mlp_fc1_bias=getattr(vit_config, "mlp_fc1_bias", True),
+        mlp_fc2_bias=getattr(vit_config, "mlp_fc2_bias", True),
+        use_rms_norm=False,
+        causal=False,
+        hidden_features_scaling_factor=1.0,
+        mask_token=False,
+        learned_pos_embedding=False,
+        patch_dropout=0,
+        sinusoidal_pos_embedding=vit_config.model_type == "vit_mae"
+    )
+
+    
+class NomicAttentionPooling(nn.Module):
+    def __init__(
+        self,
+        config
+    ):
+        super().__init__()
+        self.embed_dim = config.n_embd
+        self.use_flash_attn = config.use_flash_attn
+        self.fused_bias_fc = config.fused_bias_fc
+
+        self.num_heads = config.n_head
+        self.num_heads_kv = config.num_heads_kv if getattr(config, "num_heads_kv", None) is not None else self.num_heads
+        assert self.embed_dim % self.num_heads == 0, "embed_dim must be divisible by num_heads"
+        self.head_dim = self.embed_dim // self.num_heads
+        # we don't really support mqa / gqa for now
+        kv_dim = 2 * self.head_dim * self.num_heads_kv
+
+        self.register_buffer(
+            "norm_factor",
+            torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32)).to(torch.get_default_dtype()),
+            persistent=False,
+        )
+
+        self.Wq = nn.Linear(self.embed_dim, self.embed_dim, bias=config.qkv_proj_bias)
+        self.Wkv = nn.Linear(self.embed_dim, kv_dim, bias=config.qkv_proj_bias)
+        
+        self.latent = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.qkv_proj_bias)
+        self.causal = config.causal
+        self.drop = nn.Dropout(config.attn_pdrop)
+
+    def init_weights(self):
+        trunc_normal_tf_(self.latent, std=self.embed_dim ** -0.5)
+
+    def forward(
+        self,
+        kv,
+        attention_mask=None,
+        cu_seqlens_k=None,
+        max_seqlen_k=None,
+        is_padded_inputs: Optional[bool] = True,
+        output_attentions: bool = False,
+    ):
+        """Implements the multihead softmax attention.
+        Arguments
+        ---------
+            q: The tensor containing the query. (B, Sq, H, D)
+            kv: The tensor containing the key and value. (B, Sk, 2, H_k, D)
+            causal: if passed, will override self.causal
+            cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+                of the sequences in the batch, used to index into q.
+            max_seqlen: int. Maximum sequence length in the batch of q.
+            cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
+                of the sequences in the batch, used to index into kv.
+            max_seqlen_k: int. Maximum sequence length in the batch of k and v.
+        """
+        q_latent = self.latent.expand(kv.size(0), -1, -1)
+        q = self.Wq(q_latent)
+        bsz, q_len, h_size = q.shape
+        kv = self.Wkv(kv)
+        query = rearrange(q, "... (h d) -> ... h d", d=self.head_dim)
+        kv = rearrange(kv, "... (two hkv d) -> ... two hkv d", two=2, d=self.head_dim)
+
+        key, value = kv[:, :, 0], kv[:, :, 1]
+
+        query = query.permute(0, 2, 1, 3)
+        key = key.permute(0, 2, 1, 3)
+        value = value.permute(0, 2, 1, 3)
+
+        attention_scores = torch.matmul(query, key.transpose(-1, -2)) / self.norm_factor
+        if attention_mask is not None:
+            attention_scores = attention_scores + attention_mask
+
+        attentions_probs = F.softmax(attention_scores, dim=-1)
+        attentions_probs = self.drop(attentions_probs)
+
+        attn_output = torch.matmul(attentions_probs, value)
+        attn_output = rearrange(attn_output.permute(0, 2, 1, 3), "... h d -> ... (h d)")
+
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output
+
+        
+class NomicMultiHeadAttentionPooling(nn.Module):
+    def __init__(
+        self,
+        config,
+    ):
+        super().__init__()
+        self.prenorm = config.prenorm
+        self.fused_dropout_add_ln = config.fused_dropout_add_ln
+
+        self.attn = NomicAttentionPooling(config)
+        activation = (
+            F.sigmoid
+            if config.activation_function == "glu"
+            else (F.silu if config.activation_function == "swiglu" else F.gelu)
+        )
+        if config.activation_function in ["glu", "swiglu", "geglu"]:
+            self.mlp = NomciBertGatedMLP(
+                config.n_embd,
+                hidden_features=config.n_inner,
+                bias1=config.mlp_fc1_bias,
+                bias2=config.mlp_fc2_bias,
+                activation=activation,
+                fused_bias_fc=config.fused_bias_fc,
+            )
+        else:
+            self.mlp = NomicBertMLP(
+                config.n_embd,
+                hidden_features=config.n_inner,
+                bias1=config.mlp_fc1_bias,
+                bias2=config.mlp_fc2_bias,
+                activation=activation,
+                fused_bias_fc=config.fused_bias_fc,
+            )
+
+        self.dropout1 = nn.Dropout(config.resid_pdrop)
+        self.norm1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.dropout2 = nn.Dropout(config.resid_pdrop)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ):
+        r"""Pass the input through the encoder layer.
+
+        Args:
+            hidden_states: the sequence to the encoder layer (required).
+            residual: if postnorm, residual=None, If prenorm, hidden_states = Attn/MLP(LN(residual))
+            mixer_subset: for cross-attention only. If not None, will take a subset of x
+                before applying the query projection. Useful for e.g., ViT where we only care
+                about the CLS token in the last layer.
+        """
+
+        attn_outputs = self.attn(
+                hidden_states,
+                attention_mask=attention_mask,
+            )
+
+        normed = self.norm1(attn_outputs)
+        hidden_states = hidden_states + self.mlp(normed)
+
+        return hidden_states
+
+class NomicVisionPreTrainedModel(PreTrainedModel):
+    """An abstract class to handle weights initialization and
+    a simple interface for dowloading and loading pretrained models.
+    """
+
+    config_class = NomicBertConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Block"]
+    _skip_keys_device_placement = "past_key_values"
+
+    def __init__(self, config, *inputs, **kwargs):
+        super().__init__(config)
+        if not isinstance(config, GPT2Config):
+            raise ValueError(
+                "Parameter config in `{}(config)` should be an instance of class `GPT2Config`. "
+                "To create a model from a Google pretrained model use "
+                "`model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`".format(
+                    self.__class__.__name__, self.__class__.__name__
+                )
+            )
+        self.config = config
+
+class NomicVisionModel(NomicVisionPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.embeddings = NomicVisionPatchEmbeddings(config)
+        self.layers = nn.ModuleList([NomicBertBlock(config) for _ in range(config.n_layer)])
+
+        self.selector = NomicMultiHeadAttentionPooling(config)
+
+        self.global_pool = getattr(config, "global_pool", None)
+        self.num_prefix_tokens = (1 if not getattr(config, "no_cls_token", False) else 0) + getattr(config, "register_tokens", 0)
+
+        self.apply(partial(_init_weights, initializer_range=config.initializer_range))
+
+    def forward(
+        self,
+        pixel_values,
+        attention_mask=None,
+        position_ids=None,
+        token_type_ids=None,
+        return_dict=None,
+        matryoshka_dim=None,
+    ):
+        embeddings, rope = self.embeddings(pixel_values)
+
+        original_dtype = embeddings.dtype
+
+        hidden_states = embeddings 
+        # unused but easier to pass to gradient checkpointing as words
+        residual = None
+        for layer in self.layers:
+            # need to pass none for backwards compatability
+            hidden_states, _, residual = layer(hidden_states, None, residual=residual, is_padded_inputs=False, rope=rope)
+
+        hidden_states = hidden_states + residual
+        if self.global_pool == "avg":
+            hidden_states = hidden_states[:, self.num_prefix_tokens:].mean(dim=1)
+
+        pooled_output = self.selector(hidden_states)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=pooled_output,
+            hidden_states=hidden_states,
+        )