Add distributed end-to-end training

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
 `<exposure_cache_dir>/train_channel_stats.npz`; use
 `--recompute_channel_stats` only when a forced refresh is needed.

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,
+    model,
-    readout,
     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)
+    logits, current_times, output_readout_mask = model(**model_kwargs)
-    if not isinstance(hidden, torch.Tensor):
+    non_blocking = device.type == "cuda"
-        raise TypeError("DeepHealth forward must return a hidden-state tensor")
+    targets = {
-    model_out = DeepHealthOutput(
+        "target_event_seq": batch_cpu["target_event_seq"].to(
-        hidden=hidden,
+            device, non_blocking=non_blocking
-        time_seq=batch["time_seq"][:, : hidden.size(1)],
+        ),
-        padding_mask=batch["padding_mask"][:, : hidden.size(1)],
+        "target_time_seq": batch_cpu["target_time_seq"].to(
-        event_len=int(hidden.size(1)),
+            device, non_blocking=non_blocking
-    )
+        ),
-    targets = build_augmented_next_step_targets(
+    }
-        batch_cpu=batch_cpu,
+    if args.target_mode == "uts":
-        model_out=model_out,
+        targets["target_dt_unique"] = batch_cpu["target_dt_unique"].to(
-        include_uts_targets=args.target_mode == "uts",
+            device, non_blocking=non_blocking
-    )
+        )
-    readout_out = readout(
+        targets["target_multi_hot"] = batch_cpu["target_multi_hot"].to(
-        hidden=model_out.hidden,
+            device, non_blocking=non_blocking
-        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)
 
     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,
+            current_times=current_times,
-            padding_mask=readout_out.readout_mask,
+            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,
+    model,
-    readout,
     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)
+    totals = torch.zeros(3, device=device, dtype=torch.float64)
-    total = torch.zeros((), device=device)
-    n_batches = 0
-    skipped = 0
     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:
+        with torch.autocast(
-            loss, parts = compute_next_step_loss(args, model, readout, criterion, batch, device)
+            device_type=device.type,
-            if is_train:
+            dtype=torch.float16,
-                if optimizer is None:
+            enabled=amp_enabled,
-                    raise ValueError("optimizer is required for training")
+        ):
-                optimizer.zero_grad(set_to_none=True)
+            loss, parts = compute_next_step_loss(
-                loss.backward()
+                args, model, criterion, batch, device
-                if args.grad_clip > 0:
+            )
-                    clip_grad_norm_(model.parameters(), args.grad_clip)
+        finite = torch.isfinite(loss).to(dtype=torch.int32)
-                optimizer.step()
+        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()
+        totals[0] += loss.detach().double()
-            n_batches += 1
+        totals[1] += 1
-            for name, value in parts.items():
+        for name, value in parts.items():
-                parts_sum[name] = parts_sum.get(name, torch.zeros((), device=device)) + value.detach()
+            parts_sum[name] = (
-            if (batch_idx + 1) % progress_interval == 0:
+                parts_sum.get(name, torch.zeros((), device=device))
-                avg = total / max(1, n_batches)
+                + value.detach()
-                postfix = {
+            )
-                    "loss": f"{float(loss.detach().cpu()):.4f}",
+        if rank == 0 and (batch_idx + 1) % progress_interval == 0:
-                    "avg": f"{float(avg.detach().cpu()):.4f}",
+            progress.set_postfix(
-                    "skipped": skipped,
+                loss=f"{loss.detach().item():.4f}",
-                }
+                avg=f"{(totals[0] / totals[1].clamp_min(1)).item():.4f}",
-                for name, value in parts_sum.items():
+                skipped=int(totals[2].item()),
-                    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]}")
 
+    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")
+        logger.info(f"Skipped {skipped} rank-batches due to non-finite loss")
-    return float((total / max(1, n_batches)).detach().cpu()) if n_batches else float("inf")
+    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, rank, local_rank, world_size = init_distributed(args)
-    device = resolve_device(args.device)
+    set_seed(args.seed + rank)
     configure_torch_for_training(device)
 
-    run_dir, run_name = create_unique_run_dir(
+    run_dir, run_name = distributed_run_dir(args, rank, world_size)
-        lambda timestamp: (
+    logger = rank_logger(rank, run_dir)
-            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)
 
     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(
+        if world_size > 1:
-            train_subset,
+            train_batch_sampler = DistributedExposureLocalityBatchSampler(
-            batch_sampler=ExposureLocalityBatchSampler(
                 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_sampler = (
-            train_subset,
+            DistributedSampler(
-            batch_size=args.batch_size,
+                train_subset,
-            sampler=RandomSampler(
+                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)
+    backbone = build_model(args, dataset)
-    model = maybe_wrap_data_parallel(model, args, device, logger)
     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(
+    if rank == 0:
-        args,
+        save_config(
-        run_dir / "train_config.json",
+            args,
-        extra=build_metadata(args, dataset, run_name, train_subset, val_subset, test_subset),
+            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:
+    if rank == 0:
-        json.dump(history, f, indent=2)
+        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__":