"""Pretrain a lightweight TimesNet autoencoder on training-set exposure.""" from __future__ import annotations import argparse import hashlib import json import logging import math import os from pathlib import Path import numpy as np import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel from torch.optim import AdamW from torch.utils.data import DataLoader, Dataset, DistributedSampler from tqdm import tqdm from backbones import TimesNetExposureAutoencoder from dataset import ExposureCache from train_util import ( configure_torch_for_training, create_unique_run_dir, load_eid_file, resolve_device, set_seed, setup_logging, ) class ExposureWindowDataset(Dataset): def __init__(self, cache: ExposureCache, row_indices: np.ndarray): self.cache = cache self.row_indices = np.asarray(row_indices, dtype=np.int64) def __len__(self) -> int: return len(self.row_indices) def __getitem__(self, index: int) -> dict[str, torch.Tensor]: row = int(self.row_indices[index]) return { "daily": torch.from_numpy( np.array(self.cache.daily[row], dtype=np.float32, copy=True) ), "monthly": torch.from_numpy( np.array(self.cache.monthly[row], dtype=np.float32, copy=True) ), } def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Pretrain a lightweight TimesNet exposure autoencoder" ) parser.add_argument("--exposure_cache_dir", default=None) parser.add_argument("--train_eid_file", default="ukb_train_eid.csv") parser.add_argument("--val_eid_file", default="ukb_val_eid.csv") parser.add_argument( "--channel_stats_file", default=None, help=( "Cached channel statistics .npz file. Defaults to " "/train_channel_stats.npz." ), ) parser.add_argument( "--recompute_channel_stats", action="store_true", help="Ignore a compatible statistics cache and recompute it.", ) parser.add_argument("--runs_root", default="runs") parser.add_argument( "--resume_checkpoint", default=None, help=( "Resume training from a run directory or checkpoint. A directory " "uses last.pt when present, otherwise best.pt." ), ) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--n_embd", type=int, default=120) parser.add_argument("--d_model", type=int, default=64) parser.add_argument("--n_layers", type=int, default=2) parser.add_argument("--top_k", type=int, default=2) parser.add_argument("--n_backbone_blocks", type=int, default=1) parser.add_argument("--backbone_kernel_size", type=int, default=5) parser.add_argument("--backbone_expansion", type=float, default=2.0) parser.add_argument("--dropout", type=float, default=0.0) parser.add_argument("--mask_ratio", type=float, default=0.25) parser.add_argument("--batch_size", type=int, default=16) parser.add_argument("--base_lr", type=float, default=3e-4) parser.add_argument("--weight_decay", type=float, default=0.05) parser.add_argument("--max_epochs", type=int, default=100) parser.add_argument("--warmup_epochs", type=int, default=5) parser.add_argument("--patience", type=int, default=12) parser.add_argument("--grad_clip", type=float, default=1.0) parser.add_argument("--num_workers", type=int, default=4) parser.add_argument("--device", default="cuda") parser.add_argument("--amp", action=argparse.BooleanOptionalAction, default=True) parser.add_argument( "--data_parallel", action="store_true", help="Use torch.nn.DataParallel across multiple CUDA devices.", ) parser.add_argument( "--gpu_ids", 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 on CUDA and gloo otherwise.", ) parser.add_argument( "--prefetch_factor", type=int, default=4, help="DataLoader batches prefetched by each worker.", ) return parser.parse_args() def validate_args(args: argparse.Namespace) -> None: if not args.exposure_cache_dir: raise ValueError( "--exposure_cache_dir is required unless loaded from resume config" ) if not 0.0 <= args.mask_ratio < 1.0: raise ValueError("--mask_ratio must be in [0, 1)") if args.num_workers > 0 and args.prefetch_factor <= 0: raise ValueError("--prefetch_factor must be positive") if isinstance(args.gpu_ids, str) and 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 def resolve_resume_checkpoint(resume_path: str | None) -> Path | None: if not resume_path: return None path = Path(resume_path) if path.is_dir(): last_path = path / "last.pt" if last_path.is_file(): return last_path best_path = path / "best.pt" if best_path.is_file(): return best_path raise FileNotFoundError( f"Resume run directory has neither last.pt nor best.pt: {path}" ) if not path.is_file(): raise FileNotFoundError(f"--resume_checkpoint does not exist: {path}") return path def apply_resume_config(args: argparse.Namespace, resume_checkpoint: Path) -> None: config_path = resume_checkpoint.parent / "train_config.json" if not config_path.is_file(): raise FileNotFoundError( f"Resume requires train_config.json next to checkpoint: {config_path}" ) config = json.loads(config_path.read_text(encoding="utf-8")) resume_value = str(resume_checkpoint) for key, value in config.items(): if hasattr(args, key): setattr(args, key, value) args.resume_checkpoint = resume_value def select_rows(cache: ExposureCache, eids: set[int], split: str) -> np.ndarray: valid_row = np.asarray(cache.row_index, dtype=np.int64) >= 0 selected_events = valid_row & np.isin(cache.eids, np.fromiter(eids, np.int64)) rows = np.unique( np.asarray(cache.row_index[selected_events], dtype=np.int64) ) if len(rows) == 0: raise ValueError(f"{split} exposure rows are empty after EID filtering") return rows def maybe_wrap_data_parallel( model: TimesNetExposureAutoencoder, args: argparse.Namespace, device: torch.device, 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 = ( int(device.index) if device.index is not None else int(torch.cuda.current_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") if any(idx < 0 or idx >= torch.cuda.device_count() for idx in device_ids): raise ValueError(f"CUDA device id is out of range: {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 unwrap_model(model) -> TimesNetExposureAutoencoder: if isinstance(model, (torch.nn.DataParallel, DistributedDataParallel)): return model.module return model 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") local_rank = int(os.environ["LOCAL_RANK"]) rank = int(os.environ["RANK"]) if not torch.cuda.is_available(): raise ValueError("Multi-process exposure training requires CUDA") torch.cuda.set_device(local_rank) backend = args.ddp_backend or "nccl" dist.init_process_group(backend=backend, init_method="env://") return torch.device("cuda", local_rank), rank, local_rank, world_size def rank_logger(rank: int, run_dir: Path): 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 distributed_run_dir( args: argparse.Namespace, rank: int, world_size: int ) -> tuple[Path, str]: payload: list[str | None] = [None, None] if rank == 0: if args.resume_checkpoint: resume_path = Path(args.resume_checkpoint) run_dir = resume_path.parent run_name = run_dir.name else: run_dir, run_name = create_unique_run_dir( lambda stamp: f"exposure_ae_{stamp}", Path(args.runs_root) ) 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 channel_stats( cache: ExposureCache, rows: np.ndarray, chunk_size: int = 256 ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: results = [] for source in (cache.daily, cache.monthly): sums = np.zeros(source.shape[-1], dtype=np.float64) squares = np.zeros_like(sums) counts = np.zeros_like(sums) for start in tqdm(range(0, len(rows), chunk_size), desc="Channel statistics"): values = np.asarray(source[rows[start:start + chunk_size]], dtype=np.float64) finite = np.isfinite(values) clean = np.where(finite, values, 0.0) sums += clean.sum(axis=(0, 1)) squares += np.square(clean).sum(axis=(0, 1)) counts += finite.sum(axis=(0, 1)) mean = sums / np.maximum(counts, 1.0) variance = squares / np.maximum(counts, 1.0) - np.square(mean) std = np.sqrt(np.maximum(variance, 1e-12)) results.extend([mean.astype(np.float32), std.astype(np.float32)]) return tuple(results) def eid_set_hash(eids: set[int]) -> str: digest = hashlib.sha256() for eid in sorted(eids): digest.update(f"{eid}\n".encode("ascii")) return digest.hexdigest() def load_or_compute_channel_stats( cache: ExposureCache, rows: np.ndarray, train_eids: set[int], stats_path: Path, recompute: bool, logger, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: eid_hash = eid_set_hash(train_eids) if stats_path.is_file() and not recompute: try: with np.load(stats_path, allow_pickle=False) as saved: compatible = ( str(saved["train_eid_sha256"].item()) == eid_hash and int(saved["cache_event_rows"].item()) == len(cache.eids) and int(saved["train_window_rows"].item()) == len(rows) ) if compatible: logger.info(f"Loading channel statistics from {stats_path}") return ( saved["daily_mean"].astype(np.float32), saved["daily_std"].astype(np.float32), saved["monthly_mean"].astype(np.float32), saved["monthly_std"].astype(np.float32), ) logger.info("Channel statistics cache is stale; recomputing") except (KeyError, OSError, ValueError) as exc: logger.warning( f"Could not read channel statistics cache ({exc}); recomputing" ) logger.info("Computing channel statistics from training exposure") stats = channel_stats(cache, rows) stats_path.parent.mkdir(parents=True, exist_ok=True) np.savez( stats_path, daily_mean=stats[0], daily_std=stats[1], monthly_mean=stats[2], monthly_std=stats[3], train_eid_sha256=np.asarray(eid_hash), cache_event_rows=np.asarray(len(cache.eids), dtype=np.int64), train_window_rows=np.asarray(len(rows), dtype=np.int64), ) logger.info(f"Saved channel statistics to {stats_path}") return stats def masked_mse( prediction: torch.Tensor, target: torch.Tensor, mask: torch.Tensor ) -> torch.Tensor: error = (prediction - target).square() * mask return error.sum() / mask.sum().clamp_min(1.0) def run_epoch( model, loader: DataLoader, device: torch.device, stats: tuple[torch.Tensor, ...], mask_ratio: float, optimizer: AdamW | None, scaler: torch.amp.GradScaler, grad_clip: float, amp_enabled: bool, show_progress: bool, ) -> float: training = optimizer is not None model.train(training) loss_accumulator = torch.zeros(2, device=device, dtype=torch.float64) daily_mean, daily_std, monthly_mean, monthly_std = stats context = torch.enable_grad if training else torch.no_grad with context(): for batch in tqdm( loader, desc="train" if training else "val", leave=False, disable=not show_progress, ): daily = batch["daily"].to(device, non_blocking=True) monthly = batch["monthly"].to(device, non_blocking=True) daily_observed = torch.isfinite(daily) monthly_observed = torch.isfinite(monthly) daily = (torch.nan_to_num(daily) - daily_mean) / daily_std monthly = (torch.nan_to_num(monthly) - monthly_mean) / monthly_std daily = daily * daily_observed monthly = monthly * monthly_observed if training and mask_ratio > 0: daily_input_mask = daily_observed & ( torch.rand_like(daily) >= mask_ratio ) monthly_input_mask = monthly_observed & ( torch.rand_like(monthly) >= mask_ratio ) else: daily_input_mask = daily_observed monthly_input_mask = monthly_observed daily_input = daily * daily_input_mask monthly_input = monthly * monthly_input_mask if training: optimizer.zero_grad(set_to_none=True) with torch.autocast( device_type=device.type, dtype=torch.float16, enabled=amp_enabled, ): daily_hat, monthly_hat, _ = model( daily_input, monthly_input, daily_input_mask, monthly_input_mask, ) loss = ( masked_mse(daily_hat, daily, daily_observed) + masked_mse(monthly_hat, monthly, monthly_observed) ) if training: scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip) scaler.step(optimizer) scaler.update() batch_size = daily.size(0) loss_accumulator[0] += loss.detach().double() * batch_size loss_accumulator[1] += batch_size if dist.is_initialized(): dist.all_reduce(loss_accumulator, op=dist.ReduceOp.SUM) return float((loss_accumulator[0] / loss_accumulator[1].clamp_min(1)).item()) def learning_rate(epoch: int, args: argparse.Namespace) -> float: if epoch < args.warmup_epochs: return args.base_lr * (epoch + 1) / max(args.warmup_epochs, 1) progress = (epoch - args.warmup_epochs) / max( args.max_epochs - args.warmup_epochs - 1, 1 ) return args.base_lr * 0.5 * (1.0 + math.cos(math.pi * progress)) def autoencoder_checkpoint_payload( model, optimizer: AdamW, scaler: torch.amp.GradScaler, config: dict, raw_stats: tuple[np.ndarray, ...], epoch: int, val_loss: float, best_loss: float, stale_epochs: int, history: list[dict], include_training_state: bool, ) -> dict: checkpoint_model = unwrap_model(model) payload = { "model_state_dict": checkpoint_model.state_dict(), "encoder_state_dict": checkpoint_model.encoder.state_dict(), "model_config": { key: config[key] for key in ( "n_embd", "d_model", "n_layers", "top_k", "n_backbone_blocks", "backbone_kernel_size", "backbone_expansion", "dropout", ) }, "normalization": { "daily_mean": raw_stats[0], "daily_std": raw_stats[1], "monthly_mean": raw_stats[2], "monthly_std": raw_stats[3], }, "epoch": epoch, "val_loss": val_loss, "best_loss": best_loss, "stale_epochs": stale_epochs, "history": history, } if include_training_state: payload["optimizer_state_dict"] = optimizer.state_dict() payload["scaler_state_dict"] = scaler.state_dict() return payload def save_autoencoder_checkpoint(payload: dict, checkpoint_path: Path) -> None: tmp_path = checkpoint_path.with_suffix(checkpoint_path.suffix + ".tmp") torch.save(payload, tmp_path) tmp_path.replace(checkpoint_path) def torch_load_checkpoint(checkpoint_path: Path, map_location): try: return torch.load( checkpoint_path, map_location=map_location, weights_only=False ) except TypeError: return torch.load(checkpoint_path, map_location=map_location) def load_resume_checkpoint( checkpoint_path: Path, model, optimizer: AdamW, scaler: torch.amp.GradScaler, val_loader: DataLoader, stats: tuple[torch.Tensor, ...], device: torch.device, grad_clip: float, amp_enabled: bool, show_progress: bool, logger, ) -> tuple[int, float, int, list[dict]]: checkpoint = torch_load_checkpoint(checkpoint_path, map_location=device) if "model_state_dict" not in checkpoint: raise KeyError( f"Checkpoint does not contain model_state_dict: {checkpoint_path}" ) unwrap_model(model).load_state_dict(checkpoint["model_state_dict"]) has_training_state = "optimizer_state_dict" in checkpoint if has_training_state: optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) else: logger.warning( "Resume checkpoint has no optimizer state; continuing with a fresh " "optimizer" ) if "scaler_state_dict" in checkpoint: scaler.load_state_dict(checkpoint["scaler_state_dict"]) elif scaler.is_enabled(): logger.warning( "Resume checkpoint has no AMP scaler state; continuing with a fresh " "scaler" ) start_epoch = int(checkpoint.get("epoch", 0)) history = list(checkpoint.get("history", [])) history_path = checkpoint_path.parent / "history.json" if not history and history_path.is_file(): history = json.loads(history_path.read_text(encoding="utf-8")) if has_training_state: best_loss = float( checkpoint.get("best_loss", checkpoint.get("val_loss", float("inf"))) ) stale_epochs = int(checkpoint.get("stale_epochs", 0)) else: current_val_loss = run_epoch( model, val_loader, device, stats, 0.0, None, scaler, grad_clip, amp_enabled, show_progress, ) history_best_epoch = start_epoch historical_best = float(checkpoint.get("val_loss", float("inf"))) for entry in history: if "val_loss" not in entry: continue entry_loss = float(entry["val_loss"]) if entry_loss < historical_best: historical_best = entry_loss history_best_epoch = int(entry.get("epoch", len(history))) if math.isfinite(historical_best): tolerance = max(1e-4, abs(historical_best) * 1e-3) if abs(current_val_loss - historical_best) > tolerance: logger.warning( "Legacy best.pt validation loss differs from history: " f"recomputed={current_val_loss:.6f}, " f"history_best={historical_best:.6f}" ) best_loss = min(historical_best, current_val_loss) if history: start_epoch = max(start_epoch, int(history[-1].get("epoch", len(history)))) if current_val_loss < historical_best: stale_epochs = 0 else: stale_epochs = max(0, start_epoch - history_best_epoch) logger.info( f"Validated legacy checkpoint {checkpoint_path.name}: " f"val={current_val_loss:.6f}, historical_best={historical_best:.6f}" ) logger.info( f"Resumed from {checkpoint_path} at epoch {start_epoch}; " f"best_val={best_loss:.6f}, stale_epochs={stale_epochs}" ) return start_epoch, best_loss, stale_epochs, history def run_training( args: argparse.Namespace, device: torch.device, rank: int, local_rank: int, world_size: int, ) -> None: 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) cache = ExposureCache(args.exposure_cache_dir) train_eids = load_eid_file(args.train_eid_file) val_eids = load_eid_file(args.val_eid_file) if train_eids & val_eids: raise ValueError("train and validation EID files must be disjoint") train_rows = select_rows(cache, train_eids, "Training") val_rows = select_rows(cache, val_eids, "Validation") stats_path = ( Path(args.channel_stats_file) if args.channel_stats_file else Path(args.exposure_cache_dir) / "train_channel_stats.npz" ) if rank == 0: raw_stats = load_or_compute_channel_stats( cache, train_rows, train_eids, stats_path, args.recompute_channel_stats, logger, ) if world_size > 1: dist.barrier() if rank != 0: raw_stats = load_or_compute_channel_stats( cache, train_rows, train_eids, stats_path, False, logger ) stats = tuple( torch.as_tensor(value, device=device).view(1, 1, -1) for value in raw_stats ) 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 loader_kwargs = dict( batch_size=local_batch_size, num_workers=args.num_workers, pin_memory=device.type == "cuda", persistent_workers=args.num_workers > 0, ) if args.num_workers > 0: loader_kwargs["prefetch_factor"] = args.prefetch_factor train_dataset = ExposureWindowDataset(cache, train_rows) val_dataset = ExposureWindowDataset(cache, val_rows) train_sampler = ( DistributedSampler( train_dataset, num_replicas=world_size, rank=rank, shuffle=True, seed=args.seed, ) if world_size > 1 else None ) val_sampler = ( DistributedSampler( val_dataset, num_replicas=world_size, rank=rank, shuffle=False ) if world_size > 1 else None ) train_loader = DataLoader( train_dataset, sampler=train_sampler, shuffle=train_sampler is None, **loader_kwargs ) val_loader = DataLoader( val_dataset, sampler=val_sampler, shuffle=False, **loader_kwargs ) model = TimesNetExposureAutoencoder( n_embd=args.n_embd, d_model=args.d_model, n_layers=args.n_layers, top_k=args.top_k, n_backbone_blocks=args.n_backbone_blocks, backbone_kernel_size=args.backbone_kernel_size, backbone_expansion=args.backbone_expansion, dropout=args.dropout, ).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) optimizer = AdamW( model.parameters(), lr=args.base_lr, weight_decay=args.weight_decay, betas=(0.9, 0.95), ) amp_enabled = bool(args.amp and device.type == "cuda") scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled) logger.info( f"Run {run_name}: device={device}, train_rows={len(train_rows):,}, " f"val_rows={len(val_rows):,}" ) config = vars(args) | { "train_rows": len(train_rows), "val_rows": len(val_rows), "daily_mean": raw_stats[0].tolist(), "daily_std": raw_stats[1].tolist(), "monthly_mean": raw_stats[2].tolist(), "monthly_std": raw_stats[3].tolist(), } if rank == 0 and not args.resume_checkpoint: (run_dir / "train_config.json").write_text( json.dumps(config, indent=2), encoding="utf-8" ) best_loss = float("inf") stale_epochs = 0 history = [] start_epoch = 0 if args.resume_checkpoint: start_epoch, best_loss, stale_epochs, history = load_resume_checkpoint( Path(args.resume_checkpoint), model, optimizer, scaler, val_loader, stats, device, args.grad_clip, amp_enabled, rank == 0, logger, ) if start_epoch >= args.max_epochs: logger.info( f"Resume checkpoint is already at epoch {start_epoch}; " f"--max_epochs={args.max_epochs} leaves no remaining epochs" ) for epoch in range(start_epoch, args.max_epochs): if train_sampler is not None: train_sampler.set_epoch(epoch) lr = learning_rate(epoch, args) for group in optimizer.param_groups: group["lr"] = lr train_loss = run_epoch( model, train_loader, device, stats, args.mask_ratio, optimizer, scaler, args.grad_clip, amp_enabled, rank == 0, ) val_loss = run_epoch( model, val_loader, device, stats, 0.0, None, scaler, args.grad_clip, amp_enabled, rank == 0, ) logger.info( f"Epoch {epoch + 1:03d} | lr={lr:.3e} | " f"train={train_loss:.6f} | val={val_loss:.6f}" ) history.append( {"epoch": epoch + 1, "lr": lr, "train_loss": train_loss, "val_loss": val_loss} ) if val_loss < best_loss: best_loss = val_loss stale_epochs = 0 if rank == 0: save_autoencoder_checkpoint( autoencoder_checkpoint_payload( model, optimizer, scaler, config, raw_stats, epoch + 1, val_loss, best_loss, stale_epochs, history, include_training_state=False, ), run_dir / "best.pt", ) else: stale_epochs += 1 if rank == 0: save_autoencoder_checkpoint( autoencoder_checkpoint_payload( model, optimizer, scaler, config, raw_stats, epoch + 1, val_loss, best_loss, stale_epochs, history, include_training_state=True, ), run_dir / "last.pt", ) (run_dir / "history.json").write_text( json.dumps(history, indent=2), encoding="utf-8" ) if stale_epochs >= args.patience: logger.info(f"Early stopping after {epoch + 1} epochs") break logger.info(f"Best validation loss: {best_loss:.6f}") logger.info(f"Checkpoint: {run_dir / 'best.pt'}") def main() -> None: args = parse_args() resume_checkpoint = resolve_resume_checkpoint(args.resume_checkpoint) if resume_checkpoint is not None: apply_resume_config(args, resume_checkpoint) validate_args(args) init_done = False try: device, rank, local_rank, world_size = init_distributed(args) init_done = dist.is_initialized() run_training(args, device, rank, local_rank, world_size) finally: if init_done and dist.is_initialized(): dist.destroy_process_group() if __name__ == "__main__": main()