nvalchemi.distributed.OpAdapter

class nvalchemi.distributed.OpAdapter(op, arg_transforms=None, output_transforms=None, *, gather_inputs=(), neighbors_inputs=(), gather_inputs_full=(), owned_slice_inputs=(), scatter_outputs=(), all_reduce_outputs=(), slice_outputs_owned=(), install_site='')

Adapt a @torch.library.custom_op / @torch.library.triton_op kernel to ShardTensor-aware dispatch.

Two dicts declare per-position pre/post transformations:

arg_transforms (input position → ArgTransform)

GatherInputs — halo-pad owned input to (n_padded, *F) before kernel. GatherInputsFull — sharded analogue: full-gather to (n_global + 1, *F). SliceOwned — slice halo-padded input to (n_owned, *F).

output_transforms (output position → OutputTransform)

ScatterOutputs — halo_reverse_exchange + halo_forward_exchange after kernel. AllReduceSum — cross-rank SUM (autograd-symmetric). SliceOutputsOwned — slice (n_global + 1, *F) back to (n_owned + 1, *F).

The positional-role form names which I/O plays which cross-rank role by position, with no transform objects to import:

OpAdapter(torch.ops.ns.fused_op, scatter_outputs=[0])      # node scatter
OpAdapter(op, gather_inputs=[0], scatter_outputs=[0])      # neighbor read + scatter

The role keywords lower onto the same two transform dicts, so the dict form above stays valid and serialization is unchanged.

See the module docstring for the full worked example.

Parameters:

• op (Any)

• arg_transforms (dict[int, GatherInputs | GatherInputsFull | SliceOwned])

• output_transforms (dict[int, ScatterOutputs | AllReduceSum | SliceOutputsOwned])

• gather_inputs (tuple[int, ...])

• neighbors_inputs (tuple[int, ...])

• gather_inputs_full (tuple[int, ...])

• owned_slice_inputs (tuple[int, ...])

• scatter_outputs (tuple[int, ...])

• all_reduce_outputs (tuple[int, ...])

• slice_outputs_owned (tuple[int, ...])

• install_site (str)

property all_reduce_replicated_outputs: tuple[int, ...]

Subset of all_reduce_outputs whose reduced value feeds a replicated whole-system computation, so the adjoint passes the incoming gradient through instead of reducing it again.

describe(state='pending', error=None)

Return an AdapterStatus snapshot of this adapter.

Parameters:

• state (Literal['pending', 'installed', 'restored', 'failed'])

• error (str | None)

Return type:

AdapterStatus

classmethod from_dict(d)

Reconstruct an OpAdapter from to_dict() output.

Parameters:

d (dict[str, Any])

Return type:

OpAdapter

install()

Register a ShardTensor-aware dispatch handler on the op (and its overload packet, if any). Returns a memento that restore() consumes to clear the registration. A lazily-named op is resolved to its live handle here (raises if the declaring module isn’t imported).

Return type:

dict[str, Any]

restore(memento)

Clear the handler registered by install().

Parameters:

memento (dict[str, Any])

Return type:

None

to_dict()

Serialize to a JSON-roundtrippable dict.

Return type:

dict[str, Any]