nvalchemi.distributed.DomainConfig

class nvalchemi.distributed.DomainConfig(*, cutoff, skin=0.0, ghost_width=None, strategy=StrategyKind.HALO, compile=False, require_nondegenerate=False, mesh=None, mesh_dim='domain', grid_dims=None, scripted_marshal='auto', scripted_marshal_exclude=(), migration_hysteresis=None)

Configuration for one spatial domain-decomposition scope.

Parameters:

• cutoff (float) – Interaction cutoff radius used by the model.

• skin (float) – Additional skin distance for neighbor-list buffering. Default 0.

• ghost_width (float | None) – Width of the ghost (halo) region. When None, the effective width defaults to cutoff + skin via effective_ghost_width(). An explicit value must be at least cutoff + skin.

• strategy (StrategyKind) – Parallelization strategy for this scope: StrategyKind.HALO (spatial domain decomposition with a ghost halo, the default) or StrategyKind.GRAPH_PARTITION (node-partition graph parallel).

• compile (bool) – Compile intent for the distributed forward. When True the framework owns the compiled forward and pads per-rank atom/edge counts to stable shapes so the compiled graph is reused across steps; when False the padder is disabled. Default False.

• require_nondegenerate (bool) – When True, a degenerate partition — one where some rank’s halo already covers every atom (0 remote atoms) — is a hard error instead of a warning. Default False.

• mesh (DeviceMesh | None) – Optional torch.distributed.device_mesh.DeviceMesh describing the rank topology. None for single-rank runs.

• mesh_dim (str) – Name of the mesh dimension used for domain parallelism. Default "domain".

• grid_dims (tuple[int, int, int] | None) – Explicit grid dimensions for the spatial decomposition. When None, the partitioner chooses cells-per-dim from the cell matrix and cutoff.

• scripted_marshal ({"auto", "declared", "off"}) – Controls marshalling of @torch.jit.script ops across the ShardTensor boundary (a scripted kernel reading a ShardTensor’s storage-less data_ptr triggers a CUDA illegal memory access). "auto" (default): auto-discover scripted submodules and wrap them, plus install the spec’s declared JitAdapter marshallers. "declared": install only the spec’s declared adapters, no auto-discovery. "off": no marshalling at all. Overridable via NVALCHEMI_SCRIPTED_MARSHAL.

• scripted_marshal_exclude (tuple[str, ...]) – Submodule-name substrings to skip during "auto" discovery — for a scripted op that genuinely needs cross-rank data (where marshalling to local would silently give wrong numbers) or is handled via custom_ops.

• migration_hysteresis (float | None) – Migration-hysteresis margin in angstrom: an atom keeps its current owner until it is this far past a domain boundary, preventing per-step migration thrashing of boundary atoms. When None (default), the effective value is skin / 2 (see effective_migration_hysteresis()); it must be < skin so a deferred atom stays within the owner’s halo.

effective_ghost_width()

Return the ghost region width, defaulting to cutoff + skin.

Return type:

float

effective_migration_hysteresis()

Migration-hysteresis margin (angstrom). Defaults to skin/2.

An atom keeps its current owner until it is this far past a domain boundary, preventing per-step migration thrashing of boundary atoms. Must be < skin so a deferred atom stays within the owner’s halo (ghost_width = cutoff + skin).

Return type:

float

model_config = {'arbitrary_types_allowed': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].