harness: stabilize embed OOM tuning and defaults#1823
Open
jioffe502 wants to merge 4 commits intoNVIDIA:mainfrom
Open
harness: stabilize embed OOM tuning and defaults#1823jioffe502 wants to merge 4 commits intoNVIDIA:mainfrom
jioffe502 wants to merge 4 commits intoNVIDIA:mainfrom
Conversation
Port embed OOM hardening and decoupled inference microbatch controls into the graph-era harness path so bo767 runs can be tuned for stability without conflating Ray scheduling batch size. Capture OOM outlier metadata for attribution, restore graph Ray init parity knobs, and document the refactor/validation context for lead review. Signed-off-by: jioffe502 <jioffe@nvidia.com>
jioffe502
changed the title
Contributor
Greptile SummaryThis PR decouples embedding transport batch size (
|
| Filename | Overview |
|---|---|
| nemo_retriever/src/nemo_retriever/model/local/llama_nemotron_embed_1b_v2_embedder.py | Core change: adds adaptive OOM retry loop in _embed_local with _adaptive_batch_size actor state, plus diagnostic helpers. CUDNN_STATUS_INTERNAL_ERROR in the OOM-detection set is broader than pure allocation failures (P2). |
| nemo_retriever/src/nemo_retriever/harness/config.py | Adds embed_inference_batch_size: int = 32 to HarnessConfig and TUNING_FIELDS; validation and env-var wiring are consistent with the existing pattern. |
| nemo_retriever/src/nemo_retriever/harness/run.py | Correctly forwards --embed-inference-batch-size to the subprocess command; no issues found. |
| nemo_retriever/src/nemo_retriever/graph_ingestor.py | Adds _resolve_object_store_memory_bytes() for Ray object-store sizing and forwards ray_log_to_driver/debug into Ray init; new public constructor params are missing from the class docstring. |
| nemo_retriever/src/nemo_retriever/examples/graph_pipeline.py | Exposes --embed-inference-batch-size CLI flag and wires it into EmbedParams; fallback to embed_batch_size preserves backward compatibility for callers that don't set the new flag. |
| nemo_retriever/src/nemo_retriever/utils/ray_resource_hueristics.py | No functional changes; EMBED_BATCH_SIZE comment now inaccurately describes the constant as driving both transport and inference batch sizing after the decoupling introduced in this PR. |
| nemo_retriever/harness/test_configs.yaml | Only two of eight presets carry embed_inference_batch_size: 32; the six VL presets fall back to the HarnessConfig dataclass default, making the YAML incomplete as self-documentation for operators tuning those presets. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A["embed(texts, batch_size)"] --> B["_embed_local()"]
B --> C["Sort texts by length (longest first)"]
C --> D["current_bs = min(target_bs, _adaptive_batch_size)"]
D --> E{"i < len(sorted_texts)?"}
E -- No --> F["torch.cat(outs)"]
F --> G["Reorder to original indices"]
G --> H["Return CPU tensor"]
E -- Yes --> I["chunk = texts[i : i + current_bs]"]
I --> J["_embed_chunk(chunk)"]
J -- Success --> K["outs.append(result)\ni += current_bs"]
K --> L{"current_bs < target_bs?"}
L -- Yes --> M["success_streak += 1"]
M --> N{"streak >= 3?"}
N -- Yes --> O["current_bs = min(target_bs, current_bs * 2)\n_adaptive_batch_size = current_bs"]
N -- No --> E
O --> E
L -- No --> P["_adaptive_batch_size = target_bs"]
P --> E
J -- "torch.cuda.OutOfMemoryError\nor OOM-like RuntimeError" --> Q["torch.cuda.empty_cache()\nsuccess_streak = 0"]
Q --> R{"current_bs <= 1?"}
R -- Yes --> S["Re-raise"]
R -- No --> T["current_bs = current_bs // 2\n_adaptive_batch_size = current_bs\nwarn with diagnostics"]
T --> E
Prompt To Fix All With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/utils/ray_resource_hueristics.py
Line: 25
Comment:
**Stale comment after decoupling**
`EMBED_BATCH_SIZE` is now only the Ray Data transport/scheduling batch size; the model inference microbatch is controlled separately by `embed_inference_batch_size`. The trailing comment incorrectly implies this constant drives both knobs, which is the coupling this PR was designed to remove.
```suggestion
EMBED_BATCH_SIZE = 256 # Ray Data transport/scheduling batch size (inference microbatch is controlled separately)
```
How can I resolve this? If you propose a fix, please make it concise.
---
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/model/local/llama_nemotron_embed_1b_v2_embedder.py
Line: 24-34
Comment:
**`CUDNN_STATUS_INTERNAL_ERROR` is broader than OOM**
`CUDNN_STATUS_INTERNAL_ERROR` is a generic cuDNN failure covering workspace allocation failures but also unsupported operations, invalid tensor shapes, and hardware faults. Including it here means any unexpected cuDNN error (not just memory pressure) will silently trigger batch-halving retries. If the root cause is not VRAM pressure, each retry at a reduced batch size will also fail, burning through 7+ retry cycles before eventually re-raising — with misleading "retrying with batch_size=…" diagnostics. Consider tightening this to tokens that more specifically indicate allocation failures.
```suggestion
def _is_cuda_oom_like_error(exc: RuntimeError) -> bool:
msg = str(exc).upper()
return any(
token in msg
for token in (
"CUDA OUT OF MEMORY",
"CUBLAS_STATUS_ALLOC_FAILED",
"CUDNN_STATUS_ALLOC_FAILED",
)
)
```
How can I resolve this? If you propose a fix, please make it concise.
---
This is a comment left during a code review.
Path: nemo_retriever/harness/test_configs.yaml
Line: 77-118
Comment:
**`embed_inference_batch_size` missing from most presets**
Only `PE_GE_OCR_TE_DENSE` and `PE_GE_OCR_TE_HYBRID` carry the new key. The six VL presets (`PE_GE_OCR_VL_*`) rely on the `HarnessConfig` dataclass default (`32`), but an operator reading these YAML files to tune a VL run won't know that knob exists or what its effective value is. Adding `embed_inference_batch_size: 32` to all presets keeps the YAML self-documenting and makes future tuning explicit.
How can I resolve this? If you propose a fix, please make it concise.
---
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/graph_ingestor.py
Line: 172-199
Comment:
**New public constructor params missing from docstring**
`ray_log_to_driver` and `debug` were added as public constructor parameters in this PR but are not documented in the class docstring. Per the `docstrings-public-interface` standard, every public parameter should have a description entry.
Add to the class docstring Parameters section:
```
ray_log_to_driver
Forward Ray worker logs to the driver process (default True).
debug
Enable DEBUG-level logging in Ray runtime workers (default False).
```
**Rule Used:** Public modules, classes, and functions must have d... ([source](https://app.greptile.com/review/custom-context?memory=docstrings-public-interface))
How can I resolve this? If you propose a fix, please make it concise.Reviews (5): Last reviewed commit: "Merge branch 'main' into feature/embed-o..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comment on lines
+271
to
+281
| capture_path = Path(capture_path_raw).expanduser() | ||
| capture_is_dir = capture_path_raw.endswith("/") or (capture_path.exists() and capture_path.is_dir()) | ||
| if capture_is_dir: | ||
| capture_path.mkdir(parents=True, exist_ok=True) | ||
| capture_file = capture_path / f"embed_oom_outliers_pid{os.getpid()}.jsonl" | ||
| else: | ||
| capture_path.parent.mkdir(parents=True, exist_ok=True) | ||
| capture_file = capture_path | ||
|
|
||
| line = json.dumps(event, ensure_ascii=True) + "\n" | ||
| with capture_file.open("a", encoding="utf-8") as f: |
Contributor
There was a problem hiding this comment.
_capture_oom_outlier_event can abort OOM retry
_capture_oom_outlier_event is called from inside except torch.cuda.OutOfMemoryError: and except RuntimeError: blocks where the intent is to halve the batch size and retry. If capture_path.mkdir() or capture_file.open(...) raises (disk full, permission denied, race on directory creation) the new I/O exception propagates out of the except-handler, replacing the OOM and killing the embedding job instead of retrying. The whole function must be wrapped in a top-level try/except so that any failure is silently dropped rather than breaking recovery.
try:
capture_path = Path(capture_path_raw).expanduser()
capture_is_dir = capture_path_raw.endswith("/") or (capture_path.exists() and capture_path.is_dir())
if capture_is_dir:
capture_path.mkdir(parents=True, exist_ok=True)
capture_file = capture_path / f"embed_oom_outliers_pid{os.getpid()}.jsonl"
else:
capture_path.parent.mkdir(parents=True, exist_ok=True)
capture_file = capture_path
line = json.dumps(event, ensure_ascii=True) + "\n"
with capture_file.open("a", encoding="utf-8") as f:
try:
import fcntl
fcntl.flock(f.fileno(), fcntl.LOCK_EX)
f.write(line)
f.flush()
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
except Exception:
f.write(line)
f.flush()
self._oom_capture_event_count += 1
except Exception as capture_exc:
logger.debug("OOM capture write failed (non-fatal): %s", capture_exc)Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/model/local/llama_nemotron_embed_1b_v2_embedder.py
Line: 271-281
Comment:
**Unguarded I/O in `_capture_oom_outlier_event` can abort OOM retry**
`_capture_oom_outlier_event` is called from inside `except torch.cuda.OutOfMemoryError:` and `except RuntimeError:` blocks where the intent is to halve the batch size and retry. If `capture_path.mkdir()` or `capture_file.open(...)` raises (disk full, permission denied, race on directory creation) the new I/O exception propagates out of the except-handler, replacing the OOM and killing the embedding job instead of retrying. The whole function must be wrapped in a top-level `try/except` so that any failure is silently dropped rather than breaking recovery.
```python
try:
capture_path = Path(capture_path_raw).expanduser()
capture_is_dir = capture_path_raw.endswith("/") or (capture_path.exists() and capture_path.is_dir())
if capture_is_dir:
capture_path.mkdir(parents=True, exist_ok=True)
capture_file = capture_path / f"embed_oom_outliers_pid{os.getpid()}.jsonl"
else:
capture_path.parent.mkdir(parents=True, exist_ok=True)
capture_file = capture_path
line = json.dumps(event, ensure_ascii=True) + "\n"
with capture_file.open("a", encoding="utf-8") as f:
try:
import fcntl
fcntl.flock(f.fileno(), fcntl.LOCK_EX)
f.write(line)
f.flush()
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
except Exception:
f.write(line)
f.flush()
self._oom_capture_event_count += 1
except Exception as capture_exc:
logger.debug("OOM capture write failed (non-fatal): %s", capture_exc)
```
How can I resolve this? If you propose a fix, please make it concise.
Comment on lines
+83
to
+93
| try: | ||
| tokenized = tokenizer( | ||
| list(chunk), | ||
| padding=False, | ||
| truncation=True, | ||
| max_length=max(1, int(max_length)), | ||
| return_length=True, | ||
| ) | ||
| token_lengths = tokenized.get("length") | ||
| except Exception: | ||
| token_lengths = None |
Contributor
There was a problem hiding this comment.
The bare except Exception: on the tokenizer call in _batch_length_summary swallows tokenization errors without logging. Since this helper is called during OOM recovery to build diagnostic information, a silent failure means the warning message emitted to the user will have tok_max=None, tok_p95=None with no indication that tokenization itself failed. Consider adding a logger.debug(...) at minimum so OOM diagnostic information gaps are traceable.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/model/local/llama_nemotron_embed_1b_v2_embedder.py
Line: 83-93
Comment:
**Silent exception swallow in tokenization path**
The bare `except Exception:` on the tokenizer call in `_batch_length_summary` swallows tokenization errors without logging. Since this helper is called during OOM recovery to build diagnostic information, a silent failure means the warning message emitted to the user will have `tok_max=None, tok_p95=None` with no indication that tokenization itself failed. Consider adding a `logger.debug(...)` at minimum so OOM diagnostic information gaps are traceable.
How can I resolve this? If you propose a fix, please make it concise.Drop per-item metadata and JSONL OOM outlier capture paths from the local embed runtime while retaining adaptive OOM retry and decoupled inference batch controls. This keeps the PR focused on core stability behavior and reduces instrumentation overhead. Signed-off-by: jioffe502 <jioffe@nvidia.com>
|
|
||
| outs: List[torch.Tensor] = [] | ||
| target_bs = max(1, int(batch_size)) | ||
| current_bs = min(target_bs, self._adaptive_batch_size) if self._adaptive_batch_size is not None else target_bs |
Contributor
There was a problem hiding this comment.
self._adaptive_batch_size is actor-level state that persists between embed() calls. On a fresh call with a new (larger) batch_size argument, line 224 initialises current_bs = min(target_bs, self._adaptive_batch_size), which silently caps the first batch at whatever the worst historical OOM level was, even if VRAM pressure has since been relieved. The streak-based growth inside the loop never propagates back across the call boundary. Consider resetting when target_bs exceeds the stored value, or document the cross-call persistence explicitly.
current_bs = (
min(target_bs, self._adaptive_batch_size)
if self._adaptive_batch_size is not None and self._adaptive_batch_size < target_bs
else target_bs
)Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/model/local/llama_nemotron_embed_1b_v2_embedder.py
Line: 224
Comment:
**Adaptive batch size reads stale state across calls**
`self._adaptive_batch_size` is actor-level state that persists between `embed()` calls. On a fresh call with a new (larger) `batch_size` argument, line 224 initialises `current_bs = min(target_bs, self._adaptive_batch_size)`, which silently caps the first batch at whatever the worst historical OOM level was, even if VRAM pressure has since been relieved. The streak-based growth inside the loop never propagates back across the call boundary. Consider resetting when `target_bs` exceeds the stored value, or document the cross-call persistence explicitly.
```python
current_bs = (
min(target_bs, self._adaptive_batch_size)
if self._adaptive_batch_size is not None and self._adaptive_batch_size < target_bs
else target_bs
)
```
How can I resolve this? If you propose a fix, please make it concise.
Comment on lines
+284
to
+293
| if not outs: | ||
| return torch.empty((0, 0), dtype=torch.float32) | ||
|
|
||
| sorted_embeddings = torch.cat(outs, dim=0) | ||
| reordered_embeddings: List[torch.Tensor | None] = [None] * len(sorted_to_original) | ||
| for sorted_idx, original_idx in enumerate(sorted_to_original): | ||
| reordered_embeddings[original_idx] = sorted_embeddings[sorted_idx] | ||
| if any(emb is None for emb in reordered_embeddings): | ||
| raise RuntimeError("Failed to reconstruct embedding order after length sorting.") | ||
| return torch.stack([emb for emb in reordered_embeddings if emb is not None], dim=0) |
Contributor
There was a problem hiding this comment.
sorted_embeddings and sorted_to_original when OOM leaves partial results
If OOM is raised when current_bs <= 1 (re-raised, not retried), execution exits the while loop with outs containing embeddings for only the texts processed before the fatal OOM. torch.cat(outs) produces a tensor shorter than len(sorted_to_original), and the loop for sorted_idx, original_idx in enumerate(sorted_to_original) will raise IndexError when sorted_idx exceeds that length. The if any(emb is None …) guard never executes because the IndexError fires first. Currently unreachable because the fatal re-raise exits before this code, but it is a correctness landmine for future refactors. Add an explicit length pre-check before the indexing loop.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/model/local/llama_nemotron_embed_1b_v2_embedder.py
Line: 284-293
Comment:
**Length mismatch between `sorted_embeddings` and `sorted_to_original` when OOM leaves partial results**
If OOM is raised when `current_bs <= 1` (re-raised, not retried), execution exits the while loop with `outs` containing embeddings for only the texts processed before the fatal OOM. `torch.cat(outs)` produces a tensor shorter than `len(sorted_to_original)`, and the loop `for sorted_idx, original_idx in enumerate(sorted_to_original)` will raise `IndexError` when `sorted_idx` exceeds that length. The `if any(emb is None …)` guard never executes because the IndexError fires first. Currently unreachable because the fatal re-raise exits before this code, but it is a correctness landmine for future refactors. Add an explicit length pre-check before the indexing loop.
How can I resolve this? If you propose a fix, please make it concise.
Collaborator
Author
DGX bo767 sweep evidenceRan the dedicated 6-run
Observed results:
Recommendation:
|
Collaborator
Author
bo767 repeat sweep follow-up at
|
embed_inference_batch_size |
Passes | Mean PPS | PPS stdev | PPS range | Mean recall@5 |
recall@5 stdev |
OOM evidence |
|---|---|---|---|---|---|---|---|
1 |
10/10 |
152.82 |
1.38 |
151.46-156.59 |
0.8355 |
0.0006 |
none |
32 |
10/10 |
159.34 |
1.40 |
156.94-162.11 |
0.8354 |
0.0014 |
none |
64 |
10/10 |
158.64 |
1.67 |
156.68-163.34 |
0.8309 |
0.0108 |
none |
256 |
9/10 |
158.51* |
1.59* |
156.12-161.49* |
0.8124* |
0.0246* |
1 failed run, 8 OOM retries in the failed run |
* computed over the 9 successful mb=256 runs only.
Takeaways:
32is the best default candidate in this repeat sweep: highest mean PPS, no OOM retries, and stable recall.64is within observed run-to-run noise on ingest (-0.70 PPS,-0.44%vs32) and shows worse recall stability, so this does not support moving the default from32to64.1is clearly slower (-6.52 PPS,-4.09%vs32) with no stability upside.256is not a viable default candidate: one run failed, the failed run logged8CUDA OOM during embedding; retryingevents, and the successful runs still showed materially worse recall stability.- This is directionally consistent with the current HF/local embedding path not realizing reliable end-to-end gains from larger microbatches in this workload.
- The VDB tail is already de-bottlenecked (
VDBUploadOperator batch_size=64), so pushing embed microbatch higher is not buying a downstream write-path win here; in practice it mainly adds memory pressure and instability.
Recommendation / next step:
- If we want one default from the current evidence, this repeat sweep supports
embed_inference_batch_size=32atgpu_embed=0.25. - From here, I think there are two reasonable paths:
- Land this change with
32as the default and treat broader SKU validation as follow-up work. - Keep testing in this PR, but narrow it to cross-SKU validation focused on
32vs64rather than continuing to probe larger batch sizes.
- Land this change with
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
TLDR
This PR hardens embedding stability by separating two different control knobs:
embed_batch_sizefor Ray Data transport/scheduling, andembed_inference_batch_sizefor local model forward-pass VRAM pressure.It keeps the change focused on core runtime behavior: decoupled tuning controls plus adaptive OOM retry.
Problem
During bo767 validation, we observed long-sequence tail batches driving CUDA OOM-like failures and throughput volatility. The key discovery is that one batch knob is doing two jobs:
embed_batch_size), andinference_batch_size).That coupling makes stability tuning hard, especially when sequence lengths are heterogeneous.
What Changed
1) Decoupled controls in harness/CLI
embed_inference_batch_sizeto harness config/defaults/validation.--embed-inference-batch-sizethrough to runtime.embed_inference_batch_sizeoverembed_batch_sizefor model forward-pass batching.2) Adaptive OOM stabilization in local embed runtime
3) Runtime parity + operational controls
RAY_DEFAULT_OBJECT_STORE_MEMORY_PROPORTION/RAY_OBJECT_STORE_MEMORY_BYTES.Quantitative Evidence (bo767)
All runs below use
embed_inference_batch_size=32.nemo_retriever/artifacts/bo767_20260406_193158_UTC/results.jsonnemo_retriever/artifacts/bo767_20260406_202408_UTC/results.jsonnemo_retriever/artifacts/bo767_20260406_204929_UTC/results.jsonnemo_retriever/artifacts/bo767_20260408_154723_UTC/results.jsonnemo_retriever/artifacts/bo767_20260408_160637_UTC/results.jsonnemo_retriever/artifacts/bo767_20260408_172412_UTC/results.jsonnemo_retriever/artifacts/bo767_20260408_161941_UTC/results.jsonnemo_retriever/artifacts/bo767_20260408_173706_UTC/results.jsonKey Takeaways
embed_inference_batch_size=32) is stable and repeatable across runs.gpu_embedpolicy in these validations:gpu_embed=1.0cohort: ~146 PPSgpu_embed=0.25cohort: >150 PPSRecommendation
embed_inference_batch_size=32in harness presets.gpu_embedas an explicit policy knob for follow-up performance decisioning (isolation vs throughput).Why This Is Safe
32) aligns with observed stability.When To Override
embed_inference_batch_sizeonly with sustained VRAM headroom and zero OOM retry pressure.gpu_embedaccording to target objective:1.00.25) with recall guardrails.Rollout
embed_inference_batch_size=32defaults.gpu_embedperformance policy in a dedicated follow-up decision.Rollback
embed_inference_batch_sizeand/orgpu_embedvia harness CLI.Open Questions For Lead Decision
gpu_embedpolicy should define productized performance baselines?Test Plan
uv run pytest tests/test_harness_config.py tests/test_harness_run.py tests/test_resource_heuristics.pyuv run pytest tests/test_ingest_interface.py tests/test_graph_pipeline_registry.pyuv run pytest tests/test_create_local_embedder.py tests/test_multimodal_embed.py tests/test_operator_flags_and_cpu_actors.pyembed_inference_batch_size=32(artifacts listed above)