diff --git a/dingo/core/posterior_models/base_model.py b/dingo/core/posterior_models/base_model.py index 68508d199..2bb867aae 100755 --- a/dingo/core/posterior_models/base_model.py +++ b/dingo/core/posterior_models/base_model.py @@ -3,25 +3,43 @@ as well as functions for training and testing across an epoch. """ -from abc import abstractmethod, ABC +import ctypes +import json import os +import time +from abc import ABC, abstractmethod +from collections import OrderedDict +from multiprocessing import Value +from collections.abc import Sized from os.path import join -import h5py +from typing import Optional, Tuple -import torch -import dingo.core.utils as utils -from torch.utils.data import Dataset -import time +import h5py import numpy as np +import torch +import torch.distributed as dist from threadpoolctl import threadpool_limits +from torch.amp import autocast +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.utils.data import Dataset + +try: + from torch.amp import GradScaler +except ImportError: + # PyTorch < 2.3: GradScaler is not yet in torch.amp; use torch.cuda.amp. + # Wrap it to accept the same device-string call signature as the new API. + from torch.cuda.amp import GradScaler as _CudaGradScaler + + class GradScaler: # type: ignore[no-redef] + def __new__(cls, device="cuda", **kwargs): + return _CudaGradScaler(**kwargs) + + +import dingo.core.utils as utils import dingo.core.utils.trainutils -import json -from collections import OrderedDict -from typing import Optional from dingo.core.utils.backward_compatibility import update_model_config from dingo.core.utils.misc import get_version - -from dingo.core.utils.trainutils import EarlyStopping +from dingo.core.utils.trainutils import EarlyStopping, RuntimeLimits class BasePosteriorModel(ABC): @@ -60,6 +78,7 @@ def __init__( self.version = f"dingo={get_version()}" # dingo version self.device = None + self.rank = None self.optimizer_kwargs = None self.network_kwargs = None self.scheduler_kwargs = None @@ -72,6 +91,7 @@ def __init__( # separately, and before calling initialize_optimizer_and_scheduler(). self.epoch = 0 + self.iteration = 0 self.network = None self.optimizer = None self.scheduler = None @@ -185,19 +205,19 @@ def loss(self, theta: torch.Tensor, *context: torch.Tensor): """ pass - def network_to_device(self, device): + def network_to_device(self, device: str) -> None: """ - Put model to device, and set self.device accordingly. + Put model to device and set ``self.device`` accordingly. + + Accepts plain device strings (``"cpu"``, ``"cuda"``) as well as + rank-qualified CUDA strings (``"cuda:0"``, ``"cuda:1"``, …). In the + latter case ``self.rank`` is set to the integer rank index. """ - if device not in ("cpu", "cuda"): - raise ValueError(f"Device should be either cpu or cuda, got {device}.") + if "cpu" not in device and "cuda" not in device: + raise ValueError(f"Device should contain 'cpu' or 'cuda', got {device}.") + if ":" in device: + self.rank = int(device.split(":")[1]) self.device = torch.device(device) - # Commented below so that code runs on first cuda device in the case of multiple. - # if device == 'cuda' and torch.cuda.device_count() > 1: - # print("Using", torch.cuda.device_count(), "GPUs.") - # raise NotImplementedError('This needs testing!') - # # dim = 0 [512, ...] -> [256, ...], [256, ...] on 2 GPUs - # self.network = torch.nn.DataParallel(self.network) print(f"Putting posterior model to device {self.device}.") self.network.to(self.device) @@ -232,10 +252,17 @@ def save_model( saved, e.g. optimizer state dict """ + # Strip the DDP wrapper so the checkpoint can be loaded on any number of GPUs. + if isinstance(self.network, DDP): + model_state_dict = self.network.module.state_dict() + else: + model_state_dict = self.network.state_dict() + model_dict = { "model_kwargs": self.model_kwargs, - "model_state_dict": self.network.state_dict(), + "model_state_dict": model_state_dict, "epoch": self.epoch, + "iteration": self.iteration, "version": self.version, } @@ -328,9 +355,9 @@ def load_model( self.model_kwargs = d["model_kwargs"] update_model_config(self.model_kwargs) # For backward compatibility - self.initialize_network() self.epoch = d["epoch"] + self.iteration = d.get("iteration", 0) self.metadata = d["metadata"] @@ -341,6 +368,7 @@ def load_model( self.event_metadata = d["event_metadata"] if device != "meta": + self.initialize_network() self.network.load_state_dict(d["model_state_dict"]) self.network_to_device(device) @@ -366,70 +394,114 @@ def train( train_loader: torch.utils.data.DataLoader, test_loader: torch.utils.data.DataLoader, train_dir: str, - runtime_limits: object = None, - checkpoint_epochs: int = None, - use_wandb=False, - test_only=False, + train_sampler: Optional[torch.utils.data.DistributedSampler] = None, + runtime_limits: Optional[RuntimeLimits] = None, + checkpoint_epochs: Optional[int] = None, + use_wandb: bool = False, + test_only: bool = False, early_stopping: Optional[EarlyStopping] = None, + gradient_updates_per_optimizer_step: int = 1, + automatic_mixed_precision: bool = False, + world_size: int = 1, + global_epoch: ctypes.c_int = Value(ctypes.c_int, 1), ): """ + Train the network for one or more epochs. Parameters ---------- - train_loader - test_loader - train_dir - runtime_limits - checkpoint_epochs - use_wandb - test_only: bool = False - if True, training is skipped - early_stopping: EarlyStopping - Optional EarlyStopping instance. - - Returns - ------- - + train_loader : DataLoader + test_loader : DataLoader + train_dir : str + Directory for saving models and history. + train_sampler : DistributedSampler, optional + Required for DDP training to re-shuffle data each epoch. + runtime_limits : RuntimeLimits, optional + checkpoint_epochs : int, optional + Save a named checkpoint every this many epochs. + use_wandb : bool + test_only : bool + If True, skip training and only evaluate on the test set. + early_stopping : EarlyStopping, optional + gradient_updates_per_optimizer_step : int + Accumulate gradients over this many batches before calling + ``optimizer.step()``. Use >1 to simulate a larger effective batch + size without increasing per-GPU memory usage. + automatic_mixed_precision : bool + Train with ``torch.amp`` mixed precision (FP16/BF16 forward pass, + FP32 parameter updates). + world_size : int + Number of GPUs (used only for logging the effective batch size). + global_epoch : multiprocessing.Value + Shared counter updated so that the WaveformDataset can query the + current epoch from any worker process. """ + is_primary = self.rank is None or self.rank == 0 if test_only: - test_loss = test_epoch(self, test_loader) - print(f"test loss: {test_loss:.3f}") - - else: - while not runtime_limits.limits_exceeded(self.epoch): - self.epoch += 1 - - # Training - lr = utils.get_lr(self.optimizer) - with threadpool_limits(limits=1, user_api="blas"): + test_loss = test_epoch(self, dataloader=test_loader) + if is_primary: + print(f"test loss: {test_loss:.3f}") + return + + while not runtime_limits.limits_exceeded(self.epoch): + self.epoch += 1 + global_epoch.value = self.epoch + + lr = utils.get_lr(self.optimizer) + with threadpool_limits(limits=1, user_api="blas"): + if train_sampler is not None: + # Ensure each epoch sees a different random shuffle. + train_sampler.set_epoch(self.epoch) + + if is_primary: print(f"\nStart training epoch {self.epoch} with lr {lr}") - time_start = time.time() - train_loss = train_epoch(self, train_loader) - train_time = time.time() - time_start + time_start = torch.tensor(time.time(), device=self.device, dtype=torch.float64) + + train_loss, n_iter = train_epoch( + self, + dataloader=train_loader, + gradient_updates_per_optimizer_step=gradient_updates_per_optimizer_step, + automatic_mixed_precision=automatic_mixed_precision, + world_size=world_size, + ) + self.iteration += n_iter + if self.rank is not None: + dist.barrier() + dist.all_reduce(time_start, op=dist.ReduceOp.MIN) + train_time = time.time() - time_start.item() + + if is_primary: print( "Done. This took {:2.0f}:{:2.0f} min.".format( *divmod(train_time, 60) ) ) - - # Testing print(f"Start testing epoch {self.epoch}") - time_start = time.time() - test_loss = test_epoch(self, test_loader) - test_time = time.time() - time_start + time_start = torch.tensor(time.time(), device=self.device, dtype=torch.float64) + test_loss = test_epoch( + self, + dataloader=test_loader, + gradient_updates_per_optimizer_step=gradient_updates_per_optimizer_step, + world_size=world_size, + ) + if self.rank is not None: + dist.barrier() + dist.all_reduce(time_start, op=dist.ReduceOp.MIN) + test_time = time.time() - time_start.item() + + if is_primary: print( "Done. This took {:2.0f}:{:2.0f} min.".format( - *divmod(time.time() - time_start, 60) + *divmod(test_time, 60) ) ) - # scheduler step for learning rate - utils.perform_scheduler_step(self.scheduler, test_loss) + utils.perform_scheduler_step(self.scheduler, test_loss) - # write history and save model + if is_primary: utils.write_history(train_dir, self.epoch, train_loss, test_loss, lr) utils.save_model(self, train_dir, checkpoint_epochs=checkpoint_epochs) if use_wandb: @@ -451,70 +523,179 @@ def train( except ImportError: print("wandb not installed. Skipping logging to wandb.") - if early_stopping is not None: - # Whether to use train or test loss - early_stopping_loss = ( - test_loss - if early_stopping.metric == "validation" - else train_loss + if early_stopping is not None: + early_stopping_loss = ( + test_loss if early_stopping.metric == "validation" else train_loss + ) + is_best_model = early_stopping(early_stopping_loss) + if is_best_model and is_primary: + self.save_model( + join(train_dir, "best_model.pt"), save_training_info=False ) - is_best_model = early_stopping(early_stopping_loss) - if is_best_model: - self.save_model( - join(train_dir, "best_model.pt"), save_training_info=False - ) - if early_stopping.early_stop: + if early_stopping.early_stop: + if is_primary: print("Early stopping") - break + break + + if is_primary: print(f"Finished training epoch {self.epoch}.\n") -def train_epoch(pm, dataloader): +def _dataset_len(dataloader: torch.utils.data.DataLoader) -> int: + """Return the number of samples in a DataLoader's dataset. + + ``DataLoader.dataset`` is typed as ``Dataset``, which does not inherit from + ``Sized`` in PyTorch's stubs even though all concrete datasets implement + ``__len__``. This helper asserts the contract and returns the length. + """ + dataset = dataloader.dataset + if not isinstance(dataset, Sized): + raise TypeError( + f"DataLoader dataset of type {type(dataset).__name__} does not " + "implement __len__. Cannot determine dataset length." + ) + return len(dataset) + + +def train_epoch( + pm: BasePosteriorModel, + dataloader: torch.utils.data.DataLoader, + gradient_updates_per_optimizer_step: int = 1, + automatic_mixed_precision: bool = False, + world_size: int = 1, +) -> Tuple[float, int]: + """ + Train the network for one epoch. + + Parameters + ---------- + pm : BasePosteriorModel + dataloader : DataLoader + gradient_updates_per_optimizer_step : int + Accumulate gradients over this many mini-batches before stepping. + Values >1 simulate a larger effective batch without extra GPU memory. + automatic_mixed_precision : bool + Use ``torch.amp`` (FP16 forward pass, FP32 updates). + world_size : int + Number of GPUs, used only for logging the effective batch size. + + Returns + ------- + Tuple[float, int] + Average loss over the epoch and number of optimizer steps performed. + """ pm.network.train() + + if pm.rank is None: + effective_bs = dataloader.batch_size + else: + effective_bs = dataloader.batch_size * world_size + loss_info = dingo.core.utils.trainutils.LossInfo( - pm.epoch, - len(dataloader.dataset), - dataloader.batch_size, + epoch=pm.epoch, + len_dataset=_dataset_len(dataloader), + batch_size_per_grad_update=effective_bs, mode="Train", print_freq=1, + device=pm.device, ) + scaler = GradScaler("cuda") if automatic_mixed_precision else None + for batch_idx, data in enumerate(dataloader): - loss_info.update_timer() - pm.optimizer.zero_grad() - # data to device + loss_info.update_timer("Dataloader") + + if batch_idx % gradient_updates_per_optimizer_step == 0: + pm.optimizer.zero_grad(set_to_none=True) + data = [d.to(pm.device, non_blocking=True) for d in data] - # compute loss - loss = pm.loss(data[0], *data[1:]) - # backward pass and optimizer step - loss.backward() - pm.optimizer.step() - # update loss for history and logging - loss_info.update(loss.detach().item(), len(data[0])) - loss_info.print_info(batch_idx) - return loss_info.get_avg() + if automatic_mixed_precision: + with autocast("cuda"): + result = pm.loss(data[0], *data[1:]) + loss = result[0] if isinstance(result, tuple) else result + scaler.scale(loss).backward() + else: + result = pm.loss(data[0], *data[1:]) + loss = result[0] if isinstance(result, tuple) else result + loss.backward() + + loss_info.cache_loss(loss=loss, n=len(data[0])) + + if (batch_idx + 1) % gradient_updates_per_optimizer_step == 0: + if automatic_mixed_precision: + scaler.step(pm.optimizer) + scaler.update() + else: + pm.optimizer.step() + + loss_info.update() + if pm.rank is None or pm.rank == 0: + loss_info.print_info(batch_idx) + return loss_info.get_avg(), loss_info.get_iteration() -def test_epoch(pm, dataloader): + +def test_epoch( + pm: BasePosteriorModel, + dataloader: torch.utils.data.DataLoader, + gradient_updates_per_optimizer_step: int = 1, + world_size: int = 1, +) -> float: + """ + Evaluate the network on the test set. + + Parameters + ---------- + pm : BasePosteriorModel + dataloader : DataLoader + gradient_updates_per_optimizer_step : int + Used to match the effective batch size of the training loop for + comparable loss values. + world_size : int + Number of GPUs, used only for logging. + + Returns + ------- + float + Average loss over the test set. + """ with torch.no_grad(): pm.network.eval() + + if pm.rank is None: + effective_bs = dataloader.batch_size + else: + effective_bs = dataloader.batch_size * world_size + + if _dataset_len(dataloader) < effective_bs: + if pm.rank is None or pm.rank == 0: + print( + f"Warning: test dataset (len {_dataset_len(dataloader)}) is smaller " + f"than effective_batch_size={effective_bs}. " + "Test loss computed over full dataset; may not be comparable to train loss." + ) + effective_bs = _dataset_len(dataloader) + gradient_updates_per_optimizer_step = 1 + loss_info = dingo.core.utils.trainutils.LossInfo( - pm.epoch, - len(dataloader.dataset), - dataloader.batch_size, + epoch=pm.epoch, + len_dataset=_dataset_len(dataloader), + batch_size_per_grad_update=effective_bs, mode="Test", print_freq=1, + device=pm.device, ) for batch_idx, data in enumerate(dataloader): loss_info.update_timer() - # data to device data = [d.to(pm.device, non_blocking=True) for d in data] - # compute loss - loss = pm.loss(data[0], *data[1:]) - # update loss for history and logging - loss_info.update(loss.item(), len(data[0])) - loss_info.print_info(batch_idx) + result = pm.loss(data[0], *data[1:]) + loss = result[0] if isinstance(result, tuple) else result + loss_info.cache_loss(loss, len(data[0])) + if (batch_idx + 1) % gradient_updates_per_optimizer_step == 0: + loss_info.update() + if pm.rank is None or pm.rank == 0: + loss_info.print_info(batch_idx) return loss_info.get_avg() diff --git a/dingo/core/utils/torchutils.py b/dingo/core/utils/torchutils.py index 3b37bca9f..e326695be 100644 --- a/dingo/core/utils/torchutils.py +++ b/dingo/core/utils/torchutils.py @@ -1,21 +1,135 @@ +import os +from socket import gethostname +from typing import Any, Iterable, Optional, Tuple, Union + +import bilby import numpy as np import torch +import torch.distributed as dist import torch.nn as nn from torch.nn import functional as F -from torch.utils.data import DataLoader -from typing import Union, Tuple, Iterable -import bilby +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.utils.data import DataLoader, DistributedSampler def fix_random_seeds(_): """Utility function to set random seeds when using multiple workers for DataLoader.""" - np.random.seed(int(torch.initial_seed()) % (2 ** 32 - 1)) + np.random.seed(int(torch.initial_seed()) % (2**32 - 1)) try: - bilby.core.utils.random.seed(int(torch.initial_seed()) % (2 ** 32 - 1)) + bilby.core.utils.random.seed(int(torch.initial_seed()) % (2**32 - 1)) except AttributeError: # In case using an old version of Bilby. pass +def get_cuda_info() -> dict[str, Any]: + """Get information about the CUDA devices available in the system.""" + if not torch.cuda.is_available(): + return {} + return { + "cuDNN version": torch.backends.cudnn.version(), + "CUDA version": torch.version.cuda, + "device count": torch.cuda.device_count(), + "device name": torch.cuda.get_device_name(0), + "memory (GB)": round( + torch.cuda.get_device_properties(0).total_memory / 1024**3, 1 + ), + } + + +def document_gpus(target_dir: str) -> None: + """ + Document the current GPU resources to an ``info_gpus.txt`` file inside + *target_dir*. + """ + cuda_info = get_cuda_info() + with open(os.path.join(target_dir, "info_gpus.txt"), "w") as f: + f.write(f"# Running on host:\n{gethostname()}\n") + f.write("# CUDA information:\n") + for k, v in cuda_info.items(): + f.write(f"{k}: {v}\n") + + +def set_seed_based_on_rank(rank: int) -> None: + """ + Set NumPy and Torch seeds for a DDP worker process based on *rank* so that + each process draws different random samples. + """ + initial_torch_seed = torch.initial_seed() + torch.manual_seed(initial_torch_seed + rank) + if torch.cuda.is_available(): + torch.cuda.manual_seed(initial_torch_seed + rank) + torch.backends.cudnn.deterministic = True + # NumPy expects seeds in [0, 2**32). + reduced_seed = int(initial_torch_seed) % (2**32 - 1) + np.random.seed(reduced_seed + rank) + + +def setup_ddp(rank: int, world_size: int, port: int = 12355) -> None: + """ + Initialise the NCCL process group for DDP training. + + Parameters + ---------- + rank : int + Rank of this process within the group. + world_size : int + Total number of processes (= number of GPUs). + port : int + Port used for the ``MASTER_ADDR`` rendezvous. When running multiple + experiments on the same node, choose a different port for each to avoid + collisions. + """ + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + + if dist.is_nccl_available(): + dist.init_process_group(backend="nccl", rank=rank, world_size=world_size) + else: + raise RuntimeError( + "NCCL backend is not available. " + "Fall back to single-GPU training or install a CUDA-enabled PyTorch build." + ) + torch.cuda.set_device(rank) + print( + f"Process group initialised: backend={dist.get_backend()}, " + f"rank={dist.get_rank()}, world_size={dist.get_world_size()}." + ) + + +def cleanup_ddp() -> None: + """Tear down the distributed process group.""" + dist.destroy_process_group() + print("Destroyed process group.") + + +def replace_BatchNorm_with_SyncBatchNorm(network: nn.Module) -> nn.Module: + """Replace all BatchNorm layers with SyncBatchNorm for DDP training.""" + return nn.SyncBatchNorm.convert_sync_batchnorm(network) + + +def print_number_of_model_parameters(network: nn.Module) -> None: + """ + Print the number of fixed and learnable parameters of *network*. + Handles DDP-wrapped networks transparently. + """ + bare = network.module if isinstance(network, DDP) else network + + n_grad = get_number_of_model_parameters(network, (True,)) + n_nograd = get_number_of_model_parameters(network, (False,)) + print(f"Fixed parameters: {n_nograd}\nLearnable parameters: {n_grad}") + + try: + if bare.name == "FlowWrapper": + n_emb = get_number_of_model_parameters(bare.embedding_net, (True,)) + n_flow = get_number_of_model_parameters(bare.flow, (True,)) + print( + f" - learnable embedding network parameters: {n_emb} ({n_emb / n_grad * 100:.2f}%)\n" + f" - learnable flow parameters: {n_flow} ({n_flow / n_grad * 100:.2f}%)" + ) + except Exception: + pass + + def get_activation_function_from_string(activation_name: str): """ Returns an activation function, based on the name provided. @@ -191,11 +305,16 @@ def build_train_and_test_loaders( train_fraction: float, batch_size: int, num_workers: int, -): + world_size: Optional[int] = None, + rank: Optional[int] = None, +) -> Tuple[DataLoader, DataLoader, Optional[DistributedSampler]]: """ Split the dataset into train and test sets, and build corresponding DataLoaders. The random split uses a fixed seed for reproducibility. + When *world_size* and *rank* are given, ``DistributedSampler`` instances are + created so that each GPU processes a non-overlapping shard of the data. + Parameters ---------- dataset : torch.utils.data.Dataset @@ -203,11 +322,17 @@ def build_train_and_test_loaders( Fraction of dataset to use for training. The remainder is used for testing. Should lie between 0 and 1. batch_size : int + Batch size *per GPU*. num_workers : int + world_size : int, optional + Total number of DDP processes (GPUs). + rank : int, optional + Rank of the current DDP process. Returns ------- - (train_loader, test_loader) + (train_loader, test_loader, train_sampler) + *train_sampler* is ``None`` for single-GPU training. """ # Split the dataset. This function uses a fixed seed for reproducibility. @@ -215,25 +340,56 @@ def build_train_and_test_loaders( dataset, train_fraction ) - # Build DataLoaders - train_loader = DataLoader( - train_dataset, - batch_size=batch_size, - shuffle=True, - pin_memory=True, - num_workers=num_workers, - worker_init_fn=fix_random_seeds, - ) - test_loader = DataLoader( - test_dataset, - batch_size=batch_size, - shuffle=False, - pin_memory=True, - num_workers=num_workers, - worker_init_fn=fix_random_seeds, - ) - - return train_loader, test_loader + persistent_workers = num_workers > 0 + + if rank is not None and world_size is not None: + # DDP path: shuffle is handled by the sampler. + train_sampler = DistributedSampler( + train_dataset, shuffle=True, num_replicas=world_size, rank=rank + ) + test_sampler = DistributedSampler( + test_dataset, shuffle=False, num_replicas=world_size, rank=rank + ) + train_loader = DataLoader( + train_dataset, + batch_size=batch_size, + sampler=train_sampler, + pin_memory=False, + num_workers=num_workers, + worker_init_fn=fix_random_seeds, + persistent_workers=persistent_workers, + ) + test_loader = DataLoader( + test_dataset, + batch_size=batch_size, + sampler=test_sampler, + pin_memory=False, + num_workers=num_workers, + worker_init_fn=fix_random_seeds, + persistent_workers=persistent_workers, + ) + else: + train_sampler = None + train_loader = DataLoader( + train_dataset, + batch_size=batch_size, + shuffle=True, + pin_memory=True, + num_workers=num_workers, + worker_init_fn=fix_random_seeds, + persistent_workers=persistent_workers, + ) + test_loader = DataLoader( + test_dataset, + batch_size=batch_size, + shuffle=False, + pin_memory=True, + num_workers=num_workers, + worker_init_fn=fix_random_seeds, + persistent_workers=persistent_workers, + ) + + return train_loader, test_loader, train_sampler def set_requires_grad_flag( diff --git a/dingo/core/utils/trainutils.py b/dingo/core/utils/trainutils.py index bc7d12154..d7fdde496 100644 --- a/dingo/core/utils/trainutils.py +++ b/dingo/core/utils/trainutils.py @@ -1,9 +1,12 @@ -import time +import csv import os +import time +from os.path import isfile, join +from typing import Literal, Optional + import numpy as np -from os.path import join, isfile -import csv -from typing import Literal +import torch +import torch.distributed as dist class AvgTracker: @@ -96,57 +99,147 @@ def __call__(self, val_loss: float): class LossInfo: - def __init__(self, epoch, len_dataset, batch_size, mode="Train", print_freq=1): + def __init__( + self, + epoch: int, + len_dataset: int, + batch_size_per_grad_update: int, + mode: str = "Train", + print_freq: int = 1, + device: torch.device = torch.device("cuda"), + ): # data for print statements self.epoch = epoch + self.iteration = 0 self.len_dataset = len_dataset - self.batch_size = batch_size + self.batch_size_per_grad_update = batch_size_per_grad_update self.mode = mode self.print_freq = print_freq + self.device = device # track loss self.loss_tracker = AvgTracker() self.loss = None - # track computation times - self.times = {"Dataloader": AvgTracker(), "Network": AvgTracker()} - self.t = time.time() - - def update_timer(self, timer_mode="Dataloader"): - self.times[timer_mode].update(time.time() - self.t) + self.cached_losses: list = [] + self.cached_n: list = [] + # Use dist.is_initialized() so that LossInfo is DDP-aware only when a + # process group has actually been set up. + self.is_ddp = dist.is_initialized() + if self.is_ddp: + self.num_gpus = dist.get_world_size() + self.times = { + "Dataloader": AvgTracker(), + "Network": AvgTracker(), + "Aggregation": AvgTracker(), + } + else: + self.times = {"Dataloader": AvgTracker(), "Network": AvgTracker()} self.t = time.time() - def update(self, loss, n): - self.loss = loss - self.loss_tracker.update(loss * n, n) + def cache_loss(self, loss: torch.Tensor, n: int) -> None: + """Cache *loss* from one gradient-accumulation step.""" + self.cached_losses.append(loss.detach()) + self.cached_n.append(n) self.update_timer(timer_mode="Network") - def get_avg(self): + def _reset_cached_losses(self) -> None: + self.cached_losses = [] + self.cached_n = [] + + def update_timer(self, timer_mode: str = "Dataloader") -> None: + if self.is_ddp: + dt = torch.tensor(time.time() - self.t, device=self.device) + dist.barrier() + dist.reduce(dt, dst=0, op=dist.ReduceOp.MAX) + dt = dt.item() + else: + dt = time.time() - self.t + self.times[timer_mode].update(dt) + self.t = time.time() + + def update(self) -> None: + """Aggregate cached losses across gradient-accumulation steps (and across GPUs in DDP).""" + self.iteration += 1 + loss = torch.mean(torch.tensor(self.cached_losses, device=self.device)) + n = torch.tensor(sum(self.cached_n), device=self.device, dtype=torch.float32) + + if self.is_ddp: + # Reduce absolute loss across GPUs so that normalization is correct + # even when GPUs process different numbers of samples. + abs_loss = loss * n + dist.barrier() + dist.reduce(abs_loss, dst=0) + dist.reduce(n, dst=0) + loss = abs_loss / n + self.update_timer(timer_mode="Aggregation") + + self.loss = loss.item() + self.loss_tracker.update(self.loss * n.item(), n.item()) + self._reset_cached_losses() + + def get_avg(self) -> float: return self.loss_tracker.get_avg() - def print_info(self, batch_idx): + def get_iteration(self) -> int: + """Return the number of optimizer steps performed this epoch.""" + if self.is_ddp: + dist.barrier() + iteration = torch.tensor( + self.iteration, device=self.device, dtype=torch.int64 + ) + # all_gather requires output tensors to match the shape of the input. + gathered = [torch.zeros_like(iteration) for _ in range(self.num_gpus)] + dist.all_gather(gathered, iteration) + if not all(torch.equal(gathered[0], t) for t in gathered): + raise ValueError( + f"DDP ranks disagree on iteration count: " + f"{[t.item() for t in gathered]}" + ) + return self.iteration + + def print_info(self, batch_idx: int) -> None: if batch_idx % self.print_freq == 0: print( "{} Epoch: {} [{}/{} ({:.0f}%)]".format( self.mode, self.epoch, - min(batch_idx * self.batch_size, self.len_dataset), + min( + (batch_idx + 1) * self.batch_size_per_grad_update, + self.len_dataset, + ), self.len_dataset, - 100.0 * batch_idx * self.batch_size / self.len_dataset, + min( + 100.0 + * (batch_idx + 1) + * self.batch_size_per_grad_update + / self.len_dataset, + 100, + ), ), end="\t\t", ) - # print loss print(f"Loss: {self.loss:.3f} ({self.get_avg():.3f})", end="\t\t") - # print computation times td, td_avg = self.times["Dataloader"].x, self.times["Dataloader"].get_avg() tn, tn_avg = self.times["Network"].x, self.times["Network"].get_avg() print(f"Time Dataloader: {td:.3f} ({td_avg:.3f})", end="\t\t") - print(f"Time Network: {tn:.3f} ({tn_avg:.3f})") + if self.is_ddp: + ta, ta_avg = ( + self.times["Aggregation"].x, + self.times["Aggregation"].get_avg(), + ) + print(f"Time Network: {tn:.3f} ({tn_avg:.3f})", end="\t\t") + print(f"Time Loss Aggregation: {ta:.3f} ({ta_avg:.3f})") + else: + print(f"Time Network: {tn:.3f} ({tn_avg:.3f})") class RuntimeLimits: """ Keeps track of the runtime limits (time limit, epoch limit, max. number of epochs for model). + + In DDP training, ``limits_exceeded`` broadcasts the result across all ranks + so that every process stops at the same epoch even when only one rank (e.g. + rank 0) detects that the wall-clock limit has been reached. """ def __init__( @@ -155,6 +248,7 @@ def __init__( max_epochs_per_run: int = None, max_epochs_total: int = None, epoch_start: int = None, + device: torch.device = torch.device("cuda"), ): """ @@ -169,19 +263,26 @@ def __init__( maximum total number of epochs for model epoch_start: int = None start epoch of run + device: torch.device + Device used for the DDP all-reduce broadcast. """ self.max_time_per_run = max_time_per_run self.max_epochs_per_run = max_epochs_per_run self.max_epochs_total = max_epochs_total self.epoch_start = epoch_start self.time_start = time.time() + self.device = device + self.is_ddp = dist.is_initialized() if max_epochs_per_run is not None and epoch_start is None: - raise ValueError("epoch_start required to check " "max_epochs_per_run.") + raise ValueError("epoch_start required to check max_epochs_per_run.") - def limits_exceeded(self, epoch: int = None): + def limits_exceeded(self, epoch: Optional[int] = None) -> bool: """ Check whether any of the runtime limits are exceeded. + In DDP mode the boolean is broadcast via ``all_reduce`` so all ranks + agree on whether to stop. + Parameters ---------- epoch: int = None @@ -192,13 +293,12 @@ def limits_exceeded(self, epoch: int = None): flag whether runtime limits are exceeded and run should be stopped; if limits_exceeded = True, this prints a message for the reason """ + exceeded = False # check time limit for run if self.max_time_per_run is not None: if time.time() - self.time_start >= self.max_time_per_run: - print( - f"Stop run: Time limit of {self.max_time_per_run} s " f"exceeded." - ) - return True + print(f"Stop run: Time limit of {self.max_time_per_run} s exceeded.") + exceeded = True # check epoch limit for run if self.max_epochs_per_run is not None: if epoch is None: @@ -207,18 +307,24 @@ def limits_exceeded(self, epoch: int = None): print( f"Stop run: Epoch limit of {self.max_epochs_per_run} per run reached." ) - return True + exceeded = True # check total epoch limit if self.max_epochs_total is not None: if epoch >= self.max_epochs_total: print( f"Stop run: Total epoch limit of {self.max_epochs_total} reached." ) - return True - # return False if none of the limits is exceeded - return False + exceeded = True + + if self.is_ddp: + flag = torch.tensor(exceeded, device=self.device, dtype=torch.bool) + dist.barrier() + dist.all_reduce(flag, op=dist.ReduceOp.MAX) + exceeded = flag.item() + + return exceeded - def local_limits_exceeded(self, epoch: int = None): + def local_limits_exceeded(self, epoch: Optional[int] = None) -> bool: """ Check whether any of the local runtime limits are exceeded. Local runtime limits include max_epochs_per_run and max_time_per_run, but not max_epochs_total. diff --git a/dingo/gw/training/train_builders.py b/dingo/gw/training/train_builders.py index ca62856bb..7a0c7dacc 100755 --- a/dingo/gw/training/train_builders.py +++ b/dingo/gw/training/train_builders.py @@ -1,32 +1,32 @@ -from typing import List, Optional import copy +from typing import List, Optional +import numpy as np import torch.multiprocessing import torchvision -from threadpoolctl import threadpool_limits from bilby.gw.detector import InterferometerList +from threadpoolctl import threadpool_limits -from dingo.gw.SVD import SVDBasis - +from dingo.core.utils import * from dingo.gw.dataset.waveform_dataset import WaveformDataset from dingo.gw.domains import build_domain +from dingo.gw.gwutils import * +from dingo.gw.noise.asd_dataset import ASDDataset +from dingo.gw.prior import default_inference_parameters +from dingo.gw.SVD import SVDBasis from dingo.gw.transforms import ( + AddWhiteNoiseComplex, + CropMaskStrainRandom, + GetDetectorTimes, + GNPECoalescenceTimes, ProjectOntoDetectors, + RepackageStrainsAndASDS, + SampleExtrinsicParameters, SampleNoiseASD, - WhitenAndScaleStrain, - AddWhiteNoiseComplex, SelectStandardizeRepackageParameters, - RepackageStrainsAndASDS, UnpackDict, - GNPECoalescenceTimes, - SampleExtrinsicParameters, - GetDetectorTimes, - CropMaskStrainRandom, + WhitenAndScaleStrain, ) -from dingo.gw.noise.asd_dataset import ASDDataset -from dingo.gw.prior import default_inference_parameters -from dingo.gw.gwutils import * -from dingo.core.utils import * def build_dataset( @@ -63,7 +63,13 @@ def build_dataset( return wfd -def set_train_transforms(wfd, data_settings, asd_dataset_path, omit_transforms=None): +def set_train_transforms( + wfd: WaveformDataset, + data_settings: dict, + asd_dataset_path: str, + omit_transforms: Optional[List[type]] = None, + print_output: bool = True, +) -> None: """ Set the transform attribute of a waveform dataset based on a settings dictionary. The transform takes waveform polarizations, samples random extrinsic parameters, @@ -80,11 +86,15 @@ def set_train_transforms(wfd, data_settings, asd_dataset_path, omit_transforms=N Path corresponding to the ASD dataset used to generate noise. omit_transforms : List of sub-transforms to omit from the full composition. + print_output : bool + Whether to write informational messages to stdout. Set to False for + non-primary DDP ranks. """ - print(f"Setting train transforms.") - if omit_transforms is not None: - print("Omitting \n\t" + "\n\t".join([t.__name__ for t in omit_transforms])) + if print_output: + print(f"Setting train transforms.") + if omit_transforms is not None: + print("Omitting \n\t" + "\n\t".join([t.__name__ for t in omit_transforms])) # By passing the wfd domain when instantiating the noise dataset, this ensures the # domains will match. In particular, it truncates the ASD dataset beyond the new @@ -142,9 +152,11 @@ def set_train_transforms(wfd, data_settings, asd_dataset_path, omit_transforms=N # parameters. try: standardization_dict = data_settings["standardization"] - print("Using previously-calculated parameter standardizations.") + if print_output: + print("Using previously-calculated parameter standardizations.") except KeyError: - print("Calculating new parameter standardizations.") + if print_output: + print("Calculating new parameter standardizations.") standardization_dict = get_standardization_dict( extrinsic_prior_dict, wfd, @@ -201,7 +213,7 @@ def build_svd_for_embedding_network( num_workers: int = 0, batch_size: int = 1000, out_dir: Optional[str] = None, -) -> List: +) -> List[np.ndarray]: """ Construct SVD matrices V based on clean waveforms in each interferometer. These will be used to seed the weights of the initial projection part of the embedding @@ -277,7 +289,7 @@ def build_svd_for_embedding_network( loader = DataLoader( wfd, batch_size=batch_size, - num_workers= 0, + num_workers=0, worker_init_fn=fix_random_seeds, ) with threadpool_limits(limits=1, user_api="blas"): diff --git a/dingo/gw/training/train_pipeline.py b/dingo/gw/training/train_pipeline.py index 57ca57061..8fc2a0b44 100644 --- a/dingo/gw/training/train_pipeline.py +++ b/dingo/gw/training/train_pipeline.py @@ -1,38 +1,53 @@ -from typing import Optional, Tuple -import os - -import numpy as np -import yaml import argparse +import ctypes +import os +import queue import shutil import textwrap import time from copy import deepcopy +from multiprocessing import Value +from typing import Optional, Tuple +import numpy as np +import torch +import torch.multiprocessing as mp +import yaml from threadpoolctl import threadpool_limits +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.utils.data import DataLoader, DistributedSampler +from dingo.core.posterior_models.base_model import BasePosteriorModel from dingo.core.posterior_models.build_model import ( autocomplete_model_kwargs, build_model_from_kwargs, ) -from dingo.gw.training.train_builders import ( - build_dataset, - set_train_transforms, - build_svd_for_embedding_network, -) -from dingo.core.utils.trainutils import RuntimeLimits from dingo.core.utils import ( - set_requires_grad_flag, - get_number_of_model_parameters, build_train_and_test_loaders, + print_number_of_model_parameters, + set_requires_grad_flag, ) -from dingo.core.utils.trainutils import EarlyStopping +from dingo.core.utils.torchutils import ( + cleanup_ddp, + document_gpus, + replace_BatchNorm_with_SyncBatchNorm, + set_seed_based_on_rank, + setup_ddp, +) +from dingo.core.utils.trainutils import EarlyStopping, RuntimeLimits from dingo.gw.dataset import WaveformDataset -from dingo.core.posterior_models import BasePosteriorModel +from dingo.gw.training.train_builders import ( + build_dataset, + build_svd_for_embedding_network, + set_train_transforms, +) def copy_files_to_local( - file_path: str, local_dir: Optional[str], leave_keys_on_disk: bool, is_condor: bool = False, + file_path: str, + local_dir: Optional[str], + leave_keys_on_disk: bool, + is_condor: bool = False, ) -> str: """ Copy files to local node if local_dir is provided to minimize network traffic during training. @@ -83,7 +98,7 @@ def prepare_training_new( """ Based on a settings dictionary, initialize a WaveformDataset and PosteriorModel. - For model type 'nsf+embedding' (the only acceptable type at this point) this also + For the standard DenseResNet embedding network type (the only acceptable type at this point) this also initializes the embedding network projection stage with SVD V matrices based on clean detector waveforms. @@ -100,47 +115,10 @@ def prepare_training_new( ------- (BasePosteriorModel, WaveformDataset) """ - data_settings = deepcopy(train_settings["data"]) - # Optionally copy files to local and update path - data_settings["waveform_dataset_path"] = copy_files_to_local( - file_path=data_settings["waveform_dataset_path"], - local_dir=local_settings.get("local_cache_path", None), - leave_keys_on_disk=local_settings.get("leave_waveforms_on_disk", True), - is_condor=True if "condor" in local_settings else False, - ) - wfd = build_dataset( - data_settings=data_settings, - leave_waveforms_on_disk=local_settings.get("leave_waveforms_on_disk", True), - ) # No transforms yet - initial_weights = {} - - # The embedding network is assumed to have an SVD projection layer. If other types - # of embedding networks are added in the future, update this code. - - if train_settings["model"].get("embedding_kwargs", None): - # First, build the SVD for seeding the embedding network. - print("\nBuilding SVD for initialization of embedding network.") - initial_weights["V_rb_list"] = build_svd_for_embedding_network( - wfd, - train_settings["data"], - train_settings["training"]["stage_0"]["asd_dataset_path"], - num_workers=local_settings["num_workers"], - batch_size=train_settings["training"]["stage_0"]["batch_size"], - out_dir=train_dir, - **train_settings["model"]["embedding_kwargs"]["svd"], - ) - - # Now set the transforms for training. We need to do this here so that we can (a) - # get the data dimensions to configure the network, and (b) save the - # parameter standardization dict in the PosteriorModel. In principle, (a) could - # be done without generating data (by careful calculation) and (b) could also - # be done outside the transform setup. But for now, this is convenient. The - # transforms will be reset later by initialize_stage(). - - set_train_transforms( - wfd, - train_settings["data"], - train_settings["training"]["stage_0"]["asd_dataset_path"], + wfd, initial_weights = _prepare_wfd_and_initial_weights( + train_settings=train_settings, + train_dir=train_dir, + local_settings=local_settings, ) # This modifies the model settings in-place. @@ -233,12 +211,14 @@ def initialize_stage( wfd: WaveformDataset, stage: dict, num_workers: int, + world_size: Optional[int] = None, + rank: Optional[int] = None, resume: bool = False, -): +) -> Tuple[DataLoader, DataLoader, Optional[DistributedSampler]]: """ Initializes training based on PosteriorModel metadata and current stage: * Builds transforms (based on noise settings for current stage); - * Builds DataLoaders; + * Builds DataLoaders (with DistributedSampler for DDP); * At the beginning of a stage (i.e., if not resuming mid-stage), initializes a new optimizer and scheduler; * Freezes / unfreezes SVD layer of embedding network @@ -250,32 +230,57 @@ def initialize_stage( stage : dict Settings specific to current stage of training num_workers : int + world_size : int, optional + Total number of DDP processes (GPUs). + rank : int, optional + Rank of this DDP process. resume : bool Whether training is resuming mid-stage. This controls whether the optimizer and scheduler should be re-initialized based on contents of stage dict. Returns ------- - (train_loader, test_loader) + (train_loader, test_loader, train_sampler) + *train_sampler* is ``None`` in single-GPU mode. """ train_settings = pm.metadata["train_settings"] + print_output = rank is None or rank == 0 # Rebuild transforms based on possibly different noise. - set_train_transforms(wfd, train_settings["data"], stage["asd_dataset_path"]) - - # Allows for changes in batch size between stages. - train_loader, test_loader = build_train_and_test_loaders( + set_train_transforms( wfd, - train_settings["data"]["train_fraction"], - stage["batch_size"], - num_workers, + train_settings["data"], + stage["asd_dataset_path"], + print_output=print_output, + ) + + # Convert total batch size to per-GPU batch size for DDP. + if world_size is not None and world_size > 1: + total_batch_size = stage["batch_size"] + if total_batch_size % world_size != 0: + raise ValueError( + f"Total batch size {total_batch_size} is not divisible by " + f"the number of GPUs {world_size}." + ) + batch_size_per_gpu = total_batch_size // world_size + else: + batch_size_per_gpu = stage["batch_size"] + + train_loader, test_loader, train_sampler = build_train_and_test_loaders( + dataset=wfd, + train_fraction=train_settings["data"]["train_fraction"], + batch_size=batch_size_per_gpu, + num_workers=num_workers, + world_size=world_size, + rank=rank, ) if not resume: # New optimizer and scheduler. If we are resuming, these should have been # loaded from the checkpoint. - print("Initializing new optimizer and scheduler.") + if print_output: + print("Initializing new optimizer and scheduler.") pm.optimizer_kwargs = stage["optimizer"] pm.scheduler_kwargs = stage["scheduler"] pm.initialize_optimizer_and_scheduler() @@ -283,6 +288,11 @@ def initialize_stage( # Freeze/unfreeze RB layer if necessary if "freeze_rb_layer" in stage: if stage["freeze_rb_layer"]: + if world_size is not None and world_size > 1: + raise ValueError( + "Cannot freeze the RB layer during DDP training — modify " + "the network before wrapping it with DDP." + ) set_requires_grad_flag( pm.network, name_contains="layers_rb", requires_grad=False ) @@ -290,16 +300,20 @@ def initialize_stage( set_requires_grad_flag( pm.network, name_contains="layers_rb", requires_grad=True ) - n_grad = get_number_of_model_parameters(pm.network, (True,)) - n_nograd = get_number_of_model_parameters(pm.network, (False,)) - print(f"Fixed parameters: {n_nograd}\nLearnable parameters: {n_grad}\n") - return train_loader, test_loader + if print_output: + print_number_of_model_parameters(pm.network) + + return train_loader, test_loader, train_sampler def train_stages( - pm: BasePosteriorModel, wfd: WaveformDataset, train_dir: str, local_settings: dict -) -> bool: + pm: BasePosteriorModel, + wfd: WaveformDataset, + train_dir: str, + local_settings: dict, + global_epoch: ctypes.c_int = Value(ctypes.c_int, 1), +) -> Tuple[bool, bool]: """ Train the network, iterating through the sequence of stages. Stages can change certain settings such as the noise characteristics, optimizer, and scheduler settings. @@ -311,19 +325,25 @@ def train_stages( train_dir : str Directory for saving checkpoints and train history. local_settings : dict + global_epoch : multiprocessing.Value + Shared epoch counter, updated each epoch and forwarded to ``pm.train``. Returns ------- - bool - True if all stages are complete - False otherwise + (training_finished, resume_training) : (bool, bool) """ train_settings = pm.metadata["train_settings"] runtime_limits = RuntimeLimits( - epoch_start=pm.epoch, **local_settings["runtime_limits"] + epoch_start=pm.epoch, + device=pm.device, + **local_settings["runtime_limits"], ) + rank = local_settings.get("rank", None) + world_size = local_settings.get("world_size", None) + print_primary = rank is None or rank == 0 + # Extract list of stages from settings dict stages = [] num_stages = 0 @@ -335,22 +355,38 @@ def train_stages( break end_epochs = list(np.cumsum([stage["epochs"] for stage in stages])) + resume_training = False num_starting_stage = np.searchsorted(end_epochs, pm.epoch + 1) for n in range(num_starting_stage, num_stages): stage = stages[n] if pm.epoch == end_epochs[n] - stage["epochs"]: - print(f"\nBeginning training stage {n}. Settings:") - print(yaml.dump(stage, default_flow_style=False, sort_keys=False)) - train_loader, test_loader = initialize_stage( - pm, wfd, stage, local_settings["num_workers"], resume=False + if print_primary: + print(f"\nBeginning training stage {n}. Settings:") + print(yaml.dump(stage, default_flow_style=False, sort_keys=False)) + train_loader, test_loader, train_sampler = initialize_stage( + pm, + wfd, + stage, + local_settings["num_workers"], + world_size=world_size, + rank=rank, + resume=False, ) else: - print(f"\nResuming training in stage {n}. Settings:") - print(yaml.dump(stage, default_flow_style=False, sort_keys=False)) - train_loader, test_loader = initialize_stage( - pm, wfd, stage, local_settings["num_workers"], resume=True + if print_primary: + print(f"\nResuming training in stage {n}. Settings:") + print(yaml.dump(stage, default_flow_style=False, sort_keys=False)) + train_loader, test_loader, train_sampler = initialize_stage( + pm, + wfd, + stage, + local_settings["num_workers"], + world_size=world_size, + rank=rank, + resume=True, ) + early_stopping = None if stage.get("early_stopping"): try: @@ -363,31 +399,40 @@ def train_stages( runtime_limits.max_epochs_total = end_epochs[n] pm.train( - train_loader, - test_loader, + train_loader=train_loader, + test_loader=test_loader, train_dir=train_dir, + train_sampler=train_sampler, runtime_limits=runtime_limits, checkpoint_epochs=local_settings["checkpoint_epochs"], use_wandb=local_settings.get("wandb", False), test_only=local_settings.get("test_only", False), early_stopping=early_stopping, + gradient_updates_per_optimizer_step=stage.get( + "gradient_updates_per_optimizer_step", 1 + ), + automatic_mixed_precision=stage.get("automatic_mixed_precision", False), + world_size=world_size if world_size is not None else 1, + global_epoch=global_epoch, ) + # if test_only, model should not be saved, and run is complete if local_settings.get("test_only", False): - return True + return True, False - if pm.epoch == end_epochs[n]: + if pm.epoch == end_epochs[n] and print_primary: save_file = os.path.join(train_dir, f"model_stage_{n}.pt") print(f"Training stage complete. Saving to {save_file}.") pm.save_model(save_file, save_training_info=True) + if runtime_limits.local_limits_exceeded(pm.epoch): - print("Local runtime limits reached. Ending program.") + if print_primary: + print("Local runtime limits reached. Ending program.") + resume_training = True break - if pm.epoch == end_epochs[-1]: - return True - else: - return False + training_finished = pm.epoch == end_epochs[-1] + return training_finished, resume_training def parse_args(): @@ -434,20 +479,393 @@ def parse_args(): return args +def get_num_gpus(local_settings: dict) -> int: + """ + Return the number of GPUs to use for training. + + The canonical location is ``local_settings["num_gpus"]`` (default: 1). + + For backward compatibility, ``local_settings["condor"]["num_gpus"]`` is + accepted as a fallback when ``local_settings["num_gpus"]`` is absent, but + its use is deprecated and will be removed in a future release. + """ + if "num_gpus" in local_settings: + return int(local_settings["num_gpus"]) + condor_num_gpus = local_settings.get("condor", {}).get("num_gpus") + if condor_num_gpus is not None: + import warnings + + warnings.warn( + "Specifying 'num_gpus' under 'local.condor' is deprecated. " + "Move it to 'local.num_gpus' instead.", + DeprecationWarning, + stacklevel=2, + ) + return int(condor_num_gpus) + return 1 + + +def run_training( + train_settings: Optional[dict], + local_settings: dict, + train_dir: str, + ckpt_file: Optional[str], + resume: bool, +) -> Tuple[bool, bool, int]: + """ + Single-GPU training entry point. + + Parameters + ---------- + train_settings : dict or None + Full training settings (from YAML). ``None`` when resuming. + local_settings : dict + train_dir : str + ckpt_file : str or None + Checkpoint path used when *resume* is True. + resume : bool + + Returns + ------- + (complete, resume, epoch) : (bool, bool, int) + """ + if not resume: + pm, wfd = prepare_training_new(train_settings, train_dir, local_settings) + else: + pm, wfd = prepare_training_resume(ckpt_file, local_settings, train_dir) + + global_epoch = Value(ctypes.c_int, pm.epoch) + wfd.epoch = global_epoch + + with threadpool_limits(limits=1, user_api="blas"): + complete, resume_flag = train_stages( + pm=pm, + wfd=wfd, + train_dir=train_dir, + local_settings=local_settings, + global_epoch=global_epoch, + ) + + return complete, resume_flag, pm.epoch + + +def run_training_ddp( + rank: int, + world_size: int, + train_settings: Optional[dict], + local_settings: dict, + train_dir: str, + wfd: WaveformDataset, + initial_weights: Optional[dict], + ckpt_file: Optional[str], + resume: bool, + result_queue: mp.Queue, + global_epoch: ctypes.c_int, +) -> None: + """ + Worker function executed by each DDP process. + + Each GPU runs this function in its own process, identified by *rank*. + Results (complete flag, resume flag, epoch) are put into *result_queue* + so the parent process can collect them. + + Parameters + ---------- + rank : int + world_size : int + train_settings : dict or None + local_settings : dict + train_dir : str + wfd : WaveformDataset + Dataset object pre-built in the parent process and shared with workers. + initial_weights : dict or None + SVD-based initial weights for the embedding network. + ckpt_file : str or None + resume : bool + result_queue : mp.Queue + global_epoch : multiprocessing.Value + """ + try: + setup_ddp(rank, world_size) + set_seed_based_on_rank(rank) + + if rank == 0: + document_gpus(train_dir) + + local_settings = dict(local_settings) # avoid mutating the original + local_settings["device"] = f"cuda:{rank}" + local_settings["rank"] = rank + local_settings["world_size"] = world_size + + if not resume: + full_settings = { + "dataset_settings": wfd.settings, + "train_settings": train_settings, + } + print_output = rank == 0 + if print_output: + print("\nInitializing new posterior model.") + print("Complete settings:") + print( + yaml.dump(full_settings, default_flow_style=False, sort_keys=False) + ) + + pm = build_model_from_kwargs( + settings=full_settings, + initial_weights=initial_weights, + device=local_settings["device"], + ) + + if rank == 0 and local_settings.get("wandb", False): + try: + import wandb + + wandb.init( + config=full_settings, + dir=train_dir, + **local_settings["wandb"], + ) + except ImportError: + print("WandB is enabled but not installed.") + else: + pm = build_model_from_kwargs( + filename=ckpt_file, device=local_settings["device"] + ) + if rank == 0 and local_settings.get("wandb", False): + try: + import wandb + + wandb.init( + resume="must", + dir=train_dir, + **local_settings["wandb"], + ) + except ImportError: + print("WandB is enabled but not installed.") + + pm.network = replace_BatchNorm_with_SyncBatchNorm(pm.network) + pm.network = DDP(pm.network, device_ids=[rank]) + + global_epoch.value = pm.epoch + wfd.epoch = global_epoch + + with threadpool_limits(limits=1, user_api="blas"): + complete, resume_flag = train_stages( + pm=pm, + wfd=wfd, + train_dir=train_dir, + local_settings=local_settings, + global_epoch=global_epoch, + ) + + if complete and local_settings.get("wandb", False) and rank == 0: + try: + import wandb + + wandb.finish() + except ImportError: + pass + + cleanup_ddp() + result_queue.put((rank, complete, resume_flag, pm.epoch, None)) + + except Exception as exc: + import traceback + + result_queue.put((rank, False, False, 0, traceback.format_exc())) + raise + + +def run_multi_gpu_training( + world_size: int, + train_settings: Optional[dict], + local_settings: dict, + train_dir: str, + ckpt_file: Optional[str], + resume: bool, +) -> Tuple[bool, bool, int]: + """ + Multi-GPU DDP training entry point. + + Spawns one process per GPU, each running :func:`run_training_ddp`. + + Parameters + ---------- + world_size : int + Number of GPUs. + train_settings : dict or None + local_settings : dict + train_dir : str + ckpt_file : str or None + resume : bool + + Returns + ------- + (complete, resume, epoch) : (bool, bool, int) + """ + initial_weights = None + + if not resume: + # Build the dataset and compute SVD-based initialisation *once* in the + # parent process so it can be shared across all worker processes. + wfd, initial_weights = _prepare_wfd_and_initial_weights( + train_settings=train_settings, + train_dir=train_dir, + local_settings=local_settings, + ) + autocomplete_model_kwargs( + model_kwargs=train_settings["model"], + data_sample=wfd[0], + ) + else: + d = torch.load(ckpt_file, map_location="cpu") + train_settings = d["metadata"]["train_settings"] + data_settings = deepcopy(train_settings["data"]) + data_settings["waveform_dataset_path"] = copy_files_to_local( + file_path=data_settings["waveform_dataset_path"], + local_dir=local_settings.get("local_cache_path", None), + leave_keys_on_disk=local_settings.get("leave_waveforms_on_disk", True), + is_condor="condor" in local_settings, + ) + wfd = build_dataset( + data_settings=data_settings, + leave_waveforms_on_disk=local_settings.get("leave_waveforms_on_disk", True), + ) + + global_epoch = Value(ctypes.c_int, 0) + result_queue = mp.Queue() + processes = [] + + for rank in range(world_size): + p = mp.Process( + target=run_training_ddp, + args=( + rank, + world_size, + train_settings, + local_settings, + train_dir, + wfd, + initial_weights, + ckpt_file, + resume, + result_queue, + global_epoch, + ), + ) + p.start() + processes.append(p) + + error_occurred = False + for p in processes: + p.join() + if p.exitcode != 0: + error_occurred = True + print(f"Process {p.pid} exited with code {p.exitcode}.") + for proc in processes: + if proc.is_alive(): + proc.terminate() + + results = [] + try: + while True: + item = result_queue.get_nowait() + results.append(item) + if item[-1] is not None: + print(f"Rank {item[0]} failed:\n{item[-1]}") + else: + print(f"Rank {item[0]} finished successfully.") + except queue.Empty: + pass + + if error_occurred or len(results) != world_size: + raise RuntimeError( + f"DDP training failed: {len(results)}/{world_size} processes reported results." + ) + + complete_flags = [r[1] for r in results] + resume_flags = [r[2] for r in results] + epochs = [r[3] for r in results] + + assert ( + len(set(complete_flags)) == 1 + ), f"Inconsistent 'complete' flags: {complete_flags}" + assert len(set(resume_flags)) == 1, f"Inconsistent 'resume' flags: {resume_flags}" + assert len(set(epochs)) == 1, f"Inconsistent epoch counts: {epochs}" + + return complete_flags[0], resume_flags[0], epochs[0] + + +def _prepare_wfd_and_initial_weights( + train_settings: dict, + train_dir: str, + local_settings: dict, +) -> Tuple[WaveformDataset, Optional[dict]]: + """ + Build the WaveformDataset and, if applicable, compute SVD-based initial + weights for the embedding network. + + This is called once in the parent process before spawning DDP workers so + that the expensive SVD computation and dataset loading happen only once. + + Returns + ------- + (wfd, initial_weights) + """ + data_settings = deepcopy(train_settings["data"]) + data_settings["waveform_dataset_path"] = copy_files_to_local( + file_path=data_settings["waveform_dataset_path"], + local_dir=local_settings.get("local_cache_path", None), + leave_keys_on_disk=local_settings.get("leave_waveforms_on_disk", True), + is_condor="condor" in local_settings, + ) + wfd = build_dataset( + data_settings=data_settings, + leave_waveforms_on_disk=local_settings.get("leave_waveforms_on_disk", True), + ) + + initial_weights = {} + model_kwargs = train_settings["model"] + + if ( + model_kwargs.get("embedding_kwargs") + and "svd" in model_kwargs["embedding_kwargs"] + and not model_kwargs["embedding_kwargs"]["svd"].get("no_init", False) + ): + batch_size = train_settings["training"]["stage_0"]["batch_size"] + print("\nBuilding SVD for initialization of embedding network.") + initial_weights["V_rb_list"] = build_svd_for_embedding_network( + wfd=wfd, + data_settings=train_settings["data"], + asd_dataset_path=train_settings["training"]["stage_0"]["asd_dataset_path"], + batch_size=batch_size, + out_dir=train_dir, + **model_kwargs["embedding_kwargs"]["svd"], + ) + else: + initial_weights = None + + set_train_transforms( + wfd, + train_settings["data"], + train_settings["training"]["stage_0"]["asd_dataset_path"], + ) + + return wfd, initial_weights + + def train_local(): args = parse_args() os.makedirs(args.train_dir, exist_ok=True) - if args.settings_file is not None: + resume = args.checkpoint is not None + + if not resume: print("Beginning new training run.") with open(args.settings_file, "r") as fp: train_settings = yaml.safe_load(fp) - # Extract the local settings from train settings file, save it separately. This - # file can later be modified, and the settings take effect immediately upon - # resuming. - local_settings = train_settings.pop("local") with open(os.path.join(args.train_dir, "local_settings.yaml"), "w") as f: if ( @@ -461,19 +879,32 @@ def train_local(): except ImportError: print("wandb not installed, cannot generate run id.") yaml.dump(local_settings, f, default_flow_style=False, sort_keys=False) - - pm, wfd = prepare_training_new(train_settings, args.train_dir, local_settings) - else: print("Resuming training run.") + train_settings = None with open(os.path.join(args.train_dir, "local_settings.yaml"), "r") as f: local_settings = yaml.safe_load(f) - pm, wfd = prepare_training_resume( - args.checkpoint, local_settings, args.train_dir - ) - with threadpool_limits(limits=1, user_api="blas"): - complete = train_stages(pm, wfd, args.train_dir, local_settings) + num_gpus = get_num_gpus(local_settings) + if local_settings.get("device") == "cuda" and num_gpus > 1: + complete, resume_flag, _ = run_multi_gpu_training( + world_size=num_gpus, + train_settings=train_settings, + local_settings=local_settings, + train_dir=args.train_dir, + ckpt_file=args.checkpoint, + resume=resume, + ) + else: + if local_settings.get("device") == "cuda": + document_gpus(args.train_dir) + complete, resume_flag, _ = run_training( + train_settings=train_settings, + local_settings=local_settings, + train_dir=args.train_dir, + ckpt_file=args.checkpoint, + resume=resume, + ) if complete: if args.exit_command: diff --git a/dingo/gw/training/train_pipeline_condor.py b/dingo/gw/training/train_pipeline_condor.py index 0ea970f26..7141fae3e 100644 --- a/dingo/gw/training/train_pipeline_condor.py +++ b/dingo/gw/training/train_pipeline_condor.py @@ -1,13 +1,15 @@ +import argparse import os import sys -from os.path import join, isfile +from os.path import isfile, join + import yaml -import argparse -from dingo.gw.training import ( - prepare_training_new, - prepare_training_resume, - train_stages, +from dingo.core.utils.torchutils import document_gpus +from dingo.gw.training.train_pipeline import ( + get_num_gpus, + run_multi_gpu_training, + run_training, ) @@ -30,16 +32,39 @@ def create_submission_file( # getenv required for GPU training because wandb needs $HOME to be defined lines.append(f"getenv = True\n") lines.append(f'executable = {condor_settings["executable"]}\n') - lines.append(f'request_cpus = {condor_settings["num_cpus"]}\n') - lines.append(f'request_memory = {condor_settings["memory_cpus"]}\n') - lines.append(f'request_gpus = {condor_settings["num_gpus"]}\n') - lines.append( - f"requirements = TARGET.CUDAGlobalMemoryMb > " - f'{condor_settings["memory_gpus"]}\n\n' - ) + if "request_disk" in condor_settings: lines.append(f'request_disk = {condor_settings["request_disk"]}\n') - lines.append(f'arguments = "{condor_settings["arguments"]}"\n') + if "num_cpus" in condor_settings: + lines.append(f'request_cpus = {condor_settings["num_cpus"]}\n') + if "memory_cpus" in condor_settings: + lines.append(f'request_memory = {condor_settings["memory_cpus"]}\n') + + requirements = condor_settings.get("requirements", "") + if "memory_gpus" in condor_settings: + gpu_req = f"TARGET.CUDAGlobalMemoryMb > {condor_settings['memory_gpus']}" + if requirements: + if ")" not in requirements: + requirements = f"({requirements})" + requirements = f"{requirements} && ({gpu_req})" + else: + requirements = gpu_req + if requirements: + lines.append(f"requirements = {requirements}\n") + + if "num_gpus" in condor_settings: + num_gpus = condor_settings["num_gpus"] + lines.append(f"request_gpus = {num_gpus}\n") + # Cluster-specific full-node templates (MPI-IS). + if num_gpus == 8: + lines.append("use template : FullNode\n") + elif num_gpus >= 6: + lines.append(f"use template : FullNode({num_gpus})\n") + + if "arguments" in condor_settings: + lines.append(f'arguments = "{condor_settings["arguments"]}"\n') + + lines.append("\n") lines.append(f'error = {join(train_dir, "info.err")}\n') lines.append(f'output = {join(train_dir, "info.out")}\n') lines.append(f'log = {join(train_dir, "info.log")}\n') @@ -79,20 +104,16 @@ def train_condor(): ) args = parser.parse_args() - # For condor settings, first try looking for a local settings file. Otherwise, - # defer to train_settings.yaml. - # if isfile(join(args.train_dir, 'local_settings.yaml')): - # with open(join(args.train_dir, 'local_settings.yaml')) as f: - # condor_settings = yaml.safe_load(f)['condor'] - # else: - if not args.start_submission: # # TRAIN # - if not isfile(join(args.train_dir, args.checkpoint)): + ckpt_path = join(args.train_dir, args.checkpoint) + resume = isfile(ckpt_path) + + if not resume: print("Beginning new training run.") with open(join(args.train_dir, "train_settings.yaml"), "r") as f: train_settings = yaml.safe_load(f) @@ -100,7 +121,6 @@ def train_condor(): # Extract the local settings from train settings file, save it separately. # This file can later be modified, and the settings take effect immediately # upon resuming. - local_settings = train_settings.pop("local") with open(os.path.join(args.train_dir, "local_settings.yaml"), "w") as f: if ( @@ -114,25 +134,35 @@ def train_condor(): except ImportError: print("wandb not installed, cannot generate run id.") yaml.dump(local_settings, f, default_flow_style=False, sort_keys=False) - - pm, wfd = prepare_training_new( - train_settings, args.train_dir, local_settings - ) - else: print("Resuming training run.") + train_settings = None with open(os.path.join(args.train_dir, "local_settings.yaml"), "r") as f: local_settings = yaml.safe_load(f) - pm, wfd = prepare_training_resume( - join(args.train_dir, args.checkpoint), - local_settings, + + num_gpus = get_num_gpus(local_settings) + if local_settings.get("device") == "cuda" and num_gpus > 1: + complete, _, pm_epoch = run_multi_gpu_training( + world_size=num_gpus, + train_settings=train_settings, + local_settings=local_settings, train_dir=args.train_dir, + ckpt_file=ckpt_path, + resume=resume, + ) + else: + if local_settings.get("device") == "cuda": + document_gpus(args.train_dir) + complete, _, pm_epoch = run_training( + train_settings=train_settings, + local_settings=local_settings, + train_dir=args.train_dir, + ckpt_file=ckpt_path, + resume=resume, ) - - complete = train_stages(pm, wfd, args.train_dir, local_settings) print("Copying log files") - copy_logfiles(args.train_dir, epoch=pm.epoch) + copy_logfiles(args.train_dir, epoch=pm_epoch) # # PREPARE NEXT SUBMISSION @@ -140,7 +170,8 @@ def train_condor(): if complete: print( - f"Training complete, job will not be resubmitted. Executing exit command: {args.exit_command}." + f"Training complete, job will not be resubmitted. " + f"Executing exit command: {args.exit_command}." ) if args.exit_command: os.system(args.exit_command) @@ -164,7 +195,11 @@ def train_condor(): submission_file = "submission_file.sub" with open(join(args.train_dir, "train_settings.yaml"), "r") as f: - condor_settings = yaml.safe_load(f)["local"]["condor"] + local_settings_from_file = yaml.safe_load(f)["local"] + condor_settings = local_settings_from_file.get("condor", {}) + # Inject num_gpus from local.num_gpus so the submission file requests the + # correct number of GPUs without the user having to duplicate it under condor:. + condor_settings["num_gpus"] = get_num_gpus(local_settings_from_file) condor_settings["arguments"] = condor_arguments condor_settings["executable"] = join( os.path.dirname(sys.executable), "dingo_train_condor" diff --git a/docs/source/training.md b/docs/source/training.md index e48413fee..4b0bb792d 100644 --- a/docs/source/training.md +++ b/docs/source/training.md @@ -75,7 +75,9 @@ training: scheduler: type: cosine T_max: 300 - batch_size: 64 + batch_size: 512 # Total effective batch size across all GPUs. +# gradient_updates_per_optimizer_step: 1 +# automatic_mixed_precision: False stage_1: epochs: 150 @@ -87,15 +89,19 @@ training: scheduler: type: cosine T_max: 150 - batch_size: 64 + batch_size: 512 +# gradient_updates_per_optimizer_step: 1 +# automatic_mixed_precision: False # Local settings that have no impact on the final trained network. local: - device: cpu # Change this to 'cuda' for training on a GPU. + device: cuda # Change this to 'cpu' for training without a GPU. num_workers: 6 +# num_gpus: 1 # Set to >1 to enable multi-GPU (DDP) training. When using + # dingo_train_condor, request_gpus is set automatically. # wandb: # project: dingo -# group: O4 +# group: my_project runtime_limits: max_time_per_run: 36000 max_epochs_per_run: 500 @@ -103,10 +109,8 @@ local: leave_waveforms_on_disk: True local_cache_path: tmp # condor: -# bid: 100 # num_cpus: 16 # memory_cpus: 128000 -# num_gpus: 1 # memory_gpus: 8000 # request_disk: 50GB ``` @@ -132,7 +136,7 @@ asd_dataset_path : Points to an `ASDDataset` file. Each stage can have its own ASD dataset, which is useful for implementing a pre-training stage with fixed ASD and a fine-tuning stage with variable ASD. freeze_rb_layer -: Whether to freeze the first layer of the embedding network in `nsf+embedding` models. This layer is seeded with reduced (SVD) basis vectors, so freezing this layer during pre-training simply projects data onto the basis coefficients. In the fine-tuning stage, when other weights are more stable, unfreezing this can be useful. +: Whether to freeze the first layer of the embedding network. This layer is seeded with reduced (SVD) basis vectors, so freezing this layer during pre-training simply projects data onto the basis coefficients. In the fine-tuning stage, when other weights are more stable, unfreezing this can be useful. optimizer : Specify [optimizer](https://pytorch.org/docs/stable/optim.html) type and parameters such as initial learning rate. @@ -141,7 +145,13 @@ scheduler : Use a [learning rate scheduler](https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate) to reduce the learning rate over time. This can improve overall optimization. batch_size -: Number of training samples per mini-batch. For a training dataset of size $N$, then each epoch will consist of $N / \text{batch_size}$ batches. Generally training will be faster for a larger batch size, but will require additional iterations. +: Total number of training samples per optimizer step. For a training dataset of size $N$, each epoch consists of $N /$ `batch_size` optimizer steps. When using multiple GPUs, this is the *effective* batch size across all GPUs; each GPU processes `batch_size` / `num_gpus` samples per step. + +gradient_updates_per_optimizer_step +: (Optional, default 1) Number of forward–backward passes to accumulate before calling the optimizer. Setting this to $k$ simulates an effective batch size of $k \times$ `batch_size` without increasing GPU memory usage. This is useful when the desired batch size does not fit in GPU memory. + +automatic_mixed_precision +: (Optional, default `False`) Enable automatic mixed precision (AMP) training. With AMP, the forward pass runs in FP16, while optimizer state and parameter updates remain in FP32. This can roughly halve GPU memory usage and increase throughput on GPUs with Tensor Core hardware (e.g. NVIDIA A100, V100). Requires PyTorch >= 2.0 and a CUDA device. ```{important} The stage-training framework allows for separate pre-training and fine-tuning stages. We found that having a pre-training stage where we freeze certain network weights and fix the noise ASD improves overall training results. @@ -157,6 +167,9 @@ device num_workers : Number of CPU worker processes to use for pre-processing training data before copying to the GPU. Data pre-processing (inluding decompression, projection to detectors, and noise generation) is quite expensive, so using 16 or 32 processes is recommended, otherwise this can become a bottleneck. We recommend monitoring the GPU utilization percentage as well as time spent on pre-processing (output during training) to fine-tune this number. +num_gpus +: (Optional, default 1) Number of GPUs to use for training. Setting this to more than 1 enables data-parallel multi-GPU training (PyTorch DDP). The `batch_size` specified in each training stage is the total effective batch size; it is divided equally across GPUs. When using `dingo_train_condor`, the HTCondor `request_gpus` directive is set automatically from this value. See [](#multi-gpu-training) for details. + wandb : Settings for [Weights & Biases](https://wandb.ai/site). If you have an account, you can use this to track your training progress and compare different runs. @@ -173,7 +186,104 @@ local_cache_path : When training on a cluster and loading waveforms during training (i.e., `leave_waveforms_on_disk=True`), the waveform dataset should be copied to the disk storage of the local node at the beginning of training. This prevents unexpected long data loading times during training due to network traffic. Usually, paths for local storage are `tmp` or `dev/shm`. When submitting the job with `condor`, `request_disk: 50GB` should be included in the `condor` settings with the requested disk space larger than the size of the waveform dataset used for training. condor -: Settings for [HTCondor](https://htcondor.readthedocs.io/en/latest/index.html). The condor script will (re)submit itself according to these options. +: Settings for [HTCondor](https://htcondor.readthedocs.io/en/latest/index.html). The condor script will (re)submit itself according to these options. Available keys are `bid`, `num_cpus`, `memory_cpus`, `memory_gpus`, `request_disk`, and `requirements`. The number of requested GPUs (`request_gpus`) is derived automatically from `num_gpus` above and does not need to be specified here. + +## Multi-GPU training + +Dingo supports data-parallel training across multiple GPUs using +[PyTorch DDP](https://pytorch.org/docs/stable/notes/ddp.html) (DistributedDataParallel). +Each GPU processes a different shard of the mini-batch simultaneously, effectively +scaling throughput with the number of GPUs. + +For a detailed practical guide — including how to scale `batch_size` and learning +rate for actual speedup, how to interpret the training log, and troubleshooting +tips — see [](training_multi_gpu.md). + +### Enabling multi-GPU training + +Set `num_gpus` in the `local` section of your settings file: + +```yaml +local: + device: cuda + num_gpus: 4 + num_workers: 8 # Recommended: num_workers per GPU, e.g. 2–8 + ... +``` + +The `batch_size` in each training stage is always the **total effective batch size** +across all GPUs. Dingo divides it equally, so `batch_size` must be divisible by +`num_gpus`. Each GPU therefore processes `batch_size / num_gpus` samples per step, +which keeps the gradient statistics and learning dynamics identical to single-GPU +training with the same `batch_size`. + +```{important} +`batch_size` is the total effective batch size across all GPUs, not the per-GPU +batch size. Increasing `num_gpus` does not change the effective batch size or +require adjusting any other hyperparameters. +``` + +```note +`freeze_rb_layer = True` is currently not allowed since additional changes +would be required to exclude the frozen RB layer from the DDP model wrapper. +``` + +### Gradient accumulation + +For very large effective batch sizes that do not fit in GPU memory, use gradient +accumulation alongside multi-GPU training: + +```yaml +training: + stage_0: + batch_size: 4096 + gradient_updates_per_optimizer_step: 4 # effective batch = 4096 * 4 = 16384 +``` + +Each optimizer step accumulates gradients over `gradient_updates_per_optimizer_step` +forward–backward passes. Multi-GPU and gradient accumulation can be combined freely. + +### Automatic mixed precision + +Enable AMP to reduce GPU memory usage and increase throughput on modern GPUs: + +```yaml +training: + stage_0: + batch_size: 4096 + automatic_mixed_precision: True +``` + +AMP runs the forward pass in FP16 and the optimizer step in FP32, typically halving +memory usage with negligible impact on model quality. It requires a CUDA device and +PyTorch >= 2.0. + +### Using multiple GPUs with `dingo_train_condor` + +When submitting via HTCondor, the `request_gpus` directive is set automatically from +`local.num_gpus`. No changes to the `condor:` block are needed: + +```yaml +local: + device: cuda + num_gpus: 4 + condor: + num_cpus: 64 + memory_cpus: 256000 + memory_gpus: 24000 # Memory per GPU in MB — used for node selection only. + request_disk: 50GB +``` + +```{note} +On clusters with full-node GPU allocations (e.g., MPI-IS), requesting 6 or more +GPUs automatically applies the appropriate FullNode HTCondor template. +``` + +### Requirements + +- PyTorch >= 2.0 (required for `torch.amp`) +- An NCCL-capable build of PyTorch (standard for CUDA-enabled installations) +- One CUDA GPU per process ## Command-line scripts diff --git a/docs/source/training_multi_gpu.md b/docs/source/training_multi_gpu.md new file mode 100644 index 000000000..5a31deb23 --- /dev/null +++ b/docs/source/training_multi_gpu.md @@ -0,0 +1,196 @@ +# Multi-GPU Training + +Dingo supports data-parallel training across multiple GPUs using +[PyTorch DDP](https://pytorch.org/docs/stable/notes/ddp.html) +(DistributedDataParallel). This tutorial explains how DDP works in Dingo, why +simply setting `num_gpus: 8` is not enough to achieve a speedup, how to read the +extended training log, and which hyperparameters require tuning. + +## How DDP works in Dingo + +The goal of DDP is to increase the data throughput during training by processing +subsets of the mini-batch in parallel on separate GPUs. +Each GPU holds a full copy of the DINGO model, but each model sees a different, random subset of the waveform data set. +After every backward pass, DDP performs a gradient all-reduce: it +aggregates gradients across all GPUs resulting in a gradient update based on the full effective batch size, +not the per-GPU batch size. +The resulting averaged gradient is used to update all models equivalently such that every model replica stays the same. +The optimizer step is identical to single-GPU training with the effective batch size. + +Through DDP, it is possible to increase the effective batch size beyond the single-GPU memory limit. +Now, the limiting factor for training throughput is the memory limit of a single GPU which processes the effective +batch size divided by the number of GPUs. Out-of-memory errors can be addressed by either increasing the number of GPUs +or by reducing the effective batch size. + +There are certain caveats to DDP: +* Since the gradients have to be synced across GPUs, the interconnect between GPUs becomes highly relevant. + Slow interconnects can significantly slow down training. Therefore, it is recommended to only run DDP on GPUs + that are located on the same node where fast interconnect is independent of network traffic. +* In addition to the full model replica (with its gradients and optimizer state) and the batch subset, + additional GPU memory is required for syncing and aggregating gradients. + Therefore, the per-GPU batch size can be lower than in single-GPU training. + +## Compute requirements + +- A single node with N GPUs connected by fast interconnect (tested on A100 and H100 nodes). +- An NCCL-capable PyTorch build (the default for any CUDA-enabled installation). +- When submitting via HTCondor, `request_gpus` is derived automatically from + `num_gpus` — no changes to the `condor:` block are needed. + +## Settings changes + +Multi-GPU training is enabled within Dingo by changing `num_gpus: 1` to the number of available GPUs. +Additionally, it is recommended to scale `batch_size` and `lr` in each training stage (see explanations below): + +```yaml +# single-GPU baseline +local: + device: cuda + num_workers: 16 + num_gpus: 1 +training: + stage_0: + batch_size: 4096 + optimizer: + lr: 5.0e-5 + stage_1: + batch_size: 4096 + optimizer: + lr: 1.0e-5 +``` + +```yaml +# 8-GPU DDP — scale batch_size and lr by num_gpus +local: + device: cuda + num_workers: 32 # scale with num_gpus, e.g. 4–8 per GPU + num_gpus: 8 +training: + stage_0: + batch_size: 32768 # = 4096 × 8 + optimizer: + lr: 4.0e-4 # = 5e-5 × 8 + stage_1: + batch_size: 32768 + optimizer: + lr: 8.0e-5 # = 1e-5 × 8 +``` + +### Increasing the `batch_size` + +Dingo interprets `batch_size` as the **total effective batch size across all GPUs** and divides it equally: +``` +per-GPU batch size = batch_size / num_gpus +``` + +Pytorch's `DistributedSampler` also splits the waveform dataset equally, so each GPU sees `1/N` of +the data per epoch. Putting both together, the number of optimizer steps per epoch is: +``` +steps / epoch = dataset_size × train_fraction / batch_size +``` + +This formula is **independent of `num_gpus`**. The table below shows the +consequences: + +| `num_gpus` | `batch_size` | per-GPU batch size | steps/epoch | outcome | +|:---:|:---:|:------------------:|:---:|:-----------------------------------------------------------------| +| 1 | 4 096 | 4 096 | N | baseline | +| 8 | 4 096 | 512 | N | same steps + all-reduce overhead → **slower** than the baseline | +| 8 | 32 768 | 4 096 | N/8 | fewer steps, good GPU utilisation → **faster** than the baseline | + +When the `batch_size` is scaled by `num_gpus`, the actual wall-clock speedup is somewhat less than 8× because the +gradient all-reduce adds overhead per step (visible as inflated *Time Network* in the log); a factor of 4–6× is typical. + +In general, it is recommended to increase the batch size as much as possible to fully utilize the GPU memory. + +### Learning rate scaling + +When the effective batch size increases by a factor of N, the learning rate should be increased +by the same factor (linear scaling rule, {cite:p}`Goyal:2017`): + +``` +lr_multi = lr_single × num_gpus +``` + +**Intuition**: with N× more samples per gradient update the gradient estimate has +lower variance, so a proportionally larger step can be taken without destabilising +training. + +Every training stage has to be adapted independently. + +### Freezing layers + +It is currently not possible to set `freeze_rb_layer: True` in DDP. The reason is that when starting the separate +DDP processes, it is fixed which network parameters have to be synced across GPUs and GPU memory is allocated +accordingly. In the current implementation, the stages are initialized (and therefore `freeze_rb_layer: True` is set) +afterward. If the layer is frozen within a DDP process, the all-reduce fails since the reduced tensors have a different +shape than expected. +Allowing `freeze_rb_layer: True` with DDP would require significant restructuring of the code. + +## Reading the training log + +Within Dingo, the separate processes for DDP are spawned after loading the waveform dataset, building the Dingo model, +and computing the SVD which initializes the first layer of the embedding network. +At this point, a statement like +`Process group initialised: backend=nccl, rank=0, world_size=8` is printed for every GPU +into `info.out`. +Afterward, every `print` statement within the Dingo code is executed by every GPU process (if this is not specifically +prevented). As a result, duplicate information can appear in `info.out`. Since each GPU process prints outputs to +`info.out` in parallel, the order of these statements can be mixed up compared to single-GPU training. +If information is not duplicated, it is only printed by `rank = 0`. + +During training, each printed line looks like the following in single-GPU mode: +``` +Train Epoch: 1 [4096/4750000 (0%)] Loss: -2.771 (-2.799) Time Dataloader: 0.043 (0.043) Time Network: 0.415 (0.415) +``` + +In multi-GPU mode, rank 0 additionally prints a third timing column: +``` +Train Epoch: 1 [4096/4750000 (0%)] Loss: -2.771 (-2.799) Time Dataloader: 0.018 (0.018) Time Network: 0.700 (0.700) Time Loss Aggregation: 0.090 (0.090) +``` + +Each value is the current step; the value in parentheses is the running average +over the epoch. **In DDP mode all three timings are reported as the maximum across +all GPUs** (`dist.ReduceOp.MAX`), i.e. they reflect the slowest rank. + +### Time Dataloader +Wall-clock time spent loading and pre-processing one batch (waveform +decompression, projection to detectors, noise generation). Expect this to be +smaller than in single-GPU mode because each GPU loads a smaller per-GPU batch. + +### Time Network +Forward pass + backward pass + **DDP gradient all-reduce** + optimizer step. +Unlike single-GPU training this includes gradient synchronisation, so *Time +Network* per step will typically be *larger* than in single-GPU mode even though +each GPU processes fewer samples. The speedup comes from fewer steps per epoch, +not from faster individual steps. + +### Time Loss Aggregation *(DDP only)* +Time required to reduce per-GPU losses to rank 0 and compute the global average. +Includes a `dist.barrier()` that waits for the slowest GPU to finish its network +pass before proceeding. Occasional spikes can indicate increased wait time at the barrier. +Sustained high values could indicate a load imbalance between GPUs. + +### Guidelines for tuning `num_workers` + +- *Time Dataloader* should be clearly smaller than *Time Network*. If dataloader + time dominates, increase `num_workers` (aim for 4–8 per GPU). +- *Time Network* being larger per step than single-GPU is expected; this is the + all-reduce cost. +- *Time Loss Aggregation* should be small relative to *Time Network*. + +## Monitoring with WandB + +Only rank 0 writes to WandB, so the run appears in the dashboard exactly as a +single-GPU run. The loss logged is the globally reduced value (weighted average +over all GPU contributions), making curves directly comparable across single- and +multi-GPU runs. + +Enable WandB in the `local` section: + +```yaml +local: + wandb: + project: dingo + group: O4 # optional: group related runs for comparison +``` diff --git a/examples/gnpe_model/train_settings.yaml b/examples/gnpe_model/train_settings.yaml index 747352d69..4904c5a72 100644 --- a/examples/gnpe_model/train_settings.yaml +++ b/examples/gnpe_model/train_settings.yaml @@ -106,11 +106,13 @@ local: max_epochs_per_run: 500 checkpoint_epochs: 10 leave_waveforms_on_disk: True - # local_cache_path: tmp # uncomment to avoid slow data loading during training due to network traffic on cluster - # Local settings related to condor, remove if not used on cluster + # num_gpus: 1 # Set to >1 to enable multi-GPU (DDP) training. batch_size above is the + # total effective batch size; it is divided equally across GPUs. + # When using Condor, request_gpus in the submission file is set automatically. + # local_cache_path: /tmp # uncomment to avoid slow data loading during training due to network traffic on cluster + # Settings for HTCondor job submission (remove if not running on a Condor cluster). condor: - num_cpus: 16 + num_cpus: 32 memory_cpus: 128000 - num_gpus: 1 memory_gpus: 8000 request_disk: 50GB \ No newline at end of file diff --git a/examples/gnpe_model/train_settings_init.yaml b/examples/gnpe_model/train_settings_init.yaml index 7bdf87ee5..f9a6ab096 100644 --- a/examples/gnpe_model/train_settings_init.yaml +++ b/examples/gnpe_model/train_settings_init.yaml @@ -77,7 +77,7 @@ training: # Local settings for training that have no impact on the final trained network. local: device: cpu # Change this to 'cuda' for training on a GPU. - num_workers: 6 # num_workers >0 does not work on Mac, see https://stackoverflow.com/questions/64772335/pytorch-w-parallelnative-cpp206 + num_workers: 16 # num_workers >0 does not work on Mac, see https://stackoverflow.com/questions/64772335/pytorch-w-parallelnative-cpp206 use_wandb: True wandb_api_key: # insert_here for monitoring many runs. Can be obtained from https://wandb.ai/settings runtime_limits: @@ -85,10 +85,13 @@ local: max_epochs_per_run: 500 checkpoint_epochs: 10 leave_waveforms_on_disk: True - # local_cache_path: tmp # uncomment to avoid slow data loading during training due to network traffic on cluster + # num_gpus: 1 # Set to >1 to enable multi-GPU (DDP) training. batch_size above is the + # total effective batch size; it is divided equally across GPUs. + # When using Condor, request_gpus in the submission file is set automatically. + # local_cache_path: /tmp # uncomment to avoid slow data loading during training due to network traffic on cluster + # Settings for HTCondor job submission (remove if not running on a Condor cluster). condor: num_cpus: 16 memory_cpus: 128000 - num_gpus: 1 memory_gpus: 8000 request_disk: 50GB \ No newline at end of file diff --git a/examples/npe_model/train_settings.yaml b/examples/npe_model/train_settings.yaml index 1095f6d54..0edbb859c 100644 --- a/examples/npe_model/train_settings.yaml +++ b/examples/npe_model/train_settings.yaml @@ -100,16 +100,18 @@ training: local: device: cuda # Change this to 'cuda' for training on a GPU. num_workers: 32 # num_workers >0 does not work on Mac, see https://stackoverflow.com/questions/64772335/pytorch-w-parallelnative-cpp206 + # num_gpus: 1 # Set to >1 to enable multi-GPU (DDP) training. batch_size above is the + # total effective batch size; it is divided equally across GPUs. + # When using Condor, request_gpus in the submission file is set automatically. runtime_limits: max_time_per_run: 3600000 max_epochs_per_run: 500 checkpoint_epochs: 50 leave_waveforms_on_disk: True - # local_cache_path: tmp # uncomment to avoid slow data loading during training due to network traffic on cluster - # Local settings related to condor, remove if not used on cluster + # local_cache_path: /tmp # uncomment to avoid slow data loading during training due to network traffic on cluster + # Settings for HTCondor job submission (remove if not running on a Condor cluster). condor: - num_cpus: 16 + num_cpus: 32 memory_cpus: 128000 - num_gpus: 1 memory_gpus: 8000 request_disk: 50GB \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index fa510cd77..fa9bdb603 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,7 +42,7 @@ dependencies = [ "scikit-learn", "scipy", "threadpoolctl", - "torch", + "torch>=2.0", "torchdiffeq", "torchvision", "tqdm", diff --git a/tests/core/test_amp.py b/tests/core/test_amp.py new file mode 100644 index 000000000..6f0f1c601 --- /dev/null +++ b/tests/core/test_amp.py @@ -0,0 +1,128 @@ +"""Tests for automatic mixed precision (AMP) and gradient accumulation.""" + +import types + +import pytest +import torch +import torch.nn as nn +from torch.utils.data import DataLoader, TensorDataset + +from dingo.core.posterior_models.base_model import train_epoch + + +def _make_mock_pm(device="cpu"): + """Create a minimal mock posterior model for testing train_epoch.""" + net = nn.Linear(4, 1) + net.to(device) + + pm = types.SimpleNamespace() + pm.network = net + pm.device = torch.device(device) + pm.rank = None + pm.epoch = 1 + pm.optimizer = torch.optim.SGD(net.parameters(), lr=0.01) + pm.scheduler = None + pm.scheduler_kwargs = None + + def loss_fn(theta, *context): + return net(theta).mean() + + pm.loss = loss_fn + return pm + + +def _make_dataloader(n_samples=16, batch_size=8, input_dim=4): + theta = torch.randn(n_samples, input_dim) + return DataLoader(TensorDataset(theta), batch_size=batch_size) + + +class TestAmpImports: + def test_gradscaler_importable(self): + try: + from torch.amp import GradScaler + except ImportError: + from torch.cuda.amp import GradScaler + assert GradScaler is not None + + def test_autocast_importable(self): + from torch.amp import autocast + + assert autocast is not None + + def test_base_model_uses_torch_amp(self): + import inspect + + from dingo.core.posterior_models import base_model + + source = inspect.getsource(base_model) + assert "from torch.amp import" in source + + +class TestTrainEpochWithoutAmp: + def test_returns_loss_and_iteration(self): + pm = _make_mock_pm() + dl = _make_dataloader() + avg_loss, n_iter = train_epoch(pm, dl, automatic_mixed_precision=False) + assert isinstance(avg_loss, float) + assert n_iter > 0 + + def test_updates_parameters(self): + pm = _make_mock_pm() + dl = _make_dataloader() + params_before = [p.clone() for p in pm.network.parameters()] + train_epoch(pm, dl, automatic_mixed_precision=False) + params_after = list(pm.network.parameters()) + assert any( + not torch.equal(b, a) for b, a in zip(params_before, params_after) + ), "Parameters should be updated after training." + + def test_gradient_accumulation(self): + """With 16 samples, batch=8, accum=2 → 2 batches → 1 optimizer step.""" + pm = _make_mock_pm() + dl = _make_dataloader(n_samples=16, batch_size=8) + _, n_iter = train_epoch( + pm, + dl, + gradient_updates_per_optimizer_step=2, + automatic_mixed_precision=False, + ) + assert n_iter == 1 + + def test_gradient_accumulation_no_step_on_incomplete_tail(self): + """3 batches with accum=2 → only 1 complete step (last batch dropped).""" + pm = _make_mock_pm() + dl = _make_dataloader(n_samples=24, batch_size=8) # 3 batches + _, n_iter = train_epoch( + pm, + dl, + gradient_updates_per_optimizer_step=2, + automatic_mixed_precision=False, + ) + assert n_iter == 1 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +class TestTrainEpochWithAmp: + def test_amp_returns_loss(self): + pm = _make_mock_pm(device="cuda") + dl = _make_dataloader() + avg_loss, n_iter = train_epoch(pm, dl, automatic_mixed_precision=True) + assert isinstance(avg_loss, float) + assert n_iter > 0 + + def test_amp_loss_is_finite(self): + pm = _make_mock_pm(device="cuda") + dl = _make_dataloader() + avg_loss, _ = train_epoch(pm, dl, automatic_mixed_precision=True) + assert torch.isfinite(torch.tensor(avg_loss)) + + def test_amp_with_gradient_accumulation(self): + pm = _make_mock_pm(device="cuda") + dl = _make_dataloader(n_samples=16, batch_size=8) + _, n_iter = train_epoch( + pm, + dl, + gradient_updates_per_optimizer_step=2, + automatic_mixed_precision=True, + ) + assert n_iter == 1 diff --git a/tests/core/test_multi_gpu.py b/tests/core/test_multi_gpu.py new file mode 100644 index 000000000..fe58f54b5 --- /dev/null +++ b/tests/core/test_multi_gpu.py @@ -0,0 +1,367 @@ +""" +Tests for multi-GPU DDP utilities. + +Most tests run without a real GPU by: + * testing individual components in isolation, or + * using the CPU-compatible ``gloo`` backend for process-group tests. +""" + +import os +import warnings + +import numpy as np +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn as nn +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.utils.data import TensorDataset + +from dingo.core.utils.torchutils import ( + build_train_and_test_loaders, + get_cuda_info, + replace_BatchNorm_with_SyncBatchNorm, + set_seed_based_on_rank, +) +from dingo.core.utils.trainutils import LossInfo, RuntimeLimits +from dingo.gw.training.train_pipeline import get_num_gpus + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class _TinyNet(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 1) + self.bn = nn.BatchNorm1d(4) + self.name = "FlowWrapper" + + def forward(self, x): + return self.linear(self.bn(x)) + + +def _setup_gloo(rank: int, world_size: int, port: int): + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + + +def _cleanup(): + dist.destroy_process_group() + + +# --------------------------------------------------------------------------- +# Utility-function tests (no GPU, no process group required) +# --------------------------------------------------------------------------- + + +class TestGetNumGpus: + def test_reads_from_local_num_gpus(self): + assert get_num_gpus({"num_gpus": 4}) == 4 + + def test_defaults_to_one(self): + assert get_num_gpus({}) == 1 + + def test_condor_num_gpus_deprecated_fallback(self): + """condor.num_gpus is still honoured but raises a DeprecationWarning.""" + with pytest.warns(DeprecationWarning, match="local.condor"): + result = get_num_gpus({"condor": {"num_gpus": 4}}) + assert result == 4 + + def test_local_num_gpus_takes_precedence(self): + """local.num_gpus wins and no warning is raised.""" + with warnings.catch_warnings(): + warnings.simplefilter("error") + assert get_num_gpus({"num_gpus": 4, "condor": {"num_gpus": 2}}) == 4 + + def test_returns_int(self): + assert isinstance(get_num_gpus({"num_gpus": "2"}), int) + + +class TestGetCudaInfo: + def test_returns_dict(self): + info = get_cuda_info() + assert isinstance(info, dict) + if torch.cuda.is_available(): + assert "device count" in info + else: + assert info == {} + + +class TestSetSeedBasedOnRank: + def test_different_seeds_per_rank(self): + torch.manual_seed(42) + set_seed_based_on_rank(0) + x0 = torch.rand(5).tolist() + + torch.manual_seed(42) + set_seed_based_on_rank(1) + x1 = torch.rand(5).tolist() + + assert x0 != x1, "Different ranks should produce different random draws." + + def test_numpy_seed_set(self): + set_seed_based_on_rank(0) + a = np.random.rand(3) + set_seed_based_on_rank(0) + b = np.random.rand(3) + np.testing.assert_array_equal(a, b) + + +class TestReplaceBatchNorm: + def test_converts_batchnorm(self): + net = _TinyNet() + assert isinstance(net.bn, nn.BatchNorm1d) + net = replace_BatchNorm_with_SyncBatchNorm(net) + assert isinstance(net.bn, nn.SyncBatchNorm) + + +class TestDDPStateDictStripping: + """Test that model_dict() strips DDP 'module.' prefix when saving.""" + + def test_save_model_strips_ddp_prefix(self, tmp_path): + from dingo.core.posterior_models.normalizing_flow import ( + NormalizingFlowPosteriorModel, + ) + + model_kwargs = { + "posterior_model_type": "normalizing_flow", + "posterior_kwargs": { + "input_dim": 2, + "context_dim": 4, + "num_flow_steps": 2, + "base_transform_kwargs": { + "hidden_dim": 8, + "num_transform_blocks": 1, + "activation": "elu", + "dropout_probability": 0.0, + "batch_norm": False, + "num_bins": 4, + "base_transform_type": "rq-coupling", + }, + }, + "embedding_type": None, + "embedding_kwargs": None, + } + metadata = {"train_settings": {"model": model_kwargs}} + + pm = NormalizingFlowPosteriorModel(metadata=metadata, device="cpu") + + # Simulate DDP wrapping by injecting "module." prefixes into state dict. + original_sd = pm.network.state_dict() + ddp_sd = {f"module.{k}": v for k, v in original_sd.items()} + + # Patch network so isinstance(network, DDP) is False but keys have "module.". + # We test the isinstance(network, DDP) branch explicitly below. + class _FakeDDP(nn.Module): + def __init__(self, mod): + super().__init__() + self.module = mod + + def state_dict(self, **kw): + return { + f"module.{k}": v for k, v in self.module.state_dict(**kw).items() + } + + pm.network = _FakeDDP( + pm.network.module if hasattr(pm.network, "module") else pm.network + ) + + # Manually test the save logic: DDP branch should strip "module." prefix. + if isinstance(pm.network, DDP): + saved_sd = pm.network.module.state_dict() + else: + # Simulate detecting wrapper by checking for the attribute. + saved_sd = ( + pm.network.module.state_dict() + if hasattr(pm.network, "module") + else pm.network.state_dict() + ) + + assert not any( + "module." in k for k in saved_sd.keys() + ), "Saved state dict should not contain 'module.' prefix." + + +class TestBuildTrainAndTestLoaders: + def test_single_gpu_returns_none_sampler(self): + dataset = TensorDataset(torch.randn(100, 4)) + train_loader, test_loader, sampler = build_train_and_test_loaders( + dataset=dataset, + train_fraction=0.8, + batch_size=10, + num_workers=0, + ) + assert sampler is None + assert len(train_loader.dataset) == 80 + assert len(test_loader.dataset) == 20 + + def test_ddp_returns_distributed_sampler(self): + from torch.utils.data import DistributedSampler + + dataset = TensorDataset(torch.randn(100, 4)) + train_loader, test_loader, sampler = build_train_and_test_loaders( + dataset=dataset, + train_fraction=0.8, + batch_size=5, + num_workers=0, + world_size=2, + rank=0, + ) + assert isinstance(sampler, DistributedSampler) + # Each rank sees ceil(80/2) = 40 samples. + assert len(train_loader.dataset) == 80 + + def test_ddp_splits_data_across_ranks(self): + from torch.utils.data import DistributedSampler + + dataset = TensorDataset(torch.arange(100).float().unsqueeze(1)) + _, _, sampler0 = build_train_and_test_loaders( + dataset, 1.0, 5, 0, world_size=2, rank=0 + ) + _, _, sampler1 = build_train_and_test_loaders( + dataset, 1.0, 5, 0, world_size=2, rank=1 + ) + # The two samplers should produce disjoint index sets. + indices0 = set(sampler0) + indices1 = set(sampler1) + assert len(indices0 & indices1) == 0, "DDP samplers should partition the data." + + +class TestLossInfoSingleGPU: + """LossInfo tests without a process group.""" + + def test_basic_update(self): + info = LossInfo( + epoch=1, + len_dataset=100, + batch_size_per_grad_update=10, + device=torch.device("cpu"), + ) + loss = torch.tensor(2.0) + info.update_timer("Dataloader") + info.cache_loss(loss, n=10) + info.update() + assert abs(info.get_avg() - 2.0) < 1e-5 + + def test_iteration_count(self): + info = LossInfo(1, 100, 10, device=torch.device("cpu")) + for _ in range(3): + info.update_timer("Dataloader") + info.cache_loss(torch.tensor(1.0), n=10) + info.update() + assert info.get_iteration() == 3 + + +class TestRuntimeLimitsSingleGPU: + def test_epoch_limit(self): + rl = RuntimeLimits( + max_epochs_total=5, epoch_start=0, device=torch.device("cpu") + ) + rl.max_epochs_total = 5 + assert not rl.limits_exceeded(4) + assert rl.limits_exceeded(5) + + +# --------------------------------------------------------------------------- +# Process-group tests using gloo (CPU, no NCCL required) +# --------------------------------------------------------------------------- + + +def _worker_loss_info(rank, world_size, port, result_queue): + """Worker function for testing LossInfo in DDP mode with gloo.""" + _setup_gloo(rank, world_size, port) + device = torch.device("cpu") + info = LossInfo(1, 100, 10, device=device) + # Simulate two optimizer steps with different losses per rank. + for i in range(2): + info.update_timer("Dataloader") + loss = torch.tensor(float(rank + 1) * (i + 1)) + info.cache_loss(loss, n=10) + info.update() + avg = info.get_avg() + n_iter = info.get_iteration() + result_queue.put((rank, avg, n_iter)) + _cleanup() + + +def _worker_runtime_limits(rank, world_size, port, result_queue): + """Worker function for testing RuntimeLimits broadcast with gloo.""" + _setup_gloo(rank, world_size, port) + device = torch.device("cpu") + rl = RuntimeLimits(max_epochs_total=5, epoch_start=0, device=device) + # Rank 0 sees epoch=6 (over limit); rank 1 sees epoch=4 (under). + epoch = 6 if rank == 0 else 4 + exceeded = rl.limits_exceeded(epoch) + result_queue.put((rank, exceeded)) + _cleanup() + + +@pytest.fixture() +def free_port(): + """Find a free port for the gloo process group.""" + import socket + + with socket.socket() as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +class TestGlooDDP: + """Integration tests using the CPU gloo backend — no GPU required.""" + + def test_loss_info_aggregation(self, free_port): + """Both ranks should observe the same averaged loss.""" + result_queue = mp.Queue() + world_size = 2 + processes = [ + mp.Process( + target=_worker_loss_info, + args=(rank, world_size, free_port, result_queue), + ) + for rank in range(world_size) + ] + for p in processes: + p.start() + for p in processes: + p.join() + assert p.exitcode == 0, f"Worker process failed with code {p.exitcode}" + + results = {} + while not result_queue.empty(): + rank, avg, n_iter = result_queue.get() + results[rank] = (avg, n_iter) + + assert len(results) == world_size + # Both ranks should report the same number of iterations. + assert results[0][1] == results[1][1] == 2 + + def test_runtime_limits_broadcast(self, free_port): + """When rank 0 hits the epoch limit, rank 1 should also stop.""" + result_queue = mp.Queue() + world_size = 2 + processes = [ + mp.Process( + target=_worker_runtime_limits, + args=(rank, world_size, free_port, result_queue), + ) + for rank in range(world_size) + ] + for p in processes: + p.start() + for p in processes: + p.join() + assert p.exitcode == 0 + + results = {} + while not result_queue.empty(): + rank, exceeded = result_queue.get() + results[rank] = exceeded + + assert len(results) == world_size + # Both ranks must agree: limit is exceeded because rank 0 triggered it. + assert results[0] is True + assert results[1] is True