Buckets:
| import unittest | |
| import torch | |
| from shared.supervised import pack_visual_codes, unpack_visual_codes | |
| from sft.config import InformationGainConfig | |
| from sft.data import ( | |
| build_transition_events, | |
| cached_sample, | |
| collate_samples, | |
| plan_epoch_batches, | |
| plan_evaluation_batches, | |
| ) | |
| from uniar import SPECIAL_TOKEN_IDS | |
| class SFTDataTest(unittest.TestCase): | |
| def test_shared_visual_code_packing_roundtrips_little_endian_bits(self): | |
| codes = torch.randint(0, 2, (3, 2, 4, 7), dtype=torch.uint8) | |
| packed, shape = pack_visual_codes(codes) | |
| restored = unpack_visual_codes(packed, shape) | |
| self.assertTrue(torch.equal(restored, codes)) | |
| def _row(index: int, length: int = 8) -> dict: | |
| return { | |
| "sample_id": f"sample-{index}", | |
| "sequence_length": length, | |
| "cache_shard": "shard.pt", | |
| } | |
| def test_adjacent_images_keep_the_intervening_scaffold_text_target(self): | |
| token = SPECIAL_TOKEN_IDS | |
| ids = [99, 98] | |
| ids += [token["<think>"], 10, token["</think>"]] | |
| ids += [77, token["<|vision_start|>"], 151655, 151655, token["<|vision_end|>"]] | |
| ids += [ | |
| 77, | |
| token["<|vision_start|>"], | |
| 151655, | |
| 151655, | |
| 151655, | |
| token["<|vision_end|>"], | |
| ] | |
| ids += [ | |
| token["<think>"], | |
| 20, | |
| 21, | |
| token["</think>"], | |
| 30, | |
| 31, | |
| token["<|im_end|>"], | |
| ] | |
| generated = torch.zeros(len(ids), dtype=torch.bool) | |
| first = [index for index, value in enumerate(ids) if value == 151655] | |
| generated[first] = True | |
| item = { | |
| "sample_id": "fixture", | |
| "input_ids": torch.tensor(ids), | |
| "prefix_length": 2, | |
| "generated_sequence_mask": generated, | |
| "generated_flat_mask": torch.tensor([False, True, True, True, True, True]), | |
| "generated_image_token_counts": torch.tensor([2, 3]), | |
| "trajectory": { | |
| "assistant": [ | |
| {"role": "thought"}, | |
| {"role": "generated_state"}, | |
| {"role": "generated_state"}, | |
| {"role": "thought"}, | |
| {"role": "answer"}, | |
| ] | |
| }, | |
| } | |
| events = build_transition_events(item, InformationGainConfig()) | |
| self.assertEqual( | |
| [event.direction for event in events], | |
| ["text_to_visual", "visual_to_text", "visual_to_text"], | |
| ) | |
| self.assertEqual( | |
| [span.length for span in events[0].target_sequence_spans], [2] | |
| ) | |
| self.assertEqual( | |
| [span.length for span in events[0].target_visual_spans], [2] | |
| ) | |
| self.assertEqual(events[0].target_steps, 2) | |
| self.assertEqual(len(events[1].source_sequence_spans), 1) | |
| self.assertEqual(len(events[1].target_sequence_spans), 1) | |
| self.assertEqual( | |
| [span.length for span in events[1].source_sequence_spans], [2] | |
| ) | |
| self.assertEqual( | |
| [span.length for span in events[1].target_sequence_spans], [2] | |
| ) | |
| self.assertEqual(events[1].target_steps, 2) | |
| self.assertEqual( | |
| [ | |
| ids[position] | |
| for span in events[1].target_sequence_spans | |
| for position in range(span.start, span.end) | |
| ], | |
| [77, token["<|vision_start|>"]], | |
| ) | |
| self.assertEqual(len(events[2].source_sequence_spans), 1) | |
| self.assertEqual( | |
| [span.length for span in events[2].source_sequence_spans], [3] | |
| ) | |
| self.assertEqual( | |
| [span.length for span in events[2].target_sequence_spans], [4, 3] | |
| ) | |
| self.assertEqual(events[2].target_steps, 7) | |
| self.assertEqual( | |
| ids[events[2].target_sequence_spans[-1].end - 1], | |
| token["<|im_end|>"], | |
| ) | |
| def test_intermediate_text_target_owns_autonomous_image_scaffold(self): | |
| token = SPECIAL_TOKEN_IDS | |
| ids = [99, 98] | |
| ids += [token["<think>"], 10, token["</think>"], 77] | |
| ids += [token["<|vision_start|>"], 151655, token["<|vision_end|>"]] | |
| ids += [22, token["<think>"], 20, token["</think>"], 78, 79] | |
| ids += [token["<|vision_start|>"], 151655, token["<|vision_end|>"]] | |
| ids += [30, token["<|im_end|>"]] | |
| generated = torch.tensor([value == 151655 for value in ids]) | |
| item = { | |
| "sample_id": "window", | |
| "input_ids": torch.tensor(ids), | |
| "prefix_length": 2, | |
| "generated_sequence_mask": generated, | |
| "generated_flat_mask": torch.tensor([True, True]), | |
| "generated_image_token_counts": torch.tensor([1, 1]), | |
| "trajectory": { | |
| "assistant": [ | |
| {"role": "thought"}, | |
| {"role": "generated_state"}, | |
| {"role": "thought"}, | |
| {"role": "generated_state"}, | |
| {"role": "answer"}, | |
| ] | |
| }, | |
| } | |
| events = build_transition_events(item, InformationGainConfig()) | |
| visual_to_text = events[1] | |
| target_ids = [ | |
| value | |
| for span in visual_to_text.target_sequence_spans | |
| for value in ids[span.start : span.end] | |
| ] | |
| self.assertEqual( | |
| target_ids, | |
| [ | |
| 22, | |
| token["<think>"], | |
| 20, | |
| token["</think>"], | |
| 78, | |
| 79, | |
| token["<|vision_start|>"], | |
| ], | |
| ) | |
| self.assertEqual(target_ids[-1], token["<|vision_start|>"]) | |
| self.assertNotIn(token["<|vision_end|>"], target_ids) | |
| text_to_visual = events[2] | |
| visual_target_ids = [ | |
| ids[i] | |
| for span in text_to_visual.target_sequence_spans | |
| for i in range(span.start, span.end) | |
| ] | |
| self.assertEqual(visual_target_ids, [151655]) | |
| self.assertEqual( | |
| visual_to_text.target_sequence_spans[0].end, | |
| text_to_visual.target_sequence_spans[0].start, | |
| ) | |
| # IG and EMA consume the same TransitionEvent: only semantic thought | |
| # payload is blinded, while the student-owned delimiter is scored. | |
| self.assertEqual( | |
| [ | |
| ids[i] | |
| for span in text_to_visual.source_sequence_spans | |
| for i in range(span.start, span.end) | |
| ], | |
| [20], | |
| ) | |
| def test_cached_text_labels_follow_generated_delimiter_ownership(self): | |
| token = SPECIAL_TOKEN_IDS | |
| ids = [99, 98] | |
| ids += [token["<think>"], 10, token["</think>"], 77] | |
| ids += [token["<|vision_start|>"], 151655, token["<|vision_end|>"]] | |
| ids += [30, token["<|im_end|>"]] | |
| vision_start = ids.index(token["<|vision_start|>"]) | |
| visual = ids.index(151655) | |
| vision_end = ids.index(token["<|vision_end|>"]) | |
| answer = ids.index(30) | |
| im_end = ids.index(token["<|im_end|>"]) | |
| raw_labels = torch.tensor(ids) | |
| raw_labels[:2] = -100 | |
| raw_labels[visual] = -100 | |
| codes = torch.zeros(1, 4, 4, 64, dtype=torch.uint8) | |
| packed, shape = pack_visual_codes(codes) | |
| item = { | |
| "schema": "zebra_uniar_cache_v1", | |
| "sample_id": "label-ownership", | |
| "config": "fixture", | |
| "source": {}, | |
| "input_ids": torch.tensor(ids), | |
| "labels": raw_labels, | |
| "prefix_length": 2, | |
| "image_grid_thw": torch.tensor([[1, 2, 2]]), | |
| "generated_sequence_mask": torch.tensor( | |
| [value == 151655 for value in ids] | |
| ), | |
| "generated_flat_mask": torch.tensor([True]), | |
| "generated_image_token_counts": torch.tensor([1]), | |
| "visual_codes_packed": packed, | |
| "visual_codes_shape": shape, | |
| "trajectory": { | |
| "assistant": [ | |
| {"role": "thought"}, | |
| {"role": "generated_state"}, | |
| {"role": "answer"}, | |
| ] | |
| }, | |
| } | |
| sample = cached_sample(item, InformationGainConfig()) | |
| batch = collate_samples([sample], pad_token_id=0) | |
| self.assertEqual(int(batch.labels[0, vision_start]), ids[vision_start]) | |
| self.assertEqual(int(batch.labels[0, visual]), -100) | |
| self.assertEqual(int(batch.labels[0, vision_end]), -100) | |
| self.assertEqual(int(batch.labels[0, answer]), ids[answer]) | |
| self.assertEqual(int(batch.labels[0, im_end]), ids[im_end]) | |
| self.assertEqual(int(raw_labels[vision_start]), ids[vision_start]) | |
| self.assertEqual(int(raw_labels[vision_end]), ids[vision_end]) | |
| def test_repeat_tail_cycles_even_when_batches_are_fewer_than_ranks(self): | |
| groups = plan_epoch_batches( | |
| [self._row(0)], | |
| seed=0, | |
| epoch=0, | |
| micro_batch_size=1, | |
| batch_padded_token_budget=16, | |
| max_sequence_tokens=32, | |
| world_size=4, | |
| distributed_tail="repeat", | |
| shuffle=False, | |
| ) | |
| self.assertEqual(len(groups), 1) | |
| self.assertEqual(len(groups[0]), 4) | |
| self.assertTrue(all(group[0]["sample_id"] == "sample-0" for group in groups[0])) | |
| def test_hard_context_limit_and_soft_batch_budget_are_distinct(self): | |
| groups = plan_epoch_batches( | |
| [self._row(0, length=24)], | |
| seed=0, | |
| epoch=0, | |
| micro_batch_size=2, | |
| batch_padded_token_budget=16, | |
| max_sequence_tokens=32, | |
| world_size=1, | |
| distributed_tail="drop", | |
| shuffle=False, | |
| ) | |
| self.assertEqual(groups[0][0][0]["sequence_length"], 24) | |
| with self.assertRaisesRegex(ValueError, "max_sequence_tokens"): | |
| plan_epoch_batches( | |
| [self._row(0, length=33)], | |
| seed=0, | |
| epoch=0, | |
| micro_batch_size=1, | |
| batch_padded_token_budget=16, | |
| max_sequence_tokens=32, | |
| world_size=1, | |
| distributed_tail="drop", | |
| shuffle=False, | |
| ) | |
| def test_shuffle_preserves_shard_locality(self): | |
| rows = [] | |
| for shard in range(4): | |
| for index in range(5): | |
| row = self._row(shard * 10 + index) | |
| row["cache_shard"] = f"shard-{shard}.pt" | |
| rows.append(row) | |
| groups = plan_epoch_batches( | |
| rows, | |
| seed=31, | |
| epoch=2, | |
| micro_batch_size=1, | |
| batch_padded_token_budget=16, | |
| max_sequence_tokens=32, | |
| world_size=1, | |
| distributed_tail="drop", | |
| shuffle=True, | |
| ) | |
| sequence = [group[0][0]["cache_shard"] for group in groups] | |
| runs = [ | |
| shard | |
| for index, shard in enumerate(sequence) | |
| if index == 0 or shard != sequence[index - 1] | |
| ] | |
| self.assertEqual(len(runs), 4) | |
| self.assertEqual(set(runs), {f"shard-{index}.pt" for index in range(4)}) | |
| def test_evaluation_batches_preserve_rows_and_assign_whole_shards(self): | |
| rows = [] | |
| for shard, lengths in enumerate(((3, 4, 5), (6, 7), (2, 8, 9))): | |
| for index, length in enumerate(lengths): | |
| row = self._row(shard * 10 + index, length=length) | |
| row["cache_shard"] = f"shard-{shard}.pt" | |
| rows.append(row) | |
| ranks = [ | |
| plan_evaluation_batches( | |
| rows, | |
| micro_batch_size=2, | |
| batch_padded_token_budget=16, | |
| rank=rank, | |
| world_size=2, | |
| ) | |
| for rank in range(2) | |
| ] | |
| selected = [row for batches in ranks for batch in batches for row in batch] | |
| self.assertEqual( | |
| sorted(row["sample_id"] for row in selected), | |
| sorted(row["sample_id"] for row in rows), | |
| ) | |
| shard_ranks = {} | |
| for rank, batches in enumerate(ranks): | |
| for batch in batches: | |
| self.assertEqual(len({row["cache_shard"] for row in batch}), 1) | |
| self.assertLessEqual(len(batch), 2) | |
| self.assertLessEqual( | |
| max(row["sequence_length"] for row in batch) * len(batch), 16 | |
| ) | |
| for row in batch: | |
| shard_ranks.setdefault(row["cache_shard"], set()).add(rank) | |
| self.assertTrue(all(len(values) == 1 for values in shard_ranks.values())) | |
| if __name__ == "__main__": | |
| unittest.main() | |
Xet Storage Details
- Size:
- 12.8 kB
- Xet hash:
- 27c163fa452c16faa614865f6f3b8f0dde182232d6aef0c4e272c0e2dd4a0c4b
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.