diff --git a/deepspec/data/target_cache_dataset.py b/deepspec/data/target_cache_dataset.py index d6e7b98..523a5d6 100644 --- a/deepspec/data/target_cache_dataset.py +++ b/deepspec/data/target_cache_dataset.py @@ -59,18 +59,20 @@ def expected_target_cache_tensor_nbytes( seq_len: int, hidden_size: int, num_target_layers: int, + hidden_dtype: str = "bfloat16", ): numel = expected_target_cache_tensor_numel( seq_len=seq_len, hidden_size=hidden_size, num_target_layers=num_target_layers, ) + hidden_element_size = 2 if hidden_dtype == "bfloat16" else 1 return { "input_ids": numel["input_ids"] * 4, "attention_mask": numel["attention_mask"], "loss_mask": numel["loss_mask"], - "target_hidden_states": numel["target_hidden_states"] * 2, - "target_last_hidden_states": numel["target_last_hidden_states"] * 2, + "target_hidden_states": numel["target_hidden_states"] * hidden_element_size, + "target_last_hidden_states": numel["target_last_hidden_states"] * hidden_element_size, } @@ -151,9 +153,10 @@ def validate_target_cache_manifest(*, cache_dir: str, manifest): "Unsupported target cache manifest version: " f"{manifest['version']} != {TARGET_CACHE_VERSION}" ) - assert manifest["hidden_dtype"] == TARGET_CACHE_HIDDEN_DTYPE, ( - "Unsupported hidden_dtype in target cache manifest: " - f"{manifest['hidden_dtype']}" + supported_hidden_dtypes = ("bfloat16", "float8_e4m3fn", "float8_e5m2") + assert manifest["hidden_dtype"] in supported_hidden_dtypes, ( + f"Unsupported hidden_dtype in target cache manifest: {manifest['hidden_dtype']}. " + f"Supported dtypes are {supported_hidden_dtypes}." ) assert manifest["token_dtype"] == TARGET_CACHE_TOKEN_DTYPE, ( "Unsupported token_dtype in target cache manifest: " @@ -228,6 +231,11 @@ def _tensor_to_bfloat16_bytes(tensor: torch.Tensor): return cpu_tensor.view(torch.uint16).numpy().tobytes() +def _tensor_to_float8_bytes(tensor: torch.Tensor, dtype: torch.dtype): + cpu_tensor = tensor.detach().to(device="cpu", dtype=dtype).contiguous() + return cpu_tensor.view(torch.uint8).numpy().tobytes() + + def compute_local_sample_range(*, num_samples: int, rank: int, world_size: int): base = int(num_samples) // int(world_size) remainder = int(num_samples) % int(world_size) @@ -288,24 +296,36 @@ def build_target_cache_sample_bytes( loss_mask: torch.Tensor, target_hidden_states: torch.Tensor, target_last_hidden_states: torch.Tensor, + hidden_dtype: str = "bfloat16", ): + if hidden_dtype == "bfloat16": + th_bytes = _tensor_to_bfloat16_bytes(target_hidden_states) + tlh_bytes = _tensor_to_bfloat16_bytes(target_last_hidden_states) + elif hidden_dtype == "float8_e4m3fn": + th_bytes = _tensor_to_float8_bytes(target_hidden_states, torch.float8_e4m3fn) + tlh_bytes = _tensor_to_float8_bytes(target_last_hidden_states, torch.float8_e4m3fn) + elif hidden_dtype == "float8_e5m2": + th_bytes = _tensor_to_float8_bytes(target_hidden_states, torch.float8_e5m2) + tlh_bytes = _tensor_to_float8_bytes(target_last_hidden_states, torch.float8_e5m2) + else: + raise ValueError(f"Unsupported hidden_dtype for writing bytes: {hidden_dtype}") + return TargetCacheSampleBytes( sample_id=int(sample_id), seq_len=int(input_ids.shape[0]), input_ids=_tensor_to_bytes(input_ids, torch.int32), attention_mask=_tensor_to_bytes(attention_mask, torch.uint8), loss_mask=_tensor_to_bytes(loss_mask, torch.uint8), - target_hidden_states=_tensor_to_bfloat16_bytes(target_hidden_states), - target_last_hidden_states=_tensor_to_bfloat16_bytes( - target_last_hidden_states - ), + target_hidden_states=th_bytes, + target_last_hidden_states=tlh_bytes, ) class LocalTargetCacheWriter: - def __init__(self, *, rank_dir: str, max_shard_bytes: int): + def __init__(self, *, rank_dir: str, max_shard_bytes: int, hidden_dtype: str = "bfloat16"): self.rank_dir = rank_dir self.max_shard_bytes = int(max_shard_bytes) + self.hidden_dtype = hidden_dtype self.local_index_path = os.path.join(rank_dir, "samples.local.idx") self.index_handle = open(self.local_index_path, "wb") self.current_shard_id = -1 @@ -403,6 +423,7 @@ def write_sample( loss_mask=loss_mask, target_hidden_states=target_hidden_states, target_last_hidden_states=target_last_hidden_states, + hidden_dtype=self.hidden_dtype, ) self.write_sample_bytes(sample) @@ -414,10 +435,13 @@ def __init__( rank_dir: str, max_shard_bytes: int, max_queue_size: int = 128, + hidden_dtype: str = "bfloat16", ): + self.hidden_dtype = hidden_dtype self.writer = LocalTargetCacheWriter( rank_dir=rank_dir, max_shard_bytes=max_shard_bytes, + hidden_dtype=hidden_dtype, ) # Queue CPU byte records only; never hold CUDA tensor references here. self.queue = queue.Queue(maxsize=int(max_queue_size)) @@ -483,6 +507,7 @@ def write_sample( loss_mask=loss_mask, target_hidden_states=target_hidden_states, target_last_hidden_states=target_last_hidden_states, + hidden_dtype=self.hidden_dtype, ) self._put(sample) self.num_local_samples += 1 @@ -583,6 +608,7 @@ def build_target_cache_manifest( shards, target_layer_ids, hidden_size: int, + hidden_dtype: str = "bfloat16", extra_fields=None, ): manifest = { @@ -590,7 +616,7 @@ def build_target_cache_manifest( "num_samples": int(num_samples), "num_shards": len(shards), "target_layer_ids": [int(layer_id) for layer_id in target_layer_ids], - "hidden_dtype": TARGET_CACHE_HIDDEN_DTYPE, + "hidden_dtype": hidden_dtype, "token_dtype": TARGET_CACHE_TOKEN_DTYPE, "mask_dtype": TARGET_CACHE_MASK_DTYPE, "index_record_size": INDEX_RECORD_SIZE, @@ -621,6 +647,7 @@ def __init__(self, cache_dir: str, max_open_shards: int = 4): self.hidden_size = int(self.manifest["hidden_size"]) self.target_layer_ids = [int(layer_id) for layer_id in self.manifest["target_layer_ids"]] self.num_target_layers = len(self.target_layer_ids) + self.hidden_dtype = self.manifest.get("hidden_dtype", "bfloat16") self.index_path = os.path.join(self.cache_dir, "samples.idx") self.index_file = None self.index_mmap = None @@ -750,6 +777,28 @@ def _read_bfloat16_tensor_from_shard( tensor = torch.from_numpy(array).view(torch.bfloat16) return tensor.view(*shape) + def _read_float8_tensor_from_shard( + self, + *, + shard_mmap, + offset: int, + shape, + nbytes: int, + torch_dtype, + ): + assert int(offset) + int(nbytes) <= shard_mmap.size(), ( + "Target cache tensor extends beyond shard size: " + f"offset={offset}, nbytes={nbytes}, shard_size={shard_mmap.size()}" + ) + array = np.frombuffer( + shard_mmap, + dtype=np.uint8, + count=int(np.prod(shape)), + offset=int(offset), + ).copy() + tensor = torch.from_numpy(array).view(torch_dtype) + return tensor.view(*shape) + def __getitem__(self, index: int): if not (0 <= int(index) < self.num_samples): raise IndexError(index) @@ -761,6 +810,7 @@ def __getitem__(self, index: int): seq_len=seq_len, hidden_size=self.hidden_size, num_target_layers=self.num_target_layers, + hidden_dtype=self.hidden_dtype, ) input_ids = self._read_tensor_from_shard( shard_mmap=shard_mmap, @@ -778,18 +828,51 @@ def __getitem__(self, index: int): torch_dtype=torch.uint8, nbytes=nbytes["loss_mask"], ) - target_hidden_states = self._read_bfloat16_tensor_from_shard( - shard_mmap=shard_mmap, - offset=record["target_hidden_states_offset"], - shape=(seq_len, self.num_target_layers * self.hidden_size), - nbytes=nbytes["target_hidden_states"], - ) - target_last_hidden_states = self._read_bfloat16_tensor_from_shard( - shard_mmap=shard_mmap, - offset=record["target_last_hidden_states_offset"], - shape=(seq_len, self.hidden_size), - nbytes=nbytes["target_last_hidden_states"], - ) + if self.hidden_dtype == "bfloat16": + target_hidden_states = self._read_bfloat16_tensor_from_shard( + shard_mmap=shard_mmap, + offset=record["target_hidden_states_offset"], + shape=(seq_len, self.num_target_layers * self.hidden_size), + nbytes=nbytes["target_hidden_states"], + ) + target_last_hidden_states = self._read_bfloat16_tensor_from_shard( + shard_mmap=shard_mmap, + offset=record["target_last_hidden_states_offset"], + shape=(seq_len, self.hidden_size), + nbytes=nbytes["target_last_hidden_states"], + ) + elif self.hidden_dtype == "float8_e4m3fn": + target_hidden_states = self._read_float8_tensor_from_shard( + shard_mmap=shard_mmap, + offset=record["target_hidden_states_offset"], + shape=(seq_len, self.num_target_layers * self.hidden_size), + nbytes=nbytes["target_hidden_states"], + torch_dtype=torch.float8_e4m3fn, + ).to(dtype=torch.bfloat16) + target_last_hidden_states = self._read_float8_tensor_from_shard( + shard_mmap=shard_mmap, + offset=record["target_last_hidden_states_offset"], + shape=(seq_len, self.hidden_size), + nbytes=nbytes["target_last_hidden_states"], + torch_dtype=torch.float8_e4m3fn, + ).to(dtype=torch.bfloat16) + elif self.hidden_dtype == "float8_e5m2": + target_hidden_states = self._read_float8_tensor_from_shard( + shard_mmap=shard_mmap, + offset=record["target_hidden_states_offset"], + shape=(seq_len, self.num_target_layers * self.hidden_size), + nbytes=nbytes["target_hidden_states"], + torch_dtype=torch.float8_e5m2, + ).to(dtype=torch.bfloat16) + target_last_hidden_states = self._read_float8_tensor_from_shard( + shard_mmap=shard_mmap, + offset=record["target_last_hidden_states_offset"], + shape=(seq_len, self.hidden_size), + nbytes=nbytes["target_last_hidden_states"], + torch_dtype=torch.float8_e5m2, + ).to(dtype=torch.bfloat16) + else: + raise ValueError(f"Unsupported hidden_dtype: {self.hidden_dtype}") return { "input_ids": input_ids, "loss_mask": loss_mask, diff --git a/scripts/data/prepare_target_cache.py b/scripts/data/prepare_target_cache.py index d072582..5907e52 100644 --- a/scripts/data/prepare_target_cache.py +++ b/scripts/data/prepare_target_cache.py @@ -149,6 +149,13 @@ def parse_args(): parser.add_argument("--max-shard-bytes", type=int, default=64 * 1024**3) parser.add_argument("--local-batch-size", type=int, default=32) parser.add_argument("--num-workers", type=int, default=4) + parser.add_argument( + "--hidden-dtype", + type=str, + default="bfloat16", + choices=["bfloat16", "float8_e4m3fn", "float8_e5m2"], + help="Data type for saving target cached hidden states." + ) cli_args = parser.parse_args() config = parse_opts_to_config(cli_args.opts, load_config(cli_args.config)) return cli_args, config @@ -163,6 +170,7 @@ def _write_manifest( hidden_size: int, min_loss_tokens: int, shards, + hidden_dtype: str = "bfloat16", ): num_samples = sum( int( @@ -177,6 +185,7 @@ def _write_manifest( shards=shards, target_layer_ids=target_layer_ids, hidden_size=hidden_size, + hidden_dtype=hidden_dtype, extra_fields={ "target_model_name_or_path": str(config.model.target_model_name_or_path), "source_jsonl_paths": [str(path) for path in train_data_paths], @@ -275,6 +284,7 @@ def main(local_rank: int): rank_dir=rank_dir, max_shard_bytes=int(cli_args.max_shard_bytes), max_queue_size=int(cli_args.local_batch_size) * 4, + hidden_dtype=cli_args.hidden_dtype, ) processed_local_samples = 0 @@ -385,6 +395,7 @@ def main(local_rank: int): hidden_size=target_hidden_size, min_loss_tokens=min_loss_tokens, shards=shards, + hidden_dtype=cli_args.hidden_dtype, ) cleanup_target_cache_tmp_dir(output_dir) print_on_global_main( diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..6e99e26 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +# Empty init file to mark tests directory as a python package diff --git a/tests/test_quantization.py b/tests/test_quantization.py new file mode 100644 index 0000000..1487dab --- /dev/null +++ b/tests/test_quantization.py @@ -0,0 +1,327 @@ +import os +import sys +import json +import shutil +import tempfile +import unittest +from unittest.mock import patch, MagicMock +import torch +import numpy as np + +from deepspec.data.target_cache_dataset import ( + AsyncTargetCacheWriter, + CacheDataset, + CacheCollator, + build_target_cache_manifest, + write_target_cache_manifest, + finalize_target_cache_index, + LocalCacheWriteSummary, +) + +# Mock classes for integration test +class MockConfigObj: + def __init__(self): + self.model_type = "qwen3" + self.hidden_size = 64 + +class MockTokenizer: + def __init__(self): + self.eos_token_id = 0 + self.pad_token_id = 0 + + def apply_chat_template(self, messages, **kwargs): + return "<|im_start|>user\nhello<|im_end|>\n<|im_start|>assistant\nworld<|im_end|>\n" + + def __call__(self, text, **kwargs): + class MockEncoding: + def __init__(self, seq_len): + self.input_ids = torch.randint(0, 1000, (1, seq_len), dtype=torch.int32) + self.attention_mask = torch.ones(1, seq_len, dtype=torch.uint8) + return MockEncoding(128) + + def encode(self, text, **kwargs): + length = max(1, len(text) // 2) + return list(range(length)) + +class MockModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = MockConfigObj() + self.embed_tokens = torch.nn.Embedding(1000, 64) + self.layers = torch.nn.ModuleList([torch.nn.Linear(64, 64) for _ in range(2)]) + + def forward(self, input_ids, attention_mask=None, **kwargs): + x = self.embed_tokens(input_ids) + for layer in self.layers: + x = layer(x) + + class MockOutput: + def __init__(self, last_hidden): + self.last_hidden_state = last_hidden + return MockOutput(x.to(dtype=torch.bfloat16)) + + +class TestTargetCacheQuantization(unittest.TestCase): + + def setUp(self): + self.tmp_dir = tempfile.mkdtemp() + self.seq_len = 128 + self.hidden_size = 64 + self.num_layers = 2 + + # Generate some mock target data + torch.manual_seed(42) + self.mock_hidden_states = torch.randn(self.seq_len, self.num_layers * self.hidden_size) * 0.5 + self.mock_last_hidden_states = torch.randn(self.seq_len, self.hidden_size) * 0.5 + + self.mock_input_ids = torch.randint(0, 1000, (self.seq_len,), dtype=torch.int32) + self.mock_attention_mask = torch.ones(self.seq_len, dtype=torch.uint8) + self.mock_loss_mask = torch.ones(self.seq_len, dtype=torch.uint8) + + def tearDown(self): + shutil.rmtree(self.tmp_dir) + + def _write_mock_cache(self, output_dir: str, dtype: str): + os.makedirs(output_dir, exist_ok=True) + rank_dir = os.path.join(output_dir, "_tmp", "rank_0") + os.makedirs(rank_dir, exist_ok=True) + + writer = AsyncTargetCacheWriter( + rank_dir=rank_dir, + max_shard_bytes=10 * 1024 * 1024, # 10MB + hidden_dtype=dtype, + ) + writer.write_sample( + input_ids=self.mock_input_ids, + attention_mask=self.mock_attention_mask, + loss_mask=self.mock_loss_mask, + target_hidden_states=self.mock_hidden_states, + target_last_hidden_states=self.mock_last_hidden_states, + ) + writer.close() + + summary = LocalCacheWriteSummary( + global_rank=0, + source_sample_start=0, + source_sample_end=1, + num_local_samples=1, + num_local_shards=1, + local_shard_files=list(writer.local_shard_files), + ) + + src_shard = os.path.join(rank_dir, writer.local_shard_files[0]) + dst_shard = os.path.join(output_dir, "shard-00000.bin") + shutil.copy(src_shard, dst_shard) + + src_idx = os.path.join(rank_dir, "samples.local.idx") + dst_idx = os.path.join(output_dir, "samples.idx") + shutil.copy(src_idx, dst_idx) + + shards_metadata = [{"shard_id": 0, "file_name": "shard-00000.bin"}] + manifest = build_target_cache_manifest( + num_samples=1, + shards=shards_metadata, + target_layer_ids=[0, 1], + hidden_size=self.hidden_size, + hidden_dtype=dtype, + extra_fields={"target_model_name_or_path": "mock-model"}, + ) + write_target_cache_manifest(output_dir=output_dir, manifest=manifest) + shutil.rmtree(os.path.join(output_dir, "_tmp")) + + def test_bfloat16_correctness(self): + cache_path = os.path.join(self.tmp_dir, "bf16_cache") + self._write_mock_cache(cache_path, "bfloat16") + + # Verify sizes + shard_path = os.path.join(cache_path, "shard-00000.bin") + bf16_size = os.path.getsize(shard_path) + + expected_size = ( + self.seq_len * 4 + + self.seq_len * 1 + + self.seq_len * 1 + + self.seq_len * self.num_layers * self.hidden_size * 2 + + self.seq_len * self.hidden_size * 2 + ) + self.assertEqual(bf16_size, expected_size) + + # Read back and compare + dataset = CacheDataset(cache_path) + sample = dataset[0] + self.assertEqual(sample["target_hidden_states"].dtype, torch.bfloat16) + self.assertEqual(sample["target_last_hidden_states"].dtype, torch.bfloat16) + + self.assertTrue(torch.allclose(sample["target_hidden_states"], self.mock_hidden_states.to(torch.bfloat16))) + self.assertTrue(torch.allclose(sample["target_last_hidden_states"], self.mock_last_hidden_states.to(torch.bfloat16))) + dataset.close() + + def test_float8_e4m3fn_quantization(self): + cache_path = os.path.join(self.tmp_dir, "fp8_e4m3fn_cache") + self._write_mock_cache(cache_path, "float8_e4m3fn") + + # Verify sizes (should be reduced) + shard_path = os.path.join(cache_path, "shard-00000.bin") + fp8_size = os.path.getsize(shard_path) + + expected_size = ( + self.seq_len * 4 + + self.seq_len * 1 + + self.seq_len * 1 + + self.seq_len * self.num_layers * self.hidden_size * 1 + + self.seq_len * self.hidden_size * 1 + ) + self.assertEqual(fp8_size, expected_size) + + # Read back and compare + dataset = CacheDataset(cache_path) + sample = dataset[0] + self.assertEqual(sample["target_hidden_states"].dtype, torch.bfloat16) + self.assertEqual(sample["target_last_hidden_states"].dtype, torch.bfloat16) + + diff_states = (sample["target_hidden_states"] - self.mock_hidden_states.to(torch.bfloat16)).abs().mean().item() + diff_last = (sample["target_last_hidden_states"] - self.mock_last_hidden_states.to(torch.bfloat16)).abs().mean().item() + + print(f"FP8 E4M3FN target_hidden_states MAE: {diff_states}") + print(f"FP8 E4M3FN target_last_hidden_states MAE: {diff_last}") + + self.assertLess(diff_states, 0.05) + self.assertLess(diff_last, 0.05) + dataset.close() + + def test_float8_e5m2_quantization(self): + cache_path = os.path.join(self.tmp_dir, "fp8_e5m2_cache") + self._write_mock_cache(cache_path, "float8_e5m2") + + # Verify sizes (should be reduced) + shard_path = os.path.join(cache_path, "shard-00000.bin") + fp8_size = os.path.getsize(shard_path) + + expected_size = ( + self.seq_len * 4 + + self.seq_len * 1 + + self.seq_len * 1 + + self.seq_len * self.num_layers * self.hidden_size * 1 + + self.seq_len * self.hidden_size * 1 + ) + self.assertEqual(fp8_size, expected_size) + + # Read back and compare + dataset = CacheDataset(cache_path) + sample = dataset[0] + self.assertEqual(sample["target_hidden_states"].dtype, torch.bfloat16) + self.assertEqual(sample["target_last_hidden_states"].dtype, torch.bfloat16) + + diff_states = (sample["target_hidden_states"] - self.mock_hidden_states.to(torch.bfloat16)).abs().mean().item() + diff_last = (sample["target_last_hidden_states"] - self.mock_last_hidden_states.to(torch.bfloat16)).abs().mean().item() + + print(f"FP8 E5M2 target_hidden_states MAE: {diff_states}") + print(f"FP8 E5M2 target_last_hidden_states MAE: {diff_last}") + + self.assertLess(diff_states, 0.1) + self.assertLess(diff_last, 0.1) + dataset.close() + + def test_cache_collator_and_dataset(self): + cache_path = os.path.join(self.tmp_dir, "collator_cache") + self._write_mock_cache(cache_path, "float8_e4m3fn") + + dataset = CacheDataset(cache_path) + collator = CacheCollator() + + # Test collating a single sample list + batch = collator([dataset[0], dataset[0]]) + + self.assertIn("input_ids", batch) + self.assertIn("loss_mask", batch) + self.assertIn("attention_mask", batch) + self.assertIn("target_hidden_states", batch) + self.assertIn("target_last_hidden_states", batch) + + # Verify shapes + self.assertEqual(batch["input_ids"].shape, (2, self.seq_len)) + self.assertEqual(batch["target_hidden_states"].shape, (2, self.seq_len, self.num_layers * self.hidden_size)) + self.assertEqual(batch["target_last_hidden_states"].shape, (2, self.seq_len, self.hidden_size)) + + # Verify datatypes after collation + self.assertEqual(batch["target_hidden_states"].dtype, torch.bfloat16) + self.assertEqual(batch["target_last_hidden_states"].dtype, torch.bfloat16) + dataset.close() + + @patch("scripts.data.prepare_target_cache.dist.destroy_process_group") + @patch("scripts.data.prepare_target_cache.dist.barrier") + @patch("scripts.data.prepare_target_cache.dist.broadcast_object_list") + @patch("scripts.data.prepare_target_cache.dist.get_world_size", return_value=1) + @patch("scripts.data.prepare_target_cache.dist.get_rank", return_value=0) + @patch("deepspec.utils.distributed.is_global_main_process", return_value=True) + @patch("deepspec.utils.distributed.is_local_main_process", return_value=True) + @patch("scripts.data.prepare_target_cache.init_dist") + @patch("scripts.data.prepare_target_cache.AutoTokenizer.from_pretrained") + @patch("scripts.data.prepare_target_cache.AutoModel.from_pretrained") + def test_prepare_target_cache_script_integration(self, mock_from_pretrained_model, mock_from_pretrained_tokenizer, mock_init_dist, mock_is_local, mock_is_global, mock_get_rank, mock_get_world_size, mock_broadcast, mock_barrier, mock_destroy): + # Configure mocks + mock_init_dist.return_value = (torch.device("cpu"), 0, 1) + mock_from_pretrained_tokenizer.return_value = MockTokenizer() + mock_from_pretrained_model.return_value = MockModel() + + # Create temporary mock config and dataset files + config_path = os.path.join(self.tmp_dir, "mock_config.py") + with open(config_path, "w", encoding="utf-8") as f: + f.write(""" +import os +model = dict( + target_model_name_or_path="mock-model", + target_layer_ids=[0, 1], +) +seed = 42 +data = dict( + chat_template="qwen", + max_length=128, +) +""") + + dataset_path = os.path.join(self.tmp_dir, "mock_dataset.jsonl") + with open(dataset_path, "w", encoding="utf-8") as f: + f.write(json.dumps({"conversations": [{"role": "user", "content": "hello"}, {"role": "assistant", "content": "world"}]}) + "\n") + + output_dir = os.path.join(self.tmp_dir, "integration_output_cache") + + # Mock argv parameters (set min-loss-tokens to 0) + test_argv = [ + "prepare_target_cache.py", + "--config", config_path, + "--train-data-path", dataset_path, + "--output-dir", output_dir, + "--hidden-dtype", "float8_e4m3fn", + "--local-batch-size", "1", + "--min-loss-tokens", "0", + ] + + from scripts.data.prepare_target_cache import main as prepare_cache_main + + with patch.object(sys, "argv", test_argv): + # Run the preparation script's main + prepare_cache_main(0) + + # Verify output exists and is readable + self.assertTrue(os.path.exists(output_dir)) + self.assertTrue(os.path.exists(os.path.join(output_dir, "manifest.json"))) + self.assertTrue(os.path.exists(os.path.join(output_dir, "samples.idx"))) + self.assertTrue(os.path.exists(os.path.join(output_dir, "shard-00000.bin"))) + + # Check manifest contents + with open(os.path.join(output_dir, "manifest.json"), "r") as f: + manifest = json.load(f) + self.assertEqual(manifest["hidden_dtype"], "float8_e4m3fn") + self.assertEqual(manifest["num_samples"], 1) + + # Load using dataset and verify it works + dataset = CacheDataset(output_dir) + self.assertEqual(len(dataset), 1) + sample = dataset[0] + self.assertEqual(sample["target_hidden_states"].shape, (128, 2 * 64)) + self.assertEqual(sample["target_hidden_states"].dtype, torch.bfloat16) + dataset.close() + +if __name__ == '__main__': + unittest.main()