agilerl 2.8.0.dev0__tar.gz → 2.8.0.dev2__tar.gz

{agilerl-2.8.0.dev0 → agilerl-2.8.0.dev2}/.gitignore

@@ -1,5 +1,5 @@
-# Symlink for Pylance/Pyright type checking of namespace package
-agilerl/arena
+# Symlink for Pylance/Pyright type checking of namespace package (repo root only)
+/agilerl/arena
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
@@ -162,5 +162,5 @@ wandb/
 CLAUDE.md
 AGENTS.md
 
-# Simlink for type-checking
-agilerl/arena/
+# Symlink for type-checking (repo root only)
+/agilerl/arena/

{agilerl-2.8.0.dev0 → agilerl-2.8.0.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agilerl
-Version: 2.8.0.dev0
+Version: 2.8.0.dev2
 Summary: AgileRL is a deep reinforcement learning library focused on improving RL development through RLOps.
 Author-email: Nick Ustaran-Anderegg <dev@agilerl.com>
 License-Expression: Apache-2.0
@@ -87,7 +87,7 @@ AgileRL is a Deep Reinforcement Learning library focused on improving developmen
 
 This library is initially focused on reducing the time taken for training models and hyperparameter optimization (HPO) by pioneering [evolutionary HPO techniques](https://docs.agilerl.com/en/latest/evo_hyperparam_opt/index.html) for reinforcement learning.<br>
 Evolutionary HPO has been shown to drastically reduce overall training times by automatically converging on optimal hyperparameters, without requiring numerous training runs.<br>
-We are constantly adding more algorithms and features. AgileRL already includes state-of-the-art evolvable [on-policy](https://docs.agilerl.com/en/latest/on_policy/index.html), [off-policy](https://docs.agilerl.com/en/latest/off_policy/index.html), [offline](https://docs.agilerl.com/en/latest/offline_training/index.html), [multi-agent](https://docs.agilerl.com/en/latest/multi_agent_training/index.html) and [contextual multi-armed bandit](https://docs.agilerl.com/en/latest/bandits/index.html) reinforcement learning algorithms with [distributed training](https://docs.agilerl.com/en/latest/distributed_training/index.html).
+We are constantly adding more algorithms and features. AgileRL already includes state-of-the-art evolvable [LLM fine-tuning](https://docs.agilerl.com/en/latest/llm_finetuning/index.html), [on-policy](https://docs.agilerl.com/en/latest/on_policy/index.html), [off-policy](https://docs.agilerl.com/en/latest/off_policy/index.html), [offline](https://docs.agilerl.com/en/latest/offline_training/index.html), [multi-agent](https://docs.agilerl.com/en/latest/multi_agent_training/index.html) and [contextual multi-armed bandit](https://docs.agilerl.com/en/latest/bandits/index.html) reinforcement learning algorithms with [distributed training](https://docs.agilerl.com/en/latest/distributed_training/index.html).
 
 <p align="center">
   <img src=https://user-images.githubusercontent.com/47857277/236407686-21363eb3-ffcf-419f-b019-0be4ddf1ed4a.gif width="100%" max-width="900">
@@ -105,6 +105,16 @@ We are constantly adding more algorithms and features. AgileRL already includes
 
 ## Benchmarks
 
+### LLM Fine-tuning
+
+AgileRL's multi-turn LLM training enables state-of-the-art performance on long-horizon tasks with small models. In the following example, AgileRL's CISPO was benchmarked against ART and TRL on the <a href="https://github.com/axon-rl/gem">GEM</a> Sudoku Hard task. This is a difficult multi-turn problem, which requires a context length of 32k tokens and up to 50 turns per rollout. The sync AgileRL run is a single agent using the AgileRL framework. The async and HPO runs were performed on <a href="https://arena.agilerl.com">Arena</a>, AgileRL's RLOps platform. All runs used the same starting hyperparameters. AgileRL runs were run on A100 40GB nodes, whereas ART and TRL required A100 80GB nodes due to a lack of optimizations. AgileRL runs significantly outperformed those using the ART and TRL frameworks.
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/AgileRL/AgileRL/main/docs/_static/multi_turn_llm_benchmarks.png" min-width="100%" width="700">
+</p>
+
+### Classic RL
+
 Reinforcement learning algorithms and libraries are usually benchmarked once the optimal hyperparameters for training are known, but it often takes hundreds or thousands of experiments to discover these. This is unrealistic and does not reflect the true, total time taken for training. What if we could remove the need to conduct all these prior experiments?
 
 In the charts below, a single AgileRL run, which automatically tunes hyperparameters, is benchmarked against Optuna's multiple training runs traditionally required for hyperparameter optimization, demonstrating the real time savings possible. Global steps is the sum of every step taken by any agent in the environment, including across an entire population.
@@ -402,10 +412,10 @@ Use the `ArenaClient` to interact with Arena programmatically from scripts or no
 from agilerl.arena import ArenaClient
 
 client = ArenaClient()
-client.login()
+client.login() # OAuth2 device-flow authentication
 
-# Register and validate a custom environment
-client.validate_environment(source="path/to/my_env.py")
+# Upload and validate a custom environment
+client.validate_environment(name="my-custom-env", source="path/to/my_env.py")
 
 # Train on validated custom environment
 client.submit_experiment(
@@ -423,13 +433,13 @@ The same operations are available from the command line:
 arena login
 
 # Upload and validate
-arena env validate --source path/to/my_env.py
+arena env validate my-custom-env --source path/to/my_env.py
 
 # Train on validated custom environment
 arena experiments submit path/to/manifest.yaml --project my-project
 ```
 
-For the full CLI and Python SDK reference—including authentication, environment validation, experiments, and deployment—see the [Arena Client](https://docs.agilerl.com/en/latest/arena/index.html) documentation.
+For the full CLI and Python SDK reference, including authentication, environment validation, experiments, and deployment, see the [Arena Client](https://docs.agilerl.com/en/latest/arena/index.html) documentation.
 
 ## Tutorials
 

{agilerl-2.8.0.dev0 → agilerl-2.8.0.dev2}/README.md

@@ -23,7 +23,7 @@ AgileRL is a Deep Reinforcement Learning library focused on improving developmen
 
 This library is initially focused on reducing the time taken for training models and hyperparameter optimization (HPO) by pioneering [evolutionary HPO techniques](https://docs.agilerl.com/en/latest/evo_hyperparam_opt/index.html) for reinforcement learning.<br>
 Evolutionary HPO has been shown to drastically reduce overall training times by automatically converging on optimal hyperparameters, without requiring numerous training runs.<br>
-We are constantly adding more algorithms and features. AgileRL already includes state-of-the-art evolvable [on-policy](https://docs.agilerl.com/en/latest/on_policy/index.html), [off-policy](https://docs.agilerl.com/en/latest/off_policy/index.html), [offline](https://docs.agilerl.com/en/latest/offline_training/index.html), [multi-agent](https://docs.agilerl.com/en/latest/multi_agent_training/index.html) and [contextual multi-armed bandit](https://docs.agilerl.com/en/latest/bandits/index.html) reinforcement learning algorithms with [distributed training](https://docs.agilerl.com/en/latest/distributed_training/index.html).
+We are constantly adding more algorithms and features. AgileRL already includes state-of-the-art evolvable [LLM fine-tuning](https://docs.agilerl.com/en/latest/llm_finetuning/index.html), [on-policy](https://docs.agilerl.com/en/latest/on_policy/index.html), [off-policy](https://docs.agilerl.com/en/latest/off_policy/index.html), [offline](https://docs.agilerl.com/en/latest/offline_training/index.html), [multi-agent](https://docs.agilerl.com/en/latest/multi_agent_training/index.html) and [contextual multi-armed bandit](https://docs.agilerl.com/en/latest/bandits/index.html) reinforcement learning algorithms with [distributed training](https://docs.agilerl.com/en/latest/distributed_training/index.html).
 
 <p align="center">
   <img src=https://user-images.githubusercontent.com/47857277/236407686-21363eb3-ffcf-419f-b019-0be4ddf1ed4a.gif width="100%" max-width="900">
@@ -41,6 +41,16 @@ We are constantly adding more algorithms and features. AgileRL already includes
 
 ## Benchmarks
 
+### LLM Fine-tuning
+
+AgileRL's multi-turn LLM training enables state-of-the-art performance on long-horizon tasks with small models. In the following example, AgileRL's CISPO was benchmarked against ART and TRL on the <a href="https://github.com/axon-rl/gem">GEM</a> Sudoku Hard task. This is a difficult multi-turn problem, which requires a context length of 32k tokens and up to 50 turns per rollout. The sync AgileRL run is a single agent using the AgileRL framework. The async and HPO runs were performed on <a href="https://arena.agilerl.com">Arena</a>, AgileRL's RLOps platform. All runs used the same starting hyperparameters. AgileRL runs were run on A100 40GB nodes, whereas ART and TRL required A100 80GB nodes due to a lack of optimizations. AgileRL runs significantly outperformed those using the ART and TRL frameworks.
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/AgileRL/AgileRL/main/docs/_static/multi_turn_llm_benchmarks.png" min-width="100%" width="700">
+</p>
+
+### Classic RL
+
 Reinforcement learning algorithms and libraries are usually benchmarked once the optimal hyperparameters for training are known, but it often takes hundreds or thousands of experiments to discover these. This is unrealistic and does not reflect the true, total time taken for training. What if we could remove the need to conduct all these prior experiments?
 
 In the charts below, a single AgileRL run, which automatically tunes hyperparameters, is benchmarked against Optuna's multiple training runs traditionally required for hyperparameter optimization, demonstrating the real time savings possible. Global steps is the sum of every step taken by any agent in the environment, including across an entire population.
@@ -338,10 +348,10 @@ Use the `ArenaClient` to interact with Arena programmatically from scripts or no
 from agilerl.arena import ArenaClient
 
 client = ArenaClient()
-client.login()
+client.login() # OAuth2 device-flow authentication
 
-# Register and validate a custom environment
-client.validate_environment(source="path/to/my_env.py")
+# Upload and validate a custom environment
+client.validate_environment(name="my-custom-env", source="path/to/my_env.py")
 
 # Train on validated custom environment
 client.submit_experiment(
@@ -359,13 +369,13 @@ The same operations are available from the command line:
 arena login
 
 # Upload and validate
-arena env validate --source path/to/my_env.py
+arena env validate my-custom-env --source path/to/my_env.py
 
 # Train on validated custom environment
 arena experiments submit path/to/manifest.yaml --project my-project
 ```
 
-For the full CLI and Python SDK reference—including authentication, environment validation, experiments, and deployment—see the [Arena Client](https://docs.agilerl.com/en/latest/arena/index.html) documentation.
+For the full CLI and Python SDK reference, including authentication, environment validation, experiments, and deployment, see the [Arena Client](https://docs.agilerl.com/en/latest/arena/index.html) documentation.
 
 ## Tutorials
 

{agilerl-2.8.0.dev0 → agilerl-2.8.0.dev2}/agilerl/algorithms/cispo.py

@@ -4,9 +4,11 @@ from __future__ import annotations
 
 from typing import Any
 
-from agilerl.algorithms.grpo import GRPO, _signatures_without_loss_type
+from agilerl.algorithms.grpo import GRPO
+from agilerl.utils.algo_utils import inherit_init_signature
 
 
+@inherit_init_signature(GRPO, fixed={"loss_type"})
 class CISPO(GRPO):
     """CISPO loss variant of :class:`agilerl.algorithms.grpo.GRPO`
 
@@ -16,8 +18,3 @@ class CISPO(GRPO):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         """Initialize a CISPO agent with fixed ``loss_type``."""
         super().__init__(*args, loss_type="cispo", **kwargs)
-
-
-_CISPO_CLASS_SIG, _CISPO_INIT_SIG = _signatures_without_loss_type()
-CISPO.__signature__ = _CISPO_CLASS_SIG
-CISPO.__init__.__signature__ = _CISPO_INIT_SIG

{agilerl-2.8.0.dev0 → agilerl-2.8.0.dev2}/agilerl/algorithms/core/base.py

@@ -149,7 +149,6 @@ if TYPE_CHECKING or HAS_LLM_DEPENDENCIES:
         build_vllm_llm_init_kwargs,
         build_vllm_rollout_lora_request,
         create_model_from_name_or_path,
-        format_colocated_vllm_oom_hint,
         gather_if_zero3,
         get_model_name_or_path,
         get_state_dict,
@@ -164,8 +163,12 @@ if TYPE_CHECKING or HAS_LLM_DEPENDENCIES:
 if TYPE_CHECKING or HAS_DEEPSPEED:
     from deepspeed.checkpoint.utils import clone_tensors_for_torch_save
 
-if TYPE_CHECKING or HAS_VLLM:
+if TYPE_CHECKING:
+    from vllm import LLM, CompletionOutput, SamplingParams
+elif HAS_VLLM:
     from vllm import LLM, CompletionOutput, SamplingParams
+else:
+    LLM = CompletionOutput = SamplingParams = None
 
 __all__ = ["ActionResult", "EvolvableAlgorithm", "MultiAgentRLAlgorithm", "RLAlgorithm"]
 
@@ -173,8 +176,6 @@ logger = logging.getLogger(__name__)
 
 SelfAgentWrapper = TypeVar("SelfAgentWrapper", bound=AgentWrapperProtocol)
 
-logger = logging.getLogger(__name__)
-
 
 class _RegistryMeta(type):
     """Metaclass to wrap registry information after algorithm is done
@@ -2239,18 +2240,10 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
     :param lora_target_scope: Optional PEFT LoRA path scope for multimodal models
         (e.g. ``"language_model"``). Passed to :func:`adapt_lora_config_for_model`.
     :type lora_target_scope: str | None, optional
-    :param fused_logprobs_chunk_rows: Standard (non-Liger) path only. Rows
-        (tokens) per ``(chunk_rows, vocab)`` logit tile when computing per-token
-        log-probs via the fused-linear-logprob path. Peak logits memory is
-        ``O(chunk_rows * vocab)`` regardless of batch/sequence length. ``None``
-        (default) auto-tunes to a ~256 MB fp32 tile.
-    :type fused_logprobs_chunk_rows: int | None, optional
-    :param fused_loss_chunk_rows: Rows per ``(chunk_rows, vocab)`` logit tile in
-        the token-level Liger fused policy loss. ``None`` (default) auto-tunes to
-        a ~256 MB fp32 logit workspace — the same heuristic as
-        ``fused_logprobs_chunk_rows`` on the standard path; pass an int to
-        override.
-    :type fused_loss_chunk_rows: int | None, optional
+    :param chunk_rows: Primary chunk-size knob for fused logit tiles used by
+        both the standard fused-logprob path and the Liger fused-loss path.
+        ``None`` (default) preserves each path's auto-tuned behavior.
+    :type chunk_rows: int | None, optional
     :param vllm_importance_sampling_correction: When ``True`` (default) and
         ``use_vllm=True``, correct the rollout/trainer log-prob mismatch by
         weighting each training token by ``clamp(exp(trainer - sampling),
@@ -2339,8 +2332,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
         activation_offload: bool = False,
         use_sequence_packing: bool = False,
         lora_target_scope: str | None = None,
-        fused_logprobs_chunk_rows: int | None = None,
-        fused_loss_chunk_rows: int | None = None,
+        chunk_rows: int | None = None,
         vllm_importance_sampling_correction: bool = True,
         vllm_importance_sampling_cap: float = 2.0,
     ) -> None:
@@ -2503,10 +2495,10 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
         self.wrap = wrap
         self.use_separate_reference_adapter = use_separate_reference_adapter
         self.cast_logprobs_to_fp32 = cast_logprobs_to_fp32
-        # Per-chunk row count for the fused-linear-logprob workspace. ``None``
-        # uses a vocab-aware ~256 MB-workspace heuristic;
-        # set explicitly to trade kernel-launch count vs ``rows * vocab`` peak.
-        self.fused_logprobs_chunk_rows = fused_logprobs_chunk_rows
+        if chunk_rows is not None and chunk_rows <= 0:
+            msg = f"chunk_rows must be a positive int or None, got {chunk_rows}."
+            raise ValueError(msg)
+        self.chunk_rows = chunk_rows
         # vLLM sampling-mismatch correction (truncated importance sampling).
         # The rollout is drawn from vLLM but the loss treats the trainer's
         # recomputed ``old_log_probs`` as the behaviour policy; the two differ
@@ -2522,13 +2514,6 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
         # Kept on even when use_vllm=False: decoupled rollouts still sample
         # from a separate vLLM engine.
         self._is_correction_liger_warned = False
-        if fused_loss_chunk_rows is not None and fused_loss_chunk_rows <= 0:
-            msg = (
-                f"fused_loss_chunk_rows must be a positive int or None, "
-                f"got {fused_loss_chunk_rows}."
-            )
-            raise ValueError(msg)
-        self.fused_loss_chunk_rows = fused_loss_chunk_rows
         # Warn-once flag for the canonical Liger + non-token importance-sampling
         # "not memory-bounded" warning (see :meth:`_warn_liger_non_token_is`).
         self._liger_non_token_warned = False
@@ -2553,6 +2538,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
         self._vllm_moved = False
         self._vllm_lora_loaded = False
         self._vllm_lora_staging_dir: Path | None = None
+        self._vllm_lora_staging_dir_is_temp = True
         self._vllm_rollout_lora_request: Any | None = None
         # Colocated vLLM (use_vllm=True) runs the rollout engine and the HF
         # trainer in one process. Each holds its OWN base: vLLM cycles its base
@@ -4201,7 +4187,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
                 fused_ids[start:end, 1:],
                 temperature=self.temperature,
                 cast_to_fp32=self.cast_logprobs_to_fp32,
-                _chunk_rows=self.fused_logprobs_chunk_rows,
+                chunk_rows=self.chunk_rows,
             )
             del first
 
@@ -4376,7 +4362,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
             packed.input_ids[:, 1:],
             temperature=self.temperature,
             cast_to_fp32=self.cast_logprobs_to_fp32,
-            _chunk_rows=self.fused_logprobs_chunk_rows,
+            chunk_rows=self.chunk_rows,
         )
         log_probs = unpack_logprobs(packed_lp, packed)
 
@@ -4600,7 +4586,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
                         packed.input_ids[:, 1:],
                         temperature=self.temperature,
                         cast_to_fp32=self.cast_logprobs_to_fp32,
-                        _chunk_rows=self.fused_logprobs_chunk_rows,
+                        chunk_rows=self.chunk_rows,
                     )
                     # Map back to the dense (mb, T-1) frame so the loss path is
                     # unchanged; cross-segment boundary predictions are dropped.
@@ -4614,7 +4600,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
                         batch_ids[:, 1:],
                         temperature=self.temperature,
                         cast_to_fp32=self.cast_logprobs_to_fp32,
-                        _chunk_rows=self.fused_logprobs_chunk_rows,
+                        chunk_rows=self.chunk_rows,
                     )
 
                 first = None
@@ -4696,13 +4682,17 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
     def _ensure_vllm_lora_staging_dir(self) -> Path:
         """Resolve (once) the dir the rollout LoRA adapter is exported to.
 
-        Honours ``VLLMConfig.lora_staging_dir`` when set — e.g. a shared/NFS
-        path that a colocated Ray rollout worker reads the adapter from — by
-        creating it (parents included) and marking it non-temporary so
-        ``clean_up`` never deletes it. Otherwise falls back to a process-private
-        ``mkdtemp`` that ``clean_up`` removes. Idempotent: both the colocated
-        init (``_configure_vllm``) and every adapter sync (``_move_lora_to_vllm``)
-        call this, so the same directory is used throughout the agent's life.
+        The staging dir is always process-private: each rank exports its own
+        adapter copy and reads it back locally, so ranks never race on shared
+        files. Honours ``VLLMConfig.lora_staging_dir`` when set — e.g. a known
+        path that orchestrated deployments expect the adapter under — staging
+        in a ``rank_<process_index>`` subdirectory of that root when
+        distributed. The dir is created (parents included) and marked
+        non-temporary so ``clean_up`` never deletes it. Otherwise falls back
+        to a process-private ``mkdtemp`` that ``clean_up`` removes.
+        Idempotent: both the colocated init (``_configure_vllm``) and every
+        adapter sync (``_move_lora_to_vllm``) call this, so the same directory
+        is used throughout the agent's life.
 
         :return: The resolved staging directory.
         :rtype: pathlib.Path
@@ -4710,8 +4700,11 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
         if self._vllm_lora_staging_dir is None:
             configured = getattr(self.vllm_config, "lora_staging_dir", None)
             if configured is not None:
-                self._vllm_lora_staging_dir = Path(configured)
-                self._vllm_lora_staging_dir.mkdir(parents=True, exist_ok=True)
+                staging_dir = Path(configured)
+                if self.accelerator is not None and self.accelerator.num_processes > 1:
+                    staging_dir = staging_dir / f"rank_{self.accelerator.process_index}"
+                staging_dir.mkdir(parents=True, exist_ok=True)
+                self._vllm_lora_staging_dir = staging_dir
                 self._vllm_lora_staging_dir_is_temp = False
             else:
                 self._vllm_lora_staging_dir = Path(
@@ -4734,12 +4727,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
         peft_ref = self._get_peft_model_for_vllm_sync()
         peft_ref.set_adapter(self._vllm_rollout_adapter)
 
-        # Export to a fixed staging dir + id and refresh the resident rollout
-        # slot in place: ``load_inplace`` (2nd sync onward) re-reads the updated
-        # weights from disk, and the fixed id avoids per-sync CUDA-graph
-        # accumulation that would grow GPU memory across iterations.
         staging_dir = self._ensure_vllm_lora_staging_dir()
-        is_main_process = self.accelerator is None or self.accelerator.is_main_process
         with gather_if_zero3(self.zero_stage, list(peft_ref.parameters())):
             if self.lora_config is None:
                 msg = "lora_config is required for vLLM LoRA adapter export."
@@ -4749,33 +4737,9 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
                 staging_dir,
                 self._vllm_rollout_adapter,
                 target_modules=self.lora_config.target_modules,
-                is_main_process=is_main_process,
             )
         if self.accelerator is not None:
             self.accelerator.wait_for_everyone()
-        if not adapter_path.is_dir():
-            msg = (
-                f"PEFT adapter export for {self._vllm_rollout_adapter!r} not found under "
-                f"{staging_dir}. Expected {adapter_path} or adapter_config.json in "
-                f"{staging_dir}."
-            )
-            raise FileNotFoundError(msg)
-
-        if is_main_process and logger.isEnabledFor(logging.DEBUG):
-            # Sum of L2 norms of the trained-from-zero LoRA-B weights; a value
-            # that changes across syncs confirms the trainer is exporting
-            # updated weights into the rollout adapter. Gated on the standard
-            # logging level (the norm reduction costs a GPU sync).
-            lora_b_sq = sum(
-                float(p.detach().float().pow(2).sum().item())
-                for n, p in peft_ref.named_parameters()
-                if "lora_B" in n and self._vllm_rollout_adapter in n
-            )
-            logger.debug(
-                "lora-sync: actor lora_B L2=%.6f path=%s",
-                lora_b_sq**0.5,
-                adapter_path,
-            )
 
         # One-shot refresh of the resident slot. ``load_inplace`` forces vLLM to
         # re-read the (updated) adapter weights from disk; required from the second
@@ -4784,11 +4748,19 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
             adapter_path,
             load_inplace=self._vllm_lora_loaded,
         )
-        loaded = self.llm.llm_engine.add_lora(refresh_request)
+        lora_device = torch.device(self.device)
+        if lora_device.type == "cuda":
+            # Pin the CUDA context to this agent's device: vLLM's LoRA copy
+            # kernels otherwise launch on the process-default device.
+            with torch.cuda.device(lora_device):
+                loaded = self.llm.llm_engine.add_lora(refresh_request)
+        else:
+            loaded = self.llm.llm_engine.add_lora(refresh_request)
         if not loaded:
             msg = (
-                f"vLLM failed to load LoRA adapter from {adapter_path}. "
-                "Check max_lora_rank / target module names match the trainer."
+                "vLLM failed to load LoRA adapter from "
+                f"{adapter_path}. Check max_lora_rank / target module "
+                "names match the trainer."
             )
             raise RuntimeError(msg)
 
@@ -5077,13 +5049,13 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
         logits: torch.Tensor,
         index: torch.Tensor,
         cast_to_fp32: bool = True,
-        _chunk_rows: int = 1,
+        chunk_rows: int = 1,
     ) -> torch.Tensor:
         """Calculate log probabilities for previously generated token ids.
 
-        Processes ``_chunk_rows`` rows at a time so peak memory stays bounded to
-        ``(_chunk_rows, seq_len, vocab_size)`` rather than the full batch, avoiding
-        OOM on large-vocabulary models. Default ``_chunk_rows=1`` minimizes the
+        Processes ``chunk_rows`` rows at a time so peak memory stays bounded to
+        ``(chunk_rows, seq_len, vocab_size)`` rather than the full batch, avoiding
+        OOM on large-vocabulary models. Default ``chunk_rows=1`` minimizes the
         fp32 workspace at the cost of more kernel launches; raise to amortize
         launch overhead when memory headroom allows.
 
@@ -5120,12 +5092,12 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
             result = target - log_z
             return result.to(orig_dtype) if cast_to_fp32 else result
 
-        if B <= _chunk_rows:
+        if B <= chunk_rows:
             return _logprobs_chunk(logits, index)
 
         per_token_logps = []
-        for start in range(0, B, _chunk_rows):
-            end = min(start + _chunk_rows, B)
+        for start in range(0, B, chunk_rows):
+            end = min(start + chunk_rows, B)
             per_token_logps.append(
                 _logprobs_chunk(logits[start:end], index[start:end]),
             )
@@ -5135,10 +5107,9 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
     def _resolve_fused_chunk_rows(vocab_size: int, explicit: int | None = None) -> int:
         """Rows per fused ``(chunk_rows, vocab)`` logit tile.
 
-        Shared by the fused-linear-logprob (standard) path
-        (``fused_logprobs_chunk_rows``) and the Liger fused-loss path
-        (``fused_loss_chunk_rows``) so both bound their per-chunk logit
-        workspace identically. A positive ``explicit`` overrides; ``None``
+        Shared by the fused-linear-logprob (standard) path and the Liger
+        fused-loss path so both bound their per-chunk logit workspace
+        identically. A positive ``explicit`` overrides; ``None``
         auto-tunes to a ~256 MB fp32 logit workspace (fewer rows at larger
         vocab), clamped to ``[128, 4096]``.
 
@@ -5162,12 +5133,12 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
         target_ids: torch.Tensor,
         temperature: float = 1.0,
         cast_to_fp32: bool = True,
-        _chunk_rows: int | None = None,
+        chunk_rows: int | None = None,
     ) -> torch.Tensor:
         """Per-token target logprobs without materializing the full ``(B, T, V)``
         logits tensor.
 
-        Tiles flat over ``(B*T)`` with workspace bounded to ``(_chunk_rows, V)``
+        Tiles flat over ``(B*T)`` with workspace bounded to ``(chunk_rows, V)``
         per iteration. Counterpart of :meth:`_logprobs_from_logits` for
         callers that hold hidden states and the lm_head separately. **No-grad
         only** — gradients won't flow to ``lm_head_weight`` from this fn. The
@@ -5195,16 +5166,16 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
             in fp32 then cast back. Same semantics as
             :meth:`_logprobs_from_logits`.
         :type cast_to_fp32: bool, optional
-        :param _chunk_rows: rows of the flattened ``(B*T)`` workspace per
-            iteration; trades launch count vs ``_chunk_rows * V`` peak. When
+        :param chunk_rows: rows of the flattened ``(B*T)`` workspace per
+            iteration; trades launch count vs ``chunk_rows * V`` peak. When
             ``None`` (default) it is resolved from the vocab size via
             a ~256 MB fp32 workspace heuristic.
-        :type _chunk_rows: int | None, optional
+        :type chunk_rows: int | None, optional
         :return: ``(B, T)`` per-token logprobs in ``hidden.dtype``.
         :rtype: torch.Tensor
         """
-        _chunk_rows = LLMAlgorithm._resolve_fused_chunk_rows(
-            lm_head_weight.shape[0], _chunk_rows
+        chunk_rows = LLMAlgorithm._resolve_fused_chunk_rows(
+            lm_head_weight.shape[0], chunk_rows
         )
         return fused_linear_logprobs_chunked(
             hidden,
@@ -5213,7 +5184,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
             target_ids,
             temperature=temperature,
             cast_to_fp32=cast_to_fp32,
-            chunk_rows=_chunk_rows,
+            chunk_rows=chunk_rows,
         )
 
     @staticmethod
@@ -5224,7 +5195,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
         target_ids: torch.Tensor,
         temperature: float = 1.0,
         cast_to_fp32: bool = True,
-        _chunk_rows: int | None = None,
+        chunk_rows: int | None = None,
     ) -> torch.Tensor:
         """Gradient-aware version of :meth:`_logprobs_from_hidden_fused`.
 
@@ -5250,15 +5221,15 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
         :type temperature: float, optional
         :param cast_to_fp32: run the per-chunk reduction in fp32.
         :type cast_to_fp32: bool, optional
-        :param _chunk_rows: rows of the flattened ``(B*T)`` workspace per chunk.
+        :param chunk_rows: rows of the flattened ``(B*T)`` workspace per chunk.
             When ``None`` (default) it is resolved from the vocab size via
             a ~256 MB fp32 workspace heuristic.
-        :type _chunk_rows: int | None, optional
+        :type chunk_rows: int | None, optional
         :return: ``(B, T)`` per-token logprobs in ``hidden.dtype``.
         :rtype: torch.Tensor
         """
-        _chunk_rows = LLMAlgorithm._resolve_fused_chunk_rows(
-            lm_head_weight.shape[0], _chunk_rows
+        chunk_rows = LLMAlgorithm._resolve_fused_chunk_rows(
+            lm_head_weight.shape[0], chunk_rows
         )
         return FusedLinearLogProbsFunction.apply(
             hidden,
@@ -5267,7 +5238,7 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
             target_ids,
             temperature,
             cast_to_fp32,
-            _chunk_rows,
+            chunk_rows,
         )
 
     def _configure_batch_size_per_process(
@@ -5729,10 +5700,11 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
     def _sleep_vllm_after_init(self) -> None:
         """Put the colocated engine to sleep once after construction.
 
-        Native ``sleep(level=1)``: vLLM backs its base up to host RAM and frees
-        the KV cache; ``wake_up()`` restores the base (dense or bnb 4-bit).
+        Native ``sleep(level=sleep_mode_level)``: vLLM cycles its allocator
+        state based on the configured sleep level; ``wake_up()`` restores the
+        engine allocations.
         """
-        self.llm.sleep(level=1)
+        self.llm.sleep(level=self.vllm_config.sleep_mode_level)
         self._vllm_awake = False
         if self.accelerator is None or self.accelerator.is_main_process:
             log_cuda_memory_snapshot("vLLM sleep complete")
@@ -5745,8 +5717,6 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
             return
 
         ds_plugin = getattr(self.accelerator.state, "deepspeed_plugin", None)
-        if ds_plugin is None:
-            return
 
         ds_config = ds_plugin.deepspeed_config
         if "gradient_clipping" not in ds_config:
@@ -5860,15 +5830,16 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
 
     def _prepare_vllm_for_training(self) -> None:
         """Prepare vLLM for learning."""
-        if self._vllm_awake and (
-            self.accelerator is None or self.accelerator.is_main_process
-        ):
+        if not self.use_vllm:
+            return
+        # Every rank holds its own colocated engine (external_launcher), so
+        # every rank must sleep it — not just the main process.
+        if self.vllm_config.sleep_mode and self._vllm_awake:
             torch.cuda.empty_cache()
-            self.llm.sleep(level=1)
+            self.llm.sleep(level=self.vllm_config.sleep_mode_level)
             self._vllm_awake = False
 
-        if self.use_vllm:
-            self._vllm_moved = False
+        self._vllm_moved = False
 
     def _prepare_vllm_for_generation(self) -> None:
         if self.use_memory_efficient_params:
@@ -5881,35 +5852,11 @@ class LLMAlgorithm(EvolvableAlgorithm, ABC):
                 log_cuda_memory_snapshot(
                     "trainer base offloaded to CPU (before vLLM wake)"
                 )
-        if not self._vllm_awake and (
-            self.accelerator is None or self.accelerator.is_main_process
-        ):
+        # Every rank holds its own colocated engine, and _sleep_vllm_after_init
+        # slept them all; we wake them all here.
+        if self.vllm_config.sleep_mode and not self._vllm_awake:
             torch.cuda.empty_cache()
-            device_index = (
-                self.accelerator.local_process_index
-                if self.accelerator is not None
-                else 0
-            )
-            try:
-                self.llm.wake_up()
-            except RuntimeError as err:  # pragma: no cover
-                err_text = str(err).lower()
-                if "out of memory" in err_text or "cuda error" in err_text:
-                    vcfg = self.vllm_config
-                    hint = format_colocated_vllm_oom_hint(
-                        device_index,
-                        kv_cache_memory_bytes=(
-                            vcfg.kv_cache_memory_bytes if vcfg is not None else None
-                        ),
-                        gpu_memory_utilization=(
-                            vcfg.gpu_memory_utilization if vcfg is not None else None
-                        ),
-                        max_model_len=getattr(self, "max_model_len", None),
-                        trainer_on_gpu=not self.use_memory_efficient_params,
-                    )
-                    msg = f"vLLM wake_up failed (GPU OOM).\n{hint}"
-                    raise RuntimeError(msg) from err
-                raise
+            self.llm.wake_up()
             self._vllm_awake = True
             if self.accelerator is None or self.accelerator.is_main_process:
                 log_cuda_memory_snapshot("vLLM base restored on GPU (after wake)")