[Generator] Support non-remote (e.g. colocated) SGLang engine

skyrl-train/examples/gsm8k/run_gsm8k.sh

@@ -12,7 +12,10 @@ DATA_DIR="$HOME/data/gsm8k"
 NUM_GPUS=4
 LOGGER="wandb"  # change to "console" to print to stdout
 
-uv run --isolated --extra vllm -m skyrl_train.entrypoints.main_base \
+INFERENCE_BACKEND="vllm"
+# INFERENCE_BACKEND="sglang"
+
+uv run --isolated --extra $INFERENCE_BACKEND -m skyrl_train.entrypoints.main_base \
   data.train_data="['$DATA_DIR/train.parquet']" \
   data.val_data="['$DATA_DIR/validation.parquet']" \
   trainer.algorithm.advantage_estimator="grpo" \
@@ -37,7 +40,7 @@ uv run --isolated --extra vllm -m skyrl_train.entrypoints.main_base \
   generator.sampling_params.max_generate_length=1024 \
   trainer.policy.optimizer_config.lr=1.0e-6 \
   trainer.algorithm.use_kl_loss=true \
-  generator.backend=vllm \
+  generator.backend=$INFERENCE_BACKEND \
   generator.run_engines_locally=true \
   generator.weight_sync_backend=nccl \
   generator.async_engine=true \

skyrl-train/pyproject.toml

@@ -63,7 +63,7 @@ dev = [
     "black==24.10.0",
     "pytest>=6.2.5",
     "pytest-asyncio",
-    "pre-commit"
+    "pre-commit",
 ]
 docs = [
     "sphinx>=7.0.0",
@@ -82,7 +82,7 @@ vllm = [
     "torchvision"
 ]
 sglang = [
-    "sglang[srt,openai,torch_memory_saver]==0.4.8.post1",
+    "sglang[srt,openai,torch_memory_saver]==0.4.8.post1",  # 0.4.9.post1 causes non-colocate weight broadcast to hang
     # The version is pinned to 0.2.5 because sglang requires this
     # NOTE (sumanthrh): This can be made a common dependency, but then different inference engines can pin different compatible flashinfer versions and it might quickly break.
     "flashinfer-python@https://download.pytorch.org/whl/cu128/flashinfer/flashinfer_python-0.2.6.post1%2Bcu128torch2.7-cp39-abi3-linux_x86_64.whl",

skyrl-train/skyrl_train/entrypoints/main_base.py

@@ -51,12 +51,13 @@ def create_ray_wrapped_inference_engines_from_config(cfg: DictConfig, colocate_p
         max_model_len=cfg.generator.max_input_length + cfg.generator.sampling_params.max_generate_length,
         shared_pg=colocate_pg,
         gpu_memory_utilization=cfg.generator.gpu_memory_utilization,
-        vllm_enable_sleep=cfg.trainer.placement.colocate_all,
+        inference_engine_enable_sleep=cfg.trainer.placement.colocate_all,
         async_engine=cfg.generator.async_engine,
         max_num_batched_tokens=cfg.generator.max_num_batched_tokens,
         max_num_seqs=cfg.generator.max_num_seqs,
         sampling_params=get_sampling_params_for_backend(cfg.generator.backend, cfg.generator.sampling_params),
         tokenizer=tokenizer,
+        backend=cfg.generator.backend,
     )
 
 

skyrl-train/skyrl_train/inference_engines/ray_wrapped_inference_engine.py

@@ -25,7 +25,7 @@ def tp_size(self):
         return ray.get(self.inference_engine_actor.tp_size.remote())
 
     async def generate(self, input_batch: InferenceEngineInput) -> InferenceEngineOutput:
-        return await self.inference_engine_actor.generate.remote(input_batch)
+        return await self.inference_engine_actor.generate.remote(input_batch=input_batch)
 
     async def wake_up(self, *args: Any, **kwargs: Any):
         return await self.inference_engine_actor.wake_up.remote(*args, **kwargs)
@@ -62,21 +62,29 @@ def create_ray_wrapped_inference_engines(
     max_model_len: int,
     shared_pg=None,
     gpu_memory_utilization=None,
-    vllm_enable_sleep=False,
+    inference_engine_enable_sleep=False,
     async_engine=False,
     max_num_batched_tokens=8192,
     max_num_seqs=1024,
     sampling_params: Optional[Dict[str, Any]] = None,
     tokenizer=None,
+    backend="vllm",
 ) -> List[InferenceEngineInterface]:
     """
     Create a list of RayWrappedInferenceEngine instances wrapping Ray actor handles to InferenceEngineInterface instances.
     """
-    import vllm
-    from skyrl_train.inference_engines.vllm.vllm_engine import VLLMRayActor, AsyncVLLMRayActor
     from skyrl_train.utils import ray_noset_visible_devices, get_all_env_variables, get_ray_pg_ready_with_timeout
 
-    assert vllm.__version__ >= "0.8.3", "SkyTrainer only supports vLLM >= 0.8.3"
+    if backend == "vllm":
+        import vllm
+        from skyrl_train.inference_engines.vllm.vllm_engine import VLLMRayActor, AsyncVLLMRayActor
+
+        assert vllm.__version__ >= "0.8.3", "SkyTrainer only supports vLLM >= 0.8.3"
+    elif backend == "sglang":
+        # We import SGLang later to avoid importing vllm. See `get_sglang_engine` for more.
+        pass
+    else:
+        raise ValueError(f"Unsupported backend: {backend}")
     inference_engine_actors = []
     noset_visible_devices = ray_noset_visible_devices(ray.get(get_all_env_variables.remote()))
     # NOTE: we use the ray backend for tensor parallel size > 1 to explicitly manage resource allocation
@@ -106,42 +114,92 @@ def create_ray_wrapped_inference_engines(
             placement_group_bundle_index=i * tensor_parallel_size,
         )
 
-        if async_engine:
-            actor_class = AsyncVLLMRayActor
-        else:
-            actor_class = VLLMRayActor
-
-        vllm_engine = actor_class.options(
-            num_cpus=num_gpus,
-            num_gpus=num_gpus,
-            scheduling_strategy=scheduling_strategy,
-        ).remote(
-            model=pretrain,
-            enforce_eager=enforce_eager,
-            worker_extension_cls="skyrl_train.inference_engines.vllm.vllm_engine.WorkerWrap",
-            tensor_parallel_size=tensor_parallel_size,
-            seed=seed + i,
-            distributed_executor_backend=distributed_executor_backend,
-            max_model_len=max_model_len,
-            enable_prefix_caching=enable_prefix_caching,
-            dtype=model_dtype,
-            trust_remote_code=True,
-            vllm_v1_disable_multiproc=vllm_v1_disable_multiproc,
-            gpu_memory_utilization=gpu_memory_utilization,
-            bundle_indices=bundle_indices,
-            num_gpus=0.2 if use_hybrid_engine else 1,
-            enable_sleep_mode=vllm_enable_sleep,
-            noset_visible_devices=noset_visible_devices,
-            max_num_batched_tokens=max_num_batched_tokens,
-            max_num_seqs=max_num_seqs,
-            sampling_params=sampling_params,
-            tokenizer=tokenizer,
-        )
-        inference_engine_actors.append(vllm_engine)
+        if backend == "vllm":
+            if async_engine:
+                actor_class = AsyncVLLMRayActor
+            else:
+                actor_class = VLLMRayActor
+
+            engine = actor_class.options(
+                num_cpus=num_gpus,
+                num_gpus=num_gpus,
+                scheduling_strategy=scheduling_strategy,
+            ).remote(
+                model=pretrain,
+                enforce_eager=enforce_eager,
+                worker_extension_cls="skyrl_train.inference_engines.vllm.vllm_engine.WorkerWrap",
+                tensor_parallel_size=tensor_parallel_size,
+                seed=seed + i,
+                distributed_executor_backend=distributed_executor_backend,
+                max_model_len=max_model_len,
+                enable_prefix_caching=enable_prefix_caching,
+                dtype=model_dtype,
+                trust_remote_code=True,
+                vllm_v1_disable_multiproc=vllm_v1_disable_multiproc,
+                gpu_memory_utilization=gpu_memory_utilization,
+                bundle_indices=bundle_indices,
+                num_gpus=0.2 if use_hybrid_engine else 1,
+                enable_sleep_mode=inference_engine_enable_sleep,
+                noset_visible_devices=noset_visible_devices,
+                max_num_batched_tokens=max_num_batched_tokens,
+                max_num_seqs=max_num_seqs,
+                sampling_params=sampling_params,
+                tokenizer=tokenizer,
+            )
+        elif backend == "sglang":
+            # NOTE: there is no async / sync engine distinction in SGLang
+
+            # NOTE(Charlie): We need `torch.cuda.is_available()` to be True to import SGLang. Otherwise, it requires
+            # importing vllm. See https://github.com/sgl-project/sglang/blob/v0.4.8.post1/python/sglang/srt/layers/quantization/utils.py#L11-L17
+            # Similar comment: https://github.com/volcengine/verl/blob/9cc307767b0c787e8f5ef581dac929f7bde044ef/verl/workers/fsdp_workers.py#L520-L527
+            @ray.remote
+            def get_sglang_engine():
+                # A workaround to avoid importing vllm is to give this task a GPU.
+                import os
+
+                before_cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "")
+                os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+                from skyrl_train.inference_engines.sglang.sglang_engine import SGLangRayActor
+
+                os.environ["CUDA_VISIBLE_DEVICES"] = before_cuda_visible_devices
+
+                actor_class = SGLangRayActor
+                engine = actor_class.options(
+                    num_cpus=num_gpus,
+                    num_gpus=num_gpus,
+                    scheduling_strategy=scheduling_strategy,
+                ).remote(
+                    model_path=pretrain,
+                    tp_size=tensor_parallel_size,
+                    mem_fraction_static=gpu_memory_utilization,
+                    random_seed=seed + i,
+                    context_length=max_model_len,
+                    disable_radix_cache=not enable_prefix_caching,
+                    dtype=model_dtype,
+                    trust_remote_code=True,
+                    max_prefill_tokens=max_num_batched_tokens,
+                    max_running_requests=max_num_seqs,
+                    # Borrowed from veRL's SGLang rollout
+                    mm_attention_backend="fa3",
+                    attention_backend="fa3",
+                    enable_memory_saver=inference_engine_enable_sleep,
+                    # Will be popped before instantiating sgl.Engine
+                    distributed_executor_backend=distributed_executor_backend,
+                    noset_visible_devices=noset_visible_devices,
+                    bundle_indices=bundle_indices,
+                    num_gpus=0.2 if use_hybrid_engine else 1,
+                    sampling_params=sampling_params,
+                    tokenizer=tokenizer,
+                )
+                return engine
+
+            engine = ray.get(get_sglang_engine.remote())
+
+        inference_engine_actors.append(engine)
 
     engines = [RayWrappedInferenceEngine(actor_handle) for actor_handle in inference_engine_actors]
 
-    if vllm_enable_sleep:
+    if inference_engine_enable_sleep:
         sleep_refs = [engine.inference_engine_actor.sleep.remote() for engine in engines]
         ray.get(sleep_refs)