diff --git a/README.md b/README.md index 655755d..f1e9434 100644 --- a/README.md +++ b/README.md @@ -83,6 +83,14 @@ torchrun --standalone --nproc_per_node=4 train_exposure_autoencoder.py \ In DDP mode, `--batch_size` is the global batch size and must be divisible by the number of processes. + +The end-to-end next-step trainer supports the same DDP launch pattern: + +```bash +torchrun --standalone --nproc_per_node=4 train_next_step.py \ + --exposure_cache_dir ukb_exposure_cache \ + --batch_size 128 +``` Training-channel statistics are cached at `/train_channel_stats.npz`; use `--recompute_channel_stats` only when a forced refresh is needed. diff --git a/train_next_step.py b/train_next_step.py index 7d13644..6f6fa84 100644 --- a/train_next_step.py +++ b/train_next_step.py @@ -11,19 +11,29 @@ import argparse import json import logging import math +import os import time +from pathlib import Path from typing import Any, Dict, Iterator, List import numpy as np import torch +import torch.distributed as dist +import torch.nn as nn +from torch.nn.parallel import DistributedDataParallel from torch.nn.utils import clip_grad_norm_ from torch.optim import AdamW -from torch.utils.data import DataLoader, RandomSampler, Sampler +from torch.utils.data import ( + DataLoader, + DistributedSampler, + RandomSampler, + Sampler, +) from tqdm.auto import tqdm from dataset import HealthDataset, collate_fn from losses import build_loss -from models import DeepHealth, DeepHealthOutput +from models import DeepHealth from readouts import build_readout from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX from train_util import ( @@ -106,6 +116,92 @@ class ExposureLocalityBatchSampler(Sampler[List[int]]): return (block_id, block_offset, raw_idx) +class DistributedExposureLocalityBatchSampler(Sampler[List[int]]): + """Shard samples across ranks, then batch each shard by exposure locality.""" + + def __init__( + self, + data_source, + batch_size: int, + buffer_size: int, + seed: int, + rank: int, + world_size: int, + ) -> None: + self.data_source = data_source + self.batch_size = int(batch_size) + self.buffer_size = max(int(buffer_size), self.batch_size) + self.distributed_sampler = DistributedSampler( + data_source, + num_replicas=world_size, + rank=rank, + shuffle=True, + seed=seed, + drop_last=False, + ) + self._key_sampler = ExposureLocalityBatchSampler( + data_source, batch_size, buffer_size, seed + ) + + def set_epoch(self, epoch: int) -> None: + self.distributed_sampler.set_epoch(epoch) + + def __iter__(self) -> Iterator[List[int]]: + indices = list(self.distributed_sampler) + for start in range(0, len(indices), self.buffer_size): + buffer = indices[start:start + self.buffer_size] + buffer.sort(key=self._key_sampler._locality_key) + for batch_start in range(0, len(buffer), self.batch_size): + yield buffer[batch_start:batch_start + self.batch_size] + + def __len__(self) -> int: + return math.ceil(len(self.distributed_sampler) / self.batch_size) + + +class NextStepTrainingModel(nn.Module): + """Keep backbone, readout, and risk head inside the DDP forward boundary.""" + + def __init__(self, model: DeepHealth, readout: nn.Module, readout_name: str): + super().__init__() + self.model = model + self.readout = readout + self.readout_name = readout_name + + def forward( + self, + event_seq: torch.Tensor, + time_seq: torch.Tensor, + sex: torch.Tensor, + padding_mask: torch.Tensor, + readout_mask: torch.Tensor | None = None, + exposure_daily: torch.Tensor | None = None, + exposure_monthly: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + hidden = self.model( + event_seq=event_seq, + time_seq=time_seq, + sex=sex, + padding_mask=padding_mask, + target_mode="next_token", + exposure_daily=exposure_daily, + exposure_monthly=exposure_monthly, + ) + if not isinstance(hidden, torch.Tensor): + raise TypeError("DeepHealth forward must return a hidden-state tensor") + current_times = time_seq[:, :hidden.size(1)] + current_padding = padding_mask[:, :hidden.size(1)] + readout_out = self.readout( + hidden=hidden, + time_seq=current_times, + padding_mask=current_padding, + readout_mask=readout_mask + if self.readout_name == "same_time_group_end" + else None, + ) + logits = self.model.calc_risk(readout_out.hidden) + return logits, current_times, readout_out.readout_mask + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Train DeepHealth with next-token/point supervision") @@ -174,6 +270,12 @@ def parse_args() -> argparse.Namespace: ), ) parser.add_argument("--device", type=str, default="cuda") + parser.add_argument( + "--amp", + action=argparse.BooleanOptionalAction, + default=True, + help="Use CUDA automatic mixed precision.", + ) parser.add_argument( "--data_parallel", action="store_true", @@ -185,6 +287,12 @@ def parse_args() -> argparse.Namespace: default=None, help="Comma-separated CUDA device ids for --data_parallel, e.g. 0,1,2,3.", ) + parser.add_argument( + "--ddp_backend", + default=None, + choices=["nccl", "gloo"], + help="DDP backend. Defaults to nccl for torchrun multi-GPU training.", + ) parser.add_argument("--progress_interval", type=int, default=20) args = parser.parse_args() @@ -242,7 +350,9 @@ def _cuda_device_index(device: torch.device) -> int: def unwrap_model(model): - return model.module if isinstance(model, torch.nn.DataParallel) else model + if isinstance(model, (torch.nn.DataParallel, DistributedDataParallel)): + return model.module + return model def maybe_wrap_data_parallel( @@ -259,14 +369,61 @@ def maybe_wrap_data_parallel( raise ValueError("--data_parallel requires at least two CUDA devices") primary = _cuda_device_index(device) device_ids = args.gpu_ids if args.gpu_ids else list(range(torch.cuda.device_count())) - if primary not in device_ids: - device_ids = [primary, *[idx for idx in device_ids if idx != primary]] + device_ids = [primary, *[idx for idx in device_ids if idx != primary]] if len(device_ids) < 2: raise ValueError("--data_parallel needs at least two device ids") logger.info(f"Using DataParallel on CUDA devices: {device_ids}") return torch.nn.DataParallel(model, device_ids=device_ids, output_device=primary) +def init_distributed( + args: argparse.Namespace, +) -> tuple[torch.device, int, int, int]: + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if world_size == 1: + return resolve_device(args.device), 0, 0, 1 + if args.data_parallel: + raise ValueError("--data_parallel cannot be combined with torchrun/DDP") + if not torch.cuda.is_available(): + raise ValueError("Multi-process next-step training requires CUDA") + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group( + backend=args.ddp_backend or "nccl", + init_method="env://", + ) + return torch.device("cuda", local_rank), rank, local_rank, world_size + + +def distributed_run_dir( + args: argparse.Namespace, rank: int, world_size: int +) -> tuple[Path, str]: + payload: list[str | None] = [None, None] + if rank == 0: + run_dir, run_name = create_unique_run_dir( + lambda timestamp: ( + f"absolute_exponential_next_token_{args.target_mode}_" + f"gap_{args.no_event_interval_years:g}y_" + f"{'exposure' if args.exposure_cache_dir else 'noexposure'}_" + f"{timestamp}" + ) + ) + payload = [str(run_dir), run_name] + if world_size > 1: + dist.broadcast_object_list(payload, src=0) + return Path(str(payload[0])), str(payload[1]) + + +def rank_logger(rank: int, run_dir: Path) -> logging.Logger: + if rank == 0: + return setup_logging(run_dir) + logger = logging.getLogger(f"DeepHealth.rank{rank}") + logger.handlers.clear() + logger.addHandler(logging.NullHandler()) + return logger + + def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth: return DeepHealth( vocab_size=dataset.vocab_size, @@ -314,38 +471,15 @@ def build_next_step_loss(args: argparse.Namespace): ) -def build_augmented_next_step_targets( - batch_cpu: Dict[str, torch.Tensor], - model_out: DeepHealthOutput, - include_uts_targets: bool, -) -> Dict[str, torch.Tensor]: - device = model_out.hidden.device - non_blocking = device.type == "cuda" - targets = { - "target_event_seq": batch_cpu["target_event_seq"].to(device, non_blocking=non_blocking), - "target_time_seq": batch_cpu["target_time_seq"].to(device, non_blocking=non_blocking), - "readout_mask": batch_cpu["readout_mask"].to(device, non_blocking=non_blocking), - } - if include_uts_targets: - targets["target_dt_unique"] = batch_cpu["target_dt_unique"].to( - device, non_blocking=non_blocking - ) - targets["target_multi_hot"] = batch_cpu["target_multi_hot"].to( - device, non_blocking=non_blocking - ) - return targets - - def compute_next_step_loss( args: argparse.Namespace, - model: DeepHealth, - readout, + model, criterion, batch: Dict[str, torch.Tensor], device: torch.device, ) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]: batch_cpu = batch - input_keys = list(MODEL_INPUT_KEYS) + input_keys = [*MODEL_INPUT_KEYS, "readout_mask"] input_keys.extend(key for key in EXPOSURE_INPUT_KEYS if key in batch_cpu) batch = move_batch_to_device( {key: batch_cpu[key] for key in input_keys}, @@ -356,42 +490,36 @@ def compute_next_step_loss( "time_seq": batch["time_seq"], "sex": batch["sex"], "padding_mask": batch["padding_mask"], - "target_mode": "next_token", + "readout_mask": batch["readout_mask"], } if "exposure_daily" in batch: model_kwargs["exposure_daily"] = batch["exposure_daily"] model_kwargs["exposure_monthly"] = batch["exposure_monthly"] - hidden = model(**model_kwargs) - if not isinstance(hidden, torch.Tensor): - raise TypeError("DeepHealth forward must return a hidden-state tensor") - model_out = DeepHealthOutput( - hidden=hidden, - time_seq=batch["time_seq"][:, : hidden.size(1)], - padding_mask=batch["padding_mask"][:, : hidden.size(1)], - event_len=int(hidden.size(1)), - ) - targets = build_augmented_next_step_targets( - batch_cpu=batch_cpu, - model_out=model_out, - include_uts_targets=args.target_mode == "uts", - ) - readout_out = readout( - hidden=model_out.hidden, - time_seq=model_out.time_seq, - padding_mask=model_out.padding_mask, - readout_mask=targets["readout_mask"] - if args.readout_name == "same_time_group_end" - else None, - ) - logits = unwrap_model(model).calc_risk(readout_out.hidden) + logits, current_times, output_readout_mask = model(**model_kwargs) + non_blocking = device.type == "cuda" + targets = { + "target_event_seq": batch_cpu["target_event_seq"].to( + device, non_blocking=non_blocking + ), + "target_time_seq": batch_cpu["target_time_seq"].to( + device, non_blocking=non_blocking + ), + } + if args.target_mode == "uts": + targets["target_dt_unique"] = batch_cpu["target_dt_unique"].to( + device, non_blocking=non_blocking + ) + targets["target_multi_hot"] = batch_cpu["target_multi_hot"].to( + device, non_blocking=non_blocking + ) if args.target_mode == "delphi2m": loss, parts = criterion( logits=logits, target_events=targets["target_event_seq"], target_times=targets["target_time_seq"], - current_times=model_out.time_seq, - padding_mask=readout_out.readout_mask, + current_times=current_times, + padding_mask=output_readout_mask, return_components=True, ) else: @@ -399,70 +527,82 @@ def compute_next_step_loss( logits=logits, target_multi_hot=targets["target_multi_hot"], target_dt_unique=targets["target_dt_unique"], - readout_mask=readout_out.readout_mask, + readout_mask=output_readout_mask, return_components=True, ) - if not torch.isfinite(loss): - raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}") return loss, parts def run_epoch( logger: logging.Logger, args: argparse.Namespace, - model: DeepHealth, - readout, + model, criterion, loader: DataLoader, optimizer: AdamW | None, device: torch.device, is_train: bool, + rank: int, + scaler: torch.amp.GradScaler, + amp_enabled: bool, ) -> float: model.train(is_train) - readout.train(is_train) - total = torch.zeros((), device=device) - n_batches = 0 - skipped = 0 + totals = torch.zeros(3, device=device, dtype=torch.float64) parts_sum: Dict[str, torch.Tensor] = {} desc = "train" if is_train else "val" progress_interval = max(1, int(args.progress_interval)) - progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True) + progress = tqdm( + loader, desc=desc, leave=False, dynamic_ncols=True, disable=rank != 0 + ) for batch_idx, batch in enumerate(progress): - try: - loss, parts = compute_next_step_loss(args, model, readout, criterion, batch, device) - if is_train: - if optimizer is None: - raise ValueError("optimizer is required for training") - optimizer.zero_grad(set_to_none=True) - loss.backward() - if args.grad_clip > 0: - clip_grad_norm_(model.parameters(), args.grad_clip) - optimizer.step() + with torch.autocast( + device_type=device.type, + dtype=torch.float16, + enabled=amp_enabled, + ): + loss, parts = compute_next_step_loss( + args, model, criterion, batch, device + ) + finite = torch.isfinite(loss).to(dtype=torch.int32) + if dist.is_initialized(): + dist.all_reduce(finite, op=dist.ReduceOp.MIN) + if not bool(finite.item()): + totals[2] += 1 + continue + if is_train: + if optimizer is None: + raise ValueError("optimizer is required for training") + optimizer.zero_grad(set_to_none=True) + scaler.scale(loss).backward() + scaler.unscale_(optimizer) + if args.grad_clip > 0: + clip_grad_norm_(model.parameters(), args.grad_clip) + scaler.step(optimizer) + scaler.update() - total = total + loss.detach() - n_batches += 1 - for name, value in parts.items(): - parts_sum[name] = parts_sum.get(name, torch.zeros((), device=device)) + value.detach() - if (batch_idx + 1) % progress_interval == 0: - avg = total / max(1, n_batches) - postfix = { - "loss": f"{float(loss.detach().cpu()):.4f}", - "avg": f"{float(avg.detach().cpu()):.4f}", - "skipped": skipped, - } - for name, value in parts_sum.items(): - postfix[name] = f"{float((value / max(1, n_batches)).detach().cpu()):.4f}" - progress.set_postfix(postfix) - except RuntimeError as exc: - if "Loss is not finite" not in str(exc): - raise - skipped += 1 - logger.warning(f"Batch {batch_idx} skipped: {str(exc)[:120]}") + totals[0] += loss.detach().double() + totals[1] += 1 + for name, value in parts.items(): + parts_sum[name] = ( + parts_sum.get(name, torch.zeros((), device=device)) + + value.detach() + ) + if rank == 0 and (batch_idx + 1) % progress_interval == 0: + progress.set_postfix( + loss=f"{loss.detach().item():.4f}", + avg=f"{(totals[0] / totals[1].clamp_min(1)).item():.4f}", + skipped=int(totals[2].item()), + ) + if dist.is_initialized(): + dist.all_reduce(totals, op=dist.ReduceOp.SUM) + skipped = int(totals[2].item()) if skipped: - logger.info(f"Skipped {skipped} batches due to non-finite loss") - return float((total / max(1, n_batches)).detach().cpu()) if n_batches else float("inf") + logger.info(f"Skipped {skipped} rank-batches due to non-finite loss") + if totals[1].item() == 0: + return float("inf") + return float((totals[0] / totals[1]).item()) def build_metadata( @@ -499,6 +639,7 @@ def build_metadata( "exposure_locality_buffer_size": int(args.exposure_locality_buffer_size), "data_parallel": bool(args.data_parallel), "gpu_ids": args.gpu_ids, + "ddp_world_size": int(os.environ.get("WORLD_SIZE", "1")), "split_sizes": { "train": int(len(train_subset)), "val": int(len(val_subset)), @@ -511,19 +652,12 @@ def build_metadata( def main() -> None: args = parse_args() - set_seed(args.seed) - device = resolve_device(args.device) + device, rank, local_rank, world_size = init_distributed(args) + set_seed(args.seed + rank) configure_torch_for_training(device) - run_dir, run_name = create_unique_run_dir( - lambda timestamp: ( - f"absolute_exponential_next_token_{args.target_mode}_" - f"gap_{args.no_event_interval_years:g}y_" - f"{'exposure' if args.exposure_cache_dir else 'noexposure'}_" - f"{timestamp}" - ) - ) - logger = setup_logging(run_dir) + run_dir, run_name = distributed_run_dir(args, rank, world_size) + logger = rank_logger(rank, run_dir) logger.info(f"Starting next-step training run: {run_name}") logger.info(f"Device: {device}") @@ -571,6 +705,12 @@ def main() -> None: f"Samples: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}" ) + if args.batch_size % world_size != 0: + raise ValueError( + f"--batch_size={args.batch_size} must be divisible by " + f"DDP world size {world_size}" + ) + local_batch_size = args.batch_size // world_size dataloader_kwargs = { "collate_fn": collate_fn, "num_workers": args.num_workers, @@ -584,42 +724,90 @@ def main() -> None: ) if use_locality_sampler: logger.info("Using exposure-locality batch sampler for training") - train_loader = DataLoader( - train_subset, - batch_sampler=ExposureLocalityBatchSampler( + if world_size > 1: + train_batch_sampler = DistributedExposureLocalityBatchSampler( train_subset, - batch_size=args.batch_size, + batch_size=local_batch_size, buffer_size=args.exposure_locality_buffer_size, seed=args.seed, - ), + rank=rank, + world_size=world_size, + ) + else: + train_batch_sampler = ExposureLocalityBatchSampler( + train_subset, + batch_size=local_batch_size, + buffer_size=args.exposure_locality_buffer_size, + seed=args.seed, + ) + train_loader = DataLoader( + train_subset, + batch_sampler=train_batch_sampler, **dataloader_kwargs, ) else: - train_loader = DataLoader( - train_subset, - batch_size=args.batch_size, - sampler=RandomSampler( + train_sampler = ( + DistributedSampler( + train_subset, + num_replicas=world_size, + rank=rank, + shuffle=True, + seed=args.seed, + ) + if world_size > 1 + else RandomSampler( train_subset, generator=torch.Generator().manual_seed(args.seed), - ), + ) + ) + train_loader = DataLoader( + train_subset, + batch_size=local_batch_size, + sampler=train_sampler, **dataloader_kwargs, ) + val_sampler = ( + DistributedSampler( + val_subset, num_replicas=world_size, rank=rank, shuffle=False + ) + if world_size > 1 else None + ) + test_sampler = ( + DistributedSampler( + test_subset, num_replicas=world_size, rank=rank, shuffle=False + ) + if world_size > 1 else None + ) val_loader = DataLoader( val_subset, - batch_size=args.batch_size, + batch_size=local_batch_size, + sampler=val_sampler, shuffle=False, **dataloader_kwargs, ) test_loader = DataLoader( test_subset, - batch_size=args.batch_size, + batch_size=local_batch_size, + sampler=test_sampler, shuffle=False, **dataloader_kwargs, ) - model = build_model(args, dataset).to(device) - model = maybe_wrap_data_parallel(model, args, device, logger) + backbone = build_model(args, dataset) readout = build_next_step_readout(args).to(device) + model = NextStepTrainingModel( + backbone, readout, args.readout_name + ).to(device) + if world_size > 1: + model = DistributedDataParallel( + model, device_ids=[local_rank], output_device=local_rank + ) + logger.info( + f"Using DDP with {world_size} processes; " + f"global_batch={args.batch_size}, per_gpu_batch={local_batch_size}" + ) + else: + model = maybe_wrap_data_parallel(model, args, device, logger) criterion = build_next_step_loss(args) optimizer = AdamW( model.parameters(), @@ -627,13 +815,18 @@ def main() -> None: betas=tuple(args.betas), weight_decay=args.weight_decay, ) + amp_enabled = bool(args.amp and device.type == "cuda") + scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled) adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128) - save_config( - args, - run_dir / "train_config.json", - extra=build_metadata(args, dataset, run_name, train_subset, val_subset, test_subset), - ) + if rank == 0: + save_config( + args, + run_dir / "train_config.json", + extra=build_metadata( + args, dataset, run_name, train_subset, val_subset, test_subset + ), + ) best_val = float("inf") patience = 0 @@ -642,18 +835,29 @@ def main() -> None: start = time.time() for epoch in range(args.max_epochs): + if hasattr(train_loader.batch_sampler, "set_epoch"): + train_loader.batch_sampler.set_epoch(epoch) + elif hasattr(train_loader.sampler, "set_epoch"): + train_loader.sampler.set_epoch(epoch) lr = get_lr(epoch, args, adaptive_lr) set_optimizer_lr(optimizer, lr) - train_loss = run_epoch(logger, args, model, readout, criterion, train_loader, optimizer, device, True) + train_loss = run_epoch( + logger, args, model, criterion, train_loader, optimizer, + device, True, rank, scaler, amp_enabled, + ) with torch.no_grad(): - val_loss = run_epoch(logger, args, model, readout, criterion, val_loader, None, device, False) + val_loss = run_epoch( + logger, args, model, criterion, val_loader, None, + device, False, rank, scaler, amp_enabled, + ) is_best = val_loss < best_val if is_best: best_val = val_loss patience = 0 - save_checkpoint(unwrap_model(model), best_model_path) + if rank == 0: + save_checkpoint(unwrap_model(model).model, best_model_path) else: patience += 1 @@ -675,15 +879,25 @@ def main() -> None: logger.info(f"Early stopping triggered at epoch {epoch + 1}") break - with (run_dir / "history.json").open("w", encoding="utf-8") as f: - json.dump(history, f, indent=2) + if rank == 0: + with (run_dir / "history.json").open("w", encoding="utf-8") as f: + json.dump(history, f, indent=2) logger.info("Evaluating best model on next-step test split...") - unwrap_model(model).load_state_dict(torch.load(best_model_path, map_location=device)) + if world_size > 1: + dist.barrier() + unwrap_model(model).model.load_state_dict( + torch.load(best_model_path, map_location=device) + ) with torch.no_grad(): - test_loss = run_epoch(logger, args, model, readout, criterion, test_loader, None, device, False) + test_loss = run_epoch( + logger, args, model, criterion, test_loader, None, + device, False, rank, scaler, amp_enabled, + ) logger.info(f"Test loss: {test_loss:.6f}") logger.info(f"Best checkpoint: {best_model_path}") + if dist.is_initialized(): + dist.destroy_process_group() if __name__ == "__main__":