Distributed Simulation and Training

Distributed Simulation and Training

Understand how one quantum workload can be partitioned across processes, prepare a distributed runtime, and launch a small statevector training example with explicit ownership and correctness checks.

Version: FlagQuantum 0.3.1. Both examples were verified with two local CPU processes, Gloo, Python 3.12.14, PyTorch 2.13.0, and macOS Apple Silicon. GPU and multi-node commands below are configuration examples requiring the corresponding hardware and network; they were not executed on this machine.

Formula/code review — October 20, 2026: The equations are paired with code excerpts, and the local CPU examples were re-run with Python 3.12.14 and PyTorch 2.13.0. See the review record for executed commands, distributed checks, and untested hardware/provider paths. Earlier environment and output records are retained.

1. What Is Distributed Quantum Simulation?

Distributing a workload means assigning parts of its computation or state to different processes, called ranks, and coordinating them through communication.

Pattern	What is distributed?	Interpretation
Independent circuit jobs	Different circuits, parameter choices, or shots	Useful task parallelism; not evidence that one larger state fits.
Replicated/data-parallel training	Separate full model/state replicas process different data	Can improve throughput; each replica still needs its own state storage.
Sharded statevector	One logical amplitude vector is partitioned across ranks	Can expand state capacity, with communication and working-space costs.
Distributed MPS	Site tensors or related work have rank owners	Costs depend on bond dimensions, ordering, and boundary communication.
Distributed tensor contraction	Contraction work or slices are assigned across ranks	Requires a defined work partition and combination of partial results.

Launching several copies of a program is not, by itself, sharding. The training example here partitions one logical statevector and reports distribution_semantics="sharded_across_ranks".

For n qubits and W evenly participating ranks, a statevector shard contains roughly 2**n / W amplitudes. This is the state array only: communication buffers, backward computation, parameters, optimizer state, and other working storage also consume memory. Gates that cross partition boundaries require communication, while some operations can remain local.

Training Adds Ownership and Synchronization

A distributed training step needs consistent forward values, gradients, and parameter updates. Depending on the workflow, gradients can be combined with an all-reduce or reduced onto owners with a reduce-scatter. The optimizer updates the parameters it owns, and updated values are made available where subsequent execution needs them.

The example's ranks construct identical small circuit definitions and initial parameters. They then cooperate on a sharded quantum state. Replicating this circuit metadata is different from replicating the full amplitude vector.

2. Understand the Process Topology

Term	Meaning
RANK	Global process index, from 0 through WORLD_SIZE-1
WORLD_SIZE	Total number of worker processes
LOCAL_RANK	Process index within its host; used to choose a visible GPU
LOCAL_WORLD_SIZE	Workers launched on the local host
Node count	Number of participating hosts
Rendezvous	How workers find the same distributed job and establish membership
Process-group backend	Communication implementation, such as Gloo or NCCL

In the verified run, two processes share one host: WORLD_SIZE=2, LOCAL_WORLD_SIZE=2, node count 1. Two hosts with one process each would instead have WORLD_SIZE=2, LOCAL_WORLD_SIZE=1, node count 2.

The scripts use a uniform number of workers per host. torchrun, invoked here as python -m torch.distributed.run, supplies the rank environment variables. Run them through the launcher rather than as ordinary single-process scripts.

3. Prepare the Environment

Use Python 3.10–3.12 and the 0.3.1 core PyTorch requirement >=2.10,<2.15. The release recommends PyTorch 2.13.x or 2.14.x for new upstream CPU/CUDA environments.

CPU and Gloo

Follow Installation and Compatibility for virtual-environment creation and platform-specific commands. On Linux, for example:

python -m pip install "torch>=2.13,<2.15" --index-url https://download.pytorch.org/whl/cpu
python -m pip install "flagquantum==0.3.1" numpy
python -m pip check

On Apple Silicon, use the default PyPI index for the PyTorch command. The tested macOS PyTorch build provides Gloo. Linux is the intended environment for the NCCL examples below. A Windows package installation alone does not establish support for this distributed workflow and its timeout mechanisms; native Windows execution was not verified here.

Check the distributed backends in the selected interpreter:

python -c "import torch, torch.distributed as d; print(torch.__version__); print('Distributed:', d.is_available()); print('Gloo:', d.is_gloo_available()); print('NCCL:', d.is_nccl_available())"

The CPU verification environment reported PyTorch 2.13.0, distributed availability True, Gloo True, and NCCL False. That is sufficient for the local CPU examples.

NVIDIA GPUs and NCCL

Use Accelerator Setup to install a driver-compatible CUDA PyTorch build. For one process per GPU, every local rank must have a distinct visible CUDA device. Both scripts call torch.cuda.set_device(LOCAL_RANK) before initializing NCCL.

For two hosts, also provide matching package versions, identical scripts, reachable rendezvous and backend communication, and consistent device visibility. Record the actual network route; the presence of a network interface does not establish RDMA use. Use the release multi-node runbook for qualified workflows and their additional provisioning, staging, and checkpoint requirements.

4. Verify Communication Before Quantum Execution

First check that the ranks can initialize a process group and perform a collective. This checks the launcher and communication setup; it does not test quantum sharding.

Save as dist_preflight.py:

import argparse
import os
from datetime import timedelta

parser = argparse.ArgumentParser()
parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
args = parser.parse_args()

import torch
import torch.distributed as dist

local_rank = int(os.environ["LOCAL_RANK"])
local_world_size = int(os.environ["LOCAL_WORLD_SIZE"])
if args.device == "cuda":
    if not torch.cuda.is_available() or local_rank >= torch.cuda.device_count():
        raise RuntimeError("Each local rank requires its own visible CUDA device.")
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    backend = "nccl"
else:
    device = torch.device("cpu")
    backend = "gloo"

dist.init_process_group(backend=backend, timeout=timedelta(seconds=60))
try:
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    if world_size < 2:
        raise RuntimeError("Launch at least two workers with torchrun.")
    value = torch.tensor(rank + 1, dtype=torch.int64, device=device)
    dist.all_reduce(value, op=dist.ReduceOp.SUM)
    expected = world_size * (world_size + 1) // 2
    assert value.item() == expected
    if rank == 0:
        print("Backend:", backend)
        print("World size:", world_size)
        print("Local world size:", local_world_size)
        print("Node count:", world_size // local_world_size)
        print("Rank-sum result:", value.item())
        print("Collective preflight passed.")
finally:
    dist.destroy_process_group()

Choose an unused port and run two workers on one host. The verified launch uses an explicit IPv4 loopback address:

python -m torch.distributed.run --nnodes=1 --nproc-per-node=2 --master-addr=127.0.0.1 --master-port=29531 dist_preflight.py --device cpu

Expected rank-0 application output:

Backend: gloo
World size: 2
Local world size: 2
Node count: 1
Rank-sum result: 3
Collective preflight passed.

The two ranks contribute 1 and 2, and all-reduce returns 3. Launcher notices, networking diagnostics, or backend warnings can appear in addition to the application output.

--standalone --nnodes=1 --nproc-per-node=2 is another common single-host launcher configuration. On the verification host it encountered an IPv6 reverse-name resolution issue; the explicit IPv4 launch above passed. Use a different free port if the selected one is occupied. The process-group timeout applies after workers reach initialization; rendezvous has its own launcher-level timeout behavior.

5. Run a Truly Sharded Training Workflow

The user-facing workflow in this example is:

flagquantum.experimental.distributed.train_distributed_statevector

It is an experimental-namespace workflow with a version-specific signature. Unlike the local fq.train helper, this entry point takes a circuit, a selected Z-observable qubit, an optimizer name, and a learning rate. It does not accept an arbitrary objective= callback in 0.3.1.

We prepare three qubits using RY angles alpha and beta on qubits 0 and 1, then CX gates from each onto qubit 2. The selected expectation has a simple analytical reference:

$$\langle Z_2\rangle = \cos(\alpha)\cos(\beta).$$

Formula / mathematical operation	Corresponding code excerpt	Check / interpretation
$E(\alpha,\beta)=\cos\alpha\cos\beta$	loss = torch.cos(reference[0]) * torch.cos(reference[1])	This independent reference appears in dist_train.py.
$\nabla E=(-\sin\alpha\cos\beta,-\cos\alpha\sin\beta)$	loss.backward()	The local analytical graph provides a reference for sharded gradients and updates.
$\theta_{t+1}=\theta_t-0.1\nabla E$	optimizer.step()	The reference uses SGD; the sharded helper selects optimizer="sgd", lr=0.1.
$\theta_{\mathrm{sharded}}=\theta_{\mathrm{reference}}$	torch.allclose(learned, reference.detach(), atol=1e-5)	Both parameters and all three pre-update losses are checked. Launch with two workers as documented below.

Each training step minimizes this expectation with SGD. Two independent leaf tensors provide two optimizer-owned parameters, assigned across the two ranks. The script compares the distributed loss history and updated parameters with an independent PyTorch calculation of the analytical function. That reference does not allocate another quantum statevector.

Example 2: Three Sharded Forward/Backward/Optimizer Steps

Save as dist_train.py:

import argparse
import json
import os
from datetime import timedelta
from pathlib import Path

parser = argparse.ArgumentParser()
parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
parser.add_argument("--report-dir", default="distributed-reports")
args = parser.parse_args()

import torch
import torch.distributed as dist
import flagquantum as fq
from flagquantum.experimental.distributed import train_distributed_statevector

local_rank = int(os.environ["LOCAL_RANK"])
if args.device == "cuda":
    if not torch.cuda.is_available() or local_rank >= torch.cuda.device_count():
        raise RuntimeError("Each local rank requires its own visible CUDA device.")
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    backend = "nccl"
else:
    device = torch.device("cpu")
    backend = "gloo"

dist.init_process_group(backend=backend, timeout=timedelta(seconds=60))
try:
    if dist.get_world_size() != 2:
        raise RuntimeError("This bounded example requires exactly two workers.")
    alpha = torch.tensor(0.6, device=device, requires_grad=True)
    beta = torch.tensor(0.4, device=device, requires_grad=True)
    circuit = fq.Circuit(3, device=device).ry(0, alpha).ry(1, beta).cx(0, 2).cx(1, 2)
    training = train_distributed_statevector(
        circuit, steps=3, observable_qubit=2, optimizer="sgd", lr=0.1,
        timeout_seconds=45.0,
    )
    summary = training.summary()

    # Independent analytic reference: <Z2> = cos(alpha) * cos(beta).
    reference = torch.nn.Parameter(torch.tensor([0.6, 0.4]))
    optimizer = torch.optim.SGD([reference], lr=0.1)
    expected_losses = []
    for _ in range(3):
        optimizer.zero_grad()
        loss = torch.cos(reference[0]) * torch.cos(reference[1])
        expected_losses.append(loss.detach().item())
        loss.backward()
        optimizer.step()
    learned = torch.stack((alpha, beta)).detach().cpu()
    assert torch.allclose(torch.tensor(training.losses), torch.tensor(expected_losses), atol=1e-5)
    assert torch.allclose(learned, reference.detach(), atol=1e-5)
    assert summary["distribution_semantics"] == "sharded_across_ranks"
    assert training.completed_steps == 3
    assert training.communication_events > 0 and training.communication_bytes > 0
    after = (torch.cos(learned[0]) * torch.cos(learned[1])).item()

    report = dict(summary, device=str(device), trained_parameters=learned.tolist(), post_update_expectation=after)
    report_dir = Path(args.report_dir)
    report_dir.mkdir(parents=True, exist_ok=True)
    (report_dir / f"rank-{training.rank}.json").write_text(json.dumps(report, indent=2))
    dist.barrier()
    if training.rank == 0:
        print("Backend:", backend)
        print("World size:", training.world_size)
        print("Local world size:", training.local_world_size)
        print("Node count:", training.node_count)
        print("Distribution:", summary["distribution_semantics"])
        print("Gradient ownership:", summary["gradient_ownership_semantics"])
        print("Parameter owners:", [item["owner_rank"] for item in summary["optimizer_ownership"]])
        print("Completed steps:", training.completed_steps)
        print("Local state bytes:", training.local_state_bytes)
        print("Loss history:", [round(value, 6) for value in training.losses])
        print("Trained parameters:", [round(value.item(), 6) for value in learned])
        print(f"Post-update <Z2>: {after:.6f}")
        print("Communication observed:", training.communication_events > 0)
        print("Scalability claim allowed:", summary["scalability_claim_allowed"])
        print("Blockers:", list(summary["blockers"]))
        print("Two per-rank reports written.")
        print("Sharded training and analytic-reference checks passed.")
finally:
    dist.destroy_process_group()

Run with a fresh report directory:

python -m torch.distributed.run --nnodes=1 --nproc-per-node=2 --master-addr=127.0.0.1 --master-port=29532 dist_train.py --device cpu --report-dir distributed-reports

Expected rank-0 application output from the verified CPU run:

Backend: gloo
World size: 2
Local world size: 2
Node count: 1
Distribution: sharded_across_ranks
Gradient ownership: sharded_across_ranks
Parameter owners: [0, 1]
Completed steps: 3
Local state bytes: 32
Loss history: [0.760185, 0.721797, 0.679376]
Trained parameters: [0.765161, 0.499551]
Post-update <Z2>: 0.633130
Communication observed: True
Scalability claim allowed: False
Blockers: ['release_capacity_and_multinode_acceptance_pending']
Two per-rank reports written.
Sharded training and analytic-reference checks passed.

The scripts run in separate process groups and can be tested independently. This training script deliberately requires exactly two workers. Increasing the worker count is not a supported edit to this particular bounded example.

Numerical final digits and the available gradient-reduction route can differ with the runtime. The correctness checks use a tolerance of 1e-5, and the JSON reports preserve the actual gradient scheme. A tiny example verifies execution and semantics, not speedup or increased usable capacity.

Read the Reports Correctly

distributed-reports/rank-0.json and rank-1.json contain the actual per-rank summaries, device, trained parameters, and final analytical expectation. Use a fresh directory to avoid overwriting earlier reports.

Field	Interpretation
world_size, local_world_size, node_count	Actual job topology
distribution_semantics	Whether one logical state was sharded
optimizer_ownership	Parameter owner, reduction route, and optimizer-update metadata
gradient_ownership_semantics, gradient_reduction	Where gradients reside and how they were combined
local_state_bytes	Rank-local state-storage estimate, not total process memory
peak_memory_bytes	Reported memory metric; interpretation depends on device/runtime
communication_events, communication_bytes	Runtime-reported communication diagnostics, not an independent network profiler trace
losses	Expectation values from training forward passes, before each corresponding update
scalability_claim_allowed, blockers	Explicit claim boundary and unresolved generic workflow gates

The 32-byte local state figure corresponds to half of an eight-amplitude complex64 state. It excludes working buffers and other allocations. On CPU, the peak-memory metric reflects process memory rather than an isolated shard allocation.

The last recorded loss describes the forward pass before the last update. Post-update <Z2> evaluates the learned parameters using the analytical function after all three updates. It is not a fourth distributed forward execution.

The summary explicitly keeps scalability_claim_allowed=False and the displayed blockers. Preserve those fields; do not rewrite them because a small correctness run succeeded or because another qualified workflow has stronger evidence.

6. Launch on Two NVIDIA GPUs

After confirming that two CUDA devices are visible, run the communication check first:

python -m torch.distributed.run --nnodes=1 --nproc-per-node=2 --master-addr=127.0.0.1 --master-port=29531 dist_preflight.py --device cuda

Then run the bounded training example:

python -m torch.distributed.run --nnodes=1 --nproc-per-node=2 --master-addr=127.0.0.1 --master-port=29532 dist_train.py --device cuda --report-dir cuda-distributed-reports

For a successful run, the backend is nccl, each local rank uses its assigned GPU, the topology remains two workers on one node, and the analytical correctness checks pass. These GPU commands were not locally executed. Do not expect CPU memory metrics or a particular communication implementation to be identical on CUDA.

For the release's larger CUDA evaluation workflows, see distributed statevector topologies. Those workloads have their own sizes, prerequisites, and evidence; start with the small preflight above rather than treating a large benchmark launcher as an installation check.

7. Launch Across Two Hosts

The following manual launch uses one visible GPU per host, so total world size remains two. Put the same scripts in the working directory on both hosts and select the same compatible environment. Replace NODE0_ADDRESS with an address reachable from node 1, and choose a free rendezvous port reachable by both workers.

On node 0:

python -m torch.distributed.run --nnodes=2 --nproc-per-node=1 --node-rank=0 --master-addr=NODE0_ADDRESS --master-port=29531 dist_preflight.py --device cuda

On node 1, while node 0 is waiting:

python -m torch.distributed.run --nnodes=2 --nproc-per-node=1 --node-rank=1 --master-addr=NODE0_ADDRESS --master-port=29531 dist_preflight.py --device cuda

Successful preflight should report NCCL, world size 2, local world size 1, node count 2, and rank sum 3. After that job finishes on both hosts, launch the same pair of commands with dist_train.py in place of dist_preflight.py, using fresh report directories and a free port.

In this manual two-host layout, report files are node-local unless you choose a shared output directory. Collect rank-0.json from node 0 and rank-1.json from node 1.

These are unexecuted launch templates, not multi-node acceptance evidence. Manual launching does not replace the release runbook's qualified setup, transport checks, shared staging/checkpoint procedures, and workload-specific validation. Never use 127.0.0.1 as the address of another host. A reachable rendezvous port alone does not establish that all backend communication can pass through the network.

8. Checkpoints, Other Representations, and Limits

The JSON files above are diagnostic reports, not resumable checkpoints. The statevector training API separately provides checkpoint_dir, checkpoint_interval, and resume. Use the documented compatible topology, optimizer, and checkpoint contract, with storage accessible to the ranks that need it. Do not infer elastic rank-count changes or automatic recovery from the presence of a resume option.

FlagQuantum also exposes separately scoped train_distributed_mps and train_distributed_tensor_network workflows. Their state layouts, communication, parameter ownership, checkpoint behavior, and evidence boundaries differ. Changing mode= in a local example is not sufficient to reproduce those distributed workflows. Use their release examples and runbook sections.

The release capability catalog lists production-supported sharded statevector training and separately qualified two-node workflows. The stated scope applies to each recorded hardware/software/workload configuration, not arbitrary node counts, every gate, or all accelerators. The local CPU test on this page does not extend those claims.

Limitation or failure	What to check
More ranks do not improve runtime	Account for communication, synchronization, workload size, topology, and process overhead.
Job hangs during startup	Check that every worker launched with matching job settings, address, port, and reachable interfaces.
Backend mismatch	Use Gloo for this CPU example and a compatible NCCL-enabled Linux/CUDA environment for the GPU commands.
Several ranks use the same GPU	Set the local device from LOCAL_RANK before NCCL initialization.
Different circuits or collective order across ranks	Keep circuit structure, initialization, training settings, and communication sequence consistent.
One rank exits while others wait	Inspect all worker errors and launcher logs; timeouts and cleanup do not eliminate the underlying mismatch.
Out of memory despite sharding	Include backward state, communication buffers, and optimizer/parameter ownership in the budget.
Treating process memory as shard memory	Separate process-wide memory metrics, local-state estimates, and accelerator allocations.
A domestic card is visible	Consult its separately validated runtime and distributed transport; the --device cuda script is not generic FlagOS certification.

Continue Learning and References

• Accelerator Setup: runtime detection and single-device checks.
• Choosing a Simulation Method: representation tradeoffs.
• Training Parameterized Circuits: the local training lifecycle.
• Wiki Directory: published guides.
• Release distributed namespace, statevector training implementation, and multi-node runbook.
• Release MPS examples and capability catalog.
• PyTorch torchrun documentation and distributed communication documentation.