diff --git a/README.md b/README.md index cc4b581..6801ba8 100644 --- a/README.md +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index 7bbdfdd..282182b 100644 --- a/pyproject.toml +++ b/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"] diff --git a/scripts/train.py b/scripts/train.py index ab69586..3fbeafa 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -19,7 +19,12 @@ """ import argparse +import contextlib import json +import multiprocessing as mp +import os +import shutil +import tempfile from datetime import datetime from pathlib import Path @@ -28,13 +33,24 @@ import torch import wandb from loguru import logger +from torch.nn.parallel import DistributedDataParallel as DDP from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, StepLR from torch.utils.data import DataLoader from torch_geometric.data import HeteroData from tqdm import tqdm -from src.dataset import get_dataloader +from src.dataset import get_dataloader, ProteinWaterDataset +from src.distributed import ( + all_reduce_means, + ddp_barrier, + ddp_is_active, + ddp_rank_and_world, + is_main_process, + run_once_on_main, + setup_distributed, + teardown_distributed, +) from src.encoder_base import build_encoder from src.flow import DYNAMIC_EDGE_POLICIES, FlowMatcher, FlowWaterGVP from src.utils import ( @@ -663,9 +679,16 @@ def run_eval_sampling( run_dir: Path to run directory for saving outputs """ flow_matcher.model.eval() + + # Each rank integrates a disjoint stride of eval_indices; the metric sums are + # all-reduced below so every rank ends with identical averages. + rank, world_size = ddp_rank_and_world() results = [] for i, idx in enumerate(eval_indices): + # Shard by global position i, so plot/GIF filenames never collide. + if i % world_size != rank: + continue graph = val_loader.dataset[idx] if graph["water"].num_nodes == 0: continue @@ -726,17 +749,42 @@ def run_eval_sampling( pdb_id=graph.pdb_id, ) - if results: - avg_metrics = { - "eval/avg_rmsd": np.mean([r["rmsd"] for r in results]), - "eval/avg_precision": np.mean([r["precision"] for r in results]), - "eval/avg_recall": np.mean([r["recall"] for r in results]), - "eval/avg_f1": np.mean([r["f1"] for r in results]), - "eval/avg_auc_pr": np.mean([r["auc_pr"] for r in results]), - } - wandb.log(avg_metrics, step=global_step) - return avg_metrics - return {} + # Every rank must reach this, even one whose stride was all zero-water graphs. + avg_metrics, _ = all_reduce_means( + { + f"eval/avg_{key}": sum(r[key] for r in results) + for key in ("rmsd", "precision", "recall", "f1", "auc_pr") + }, + len(results), + device, + ) + if not avg_metrics: + return {} + wandb.log(avg_metrics, step=global_step) + return avg_metrics + + +def _collective_device(args: argparse.Namespace) -> torch.device: + """ + Device for DDP collective buffers. + + NCCL requires this rank's own CUDA device; the CPU fallback only ever runs + single-process, where the collective is a no-op. + """ + return torch.device(args.device if torch.cuda.is_available() else "cpu") + + +def _needs_grad_sync(step: int, n_batches: int, accum_steps: int) -> bool: + """ + Whether this micro-step's backward must all-reduce gradients under DDP. + + True on every accumulation boundary, and throughout the epoch's trailing + partial window: that window ends in an optimizer.step() too, and stepping on + gradients that were never all-reduced leaves the ranks permanently diverged. + """ + if (step + 1) % accum_steps == 0: + return True + return step >= n_batches - (n_batches % accum_steps) def train_epoch( @@ -763,11 +811,17 @@ def train_epoch( skipped_batches += 1 continue - metrics = flow_matcher.training_step( - batch, - use_self_conditioning=args.use_self_cond, - accumulation_steps=args.grad_accum_steps, + # Suppress the gradient all-reduce on micro-steps that are not followed by + # an optimizer.step(), keeping comms at one all-reduce per optimizer step. + no_sync = ddp_is_active() and not _needs_grad_sync( + step, len(train_loader), args.grad_accum_steps ) + with flow_matcher.model.no_sync() if no_sync else contextlib.nullcontext(): + metrics = flow_matcher.training_step( + batch, + use_self_conditioning=args.use_self_cond, + accumulation_steps=args.grad_accum_steps, + ) if metrics["per_sample_info"] is not None: per_sample_losses = metrics["per_sample_info"]["losses"].cpu() @@ -840,6 +894,14 @@ def train_epoch( final_global_step = (epoch - 1) * len(train_loader) + len(train_loader) - 1 + # One collective per epoch so metrics cover every rank's shard. Must run before + # the zero-batch check, or a rank that skipped everything would never enter it. + train_metrics, processed_batches = all_reduce_means( + {"train/epoch_loss": total_loss, "train/epoch_rmsd": total_rmsd}, + processed_batches, + _collective_device(args), + ) + if processed_batches == 0: logger.warning( f"Epoch {epoch}: skipped all {skipped_batches} train batches (no waters)." @@ -853,14 +915,7 @@ def train_epoch( logger.info( f"Epoch {epoch} [Train] processed_batches={processed_batches}, skipped_batches={skipped_batches}" ) - return ( - { - "train/epoch_loss": total_loss / processed_batches, - "train/epoch_rmsd": total_rmsd / processed_batches, - }, - final_global_step, - optimizer_step_count, - ) + return train_metrics, final_global_step, optimizer_step_count @torch.no_grad() @@ -886,6 +941,13 @@ def val_epoch( total_loss += metrics["loss"] total_rmsd += metrics["rmsd"] + # Best-checkpoint selection keys off val/loss, so ranks must agree on it. + val_metrics, processed_batches = all_reduce_means( + {"val/loss": total_loss, "val/rmsd": total_rmsd}, + processed_batches, + _collective_device(args), + ) + if processed_batches == 0: logger.warning( f"Epoch {epoch}: skipped all {skipped_batches} val batches (no waters)." @@ -895,10 +957,7 @@ def val_epoch( logger.info( f"Epoch {epoch} [Val] processed_batches={processed_batches}, skipped_batches={skipped_batches}" ) - return { - "val/loss": total_loss / processed_batches, - "val/rmsd": total_rmsd / processed_batches, - } + return val_metrics def count_parameters(model): @@ -985,27 +1044,121 @@ def build_scheduler(optimizer, args): return warmup_scheduler, main_scheduler +def _build_cache_shard( + list_file: str, processed_dir: str, dataset_kwargs: dict +) -> None: + """ + Pool worker: build the geometry cache for one shard's list. + + Already-cached entries are skipped, and shards hold disjoint keys, so workers + never write the same file. + """ + ProteinWaterDataset( + pdb_list_file=list_file, + processed_dir=processed_dir, + preprocess=True, + **dataset_kwargs, + ) + + +def build_cache(args: argparse.Namespace) -> None: + """ + Build the geometry cache for the train+val lists as the sole writer. + + Preprocessing is CPU/PyMOL only, so this runs before the DDP group exists -- + hence race-free. A warm cache is a fast no-op; a cold build is parallelized + across CPU cores over disjoint key shards. + """ + dataset_kwargs, _, _ = _build_dataset_config(args) + ids = set() + for lst in (args.train_list, args.val_list): + with open(lst) as f: + ids.update(line.strip() for line in f if line.strip()) + sorted_ids = sorted(ids) + if not sorted_ids: + return + + tmpdir = Path(tempfile.mkdtemp(prefix="wf_build_")) + try: + union = tmpdir / "union.txt" + union.write_text("\n".join(sorted_ids) + "\n") + # Parse-only probe to find which entries still need building. + probe = ProteinWaterDataset( + pdb_list_file=str(union), + processed_dir=args.processed_dir, + preprocess=False, + **dataset_kwargs, + ) + missing = [ + entry["cache_key"] + for entry in probe.entries + if not (probe.geometry_dir / f"{entry['cache_key']}.pt").is_file() + ] + if not missing: + return # warm cache: nothing to build + + logger.info(f"build_cache: preprocessing {len(missing)} missing entries") + n_shards = max(1, min(len(missing), os.cpu_count() or 1)) + if n_shards == 1: + _build_cache_shard(str(union), args.processed_dir, dataset_kwargs) + return + shard_files = [] + for i in range(n_shards): + shard = tmpdir / f"shard_{i}.txt" + shard.write_text("\n".join(missing[i::n_shards]) + "\n") + shard_files.append(str(shard)) + # spawn (not fork): safe alongside PyMOL's C extension and any threads. + ctx = mp.get_context("spawn") + with ctx.Pool(n_shards) as pool: + pool.starmap( + _build_cache_shard, + [(shard, args.processed_dir, dataset_kwargs) for shard in shard_files], + ) + finally: + shutil.rmtree(tmpdir, ignore_errors=True) + + def main(): """Run the full training pipeline.""" args = parse_args() + + # Build the cache before the NCCL group exists: a cold build can take hours and + # would trip a GPU-collective timeout. The store is reused as NCCL's rendezvous. + store = run_once_on_main(lambda: build_cache(args), key="wf_cache_ready") + + # Under torchrun each rank binds its own GPU; a plain launch yields (0, 0, 1). + rank, local_rank, world_size = setup_distributed(store=store) + main_proc = is_main_process(rank) + if ddp_is_active(): + args.device = f"cuda:{local_rank}" device = torch.device(args.device if torch.cuda.is_available() else "cpu") if args.run_name is None: args.run_name = generate_run_name(args) run_dir = Path(args.save_dir) / args.run_name - run_dir.mkdir(parents=True, exist_ok=True) - (run_dir / "checkpoints").mkdir(exist_ok=True) - (run_dir / "plots").mkdir(exist_ok=True) - (run_dir / "gifs").mkdir(exist_ok=True) - + if main_proc: + run_dir.mkdir(parents=True, exist_ok=True) + (run_dir / "checkpoints").mkdir(exist_ok=True) + (run_dir / "plots").mkdir(exist_ok=True) + (run_dir / "gifs").mkdir(exist_ok=True) + ddp_barrier() # other ranks wait until run_dir exists + + # Rank 0 owns the log file; other ranks log to console only, so N processes + # never interleave writes into it. log_file = Path(args.log_file) if args.log_file else run_dir / "train.log" - setup_logging_for_tqdm(level=args.log_level, log_file=str(log_file)) + setup_logging_for_tqdm( + level=args.log_level, log_file=str(log_file) if main_proc else None + ) logger.info("=" * 60) logger.info(f"Run name: {args.run_name}") logger.info(f"Run directory: {run_dir}") logger.info(f"Log file: {log_file}") + if ddp_is_active(): + logger.info( + f"DDP active: rank={rank} local_rank={local_rank} world_size={world_size}" + ) logger.info("=" * 60) # data loaders @@ -1022,6 +1175,7 @@ def main(): prefetch_factor=args.prefetch_factor, persistent_workers=args.persistent_workers, duplicate_single_sample=args.duplicate_single_sample, + distributed=ddp_is_active(), **dataset_kwargs, ) @@ -1035,6 +1189,7 @@ def main(): prefetch_factor=args.prefetch_factor, persistent_workers=args.persistent_workers, duplicate_single_sample=args.duplicate_single_sample, + distributed=ddp_is_active(), **dataset_kwargs, ) @@ -1047,13 +1202,14 @@ def main(): ).tolist() eval_indices_file = run_dir / "eval_indices.txt" - with open(eval_indices_file, "w") as f: - f.write("# Fixed evaluation sample indices\n") - for idx in eval_indices: - graph = val_loader.dataset[idx] - pdb_id = getattr(graph, "pdb_id", "unknown") - f.write(f"{idx}\t{pdb_id}\n") - logger.info(f"Fixed eval indices saved to: {eval_indices_file}") + if main_proc: + with open(eval_indices_file, "w") as f: + f.write("# Fixed evaluation sample indices\n") + for idx in eval_indices: + graph = val_loader.dataset[idx] + pdb_id = getattr(graph, "pdb_id", "unknown") + f.write(f"{idx}\t{pdb_id}\n") + logger.info(f"Fixed eval indices saved to: {eval_indices_file}") logger.info(f"Evaluating on {len(eval_indices)} proteins at each eval epoch") # detect input dimension and resolve encoder configuration from sample data @@ -1078,19 +1234,38 @@ def main(): config_dict["node_scalar_in"] = node_scalar_in config_dict["resolved_encoder_config"] = encoder_config config_file = run_dir / "config.json" - with open(config_file, "w") as f: - json.dump(config_dict, f, indent=2) - logger.info(f"Configuration saved to: {config_file}") - + if main_proc: + with open(config_file, "w") as f: + json.dump(config_dict, f, indent=2) + logger.info(f"Configuration saved to: {config_file}") + + # Non-main ranks run a disabled client, so every wandb.log call site stays a + # no-op without a per-call guard. None on the main rank defers to WANDB_MODE + # (default: online); an explicit mode here would override that env var. wandb.init( project=args.wandb_project, dir=args.wandb_dir, name=args.run_name, config=config_dict, + mode=None if main_proc else "disabled", ) model = build_model(args, device, encoder_config=encoder_config) - trainable_params, total_params = count_parameters(model) + if ddp_is_active(): + # broadcast_buffers=False is safe (no BatchNorm; LayerNorm has no synced + # buffers). find_unused_parameters=True because ablated edge types can + # leave the used-parameter set varying across backwards. + model = DDP( + model, + device_ids=[local_rank], + broadcast_buffers=False, + find_unused_parameters=True, + ) + # Unwrapped module: parameter access, sanity forward, sampling, and + # state_dicts. Saving the wrapper would prefix every key with "module.". + raw_model = getattr(model, "module", model) + + trainable_params, total_params = count_parameters(raw_model) logger.info("Model statistics:") logger.info(f"Trainable parameters: {trainable_params:,}") logger.info(f"Total parameters: {total_params:,}") @@ -1098,17 +1273,17 @@ def main(): # quick forward pass sanity check for cached embedding encoders if _uses_cached_embeddings(args.encoder_type): logger.info(f"Testing forward pass with {args.encoder_type.upper()}...") - model.eval() + raw_model.eval() batch = next(iter(train_loader)).to(device) with torch.no_grad(): num_graphs = int(batch["protein"].batch.max().item()) + 1 t = torch.zeros(num_graphs, device=device) - v_out = model(batch, t) + v_out = raw_model(batch, t) logger.info(f"Forward pass successful! Output shape: {v_out.shape}") logger.info(f"Output stats: mean={v_out.mean():.4f}, std={v_out.std():.4f}") if v_out.std() < 1e-6: logger.warning("Model output is constant! This indicates a problem.") - model.train() + raw_model.train() flow_matcher = FlowMatcher( model=model, @@ -1121,7 +1296,7 @@ def main(): ) optimizer = AdamW( - [p for p in model.parameters() if p.requires_grad], + [p for p in raw_model.parameters() if p.requires_grad], lr=args.lr, weight_decay=args.weight_decay, ) @@ -1131,6 +1306,11 @@ def main(): optimizer_step_count = 0 for epoch in range(1, args.epochs + 1): + # Without this every epoch replays the same shard order on every rank. + if ddp_is_active(): + train_loader.sampler.set_epoch(epoch) + val_loader.sampler.set_epoch(epoch) + train_metrics, global_step, optimizer_step_count = train_epoch( flow_matcher, train_loader, @@ -1157,22 +1337,25 @@ def main(): f"val_loss={val_metrics['val/loss']:.4f}, val_rmsd={val_metrics['val/rmsd']:.2f}" ) + # val/loss is all-reduced, so every rank agrees on the best epoch and + # updates best_val_loss identically; only rank 0 writes it out. if val_metrics["val/loss"] < best_val_loss: best_val_loss = val_metrics["val/loss"] - save_checkpoint( - model, - optimizer, - warmup_scheduler, - main_scheduler, - epoch, - optimizer_step_count, - run_dir / "checkpoints" / "best.pt", - best=True, - ) + if main_proc: + save_checkpoint( + raw_model, + optimizer, + warmup_scheduler, + main_scheduler, + epoch, + optimizer_step_count, + run_dir / "checkpoints" / "best.pt", + best=True, + ) - if epoch % args.save_every == 0: + if epoch % args.save_every == 0 and main_proc: save_checkpoint( - model, + raw_model, optimizer, warmup_scheduler, main_scheduler, @@ -1181,17 +1364,24 @@ def main(): run_dir / "checkpoints" / f"epoch_{epoch}.pt", ) + # All ranks enter: run_eval_sampling ends in a collective. Swap in the + # unwrapped module so no DDP forward machinery fires during integration. if epoch % args.eval_every == 0: - eval_metrics = run_eval_sampling( - flow_matcher, - val_loader, - args, - epoch, - device, - global_step, - eval_indices, - run_dir, - ) + wrapped = flow_matcher.model + flow_matcher.model = raw_model + try: + eval_metrics = run_eval_sampling( + flow_matcher, + val_loader, + args, + epoch, + device, + global_step, + eval_indices, + run_dir, + ) + finally: + flow_matcher.model = wrapped if eval_metrics: logger.info( f"Eval: RMSD={eval_metrics['eval/avg_rmsd']:.2f}A, " @@ -1201,7 +1391,11 @@ def main(): f"AUC-PR={eval_metrics['eval/avg_auc_pr']:.3f}" ) + # Realign ranks: rank 0 may have spent extra time writing checkpoints. + ddp_barrier() + wandb.finish() + teardown_distributed() logger.info("Training complete.") diff --git a/src/dataset.py b/src/dataset.py index 4a6ba0b..960c402 100644 --- a/src/dataset.py +++ b/src/dataset.py @@ -26,7 +26,8 @@ from loguru import logger from scipy.spatial.distance import cdist from torch import Tensor -from torch.utils.data import DataLoader, Dataset +from torch.utils.data import DataLoader, Dataset, Sampler +from torch.utils.data.distributed import DistributedSampler from torch_cluster import radius_graph from torch_geometric.data import Batch, HeteroData from tqdm import tqdm @@ -1390,6 +1391,8 @@ def get_dataloader( pin_memory: bool = True, prefetch_factor: int = 4, persistent_workers: bool = True, + sampler: Sampler | None = None, + distributed: bool = False, **dataset_kwargs, ) -> DataLoader: """ @@ -1409,6 +1412,12 @@ def get_dataloader( pin_memory: Pin memory for faster CPU-GPU transfer (default True) prefetch_factor: Number of batches to prefetch per worker (default 4) persistent_workers: Keep workers alive between epochs (default True) + sampler: Optional sampler (e.g. DistributedSampler for DDP). When + provided, `shuffle` is ignored (the sampler owns ordering). + distributed: If True (and no explicit sampler is given), build a + DistributedSampler from the env-configured process group so each + DDP rank sees a disjoint shard. Call + `loader.sampler.set_epoch(epoch)` each epoch to reshuffle. **dataset_kwargs: Additional arguments passed to ProteinWaterDataset (e.g., cutoff, include_mates, duplicate_single_sample) @@ -1427,10 +1436,18 @@ def get_dataloader( **dataset_kwargs, ) + if sampler is None and distributed: + # drop_last=False keeps every sample; the sampler pads the last shard by + # repeating a few, which double-counts them in epoch metrics. Accepted: + # dropping instead gives ranks uneven shards and stalls the all-reduce. + sampler = DistributedSampler(dataset, shuffle=shuffle, drop_last=False) + loader = DataLoader( dataset, batch_size=batch_size, - shuffle=shuffle, + # A sampler is mutually exclusive with shuffle; let the sampler own ordering. + shuffle=shuffle if sampler is None else False, + sampler=sampler, num_workers=num_workers, pin_memory=pin_memory, prefetch_factor=prefetch_factor if num_workers > 0 else None, diff --git a/src/distributed.py b/src/distributed.py new file mode 100644 index 0000000..60b5376 --- /dev/null +++ b/src/distributed.py @@ -0,0 +1,165 @@ +# distributed.py +""" +Distributed (DDP) helpers shared by the training entry points. + +DDP is activated purely by the launcher: `torchrun --nproc_per_node=N` sets +WORLD_SIZE / RANK / LOCAL_RANK. Without those env vars every helper degrades to +single-GPU behavior, so scripts run unchanged under `python -m scripts.train`. +Nothing here reads CLI arguments, which keeps a recorded config.json identical +whether the run used one GPU or eight. +""" + +import os +from collections.abc import Callable +from datetime import timedelta + +import torch +import torch.distributed as dist + + +def _ddp_world_size() -> int: + """World size from the launcher env, or 1 when not launched by torchrun.""" + return int(os.environ.get("WORLD_SIZE", "1")) + + +def ddp_is_active() -> bool: + """True when running under a multi-process launcher.""" + return _ddp_world_size() > 1 + + +def ddp_rank_and_world() -> tuple[int, int]: + """ + This rank's index and the world size, from the live process group. + + Returns (0, 1) when not distributed, so callers can shard work by + `i % world_size == rank` without branching. + """ + if not ddp_is_active(): + return 0, 1 + return dist.get_rank(), dist.get_world_size() + + +def setup_distributed(store: dist.Store | None = None) -> tuple[int, int, int]: + """ + Initialize the NCCL process group if launched under torchrun. + + Args: + store: Optional rendezvous store. Passing the CPU-side TCPStore from + `run_once_on_main` builds the NCCL group on top of it instead of + re-rendezvousing via env://, so one store serves both phases. + + Returns: + (rank, local_rank, world_size); (0, 0, 1) when not distributed. + """ + if not ddp_is_active(): + return 0, 0, 1 + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + # device_id binds this rank to its GPU eagerly and verifies the rank->GPU + # mapping at init, instead of inferring it from rank order. + if store is not None: + dist.init_process_group( + backend="nccl", + store=store, + rank=int(os.environ["RANK"]), + world_size=int(os.environ["WORLD_SIZE"]), + device_id=torch.device(f"cuda:{local_rank}"), + ) + else: + dist.init_process_group( + backend="nccl", device_id=torch.device(f"cuda:{local_rank}") + ) + return dist.get_rank(), local_rank, dist.get_world_size() + + +def is_main_process(rank: int) -> bool: + """True on the single rank that owns disk IO and W&B logging.""" + return rank == 0 + + +def ddp_barrier() -> None: + """Block until every rank arrives. No-op when not distributed.""" + if ddp_is_active(): + dist.barrier() + + +def teardown_distributed() -> None: + """Destroy the process group. No-op when not distributed.""" + if ddp_is_active(): + dist.destroy_process_group() + + +def run_once_on_main(work: Callable[[], None], key: str) -> dist.Store | None: + """ + Run `work` on rank 0 only, holding every other rank until it finishes. + Primarily for dataset processing on 0 rank to prevent race conditions. + + Coordination runs on a CPU-side TCPStore rather than a GPU collective, so a + long single-writer job never counts against an NCCL timeout. Call this before + `setup_distributed` and hand the store back to it. + + Args: + work: Executed on rank 0 only. An exception propagates on rank 0 and + leaves the others blocked -- a hang, not a half-written result. + key: Store key signalling completion. Must be unique per work item. + + Returns: + The store, for `setup_distributed(store=...)`; None when not distributed + (in which case `work` simply runs inline). + """ + if not ddp_is_active(): + work() + return None + + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + host = os.environ["MASTER_ADDR"] + port = int(os.environ["MASTER_PORT"]) + # Client of torchrun's agent store (is_master=False) -- the agent already owns + # MASTER_PORT. The timeout covers a cold build, which can take hours. + store = dist.TCPStore(host, port, world, False, timeout=timedelta(hours=24)) + + if rank == 0: + work() + store.set(key, "1") + else: + store.wait([key]) + return store + + +def all_reduce_means( + sums: dict[str, float], count: int, device: torch.device +) -> tuple[dict[str, float], int]: + """ + Average per-item metric sums across ranks. + + Reduces the metric sums and the item count in one collective, then divides. + Summing before dividing is what keeps the mean correct when ranks processed + unequal numbers of items; float64 keeps it insensitive to the order NCCL + combines ranks in. + + Call this on every rank, including ones where `count` is 0 -- a rank that + skips the collective hangs the rest. That is why an empty result is reported + through the return value rather than an early exit. + + Args: + sums: Metric name -> sum of that metric over this rank's items. + count: Number of items this rank contributed to `sums`. + device: Buffer device. Must be this rank's CUDA device under NCCL. + + Returns: + (means, total_count). `means` is empty when `total_count` is 0. + """ + keys = list(sums) + totals = torch.tensor( + [float(sums[k]) for k in keys] + [float(count)], + dtype=torch.float64, + device=device, + ) + if ddp_is_active(): + dist.all_reduce(totals, op=dist.ReduceOp.SUM) + + total_count = int(totals[-1].item()) + if total_count == 0: + return {}, 0 + return {k: (totals[i] / totals[-1]).item() for i, k in enumerate(keys)}, total_count diff --git a/src/flow.py b/src/flow.py index 0728dbc..4ccc1b6 100644 --- a/src/flow.py +++ b/src/flow.py @@ -907,7 +907,10 @@ def __init__( self.t_distort = t_distort self.sigma_distort = sigma_distort self.loss_eps = loss_eps - self.graph_cutoff = getattr(model, "cutoff", 8.0) + # DDP does not forward attribute lookups to the module it wraps, so + # reading `cutoff` off the wrapper would silently fall back to the + # default and change the water prior's sampling radius under DDP only. + self.graph_cutoff = getattr(getattr(model, "module", model), "cutoff", 8.0) self.sampling_strategy = sampling_strategy @staticmethod diff --git a/tests/test_distributed.py b/tests/test_distributed.py new file mode 100644 index 0000000..75456a7 --- /dev/null +++ b/tests/test_distributed.py @@ -0,0 +1,455 @@ +""" +Tests for the DDP helpers and the single-writer cache prebuild. + +These run single-process, so they cover the logic that decides *what* each rank +does -- sharding, the sum-then-divide reduction, sampler wiring, cache shard +disjointness -- not NCCL itself. The collective paths degrade to no-ops without a +launcher, which is exactly the property most of these assert. +""" + +import sys +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch +from torch.utils.data import SequentialSampler + +import scripts.train as train +from src.dataset import get_dataloader, ProteinWaterDataset +from src.distributed import ( + _ddp_world_size, + all_reduce_means, + ddp_barrier, + ddp_is_active, + ddp_rank_and_world, + is_main_process, + run_once_on_main, + setup_distributed, + teardown_distributed, +) +from src.flow import FlowMatcher + + +CPU = torch.device("cpu") + + +@pytest.fixture +def launcher_env(monkeypatch): + """Factory that fakes a torchrun launch by setting WORLD_SIZE.""" + + def _set(world_size): + monkeypatch.setenv("WORLD_SIZE", str(world_size)) + + return _set + + +# ============== Launcher detection ============== + + +def test_world_size_defaults_to_one_without_launcher(monkeypatch): + monkeypatch.delenv("WORLD_SIZE", raising=False) + assert _ddp_world_size() == 1 + assert ddp_is_active() is False + + +def test_single_process_launch_is_not_ddp(launcher_env): + """torchrun --nproc_per_node=1 must take the plain single-GPU path.""" + launcher_env(1) + assert ddp_is_active() is False + + +def test_world_size_above_one_activates_ddp(launcher_env): + launcher_env(4) + assert ddp_is_active() is True + + +def test_setup_distributed_is_inert_without_launcher(monkeypatch): + """No process group is created, so a plain `python -m scripts.train` works.""" + monkeypatch.delenv("WORLD_SIZE", raising=False) + assert setup_distributed() == (0, 0, 1) + + +def test_rank_and_world_default_to_solo(monkeypatch): + monkeypatch.delenv("WORLD_SIZE", raising=False) + assert ddp_rank_and_world() == (0, 1) + + +def test_barrier_and_teardown_are_noops_without_launcher(monkeypatch): + """Both would raise on an uninitialized process group if they ran for real.""" + monkeypatch.delenv("WORLD_SIZE", raising=False) + ddp_barrier() + teardown_distributed() + + +def test_is_main_process_only_rank_zero(): + assert is_main_process(0) is True + assert is_main_process(1) is False + + +# ============== Stride sharding ============== + + +@pytest.mark.parametrize("world_size", [1, 2, 3, 8]) +def test_stride_shard_covers_every_index_exactly_once(world_size): + """ + The `i % world_size == rank` rule run_eval_sampling shards on. + + A structure evaluated twice would be double-counted in the all-reduced means; + one evaluated zero times would be silently dropped. + """ + n_structures = 17 + owners = [ + [i for i in range(n_structures) if i % world_size == rank] + for rank in range(world_size) + ] + assert sorted(i for shard in owners for i in shard) == list(range(n_structures)) + + +# ============== Gradient sync scheduling ============== + + +def _sync_steps(n_batches, accum_steps): + return [ + step + for step in range(n_batches) + if train._needs_grad_sync(step, n_batches, accum_steps) + ] + + +def test_every_step_syncs_without_accumulation(): + assert _sync_steps(5, 1) == [0, 1, 2, 3, 4] + + +def test_only_boundaries_sync_when_batches_divide_evenly(): + """No trailing window, so the intermediate micro-steps can skip the collective.""" + assert _sync_steps(8, 4) == [3, 7] + + +def test_trailing_partial_window_syncs_throughout(): + """ + 10 batches at accum=4 leaves steps 8-9 in a window that still ends in a step(). + + Letting those accumulate under no_sync would apply un-reduced gradients and + drift the ranks apart for the rest of training. + """ + assert _sync_steps(10, 4) == [3, 7, 8, 9] + + +@pytest.mark.parametrize("n_batches", [1, 2, 5, 7, 10, 33]) +@pytest.mark.parametrize("accum_steps", [1, 2, 3, 4, 8]) +def test_last_batch_of_an_epoch_always_syncs(n_batches, accum_steps): + """The epoch's final optimizer.step() must never run on unsynced gradients.""" + assert train._needs_grad_sync(n_batches - 1, n_batches, accum_steps) + + +# ============== all_reduce_means ============== + + +def test_all_reduce_means_divides_sums_by_count(): + means, count = all_reduce_means({"a": 10.0, "b": 5.0}, 4, CPU) + assert count == 4 + assert means == {"a": 2.5, "b": 1.25} + + +def test_all_reduce_means_reports_empty_instead_of_dividing_by_zero(): + """A rank with no items must still return, not raise -- callers branch on this.""" + assert all_reduce_means({"a": 0.0}, 0, CPU) == ({}, 0) + + +def test_all_reduce_means_matches_a_plain_mean(): + """Sum-then-divide reproduces the single-process average it replaced.""" + values = [0.31, 2.75, 1.5, 9.125, 0.0625] + means, _ = all_reduce_means({"m": sum(values)}, len(values), CPU) + assert means["m"] == pytest.approx(sum(values) / len(values), rel=1e-12) + + +def test_all_reduce_means_weights_by_item_count(): + """ + Reducing sums (not per-rank means) is what makes unequal shards correct. + + Two ranks holding 1 item at 10.0 and 3 items at 2.0 average to 4.0, not the + 6.0 that averaging their two means would give. + """ + means, count = all_reduce_means({"m": 10.0 + 6.0}, 1 + 3, CPU) + assert count == 4 + assert means["m"] == pytest.approx(4.0) + + +def test_all_reduce_means_preserves_key_order(): + """Keys and the reduced buffer are zipped positionally.""" + sums = {"z": 3.0, "a": 6.0, "m": 9.0} + means, _ = all_reduce_means(sums, 3, CPU) + assert list(means) == ["z", "a", "m"] + assert means == {"z": 1.0, "a": 2.0, "m": 3.0} + + +# ============== Reading model config through the DDP wrapper ============== + + +def test_flow_matcher_reads_cutoff_through_a_wrapper(): + """ + DDP does not forward attribute lookups to the module it wraps. + + Reading `cutoff` off the wrapper would fall back to the 8.0 default and + silently change the water prior's sampling radius under DDP only. + """ + wrapped = SimpleNamespace(module=SimpleNamespace(cutoff=12.0)) + assert FlowMatcher(model=wrapped).graph_cutoff == 12.0 + + +def test_flow_matcher_reads_cutoff_off_a_bare_model(): + assert FlowMatcher(model=SimpleNamespace(cutoff=12.0)).graph_cutoff == 12.0 + + +# ============== run_once_on_main ============== + + +def test_run_once_on_main_runs_inline_without_launcher(monkeypatch): + monkeypatch.delenv("WORLD_SIZE", raising=False) + calls = [] + store = run_once_on_main(lambda: calls.append("built"), key="k") + assert calls == ["built"] + assert store is None + + +# ============== get_dataloader sampler wiring ============== + + +@pytest.fixture +def pdb_list_file(tmp_path): + path = tmp_path / "list.txt" + path.write_text("6eey_final\n") + return path + + +def _loader(pdb_list_file, tmp_path, pdb_base_dir, **kwargs): + return get_dataloader( + pdb_list_file=str(pdb_list_file), + processed_dir=str(tmp_path / "processed"), + base_pdb_dir=str(pdb_base_dir), + num_workers=0, + preprocess=False, + **kwargs, + ) + + +def test_dataloader_has_no_sampler_by_default(pdb_list_file, tmp_path, pdb_base_dir): + loader = _loader(pdb_list_file, tmp_path, pdb_base_dir, shuffle=True) + assert not isinstance(loader.sampler, torch.utils.data.DistributedSampler) + + +def test_explicit_sampler_overrides_shuffle(pdb_list_file, tmp_path, pdb_base_dir): + """ + DataLoader raises if both are set, so `shuffle` must yield to the sampler. + + Passing shuffle=True alongside a sampler is what a caller does by accident; + it has to be tolerated rather than crash mid-run. + """ + dataset = ProteinWaterDataset( + pdb_list_file=str(pdb_list_file), + processed_dir=str(tmp_path / "probe"), + base_pdb_dir=str(pdb_base_dir), + preprocess=False, + ) + sampler = SequentialSampler(dataset) + loader = _loader( + pdb_list_file, tmp_path, pdb_base_dir, shuffle=True, sampler=sampler + ) + assert loader.sampler is sampler + + +def test_distributed_flag_is_ignored_when_a_sampler_is_given( + pdb_list_file, tmp_path, pdb_base_dir +): + dataset = ProteinWaterDataset( + pdb_list_file=str(pdb_list_file), + processed_dir=str(tmp_path / "probe"), + base_pdb_dir=str(pdb_base_dir), + preprocess=False, + ) + sampler = SequentialSampler(dataset) + loader = _loader( + pdb_list_file, tmp_path, pdb_base_dir, distributed=True, sampler=sampler + ) + assert loader.sampler is sampler + + +def test_distributed_builds_a_distributed_sampler( + pdb_list_file, tmp_path, pdb_base_dir, monkeypatch +): + """ + distributed=True must produce a DistributedSampler with shuffle honored. + + DistributedSampler reads the process group, so stand in a world of one rather + than initializing NCCL. + """ + import torch.utils.data.distributed as dist_data + + monkeypatch.setattr(dist_data.dist, "is_available", lambda: True) + monkeypatch.setattr(dist_data.dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist_data.dist, "get_world_size", lambda: 1) + monkeypatch.setattr(dist_data.dist, "get_rank", lambda: 0) + + loader = _loader( + pdb_list_file, tmp_path, pdb_base_dir, shuffle=True, distributed=True + ) + assert isinstance(loader.sampler, torch.utils.data.DistributedSampler) + assert loader.sampler.shuffle is True + assert loader.sampler.drop_last is False + + +# ============== build_cache ============== + + +class _InlinePool: + """Stand-in for a spawn Pool that runs starmap in-process.""" + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def starmap(self, fn, argsets): + return [fn(*args) for args in argsets] + + +@pytest.fixture +def inline_pool(monkeypatch): + monkeypatch.setattr( + train.mp, + "get_context", + lambda method: SimpleNamespace(Pool=lambda n: _InlinePool()), + ) + + +@pytest.fixture +def cache_args(tmp_path, pdb_base_dir): + """Real parsed args for a train+val pair with overlapping ids.""" + train_list = tmp_path / "train.txt" + val_list = tmp_path / "val.txt" + train_list.write_text("\n".join(f"e{i}_final" for i in range(6)) + "\n") + val_list.write_text("e5_final\ne6_final\n") # e5 overlaps the train list + + argv = [ + "train.py", + "--train_list", + str(train_list), + "--val_list", + str(val_list), + "--processed_dir", + str(tmp_path / "cache"), + "--base_pdb_dir", + str(pdb_base_dir), + ] + with mock.patch.object(sys, "argv", argv): + return train.parse_args() + + +def _geometry_dir(args): + """Where build_cache's probe will look for cached geometry.""" + dataset_kwargs, _, _ = train._build_dataset_config(args) + probe = ProteinWaterDataset( + pdb_list_file=args.val_list, + processed_dir=args.processed_dir, + preprocess=False, + **dataset_kwargs, + ) + return probe.geometry_dir + + +@pytest.fixture +def recorded_shards(monkeypatch): + """ + Capture the cache keys each shard was handed, one list per worker call. + + Read eagerly: build_cache deletes its tmpdir before returning, so the shard + files are gone by the time the test body runs. + """ + recorded = [] + + def _record(list_file, processed_dir, dataset_kwargs): + recorded.append([line for line in Path(list_file).read_text().split() if line]) + + monkeypatch.setattr(train, "_build_cache_shard", _record) + return recorded + + +def test_build_cache_shards_are_disjoint_and_complete( + cache_args, recorded_shards, inline_pool, monkeypatch +): + """Every missing key is built exactly once; workers never race on a file.""" + monkeypatch.setattr(train.os, "cpu_count", lambda: 3) + train.build_cache(cache_args) + + shards = recorded_shards + assert len(shards) == 3 + built = [key for shard in shards for key in shard] + assert sorted(built) == [f"e{i}_final" for i in range(7)] + assert len(built) == len(set(built)) + + +def test_build_cache_deduplicates_ids_across_train_and_val( + cache_args, recorded_shards, inline_pool, monkeypatch +): + """e5 appears in both lists but must only be built once.""" + monkeypatch.setattr(train.os, "cpu_count", lambda: 3) + train.build_cache(cache_args) + + built = [key for shard in recorded_shards for key in shard] + assert built.count("e5_final") == 1 + + +def test_build_cache_skips_already_cached_entries( + cache_args, recorded_shards, inline_pool, monkeypatch +): + monkeypatch.setattr(train.os, "cpu_count", lambda: 2) + geometry_dir = _geometry_dir(cache_args) + geometry_dir.mkdir(parents=True, exist_ok=True) + for i in range(5): + (geometry_dir / f"e{i}_final.pt").touch() + + train.build_cache(cache_args) + + built = [key for shard in recorded_shards for key in shard] + assert sorted(built) == ["e5_final", "e6_final"] + + +def test_build_cache_is_a_noop_on_a_warm_cache( + cache_args, recorded_shards, inline_pool +): + """A warm cache must not spawn a pool at all -- this runs every job's startup.""" + geometry_dir = _geometry_dir(cache_args) + geometry_dir.mkdir(parents=True, exist_ok=True) + for i in range(7): + (geometry_dir / f"e{i}_final.pt").touch() + + train.build_cache(cache_args) + assert recorded_shards == [] + + +def test_build_cache_runs_single_shard_without_a_pool( + cache_args, recorded_shards, monkeypatch +): + """ + With one usable core the pool is skipped entirely. + + No inline_pool fixture here: reaching mp.get_context would spawn real workers. + """ + monkeypatch.setattr(train.os, "cpu_count", lambda: 1) + train.build_cache(cache_args) + + assert len(recorded_shards) == 1 + assert sorted(recorded_shards[0]) == [f"e{i}_final" for i in range(7)] + + +def test_build_cache_handles_empty_lists(cache_args, recorded_shards, tmp_path): + empty = tmp_path / "empty.txt" + empty.write_text("") + cache_args.train_list = str(empty) + cache_args.val_list = str(empty) + + train.build_cache(cache_args) + assert recorded_shards == []