Update modeling_hf_nomic_bert.py

modeling_hf_nomic_bert.py

@@ -119,7 +119,7 @@ def filter_shapes(state_dict, model):
     return filtered_state_dict
 
     
-def remap_bert_state_dict(state_dict, config,  remove_bert=False, remove_cls_weights=False, add_pooling_layer=False):
+def remap_bert_state_dict(state_dict, config, remove_bert=False, remove_cls_weights=False, add_pooling_layer=False):
     """
     Map the state_dict of a Huggingface BERT model to be flash_attn compatible.
     """
@@ -225,6 +225,16 @@ def remap_bert_state_dict(state_dict, config,  remove_bert=False, remove_cls_wei
 
     state_dict = OrderedDict((key_mapping_decoder_bias(k), v) for k, v in state_dict.items())
 
+    if remove_cls_weights:
+        cls_weights = ["cls.predictions.decoder.bias",
+                       "cls.predictions.transform.dense.weight",
+                       "cls.predictions.transform.dense.bias", 
+                       "cls.predictions.transform.layer_norm.weight", 
+                       "cls.predictions.transform.layer_norm.bias", 
+                       "cls.predictions.decoder.weight"]
+        for weight in cls_weights:
+            state_dict.pop(weight, None)
+
     # Word embedding
     pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
     if pad_vocab_size_multiple > 1:
@@ -232,18 +242,19 @@ def remap_bert_state_dict(state_dict, config,  remove_bert=False, remove_cls_wei
         state_dict["bert.embeddings.word_embeddings.weight"] = F.pad(
             word_embeddings, (0, 0, 0, config.vocab_size - word_embeddings.shape[0])
         )
-        decoder_weight = state_dict["cls.predictions.decoder.weight"]
-        state_dict["cls.predictions.decoder.weight"] = F.pad(
-            decoder_weight, (0, 0, 0, config.vocab_size - decoder_weight.shape[0])
-        )
-        # If the vocab was padded, we want to set the decoder bias for those padded indices to be
-        # strongly negative (i.e. the decoder shouldn't predict those indices).
-        # TD [2022-05-09]: I don't think it affects the MLPerf training.
-        if "cls.predictions.decoder.bias" in state_dict:
-            decoder_bias = state_dict["cls.predictions.decoder.bias"]
-            state_dict["cls.predictions.decoder.bias"] = F.pad(
-                decoder_bias, (0, config.vocab_size - decoder_bias.shape[0]), value=-100.0
+        if not remove_cls_weights:
+            decoder_weight = state_dict["cls.predictions.decoder.weight"]
+            state_dict["cls.predictions.decoder.weight"] = F.pad(
+                decoder_weight, (0, 0, 0, config.vocab_size - decoder_weight.shape[0])
             )
+            # If the vocab was padded, we want to set the decoder bias for those padded indices to be
+            # strongly negative (i.e. the decoder shouldn't predict those indices).
+            # TD [2022-05-09]: I don't think it affects the MLPerf training.
+            if "cls.predictions.decoder.bias" in state_dict:
+                decoder_bias = state_dict["cls.predictions.decoder.bias"]
+                state_dict["cls.predictions.decoder.bias"] = F.pad(
+                    decoder_bias, (0, config.vocab_size - decoder_bias.shape[0]), value=-100.0
+                )
 
     if add_pooling_layer is False:
         pooler_weights = ["bert.pooler.dense.weight",
@@ -252,16 +263,6 @@ def remap_bert_state_dict(state_dict, config,  remove_bert=False, remove_cls_wei
         for key in pooler_weights:
             state_dict.pop(key, None)
 
-    if remove_cls_weights:
-        cls_weights = ["cls.predictions.decoder.bias",
-                       "cls.predictions.transform.dense.weight",
-                       "cls.predictions.transform.dense.bias", 
-                       "cls.predictions.transform.layer_norm.weight", 
-                       "cls.predictions.transform.layer_norm.bias", 
-                       "cls.predictions.decoder.weight"]
-        for weight in cls_weights:
-            state_dict.pop(weight, None)
-
     if remove_bert:
         def remove_bert_prefix(key):
             key = re.sub(r"^bert.", "", key)
@@ -319,9 +320,21 @@ class NomicBertPreTrainedModel(PreTrainedModel):
         remove_bert_prefix = cls != NomicBertForPreTraining
         ignore_mismatched_shapes = kwargs.pop("ignore_mismatched_sizes", False)
         num_labels = kwargs.pop("num_labels", None)
+        rotary_scaling_factor = kwargs.pop("rotary_scaling_factor", None)
+        if rotary_scaling_factor:
+            config.rotary_scaling_factor = rotary_scaling_factor
+        if config.n_positions <= 0 and config.rotary_emb_fraction > 0:
+            config.n_positions = 2048
         if num_labels:
             config.num_labels = num_labels
-        model = cls(config, *inputs)
+
+        if "add_pooling_layer" in kwargs:
+            model = cls(config, *inputs, add_pooling_layer=kwargs.pop("add_pooling_layer"))
+        else:
+            if cls == NomicBertModel:
+                model = cls(config, *inputs, add_pooling_layer=False)
+            else:
+                model = cls(config, *inputs)
         # TODO: fix this
         # Assuming we know what we're doing when loading from disk
         # Prob a bad assumption but i'm tired and want to train this asap
@@ -379,9 +392,9 @@ class NomicBertEmbeddings(nn.Module):
         self.word_embeddings = nn.Embedding(
             config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id
         )
-        self.max_position_embeddings = config.max_position_embeddings
+        self.max_position_embeddings = config.max_position_embeddings if config.rotary_emb_fraction <= 0 else 0
         self.type_vocab_size = config.type_vocab_size
-        if self.max_position_embeddings > 0:
+        if self.max_position_embeddings > 0 and config.rotary_emb_fraction <= 0:
             self.position_embeddings = nn.Embedding(
                 config.max_position_embeddings, config.hidden_size,
             )
@@ -542,6 +555,12 @@ class NomicBertRotaryEmbedding(nn.Module):
         self.register_buffer("inv_freq", inv_freq, persistent=False)
         self.interleaved = interleaved
         self.scale_base = scale_base
+        scale = (
+            (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)
+            if scale_base is not None
+            else None
+        )
+        self.register_buffer("scale", scale, persistent=False)
 
         self._seq_len_cached = 0
         self._cos_cached = None
@@ -607,7 +626,9 @@ class NomicBertRotaryEmbedding(nn.Module):
         Apply rotary embedding *inplace* to qkv and / or kv.
         """
         seqlen = qkv.shape[1]
-        if max_seqlen is not None:
+        if seqlen > self._seq_len_cached:
+            self._update_cos_sin_cache(seqlen, device=qkv.device, dtype=qkv.dtype)
+        elif max_seqlen is not None:
             self._update_cos_sin_cache(max_seqlen, device=qkv.device, dtype=qkv.dtype)
         elif isinstance(seqlen_offset, int):
             self._update_cos_sin_cache(seqlen + seqlen_offset, device=qkv.device, dtype=qkv.dtype)
@@ -617,6 +638,79 @@ class NomicBertRotaryEmbedding(nn.Module):
         return torch.stack((q_rot, k_rot, qkv[:, :, 2]), dim=2)
 
 
+class NomicBertDynamicNTKRotaryEmbedding(NomicBertRotaryEmbedding):
+    def __init__(self, rotary_scaling_factor, max_position_embeddings, **kwargs):
+        super().__init__(**kwargs)
+        self.rotary_scaling_factor = rotary_scaling_factor
+        self.max_position_embeddings = max_position_embeddings
+
+        
+    def _compute_inv_freq(self, base=None, device=None):
+        if base is None:
+            base = self.base
+        return 1.0 / (
+            base 
+            ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim)
+        )
+
+    def _update_cos_sin_cache(self, seqlen, device=None, dtype=None):
+        # Reset the tables if the sequence length has changed,
+        # if we're on a new device (possibly due to tracing for instance),
+        # or if we're switching from inference mode to training
+        if seqlen > self.max_position_embeddings:
+            base = self.base * (
+                (self.rotary_scaling_factor * seqlen / self.max_position_embeddings) - (self.rotary_scaling_factor - 1)
+            ) ** (self.dim / (self.dim - 2))
+            inv_freq = self._compute_inv_freq(base=base, device=device)
+            self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        if (
+            seqlen > self._seq_len_cached
+            or self._cos_cached is None
+            or self._cos_cached.device != device
+            or self._cos_cached.dtype != dtype
+            or (self.training and self._cos_cached.is_inference())
+        ):
+            self._seq_len_cached = seqlen
+            # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16
+            # And the output of arange can be quite large, so bf16 would lose a lot of precision.
+            # However, for compatibility reason, we add an option to use the dtype of self.inv_freq.
+            if self.pos_idx_in_fp32:
+                t = torch.arange(seqlen, device=device, dtype=torch.float32)
+                # We want fp32 here as well since inv_freq will be multiplied with t, and the output
+                # will be large. Having it in bf16 will lose a lot of precision and cause the
+                # cos & sin output to change significantly.
+                # We want to recompute self.inv_freq if it was not loaded in fp32
+                if self.inv_freq.dtype != torch.float32:
+                    if seqlen > self.max_position_embeddings:
+                        base = self.base * (
+                            (self.scaling_factor * seqlen / self.max_position_embeddings) - (self.scaling_factor - 1)
+                        ) ** (self.dim / (self.dim - 2))
+                    else:
+                        base = self.base
+                    inv_freq = self._compute_inv_freq(device=device, base=base)
+                else:
+                    inv_freq = self.inv_freq
+            else:
+                t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
+                inv_freq = self.inv_freq
+            # Don't do einsum, it converts fp32 to fp16 under AMP
+            # freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            freqs = torch.outer(t, inv_freq)
+            if self.scale is None:
+                self._cos_cached = torch.cos(freqs).to(dtype)
+                self._sin_cached = torch.sin(freqs).to(dtype)
+            else:
+                power = (
+                    torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device)
+                    - seqlen // 2
+                ) / self.scale_base
+                scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1")
+                # We want the multiplication by scale to happen in fp32
+                self._cos_cached = (torch.cos(freqs) * scale).to(dtype)
+                self._sin_cached = (torch.sin(freqs) * scale).to(dtype)
+                self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
+                self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)
 
 class NomicBertAttention(nn.Module):
     """Multi-head self-attention and cross-attention"""
@@ -651,12 +745,22 @@ class NomicBertAttention(nn.Module):
 
         self.rotary_emb_dim = self.head_dim * config.rotary_emb_fraction
         if self.rotary_emb_dim > 0:
-            self.rotary_emb = NomicBertRotaryEmbedding(
-                self.rotary_emb_dim,
-                base=config.rotary_emb_base,
-                scale_base=config.rotary_emb_scale_base,
-                interleaved=config.rotary_emb_interleaved,
-            )
+            if config.rotary_scaling_factor:
+                self.rotary_emb = NomicBertDynamicNTKRotaryEmbedding(
+                    dim=self.rotary_emb_dim,
+                    base=config.rotary_emb_base,
+                    scale_base=config.rotary_emb_scale_base,
+                    interleaved=config.rotary_emb_interleaved,
+                    rotary_scaling_factor=config.rotary_scaling_factor,
+                    max_position_embeddings=config.n_positions,
+                ) 
+            else:
+                self.rotary_emb = NomicBertRotaryEmbedding(
+                    dim=self.rotary_emb_dim,
+                    base=config.rotary_emb_base,
+                    scale_base=config.rotary_emb_scale_base,
+                    interleaved=config.rotary_emb_interleaved,
+                )
             # bug in xformers: https://github.com/facebookresearch/xformers/issues/841
             # uses the head dimension instead of the sequence dimension
             self.rotary_head_dim = getattr(config, "rotary_head_dim", False)