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

[Abecid]

Extracted from #1391.

GenRL-Stack: 2/6

Purpose

Fix GenRL prompt group placement so repeated samples from the same prompt stay together on the same data-parallel rank.

This matches the DistributedKRepeatSampler invariant and avoids per-rank prompt groups being split/shuffled across ranks.

Fixes #

Changes

• Require batch_size % num_video_per_prompt == 0 in DistributedKRepeatSampler.
• Assign whole prompt groups to each rank before repeating samples.
• Preserve epoch-tagged prompt samples for deterministic per-epoch behavior.

Test Plan

python -m py_compile fastvideo/train/methods/rl/utils/data.py

Test Results

Test output

py_compile passed locally.

[mergify]

⚠️ PR title format required

Your PR title must start with a type tag in brackets. Examples:

• [feat] Add new model support
• [bugfix] Fix VAE tiling corruption
• [refactor] Restructure training pipeline
• [perf] Optimize attention kernel
• [ci] Update test infrastructure
• [infra] Add activation trace hooks
• [docs] Add inference guide
• [misc] Clean up configs
• [new-model] Port Flux2 to FastVideo
• [skill] Add add-model agent skill

Valid tags: feat, feature, bugfix, fix, refactor, perf, ci, infra, doc, docs, misc, chore, kernel, new-model, skill, skills

Please update your PR title and the merge protection check will pass automatically.

[gemini-code-assist]

Code Review

This pull request refactors the distributed data sampler in fastvideo/train/methods/rl/utils/data.py to simplify the indexing logic and ensure each rank receives whole prompt groups. It adds a validation check to ensure batch_size is divisible by k. The reviewer suggested adding additional defensive checks to verify that k is a positive integer and that the dataset contains enough samples (len(self.dataset) >= self.m) to prevent potential runtime errors or silent bugs during distributed training.

[gemini-code-assist]

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