fix: rewrite GRPO data pipeline for offline record-level access

- process_list_field returns List[List[int]] preserving per-response boundaries
- GRPODataset rewritten to record-level __getitem__ (no windowing/stride)
- grpo_collate_fn pads variable-length responses into [B, G, R] tensors
- JsonlStore detects nested List[List[int]] and stores List[Tensor] per record
- Store._normalize skips nested-list keys from cumsum bookkeeping
- Pipeline._flush handles nested lists without cross-record flattening
- Export grpo_collate_fn from astrai.dataset
- 6 new GRPO tests + 2 updated builder tests, 114 total pass

astrai/dataset/__init__.py

@@ -1,6 +1,7 @@
 from astrai.dataset.dataset import (
     BaseDataset,
     DatasetFactory,
+    grpo_collate_fn,
 )
 from astrai.dataset.sampler import ResumableDistributedSampler
 from astrai.dataset.storage import (
@@ -21,6 +22,7 @@ from astrai.serialization import (
 __all__ = [
     "BaseDataset",
     "DatasetFactory",
+    "grpo_collate_fn",
     "Store",
     "StoreFactory",
     "H5Store",

astrai/dataset/dataset.py

@@ -15,6 +15,49 @@ from astrai.dataset.storage import (
 from astrai.factory import BaseFactory
 
 
+def grpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
+    """Collate variable-length GRPO samples into padded 3-D tensors.
+
+    Input: list of dicts, each with:
+      - prompts:  [P_i]
+      - responses: list of G tensors, each [R_ij]
+      - masks:     list of G tensors, each [R_ij]
+      - rewards:  [G]
+
+    Output:
+      - prompts:   [B, P_max]
+      - responses: [B, G, R_max]
+      - masks:     [B, G, R_max]
+      - rewards:   [B, G]
+    """
+    B = len(batch)
+    G = len(batch[0]["responses"])
+    P_max = max(b["prompts"].size(0) for b in batch)
+    R_max = max(r.size(0) for b in batch for r in b["responses"])
+
+    prompts = torch.zeros(B, P_max, dtype=torch.long)
+    responses = torch.zeros(B, G, R_max, dtype=torch.long)
+    masks = torch.zeros(B, G, R_max, dtype=torch.bool)
+    rewards = torch.zeros(B, G, dtype=torch.float32)
+
+    for i, b in enumerate(batch):
+        p_len = b["prompts"].size(0)
+        prompts[i, :p_len] = b["prompts"]
+        rewards[i, : b["rewards"].size(0)] = b["rewards"]
+        for g in range(min(G, len(b["responses"]))):
+            r_len = b["responses"][g].size(0)
+            responses[i, g, :r_len] = b["responses"][g]
+            if g < len(b["masks"]):
+                masks[i, g, :r_len] = b["masks"][g]
+
+    return {
+        "prompts": prompts,
+        "responses": responses,
+        "masks": masks,
+        "rewards": rewards,
+    }
+
+
 class BaseDataset(Dataset, ABC):
     """Abstract base class for all dataset types.
 
@@ -250,28 +293,85 @@ class DPODataset(BaseDataset):
 
 @DatasetFactory.register("grpo")
 class GRPODataset(BaseDataset):
-    """Dataset for Group Relative Policy Optimization training."""
+    """Dataset for offline Group Relative Policy Optimization.
+
+    Unlike the window-based datasets (SEQ/SFT/DPO), GRPO data is
+    record-structured: each sample is one prompt with its group of
+    responses and scalar rewards.  There is no windowing or stride —
+    every record is an independent training unit.
+
+    Expected storage layout (produced by JsonlStore or pre-tokenized):
+
+    - ``prompts``:   List[Tensor]  — one 1-D token tensor per record
+    - ``responses``: List[List[Tensor]] — G response tensors per record
+    - ``masks``:     List[List[Tensor]] — G mask tensors per record
+    - ``rewards``:   List[Tensor]  — one 1-D float tensor (len G) per record
+    """
+
+    def __init__(self, window_size: int = 0, stride: int = 0, **kwargs):
+        super().__init__(window_size=window_size, stride=stride or window_size)
+        self._records: List[dict] = []
 
     @property
     def required_keys(self) -> List[str]:
         return ["prompts", "responses", "masks", "rewards"]
 
-    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
+    def load(self, load_path: str, storage_type: Optional[str] = None, **kwargs):
-        return self.storage.fetch(begin_idx, end_idx, key)
+        if storage_type is None:
+            storage_type = detect_format(load_path)
+        self.storage = StoreFactory.create(storage_type, **kwargs)
+        self._load_path = load_path
+        self.storage.load(load_path, **kwargs)
+        self._validate_keys()
+        self._build_records()
+
+    def _validate_keys(self):
+        actual_keys = set(self.storage.keys)
+        missing = [k for k in self.required_keys if k not in actual_keys]
+        if missing:
+            raise KeyError(
+                f"GRPODataset requires keys {self.required_keys}, "
+                f"but storage only has {sorted(actual_keys)}. Missing: {missing}"
+            )
+
+    def _build_records(self):
+        """Unfold segmented storage into per-record lists.
+
+        ``prompts`` is a flat list of 1-D tensors (one per record).
+        ``responses`` / ``masks`` are nested lists (G tensors per record).
+        ``rewards`` is a flat list of 1-D tensors (len G per record).
+        """
+        prompt_segs = self.storage._data.get("prompts", [])
+        response_segs = self.storage._data.get("responses", [])
+        mask_segs = self.storage._data.get("masks", [])
+        reward_segs = self.storage._data.get("rewards", [])
+
+        n_records = len(prompt_segs)
+        self._records = []
+        for i in range(n_records):
+            self._records.append(
+                {
+                    "prompts": prompt_segs[i],
+                    "responses": response_segs[i] if i < len(response_segs) else [],
+                    "masks": mask_segs[i] if i < len(mask_segs) else [],
+                    "rewards": reward_segs[i]
+                    if i < len(reward_segs)
+                    else torch.tensor([]),
+                }
+            )
+
+    @property
+    def count(self) -> int:
+        return len(self._records)
+
+    def __len__(self) -> int:
+        return len(self._records)
 
     def __getitem__(self, index: int) -> Dict[str, Tensor]:
-        begin_idx, end_idx = self.get_index(index)
+        rec = self._records[index]
-
-        prompts = self._fetch_data(begin_idx, end_idx, "prompts").to(dtype=torch.long)
-        responses = self._fetch_data(begin_idx, end_idx, "responses").to(
-            dtype=torch.long
-        )
-        masks = self._fetch_data(begin_idx, end_idx, "masks").to(dtype=torch.bool)
-        rewards = self._fetch_data(begin_idx, end_idx, "rewards")
-
         return {
-            "prompts": prompts,
+            "prompts": rec["prompts"].to(dtype=torch.long),
-            "responses": responses,
+            "responses": [r.to(dtype=torch.long) for r in rec["responses"]],
-            "masks": masks,
+            "masks": [m.to(dtype=torch.bool) for m in rec["masks"]],
-            "rewards": rewards,
+            "rewards": rec["rewards"].to(dtype=torch.float32),
         }

astrai/dataset/storage.py

@@ -148,26 +148,37 @@ class Store(ABC):
 
         return results[0] if len(results) == 1 else torch.cat(results, dim=0)
 
-    def _normalize(self, raw: Dict[str, List[Tensor]]):
+    def _normalize(self, raw: Dict[str, list]):
         """Register segments and pre-compute cumulative lengths.
 
         Does NOT concatenate — segments are kept as-is to avoid OOM on
         large datasets.  Sets ``self._length`` to the minimum total
-        element count across all keys.
+        element count across all flat-tensor keys.
+
+        For GRPO multi-response keys, values may be ``List[List[Tensor]]``
+        (one list of G tensors per record).  These are stored as-is and
+        excluded from the cumulative-length bookkeeping since they are
+        accessed record-by-record via ``_data`` rather than via ``fetch``.
         """
+        flat_lengths = []
         for key, tensors in raw.items():
             self._data[key] = tensors
+            if not tensors:
+                self._cum[key] = []
+                flat_lengths.append(0)
+                continue
+            # Skip nested lists (GRPO responses/masks) — record-level access
+            if isinstance(tensors[0], list):
+                self._cum[key] = []
+                continue
             cum = []
             total = 0
             for t in tensors:
                 total += t.shape[0]
                 cum.append(total)
             self._cum[key] = cum
-        self._length = (
+            flat_lengths.append(cum[-1] if cum else 0)
-            min((cum[-1] if cum else 0) for cum in self._cum.values())
+        self._length = min(flat_lengths) if flat_lengths else 0
-            if self._cum
-            else 0
-        )
 
 
 class StoreFactory(BaseFactory["Store"]):
@@ -274,6 +285,12 @@ class JsonlStore(Store):
             for key, ids in result.items():
                 if key not in raw:
                     raw[key] = []
+                if ids and isinstance(ids[0], list):
+                    # GRPO multi-response: List[List[int]] → List[Tensor]
+                    raw[key].append(
+                        [torch.tensor(sub, dtype=self._infer_dtype(sub)) for sub in ids]
+                    )
+                else:
                     raw[key].append(torch.tensor(ids, dtype=self._infer_dtype(ids)))
 
         for jsonl_path in sorted(root.glob("*.jsonl")):
@@ -314,7 +331,7 @@ class JsonlStore(Store):
 
     @staticmethod
     def _primary_ids(result: dict) -> List[int]:
-        """Return the first integer list in *result* as the primary id sequence."""
+        """Return the first flat integer list in *result* as the primary id sequence."""
         for val in result.values():
             if isinstance(val, list) and val and isinstance(val[0], int):
                 return val

astrai/preprocessing/builder.py

@@ -95,8 +95,15 @@ class SectionRenderer:
         return all_ids, loss_mask
 
     def process_list_field(self, item: dict, sections: list, config, tokenizer):
-        all_ids: list[int] = []
+        """Tokenize a list-valued field, preserving per-element boundaries.
-        loss_mask: list[int] = []
+
+        Returns ``(list_of_id_lists, list_of_mask_lists)`` where each
+        inner list corresponds to one element of the source list.  This
+        is critical for GRPO where each response must stay a separate
+        sequence so the strategy can form a ``[G, R]`` tensor.
+        """
+        per_item_ids: list[list[int]] = []
+        per_item_masks: list[list[int]] = []
 
         for sec in sections:
             field = sec["field"]
@@ -108,17 +115,13 @@ class SectionRenderer:
                 continue
 
             for val in values:
+                ids: list[int] = []
+                mask: list[int] = []
                 if use_template:
                     if isinstance(val, list):
                         wrapper = {field: val}
                         self._append_template(
-                            wrapper,
+                            wrapper, field, action, tokenizer, config, ids, mask
-                            field,
-                            action,
-                            tokenizer,
-                            config,
-                            all_ids,
-                            loss_mask,
                         )
                 else:
                     wrapper = {field: str(val)}
@@ -130,17 +133,19 @@ class SectionRenderer:
                         False,
                         False,
                         config,
-                        all_ids,
+                        ids,
-                        loss_mask,
+                        mask,
                     )
-
+                if ids:
                     max_len = config.preprocessing.max_seq_len
-        all_ids = all_ids[:max_len]
+                    ids = ids[:max_len]
-        loss_mask = loss_mask[: len(all_ids)]
+                    mask = mask[: len(ids)]
+                    per_item_ids.append(ids)
+                    per_item_masks.append(mask)
 
-        if not all_ids:
+        if not per_item_ids:
             return None, None
-        return all_ids, loss_mask
+        return per_item_ids, per_item_masks
 
     @staticmethod
     def is_value_section(sections: list) -> bool:
@@ -282,7 +287,15 @@ class MultiOutputMaskBuilder(BaseMaskBuilder):
                 ids, mask = self.renderer.process_list_field(
                     item, sections, config, tokenizer
                 )
-            else:
+                if ids is None:
+                    continue
+                # ids is List[List[int]] — preserve per-response structure
+                result[output_key] = ids
+                if mask is not None:
+                    result[mask_key] = mask
+                any_output = True
+                continue
+
             ids, mask = self.renderer.process_sections(
                 item, sections, config, tokenizer, is_top_level=True
             )

astrai/preprocessing/pipeline.py

@@ -180,6 +180,21 @@ class Pipeline:
                 dt = _STR_TO_DTYPE.get(
                     self.config.output.dtype.get(key, "int32"), torch.int32
                 )
+                # GRPO multi-response keys store List[List[int]] per record
+                # (responses/masks).  Rewards store List[float] per record.
+                # Both produce List[Tensor] (one tensor per record), but
+                # responses need inner flattening while rewards do not.
+                if ids_list and isinstance(ids_list[0], list):
+                    tensors[key] = [
+                        torch.tensor(
+                            list(chain.from_iterable(ids))
+                            if ids and isinstance(ids[0], list)
+                            else ids,
+                            dtype=dt,
+                        )
+                        for ids in ids_list
+                    ]
+                else:
                     tensors[key] = [
                         torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt)
                     ]

tests/data/test_dataset.py

@@ -327,43 +327,82 @@ def test_normalize_mixed_empty_key():
 
 
 def test_grpo_dataset_dtype(base_test_env):
+    """GRPO dataset returns correct dtypes for per-record structured data."""
+    from astrai.dataset.dataset import GRPODataset
+
     test_dir = base_test_env["test_dir"]
-    dummy_data = {
+    G = 4
-        "prompts": [torch.randint(0, 100, (100,), dtype=torch.int32)],
+    dataset = GRPODataset()
-        "responses": [torch.randint(0, 100, (100,), dtype=torch.int32)],
+    dataset.storage = type(
-        "masks": [torch.ones(100, dtype=torch.int32)],
+        "FakeStore",
-        "rewards": [torch.ones(100, dtype=torch.float32)],
+        (),
-    }
+        {
-    dataset = _make_seq_dataset(
+            "keys": ["prompts", "responses", "masks", "rewards"],
-        test_dir, "grpo_dtype", train_type="grpo", data=dummy_data, window_size=32
+            "_data": {
-    )
+                "prompts": [torch.randint(0, 100, (10,), dtype=torch.int32)],
+                "responses": [
+                    [torch.randint(0, 100, (5,), dtype=torch.int32) for _ in range(G)]
+                ],
+                "masks": [[torch.ones(5, dtype=torch.int32) for _ in range(G)]],
+                "rewards": [torch.rand(G, dtype=torch.float32)],
+            },
+        },
+    )()
+    dataset._build_records()
     item = dataset[0]
 
     assert item["prompts"].dtype == torch.long
-    assert item["responses"].dtype == torch.long
+    assert all(r.dtype == torch.long for r in item["responses"])
-    assert item["masks"].dtype == torch.bool
+    assert all(m.dtype == torch.bool for m in item["masks"])
     assert item["rewards"].dtype == torch.float32
 
 
 def test_grpo_dataset_load(base_test_env):
+    """GRPO dataset loads record-structured data with per-response boundaries."""
+    from astrai.dataset.dataset import GRPODataset
+
     test_dir = base_test_env["test_dir"]
-    dummy_data = {
+    G = 3
-        "prompts": [_rand_seq(200)],
+    prompt_len = 8
-        "responses": [_rand_seq(200)],
+    resp_lens = [5, 7, 4]
-        "masks": [torch.ones(200, dtype=torch.int64)],
+    dataset = GRPODataset()
-        "rewards": [torch.rand(200, dtype=torch.float32)],
+    dataset.storage = type(
-    }
+        "FakeStore",
-    dataset = _make_seq_dataset(
+        (),
-        test_dir, "grpo_test", train_type="grpo", data=dummy_data
+        {
-    )
+            "keys": ["prompts", "responses", "masks", "rewards"],
-    assert len(dataset) > 0
+            "_data": {
+                "prompts": [torch.randint(0, 100, (prompt_len,))],
+                "responses": [[torch.randint(0, 100, (rl,)) for rl in resp_lens]],
+                "masks": [[torch.ones(rl, dtype=torch.int64) for rl in resp_lens]],
+                "rewards": [torch.tensor([0.9, 0.3, 0.7], dtype=torch.float32)],
+            },
+        },
+    )()
+    dataset._build_records()
+
+    assert len(dataset) == 1
     item = dataset[0]
     assert "prompts" in item
     assert "responses" in item
     assert "masks" in item
     assert "rewards" in item
-    assert item["prompts"].shape[0] == 64
+
-    assert item["responses"].shape[0] == 64
+    # Prompts is 1-D
+    assert item["prompts"].shape == (prompt_len,)
+
+    # Responses is a list of G tensors with correct lengths
+    assert len(item["responses"]) == G
+    for i, r in enumerate(item["responses"]):
+        assert r.shape == (resp_lens[i],)
+
+    # Masks align with responses
+    assert len(item["masks"]) == G
+    for i, m in enumerate(item["masks"]):
+        assert m.shape == (resp_lens[i],)
+
+    # Rewards has G elements
+    assert item["rewards"].shape == (G,)
 
 
 def test_detect_format_bin_dir(base_test_env):
@@ -621,3 +660,231 @@ def test_jsonl_store_pipeline_config_roundtrip(base_test_env):
     config = PipelineConfig.from_dict(raw)
     assert config.output.position_ids_mode == "doc_reset"
     assert config.preprocessing.max_seq_len == 64
+
+
+# ---------------------------------------------------------------------------
+# GRPO end-to-end: builder → JsonlStore → GRPODataset → collate_fn
+# ---------------------------------------------------------------------------
+
+
+def _write_grpo_jsonl(test_dir, tokenizer_path, records):
+    """Write a GRPO JSONL dataset directory with config."""
+    data_dir = os.path.join(test_dir, "grpo_jsonl")
+    os.makedirs(data_dir, exist_ok=True)
+
+    with open(os.path.join(data_dir, "data.jsonl"), "w", encoding="utf-8") as f:
+        for rec in records:
+            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
+
+    config = {
+        "tokenizer_path": tokenizer_path,
+        "version": 1,
+        "input": {
+            "sources": {
+                "prompts": {
+                    "sections": [
+                        {
+                            "field": "prompt",
+                            "action": "mask",
+                            "add_special_tokens": True,
+                        }
+                    ]
+                },
+                "responses": {
+                    "sections": [{"field": "responses", "action": "train"}],
+                    "list_field": True,
+                    "mask_key": "masks",
+                },
+                "rewards": {
+                    "sections": [{"field": "rewards", "action": "value"}],
+                },
+            }
+        },
+        "mask": {"user": "mask", "assistant": "train"},
+        "mask_default": "mask",
+        "preprocessing": {"max_seq_len": 128},
+        "output": {"position_ids_mode": "none"},
+    }
+
+    with open(
+        os.path.join(data_dir, "dataset_config.json"), "w", encoding="utf-8"
+    ) as f:
+        json.dump(config, f, ensure_ascii=False, indent=2)
+
+    return data_dir
+
+
+def test_grpo_builder_preserves_response_boundaries(base_test_env):
+    """MultiOutputMaskBuilder with list_field returns List[List[int]] for responses."""
+    from astrai.preprocessing.builder import SectionedMaskBuilder
+    from tests.data.conftest import make_grpo_no_template_config
+
+    tokenizer = base_test_env["tokenizer"]
+    tokenizer_path = _save_test_tokenizer(base_test_env["test_dir"], tokenizer)
+
+    builder = SectionedMaskBuilder()
+    config = make_grpo_no_template_config()
+    config.preprocessing.max_seq_len = 128
+
+    item = {
+        "prompt": "What is 2+2?",
+        "responses": ["4", "four", "2+2=4"],
+        "rewards": [0.9, 0.1, 0.5],
+    }
+
+    result = builder.build(item, config, tokenizer)
+    assert result is not None
+
+    # prompts should be flat list of ints
+    assert isinstance(result["prompts"], list)
+    assert isinstance(result["prompts"][0], int)
+
+    # responses should be list of lists (one per response)
+    assert isinstance(result["responses"], list)
+    assert isinstance(result["responses"][0], list)
+    assert isinstance(result["responses"][0][0], int)
+    assert len(result["responses"]) == 3
+
+    # masks should match responses structure
+    assert isinstance(result["masks"], list)
+    assert len(result["masks"]) == 3
+    for i in range(3):
+        assert len(result["masks"][i]) == len(result["responses"][i])
+
+    # rewards should be flat list of floats
+    assert isinstance(result["rewards"], list)
+    assert all(isinstance(r, float) for r in result["rewards"])
+    assert len(result["rewards"]) == 3
+
+
+def test_grpo_end_to_end_jsonl(base_test_env):
+    """Full GRPO pipeline: JSONL → JsonlStore → GRPODataset → collate_fn."""
+    from astrai.dataset.dataset import grpo_collate_fn
+
+    test_dir = base_test_env["test_dir"]
+    tokenizer = base_test_env["tokenizer"]
+    tokenizer_path = _save_test_tokenizer(test_dir, tokenizer)
+
+    records = [
+        {
+            "prompt": "What is 2+2?",
+            "responses": ["4", "four", "The answer is 4"],
+            "rewards": [0.9, 0.1, 0.5],
+        },
+        {
+            "prompt": "Write a haiku",
+            "responses": ["Leaves fall", "Cherry blossoms bloom in spring"],
+            "rewards": [0.3, 0.8],
+        },
+    ]
+
+    data_dir = _write_grpo_jsonl(test_dir, tokenizer_path, records)
+
+    dataset = DatasetFactory.load("grpo", data_dir, window_size=0)
+    assert len(dataset) == 2
+
+    # Item 0: 3 responses
+    item0 = dataset[0]
+    assert item0["prompts"].ndim == 1
+    assert len(item0["responses"]) == 3
+    assert len(item0["masks"]) == 3
+    assert item0["rewards"].shape == (3,)
+    for r, m in zip(item0["responses"], item0["masks"]):
+        assert r.shape == m.shape
+
+    # Item 1: 2 responses (different group size)
+    item1 = dataset[1]
+    assert len(item1["responses"]) == 2
+    assert item1["rewards"].shape == (2,)
+
+    # Collate: batch records with same G (item0 has G=3)
+    batch = grpo_collate_fn([item0, item0])
+    assert batch["prompts"].shape[0] == 2
+    assert batch["responses"].ndim == 3
+    assert batch["responses"].shape[0] == 2
+    assert batch["responses"].shape[1] == 3  # G=3
+    assert batch["masks"].shape == batch["responses"].shape
+    assert batch["rewards"].shape == (2, 3)
+
+
+def test_grpo_collate_variable_lengths():
+    """collate_fn pads variable-length responses to [B, G, R_max]."""
+    from astrai.dataset.dataset import grpo_collate_fn
+
+    batch = [
+        {
+            "prompts": torch.tensor([1, 2, 3]),
+            "responses": [torch.tensor([4, 5]), torch.tensor([6, 7, 8, 9])],
+            "masks": [torch.tensor([1, 1]), torch.tensor([1, 1, 1, 1])],
+            "rewards": torch.tensor([0.9, 0.1]),
+        },
+        {
+            "prompts": torch.tensor([10, 11]),
+            "responses": [torch.tensor([12]), torch.tensor([13, 14, 15])],
+            "masks": [torch.tensor([1]), torch.tensor([1, 1, 1])],
+            "rewards": torch.tensor([0.5, 0.5]),
+        },
+    ]
+
+    result = grpo_collate_fn(batch)
+
+    assert result["prompts"].shape == (2, 3)  # B=2, P_max=3
+    assert result["responses"].shape == (2, 2, 4)  # B=2, G=2, R_max=4
+    assert result["masks"].shape == (2, 2, 4)
+    assert result["rewards"].shape == (2, 2)
+
+    # Check padding: item 1 prompt is length 2, padded to 3
+    assert result["prompts"][1, 2] == 0
+
+    # Check response content: item 0, response 0 is [4,5] padded to 4
+    assert result["responses"][0, 0, 0] == 4
+    assert result["responses"][0, 0, 1] == 5
+    assert result["responses"][0, 0, 2] == 0  # padded
+    assert result["masks"][0, 0, 2] == False  # padded
+
+    # Check response content: item 0, response 1 is [6,7,8,9] no padding
+    assert result["responses"][0, 1, 3] == 9
+    assert result["masks"][0, 1, 3] == True
+
+
+def test_grpo_multiple_records(base_test_env):
+    """GRPODataset loads multiple records with correct structure."""
+    from astrai.dataset.dataset import GRPODataset
+
+    G = 4
+    n_records = 5
+
+    dummy_responses = [
+        [torch.randint(0, 100, (np.random.randint(3, 8),)) for _ in range(G)]
+        for _ in range(n_records)
+    ]
+    dataset = GRPODataset()
+    dataset.storage = type(
+        "FakeStore",
+        (),
+        {
+            "keys": ["prompts", "responses", "masks", "rewards"],
+            "_data": {
+                "prompts": [torch.randint(0, 100, (10,)) for _ in range(n_records)],
+                "responses": dummy_responses,
+                "masks": [
+                    [torch.ones(r.shape[0], dtype=torch.int64) for r in resps]
+                    for resps in dummy_responses
+                ],
+                "rewards": [
+                    torch.rand(G, dtype=torch.float32) for _ in range(n_records)
+                ],
+            },
+        },
+    )()
+    dataset._build_records()
+
+    assert len(dataset) == n_records
+
+    for i in range(n_records):
+        item = dataset[i]
+        assert len(item["responses"]) == G
+        assert len(item["masks"]) == G
+        assert item["rewards"].shape == (G,)
+        for g in range(G):
+            assert item["responses"][g].shape == item["masks"][g].shape

tests/data/test_preprocess_builder.py

@@ -349,7 +349,17 @@ def test_grpo_basic(chat_tokenizer, builder):
     assert "responses" in result
     assert "masks" in result
     assert "rewards" in result
-    assert len(result["responses"]) == len(result["masks"])
+
+    # responses is List[List[int]] — one per response
+    assert len(result["responses"]) == 4
+    assert all(isinstance(r, list) for r in result["responses"])
+    assert all(isinstance(r[0], int) for r in result["responses"])
+
+    # masks is List[List[int]] — one per response, matching length
+    assert len(result["masks"]) == 4
+    for i in range(4):
+        assert len(result["masks"][i]) == len(result["responses"][i])
+
     assert result["rewards"] == [1.0, 0.5, 0.8, 0.2]
 
 
@@ -362,8 +372,11 @@ def test_grpo_response_tokens_all_trained(chat_tokenizer, builder):
     }
     result = builder.build(item, config, chat_tokenizer)
     masks = result["masks"]
-    assert all(m == 1 for m in masks)
+    # masks is List[List[int]] — each response's mask should be all 1s
-    assert len(masks) == len(result["responses"])
+    assert len(masks) == 2
+    for m in masks:
+        assert all(v == 1 for v in m)
+        assert len(m) == len(result["responses"][masks.index(m)])
 
 
 def test_grpo_single_reward(chat_tokenizer, builder):