add opt

config.json

@@ -0,0 +1,26 @@
+{
+  "activation_dropout": 0.0,
+  "activation_function": "relu",
+  "architectures": [
+    "OPTModel"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "d_model": 1024,
+  "dropout": 0.1,
+  "eos_token_id": 2,
+  "ffn_dim": 4096,
+  "init_std": 0.02,
+  "layerdrop": 0.0,
+  "max_position_embeddings": 2048,
+  "model_type": "opt",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 24,
+  "output_projection": true,
+  "pad_token_id": 1,
+  "torch_dtype": "float16",
+  "transformers_version": "4.19.0.dev0",
+  "vocab_size": 50272,
+  "word_embed_proj_dim": 512,
+  "do_layer_norm_before": false
+}

run.sh

@@ -1,2 +1,2 @@
 #!/usr/bin/env bash
-CUDA_VISIBLE_DEVICES="0" torchrun run_model.py --pipeline-model-parallel-size 1 --tensor-model-parallel-size 1
+CUDA_VISIBLE_DEVICES="0,3" torchrun run_model.py --pipeline-model-parallel-size 1 --tensor-model-parallel-size 1

run_model.py

@@ -1,9 +1,9 @@
 #!/usr/bin/env python3
-#!/usr/bin/env python3
 import os
 from transformers import AutoTokenizer, GPT2Tokenizer
 from megatron.initialize import initialize_megatron
 from metaseq import checkpoint_utils
+from transformers import OPTForCausalLM
 import torch
 
 path = "./model"
@@ -34,32 +34,56 @@ checkpoint = checkpoint_utils.load_model_ensemble_and_task(
 )
 
 model = checkpoint[0][0].eval()
-model = model.cuda().half()
+model = model.to("cuda:0").half()
+
+hf_model = OPTForCausalLM.from_pretrained("../opt-350m").to("cuda:1").half()
 
 
 # forward passes
 def single_batch_forward_logits(prompts):
     input_ids = tokenizer(prompts, return_tensors="pt").input_ids
     input_ids = torch.cat([torch.tensor([[0]]), input_ids], dim=-1)
-    input_ids = input_ids.cuda()
+    input_ids = input_ids.to("cuda:0")
     with torch.no_grad():
         logits = model(input_ids)[0]
     return logits
 
+# forward hf
+def forward_hf(prompts):
+    input_ids = tokenizer(prompts, return_tensors="pt").input_ids
+    input_ids = torch.cat([torch.tensor([[0]]), input_ids], dim=-1)
+    input_ids = input_ids.to("cuda:1")
+    with torch.no_grad():
+        logits = hf_model(input_ids)[0]
+    return logits
+
 prompts = [
-    "Today is a beautiful day and I want to",
-    "In the city of",
-    "Paris is the capital of France and",
-    "Computers and mobile phones have taken",
+   "Today is a beautiful day and I want to",
+   "In the city of",
+   "Paris is the capital of France and",
+   "Computers and mobile phones have taken",
 ]
 
+prompts = [
+   "Today is a beautiful day and I want to",
+]
+
+#import ipdb; ipdb.set_trace()
 print("Next word generation")
 for prompt in prompts:
     print("-------------")
     print(f"Prompt: {prompt}...\n")
-    logits = single_batch_forward_logits(prompt)
+    logits_fsq = single_batch_forward_logits(prompt)
+    pred_next_token = torch.argmax(logits_fsq[0, -1], -1)
+    next_token = tokenizer.convert_ids_to_tokens([pred_next_token])
+    next_token = next_token[0].replace("Ġ", "")
+    print(f"Next word: {next_token}")
+    print("-------------")
+    logits = forward_hf(prompt)
     pred_next_token = torch.argmax(logits[0, -1], -1)
     next_token = tokenizer.convert_ids_to_tokens([pred_next_token])
     next_token = next_token[0].replace("Ġ", "")
     print(f"Next word: {next_token}")
     print("-------------")
+
+torch.allclose(logits_fsq.cpu(), logits.cpu(), atol=1e-3)