Version 4.3.0 (#103)

Author: sansiro77

README.md

@@ -177,6 +177,63 @@ print(pattern().full_state)
 print(cir() / pattern().full_state)
 ```
 
+- Distributed simulation of quantum circuit
+
+```python
+import torch
+# OMP_NUM_THREADS=2 torchrun --nproc_per_node=4 main.py
+backend = 'gloo' # for CPU
+# torchrun --nproc_per_node=4 main.py
+backend = 'nccl' # for GPU
+rank, world_size, local_rank = dq.setup_distributed(backend)
+data = torch.arange(4, dtype=torch.float, requires_grad=True)
+cir = dq.DistributedQubitCircuit(4)
+cir.rylayer(encode=True)
+cir.cnot_ring()
+cir.observable(0)
+cir.observable(1, 'x')
+if backend == 'nccl':
+    data = data.to(f'cuda:{local_rank}')
+    cir.to(f'cuda:{local_rank}')
+state = cir(data).amps
+result = cir.measure(with_prob=True)
+exp = cir.expectation().sum()
+exp.backward()
+if rank == 0:
+    print(state)
+    print(result)
+    print(exp)
+    print(data.grad)
+dq.cleanup_distributed()
+```
+
+- Distributed simulation of photonic quantum circuit
+
+```python
+# OMP_NUM_THREADS=2 torchrun --nproc_per_node=4 main.py
+backend = 'gloo' # for CPU
+# torchrun --nproc_per_node=4 main.py
+backend = 'nccl' # for GPU
+rank, world_size, local_rank = dq.setup_distributed(backend)
+nmode = 4
+cutoff = 4
+data = torch.arange(14, dtype=torch.float) / 10
+cir = dq.DistributedQumodeCircuit(nmode, [0] * nmode, cutoff)
+for i in range(nmode):
+    cir.s(i, encode=True)
+for i in range(nmode - 1):
+    cir.bs([i, i + 1], encode=True)
+if backend == 'nccl':
+    data = data.to(f'cuda:{local_rank}')
+    cir.to(f'cuda:{local_rank}')
+state = cir(data).amps
+result = cir.measure(with_prob=True)
+if rank == 0:
+    print(state)
+    print(result)
+dq.cleanup_distributed()
+```
+
 # License
 
 DeepQuantum is open source, released under the Apache License, Version 2.0.

src/deepquantum/__init__.py

@@ -3,7 +3,7 @@
 DeepQuantum can be directly imported.
 """
 
-__version__ = '4.2.0'
+__version__ = '4.3.0'
 
 
 from . import adjoint
@@ -42,6 +42,7 @@
 from .channel import BitFlip, PhaseFlip, Depolarizing, Pauli, AmplitudeDamping, PhaseDamping
 from .channel import GeneralizedAmplitudeDamping
 from .circuit import QubitCircuit, DistributedQubitCircuit
+from .communication import setup_distributed, cleanup_distributed
 from .gate import U3Gate, PhaseShift, Identity, PauliX, PauliY, PauliZ, Hadamard
 from .gate import SGate, SDaggerGate, TGate, TDaggerGate
 from .gate import Rx, Ry, Rz, ProjectionJ, CombinedSingleGate

src/deepquantum/communication.py

@@ -9,7 +9,7 @@
 import torch.distributed as dist
 
 
-def setup_distributed(port = '29500', backend = 'nccl') -> Tuple[int, int, int]:
+def setup_distributed(backend: str = 'nccl', port: str = '29500') -> Tuple[int, int, int]:
     """Initialize torch.distributed."""
     try:
         # These should be set by the launch script (e.g., torchrun)
@@ -24,16 +24,18 @@ def setup_distributed(port = '29500', backend = 'nccl') -> Tuple[int, int, int]:
         local_rank = 0
         os.environ['MASTER_ADDR'] = 'localhost'
         os.environ['MASTER_PORT'] = port
-
-    print(f'Initializing distributed setup: Rank {rank}/{world_size}, Local Rank (GPU): {local_rank}')
-
+    if backend == 'nccl':
+        print(f'Initializing distributed setup: Rank {rank}/{world_size}, Local Rank (GPU): {local_rank}')
+    elif backend == 'gloo':
+        print(f'Initializing distributed setup: Rank {rank}/{world_size}, Local Rank (CPU): {local_rank}')
     # Initialize the process group
     dist.init_process_group(backend, world_size=world_size, rank=rank)
-
-    # Pin the current process to a specific GPU
-    torch.cuda.set_device(local_rank)
-
-    print(f'Rank {rank} initialized, using GPU {local_rank}.')
+    if backend == 'nccl':
+        # Pin the current process to a specific GPU
+        torch.cuda.set_device(local_rank)
+        print(f'Rank {rank} initialized, using GPU {local_rank}.')
+    elif backend == 'gloo':
+        print(f'Rank {rank} initialized.')
     return rank, world_size, local_rank