multi-GPU robustness: translator.py

src/translator.py

@@ -10,14 +10,20 @@ from functools import lru_cache
 from pathlib import Path
 from typing import Callable, Iterator
 
-import ctranslate2
 import sentencepiece as spm
 from huggingface_hub import snapshot_download
 
 from chunker import split_chunks
-from hardware import HardwareProfile, detect_hardware_profile, resolve_gpu_indices
+from hardware import (
+    HardwareProfile,
+    auto_all_gpus_by_default,
+    detect_hardware_profile,
+    resolve_gpu_indices,
+)
 from token_chunker import source_token_ids, split_for_translation
 
+import ctranslate2
+
 ROOT = Path(__file__).resolve().parent.parent
 MODELS_DIR = Path(os.environ.get("HACHIMIMT_MODELS_DIR", ROOT / "models"))
 SPECIAL_ID_TO_TOKEN = {0: "<pad>", 1: "<s>", 2: "</s>", 3: "<unk>"}
@@ -155,6 +161,47 @@ def default_ct2_compute_type(device: str) -> str:
     return "int8_float16" if device == "cuda" else "int8_float32"
 
 
+def _ct2_gpu_index_attempts(gpu_indices: list[int]) -> list[list[int]]:
+    """Try all requested GPUs first, then one GPU before giving up to CPU."""
+    if len(gpu_indices) <= 1:
+        return [list(gpu_indices)]
+    return [list(gpu_indices), [gpu_indices[0]]]
+
+
+def _ct2_translator_kwargs(
+    *,
+    device: str,
+    compute_type: str,
+    intra_threads: int,
+    inter_threads: int,
+    gpu_indices: list[int] | None = None,
+) -> tuple[dict[str, object], int, str | None]:
+    kwargs: dict[str, object] = dict(
+        device=device,
+        compute_type=compute_type,
+        intra_threads=intra_threads,
+        inter_threads=max(1, int(inter_threads)),
+    )
+    if device != "cuda":
+        return kwargs, max(1, int(inter_threads)), None
+
+    if not gpu_indices:
+        raise RuntimeError("Không có GPU CUDA khả dụng.")
+
+    selected = list(dict.fromkeys(gpu_indices))
+    if len(selected) == 1:
+        # Luôn truyền device_index, kể cả single GPU: env "1" phải dùng GPU 1.
+        kwargs["device_index"] = selected[0]
+    else:
+        # CT2 inter_threads = replica trên MỖI device. Với nhiều GPU, giữ 1
+        # replica/GPU để tránh nhân VRAM và đã nhanh hơn trong benchmark T4x2.
+        kwargs["device_index"] = selected
+        kwargs["inter_threads"] = 1
+
+    actual_inter_threads = int(kwargs["inter_threads"])
+    return kwargs, len(selected) * actual_inter_threads, ",".join(str(i) for i in selected)
+
+
 @lru_cache(maxsize=1)
 def _optional_torch():
     try:
@@ -386,6 +433,9 @@ class HachimiTranslator:
         batch_type = os.environ.get("HACHIMIMT_CT2_BATCH_TYPE", "tokens").strip().lower()
         self._ct2_batch_type = batch_type if batch_type in {"examples", "tokens"} else "tokens"
         self._ct2_compute_type: str | None = None
+        self._ct2_actual_inter_threads = self._ct2_inter_threads
+        self._ct2_worker_count = 1
+        self._ct2_device_indices_label: str | None = None
         self._batch_size = self._profile.batch_size
         self._tokenize_workers = self._profile.tokenize_workers
         self._tokenize_pool: ThreadPoolExecutor | None = None
@@ -486,11 +536,19 @@ class HachimiTranslator:
         msg = f"{prefix} {config.label} · {engine} · {self.device_label()}"
         if backend == Backend.CT2 and self._ct2_compute_type:
             msg += f" · compute={self._ct2_compute_type}"
+            window_multiplier = self._ct2_effective_window_multiplier()
+            window_part = f"window={self._ct2_window_multiplier}x"
+            if window_multiplier != self._ct2_window_multiplier:
+                window_part += f"/{window_multiplier}x"
             msg += (
                 f" · batch_type={self._ct2_batch_type}"
-                f" · window={self._ct2_window_multiplier}x"
-                f" · inter={self._ct2_inter_threads}"
+                f" · {window_part}"
+                f" · inter={self._ct2_actual_inter_threads}"
             )
+            if self._ct2_worker_count > 1:
+                msg += f" · workers={self._ct2_worker_count}"
+            if self._ct2_device_indices_label:
+                msg += f" · gpu={self._ct2_device_indices_label}"
         if beam_size is not None:
             msg += f" · beam={beam_size}"
         return msg
@@ -505,6 +563,9 @@ class HachimiTranslator:
         self._tokenizer = None
         self._model_path = None
         self._ct2_compute_type = None
+        self._ct2_actual_inter_threads = self._ct2_inter_threads
+        self._ct2_worker_count = 1
+        self._ct2_device_indices_label = None
         if self._tokenize_pool is not None:
             self._tokenize_pool.shutdown(wait=False, cancel_futures=True)
             self._tokenize_pool = None
@@ -554,50 +615,48 @@ class HachimiTranslator:
 
         env_compute_type = os.environ.get("HACHIMIMT_COMPUTE_TYPE", "").strip()
         ct2_device = "cuda" if self._profile.has_cuda else "cpu"
-        attempts = [(ct2_device, default_ct2_compute_type(ct2_device))]
-        if not env_compute_type and ct2_device == "cuda":
-            attempts.extend([("cuda", "int8_float32"), ("cpu", "int8_float32")])
-        elif env_compute_type and ct2_device == "cuda":
-            attempts.append(("cpu", "int8_float32"))
-
-        # Multi-GPU: nếu có >1 GPU (vd Kaggle T4x2) → CT2 chia batch ra các GPU
-        # (~1.69× compute đo thực). Chỉ áp cho device cuda; CPU bỏ qua device_index.
-        try:
-            cuda_count = ctranslate2.get_cuda_device_count()
-        except Exception:
-            cuda_count = 1
-        gpu_indices = resolve_gpu_indices(
-            cuda_count, os.environ.get("HACHIMIMT_GPU_INDICES")
-        )
+        attempts: list[tuple[str, str, list[int] | None]] = []
+        if ct2_device == "cuda":
+            try:
+                cuda_count = ctranslate2.get_cuda_device_count()
+            except Exception:
+                cuda_count = 0
+            gpu_indices = resolve_gpu_indices(
+                cuda_count,
+                os.environ.get("HACHIMIMT_GPU_INDICES"),
+                auto_all=auto_all_gpus_by_default(),
+            )
+            compute_types = [default_ct2_compute_type("cuda")]
+            if not env_compute_type and "int8_float32" not in compute_types:
+                compute_types.append("int8_float32")
+            for compute_type in compute_types:
+                for candidate_indices in _ct2_gpu_index_attempts(gpu_indices):
+                    attempts.append(("cuda", compute_type, candidate_indices))
+            attempts.append(("cpu", "int8_float32", None))
+        else:
+            cpu_compute_type = default_ct2_compute_type("cpu")
+            attempts.append(("cpu", cpu_compute_type, None))
+            if cpu_compute_type != "int8_float32":
+                attempts.append(("cpu", "int8_float32", None))
 
         translator = None
         last_error: Exception | None = None
-        for device, compute_type in attempts:
+        for device, compute_type, gpu_indices in attempts:
             try:
-                kwargs = dict(
+                kwargs, worker_count, device_indices_label = _ct2_translator_kwargs(
                     device=device,
                     compute_type=compute_type,
                     intra_threads=self._ct2_threads,
                     inter_threads=self._ct2_inter_threads,
+                    gpu_indices=gpu_indices,
                 )
-                if device == "cuda":
-                    # LUÔN truyền device_index khi cuda — kể cả 1 GPU (vd "1" để
-                    # chọn GPU 1; nếu bỏ qua thì CT2 mặc định GPU 0 = chạy sai).
-                    if not gpu_indices:
-                        raise RuntimeError("Không có GPU CUDA khả dụng.")
-                    if len(gpu_indices) == 1:
-                        kwargs["device_index"] = gpu_indices[0]   # int
-                    else:
-                        # Đa GPU: list device + inter_threads=1 = 1 REPLICA/GPU.
-                        # (CT2 inter_threads = replica MỖI device, KHÔNG phải tổng.
-                        # Benchmark T4x2: 2GPU·1replica = 1.64× & NHANH NHẤT; 2
-                        # replica/GPU = 4 worker CHẬM HƠN -17% + 4× VRAM → để =1.)
-                        kwargs["device_index"] = gpu_indices
-                        kwargs["inter_threads"] = 1
                 translator = ctranslate2.Translator(
                     str(model_path / config.ct2_subdir), **kwargs
                 )
                 self._ct2_compute_type = compute_type
+                self._ct2_actual_inter_threads = int(kwargs["inter_threads"])
+                self._ct2_worker_count = worker_count
+                self._ct2_device_indices_label = device_indices_label
                 break
             except Exception as exc:
                 last_error = exc
@@ -690,9 +749,17 @@ class HachimiTranslator:
     def _runtime_window_size(self, beam_size: int) -> int:
         batch_size = self._runtime_batch_size(beam_size)
         if self._backend == Backend.CT2:
-            return max(batch_size, batch_size * self._ct2_window_multiplier)
+            return max(batch_size, batch_size * self._ct2_effective_window_multiplier())
         return batch_size
 
+    def _ct2_effective_window_multiplier(self) -> int:
+        multiplier = self._ct2_window_multiplier
+        if self._ct2_batch_type == "tokens" and self._ct2_worker_count > 1:
+            # Multi-GPU needs enough queued chunks for CT2 to split into multiple
+            # token sub-batches; cap keeps large files from over-buffering.
+            multiplier *= min(self._ct2_worker_count * 2, 8)
+        return max(1, min(32, multiplier))
+
     def _ct2_max_batch_size(self, config: ModelConfig) -> int:
         if self._ct2_batch_type == "tokens":
             return self._batch_size * config.ct2_max_input_tokens
@@ -824,8 +891,13 @@ class HachimiTranslator:
 
         for start in range(0, total, window_size):
             end = min(start + window_size, total)
+            worker_label = (
+                f", workers {self._ct2_worker_count}"
+                if self._backend == Backend.CT2 and self._ct2_worker_count > 1
+                else ""
+            )
             batch_label = (
-                f"window {window_size}, batch {batch_size}, {self._ct2_batch_type}"
+                f"window {window_size}, batch {batch_size}, {self._ct2_batch_type}{worker_label}"
                 if self._backend == Backend.CT2
                 else f"batch {batch_size}"
             )