Adding DDP to the repo for multi-GPU training

README.md

@@ -12,6 +12,7 @@ WaterFlow/
 │   ├── gvp.py              # Geometric Vector Perceptron layers
 │   ├── gvp_encoder.py      # GVP-based protein encoder
 │   ├── encoder_base.py     # Encoder registry and factory (includes ESM/SLAE)
+│   ├── distributed.py      # DDP helpers (rank discovery, collectives, barriers)
 │   ├── constants.py        # Shared constants (RBF bins, etc.)
 │   └── utils.py            # Metrics, plotting, logging utilities
 ├── scripts/                # Executable scripts
@@ -21,6 +22,7 @@ WaterFlow/
 │   └── generate_slae_embeddings.py  # Precompute SLAE embeddings
 ├── tests/                  # Test suite
 │   ├── test_dataset.py     # Dataset and preprocessing tests
+│   ├── test_distributed.py # DDP helper and cache prebuild tests
 │   ├── test_flow.py        # Flow matching tests
 │   ├── test_encoder.py     # Encoder tests
 │   ├── test_forward.py     # End-to-end forward pass tests
@@ -276,6 +278,27 @@ uv run python -m scripts.train \
     --processed_dir ~/flow_cache/
 ```
 
+### Multi-GPU Training (DDP)
+
+No DDP flag — `torchrun`'s env vars are the only switch; a plain
+`python -m scripts.train` runs single-GPU as before.
+
+```bash
+uv run torchrun --nproc_per_node=4 -m scripts.train \
+    --train_list splits/train_list_0.95.txt \
+    --val_list splits/valid_list_0.05.txt \
+    --encoder_type gvp \
+    --batch_size 4  # per rank -> effective 16
+```
+
+- Each rank trains on a disjoint `DistributedSampler` shard, reshuffled per epoch.
+- Gradients all-reduce once per optimizer step; train/val/eval metrics are
+  all-reduced, so every rank agrees on the best epoch.
+- Rank 0 owns disk and W&B (config, checkpoints, logs). Checkpoints hold the
+  unwrapped `state_dict`, so `inference.py` loads them unchanged.
+- The geometry cache is built by rank 0 before the NCCL group exists,
+  coordinated on a CPU-side store — a cold build can't trip a collective timeout.
+
 ### Resuming from Checkpoints
 
 To resume training from a checkpoint, you can load the model weights and optimizer state:

pyproject.toml

@@ -65,6 +65,15 @@ missing-argument = "ignore"
 # Dict.get() returns union types that don't narrow well
 invalid-argument-type = "ignore"
 
+# torch.distributed declares its API under `if is_available():`, so every member
+# reads as conditionally defined. Scoped to the one module that uses it -- the
+# training scripts reach DDP through src/distributed.py, not torch.distributed.
+[[tool.ty.overrides]]
+include = ["src/distributed.py"]
+
+[tool.ty.overrides.rules]
+possibly-missing-attribute = "ignore"
+
 [tool.ruff.lint]
 fixable = ["I001", "F401", "UP"]
 ignore = ["E402", "E501", "E721", "E731", "E741", "F722", "F821", "UP015", "UP037"]