torchrl 0.11.0__cp314-cp314-manylinux_2_28_aarch64.whl

benchmarks/test_compressed_storage_benchmark.py

@@ -0,0 +1,145 @@
+# 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.
+import io
+import pickle
+
+import pytest
+import torch
+
+
+try:
+    from safetensors.torch import save
+except ImportError:
+    save = None
+
+
+class TestCompressedStorageBenchmark:
+    """Benchmark tests for CompressedListStorage."""
+
+    @staticmethod
+    def make_compressible_mock_data(num_experiences: int, device=None) -> dict:
+        """Easily compressible data for testing."""
+        if device is None:
+            device = torch.device("cpu")
+
+        return {
+            "observations": torch.zeros(
+                (num_experiences, 4, 84, 84),
+                dtype=torch.uint8,
+                device=device,
+            ),
+            "actions": torch.zeros((num_experiences,), device=device),
+            "rewards": torch.zeros((num_experiences,), device=device),
+            "next_observations": torch.zeros(
+                (num_experiences, 4, 84, 84),
+                dtype=torch.uint8,
+                device=device,
+            ),
+            "terminations": torch.zeros(
+                (num_experiences,), dtype=torch.bool, device=device
+            ),
+            "truncations": torch.zeros(
+                (num_experiences,), dtype=torch.bool, device=device
+            ),
+            "batch_size": [num_experiences],
+        }
+
+    @staticmethod
+    def make_uncompressible_mock_data(num_experiences: int, device=None) -> dict:
+        """Uncompressible data for testing."""
+        if device is None:
+            device = torch.device("cpu")
+        return {
+            "observations": torch.randn(
+                (num_experiences, 4, 84, 84),
+                dtype=torch.float32,
+                device=device,
+            ),
+            "actions": torch.randint(0, 10, (num_experiences,), device=device),
+            "rewards": torch.randn(
+                (num_experiences,), dtype=torch.float32, device=device
+            ),
+            "next_observations": torch.randn(
+                (num_experiences, 4, 84, 84),
+                dtype=torch.float32,
+                device=device,
+            ),
+            "terminations": torch.rand((num_experiences,), device=device)
+            < 0.2,  # ~20% True
+            "truncations": torch.rand((num_experiences,), device=device)
+            < 0.1,  # ~10% True
+            "batch_size": [num_experiences],
+        }
+
+    @pytest.mark.benchmark(
+        group="tensor_serialization_speed",
+        min_time=0.1,
+        max_time=0.5,
+        min_rounds=5,
+        disable_gc=True,
+        warmup=False,
+    )
+    @pytest.mark.parametrize(
+        "serialization_method",
+        ["pickle", "torch.save", "untyped_storage", "numpy", "safetensors"],
+    )
+    def test_tensor_to_bytestream_speed(self, benchmark, serialization_method: str):
+        """Benchmark the speed of different tensor serialization methods.
+
+        TODO: we might need to also test which methods work on the gpu.
+        pytest benchmarks/test_compressed_storage_benchmark.py::TestCompressedStorageBenchmark::test_tensor_to_bytestream_speed -v --benchmark-only --benchmark-sort='mean' --benchmark-columns='mean, ops'
+
+        ------------------------ benchmark 'tensor_to_bytestream_speed': 5 tests -------------------------
+        Name (time in us)                                           Mean (smaller is better)   OPS (bigger is better)
+        --------------------------------------------------------------------------------------------------
+        test_tensor_serialization_speed[numpy]                    2.3520 (1.0)      425,162.1779 (1.0)
+        test_tensor_serialization_speed[safetensors]             14.7170 (6.26)      67,948.7129 (0.16)
+        test_tensor_serialization_speed[pickle]                  19.0711 (8.11)      52,435.3333 (0.12)
+        test_tensor_serialization_speed[torch.save]              32.0648 (13.63)     31,186.8261 (0.07)
+        test_tensor_serialization_speed[untyped_storage]     38,227.0224 (>1000.0)       26.1595 (0.00)
+        --------------------------------------------------------------------------------------------------
+        """
+
+        def serialize_with_pickle(data: torch.Tensor) -> bytes:
+            """Serialize tensor using pickle."""
+            buffer = io.BytesIO()
+            pickle.dump(data, buffer)
+            return buffer.getvalue()
+
+        def serialize_with_untyped_storage(data: torch.Tensor) -> bytes:
+            """Serialize tensor using torch's built-in method."""
+            return bytes(data.untyped_storage())
+
+        def serialize_with_numpy(data: torch.Tensor) -> bytes:
+            """Serialize tensor using numpy."""
+            return data.numpy().tobytes()
+
+        def serialize_with_safetensors(data: torch.Tensor) -> bytes:
+            return save({"0": data})
+
+        def serialize_with_torch(data: torch.Tensor) -> bytes:
+            """Serialize tensor using torch's built-in method."""
+            buffer = io.BytesIO()
+            torch.save(data, buffer)
+            return buffer.getvalue()
+
+        # Benchmark each serialization method
+        if serialization_method == "pickle":
+            serialize_fn = serialize_with_pickle
+        elif serialization_method == "torch.save":
+            serialize_fn = serialize_with_torch
+        elif serialization_method == "untyped_storage":
+            serialize_fn = serialize_with_untyped_storage
+        elif serialization_method == "numpy":
+            serialize_fn = serialize_with_numpy
+        elif serialization_method == "safetensors":
+            serialize_fn = serialize_with_safetensors
+        else:
+            raise ValueError(f"Unknown serialization method: {serialization_method}")
+
+        data = self.make_compressible_mock_data(1).get("observations")
+
+        # Run the actual benchmark
+        benchmark(serialize_fn, data)

benchmarks/test_envs_benchmark.py

@@ -0,0 +1,133 @@
+# 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.
+import argparse
+
+import pytest
+import torch
+
+from tensordict import TensorDict
+from torchrl.envs import ParallelEnv, SerialEnv, step_mdp, StepCounter, TransformedEnv
+from torchrl.envs.libs.dm_control import DMControlEnv
+
+
+def make_simple_env():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = DMControlEnv("cheetah", "run", device=device)
+    env.rollout(3)
+    return ((env,), {})
+
+
+def make_transformed_env():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = TransformedEnv(DMControlEnv("cheetah", "run", device=device), StepCounter(50))
+    env.rollout(3)
+    return ((env,), {})
+
+
+def make_serial_env():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = SerialEnv(3, lambda: DMControlEnv("cheetah", "run", device=device))
+    env.rollout(3)
+    return ((env,), {})
+
+
+def make_parallel_env():
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    env = ParallelEnv(3, lambda: DMControlEnv("cheetah", "run", device=device))
+    env.rollout(3)
+    return ((env,), {})
+
+
+def make_nested_td():
+    return TensorDict(
+        {
+            ("agent", "action"): 0,
+            ("agent", "done"): 0,
+            ("agent", "obs"): 0,
+            ("agent", "other"): 0,
+            ("next", "agent", "action"): 1,
+            ("next", "agent", "reward"): 1,
+            ("next", "agent", "done"): 1,
+            ("next", "agent", "obs"): 1,
+        },
+        [],
+    )
+
+
+def make_flat_td():
+    return TensorDict(
+        {
+            "action": 0,
+            "done": 0,
+            "obs": 0,
+            "other": 0,
+            ("next", "action"): 1,
+            ("next", "reward"): 1,
+            ("next", "done"): 1,
+            ("next", "obs"): 1,
+        },
+        [],
+    )
+
+
+def execute_env(env):
+    env.rollout(1000, break_when_any_done=False)
+
+
+def test_simple(benchmark):
+    (c,), _ = make_simple_env()
+    benchmark(execute_env, c)
+
+
+def test_transformed(benchmark):
+    (c,), _ = make_transformed_env()
+    benchmark(execute_env, c)
+
+
+def test_serial(benchmark):
+    (c,), _ = make_serial_env()
+    benchmark(execute_env, c)
+
+
+def test_parallel(benchmark):
+    (c,), _ = make_parallel_env()
+    benchmark(execute_env, c)
+
+
+@pytest.mark.parametrize("nested", [True, False])
+@pytest.mark.parametrize("keep_other", [True, False])
+@pytest.mark.parametrize("exclude_reward", [True, False])
+@pytest.mark.parametrize("exclude_done", [True, False])
+@pytest.mark.parametrize("exclude_action", [True, False])
+def test_step_mdp_speed(
+    benchmark, nested, keep_other, exclude_reward, exclude_done, exclude_action
+):
+    if nested:
+        td = make_nested_td()
+        reward_key = ("agent", "reward")
+        done_key = ("agent", "done")
+        action_key = ("agent", "action")
+    else:
+        td = make_flat_td()
+        reward_key = "reward"
+        done_key = "done"
+        action_key = "action"
+
+    benchmark(
+        step_mdp,
+        td,
+        action_keys=action_key,
+        reward_keys=reward_key,
+        done_keys=done_key,
+        keep_other=keep_other,
+        exclude_reward=exclude_reward,
+        exclude_done=exclude_done,
+        exclude_action=exclude_action,
+    )
+
+
+if __name__ == "__main__":
+    args, unknown = argparse.ArgumentParser().parse_known_args()
+    pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

benchmarks/test_llm.py

@@ -0,0 +1,101 @@
+# 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
+
+import pytest
+import torch
+from tensordict import set_list_to_stack, TensorDict
+from torchrl.data.llm import History
+from torchrl.modules.llm.policies.common import ChatHistory
+from torchrl.modules.llm.policies.transformers_wrapper import TransformersWrapper
+
+_has_transformers = importlib.import_module("transformers") is not None
+
+# Skip all these tests if gpu is not available
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="GPU not available"
+)
+
+
+@pytest.fixture(scope="module")
+def transformers_wrapper():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    with torch.device(device):
+        model = TransformersWrapper(
+            model="Qwen/Qwen2.5-0.5B",
+            tokenizer="Qwen/Qwen2.5-0.5B",
+            pad_model_input=False,
+            generate=False,
+        )
+        return model
+
+
+@pytest.mark.skipif(not _has_transformers, reason="transformers not installed")
+class TestWrappers:
+    @pytest.mark.parametrize("packing", [True, False])
+    @set_list_to_stack(True)
+    def test_packing(self, benchmark, transformers_wrapper, packing: bool):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        with torch.device(device):
+            transformers_wrapper = TransformersWrapper(
+                model=transformers_wrapper.model,
+                tokenizer=transformers_wrapper.tokenizer,
+                pad_model_input=not packing,
+                generate=False,
+                pad_output=False,
+            )
+            data = TensorDict(
+                {
+                    "history": ChatHistory(
+                        full=History(
+                            role=[
+                                ["user", "assistant"],
+                                ["user", "assistant"],
+                                ["user", "assistant"],
+                                ["user", "assistant"],
+                            ],
+                            content=[
+                                [
+                                    "Lorem ipsum dolor sit amet",
+                                    "consectetur adipiscing elit",
+                                ],
+                                [
+                                    "sed do eiusmod tempor incididunt ut labore et dolore magna aliqua",
+                                    "Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat",
+                                ],
+                                [
+                                    "Lorem ipsum dolor sit amet",
+                                    "consectetur adipiscing elit",
+                                ],
+                                [
+                                    "sed do eiusmod tempor incididunt ut labore et dolore magna aliqua",
+                                    "Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat",
+                                ],
+                            ],
+                            batch_size=(4, 2),
+                            device=device,
+                        ),
+                        batch_size=(4,),
+                        device=device,
+                    )
+                },
+                batch_size=(4,),
+                device=device,
+            ).to_lazystack()
+
+            def setup():
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+
+            benchmark.pedantic(
+                transformers_wrapper,
+                (data,),
+                rounds=10,
+                warmup_rounds=3,
+                setup=setup,
+            )

benchmarks/test_non_tensor_env_benchmark.py

@@ -0,0 +1,70 @@
+# 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 gc
+import time
+
+import pytest
+from tensordict import set_capture_non_tensor_stack
+from torchrl.envs import ParallelEnv, SerialEnv
+from torchrl.testing.mocking_classes import EnvWithMetadata
+
+
+def _rollout(env, n_steps: int, break_when_any_done: bool) -> None:
+    env.rollout(n_steps, break_when_any_done=break_when_any_done)
+
+
+@pytest.mark.parametrize("break_when_any_done", [True, False])
+@pytest.mark.parametrize(
+    "kind,use_buffers",
+    [
+        pytest.param("single", None, id="single"),
+        pytest.param("serial", False, id="serial-no-buffers"),
+        pytest.param("serial", True, id="serial-buffers"),
+        pytest.param("parallel", False, id="parallel-no-buffers"),
+        pytest.param("parallel", True, id="parallel-buffers"),
+    ],
+)
+@pytest.mark.parametrize("n_steps", [1000])
+def test_non_tensor_env_rollout_speed(
+    benchmark,
+    break_when_any_done: bool,
+    kind: str,
+    use_buffers: bool | None,
+    n_steps: int,
+):
+    """Benchmarks a single rollout, after a warmup rollout, for non-tensor stacking envs.
+
+    Mirrors `test/test_envs.py::TestNonTensorEnv`'s option matrix (single/serial/parallel,
+    break_when_any_done, use_buffers).
+    """
+    with set_capture_non_tensor_stack(False):
+        if kind == "single":
+            env = EnvWithMetadata()
+        elif kind == "serial":
+            env = SerialEnv(2, EnvWithMetadata, use_buffers=use_buffers)
+        elif kind == "parallel":
+            env = ParallelEnv(2, EnvWithMetadata, use_buffers=use_buffers)
+        else:
+            raise RuntimeError(f"Unknown kind={kind}")
+
+        env.set_seed(0)
+        env.reset()
+
+        try:
+            # Warmup run (not timed)
+            _rollout(env, n_steps=n_steps, break_when_any_done=break_when_any_done)
+
+            # Timed run(s)
+            benchmark(
+                _rollout, env, n_steps=n_steps, break_when_any_done=break_when_any_done
+            )
+        finally:
+            env.close(raise_if_closed=False)
+            del env
+            # Give multiprocessing envs a brief chance to terminate cleanly.
+            time.sleep(0.05)
+            gc.collect()