diff --git a/README.md b/README.md index 1ae8091..655755d 100644 --- a/README.md +++ b/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 +`/train_channel_stats.npz`; use +`--recompute_channel_stats` only when a forced refresh is needed. diff --git a/backbones.py b/backbones.py index 0c42409..689a580 100644 --- a/backbones.py +++ b/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: diff --git a/train_exposure_autoencoder.py b/train_exposure_autoencoder.py index afdd203..55f9b76 100644 --- a/train_exposure_autoencoder.py +++ b/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 " + "/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 - total_samples = 0 + loss_accumulator = torch.zeros(2, device=device, dtype=torch.float64) 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 - total_samples += batch_size - return total_loss / max(total_samples, 1) + loss_accumulator[0] += loss.detach().double() * batch_size + loss_accumulator[1] += batch_size + 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 = resolve_device(args.device) + device, rank, local_rank, world_size = init_distributed(args) + set_seed(args.seed + rank) configure_torch_for_training(device) - run_dir, run_name = create_unique_run_dir( - lambda stamp: f"exposure_ae_{stamp}", Path(args.runs_root) - ) - logger = setup_logging(run_dir) + run_dir, run_name = distributed_run_dir(args, rank, world_size) + logger = rank_logger(rank, 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,7 +471,16 @@ def main() -> None: backbone_expansion=args.backbone_expansion, dropout=args.dropout, ).to(device) - model = maybe_wrap_data_parallel(model, args, device, logger) + 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, weight_decay=args.weight_decay, betas=(0.9, 0.95), @@ -312,24 +499,27 @@ def main() -> None: "monthly_mean": raw_stats[2].tolist(), "monthly_std": raw_stats[3].tolist(), } - (run_dir / "train_config.json").write_text( - json.dumps(config, indent=2), encoding="utf-8" - ) + if rank == 0: + (run_dir / "train_config.json").write_text( + json.dumps(config, indent=2), encoding="utf-8" + ) best_loss = float("inf") 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,39 +532,43 @@ def main() -> None: if val_loss < best_loss: best_loss = val_loss stale_epochs = 0 - checkpoint_model = unwrap_model(model) - torch.save( - { - "model_state_dict": checkpoint_model.state_dict(), - "encoder_state_dict": checkpoint_model.encoder.state_dict(), - "model_config": { - key: config[key] for key in ( - "n_embd", "d_model", "n_layers", "top_k", - "n_backbone_blocks", "backbone_kernel_size", - "backbone_expansion", "dropout", - ) + if rank == 0: + checkpoint_model = unwrap_model(model) + torch.save( + { + "model_state_dict": checkpoint_model.state_dict(), + "encoder_state_dict": checkpoint_model.encoder.state_dict(), + "model_config": { + key: config[key] for key in ( + "n_embd", "d_model", "n_layers", "top_k", + "n_backbone_blocks", "backbone_kernel_size", + "backbone_expansion", "dropout", + ) + }, + "normalization": { + "daily_mean": raw_stats[0], + "daily_std": raw_stats[1], + "monthly_mean": raw_stats[2], + "monthly_std": raw_stats[3], + }, + "epoch": epoch + 1, + "val_loss": val_loss, }, - "normalization": { - "daily_mean": raw_stats[0], - "daily_std": raw_stats[1], - "monthly_mean": raw_stats[2], - "monthly_std": raw_stats[3], - }, - "epoch": epoch + 1, - "val_loss": val_loss, - }, - run_dir / "best.pt", - ) + run_dir / "best.pt", + ) else: stale_epochs += 1 - (run_dir / "history.json").write_text( - json.dumps(history, indent=2), encoding="utf-8" - ) + if rank == 0: + (run_dir / "history.json").write_text( + json.dumps(history, indent=2), encoding="utf-8" + ) if stale_epochs >= args.patience: logger.info(f"Early stopping after {epoch + 1} epochs") 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__":