Implementation:Hpcaitech ColossalAI NaiveExperienceBuffer

Knowledge Sources	Hpcaitech_ColossalAI
Domains	Reinforcement Learning, RLHF, Experience Replay
Last Updated	2026-02-09 00:00 GMT

Overview

A simple experience buffer implementation with CPU offloading and random sampling for PPO training.

Description

NaiveExperienceBuffer is a concrete implementation of ExperienceBuffer that stores experience data as individual BufferItem entries. It supports optional CPU offloading to reduce GPU memory pressure during PPO training, automatically moving experience data to CPU on append and back to the target CUDA device on sample. The buffer maintains a shuffled random number sequence for pseudo-random sampling without replacement within each epoch, and supports an optional capacity limit that evicts the oldest samples when exceeded.

Usage

Use this buffer in PPO training loops when you need a straightforward experience replay mechanism. It is the default experience buffer for ColossalChat RLHF training and is suitable for single-node or multi-node setups where experience data fits in CPU memory.

Code Reference

Source Location

• Repository: Hpcaitech_ColossalAI
• File: applications/ColossalChat/coati/experience_buffer/naive.py
• Lines: 1-77

Signature

class NaiveExperienceBuffer(ExperienceBuffer):
    def __init__(self, sample_batch_size: int, limit: int = 0, cpu_offload: bool = True) -> None:

    @torch.no_grad()
    def append(self, experience: Experience) -> None:

    def clear(self) -> None:

    @torch.no_grad()
    def sample(self) -> Experience:

    def __len__(self) -> int:

    def __getitem__(self, idx: int) -> BufferItem:

    def collate_fn(self, batch) -> Experience:

Import

from coati.experience_buffer.naive import NaiveExperienceBuffer

I/O Contract

Inputs (__init__)

Name	Type	Required	Description
sample_batch_size	int	Yes	Batch size when sampling from the buffer
limit	int	No	Maximum number of stored samples; <= 0 means unlimited, defaults to 0
cpu_offload	bool	No	Whether to offload experience to CPU, defaults to True

Outputs (sample)

Name	Type	Description
return	Experience	A batch of sampled experiences moved to the target CUDA device

Usage Examples

from coati.experience_buffer.naive import NaiveExperienceBuffer

# Create buffer with batch size 8, limit 1000 samples, CPU offloading enabled
buffer = NaiveExperienceBuffer(sample_batch_size=8, limit=1000, cpu_offload=True)

# Append experience from the experience maker
buffer.append(experience)

# Sample a batch for PPO training
batch = buffer.sample()

# Check buffer size
print(f"Buffer contains {len(buffer)} items")

# Clear after epoch
buffer.clear()

Related Pages

• Environment:Hpcaitech_ColossalAI_CUDA_GPU_Environment