""" fp8_cast_bf16_fallback.py ------------------------- Standalone FP8 E4M3 → BF16 dequantisation for DeepSeek-V3/Math-V2 style safetensors checkpoints. Processes ONE SHARD AT A TIME with explicit memory cleanup so that peak RAM is bounded by the size of the largest single shard (~5-15 GB) rather than the full model (~800 GB). Safe to run on standard HF Spaces hardware. Usage: python fp8_cast_bf16_fallback.py \ --input-fp8-hf-path /workspace/model-fp8 \ --output-bf16-hf-path /workspace/model-bf16 # Resume a partial run (already-converted shards are skipped): python fp8_cast_bf16_fallback.py \ --input-fp8-hf-path /workspace/model-fp8 \ --output-bf16-hf-path /workspace/model-bf16 \ --resume """ import argparse import gc import json import os import shutil from pathlib import Path import torch from safetensors import safe_open from safetensors.torch import save_file from tqdm import tqdm # ── Tile dequantisation ────────────────────────────────────────────────────── def dequantize_fp8_tile(weight_fp8: torch.Tensor, scale_inv: torch.Tensor) -> torch.Tensor: """ weight_fp8 : (rows, cols) in float8_e4m3fn scale_inv : (ceil(rows/128), ceil(cols/128)) in float32 Returns : (rows, cols) in bfloat16 Processes in 128-row strips to avoid materialising a full (rows, cols) float32 intermediate — halves transient memory usage for large tensors. """ 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 # index into scale_inv rows strip_fp8 = weight_fp8[row_start:row_end, :] # (≤128, cols) # Build per-element scale for this strip scale_strip = ( scale_inv[tile_row : tile_row + 1, :] # (1, tile_cols) .repeat_interleave(128, dim=1)[:, :cols] # (1, cols) .expand(row_end - row_start, -1) # (strip_rows, cols) ) out[row_start:row_end, :] = ( strip_fp8.to(torch.float32).mul_(scale_strip).to(torch.bfloat16) ) return out # ── Shard conversion ───────────────────────────────────────────────────────── def convert_shard(shard_path: Path, out_path: Path, fp8_key_names: set[str] | None = None) -> dict[str, str]: """ Convert a single safetensors shard in-place (one tensor at a time). Returns a dict mapping tensor_name → shard filename for index rebuilding. """ # ── First pass: inventory the shard ────────────────────────────────────── key_list: list[str] = [] metadata: dict = {} with safe_open(str(shard_path), framework="pt", device="cpu") as f: key_list = list(f.keys()) try: metadata = f.metadata() or {} except Exception: metadata = {} # Detect FP8 keys in *this shard* (a key is FP8 if its _scale_inv is present) all_keys_set = set(key_list) fp8_keys_in_shard = { k for k in key_list if k.endswith(".weight") and f"{k}_scale_inv" in all_keys_set # Also accept an external hint (fp8_key_names) for cross-shard scales or (fp8_key_names and k in fp8_key_names) } scale_keys = {f"{k}_scale_inv" for k in fp8_keys_in_shard} # ── Second pass: convert tensor by tensor ───────────────────────────────── tensors_out: dict[str, torch.Tensor] = {} with safe_open(str(shard_path), framework="pt", device="cpu") as f: for key in key_list: if key in scale_keys: # Drop scale tensors from output continue tensor = f.get_tensor(key) if key in fp8_keys_in_shard and tensor.dtype == torch.float8_e4m3fn: # Load companion scale and dequantise 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(out_path), metadata=metadata) weight_map = {k: shard_path.name for k in tensors_out} # Explicit cleanup del tensors_out gc.collect() return weight_map # ── Directory conversion ───────────────────────────────────────────────────── def convert_directory(fp8_dir: Path, bf16_dir: Path, resume: bool = False) -> None: bf16_dir.mkdir(parents=True, exist_ok=True) # Collect shards index_path = fp8_dir / "model.safetensors.index.json" if index_path.exists(): with open(index_path) as f: index = json.load(f) shard_names = sorted(set(index["weight_map"].values())) shard_paths = [fp8_dir / s for s in shard_names] original_weight_map: dict[str, str] = index.get("weight_map", {}) else: single = fp8_dir / "model.safetensors" if not single.exists(): raise FileNotFoundError(f"No safetensors found in {fp8_dir}") shard_paths = [single] original_weight_map = {} print(f"Found {len(shard_paths)} shard(s).") print(f"Peak RAM per shard: ~{max(p.stat().st_size for p in shard_paths) * 2 / 1e9:.1f} GB " f"(2× largest shard on disk)") combined_weight_map: dict[str, str] = {} for shard_path in tqdm(shard_paths, desc="Converting shards", unit="shard"): out_path = bf16_dir / shard_path.name if resume and out_path.exists() and out_path.stat().st_size > 0: tqdm.write(f" ↷ Skipping {shard_path.name} (already converted)") # Still need to rebuild the weight map entry for the index with safe_open(str(out_path), framework="pt", device="cpu") as f: for key in f.keys(): combined_weight_map[key] = shard_path.name continue tqdm.write(f" → {shard_path.name} ({shard_path.stat().st_size / 1e9:.1f} GB on disk)") shard_map = convert_shard(shard_path, out_path) combined_weight_map.update(shard_map) tqdm.write(f" ✔ written → {out_path.name} ({out_path.stat().st_size / 1e9:.1f} GB)") # ── Rebuild index ──────────────────────────────────────────────────────── if index_path.exists(): # Exclude _scale_inv keys from the new weight map new_weight_map = { k: v for k, v in original_weight_map.items() if not k.endswith("_scale_inv") } # Merge with actually-written map (handles edge cases) new_weight_map.update(combined_weight_map) 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) print("Index written.") # ── Copy config / tokeniser files ──────────────────────────────────────── for pattern in [ "config.json", "generation_config.json", "tokenizer*.json", "tokenizer*.model", "special_tokens_map.json", "*.py", ]: for src in fp8_dir.glob(pattern): dst = bf16_dir / src.name if not dst.exists(): shutil.copy2(src, dst) print(f"Copied {src.name}") print(f"\n✔ BF16 model written to: {bf16_dir}") # ── CLI ────────────────────────────────────────────────────────────────────── def main() -> None: parser = argparse.ArgumentParser( description="Dequantise DeepSeek FP8 safetensors → BF16 (low RAM, shard-by-shard)" ) parser.add_argument("--input-fp8-hf-path", required=True, type=Path) parser.add_argument("--output-bf16-hf-path", required=True, type=Path) parser.add_argument( "--resume", action="store_true", help="Skip shards whose output file already exists (useful after a crash)", ) args = parser.parse_args() convert_directory(args.input_fp8_hf_path, args.output_bf16_hf_path, args.resume) if __name__ == "__main__": main()