Fix various issues related to OPT models

mtj_softtuner/core.py

@@ -8,6 +8,7 @@
 import progressbar
 from tqdm.auto import tqdm
 import multiprocessing
+import contextlib
 import functools
 import time
 import json
@@ -167,11 +168,6 @@ def __call__(self, x: jnp.DeviceArray, **kwargs) -> jnp.DeviceArray:
         )
         mask = jnp.broadcast_to((x < self.in_dim)[:, jnp.newaxis], proj_out.shape)
         proj_out = jnp.where(mask, proj_out, 0)
-        if (
-            not kwargs.get("mtj_softtuner_disable_pe", False)
-            and getattr(self, "project_in") is not None
-        ):
-            proj_out @= self.project_in
         soft_shard_start_index = (
             jax.lax.axis_index("shard") * self.softtune_in_dim_per_shard
         )
@@ -186,6 +182,11 @@ def __call__(self, x: jnp.DeviceArray, **kwargs) -> jnp.DeviceArray:
             self, "has_sqrt_embed_scale", False
         ):
             proj_out *= jnp.sqrt(self.out_dim).astype(proj_out.dtype)
+        if (
+            not kwargs.get("mtj_softtuner_disable_pe", False)
+            and getattr(self, "project_in") is not None
+        ):
+            proj_out @= self.project_in
         if (
             not kwargs.get("mtj_softtuner_disable_pe", False)
             and not getattr(self, "post_embed", False)
@@ -482,6 +483,8 @@ def hf_conversion_callback(model_dict, f, **_):
             try:
                 last_storage_key = None
                 f = None
+                if utils.num_shards is not None:
+                    utils.current_shard += 1
                 for key in sorted(
                     model_dict.keys(),
                     key=lambda k: (model_dict[k].key, model_dict[k].seek_offset),
@@ -685,9 +688,29 @@ def initialize_thread_resources(shards: int, backend=None):
 def get_tokenizer(
     params: dict, tokenizer_id: Optional[str] = None
 ) -> transformers.PreTrainedTokenizerBase:
-    return transformers.GPT2TokenizerFast.from_pretrained(
-        params.get("tokenizer_id", "gpt2")
-    )
+    try:
+        tokenizer = transformers.AutoTokenizer.from_pretrained(
+            params.get("tokenizer_id", "gpt2")
+        )
+    except ValueError:
+        tokenizer = transformers.AutoTokenizer.from_pretrained(
+            params.get("tokenizer_id", "gpt2"), use_fast=False
+        )
+
+    @contextlib.contextmanager
+    def _mtj_softtuner_no_prefix():
+        add_bos_token = getattr(tokenizer, "add_bos_token", False)
+        add_prefix_space = getattr(tokenizer, "add_prefix_space", False)
+        tokenizer.add_bos_token = False
+        tokenizer.add_prefix_space = False
+        try:
+            yield
+        finally:
+            tokenizer.add_bos_token = add_bos_token
+            tokenizer.add_prefix_space = add_prefix_space
+
+    tokenizer._mtj_softtuner_no_prefix = _mtj_softtuner_no_prefix
+    return tokenizer
 
 
 def get_hf_checkpoint_metadata(ckpt_path: str):

mtj_softtuner/trainers/basic.py

@@ -81,12 +81,13 @@ def get_initial_soft_embeddings(
                 )
             )
             tokenizer = self.get_tokenizer()
-            special_tokens = set(
-                itertools.chain.from_iterable(
-                    tokenizer.encode(v)
-                    for v in tokenizer.special_tokens_map_extended.values()
+            with tokenizer._mtj_softtuner_no_prefix():
+                special_tokens = set(
+                    itertools.chain.from_iterable(
+                        tokenizer.encode(str(v))
+                        for v in tokenizer.special_tokens_map_extended.values()
+                    )
                 )
-            )
             sample_space = [
                 k for k in range(self.data.params["n_vocab"]) if k not in special_tokens
             ]
@@ -105,4 +106,5 @@ def tokenize_dataset_callback(
     ) -> List[int]:
         if self.data.newlinemode == "s":
             text = text.replace("\n", "</s>")
-        return tokenizer.encode(text) + self.data.params["eos_token"]
+        with tokenizer._mtj_softtuner_no_prefix():
+            return tokenizer.encode(text) + self.data.params["eos_token"]