diff --git a/arealite/README.md b/arealite/README.md index baa58b9f12..d70b86ebca 100644 --- a/arealite/README.md +++ b/arealite/README.md @@ -83,7 +83,6 @@ def main_grpo(): # or asynchronous rollout with filtering and off-policyness control # rollout_batch = rollout.prepare_batch(batch, # workflow=MyRolloutWorkflow(rollout_config.workflow), - # offpolicyness=4, # should_accept=lambda x: x['rewards'].mean() > 0) # In the single-controller mode diff --git a/arealite/api/cli_args.py b/arealite/api/cli_args.py index 39a471e963..65ff7cbd92 100644 --- a/arealite/api/cli_args.py +++ b/arealite/api/cli_args.py @@ -1,5 +1,5 @@ from dataclasses import asdict, dataclass, field -from typing import List +from typing import Dict, List, Optional @dataclass @@ -70,8 +70,160 @@ def new(self, **kwargs): return GenerationHyperparameters(**args) +@dataclass +class SGLangConfig: + """Configuration for SGLang runtime. Refer to: + https://github.com/sgl-project/sglang for detailed documentation. + """ + + disable_cuda_graph: bool = False + disable_radix_cache: bool = False + disable_cuda_graph_padding: bool = False + enable_nccl_nvls: bool = False + disable_outlines_disk_cache: bool = False + disable_custom_all_reduce: bool = False + disable_overlap_schedule: bool = False + enable_mixed_chunk: bool = False + enable_dp_attention: bool = False + enable_ep_moe: bool = False + enable_torch_compile: bool = False + torch_compile_max_bs: int = 32 + cuda_graph_max_bs: Optional[int] = None + cuda_graph_bs: Optional[List[int]] = None + torchao_config: str = "" + enable_nan_detection: bool = False + enable_p2p_check: bool = False + triton_attention_reduce_in_fp32: bool = False + triton_attention_num_kv_splits: int = 8 + num_continuous_decode_steps: int = 1 + enable_memory_saver: bool = False + allow_auto_truncate: bool = False + # NOTE: to avoid the illegal memory access error + attention_backend: Optional[str] = "flashinfer" + sampling_backend: Optional[str] = None + context_length: Optional[int] = 32768 + mem_fraction_static: Optional[float] = 0.9 + max_running_requests: Optional[int] = None + # NOTE: chunked_prefill_size is by default 8192 on GPUs with 80GB mem in SGLang, + # but we disable it to avoid precision issues + chunked_prefill_size: Optional[int] = -1 + max_prefill_tokens: int = 32768 + schedule_policy: str = "lpm" + schedule_conservativeness: float = 1.0 + cpu_offload_gb: int = 0 + + dtype: str = "float16" + kv_cache_dtype: str = "auto" + + # logging + log_level: str = "warning" + log_level_http: Optional[str] = "warning" + log_requests: bool = False + log_requests_level: int = 0 + show_time_cost: bool = False + enable_metrics: bool = True # Exports Prometheus-like metrics + # The interval (in decoding iterations) to log throughput + # and update prometheus metrics + decode_log_interval: int = 1 + + # Use staticmethod to make OmegaConf happy. + @staticmethod + def build_cmd( + sglang_config: "SGLangConfig", + model_path, + tp_size, + base_gpu_id, + dist_init_addr: Optional[str] = None, + served_model_name: Optional[str] = None, + skip_tokenizer_init: bool = True, + ): + from realhf.base import network, pkg_version, seeding + from realhf.experiments.common.utils import asdict as conf_as_dict + + args: Dict = conf_as_dict(sglang_config) + args["random_seed"] = seeding.get_seed() + + if served_model_name is None: + served_model_name = model_path + host_ip = network.gethostip() + host = "localhost" if not sglang_config.enable_metrics else host_ip + args = dict( + host=host, + model_path=model_path, + # Model and tokenizer + tokenizer_path=model_path, + tokenizer_mode="auto", + load_format="auto", + trust_remote_code=True, + device="cuda", + served_model_name=served_model_name, + is_embedding=False, + skip_tokenizer_init=skip_tokenizer_init, + # Other runtime options + tp_size=tp_size, + # Because we have set CUDA_VISIBLE_DEVICES to a single GPU in each process + base_gpu_id=base_gpu_id, + nnodes=1, + node_rank=0, + dist_init_addr=dist_init_addr, + **args, + ) + + if pkg_version.is_version_less("sglang", "0.4.4"): + args.pop("log_requests_level") + if pkg_version.is_version_less("sglang", "0.4.3"): + args.pop("enable_nccl_nvls") + args.pop("triton_attention_num_kv_splits") + args.pop("cuda_graph_bs") + args.pop("enable_memory_saver") + args.pop("allow_auto_truncate") + args.pop("file_storage_path") + + flags = [] + for k, v in args.items(): + if v is None or v is False or v == "": + continue + if v is True: + flags.append(f"--{k.replace('_','-')} ") + continue + if isinstance(v, list): + values = " ".join(map(str, v)) + flags.append(f"--{k.replace('_','-')} {values}") + continue + flags.append(f"--{k.replace('_','-')} {v}") + flags = " ".join(flags) + return f"python3 -m sglang.launch_server {flags}" + + @dataclass class InferenceEngineConfig: + experiment_name: str + trial_name: str + max_concurrent_rollouts: None | int = field( + default=None, + metadata={ + "help": "Maximum number of concurrent requests to the inference engine." + }, + ) + queue_size: None | int = field( + default=None, + metadata={"help": "Input/Output queue size for async rollout."}, + ) + consumer_batch_size: int = field( + default=1, + metadata={ + "help": "Batch size for consuming requests from the queue. " + "If None, it will be set to max_concurrent_requests.", + }, + ) + max_head_offpolicyness: int = field( + default=0, + metadata={ + "help": "Maximum off-policyness for the head. " + "If the current version is more than this many versions behind, " + "the request will not be accepted.", + }, + ) # Used by remote inference engines. server_addrs: List[str] = field( default_factory=list, diff --git a/arealite/api/engine_api.py b/arealite/api/engine_api.py index 7badfce7d7..c9a4b7ede3 100644 --- a/arealite/api/engine_api.py +++ b/arealite/api/engine_api.py @@ -109,6 +109,10 @@ def initialize(self, addr: str | None, ft_spec): """Initialize environments for distributed inference and load models.""" raise NotImplementedError() + def destroy(self): + """Destroy the engine and release GPU memory.""" + pass + def update_weights(self, meta: WeightUpdateMeta) -> Future: """Update weights in the inference engine.""" raise NotImplementedError() @@ -121,7 +125,7 @@ def submit(self, data: Dict[str, Any], workflow: "RolloutWorkflow") -> None: """Asynchronously submit a request to the inference engine. Exits immediately.""" raise NotImplementedError() - def wait(self, count: int, timeout: int) -> TensorDict: + def wait(self, count: int, timeout: float) -> TensorDict: """Wait for a specified number of requests to complete, with a timeout.""" raise NotImplementedError() diff --git a/arealite/api/workflow_api.py b/arealite/api/workflow_api.py index d476ed45b3..9141399c1d 100644 --- a/arealite/api/workflow_api.py +++ b/arealite/api/workflow_api.py @@ -9,7 +9,7 @@ class RolloutWorkflow: async def arun_episode( - self, engine: InferenceEngine, data: Dict[str, Any] + self, engine: "InferenceEngine", data: Dict[str, Any] ) -> TensorDict: """Run a single episode of the workflow. diff --git a/arealite/engine/sglang_remote.py b/arealite/engine/sglang_remote.py index 453f659ee1..036b4781eb 100644 --- a/arealite/engine/sglang_remote.py +++ b/arealite/engine/sglang_remote.py @@ -3,8 +3,8 @@ import time import traceback from concurrent.futures import ThreadPoolExecutor -from queue import Empty, Queue -from typing import TYPE_CHECKING, Any, Callable, Dict, Optional +from queue import Empty, Full, Queue +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional import aiohttp import torch.distributed as dist @@ -18,7 +18,7 @@ RolloutStat, WeightUpdateMeta, ) -from realhf.base import logging, name_resolve, names, pkg_version +from realhf.base import logging, pkg_version if TYPE_CHECKING: from arealite.api.workflow_api import RolloutWorkflow @@ -31,19 +31,27 @@ SGLANG_TOKEN_OUTPUT_IDENTIFIER = "token_ids" ROLLOUT_POLL_WAIT_TIME = 0.4 +RID_CACHE_SIZE = 128 class RemoteSGLangEngine(InferenceEngine): def __init__(self, config: InferenceEngineConfig): + config.max_concurrent_rollouts = ( + config.max_concurrent_rollouts or config.consumer_batch_size + ) self.config = config self.rid_to_address = {} + # Maintain the addresses for the recent 128 requests + self.rid_queue = [] + self.addresses = config.server_addrs self.server_idx = 0 - self.input_queue = Queue(maxsize=config.max_concurrent_rollouts) - self.output_queue = Queue(maxsize=config.max_concurrent_rollouts) + qsize = config.queue_size or config.max_concurrent_rollouts * 10 + self.input_queue = Queue(maxsize=qsize) + self.output_queue = Queue(maxsize=qsize) self.result_cache = [] self.exiting = threading.Event() @@ -51,32 +59,35 @@ def __init__(self, config: InferenceEngineConfig): self.rollout_stat = RolloutStat() - def _get_model_version(self) -> int: - name = names.model_version( - self.config.experiment_name, - self.config.trial_name, - "actor", - ) - try: - return int(name_resolve.get(name)) - except name_resolve.NameEntryNotFoundError: - return 0 + self._version = 0 def initialize(self, addr: str | None, ft_spec: Optional[Dict[str, Any]] = None): self.rollout_thread = threading.Thread(target=self._rollout_thread) self.rollout_thread.start() + def destroy(self): + self.exiting.set() + self.rollout_thread.join() + + def set_version(self, version): + with self.lock: + self._version = version + + def get_version(self): + with self.lock: + return self._version + def _rollout_thread(self): """Thread that runs the rollout loop.""" try: - asyncio.run_coroutine_threadsafe(self._rollout_thread_async()) - finally: - self.exiting.set() + asyncio.run(self._rollout_thread_async()) + except Exception as e: + traceback.print_exc() async def _rollout_thread_async(self): data = None - rollout_tasks: Dict[int, asyncio.Task] = {} + rollout_tasks: Dict[str, asyncio.Task] = {} rid = 0 try: @@ -85,7 +96,7 @@ async def _rollout_thread_async(self): if data is None: try: data, workflow = self.input_queue.get_nowait() - logger.debug(f"Get data from puller: {data}") + logger.info(f"Get data from puller: {data}") except Empty: logger.debug(f"No data from puller stream.") @@ -104,17 +115,17 @@ async def _rollout_thread_async(self): ) # Staleness control - version = self._get_model_version() + version = self.get_version() ofp = self.config.max_head_offpolicyness with self.lock: sample_cnt = self.rollout_stat.accepted + self.rollout_stat.running - expected_version = sample_cnt // self.train_batch_size + expected_version = sample_cnt // self.config.consumer_batch_size not_staled = expected_version <= ofp + version can_rollout &= not_staled if not not_staled: cannot_rollout_reason.append( f"Staled: expected version ({expected_version}) = " - f"global sample cnt ({sample_cnt}) // batch size ({self.train_batch_size}), " + f"global sample cnt ({sample_cnt}) // batch size ({self.config.consumer_batch_size}), " f"current latest version {version}, " f"offpolicyness {self.config.max_head_offpolicyness}." ) @@ -130,12 +141,12 @@ async def _rollout_thread_async(self): task = asyncio.create_task( workflow.arun_episode(self, data), name=str(rid) ) - rollout_tasks[rid] = task + rollout_tasks[str(rid)] = task with self.lock: self.rollout_stat.submitted += 1 self.rollout_stat.running += 1 - logger.debug( + logger.info( f"Submit rollout rid {rid}. " f"Submit: {self.rollout_stat.submitted}, " f"running: {self.rollout_stat.running}, " @@ -163,12 +174,18 @@ async def _rollout_thread_async(self): traj: TensorDict task_rid = task.get_name() rollout_tasks.pop(task_rid) + self.rollout_stat.accepted += 1 - self.output_queue.put(traj) + try: + self.output_queue.put_nowait(traj) + except Full: + raise RuntimeError( + "Output queue full. Please increase queue_size." + ) with self.lock: self.rollout_stat.running -= 1 - logger.debug( + logger.info( f"Finish rollout {task_rid}. " f"Submit: {self.rollout_stat.submitted}, " f"running: {self.rollout_stat.running}, " @@ -256,7 +273,11 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: gconfig = req.gconfig stop_token_ids = gconfig.stop_token_ids - assert gconfig.n_samples == 1 + if gconfig.n_samples != 1: + raise ValueError( + "RemoteSGLangEngine does not support n_samples > 1. " + "Please call generate for multiple times with n_samples = 1." + ) sample_params = { "top_p": gconfig.top_p, "top_k": gconfig.top_k, @@ -265,8 +286,8 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: "stop_token_ids": stop_token_ids, } + # NOTE: rid should NOT be passed in payload payload = { - "rid": req.rid, "text": req.text, "sampling_params": sample_params, "return_logprob": True, @@ -287,6 +308,17 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: completions = "" stop_reason = "length" + if req.rid in self.rid_to_address: + server_addr = self.rid_to_address[req.rid] + else: + server_addr = self.choose_server() + if len(self.rid_queue) >= RID_CACHE_SIZE: + # Remove the oldest entry if cache is full + oldest_rid = self.rid_queue.pop(0) + self.rid_to_address.pop(oldest_rid, None) + self.rid_to_address[req.rid] = server_addr + self.rid_queue.append(req.rid) + while ( stop_reason != "stop" and len(accumulated_output_tokens) < gconfig.max_new_tokens @@ -298,6 +330,7 @@ async def agenerate(self, req: LLMRequest) -> LLMResponse: method="POST", max_retries=3, timeout=self.config.request_timeout, + target_addr=server_addr, ) result = await response.json() @@ -369,9 +402,12 @@ async def aupdate_weights_from_disk(self, addr, path: str): ) def submit(self, data: Dict[str, Any], workflow: "RolloutWorkflow") -> None: - self.input_queue.put((workflow, data)) + try: + self.input_queue.put_nowait((data, workflow)) + except Full: + raise RuntimeError("Input queue full. Please increase queue_size.") - def wait(self, count: int, timeout: int, should_accept: Callable) -> TensorDict: + def wait(self, count: int, timeout: float, should_accept: Callable) -> TensorDict: tik = time.perf_counter() accepted = len(self.result_cache) while ( @@ -384,8 +420,9 @@ def wait(self, count: int, timeout: int, should_accept: Callable) -> TensorDict: if should_accept(result): self.result_cache.append(result) accepted += 1 + else: with self.lock: - self.rollout_stat.accepted += 1 + self.rollout_stat.accepted -= 1 except Empty: time.sleep(ROLLOUT_POLL_WAIT_TIME) if self.exiting.is_set(): @@ -399,3 +436,15 @@ def wait(self, count: int, timeout: int, should_accept: Callable) -> TensorDict: self.result_cache[count:], ) return TensorDict.cat(results, dim=0) + + def rollout( + self, data: List[Dict[str, Any]], workflow: "RolloutWorkflow" + ) -> TensorDict: + """Submit a batch of requests to the inference engine and wait for the results.""" + for item in data: + self.submit(item, workflow) + return self.wait( + count=len(data), + timeout=self.config.request_timeout, + should_accept=lambda x: True, + ) diff --git a/arealite/tests/test_sglang_engine.py b/arealite/tests/test_sglang_engine.py index e69de29bb2..493f33a7d3 100644 --- a/arealite/tests/test_sglang_engine.py +++ b/arealite/tests/test_sglang_engine.py @@ -0,0 +1,188 @@ +import os +import subprocess +import sys +import time +import uuid + +import pytest +import requests +import torch +from tensordict import TensorDict + +from arealite.api.cli_args import ( + GenerationHyperparameters, + InferenceEngineConfig, + SGLangConfig, +) +from arealite.api.io_struct import FinetuneSpec, LLMRequest, LLMResponse +from realhf.api.core.data_api import load_hf_tokenizer +from realhf.base import name_resolve, network, seeding + +EXPR_NAME = "test_sglang_engine" +TRIAL_NAME = "trial_0" +MODEL_PATH = "/storage/testing/models/Qwen__Qwen3-1.7B/" +if not os.path.exists(MODEL_PATH): + MODEL_PATH = "Qwen/Qwen2-0.5B" +PORT = 13887 +DIST_PORT = 15887 +HOST = network.gethostip() + + +def check_server_health(base_url): + # Check server endpoint + try: + response = requests.get( + f"{base_url}/metrics", + timeout=30, + ) + return response.status_code == 200 + except requests.exceptions.RequestException: + return False + + +@pytest.fixture(scope="module") +def sglang_server(): + from realhf.base import seeding + + seeding.set_random_seed(1, EXPR_NAME) + cmd = SGLangConfig.build_cmd( + sglang_config=SGLangConfig(mem_fraction_static=0.3), + model_path=MODEL_PATH, + tp_size=1, + base_gpu_id=0, + dist_init_addr=f"{HOST}:{DIST_PORT}", + served_model_name=MODEL_PATH, + skip_tokenizer_init=False, + ) + # Launch process + full_command = f"{cmd} --port {PORT}" + full_command = full_command.replace("\\\n", " ").replace("\\", " ") + process = subprocess.Popen( + full_command.split(), + text=True, + stdout=sys.stdout, + stderr=sys.stdout, + ) + base_url = f"http://{HOST}:{PORT}" + tik = time.time() + while time.time() - tik < 90: + if check_server_health(base_url): + break + time.sleep(1) + if time.time() - tik > 90: + raise RuntimeError("server launch failed") + yield + process.terminate() + + +@pytest.mark.skip("") +@pytest.mark.asyncio +async def test_remote_sglang_generate(sglang_server): + from arealite.engine.sglang_remote import RemoteSGLangEngine + + config = InferenceEngineConfig(experiment_name=EXPR_NAME, trial_name=TRIAL_NAME) + config.server_addrs = [f"{HOST}:{PORT}"] + engine = RemoteSGLangEngine(config) + req = LLMRequest( + rid=str(uuid.uuid4()), + text="hello! how are you today", + gconfig=GenerationHyperparameters(max_new_tokens=16), + ) + resp = await engine.agenerate(req) + assert isinstance(resp, LLMResponse) + assert resp.input_tokens == req.input_ids + assert ( + len(resp.output_logprobs) + == len(resp.output_tokens) + == len(resp.output_versions) + ) + assert isinstance(resp.completions, str) + + +@pytest.mark.skip("") +@pytest.mark.parametrize("n_samples", [1, 2, 4]) +def test_remote_sglang_rollout(sglang_server, n_samples): + from arealite.engine.sglang_remote import RemoteSGLangEngine + from arealite.workflow.rlvr import RLVRWorkflow + + config = InferenceEngineConfig( + experiment_name=EXPR_NAME, + trial_name=TRIAL_NAME, + max_concurrent_rollouts=2, + consumer_batch_size=2, + ) + config.server_addrs = [f"{HOST}:{PORT}"] + engine = RemoteSGLangEngine(config) + engine.initialize(None, None) + + gconfig = GenerationHyperparameters( + max_new_tokens=16, greedy=False, n_samples=n_samples + ) + tokenizer = load_hf_tokenizer(MODEL_PATH) + + workflow = RLVRWorkflow( + reward_fn=lambda **kwargs: 1.0, # Dummy reward function + gconfig=gconfig, + tokenizer=tokenizer, + ) + + data = { + "messages": [{"role": "user", "content": "Hello, how are you?"}], + } + result = engine.rollout([data] * 2, workflow=workflow) + assert isinstance(result, TensorDict) + bs = result.batch_size + assert bs == torch.Size([2 * n_samples]) + engine.destroy() + + +@pytest.mark.parametrize("ofp", [1, 2, 4, 8, 16]) +@pytest.mark.parametrize("bs", [2, 4]) +@pytest.mark.parametrize("n_samples", [2, 1]) +def test_remote_sglang_staleness_control(sglang_server, bs, ofp, n_samples): + from arealite.engine.sglang_remote import RemoteSGLangEngine + from arealite.workflow.rlvr import RLVRWorkflow + + config = InferenceEngineConfig( + experiment_name=EXPR_NAME, + trial_name=TRIAL_NAME, + consumer_batch_size=bs, + max_head_offpolicyness=ofp, + ) + config.server_addrs = [f"{HOST}:{PORT}"] + engine = RemoteSGLangEngine(config) + engine.initialize(None, None) + + gconfig = GenerationHyperparameters( + max_new_tokens=16, greedy=False, n_samples=n_samples + ) + tokenizer = load_hf_tokenizer(MODEL_PATH) + + workflow = RLVRWorkflow( + reward_fn=lambda **kwargs: 1.0, # Dummy reward function + gconfig=gconfig, + tokenizer=tokenizer, + ) + data = { + "messages": [{"role": "user", "content": "Hello, how are you?"}], + } + for _ in range(bs * 2): + engine.submit(data, workflow=workflow) + + # wait for some time + time.sleep(15) + assert engine.output_queue.qsize() == min(bs * 2, bs * (ofp + 1)) + + # Update model version + engine.set_version(1) + print("Updated model version", flush=True) + + # submit again + for _ in range(bs * 2): + engine.submit(data, workflow=workflow) + # wait for some time + time.sleep(15) + assert engine.output_queue.qsize() == min(bs * 4, bs * (ofp + 2)) + + # exit + engine.destroy() diff --git a/arealite/utils/padding.py b/arealite/utils/padding.py new file mode 100644 index 0000000000..ac12d60ac8 --- /dev/null +++ b/arealite/utils/padding.py @@ -0,0 +1,57 @@ +from typing import Any, Dict, List + +import torch +from tensordict import TensorDict + + +def concat_padded_tensors( + tensor_dicts: List[TensorDict], pad_value: float = 0.0 +) -> TensorDict: + """Concatenate and pad tensors from multiple padded tensor dictionaries.""" + if not tensor_dicts: + return TensorDict() + + batch_sizes = [tuple(d.batch_size) for d in tensor_dicts] + new_batch_size = [sum(x[0] for x in batch_sizes), *batch_sizes[0][1:]] + + # Find max sequence length across all dictionaries + lens = [] + for tensor_dict in tensor_dicts: + for key, tensor in tensor_dict.items(): + if key != "attention_mask" and len(tensor.shape) == 2: + lens.append(tensor.shape[1]) + break + max_length = max(lens) + attn_mask = torch.arange(max_length).unsqueeze(0) < torch.tensor(lens).unsqueeze(1) + + result = {} + # Process each key + for key in tensor_dicts[0].keys(): + tensors_to_concat = [] + for tensor_dict in tensor_dicts: + tensor = tensor_dict[key] + # Skip 1D tensors like rewards + if len(tensor.shape) == 1: + tensors_to_concat.append(tensor) + continue + current_length = tensor.shape[1] + if current_length < max_length: + # Pad tensor to max_length + pad_width = max_length - current_length + if key == "attention_mask": + # Pad attention mask with 0s + padding = torch.zeros( + (tensor.shape[0], pad_width), dtype=tensor.dtype + ) + else: + # Pad feature tensors with pad_value + padding = torch.full( + (tensor.shape[0], pad_width), pad_value, dtype=tensor.dtype + ) + tensor = torch.cat([tensor, padding], dim=1) + tensors_to_concat.append(tensor) + + result[key] = torch.cat(tensors_to_concat, dim=0) + if "attention_mask" not in result: + result["attention_mask"] = attn_mask + return TensorDict(result, batch_size=new_batch_size) diff --git a/arealite/workflow/rlvr.py b/arealite/workflow/rlvr.py index a260543a0f..5479f09f31 100644 --- a/arealite/workflow/rlvr.py +++ b/arealite/workflow/rlvr.py @@ -1,3 +1,4 @@ +import asyncio import uuid import torch @@ -7,6 +8,7 @@ from arealite.api.cli_args import GenerationHyperparameters from arealite.api.io_struct import LLMRequest from arealite.api.workflow_api import RolloutWorkflow +from arealite.utils.padding import concat_padded_tensors class RLVRWorkflow(RolloutWorkflow): @@ -24,33 +26,37 @@ async def arun_episode(self, engine, data): text = self.tokenizer.apply_chat_template( data["messages"], tokenize=False, add_generation_prompt=True ) + n_samples = self.gconfig.n_samples req = LLMRequest( rid=uuid.uuid4().hex, text=text, - gconfig=self.gconfig, + gconfig=self.gconfig.new(n_samples=1), ) - resp = await engine.agenerate(req) - - seq = resp.input_tokens + resp.output_tokens - logprobs = [0] * resp.input_len + resp.output_logprobs - prompt_mask = [1] * resp.input_len + [0] * resp.output_len - versions = [-1] * resp.input_len + resp.output_versions - - reward = self.reward_fn( - prompt=req.text, - completions=resp.completions, - prompt_ids=resp.input_tokens, - completion_ids=resp.output_tokens, - **data, - ) - res = dict( - # unsqueeze to add an additional batch dimension - input_ids=torch.tensor(seq).unsqueeze(0), - prompt_mask=torch.tensor(prompt_mask).unsqueeze(0), - logprobs=torch.tensor(logprobs).unsqueeze(0), - versions=torch.tensor(versions).unsqueeze(0), - # reward - rewards=torch.tensor([reward]), - ) - - return TensorDict(res, batch_size=[1]) + resps = await asyncio.gather(*[engine.agenerate(req) for _ in range(n_samples)]) + + results = [] + for resp in resps: + seq = resp.input_tokens + resp.output_tokens + logprobs = [0] * resp.input_len + resp.output_logprobs + prompt_mask = [1] * resp.input_len + [0] * resp.output_len + versions = [-1] * resp.input_len + resp.output_versions + + reward = self.reward_fn( + prompt=req.text, + completions=resp.completions, + prompt_ids=resp.input_tokens, + completion_ids=resp.output_tokens, + **data, + ) + res = dict( + # unsqueeze to add an additional batch dimension + input_ids=torch.tensor(seq).unsqueeze(0), + prompt_mask=torch.tensor(prompt_mask).unsqueeze(0), + logprobs=torch.tensor(logprobs).unsqueeze(0), + versions=torch.tensor(versions).unsqueeze(0), + # reward + rewards=torch.tensor([reward]), + ) + results.append(TensorDict(res, batch_size=[1])) + + return concat_padded_tensors(results) diff --git a/pyproject.toml b/pyproject.toml index 7597e137a0..3ec0a71ea8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,9 @@ dependencies = [ "hydra-core==1.4.0.dev1", "packaging", "tabulate", + "torchdata", + "gymnasium", + "tensordict", # Monitoring and logging "wandb", diff --git a/requirements.txt b/requirements.txt index 4ffab7e39b..b719cf7636 100644 --- a/requirements.txt +++ b/requirements.txt @@ -69,4 +69,7 @@ word2number Pebble timeout-decorator prettytable -swanlab[dashboard] \ No newline at end of file +swanlab[dashboard] +torchdata +gymnasium +tensordict \ No newline at end of file