scripts/streaming_cast.py

"""
streaming_cast.py
-----------------
Downloads DeepSeek-Math-V2 FP8 shards ONE AT A TIME, casts to BF16,
and writes each shard to a local directory (the HF Storage Bucket mounted
at /data). Deletes the FP8 shard immediately after casting.

Peak local ephemeral disk: ~0 GB (everything goes to /data).
Peak bucket disk per shard: ~13 GB (4 GB FP8 + 9 GB BF16 simultaneously),
dropping to ~9 GB once FP8 is deleted.

Usage:
    python streaming_cast.py \
        --src-repo  deepseek-ai/DeepSeek-Math-V2 \
        --dst-repo  memmywinks/DeepSeek_Math_V2 \
        --bf16-dir  /data/bf16 \
        --work-dir  /data/scratch \
        --resume
"""

import argparse
import gc
import json
import os
import shutil
import time
from pathlib import Path

import torch
from huggingface_hub import HfApi, hf_hub_download
from safetensors import safe_open
from safetensors.torch import save_file


# ── FP8 dequantisation ────────────────────────────────────────────────────────

def dequantize_fp8_tile(weight_fp8: torch.Tensor,
                         scale_inv: torch.Tensor) -> torch.Tensor:
    rows, cols = weight_fp8.shape
    out = torch.empty((rows, cols), dtype=torch.bfloat16)
    for row_start in range(0, rows, 128):
        row_end  = min(row_start + 128, rows)
        tile_row = row_start // 128
        strip    = weight_fp8[row_start:row_end, :]
        scale    = (
            scale_inv[tile_row:tile_row+1, :]
            .repeat_interleave(128, dim=1)[:, :cols]
            .expand(row_end - row_start, -1)
        )
        out[row_start:row_end, :] = strip.to(torch.float32).mul_(scale).to(torch.bfloat16)
    return out


def cast_shard(src_path: Path, dst_path: Path) -> None:
    with safe_open(str(src_path), framework="pt", device="cpu") as f:
        keys = list(f.keys())
        try:
            metadata = f.metadata() or {}
        except Exception:
            metadata = {}

    all_keys = set(keys)
    fp8_keys = {k for k in keys if k.endswith(".weight") and f"{k}_scale_inv" in all_keys}
    scale_keys = {f"{k}_scale_inv" for k in fp8_keys}

    tensors_out = {}
    with safe_open(str(src_path), framework="pt", device="cpu") as f:
        for key in keys:
            if key in scale_keys:
                continue
            tensor = f.get_tensor(key)
            if key in fp8_keys and tensor.dtype == torch.float8_e4m3fn:
                scale_inv = f.get_tensor(f"{key}_scale_inv").to(torch.float32)
                tensor = dequantize_fp8_tile(tensor, scale_inv)
                del scale_inv
            tensors_out[key] = tensor

    save_file(tensors_out, str(dst_path), metadata=metadata)
    del tensors_out
    gc.collect()


# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--src-repo",  default="deepseek-ai/DeepSeek-Math-V2")
    parser.add_argument("--dst-repo",  required=True,
                        help="HF dataset repo (used only to fetch index + upload final index/configs)")
    parser.add_argument("--bf16-dir",  default="/data/bf16",
                        help="Local directory to write BF16 shards (should be on bucket mount)")
    parser.add_argument("--work-dir",  default="/data/scratch")
    parser.add_argument("--resume",    action="store_true",
                        help="Skip shards already present in --bf16-dir")
    args = parser.parse_args()

    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    api   = HfApi(token=token)

    bf16_dir = Path(args.bf16_dir)
    work_dir = Path(args.work_dir)
    fp8_dir  = work_dir / "shard_fp8"
    bf16_dir.mkdir(parents=True, exist_ok=True)
    fp8_dir.mkdir(parents=True, exist_ok=True)

    # ── Fetch index ───────────────────────────────────────────────────────────
    print(f"Fetching model index from {args.src_repo}...")
    index_path = hf_hub_download(
        repo_id=args.src_repo,
        filename="model.safetensors.index.json",
        token=token,
    )
    with open(index_path) as f:
        index = json.load(f)

    shard_names = sorted(set(index["weight_map"].values()))
    print(f"Found {len(shard_names)} shards.")
    print(f"BF16 output dir: {bf16_dir}")

    # ── Resume: skip shards already on disk ───────────────────────────────────
    if args.resume:
        already_done = {s for s in shard_names if (bf16_dir / s).exists()}
        if already_done:
            print(f"Resume: {len(already_done)}/{len(shard_names)} shards already on disk, skipping.")
        shard_names = [s for s in shard_names if s not in already_done]

    if not shard_names:
        print("All shards already cast. Nothing to do.")
    else:
        total = len(shard_names)
        for i, shard_name in enumerate(shard_names, 1):
            print(f"\n[{i}/{total}] {shard_name}")
            fp8_path  = fp8_dir  / shard_name
            bf16_path = bf16_dir / shard_name

            # Download FP8 shard
            t0 = time.time()
            print(f"  Downloading to {fp8_path}...")
            hf_hub_download(
                repo_id=args.src_repo,
                filename=shard_name,
                local_dir=str(fp8_dir),
                token=token,
            )
            print(f"  Downloaded in {time.time()-t0:.0f}s  ({fp8_path.stat().st_size/1e9:.1f} GB)")

            # Cast FP8 → BF16, write directly to bucket
            t0 = time.time()
            print(f"  Casting → {bf16_path}...")
            cast_shard(fp8_path, bf16_path)
            print(f"  Cast in {time.time()-t0:.0f}s  ({bf16_path.stat().st_size/1e9:.1f} GB)")

            # Delete FP8 immediately to free bucket space
            fp8_path.unlink(missing_ok=True)
            data_free = shutil.disk_usage(str(bf16_dir)).free / 1e9
            print(f"  FP8 deleted. Bucket free: {data_free:.0f} GB")

    # ── Write BF16 index + copy config files ──────────────────────────────────
    print("\nWriting model index and config files to BF16 dir...")
    new_weight_map = {
        k: v for k, v in index["weight_map"].items()
        if not k.endswith("_scale_inv")
    }
    new_index = {"metadata": index.get("metadata", {}), "weight_map": new_weight_map}
    with open(bf16_dir / "model.safetensors.index.json", "w") as f:
        json.dump(new_index, f, indent=2)

    for fname in ["config.json", "generation_config.json", "tokenizer.json",
                  "tokenizer_config.json", "special_tokens_map.json"]:
        try:
            local = hf_hub_download(repo_id=args.src_repo, filename=fname, token=token)
            shutil.copy2(local, bf16_dir / fname)
            print(f"  Copied {fname}")
        except Exception as e:
            print(f"  Skipping {fname}: {e}")

    print(f"\n✔ All shards cast. BF16 model ready at: {bf16_dir}")
    print("Next: convert_hf_to_gguf.py will read from this directory.")


if __name__ == "__main__":
    main()