[Feature Request] Compressing data stored in the Replay Buffer

[AdrianOrenstein]

Motivation

Replay buffers are used to store a lot of data and ideally we could put this data on a GPU.
Often raw sensory observations are stored in these buffers, such as images or audio, which consumes many gigabytes of precious VRAM.
In torchRL, memmapping and prefetching could be considered the fix for this issue but I have instead pursued compression instead.

I've found that zstd is effective in compressing atari frames (reducing storage use by 150-250x) and my DQN implementation has gone from storing ~4M frames consuming ~20ish GB to less than a GB.
I would like some pointers on where it might be sensible to integrate compression into the torch rl replay buffer in the hopes that compressing raw sensory observations could be more widely adopted.

Solution

I would like feedback on my proposed solutions to implementing compression in the replay buffer.

A CompressedReplayBuffer class that inherents from the ReplayBuffer class could be one option.
This class would compress data "as it comes in", by modifying the append and extend functions.
The compressed data is then passed to the parent ReplayBuffer class as a uint8 byte stream for storage.
The compressed data is then decompressed in the sample function.
Although, the issue in this approach is that the compressed byte stream can be of variable shape.
Some images might be compressed into ~150 bytes, but others might be ~188 bytes.
The TensorStorage object expects to receive objects of a consistent shape, so perhaps this is an inconvenient place to implement compression..

A compressed storage type might make more sense, as ultimately we want a storage type that can tolerate storing data of different lengths.
One approach could be to use the standard replay buffer but with a ListStorage type.
By modifying the get and set functions we could compress and decompress as data comes in and out.
Alternatively, TensorStorage could be modified to use a nested & jagged layout, which might be more efficient.

Additional context

Although I have used zstd as the compression algorithm in this issue, ideally we don't lock someone into using it.
Compression is just converting any arbitrary tensor into a uint8 1d tensor (the byte stream) and then decompression is rebuilding the prior tensor and recovering the shape and datatype.
If the data is sensory observations from a gym-like environment, we can reconstruct the observation's original shape and datatype from the gym environment metadata.

Checklist

• I have checked that there is no similar issue in the repo (required)

[AdrianOrenstein]

Tensordict - Feature Request - Support for nested tensors Might be relevant.

[vmoens]

I like the idea!
We also experimented with torchcodec to store videos instead of large uint8 tensors and it worked very well.
I'd like to get these two things right before committing to the feature:

• extendable and clean API: we don't want to overfit to a specific solution/backend
• how do we specify the compressing format ahead of time? Lazy init etc
• What if the compressed format is read only?
• serialization / checkpointing
• strange data formats: non tensors, ragged dimensions
• torch.compile compatibility

[AdrianOrenstein]

Thanks. I'll have a think. Also I've found a combination of compression being done on the CPU, and batched decompression on the GPU to be the fastest pipeline for my application. I used NVComp as my gpu backend. So some additional api might be setting what device to do the compression and decompression on, sensible defaults being the device the data is initially on, and where the data is being sent to upon sampling the buffer.

What might make gpu compression more effective is if the underlying storage type waits until enough objects need compression and schedules them all at once.

[vmoens]

Can you share some (pseudo-)code of how you're using zstd in practice?
I wonder how many files / compressed file we put, if we dynamically create / delete these files as the buffer is circularly filled etc.

So some additional api might be setting what device to do the compression and decompression on

Yes I agree!

[AdrianOrenstein]

Yeah, here you go:

CPU compression / decompression:

pip install zstd

import zstd
import torch
import numpy as np
obs, _ = env.reset(seed=0)
compressed_obs = zstd.compress(obs.tobytes())
decompressed_obs = zstd.decompress(compressed_obs)
pt_obs = torch.frombuffer(decompressed_obs, dtype=torch.uint8).view(4, 84, 84)
assert np.allclose(obs, pt_obs.cpu().numpy())
print(f"passed correctness check for zstd: from {len(obs.tobytes())} bytes to {len(compressed_obs)} @ compression ratio of {len(obs.tobytes()) / len(compressed_obs):0.0f}x")

GPU compression / decompression with nvComp with:

I'm using nvComp as the compress/decompress library. They only provide wheels for python 3.12 and below NVIDIA/CUDALibrarySamples#259 (comment), so don't use 3.13 just yet. Unfortunately they made the library closed-source so we need to wait for the next public release of nvComp 5.0 (maybe happening in July).

I'm using a docker image based from nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04 and python 3.12. As I'm using the cuda 12 image, we install cu12 nvcomp:

pip install nvidia-nvcomp-cu12

import nvcomp
codec = nvcomp.Codec(algorithm="Zstd", bitstream_kind=nvcomp.BitstreamKind.RAW)
obs, _ = env.reset(seed=0)
nv_obs = nvcomp.as_array(obs).cuda()
compressed_obs = codec.encode(nv_obs)
decompressed_obs = codec.decode(compressed_obs, data_type="|u1")
pt_obs = torch.from_dlpack(decompressed_obs).clone().view(4, 84, 84)
assert np.allclose(obs, pt_obs.cpu().numpy())

# Some measures of performance:
# rollout with Zstd and BitstreamKind.RAW, avg_compression_ratio=97 @ obs/s=1173
# batch sampling and decompression with Zstd and nvcomp.BitstreamKind.RAW @ obs/s=100509

CPU compression, batched GPU decompression:

pip install zstd nvidia-nvcomp-cu12

import zstd
import torch
import nvcomp

codec = nvcomp.Codec(algorithm="Zstd", bitstream_kind=nvcomp.BitstreamKind.RAW)

obs, _ = env.reset(seed=0)
compressed_obs: bytes = zstd.compress(obs.tobytes())

nv_compressed_obs = nvcomp.as_array(compressed_obs).cuda()
nv_decompressed_obs = codec.decode(nv_compressed_obs, data_type="|u1")
pt_obs = torch.from_dlpack(nv_decompressed_obs).clone().view(4, 84, 84)
assert np.allclose(obs, pt_obs.cpu().numpy())

# Some measures of performance:
# rollout with cpu zstd, avg_compression_ratio=94 @ obs/s=1861
# batch sampling and decompression with Zstd and nvcomp.BitstreamKind.RAW @ obs/s=100509

Despite my initial benchmarks showing that (batch size 1) cpu compression is faster than gpu compression, I still think some kind of lazy batched compression on the gpu will be the fastest. So accumulating a batch-size of observations and compressing them all at once is probably more ideal behaviour for a replay buffer as opposed to compressing each observation immediately as it comes in.

I used to use lz4 as well, so if you want examples of that use pip install lz4 and for compatability with NvComp I think you'd need to use lz4.block (as per https://docs.nvidia.com/cuda/nvcomp/lz4.html).

[AdrianOrenstein]

• extendable and clean API: we don't want to overfit to a specific solution/backend
  I was thinking we could pass in encode/decode functions into the storage object? I think the storage object should be the thing that is responsible for compression/decompression, and the replay buffer shouldn't really care.

I feel like compressing the searlalized bytestream that memmap would need to produce would give us most of the benefits. I'm not sure how torchrl serializes every object for memmap without pickle, is there anywhere I could read up on it?

The encode function would ingest any type objects and return a bytearray, and a decode function that accepts bytearrays and reconstructs the original objects. I think we could heavily reuse the serialisation code I imagine is used by memmap.

[AdrianOrenstein]

A very quick and dirty example of compressing and decompressing a tensordict.

import torch
import zstd
import pickle
import io
from tensordict import tensorclass

device = "cpu"


@tensorclass
class PixelExperiences:
    observations: torch.Tensor
    actions: torch.Tensor
    rewards: torch.Tensor
    next_observations: torch.Tensor
    terminations: torch.Tensor
    truncations: torch.Tensor


def make_compressible_mock_data(num_experiences: int) -> PixelExperiences:
    """Easily compressible data for testing."""
    return PixelExperiences(
        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],
    )


data = make_compressible_mock_data(1).consolidate()
serialized_data = data.state_dict(flatten=True)

Compressing the whole tensordict:

# Compress
whole_buff = io.BytesIO()
pickle.dump(serialized_data, whole_buff)
whole_buff.seek(0)
compressed_buff = zstd.compress(whole_buff.getvalue(), 16)

# Decompress
decompressed_buff = zstd.decompress(compressed_buff)
deserialised_state_dict = pickle.loads(decompressed_buff)
data_2 = make_compressible_mock_data(1)
data_consolidated_2 = data.load_state_dict(state_dict=deserialised_state_dict, from_flatten=True)
assert data is data_consolidated_2
print(
    "from",
    len(whole_buff.getvalue()),
    "to",
    len(compressed_buff),
    f", compression ratio: {len(whole_buff.getvalue()) / len(compressed_buff):0.0f}x",
)
# from 58436 to 580 , compression ratio: 101x

Compressing just the values in the tensordict

# Compress
metadata = tuple(serialized_data.keys())
buff = io.BytesIO()
pickle.dump(tuple(serialized_data.values()), buff)
buff.seek(0)
compressed_buff = zstd.compress(buff.getvalue(), 16)

# Decompress
decompressed_buff = zstd.decompress(compressed_buff)
deserialised_state_dict = pickle.loads(decompressed_buff)
data_2 = make_compressible_mock_data(1)
data_consolidated_2 = data.load_state_dict(state_dict=dict(zip(metadata, deserialised_state_dict)), from_flatten=True)
assert data is data_consolidated_2
print(
    "from",
    len(whole_buff.getvalue()),
    "to",
    len(compressed_buff),
    f", compression ratio: {len(whole_buff.getvalue()) / len(compressed_buff):0.0f}x",
)
# from 58436 to 563 , compression ratio: 104x

strange data formats: non tensors, ragged dimensions

I guess this satisfies the strange data formats via pickle. But if we can see that all of the tensordict is populated with tensors, we should just be able to read out the bytes from the tensor object.
I'll look into how we might compress and decompress tensors on the gpu.

[AdrianOrenstein]

The compressed byte values still appear to include unnecessary metadata needed to reconstruct the whole object. As stored objects keep metadata about the expected tensordicts it should receive, I think it is sensible to assume that we can just store the underlying bytestream of the incoming objects and reconstruct the pytorch tensor upon retrieval. So we can just compress & store the byte stream of the tensor rather than serialising the whole object.

So to get the actual data of each tensor, we can usebytes(the_tensor.untyped_storage()). This is similar to numpy's .to_bytes():

t = torch.HalfTensor([[1.23], [3.45], [1.2445]])
torch_bytes = bytes(t.untyped_storage())
np_bytes = t.numpy().tobytes()
assert torch_bytes == np_bytes

Compressing the underlying byte store

mock_data = make_compressible_mock_data(1) + 2 # +2 to sanity check we decompress correctly

def to_bytestream(data_to_bytestream: Union[torch.Tensor, np.array, Any]) -> bytes:
    if isinstance(data_to_bytestream, torch.Tensor):
        if data_to_bytestream.is_nested:
            byte_stream = bytes(data_to_bytestream.values().untyped_storage())
        else:
            byte_stream = bytes(data_to_bytestream.untyped_storage())

    elif isinstance(data_to_bytestream, np.array):
        byte_stream = bytes(data_to_bytestream.tobytes())

    else:
        buffer = io.BytesIO()
        pickle.dump(serialized_data, buffer)
        buffer.seek(0)
        byte_stream = bytes(byte_stream)

    return byte_stream


def compress(
    data: Any,
    compress_fn: Union[str, Callable[[bytes, int], bytes]] = "zstd",
    level: int = 16,
    **compress_fn_kwargs,
) -> torch.Tensor:
    """
    Compress data using a specified compression function.

    Args:
        data: The data to compress. Can be any type that can be converted to bytes.
        compress_fn: Either a string identifier for a compression algorithm or a callable that takes bytes and compression level.
        level: Compression level (default: 16). Higher values give better compression but are slower.

    Returns:
        torch.Tensor: A tensor containing the compressed data as uint8 bytes.

    Example:
        >>> compressed_tn = compress(torch.randn(1000), "zstd", level=16)
        >>> compresse_np = compress(np.array([10, 12, 13]), "zstd", level=16)
        >>> compressed_py = compress(dict("test": tuple(1, 2, 3))), "zstd", level=16)
    """
    byte_stream = to_bytestream(data)

    writable_byte_stream = None

    if isinstance(compress_fn, str):
        if compress_fn.lower() == "zstd":
            import zstd

            # Use positional arguments instead of keyword arguments
            writable_byte_stream = bytearray(zstd.compress(byte_stream, level))

    # our default
    if writable_byte_stream is None or (isinstance(compress_fn, str) and compress_fn.lower() == "zlib"):
        import zlib

        assert (0 <= level <= 9) or (level == -1), "zlib expects a compression level in 0-9 or -1"
        writable_byte_stream = bytearray(
            bytearray(zlib.compress(byte_stream, level=level, wbits=compress_fn_kwargs.get("wbits", 15)))
        )

    # TODO: we should handle function calls

    return torch.frombuffer(writable_byte_stream, dtype=torch.uint8)


compressed_mock_data = mock_data.apply(
    lambda d: torch.nested.nested_tensor(
        list(map(lambda x: compress(x, compress_fn="zstd", level=16), d)), layout=torch.jagged
    )
)
print(compressed_mock_data)
# PixelExperiences(
#     actions=NestedTensor(shape=torch.Size([1, j3]), device=cpu, dtype=torch.uint8, is_shared=False),
#     next_observations=NestedTensor(shape=torch.Size([1, j5]), device=cpu, dtype=torch.uint8, is_shared=False),
#     observations=NestedTensor(shape=torch.Size([1, j2]), device=cpu, dtype=torch.uint8, is_shared=False),
#     rewards=NestedTensor(shape=torch.Size([1, j4]), device=cpu, dtype=torch.uint8, is_shared=False),
#     terminations=NestedTensor(shape=torch.Size([1, j6]), device=cpu, dtype=torch.uint8, is_shared=False),
#     truncations=NestedTensor(shape=torch.Size([1, j7]), device=cpu, dtype=torch.uint8, is_shared=False),
#     batch_size=torch.Size([1]),
#     device=None,
#     is_shared=False
# )
print(
    "alg=zstd level=16",
    "from",
    mock_data.bytes(),
    "to",
    compressed_mock_data.bytes(),
    f", compression ratio: {data.bytes() / compressed_mock_data.bytes():0.0f}x",
)
# alg=zstd level=16 from 56458 to 178, compression ratio: 317x

Decompressing the byte store into the object

from typing import Literal, Optional

def decompress(
    data: Any,
    decompress_fn: Union[str, Callable[[bytes, int], bytes]] = "zstd",
    **decompress_fn_kwargs,
) -> torch.Tensor:
    byte_stream = to_bytestream(data)

    writable_byte_stream = None

    if isinstance(decompress_fn, str):
        if decompress_fn.lower() == "zstd":
            import zstd

            # Use positional arguments instead of keyword arguments
            writable_byte_stream = bytearray(zstd.decompress(byte_stream))

    # our default
    if writable_byte_stream is None or (isinstance(decompress_fn, str) and decompress_fn.lower() == "zlib"):
        import zlib

        writable_byte_stream = bytearray(
            bytearray(zlib.decompress(byte_stream, wbits=decompress_fn_kwargs.get("wbits", 15)))
        )

    # TODO: we should handle function calls

    return torch.frombuffer(writable_byte_stream, dtype=torch.uint8)

decompressed_mock_data = compressed_mock_data.apply(
    lambda d: torch.stack(list(map(lambda x: decompress(x, decompress_fn="zstd"), d)))
).to_tensordict()

print(decompressed_mock_data)
# TensorDict(
#     fields={
#         actions: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.uint8, is_shared=False),
#         next_observations: Tensor(shape=torch.Size([1, 28224]), device=cpu, dtype=torch.uint8, is_shared=False),
#         observations: Tensor(shape=torch.Size([1, 28224]), device=cpu, dtype=torch.uint8, is_shared=False),
#         rewards: Tensor(shape=torch.Size([1, 4]), device=cpu, dtype=torch.uint8, is_shared=False),
#         terminations: Tensor(shape=torch.Size([1, 8]), device=cpu, dtype=torch.uint8, is_shared=False),
#         truncations: Tensor(shape=torch.Size([1, 8]), device=cpu, dtype=torch.uint8, is_shared=False)},
#     batch_size=torch.Size([1]),
#     device=None,
#     is_shared=False
# )

def rebuild_object(
    bytestream_to_data: torch.Tensor,
    dummy_data: Optional[TensorDict] = None,
    output_type: Literal["tensor", "object"] = "tensor",
) -> TensorDict:
    if output_type == "tensor":
        assert dummy_data is not None, f"For now we rebuild the object by replacing the byte store"
        out_td = dummy_data.clone()

        for k_zd, v_zd in out_td.items():
            v_zd.set_(
                bytestream_to_data[k_zd],
                storage_offset=v_zd.storage_offset(),
                stride=v_zd.stride(),
                size=v_zd.size(),
            )

        return out_td

    return pickle.loads(io.BytesIO(to_bytestream(bytestream_to_data)))


new_object = rebuild_object(
    bytestream_to_data=decompressed_mock_data,
    dummy_data=make_compressible_mock_data(1).to_tensordict(),
)
print(new_object)
# TensorDict(
#     fields={
#         actions: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.uint8, is_shared=False),
#         next_observations: Tensor(shape=torch.Size([1, 4, 84, 84]), device=cpu, dtype=torch.uint8, is_shared=False),
#         observations: Tensor(shape=torch.Size([1, 4, 84, 84]), device=cpu, dtype=torch.uint8, is_shared=False),
#         rewards: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
#         terminations: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, is_shared=False),
#         truncations: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, is_shared=False)},
#     batch_size=torch.Size([1]),
#     device=None,
#     is_shared=False
# )

So above, I showed: to_bytestream, compress to a nested tensordict, decompress, rebuild the object on the device with the right types & shapes. Using this to compress tensors in a tensordict we get a better compression ratio: from 104x to 317x, which we can expect from just using tensors filled with zeros. I'll make a branch and start implementing a storage object that stores the byte values of a tensor object.

[AdrianOrenstein]

I've written my thoughts for each of your points @vmoens, let me know if you have any comments. I'll work on tracking my progress on making a CompressedTensorStorage class in a draft PR.

• extendable and clean API: we don't want to overfit to a specific solution/backend
  • So I think the main breakdown are these features:
  • to_bytestream: implements taking Union[torch.Tensor, np.array, Any] and outputs a batch of torch.tensor(..., dtype=torch.uint8);
  • compress: takes a batch of torch.tensor(..., dtype=torch.uint8) and outputs another byte stream but in the nested format torch.tensor(..., dtype=torch.uint8, layout=torch.nested) as each sample may have different byte stream lengths.
• how do we specify the compressing format ahead of time? Lazy init etc
  • In my example I have a few preset options, "zstd" or the default "zlib" available as a standard library. I'd extend this to have nvcomp and potentially torchvision.jpeg.
• What if the compressed format is read only?
  • The compressed format I've used is a torch.nested.tensor of bytes.
• serialization. See to_bytestream above, but essentially we can either use pickle or extract the byte storage directly.
• checkpointing. We should be able to memmap the tensordict containing compressed data.
• strange data formats: non tensors, ragged dimensions. Non tensors can be captured with pickle. We can access the byte storage directly from ragged/nested tensors, and we can reconstruct the decompressed object by keeping around the metadata of the original data.
• torch.compile compatibility. I'm assuming if we call zlib/zstd we will break the graph by calling out to an external library. Additionally, I don't think nvcomp's cuda calls are captured, leading to another graph break.

[vmoens]

Hey just came back from working on the new release, now I can give some time to look into this!