github/AMD-SHARK-Studio
1
1
Fork 0
mirror of https://github.com/nod-ai/AMD-SHARK-Studio.git synced 2026-04-03 03:00:17 -04:00
Code Issues Packages Projects Releases Wiki Activity

improved sharded performance and fixed issue with lmhead on rocm

Browse Source
This commit is contained in:
Elias Joseph
2023-12-04 18:26:46 -08:00
committed by PhaneeshB
parent d72da3801f
commit e43876cff5
2 changed files with 88 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
apps/language_models/scripts/vicuna.py
View File

@@ -12,6 +12,8 @@ import sys
import time
from dataclasses import dataclass
from os import environ
from dataclasses import dataclass
from os import environ
import torch
import torch_mlir
@@ -510,6 +512,8 @@ class ShardedVicuna(VicunaBase):
n_devices=None,
) -> None:
self.hf_auth_token = hf_auth_token
self.hidden_state_size_dict = {"vicuna": 4096, "llama2_7b": 4096, "llama2_13b" : 5120}
self.n_layers_dict = {"vicuna": 32, "llama2_7b": 32, "llama2_13b" : 40}
super().__init__(
model_name,
hf_model_path,
@@ -711,6 +715,27 @@ class ShardedVicuna(VicunaBase):
device_idx = max(idx_votes, key=idx_votes.get)
return device_idx
def write_dynamic_inputs_lmhead(self, ir, sample_input_length):
if self.precision in ["fp16", "int4"]:
precision_str = "f16"
else:
precision_str = "f32"
lines = ir.splitlines()
new_lines = []
for line in lines:
if f"%cst_0 =" in line:
new_lines.append(line)
new_lines.append("%c1 = arith.constant 1 : index")
new_lines.append(f"%dim = tensor.dim %arg0, %c1 : tensor<1x?x{self.hidden_state_size_dict[self.model_name]}x{precision_str}>")
else:
line = re.sub(f"{sample_input_length}x", "?x", line)
if "?x" in line:
line = re.sub("tensor.empty\(\)", "tensor.empty(%dim)", line)
new_lines.append(line)
return "\n".join(new_lines)
def compile_lmhead(
self,
lmh,
@@ -775,14 +800,21 @@ class ShardedVicuna(VicunaBase):
use_tracing=False,
verbose=False,
)
"""
bytecode_stream = BytesIO()
module.operation.write_bytecode(bytecode_stream)
bytecode = bytecode_stream.getvalue()
f_ = open(mlir_path, "wb")
f_.write(bytecode)
f_.close()
"""
module = str(module)
if self.precision in ["int4", "fp16"]:
module = self.write_dynamic_inputs_lmhead(module, 137)
filepath = Path(f"{self.dir_name}/lmhead.mlir")
f_ = open(mlir_path, "w+")
f_.write(module)
f_.close()
# download_public_file(
# "gs://shark_tank/elias/compressed_sv/lmhead.mlir",
# filepath.absolute(),
@@ -1163,6 +1195,7 @@ class ShardedVicuna(VicunaBase):
device_idx = idx % self.n_devices
else:
device_idx = None
print(device_idx, self.n_devices)
module = SharkInference(
None,
device=device,
@@ -1180,7 +1213,7 @@ class ShardedVicuna(VicunaBase):
if self.n_devices is not None:
device_idx = idx % self.n_devices
else:
device_idx = 0
device_idx = None
module = SharkInference(
mlirs[idx],
device=device,
@@ -1238,7 +1271,7 @@ class ShardedVicuna(VicunaBase):
if self.n_devices is not None:
device_idx = idx % self.n_devices
else:
device_idx = 0
device_idx = None
module = SharkInference(
None,
device=device,
@@ -1256,7 +1289,7 @@ class ShardedVicuna(VicunaBase):
if self.n_devices is not None:
device_idx = idx % self.n_devices
else:
device_idx = 0
device_idx = None
module = SharkInference(
mlirs[idx],
device=device,
@@ -1320,41 +1353,40 @@ class ShardedVicuna(VicunaBase):
placeholder_pkv_segment = tuple(
(
torch.zeros([1, 32, SAMPLE_INPUT_LEN, 128]),
torch.zeros([1, 32, SAMPLE_INPUT_LEN, 128]),
torch.zeros([1, self.n_layers_dict[self.model_name], SAMPLE_INPUT_LEN, 128]),
torch.zeros([1, self.n_layers_dict[self.model_name], SAMPLE_INPUT_LEN, 128]),
)
for _ in range(8)
)
placeholder_pkv_full = tuple(
(
torch.zeros([1, 32, SAMPLE_INPUT_LEN, 128]),
torch.zeros([1, 32, SAMPLE_INPUT_LEN, 128]),
torch.zeros([1, self.n_layers_dict[self.model_name], SAMPLE_INPUT_LEN, 128]),
torch.zeros([1, self.n_layers_dict[self.model_name], SAMPLE_INPUT_LEN, 128]),
)
for _ in range(32)
for _ in range(self.n_layers_dict[self.model_name])
)
placeholder_input0 = (
torch.zeros([1, SAMPLE_INPUT_LEN, 4096]),
torch.zeros([1, SAMPLE_INPUT_LEN, self.hidden_state_size_dict[self.model_name]]),
torch.zeros([1, 1, SAMPLE_INPUT_LEN, SAMPLE_INPUT_LEN]),
torch.zeros([1, SAMPLE_INPUT_LEN], dtype=torch.int64),
)
placeholder_input1 = (
torch.zeros([1, 1, 4096]),
torch.zeros([1, 1, self.hidden_state_size_dict[self.model_name]]),
torch.zeros([1, 1, 1, SAMPLE_INPUT_LEN + 1]),
torch.zeros([1, 1], dtype=torch.int64),
torch.zeros([1, 32, SAMPLE_INPUT_LEN, 128]),
torch.zeros([1, 32, SAMPLE_INPUT_LEN, 128]),
torch.zeros([1, self.n_layers_dict[self.model_name], SAMPLE_INPUT_LEN, 128]),
torch.zeros([1, self.n_layers_dict[self.model_name], SAMPLE_INPUT_LEN, 128]),
)
norm = VicunaNorm(vicuna_model.model.norm)
device_idx = self.get_device_index(
r"vicuna\.model\.model\.norm(?:\.|\s|$)"
)
print(device_idx)
norm = self.compile_norm(
norm,
torch.zeros([1, SAMPLE_INPUT_LEN, 4096]),
torch.zeros([1, SAMPLE_INPUT_LEN, self.hidden_state_size_dict[self.model_name]]),
device=self.device,
device_idx=device_idx,
)
@@ -1363,7 +1395,6 @@ class ShardedVicuna(VicunaBase):
device_idx = self.get_device_index(
r"vicuna\.model\.model\.embed_tokens(?:\.|\s|$)"
)
print(device_idx)
embeddings = self.compile_embedding(
embeddings,
(torch.zeros([1, SAMPLE_INPUT_LEN], dtype=torch.int64)),
@@ -1375,10 +1406,9 @@ class ShardedVicuna(VicunaBase):
device_idx = self.get_device_index(
r"vicuna\.model\.lm_head(?:\.|\s|$)"
)
print(device_idx)
lmhead = self.compile_lmhead(
lmhead,
torch.zeros([1, SAMPLE_INPUT_LEN, 4096]),
torch.zeros([1, SAMPLE_INPUT_LEN, self.hidden_state_size_dict[self.model_name]]),
device=self.device,
device_idx=device_idx,
)
@@ -1667,12 +1697,7 @@ class UnshardedVicuna(VicunaBase):
new_lines = []
# Using a while loop and the pop method to avoid creating a copy of module
if "llama2_13b" in self.model_name:
pkv_tensor_shape = "tensor<1x40x?x128x"
elif "llama2_70b" in self.model_name:
pkv_tensor_shape = "tensor<1x8x?x128x"
else:
pkv_tensor_shape = "tensor<1x32x?x128x"
pkv_tensor_shape = f"tensor<1x{self.n_layers_dict[self.model_name]}x?x128x"
if self.precision in ["fp16", "int4", "int8"]:
pkv_tensor_shape += "f16>"
else:

56
apps/language_models/src/model_wrappers/vicuna_sharded_model.py
View File

@@ -1,4 +1,5 @@
import torch
import time
class FirstVicunaLayer(torch.nn.Module):
@@ -110,9 +111,12 @@ class LMHeadCompiled(torch.nn.Module):
self.model = shark_module
def forward(self, hidden_states):
hidden_states = hidden_states.detach()
hidden_states_sample = hidden_states.detach()
output = self.model("forward", (hidden_states,))
output = torch.tensor(output)
return output
@@ -132,12 +136,14 @@ class VicunaNormCompiled(torch.nn.Module):
self.model = shark_module
def forward(self, hidden_states):
try:
hidden_states.detach()
except:
pass
output = self.model("forward", (hidden_states,))
output = self.model("forward", (hidden_states,), send_to_host=True)
output = torch.tensor(output)
return output
@@ -157,9 +163,11 @@ class VicunaEmbeddingCompiled(torch.nn.Module):
self.model = shark_module
def forward(self, input_ids):
input_ids.detach()
output = self.model("forward", (input_ids,))
output = self.model("forward", (input_ids,), send_to_host=True)
output = torch.tensor(output)
return output
@@ -177,10 +185,12 @@ class CompiledVicunaLayer(torch.nn.Module):
output_attentions=False,
use_cache=True,
):
if past_key_value is None:
hidden_states = hidden_states.detach()
attention_mask = attention_mask.detach()
position_ids = position_ids.detach()
#hidden_states = hidden_states.detach()
#attention_mask = attention_mask.detach()
#position_ids = position_ids.detach()
output = self.model(
"first_vicuna_forward",
(
@@ -188,11 +198,17 @@ class CompiledVicunaLayer(torch.nn.Module):
attention_mask,
position_ids,
),
send_to_host=False,
)
output0 = torch.tensor(output[0])
output1 = torch.tensor(output[1])
output2 = torch.tensor(output[2])
#output0 = torch.tensor(output[0])
#output1 = torch.tensor(output[1])
#output2 = torch.tensor(output[2])
output0 = output[0]
output1 = output[1]
output2 = output[2]
return (
output0,
@@ -202,11 +218,12 @@ class CompiledVicunaLayer(torch.nn.Module):
),
)
else:
hidden_states = hidden_states.detach()
attention_mask = attention_mask.detach()
position_ids = position_ids.detach()
pkv0 = past_key_value[0].detach()
pkv1 = past_key_value[1].detach()
#hidden_states = hidden_states.detach()
#attention_mask = attention_mask.detach()
#position_ids = position_ids.detach()
#pkv0 = past_key_value[0].detach()
pkv0 = past_key_value[0]
pkv1 = past_key_value[1]
output = self.model(
"second_vicuna_forward",
(
@@ -216,11 +233,16 @@ class CompiledVicunaLayer(torch.nn.Module):
pkv0,
pkv1,
),
send_to_host=False,
)
output0 = torch.tensor(output[0])
output1 = torch.tensor(output[1])
output2 = torch.tensor(output[2])
#output0 = torch.tensor(output[0])
#output1 = torch.tensor(output[1])
#output2 = torch.tensor(output[2])
output0 = output[0]
output1 = output[1]
output2 = output[2]
return (
output0,