#!/usr/bin/env python # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ Restartable OOMptimizer for distributed speechlm2 training. This script intentionally lives next to, rather than inside, ``oomptimizer.py``. The original OOMptimizer is a single-process calibration tool that relies on catching CUDA OOM exceptions in the same Python process, emptying enough state to keep going, and then continuing a binary search over synthetic batch sizes. That model is useful for one GPU and for simple DDP-style memory estimates, but it becomes unreliable once the model is truly distributed. With FSDP2, EP, and multi-node NCCL process groups, a single rank hitting CUDA OOM can leave other ranks blocked in collectives, can poison the process group, and often prevents the training process from reaching Python exception handling on every rank. In practice, trying to recover from those errors in-process is exactly the failure mode we want to avoid. The distributed OOMptimizer uses a different unit of recovery: a whole ``torchrun`` child job. A lightweight supervisor process owns the search state and launches short-lived probe jobs. Each probe instantiates the real model and optimizer from the provided config, creates synthetic batches through the model's OOMptimizer schema, runs one or more candidate batch sizes, records the observed peak CUDA memory, and exits. If a candidate succeeds, rank 0 writes a JSONL record with the batch size, bucket, peak allocated memory, peak reserved memory, and target memory. If a candidate reaches the requested memory fraction, the worker records ``memory_target`` and stops probing that session. If a candidate OOMs and records that fact, the supervisor marks it as the first bad candidate. If a child job dies without a result record, the supervisor treats the candidate as indeterminate, retries it, and refuses to use that failure as a memory bound until there is enough evidence to avoid corrupting the search. Example single-node invocation using one ``torchrun`` supervisor process that launches 8-rank child probes:: torchrun --standalone --nproc-per-node=1 scripts/speechlm2/distributed_oomptimizer.py \ --module-name nemo.collections.speechlm2.models.SALMAutomodel \ --config-path /path/to/experiment.yaml \ --buckets '[128,256,512,1024]' \ --memory-fraction 0.9 \ --nproc-per-node 8 Example four-node SLURM invocation:: srun --nodes=4 --ntasks-per-node=1 --gpus-per-node=8 \ python scripts/speechlm2/distributed_oomptimizer.py \ --module-name nemo.collections.speechlm2.models.SALMAutomodel \ --config-path /path/to/experiment.yaml \ --buckets '[128,256,512,1024]' \ --supervisor-nnodes 4 \ --rdzv-endpoint "${MASTER_ADDR}:29500" \ --nproc-per-node 8 The main control flow is: 1. The supervisor reads bucket boundaries and model config, then converts each bucket into synthetic input/output sequence lengths. SALM-style audio-locator models have their own conversion because a single token bucket represents both audio-equivalent tokens and text tokens; the ``--salm-audio-token-ratio`` option controls that split. 2. Buckets are processed from largest to smallest to preserve the memory-fragmentation behavior expected during real training. The next smaller bucket starts near the previous bucket's discovered batch size instead of starting from scratch. 3. For each bucket, the supervisor proposes one or more candidate batch sizes. Early probes expand quickly; later probes use the observed memory slope when possible, otherwise they fall back to doubling or bisection. 4. A probe session is launched with ``torchrun --max-restarts=0`` and a short process-group timeout. On a single node the supervisor uses ``--standalone``. On multiple nodes, one supervisor per node coordinates through a shared filesystem barrier and a rendezvous endpoint. 5. Probe workers run the actual model training step under the requested dtype. In distributed mode, workers reduce the maximum observed CUDA memory across ranks so the profile reflects the most memory-constrained rank. 6. The supervisor merges successful, memory-target, and explicit OOM observations into the same search state: ``max_ok`` tracks the largest usable batch size and ``min_err`` tracks the smallest known bad batch size. No-record child failures are retried instead of being interpreted as memory bounds. The search finishes when the relative gap between those bounds is below ``--threshold`` or the bounds differ by one. 7. The primary supervisor emits the same style of final ``bucket_duration_bins`` and ``bucket_batch_size`` output as the original tool, while preserving the per-probe logs for debugging. """ import importlib import json import math import os import signal import subprocess import sys import time from dataclasses import dataclass, field from datetime import timedelta from pathlib import Path import click import lightning.pytorch as pl import torch from omegaconf import OmegaConf from torch.utils.data import DataLoader, IterableDataset from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, MaskType, NeuralType from nemo.utils import logging from nemo.utils.oomptimizer import SequenceLengthResolver from nemo.utils.oomptimizer import is_2d_bucketing as _is_2d_bucketing from nemo.utils.trainer_utils import resolve_trainer_cfg class ProfilingBatchGenerator: """ ProfilingBatchGenerator is used to generate artificial mini-batches for model training and tracking the progress of batch size optimization. The high-level usage API is the following:: >>> gen = ProfilingBatchGenerator(schema) ... finished = False ... while not finished: ... batch = gen(input_seq_len, output_seq_len) ... try: ... training_step(model, batch) ... oom = False ... except torch.cuda.OutOfMemoryError: ... oom = True ... finished = gen.advance(oom) ... solution = gen.max_batch_size # The solution of the search problem. ... gen.reset() # Can re-use for other sequence lengths now. The search terminates once the difference between max working batch size and min OOM batch size divided by the latter is smaller than ``rel_gap_thresh`` that difference amounts to a single element. For example, a max working batch size is 96 and min OOM batch size is 100 indicates a gap of 0.04, which would terminate the search with threshold of 0.05. In order to generate mini-batches compatible with a given model, the generator: * accepts a ``schema`` argument in its constructor, and * accepts input/output sequence lengths in each call to generate a mini-batch. ``schema`` has the following structure:: >>> { ... "cls": tuple | MyBatchType, ... "inputs": [ ... { ... "type": NeuralType(...) | Literal["dummy"], ... "seq_length": Literal["input", "output"], ... "vocab_size": int, # optional, required only for LabelsType ... "name": str, # optional, indicates kwarg ... }, ... ..., ... ] ... } ``cls`` indicates how we should construct the mini-batch. Typically you can just use ``tuple`` for most batch schemas. However, if the model expects a specific, e.g., dataclass, you can tell ``ProfilingBatchGenerator`` to use it. The mini-batch object will be constructed using the items in ``inputs``. Each element of ``inputs`` specifies a NeMo NeuralType which needs to have a defined ``elements_type``. The supported types are ``AudioSignal``, ``LengthsType`` and ``LabelsType``. If "type" is not a NeuralType, we interpret that as a placeholder tensor that's not relevant but expected by the model/batch constructor. In addition, ``"seq_length"`` key is used to determine whether we should apply input or output sequence length to a given tensor. Optional keys: * ``vocab_size`` is required for ``LabelsType`` so that we can generate proper label values. * ``name`` is required if objects of ``cls`` have to be constructed using keyword arguments. A simple schema example for a model using audio/lengths tensor pair (unsupervised/self-supervised):: >>> { ... "cls": tuple, ... "inputs": [ ... {"type": NeuralType(("B", "T"), AudioSignal()), "seq_length": "input"}, ... {"type": NeuralType(("B"), LengthsType()), "seq_length": "input"}, ... ] ... } """ def __init__( self, schema: dict, start_batch_size: int = 32, rel_gap_thresh: float = 0.05, device: str = "cuda", float_dtype: torch.dtype = torch.float32, ): self.schema = schema self.start_batch_size = start_batch_size self.rel_gap_thresh = rel_gap_thresh self.device = device self.float_dtype = float_dtype self.reset() def __call__(self, input_seq_length: int, output_seq_length: int): B = self._current select_seq_length = {"input": input_seq_length, "output": output_seq_length} batch = [] names = [] for item in self.schema["inputs"]: nt = item["type"] if isinstance(nt, str) and nt == "constant": if isinstance(val := item["value"], str) and val == "batch": tnsr = torch.tensor([B], dtype=torch.long, device=self.device) else: tnsr = torch.tensor([val], dtype=torch.long, device=self.device) elif not isinstance(nt, NeuralType): # placeholder tnsr = torch.tensor([]) elif isinstance(nt.elements_type, AudioSignal): seq_length = select_seq_length[item["seq_length"]] tnsr = torch.randn(B, seq_length, dtype=self.float_dtype, device=self.device) elif isinstance(nt.elements_type, LengthsType): seq_length = select_seq_length[item["seq_length"]] tnsr = torch.ones(B, dtype=torch.long, device=self.device) * seq_length elif isinstance(nt.elements_type, MaskType): seq_length = select_seq_length[item["seq_length"]] tnsr = torch.ones(B, seq_length, device=self.device, dtype=torch.bool) elif isinstance(nt.elements_type, LabelsType): seq_length = select_seq_length[item["seq_length"]] tnsr = torch.randint(0, item["vocab_size"], size=(B, seq_length), device=self.device) replacement_id = int(item.get("excluded_token_replacement_id", 0)) for token_id in item.get("excluded_token_ids", []): tnsr.masked_fill_(tnsr == token_id, replacement_id) for position, token_id in item.get("forced_token_ids", {}).items(): position = int(position) if position < 0: position += seq_length if 0 <= position < seq_length: tnsr[:, position] = token_id else: raise RuntimeError("Unexpected item in oomptimizer schema: {item}") batch.append(tnsr) names.append(item.get("name")) args = [elem for name, elem in zip(names, batch) if name is None] kwargs = {name: elem for name, elem in zip(names, batch) if name is not None} if not kwargs and self.schema["cls"] == tuple: return tuple(args) return self.schema["cls"](*args, **kwargs) @property def max_batch_size(self) -> int | None: """ Return the solution of the batch size search problem. It will keep returning None until the search is done. """ if ( self._max_ok is not None and self._min_err is not None and (self.current_rel_gap <= self.rel_gap_thresh or self._min_err - self._max_ok <= 1) ): return self._max_ok return None @property def current_rel_gap(self) -> float | None: """ Return the current gap between the largest batch that works and the smallest batch that triggers OOM. The gap is defined as the batch size difference divided by the larger element. E.g., if the best found batch size is 95 and the smallest that triggers OOM is 100, the gap is 0.05. """ if self._min_err is None or self._max_ok is None: return None return (self._min_err - self._max_ok) / self._min_err def reset(self): """Reset the generator to prepare it for a new search.""" self._current = self.start_batch_size self._max_ok = None # max batch size that works self._min_err = None # min batch size that doesn't work def advance(self, oom: bool) -> bool: """ Adjusts the current batch size based on the outcome. Returns a bool indicating whether the calibration is complete. """ if self.max_batch_size is not None: return True if oom: # Training step failed with OOM. # Update the minimum known batch size that causes an error. self._min_err = min(float("inf") if self._min_err is None else self._min_err, self._current) # Training step failed on OOM if self._max_ok is None: # We haven't found a batch size that works yet, keep going 2x down. self._current = round(self._current / 2) else: # Try the middle-point between the known extremes. self._current = round((self._max_ok + self._min_err) / 2) else: # Training step successful. # Update the maximum known batch size that works. self._max_ok = max(-1 if self._max_ok is None else self._max_ok, self._current) if self._min_err is None: # We haven't found a batch size that causes an error yet, keep going 2x higher self._current *= 2 else: # Try the middle-point between the known extremes. self._current = round((self._max_ok + self._min_err) / 2) return False class FloatList(click.Option): """Support passing bucket duration bins as [1.1,2.5,5.6,...]""" name = "list[float]" def type_cast_value(self, ctx, value): if isinstance(value, list) and all(isinstance(v, float) for v in value): return value try: import ast ans = ast.literal_eval(value) if isinstance(ans[0], list): ans = [tuple(item) for item in ans] return ans except ValueError: raise click.BadParameter(value) def _parse_int_list(value: str) -> list[int]: if value.startswith("["): import ast parsed = ast.literal_eval(value) return [int(item) for item in parsed] return [int(item) for item in value.split(",") if item] class GpuMemoryMonitor: @staticmethod def count_visible_devices() -> int: visible = os.environ.get("CUDA_VISIBLE_DEVICES") if visible: return len([item for item in visible.split(",") if item.strip()]) try: result = subprocess.run( ["nvidia-smi", "--query-gpu=index", "--format=csv,noheader,nounits"], text=True, capture_output=True, check=True, ) except (OSError, subprocess.CalledProcessError): return 1 return max(1, len([line for line in result.stdout.splitlines() if line.strip()])) @staticmethod def trainer_devices_to_int(devices) -> int: if isinstance(devices, int): return devices if isinstance(devices, (list, tuple)): return len(devices) if isinstance(devices, str): if devices.isdigit(): return int(devices) if devices in ("auto", "-1"): return GpuMemoryMonitor.count_visible_devices() return 1 @staticmethod def query_memory_mib() -> list[tuple[int, int]]: result = subprocess.run( ["nvidia-smi", "--query-gpu=memory.used,memory.total", "--format=csv,noheader,nounits"], text=True, capture_output=True, check=True, ) memory = [] for line in result.stdout.splitlines(): if not line.strip(): continue used, total = line.split(",") memory.append((int(used.strip()), int(total.strip()))) return memory @staticmethod def target_memory_bytes(memory_fraction: float) -> float: gpu_memory = GpuMemoryMonitor.query_memory_mib() if not gpu_memory: raise click.ClickException("Could not query GPU memory via nvidia-smi.") return memory_fraction * min(total for _, total in gpu_memory) * 1024 * 1024 @staticmethod def wait_for_reclaim(timeout_seconds: float, tolerance_mb: int, poll_interval_seconds: float = 2.0) -> None: if timeout_seconds <= 0: return deadline = time.monotonic() + timeout_seconds while time.monotonic() < deadline: try: memory = GpuMemoryMonitor.query_memory_mib() except (OSError, subprocess.CalledProcessError): return if memory and max(used for used, _ in memory) <= tolerance_mb: return time.sleep(poll_interval_seconds) @dataclass class FileBarrier: log_dir: Path node_rank: int nnodes: int timeout_seconds: float = 300.0 run_id: str | None = None @property def barrier_root(self) -> Path: root = self.log_dir / ".supervisor_barriers" if self.run_id: root = root / self.safe_path_part(self.run_id) return root def wait(self, name: str) -> None: if self.nnodes <= 1: return barrier_dir = self.barrier_root / self.safe_path_part(name) barrier_dir.mkdir(parents=True, exist_ok=True) marker = barrier_dir / f"rank_{self.node_rank}.ready" marker.write_text(str(os.getpid())) deadline = time.monotonic() + self.timeout_seconds while time.monotonic() < deadline: if len(list(barrier_dir.glob("rank_*.ready"))) >= self.nnodes: return time.sleep(1.0) raise TimeoutError(f"Timed out in OOMptimizer supervisor barrier {name}.") @staticmethod def safe_path_part(value: str) -> str: allowed = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_." safe = "".join(ch if ch in allowed else "_" for ch in str(value)) return safe or "run" @staticmethod def wait_for_path(path: Path, timeout_seconds: float, description: str) -> None: deadline = time.monotonic() + timeout_seconds while time.monotonic() < deadline: if path.exists(): return time.sleep(1.0) raise TimeoutError(f"Timed out waiting for {description}: {path}") @dataclass class ProbeOutcome: records: list[dict] = field(default_factory=list) failed_candidate: int | None = None indeterminate_candidate: int | None = None failure_kind: str | None = None failure_summary: str | None = None log_path: Path = Path() returncode: int | None = None def to_json(self) -> dict: return { "records": self.records, "failed_candidate": self.failed_candidate, "indeterminate_candidate": self.indeterminate_candidate, "failure_kind": self.failure_kind, "failure_summary": self.failure_summary, "log_path": str(self.log_path), "returncode": self.returncode, } @classmethod def from_json(cls, data: dict) -> "ProbeOutcome": return cls( records=data["records"], failed_candidate=data["failed_candidate"], indeterminate_candidate=data.get("indeterminate_candidate"), failure_kind=data.get("failure_kind"), failure_summary=data.get("failure_summary"), log_path=Path(data["log_path"]), returncode=data["returncode"], ) @dataclass class ProbeStore: log_dir: Path probe_index: int bucket: object batch_sizes: list[int] node_rank: int nnodes: int @property def safe_bucket(self) -> str: return str(self.bucket).replace("/", "_").replace("[", "").replace("]", "").replace(",", "_") @property def result_path(self) -> Path: return self.log_dir / f"probe_{self.probe_index:04d}_bucket_{self.safe_bucket}_bs_{self.batch_sizes[0]}.jsonl" @property def log_path(self) -> Path: log_suffix = "" if self.nnodes <= 1 else f"_node{self.node_rank}" return ( self.log_dir / f"probe_{self.probe_index:04d}_bucket_{self.safe_bucket}_bs_{self.batch_sizes[0]}{log_suffix}.log" ) @property def outcome_path(self) -> Path: return ( self.log_dir / f"probe_{self.probe_index:04d}_bucket_{self.safe_bucket}_bs_{self.batch_sizes[0]}_outcome.json" ) def log_paths(self) -> list[Path]: if self.nnodes <= 1: return [self.log_path] if self.log_path.exists() else [] pattern = f"probe_{self.probe_index:04d}_bucket_{self.safe_bucket}_bs_{self.batch_sizes[0]}_node*.log" paths = sorted(self.log_dir.glob(pattern)) if self.log_path.exists() and self.log_path not in paths: paths.append(self.log_path) return paths def prepare(self) -> None: self.log_dir.mkdir(parents=True, exist_ok=True) if self.node_rank == 0: self.result_path.unlink(missing_ok=True) self.outcome_path.unlink(missing_ok=True) self.log_path.unlink(missing_ok=True) def read_records(self) -> list[dict]: return self.read_records_from_path(self.result_path) def first_unreported_candidate(self, records: list[dict]) -> int | None: reported = {int(record["batch_size"]) for record in records} for candidate in self.batch_sizes: if candidate not in reported: return candidate return None def write_outcome(self, outcome: ProbeOutcome) -> None: self.write_json_atomic(self.outcome_path, outcome.to_json()) def read_outcome(self) -> ProbeOutcome: return ProbeOutcome.from_json(self.read_json(self.outcome_path)) @staticmethod def read_records_from_path(path: Path) -> list[dict]: records = [] seen = set() if not path.exists(): return records with path.open() as f: for line in f: line = line.strip() if line: record = json.loads(line) key = json.dumps(record, sort_keys=True) if key not in seen: records.append(record) seen.add(key) return records @staticmethod def append_record_to_path(path: Path, record: dict) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("a") as f: f.write(json.dumps(record, sort_keys=True) + "\n") f.flush() os.fsync(f.fileno()) @staticmethod def read_json(path: Path) -> dict: with path.open() as f: return json.load(f) @staticmethod def write_json_atomic(path: Path, data: dict) -> None: path.parent.mkdir(parents=True, exist_ok=True) tmp_path = path.with_suffix(path.suffix + f".tmp.{os.getpid()}") with tmp_path.open("w") as f: json.dump(data, f, sort_keys=True) f.flush() os.fsync(f.fileno()) tmp_path.replace(path) @dataclass class DistributedSearchState: current: int threshold: float target_memory: float max_ok: int | None = None min_err: int | None = None ok_points: list[tuple[int, int]] = field(default_factory=list) @property def finished(self) -> bool: if self.max_ok is None or self.min_err is None: return False return (self.min_err - self.max_ok) / self.min_err <= self.threshold or self.min_err - self.max_ok <= 1 def make_plan(self) -> list[int]: current = max(1, int(self.current)) if self.min_err is not None: plan = [min(current, max(1, self.min_err - 1))] if self.ok_points: while len(plan) < 3 and plan[-1] < self.min_err - 1: next_candidate = plan[-1] + max(1, round((self.min_err - plan[-1]) / 2)) next_candidate = min(next_candidate, self.min_err - 1) if next_candidate in plan: break plan.append(next_candidate) return plan plan = [current] if len(self.ok_points) < 2: while len(plan) < 3: plan.append(plan[-1] * 2) else: predicted = self._predict_batch_for_target() plan.append(predicted if predicted is not None else plan[-1] * 2) deduped = [] for item in plan: if item not in deduped: deduped.append(item) return deduped def apply_records(self, records: list[dict]) -> list[dict]: events = [] for record in records: batch_size = int(record["batch_size"]) peak_allocated = int(record.get("peak_allocated", 0)) status = record["status"] if status == "ok": self.max_ok = max(batch_size, -1 if self.max_ok is None else self.max_ok) self.ok_points.append((batch_size, peak_allocated)) events.append( { "kind": "ok", "batch_size": batch_size, "peak_allocated": peak_allocated, } ) elif status == "memory_target": self.max_ok = max(batch_size, -1 if self.max_ok is None else self.max_ok) self.min_err = min(batch_size + 1, int(1e18) if self.min_err is None else self.min_err) self.ok_points.append((batch_size, peak_allocated)) events.append( { "kind": "memory_target", "batch_size": batch_size, "peak_allocated": peak_allocated, } ) else: self.min_err = min(batch_size, int(1e18) if self.min_err is None else self.min_err) events.append({"kind": "failed", "batch_size": batch_size, "status": status}) return events def mark_failed_candidate(self, failed_candidate: int) -> None: self.min_err = min(failed_candidate, int(1e18) if self.min_err is None else self.min_err) def record_batch_size_one_failure(self) -> bool: if self.max_ok is None and self.min_err is not None and self.min_err <= 1: self.max_ok = 0 return True return False def advance(self) -> None: if not self.finished: self.current = self._next_batch_size() def _next_batch_size(self) -> int: if self.max_ok is None: assert self.min_err is not None return max(1, self.min_err // 2) if self.min_err is not None: return self.max_ok + max(1, round((self.min_err - self.max_ok) / 2)) predicted = self._predict_batch_for_target() return predicted if predicted is not None else max(1, self.max_ok * 2) def _predict_batch_for_target(self) -> int | None: points_by_batch = {} for batch_size, peak_allocated in self.ok_points: if peak_allocated > 0: points_by_batch[int(batch_size)] = int(peak_allocated) points = sorted(points_by_batch.items()) if len(points) < 2: return None b1, p1 = points[-2] b2, p2 = points[-1] if b2 <= b1 or p2 <= p1: return None slope = (p2 - p1) / (b2 - b1) intercept = p2 - slope * b2 predicted = math.floor((self.target_memory - intercept) / slope) if predicted <= b2: return None return int(min(predicted, max(b2 + 1, b2 * 2))) @dataclass class TorchrunProbeLauncher: module_name: str config_path: str nproc_per_node: int nnodes: int node_rank: int rdzv_endpoint: str | None memory_fraction: float dtype: str ddp: bool salm_audio_token_ratio: float distributed_timeout_seconds: float probe_timeout_seconds: float log_dir: Path run_id: str def run( self, *, bucket, seq_len_in: int, seq_len_out: int, batch_sizes: list[int], probe_index: int, rdzv_id: str, ) -> ProbeOutcome: store = ProbeStore( log_dir=self.log_dir, probe_index=probe_index, bucket=bucket, batch_sizes=batch_sizes, node_rank=self.node_rank, nnodes=self.nnodes, ) store.prepare() cmd = [ *self.torchrun_launcher(), f"--nnodes={self.nnodes}", f"--nproc-per-node={self.nproc_per_node}", ] if self.nnodes <= 1: cmd.append("--standalone") else: if not self.rdzv_endpoint: raise click.ClickException("--rdzv-endpoint is required when supervisor nnodes > 1.") cmd.extend( [ f"--node-rank={self.node_rank}", "--rdzv-backend=c10d", f"--rdzv-endpoint={self.rdzv_endpoint}", f"--rdzv-id={rdzv_id}", ] ) cmd.extend( [ "--max-restarts=0", "--monitor-interval=1", str(Path(__file__).resolve()), "--module-name", self.module_name, "--config-path", self.config_path, "--memory-fraction", str(self.memory_fraction), "--dtype", self.dtype, "--salm-audio-token-ratio", str(self.salm_audio_token_ratio), "--distributed-timeout-seconds", str(self.distributed_timeout_seconds), "--probe-batch-sizes", ",".join(str(item) for item in batch_sizes), "--probe-seq-len-in", str(seq_len_in), "--probe-seq-len-out", str(seq_len_out), "--probe-result-path", str(store.result_path), "--probe-bucket", str(bucket), "--no-distributed-supervisor", ] ) cmd.append("--ddp" if self.ddp else "--no-ddp") env = self.clean_launcher_env(os.environ) env.setdefault("NCCL_CUMEM_ENABLE", "0") env.setdefault("TORCH_NCCL_ASYNC_ERROR_HANDLING", "1") env.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") barrier = FileBarrier(self.log_dir, self.node_rank, self.nnodes, run_id=self.run_id) barrier.wait(f"probe_{probe_index:04d}_start") with store.log_path.open("w") as log_f: log_f.write(f"COMMAND: {' '.join(cmd)}\n") log_f.flush() proc = subprocess.Popen( cmd, stdout=log_f, stderr=subprocess.STDOUT, text=True, env=env, preexec_fn=os.setsid, ) timed_out = False try: returncode = proc.wait(timeout=self.probe_timeout_seconds) except subprocess.TimeoutExpired: timed_out = True self.terminate_process_group(proc) returncode = proc.returncode log_f.write(f"\nOOMPTIMIZER_PROBE_TIMEOUT after {self.probe_timeout_seconds}s\n") log_f.flush() if self.nnodes <= 1 or self.node_rank == 0: records = store.read_records() failed_candidate = None indeterminate_candidate = None failure_kind = None failure_summary = None if timed_out or returncode != 0: failure_kind, failure_summary = self.classify_failure(store.log_paths(), timed_out=timed_out) if records and records[-1].get("status") not in ("ok", "memory_target"): failed_candidate = None else: candidate = store.first_unreported_candidate(records) if candidate is None and records: candidate = int(records[-1]["batch_size"]) if candidate is not None: if failure_kind == "oom": failed_candidate = candidate else: indeterminate_candidate = candidate outcome = ProbeOutcome( records=records, failed_candidate=failed_candidate, indeterminate_candidate=indeterminate_candidate, failure_kind=failure_kind, failure_summary=failure_summary, log_path=store.log_path, returncode=returncode, ) if self.nnodes > 1: store.write_outcome(outcome) else: FileBarrier.wait_for_path( store.outcome_path, self.probe_timeout_seconds + 60.0, "multi-node probe outcome" ) outcome = store.read_outcome() barrier.wait(f"probe_{probe_index:04d}_done") return outcome @staticmethod def torchrun_launcher() -> list[str]: torchrun = Path(sys.executable).with_name("torchrun") if torchrun.exists() and os.access(torchrun, os.X_OK): return [str(torchrun)] return [sys.executable, "-m", "torch.distributed.run"] @staticmethod def clean_launcher_env(env: dict[str, str]) -> dict[str, str]: env = dict(env) # Supervisors may be launched by srun with one task per node. If these rank variables leak into the torchrun # workers, Lightning prefers SLURMEnvironment over TorchElastic and sees only the supervisor task world. for name in ( "SLURM_PROCID", "SLURM_LOCALID", "SLURM_NODEID", "SLURM_NTASKS", "SLURM_TASKS_PER_NODE", "SLURM_GTIDS", "SLURM_STEP_TASKS_PER_NODE", ): env.pop(name, None) for name in ( "WORLD_SIZE", "RANK", "LOCAL_RANK", "GROUP_RANK", "ROLE_RANK", "ROLE_WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", ): env.pop(name, None) return env @staticmethod def terminate_process_group(proc: subprocess.Popen, grace_seconds: float = 10.0) -> None: try: os.killpg(proc.pid, signal.SIGTERM) except ProcessLookupError: return try: proc.wait(timeout=grace_seconds) except subprocess.TimeoutExpired: try: os.killpg(proc.pid, signal.SIGKILL) except ProcessLookupError: # The process group already exited between timeout handling and SIGKILL. # This is expected in racey shutdown paths and requires no further action. return proc.wait() @staticmethod def classify_failure(log_paths: list[Path], timed_out: bool) -> tuple[str | None, str | None]: text = "\n".join(TorchrunProbeLauncher.read_log_tail(path) for path in log_paths) if any( pattern in text for pattern in ( "CUDA out of memory", "CUDACachingAllocator", "cuFFT error: CUFFT_INTERNAL_ERROR", ) ): return "oom", TorchrunProbeLauncher.first_matching_line( text, ("CUDA out of memory", "CUDACachingAllocator") ) if "OOMPTIMIZER_PROBE_TIMEOUT" in text or timed_out: return "timeout", TorchrunProbeLauncher.first_matching_line(text, ("OOMPTIMIZER_PROBE_TIMEOUT",)) if "Watchdog caught collective operation timeout" in text or "ProcessGroupNCCL" in text: return "distributed", TorchrunProbeLauncher.first_matching_line( text, ("Watchdog caught collective operation timeout", "ProcessGroupNCCL") ) if "ChildFailedError" in text or "Traceback (most recent call last)" in text: return "child_error", TorchrunProbeLauncher.first_matching_line( text, ("ChildFailedError", "Traceback (most recent call last)") ) if text.strip(): return "child_error", text.strip().splitlines()[-1][:500] return None, None @staticmethod def read_log_tail(path: Path, max_bytes: int = 512_000) -> str: try: with path.open("rb") as f: try: f.seek(-max_bytes, os.SEEK_END) except OSError: pass return f.read().decode("utf-8", errors="replace") except OSError: return "" @staticmethod def first_matching_line(text: str, patterns: tuple[str, ...]) -> str | None: for line in text.splitlines(): if any(pattern in line for pattern in patterns): return line[:500] return None def _is_torchrun_worker() -> bool: return bool( os.environ.get("TORCHELASTIC_RUN_ID") or ("LOCAL_RANK" in os.environ and "RANK" in os.environ and "GROUP_RANK" in os.environ) ) def _supervisor_run_id() -> str: explicit = os.environ.get("OOMPTIMIZER_SUPERVISOR_RUN_ID") if explicit: return explicit slurm_id = "_".join(part for part in (os.environ.get("SLURM_JOB_ID"), os.environ.get("SLURM_STEP_ID")) if part) return slurm_id or str(os.getpid()) def _run_distributed_supervisor( *, pretrained_name: str | None, module_name: str | None, config_path: str | None, buckets, threshold: float, start_batch_size: int, ratio: float, memory_fraction: float, dtype: str, ddp: bool, salm_audio_token_ratio: float, distributed_timeout_seconds: float, nproc_per_node: int | None, supervisor_nnodes: int | None, supervisor_node_rank: int | None, rdzv_endpoint: str | None, probe_log_dir: str | None, probe_timeout_seconds: float, probe_memory_reclaim_timeout_seconds: float, probe_memory_tolerance_mb: int, max_probe_retries: int, ) -> None: assert pretrained_name is None, "--pretrained-name is not supported yet for Duplex S2S" assert config_path is not None, "--module-name requires --config-path to be specified as well." assert module_name is not None, "--config-path requires --module-name to be specified as well." cfg = OmegaConf.load(config_path) requested_devices = GpuMemoryMonitor.trainer_devices_to_int(OmegaConf.select(cfg, "trainer.devices", default=1)) nproc_per_node = int(nproc_per_node or requested_devices) if nproc_per_node <= 1: raise click.ClickException("Distributed supervisor requires nproc_per_node > 1.") supervisor_nnodes = int( supervisor_nnodes or os.environ.get("OOMPTIMIZER_SUPERVISOR_NNODES") or os.environ.get("SLURM_NNODES") or 1 ) supervisor_node_rank = int( supervisor_node_rank if supervisor_node_rank is not None else os.environ.get("OOMPTIMIZER_SUPERVISOR_NODE_RANK") or os.environ.get("SLURM_NODEID") or os.environ.get("SLURM_PROCID") or 0 ) rdzv_endpoint = rdzv_endpoint or os.environ.get("OOMPTIMIZER_RDZV_ENDPOINT") if supervisor_nnodes > 1 and not rdzv_endpoint: master_addr = os.environ.get("MASTER_ADDR") or os.environ.get("SLURM_MASTER_NODE") master_port = os.environ.get("MASTER_PORT") or os.environ.get("OOMPTIMIZER_RDZV_PORT") or "29500" if master_addr: rdzv_endpoint = f"{master_addr}:{master_port}" if supervisor_nnodes > 1 and not rdzv_endpoint: raise click.ClickException( "Multi-node distributed supervisor requires --rdzv-endpoint or OOMPTIMIZER_RDZV_ENDPOINT." ) if not 0 <= supervisor_node_rank < supervisor_nnodes: raise click.ClickException( f"Supervisor node rank must be in [0, {supervisor_nnodes}); got {supervisor_node_rank}." ) is_primary_supervisor = supervisor_node_rank == 0 length_resolver = SequenceLengthResolver( cfg=cfg, ratio=ratio, salm_audio_token_ratio=salm_audio_token_ratio, module_name=module_name, ) max_seq_lens = length_resolver.resolve_many(buckets) target_memory = GpuMemoryMonitor.target_memory_bytes(memory_fraction) log_dir = ( Path(probe_log_dir) if probe_log_dir else Path(config_path).with_suffix("").parent / (Path(config_path).stem + "_oomptimizer_probes") ) supervisor_run_id = _supervisor_run_id() if is_primary_supervisor: click.echo("Starting restartable distributed profiling.") click.echo(f"Probe logs: {log_dir}") click.echo(f"Supervisor run id: {supervisor_run_id}") click.echo( f"Using nnodes={supervisor_nnodes}, nproc_per_node={nproc_per_node}; " f"target allocated memory={target_memory / (1024 ** 3):.2f}GiB" ) launcher = TorchrunProbeLauncher( module_name=module_name, config_path=config_path, nproc_per_node=nproc_per_node, nnodes=supervisor_nnodes, node_rank=supervisor_node_rank, rdzv_endpoint=rdzv_endpoint, memory_fraction=memory_fraction, dtype=dtype, ddp=ddp, salm_audio_token_ratio=salm_audio_token_ratio, distributed_timeout_seconds=distributed_timeout_seconds, probe_timeout_seconds=probe_timeout_seconds, log_dir=log_dir, run_id=supervisor_run_id, ) profile = {} next_start = max(1, start_batch_size) probe_index = 0 indeterminate_retries: dict[tuple[str, int], int] = {} for bucket, (seq_len_in, seq_len_out) in reversed(list(zip(buckets, max_seq_lens))): if is_primary_supervisor: click.echo(f"The current sequence lengths are: input={seq_len_in} output={seq_len_out}.") search = DistributedSearchState(current=next_start, threshold=threshold, target_memory=target_memory) while not search.finished: plan = search.make_plan() if is_primary_supervisor: click.echo( f"\tProbe plan for bucket={bucket}: {plan} " f"(max_ok={search.max_ok}, min_err={search.min_err}, ok_points={len(search.ok_points)})" ) outcome = launcher.run( bucket=bucket, seq_len_in=seq_len_in, seq_len_out=seq_len_out, batch_sizes=plan, rdzv_id=f"{supervisor_run_id}_{probe_index:04d}", probe_index=probe_index, ) probe_index += 1 for event in search.apply_records(outcome.records): match event["kind"]: case "ok": batch_size = event["batch_size"] peak_allocated = event["peak_allocated"] if is_primary_supervisor: click.echo( f"\tOK batch={batch_size}; peak={peak_allocated / (1024 ** 3):.2f}GiB " f"({peak_allocated / target_memory:.1%} of target)" ) case "memory_target": batch_size = event["batch_size"] peak_allocated = event["peak_allocated"] if is_primary_supervisor: click.echo( f"\tMEMORY TARGET batch={batch_size}; peak={peak_allocated / (1024 ** 3):.2f}GiB " f"({peak_allocated / target_memory:.1%} of target)" ) case "failed": batch_size = event["batch_size"] status = event["status"] if is_primary_supervisor: click.echo(f"\tFAILED batch={batch_size}; status={status}") if outcome.indeterminate_candidate is not None: candidate = int(outcome.indeterminate_candidate) key = (str(bucket), candidate) indeterminate_retries[key] = indeterminate_retries.get(key, 0) + 1 GpuMemoryMonitor.wait_for_reclaim( probe_memory_reclaim_timeout_seconds, probe_memory_tolerance_mb, ) if indeterminate_retries[key] <= max_probe_retries: search.current = candidate if is_primary_supervisor: click.secho( f"\tINDETERMINATE batch={candidate}; " f"child_returncode={outcome.returncode}; kind={outcome.failure_kind}; " f"retry={indeterminate_retries[key]}/{max_probe_retries}; log={outcome.log_path}", fg="yellow", ) if outcome.failure_summary: click.secho(f"\t {outcome.failure_summary}", fg="yellow") continue if search.max_ok is None: raise click.ClickException( f"Probe for bucket={bucket} batch={candidate} failed {indeterminate_retries[key]} times " "without producing a usable result before any lower bound was found. " "Refusing to treat this as OOM because it would corrupt the search. " f"Last failure kind={outcome.failure_kind}; log={outcome.log_path}" ) search.mark_failed_candidate(candidate) if is_primary_supervisor: click.secho( f"\tFAILED batch={candidate}; child_returncode={outcome.returncode}; " f"kind={outcome.failure_kind}; retries_exhausted={indeterminate_retries[key]}; " f"log={outcome.log_path}", fg="yellow", ) if outcome.failed_candidate is not None: search.mark_failed_candidate(outcome.failed_candidate) if is_primary_supervisor: click.echo( f"\tFAILED batch={outcome.failed_candidate}; " f"child_returncode={outcome.returncode}; log={outcome.log_path}" ) GpuMemoryMonitor.wait_for_reclaim( probe_memory_reclaim_timeout_seconds, probe_memory_tolerance_mb, ) if search.record_batch_size_one_failure(): if is_primary_supervisor: click.secho( f"\tBatch size 1 failed for bucket={bucket}; recording max_batch_size=0 and continuing.", fg="yellow", ) search.advance() if is_primary_supervisor: click.secho( f"=> Optimal setting for bucket={bucket} (input={seq_len_in} output={seq_len_out}) " f"is max_batch_size={search.max_ok}", fg="green", ) profile[(bucket, seq_len_in, seq_len_out)] = search.max_ok next_start = max(search.max_ok + 1, int(math.ceil(search.max_ok * 1.5))) if is_primary_supervisor: _emit_profile(profile, buckets, memory_fraction, ddp, dtype) def _emit_profile(profile: dict, buckets, memory_fraction: float, ddp: bool, dtype: str) -> None: profile = dict(reversed(list(profile.items()))) click.echo("The 1st stage profile is:") for (bucket, seq_len_in, seq_len_out), bs in profile.items(): click.echo(f"Bucket={bucket} (input={seq_len_in} output={seq_len_out}) => max_batch_size={bs}") if _is_2d_bucketing(buckets): final_profile = [["[" + ",".join(map(str, b)) + "]", bs] for (b, _, __), bs in profile.items()] else: click.echo("Bucket merging stage...") final_profile = [] for idx, ((bucket, seq_len_in, seq_len_out), bs) in enumerate(profile.items()): if idx == 0: final_profile.append([bucket, bs]) continue if bs == final_profile[-1][1]: click.echo(f"Merging bucket {idx} with bucket {idx-1} due to identical batch sizes.") final_profile[-1][0] = bucket continue final_profile.append([bucket, bs]) click.secho(f"The profile was created with the following settings:") click.secho(f"* using {memory_fraction:.1%} of available GPU RAM.") click.secho(f"* {'' if ddp else 'not '}simulating DDP memory overhead.") click.secho(f"* using AMP with dtype={dtype}.") click.secho("The final profile is:", bold=True) click.secho("\tbucket_duration_bins=[" + ",".join(str(seqlen) for seqlen, bs in final_profile) + "]", bold=True) click.secho("\tbucket_batch_size=[" + ",".join(str(bs) for seqlen, bs in final_profile) + "]", bold=True) def _is_oom_like(error: RuntimeError) -> bool: error_msg = str(error) return ( "cuFFT error: CUFFT_INTERNAL_ERROR" in error_msg or "CUDA out of memory" in error_msg or "CUDACachingAllocator" in error_msg ) @click.command(context_settings={'show_default': True}) @click.option( "-n", "--pretrained-name", type=str, default=None, help="Name of a pretrained model to use, e.g. 'nvidia/canary-1b'.", ) @click.option( "-m", "--module-name", type=str, default=None, help="Full path to NeMo's module corresponding to CONFIG_PATH, e.g. 'nemo.collections.asr.models.EncDecMultiTaskModel'.", ) @click.option( "-c", "--config-path", type=str, default=None, help="Path to the training configuration file for MODULE_NAME." ) @click.option( "-b", "--buckets", cls=FloatList, default=[5.0, 10.0, 15.0, 20.0, 25.0, 30.0], help="List of upper-bound bucket bins (i.e. first bucket is [0.0 - item0), second bucket is [item0 - item1), etc.). " "We also support a nested list for 2D bucketing, e.g. [[2.0, 10],[2.0,20],[4.5,15],[4.5,30],...], " "where each item is a pair of (max_input_seq_len, max_output_seq_len) for a given bucket.", ) @click.option( "-t", "--threshold", type=float, default=0.05, help="Search stopping criterion in range [0, 1], lower is more precise. Interpret as the uncerainty gap, i.e. (min_oom_batch_size - max_ok_batch_size) / min_oom_batch_size.", ) @click.option("-s", "--start-batch-size", type=int, default=32, help="Initial batch size to start the search from.") @click.option( "-r", "--ratio", type=float, default=12, # conservative estimate towards longer transcripts help="The output_sequence_length to input_sequence_length ratio for the purpose of determing the maximum output sequence lengths. " "The interpretation depends on input and output modalities. Examples: for audio->text it's tokens per second. " "For text->audio it's seconds per token. For audio->audio it's output seconds per input second. " "For text->text it's output tokens per input token. " "In general larger ratio means longer output sequences and increased memory consumption. " "The default value is set adequately for automatic speech recognition. " "This argument is ignored when 2D buckets are provided to --buckets option.", ) @click.option( "-f", "--memory-fraction", type=float, default=0.9, help="Limits the use of CUDA memory for this process to MEMORY_FRACTION of the total device memory. " "By default we force 5% memory to be unused to account for non-training-loop related CUDA memory usage" "in actual training scripts.", ) @click.option( "-y", "--dtype", default="bfloat16", help="Float precision to use for computation (used together with autocast).", ) @click.option( "--ddp/--no-ddp", type=bool, default=True, help="Whether we should simulate DDP GPU RAM usage. Stores an extra copy of the model in GPU memory. Enabled by default.", ) @click.option( "--salm-audio-token-ratio", type=float, default=0.75, help="For SALM-style 1D token buckets, fraction of the bucket represented by audio-equivalent tokens.", ) @click.option( "--distributed-timeout-seconds", type=float, default=15.0, help="Process-group timeout used for distributed profiling so collective failures surface quickly.", ) @click.option( "--distributed-supervisor/--no-distributed-supervisor", type=bool, default=True, help="Use restartable torchrun child probes for multi-GPU configs instead of in-process OOM recovery.", ) @click.option( "--nproc-per-node", type=int, default=None, help="Number of local workers used by the distributed supervisor. Defaults to trainer.devices.", ) @click.option( "--supervisor-nnodes", type=int, default=None, help="Number of nodes coordinated by the distributed supervisor. Defaults to OOMPTIMIZER_SUPERVISOR_NNODES or SLURM_NNODES.", ) @click.option( "--supervisor-node-rank", type=int, default=None, help="Node rank for the distributed supervisor. Defaults to OOMPTIMIZER_SUPERVISOR_NODE_RANK, SLURM_NODEID, or SLURM_PROCID.", ) @click.option( "--rdzv-endpoint", type=str, default=None, help="Torchrun rendezvous endpoint for multi-node supervisor probes. Defaults to OOMPTIMIZER_RDZV_ENDPOINT.", ) @click.option( "--probe-log-dir", type=str, default=None, help="Directory where distributed supervisor probe logs and JSONL results are written.", ) @click.option( "--probe-timeout-seconds", type=float, default=900.0, help="Wall-clock timeout for one distributed probe session.", ) @click.option( "--probe-memory-reclaim-timeout-seconds", type=float, default=60.0, help="How long the supervisor waits for GPU memory to be reclaimed after a failed child probe.", ) @click.option( "--probe-memory-tolerance-mb", type=int, default=1024, help="GPU memory threshold used by the supervisor reclaim wait.", ) @click.option( "--max-probe-retries", type=int, default=2, help="Number of retries for child probes that fail without an explicit OOM or memory result.", ) @click.option("--probe-batch-sizes", type=str, default=None, hidden=True) @click.option("--probe-seq-len-in", type=int, default=None, hidden=True) @click.option("--probe-seq-len-out", type=int, default=None, hidden=True) @click.option("--probe-result-path", type=str, default=None, hidden=True) @click.option("--probe-bucket", type=str, default=None, hidden=True) def oomptimizer( pretrained_name: str | None, module_name: str | None, config_path: str | None, buckets: list[float], threshold: float, start_batch_size: int, ratio: float, memory_fraction: float, dtype: str, ddp: bool, salm_audio_token_ratio: float, distributed_timeout_seconds: float, distributed_supervisor: bool, nproc_per_node: int | None, supervisor_nnodes: int | None, supervisor_node_rank: int | None, rdzv_endpoint: str | None, probe_log_dir: str | None, probe_timeout_seconds: float, probe_memory_reclaim_timeout_seconds: float, probe_memory_tolerance_mb: int, max_probe_retries: int, probe_batch_sizes: str | None, probe_seq_len_in: int | None, probe_seq_len_out: int | None, probe_result_path: str | None, probe_bucket: str | None, ): """ OOMptimizer finds the optimal batch sizes for training your model with bucketing dataloading. It performs a search over batch sizes until it converges by measuring the GPU memory usage for a model's training step and optimizer update. \b There are two main usage patterns: for using a pretrained model or an untrained model configuration. The latter is more flexible but requires the user to provide two separate arguments. Examples: * python oomptimizer.py --pretrained-name nvidia/canary-1b * python oomptimizer.py --module-name nemo.collections.asr.models.EncDecMultiTaskModel \ --config-path examples/asr/conf/speech_multitask/fast-conformer_aed.yaml Dynamic bucketing is notoriously difficult to tune as you risk running into CUDA OOM many steps into the training. In order to simplify finding the optimal settings, OOMptimizer scans each bucket to find the maximum possible batch size that doesn't trigger a CUDA OOM. \b The suggested workflow is the following: 1) Run scripts/speech_recognition/estimate_duration_bins.py to get the duration distribution of your data. (consider running estimate_duration_bins_2d.py for models with a strong dependency on output sequence length such as attention-encoder-decoder models). 2) Run OOMptimizer to find the optimal batch sizes for your specific model, optimizer, and GPU. 3) Use these optimal settings in your actual training script and enjoy optimal GPU utilization OOM-free. In the unlikely event that OOMptimizer bucket batch sizes are still leading to OOMs, please try a lower setting of the MEMORY_FRACTION option, e.g. 0.75 (75% of GPU memory). This may be required in very complex setups where there are additional GPU RAM loads that can't be anticipated through the combination of training_step and optimizer update. """ assert pretrained_name is None, "--pretrained-name is not supported yet for Duplex S2S" if all(opt is None for opt in (pretrained_name, module_name, config_path)): click.secho( "You need to provide either PRETRAINED_NAME or the pair of MODULE_NAME and CONFIG_PATH.", fg="yellow" ) sys.exit(1) world_size = int(os.environ.get("WORLD_SIZE", "1")) is_outer_torchrun_worker = _is_torchrun_worker() if ( probe_batch_sizes is None and distributed_supervisor and (not is_outer_torchrun_worker or world_size == 1) and config_path is not None ): cfg_for_supervisor = OmegaConf.load(config_path) requested_devices = GpuMemoryMonitor.trainer_devices_to_int( OmegaConf.select(cfg_for_supervisor, "trainer.devices", default=1) ) requested_devices = int(nproc_per_node or requested_devices) if requested_devices > 1: _run_distributed_supervisor( pretrained_name=pretrained_name, module_name=module_name, config_path=config_path, buckets=buckets, threshold=threshold, start_batch_size=start_batch_size, ratio=ratio, memory_fraction=memory_fraction, dtype=dtype, ddp=ddp, salm_audio_token_ratio=salm_audio_token_ratio, distributed_timeout_seconds=distributed_timeout_seconds, nproc_per_node=nproc_per_node, supervisor_nnodes=supervisor_nnodes, supervisor_node_rank=supervisor_node_rank, rdzv_endpoint=rdzv_endpoint, probe_log_dir=probe_log_dir, probe_timeout_seconds=probe_timeout_seconds, probe_memory_reclaim_timeout_seconds=probe_memory_reclaim_timeout_seconds, probe_memory_tolerance_mb=probe_memory_tolerance_mb, max_probe_retries=max_probe_retries, ) return logging.setLevel(logging.CRITICAL) local_rank = int(os.environ.get("LOCAL_RANK", "0")) distributed = world_size > 1 torch.cuda.set_device(local_rank) device = torch.device(f"cuda:{local_rank}") dtype = getattr(torch, dtype) # Distributed profiling stops on allocated memory. Leave extra reservation headroom for FSDP all-gathers and # allocator cache so the artificial cap does not reject candidates before the target is reached. memory_cap = memory_fraction if not distributed else min(0.99, memory_fraction + 0.10) torch.cuda.set_per_process_memory_fraction(memory_cap, device) if distributed: torch.distributed.init_process_group(backend="nccl", timeout=timedelta(seconds=distributed_timeout_seconds)) torch.set_float32_matmul_precision("medium") torch.backends.cudnn.allow_tf32 = True assert config_path is not None, "--module-name requires --config-path to be specified as well." assert module_name is not None, "--config-path requires --module-name to be specified as well." cfg = OmegaConf.load(config_path) namespace, name = module_name.rsplit('.', maxsplit=1) model_cls = getattr(importlib.import_module(namespace), name) trainer_cfg = resolve_trainer_cfg(cfg.trainer) if not distributed: trainer_cfg = {**trainer_cfg, "devices": 1, "num_nodes": 1} trainer_cfg.pop("strategy", None) trainer = pl.Trainer( **{ **trainer_cfg, "max_steps": 1, "max_epochs": 1, "limit_val_batches": 0.0, "val_check_interval": 0.0, } ) with trainer.init_module(): model = model_cls(OmegaConf.to_container(cfg.model, resolve=True)) model = model.to(device) if not hasattr(model, "oomptimizer_schema"): click.secho( f"We read model of type {type(model)} which doesn't seem to support OOMptimizer " f"(we could not find the property .oomptimizer_schema).", fg="red", ) sys.exit(1) schema = model.oomptimizer_schema length_resolver = SequenceLengthResolver( cfg=cfg, ratio=ratio, salm_audio_token_ratio=salm_audio_token_ratio, module_name=module_name, model=model, schema=schema, ) click.echo("Starting profiling.") max_seq_lens = length_resolver.resolve_many(buckets) target_memory = memory_fraction * torch.cuda.get_device_properties(device).total_memory profile_by_memory = distributed gen = ProfilingBatchGenerator( schema=schema, start_batch_size=start_batch_size, rel_gap_thresh=threshold, device=device, float_dtype=dtype ) profile = {} class _GenDataset(IterableDataset): def __iter__(self): gen.reset() gen._current = 1 yield gen(*length_resolver.resolve_one(33)) gen.reset() def __len__(self): return 1 # initialize everything PTL needs trainer.fit(model, DataLoader(_GenDataset(), batch_size=None)) model = model.to(device) optimizer = model.configure_optimizers()["optimizer"] model.log = lambda *args, **kwargs: None # no logging if probe_batch_sizes is not None: if probe_seq_len_in is None or probe_seq_len_out is None or probe_result_path is None: raise click.ClickException("--probe-batch-sizes requires probe sequence lengths and result path.") ProbeWorkerRunner( gen=gen, model=model, optimizer=optimizer, seq_len_in=probe_seq_len_in, seq_len_out=probe_seq_len_out, batch_sizes=_parse_int_list(probe_batch_sizes), result_path=Path(probe_result_path), target_memory=target_memory, bucket=probe_bucket, distributed=distributed, device=device, ).run() return # Iterate buckets from the largest to the smallest sequences. This usually ends up creating # a tiny bit smaller batches, likely due to worse memory fragmentation. with torch.autocast("cuda", dtype=None, enabled=False): for bucket, (seq_len_in, seq_len_out) in reversed(list(zip(buckets, max_seq_lens))): click.echo(f"The current sequence lengths are: input={seq_len_in} output={seq_len_out}.") gen.reset() batch_idx = 0 def step(): click.echo( f"\t[BEGIN step] [CUDA RAM CURRENT: {torch.cuda.memory_allocated() / (1024 * 1024):.1f}MB] [CUDA RAM MAX: {torch.cuda.max_memory_allocated() / (1024*1024):.1f}MB]" ) batch = gen(seq_len_in, seq_len_out) oom = False peak_allocated = 0 status = "OK" try: click.echo(f"\tCurrent gap: {gen.current_rel_gap}... ", nl=False) optimizer.zero_grad() out = model.training_step(batch, batch_idx) out['loss'].sum().backward() optimizer.step() peak_allocated = torch.cuda.max_memory_allocated() except torch.cuda.OutOfMemoryError as e: oom = True status = "OOM!" except RuntimeError as e: error_msg = str(e) oom_like = ( "cuFFT error: CUFFT_INTERNAL_ERROR" in error_msg or "CUDA out of memory" in error_msg or "CUDACachingAllocator" in error_msg ) if not oom_like: raise oom = True status = "OOM!" else: status = "OK!" finally: if distributed: oom_t = torch.tensor([int(oom)], dtype=torch.int32, device=device) try: torch.distributed.all_reduce(oom_t, op=torch.distributed.ReduceOp.MAX) oom = bool(oom_t.item()) except RuntimeError: oom = True if not oom and profile_by_memory: peak_t = torch.tensor([peak_allocated], dtype=torch.float64, device=device) torch.distributed.all_reduce(peak_t, op=torch.distributed.ReduceOp.MAX) peak_allocated = int(peak_t.item()) if peak_allocated >= target_memory: oom = True status = f"MEMORY TARGET ({peak_allocated / (1024 * 1024):.1f}MB)!" elif oom: status = "OOM!" click.secho(status, fg="yellow" if oom else "green") click.echo( f"\t[END step] [CUDA RAM CURRENT: {torch.cuda.memory_allocated() / (1024 * 1024):.1f}MB] [CUDA RAM MAX: {torch.cuda.max_memory_allocated() / (1024*1024):.1f}MB]" ) del batch if oom: optimizer.zero_grad(set_to_none=True) torch.cuda.empty_cache() # Note: We could call empty_cache() to free up some more memory on the GPU, # but we have found out empirically that this causes a mismatched condition # between OOMptimizer and the actual training. During training, there is some # degree of memory fragmentation and it's better to simulate that in OOMptimizer. # torch.cuda.memory.empty_cache() torch.cuda.reset_peak_memory_stats() return oom oom = step() while not (finished := gen.advance(oom)): click.echo("\t" + "=" * 80) oom = step() click.secho( f"=> Optimal setting for bucket={bucket} (input={seq_len_in} output={seq_len_out}) is max_batch_size={gen.max_batch_size}", fg="green", ) profile[(bucket, seq_len_in, seq_len_out)] = gen.max_batch_size gen.start_batch_size = gen.max_batch_size * 2 _emit_profile(profile, buckets, memory_fraction, ddp, dtype) @dataclass class ProbeWorkerRunner: gen: ProfilingBatchGenerator model: pl.LightningModule optimizer: torch.optim.Optimizer seq_len_in: int seq_len_out: int batch_sizes: list[int] result_path: Path target_memory: float bucket: str | None distributed: bool device: torch.device def run(self) -> None: global_rank = int(os.environ.get("RANK", "0")) batch_idx = 0 with torch.autocast("cuda", dtype=None, enabled=False): for batch_size in self.batch_sizes: click.echo( f"OOMPTIMIZER_PROBE bucket={self.bucket} batch_size={batch_size} " f"input={self.seq_len_in} output={self.seq_len_out}" ) self.gen.reset() self.gen._current = batch_size torch.cuda.reset_peak_memory_stats() batch = None try: self.optimizer.zero_grad() batch = self.gen(self.seq_len_in, self.seq_len_out) out = self.model.training_step(batch, batch_idx) out['loss'].sum().backward() self.optimizer.step() torch.cuda.synchronize(self.device) peak_allocated = torch.cuda.max_memory_allocated() peak_reserved = torch.cuda.max_memory_reserved() except torch.cuda.OutOfMemoryError as e: click.echo(f"OOMPTIMIZER_PROBE_OOM batch_size={batch_size}: {e}") self._record_oom(global_rank, batch_size, e) os._exit(42) except RuntimeError as e: if not _is_oom_like(e): raise click.echo(f"OOMPTIMIZER_PROBE_OOM_LIKE batch_size={batch_size}: {e}") self._record_oom(global_rank, batch_size, e) os._exit(43) finally: if batch is not None: del batch if self.distributed: try: peak_t = torch.tensor([peak_allocated, peak_reserved], dtype=torch.float64, device=self.device) torch.distributed.all_reduce(peak_t, op=torch.distributed.ReduceOp.MAX) peak_allocated = int(peak_t[0].item()) peak_reserved = int(peak_t[1].item()) except RuntimeError as e: click.echo(f"OOMPTIMIZER_PROBE_COLLECTIVE_FAILED batch_size={batch_size}: {e}") os._exit(44) status = "memory_target" if peak_allocated >= self.target_memory else "ok" if global_rank == 0: ProbeStore.append_record_to_path( self.result_path, { "batch_size": batch_size, "bucket": self.bucket, "status": status, "peak_allocated": peak_allocated, "peak_reserved": peak_reserved, "target_memory": self.target_memory, }, ) click.echo( f"OOMPTIMIZER_PROBE_RESULT batch_size={batch_size} status={status} " f"peak_allocated={peak_allocated / (1024 ** 3):.2f}GiB" ) if status == "memory_target": break def _record_oom(self, global_rank: int, batch_size: int, error: RuntimeError) -> None: if global_rank == 0: ProbeStore.append_record_to_path( self.result_path, { "batch_size": batch_size, "bucket": self.bucket, "status": "oom", "message": str(error), }, ) if __name__ == "__main__": oomptimizer()