b278557935
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Signed-off-by: Abatom <abzhonghua@gmail.com> Co-authored-by: Zhonghua Deng <abatom@163.com>
340 lines
10 KiB
Python
340 lines
10 KiB
Python
from typing import Dict, List, Optional
|
|
|
|
import torch
|
|
|
|
from vllm.lora.lora import LoRALayerWeights, PackedLoRALayerWeights
|
|
|
|
|
|
class DummyLoRAManager:
|
|
|
|
def __init__(self, device: torch.device = "cuda:0"):
|
|
super().__init__()
|
|
self._loras: Dict[str, LoRALayerWeights] = {}
|
|
self._device = device
|
|
|
|
def set_module_lora(self, module_name: str, lora: LoRALayerWeights):
|
|
self._loras[module_name] = lora
|
|
|
|
def get_module_lora(self, module_name: str) -> LoRALayerWeights:
|
|
return self._loras[module_name]
|
|
|
|
def init_random_lora(
|
|
self,
|
|
module_name: str,
|
|
weight: torch.Tensor,
|
|
rank: int = 8,
|
|
generate_embeddings_tensor: int = 0,
|
|
):
|
|
lora = LoRALayerWeights(
|
|
module_name,
|
|
rank=rank,
|
|
lora_alpha=1,
|
|
lora_a=torch.rand([weight.shape[1], rank],
|
|
dtype=weight.dtype,
|
|
device=self._device),
|
|
lora_b=torch.rand([rank, weight.shape[0]],
|
|
dtype=weight.dtype,
|
|
device=self._device),
|
|
)
|
|
if generate_embeddings_tensor:
|
|
lora.embeddings_tensor = torch.rand(
|
|
5,
|
|
generate_embeddings_tensor,
|
|
dtype=weight.dtype,
|
|
device=self._device,
|
|
)
|
|
self.set_module_lora(module_name, lora)
|
|
|
|
return lora
|
|
|
|
def init_lora(
|
|
self,
|
|
module_name: str,
|
|
input_dim: int,
|
|
output_dim: int,
|
|
rank=8,
|
|
noop=False,
|
|
embeddings_tensor=None,
|
|
):
|
|
lora = LoRALayerWeights(
|
|
module_name,
|
|
rank=rank,
|
|
lora_alpha=1,
|
|
lora_a=torch.rand([input_dim, rank], device="cuda"),
|
|
lora_b=torch.rand([rank, output_dim], device="cuda"),
|
|
embeddings_tensor=embeddings_tensor,
|
|
)
|
|
self.set_module_lora(module_name, lora)
|
|
return lora
|
|
|
|
def reset_lora(self):
|
|
self._loras = {}
|
|
|
|
def init_packed_lora(
|
|
self,
|
|
module_name: str,
|
|
input_dim: int,
|
|
output_dims: List[int],
|
|
noop_lora_index: Optional[List[int]] = None,
|
|
rank: int = 8,
|
|
):
|
|
base_loras: List[LoRALayerWeights] = []
|
|
noop_lora_index_set = set(noop_lora_index or [])
|
|
|
|
for i, out_dim in enumerate(output_dims):
|
|
base_lora = self.init_lora(
|
|
module_name + "_000_" + str(i),
|
|
input_dim,
|
|
out_dim,
|
|
rank=rank,
|
|
noop=i in noop_lora_index_set,
|
|
)
|
|
base_loras.append(base_lora)
|
|
packed_lora = PackedLoRALayerWeights.pack(base_loras)
|
|
self.set_module_lora(module_name, packed_lora)
|
|
return packed_lora
|
|
|
|
|
|
def assert_close(a, b):
|
|
rtol, atol = {
|
|
torch.float16: (6e-2, 6e-2),
|
|
torch.bfloat16: (6e-2, 6e-2),
|
|
torch.float32: (1e-2, 1e-2),
|
|
}[a.dtype]
|
|
torch.testing.assert_close(a, b, rtol=rtol, atol=atol)
|
|
|
|
|
|
def ref_torch_groupgemm(
|
|
out_tensor,
|
|
inputs,
|
|
lora_weights,
|
|
lora_indices_tensor,
|
|
seq_len_tensor,
|
|
batches,
|
|
scaling,
|
|
op_type,
|
|
) -> torch.Tensor:
|
|
out_list = []
|
|
current_offset = 0
|
|
for lora_index, b_length in zip(range(batches), seq_len_tensor):
|
|
input_weight = inputs[current_offset:b_length + current_offset, :]
|
|
current_offset += b_length
|
|
lora_weight = lora_weights[lora_indices_tensor[lora_index]]
|
|
result = torch.nn.functional.linear(input_weight, lora_weight)
|
|
result *= scaling
|
|
out_list.append(result)
|
|
cat_result = torch.cat(out_list, dim=0)
|
|
if op_type == "expand":
|
|
out_tensor += cat_result
|
|
else:
|
|
out_tensor.copy_(cat_result)
|
|
return
|
|
|
|
|
|
def generate_data(
|
|
batches,
|
|
hidden_size,
|
|
lora_nums,
|
|
max_rank,
|
|
seq_length,
|
|
dtype,
|
|
op_type,
|
|
device,
|
|
):
|
|
seq_len_tensor = torch.randint(seq_length, seq_length + 1,
|
|
(batches, )).to(device)
|
|
b_seq_start_loc = torch.cumsum(
|
|
torch.tensor([0] + seq_len_tensor[:-1].tolist(), dtype=torch.long),
|
|
dim=0,
|
|
).to(device)
|
|
total_tokens = seq_len_tensor.sum()
|
|
if op_type == "shrink":
|
|
inputs_tensor = torch.rand((total_tokens, hidden_size),
|
|
dtype=dtype).to(device)
|
|
lora_weights = torch.rand(
|
|
(lora_nums, max_rank, hidden_size), # col-major
|
|
dtype=dtype,
|
|
).to(device)
|
|
# shrink op need atomic_add, so output is initinized by 0
|
|
ref_out_tensor = torch.zeros((total_tokens, max_rank),
|
|
dtype=dtype,
|
|
device=inputs_tensor.device)
|
|
# NOTE shrink kernel using torch.float32 as output type
|
|
our_out_tensor = torch.zeros((total_tokens, max_rank),
|
|
dtype=torch.float32).to(device)
|
|
else:
|
|
inputs_tensor = torch.rand(
|
|
(total_tokens, max_rank),
|
|
dtype=dtype,
|
|
).to(device)
|
|
lora_weights = torch.rand(
|
|
(lora_nums, hidden_size, max_rank), # col-major
|
|
dtype=dtype,
|
|
).to(device)
|
|
# expand op needs to complete y+=a@lora_b, so output is
|
|
# initinized randomly
|
|
ref_out_tensor = torch.rand(
|
|
(total_tokens, hidden_size),
|
|
dtype=dtype,
|
|
).to(device)
|
|
# Ensure the same input.
|
|
our_out_tensor = ref_out_tensor.clone()
|
|
lora_indices_tensor = torch.randint(0,
|
|
lora_nums - 1 if lora_nums > 1 else 1,
|
|
(batches, )).to(device)
|
|
indices = torch.zeros((total_tokens), dtype=torch.long).to(device)
|
|
current_offset = 0
|
|
for b_id in range(batches):
|
|
lora_index = lora_indices_tensor[b_id]
|
|
indices[current_offset:current_offset +
|
|
seq_len_tensor[b_id]].copy_(lora_index)
|
|
current_offset += seq_len_tensor[b_id].item()
|
|
return (
|
|
inputs_tensor,
|
|
lora_weights,
|
|
our_out_tensor,
|
|
ref_out_tensor,
|
|
b_seq_start_loc,
|
|
lora_indices_tensor,
|
|
seq_len_tensor,
|
|
indices,
|
|
)
|
|
|
|
|
|
def generate_data_for_expand_nslices(
|
|
batches,
|
|
hidden_size,
|
|
lora_nums,
|
|
max_rank,
|
|
seq_length,
|
|
dtype,
|
|
nslices,
|
|
device,
|
|
):
|
|
seq_len_tensor = torch.randint(seq_length, seq_length + 1,
|
|
(batches, )).to(device)
|
|
b_seq_start_loc = torch.cumsum(
|
|
torch.tensor([0] + seq_len_tensor[:-1].tolist(), dtype=torch.long),
|
|
dim=0,
|
|
).to(device)
|
|
total_tokens = seq_len_tensor.sum()
|
|
inputs_tensor = torch.rand(
|
|
(total_tokens, max_rank),
|
|
dtype=dtype,
|
|
).to(device)
|
|
lora_weights_lst = []
|
|
for _ in range(nslices):
|
|
lora_weights_lst.append(
|
|
torch.rand(
|
|
(lora_nums, hidden_size, max_rank), # col-major
|
|
dtype=dtype,
|
|
).to(device))
|
|
# expand op needs to complete y+=a@lora_b, so output is
|
|
# initinized randomly
|
|
ref_out_tensor = torch.rand((total_tokens, hidden_size * nslices),
|
|
dtype=dtype).to(device)
|
|
# Ensure the same input.
|
|
our_out_tensor = ref_out_tensor.clone()
|
|
lora_indices_tensor = torch.randint(0,
|
|
lora_nums - 1 if lora_nums > 1 else 1,
|
|
(batches, ))
|
|
indices = torch.zeros((total_tokens), dtype=torch.long).to(device)
|
|
current_offset = 0
|
|
for b_id in range(batches):
|
|
lora_index = lora_indices_tensor[b_id]
|
|
indices[current_offset:current_offset +
|
|
seq_len_tensor[b_id]] = (lora_index.item())
|
|
current_offset += seq_len_tensor[b_id].item()
|
|
|
|
lora_indices_tensor = lora_indices_tensor.to(device)
|
|
return (
|
|
inputs_tensor,
|
|
lora_weights_lst,
|
|
our_out_tensor,
|
|
ref_out_tensor,
|
|
b_seq_start_loc,
|
|
lora_indices_tensor,
|
|
seq_len_tensor,
|
|
indices,
|
|
)
|
|
|
|
|
|
def generate_data_for_nslices(
|
|
batches,
|
|
hidden_size,
|
|
lora_nums,
|
|
max_rank,
|
|
seq_length,
|
|
nslices,
|
|
dtype,
|
|
op_type,
|
|
device,
|
|
):
|
|
seq_len_tensor = torch.randint(seq_length, seq_length + 1,
|
|
(batches, )).to(device)
|
|
b_seq_start_loc = torch.cumsum(
|
|
torch.tensor([0] + seq_len_tensor[:-1].tolist(), dtype=torch.long),
|
|
dim=0,
|
|
).to(device)
|
|
total_tokens = seq_len_tensor.sum()
|
|
|
|
lora_weights_lst = []
|
|
if op_type == "shrink":
|
|
|
|
inputs_tensor = torch.rand((total_tokens, hidden_size),
|
|
dtype=dtype).to(device)
|
|
|
|
for _ in range(nslices):
|
|
if op_type == "shrink":
|
|
lora_weights_lst.append(
|
|
torch.rand(
|
|
(lora_nums, max_rank, hidden_size), # col-major
|
|
dtype=dtype,
|
|
).to(device))
|
|
# NOTE shrink kernel using torch.float32 as output type
|
|
# shrink op need atomic_add, so output is initinized by 0
|
|
our_out_tensor = torch.zeros(
|
|
(nslices, total_tokens, max_rank),
|
|
dtype=torch.float32,
|
|
).to(device)
|
|
else:
|
|
inputs_tensor = torch.rand(
|
|
(nslices, total_tokens, max_rank),
|
|
dtype=dtype,
|
|
).to(device)
|
|
for _ in range(nslices):
|
|
lora_weights_lst.append(
|
|
torch.rand(
|
|
(lora_nums, hidden_size, max_rank), # col-major
|
|
dtype=dtype,
|
|
).to(device))
|
|
# expand op needs to complete y+=a@lora_b, so output is
|
|
# initinized randomly
|
|
our_out_tensor = torch.rand((total_tokens, hidden_size * nslices),
|
|
dtype=dtype).to(device)
|
|
|
|
# Ensure the same input.
|
|
ref_out_tensor = our_out_tensor.clone()
|
|
lora_indices_tensor = torch.randint(0,
|
|
lora_nums - 1 if lora_nums > 1 else 1,
|
|
(batches, ))
|
|
indices = torch.zeros((total_tokens), dtype=torch.long).to(device)
|
|
current_offset = 0
|
|
for b_id in range(batches):
|
|
lora_index = lora_indices_tensor[b_id]
|
|
indices[current_offset:current_offset +
|
|
seq_len_tensor[b_id]] = (lora_index.item())
|
|
current_offset += seq_len_tensor[b_id].item()
|
|
|
|
lora_indices_tensor = lora_indices_tensor.to(device)
|
|
return (
|
|
inputs_tensor,
|
|
lora_weights_lst,
|
|
our_out_tensor,
|
|
ref_out_tensor,
|
|
b_seq_start_loc,
|
|
lora_indices_tensor,
|
|
seq_len_tensor,
|
|
indices,
|
|
)
|