Upload 2 files

split-torch-model-v1.py

@@ -0,0 +1,59 @@
+import torch
+from transformers import AutoTokenizer
+from torch import nn
+
+device1 = torch.device("cuda:0")
+device2 = torch.device("cuda:1")
+
+class SplitModel(nn.Module):
+    def __init__(self, embedding_layer, dropout_layer, gptj_blocks, layer_norm, lm_head):
+        super(SplitModel, self).__init__()
+        self.embedding_layer = embedding_layer
+        self.dropout_layer = dropout_layer
+        self.gptj_blocks = gptj_blocks
+        self.layer_norm = layer_norm
+        self.lm_head = lm_head
+    
+    def forward(self, input_ids, attention_mask):
+        tensor_ids = self.dropout_layer(self.embedding_layer(input_ids))
+        # GPTJBlock is missing the embedding positions that are necessary for self-attention.
+        # To fix this issue, you need to ensure that the position_ids are passed to each GPTJBlock during the forward pass. 
+        position_ids = torch.arange(tensor_ids.shape[1], dtype=torch.long, device=tensor_ids.device)
+        for block in self.gptj_blocks:
+            tensor_ids = block(tensor_ids, attention_mask=attention_mask, position_ids=position_ids)[0]
+        tensor_ids = tensor_ids.to(device2)
+        tensor_ids = self.layer_norm(tensor_ids)
+        logits = self.lm_head(tensor_ids)
+        logits = logits.to(device1)
+        return logits
+
+model_dir = "pt_fp32"
+model_path = f"{model_dir}/torch_model.pt"
+
+tokenizer = AutoTokenizer.from_pretrained(model_dir)
+full_model = torch.load(model_path)
+
+embedding_layer = full_model.transformer.wte.to(device1)
+dropout_layer = full_model.transformer.drop.to(device1)
+gptj_blocks = full_model.transformer.h.to(device1)
+layer_norm = full_model.transformer.ln_f.to(device2)
+lm_head = full_model.lm_head.to(device2)
+
+split_model = SplitModel(embedding_layer, dropout_layer, gptj_blocks, layer_norm, lm_head)
+
+input_text = "Hi I am Jade and I love"
+input("Press enter please")
+input_tokens = tokenizer.encode_plus(input_text, return_tensors="pt").to(device1)
+input_ids = input_tokens["input_ids"]
+temperature = 0.8
+max_new_tokens = 50
+with torch.no_grad():
+    for _ in range(max_new_tokens):
+        attention_mask = torch.ones_like(input_ids).to(device1)
+        logits = split_model(input_ids, attention_mask)[:, -1] / temperature
+        probabilities = torch.softmax(logits, dim=-1)
+        sampled_token_ids = torch.multinomial(probabilities, num_samples=1)
+        input_ids = torch.cat((input_ids, sampled_token_ids), dim=-1)
+    generated_ids = input_ids.squeeze().tolist()
+output = tokenizer.decode(generated_ids, skip_special_tokens=True)
+print(output)

split-torch-model-v2.py

@@ -0,0 +1,54 @@
+from transformers import AutoTokenizer
+import torch
+
+device1 = torch.device("cuda:0")
+device2 = torch.device("cuda:1")
+
+class SplitModel(torch.nn.Module):
+    def __init__(self, base_model):
+        super(SplitModel, self).__init__()
+        self.embedding_layer = base_model.transformer.wte.to(device1)
+        # self.dropout_layer = base_model.transformer.drop.to(device1)
+        self.gptj_blocks1 = torch.nn.ModuleList(base_model.transformer.h[:14]).to(device1)
+        self.gptj_blocks2 = torch.nn.ModuleList(base_model.transformer.h[14:]).to(device2)
+        self.layer_norm = base_model.transformer.ln_f.to(device2)
+        self.lm_head = base_model.lm_head.to(device2)
+    
+    def forward(self, input_ids, attention_mask):
+        # tensor_ids = self.dropout_layer(self.embedding_layer(input_ids))
+        tensor_ids = self.embedding_layer(input_ids)
+        position_ids = torch.arange(tensor_ids.shape[1], dtype=torch.long, device=tensor_ids.device)
+        for block in self.gptj_blocks1:
+            tensor_ids = block(tensor_ids, attention_mask=attention_mask, position_ids=position_ids)[0]
+        tensor_ids = tensor_ids.to(device2)
+        position_ids = position_ids.to(device2)
+        attention_mask = attention_mask.to(device2)
+        for block in self.gptj_blocks2:
+            tensor_ids = block(tensor_ids, attention_mask=attention_mask, position_ids=position_ids)[0]
+        tensor_ids = self.layer_norm(tensor_ids)
+        logits = self.lm_head(tensor_ids)
+        return logits.to(device1)
+
+model_dir = "pt_fp32"
+model_path = f"{model_dir}/torch_model.pt"
+
+tokenizer = AutoTokenizer.from_pretrained(model_dir)
+split_model = SplitModel(torch.load(model_path))
+
+input_text = "Hi I am Jade and I love"
+input_tokens = tokenizer.encode_plus(input_text, return_tensors="pt").to(device1)
+input_ids = input_tokens["input_ids"]
+temperature = 0.5
+max_new_tokens = 50
+with torch.no_grad():
+    # split_model.eval()
+    for _ in range(max_new_tokens):
+        attention_mask = torch.ones_like(input_ids).to(device1)
+        logits = split_model(input_ids, attention_mask)[:, -1] / temperature
+        probabilities = torch.softmax(logits, dim=-1)
+        sampled_token_ids = torch.multinomial(probabilities, num_samples=1)
+        input_ids = torch.cat((input_ids, sampled_token_ids), dim=-1)
+        del logits, probabilities, sampled_token_ids
+    generated_ids = input_ids.squeeze().tolist()
+output = tokenizer.decode(generated_ids, skip_special_tokens=True)
+print(output)