fix gramma & add more check

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>
Co-authored-by: zhabuye <2947436155@qq.com>

Author: 0oshowero0

tests/test_samplers.py

@@ -499,7 +499,7 @@ def test_rank_aware_sampler_multiple_dp_groups(self):
         sampled0_g0, consumed0_g0 = sampler.sample(
             ready_indexes, batch_size, dp_group=0, dp_world_size=dp_world_size, world_size=world_size
         )
-        # minic the consumption status update managed in TransferQueueController
+        # mimic the consumption status update managed in TransferQueueController
         ready_indexes = [i for i in ready_indexes if i not in consumed0_g0]
 
         # DP group 1: rank 0 samples first

transfer_queue/sampler/base.py

@@ -34,9 +34,8 @@ class BaseSampler(ABC):
     - **SequentialSampler**: Default sampler, selects samples sequentially without replacement
     - **GRPOGroupNSampler**: A sampler that performs sampling on continuous N samples only when all of them are ready.
                             It assumes the N samples associated with the same prompt are stored contiguously
-    - **RankAwareSampler**: Rank-aware sampling for distributed training where each ranks independently retrieve data
-                            by themselves. This sampler will guarantee ranks of the same DP group consume identical
-                            samples.
+    - **RankAwareSampler**: Rank-aware sampling for distributed training where each rank retrieves data independently.
+                            This sampler will guarantee ranks of the same DP group consume identical samples.
 
     NOTE: Always return both sampled and consumed indexes (may be identical).
     """

transfer_queue/sampler/rank_aware_sampler.py

@@ -22,10 +22,9 @@ class RankAwareSampler(BaseSampler):
     """Rank-aware sampler for distributed training with TransferQueue.
 
     This sampler is designed for distributed data parallel training scenarios
-    where each ranks independently retrieve data by themselves.
+    where each rank retrieves data independently.
 
-    Each rank independently calls the sampler, passing its own rank information,
-    and the sampler guarantees that all ranks within the same DP group receive
+    This sampler guarantees that all ranks within the same DP group receive
     the same sample indices.
 
     The sampler maintains per-DP-group state to coordinate sampling across ranks:
@@ -72,7 +71,7 @@ def sample(
             ready_indexes: List of global indices for which all required fields of the
                 corresponding samples have been produced, and the samples are not labeled
                 as consumed in the corresponding task.
-            batch_size: batch_size: Number of samples to select. If larger than available
+            batch_size: Number of samples to select. If larger than available
                 ready samples, all available samples will be returned.
             dp_group: The group id of current data parallel group. Used to
                 identify which DP group this rank belongs to.
@@ -89,9 +88,8 @@ def sample(
             List of global indices of length batch_size that should be labeled as consumed
             (will never be retrieved in the future)
 
-        Raises:
-            RuntimeError: If the fetch count exceeds the expected number of
-                fetches per DP group.
+        Raise:
+            RuntimeError: If ``world_size`` is not divisible by ``dp_world_size``.
 
         Note:
             The ``world_size // dp_world_size`` calculation determines how many
@@ -102,6 +100,9 @@ def sample(
         data_for_dp_group = self._states.get(dp_group, None)
 
         # Calculate how many times this batch should be fetched across all ranks
+        if world_size % dp_world_size != 0:
+            raise RuntimeError(f"world_size ({world_size}) is not divisible by dp_world_size ({dp_world_size})")
+
         fetches_per_batch = world_size // dp_world_size
 
         if data_for_dp_group is None: