nvalchemi.distributed.trace_and_validate#

nvalchemi.distributed.trace_and_validate(model_factory, sample_batch, *, world_size=2, device='cuda:0', atol=1e-5, rtol=1e-4, auto_fix=True, max_fix_attempts=8, backend='auto', timeout_sec=120.0, watched_helper_packages=None, helper_sample_every=8, layer_diagnostic=True)[source]#

Infer a distribution spec, validate it on a single-GPU multi-process run, and (optionally) auto-fix when validation fails.

Parameters:

model_factory (Callable[[], Any]) – Callable returning a freshly-constructed wrapper. Called once in the launcher process for the reference run, and once per rank in each spawned worker. Pristine state every time — no shared module graph between processes.
sample_batch (Any) – A Batch (or compatible) carrying positions / cell / pbc on the target device. Small enough that world_size copies fit in memory at once.
world_size (int) – Virtual ranks to spawn on the same GPU. The default (2) is sufficient to flush the dispatch logic; larger values catch partition-dependent bugs but cost spawn overhead linearly.
device (str | device) – CUDA device all ranks bind to. Default "cuda:0". CPU validation is not supported by this entry point — CPU/GPU numerical drift makes it unreliable; if you need it, call the harness in test_dispatch_trace_gloo.py directly.
atol (float) – Per-output absolute tolerance. Pass criterion (per output) is abs_diff <= atol OR rel_diff <= rtol — same convention torch.testing.assert_close() uses, so extensive quantities (energy scales linearly with atom count) compare correctly across system sizes.
rtol (float) – Per-output relative tolerance. Default 1e-4 covers fp32 round-off accumulation across collective reductions on the cpu:gloo,cuda:gloo backend; tighten to e.g. 1e-5 when running NCCL or fp64.
auto_fix (bool) – When the initial inferred spec fails validation, try rule-based mutations. Disable to get a single-attempt report.
max_fix_attempts (int) – Cap on the number of distinct specs auto-fix will try.
backend (str) – "nccl", "gloo", or "auto" (NCCL when CUDA is available, else Gloo). Both correctly route over CUDA tensors; NCCL is faster.
timeout_sec (float) – Per-spawn join timeout.
watched_helper_packages (Sequence[str] | None) – Fully-qualified module paths whose top-level Python helpers get instrumented during the reference and per-rank runs. The helper_trace proxy records each call’s input / output shapes + sums; the helper_diagnosis classifier then flags helpers that look like distribution gaps (per-system reductions whose per-rank outputs sum to the reference output but aren’t declared in spec.distribution.third_party_helpers). Defaults to ("aimnet.nbops",). Pass an explicit empty tuple to disable. Unimportable packages are skipped silently.
helper_sample_every (int) – Record every Nth call after the first call per (module, function). Default 8 keeps overhead bounded for hot helpers (mol_sum runs multiple times per layer); set to 1 for exhaustive recording (debug only).
layer_diagnostic (bool)

Returns:

Carries the working (or best-guess) spec, every attempt’s diff/handler-counts, and a one-line next_action.

Return type:

TraceReport