File size: 7,989 Bytes
e791b16 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 | # Rollout engine implementations: StubEngine, HFEngine, and VLLMEngineWrapper.
# Supports CPU and CUDA GPU execution, pure-Python async testing,
# and high-performance tensor-native execution on NVIDIA GPUs.
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:
# Fallback to lightweight linear language model for testing/prototyping
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,
)
|