| |
| |
| |
|
|
| import asyncio |
| import random |
| from dataclasses import dataclass |
| from typing import AsyncIterator, List, Optional, Tuple, Union |
|
|
| import torch |
|
|
|
|
| @dataclass |
| class GenerationResult: |
| """Result of an asynchronous generation request.""" |
| input_ids: torch.Tensor |
| ids: torch.Tensor |
| logprobs: torch.Tensor |
| gt_ids: Optional[torch.Tensor] = None |
| eos_token_id: int = 99 |
|
|
|
|
| class StubEngine: |
| """Minimal async generation engine for testing without vLLM or GPU. |
| |
| Generates tokens by sampling uniformly from [0, vocab_size) until EOS |
| or max_new_tokens is reached. Uses asyncio.sleep(0.001) to force real |
| coroutine yielding so that concurrent callers actually interleave. |
| """ |
|
|
| def __init__(self, vocab_size: int = 100, max_new_tokens: int = 8, eos_token_id: int = 99): |
| self.vocab_size = vocab_size |
| self.max_new_tokens = max_new_tokens |
| self.eos_token_id = eos_token_id |
|
|
| async def async_generate( |
| self, prompt_ids: torch.Tensor |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| """Generate tokens asynchronously for a single prompt.""" |
| await asyncio.sleep(0.001) |
|
|
| generated_ids = [] |
| log_probs = [] |
|
|
| device = prompt_ids.device if isinstance(prompt_ids, torch.Tensor) else torch.device("cpu") |
| uniform_log_prob = -torch.log(torch.tensor(float(self.vocab_size), device=device)).item() |
|
|
| for _ in range(self.max_new_tokens): |
| token_id = random.randint(0, self.vocab_size - 1) |
| generated_ids.append(token_id) |
| log_probs.append(uniform_log_prob) |
| if token_id == self.eos_token_id: |
| break |
|
|
| return ( |
| torch.tensor(generated_ids, dtype=torch.long, device=device), |
| torch.tensor(log_probs, dtype=torch.float32, device=device), |
| ) |
|
|
|
|
| class HFEngine: |
| """Async generation engine using PyTorch / Transformers on CUDA or CPU. |
| |
| Runs autoregressive generation directly on the target GPU without CPU roundtrips, |
| extracting exact log-probabilities for RLHF/GRPO. |
| """ |
|
|
| def __init__( |
| self, |
| model_or_path: Union[str, torch.nn.Module], |
| tokenizer=None, |
| device: Optional[str] = None, |
| max_new_tokens: int = 32, |
| eos_token_id: Optional[int] = None, |
| vocab_size: int = 1000, |
| ): |
| self.device = torch.device( |
| device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu") |
| ) |
| self.max_new_tokens = max_new_tokens |
| self.tokenizer = tokenizer |
| self.eos_token_id = eos_token_id if eos_token_id is not None else 99 |
| self.vocab_size = vocab_size |
|
|
| if isinstance(model_or_path, torch.nn.Module): |
| self.model = model_or_path.to(self.device) |
| self.model.eval() |
| elif isinstance(model_or_path, str) and model_or_path: |
| try: |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| self.model = AutoModelForCausalLM.from_pretrained(model_or_path).to(self.device) |
| self.model.eval() |
| if self.tokenizer is None: |
| self.tokenizer = AutoTokenizer.from_pretrained(model_or_path) |
| if hasattr(self.model.config, "vocab_size"): |
| self.vocab_size = self.model.config.vocab_size |
| if hasattr(self.model.config, "eos_token_id") and self.model.config.eos_token_id is not None: |
| self.eos_token_id = self.model.config.eos_token_id |
| except Exception: |
| |
| self.model = torch.nn.Sequential( |
| torch.nn.Embedding(self.vocab_size, 128), |
| torch.nn.Linear(128, self.vocab_size), |
| ).to(self.device) |
| self.model.eval() |
| else: |
| self.model = torch.nn.Sequential( |
| torch.nn.Embedding(self.vocab_size, 128), |
| torch.nn.Linear(128, self.vocab_size), |
| ).to(self.device) |
| self.model.eval() |
|
|
| async def async_generate( |
| self, prompt_ids: torch.Tensor |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| """Asynchronously generate tokens on GPU and return (generated_ids, logprobs).""" |
| await asyncio.sleep(0.0001) |
|
|
| input_ids = prompt_ids.to(self.device) |
| if input_ids.ndim == 1: |
| curr_tokens = input_ids.clone() |
| else: |
| curr_tokens = input_ids[0].clone() |
|
|
| generated_ids = [] |
| log_probs = [] |
|
|
| with torch.no_grad(): |
| for _ in range(self.max_new_tokens): |
| logits = self.model(curr_tokens.unsqueeze(0)) |
| if hasattr(logits, "logits"): |
| logits = logits.logits |
| next_token_logits = logits[0, -1, :] |
| probs = torch.softmax(next_token_logits, dim=-1) |
| log_p = torch.log_softmax(next_token_logits, dim=-1) |
|
|
| next_token = torch.multinomial(probs, num_samples=1) |
| token_val = next_token.item() |
| token_logp = log_p[token_val].item() |
|
|
| generated_ids.append(token_val) |
| log_probs.append(token_logp) |
| curr_tokens = torch.cat([curr_tokens, next_token]) |
|
|
| if token_val == self.eos_token_id: |
| break |
|
|
| return ( |
| torch.tensor(generated_ids, dtype=torch.long, device=self.device), |
| torch.tensor(log_probs, dtype=torch.float32, device=self.device), |
| ) |
|
|
|
|
| class VLLMEngineWrapper: |
| """Wrapper that supports vLLM if installed, or falls back to HFEngine or StubEngine.""" |
|
|
| def __init__(self, model_path: Optional[str] = None, engine=None, **engine_kwargs): |
| self.model_path = model_path |
| self.engine = engine |
|
|
| if self.engine is None: |
| vllm_available = False |
| try: |
| import vllm |
| vllm_available = True |
| except ImportError: |
| vllm_available = False |
|
|
| if vllm_available and model_path: |
| from vllm import AsyncEngineArgs, AsyncLLMEngine |
| args = AsyncEngineArgs(model=model_path, **engine_kwargs) |
| self.engine = AsyncLLMEngine.from_engine_args(args) |
| self.is_vllm = True |
| else: |
| self.is_vllm = False |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| if model_path and model_path != "stub": |
| self.engine = HFEngine(model_path, device=device, **engine_kwargs) |
| else: |
| self.engine = StubEngine(**engine_kwargs) if engine_kwargs else StubEngine() |
| else: |
| self.is_vllm = False |
|
|
| async def async_generate( |
| self, prompts: Union[List[dict], torch.Tensor] |
| ) -> AsyncIterator[GenerationResult]: |
| """Stream generated results for a list of prompt dicts or a single prompt.""" |
| if isinstance(prompts, torch.Tensor): |
| prompts = [{"input_ids": prompts}] |
|
|
| for p in prompts: |
| input_ids = p.get("input_ids") |
| gt_ids = p.get("gt_ids", None) |
| eos_token_id = p.get("eos_token_id", 99) |
|
|
| if isinstance(input_ids, list): |
| input_ids = torch.tensor(input_ids) |
|
|
| if input_ids is None: |
| input_ids = torch.tensor([1, 2, 3]) |
|
|
| ids, logprobs = await self.engine.async_generate(input_ids) |
| yield GenerationResult( |
| input_ids=input_ids, |
| ids=ids, |
| logprobs=logprobs, |
| gt_ids=gt_ids, |
| eos_token_id=eos_token_id, |
| ) |
|
|
