diff --git a/astrai/dataset/dataset.py b/astrai/dataset/dataset.py index 08a4566..4db6c7d 100644 --- a/astrai/dataset/dataset.py +++ b/astrai/dataset/dataset.py @@ -492,9 +492,6 @@ class DPODataset(BaseDataset): required_keys = ["chosen", "rejected", "chosen_mask", "rejected_mask"] - def make_processor(self, tokenizer, max_len: int): - return partial(dpo_processor, tokenizer=tokenizer, max_len=max_len) - def __getitem__(self, index: int) -> Dict[str, Tensor]: return { "chosen": self.store.fetch_record(index, "chosen").to(dtype=torch.long), diff --git a/astrai/inference/cache/buffer.py b/astrai/inference/cache/buffer.py index de89892..b43a1aa 100644 --- a/astrai/inference/cache/buffer.py +++ b/astrai/inference/cache/buffer.py @@ -72,16 +72,6 @@ class KVStorage: (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype ) - def get_key_buffer(self, layer_id: int) -> Tensor: - return self.k_buffer[layer_id] - - def get_value_buffer(self, layer_id: int) -> Tensor: - return self.v_buffer[layer_id] - - def set_kv_buffer(self, layer_id: int, loc: Tensor, k: Tensor, v: Tensor) -> None: - self.k_buffer[layer_id, loc] = k - self.v_buffer[layer_id, loc] = v - @dataclass class KVCache: diff --git a/astrai/inference/metrics.py b/astrai/inference/metrics.py index 5aaba38..b6abb6f 100644 --- a/astrai/inference/metrics.py +++ b/astrai/inference/metrics.py @@ -148,17 +148,6 @@ class MetricsCollector: self._completed.append(timing) self._accumulate(timing) - def clear(self): - """Reset all state (e.g. on engine shutdown).""" - self._timings.clear() - self._completed.clear() - self._ttft_ms_sum = 0.0 - self._ttft_ms_count = 0 - self._decode_tps_sum = 0.0 - self._decode_tps_count = 0 - self._e2e_ms_sum = 0.0 - self._e2e_ms_count = 0 - # timing scopes @contextmanager @@ -180,17 +169,6 @@ class MetricsCollector: t._decode_steps += 1 t._decode_total_s += dt - # access - - def get_timing(self, task_id: str) -> Optional[TaskTiming]: - """Return the timing record for *task_id* (active or completed).""" - if task_id in self._timings: - return self._timings[task_id] - for t in self._completed: - if t.task_id == task_id: - return t - return None - # aggregate stats def get_stats(self) -> Dict[str, Any]: diff --git a/astrai/inference/scheduler.py b/astrai/inference/scheduler.py index e947b54..56061b0 100644 --- a/astrai/inference/scheduler.py +++ b/astrai/inference/scheduler.py @@ -79,29 +79,21 @@ class InferenceScheduler: if backend is None: self._backend = None - default_backend = get_backend() - self._backend_name = type(default_backend).__name__ - with attn_backend(default_backend): - self._executor = Executor( - model=model, - kv_cache=self._cache, - task_cache=self._task_cache, - device=self.device, - dtype=self.dtype, - enable_cuda_graph=enable_cuda_graph, - ) + active_backend = get_backend() else: - with attn_backend(backend): + active_backend = backend + with attn_backend(active_backend): + if backend is not None: self._backend = get_backend() - self._backend_name = type(self._backend).__name__ - self._executor = Executor( - model=model, - kv_cache=self._cache, - task_cache=self._task_cache, - device=self.device, - dtype=self.dtype, - enable_cuda_graph=enable_cuda_graph, - ) + self._backend_name = type(get_backend()).__name__ + self._executor = Executor( + model=model, + kv_cache=self._cache, + task_cache=self._task_cache, + device=self.device, + dtype=self.dtype, + enable_cuda_graph=enable_cuda_graph, + ) self._stop_event = threading.Event() self._loop_thread: Optional[threading.Thread] = None @@ -283,12 +275,7 @@ class InferenceScheduler: except Exception as e: self._stop_event.set() logger.error(f"Scheduler loop crashed: {e}", exc_info=True) - for task in self._task_mgr.get_active_tasks(): - self._task_mgr.invoke_callback(task.task_id, STOP) - self._task_cache.task_free(task.task_id) - for task in self._task_mgr.get_waiting_tasks(): - self._task_mgr.invoke_callback(task.task_id, STOP) - self._task_mgr.clear_queues() + self._abort_and_clear(free_waiting=False) def start(self): if self._loop_thread is not None and self._loop_thread.is_alive(): @@ -304,15 +291,20 @@ class InferenceScheduler: if self._loop_thread is not None: self._loop_thread.join(timeout=2.0) self._loop_thread = None + self._abort_and_clear(free_waiting=True) + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def _abort_and_clear(self, free_waiting: bool): + """Invoke STOP callbacks, release cache slots, and clear task queues.""" for task in self._task_mgr.get_active_tasks(): self._task_mgr.invoke_callback(task.task_id, STOP) self._task_cache.task_free(task.task_id) for task in self._task_mgr.get_waiting_tasks(): self._task_mgr.invoke_callback(task.task_id, STOP) - self._task_cache.task_free(task.task_id) + if free_waiting: + self._task_cache.task_free(task.task_id) self._task_mgr.clear_queues() - if torch.cuda.is_available(): - torch.cuda.empty_cache() def run_batch( self, diff --git a/astrai/inference/workspace.py b/astrai/inference/workspace.py index eaf4287..6e4cc76 100644 --- a/astrai/inference/workspace.py +++ b/astrai/inference/workspace.py @@ -124,13 +124,6 @@ class InferenceWorkspace: device=device, ) - def decode_buffers(self, batch: int, q_heads: int): - """Return ``(o_part, ml_part)`` view sliced to live dimensions.""" - return ( - self.decode_o_part[:batch, :q_heads], - self.decode_ml_part[:batch, :q_heads], - ) - def fill_input_ids(self, ids: "list[int]") -> Tensor: """Write ``ids`` into the device buffer and return ``[B]``. diff --git a/astrai/model/automodel.py b/astrai/model/automodel.py index 6d79a3e..99b1d7d 100644 --- a/astrai/model/automodel.py +++ b/astrai/model/automodel.py @@ -4,7 +4,7 @@ AutoModel base class for model loading and saving. from contextlib import contextmanager from pathlib import Path -from typing import Self, Union +from typing import Union import torch.nn as nn @@ -90,7 +90,3 @@ class AutoModel(nn.Module): state_dict=self.state_dict(), save_directory=str(save_directory), ) - - def to(self, *args, **kwargs) -> Self: - """Move model to device/dtype.""" - return super().to(*args, **kwargs) diff --git a/astrai/trainer/schedule.py b/astrai/trainer/schedule.py index c9db9bd..d0e373d 100644 --- a/astrai/trainer/schedule.py +++ b/astrai/trainer/schedule.py @@ -2,7 +2,7 @@ import math from abc import ABC, abstractmethod -from typing import Any, Dict, List +from typing import List from torch.optim.lr_scheduler import LRScheduler @@ -20,12 +20,6 @@ class BaseScheduler(LRScheduler, ABC): """Calculate the current learning rate.""" raise NotImplementedError - def state_dict(self) -> Dict[str, Any]: - return super().state_dict() - - def load_state_dict(self, state_dict: Dict[str, Any]): - super().load_state_dict(state_dict) - class SchedulerFactory(BaseFactory["BaseScheduler"]): """Factory class for creating learning rate schedulers. diff --git a/tests/inference/test_cache.py b/tests/inference/test_cache.py index 7435026..6826ef6 100644 --- a/tests/inference/test_cache.py +++ b/tests/inference/test_cache.py @@ -201,23 +201,6 @@ def test_req_to_token_pool_write(): # ---- KVStorage ---- -def test_kv_storage_set_and_get(): - storage = KVStorage( - size=16, - n_layers=2, - n_kv_heads=4, - head_dim=8, - device=torch.device("cpu"), - dtype=torch.float32, - ) - loc = torch.tensor([[0, 1]], dtype=torch.long) - k = torch.randn(1, 2, 4, 8) - v = torch.randn(1, 2, 4, 8) - storage.set_kv_buffer(0, loc, k, v) - assert torch.allclose(storage.get_key_buffer(0)[loc], k) - assert torch.allclose(storage.get_value_buffer(0)[loc], v) - - def test_kv_storage_buffer_shape(): storage = KVStorage( size=32,