[Feature] Compressed storage gpu #3062
Open
+1,793
−66
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3062
Note: Links to docs will display an error until the docs builds have been completed. ❌ 13 New Failures, 1 Pending, 3 Unrelated FailuresAs of commit 0dbd233 with merge base db0e30d ( NEW FAILURES - The following jobs have failed:
FLAKY - The following job failed but was likely due to flakiness present on trunk:
BROKEN TRUNK - The following jobs failed but was present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
|
When the tensor is on the CPU |
… cursor logic to a view class. Passing all tests now.
vmoens
force-pushed
the
compressed-storage-gpu
branch
2 times, most recently
from
74d85fa to
95f532e
Compare
…o_bytestream speed test.
vmoens
force-pushed
the
compressed-storage-gpu
branch
from
95f532e to
5581cf6
Compare
|
Added some examples of compressing on the cpu, batched decompression on the gpu. I noticed that in my example of an Atari rollout, the Compressing on the CPU first, then transferring, and re-using this compressed observation for the next transition gets about double the Atari transitions per second. |
|
@vmoens I think we're essentially done with this PR, except for a cleanup pass. Do we want CompressedListStorage to be mentioned in the documentation? Maybe have a page on compression to showcase the VRAM storage savings on the gpu? |
vmoens
approved these changes
Jul 23, 2025
| @@ -20,7 +20,7 @@ repos: | |||
| - libcst == 0.4.7 | |||
|
|
|||
| - repo: https://github.com/pycqa/flake8 | |||
| rev: 4.0.1 | |||
| rev: 6.0.0 | |||
There was a problem hiding this comment.
Not against upgrades but can you comment why we need it?
There was a problem hiding this comment.
4.0.1 has an issue that causes the error below:
flake8...................................................................Failed
- hook id: flake8
- exit code: 1
Traceback (most recent call last):
File "/home/adrian/.cache/pre-commit/repoxdgdrlah/py_env-python3.12/bin/flake8", line 8, in <module>
sys.exit(main())
^^^^^^
File "/home/adrian/.cache/pre-commit/repoxdgdrlah/py_env-python3.12/lib/python3.12/site-packages/flake8/main/cli.py", line 22, in main
app.run(argv)
File "/home/adrian/.cache/pre-commit/repoxdgdrlah/py_env-python3.12/lib/python3.12/site-packages/flake8/main/application.py", line 375, in run
self._run(argv)
File "/home/adrian/.cache/pre-commit/repoxdgdrlah/py_env-python3.12/lib/python3.12/site-packages/flake8/main/application.py", line 363, in _run
self.initialize(argv)
File "/home/adrian/.cache/pre-commit/repoxdgdrlah/py_env-python3.12/lib/python3.12/site-packages/flake8/main/application.py", line 343, in initialize
self.find_plugins(config_finder)
File "/home/adrian/.cache/pre-commit/repoxdgdrlah/py_env-python3.12/lib/python3.12/site-packages/flake8/main/application.py", line 157, in find_plugins
self.check_plugins = plugin_manager.Checkers(local_plugins.extension)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/adrian/.cache/pre-commit/repoxdgdrlah/py_env-python3.12/lib/python3.12/site-packages/flake8/plugins/manager.py", line 363, in __init__
self.manager = PluginManager(
^^^^^^^^^^^^^^
File "/home/adrian/.cache/pre-commit/repoxdgdrlah/py_env-python3.12/lib/python3.12/site-packages/flake8/plugins/manager.py", line 243, in __init__
self._load_entrypoint_plugins()
File "/home/adrian/.cache/pre-commit/repoxdgdrlah/py_env-python3.12/lib/python3.12/site-packages/flake8/plugins/manager.py", line 261, in _load_entrypoint_plugins
eps = importlib_metadata.entry_points().get(self.namespace, ())
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'EntryPoints' object has no attribute 'get'
From what I understand from the issues that recommend upgrading, there was a depreciation by importlib-metadata which is fixed in 5.x.x or greater.
As an aside, you might like the speed of Ruff, I use it for my own projects.
| - **Data Integrity**: Maintains full data fidelity through lossless compression | ||
| - **Flexible Compression**: Supports custom compression algorithms or uses zstd by default | ||
| - **TensorDict Support**: Seamlessly works with TensorDict structures | ||
| - **Checkpointing**: Full support for saving and loading compressed data |
There was a problem hiding this comment.
Maybe add GPU support here
Comment on lines
+216
to
+231
| >>> import torch | ||
| >>> from torchrl.data import ReplayBuffer, CompressedListStorage | ||
| >>> from tensordict import TensorDict | ||
| >>> | ||
| >>> # Create a compressed storage for image data | ||
| >>> storage = CompressedListStorage(max_size=1000, compression_level=3) | ||
| >>> rb = ReplayBuffer(storage=storage, batch_size=32) | ||
| >>> | ||
| >>> # Add image data | ||
| >>> images = torch.randn(100, 3, 84, 84) # Atari-like frames | ||
| >>> data = TensorDict({"obs": images}, batch_size=[100]) | ||
| >>> rb.extend(data) | ||
| >>> | ||
| >>> # Sample data (automatically decompressed) | ||
| >>> sample = rb.sample(16) | ||
| >>> print(sample["obs"].shape) # torch.Size([16, 3, 84, 84]) |
There was a problem hiding this comment.
@AdrianOrenstein Can you check that this still makes sense and runs?
| @@ -0,0 +1,199 @@ | |||
| #!/usr/bin/env python3 | |||
|
|
||
| from torchrl.data.replay_buffers.storages import ListStorage | ||
|
|
||
| gym.register_envs(ale_py) |
| @@ -0,0 +1,182 @@ | |||
| #!/usr/bin/env python3 | |||
| """ | |||
| Example demonstrating the use of CompressedStorage for memory-efficient replay buffers on the GPU. | |||
There was a problem hiding this comment.
Is it worth having 2 distinct examples? Not particularly unhappy with this, just more repeated code and could be harder for people to spot what the differences are between the two (hard to know where to focus your attention)
| # === CompressedListStorage + ReplayBuffer with GPU compression === | ||
| print("\n=== ListStorage + ReplayBuffer (GPU) Example ===\n") | ||
|
|
||
| codec = nvcomp.Codec(algorithm=algorithm, bitstream_kind=bitstream_kind) |
There was a problem hiding this comment.
I think this deserves a bit of explanation :)
Also I would check if the lib is installed beforehand and raise a warning /exception if it cannot be found
has_nvcomp = importlib.util.findspec("nvcomp", None) is not None
if not has_nvcomp:
raise ImportError(...)
| storage = CompressedListStorage(max_size=1000, compression_level=6) | ||
|
|
||
| # Create some sample data with different shapes and types | ||
| print("Creating sample data...") |
There was a problem hiding this comment.
we should remove the prints in favor of logger
|
Thanks for the code review! |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
Replay buffers are used to store a lot of data and are used to feed neural networks with batched samples to learn from. So then ideally we could put this data as close to where the network is being updated. Often raw sensory observations are stored in these buffers, such as images, audio, or text, which consumes many gigabytes of precious memory. CPU memory and accelerator VRAM may be limited, or memory transfer between these devices may be costly. So this PR aims to streamline data compression to aid in efficient storage and memory transfer.
Mainly, creating a compressed storage object will aid in training state-of-the-art RL methods on benchmarks such as the Atari Learning Environment. The
~torchrl.data.replay_buffers.storages.CompressedStorageclass provides the memory savings through compression.closes #3058
closes #2983
Types of changes
What types of changes does your code introduce? Remove all that do not apply:
Checklist
Go over all the following points, and put an
xin all the boxes that apply.If you are unsure about any of these, don't hesitate to ask. We are here to help!