Skip to content

[feat] GenRL: keep repeated prompt samples on one rank #1401

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Davids048 merged 2 commits into from
Jun 4, 2026
Merged

[feat] GenRL: keep repeated prompt samples on one rank #1401

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cab9b30
[genrl]: keep repeated prompt samples on one rank
Abecid May 27, 2026
14ac936
[patch]: validate GenRL prompt sampler inputs
Davids048 Jun 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 25 additions & 22 deletions fastvideo/train/methods/rl/utils/data.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,17 +149,34 @@ def __init__(
):
self.dataset = dataset
self.batch_size = batch_size
self.k = k
self.k = k # Repeats/videos per prompt.
self.num_replicas = num_replicas
self.rank = rank
self.seed = seed
self.total_samples = num_replicas * batch_size
if self.k <= 0:
raise ValueError(
f"k must be a positive integer. Got k={k}."
)
if self.batch_size % self.k != 0:
raise ValueError(
"batch_size must be divisible by k so each rank receives "
"whole prompt groups. Got "
f"batch_size={batch_size}, k={k}."
)
assert self.total_samples % self.k == 0, (
f"k cannot divide n*b: k={k} "
f"num_replicas={num_replicas} "
f"batch_size={batch_size}"
)
self.m = self.total_samples // self.k
self.m = self.total_samples // self.k # Unique prompts across ranks.
if len(self.dataset) < self.m:
raise ValueError(
"dataset must contain at least one prompt per global "
"prompt group. Got "
f"dataset_size={len(self.dataset)}, required={self.m}."
)
self.groups_per_rank = self.batch_size // self.k
Comment on lines +161 to +179

@gemini-code-assist gemini-code-assist Bot May 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

To prevent potential runtime errors or silent bugs during distributed training, we should add defensive checks to ensure that k is a positive integer and that the dataset contains enough samples (len(self.dataset) >= self.m). If the dataset is smaller than self.m, some ranks will receive fewer or empty prompt groups, which can cause shape mismatches and hang or crash distributed training during collective operations (like all_gather).

        if self.k <= 0:
            raise ValueError(f"k must be a positive integer. Got k={k}.")
        if self.batch_size % self.k != 0:
            raise ValueError(
                "batch_size must be divisible by k so each rank receives "
                "whole prompt groups. Got "
                f"batch_size={batch_size}, k={k}."
            )
        assert self.total_samples % self.k == 0, (
            f"k cannot divide n*b: k={k} "
            f"num_replicas={num_replicas} "
            f"batch_size={batch_size}"
        )
        self.m = self.total_samples // self.k
        if len(self.dataset) < self.m:
            raise ValueError(
                f"Dataset length ({len(self.dataset)}) must be at least "
                f"required samples per epoch ({self.m}) to support the requested batch size and replicas."
            )
        self.groups_per_rank = self.batch_size // self.k

self.epoch = 0

def __iter__(self):
Expand All @@ -169,28 +186,14 @@ def __iter__(self):
indices = torch.randperm(
len(self.dataset), generator=g
)[: self.m].tolist()
repeated = [
idx
for idx in indices
start = self.rank * self.groups_per_rank
end = start + self.groups_per_rank
rank_groups = indices[start:end]
yield [
(self.epoch, idx)
for idx in rank_groups
for _ in range(self.k)
]
shuffled_idx = torch.randperm(
len(repeated), generator=g
).tolist()
shuffled = [
repeated[i] for i in shuffled_idx
]
per_card = []
for i in range(self.num_replicas):
start = i * self.batch_size
end = start + self.batch_size
per_card.append(
[
(self.epoch, idx)
for idx in shuffled[start:end]
]
)
yield per_card[self.rank]

def set_epoch(self, epoch: int):
self.epoch = epoch
Expand Down
Loading