File size: 2,491 Bytes
4798617 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 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) |