[Feature] Restructure LLM pip extras for backend flexibility

- Add llm-vllm, llm-sglang, llm-all extras for backend selection
- Base llm extra no longer includes inference backend
- Update sglang_nccl.py to use SGLang's native NCCL utilities
- Remove vLLM dependency from SGLang weight sync code

Users can now:
- pip install torchrl[llm-vllm] for vLLM backend
- pip install torchrl[llm-sglang] for SGLang backend
- pip install torchrl[llm-all] for both backends

Co-authored-by: Cursor <cursoragent@cursor.com>
ghstack-source-id: d53065a
Pull-Request: #3436

Author: vmoens

pyproject.toml

@@ -93,9 +93,9 @@ marl = ["vmas>=1.2.10", "pettingzoo>=1.24.1", "dm-meltingpot; python_version>='3
 open_spiel = ["open_spiel>=1.5"]
 brax = ["jax>=0.7.0; python_version>='3.11'", "brax; python_version>='3.11'"]
 procgen = ["procgen"]
+# Base LLM dependencies (no inference backend - use llm-vllm or llm-sglang)
 llm = [
     "transformers",
-    "vllm",
     "playwright",
     "datasets",
     "langdetect",
@@ -107,12 +107,27 @@ llm = [
     "einops",
     "safetensors",
 ]
+# LLM with vLLM backend
+llm-vllm = [
+    "torchrl[llm]",
+    "vllm",
+]
+# LLM with SGLang backend
+llm-sglang = [
+    "torchrl[llm]",
+    "sglang[all]",
+]
+# LLM with both backends
+llm-all = [
+    "torchrl[llm]",
+    "vllm",
+    "sglang[all]",
+]
 grpo = [
     # Core dependencies for GRPO training
-    "datasets",
+    "torchrl[llm-vllm]",
     "peft",
     "wandb",
-    "vllm",
     "transformers",
     "accelerate",
     "ray",
@@ -121,10 +136,6 @@ grpo = [
     "flash-attn",
     "bitsandbytes",
     "xformers",
-    # LLM-related dependencies that may be needed
-    "nltk",
-    "langdetect", 
-    "immutabledict",
 ]
 dev = [
     "pre-commit",

torchrl/weight_update/llm/sglang_nccl.py

@@ -191,33 +191,21 @@ def init_all_workers_group(
         # Step 2: Initialize trainer's NCCL communicator
         torch.cuda.set_device(self.device)
 
-        try:
-            from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
-            from vllm.distributed.utils import StatelessProcessGroup
-
-            pg = StatelessProcessGroup.create(
-                host=self.master_address,
-                port=self.master_port,
-                rank=0,
-                world_size=self.world_size,
-            )
-            self._comm_group = PyNcclCommunicator(
-                pg, device=torch.device(f"cuda:{self.device}")
-            )
-        except ImportError:
-            # Fallback to torch.distributed if vLLM not available
-            torchrl_logger.warning(
-                "vLLM not available, falling back to torch.distributed for NCCL init. "
-                "This may not be compatible with SGLang's NCCL implementation."
-            )
-            if not torch.distributed.is_initialized():
-                torch.distributed.init_process_group(
-                    backend="nccl",
-                    init_method=f"tcp://{self.master_address}:{self.master_port}",
-                    rank=0,
-                    world_size=self.world_size,
-                )
-            self._comm_group = torch.distributed.distributed_c10d._get_default_group()
+        # Use SGLang's native NCCL utilities (no vLLM dependency)
+        from sglang.srt.distributed.device_communicators.pynccl import (
+            PyNcclCommunicator,
+        )
+        from sglang.srt.distributed.utils import StatelessProcessGroup
+
+        pg = StatelessProcessGroup.create(
+            host=self.master_address,
+            port=self.master_port,
+            rank=0,
+            world_size=self.world_size,
+        )
+        self._comm_group = PyNcclCommunicator(
+            pg, device=torch.device(f"cuda:{self.device}")
+        )
 
         torchrl_logger.info("NCCL group initialized successfully")