diff --git a/astrai/dataset/__init__.py b/astrai/dataset/__init__.py index 562fa6f..51e4dc4 100644 --- a/astrai/dataset/__init__.py +++ b/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", diff --git a/astrai/dataset/dataset.py b/astrai/dataset/dataset.py index 4f07010..e996fab 100644 --- a/astrai/dataset/dataset.py +++ b/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: - return self.storage.fetch(begin_idx, end_idx, key) + def load(self, load_path: str, storage_type: Optional[str] = None, **kwargs): + 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) - - 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") - + rec = self._records[index] return { - "prompts": prompts, - "responses": responses, - "masks": masks, - "rewards": rewards, + "prompts": rec["prompts"].to(dtype=torch.long), + "responses": [r.to(dtype=torch.long) for r in rec["responses"]], + "masks": [m.to(dtype=torch.bool) for m in rec["masks"]], + "rewards": rec["rewards"].to(dtype=torch.float32), } diff --git a/astrai/dataset/storage.py b/astrai/dataset/storage.py index e78575d..ac68efd 100644 --- a/astrai/dataset/storage.py +++ b/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 = ( - min((cum[-1] if cum else 0) for cum in self._cum.values()) - if self._cum - else 0 - ) + flat_lengths.append(cum[-1] if cum else 0) + self._length = min(flat_lengths) if flat_lengths else 0 class StoreFactory(BaseFactory["Store"]): @@ -274,7 +285,13 @@ class JsonlStore(Store): for key, ids in result.items(): if key not in raw: raw[key] = [] - raw[key].append(torch.tensor(ids, dtype=self._infer_dtype(ids))) + 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")): with open(jsonl_path, "r", encoding="utf-8") as f: @@ -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 diff --git a/astrai/preprocessing/builder.py b/astrai/preprocessing/builder.py index b798e3e..17da350 100644 --- a/astrai/preprocessing/builder.py +++ b/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] = [] - loss_mask: list[int] = [] + """Tokenize a list-valued field, preserving per-element boundaries. + + 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, - field, - action, - tokenizer, - config, - all_ids, - loss_mask, + wrapper, field, action, tokenizer, config, ids, mask ) else: wrapper = {field: str(val)} @@ -130,17 +133,19 @@ class SectionRenderer: False, False, config, - all_ids, - loss_mask, + ids, + mask, ) + if ids: + max_len = config.preprocessing.max_seq_len + ids = ids[:max_len] + mask = mask[: len(ids)] + per_item_ids.append(ids) + per_item_masks.append(mask) - max_len = config.preprocessing.max_seq_len - all_ids = all_ids[:max_len] - loss_mask = loss_mask[: len(all_ids)] - - 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,10 +287,18 @@ class MultiOutputMaskBuilder(BaseMaskBuilder): ids, mask = self.renderer.process_list_field( item, sections, config, tokenizer ) - else: - ids, mask = self.renderer.process_sections( - item, sections, config, tokenizer, is_top_level=True - ) + 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 + ) if ids is None: continue diff --git a/astrai/preprocessing/pipeline.py b/astrai/preprocessing/pipeline.py index a351fa9..fca8c21 100644 --- a/astrai/preprocessing/pipeline.py +++ b/astrai/preprocessing/pipeline.py @@ -180,9 +180,24 @@ class Pipeline: dt = _STR_TO_DTYPE.get( self.config.output.dtype.get(key, "int32"), torch.int32 ) - tensors[key] = [ - torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt) - ] + # 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) + ] if mode == "continuous" and original_sequences: pos_ids = self._position_id.generate(keys.get("sequence", [])) diff --git a/tests/data/test_dataset.py b/tests/data/test_dataset.py index cd68e41..41c8a1f 100644 --- a/tests/data/test_dataset.py +++ b/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 = { - "prompts": [torch.randint(0, 100, (100,), dtype=torch.int32)], - "responses": [torch.randint(0, 100, (100,), dtype=torch.int32)], - "masks": [torch.ones(100, dtype=torch.int32)], - "rewards": [torch.ones(100, dtype=torch.float32)], - } - dataset = _make_seq_dataset( - test_dir, "grpo_dtype", train_type="grpo", data=dummy_data, window_size=32 - ) + G = 4 + dataset = GRPODataset() + dataset.storage = type( + "FakeStore", + (), + { + "keys": ["prompts", "responses", "masks", "rewards"], + "_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 item["masks"].dtype == torch.bool + assert all(r.dtype == torch.long for r in item["responses"]) + 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 = { - "prompts": [_rand_seq(200)], - "responses": [_rand_seq(200)], - "masks": [torch.ones(200, dtype=torch.int64)], - "rewards": [torch.rand(200, dtype=torch.float32)], - } - dataset = _make_seq_dataset( - test_dir, "grpo_test", train_type="grpo", data=dummy_data - ) - assert len(dataset) > 0 + G = 3 + prompt_len = 8 + resp_lens = [5, 7, 4] + dataset = GRPODataset() + dataset.storage = type( + "FakeStore", + (), + { + "keys": ["prompts", "responses", "masks", "rewards"], + "_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 diff --git a/tests/data/test_preprocess_builder.py b/tests/data/test_preprocess_builder.py index 03a6fc7..abc8aa7 100644 --- a/tests/data/test_preprocess_builder.py +++ b/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) - assert len(masks) == len(result["responses"]) + # masks is List[List[int]] — each response's mask should be all 1s + 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):