Further logging

Author: remi-or

src/transformers/generation/continuous_batching/cache.py

@@ -208,8 +208,8 @@ class PagedAttentionMemoryHandler:
     _activation_dtype = torch.bfloat16
     _activation_safety_factor = 2
     _input_dtype = torch.int32
-    _upper_bound_max_batch_tokens = 2048
-    _upper_bound_num_blocks = 16384
+    _upper_bound_max_batch_tokens = 256
+    _upper_bound_num_blocks = 4096
 
     def __init__(
         self,
@@ -271,14 +271,20 @@ def compute_num_blocks_and_max_batch_tokens(
         m: float = 0.1,
     ) -> tuple[int, int]:
         cache_memory = self.get_available_memory(max_memory_percent)
+        logger.info(f"Cache memory: {cache_memory}")
+
+        # Compute memory footprints # TODO: check and explain better
+        mem_per_activation_token = self._activation_dtype.itemsize * (self.hidden_size + self.vocab_size) * self._activation_safety_factor
+        mem_per_cache_token = 2 * self.num_heads * self.head_dim * self.num_layers * cache_dtype.itemsize
+        mem_per_input_token = 8 * m * self._input_dtype.itemsize
+        logger.info(f"Memory per activation token: {mem_per_activation_token}")
+        logger.info(f"Memory per cache token: {mem_per_cache_token}")
+        logger.info(f"Memory per input token: {mem_per_input_token}")
 
         # Compute second-degree polynomial coefficients
         a = m * self._activation_dtype.itemsize
-        b = 8 * m * self._input_dtype.itemsize
-        b += 2 * self.num_heads * self.head_dim * self.num_layers * cache_dtype.itemsize
-        c = self._activation_dtype.itemsize * (self.hidden_size + self.vocab_size) * self._activation_safety_factor
-        c += 2 * self._input_dtype.itemsize
-        c -= cache_memory
+        b = mem_per_input_token + mem_per_cache_token
+        c = mem_per_activation_token + 2 * self._input_dtype.itemsize - cache_memory
 
         # Compute discriminant and greatest solution
         discriminant = b**2 - 4 * a * c

src/transformers/generation/continuous_batching/continuous_api.py

@@ -262,8 +262,10 @@ def prepare_next_batch(self):
             self.max_seqlen_k = max(self.max_seqlen_k, key_length)
             state.position_offset += query_length
 
-        logger.info(
-            f"Scheduled: {len(self.requests_in_batch)}, Waiting: {len(self.scheduler.waiting_requests)}, Active: {len(self.scheduler.active_requests)}. cum Q: {cumulative_seqlens_q[-1]}. cum KV: {cumulative_seqlens_k[-1]}, free blocks: {self.cache.get_num_free_blocks()}"
+        logger.debug(
+            f"Scheduled: {len(self.requests_in_batch)}, Waiting: {len(self.scheduler.waiting_requests)}, "
+            f"Active: {len(self.scheduler.active_requests)}. cum Q: {cumulative_seqlens_q[-1]}. "
+            f"cum KV: {cumulative_seqlens_k[-1]}, free blocks: {self.cache.get_num_free_blocks()}"
         )
         self._build_tensors(
             input_ids,
@@ -666,7 +668,7 @@ def _inner_generation_loop(self, batch_processor: ContinuousBatchProcessor):
             torch.cuda.synchronize()
         batch_processor.prepare_next_batch()
         device, total, reserved, allocated = get_device_and_memory_breakdown()
-        logger.info(f"[Memory] Device: {device}, Total: {total}, Reserved: {reserved}, Allocated: {allocated}")
+        logger.debug(f"[Memory] Device: {device}, Total: {total}, Reserved: {reserved}, Allocated: {allocated}")
         if torch.cuda.is_available() and self.use_cuda_graph:
             if self.current_batch == 0:
                 self.warmup(batch_processor)
@@ -780,8 +782,8 @@ def generate_batch(
         """
         if not inputs:
             return []
-        if logger.getEffectiveLevel() <= logging.INFO:
-            logger.warning("Progress bar is disabled when logger level is less than INFO")
+        if logger.getEffectiveLevel() <= logging.DEBUG:
+            logger.warning("Progress bar is disabled when logger level is less than DEBUG")
             progress_bar = False
 
         # Initialize manager with the batch inputs

src/transformers/generation/continuous_batching/core.py

@@ -11,6 +11,7 @@
 
 # We centralize the logger here to coordinate between logging and progress bar
 logger = logging.getLogger("ContinuousBatchingLogger")
+logger.setLevel(logging.INFO)
 
 
 @staticmethod
@@ -102,12 +103,35 @@ class RequestState:
     static_outputs: list[int] = field(default_factory=list)  # Generated tokens
     allocated_blocks: list[int] = field(default_factory=list)  # Block IDs allocated to the request
     position_offset: int = 0  # Current position in the sequence for position_ids
-    status: RequestStatus = RequestStatus.PENDING  # Status of the request
+    _status: RequestStatus = RequestStatus.PENDING  # Status of the request, hidden behind a property
     max_new_tokens: int = 20  # Maximum number of new tokens to generate
     eos_token_id: int = -1  # ID of the end-of-sequence token
     created_time: float = field(default_factory=time.time)  # Time the request was created
     error: Optional[str] = None  # Error message if the request failed
     next_token: Optional[str] = None  # Next token to be generated
+    lifespan: tuple[float, float] = (-1, -1)  # (time request was no longer pending, time request finished)
+
+    @property
+    def status(self) -> RequestStatus:
+        return self._status
+
+    @status.setter
+    def status(self, value: RequestStatus):
+        if self._status == RequestStatus.PENDING:
+            self.lifespan = (time.time(), -1)
+        elif value == RequestStatus.FINISHED:
+            self.lifespan = (self.lifespan[0], time.time())
+            self.log_end_of_request()
+        self._status = value
+
+    def log_end_of_request(self):
+        prefill_len = len(self.full_prompt_ids)
+        decode_len = self.generated_len()
+        start_time = self.lifespan[0] - self.created_time
+        end_time = self.lifespan[1] - self.created_time
+        logger.info(
+            f"Request {self.request_id} finished: {prefill_len = } {decode_len = } {start_time = } {end_time = }"
+        )
 
     def current_len(self) -> int:
         """Get the current length of the sequence (prompt + generated tokens)."""
@@ -148,7 +172,7 @@ def update_with_token(self, token_id: int) -> bool:
     def __repr__(self):
         msg = [
             f"request_id={self.request_id}",
-            f"status={self.status}",
+            f"status={self._status}",
             f"out_tokens={self.generated_len()}",
             f"query_length={len(self.prompt_ids)}",
             f"remaining_tokens={len(self.remaining_prompt_ids)}",