| """Auto-batch sizing for GPU inference. |
| |
| ControlMT's `translate()` runs beam search sentence-by-sentence, but we |
| ThreadPoolExecutor across N concurrent calls. N is bounded by free VRAM. |
| |
| Memory per concurrent translation at beam=2, max_len=256: |
| fp16: ~25 MB (encoder + decoder state + beam tensors + kv-cache) |
| bf16: ~25 MB |
| fp32: ~50 MB |
| """ |
|
|
| from __future__ import annotations |
| import warnings |
|
|
|
|
| _MEM_PER_SENT_MB = { |
| "float16": 25, |
| "bfloat16": 25, |
| "float32": 50, |
| "int8-dynamic": 12, |
| } |
| _MAX_BATCH = 64 |
| _TARGET_VRAM_FRAC = 0.8 |
|
|
|
|
| def auto_batch_size(device: str, dtype: str, quant: str = "none", |
| free_vram_mb: float | None = None) -> int: |
| """Estimate batch size that fits in ~80% of free VRAM. |
| |
| Returns 1 on CPU (no auto-batching there) with a warning. |
| On GPU returns 1..MAX_BATCH. |
| """ |
| if device != "cuda": |
| warnings.warn( |
| "auto_batch=True ignored on CPU (no VRAM signal to size against). " |
| "Pass an explicit batch_size for batched CPU translation.", |
| RuntimeWarning, stacklevel=2) |
| return 1 |
|
|
| if free_vram_mb is None: |
| import torch |
| free, _total = torch.cuda.mem_get_info() |
| free_vram_mb = free / 1024**2 |
|
|
| per_sent_mb = _MEM_PER_SENT_MB.get(quant if quant != "none" else dtype, 50) |
| n = int((free_vram_mb * _TARGET_VRAM_FRAC) // per_sent_mb) |
| return max(1, min(_MAX_BATCH, n)) |
|
|
