Upload load_fp8.py with huggingface_hub

load_fp8.py

@@ -0,0 +1,122 @@
+"""Load FP8-quantized cohere-transcribe model from safetensors.
+
+Usage:
+    from load_fp8 import load_model
+    model, processor = load_model("path/to/repo", compile=True)
+"""
+import torch
+import torch._dynamo
+import torch._inductor.config
+import json
+import os
+from safetensors.torch import load_file
+from transformers import CohereAsrForConditionalGeneration, AutoProcessor, AutoConfig
+
+torch._inductor.config.triton.cudagraphs = False
+
+
+def load_model(model_path, device="auto", native_fp8=True, compile=False):
+    """Load FP8-quantized cohere-transcribe model.
+    
+    Args:
+        model_path: Path to the FP8 model directory
+        device: Device for model placement
+        native_fp8: Reconstruct torchao Float8Tensors for FP8 compute (requires torchao)
+        compile: Apply torch.compile to encoder layers for max throughput
+    
+    Returns:
+        (model, processor) tuple
+    """
+    with open(os.path.join(model_path, "fp8_config.json")) as f:
+        fp8_config = json.load(f)
+    fp8_layers = set(fp8_config["fp8_layers"])
+    raw = load_file(os.path.join(model_path, "model.safetensors"))
+
+    # Build BF16 state dict (dequantize FP8 layers)
+    state_dict = {}
+    for layer_name in fp8_layers:
+        qdata = raw[layer_name + ".qdata"]
+        scale = raw[layer_name + ".scale"]
+        state_dict[layer_name] = (qdata.float() * scale).to(torch.bfloat16)
+    for key, tensor in raw.items():
+        if not key.endswith(".qdata") and not key.endswith(".scale"):
+            state_dict[key] = tensor
+
+    config = AutoConfig.from_pretrained(model_path)
+    model = CohereAsrForConditionalGeneration(config)
+    model.to(torch.bfloat16)
+    model.load_state_dict(state_dict, strict=False)
+    
+    target_device = "cuda" if (device == "auto" and torch.cuda.is_available()) else device
+    model = model.to(target_device)
+    model.eval()
+
+    # Re-quantize to native FP8
+    if native_fp8:
+        from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig
+        quantize_(model.model.encoder, Float8DynamicActivationFloat8WeightConfig())
+
+    # Compile encoder per-layer
+    if compile:
+        encoder = model.model.encoder
+        torch._dynamo.config.cache_size_limit = len(encoder.layers) + 4
+        for layer in encoder.layers:
+            layer.forward = torch.compile(layer.forward, dynamic=True, mode="max-autotune-no-cudagraphs")
+
+    processor = AutoProcessor.from_pretrained(model_path)
+    return model, processor
+
+
+if __name__ == "__main__":
+    import argparse, time
+    import numpy as np
+    from transformers.audio_utils import load_audio
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", required=True)
+    parser.add_argument("--audio", default=None)
+    parser.add_argument("--language", default="en")
+    parser.add_argument("--no-fp8", dest="native_fp8", action="store_false")
+    parser.add_argument("--compile", action="store_true")
+    args = parser.parse_args()
+
+    mode = []
+    if args.native_fp8: mode.append("FP8")
+    else: mode.append("BF16")
+    if args.compile: mode.append("+ compile")
+    print(f"Loading model ({' '.join(mode)})...")
+    t0 = time.time()
+    model, processor = load_model(args.model_path, native_fp8=args.native_fp8, compile=args.compile)
+    load_time = time.time() - t0
+
+    if args.audio:
+        audio = load_audio(args.audio, sampling_rate=16000)
+    else:
+        from huggingface_hub import hf_hub_download
+        audio_file = hf_hub_download(
+            repo_id="CohereLabs/cohere-transcribe-03-2026",
+            filename="demo/voxpopuli_test_en_demo.wav",
+        )
+        audio = load_audio(audio_file, sampling_rate=16000)
+
+    # Warmup (triggers compile if enabled)
+    inputs = processor(audio, sampling_rate=16000, return_tensors="pt", language=args.language)
+    inputs.to(model.device, dtype=model.dtype)
+    _ = model.generate(**inputs, max_new_tokens=256)
+    print(f"Loaded + warmup in {time.time()-t0:.1f}s")
+
+    dur = len(audio) / 16000
+    times = []
+    for _ in range(5):
+        inputs = processor(audio, sampling_rate=16000, return_tensors="pt", language=args.language)
+        inputs.to(model.device, dtype=model.dtype)
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        outputs = model.generate(**inputs, max_new_tokens=256)
+        torch.cuda.synchronize()
+        times.append(time.perf_counter() - t0)
+
+    text = processor.decode(outputs, skip_special_tokens=True)
+    avg = np.mean(times)
+    print(f"RTFx: {dur/avg:.1f} ({avg*1000:.1f}ms)")
+    print(f"Transcription: {text}")