Merge pull request #175 from VE-FORBRYDERNE/gptj-patch

Fix GPT-J model loading in TPU Colab when `vocab_size` is not divisible by 8

Author: henk717

maps/gptj.json

@@ -9,11 +9,11 @@
   },
   "static_weights": {
     "transformer.wte.weight": {"mtj": {"module": "embedding_shard/~/linear", "param": "w", "transforms": ["no_transpose", "vocab_pad"]}},
-    "transformer.wte.bias": {"mtj": {"module": "embedding_shard/~/linear", "param": "b"}},
+    "transformer.wte.bias": {"mtj": {"module": "embedding_shard/~/linear", "param": "b", "transforms": ["vocab_pad"]}},
     "transformer.ln_f.weight": {"mtj": {"module": "projection_shard/~/replicated_layer_norm", "param": "scale"}},
     "transformer.ln_f.bias": {"mtj": {"module": "projection_shard/~/replicated_layer_norm", "param": "offset"}},
     "lm_head.weight": {"mtj": {"module": "projection_shard/~/linear", "param": "w", "transforms": ["vocab_pad"]}},
-    "lm_head.bias": {"mtj": {"module": "projection_shard/~/linear", "param": "b"}}
+    "lm_head.bias": {"mtj": {"module": "projection_shard/~/linear", "param": "b", "transforms": ["vocab_pad"]}}
   },
   "layer_weights": {
     "transformer.h.{layer}.attn.bias": {},

tpu_mtj_backend.py

@@ -1304,7 +1304,7 @@ def callback(model_dict, f, **_):
                     if "divide_by_shards" in transforms:
                         tensor /= params["cores_per_replica"]
                     if "vocab_pad" in transforms:
-                        tensor = torch.nn.functional.pad(tensor, (0, 0, 0, params["n_vocab_padding"]))
+                        tensor = torch.nn.functional.pad(tensor, (0,) * (tensor.ndim * 2 - 1) + (params["n_vocab_padding"],))
                     if "no_transpose" not in transforms and tensor.ndim == 2:
                         tensor = tensor.T
                     tensor.unsqueeze_(0)