torchrl 0.11.0__cp314-cp314-win_amd64.whl

torchrl/collectors/llm/ray_collector.py

@@ -0,0 +1,275 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import copy
+
+import warnings
+from collections.abc import Callable, Iterator
+from typing import Any
+
+import torch
+from tensordict import TensorDictBase
+from tensordict.nn import TensorDictModuleBase
+from torchrl.collectors.llm import LLMCollector
+from torchrl.collectors.weight_update import WeightUpdaterBase
+from torchrl.data.replay_buffers.replay_buffers import ReplayBuffer
+from torchrl.envs import EnvBase
+from torchrl.envs.llm.transforms.policy_version import PolicyVersion
+
+RAY_ERR = None
+try:
+    import ray
+
+    _has_ray = True
+except ImportError as err:
+    _has_ray = False
+    RAY_ERR = err
+
+
+class RayLLMCollector(LLMCollector):
+    """A lightweight Ray implementation of the LLM Collector that can be extended and sampled remotely.
+
+    Args:
+        env (EnvBase or EnvBase constructor): the environment to be used for data collection.
+
+    Keyword Args:
+        policy (Callable[[TensorDictBase], TensorDictBase]): the policy to be used for data collection.
+        policy_factory (Callable[[], Callable], optional): a callable that returns
+            a policy instance. This is exclusive with the `policy` argument.
+        dialog_turns_per_batch (int): A keyword-only argument representing the total
+            number of elements in a batch.
+        total_dialog_turns (int): A keyword-only argument representing the total
+            number of dialog turns returned by the collector during its lifespan.
+        yield_only_last_steps (bool, optional): whether to yield every step of a trajectory, or only the
+            last (done) steps.
+        yield_completed_trajectories (bool, optional): whether to yield batches of rollouts with a given number of steps
+            or single, completed trajectories.
+        postproc (Callable, optional): A post-processing transform.
+        async_envs (bool, optional): if True, the environment will be run asynchronously.
+        replay_buffer (ReplayBuffer, optional): if provided, the collector will not yield tensordicts
+            but populate the buffer instead.
+        reset_at_each_iter (bool, optional): if True, the environment will be reset at each iteration.
+        flatten_data (bool, optional): if True, the collector will flatten the collected data
+            before returning it.
+        weight_updater (WeightUpdaterBase or constructor, optional): An instance of WeightUpdaterBase
+            or its subclass, responsible for updating the policy weights on remote inference workers.
+        ray_init_config (dict[str, Any], optional): keyword arguments to pass to ray.init().
+        remote_config (dict[str, Any], optional): keyword arguments to pass to cls.as_remote().
+        num_cpus (int, optional): Number of CPUs of the actor. Defaults to `None` (taken from remote_config).
+        num_gpus (int, optional): Number of GPUs of the actor. Defaults to `None` (taken from remote_config).
+        sync_iter (bool, optional): if `True`, items yeilded by the collector will be synced to the local process.
+            If `False`, the collector will collect the next batch of data in between yielding.
+            This has no effect when data is collected through the :meth:`start` method.
+            For example:
+
+               >>> collector = RayLLMCollector(..., sync_iter=True)
+               >>> for data in collector:  # blocking
+               ...     # expensive operation - collector is idle
+               >>> collector = RayLLMCollector(..., sync_iter=False)
+               >>> for data in collector:  # non-blocking
+               ...     # expensive operation - collector is collecting data
+
+            This is somehwat equivalent to using :class:`~torchrl.collectors.MultiSyncCollector` (`sync_iter=True`) or
+            :class:`~torchrl.collectors.MultiAsyncDataCollector` (`sync_iter=False`).
+            Defaults to `True`.
+        verbose (bool, optional): if ``True``, the collector will print progress information.
+            Defaults to `False`.
+    """
+
+    def __init__(
+        self,
+        env: EnvBase | Callable[[], EnvBase],
+        *,
+        policy: Callable[[TensorDictBase], TensorDictBase] | None = None,
+        policy_factory: Callable[[], Callable[[TensorDictBase], TensorDictBase]]
+        | None = None,
+        dialog_turns_per_batch: int,
+        total_dialog_turns: int = -1,
+        yield_only_last_steps: bool | None = None,
+        yield_completed_trajectories: bool | None = None,
+        postproc: Callable[[TensorDictBase], TensorDictBase] | None = None,
+        async_envs: bool | None = None,
+        replay_buffer: ReplayBuffer | None = None,
+        reset_at_each_iter: bool = False,
+        flatten_data: bool | None = None,
+        weight_updater: WeightUpdaterBase
+        | Callable[[], WeightUpdaterBase]
+        | None = None,
+        ray_init_config: dict[str, Any] | None = None,
+        remote_config: dict[str, Any] | None = None,
+        track_policy_version: bool | PolicyVersion = False,
+        sync_iter: bool = True,
+        verbose: bool = False,
+        num_cpus: int | None = None,
+        num_gpus: int | None = None,
+    ) -> None:
+        if not _has_ray:
+            raise RuntimeError(
+                "ray library not found, unable to create a RayLLMCollector. "
+            ) from RAY_ERR
+        if not ray.is_initialized():
+            if ray_init_config is None:
+                from torchrl.collectors.distributed.ray import DEFAULT_RAY_INIT_CONFIG
+
+                ray_init_config = DEFAULT_RAY_INIT_CONFIG
+            ray.init(**ray_init_config)
+        if not sync_iter:
+            remote_config = copy.copy(remote_config)
+            remote_config.setdefault("max_concurrency", 2)
+        if num_cpus is not None:
+            remote_config.setdefault("num_cpus", num_cpus)
+        if num_gpus is not None:
+            remote_config.setdefault("num_gpus", num_gpus)
+        remote_cls = LLMCollector.as_remote(remote_config).remote
+        self.sync_iter = sync_iter
+        self._collector = remote_cls(
+            env=env,
+            policy=policy,
+            policy_factory=policy_factory,
+            dialog_turns_per_batch=dialog_turns_per_batch,
+            total_dialog_turns=total_dialog_turns,
+            yield_only_last_steps=yield_only_last_steps,
+            yield_completed_trajectories=yield_completed_trajectories,
+            postproc=postproc,
+            async_envs=async_envs,
+            replay_buffer=replay_buffer,
+            reset_at_each_iter=reset_at_each_iter,
+            flatten_data=flatten_data,
+            weight_updater=weight_updater,
+            track_policy_version=track_policy_version,
+            verbose=verbose,
+        )
+
+    def set_postproc(self, postproc: Callable[[TensorDictBase], TensorDictBase]):
+        return ray.get(self._collector.set_postproc.remote(postproc))
+
+    def _next_remote(self) -> None:
+        return self._collector.next.remote()
+
+    def next(self) -> None:
+        """Get the next batch of data from the collector.
+
+        Returns:
+            None as the data is written directly to the replay buffer.
+        """
+        return ray.get(self._next_remote())
+
+    def __iter__(self) -> Iterator[None]:
+        """Returns an iterator that yields None as the collector writes directly to the replay buffer."""
+        if not self.sync_iter:
+            future = self._next_remote()
+        else:
+            future = None
+        while True:
+            try:
+                if self.sync_iter:
+                    yield self.next()
+                else:
+                    result = ray.get(future)
+                    future = self._next_remote()
+                    yield result
+            except StopIteration:
+                break
+
+    def start(self):
+        """Starts the collector in a background thread."""
+        pending_task = self._collector.start.remote()
+        return ray.get(pending_task)
+
+    def is_running(self):
+        return ray.get(self._collector.is_running.remote())
+
+    def shutdown(self):
+        """Shuts down the collector."""
+        pending_task = self._collector.shutdown.remote()
+        return ray.get(pending_task)
+
+    def async_shutdown(self, timeout=None):
+        """Shuts down the collector asynchronously."""
+        pending_task = self._collector.async_shutdown.remote(timeout=timeout)
+        return ray.get(pending_task)
+
+    def update_policy_weights_(
+        self,
+        policy_or_weights: TensorDictBase | TensorDictModuleBase | dict | None = None,
+        *,
+        worker_ids: torch.device | int | list[int] | list[torch.device] | None = None,
+        **kwargs,
+    ):
+        """Updates the policy weights on remote workers.
+
+        Args:
+            policy_or_weights: The weights to update with. Can be:
+                - TensorDictModuleBase: A policy module whose weights will be extracted
+                - TensorDictBase: A TensorDict containing weights
+                - dict: A regular dict containing weights
+                - None: Will try to get weights from server using _get_server_weights()
+            worker_ids: The workers to update. If None, updates all workers.
+        """
+        if "policy_weights" in kwargs:
+            warnings.warn(
+                "`policy_weights` is deprecated. Use `policy_or_weights` instead.",
+                DeprecationWarning,
+            )
+            policy_or_weights = kwargs.pop("policy_weights")
+
+        pending_task = self._collector.update_policy_weights_.remote(
+            policy_or_weights=policy_or_weights, worker_ids=worker_ids
+        )
+        return ray.get(pending_task)
+
+    @property
+    def total_dialog_turns(self):
+        """Total number of dialog turns to collect."""
+        return ray.get(self._collector.total_dialog_turns.remote)
+
+    @property
+    def dialog_turns_per_batch(self) -> int:
+        """Number of dialog turns per batch."""
+        return ray.get(self._collector.dialog_turns_per_batch.remote)
+
+    @property
+    def rollout(self) -> Callable[[], TensorDictBase]:
+        """Returns the rollout function."""
+        return ray.get(self._collector.rollout.remote())
+
+    def init_updater(self, *args, **kwargs):
+        """Initialize the weight updater with custom arguments.
+
+        This method calls init_updater on the remote collector.
+
+        Args:
+            *args: Positional arguments for weight updater initialization
+            **kwargs: Keyword arguments for weight updater initialization
+        """
+        ray.get(self._collector.init_updater.remote(*args, **kwargs))
+
+    @property
+    def policy_version(self) -> str | int | None:
+        """The current version of the policy.
+
+        Returns:
+            The current version number (int) or UUID (str), or None if version tracking is disabled.
+        """
+        return ray.get(self._collector.get_policy_version.remote())
+
+    @property
+    def weight_updater(self) -> WeightUpdaterBase:
+        """The weight updater instance.
+
+        We can pass the weight updater because it's stateless, hence serializable.
+        """
+        return ray.get(self._collector.weight_updater.remote)
+
+    @weight_updater.setter
+    def weight_updater(self, weight_updater: WeightUpdaterBase):
+        """Set the weight updater instance."""
+        ray.get(self._collector.set_weight_updater.remote(weight_updater))
+        weight_updater.register_collector(self)
+
+    def increment_version(self):
+        """Increment the policy version."""
+        return ray.get(self._collector.increment_version.remote())

torchrl/collectors/llm/utils.py

@@ -0,0 +1,36 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import importlib.util
+
+from queue import Full as QueueFull, Queue
+
+from tensordict import TensorDictBase
+
+from torchrl._utils import logger as torchrl_logger
+
+_has_ray = importlib.util.find_spec("ray") is not None
+
+
+class _QueueAsRB:
+    def __init__(self, queue: Queue | ray.util.queue.Queue):  # noqa
+        if not _has_ray:
+            raise ImportError("Ray not installed.")
+        self.queue = queue
+
+    def extend(self, data: TensorDictBase):
+        from ray.util.queue import Full as RayQueueFull
+
+        # unbind the data and put in the queue
+        for item in data.unbind(0):
+            while True:
+                try:
+                    self.queue.put_nowait(item)
+                    break
+                except (QueueFull, RayQueueFull):
+                    self.queue.get()  # Remove the oldest item to make space
+                    torchrl_logger.warn("rollout queue full. Discarding data.")
+        return

torchrl/collectors/llm/weight_update/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+from .vllm import vLLMUpdater
+from .vllm_v2 import vLLMUpdaterV2
+
+__all__ = ["vLLMUpdater", "vLLMUpdaterV2"]

torchrl/collectors/llm/weight_update/vllm.py

@@ -0,0 +1,348 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import abc
+
+import importlib.util
+
+import torch
+import torch.cuda
+import torch.distributed
+
+from tensordict import TensorDict, TensorDictBase
+from tensordict.nn import TensorDictModuleBase
+
+from torchrl._utils import logger as torchrl_logger
+
+from torchrl.collectors.weight_update import WeightUpdaterBase
+from torchrl.modules.llm.backends import stateless_init_process_group
+
+_has_vllm = importlib.util.find_spec("vllm") is not None
+
+_has_ray = importlib.util.find_spec("ray") is not None
+
+
+class vLLMUpdaterMeta(abc.ABCMeta):
+    """Metaclass for vLLMUpdater that allows switching between V1 and V2 implementations.
+
+    When instantiating vLLMUpdater with v2=True, returns a vLLMUpdaterV2 instance instead.
+    This provides a unified entry point for both updater versions while maintaining
+    backward compatibility.
+    """
+
+    def __call__(cls, *args, v2=False, **kwargs):
+        if v2:
+            # Import V2 here to avoid circular imports
+            from .vllm_v2 import vLLMUpdaterV2
+
+            # V2 has a different signature - it expects a vllm_engine parameter
+            # If the user is providing the old signature, we need to handle this gracefully
+            if args or any(
+                k in kwargs
+                for k in [
+                    "master_address",
+                    "master_port",
+                    "model_metadata",
+                    "vllm_tp_size",
+                ]
+            ):
+                # Old signature detected - we can't auto-convert, user needs to update their code
+                raise TypeError(
+                    "When using v2=True, you must provide a vllm_engine parameter instead of "
+                    "the v1 parameters (master_address, master_port, model_metadata, vllm_tp_size). "
+                    "See vLLMUpdaterV2 documentation for details."
+                )
+
+            # Forward to V2 constructor
+            return vLLMUpdaterV2(*args, **kwargs)
+        else:
+            # Use original V1 constructor
+            return super().__call__(*args, **kwargs)
+
+
+class vLLMUpdater(WeightUpdaterBase, metaclass=vLLMUpdaterMeta):
+    """A class that sends weights to vLLM workers.
+
+    This class handles synchronizing weights between a training policy and vLLM inference workers.
+    It supports both local vLLM instances and remote Ray actors.
+
+    Args:
+        master_address (str, optional): The master address for distributed training. Defaults to localhost.
+        master_port (int, optional): The master port for distributed training. If None, will auto-assign.
+        model_metadata (dict[str, tuple[torch.dtype, torch.Size]], optional): Model metadata mapping
+            parameter names to their dtype and shape. If not provided, will be extracted from policy.
+        vllm_tp_size (int, optional): vLLM tensor parallel size. Defaults to 1.
+        v2 (bool, optional): If True, returns a vLLMUpdaterV2 instance instead. This is an experimental
+            feature that provides better integration with AsyncVLLM engines. When using v2=True, you must
+            provide a vllm_engine parameter instead of the above parameters. Defaults to False.
+
+    Methods:
+        init: Initialize the updater with model metadata and initialize the group.
+        _sync_weights_with_worker: Synchronize weights with a vLLM worker.
+        _get_server_weights: Not used - weights must be passed directly.
+        _maybe_map_weights: No mapping needed.
+        all_worker_ids: Returns [0] since we only have one worker.
+
+    .. note::
+        This class assumes the policy is a transformers model that can be loaded by vLLM.
+        The policy must have a state_dict() method that returns the model weights.
+
+    .. warning::
+        The v2=True option is experimental and may have backward-compatibility breaking changes
+        in future releases. However, it is generally considered a better option for working with
+        AsyncVLLM engines and provides improved performance and reliability.
+    """
+
+    def __init__(
+        self,
+        master_address: str | None = None,
+        master_port: int | None = None,
+        model_metadata: dict[str, tuple[torch.dtype, torch.Size]] | None = None,
+        vllm_tp_size: int | None = None,
+    ):
+        torchrl_logger.debug(f"=> in {type(self).__name__}.__init__")
+        self.master_address = master_address
+        self.master_port = master_port
+        self.model_metadata = model_metadata
+        self.initialized_group = None
+        if vllm_tp_size is None:
+            vllm_tp_size = 1
+        self.vllm_tp_size = vllm_tp_size
+        self._model_ref = None
+
+    def init(self, model_metadata: dict[str, tuple[torch.dtype, torch.Size]]) -> None:
+        """Initialize the updater with model metadata and initialize the group.
+
+        Args:
+            model_metadata (dict[str, tuple[torch.dtype, torch.Size]]): The model metadata mapping
+                parameter names to their dtype and shape.
+        """
+        self.model_metadata = model_metadata
+        self.maybe_init_group()
+
+    @property
+    def master_address(self):
+        if self._master_address is None:
+            self._master_address = "localhost"  # get_ip()
+        return self._master_address
+
+    @master_address.setter
+    def master_address(self, value):
+        self._master_address = value
+
+    @property
+    def master_port(self):
+        if self._master_port is None:
+            try:
+                from vllm.utils import get_open_port
+
+                self._master_port = get_open_port()
+            except ImportError:
+                # Fallback if vLLM get_open_port not available
+                import socket
+
+                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                    s.bind(("", 0))
+                    self._master_port = s.getsockname()[1]
+        return self._master_port
+
+    @master_port.setter
+    def master_port(self, value):
+        self._master_port = value
+
+    def _get_model_ref(self):
+        """Get a reference to the model actor, either directly or through Ray.
+
+        Returns:
+            The model reference that can be used for weight updates
+        """
+        if self._model_ref is not None:
+            return self._model_ref
+
+        if hasattr(self.collector, "_collector"):
+            # We're dealing with a RayLLMCollector
+            import ray
+
+            # Get direct reference to the model actor
+            self._model_ref = ray.get(
+                self.collector._collector.get_policy_model.remote()
+            )
+        else:
+            # We're dealing with a local collector
+            self._model_ref = self.collector.policy.model
+
+        return self._model_ref
+
+    def _init_group(self):
+        import ray
+
+        torchrl_logger.debug(f"=> in {type(self).__name__}._init_group")
+        weight_sync_world_size = self.vllm_tp_size + 1
+        torchrl_logger.debug(f"initializing group with {weight_sync_world_size=}...")
+        torchrl_logger.debug(f"vllm_tp_size={self.vllm_tp_size}")
+
+        model_ref = self._get_model_ref()
+
+        torchrl_logger.debug(f"model_ref: {model_ref}")
+        # Initialize the weight update group
+        torchrl_logger.debug("Calling init_weight_update_group...")
+        init_weight_update_group_getter = model_ref.collective_rpc.remote(
+            "init_weight_update_group",
+            args=(self.master_address, self.master_port, 1, weight_sync_world_size),
+        )
+        torchrl_logger.debug("init_weight_update_group remote call succeeded")
+
+        torchrl_logger.debug("Calling stateless_init_process_group within updater...")
+        self.vllm_comm_group = stateless_init_process_group(
+            self.master_address,
+            self.master_port,
+            0,
+            weight_sync_world_size,
+            torch.device("cuda:0"),
+        )
+
+        ray.get(init_weight_update_group_getter)
+        torchrl_logger.debug("init_weight_update_group getter succeeded")
+
+        torchrl_logger.debug("group initialized")
+        self.initialized_group = True
+
+    def maybe_init_group(self):
+        if self.initialized_group is None and self.collector is not None:
+            self._init_group()
+
+    def _sync_weights_with_worker(
+        self,
+        *,
+        worker_id: int | torch.device | None = None,
+        server_weights: TensorDictBase | TensorDictModuleBase | dict | None = None,
+    ) -> None:
+        """Synchronize weights with a vLLM worker.
+
+        Args:
+            worker_id: Not used - we only have one worker.
+            server_weights: The weights to sync. Can be:
+                - TensorDictModuleBase: A policy module whose weights will be extracted
+                - TensorDictBase: A TensorDict containing weights
+                - dict: A regular dict containing weights
+                - None: Will try to get weights from server using _get_server_weights()
+        """
+        if server_weights is None:
+            raise ValueError("server_weights cannot be None for vLLM updater")
+
+        if self.initialized_group is None:
+            raise RuntimeError(
+                "Failed to update weights because sender is not initialized."
+            )
+        if self.model_metadata is None:
+            raise RuntimeError(
+                "Failed to update weights because model metadata is not set. "
+                "In async mode, you must set the model metadata in the training actor "
+                "before any weight updates."
+            )
+
+        model_ref = self._get_model_ref()
+
+        # First broadcast metadata
+        torchrl_logger.debug("broadcasting with update_weight_broadcast")
+        remotes = []
+        for k, (dtype, shape) in self.model_metadata.items():
+            remotes.append(
+                model_ref.collective_rpc.remote(
+                    "update_weight_broadcast", args=(k, dtype, shape)
+                )
+            )
+
+        # # Then update weights
+        # remotes = []
+        # pbar = tqdm.tqdm(server_weights.items(), desc="Updating weights", total=len(server_weights))
+        # for k, val in pbar:
+        #     pbar.set_description(f"Updating {k}")
+        #     remotes.append(model_ref.collective_rpc.remote("update_weight", args=(k, val)))
+        # # ray.get(remotes)
+
+        # if self.vllm_comm_group is not True:
+        torchrl_logger.debug("broadcasting...")
+        for k in self.model_metadata:
+            val = server_weights[k].to(torch.device("cuda:0"))
+            self.vllm_comm_group.broadcast(
+                val,
+                src=0,
+                stream=torch.cuda.current_stream(),
+            )
+            del val
+        import ray
+
+        ray.get(remotes)
+        torchrl_logger.debug("done broadcasting")
+        torch.cuda.synchronize()
+
+    def _get_server_weights(self) -> TensorDictBase | None:
+        """Not used - weights must be passed directly via policy."""
+        return None
+
+    def _maybe_map_weights(
+        self, server_weights: TensorDictBase | TensorDictModuleBase | dict
+    ) -> TensorDictBase:
+        """Map weights from any format to the format expected by vLLM.
+
+        Args:
+            server_weights: The weights to map. Can be:
+                - TensorDictModuleBase: A policy module whose weights will be extracted
+                - TensorDictBase: A TensorDict containing weights
+                - dict: A regular dict containing weights
+
+        Returns:
+            TensorDictBase: The mapped weights in TensorDict format
+        """
+        if isinstance(server_weights, TensorDictModuleBase):
+            # Extract weights from policy module using merge_and_unload for LLMs
+            if not hasattr(server_weights, "model"):
+                raise ValueError("TensorDictModuleBase must have a 'model' attribute")
+            # Check if it's a LoRA model
+            if hasattr(server_weights.model, "merge_and_unload"):
+                state_dict = server_weights.model.merge_and_unload().state_dict()
+            else:
+                state_dict = server_weights.model.state_dict()
+            return TensorDict(state_dict, [])
+        elif isinstance(server_weights, TensorDictBase):
+            return server_weights
+        elif isinstance(server_weights, dict):
+            return TensorDict(server_weights, [])
+        else:
+            raise TypeError(
+                f"server_weights must be TensorDictModuleBase, TensorDictBase or dict, got {type(server_weights)}"
+            )
+
+    @classmethod
+    def get_model_metadata(
+        cls, model: TensorDictModuleBase
+    ) -> dict[str, tuple[torch.dtype, torch.Size]]:
+        """Get the model metadata from a model.
+
+        Args:
+            model (TensorDictModuleBase): The model to get the metadata from.
+                Must be a TransformersWrapper or equivalent.
+
+        Returns:
+            dict[str, tuple[torch.dtype, torch.Size]]: The model metadata.
+        """
+        # Check if the model has a LoRA adapter
+        if hasattr(model.model, "merge_and_unload"):
+            sd = model.model.merge_and_unload().state_dict()
+        else:
+            sd = model.model.state_dict()
+        model_metadata = {k: (v.dtype, v.shape) for k, v in sd.items()}
+        return model_metadata
+
+    def all_worker_ids(self) -> list[int]:
+        """Returns [0] since we only have one worker."""
+        return [0]
+
+    def register_collector(self, collector: BaseCollector):  # noqa: F821
+        result = super().register_collector(collector)
+        self.register_post_hook(collector.increment_version)
+        return result