From 7455ba102daed60591a3f654bf2b6be7c5cc9676 Mon Sep 17 00:00:00 2001 From: Gyeongseo Park Date: Fri, 23 Oct 2026 10:11:46 +0900 Subject: [PATCH 1/3] fix: KV-cache generation at batch size >= 2 (#164) (#187) The exported graph ties position_ids' batch dim to input_ids'. Prefill passed position_ids [1, S] and decode [1, 1], and decode fed the [B] sampled tokens as [1, B]. Prefill now builds [B, S] positions, decode [B, 1] positions and a [B, 1] token input. Adds tests/test_kv_cache_batch.py: greedy KV-cache generation matches HF at B=1 and B=2 (GPU only). Co-authored-by: Claude Opus 5.5 (cherry picked from commit d46b8ffcada3f0839340689afb890626d291bba7) --- opt_prime/opt_prime/inference.py | 4 +-- opt_prime/opt_prime/schedule_infer.py | 9 ++++--- opt_prime/tests/test_kv_cache_batch.py | 37 ++++++++++++++++++++++++++ 3 files changed, 43 insertions(+), 7 deletions(-) create mode 100644 opt_prime/tests/test_kv_cache_batch.py diff --git a/opt_prime/opt_prime/inference.py b/opt_prime/opt_prime/inference.py index c09ae220..73f6cb94 100644 --- a/opt_prime/opt_prime/inference.py +++ b/opt_prime/opt_prime/inference.py @@ -1047,9 +1047,7 @@ def _generate_with_kv_cache( # Forward single token via decode_step if self.tpl.is_first_stage(): - token_input = next_token.unsqueeze(0) if next_token.dim() == 1 else next_token - if token_input.dim() == 1: - token_input = token_input.unsqueeze(0) # [1, 1] + token_input = next_token.reshape(-1, 1) # [B] -> [B, 1] logits = self._schedule_gen.decode_step(token_input, cur_pos) else: logits = self._schedule_gen.decode_step(None, cur_pos) diff --git a/opt_prime/opt_prime/schedule_infer.py b/opt_prime/opt_prime/schedule_infer.py index 10bf1576..e68f79ed 100644 --- a/opt_prime/opt_prime/schedule_infer.py +++ b/opt_prime/opt_prime/schedule_infer.py @@ -370,7 +370,7 @@ def prefill( Args: input_ids: Input token IDs tensor of shape [batch_size, seq_len] - position_ids: Position IDs tensor of shape [1, seq_len] for RoPE. + position_ids: Position IDs tensor of shape [batch_size, seq_len] for RoPE. If None, auto-generated as arange(seq_len). Returns: @@ -383,9 +383,10 @@ def prefill( # Set position_ids for RoPE (critical for LLaMA) if position_ids is None and input_ids is not None: seq_len = input_ids.size(1) if input_ids.dim() > 1 else input_ids.size(0) + # The export ties position_ids' batch dim to input_ids' position_ids = torch.arange( seq_len, device=input_ids.device - ).unsqueeze(0) + ).expand(input_ids.size(0) if input_ids.dim() > 1 else 1, -1) if position_ids is not None: self.optimus.run_info.env[0]["position_ids"] = \ position_ids.to(self.optimus.run_info.device) @@ -427,8 +428,8 @@ def decode_step( if self.optimus.tpl.is_first_stage(): self.get_input(input_token) # Set position_ids for RoPE — must be the absolute position - position_ids = torch.tensor( - [[position]], device=input_token.device + position_ids = torch.full( + (input_token.size(0), 1), position, device=input_token.device ) self.optimus.run_info.env[0]["position_ids"] = \ position_ids.to(self.optimus.run_info.device) diff --git a/opt_prime/tests/test_kv_cache_batch.py b/opt_prime/tests/test_kv_cache_batch.py new file mode 100644 index 00000000..d1a094c3 --- /dev/null +++ b/opt_prime/tests/test_kv_cache_batch.py @@ -0,0 +1,37 @@ +# +# Copyright (c) 2026-present, ETRI, All rights reserved. +# +# opt_prime.inference: greedy generation with the KV cache (torch.export path) matches HF at batch size 1 and 2. +import copy +import os +import sys + +import pytest +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Optimus_Inference needs a GPU") + + +@pytest.fixture(scope="module") +def tiny(): + from transformers import AutoConfig, AutoModelForCausalLM + torch.manual_seed(0) + cfg = AutoConfig.for_model("llama", hidden_size=64, intermediate_size=128, num_hidden_layers=2, + num_attention_heads=4, num_key_value_heads=2, vocab_size=128, + attn_implementation="sdpa") + return AutoModelForCausalLM.from_config(cfg).eval() + + +@pytest.mark.parametrize("batch", [1, 2]) +def test_kv_cache_generation_matches_hf(tiny, batch): + from opt_prime.inference import Optimus_Inference + ids = torch.randint(0, 128, (batch, 6), generator=torch.Generator().manual_seed(batch)) + with torch.no_grad(): + ref = copy.deepcopy(tiny).cuda().generate(ids.cuda(), max_new_tokens=5, do_sample=False, + pad_token_id=0).cpu() + eng = Optimus_Inference(copy.deepcopy(tiny), use_kv_cache=True, dynamo_capture=True, dtype=torch.float32) + eng.eval() + out = eng.generate(ids, max_new_tokens=5, do_sample=False, verbose=False).cpu() + assert torch.equal(out, ref) From 3d07128a9b0059e5609ac20883b0d1372cdf0a5d Mon Sep 17 00:00:00 2001 From: Gyeongseo Park Date: Fri, 23 Oct 2026 10:11:54 +0900 Subject: [PATCH 2/3] fix: per-rank checkpoint names for TP/DP (#165) (#188) save_ckpt named files by pipeline stage only, so with tp > 1 the TP ranks overwrote each other's shards, and all DP replicas wrote and pruned the same files. save_lora_ckpt already named TP shards but let every DP replica write. Following save_hf_ckpt, only DP rank 0 writes and save_ckpt/load_ckpt add a _tp{rank} suffix when tp > 1 (PP-only names unchanged). The checkpoint directory no longer becomes checkpoint_None_ without a postfix (the "or" fallback never triggered). Adds tests/test_ckpt_names.py. Co-authored-by: Claude Opus 5.5 (cherry picked from commit fb09a618fed5ec20e511afd9b945da3461108f1f) --- opt_prime/opt_prime/opti_pri.py | 17 +++++-- opt_prime/tests/test_ckpt_names.py | 76 ++++++++++++++++++++++++++++++ 2 files changed, 89 insertions(+), 4 deletions(-) create mode 100644 opt_prime/tests/test_ckpt_names.py diff --git a/opt_prime/opt_prime/opti_pri.py b/opt_prime/opt_prime/opti_pri.py index 4dd36604..204b40ca 100644 --- a/opt_prime/opt_prime/opti_pri.py +++ b/opt_prime/opt_prime/opti_pri.py @@ -461,7 +461,7 @@ def __init__(self, module:nn.Module, num_mb, use_gpu=True, pp_size=1, dp_size=1, self.tpl = Topology(rank, local_rank, world_size, pp_size, dp_size, tp_size) self.checkpoint = checkpoint - self.ckpt_dir = f"checkpoint_{ckpt_dir_postfix}_{self.tpl.stage}" or f"checkpoint_{self.tpl.stage}" + self.ckpt_dir = f"checkpoint_{ckpt_dir_postfix}_{self.tpl.stage}" if ckpt_dir_postfix else f"checkpoint_{self.tpl.stage}" if self.checkpoint == True: os.makedirs(self.ckpt_dir, exist_ok=True) @@ -1165,6 +1165,10 @@ def save_lora_ckpt(self, step, epoch, lora_dir=None): print("[LoRA] No LoRA config set, skipping save.") return + # DP replicas are identical → only DP rank 0 saves + if (self.get_rank() // self.tpl.tp_size) % self.tpl.dp_size != 0: + return + if lora_dir is None: lora_dir = f"lora_checkpoint_stage_{self.tpl.stage}" else: @@ -1234,7 +1238,11 @@ def save_ckpt(self, step, epoch): if self.checkpoint != True: return - ckpt_path = os.path.join(self.ckpt_dir, f"ckpt_step_{step}_epoch_{epoch}.pt") + # DP replicas are identical → only DP rank 0 saves, one file per TP shard + if (self.get_rank() // self.tpl.tp_size) % self.tpl.dp_size != 0: + return + tp = f"_tp{self.get_rank() % self.tpl.tp_size}" if self.tpl.tp_size > 1 else "" + ckpt_path = os.path.join(self.ckpt_dir, f"ckpt_step_{step}_epoch_{epoch}{tp}.pt") torch.save({ 'epoch': epoch, @@ -1245,7 +1253,7 @@ def save_ckpt(self, step, epoch): print(f"Checkpoint saved: {ckpt_path}") ckpt_files = sorted( - glob.glob(os.path.join(self.ckpt_dir, "ckpt_step_*.pt")), + glob.glob(os.path.join(self.ckpt_dir, f"ckpt_step_*{tp}.pt")), key=os.path.getmtime, reverse=True ) @@ -1257,7 +1265,8 @@ def load_ckpt(self, step, epoch): if self.checkpoint != True: return - ckpt_path = os.path.join(self.ckpt_dir, f"ckpt_step_{step}_epoch_{epoch}.pt") + tp = f"_tp{self.get_rank() % self.tpl.tp_size}" if self.tpl.tp_size > 1 else "" + ckpt_path = os.path.join(self.ckpt_dir, f"ckpt_step_{step}_epoch_{epoch}{tp}.pt") ckpt = torch.load(ckpt_path) diff --git a/opt_prime/tests/test_ckpt_names.py b/opt_prime/tests/test_ckpt_names.py new file mode 100644 index 00000000..1b4d385a --- /dev/null +++ b/opt_prime/tests/test_ckpt_names.py @@ -0,0 +1,76 @@ +# +# Copyright (c) 2026-present, ETRI, All rights reserved. +# +# opt_prime.opti_pri checkpoint naming: TP shards get their own files, only DP rank 0 writes, and the directory +# name never contains "None". Ranks are laid out as the (pp, dp, tp) device mesh: rank = (pp*dp + dp_i)*tp + tp_i. +import os +import sys +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from opt_prime import opti_pri # noqa: E402 +from opt_prime.opti_pri import Optimus_p # noqa: E402 + + +def _rank(rank, ckpt_dir, tp_size=2, dp_size=2): + o = Optimus_p.__new__(Optimus_p) # no process group: just the state the save/load paths read + o.tpl = SimpleNamespace(rank=rank, stage=0, tp_size=tp_size, dp_size=dp_size) + o.checkpoint, o.ckpt_dir, o.lora_config, o.device = True, str(ckpt_dir), "lora", torch.device("cpu") + torch.manual_seed(rank) # every rank holds different weights + submod = nn.ModuleDict({"lora_A": nn.Linear(2, 2)}) + o.run_info = SimpleNamespace(submod=submod) + o.optimizer = torch.optim.SGD(submod.parameters(), lr=0.1) + return o + + +def _weight(o): + return o.run_info.submod["lora_A"].weight.detach().clone() + + +def test_save_ckpt_one_file_per_tp_shard_from_dp_rank_0(tmp_path): + ranks = [_rank(r, tmp_path) for r in range(4)] # tp_i = r % 2, dp_i = r // 2 + for o in ranks: + o.save_ckpt(1, 0) + assert sorted(os.listdir(tmp_path)) == ["ckpt_step_1_epoch_0_tp0.pt", "ckpt_step_1_epoch_0_tp1.pt"] + expected = _weight(ranks[1]) # TP rank 1 of DP rank 0 + reload = _rank(3, tmp_path) # TP rank 1 of DP rank 1 reloads the same shard + reload.load_ckpt(1, 0) + assert torch.equal(_weight(reload), expected) + + +def test_save_ckpt_pp_only_name_unchanged(tmp_path): + _rank(0, tmp_path, tp_size=1, dp_size=1).save_ckpt(1, 0) # existing PP-only checkpoints keep loading + assert os.listdir(tmp_path) == ["ckpt_step_1_epoch_0.pt"] + + +def test_save_lora_ckpt_only_dp_rank_0_writes(tmp_path): + ranks = [_rank(r, tmp_path) for r in range(4)] + for o in ranks: + o.save_lora_ckpt(1, 0, lora_dir=str(tmp_path / "lora")) + files = sorted(os.listdir(str(tmp_path / "lora") + "_stage_0")) + assert files == ["lora_step_1_epoch_0_tp0.pt", "lora_step_1_epoch_0_tp1.pt"] + expected = _weight(ranks[1]) + reload = _rank(3, tmp_path) + reload.load_lora_ckpt(1, 0, lora_dir=str(tmp_path / "lora")) + assert torch.equal(_weight(reload), expected) + + +class _Stop(Exception): + pass + + +@pytest.mark.parametrize("postfix, name", [(None, "checkpoint_0"), ("run1", "checkpoint_run1_0")]) +def test_ckpt_dir_name(tmp_path, monkeypatch, postfix, name): + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(opti_pri, "Comm", lambda **kw: SimpleNamespace(rank=0, world_size=1, local_rank=0)) + + def stop(*a): # __init__ names the directory just before picking the GPU + raise _Stop + monkeypatch.setattr(torch.cuda, "set_device", stop) + with pytest.raises(_Stop): + Optimus_p(nn.Linear(2, 2), 1, checkpoint=True, ckpt_dir_postfix=postfix) + assert os.listdir(tmp_path) == [name] From 0e441d2fad1f0bd311f3439d85173911f3ca8d25 Mon Sep 17 00:00:00 2001 From: Gyeongseo Park Date: Fri, 23 Oct 2026 10:12:01 +0900 Subject: [PATCH 3/3] fix: scope MPS daemon detection and cleanup to our pipe dir (#166) (#189) _mps_daemon_running fell back to a host-wide pgrep, so another job's MPS daemon (other pipe directory) looked like ours; the health probe on our pipe dir then failed and the stale cleanup ran pkill -u $UID on every MPS server and control process, killing the other job's daemon. Liveness is now our control socket only, and the stale cleanup sends quit to our pipe dir and removes our socket, without pkill. Adds tests/test_mps_manager.py with a fake pipe dir. Co-authored-by: Claude Opus 5.5 (cherry picked from commit 0e69cc9d54d19e336d33920298c39507d1415811) --- opt_prime/opt_prime/mps_manager.py | 44 +++++------------------------ opt_prime/tests/test_mps_manager.py | 36 +++++++++++++++++++++++ 2 files changed, 43 insertions(+), 37 deletions(-) create mode 100644 opt_prime/tests/test_mps_manager.py diff --git a/opt_prime/opt_prime/mps_manager.py b/opt_prime/opt_prime/mps_manager.py index 71b5cfb7..63ba6fb8 100644 --- a/opt_prime/opt_prime/mps_manager.py +++ b/opt_prime/opt_prime/mps_manager.py @@ -502,34 +502,12 @@ def cleanup_barrier_files(pipe_dir: str) -> None: # --------------------------------------------------------------------------- def _mps_daemon_running(pipe_dir: str) -> bool: - """Best-effort check whether an MPS daemon is already running. + """Whether an MPS daemon serves this pipe directory (its control socket exists). - Looks for the pipe directory's control socket and probes pgrep. - Returns True if either signal is present. + A host-wide process check would mistake another job's daemon (other pipe + directory) for ours, and the stale cleanup would then hit it. """ - # 1) control socket presence (most reliable when MPS daemon started cleanly) - if os.path.exists(os.path.join(pipe_dir, "control")): - return True - # 2) Process check - try: - out = subprocess.run( - ["pgrep", "-x", "nvidia-cuda-mps-server"], - capture_output=True, text=True, timeout=2, - ) - if out.returncode == 0 and out.stdout.strip(): - return True - except (FileNotFoundError, subprocess.TimeoutExpired, OSError): - pass - try: - out = subprocess.run( - ["pgrep", "-x", "nvidia-cuda-mps-control"], - capture_output=True, text=True, timeout=2, - ) - if out.returncode == 0 and out.stdout.strip(): - return True - except (FileNotFoundError, subprocess.TimeoutExpired, OSError): - pass - return False + return os.path.exists(os.path.join(pipe_dir, "control")) # --------------------------------------------------------------------------- @@ -786,7 +764,7 @@ def _daemon_responsive(self, timeout_s: float = 2.0) -> bool: return False def _force_cleanup_stale_daemon(self) -> None: - """Best-effort cleanup of a stale daemon: try quit, then pkill (own user).""" + """Best-effort cleanup of a stale daemon on our pipe dir: quit, drop socket.""" try: subprocess.run( ["nvidia-cuda-mps-control"], @@ -798,16 +776,8 @@ def _force_cleanup_stale_daemon(self) -> None: ) except (subprocess.TimeoutExpired, OSError): pass - # Force kill stragglers (only this user's processes) - for proc_name in ("nvidia-cuda-mps-server", "nvidia-cuda-mps-control"): - try: - subprocess.run( - ["pkill", "-u", str(os.geteuid()), "-x", proc_name], - capture_output=True, timeout=5, - ) - except (FileNotFoundError, subprocess.TimeoutExpired, OSError): - pass - # Remove control socket if present + # No pkill: other jobs' daemons (other pipe directories) share the + # process names. Remove our control socket so a new daemon can bind. ctrl = os.path.join(self.pipe_dir, "control") try: os.unlink(ctrl) diff --git a/opt_prime/tests/test_mps_manager.py b/opt_prime/tests/test_mps_manager.py new file mode 100644 index 00000000..d72d7237 --- /dev/null +++ b/opt_prime/tests/test_mps_manager.py @@ -0,0 +1,36 @@ +# +# Copyright (c) 2026-present, ETRI, All rights reserved. +# +# opt_prime.mps_manager: daemon detection and stale cleanup stay scoped to this job's pipe directory. +import os +import subprocess +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from opt_prime import mps_manager # noqa: E402 + + +def _fake_run(calls): + # Every command "succeeds" as if another job's MPS daemon were running on this host. + def run(cmd, *a, **kw): + calls.append(cmd) + return subprocess.CompletedProcess(cmd, 0, stdout="4242\n", stderr="") + return run + + +def test_daemon_running_only_from_own_control_socket(tmp_path, monkeypatch): + calls = [] + monkeypatch.setattr(subprocess, "run", _fake_run(calls)) + assert not mps_manager._mps_daemon_running(str(tmp_path)) # another job's daemon is not ours + (tmp_path / "control").touch() + assert mps_manager._mps_daemon_running(str(tmp_path)) + + +def test_stale_cleanup_does_not_kill_other_daemons(tmp_path, monkeypatch): + calls = [] + monkeypatch.setattr(subprocess, "run", _fake_run(calls)) + monkeypatch.setattr(mps_manager.time, "sleep", lambda s: None) + (tmp_path / "control").touch() + mps_manager.MPSManager(pipe_dir=str(tmp_path), log_dir=str(tmp_path / "log"))._force_cleanup_stale_daemon() + assert not any(c[0] in ("pkill", "killall") for c in calls), calls + assert not (tmp_path / "control").exists()