diff --git a/.github/workflows/rl_test.yaml b/.github/workflows/rl_test.yaml index 39d459b2ab..26e7508f0f 100644 --- a/.github/workflows/rl_test.yaml +++ b/.github/workflows/rl_test.yaml @@ -35,7 +35,8 @@ jobs: runs-on: linux.g5.12xlarge.nvidia.gpu strategy: matrix: - python-version: ['3.9', '3.10', '3.11'] + # torchrl 0.13.3 requires Python >= 3.10 + python-version: ['3.10', '3.11'] torch-version: ["stable", "nightly"] # Do not run against nightlies on PR exclude: diff --git a/pyproject.toml b/pyproject.toml index 3106a27aea..9661d01478 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,11 +65,8 @@ dev = [ "expecttest", ] async_rl = [ - # We will update this to a stable release soon! - "torchrl @ git+https://github.com/pytorch/rl@0475cbf64de0cdf5f185a77badaa6eb241c63b0b", - - # We will update this to a stable release soon! - "tensordict @ git+https://github.com/pytorch/tensordict@c61d045aaadf6c0625706a3670fc6a741f31f1b0", + "torchrl==0.13.3", + "tensordict==0.13.0", # Regular PyPI dependencies "ray", diff --git a/torchtune/dev/rl/utils/_torchrl_compat.py b/torchtune/dev/rl/utils/_torchrl_compat.py new file mode 100644 index 0000000000..bc94f83886 --- /dev/null +++ b/torchtune/dev/rl/utils/_torchrl_compat.py @@ -0,0 +1,56 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# torchrl >= 0.13 removed the collector classes the async RL data path +# subclasses. Import the real classes when available and fall back to +# import-compatible stand-ins so the data path can be imported and unit-tested +# against released torchrl. +try: + from torchrl.collectors import ( + SyncDataCollector, + WeightUpdateReceiverBase, + WeightUpdateSenderBase, + ) +except ImportError: + from torchrl.collectors import BaseCollector + + + class SyncDataCollector(BaseCollector): + """Import-compatible stand-in for ``torchrl.collectors.SyncDataCollector``. + + The stand-in accepts the constructor arguments of the original class + and sets up the attributes the subclass relies on, but real data + collection is not supported. + """ + + def __init__(self, *args, **kwargs): + super().__init__() + env = kwargs.get("create_env_fn") + self.env = env() if callable(env) else env + self.policy = kwargs.get("policy") + self.frames_per_batch = kwargs.get("frames_per_batch", -1) + self.total_frames = kwargs.get("total_frames", -1) + self.weight_update_receiver = kwargs.get("weight_update_receiver") + self.weight_update_sender = kwargs.get("weight_update_sender") + self.reset_at_each_iter = kwargs.get("reset_at_each_iter", False) + self.replay_buffer = None + self._shuttle = None + + def _setup_data(self, *args, **kwargs): + raise NotImplementedError( + "SyncDataCollector is only available with torchrl < 0.13." + ) + + def _update_traj_ids(self, data): + pass + + + class WeightUpdateReceiverBase: + pass + + + class WeightUpdateSenderBase: + pass diff --git a/torchtune/dev/rl/workers/datacollectors/sync.py b/torchtune/dev/rl/workers/datacollectors/sync.py index b9c73d7771..b7590aafb4 100644 --- a/torchtune/dev/rl/workers/datacollectors/sync.py +++ b/torchtune/dev/rl/workers/datacollectors/sync.py @@ -17,15 +17,15 @@ from tensordict import lazy_stack, NonTensorStack, TensorDictBase from torchdata.stateful_dataloader import StatefulDataLoader from torchdata.stateful_dataloader.sampler import StatefulDistributedSampler -from torchrl.collectors import ( - SyncDataCollector, - WeightUpdateReceiverBase, - WeightUpdateSenderBase, -) from torchtune import utils from torchtune.dev.rl.datatypes import Trajectory from torchtune.dev.rl.utils import stateless_init_process_group +from torchtune.dev.rl.utils._torchrl_compat import ( + SyncDataCollector, + WeightUpdateReceiverBase, + WeightUpdateSenderBase, +) from vllm import LLM from vllm.worker.worker import Worker @@ -103,7 +103,16 @@ def __init__( ) # local import below LLM call to avoid vLLM no CUDA GPUs available error - from torchrl.envs import LLMEnv + try: + from torchrl.envs import LLMEnv + except ImportError: + # torchrl >= 0.13 removed LLMEnv + class LLMEnv: + @classmethod + def from_dataloader(cls, *args, **kwargs): + raise NotImplementedError( + "LLMEnv is only available with torchrl < 0.13." + ) env = LLMEnv.from_dataloader( dataloader=dataloader, diff --git a/torchtune/dev/rl/workers/weight_updaters/weight_updater.py b/torchtune/dev/rl/workers/weight_updaters/weight_updater.py index 4633c7d1fb..b4e35b1a04 100644 --- a/torchtune/dev/rl/workers/weight_updaters/weight_updater.py +++ b/torchtune/dev/rl/workers/weight_updaters/weight_updater.py @@ -5,7 +5,7 @@ # LICENSE file in the root directory of this source tree. import ray -from torchrl.collectors import WeightUpdateReceiverBase +from torchtune.dev.rl.utils._torchrl_compat import WeightUpdateReceiverBase class VLLMHFWeightUpdateReceiver(WeightUpdateReceiverBase): diff --git a/torchtune/modules/common_utils.py b/torchtune/modules/common_utils.py index ccd0f0a236..10dd71fa8a 100644 --- a/torchtune/modules/common_utils.py +++ b/torchtune/modules/common_utils.py @@ -16,7 +16,11 @@ import torch.nn as nn from torch._subclasses.fake_tensor import FakeTensorConverter, FakeTensorMode -from torchao.quantization import NF4Tensor +try: + from torchao.quantization import NF4Tensor +except ImportError: + # torchao >= 0.17 moved NF4Tensor to torchao.dtypes.nf4tensor + from torchao.dtypes.nf4tensor import NF4Tensor _use_low_cpu_ram: bool = False diff --git a/torchtune/modules/low_precision/_register_nf4_dispatch_ops.py b/torchtune/modules/low_precision/_register_nf4_dispatch_ops.py index e1de27d172..c2395791c4 100644 --- a/torchtune/modules/low_precision/_register_nf4_dispatch_ops.py +++ b/torchtune/modules/low_precision/_register_nf4_dispatch_ops.py @@ -5,8 +5,13 @@ # LICENSE file in the root directory of this source tree. import torch -from torchao.quantization import to_nf4 -from torchao.quantization.quantize_.workflows.nf4.nf4_tensor import implements as nf4_tensor_impl +try: + from torchao.quantization import to_nf4 + from torchao.quantization.quantize_.workflows.nf4.nf4_tensor import implements as nf4_tensor_impl +except ImportError: + # torchao >= 0.17 moved these to torchao.dtypes.nf4tensor + from torchao.dtypes.nf4tensor import implements as nf4_tensor_impl + from torchao.dtypes.nf4tensor import to_nf4 @nf4_tensor_impl([torch.ops.aten.clone.default]) diff --git a/torchtune/modules/low_precision/nf4_linear.py b/torchtune/modules/low_precision/nf4_linear.py index 64951354ce..81b473fa85 100644 --- a/torchtune/modules/low_precision/nf4_linear.py +++ b/torchtune/modules/low_precision/nf4_linear.py @@ -9,8 +9,12 @@ import torch import torch.nn as nn -from torchao.quantization import to_nf4 -from torchao.quantization.quantize_.workflows.nf4.nf4_tensor import linear_nf4 +try: + from torchao.quantization import to_nf4 + from torchao.quantization.quantize_.workflows.nf4.nf4_tensor import linear_nf4 +except ImportError: + # torchao >= 0.17 moved these to torchao.dtypes.nf4tensor + from torchao.dtypes.nf4tensor import linear_nf4, to_nf4 class FrozenNF4Linear(nn.Linear): diff --git a/torchtune/modules/peft/dora.py b/torchtune/modules/peft/dora.py index 4a250a3511..5881b09621 100644 --- a/torchtune/modules/peft/dora.py +++ b/torchtune/modules/peft/dora.py @@ -12,8 +12,12 @@ from torch import nn -from torchao.quantization import to_nf4 -from torchao.quantization.quantize_.workflows.nf4.nf4_tensor import linear_nf4 +try: + from torchao.quantization import to_nf4 + from torchao.quantization.quantize_.workflows.nf4.nf4_tensor import linear_nf4 +except ImportError: + # torchao >= 0.17 moved these to torchao.dtypes.nf4tensor + from torchao.dtypes.nf4tensor import linear_nf4, to_nf4 from torchtune.modules.low_precision import _register_nf4_dispatch_ops # noqa: F401 from torchtune.modules.peft import AdapterModule diff --git a/torchtune/modules/peft/lora.py b/torchtune/modules/peft/lora.py index f115bb89b2..b7af50dd70 100644 --- a/torchtune/modules/peft/lora.py +++ b/torchtune/modules/peft/lora.py @@ -12,8 +12,12 @@ from torch import nn -from torchao.quantization import to_nf4 -from torchao.quantization.quantize_.workflows.nf4.nf4_tensor import linear_nf4 +try: + from torchao.quantization import to_nf4 + from torchao.quantization.quantize_.workflows.nf4.nf4_tensor import linear_nf4 +except ImportError: + # torchao >= 0.17 moved these to torchao.dtypes.nf4tensor + from torchao.dtypes.nf4tensor import linear_nf4, to_nf4 from torchtune.modules.low_precision import _register_nf4_dispatch_ops # noqa: F401 from torchtune.modules.peft import AdapterModule diff --git a/torchtune/training/_activation_offloading.py b/torchtune/training/_activation_offloading.py index 852cb51ca5..b14a7cf9dd 100644 --- a/torchtune/training/_activation_offloading.py +++ b/torchtune/training/_activation_offloading.py @@ -13,7 +13,11 @@ import torchao from torch import nn from torch.autograd.graph import saved_tensors_hooks -from torchao.quantization import NF4Tensor +try: + from torchao.quantization import NF4Tensor +except ImportError: + # torchao >= 0.17 moved NF4Tensor to torchao.dtypes.nf4tensor + from torchao.dtypes.nf4tensor import NF4Tensor from torchtune.modules import TiedLinear from torchtune.utils import get_logger diff --git a/torchtune/training/_distributed.py b/torchtune/training/_distributed.py index bf52889812..cdbda9115d 100644 --- a/torchtune/training/_distributed.py +++ b/torchtune/training/_distributed.py @@ -36,7 +36,11 @@ from torch.nn.attention.flex_attention import BlockMask from torch.nn.modules.module import _IncompatibleKeys from torch.optim import Optimizer -from torchao.quantization import NF4Tensor, to_nf4 +try: + from torchao.quantization import NF4Tensor, to_nf4 +except ImportError: + # torchao >= 0.17 moved these to torchao.dtypes.nf4tensor + from torchao.dtypes.nf4tensor import NF4Tensor, to_nf4 from torchtune.modules import TransformerDecoder from torchtune.modules.attention import MultiHeadAttention from torchtune.modules.model_fusion import DeepFusionModel, EarlyFusionModel