DeepHealthExpo/train_next_step.py

"""
Train DeepHealth with next-token / next-time-point supervision.

The next-step dataset uses observed event histories, including CHECKUP state
tokens, plus optional gap <NO_EVENT> 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_daily",
    "exposure_monthly",
)


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_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")

    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("--exposure_cache_dir", type=str, default=None)
    parser.add_argument("--mask_onset_exposure", action="store_true")
    parser.add_argument(
        "--d_model",
        type=int,
        default=64,
        help="Internal TimesNet channel dimension for exposure encoding.",
    )
    parser.add_argument("--exposure_n_layers", type=int, default=2)
    parser.add_argument("--exposure_top_k", type=int, default=2)
    parser.add_argument("--exposure_n_backbone_blocks", type=int, default=1)
    parser.add_argument("--exposure_backbone_kernel_size", type=int, default=5)
    parser.add_argument("--exposure_backbone_expansion", type=float, default=2.0)
    parser.add_argument("--no_exposure_gate", action="store_true")
    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.d_model <= 0:
        raise ValueError("--d_model must be positive")
    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:<id>")
    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:
        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,
        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_encoder=args.exposure_cache_dir is not None,
        exposure_d_model=args.d_model,
        exposure_n_layers=args.exposure_n_layers,
        exposure_top_k=args.exposure_top_k,
        exposure_n_backbone_blocks=args.exposure_n_backbone_blocks,
        exposure_backbone_kernel_size=args.exposure_backbone_kernel_size,
        exposure_backbone_expansion=args.exposure_backbone_expansion,
        exposure_use_gate=not args.no_exposure_gate,
    )


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_daily" in batch:
        model_kwargs["exposure_daily"] = batch["exposure_daily"]
        model_kwargs["exposure_monthly"] = batch["exposure_monthly"]
    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_encoder": args.exposure_cache_dir is not None,
        "exposure_cache_dir": args.exposure_cache_dir,
        "mask_onset_exposure": bool(args.mask_onset_exposure),
        "d_model": int(args.d_model),
        "exposure_n_layers": int(args.exposure_n_layers),
        "exposure_top_k": int(args.exposure_top_k),
        "exposure_n_backbone_blocks": int(args.exposure_n_backbone_blocks),
        "exposure_backbone_kernel_size": int(args.exposure_backbone_kernel_size),
        "exposure_backbone_expansion": float(args.exposure_backbone_expansion),
        "exposure_use_gate": not bool(args.no_exposure_gate),
        "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()
    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(
        "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,
        mask_onset_exposure=args.mask_onset_exposure,
    )
    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()