huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 12:24:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

Temporary implem of torch.compile on our stuff.

Browse Source
This commit is contained in:
Nicolas Patry 2024-03-21 18:56:40 +00:00
parent ed72e92126
commit 9f44af470c
3 changed files with 58 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
server/text_generation_server/models/flash_mistral.py
View File

@ -511,6 +511,21 @@ class BaseFlashMistral(FlashCausalLM):
cuda_graph = self.cuda_graphs.get(padded_bs, None)
if cu_seqlen_prefill is not None or cuda_graph is None:
if cu_seqlen_prefill is None:
logits, speculative_logits = self.compiled_model(
input_ids=input_ids,
position_ids=position_ids,
cu_seqlen_prefill=cu_seqlen_prefill,
kv_cache=kv_cache,
block_tables=block_tables,
slots=slots,
input_lengths=input_lengths,
max_s=max_s,
prefill_cache_indices=batch.prefill_cache_indices,
lm_head_indices=lm_head_indices,
)
else:
logits, speculative_logits = self.model.forward(
input_ids=input_ids,
position_ids=position_ids,

12
server/text_generation_server/server.py
View File

@ -154,7 +154,19 @@ class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
batch = batches[0]
concat_ns = None
torch.profiler._utils._init_for_cuda_graphs()
# prof = torch.profiler.profile()
# if self.model.rank != 0:
if True:
import contextlib
prof = contextlib.nullcontext()
else:
prof = torch.profiler.profile()
with prof:
generations, next_batch, timings = self.model.generate_token(batch)
# if self.model.rank == 0:
# prof.export_chrome_trace(f"out_rank_0.json")
self.cache.set(next_batch)
return generate_pb2.DecodeResponse(

37
server/text_generation_server/utils/layers.py
View File

@ -664,27 +664,27 @@ class TensorParallelHead(SuperLayer):
return super().forward(input)
world_size = self.process_group.size()
if len(input.shape) == 2 and isinstance(self.linear, FastLinear):
out_dim = self.linear.weight.shape[0]
# if len(input.shape) == 2 and isinstance(self.linear, FastLinear):
# out_dim = self.linear.weight.shape[0]
if input.shape[0] == 1:
world_out = input.new_empty(1, out_dim * world_size)
local_out = input.new_empty(1, out_dim)
gather_input = local_out
else:
world_out = input.new_empty(out_dim * world_size, input.shape[0])
gather_input = input.new_empty(out_dim, input.shape[0])
local_out = gather_input.T
# if input.shape[0] == 1:
# world_out = input.new_empty(1, out_dim * world_size)
# local_out = input.new_empty(1, out_dim)
# gather_input = local_out
# else:
# world_out = input.new_empty(out_dim * world_size, input.shape[0])
# gather_input = input.new_empty(out_dim, input.shape[0])
# local_out = gather_input.T
torch.mm(input, self.linear.weight.T, out=local_out)
# torch.mm(input, self.linear.weight.T, out=local_out)
torch.distributed.all_gather_into_tensor(
world_out, gather_input, group=self.process_group
)
# torch.distributed.all_gather_into_tensor(
# world_out, gather_input, group=self.process_group
# )
if input.shape[0] == 1:
return world_out
return world_out.T
# if input.shape[0] == 1:
# return world_out
# return world_out.T
output = super().forward(input)
world_output = [
@ -943,6 +943,7 @@ try:
self._sin_k_cached = None
self.scaling_factor = scaling_factor
self.dynamic_args = None
self._update_cos_sin_cache(torch.float16, inv_freq.device, seqlen=4096)
def forward(
self,
@ -1086,8 +1087,6 @@ try:
# But later on goes and cast cos/sin to float anyway: https://github.com/Dao-AILab/flash-attention/blob/017716451d446e464dde9aca3a3c1ed2209caaa9/csrc/rotary/rotary_cuda.cu#L29, which looks suboptimal.
dtype = torch.float32
self._update_cos_sin_cache(dtype, position_ids.device, max_s)
cos = torch.index_select(self._cos_cached, 0, position_ids)
sin = torch.index_select(self._sin_cached, 0, position_ids)
# Note: this unsqueeze is not necessary on RoCm + VLLM ROPE implementation, but we leave it as is to avoid yet an other controlflow.