perf(ocr_qwen): support CPU + bfloat16 AVX-512 + threads explicites

Trois ajouts pour rendre le pipeline utilisable sur CPU quand la VRAM
est saturée par d'autres process :

1. Variable QWEN_DEVICE=cpu pour forcer le device CPU. Le défaut "auto"
   choisit CUDA si dispo, fallback CPU sinon.

2. Sur CPU, détection automatique du support AVX-512 BF16 via /proc/cpuinfo
   (Zen 4/5, Intel Sapphire Rapids+). Si présent, bfloat16 au lieu de
   float32 — divise par 2 la RAM et ~2x plus rapide sur matmul.

3. Appel explicite de torch.set_num_threads(N) et set_num_interop_threads(N)
   (OMP_NUM_THREADS seul ne suffit pas). Configurable via TORCH_NUM_THREADS,
   défaut = os.cpu_count().

Mesure sur Ryzen 9 9950X (Zen 5, 16c/32t, AVX-512 BF16 natif) :
- AVANT : 645% CPU (~6.5 cores), 15 Go RAM (float32)
- APRÈS : 2433% CPU (~24 cores), 8 Go RAM (bfloat16)

Appel `torch.cuda.empty_cache()` en fin d'inférence pour réduire la
fragmentation VRAM quand d'autres process GPU tournent en parallèle.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Dom
2026-04-24 23:07:45 +02:00
parent 6c8184cc03
commit c0b0cd9b87

View File

@@ -27,22 +27,64 @@ class QwenVLOCR:
def _init_model(self): def _init_model(self):
t0 = time.time() t0 = time.time()
import os as _os
# max_pixels limite le nombre de patches visuels pour éviter l'OOM # max_pixels limite le nombre de patches visuels pour éviter l'OOM
# sur images 300 dpi (2481x3509). ~800 patches = équilibre qualité/VRAM, # sur images 300 dpi (2481x3509). ~800 patches = équilibre qualité/VRAM,
# tient confortablement dans ~5-6 Go même avec d'autres processus GPU # tient confortablement dans ~5-6 Go même avec d'autres processus GPU
# en arrière-plan. Configurable via env var QWEN_MAX_PIXELS (en patches). # en arrière-plan. Configurable via env var QWEN_MAX_PIXELS (en patches).
import os as _os
max_pixels = int(_os.environ.get("QWEN_MAX_PIXELS", 800)) * 28 * 28 max_pixels = int(_os.environ.get("QWEN_MAX_PIXELS", 800)) * 28 * 28
self.processor = AutoProcessor.from_pretrained( self.processor = AutoProcessor.from_pretrained(
MODEL_PATH, MODEL_PATH,
min_pixels=256 * 28 * 28, min_pixels=256 * 28 * 28,
max_pixels=max_pixels, max_pixels=max_pixels,
) )
# Device : "auto" par défaut (GPU si dispo), "cpu" pour forcer le CPU
# quand la VRAM est saturée par d'autres process. Configurable via
# QWEN_DEVICE=cpu.
device = _os.environ.get("QWEN_DEVICE", "auto").lower()
if device == "cpu":
# Sur CPU on cherche à maximiser le throughput :
# 1. Utiliser tous les cores via torch.set_num_threads (set_num_threads
# prime sur OMP_NUM_THREADS pour les ops PyTorch natifs).
# 2. Choisir bfloat16 si le CPU le supporte nativement (Zen 5,
# Zen 4, Intel Sapphire Rapids+ ont AVX-512 BF16). Sinon float32.
n_threads = int(_os.environ.get("TORCH_NUM_THREADS", _os.cpu_count() or 8))
torch.set_num_threads(n_threads)
try:
torch.set_num_interop_threads(n_threads)
except RuntimeError:
pass # déjà initialisé, ignorer
# Détection AVX-512 BF16 via /proc/cpuinfo (Linux)
use_bf16 = False
try:
with open("/proc/cpuinfo") as f:
flags = f.read()
use_bf16 = "avx512_bf16" in flags or "amx_bf16" in flags
except Exception:
pass
dtype = torch.bfloat16 if use_bf16 else torch.float32
self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
MODEL_PATH,
torch_dtype=dtype,
device_map={"": "cpu"},
low_cpu_mem_usage=True,
)
self.device_used = "cpu"
self.cpu_threads = n_threads
self.cpu_dtype = str(dtype).replace("torch.", "")
else:
self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained( self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
MODEL_PATH, MODEL_PATH,
torch_dtype=torch.bfloat16, torch_dtype=torch.bfloat16,
device_map="auto", device_map="auto",
) )
self.device_used = "cuda" if torch.cuda.is_available() else "cpu"
self.cpu_threads = None
self.cpu_dtype = None
self.model.eval() self.model.eval()
self.load_time = time.time() - t0 self.load_time = time.time() - t0
self.vram_gb = torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0 self.vram_gb = torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0