Tiny add Sample.group_index (#475)

Author: fzyzcjy

slime/ray/rollout_data_source.py

@@ -16,6 +16,7 @@ def __init__(self, args):
         self.args = args
 
         self.epoch_id = 0
+        self.sample_group_index = 0
         self.sample_index = 0
         self.sample_offset = 0
         # TODO remove this
@@ -66,9 +67,11 @@ def get_samples(self, num_samples):
             group = []
             for _ in range(self.args.n_samples_per_prompt):
                 sample = copy.deepcopy(prompt_sample)
+                sample.group_index = self.sample_group_index
                 sample.index = self.sample_index
                 self.sample_index += 1
                 group.append(sample)
+            self.sample_group_index += 1
             samples.append(group)
         return samples
 
@@ -82,6 +85,7 @@ def save(self, rollout_id):
         state_dict = {
             "sample_offset": self.sample_offset,
             "epoch_id": self.epoch_id,
+            "sample_group_index": self.sample_group_index,
             "sample_index": self.sample_index,
             "metadata": self.metadata,
         }
@@ -106,6 +110,7 @@ def load(self, rollout_id=None):
         state_dict = torch.load(path)
         self.sample_offset = state_dict.get("sample_offset", 0)
         self.epoch_id = state_dict.get("epoch_id", 0)
+        self.sample_group_index = state_dict.get("sample_group_index", 0)
         self.sample_index = state_dict.get("sample_index", 0)
         self.metadata = state_dict.get("metadata", {})
 

slime/utils/types.py

@@ -9,6 +9,7 @@
 class Sample:
     """The sample generated"""
 
+    group_index: Optional[int] = None
     index: Optional[int] = None
     # prompt
     prompt: Union[str, list[dict[str, str]]] = ""