deepspec/scripts/data/convert_engine_dump.py

"""Convert a Lux engine target-hidden-state dump into a DeepSpec target cache.

The Rust engine (separate workstream) dumps, per training sample:

  * ``input_ids``                 [seq]            int64/int32
  * ``target_hidden_states``      [seq, L, H]      bf16   (L captured layers)
  * ``target_last_hidden_states`` [seq, H]         bf16
  * (optional) ``loss_mask``      [seq]            any int/bool
  * (optional) ``attention_mask`` [seq]            any int/bool

alongside an ``index.jsonl`` (one JSON object per line) that points at each
per-sample ``*.safetensors`` shard via a ``file`` / ``path`` / ``safetensors``
key.

This script re-emits those samples in DeepSpec's canonical *target cache*
protocol (version 2) by driving DeepSpec's OWN writer/finalizer classes -- it
does not reimplement the binary layout. The output directory is directly
consumable by ``deepspec.data.CacheDataset`` and ``scripts/train/train.sh``.

``target_hidden_states`` is stored as ``[seq, L * H]`` with the L captured
layers concatenated along the hidden axis in ``--target-layer-ids`` order --
byte-identical to what ``run_target_forward_with_hooks`` in
``scripts/data/prepare_target_cache.py`` produces
(``torch.cat([captured[id] for id in target_layer_ids], dim=-1)``).

Single process, no torch.distributed: this is the ``world_size == 1``
specialization of ``scripts/data/prepare_target_cache.py``'s finalize flow.
"""

import argparse
import json
import os

import torch
from safetensors.torch import load_file

from deepspec.data.target_cache_dataset import (
    AsyncTargetCacheWriter,
    LocalCacheWriteSummary,
    atomic_json_dump,
    build_global_target_cache_shard_map,
    build_target_cache_manifest,
    cleanup_target_cache_tmp_dir,
    finalize_target_cache_index,
    load_local_cache_write_summary,
    load_target_cache_manifest,
    prepare_target_cache_output_dir,
    rename_local_target_cache_shards,
    write_target_cache_manifest,
)


_FILE_KEYS = ("file", "path", "safetensors", "sample", "shard")
_INPUT_IDS_KEYS = ("input_ids", "input_id", "ids", "tokens")
_HIDDEN_KEYS = ("target_hidden_states", "hidden_states", "target_hidden")
_LAST_HIDDEN_KEYS = (
    "target_last_hidden_states",
    "last_hidden_states",
    "target_last_hidden",
)
_LOSS_MASK_KEYS = ("loss_mask", "loss_masks")


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--dump-dir",
        required=True,
        help="Engine-dump directory containing index.jsonl + *.safetensors.",
    )
    parser.add_argument(
        "--output-dir",
        required=True,
        help="Destination target-cache directory (must be new/empty).",
    )
    parser.add_argument(
        "--index-name",
        default="index.jsonl",
        help="Index file name inside --dump-dir (default: index.jsonl).",
    )
    parser.add_argument(
        "--hidden-size",
        type=int,
        default=4096,
        help="Target hidden dimension H (V4-Flash: 4096).",
    )
    parser.add_argument(
        "--target-layer-ids",
        default="1,11,21,31,41",
        help="Comma-separated captured target layer ids, ascending "
        "(V4-Flash: 1,11,21,31,41).",
    )
    parser.add_argument(
        "--target-model-name-or-path",
        required=True,
        help="Recorded in the manifest; MUST equal "
        "config.model.target_model_name_or_path used at train time "
        "(validate_train_cache asserts string equality).",
    )
    parser.add_argument(
        "--loss-mask-mode",
        choices=("ones_except_first", "all_ones", "from_dump"),
        default="ones_except_first",
        help="How to build loss_mask. 'ones_except_first' (default): 1 "
        "everywhere except position 0 -- position 0 has no left context so it "
        "is never a valid anchor nor a first draft target "
        "(see build_anchor_candidate_mask). 'from_dump': use the dump's "
        "loss_mask (error if a sample lacks one).",
    )
    parser.add_argument(
        "--max-shard-bytes",
        type=int,
        default=64 * 1024**3,
        help="Max bytes per output shard before rolling to a new one.",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="Only convert the first N index entries (debug).",
    )
    return parser.parse_args()


def _first_present(mapping, keys, what, where):
    for key in keys:
        if key in mapping:
            return mapping[key]
    raise KeyError(f"{what} not found in {where}; looked for keys {keys}.")


def _resolve_sample_path(dump_dir, record, line_no):
    file_ref = None
    for key in _FILE_KEYS:
        if key in record:
            file_ref = record[key]
            break
    if file_ref is None:
        raise KeyError(
            f"index line {line_no}: no shard reference; expected one of "
            f"{_FILE_KEYS} in {record!r}."
        )
    path = str(file_ref)
    if not os.path.isabs(path):
        path = os.path.join(dump_dir, path)
    if not os.path.exists(path):
        raise FileNotFoundError(f"index line {line_no}: missing sample file {path}.")
    return path


def _build_loss_mask(mode, seq_len, tensors, device):
    if mode == "from_dump":
        loss_mask = _first_present(
            tensors, _LOSS_MASK_KEYS, "loss_mask", "sample tensors"
        )
        return loss_mask.reshape(-1).to(dtype=torch.uint8)
    loss_mask = torch.ones(seq_len, dtype=torch.uint8, device=device)
    if mode == "ones_except_first" and seq_len > 0:
        loss_mask[0] = 0
    return loss_mask


def _load_sample(path, *, hidden_size, num_target_layers, loss_mask_mode):
    tensors = load_file(path)
    input_ids = _first_present(tensors, _INPUT_IDS_KEYS, "input_ids", path).reshape(-1)
    seq_len = int(input_ids.shape[0])
    assert seq_len > 0, f"{path}: empty input_ids."

    hidden = _first_present(tensors, _HIDDEN_KEYS, "target_hidden_states", path)
    expected_hidden_numel = seq_len * num_target_layers * hidden_size
    assert hidden.numel() == expected_hidden_numel, (
        f"{path}: target_hidden_states has {hidden.numel()} elements, expected "
        f"{expected_hidden_numel} = seq({seq_len}) * layers({num_target_layers}) "
        f"* hidden({hidden_size}). Got shape {tuple(hidden.shape)}."
    )
    # [seq, L, H] or already [seq, L*H] -> [seq, L*H] (row-major keeps
    # per-layer blocks contiguous, matching target_layer_ids concat order).
    target_hidden_states = hidden.reshape(seq_len, num_target_layers * hidden_size)

    last_hidden = _first_present(
        tensors, _LAST_HIDDEN_KEYS, "target_last_hidden_states", path
    )
    assert last_hidden.numel() == seq_len * hidden_size, (
        f"{path}: target_last_hidden_states has {last_hidden.numel()} elements, "
        f"expected {seq_len * hidden_size} = seq({seq_len}) * hidden({hidden_size})."
    )
    target_last_hidden_states = last_hidden.reshape(seq_len, hidden_size)

    loss_mask = _build_loss_mask(loss_mask_mode, seq_len, tensors, input_ids.device)
    assert loss_mask.shape[0] == seq_len, (
        f"{path}: loss_mask length {loss_mask.shape[0]} != seq_len {seq_len}."
    )
    attention_mask = torch.ones(seq_len, dtype=torch.uint8, device=input_ids.device)

    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "loss_mask": loss_mask,
        "target_hidden_states": target_hidden_states,
        "target_last_hidden_states": target_last_hidden_states,
    }


def _read_index(index_path, limit):
    records = []
    with open(index_path, "r", encoding="utf-8") as handle:
        for line_no, raw in enumerate(handle):
            raw = raw.strip()
            if not raw:
                continue
            records.append((line_no, json.loads(raw)))
            if limit is not None and len(records) >= limit:
                break
    assert records, f"No records read from {index_path}."
    return records


def main():
    args = parse_args()
    dump_dir = os.path.abspath(args.dump_dir)
    output_dir = os.path.abspath(args.output_dir)
    index_path = os.path.join(dump_dir, args.index_name)
    assert os.path.exists(index_path), f"Missing index file: {index_path}"

    target_layer_ids = [int(x) for x in args.target_layer_ids.split(",") if x != ""]
    assert target_layer_ids == sorted(target_layer_ids), (
        f"--target-layer-ids must be ascending, got {target_layer_ids}."
    )
    num_target_layers = len(target_layer_ids)
    hidden_size = int(args.hidden_size)

    records = _read_index(index_path, args.limit)
    print(
        f"Converting {len(records)} samples from {dump_dir}\n"
        f"  hidden_size={hidden_size} target_layer_ids={target_layer_ids} "
        f"loss_mask_mode={args.loss_mask_mode}\n"
        f"  -> {output_dir}",
        flush=True,
    )

    prepare_target_cache_output_dir(output_dir)
    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=int(args.max_shard_bytes),
    )
    try:
        for idx, (line_no, record) in enumerate(records):
            path = _resolve_sample_path(dump_dir, record, line_no)
            sample = _load_sample(
                path,
                hidden_size=hidden_size,
                num_target_layers=num_target_layers,
                loss_mask_mode=args.loss_mask_mode,
            )
            writer.write_sample(**sample)
            if (idx + 1) % 100 == 0 or (idx + 1) == len(records):
                print(f"  wrote {idx + 1}/{len(records)} samples", flush=True)
    finally:
        writer.close()

    summary = LocalCacheWriteSummary(
        global_rank=0,
        source_sample_start=0,
        source_sample_end=len(records),
        num_local_samples=writer.num_local_samples,
        num_local_shards=len(writer.local_shard_files),
        local_shard_files=list(writer.local_shard_files),
    )
    atomic_json_dump(summary.to_json(), os.path.join(rank_dir, "summary.json"))

    summaries = [load_local_cache_write_summary(rank_dir)]
    shard_map, shards = build_global_target_cache_shard_map(summaries)
    rename_local_target_cache_shards(
        output_dir=output_dir,
        rank_dir=rank_dir,
        summary=summaries[0],
        shard_map=shard_map,
    )
    num_valid_samples = finalize_target_cache_index(
        output_dir=output_dir,
        summaries=summaries,
        shard_map=shard_map,
    )
    manifest = build_target_cache_manifest(
        num_samples=num_valid_samples,
        shards=shards,
        target_layer_ids=target_layer_ids,
        hidden_size=hidden_size,
        extra_fields={
            "target_model_name_or_path": str(args.target_model_name_or_path),
            "source_engine_dump": dump_dir,
            "loss_mask_mode": args.loss_mask_mode,
            "converter": "scripts/data/convert_engine_dump.py",
        },
    )
    write_target_cache_manifest(output_dir=output_dir, manifest=manifest)
    cleanup_target_cache_tmp_dir(output_dir)

    # Re-load through DeepSpec's own validator to prove the cache is valid.
    load_target_cache_manifest(output_dir)
    print(
        f"Wrote valid target cache: {num_valid_samples} samples, "
        f"{len(shards)} shard(s) at {output_dir}",
        flush=True,
    )


if __name__ == "__main__":
    main()