""" 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()