diff --git a/README.md b/README.md index fa97463..82c2800 100644 --- a/README.md +++ b/README.md @@ -105,7 +105,9 @@ out = src( ) ``` -See [Large problems](https://algorithmiq.github.io/src-method/features/large-problems) for the budgets, the scratch directory and how to read the logged plan. +On a node with several GPUs, `Resources(devices=4)` splits the sweep among four of them, in the same process. + +See [Large problems](https://algorithmiq.github.io/src-method/features/large-problems) for the budgets, the scratch directory, several GPUs and how to read the logged plan. ### Logging diff --git a/benches/large/README.md b/benches/large/README.md index 7cc9eee..cfe7df2 100644 --- a/benches/large/README.md +++ b/benches/large/README.md @@ -1,7 +1,8 @@ # Out-of-core benchmark `bench_large.py` compresses the stack `N . V . M . U` of Pauli-transfer-matrix MPOs -(physical legs of 4) on one GPU, where `M` has a large bond and is read from disk. +(physical legs of 4) on one GPU or several, where `M` has a large bond and is read +from disk. ```bash # Write the stack to node-local NVMe (205 GB for M at the defaults, complex128). @@ -15,9 +16,13 @@ uv run python benches/large/bench_large.py run /local/stack --chi-out 2000 \ uv run python benches/large/bench_large.py generate /local/medium --bond-m 1000 uv run python benches/large/bench_large.py compare /local/medium --chi-out 500 \ --small 40GB --large 80GB --scratch-dir /local/scratch + +# Split the sweep among the GPUs of the node: 1, 2, 4 and 8 of them, in turn. +uv run python benches/large/bench_large.py scaling /local/stack --chi-out 2000 \ + --devices 1 2 4 8 --scratch-dir /local/scratch --results scaling.jsonl ``` -`run` logs the wall time, the size of CuPy's pool at the end (its high-water mark, +`run` logs the wall time, the size of CuPy's pool on each GPU at the end (its high-water mark, since the pool keeps its blocks), the host peak (`ru_maxrss`), the planned peaks, the bytes spilled and the tiers. With `--debug` before the command, `src_method` also logs its plan, the time of each pass and the stall time (`SRC stalls`). @@ -32,6 +37,15 @@ memory and node-local NVMe: 3. The stall time is below 10% of the wall time. 4. `compare` at `D_M = 1000`, `chi_out = 500` reports a distance below `1e-10`. +`scaling` logs, for each GPU count, the wall time, the speed-up and the parallel +efficiency `T_1 / (G T_G)` against the first count, and the relative distance to +its output; `--results` appends the same as JSON lines. `--devices` with +`run` splits a single run. Acceptance of the multi-GPU sweep, at the reference +size on 8xH100 or 8xH200 (and 4xA100 on Leonardo): + +1. Every distance is below `1e-10`. +2. The parallel efficiency is at least 80% on all the GPUs of the node. + ## Leonardo The [`leonardo/`](leonardo/) folder holds the Slurm script and the logs; see the diff --git a/benches/large/bench_large.py b/benches/large/bench_large.py index 8ad4143..59b8a2e 100644 --- a/benches/large/bench_large.py +++ b/benches/large/bench_large.py @@ -1,13 +1,16 @@ """Out-of-core SRC of ``N . V . M . U`` with a large ``M``, read from disk. -Three commands: +Four commands: - ``generate``: write the four MPOs site by site as ``.npy`` files, so that ``M`` is never whole in memory. -- ``run``: compress the stack with automatic (or given) budgets and report the - plan, the wall time, the device pool size, the host peak and the stall time. +- ``run``: compress the stack with automatic (or given) budgets, on one or several + GPUs, and report the plan, the wall time, the device pool sizes, the host peak + and the stall time. - ``compare``: run twice with two GPU budgets and report the relative distance between the two outputs, which checks that the plan does not change the result. +- ``scaling``: run on 1, 2, 4, ... GPUs of the node and report the wall time, the + parallel efficiency and the distance to the single-GPU output. The plan, the per-pass times and the stall time are also logged by `src_method` itself at ``DEBUG``; pass ``--debug`` to see them. @@ -15,6 +18,7 @@ from __future__ import annotations +import json import logging import resource from pathlib import Path # noqa: TC003 (cyclopts reads the annotations at runtime) @@ -94,8 +98,12 @@ def _load(directory: Path) -> list[list[np.ndarray]]: def _run( - directory: Path, chi_out: int, resources: Resources, seed: int -) -> list[np.ndarray]: + directory: Path, + chi_out: int, + resources: Resources, + seed: int, + device: str = "gpu", +) -> tuple[list[np.ndarray], float]: layers = _load(directory) plans = [] @@ -111,21 +119,27 @@ def spy(*args: object, **kwargs: object) -> object: chi_out=chi_out, dtype=layers[0][0].dtype, seed=seed, - device="gpu", + device=device, resources=resources, ) seconds = perf_counter() - start finally: sweep_module.make_plan = make_plan - import cupy # noqa: PLC0415 (GPU-only benchmark) - (plan,) = plans tiers = [site.tier for site in plan.sites] + pools = [] + if device == "gpu": + import cupy # noqa: PLC0415 (optional dependency) + + devices = resources.devices or 1 + for gpu in range(devices) if isinstance(devices, int) else devices: + with cupy.cuda.Device(gpu): + pools.append(cupy.get_default_memory_pool().total_bytes()) logger.info( - "Run complete: %.1f s, pool %d B, host peak %d B, planned device peak %d B, " + "Run complete: %.1f s, pools %s B, host peak %d B, planned device peak %d B, " "planned host peak %d B, disk %d B, tiers %s, sketch batches %s", seconds, - cupy.get_default_memory_pool().total_bytes(), + pools, resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024, plan.device_peak, plan.host_peak, @@ -133,7 +147,7 @@ def spy(*args: object, **kwargs: object) -> object: {tier: tiers.count(tier) for tier in ("device", "host", "disk")}, sorted({site.sketch_batch for site in plan.sites[1:]}), ) - return out + return out, seconds @app.command @@ -144,6 +158,7 @@ def run( gpu_memory: str | None = None, host_memory: str | None = None, scratch_dir: Path | None = None, + devices: int = 1, seed: int = 0, ) -> None: """Compress the stack written by ``generate`` on the GPU. @@ -151,12 +166,14 @@ def run( Args: directory: The directory given to ``generate``. chi_out: The output bond dimension. - gpu_memory: GPU budget, e.g. ``36GB``; detected when omitted. + gpu_memory: GPU budget of each device, e.g. ``36GB``; detected when omitted. host_memory: Host budget; detected when omitted. scratch_dir: Where to spill environments; ``$TMPDIR`` when omitted. + devices: The number of GPUs to split the sweep among. seed: Seed of the sketch. """ - _run(directory, chi_out, Resources(gpu_memory, host_memory, scratch_dir), seed) + resources = Resources(gpu_memory, host_memory, scratch_dir, devices=devices) + _run(directory, chi_out, resources, seed) def _inner(a: list[np.ndarray], b: list[np.ndarray]) -> complex: @@ -187,11 +204,82 @@ def compare( scratch_dir: Where to spill environments; ``$TMPDIR`` when omitted. seed: Seed of the sketch, the same for both runs. """ - first = _run(directory, chi_out, Resources(small, None, scratch_dir), seed) - second = _run(directory, chi_out, Resources(large, None, scratch_dir), seed) - aa, bb, ab = _inner(first, first), _inner(second, second), _inner(first, second) - distance = np.sqrt(max((aa + bb - 2 * ab).real, 0.0) / aa.real) - logger.info("Relative distance between the runs: %.3e", distance) + first, _ = _run(directory, chi_out, Resources(small, None, scratch_dir), seed) + second, _ = _run(directory, chi_out, Resources(large, None, scratch_dir), seed) + logger.info("Relative distance between the runs: %.3e", _distance(first, second)) + + +def _distance(a: list[np.ndarray], b: list[np.ndarray]) -> float: + """Relative Frobenius distance ``|a - b| / |a|`` of two MPOs.""" + aa, bb, ab = _inner(a, a), _inner(b, b), _inner(a, b) + return float(np.sqrt(max((aa + bb - 2 * ab).real, 0.0) / aa.real)) + + +@app.command +def scaling( + directory: Path, + *, + chi_out: int = 2000, + devices: Annotated[tuple[int, ...], cyclopts.Parameter(consume_multiple=True)] = ( + 1, + 2, + 4, + 8, + ), + gpu_memory: str | None = None, + host_memory: str | None = None, + scratch_dir: Path | None = None, + results: Path | None = None, + seed: int = 0, + device: str = "gpu", +) -> None: + """Run on increasing numbers of GPUs and report the strong scaling. + + The first count is the reference: the efficiency of ``G`` GPUs is + ``T_ref * G_ref / (T_G * G)``, and every output is compared with the + reference output. Host memory holds two outputs at a time. + + Args: + directory: The directory given to ``generate``. + chi_out: The output bond dimension. + devices: The GPU counts, the reference first. + gpu_memory: GPU budget of each device; detected when omitted. + host_memory: Host budget of the node; detected when omitted. + scratch_dir: Where to spill environments; ``$TMPDIR`` when omitted. + results: A JSON-lines file to append one record per run to. + seed: Seed of the sketch, the same for every run. + device: ``gpu``, or ``cpu`` to smoke-test with simulated devices. + """ + reference: list[np.ndarray] | None = None + ref_seconds = ref_devices = 0.0 + for count in devices: + resources = Resources(gpu_memory, host_memory, scratch_dir, devices=count) + out, seconds = _run(directory, chi_out, resources, seed, device) + if reference is None: + reference, ref_seconds, ref_devices = out, seconds, count + distance = 0.0 + else: + distance = _distance(reference, out) + efficiency = ref_seconds * ref_devices / (seconds * count) + logger.info( + "GPUs %d: %.1f s, speed-up %.2f, efficiency %.0f %%, distance %.3e", + count, + seconds, + ref_seconds / seconds, + 100 * efficiency, + distance, + ) + if results is not None: + record = { + "gpus": count, + "seconds": seconds, + "efficiency": efficiency, + "distance": distance, + "chi_out": chi_out, + } + with results.open("a") as f: + f.write(json.dumps(record) + "\n") + del out @app.meta.default diff --git a/docs/content/docs/contributing/testing.mdx b/docs/content/docs/contributing/testing.mdx index 02e5b48..ce1aa0c 100644 --- a/docs/content/docs/contributing/testing.mdx +++ b/docs/content/docs/contributing/testing.mdx @@ -23,6 +23,7 @@ Always run tests through `pytest`, never as `python tests/test_.py`. | `test_plan.py` | the planner: budgets, batch sizes and environment tiers | | `test_store.py` | environment storage on the device, in host memory and on disk | | `test_sites.py` | reading the cores one site at a time, with prefetching | +| `test_group.py` | the device group: work splitting, collectives, failures | | `test_logging.py` | the package leaves the host's logging configuration alone | | `test_gpu_backend.py` | the CuPy path; skipped without CuPy or a GPU | @@ -88,6 +89,16 @@ pure function, so `tests/test_plan.py` checks batch sizes and tiers from shapes alone, and `tests/test_gpu_backend.py` derives GPU budgets from a plan made with `make_plan` to reach every tier. +## Exercising several devices + +`Resources(devices=n)` on the CPU simulates `n` devices with threads that share +the host, through the same code as on GPUs: the split of the sketch columns, the +collectives and the shared reading of the cores. `tests/test_stack.py` compares +such runs with a single device, also with uneven blocks (`chi_out` not a +multiple of `n`, or smaller than it), a cutoff and spilling. Only the peer copies +and device selection are GPU-specific; their tests in `tests/test_gpu_backend.py` +skip unless at least two GPUs are visible. + ## Logging in tests `log_cli` is on, so log records show up live while the tests run. To assert on diff --git a/docs/content/docs/features/gpu.mdx b/docs/content/docs/features/gpu.mdx index 0ac9d78..847f17c 100644 --- a/docs/content/docs/features/gpu.mdx +++ b/docs/content/docs/features/gpu.mdx @@ -34,3 +34,7 @@ result = apply(H, psi, chi_out=256, device="gpu") The inputs can be NumPy or CuPy arrays; they are moved to the device for the sweep, and the result always comes back as NumPy arrays on the host. Without CuPy installed, `device="gpu"` raises an `ImportError`. + +To split one sweep among several GPUs of a node, pass +`resources=Resources(devices=n)`; see +[Large problems](/features/large-problems#several-gpus). diff --git a/docs/content/docs/features/large-problems.mdx b/docs/content/docs/features/large-problems.mdx index 3cc4bd7..165bb01 100644 --- a/docs/content/docs/features/large-problems.mdx +++ b/docs/content/docs/features/large-problems.mdx @@ -1,6 +1,6 @@ --- title: Large problems -description: Run stacks that fit neither on the GPU nor in host memory. +description: Run stacks that fit neither on the GPU nor in host memory, on one GPU or several. --- A single SRC sweep keeps, besides the input and output trains, one sketched @@ -75,13 +75,68 @@ reference size of 50 sites, `D_M = 4000` and `chi_out = 2000` in complex128, up returns, also after an error; a process killed with `SIGKILL` leaves it behind, and its name identifies the process. +## Several GPUs + +On a node with several GPUs, one sweep can be split among them: + +```python notest +from src_method import Resources, src + +out = src( + N, V, M, U, + chi_out=2000, + dtype=np.complex128, + device="gpu", + resources=Resources(devices=4, scratch_dir="/local/scratch"), +) +``` + +`devices` is a count, for the first `n` visible GPUs, or a list of device ids; +the first one gathers the sketches and runs the QR. Everything runs in the calling +process, one thread per GPU, so nothing changes for the caller: the call returns +the output cores as with one GPU, errors are raised as usual, and the inputs are +read once and shared. + +The sketch index is split: each GPU owns a contiguous block of the `chi_out` +sketch columns and computes only those columns of every environment and of every +sketch, so the environments, the dominant memory, are split among the GPUs too. +At every site of the right-to-left pass, the first GPU gathers the sketch, runs +the QR and gives each GPU a block of rows of the output core; each GPU projects +its rows of the next projected environment, and those rows are gathered on every +GPU. The random draws are those of one GPU, so the result is that of one GPU up to +rounding. + +What is not split: + +- every GPU holds a whole site: the cores, the projected environment and the + output core of the site, so the largest site that fits is that of one GPU; +- every core is read from the source once, but every GPU needs a copy. When the + GPUs reach each other directly (NVLink, or PCIe peer-to-peer), each uploads a + slice of the core and gathers the rest from the others; otherwise each uploads + the whole core; +- the QR runs on the first GPU while the others wait. + +These costs grow more slowly with the problem than the contractions do, so large +problems, the ones this is for, scale close to linearly, while small ones may run +slower on several GPUs than on one. + +The budgets apply as follows: `gpu_memory` is the budget of each GPU (detected: +the smallest free memory of the GPUs), while `host_memory` and `scratch_dir` are +shared by all of them. Every GPU writes its environments to its own scratch +directory. + +On the CPU, `devices` simulates the GPUs with threads that share the host; the +result is the same, and the mode exists to test the multi-GPU sweep without GPUs. + ## Reading the plan With `DEBUG` logging on for the `src_method` logger (see [Logging](/features/logging)), every sweep logs its plan as `SRC plan: ...`: - `prefetch`: whether the next site is read ahead; -- `device peak`, `host peak`, `disk`: the planned peaks, in bytes; +- `devices`: the device ids the sweep is split among; +- `device peak`, `host peak`, `disk`: the planned peaks, in bytes; the device + peak is that of each device, the others are totals; - `scratch`: the scratch directory, when anything spills to disk; - `tiers`: where the environment of each site lives (`device`, `host` or `disk`); - `batches`: for each site, the batch of the environment, sketch and projection @@ -89,4 +144,5 @@ With `DEBUG` logging on for the `src_method` logger (see It then logs the time of each pass and, at the end, `SRC stalls`: the seconds the sweep waited for input cores and for environments. If they are a large -fraction of the run, the disk is too slow for the compute. +fraction of the run, the disk is too slow for the compute. With several devices, +these lines start with the device they refer to. diff --git a/src/src_method/_group.py b/src/src_method/_group.py new file mode 100644 index 0000000..8e9facb --- /dev/null +++ b/src/src_method/_group.py @@ -0,0 +1,252 @@ +"""The devices of one sweep, one thread per device, and the collectives between them. + +A sweep over several devices of a node runs in one process: `DeviceGroup.run` +starts one thread per device, each running the same per-device function on its +*rank* (its position in the group). Collectives are called by every rank in the +same order. They exchange device arrays through peer copies, ordered with the +kernels by events. A rank that fails aborts the group: every other rank waiting in +a collective, or for an input core, raises `GroupAbortedError`, and `run` re-raises the +original error. + +On the host backend the devices are simulated: the threads share the host, and the +collectives are plain copies. The result is that of a single device; this mode +exists to test the distributed logic without GPUs. +""" + +from __future__ import annotations + +import threading +from concurrent.futures import ThreadPoolExecutor +from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar + +from .utils import copy_into, current_stream, enable_peer_access, use_device + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + from types import ModuleType + + from numpy.typing import NDArray + +T = TypeVar("T") + + +class GroupAbortedError(RuntimeError): + """Raised on the ranks of a group after another rank failed.""" + + +def blocks(n: int, size: int) -> list[tuple[int, int]]: + """Split ``range(n)`` into ``size`` contiguous ``(lo, hi)`` blocks, in rank order. + + Block sizes differ by at most one, the larger blocks first. Blocks may be empty + when ``n < size``. + """ + base, extra = divmod(n, size) + out, lo = [], 0 + for rank in range(size): + hi = lo + base + (rank < extra) + out.append((lo, hi)) + lo = hi + return out + + +class _Part(NamedTuple): + """What a rank publishes in a collective: an array and the event it is ready at.""" + + array: Any + ready: Any + + +class DeviceGroup: + """The devices of one sweep and the collectives between them. + + Args: + xp: Array module (``numpy`` or ``cupy``). + devices: The device ids, in rank order; rank 0 is the root. + """ + + def __init__(self, xp: ModuleType, devices: Sequence[int]) -> None: + self.xp = xp + self.devices = tuple(devices) + self.size = len(self.devices) + # Uploading 1/size of each core and gathering the rest pays off only when + # the devices copy directly between them. + self.split_uploads = self.size > 1 and enable_peer_access(xp, self.devices) + self._barrier = threading.Barrier(self.size) + # Two sets of slots, alternating between consecutive collectives: a rank may + # publish into the next collective while slower ranks still read this one. + self._parts: list[list[Any]] = [[None] * self.size, [None] * self.size] + self._done: list[list[Any]] = [[None] * self.size, [None] * self.size] + self._calls = [0] * self.size + self._on_abort: list[Callable[[], None]] = [] + self.aborted = threading.Event() + + def on_abort(self, callback: Callable[[], None]) -> None: + """Call ``callback`` when a rank fails, to release threads outside the ranks.""" + self._on_abort.append(callback) + + def abort(self) -> None: + """Wake every rank waiting in a collective or on a registered resource.""" + if self.aborted.is_set(): + return + self.aborted.set() + self._barrier.abort() + for callback in self._on_abort: + callback() + + def run(self, fn: Callable[[int], T]) -> list[T]: + """Run ``fn(rank)`` on every device, each in its own thread. + + A group of one runs ``fn(0)`` in the calling thread. + + Args: + fn: The per-device function. + + Returns: + The results, in rank order. + + Raises: + Exception: The first error raised by a rank, after every rank stopped. + """ + if self.size == 1: + with use_device(self.xp, self.devices[0]): + return [fn(0)] + first: list[BaseException] = [] + lock = threading.Lock() + + def worker(rank: int) -> T: + try: + with use_device(self.xp, self.devices[rank]): + return fn(rank) + except BaseException as err: + secondary = isinstance( + err, GroupAbortedError | threading.BrokenBarrierError + ) + with lock: + if not first and not secondary: + first.append(err) + self.abort() + raise + + with ThreadPoolExecutor(self.size, thread_name_prefix="src-device") as pool: + futures = [pool.submit(worker, rank) for rank in range(self.size)] + results, errors = [], [] + for future in futures: + try: + results.append(future.result()) + except BaseException as err: # noqa: BLE001 (re-raised below) + errors.append(err) + if first: + raise first[0] + if errors: + raise errors[0] + return results + + # --- collectives --- + + def _wait(self) -> None: + try: + self._barrier.wait() + except threading.BrokenBarrierError: + raise GroupAbortedError from None + + def _publish(self, rank: int, array: Any) -> tuple[int, list[_Part]]: # noqa: ANN401 + """Publish ``array`` and return the parts of every rank.""" + parity = self._calls[rank] % 2 + self._calls[rank] += 1 + self._parts[parity][rank] = _Part(array, current_stream(self.xp).record()) + self._wait() + return parity, list(self._parts[parity]) + + def _settle(self, rank: int, parity: int) -> None: + """Wait, on this rank's stream, until every rank has read what it published. + + Until then the arrays published by this rank must not be overwritten; the + memory pool reuses freed blocks in stream order, so making the stream wait + is enough. + """ + stream = current_stream(self.xp) + self._done[parity][rank] = stream.record() + self._wait() + for event in self._done[parity]: + stream.wait_event(event) + self._parts[parity][rank] = None + + def _pull(self, part: _Part, dst: NDArray) -> None: + stream = current_stream(self.xp) + stream.wait_event(part.ready) + copy_into(dst, part.array, self.xp, stream) + + def allgather(self, rank: int, local: NDArray) -> NDArray: + """Concatenate the blocks of every rank along the first axis, on every rank. + + Args: + rank: The calling rank. + local: This rank's block, C-contiguous on its device; the blocks may + differ along the first axis only. + + Returns: + The concatenation in rank order, on this rank's device. + """ + if self.size == 1: + return local + parity, parts = self._publish(rank, local) + full = self._empty_like_parts(parts) + lo = 0 + for part in parts: + hi = lo + part.array.shape[0] + self._pull(part, full[lo:hi]) + lo = hi + self._settle(rank, parity) + return full + + def gather(self, rank: int, local: NDArray, root: int = 0) -> NDArray | None: + """Concatenate the blocks of every rank along the first axis, on ``root``. + + Args: + rank: The calling rank. + local: This rank's block, C-contiguous on its device. + root: The rank that receives the result. + + Returns: + The concatenation on ``root``, ``None`` on the other ranks. + """ + if self.size == 1: + return local + parity, parts = self._publish(rank, local) + full = None + if rank == root: + full = self._empty_like_parts(parts) + lo = 0 + for part in parts: + hi = lo + part.array.shape[0] + self._pull(part, full[lo:hi]) + lo = hi + self._settle(rank, parity) + return full + + def scatter(self, rank: int, full: NDArray | None, root: int = 0) -> NDArray: + """Give every rank its block (see `blocks`) of the first axis of ``full``. + + Args: + rank: The calling rank. + full: The array to split, C-contiguous on ``root``; ignored elsewhere. + root: The rank that holds ``full``. + + Returns: + This rank's block, on its device. + """ + if self.size == 1: + assert full is not None # noqa: S101 (the root holds it) + return full + parity, parts = self._publish(rank, full if rank == root else None) + source = parts[root].array + lo, hi = blocks(source.shape[0], self.size)[rank] + local = self.xp.empty((hi - lo, *source.shape[1:]), dtype=source.dtype) + self._pull(_Part(source[lo:hi], parts[root].ready), local) + self._settle(rank, parity) + return local + + def _empty_like_parts(self, parts: Sequence[_Part]) -> NDArray: + first = parts[0].array + rows = sum(part.array.shape[0] for part in parts) + return self.xp.empty((rows, *first.shape[1:]), dtype=first.dtype) diff --git a/src/src_method/_plan.py b/src/src_method/_plan.py index aa44299..8a2a099 100644 --- a/src/src_method/_plan.py +++ b/src/src_method/_plan.py @@ -1,8 +1,12 @@ """Memory planning for the SRC sweep: budgets, batch sizes and environment tiers. -`make_plan` is a pure function of the core shapes, the sketch size, the dtype and -the budgets, so the plan of a run can be inspected, and tested, without loading any -data or touching a GPU. +`make_plan` is a pure function of the core shapes, the sketch size, the dtype, the +number of devices and the budgets, so the plan of a run can be inspected, and +tested, without loading any data or touching a GPU. + +With several devices, every device runs the same plan on its block of the sketch +columns and of the rows of the projected environment; the plan is that of the +largest block, and of rank 0, which also gathers the sketch and runs the QR. """ from __future__ import annotations @@ -10,19 +14,27 @@ import re import shutil import tempfile +from collections.abc import Sequence from dataclasses import dataclass -from math import prod +from math import ceil, prod from pathlib import Path from typing import TYPE_CHECKING, Literal import numpy as np from ._kernels import equations, peak_elements -from .utils import device_memory, host_memory_available, is_host +from .utils import ( + current_device, + device_count, + device_memory, + host_memory_available, + is_host, + use_device, +) if TYPE_CHECKING: import os - from collections.abc import Callable, Sequence + from collections.abc import Callable from types import ModuleType Tier = Literal["device", "host", "disk"] @@ -100,22 +112,79 @@ class Resources: scratch_dir: Directory for environments that fit in neither budget, ideally on node-local disk. Defaults to ``tempfile.gettempdir()``, which honours ``$TMPDIR``. + devices: The devices to split the sweep among: a count ``n`` for the + first ``n`` visible devices, or a sequence of device ids. Defaults to + the current device alone. With several devices, ``gpu_memory`` is the + budget of each, and ``host_memory`` and ``scratch_dir`` are shared. On + the CPU the devices are simulated, which gives the same result and + exists to test the multi-device sweep. """ gpu_memory: int | str | None = None host_memory: int | str | None = None scratch_dir: str | os.PathLike[str] | None = None + devices: int | Sequence[int] | None = None def __post_init__(self) -> None: - """Validate the explicit budgets. + """Validate the explicit budgets and devices. Raises: - TypeError: If a budget is neither an integer nor a string. - ValueError: If a budget is not a recognised size. + TypeError: If a budget is neither an integer nor a string, or the + devices are neither a count nor a sequence of integers. + ValueError: If a budget is not a recognised size, the device count is + not positive, or the device ids are negative, repeated or none. """ for value in (self.gpu_memory, self.host_memory): if value is not None: parse_size(value) + _device_ids(self.devices) + + +def _device_ids(devices: object) -> tuple[int, ...] | None: + """Validate a device request and return its ids, or ``None`` for the default.""" + if devices is None: + return None + if isinstance(devices, int | np.integer) and not isinstance(devices, bool): + if devices < 1: + msg = f"Expected a positive device count, got {devices}." + raise ValueError(msg) + return tuple(range(int(devices))) + if isinstance(devices, str) or not isinstance(devices, Sequence): + msg = f"Expected a device count or a sequence of device ids, got {devices!r}." + raise TypeError(msg) + ids = [] + for d in devices: + if not isinstance(d, int | np.integer) or isinstance(d, bool): + msg = f"Expected integer device ids, got {devices!r}." + raise TypeError(msg) + ids.append(int(d)) + if not ids or min(ids) < 0 or len(set(ids)) != len(ids): + msg = f"Expected distinct non-negative device ids, got {devices!r}." + raise ValueError(msg) + return tuple(ids) + + +def resolve_devices(resources: Resources | None, xp: ModuleType) -> tuple[int, ...]: + """Return the device ids of a sweep, in rank order. + + Args: + resources: The requested devices, or ``None`` for the current device. + xp: Array module (``numpy`` or ``cupy``). + + Returns: + The device ids. + + Raises: + ValueError: If a device does not exist. + """ + ids = _device_ids(resources.devices if resources is not None else None) + if ids is None: + return (current_device(xp),) + count = device_count(xp) + if count is not None and max(ids) >= count: + msg = f"Requested devices {list(ids)}, but only {count} are visible." + raise ValueError(msg) + return ids @dataclass(frozen=True) @@ -141,12 +210,17 @@ class Budgets: device_cap: int | None = None -def resolve_budgets(resources: Resources | None, xp: ModuleType) -> Budgets: +def resolve_budgets( + resources: Resources | None, xp: ModuleType, devices: Sequence[int] = (0,) +) -> Budgets: """Turn `Resources` into byte budgets, detecting those left unset. Args: resources: The requested budgets, or ``None`` to detect all of them. xp: Array module (``numpy`` or ``cupy``). + devices: The devices of the sweep. The device budget is that of each of + them: the smallest detected one. Host memory and disk are those of the + node; on the CPU the simulated devices split the host budget. Returns: The budgets. @@ -164,14 +238,19 @@ def resolve_budgets(resources: Resources | None, xp: ModuleType) -> Budgets: unified = is_host(xp) cap = None if unified: - device = host + device = max(host, 0) // len(devices) else: - available, total = device_memory(xp) - margin = max(int(GPU_MARGIN_FRACTION * total), GPU_MARGIN_MIN) + caps = [] + for d in devices: + with use_device(xp, d): + available, total = device_memory(xp) + margin = max(int(GPU_MARGIN_FRACTION * total), GPU_MARGIN_MIN) + caps.append((available - margin, margin)) + margin = min(m for _, m in caps) if resources.gpu_memory is not None: device = max(parse_size(resources.gpu_memory), 0) else: - device = max(available - margin, 0) + device = max(min(a for a, _ in caps), 0) # The plan counts the bytes in use, but the pool limit applies to every # block the pool holds, including split blocks that are partly free. cap = device + margin @@ -196,6 +275,7 @@ class SitePlan: env_batch: Sketch columns per environment step (0 at the last site). sketch_batch: Sketch columns per sketch step (0 at the first site). project_batch: Rows per projection step, or per first-site step at site 0. + With several devices, batches split each device's share. tier: Where ``C_j`` is kept; site 0 stores no environment. """ @@ -209,12 +289,15 @@ class SitePlan: class Plan: """The memory plan of one sweep. + With several devices, every device runs the same plan on its share of the + sketch columns. + Attributes: sites: One entry per site. prefetch: How many sites ahead the cores are loaded (0 or 1). - device_peak: Estimated peak device bytes. - host_peak: Estimated peak host bytes, beyond the inputs. - disk_bytes: Bytes spilled to the scratch directory. + device_peak: Estimated peak bytes of each device. + host_peak: Estimated peak host bytes, beyond the inputs, over all devices. + disk_bytes: Bytes spilled to the scratch directory, over all devices. """ sites: tuple[SitePlan, ...] @@ -225,7 +308,13 @@ class Plan: class _Site: - """The memory model of one site, in bytes.""" + """The memory model of one site on one device, in bytes. + + With ``ranks`` devices, a device owns ``cols`` of the ``chi`` sketch columns + and projects ``cols`` rows of the new projected environment; rank 0 also holds + the full sketch for the QR. With one device every term is that of the + single-device sweep. + """ def __init__( self, @@ -236,6 +325,7 @@ def __init__( chi: int, itemsize: int, n_sites: int, + ranks: int = 1, ) -> None: self.j, self.shapes, self.core_bytes, self.chi, self.e = ( j, @@ -244,6 +334,10 @@ def __init__( chi, itemsize, ) + self.ranks = ranks + self.cols = ceil(chi / ranks) + # The slice of the cores a device uploads before gathering the rest. + self.chunk = ceil(core_bytes / ranks) if ranks > 1 else 0 self.eqs = equations(len(shapes)) self.left = tuple(s[0] for s in shapes) self.right = tuple(s[1] for s in shapes) @@ -252,12 +346,15 @@ def __init__( self.p = self.up * self.down # Rows of the projected environment S that enters site j right-to-left. self.eta = 1 if j == n_sites - 1 else chi - self.env_bytes = chi * self.a * itemsize + self.env_bytes = self.cols * self.a * itemsize + + def _cores(self, prefetch: int) -> int: + return self.core_bytes * (1 + prefetch) + self.chunk def env(self, b: int, prefetch: int, *, staged_in: bool, staged_out: bool) -> int: """Bytes of one environment step on ``b`` columns.""" - e, chi = self.e, self.chi - fixed = self.core_bytes * (1 + prefetch) + chi * self.p * e + e = self.e + fixed = self._cores(prefetch) + self.cols * self.p * e slices = (1 + staged_in) * b * self.a * e + staged_out * b * self.b * e peak = peak_elements( self.eqs.ltr, ((b, *self.left), (b, self.up, self.down), *self.shapes) @@ -266,11 +363,11 @@ def env(self, b: int, prefetch: int, *, staged_in: bool, staged_out: bool) -> in def sketch(self, b: int, prefetch: int, *, staged: bool) -> int: """Bytes of one sketch step on ``b`` columns.""" - e, chi = self.e, self.chi + e = self.e fixed = ( - self.core_bytes * (1 + prefetch) + self._cores(prefetch) + self.eta * self.b * e - + self.eta * self.p * chi * e + + self.eta * self.p * self.cols * e ) slices = (1 + staged) * b * self.a * e peak = ( @@ -283,22 +380,35 @@ def sketch(self, b: int, prefetch: int, *, staged: bool) -> int: return fixed + slices + peak def qr(self, prefetch: int) -> int: - """Bytes of the QR of the sketch: the sketch, ``Q`` and a workspace.""" + """Bytes of the QR of the sketch: the sketch, ``Q`` and a workspace. + + On rank 0 of a group, the local block of the sketch is alive as well while + the full sketch is gathered. + """ e = self.e + local = self.eta * self.p * self.cols * e if self.ranks > 1 else 0 return ( - self.core_bytes * (1 + prefetch) + self._cores(prefetch) + self.eta * self.b * e + 3 * self.eta * self.p * self.chi * e + + local ) + def gather(self, prefetch: int) -> int: + """Bytes of gathering the new projected environment from its row blocks.""" + if self.ranks == 1: + return 0 + e = self.e + return self._cores(prefetch) + (self.cols + self.chi) * self.a * e + def project(self, b: int, prefetch: int) -> int: """Bytes of one projection step on ``b`` rows.""" - e, chi = self.e, self.chi + e = self.e fixed = ( - self.core_bytes * (1 + prefetch) + self._cores(prefetch) + self.eta * self.b * e - + self.eta * self.p * chi * e - + chi * self.a * e + + self.eta * self.p * self.cols * e + + self.cols * self.a * e ) slices = b * self.eta * self.p * e peak = ( @@ -318,9 +428,7 @@ def first(self, b: int, prefetch: int) -> int: """Bytes of one first-site step on ``b`` rows of ``S``.""" e, chi = self.e, self.chi fixed = ( - self.core_bytes * (1 + prefetch) - + self.eta * self.b * e - + self.a * chi * self.p * e + self._cores(prefetch) + self.eta * self.b * e + self.a * chi * self.p * e ) peak = peak_elements(self.eqs.first, (*self.shapes, (b, *self.right))) * e return fixed + peak @@ -355,6 +463,8 @@ def __init__(self, sites: Sequence[_Site], budgets: Budgets, prefetch: int) -> N self.sites, self.budgets, self.prefetch = sites, budgets, prefetch self.n = len(sites) self.chi = sites[0].chi + self.cols = sites[0].cols + self.ranks = sites[0].ranks self.out_total = sum(self.chi * s.p * s.eta * s.e for s in sites) # Pinned ring of the site source, plus the host copy being read. self.site_ring = (prefetch + 2) * max(s.core_bytes for s in sites) @@ -378,14 +488,20 @@ def _costs(self, j: int, tiers: Sequence[Tier]) -> dict[str, Callable[[int], int costs["first"] = lambda b: site.first(b, pf) return costs + def _limit(self, kernel: str) -> int: + """The largest useful batch: the first site runs on rank 0 alone.""" + return self.chi if kernel == "first" else self.cols + def _batches(self, tiers: Sequence[Tier], avail: int) -> list[dict[str, int]]: batches = [] for j, site in enumerate(self.sites): if j > 0 and site.qr(self.prefetch) > avail: _fail(j, "QR", site.qr(self.prefetch), avail) + if j > 0 and site.gather(self.prefetch) > avail: + _fail(j, "gather", site.gather(self.prefetch), avail) chosen = {} for kernel, cost in self._costs(j, tiers).items(): - b = _largest_batch(cost, self.chi, avail) + b = _largest_batch(cost, self._limit(kernel), avail) if b == 0: _fail(j, kernel, cost(1), avail) chosen[kernel] = b @@ -446,10 +562,13 @@ def plan(self) -> Plan: {kernel: min(b, PREFERRED_BATCH) for kernel, b in chosen.items()} for chosen in batches ] - host_left = budgets.host - self.out_total - self.site_ring - self.env_staging + # The output and the site ring are shared by the devices; each has its own + # staging buffers and host tier. + shared = budgets.host - self.out_total - self.site_ring + host_left = shared // self.ranks - self.env_staging tiers = self._tiers(base - self._peak(staged, preferred), host_left) - disk_bytes = sum( + disk_bytes = self.ranks * sum( s.env_bytes for j, s in enumerate(self.sites) if tiers[j] == "disk" ) if disk_bytes > budgets.disk: @@ -466,7 +585,7 @@ def plan(self) -> Plan: ) host_peak = self.out_total + self.site_ring if disk_bytes or host_env: - host_peak += host_env + self.env_staging + host_peak += self.ranks * (host_env + self.env_staging) sites = tuple( SitePlan( env_batch=chosen.get("env", 0), @@ -500,6 +619,8 @@ def make_plan( chi_out: int, dtype: type | np.dtype, budgets: Budgets, + *, + devices: int = 1, ) -> Plan: """Plan the batches and the environment tiers of one sweep. @@ -509,6 +630,8 @@ def make_plan( chi_out: The sketch size. dtype: The data type of the computation. budgets: The resolved budgets. + devices: The number of devices the sweep is split among; each runs the + plan on its share of the sketch columns. Returns: The plan. Prefetching is dropped before giving up. @@ -527,6 +650,7 @@ def make_plan( chi=chi_out, itemsize=itemsize, n_sites=n_sites, + ranks=devices, ) for j, (shapes, core_bytes) in enumerate(zip(site_shapes, site_bytes)) ] diff --git a/src/src_method/_sites.py b/src/src_method/_sites.py index d59544d..be2e65c 100644 --- a/src/src_method/_sites.py +++ b/src/src_method/_sites.py @@ -1,23 +1,37 @@ """Lazy, per-site access to the cores of a stack. A core is read from its source (a NumPy array, an `np.memmap`, a zarr or HDF5 -dataset) only when the sweep reaches its site, and a background thread can read the -next site while the current one is computed. Shapes come from ``.shape`` alone, so +dataset) only when the sweep needs it, and a background thread reads ahead, in +the order the sweep visits the sites. Shapes come from ``.shape`` alone, so planning never reads any data. + +With a group of several devices the cores are read once, into one page-locked +buffer, and shared by the ranks. Where the devices copy directly between them, +each rank uploads a slice of the buffer and gathers the rest from its peers, so the +host link carries every core once rather than once per device. """ from __future__ import annotations import queue +import threading from concurrent.futures import Future, ThreadPoolExecutor from math import prod from time import perf_counter -from typing import TYPE_CHECKING, Any, Self +from typing import TYPE_CHECKING, Any, NamedTuple, Self import numpy as np +from ._group import DeviceGroup, GroupAbortedError, blocks from ._tensor_train import known_kind, pad_site, padded_shape -from .utils import NullEvent, current_stream, is_host, pinned_empty, to_device_async +from .utils import ( + NullEvent, + current_stream, + is_host, + pinned_empty, + to_device_async, + to_numpy, +) if TYPE_CHECKING: from collections.abc import Sequence @@ -27,8 +41,10 @@ from ._tensor_train import Site -# Byte alignment of each core inside a page-locked staging buffer. +# Byte alignment of each core inside a staging buffer. _ALIGN = 256 +# How often a blocked loader checks whether the source is closing. +_POLL_SECONDS = 0.1 def padded_shapes( @@ -63,16 +79,44 @@ def site_bytes(layers: Sequence[Sequence[Site]]) -> list[int]: ] +def sweep_order(n_sites: int) -> list[int]: + """Return the sites in the order the sweep reads them. + + Left to right up to the last but one site, then right to left from the last. + """ + return [*range(n_sites - 1), *range(n_sites - 1, -1, -1)] + + def _aligned(n_bytes: int) -> int: return -(-n_bytes // _ALIGN) * _ALIGN +class _Staged(NamedTuple): + """One site read into memory. + + Attributes: + cores: Per layer, a host view into ``buffer`` (``None`` for a core given as + a device array), or the padded host core on the host backend. + device: Per layer, the core given as a device array, or ``None``. + offsets: Per layer, the byte offset of its view in ``buffer``. + n_bytes: The bytes of ``buffer`` in use. + buffer: The staging buffer, or ``None`` on the host backend. + """ + + cores: list[Any] + device: list[Any] + offsets: list[int] + n_bytes: int + buffer: Any + + class SiteSource: - """Padded device cores of a stack, one site at a time. + """Padded device cores of a stack, one site at a time, in sweep order. - Reading from the source and staging into page-locked memory happen on a - background thread; the host-to-device copy is issued on the current stream when - the site is requested, so it is ordered with the kernels that use it. + A loader thread reads the sites in ``order`` into a ring of staging buffers, + ``depth`` sites ahead of the slowest rank. Each rank takes the sites in the + same order with `next`; the host-to-device copies are issued on the rank's + current stream, so they are ordered with the kernels that use them. Use as a context manager, so that the loader thread stops. @@ -80,34 +124,51 @@ class SiteSource: layers: A stack in ket form, each layer a sequence of array-likes with ``shape``, ``dtype``, ``ndim`` and ``np.asarray`` support. xp: Array module (``numpy`` or ``cupy``). - depth: How many sites `prefetch` may load ahead (0 disables the thread). + order: The sites, in the order they will be requested. + depth: How many sites the loader may read ahead of the one in use. + group: The devices that share the cores; a single device by default. """ def __init__( - self, layers: Sequence[Sequence[Site]], xp: ModuleType, *, depth: int + self, + layers: Sequence[Sequence[Site]], + xp: ModuleType, + order: Sequence[int], + *, + depth: int, + group: DeviceGroup | None = None, ) -> None: self._layers = layers self._kinds = [known_kind(layer) for layer in layers] self._last = len(layers[0]) - 1 self._xp = xp - self._pending: dict[int, Future[tuple[list[Any], Any]]] = {} - self._executor = ( - ThreadPoolExecutor(1, thread_name_prefix="src-site-loader") - if depth > 0 - else None - ) - self._buffers: queue.Queue[tuple[Any, Any]] = queue.Queue() - self.stall_seconds = 0.0 - if not is_host(xp): - n_bytes = max( - sum( - _aligned(prod(layer[j].shape) * np.dtype(layer[j].dtype).itemsize) - for layer in layers - ) - for j in range(self._last + 1) + self._group = group or DeviceGroup(xp, [0]) + self._order = list(order) + self._staged: list[Future[_Staged]] = [Future() for _ in self._order] + self._taken = [0] * len(self._order) + self._events: list[list[Any]] = [[] for _ in self._order] + self._next = [0] * self._group.size + self._lock = threading.Lock() + self._closing = threading.Event() + self.stall_seconds = [0.0] * self._group.size + # One slot per site in flight: the one in use and ``depth`` ahead. + self._slots: queue.Queue[tuple[Any, list[Any]]] = queue.Queue() + n_bytes = max( + sum( + _aligned(prod(layer[j].shape) * np.dtype(layer[j].dtype).itemsize) + for layer in layers ) - for _ in range(depth + 1): - self._buffers.put((pinned_empty(n_bytes, xp), NullEvent())) + for j in range(self._last + 1) + ) + # A single host device uses the cores as read; simulated devices stage them + # as a GPU group does, so that the host backend tests the same path. + staging = not is_host(xp) or self._group.size > 1 + for _ in range(depth + 1): + buffer = pinned_empty(n_bytes, xp) if staging else None + self._slots.put((buffer, [NullEvent()])) + self._group.on_abort(self._closing.set) + self._loader = ThreadPoolExecutor(1, thread_name_prefix="src-site-loader") + self._loader.submit(self._load) def __len__(self) -> int: """Return the number of sites.""" @@ -124,66 +185,126 @@ def __exit__( tb: TracebackType | None, ) -> None: """Stop the loader thread.""" - if self._executor is not None: - self._executor.shutdown(wait=True, cancel_futures=True) - - def prefetch(self, j: int) -> None: - """Start reading site ``j`` in the background, if there is such a site.""" - if self._executor is None or not 0 <= j <= self._last or j in self._pending: - return - self._pending[j] = self._executor.submit(self._stage, j) - - def __getitem__(self, j: int) -> tuple[NDArray, ...]: - """Return the padded cores of site ``j`` on the device.""" - future = self._pending.pop(j, None) - if future is None: - staged, buffer = self._stage(j) - else: - start = perf_counter() - staged, buffer = future.result() - self.stall_seconds += perf_counter() - start - if is_host(self._xp): - return tuple(staged) - stream = current_stream(self._xp) - cores = [ - core - if isinstance(core, self._xp.ndarray) - else to_device_async(core, self._xp, stream) - for core in staged - ] - # The buffer is refilled only once these copies are done. - self._buffers.put((buffer, stream.record())) - return tuple( - pad_site(core, kind, j, self._last) - for core, kind in zip(cores, self._kinds) - ) - - def _stage(self, j: int) -> tuple[list[Any], Any]: - """Read site ``j`` into memory. - - Returns padded arrays on the host backend, and on the GPU unpadded - page-locked copies (or the device arrays given as input) with their buffer. - """ - if is_host(self._xp): + self._closing.set() + self._loader.shutdown(wait=True, cancel_futures=True) + + # --- loading --- + + def _load(self) -> None: + for k, j in enumerate(self._order): + try: + buffer = self._acquire() + self._staged[k].set_result(self._stage(j, buffer)) + except BaseException as err: # noqa: BLE001 (handed to the ranks) + failure = GroupAbortedError() if self._closing.is_set() else err + for future in self._staged[k:]: + future.set_exception(failure) + return + + def _acquire(self) -> Any: # noqa: ANN401 (a page-locked buffer, or None) + while True: + try: + buffer, events = self._slots.get(timeout=_POLL_SECONDS) + break + except queue.Empty: + if self._closing.is_set(): + raise GroupAbortedError from None + for event in events: + event.synchronize() + return buffer + + def _stage(self, j: int, buffer: Any) -> _Staged: # noqa: ANN401 + """Read site ``j``: padded host cores, or views into the staging buffer.""" + if buffer is None: padded = [ pad_site(_read(layer[j]), kind, j, self._last) for layer, kind in zip(self._layers, self._kinds) ] - return padded, None - buffer, previous = self._buffers.get() - previous.synchronize() - staged, offset = [], 0 + return _Staged(padded, [None] * len(padded), [], 0, None) + # A core already on the device is used as it is by a single device; a group + # stages it like any other, since it lives on one device only. + keep_device = self._group.size == 1 and not is_host(self._xp) + views, device, offsets, offset = [], [], [], 0 for layer in self._layers: core = layer[j] - if not isinstance(core, self._xp.ndarray): - host = np.asarray(core) - view = buffer[offset : offset + host.nbytes] - view = view.view(host.dtype).reshape(host.shape) - np.copyto(view, host) - offset += _aligned(host.nbytes) - core = view - staged.append(core) - return staged, buffer + if keep_device and isinstance(core, self._xp.ndarray): + views.append(None) + device.append(core) + offsets.append(offset) + continue + host = np.asarray(to_numpy(core)) + view = buffer[offset : offset + host.nbytes] + view = view.view(host.dtype).reshape(host.shape) + np.copyto(view, host) + views.append(view) + device.append(None) + offsets.append(offset) + offset += _aligned(host.nbytes) + return _Staged(views, device, offsets, offset, buffer) + + def _release(self, k: int, event: Any) -> None: # noqa: ANN401 + """Record that a rank is done with its copies of site ``k``.""" + with self._lock: + self._taken[k] += 1 + self._events[k].append(event) + if self._taken[k] < self._group.size: + return + events, self._events[k] = self._events[k], [] + buffer = self._staged[k].result().buffer + self._slots.put((buffer, events)) + + # --- reading --- + + def next(self, rank: int = 0) -> tuple[NDArray, ...]: + """Return the padded cores of the next site in order, on the rank's device. + + Args: + rank: The calling rank. + + Returns: + One rank-4 ``(l, r, u, d)`` core per layer. + """ + k = self._next[rank] + self._next[rank] += 1 + j = self._order[k] + future = self._staged[k] + start = perf_counter() + while True: + try: + staged = future.result(timeout=_POLL_SECONDS) + break + except TimeoutError: + if self._closing.is_set(): + raise GroupAbortedError from None + self.stall_seconds[rank] += perf_counter() - start + if staged.buffer is None: + self._release(k, NullEvent()) + return tuple(staged.cores) + cores = self._upload(rank, k, staged) + return tuple( + pad_site(core, kind, j, self._last) + for core, kind in zip(cores, self._kinds) + ) + + def _upload(self, rank: int, k: int, staged: _Staged) -> list[NDArray]: + xp, stream = self._xp, current_stream(self._xp) + if not self._group.split_uploads: + cores = [ + device if view is None else to_device_async(view, xp, stream) + for view, device in zip(staged.cores, staged.device) + ] + # The buffer is refilled only once these copies are done. + self._release(k, stream.record()) + return cores + lo, hi = blocks(staged.n_bytes, self._group.size)[rank] + local = to_device_async(staged.buffer[lo:hi], xp, stream) + self._release(k, stream.record()) + flat = self._group.allgather(rank, local) + del local + return [ + flat[offset : offset + view.nbytes].view(view.dtype).reshape(view.shape) + for view, offset in zip(staged.cores, staged.offsets) + ] def _read(core: Site) -> NDArray: diff --git a/src/src_method/_sweep.py b/src/src_method/_sweep.py index c0dac32..975a389 100644 --- a/src/src_method/_sweep.py +++ b/src/src_method/_sweep.py @@ -13,19 +13,28 @@ Every contraction runs in batches sized by `src_method._plan.make_plan` to the memory budgets. The cores are read one site at a time (`SiteSource`) and the environments live on the device, in host memory or on disk (`EnvironmentStore`). + +With several devices (`DeviceGroup`), each owns a contiguous block of the sketch +columns. The left-to-right pass needs no communication: the columns are +independent. At every site of the right-to-left pass, the blocks of the sketch are +gathered on rank 0, which runs the QR and scatters blocks of rows of the output +core; each device projects its rows of the new ``S``, and the rows are gathered on +every device. The random draws, and hence the result, are those of one device. """ from __future__ import annotations import logging +import threading from time import perf_counter_ns from typing import TYPE_CHECKING import numpy as np +from ._group import DeviceGroup, blocks from ._kernels import SiteKernels -from ._plan import make_plan, resolve_budgets -from ._sites import SiteSource, padded_shapes, site_bytes +from ._plan import make_plan, resolve_budgets, resolve_devices +from ._sites import SiteSource, padded_shapes, site_bytes, sweep_order from ._store import EnvironmentStore from ._tensor_train import unpad from .utils import ( @@ -43,7 +52,7 @@ from numpy.typing import DTypeLike, NDArray - from ._plan import Plan, Resources + from ._plan import Budgets, Plan, Resources from ._tensor_train import Site, TrainKind logger = logging.getLogger(__name__) @@ -80,7 +89,8 @@ def sweep( xp: Array module (``numpy`` or ``cupy``). cutoff: Relative singular-value cutoff for adaptive bond truncation. dtype: The data type of the sketches. - resources: Memory budgets and scratch space; detected when ``None``. + resources: Memory budgets, scratch space and devices; detected when + ``None``. Returns: The site arrays of the compressed train in right-canonical form, as numpy @@ -90,13 +100,18 @@ def sweep( # The dtype of the environments and the output: the sketches promoted by the # cores, as the contractions would. work = np.result_type(dtype, *(site.dtype for layer in layers for site in layer)) - budgets = resolve_budgets(resources, xp) - plan = make_plan(shapes, site_bytes(layers), chi_out, work, budgets) + devices = resolve_devices(resources, xp) + group = DeviceGroup(xp, devices) + budgets = resolve_budgets(resources, xp, devices) + plan = make_plan( + shapes, site_bytes(layers), chi_out, work, budgets, devices=group.size + ) # Guarded: the per-site lists walk the whole plan, wasted work unless logged. if logger.isEnabledFor(logging.DEBUG): logger.debug( - "SRC plan: prefetch=%s, device peak=%d B, host peak=%d B, disk=%d B, " - "scratch=%s, tiers=%s, batches (env, sketch, project)=%s", + "SRC plan: devices=%s, prefetch=%s, device peak=%d B, host peak=%d B, " + "disk=%d B, scratch=%s, tiers=%s, batches (env, sketch, project)=%s", + list(devices), plan.prefetch, plan.device_peak, plan.host_peak, @@ -108,113 +123,196 @@ def sweep( for site in plan.sites ], ) - env_shapes = [(chi_out, *(s[0] for s in site)) for site in shapes] - kernels = SiteKernels(len(layers)) - with ( - device_pool_limit(xp, budgets.device_cap), - SiteSource(layers, xp, depth=plan.prefetch) as source, - EnvironmentStore( - plan, env_shapes, work, xp, budgets.scratch_dir, copy_stream=new_stream(xp) - ) as store, - ): - tms = perf_counter_ns() - _left_to_right( - kernels, source, store, plan, chi_out=chi_out, prng=prng, xp=xp, dtype=dtype - ) - logger.debug("Left-to-right sweep: %.3f s", (perf_counter_ns() - tms) * 1e-9) - tms = perf_counter_ns() - eta = _right_to_left( - kernels, - source, - store, - plan, - chi_out=chi_out, - cutoff=cutoff, - xp=xp, - dtype=work, - ) - logger.debug("Right-to-left sweep: %.3f s", (perf_counter_ns() - tms) * 1e-9) - logger.debug( - "SRC stalls: sites %.3f s, environments %.3f s", - source.stall_seconds, - store.stall_seconds, - ) - logger.debug("Device pool: %d B", device_pool_bytes(xp)) + sweep_ = _Sweep( + group, + plan, + budgets, + shapes, + chi_out=chi_out, + cutoff=cutoff, + sketches=_Sketches(prng, chi_out, dtype, group.size), + work=work, + depth=len(layers), + ) + with SiteSource( + layers, xp, sweep_order(len(shapes)), depth=plan.prefetch, group=group + ) as source: + eta = group.run(lambda rank: sweep_.run(rank, source))[0] + assert eta is not None # noqa: S101 (rank 0 returns the cores) return [to_numpy(site) for site in unpad(eta, kind)] -def _left_to_right( - kernels: SiteKernels, - source: SiteSource, - store: EnvironmentStore, - plan: Plan, - *, - chi_out: int, - prng: np.random.Generator, - xp: ModuleType, - dtype: DTypeLike, -) -> None: - """Build the environments ``C_1 .. C_{n-1}`` into the store.""" - depth = len(kernels.eqs.ltr.split(",")) - 2 - first_env = xp.ones((chi_out,) + (1,) * depth, dtype=dtype) - n_sites = len(source) - for j in range(n_sites - 1): - cores = source[j] - if plan.prefetch: - source.prefetch(j + 1) - up, down = cores[0].shape[2], cores[-1].shape[3] - omega = gaussian_sketch(prng, (chi_out, up, down), dtype, xp) - batches = list(_ranges(chi_out, plan.sites[j].env_batch)) - if j > 0: - store.prefetch(j, batches) - for lo, hi in batches: - env = first_env[lo:hi] if j == 0 else store.get(j, lo, hi) - store.put(j + 1, lo, hi, kernels.env(env, omega[lo:hi], cores)) +class _Sketches: + """The Gaussian tensors of the left-to-right pass, shared by the devices. + + Every device asks for the sites in order, and the first one at a site draws its + tensor, so the draws come in the order of a single device, whatever the + timing. A tensor is dropped once every device took it. + """ + def __init__( + self, prng: np.random.Generator, chi: int, dtype: DTypeLike, ranks: int + ) -> None: + self._prng, self._chi, self._dtype, self._ranks = prng, chi, dtype, ranks + self._drawn: dict[int, tuple[np.ndarray, list[int]]] = {} + self._lock = threading.Lock() -def _right_to_left( - kernels: SiteKernels, - source: SiteSource, - store: EnvironmentStore, - plan: Plan, - *, - chi_out: int, - cutoff: float, - xp: ModuleType, - dtype: np.dtype, -) -> list[NDArray]: - """Build the output cores, host-side, from the last site to the first. + def get(self, j: int, up: int, down: int) -> np.ndarray: + """Return the host tensor ``(chi, up, down)`` of site ``j``.""" + with self._lock: + if j not in self._drawn: + omega = gaussian_sketch( + self._prng, (self._chi, up, down), self._dtype, np + ) + self._drawn[j] = (omega, [self._ranks]) + omega, left = self._drawn[j] + left[0] -= 1 + if left[0] == 0: + del self._drawn[j] + return omega - ``dtype`` is the working dtype of the sweep, that of the environments. - """ - depth = len(kernels.eqs.ltr.split(",")) - 2 - n_sites = len(source) - eta_reversed: list[NDArray] = [] - proj = xp.ones((1,) * (depth + 1), dtype=dtype) - for j in range(n_sites - 1, 0, -1): - cores = source[j] - if plan.prefetch: - source.prefetch(j - 1) - site = plan.sites[j] - batches = list(_ranges(chi_out, site.sketch_batch)) - store.prefetch(j, batches) + +class _Sweep: + """The per-device part of a sweep; `run` executes it on one rank.""" + + def __init__( + self, + group: DeviceGroup, + plan: Plan, + budgets: Budgets, + shapes: Sequence[tuple[tuple[int, ...], ...]], + *, + chi_out: int, + cutoff: float, + sketches: _Sketches, + work: np.dtype, + depth: int, + ) -> None: + self.group, self.plan, self.budgets, self.shapes = group, plan, budgets, shapes + self.chi_out, self.cutoff, self.sketches = chi_out, cutoff, sketches + self.work, self.depth = work, depth + self.columns = blocks(chi_out, group.size) + + def _tag(self, rank: int) -> str: + if self.group.size == 1: + return "" + return f"Device {self.group.devices[rank]}: " + + def run(self, rank: int, source: SiteSource) -> list[NDArray] | None: + """Run both passes on ``rank``; rank 0 returns the output cores.""" + xp, tag = self.group.xp, self._tag(rank) + lo, hi = self.columns[rank] + env_shapes = [(hi - lo, *(s[0] for s in site)) for site in self.shapes] + kernels = SiteKernels(self.depth) + with ( + device_pool_limit(xp, self.budgets.device_cap), + EnvironmentStore( + self.plan, + env_shapes, + self.work, + xp, + self.budgets.scratch_dir, + copy_stream=new_stream(xp), + ) as store, + ): + tms = perf_counter_ns() + self._left_to_right(rank, kernels, source, store) + logger.debug( + "%sLeft-to-right sweep: %.3f s", tag, (perf_counter_ns() - tms) * 1e-9 + ) + tms = perf_counter_ns() + eta = self._right_to_left(rank, kernels, source, store) + logger.debug( + "%sRight-to-left sweep: %.3f s", tag, (perf_counter_ns() - tms) * 1e-9 + ) + logger.debug( + "%sSRC stalls: sites %.3f s, environments %.3f s", + tag, + source.stall_seconds[rank], + store.stall_seconds, + ) + logger.debug("%sDevice pool: %d B", tag, device_pool_bytes(xp)) + return eta + + def _left_to_right( + self, + rank: int, + kernels: SiteKernels, + source: SiteSource, + store: EnvironmentStore, + ) -> None: + """Build this device's columns of the environments ``C_1 .. C_{n-1}``.""" + xp, plan = self.group.xp, self.plan + col_lo, col_hi = self.columns[rank] + n_local = col_hi - col_lo + first_env = xp.ones((n_local,) + (1,) * self.depth, dtype=self.work) + for j in range(len(source) - 1): + cores = source.next(rank) + up, down = cores[0].shape[2], cores[-1].shape[3] + omega = xp.asarray(self.sketches.get(j, up, down)[col_lo:col_hi]) + batches = list(_ranges(n_local, plan.sites[j].env_batch)) + if j > 0: + store.prefetch(j, batches) + for lo, hi in batches: + env = first_env[lo:hi] if j == 0 else store.get(j, lo, hi) + store.put(j + 1, lo, hi, kernels.env(env, omega[lo:hi], cores)) + + def _right_to_left( + self, + rank: int, + kernels: SiteKernels, + source: SiteSource, + store: EnvironmentStore, + ) -> list[NDArray] | None: + """Build the output cores, host-side on rank 0, from the last site down.""" + xp, plan, group, dtype = self.group.xp, self.plan, self.group, self.work + root = rank == 0 + col_lo, col_hi = self.columns[rank] + n_local = col_hi - col_lo + eta_reversed: list[NDArray] = [] + proj = xp.ones((1,) * (self.depth + 1), dtype=dtype) + for j in range(len(source) - 1, 0, -1): + cores = source.next(rank) + site = plan.sites[j] + batches = list(_ranges(n_local, site.sketch_batch)) + store.prefetch(j, batches) + up, down = cores[0].shape[2], cores[-1].shape[3] + # Sketch columns first, so that each device's block is contiguous. + local = xp.empty((n_local, proj.shape[0], up, down), dtype=dtype) + for lo, hi in batches: + part = kernels.sketch(store.get(j, lo, hi), cores, proj) + local[lo:hi] = xp.moveaxis(part, -1, 0) + store.drop(j) + sketch = group.gather(rank, local) + del local + eta_full = None + if root: + assert sketch is not None # noqa: S101 (gathered on rank 0) + rows = proj.shape[0] * up * down + Q = truncated_qr(sketch.reshape(self.chi_out, rows).T, self.cutoff, xp) + del sketch + eta_full = xp.ascontiguousarray(Q.T).reshape( + -1, proj.shape[0], up, down + ) + del Q + eta_local = group.scatter(rank, eta_full) + if eta_full is not None: + eta_reversed.append(to_numpy(eta_full)) + del eta_full + new_local = xp.empty( + (eta_local.shape[0], *(c.shape[0] for c in cores)), dtype=dtype + ) + for lo, hi in _ranges(eta_local.shape[0], site.project_batch): + new_local[lo:hi] = kernels.project(eta_local[lo:hi].conj(), cores, proj) + del proj, eta_local + proj = group.allgather(rank, new_local) + del new_local + # Every device takes the first site, so that the source can release it. + cores = source.next(rank) + if not root: + return None up, down = cores[0].shape[2], cores[-1].shape[3] - sketch = xp.empty((proj.shape[0], up, down, chi_out), dtype=dtype) - for lo, hi in batches: - sketch[..., lo:hi] = kernels.sketch(store.get(j, lo, hi), cores, proj) - store.drop(j) - rows = sketch.shape[0] * up * down - Q = truncated_qr(sketch.reshape(rows, chi_out), cutoff, xp) - del sketch - eta_j = Q.reshape(proj.shape[0], up, down, Q.shape[1]).transpose(3, 0, 1, 2) - new_proj = xp.empty((Q.shape[1], *(c.shape[0] for c in cores)), dtype=dtype) - for lo, hi in _ranges(Q.shape[1], site.project_batch): - new_proj[lo:hi] = kernels.project(eta_j[lo:hi].conj(), cores, proj) - proj = new_proj - eta_reversed.append(to_numpy(eta_j)) - cores = source[0] - up, down = cores[0].shape[2], cores[-1].shape[3] - first = xp.empty((proj.shape[0], up, down), dtype=dtype) - for lo, hi in _ranges(proj.shape[0], plan.sites[0].project_batch): - first[lo:hi] = kernels.first(cores, proj[lo:hi]).reshape(hi - lo, up, down) - return [to_numpy(first)[None], *reversed(eta_reversed)] + first = xp.empty((proj.shape[0], up, down), dtype=dtype) + for lo, hi in _ranges(proj.shape[0], plan.sites[0].project_batch): + first[lo:hi] = kernels.first(cores, proj[lo:hi]).reshape(hi - lo, up, down) + return [to_numpy(first)[None], *reversed(eta_reversed)] diff --git a/src/src_method/utils/__init__.py b/src/src_method/utils/__init__.py index c3e661c..c631795 100644 --- a/src/src_method/utils/__init__.py +++ b/src/src_method/utils/__init__.py @@ -5,11 +5,15 @@ from ._backend import ( NullEvent, NullStream, + copy_into, + current_device, current_stream, default_rng, + device_count, device_memory, device_pool_bytes, device_pool_limit, + enable_peer_access, gaussian_sketch, get_xp, host_memory_available, @@ -20,17 +24,22 @@ to_device_async, to_host_async, to_numpy, + use_device, ) from .linalg import truncated_qr __all__ = [ "NullEvent", "NullStream", + "copy_into", + "current_device", "current_stream", "default_rng", + "device_count", "device_memory", "device_pool_bytes", "device_pool_limit", + "enable_peer_access", "gaussian_sketch", "get_xp", "host_memory_available", @@ -42,4 +51,5 @@ "to_host_async", "to_numpy", "truncated_qr", + "use_device", ] diff --git a/src/src_method/utils/_backend.py b/src/src_method/utils/_backend.py index d249412..9dd6129 100644 --- a/src/src_method/utils/_backend.py +++ b/src/src_method/utils/_backend.py @@ -168,6 +168,129 @@ def to_host_async(device: NDArray, out: np.ndarray, stream: Any) -> None: # noq device.get(stream=stream, out=out, blocking=False) +def current_device(xp: ModuleType) -> int: + """Return the id of the current device, or 0 on the host.""" + if is_host(xp): + return 0 + return xp.cuda.runtime.getDevice() + + +def device_count(xp: ModuleType) -> int | None: + """Return the number of visible devices, or ``None`` on the host. + + The host has no limit: its devices are simulated (see `use_device`). + """ + if is_host(xp): + return None + return xp.cuda.runtime.getDeviceCount() + + +@contextmanager +def use_device(xp: ModuleType, device: int) -> Iterator[None]: + """Make ``device`` current in the calling thread for the duration of the block. + + On the host backend the devices of a group are simulated: they share the host, + so this is a no-op. + + Args: + xp: Array module (``numpy`` or ``cupy``). + device: The device id. + + Yields: + Nothing. + """ + if is_host(xp): + yield + return + with xp.cuda.Device(device): + yield + + +# cudaErrorPeerAccessAlreadyEnabled: another library, or an earlier call, did it. +_PEER_ALREADY_ENABLED = 704 + + +def enable_peer_access(xp: ModuleType, devices: Sequence[int]) -> bool: + """Enable direct access between every pair of ``devices``, where supported. + + Peer copies work without it, but are then staged through the host. + + Args: + xp: Array module (``numpy`` or ``cupy``). + devices: The device ids. + + Returns: + Whether every pair has direct access; always true on the host. + """ + if is_host(xp): + return True + runtime = xp.cuda.runtime + complete = True + for device in devices: + with xp.cuda.Device(device): + for peer in devices: + if peer == device: + continue + if not runtime.deviceCanAccessPeer(device, peer): + complete = False + continue + try: + runtime.deviceEnablePeerAccess(peer) + except runtime.CUDARuntimeError as err: + if getattr(err, "status", None) != _PEER_ALREADY_ENABLED: + raise + return complete + + +def copy_into(dst: NDArray, src: NDArray, xp: ModuleType, stream: Any) -> None: # noqa: ANN401 + """Copy ``src`` into ``dst``, possibly on another device, on ``stream``. + + Both arrays must be C-contiguous with the same shape and dtype. Across devices + the copy is a peer copy: over NVLink or PCIe peer-to-peer when enabled (see + `enable_peer_access`), staged through the host otherwise. ``stream`` must belong + to the device of ``dst``, and ``src`` must be ready on it (see + ``stream.wait_event``). + + Args: + dst: The destination. + src: The source. + xp: Array module (``numpy`` or ``cupy``). + stream: The stream that performs the copy. + """ + if dst.shape != src.shape or dst.dtype != src.dtype: + msg = f"Cannot copy {src.shape} {src.dtype} into {dst.shape} {dst.dtype}." + raise ValueError(msg) + if dst.nbytes == 0: + return + if is_host(xp): + np.copyto(dst, src) + return + if not (dst.flags.c_contiguous and src.flags.c_contiguous): + msg = "Device copies need C-contiguous arrays." + raise ValueError(msg) + # CuPy arrays: ``data`` is a memory pointer and ``device`` a CUDA device. + to: Any = dst + source: Any = src + runtime = xp.cuda.runtime + if to.device.id == source.device.id: + runtime.memcpyAsync( + to.data.ptr, + source.data.ptr, + to.nbytes, + runtime.memcpyDeviceToDevice, + stream.ptr, + ) + else: + runtime.memcpyPeerAsync( + to.data.ptr, + to.device.id, + source.data.ptr, + source.device.id, + to.nbytes, + stream.ptr, + ) + + def device_memory(xp: ModuleType) -> tuple[int, int]: """Return the device bytes available to the sweep and the device total. diff --git a/tests/test_backend.py b/tests/test_backend.py index d69acd3..d467961 100644 --- a/tests/test_backend.py +++ b/tests/test_backend.py @@ -1,19 +1,25 @@ """Test the host-side behaviour of the backend helpers.""" import numpy as np +import pytest from src_method.utils import ( NullEvent, NullStream, + copy_into, + current_device, current_stream, + device_count, device_pool_bytes, device_pool_limit, + enable_peer_access, host_memory_available, is_host, new_stream, pinned_empty, to_device_async, to_host_async, + use_device, ) @@ -72,3 +78,26 @@ def test_host_memory_reads_mem_available(tmp_path): def test_host_memory_falls_back_without_meminfo(tmp_path): assert host_memory_available(str(tmp_path / "missing")) > 0 + + +def test_host_devices_are_simulated(): + assert device_count(np) is None + assert current_device(np) == 0 + assert enable_peer_access(np, [0, 1, 2]) + with use_device(np, 5): + assert current_device(np) == 0 + + +def test_copy_into_on_host(): + dst = np.empty((2, 3), dtype=complex) + + copy_into(dst, np.arange(6.0).reshape(2, 3).astype(complex), np, NullStream()) + + np.testing.assert_array_equal(dst, np.arange(6.0).reshape(2, 3)) + + +def test_copy_into_rejects_mismatches(): + with pytest.raises(ValueError, match="Cannot copy"): + copy_into(np.empty(3), np.empty(4), np, NullStream()) + with pytest.raises(ValueError, match="Cannot copy"): + copy_into(np.empty(3), np.empty(3, dtype=np.float32), np, NullStream()) diff --git a/tests/test_gpu_backend.py b/tests/test_gpu_backend.py index 6b4fdb6..a33652b 100644 --- a/tests/test_gpu_backend.py +++ b/tests/test_gpu_backend.py @@ -317,3 +317,116 @@ def spy(*_args: object) -> None: _, total = cupy.cuda.runtime.memGetInfo() margin = max(int(GPU_MARGIN_FRACTION * total), GPU_MARGIN_MIN) assert limits == [used + budget + margin] + + +# -------------------- +# --- Several GPUs --- +# -------------------- + + +@pytest.fixture +def two_gpus() -> None: + """Skip unless at least two GPUs are visible.""" + if cupy.cuda.runtime.getDeviceCount() < 2: + pytest.skip("Needs at least two GPUs.") + + +@pytest.mark.usefixtures("two_gpus") +@pytest.mark.parametrize("uploads", ["split", "full"]) +def test_two_gpus_match_one( + tmp_path, plans: list[Plan], monkeypatch: pytest.MonkeyPatch, uploads: str +) -> None: + """Two GPUs, with or without split uploads, give the operator of one. + + An odd ``chi`` gives uneven blocks, and the budgets every tier. + """ + rng = np.random.default_rng(7) + stack = [random_mpo_arrays([4, 8, 8, 8, 4], rng) for _ in range(4)] + chi = 63 + if uploads == "full": + monkeypatch.setattr( + "src_method._group.enable_peer_access", lambda _xp, _devices: False + ) + shapes, sizes = padded_shapes(stack), site_bytes(stack) + roomy = make_plan( + shapes, + sizes, + chi, + np.complex128, + Budgets(10**10, 10**10, 10**12, tmp_path, unified=False), + devices=2, + ) + env = (chi // 2 + 1) * 8**4 * 16 # one bulk environment on one GPU + # Checked with make_plan: one environment per GPU on the device, two in host + # memory, the oldest two on disk. + tight = Resources( + gpu_memory=roomy.device_peak - 2 * env, + host_memory=roomy.host_peak + 12 * env, + scratch_dir=tmp_path, + devices=2, + ) + + one = src(*stack, chi_out=chi, seed=3, dtype=np.complex128, device="gpu") + two = src( + *stack, chi_out=chi, seed=3, dtype=np.complex128, device="gpu", resources=tight + ) + + assert {site.tier for site in plans[-1].sites} == {"device", "host", "disk"} + reference = dense_mpo(one) + error = np.linalg.norm(dense_mpo(two) - reference) / np.linalg.norm(reference) + assert error < 1e-10 + assert list(tmp_path.iterdir()) == [] + + +@pytest.mark.usefixtures("two_gpus") +def test_two_gpus_take_device_inputs() -> None: + """Cores given on one GPU reach the other through the staging buffer.""" + rng = np.random.default_rng(10) + stack = [random_mpo_arrays([2, 3, 3, 2], rng) for _ in range(2)] + with cupy.cuda.Device(1): + on_device = [[cupy.asarray(t) for t in train] for train in stack] + + host = src(*stack, chi_out=8, seed=0, device="gpu") + both = src( + *on_device, chi_out=8, seed=0, device="gpu", resources=Resources(devices=2) + ) + + reference = dense_mpo(host) + assert ( + np.linalg.norm(dense_mpo(both) - reference) / np.linalg.norm(reference) < 1e-10 + ) + + +@pytest.mark.usefixtures("two_gpus") +def test_two_gpus_raise_the_first_error(monkeypatch: pytest.MonkeyPatch) -> None: + """An error on one GPU stops both, restores the pool limits and is raised.""" + rng = np.random.default_rng(11) + stack = [random_mpo_arrays([2, 3, 3, 2], rng) for _ in range(2)] + limits = [] + for d in range(2): + with cupy.cuda.Device(d): + limits.append(cupy.get_default_memory_pool().get_limit()) + + def boom(self: SiteKernels, *args: Any) -> Any: + if cupy.cuda.runtime.getDevice() == 1: + msg = "boom on GPU 1" + raise RuntimeError(msg) + return original(self, *args) + + original = SiteKernels.sketch + monkeypatch.setattr(SiteKernels, "sketch", boom) + with pytest.raises(RuntimeError, match="boom on GPU 1"): + src(*stack, chi_out=4, seed=0, device="gpu", resources=Resources(devices=2)) + for d in range(2): + with cupy.cuda.Device(d): + assert cupy.get_default_memory_pool().get_limit() == limits[d] + + +@pytest.mark.usefixtures("two_gpus") +def test_too_many_devices_raise() -> None: + count = cupy.cuda.runtime.getDeviceCount() + rng = np.random.default_rng(12) + stack = [random_mpo_arrays([2, 3, 3, 2], rng) for _ in range(2)] + + with pytest.raises(ValueError, match="visible"): + src(*stack, chi_out=4, device="gpu", resources=Resources(devices=count + 1)) diff --git a/tests/test_group.py b/tests/test_group.py new file mode 100644 index 0000000..c0bc6a0 --- /dev/null +++ b/tests/test_group.py @@ -0,0 +1,108 @@ +"""Test the device group: work splitting, collectives and failure handling.""" + +import threading + +import numpy as np +import pytest + +from src_method._group import DeviceGroup, GroupAbortedError, blocks + + +@pytest.mark.parametrize("n", [0, 1, 7, 8, 2000]) +@pytest.mark.parametrize("size", [1, 2, 3, 8]) +def test_blocks_partition_in_rank_order(n, size): + parts = blocks(n, size) + + assert len(parts) == size + assert [i for lo, hi in parts for i in range(lo, hi)] == list(range(n)) + sizes = [hi - lo for lo, hi in parts] + assert max(sizes) - min(sizes) <= 1 + assert sizes == sorted(sizes, reverse=True) + + +def test_a_group_of_one_runs_inline(): + group = DeviceGroup(np, [0]) + local = np.arange(6.0).reshape(3, 2) + + def fn(rank): + assert threading.current_thread() is threading.main_thread() + assert group.allgather(rank, local) is local + assert group.gather(rank, local) is local + assert group.scatter(rank, local) is local + return rank + + assert group.run(fn) == [0] + + +@pytest.mark.parametrize("size", [2, 3]) +@pytest.mark.parametrize("dtype", [np.float64, np.complex128]) +def test_collectives(size, dtype): + group = DeviceGroup(np, list(range(size))) + # Uneven blocks: rank r holds r + 1 rows. + parts = [ + (np.arange((r + 1) * 6) + 100 * r).reshape(r + 1, 2, 3) for r in range(size) + ] + parts = [p.astype(dtype) for p in parts] + expected = np.concatenate(parts) + + def fn(rank): + gathered = group.allgather(rank, parts[rank]) + on_root = group.gather(rank, parts[rank]) + mine = group.scatter(rank, expected if rank == 0 else None) + # A second round reuses the slots of the first. + again = group.allgather(rank, parts[rank]) + return gathered, on_root, mine, again + + results = group.run(fn) + + for rank, (gathered, on_root, mine, again) in enumerate(results): + np.testing.assert_array_equal(gathered, expected) + np.testing.assert_array_equal(again, expected) + assert (on_root is not None) == (rank == 0) + lo, hi = blocks(expected.shape[0], size)[rank] + np.testing.assert_array_equal(mine, expected[lo:hi]) + np.testing.assert_array_equal(results[0][1], expected) + + +def test_many_collectives_in_a_row(): + group = DeviceGroup(np, [0, 1, 2]) + + def fn(rank): + total = np.zeros(3) + for step in range(200): + total += group.allgather(rank, np.array([rank + step], dtype=float)) + return total + + for total in group.run(fn): + np.testing.assert_array_equal(total, [19900, 20100, 20300]) + + +def test_a_failing_rank_releases_the_others(): + group = DeviceGroup(np, [0, 1, 2]) + + def fn(rank): + if rank == 1: + msg = "boom on rank 1" + raise ValueError(msg) + # Never completes: rank 1 does not join. + return group.allgather(rank, np.zeros(1)) + + with pytest.raises(ValueError, match="boom on rank 1"): + group.run(fn) + assert group.aborted.is_set() + + +def test_abort_wakes_registered_waiters(): + group = DeviceGroup(np, [0, 1]) + released = threading.Event() + group.on_abort(released.set) + + def fn(rank): + if rank == 0: + msg = "boom" + raise RuntimeError(msg) + released.wait() + raise GroupAbortedError + + with pytest.raises(RuntimeError, match="boom"): + group.run(fn) diff --git a/tests/test_plan.py b/tests/test_plan.py index 67725c2..67fb38c 100644 --- a/tests/test_plan.py +++ b/tests/test_plan.py @@ -1,5 +1,6 @@ """Test the memory planner: sizes, budgets, batches and tiers.""" +from contextlib import contextmanager from pathlib import Path from types import ModuleType @@ -68,10 +69,29 @@ def test_resolve_budgets_detects_host_memory(monkeypatch, tmp_path): assert budgets.host == 9 * GB -def test_resolve_budgets_detects_device_memory(monkeypatch, tmp_path): - fake_xp = ModuleType("fake_xp") +@pytest.fixture +def fake_gpus(monkeypatch): + """Fake devices for the planner; returns a dict of ``id -> (free, total)``.""" + memory = {0: (30 * GB, 40 * GB)} + current = [0] + + @contextmanager + def use_device(_xp, device): + previous, current[0] = current[0], device + try: + yield + finally: + current[0] = previous + monkeypatch.setattr(plan_module, "is_host", lambda _xp: False) - monkeypatch.setattr(plan_module, "device_memory", lambda _xp: (30 * GB, 40 * GB)) + monkeypatch.setattr(plan_module, "use_device", use_device) + monkeypatch.setattr(plan_module, "device_memory", lambda _xp: memory[current[0]]) + return memory + + +@pytest.mark.usefixtures("fake_gpus") +def test_resolve_budgets_detects_device_memory(tmp_path): + fake_xp = ModuleType("fake_xp") budgets = resolve_budgets( Resources(host_memory="1GB", scratch_dir=tmp_path), fake_xp @@ -82,10 +102,9 @@ def test_resolve_budgets_detects_device_memory(monkeypatch, tmp_path): assert budgets.device_cap == 30 * GB # the pool may grow into the margin -def test_resolve_budgets_caps_explicit_device_memory_with_margin(monkeypatch, tmp_path): +@pytest.mark.usefixtures("fake_gpus") +def test_resolve_budgets_caps_explicit_device_memory_with_margin(tmp_path): fake_xp = ModuleType("fake_xp") - monkeypatch.setattr(plan_module, "is_host", lambda _xp: False) - monkeypatch.setattr(plan_module, "device_memory", lambda _xp: (30 * GB, 40 * GB)) budgets = resolve_budgets( Resources(gpu_memory="10GB", host_memory="1GB", scratch_dir=tmp_path), @@ -96,6 +115,28 @@ def test_resolve_budgets_caps_explicit_device_memory_with_margin(monkeypatch, tm assert budgets.device_cap == 14 * GB +def test_resolve_budgets_takes_the_smallest_device(fake_gpus, tmp_path): + fake_gpus[1] = (20 * GB, 40 * GB) + + budgets = resolve_budgets( + Resources(host_memory="1GB", scratch_dir=tmp_path), + ModuleType("fake_xp"), + devices=(0, 1), + ) + + assert budgets.device == 16 * GB + assert budgets.device_cap == 20 * GB + + +def test_resolve_budgets_splits_the_host_among_simulated_devices(tmp_path): + budgets = resolve_budgets( + Resources(host_memory="8GB", scratch_dir=tmp_path), np, devices=(0, 1, 2, 3) + ) + + assert budgets.host == 8 * GB + assert budgets.device == 2 * GB + + def test_resolve_budgets_leaves_the_host_uncapped(tmp_path): budgets = resolve_budgets(Resources(host_memory="1GB", scratch_dir=tmp_path), np) @@ -251,3 +292,62 @@ def test_prefetch_is_dropped_before_giving_up(): shapes, site_bytes(shapes), chi, np.complex128, budgets(lo - cores // 2) ) assert dropped.prefetch == 0 + + +def test_devices_split_the_environments_and_the_batches(): + shapes = mpo_stack_shapes(8, [4, 4, 64, 4]) + chi = 256 + env = chi * 4 * 4 * 64 * 4 * 16 # one bulk environment on one device + one = make_plan(shapes, site_bytes(shapes), chi, np.complex128, budgets(10 * GB)) + # A device budget that holds half of one device's environments on four devices + # holds all of them. + tight = one.device_peak - 3 * env + + single = make_plan( + shapes, site_bytes(shapes), chi, np.complex128, budgets(tight, host=10 * GB) + ) + four = make_plan( + shapes, + site_bytes(shapes), + chi, + np.complex128, + budgets(tight, host=10 * GB), + devices=4, + ) + + assert "host" in {s.tier for s in single.sites} + assert {s.tier for s in four.sites} == {"device"} + for site in four.sites[1:]: + assert site.sketch_batch <= chi // 4 + assert site.project_batch <= chi // 4 + assert four.sites[0].env_batch <= chi // 4 + # The first site runs on rank 0 alone, on every row. + assert four.sites[0].project_batch == chi + + +def test_devices_share_host_memory_and_disk(): + shapes = mpo_stack_shapes(8, [4, 4, 64, 4]) + chi = 256 + env = chi * 4 * 4 * 64 * 4 * 16 // 2 # one bulk environment on each of 2 devices + roomy = make_plan( + shapes, site_bytes(shapes), chi, np.complex128, budgets(10 * GB), devices=2 + ) + device = roomy.device_peak - 6 * env + # Each device stages four batches of 128 columns, one environment each. + staging = 4 * env + host = roomy.host_peak + 2 * staging + 4 * env + + plan = make_plan( + shapes, + site_bytes(shapes), + chi, + np.complex128, + budgets(device, host=host), + devices=2, + ) + + tiers = [s.tier for s in plan.sites[1:]] + # Host room for four environments in all: two per device. + assert tiers.count("host") == 2 + assert plan.disk_bytes == 2 * env * tiers.count("disk") + assert plan.host_peak <= host diff --git a/tests/test_sites.py b/tests/test_sites.py index f9a0364..1f45669 100644 --- a/tests/test_sites.py +++ b/tests/test_sites.py @@ -6,7 +6,8 @@ import pytest from src_method import src -from src_method._sites import SiteSource, padded_shapes, site_bytes +from src_method._group import DeviceGroup +from src_method._sites import SiteSource, padded_shapes, site_bytes, sweep_order from src_method._tensor_train import pad @@ -62,19 +63,57 @@ def test_padded_shapes_and_bytes_without_reading(): @pytest.mark.parametrize("depth", [0, 1]) -def test_sites_are_read_when_requested(depth): +def test_sites_come_in_the_given_order(depth): rng = np.random.default_rng(1) train = CountingTrain(random_mpo(4, 3, rng)) - with SiteSource([train], np, depth=depth) as source: - (core,) = source[2] - assert train.reads == [0, 0, 1, 0] - source.prefetch(3) - (last,) = source[3] - assert train.reads == [0, 0, 1, 1] + with SiteSource([train], np, [2, 3, 2], depth=depth) as source: + cores = [source.next() for _ in range(3)] - np.testing.assert_array_equal(core, pad(train.sites)[2]) - np.testing.assert_array_equal(last, pad(train.sites)[3]) + assert train.reads == [0, 0, 2, 1] + for (core,), j in zip(cores, [2, 3, 2]): + np.testing.assert_array_equal(core, pad(train.sites)[j]) + + +@pytest.mark.parametrize("depth", [0, 1]) +def test_simulated_devices_share_one_read(depth): + rng = np.random.default_rng(1) + train = CountingTrain(random_mpo(4, 3, rng)) + mps = [t[..., 0] for t in random_mpo(4, 2, rng)] + group = DeviceGroup(np, [0, 1, 2]) + order = sweep_order(4) + + with SiteSource([train, mps], np, order, depth=depth, group=group) as source: + got = group.run(lambda rank: [source.next(rank) for _ in order]) + + assert train.reads == [2, 2, 2, 1] + for per_rank in got: + for cores, j in zip(per_rank, order): + np.testing.assert_array_equal(cores[0], pad(train.sites)[j]) + np.testing.assert_array_equal(cores[1], pad(mps)[j]) + + +def test_a_failed_read_reaches_every_device(): + class Broken(CountingSite): + def __array__(self, dtype=None, copy=None): + msg = "unreadable" + raise OSError(msg) + + class BrokenTrain(CountingTrain): + def __getitem__(self, j): + return Broken(self, j) if j == 1 else CountingSite(self, j) + + rng = np.random.default_rng(1) + group = DeviceGroup(np, [0, 1]) + order = sweep_order(4) + + with ( + SiteSource( + [BrokenTrain(random_mpo(4, 3, rng))], np, order, depth=1, group=group + ) as source, + pytest.raises(OSError, match="unreadable"), + ): + group.run(lambda rank: [source.next(rank) for _ in order]) def test_src_reads_each_site_once_per_pass(): diff --git a/tests/test_stack.py b/tests/test_stack.py index 693f7ca..b7d0c38 100644 --- a/tests/test_stack.py +++ b/tests/test_stack.py @@ -1,5 +1,6 @@ """Test `src` over stacks of tensor trains, and the sweep behind it.""" +import threading from functools import reduce import numpy as np @@ -7,6 +8,7 @@ import src_method._sweep as sweep_module from src_method import Resources, apply, compress, src +from src_method._kernels import SiteKernels from src_method._plan import make_plan from src_method._sweep import sweep @@ -405,3 +407,94 @@ def test_compress_passes_resources(rng): with pytest.raises(MemoryError, match="working set exceeds the budget"): compress(A, chi_out=8, resources=Resources(host_memory="1kB")) + + +# ------------------------------- +# --- Several (simulated) devices --- +# ------------------------------- + + +@pytest.mark.parametrize( + ("devices", "chi"), [(2, 16), (3, 16), (3, 17), ([2, 0, 1], 16), (4, 3)] +) +def test_devices_match_one_device(devices, chi, rng): + stack = [random_mpo([3, 4, 4, 4, 3], rng) for _ in range(4)] + + one = src(*stack, chi_out=chi, seed=5, dtype=np.complex128) + many = src( + *stack, + chi_out=chi, + seed=5, + dtype=np.complex128, + resources=Resources(devices=devices), + ) + + assert [t.shape for t in many] == [t.shape for t in one] + assert rel_error(many, dense(one)) < 1e-10 + + +@pytest.mark.parametrize("spec", ["psi", "A psi", "A B", "phi A B"]) +def test_devices_are_exact_without_truncation(spec, rng): + stack = make_stack(spec, rng) + + out = src(*stack, chi_out=CHI_EXACT, resources=Resources(devices=3)) + + assert rel_error(out, dense_stack(*stack)) < 1e-10 + + +def test_devices_with_cutoff_keep_the_bonds(rng): + stack = [random_mpo([3, 4, 4, 4, 3], rng) for _ in range(3)] + + one = src(*stack, chi_out=24, cutoff=0.05, seed=2) + many = src(*stack, chi_out=24, cutoff=0.05, seed=2, resources=Resources(devices=2)) + + assert [t.shape for t in many] == [t.shape for t in one] + assert rel_error(many, dense(one)) < 1e-10 + + +def test_devices_spill_and_clean_up(rng, tmp_path, plans): + stack = [random_mpo([3, 4, 4, 4, 3], rng) for _ in range(4)] + tight = Resources(host_memory="1MB", scratch_dir=tmp_path, devices=2) + + default = src(*stack, chi_out=16, seed=3, dtype=np.complex128) + spilled = src(*stack, chi_out=16, seed=3, dtype=np.complex128, resources=tight) + + assert "disk" in {site.tier for site in plans[1].sites} + assert rel_error(spilled, dense(default)) < 1e-10 + assert list(tmp_path.iterdir()) == [] + + +def test_an_error_on_one_device_is_raised(rng, monkeypatch, tmp_path): + stack = [random_mpo([3, 4, 4, 4, 3], rng) for _ in range(2)] + calls, lock = [0], threading.Lock() + original = SiteKernels.sketch + + def flaky(self, *args): + with lock: + calls[0] += 1 + fail = calls[0] == 2 + if fail: + msg = "sketch failed" + raise RuntimeError(msg) + return original(self, *args) + + monkeypatch.setattr(SiteKernels, "sketch", flaky) + resources = Resources(host_memory="2MB", scratch_dir=tmp_path, devices=3) + with pytest.raises(RuntimeError, match="sketch failed"): + src(*stack, chi_out=16, seed=0, resources=resources) + assert list(tmp_path.iterdir()) == [] + + +def test_devices_are_validated(rng): + stack = make_stack("A B", rng) + + with pytest.raises(ValueError, match="positive device count"): + Resources(devices=0) + with pytest.raises(ValueError, match="distinct"): + Resources(devices=[1, 1]) + with pytest.raises(TypeError, match="device count or a sequence"): + Resources(devices="0,1") # ty: ignore[invalid-argument-type] + with pytest.raises(TypeError, match="integer device ids"): + Resources(devices=[0.0, 1.0]) # ty: ignore[invalid-argument-type] + # Simulated devices have no upper limit. + src(*stack, chi_out=4, resources=Resources(devices=[0, 7]))