mirror of
https://github.com/huggingface/text-generation-inference.git
synced 2025-04-21 23:12:07 +00:00
Co-authored-by: mrs303 <54661797+mrs303@users.noreply.github.com>
This commit is contained in:
Karol Damaszke
GitHub
parent
212136dff8
commit
022ce1eaaf
@ -957,6 +957,9 @@ class CausalLM(Model):
|
||||
new_input_length = input_length + 1
|
||||
|
||||
# Generated token
|
||||
if is_tokenizer_transparent(self.tokenizer) and len(stopping_criteria.stop_sequence_criterias) == 0:
|
||||
next_token_text = ''
|
||||
else:
|
||||
next_token_text, prefix_offset, read_offset = self.decode_token(
|
||||
all_input_ids[0:new_input_length, 0], prefix_offset, read_offset
|
||||
)
|
||||
@ -975,6 +978,9 @@ class CausalLM(Model):
|
||||
if i % self.world_size == self.rank:
|
||||
if stop:
|
||||
# Decode generated tokens
|
||||
if is_tokenizer_transparent(self.tokenizer):
|
||||
output_text = None
|
||||
else:
|
||||
output_text = self.decode(
|
||||
all_input_ids[new_input_length - stopping_criteria.current_tokens: new_input_length, 0]
|
||||
)
|
||||
@ -1034,6 +1040,7 @@ class CausalLM(Model):
|
||||
req.input_length = new_input_length
|
||||
req.prefix_offset = prefix_offset
|
||||
req.read_offset = read_offset
|
||||
|
||||
htorch.core.mark_step()
|
||||
self.step = self.step + 1
|
||||
if self.hb_profiler is not None:
|
||||
|
||||
@ -1,94 +0,0 @@
|
||||
import os
|
||||
import threading
|
||||
import queue
|
||||
|
||||
from contextlib import contextmanager
|
||||
import time
|
||||
import json
|
||||
|
||||
class FileWriter(threading.Thread):
|
||||
def __init__(self, filename, write_queue):
|
||||
super().__init__()
|
||||
self.filename = filename
|
||||
self.write_queue = write_queue
|
||||
self.daemon = True
|
||||
self.timer_event = threading.Event()
|
||||
|
||||
def _drain_write_queue(self):
|
||||
content = ""
|
||||
while True:
|
||||
try:
|
||||
element = self.write_queue.get_nowait()
|
||||
content += element
|
||||
except queue.Empty:
|
||||
break
|
||||
return content
|
||||
|
||||
def run(self):
|
||||
# don't check the queue too often
|
||||
while not self.timer_event.wait(1):
|
||||
# Block and wait for the next item in the queue
|
||||
content = self.write_queue.get()
|
||||
# Collect any other items in the queue
|
||||
content += self._drain_write_queue()
|
||||
|
||||
with open(self.filename, "a") as outfile:
|
||||
outfile.write(content)
|
||||
|
||||
class Profiler():
|
||||
profiling_trace_events = queue.Queue()
|
||||
event_tid = {"counter": 1, "external": 2, "internal": 3, "own": 4}
|
||||
filename = "server_events.json"
|
||||
|
||||
def __init__(self):
|
||||
self.step = 0
|
||||
self.enabled = os.getenv("TGI_PROFILER_ENABLED", "false").lower() == "true" and int(os.getenv("RANK", "0")) == 0
|
||||
if self.enabled:
|
||||
# initialize the trace file
|
||||
with open(self.filename, "w") as outfile:
|
||||
outfile.write('{"traceEvents": ')
|
||||
file_writer = FileWriter(self.filename, self.profiling_trace_events)
|
||||
file_writer.start()
|
||||
|
||||
@contextmanager
|
||||
def record_event(self, type, name, args=None, util=None, count_step=False):
|
||||
if self.enabled:
|
||||
start = time.time() * 1000000.0
|
||||
if util is not None:
|
||||
self._add_util_event(util, start)
|
||||
|
||||
if count_step:
|
||||
if args is None:
|
||||
args = {}
|
||||
args["step"] = self.step
|
||||
self.step += 1
|
||||
event = {
|
||||
"pid": 1,
|
||||
"tid": self.event_tid[type],
|
||||
"ph": "X",
|
||||
"name": name,
|
||||
"ts": start,
|
||||
"dur": None,
|
||||
"args": args
|
||||
}
|
||||
yield
|
||||
|
||||
end = time.time() * 1000000.0
|
||||
event["dur"] = end - start
|
||||
|
||||
self.profiling_trace_events.put(json.dumps([event]))
|
||||
else:
|
||||
yield
|
||||
|
||||
def _add_util_event(self, util, start):
|
||||
util_event = {
|
||||
"pid": 1,
|
||||
"tid": self.event_tid["counter"],
|
||||
"ph": "C",
|
||||
"name": "util",
|
||||
"ts": start,
|
||||
"args": {
|
||||
"util": util["util"],
|
||||
}
|
||||
}
|
||||
self.profiling_trace_events.put(json.dumps([util_event]))
|
||||
@ -16,12 +16,9 @@ from text_generation_server.models import Model, get_model
|
||||
from text_generation_server.pb import generate_pb2_grpc, generate_pb2
|
||||
from text_generation_server.tracing import UDSOpenTelemetryAioServerInterceptor
|
||||
|
||||
from .profiler import Profiler
|
||||
|
||||
class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
|
||||
def __init__(self, model: Model, cache: Cache, server_urls: List[str]):
|
||||
self.profiler = Profiler()
|
||||
with self.profiler.record_event("external", "init"):
|
||||
self.cache = cache
|
||||
self.model = model
|
||||
self.server_urls = server_urls
|
||||
@ -44,7 +41,6 @@ class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
|
||||
return generate_pb2.ServiceDiscoveryResponse(urls=self.server_urls)
|
||||
|
||||
async def ClearCache(self, request, context):
|
||||
with self.profiler.record_event("external", "clear_cache"):
|
||||
if request.HasField("id"):
|
||||
self.cache.delete(request.id)
|
||||
else:
|
||||
@ -53,12 +49,6 @@ class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
|
||||
|
||||
async def FilterBatch(self, request, context):
|
||||
batch = self.cache.pop(request.batch_id)
|
||||
with self.profiler.record_event(
|
||||
type="external",
|
||||
name="filter_batch",
|
||||
args={"batch_id": request.batch_id, "request_ids": [id for id in request.request_ids]},
|
||||
util={"util": len(batch.requests)}
|
||||
):
|
||||
if batch is None:
|
||||
raise ValueError(f"Batch ID {request.batch_id} not found in cache.")
|
||||
filtered_batch = batch.filter(request.request_ids)
|
||||
@ -72,7 +62,6 @@ class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
|
||||
batch, self.model.tokenizer, self.model.dtype, self.model.device, self.model.is_optimized_for_gaudi
|
||||
)
|
||||
|
||||
with self.profiler.record_event("external", "warmup"):
|
||||
batches = [batch_from_pb(batch) for batch in request.batches]
|
||||
self.model.warmup(batches)
|
||||
|
||||
@ -82,12 +71,6 @@ class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
|
||||
batch = self.model.batch_type.from_pb(
|
||||
request.batch, self.model.tokenizer, self.model.dtype, self.model.device, self.model.is_optimized_for_gaudi
|
||||
)
|
||||
with self.profiler.record_event(
|
||||
type="external",
|
||||
name="prefill",
|
||||
args={"batch_size": batch.batch_size, "sequence_length": batch.seq_length}
|
||||
):
|
||||
with self.profiler.record_event(type="internal", name="generate_token", count_step=True):
|
||||
generations, next_batch = self.model.generate_token([batch])
|
||||
self.cache.set(next_batch)
|
||||
|
||||
@ -97,13 +80,6 @@ class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
|
||||
)
|
||||
|
||||
async def Decode(self, request, context):
|
||||
batch0 = self.cache.cache[request.batches[0].id]
|
||||
with self.profiler.record_event(
|
||||
type="external",
|
||||
name="decode",
|
||||
args={"request_batches": [batch.id for batch in request.batches], "batch_size": batch0.batch_size},
|
||||
util={"util": len(batch0.requests)}
|
||||
):
|
||||
if len(request.batches) == 0:
|
||||
raise ValueError("Must provide at least one batch")
|
||||
|
||||
@ -117,7 +93,6 @@ class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
|
||||
if len(batches) == 0:
|
||||
raise ValueError("All batches are empty")
|
||||
|
||||
with self.profiler.record_event(type="internal", name="generate_token", count_step=True):
|
||||
generations, next_batch = self.model.generate_token(batches)
|
||||
self.cache.set(next_batch)
|
||||
|
||||
|
||||
@ -212,6 +212,11 @@ class HeterogeneousNextTokenChooser:
|
||||
scores = warper(input_ids, scores)
|
||||
|
||||
next_ids = self.choice(scores)
|
||||
# ignore logprobs if we use greedy search
|
||||
if type(self.choice) == Greedy:
|
||||
logprobs = torch.zeros_like(scores, device="cpu")
|
||||
next_logprobs = torch.zeros_like(next_ids.view(-1), device="cpu")
|
||||
else:
|
||||
logprobs = torch.log_softmax(scores, -1)
|
||||
next_logprobs = torch.gather(logprobs, 1, next_ids.view(-1, 1)).view(-1)
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user