Large contractions#
See also tensor contractions for
the basic building blocks (yastn.ncon(), yastn.tensordot(), …).
The module yastn.tensor.oe_blocksparse provides a sliced (also called
unrolled) block-sparse contraction engine on top of yastn.ncon(). It is
aimed at contractions whose single-shot intermediate tensors are too large to
fit in memory, and/or that should be spread across several GPUs. The companion
module yastn.tensor._oe_blocksparse_mp implements the multi-device
(multiprocess) dispatch of the same computation.
Motivation#
A tensor-network contraction is specified in opt_einsum’s interleaved format with an explicit output index group:
contract_with_unroll(T1, ig1, T2, ig2, ..., out_ig, unroll=..., optimize=...)
Two knobs turn this from a plain ncon into a large-contraction engine:
Slicing / unrolling. One or more index labels are unrolled: their leg is partitioned into non-overlapping
SlicedLegpieces. The contraction is then evaluated as a loop over the Cartesian product of these pieces (combos). Each iteration masks the relevant tensors down to a single slice, so every intermediate produced bynconis strictly smaller than the full-leg intermediate. This trades extra compute (the loop) for a lower peak memory footprint, and exposes the loop as an axis of parallelism.Multi-device dispatch. The independent combos of the unroll loop can be distributed round-robin across a pool of persistent worker processes pinned to different devices (e.g. several GPUs), with autograd supported through a checkpoint-style custom
torch.autograd.Function.
Both contracted labels (summed away) and output labels (kept in the result)
may be unrolled. Partials from contracted-only combos are summed with +;
partials that differ along an output-unrolled axis are grouped by their position
along that axis and reassembled into the final tensor with yastn.block().
Slicing a leg#
A SlicedLeg selects a subset of a leg’s charge sectors, each
optionally restricted to a contiguous slice inside that sector’s block
dimension. A list of non-overlapping SlicedLeg objects whose
union covers the whole leg is a valid partition; iterating over it and summing
the partial contractions reproduces the full result.
The unroll argument maps index labels to such partitions:
import yastn
# partition the leg carrying label 'k' into charge-sector-sized pieces
unroll = {'k': yastn.make_sliced_legs(A.get_legs(axis_of_k))}
# or let YASTN slice a leg uniformly into segments of at most `size`
# simply by passing an int (resolved against the network's legs):
unroll = {'k': 256} # equivalent to slice_leg_uniform(leg, 256)
Two partitioning helpers are provided:
yastn.make_sliced_legs()— oneSlicedLegper charge sector (the finest partition; the simplest, and the one an integer maps to whensizeexceeds every sector).yastn.tensor.oe_blocksparse.slice_leg_uniform()— contiguous segments of at mostsizeelements; a segment may span several charge sectors.
Entry-point API#
- yastn.get_contraction_path(*tn_to_contract, unroll=None, names: Sequence[str] = None, who: str = None, **kwargs) tuple[Sequence[tuple[int]], PathInfo][source]#
Returns optimal contraction path for tensor network contraction specified in interleaved format. Takes into account unrolled indices if any.
- Parameters:
tn_to_contract – input to einsum in interleaved format. Explicit index labeling of output is required
unroll – Mapping[Hashable,Union[Sequence[SlicedLeg],int]] indices to unroll
names – string labels for tensors used for more readable logging. The order of names has to follow order of tensors as they appear in
tn_to_contractwho – string id for logging identifying this optimal contraction path search
optimizer – str or
opt_einsum.paths.PathOptimizer, optional The optimizer to use for contraction path search. See https://optimized-einsum.readthedocs.io/en/stable/optimal_path.html for details. Options are: - None (or ‘default’, ‘dp’, ‘dynamic-programming’): use the defaultDynamicProgrammingoptimizer - user-providedPathOptimizerinstanceoptimizer_kwargs – dict, optional Additional keyword arguments to pass to the optimizer. For the default
DynamicProgrammingoptimizer, you can specifyminimize,search_outer, andcost_cap. Whereminimizecan be - ‘write’ (default) total amount of data written to memory - ‘flops’ total number of floating point operations - ‘size’ maximum size of any intermediate tensor
- Returns:
path (Sequence[tuple[int]]) – Optimal contraction path as a sequence of tuples specifying which pair of tensors to contract at each step. The path is in terms of positions in current list of tensors, which is shrinking at each step.
path_info (opt_einsum.contract.PathInfo) – Detailed information about the contraction path, including shapes and memory usage of intermediate tensors.
- yastn.contract_with_unroll(*args, **kwargs)[source]#
Extension of opt_einsum’s contract allowing for index unrolling and use of checkpointing over unrolled loop.
- Parameters:
args – input to einsum in interleaved format. Explicit index labeling of output is required
unroll – Mapping[Hashable,Sequence[SlicedLeg]] or None indices to unroll
optimize – contraction path. Optional — if omitted, computed internally via
get_contraction_path()(with the sameunrolland any path-search kwargs inkwargs).checkpoint_loop – if True, each unrolled loop iteration is wrapped in
torch.utils.checkpoint.checkpoint(), avoiding storage of masking and ncon intermediates across all iterations simultaneously.distributed – if True, dispatch the unrolled-combo sum across an already-initialised
torch.distributedprocess group (SPMD, one rank per GPU, scales across nodes via NCCL). Every rank must reach this call with a structurally identical copy of the inputs. Requires a dictunrolland a torch backend;devices/mp_workers_per_deviceare ignored. See_oe_blocksparse_dist. A single-rank group falls back to serial. Optionaldistributed_groupselects a non-default process group.
Supporting types used to build the unroll argument:
- class yastn.SlicedLeg(t, D, slices=None)[source]#
Describes a subset of a YASTN tensor leg: a selection of charge sectors, each optionally restricted to a contiguous slice within the sector’s full block dimension.
A collection of non-overlapping
SlicedLegobjects whose union covers all charge sectors of a leg forms a valid partition. Iterating over such a partition and summing the partial contractions reproduces the full result, while each individual contraction operates on a strictly smaller tensor.- Parameters:
t (Sequence[tuple | int]) – Charge sectors included in this slice. Each element must be a tuple of ints (one per symmetry component, length
NSYM), or a plain int for single-component symmetries (normalised to a 1-tuple internally).D (Sequence[int]) – Dimension of each charge sector within this slice. Must equal
len(t). Usefull_Dwhen the entire block is selected.slices (dict[tuple, slice], optional) – Maps each charge tuple to a
sliceobject into the full block dimension of that sector. If omitted, every sector usesslice(None)(i.e. the full block is selected).
- property tD#
Dict mapping charge tuple → dimension in this slice.
- yastn.make_sliced_legs(leg)[source]#
Split a YASTN
yastn.Leginto oneSlicedLegper charge sector (the simplest non-overlapping partition).Each returned
SlicedLegcovers exactly one charge sector and includes the full block dimension of that sector (slice(None)).- Parameters:
leg (yastn.Leg)
- Return type:
list[SlicedLeg]
- yastn.tensor.oe_blocksparse.slice_leg_uniform(leg: Leg, size: int)[source]#
Uniformly slice YASTN
yastn.Leginto segments of at mostsize. Resulting slices can span multiple charge sectors.The returned list of
SlicedLegobjects forms a valid partition of the leg, with each slice selecting a contiguous subset of the block dimension within each charge sector. The last slice may be smaller than the specified size.- Parameters:
leg (yastn.Leg)
size (int)
- Return type:
list[SlicedLeg]
Single-device execution#
yastn.contract_with_unroll() resolves the unroll argument (integers are
expanded to SlicedLeg partitions), optionally searches a
contraction path with yastn.get_contraction_path(), and then either falls
back to a plain yastn.ncon() (when unroll is None) or drives the sliced
loop in _contract_with_sliced_unroll.
The loop is preceded by a metadata-only prefilter (_metadata_filter_combos):
for every combo it applies the slice masks to the tensor metadata only (no GPU
work), drops combos whose masked tensors have no surviving blocks, and runs
ncon_prefilter() to record, per surviving combo, which blocks each operand
must keep (pf_trim) and — when per_combo_path is set — the effective
per-axis dimensions used to tune a combo-specific path. Only the surviving combos
enter the actual contraction loop.
Inside the loop, each combo (_contract_single_combo) masks its operands with
cached diagonal apply_mask() tensors, trims blocks according
to pf_trim (_filter_tensor_blocks()), runs yastn.ncon(), and
accumulates the partial into a bucket keyed by its position along the
output-unrolled axes. Optionally the whole loop body is wrapped in
torch.utils.checkpoint.checkpoint() (checkpoint_loop=True) so masking
and ncon intermediates are recomputed on backward instead of being kept for
every iteration. Finally, single-key results are returned directly, while
output-unrolled results are reassembled with yastn.block() and stripped of
their fusion history via drop_leg_history().
Single-device (serial) control/data flow of contract_with_unroll.#
Multi-device execution#
When devices names more than one device (or a single device with
mp_workers_per_device >= 2), _contract_with_sliced_unroll hands off to
_contract_with_sliced_unroll_mp in yastn.tensor._oe_blocksparse_mp.
A persistent pool of spawn-ed worker processes (mp_workers_per_device
per device) is created once per (devices, workers, config) and reused across
calls; each worker loops on a command queue and pins itself to its assigned
device.
The dispatcher runs the same metadata prefilter as the serial path and, in
addition, derives the output tensor’s structure directly from the input legs
(_derive_output_structs) — no data contraction is needed to know the result’s
blocks. Both results are memoized per pool, keyed on a structural fingerprint of
the inputs. The surviving combos are distributed round-robin into per-worker
assignments, and the computation is driven through a custom
torch.autograd.Function (_MultiprocSlicedUnrollFunction) that
implements a checkpoint pattern:
Forward. Input data is replicated once per unique worker device (shared via CUDA IPC), each worker contracts its assigned combos under
torch.no_grad(), zero-fills every partial to its per-key output struct, and ships the partials back. The parent sums per-key partials (gathering to the original device) and, for output-unrolled contractions, reassembles them withyastn.block().Backward. The parent re-runs only the (cheap)
yastn.block()assembly with autograd enabled to split the output gradient into per-key gradients, then asks each worker to re-run its combos with autograd on and calltorch.autograd.backward()locally. Workers return per-input gradient data, which the parent sums across workers. No live autograd graph crosses the process boundary; the forward is replayed on backward instead.
Robustness details handled by the pool: a Linux parent-death signal so orphaned workers do not leak CUDA contexts, a liveness guard that turns a dead-worker hang into an error, and multiprocess-safe logging that forwards worker log records to the parent.
Multi-device dispatch of the sliced unroll loop (forward + backward).#
Memory management#
The sliced loop keeps the peak size of any single intermediate bounded, but on
CUDA the PyTorch caching allocator can still fragment: it holds memory that
is reserved but unallocated in segments that also contain live blocks, so a
later large contiguous allocation fails even though the totals look sufficient.
torch.cuda.empty_cache() only returns segments that are entirely free, so
it cannot cure this class of fragmentation — only expandable_segments:True
(virtual-memory remapping of partial segments) can.
Two further sets of knobs matter here, documented elsewhere but worth pointing at because the unroll loop amplifies both:
Every
yastn.ncon()inside the loop fuses and unfuses legs, soYASTN_FUSE_SCATTER_THRESHandYASTN_FUSE_SCATTER_CHUNK— see GPU execution: hybrid scatter/loop — apply to each combo.YASTN_FUSE_SCATTER_CHUNKis the one to reach for if the index build itself is the memory problem, since it bounds that build’s scratch allocation.The metadata caches (Caching) are populated per process, so each spawned worker keeps its own. The fusion-index cache holds device-resident tensors;
yastn.clear_cache()releases them if the accumulated indices become significant next to the contraction’s own working set.
The knobs below tune allocator behaviour for the torch / torch_cutensor
CUDA backends (they are inert on CPU / NumPy):
Variable / argument |
Scope |
Purpose |
|---|---|---|
|
all processes, at allocator init |
PyTorch’s own global allocator config; the place to set
|
|
each spawned worker, at startup |
Runtime-mutable allocator knobs ( |
|
|
How often |
|
|
|
Ideal case — expandable_segments:True#
If the host kernel supports it, the cleanest option is PyTorch’s expandable segments, set globally through PyTorch’s own environment variable:
export PYTORCH_ALLOC_CONF=expandable_segments:True
Each process’s allocator reads this once at initialization, and spawn-ed
workers inherit the same environment — so their allocators pick it up too, with
no per-worker configuration. Expandable segments let the allocator hand back
partial segments, eliminating the fragmentation that plain empty_cache
cannot. In this case leave both YASTN_OE_ALLOC_CONF and
YASTN_OE_CUDA_CACHE_RELEASE_LEVEL unset.
Fallback — kernel without expandable_segments support#
Sharing CUDA tensors between processes with expandable_segments:True needs
the pidfd_open syscall. On kernels (or containers) that lack it, the
multi-device path crashes during CUDA IPC with an error naming
pidfd_open. This is deliberate — YASTN does not silently work around it, so
the failure clearly tells you this host cannot combine expandable segments with
multiprocess dispatch. Drop expandable_segments and use the two per-worker
knobs below instead.
Per-worker allocator settings. Apply runtime-mutable caching-allocator knobs (see Optimizing memory usage with PYTORCH_ALLOC_CONF) to each worker at startup:
export YASTN_OE_ALLOC_CONF="garbage_collection_threshold:0.8,max_split_size_mb:512"
The key knob is garbage_collection_threshold: it makes the allocator
proactively reclaim unused cached blocks once reserved memory exceeds the given
fraction, keeping the reserved-but-unallocated pile from growing and
fragmenting. max_split_size_mb additionally stops large free blocks from
being carved up for small requests, preserving big contiguous regions for the
large intermediates. Only runtime-mutable keys take effect; backend and
expandable_segments are fixed at allocator init and are ignored here. The
same string can be passed programmatically, which wins over the environment:
contract_with_unroll(..., alloc_conf="garbage_collection_threshold:0.8")
(Applies to the spawned workers of the multi-device path. A single-process
serial run has no workers; use PYTORCH_ALLOC_CONF for it, which the main
process reads at init.)
Cache-release cadence. On the torch_cutensor backend, YASTN can call the
blocking torch.cuda.empty_cache() at chosen points to return fully-free
segments to the driver. A release point fires only when the level is at least
its tag (each level adds to the ones below):
Level |
Effect |
|---|---|
|
never (default) |
|
once per contraction |
|
|
|
|
export YASTN_OE_CUDA_CACHE_RELEASE_LEVEL=2
This should be experimented with; a reasonable starting point is level 2
(clean up after every combo). Remember that empty_cache cannot defragment
free space interleaved with live allocations — that is what
expandable_segments:True is for. A malformed value is treated as 0 with a
one-time warning.