Optimize exposure autoencoder distributed training

README.md

@@ -73,3 +73,16 @@ compatible with the default gated `TimesNetExposureEncoder`, and the channel
 normalization statistics needed when the encoder is attached to DeepHealth.
 Multi-GPU pretraining follows the main trainer interface: add
 `--data_parallel --gpu_ids 0,1,2,3`.
+For efficient multi-GPU training, launch one process per GPU with DDP:
+
+```bash
+torchrun --standalone --nproc_per_node=4 train_exposure_autoencoder.py \
+  --exposure_cache_dir ukb_exposure_cache \
+  --batch_size 128
+```
+
+In DDP mode, `--batch_size` is the global batch size and must be divisible by
+the number of processes.
+Training-channel statistics are cached at
+`<exposure_cache_dir>/train_channel_stats.npz`; use
+`--recompute_channel_stats` only when a forced refresh is needed.

backbones.py

@@ -232,7 +232,8 @@ class TimesNetBlock(nn.Module):
         amplitude[0] = 0.0
         k = min(self.top_k, amplitude.numel() - 1)
         weights, indices = torch.topk(amplitude, k=k)
-        periods = [max(1, T // int(idx.item())) for idx in indices]
+        # One device synchronization per block instead of one per selected period.
+        periods = [max(1, T // int(idx)) for idx in indices.tolist()]
         return periods, weights.to(dtype=x.dtype, device=x.device)
 
     def _period_branch(self, x: torch.Tensor, period: int) -> torch.Tensor:

train_exposure_autoencoder.py

@@ -2,14 +2,19 @@
 from __future__ import annotations
 
 import argparse
+import hashlib
 import json
+import logging
 import math
+import os
 from pathlib import Path
 
 import numpy as np
 import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
 from torch.optim import AdamW
-from torch.utils.data import DataLoader, Dataset
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
 from tqdm import tqdm
 
 from backbones import TimesNetExposureAutoencoder
@@ -51,6 +56,19 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--exposure_cache_dir", required=True)
     parser.add_argument("--train_eid_file", default="ukb_train_eid.csv")
     parser.add_argument("--val_eid_file", default="ukb_val_eid.csv")
+    parser.add_argument(
+        "--channel_stats_file",
+        default=None,
+        help=(
+            "Cached channel statistics .npz file. Defaults to "
+            "<exposure_cache_dir>/train_channel_stats.npz."
+        ),
+    )
+    parser.add_argument(
+        "--recompute_channel_stats",
+        action="store_true",
+        help="Ignore a compatible statistics cache and recompute it.",
+    )
     parser.add_argument("--runs_root", default="runs")
     parser.add_argument("--seed", type=int, default=42)
     parser.add_argument("--n_embd", type=int, default=120)
@@ -82,9 +100,23 @@ def parse_args() -> argparse.Namespace:
         default=None,
         help="Comma-separated CUDA device ids for --data_parallel, e.g. 0,1,2,3.",
     )
+    parser.add_argument(
+        "--ddp_backend",
+        default=None,
+        choices=["nccl", "gloo"],
+        help="DDP backend. Defaults to nccl on CUDA and gloo otherwise.",
+    )
+    parser.add_argument(
+        "--prefetch_factor",
+        type=int,
+        default=4,
+        help="DataLoader batches prefetched by each worker.",
+    )
     args = parser.parse_args()
     if not 0.0 <= args.mask_ratio < 1.0:
         parser.error("--mask_ratio must be in [0, 1)")
+    if args.num_workers > 0 and args.prefetch_factor <= 0:
+        parser.error("--prefetch_factor must be positive")
     if args.gpu_ids:
         try:
             args.gpu_ids = [
@@ -145,7 +177,50 @@ def maybe_wrap_data_parallel(
 
 
 def unwrap_model(model) -> TimesNetExposureAutoencoder:
-    return model.module if isinstance(model, torch.nn.DataParallel) else model
+    if isinstance(model, (torch.nn.DataParallel, DistributedDataParallel)):
+        return model.module
+    return model
+
+
+def init_distributed(
+    args: argparse.Namespace,
+) -> tuple[torch.device, int, int, int]:
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    if world_size == 1:
+        return resolve_device(args.device), 0, 0, 1
+    if args.data_parallel:
+        raise ValueError("--data_parallel cannot be combined with torchrun/DDP")
+    local_rank = int(os.environ["LOCAL_RANK"])
+    rank = int(os.environ["RANK"])
+    if not torch.cuda.is_available():
+        raise ValueError("Multi-process exposure training requires CUDA")
+    torch.cuda.set_device(local_rank)
+    backend = args.ddp_backend or "nccl"
+    dist.init_process_group(backend=backend, init_method="env://")
+    return torch.device("cuda", local_rank), rank, local_rank, world_size
+
+
+def rank_logger(rank: int, run_dir: Path):
+    if rank == 0:
+        return setup_logging(run_dir)
+    logger = logging.getLogger(f"DeepHealth.rank{rank}")
+    logger.handlers.clear()
+    logger.addHandler(logging.NullHandler())
+    return logger
+
+
+def distributed_run_dir(
+    args: argparse.Namespace, rank: int, world_size: int
+) -> tuple[Path, str]:
+    payload: list[str | None] = [None, None]
+    if rank == 0:
+        run_dir, run_name = create_unique_run_dir(
+            lambda stamp: f"exposure_ae_{stamp}", Path(args.runs_root)
+        )
+        payload = [str(run_dir), run_name]
+    if world_size > 1:
+        dist.broadcast_object_list(payload, src=0)
+    return Path(str(payload[0])), str(payload[1])
 
 
 def channel_stats(
@@ -170,6 +245,61 @@ def channel_stats(
     return tuple(results)
 
 
+def eid_set_hash(eids: set[int]) -> str:
+    digest = hashlib.sha256()
+    for eid in sorted(eids):
+        digest.update(f"{eid}\n".encode("ascii"))
+    return digest.hexdigest()
+
+
+def load_or_compute_channel_stats(
+    cache: ExposureCache,
+    rows: np.ndarray,
+    train_eids: set[int],
+    stats_path: Path,
+    recompute: bool,
+    logger,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    eid_hash = eid_set_hash(train_eids)
+    if stats_path.is_file() and not recompute:
+        try:
+            with np.load(stats_path, allow_pickle=False) as saved:
+                compatible = (
+                    str(saved["train_eid_sha256"].item()) == eid_hash
+                    and int(saved["cache_event_rows"].item()) == len(cache.eids)
+                    and int(saved["train_window_rows"].item()) == len(rows)
+                )
+                if compatible:
+                    logger.info(f"Loading channel statistics from {stats_path}")
+                    return (
+                        saved["daily_mean"].astype(np.float32),
+                        saved["daily_std"].astype(np.float32),
+                        saved["monthly_mean"].astype(np.float32),
+                        saved["monthly_std"].astype(np.float32),
+                    )
+                logger.info("Channel statistics cache is stale; recomputing")
+        except (KeyError, OSError, ValueError) as exc:
+            logger.warning(
+                f"Could not read channel statistics cache ({exc}); recomputing"
+            )
+
+    logger.info("Computing channel statistics from training exposure")
+    stats = channel_stats(cache, rows)
+    stats_path.parent.mkdir(parents=True, exist_ok=True)
+    np.savez(
+        stats_path,
+        daily_mean=stats[0],
+        daily_std=stats[1],
+        monthly_mean=stats[2],
+        monthly_std=stats[3],
+        train_eid_sha256=np.asarray(eid_hash),
+        cache_event_rows=np.asarray(len(cache.eids), dtype=np.int64),
+        train_window_rows=np.asarray(len(rows), dtype=np.int64),
+    )
+    logger.info(f"Saved channel statistics to {stats_path}")
+    return stats
+
+
 def masked_mse(
     prediction: torch.Tensor, target: torch.Tensor, mask: torch.Tensor
 ) -> torch.Tensor:
@@ -187,15 +317,20 @@ def run_epoch(
     scaler: torch.amp.GradScaler,
     grad_clip: float,
     amp_enabled: bool,
+    show_progress: bool,
 ) -> float:
     training = optimizer is not None
     model.train(training)
-    total_loss = 0.0
+    loss_accumulator = torch.zeros(2, device=device, dtype=torch.float64)
-    total_samples = 0
     daily_mean, daily_std, monthly_mean, monthly_std = stats
     context = torch.enable_grad if training else torch.no_grad
     with context():
-        for batch in tqdm(loader, desc="train" if training else "val", leave=False):
+        for batch in tqdm(
+            loader,
+            desc="train" if training else "val",
+            leave=False,
+            disable=not show_progress,
+        ):
             daily = batch["daily"].to(device, non_blocking=True)
             monthly = batch["monthly"].to(device, non_blocking=True)
             daily_observed = torch.isfinite(daily)
@@ -239,9 +374,11 @@ def run_epoch(
                 scaler.step(optimizer)
                 scaler.update()
             batch_size = daily.size(0)
-            total_loss += float(loss.detach()) * batch_size
+            loss_accumulator[0] += loss.detach().double() * batch_size
-            total_samples += batch_size
+            loss_accumulator[1] += batch_size
-    return total_loss / max(total_samples, 1)
+    if dist.is_initialized():
+        dist.all_reduce(loss_accumulator, op=dist.ReduceOp.SUM)
+    return float((loss_accumulator[0] / loss_accumulator[1].clamp_min(1)).item())
 
 
 def learning_rate(epoch: int, args: argparse.Namespace) -> float:
@@ -255,13 +392,11 @@ def learning_rate(epoch: int, args: argparse.Namespace) -> float:
 
 def main() -> None:
     args = parse_args()
-    set_seed(args.seed)
+    device, rank, local_rank, world_size = init_distributed(args)
-    device = resolve_device(args.device)
+    set_seed(args.seed + rank)
     configure_torch_for_training(device)
-    run_dir, run_name = create_unique_run_dir(
+    run_dir, run_name = distributed_run_dir(args, rank, world_size)
-        lambda stamp: f"exposure_ae_{stamp}", Path(args.runs_root)
+    logger = rank_logger(rank, run_dir)
-    )
-    logger = setup_logging(run_dir)
     cache = ExposureCache(args.exposure_cache_dir)
     train_eids = load_eid_file(args.train_eid_file)
     val_eids = load_eid_file(args.val_eid_file)
@@ -269,22 +404,65 @@ def main() -> None:
         raise ValueError("train and validation EID files must be disjoint")
     train_rows = select_rows(cache, train_eids, "Training")
     val_rows = select_rows(cache, val_eids, "Validation")
-    raw_stats = channel_stats(cache, train_rows)
+    stats_path = (
+        Path(args.channel_stats_file)
+        if args.channel_stats_file
+        else Path(args.exposure_cache_dir) / "train_channel_stats.npz"
+    )
+    if rank == 0:
+        raw_stats = load_or_compute_channel_stats(
+            cache,
+            train_rows,
+            train_eids,
+            stats_path,
+            args.recompute_channel_stats,
+            logger,
+        )
+    if world_size > 1:
+        dist.barrier()
+    if rank != 0:
+        raw_stats = load_or_compute_channel_stats(
+            cache, train_rows, train_eids, stats_path, False, logger
+        )
     stats = tuple(
         torch.as_tensor(value, device=device).view(1, 1, -1)
         for value in raw_stats
     )
+    if args.batch_size % world_size != 0:
+        raise ValueError(
+            f"--batch_size={args.batch_size} must be divisible by "
+            f"DDP world size {world_size}"
+        )
+    local_batch_size = args.batch_size // world_size
     loader_kwargs = dict(
-        batch_size=args.batch_size,
+        batch_size=local_batch_size,
         num_workers=args.num_workers,
         pin_memory=device.type == "cuda",
         persistent_workers=args.num_workers > 0,
     )
+    if args.num_workers > 0:
+        loader_kwargs["prefetch_factor"] = args.prefetch_factor
+    train_dataset = ExposureWindowDataset(cache, train_rows)
+    val_dataset = ExposureWindowDataset(cache, val_rows)
+    train_sampler = (
+        DistributedSampler(
+            train_dataset, num_replicas=world_size, rank=rank,
+            shuffle=True, seed=args.seed,
+        )
+        if world_size > 1 else None
+    )
+    val_sampler = (
+        DistributedSampler(
+            val_dataset, num_replicas=world_size, rank=rank, shuffle=False
+        )
+        if world_size > 1 else None
+    )
     train_loader = DataLoader(
-        ExposureWindowDataset(cache, train_rows), shuffle=True, **loader_kwargs
+        train_dataset, sampler=train_sampler,
+        shuffle=train_sampler is None, **loader_kwargs
     )
     val_loader = DataLoader(
-        ExposureWindowDataset(cache, val_rows), shuffle=False, **loader_kwargs
+        val_dataset, sampler=val_sampler, shuffle=False, **loader_kwargs
     )
     model = TimesNetExposureAutoencoder(
         n_embd=args.n_embd, d_model=args.d_model, n_layers=args.n_layers,
@@ -293,6 +471,15 @@ def main() -> None:
         backbone_expansion=args.backbone_expansion,
         dropout=args.dropout,
     ).to(device)
+    if world_size > 1:
+        model = DistributedDataParallel(
+            model, device_ids=[local_rank], output_device=local_rank
+        )
+        logger.info(
+            f"Using DDP with {world_size} processes; "
+            f"global_batch={args.batch_size}, per_gpu_batch={local_batch_size}"
+        )
+    else:
         model = maybe_wrap_data_parallel(model, args, device, logger)
     optimizer = AdamW(
         model.parameters(), lr=args.base_lr,
@@ -312,6 +499,7 @@ def main() -> None:
         "monthly_mean": raw_stats[2].tolist(),
         "monthly_std": raw_stats[3].tolist(),
     }
+    if rank == 0:
         (run_dir / "train_config.json").write_text(
             json.dumps(config, indent=2), encoding="utf-8"
         )
@@ -320,16 +508,18 @@ def main() -> None:
     stale_epochs = 0
     history = []
     for epoch in range(args.max_epochs):
+        if train_sampler is not None:
+            train_sampler.set_epoch(epoch)
         lr = learning_rate(epoch, args)
         for group in optimizer.param_groups:
             group["lr"] = lr
         train_loss = run_epoch(
             model, train_loader, device, stats, args.mask_ratio, optimizer,
-            scaler, args.grad_clip, amp_enabled,
+            scaler, args.grad_clip, amp_enabled, rank == 0,
         )
         val_loss = run_epoch(
             model, val_loader, device, stats, 0.0, None,
-            scaler, args.grad_clip, amp_enabled,
+            scaler, args.grad_clip, amp_enabled, rank == 0,
         )
         logger.info(
             f"Epoch {epoch + 1:03d} | lr={lr:.3e} | "
@@ -342,6 +532,7 @@ def main() -> None:
         if val_loss < best_loss:
             best_loss = val_loss
             stale_epochs = 0
+            if rank == 0:
                 checkpoint_model = unwrap_model(model)
                 torch.save(
                     {
@@ -367,6 +558,7 @@ def main() -> None:
                 )
         else:
             stale_epochs += 1
+        if rank == 0:
             (run_dir / "history.json").write_text(
                 json.dumps(history, indent=2), encoding="utf-8"
             )
@@ -375,6 +567,8 @@ def main() -> None:
             break
     logger.info(f"Best validation loss: {best_loss:.6f}")
     logger.info(f"Checkpoint: {run_dir / 'best.pt'}")
+    if dist.is_initialized():
+        dist.destroy_process_group()
 
 
 if __name__ == "__main__":