feat(parakeet-nemo): speech seat moves to parakeet-unified-en under NeMo (bf16 weights)
Prime-approved switch of the fleet STT seat (fv-ml1 :8300, LiteLLM ext-stt/ whisper-1, caller talk) from the sherpa-onnx int8 seat to arm B-bf16w of the 2026-09-30 A/B (docs/pfi/parakeet-seat-ab-2026-09-30.md): p50 33/36/42/71 ms vs the old seat's 187/308/626 measured on the same card today, WER 1.965/3.026 vs the A/B floor 1.97/3.09. All three seat defects fixed: 12-min file 200s (windowed at 360 s after a GPU 0 OOM on one whole request; the A/B's own long-form method), no pause truncation, no long-form dropout. GPU 0 room: gen-small --gpu-memory-utilization 0.48 -> 0.36 (0.46 and 0.40 refuse their boot check; cyberprev+voices hold the card). Its KV is byte- pinned, so the boot log is token-identical: 670,142 tokens / 2.56x before and after. Seat rests 2,088 MiB; GPU 0 keeps ~1.9 GB free. Two runtime landmines documented in the README: NeMo's attention mask is materialised T x T even under local attention (hence the window), and httptools 0.8.0 writes a NUL into the HTTP status line that httpx — i.e. LiteLLM — rejects, so the image ships plain uvicorn with --http h11. Old seat stopped, not removed: docker stop parakeet-nemo && docker start parakeet is the rollback. License: NVIDIA Open Model License (accepted by Prime 2026-09-30); note in stacks/parakeet-nemo/README.md.
This commit is contained in:
@@ -0,0 +1,154 @@
|
||||
"""Parakeet ASR seat: nvidia/parakeet-unified-en-0.6b under NeMo torch (bf16 weights).
|
||||
|
||||
Fork of the A/B winner (services/parakeet-ab-2026-09-30 arm B-bf16w), with the three shipping
|
||||
changes that arm's doc called for and a longer warm-up. Same HTTP shape as the sherpa seat it
|
||||
replaces: model loaded at import, warm-up before traffic, async handlers over a serialised
|
||||
blocking decode, {"text": ...} responses on /transcribe and /v1/audio/transcriptions.
|
||||
|
||||
Hard-wired, because every one of these is load-bearing on GPU 0:
|
||||
bf16 weights — the encoder/decoder/joint are cast to bfloat16 (mel front end stays fp32).
|
||||
WER is identical to fp32 (399/400 utterances byte-equal in the A/B).
|
||||
CPU-then-cast — the .nemo restores on CPU, weights are cast to bf16 there, and only then move
|
||||
to the GPU. Restoring to CUDA spikes the load by ~1.5 GB; GPU 0 cannot absorb it.
|
||||
local attn — rel_pos_local_attn ±128 (NeMo's documented long-audio mode). This is what fixes
|
||||
the old seat's >400 s HTTP 500 and the long-form dropouts; memory grows linearly
|
||||
in file length. A 30-min file transcribes in ~2.6 s in one request (A/B, 2026-09-30).
|
||||
Warm-up runs ascending silent clips (1 s, 8 s, 60 s): the CUDA-graph greedy decoder and the
|
||||
encoder kernels cost ~330 ms extra on the first call at a new maximum length.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
from fastapi import FastAPI, File, HTTPException, UploadFile
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
MODEL_PATH = os.environ["MODEL_PATH"]
|
||||
WARMUP_SECONDS = [int(x) for x in os.environ.get("WARMUP_SECONDS", "1,8,60").split(",")]
|
||||
SR = 16000
|
||||
|
||||
logger = logging.getLogger("parakeet-nemo")
|
||||
logging.basicConfig(level=os.environ.get("LOG_LEVEL", "INFO"))
|
||||
|
||||
|
||||
def _load():
|
||||
import nemo.collections.asr as nemo_asr
|
||||
from omegaconf import open_dict
|
||||
|
||||
t0 = time.monotonic()
|
||||
m = nemo_asr.models.ASRModel.restore_from(MODEL_PATH, map_location="cpu")
|
||||
m.eval()
|
||||
if m.cfg.get("validation_ds") is None: # the unified .nemo ships without it; transcribe() reads it
|
||||
with open_dict(m.cfg):
|
||||
m.cfg.validation_ds = {}
|
||||
d = m.cfg.decoding
|
||||
with open_dict(d):
|
||||
d.strategy = "greedy_batch"
|
||||
d.greedy["use_cuda_graph_decoder"] = True
|
||||
m.change_decoding_strategy(d, verbose=False)
|
||||
# transcribe() sets these on entry; the direct path must match, and must not dither (dither is
|
||||
# a training-time augmentation and makes the same file decode differently on each call).
|
||||
m.preprocessor.featurizer.dither = 0.0
|
||||
m.preprocessor.featurizer.pad_to = 0
|
||||
m.change_attention_model("rel_pos_local_attn", [128, 128])
|
||||
# bf16 the serving modules BEFORE the H2D copy: halves the transfer and skips the GPU-side
|
||||
# fp32->bf16 transient entirely (the measured load spike goes from 3,194 MiB to under 2,600).
|
||||
# ⚠ Must run AFTER change_attention_model: that call rebuilds the attention modules in fp32,
|
||||
# and casting first leaves fp32 islands behind (RuntimeError: mat1 and mat2 ... BFloat16/Float,
|
||||
# hit live at the first boot of this image, 2026-10-01).
|
||||
for mod in (m.encoder, m.decoder, m.joint):
|
||||
mod.to(torch.bfloat16)
|
||||
m = m.to("cuda")
|
||||
logger.info("loaded %s (%s) bf16w local_att=128,128 in %.1fs", os.path.basename(MODEL_PATH),
|
||||
type(m).__name__, time.monotonic() - t0)
|
||||
return m
|
||||
|
||||
|
||||
model = _load()
|
||||
|
||||
|
||||
def _hyp_text(h) -> str:
|
||||
if isinstance(h, str):
|
||||
return h
|
||||
t = getattr(h, "text", None)
|
||||
if isinstance(t, str):
|
||||
return t
|
||||
return model.tokenizer.ids_to_text([int(i) for i in h.y_sequence])
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _infer(samples: np.ndarray) -> str:
|
||||
x = torch.from_numpy(samples).to("cuda", non_blocking=True).unsqueeze(0)
|
||||
xl = torch.tensor([x.shape[1]], device="cuda", dtype=torch.long)
|
||||
feats, fl = model.preprocessor(input_signal=x, length=xl)
|
||||
feats = feats.to(torch.bfloat16)
|
||||
enc, el = model.encoder(audio_signal=feats, length=fl)
|
||||
hyps = model.decoding.rnnt_decoder_predictions_tensor(encoder_output=enc, encoded_lengths=el,
|
||||
return_hypotheses=False)
|
||||
if isinstance(hyps, tuple):
|
||||
hyps = hyps[0]
|
||||
return _hyp_text(hyps[0])
|
||||
|
||||
|
||||
def _warm() -> None:
|
||||
for secs in WARMUP_SECONDS:
|
||||
t0 = time.monotonic()
|
||||
_infer(np.zeros(SR * secs, dtype=np.float32))
|
||||
torch.cuda.synchronize()
|
||||
logger.info("warmup %ss decode complete in %.1fs", secs, time.monotonic() - t0)
|
||||
|
||||
|
||||
_warm()
|
||||
app = FastAPI(title="Parakeet ASR (NeMo torch, bf16w)")
|
||||
|
||||
|
||||
def _decode(raw: bytes) -> str:
|
||||
try:
|
||||
samples, sample_rate = sf.read(io.BytesIO(raw), dtype="float32")
|
||||
except Exception as exc:
|
||||
raise HTTPException(400, f"Could not decode audio: {exc}") from exc
|
||||
if samples.ndim > 1:
|
||||
samples = samples.mean(axis=1).astype(np.float32)
|
||||
if sample_rate != SR:
|
||||
import torchaudio.functional as AF
|
||||
samples = AF.resample(torch.from_numpy(samples), sample_rate, SR).numpy()
|
||||
samples = np.ascontiguousarray(samples, dtype=np.float32)
|
||||
# Windowed long-form (> WINDOW_S): NeMo's attention mask is materialised T×T even under
|
||||
# rel_pos_local_attn, so one whole 12-min request wanted +1.1 GiB of scratch and OOMed on
|
||||
# GPU 0 (measured at acceptance, 2026-10-01). The A/B's own long-form arm used ~6-min
|
||||
# windows and lost zero clean speech on unified-en in 4/4 placements (doc § 5.4), so the
|
||||
# seat chunks at the same size: bounded memory, any length, no API change for callers.
|
||||
win = int(os.environ.get("WINDOW_S", "360")) * SR
|
||||
if len(samples) <= win:
|
||||
return _infer(samples)
|
||||
parts = [_infer(samples[i:i + win]) for i in range(0, len(samples), win)]
|
||||
return " ".join(p for p in parts if p)
|
||||
|
||||
|
||||
def _timed(raw: bytes) -> JSONResponse:
|
||||
t0 = time.perf_counter()
|
||||
text = _decode(raw)
|
||||
torch.cuda.synchronize()
|
||||
ms = (time.perf_counter() - t0) * 1000.0
|
||||
return JSONResponse({"text": text}, headers={"x-decode-ms": f"{ms:.3f}"})
|
||||
|
||||
|
||||
@app.get("/healthz")
|
||||
def healthz() -> dict[str, str]:
|
||||
return {"status": "ok"}
|
||||
|
||||
|
||||
@app.post("/transcribe")
|
||||
async def transcribe(file: UploadFile = File(...)):
|
||||
return _timed(await file.read())
|
||||
|
||||
|
||||
@app.post("/v1/audio/transcriptions")
|
||||
async def openai_transcriptions(file: UploadFile = File(...)):
|
||||
return _timed(await file.read())
|
||||
Reference in New Issue
Block a user