up

README.md

@@ -0,0 +1,43 @@
+---
+tags:
+- opt_metasq
+---
+
+# This repo let's you run the following checkpoint using facebookresearch/metaseq.
+
+Do the following:
+
+## 1. Install PyTorch
+```
+pip3 install torch==1.10.1+cu113 torchvision==0.11.2+cu113 torchaudio==0.10.1+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
+```
+
+## 2. Install Megatron
+```
+git clone https://github.com/patrickvonplaten/Megatron-LM.git
+cd Megatron-LM
+pip3 install six regex
+pip3 install -e .
+```
+## 3. Install fairscale
+```
+git clone https://github.com/facebookresearch/fairscale.git
+cd fairscale
+git checkout prefetch_fsdp_params_simple
+pip3 install -e .
+```
+## 4. Install metaseq
+```
+git clone https://github.com/patrickvonplaten/metaseq.git
+cd metaseq
+pip3 install -e .
+```
+
+## 5. Clone this repo (click top right on "How to clone")
+
+## 6. Run the following:
+
+```bash
+cd <path/to/cloned/repo>
+bash run.sh
+```

model/restored.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f76f91b201a0f295cea6cd41c09526d530559ba4ea72cd6b5c536b12d6a202d4
+size 714040635

model/special_tokens_map.json

@@ -0,0 +1 @@
+{"bos_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "eos_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "unk_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}}

model/tokenizer_config.json

@@ -0,0 +1 @@
+{"errors": "replace", "unk_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "bos_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "eos_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "add_prefix_space": false, "tokenizer_class": "GPT2Tokenizer"}

run.sh

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

run_model.py

@@ -0,0 +1,65 @@
+#!/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
+import torch
+
+path = "./model"
+
+# just need to initialize args with something, 
+# => doesn't need to correspond to the "correct" architecture for this checkpoint
+initialize_megatron(args_defaults={
+    "micro_batch_size": 1, 
+    "num_layers": 12, 
+    "hidden_size": 768, 
+    "num_attention_heads": 12,
+    "max_position_embeddings": 2048,
+    "encoder_seq_length": 2048
+})
+
+vocab_file = os.path.join(path, "gpt2-vocab.json")
+merges_file = os.path.join(path, "gpt2-merges.txt")
+
+tokenizer = GPT2Tokenizer(vocab_file, merges_file)
+tokenizer.save_pretrained(path)
+
+checkpoint = checkpoint_utils.load_model_ensemble_and_task(
+    [os.path.join(path, "restored.pt")],
+    arg_overrides={
+        "vocab_filename": vocab_file,
+        "merges_filename": merges_file,
+    }
+)
+
+model = checkpoint[0][0].eval()
+model = model.cuda().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()
+    with torch.no_grad():
+        logits = 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",
+]
+
+print("Next word generation")
+for prompt in prompts:
+    print("-------------")
+    print(f"Prompt: {prompt}...\n")
+    logits = single_batch_forward_logits(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("-------------")