torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

torchrl/trainers/algorithms/configs/utils.py

@@ -0,0 +1,252 @@
+# 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 dataclasses import dataclass
+
+from torchrl.trainers.algorithms.configs.common import ConfigBase
+
+
+@dataclass
+class AdamConfig(ConfigBase):
+    """Configuration for Adam optimizer."""
+
+    lr: float = 1e-3
+    betas: tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-4
+    weight_decay: float = 0.0
+    amsgrad: bool = False
+    _target_: str = "torch.optim.Adam"
+    _partial_: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for Adam optimizer configurations."""
+
+
+@dataclass
+class AdamWConfig(ConfigBase):
+    """Configuration for AdamW optimizer."""
+
+    lr: float = 1e-3
+    betas: tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-8
+    weight_decay: float = 1e-2
+    amsgrad: bool = False
+    maximize: bool = False
+    foreach: bool | None = None
+    capturable: bool = False
+    differentiable: bool = False
+    fused: bool | None = None
+    _target_: str = "torch.optim.AdamW"
+    _partial_: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for AdamW optimizer configurations."""
+
+
+@dataclass
+class AdamaxConfig(ConfigBase):
+    """Configuration for Adamax optimizer."""
+
+    lr: float = 2e-3
+    betas: tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-8
+    weight_decay: float = 0.0
+    _target_: str = "torch.optim.Adamax"
+    _partial_: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for Adamax optimizer configurations."""
+
+
+@dataclass
+class SGDConfig(ConfigBase):
+    """Configuration for SGD optimizer."""
+
+    lr: float = 1e-3
+    momentum: float = 0.0
+    dampening: float = 0.0
+    weight_decay: float = 0.0
+    nesterov: bool = False
+    maximize: bool = False
+    foreach: bool | None = None
+    differentiable: bool = False
+    _target_: str = "torch.optim.SGD"
+    _partial_: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for SGD optimizer configurations."""
+
+
+@dataclass
+class RMSpropConfig(ConfigBase):
+    """Configuration for RMSprop optimizer."""
+
+    lr: float = 1e-2
+    alpha: float = 0.99
+    eps: float = 1e-8
+    weight_decay: float = 0.0
+    momentum: float = 0.0
+    centered: bool = False
+    maximize: bool = False
+    foreach: bool | None = None
+    differentiable: bool = False
+    _target_: str = "torch.optim.RMSprop"
+    _partial_: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for RMSprop optimizer configurations."""
+
+
+@dataclass
+class AdagradConfig(ConfigBase):
+    """Configuration for Adagrad optimizer."""
+
+    lr: float = 1e-2
+    lr_decay: float = 0.0
+    weight_decay: float = 0.0
+    initial_accumulator_value: float = 0.0
+    eps: float = 1e-10
+    maximize: bool = False
+    foreach: bool | None = None
+    differentiable: bool = False
+    _target_: str = "torch.optim.Adagrad"
+    _partial_: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for Adagrad optimizer configurations."""
+
+
+@dataclass
+class AdadeltaConfig(ConfigBase):
+    """Configuration for Adadelta optimizer."""
+
+    lr: float = 1.0
+    rho: float = 0.9
+    eps: float = 1e-6
+    weight_decay: float = 0.0
+    foreach: bool | None = None
+    maximize: bool = False
+    differentiable: bool = False
+    _target_: str = "torch.optim.Adadelta"
+    _partial_: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for Adadelta optimizer configurations."""
+
+
+@dataclass
+class RpropConfig(ConfigBase):
+    """Configuration for Rprop optimizer."""
+
+    lr: float = 1e-2
+    etas: tuple[float, float] = (0.5, 1.2)
+    step_sizes: tuple[float, float] = (1e-6, 50.0)
+    foreach: bool | None = None
+    maximize: bool = False
+    differentiable: bool = False
+    _target_: str = "torch.optim.Rprop"
+    _partial_: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for Rprop optimizer configurations."""
+
+
+@dataclass
+class ASGDConfig(ConfigBase):
+    """Configuration for ASGD optimizer."""
+
+    lr: float = 1e-2
+    lambd: float = 1e-4
+    alpha: float = 0.75
+    t0: float = 1e6
+    weight_decay: float = 0.0
+    foreach: bool | None = None
+    maximize: bool = False
+    differentiable: bool = False
+    _target_: str = "torch.optim.ASGD"
+    _partial_: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for ASGD optimizer configurations."""
+
+
+@dataclass
+class LBFGSConfig(ConfigBase):
+    """Configuration for LBFGS optimizer."""
+
+    lr: float = 1.0
+    max_iter: int = 20
+    max_eval: int | None = None
+    tolerance_grad: float = 1e-7
+    tolerance_change: float = 1e-9
+    history_size: int = 100
+    line_search_fn: str | None = None
+    _target_: str = "torch.optim.LBFGS"
+    _partial_: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for LBFGS optimizer configurations."""
+
+
+@dataclass
+class RAdamConfig(ConfigBase):
+    """Configuration for RAdam optimizer."""
+
+    lr: float = 1e-3
+    betas: tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-8
+    weight_decay: float = 0.0
+    _target_: str = "torch.optim.RAdam"
+    _partial_: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for RAdam optimizer configurations."""
+
+
+@dataclass
+class NAdamConfig(ConfigBase):
+    """Configuration for NAdam optimizer."""
+
+    lr: float = 2e-3
+    betas: tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-8
+    weight_decay: float = 0.0
+    momentum_decay: float = 4e-3
+    foreach: bool | None = None
+    _target_: str = "torch.optim.NAdam"
+    _partial_: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for NAdam optimizer configurations."""
+
+
+@dataclass
+class SparseAdamConfig(ConfigBase):
+    """Configuration for SparseAdam optimizer."""
+
+    lr: float = 1e-3
+    betas: tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-8
+    _target_: str = "torch.optim.SparseAdam"
+    _partial_: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for SparseAdam optimizer configurations."""
+
+
+@dataclass
+class LionConfig(ConfigBase):
+    """Configuration for Lion optimizer."""
+
+    lr: float = 1e-4
+    betas: tuple[float, float] = (0.9, 0.99)
+    weight_decay: float = 0.0
+    _target_: str = "torch.optim.Lion"
+    _partial_: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for Lion optimizer configurations."""

torchrl/trainers/algorithms/configs/weight_sync_schemes.py

@@ -0,0 +1,191 @@
+# 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 dataclasses import dataclass
+from typing import Any
+
+from torchrl.trainers.algorithms.configs.common import ConfigBase
+
+
+@dataclass
+class WeightSyncSchemeConfig(ConfigBase):
+    """Base configuration for weight synchronization schemes."""
+
+    _target_: str = "torchrl.weight_update.WeightSyncScheme"
+    _partial_: bool = False
+
+    # Common argument for all schemes
+    strategy: str = "tensordict"  # "tensordict" or "state_dict"
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for weight sync scheme configurations."""
+
+
+@dataclass
+class MultiProcessWeightSyncSchemeConfig(ConfigBase):
+    """Configuration for MultiProcessWeightSyncScheme.
+
+    Weight synchronization for multiprocess operations using pipes.
+    This scheme creates transports that communicate via multiprocessing pipes.
+    """
+
+    _target_: str = "torchrl.weight_update.MultiProcessWeightSyncScheme"
+    _partial_: bool = False
+
+    strategy: str = "tensordict"  # "tensordict" or "state_dict"
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for multiprocess weight sync scheme configurations."""
+
+
+@dataclass
+class SharedMemWeightSyncSchemeConfig(ConfigBase):
+    """Configuration for SharedMemWeightSyncScheme.
+
+    Weight synchronization using shared memory for in-place weight updates.
+    Workers automatically see weight updates without explicit message passing.
+    """
+
+    _target_: str = "torchrl.weight_update.SharedMemWeightSyncScheme"
+    _partial_: bool = False
+
+    strategy: str = "tensordict"  # "tensordict" or "state_dict"
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for shared memory weight sync scheme configurations."""
+
+
+@dataclass
+class NoWeightSyncSchemeConfig(ConfigBase):
+    """Configuration for NoWeightSyncScheme.
+
+    No-op weight synchronization scheme that disables weight synchronization entirely.
+    """
+
+    _target_: str = "torchrl.weight_update.NoWeightSyncScheme"
+    _partial_: bool = False
+
+    strategy: str = "tensordict"  # Not really used, but kept for consistency
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for no weight sync scheme configurations."""
+
+
+@dataclass
+class RayWeightSyncSchemeConfig(ConfigBase):
+    """Configuration for RayWeightSyncScheme.
+
+    Weight synchronization for Ray distributed computing. Uses Ray's object store
+    and remote calls to synchronize weights across distributed workers (Ray actors).
+    """
+
+    _target_: str = "torchrl.weight_update.RayWeightSyncScheme"
+    _partial_: bool = False
+
+    strategy: str = "tensordict"  # "tensordict" or "state_dict"
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for Ray weight sync scheme configurations."""
+
+
+@dataclass
+class RayModuleTransformSchemeConfig(ConfigBase):
+    """Configuration for RayModuleTransformScheme.
+
+    Weight synchronization for RayModuleTransform actors. This scheme is designed
+    specifically for updating models hosted within Ray actors.
+    """
+
+    _target_: str = "torchrl.weight_update.RayModuleTransformScheme"
+    _partial_: bool = False
+
+    strategy: str = "tensordict"  # "tensordict" or "state_dict"
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for Ray module transform scheme configurations."""
+
+
+@dataclass
+class RPCWeightSyncSchemeConfig(ConfigBase):
+    """Configuration for RPCWeightSyncScheme.
+
+    Weight synchronization for torch.distributed.rpc. Uses RPC calls to synchronize
+    weights across distributed workers.
+    """
+
+    _target_: str = "torchrl.weight_update.RPCWeightSyncScheme"
+    _partial_: bool = False
+
+    strategy: str = "tensordict"  # "tensordict" or "state_dict"
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for RPC weight sync scheme configurations."""
+
+
+@dataclass
+class DistributedWeightSyncSchemeConfig(ConfigBase):
+    """Configuration for DistributedWeightSyncScheme.
+
+    Weight synchronization for torch.distributed. Uses torch.distributed primitives
+    (send/recv) to synchronize weights across distributed workers.
+    """
+
+    _target_: str = "torchrl.weight_update.DistributedWeightSyncScheme"
+    _partial_: bool = False
+
+    backend: str = "gloo"  # "gloo", "nccl", etc.
+    sync: bool = True
+    strategy: str = "tensordict"  # "tensordict" or "state_dict"
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for distributed weight sync scheme configurations."""
+
+
+@dataclass
+class VLLMWeightSyncSchemeConfig(ConfigBase):
+    """Configuration for VLLMWeightSyncScheme.
+
+    Weight synchronization scheme for vLLM engines using collective communication (NCCL).
+    Broadcasts weights from a trainer to vLLM inference workers with parallelism support.
+    """
+
+    _target_: str = "torchrl.weight_update.llm.VLLMWeightSyncScheme"
+    _partial_: bool = False
+
+    master_address: str | None = None  # Defaults to "localhost"
+    master_port: int | None = None  # Auto-assigned if None
+    gpus_per_replica: int = 1  # tp_size × dp_size × pp_size
+    num_replicas: int = 1
+    strategy: str = "tensordict"  # "tensordict" or "state_dict"
+    device: Any = 0  # torch.device | str | int
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for vLLM weight sync scheme configurations."""
+
+
+@dataclass
+class VLLMDoubleBufferSyncSchemeConfig(ConfigBase):
+    """Configuration for VLLMDoubleBufferSyncScheme.
+
+    Weight synchronization scheme for vLLM using double-buffered memory-mapped storage.
+    Uses TensorDict's memory-mapping capabilities to transfer weights via filesystem.
+    """
+
+    _target_: str = "torchrl.weight_update.llm.VLLMDoubleBufferSyncScheme"
+    _partial_: bool = False
+
+    remote_addr: str | None = None  # Directory path where sender writes weights
+    local_addr: str | None = None  # Directory path where receiver reads weights
+    num_threads: int = 1  # Number of threads for memmap operations
+    strategy: str = "tensordict"  # "tensordict" or "state_dict"
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for vLLM double buffer sync scheme configurations."""
+        if self.remote_addr is None:
+            raise ValueError("remote_addr is required for VLLMDoubleBufferSyncScheme")
+        if self.local_addr is None:
+            self.local_addr = self.remote_addr

torchrl/trainers/algorithms/configs/weight_update.py

@@ -0,0 +1,159 @@
+# 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 dataclasses import dataclass
+from typing import Any
+
+from torchrl.trainers.algorithms.configs.common import ConfigBase
+
+
+@dataclass
+class WeightUpdaterConfig(ConfigBase):
+    """Base configuration for weight updaters."""
+
+    _target_: str = "torchrl.collectors.WeightUpdaterBase"
+    _partial_: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for weight updater configurations."""
+
+
+@dataclass
+class VanillaWeightUpdaterConfig(ConfigBase):
+    """Configuration for VanillaWeightUpdater.
+
+    A simple implementation for updating local policy weights by directly
+    fetching them from a specified source.
+    """
+
+    _target_: str = "torchrl.collectors.VanillaWeightUpdater"
+    _partial_: bool = True
+
+    # Constructor arguments
+    weight_getter: Any = None  # Callable[[], TensorDictBase] | None
+    policy_weights: Any = None  # TensorDictBase
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for vanilla weight updater configurations."""
+
+
+@dataclass
+class MultiProcessedWeightUpdaterConfig(ConfigBase):
+    """Configuration for MultiProcessedWeightUpdater.
+
+    A remote weight updater for synchronizing policy weights across multiple
+    processes or devices in a multiprocessed environment.
+    """
+
+    _target_: str = "torchrl.collectors.MultiProcessedWeightUpdater"
+    _partial_: bool = True
+
+    # Constructor arguments
+    get_server_weights: Any = None  # Callable[[], TensorDictBase] | None
+    policy_weights: Any = None  # dict[torch.device, TensorDictBase]
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for multiprocessed weight updater configurations."""
+
+
+@dataclass
+class RayWeightUpdaterConfig(ConfigBase):
+    """Configuration for RayWeightUpdater.
+
+    A remote weight updater for synchronizing policy weights across remote
+    workers using Ray's distributed computing capabilities.
+    """
+
+    _target_: str = "torchrl.collectors.RayWeightUpdater"
+    _partial_: bool = True
+
+    # Constructor arguments
+    policy_weights: Any = None  # TensorDictBase
+    remote_collectors: Any = None  # list
+    max_interval: int = 0  # int
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for Ray weight updater configurations."""
+
+
+@dataclass
+class RPCWeightUpdaterConfig(ConfigBase):
+    """Configuration for RPCWeightUpdater.
+
+    A remote weight updater for synchronizing policy weights across remote
+    workers using RPC communication.
+    """
+
+    _target_: str = "torchrl.collectors.distributed.RPCWeightUpdater"
+    _partial_: bool = True
+
+    # Constructor arguments
+    collector_infos: Any = None
+    collector_class: Any = None
+    collector_rrefs: Any = None
+    policy_weights: Any = None  # TensorDictBase
+    num_workers: int = 0
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for RPC weight updater configurations."""
+
+
+@dataclass
+class DistributedWeightUpdaterConfig(ConfigBase):
+    """Configuration for DistributedWeightUpdater.
+
+    A remote weight updater for synchronizing policy weights across distributed
+    workers using a dictionary-like store for communication.
+    """
+
+    _target_: str = "torchrl.collectors.distributed.DistributedWeightUpdater"
+    _partial_: bool = True
+
+    # Constructor arguments
+    store: Any = None  # dict[str, str]
+    policy_weights: Any = None  # TensorDictBase
+    num_workers: int = 0
+    sync: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for distributed weight updater configurations."""
+
+
+@dataclass
+class RemoteModuleWeightUpdaterConfig(ConfigBase):
+    """Configuration for RemoteModuleWeightUpdater.
+
+    A weight updater for remote nn.Modules that requires explicit weight passing.
+    Used when the master collector doesn't have direct access to worker weights.
+    """
+
+    _target_: str = "torchrl.collectors.RemoteModuleWeightUpdater"
+    _partial_: bool = True
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for remote module weight updater configurations."""
+
+
+@dataclass
+class vLLMUpdaterConfig(ConfigBase):
+    """Configuration for vLLMUpdater.
+
+    A weight updater that sends weights to vLLM workers, supporting both local
+    vLLM instances and remote Ray actors for LLM inference.
+    """
+
+    _target_: str = "torchrl.collectors.llm.vLLMUpdater"
+    _partial_: bool = True
+
+    # Constructor arguments
+    master_address: str | None = None
+    master_port: int | None = None
+    model_metadata: Any = None  # dict[str, tuple[torch.dtype, torch.Size]] | None
+    vllm_tp_size: int | None = None
+
+    def __post_init__(self) -> None:
+        """Post-initialization hook for vLLM updater configurations."""