diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py index c5b7fbe..b436c43 100644 --- a/astrai/inference/core/executor.py +++ b/astrai/inference/core/executor.py @@ -85,9 +85,65 @@ class Executor: dtype=self.dtype, ) - def execute_prefill(self, tasks: List[Task], prompt_len: int, start_pos: int = 0): + def _sample_logits( + self, + logits: Tensor, + tasks: List[Task], + return_logprobs: bool = False, + info: Optional[SamplingBatchInfo] = None, + ): + info = info or _build_sampling_batch_info(tasks, self.device) + if info.has_freq: + history_lists = [ + t.prompt_ids[-t.rep_window :] + t.output_ids for t in tasks + ] + history_lens = [len(ids) for ids in history_lists] + max_len = max(history_lens, default=0) + padded_ids = torch.zeros( + len(tasks), max_len, dtype=torch.long, device=self.device + ) + padded_mask = torch.zeros( + len(tasks), max_len, dtype=torch.bool, device=self.device + ) + for i, ids in enumerate(history_lists): + length = len(ids) + padded_ids[i, :length] = torch.as_tensor( + ids, dtype=torch.long, device=self.device + ) + padded_mask[i, :length] = True + else: + padded_ids = None + padded_mask = None + + result = sample( + logits, + temperature=info.temperatures, + top_k=info.top_ks, + top_p=info.top_ps, + frequency_penalty=info.freq_penalties, + input_ids=padded_ids, + input_mask=padded_mask, + return_logprobs=return_logprobs, + ) + if not return_logprobs: + return result.tolist() + + tokens, logprobs = result + tokens_list = tokens.tolist() + logprobs_list = logprobs.tolist() + for task, logprob in zip(tasks, logprobs_list): + task.output_logprobs.append(float(logprob)) + return list(zip(tokens_list, logprobs_list)) + + def execute_prefill( + self, + tasks: List[Task], + prompt_len: int, + start_pos: int = 0, + return_logprobs: bool = False, + ): if start_pos >= prompt_len: - return + return [] tasks = sorted(tasks, key=lambda t: t.task_id) batch_sz = len(tasks) @@ -109,7 +165,7 @@ class Executor: ) with torch.inference_mode(): - self.model( + outputs = self.model( input_ids, input_mask=input_mask, position_ids=position_ids, @@ -119,6 +175,9 @@ class Executor: start_pos=start_pos, ), ) + logits = outputs["logits"][:, -1, :] + + return tasks, self._sample_logits(logits, tasks, return_logprobs) def execute_decode( self, tasks: List[Task], return_logprobs: bool = False @@ -167,34 +226,6 @@ class Executor: total_len = max(t.next_pos for t in tasks) + 1 input_mask = self._workspace.decode_mask(position_ids, total_len) - has_freq = info.has_freq - if has_freq: - history_lists = [] - history_lens = [] - for t in tasks: - window = t.rep_window - prompt_part = t.prompt_ids[-window:] - ids = prompt_part + t.output_ids - history_lists.append(ids) - history_lens.append(len(ids)) - - max_len = max(history_lens) if history_lens else 0 - padded_ids = torch.zeros( - len(tasks), max_len, dtype=torch.long, device=self.device - ) - padded_mask = torch.zeros( - len(tasks), max_len, dtype=torch.bool, device=self.device - ) - for i, h in enumerate(history_lists): - L = history_lens[i] - padded_ids[i, :L] = torch.as_tensor( - h, dtype=torch.long, device=self.device - ) - padded_mask[i, :L] = True - else: - padded_ids = None - padded_mask = None - with torch.inference_mode(): outputs = self.model( input_ids, @@ -207,29 +238,4 @@ class Executor: ) logits = outputs["logits"][:, -1, :] - if return_logprobs: - tokens, logprobs = sample( - logits, - temperature=info.temperatures, - top_k=info.top_ks, - top_p=info.top_ps, - frequency_penalty=info.freq_penalties, - input_ids=padded_ids, - input_mask=padded_mask, - return_logprobs=True, - ) - tokens_list = tokens.tolist() - logprobs_list = logprobs.tolist() - for t, lp in zip(tasks, logprobs_list): - t.output_logprobs.append(float(lp)) - return list(zip(tokens_list, logprobs_list)) - - return sample( - logits, - temperature=info.temperatures, - top_k=info.top_ks, - top_p=info.top_ps, - frequency_penalty=info.freq_penalties, - input_ids=padded_ids, - input_mask=padded_mask, - ).tolist() + return self._sample_logits(logits, tasks, return_logprobs, info=info) diff --git a/astrai/inference/core/scheduler.py b/astrai/inference/core/scheduler.py index aa837fe..ee323c8 100644 --- a/astrai/inference/core/scheduler.py +++ b/astrai/inference/core/scheduler.py @@ -91,11 +91,10 @@ class InferenceScheduler: Single shared primitive for both the continuous-batching loop and the synchronous ``run_batch`` path, so the two cannot drift. - Tasks must already be allocated in the KV cache. Any task that still - needs prefill (``output_tokens == 0`` and fewer cached pages than - prompt tokens) is prefilled first, grouped by ``(prompt_len, cached)`` - so a ragged batch is padded into equal-length groups. Every task is - then extended by one token and decoded. + Tasks must already be allocated in the KV cache. Tasks without output + are prefilled first and sample their first token from the final prompt + position. Tasks with output extend the cache by one position and decode + from their latest generated token. Args: tasks: Active tasks to advance by one token. @@ -109,23 +108,27 @@ class InferenceScheduler: """ cache = self._cache - to_prefill = [ - t - for t in tasks - if t.output_tokens == 0 and cache.task_cached(t.task_id) < len(t.prompt_ids) - ] + to_prefill = [t for t in tasks if t.output_tokens == 0 and t.prompt_ids] + prefilled_ids = set() + produced: List[Task] = [] if to_prefill: for t in to_prefill: t.input_tokens = len(t.prompt_ids) groups: Dict[Tuple[int, int], List[Task]] = {} for t in to_prefill: - groups.setdefault( - (len(t.prompt_ids), cache.task_cached(t.task_id)), [] - ).append(t) + start_pos = min(cache.task_cached(t.task_id), len(t.prompt_ids) - 1) + groups.setdefault((len(t.prompt_ids), start_pos), []).append(t) for (prompt_len, start_pos), group in groups.items(): - self._executor.execute_prefill(group, prompt_len, start_pos) + prefilled, step_out = self._executor.execute_prefill( + group, prompt_len, start_pos, return_logprobs=return_logprobs + ) + for t, out in zip(prefilled, step_out): + t.output_ids.append(out[0] if return_logprobs else out) + t.output_tokens += 1 + prefilled_ids.add(t.task_id) + produced.append(t) start_logical_page = start_pos // getattr(cache, "page_size", 64) for t in group: cache.task_record_hashes( @@ -135,6 +138,8 @@ class InferenceScheduler: decoded: List[Task] = [] aborted: List[Task] = [] for t in tasks: + if t.task_id in prefilled_ids: + continue if cache.task_extend(t.task_id, t.next_pos): decoded.append(t) else: @@ -148,8 +153,9 @@ class InferenceScheduler: for t, out in zip(decoded, step_out): t.output_ids.append(out[0] if return_logprobs else out) t.output_tokens += 1 + produced.append(t) - return decoded, aborted + return produced, aborted def _run_generation_loop(self): stop_ids = self._task_mgr.tokenizer.stop_ids diff --git a/astrai/inference/core/task.py b/astrai/inference/core/task.py index 9b2ec0a..c3ba3a3 100644 --- a/astrai/inference/core/task.py +++ b/astrai/inference/core/task.py @@ -105,7 +105,8 @@ class Task: @property def next_pos(self) -> int: - return self.input_tokens + len(self.output_ids) + # The first output is sampled from prefill and enters KV on the next step. + return self.input_tokens + max(0, len(self.output_ids) - 1) def is_finished(self, stop_ids: List[int]) -> bool: if self.max_tokens is not None and self.output_tokens >= self.max_tokens: diff --git a/astrai/preprocessing/builder.py b/astrai/preprocessing/builder.py index 5d5778a..74099f4 100644 --- a/astrai/preprocessing/builder.py +++ b/astrai/preprocessing/builder.py @@ -416,7 +416,11 @@ class MultiOutputMaskBuilder(BaseMaskBuilder): return None result: dict = {} - any_output = False + required_outputs = { + output_key + for output_key, spec in sources_spec.items() + if spec.get("sections") + } for output_key, spec in sources_spec.items(): sections = spec.get("sections", []) @@ -428,7 +432,6 @@ class MultiOutputMaskBuilder(BaseMaskBuilder): if ids is None: continue result[output_key] = ids - any_output = True continue list_field = spec.get("list_field", False) @@ -444,7 +447,6 @@ class MultiOutputMaskBuilder(BaseMaskBuilder): result[output_key] = ids if mask is not None: result[mask_key] = mask - any_output = True continue ids, mask = self.renderer.process_sections( @@ -460,9 +462,7 @@ class MultiOutputMaskBuilder(BaseMaskBuilder): elif "mask_key" in spec: result[mask_key] = mask - any_output = True - - if not any_output: + if not required_outputs or not required_outputs.issubset(result): return None result["domain"] = _extract_domain(item, config.output.domain_key) @@ -474,6 +474,11 @@ class MultiOutputMaskBuilder(BaseMaskBuilder): return [None] * len(items) results = [{} for _ in items] + required_outputs = { + output_key + for output_key, spec in sources_spec.items() + if spec.get("sections") + } for output_key, spec in sources_spec.items(): sections = spec.get("sections", []) if not sections: @@ -506,7 +511,7 @@ class MultiOutputMaskBuilder(BaseMaskBuilder): return [ ({**result, "domain": _extract_domain(item, config.output.domain_key)}) - if result + if required_outputs and required_outputs.issubset(result) else None for item, result in zip(items, results) ] diff --git a/tests/data/test_preprocess_builder.py b/tests/data/test_preprocess_builder.py index 4dfe327..3d0e83c 100644 --- a/tests/data/test_preprocess_builder.py +++ b/tests/data/test_preprocess_builder.py @@ -369,6 +369,19 @@ def test_dpo_missing_field_is_none(chat_tokenizer, builder): assert builder.build({"chosen": [], "rejected": []}, config, chat_tokenizer) is None +@pytest.mark.parametrize("missing", ["chosen", "rejected"]) +def test_dpo_partial_record_is_none(chat_tokenizer, builder, missing): + config = make_dpo_chat_config() + item = { + "chosen": [{"role": "assistant", "content": "Good"}], + "rejected": [{"role": "assistant", "content": "Bad"}], + } + item.pop(missing) + + assert builder.build(item, config, chat_tokenizer) is None + assert builder.build_batch([item], config, chat_tokenizer) == [None] + + def test_grpo_basic(chat_tokenizer, builder): config = make_grpo_config() item = { diff --git a/tests/inference/test_scheduler.py b/tests/inference/test_scheduler.py index 8edb6cb..dde02ff 100644 --- a/tests/inference/test_scheduler.py +++ b/tests/inference/test_scheduler.py @@ -205,6 +205,36 @@ def test_run_batch_returns_token_sequences(device): scheduler.stop() +def test_run_batch_tokens_match_full_sequence_forward(device): + scheduler, _tok, model = _make_real_scheduler(device) + prompt = [10, 20, 30, 40] + try: + expected = [] + sequence = list(prompt) + for _ in range(2): + input_ids = torch.tensor([sequence], dtype=torch.long, device=device) + position_ids = torch.arange(len(sequence), device=device).unsqueeze(0) + input_mask = torch.ones( + 1, len(sequence), len(sequence), dtype=torch.bool, device=device + ).tril() + with torch.inference_mode(): + logits = model( + input_ids, + input_mask=input_mask, + position_ids=position_ids, + )["logits"][:, -1, :] + token = logits.argmax(dim=-1).item() + expected.append(token) + sequence.append(token) + + result = scheduler.run_batch( + prompt_ids_list=[prompt], max_tokens=2, temperature=0 + ) + assert result == [expected] + finally: + scheduler.stop() + + def test_run_batch_return_logprobs_aligned(device): """return_logprobs=True gives (token_ids, logprobs) tuples with equal len.""" scheduler, _tok, _model = _make_real_scheduler(device) diff --git a/tests/inference/test_task.py b/tests/inference/test_task.py index b74886f..65de9d6 100644 --- a/tests/inference/test_task.py +++ b/tests/inference/test_task.py @@ -22,6 +22,8 @@ def test_task_next_pos(): task.input_tokens = 5 assert task.next_pos == 5 task.output_ids.append(4) + assert task.next_pos == 5 + task.output_ids.append(5) assert task.next_pos == 6