core.trainers.grpo.sampler

core.trainers.grpo.sampler

Repeat random sampler (similar to the one implemented in https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py) that adds context parallelism functionality; i.e., duplicating data across ranks in the same context parallel group.

Classes

Name	Description
ContextParallelRepeatRandomSampler	Sampler for GRPO training with context parallelism.

ContextParallelRepeatRandomSampler

core.trainers.grpo.sampler.ContextParallelRepeatRandomSampler(
    dataset,
    mini_repeat_count,
    world_size,
    rank,
    batch_size=1,
    repeat_count=1,
    context_parallel_size=1,
    shuffle=True,
    seed=0,
    drop_last=False,
)

Sampler for GRPO training with context parallelism.

This sampler ensures: - Ranks in the same context parallel (CP) group receive identical data. - Each index is repeated multiple times for sampling different completions. - Entire batches are repeated for reuse in multiple updates. - Data is properly distributed across CP groups.

In the table below, the values represent dataset indices. Each CP group has context_parallel_size = 2 GPUs working together on the same data. There are 2 CP groups (CP0 and CP1), with world_size = 4 total GPUs.

                                        Context Parallel Groups
                                |       CP0        |       CP1        |
                                |  GPU 0  |  GPU 1 |  GPU 2  |  GPU 3 |
            global_step  step    <---> mini_repeat_count=3
                                    <----------> batch_size=2 per CP group

grad_accum=2 ▲ ▲ 0 0 [0 0 0 1 1 1] [2 2 2 3 3 3] <- CP groups get different data ▼ | 0 1 [0 0 0 1 1 1] [2 2 2 3 3 3] <- Same data for each CP group GPU | | 1 2 [0 0 0 1 1 1] [2 2 2 3 3 3] <- Repeat same indices for iterations num_iterations=2 ▼ 1 3 [0 0 0 1 1 1] [2 2 2 3 3 3] <- When using gradient accumulation

                 2       4         [4 4 4  5 5 5]     [6 6 6  7 7 7]   <- New batch of data indices
                 2       5         [4 4 4  5 5 5]     [6 6 6  7 7 7]
                                    ...

Parameters

Name	Type	Description	Default
dataset	Sized	Dataset to sample from.	required
mini_repeat_count	int	How many times to repeat each sample immediately.	required
world_size	int	Total number of processes.	required
rank	int	Rank of current process.	required
batch_size	int	Number of samples per batch.	`1`
repeat_count	int	How many times to repeat the full sampling process.	`1`
context_parallel_size	int	Number of ranks in a context parallel group.	`1`
shuffle	bool	Whether to shuffle the dataset.	`True`
seed	int	Random seed for shuffling.	`0`
drop_last	bool	Whether to drop the last incomplete batch.	`False`

Methods

Name	Description
set_epoch	Sets the epoch for this sampler.

set_epoch

core.trainers.grpo.sampler.ContextParallelRepeatRandomSampler.set_epoch(epoch)

Sets the epoch for this sampler.

Parameters

Name	Type	Description	Default
epoch	int	Epoch number to use for shuffling.	required