fix(dev): support current torchao and torchrl releases

.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:

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",

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

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,

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):

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
 

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])

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):

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
 

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
 

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

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