""" Train DeepHealth with next-token / next-time-point supervision. The next-step dataset uses observed event histories, including CHECKUP state tokens, plus optional gap imputation. UTS training reads out only same-time group ends. """ from __future__ import annotations 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, DistributedSampler, RandomSampler, Sampler, ) from tqdm.auto import tqdm from dataset import HealthDataset, collate_fn from losses import build_loss from models import DeepHealth from readouts import build_readout from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX from train_util import ( configure_torch_for_training, create_unique_run_dir, resolve_device, save_checkpoint, save_config, set_optimizer_lr, set_seed, setup_logging, split_dataset, split_dataset_by_eid_files, ) MODEL_INPUT_KEYS = ( "event_seq", "time_seq", "sex", "padding_mask", ) EXPOSURE_INPUT_KEYS = ( "exposure_embedding", ) class ExposureLocalityBatchSampler(Sampler[List[int]]): """Randomized batches with within-buffer sorting by exposure parquet locality.""" def __init__( self, data_source, batch_size: int, buffer_size: int, seed: int, drop_last: bool = False, ) -> None: self.data_source = data_source self.batch_size = int(batch_size) self.buffer_size = max(int(buffer_size), self.batch_size) self.seed = int(seed) self.drop_last = bool(drop_last) self.epoch = 0 def __iter__(self) -> Iterator[List[int]]: n = len(self.data_source) generator = torch.Generator().manual_seed(self.seed + self.epoch) self.epoch += 1 shuffled = torch.randperm(n, generator=generator).tolist() for start in range(0, n, self.buffer_size): buffer = shuffled[start:start + self.buffer_size] buffer.sort(key=self._locality_key) for batch_start in range(0, len(buffer), self.batch_size): batch = buffer[batch_start:batch_start + self.batch_size] if len(batch) < self.batch_size and self.drop_last: continue yield batch def __len__(self) -> int: n = len(self.data_source) if self.drop_last: return n // self.batch_size return (n + self.batch_size - 1) // self.batch_size def _locality_key(self, local_idx: int) -> tuple[int, int, int]: dataset = self.data_source raw_idx = int(local_idx) if hasattr(dataset, "dataset") and hasattr(dataset, "indices"): raw_idx = int(dataset.indices[local_idx]) dataset = dataset.dataset sample = getattr(dataset, "samples", [])[raw_idx] exposure_index = sample.get("exposure_index") exposure_cache = getattr(dataset, "exposure_cache", None) if exposure_index is None or exposure_cache is None: return (2**31 - 1, 2**31 - 1, raw_idx) block_id, block_offset = exposure_cache.locality_key(exposure_index) 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_embedding: 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_embedding=exposure_embedding, ) 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") parser.add_argument("--data_prefix", type=str, default="ukb") parser.add_argument("--labels_file", type=str, default="labels.csv") parser.add_argument("--seed", type=int, default=42) parser.add_argument("--no_event_interval_years", type=float, default=5.0) parser.add_argument("--include_no_event_in_uts_target", action="store_true") parser.add_argument("--train_ratio", type=float, default=0.7) parser.add_argument("--val_ratio", type=float, default=0.15) parser.add_argument("--test_ratio", type=float, default=0.15) parser.add_argument("--train_eid_file", type=str, default="ukb_train_eid.csv") parser.add_argument("--val_eid_file", type=str, default="ukb_val_eid.csv") parser.add_argument("--test_eid_file", type=str, default="ukb_test_eid.csv") parser.add_argument("--n_embd", type=int, default=120) parser.add_argument("--n_head", type=int, default=10) parser.add_argument("--n_hist_layer", type=int, default=12) parser.add_argument("--dropout", type=float, default=0.0) parser.add_argument( "--input_ablation", choices=["none", "exposure_only"], default="none", help=( "Input ablation. 'exposure_only' removes disease-token embeddings " "while retaining exposure, sex, and age embeddings." ), ) parser.add_argument("--exposure_cache_dir", type=str, default=None) parser.add_argument( "--exposure_embeddings_file", type=str, default=None, help=( "Precomputed exposure embeddings. Defaults to " "/exposure_embeddings.npy." ), ) parser.add_argument("--target_mode", type=str, default="uts", choices=["delphi2m", "uts"]) parser.add_argument("--readout_name", type=str, default=None, choices=["token", "same_time_group_end", "last_valid"]) parser.add_argument("--readout_reduce", type=str, default="mean", choices=["mean", "sum"]) parser.add_argument("--t_min", type=float, default=0.0027378507871321013) parser.add_argument("--max_exp_input", type=float, default=60.0) parser.add_argument("--ce_weight", type=float, default=1.0) parser.add_argument("--time_weight", type=float, default=1.0) parser.add_argument("--ignore_no_event_in_delphi2m", action="store_true") parser.add_argument("--batch_size", type=int, default=128) parser.add_argument("--base_lr", type=float, default=3e-4) parser.add_argument("--weight_decay", type=float, default=0.1) parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99)) parser.add_argument("--grad_clip", type=float, default=1.0) parser.add_argument("--max_epochs", type=int, default=200) parser.add_argument("--warmup_epochs", type=int, default=10) parser.add_argument("--patience", type=int, default=15) parser.add_argument("--min_lr_ratio", type=float, default=0.1) parser.add_argument("--num_workers", type=int, default=4) parser.add_argument( "--prefetch_factor", type=int, default=4, help="DataLoader batches prefetched per worker when num_workers > 0.", ) parser.add_argument( "--exposure_locality_buffer_size", type=int, default=4096, help=( "Training-only shuffle buffer sorted by exposure parquet locality. " "Set 0 to use the standard RandomSampler." ), ) 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", help="Use torch.nn.DataParallel across multiple CUDA devices.", ) parser.add_argument( "--gpu_ids", type=str, 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() use_eid_split = all( getattr(args, name) for name in ("train_eid_file", "val_eid_file", "test_eid_file") ) if not use_eid_split and not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0): raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0") if args.num_workers > 0 and args.prefetch_factor <= 0: raise ValueError("prefetch_factor must be positive when num_workers > 0") if args.exposure_locality_buffer_size < 0: raise ValueError("exposure_locality_buffer_size must be non-negative") if args.exposure_embeddings_file and not args.exposure_cache_dir: raise ValueError( "--exposure_cache_dir is required with --exposure_embeddings_file" ) if args.input_ablation == "exposure_only" and not args.exposure_cache_dir: raise ValueError( "--input_ablation exposure_only requires --exposure_cache_dir" ) if args.target_mode == "uts": args.readout_name = args.readout_name or "same_time_group_end" args.include_no_event_in_uts_target = True else: args.readout_name = args.readout_name or "token" if args.gpu_ids: try: args.gpu_ids = [int(part.strip()) for part in args.gpu_ids.split(",") if part.strip()] except ValueError as exc: raise ValueError("--gpu_ids must be a comma-separated list of integers") from exc if not args.gpu_ids: raise ValueError("--gpu_ids did not contain any valid CUDA device ids") args.data_parallel = True return args def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float: if epoch < args.warmup_epochs: return adaptive_lr * (epoch + 1) / args.warmup_epochs progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs) cosine = 0.5 * (1 + math.cos(math.pi * progress)) return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio)) def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]: non_blocking = device.type == "cuda" return { key: value.to(device, non_blocking=non_blocking) if isinstance(value, torch.Tensor) else value for key, value in batch.items() } def _cuda_device_index(device: torch.device) -> int: if device.type != "cuda": raise ValueError("CUDA device is required for multi-GPU training") if device.index is not None: return int(device.index) current = torch.cuda.current_device() return int(current) def unwrap_model(model): if isinstance(model, (torch.nn.DataParallel, DistributedDataParallel)): return model.module return model def maybe_wrap_data_parallel( model: DeepHealth, args: argparse.Namespace, device: torch.device, logger: logging.Logger, ): if not args.data_parallel: return model if device.type != "cuda": raise ValueError("--data_parallel requires --device cuda or cuda:") if not torch.cuda.is_available() or torch.cuda.device_count() < 2: 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())) 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: input_label = ( args.input_ablation if args.input_ablation != "none" else ("exposure" if args.exposure_cache_dir else "noexposure") ) 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"{input_label}_" 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, n_embd=args.n_embd, n_head=args.n_head, n_hist_layer=args.n_hist_layer, target_mode="next_token", dist_mode="exponential", dropout=args.dropout, use_exposure_embeddings=args.exposure_embeddings_file is not None, input_ablation=args.input_ablation, ) def build_next_step_readout(args: argparse.Namespace): if args.readout_name == "same_time_group_end": return build_readout("same_time_group_end", reduce=args.readout_reduce) return build_readout(args.readout_name) def build_next_step_loss(args: argparse.Namespace): if args.target_mode == "delphi2m": ignored_tokens = {PAD_IDX, CHECKUP_IDX} if args.ignore_no_event_in_delphi2m: ignored_tokens.add(NO_EVENT_IDX) return build_loss( "delphi2m", ignored_tokens=ignored_tokens, t_min=args.t_min, max_exp_input=args.max_exp_input, ce_weight=args.ce_weight, time_weight=args.time_weight, ) return build_loss( "uts", ignored_idx={PAD_IDX, CHECKUP_IDX}, t_min=args.t_min, max_exp_input=args.max_exp_input, ) def compute_next_step_loss( args: argparse.Namespace, model, criterion, batch: Dict[str, torch.Tensor], device: torch.device, ) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]: batch_cpu = batch 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}, device, ) model_kwargs = { "event_seq": batch["event_seq"], "time_seq": batch["time_seq"], "sex": batch["sex"], "padding_mask": batch["padding_mask"], "readout_mask": batch["readout_mask"], } if "exposure_embedding" in batch: model_kwargs["exposure_embedding"] = batch["exposure_embedding"] 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=current_times, padding_mask=output_readout_mask, return_components=True, ) else: loss, parts = criterion( logits=logits, target_multi_hot=targets["target_multi_hot"], target_dt_unique=targets["target_dt_unique"], readout_mask=output_readout_mask, return_components=True, ) return loss, parts def run_epoch( logger: logging.Logger, args: argparse.Namespace, 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) 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, disable=rank != 0 ) for batch_idx, batch in enumerate(progress): 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() 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} 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( args: argparse.Namespace, dataset: HealthDataset, run_name: str, train_subset, val_subset, test_subset, ) -> Dict[str, Any]: return { "run_name": run_name, "dataset_class": "NextStepHealthDataset", "collate_fn": "next_step_collate_fn", "model_class": "DeepHealth", "model_target_mode": "next_token", "target_mode": args.target_mode, "dist_mode": "exponential", "dataset_metadata": { "vocab_size": int(dataset.vocab_size), }, "use_exposure_embeddings": args.exposure_embeddings_file is not None, "input_ablation": args.input_ablation, "exposure_cache_dir": args.exposure_cache_dir, "exposure_embeddings_file": args.exposure_embeddings_file, "num_workers": int(args.num_workers), "prefetch_factor": int(args.prefetch_factor), "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)), "test": int(len(test_subset)), }, "resolved_readout_name": args.readout_name, "resolved_loss_name": args.target_mode, } def main() -> None: args = parse_args() if args.exposure_cache_dir and args.exposure_embeddings_file is None: args.exposure_embeddings_file = str( Path(args.exposure_cache_dir) / "exposure_embeddings.npy" ) device, rank, local_rank, world_size = init_distributed(args) set_seed(args.seed + rank) configure_torch_for_training(device) 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}") logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}") logger.info(f"exposure_cache_dir={args.exposure_cache_dir}") logger.info(f"exposure_embeddings_file={args.exposure_embeddings_file}") logger.info(f"input_ablation={args.input_ablation}") logger.info( "DataLoader IO: " f"num_workers={args.num_workers}, " f"prefetch_factor={args.prefetch_factor if args.num_workers > 0 else None}, " f"exposure_locality_buffer_size={args.exposure_locality_buffer_size}" ) dataset = HealthDataset( data_prefix=args.data_prefix, labels_file=args.labels_file, no_event_interval_years=args.no_event_interval_years, include_no_event_in_uts_target=args.include_no_event_in_uts_target, exposure_cache_dir=args.exposure_cache_dir, exposure_embeddings_file=args.exposure_embeddings_file, ) if dataset.exposure_cache is not None: embedding_dim = int(dataset.exposure_cache.embeddings.shape[1]) if embedding_dim != args.n_embd: raise ValueError( f"Exposure embedding dim {embedding_dim} must equal " f"--n_embd={args.n_embd}" ) if args.train_eid_file and args.val_eid_file and args.test_eid_file: train_subset, val_subset, test_subset = split_dataset_by_eid_files( dataset=dataset, train_eid_file=args.train_eid_file, val_eid_file=args.val_eid_file, test_eid_file=args.test_eid_file, ) logger.info( "Using eid split files: " f"train={args.train_eid_file}, val={args.val_eid_file}, test={args.test_eid_file}" ) else: train_subset, val_subset, test_subset = split_dataset( dataset=dataset, train_ratio=args.train_ratio, val_ratio=args.val_ratio, test_ratio=args.test_ratio, seed=args.seed, ) logger.info( f"Using random ratio split: train={args.train_ratio}, " f"val={args.val_ratio}, test={args.test_ratio}, seed={args.seed}" ) logger.info( 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, "pin_memory": device.type == "cuda", "persistent_workers": args.num_workers > 0, "prefetch_factor": args.prefetch_factor if args.num_workers > 0 else None, } use_locality_sampler = ( args.exposure_cache_dir is not None and args.exposure_locality_buffer_size > 0 ) if use_locality_sampler: logger.info("Using exposure-locality batch sampler for training") if world_size > 1: train_batch_sampler = DistributedExposureLocalityBatchSampler( train_subset, 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_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=local_batch_size, sampler=val_sampler, shuffle=False, **dataloader_kwargs, ) test_loader = DataLoader( test_subset, batch_size=local_batch_size, sampler=test_sampler, shuffle=False, **dataloader_kwargs, ) 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(), lr=args.base_lr, 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) 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 history = [] best_model_path = run_dir / "best_model.pt" 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, criterion, train_loader, optimizer, device, True, rank, scaler, amp_enabled, ) with torch.no_grad(): 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 if rank == 0: save_checkpoint(unwrap_model(model).model, best_model_path) else: patience += 1 logger.info( f"Epoch {epoch + 1}/{args.max_epochs} | lr={lr:.6f} | " f"train_loss={train_loss:.6f} | val_loss={val_loss:.6f} | " f"best_val_loss={best_val:.6f} | patience={patience}/{args.patience} | " f"elapsed={time.time() - start:.1f}s" ) history.append({ "epoch": epoch + 1, "lr": lr, "train_loss": train_loss, "val_loss": val_loss, "best_val_loss": best_val, "is_best": int(is_best), }) if patience >= args.patience: logger.info(f"Early stopping triggered at epoch {epoch + 1}") break 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...") 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, 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__": try: main() finally: if dist.is_initialized(): dist.destroy_process_group()