import torch import torch.nn as nn from torch.optim import Adam from torch.utils.data import DataLoader, DistributedSampler import torch.distributed as dist import torch.multiprocessing as mp from torch.nn.parallel import DistributedDataParallel as DDP import numpy as np import math import tqdm import matplotlib.pyplot as plt import os import time from models import TimeAwareGPT2, CombinedLoss from utils import PatientEventDataset # --- Configuration --- class TrainConfig: # Data parameters train_data_path = 'ukb_real_train.bin' val_data_path = 'ukb_real_val.bin' block_length = 24 # Sequence length # Model parameters n_embd = 256 n_layer = 8 n_head = 8 pdrop = 0.1 token_pdrop = 0.1 # Training parameters max_epoch = 200 batch_size = 512 # 增大总批次大小 lr_initial = 6e-4 lr_final = 6e-5 warmup_epochs = 10 early_stopping_patience = 5 # Loss parameters ignored_token_ids = [0, 1] # Distributed training parameters world_size = torch.cuda.device_count() distributed = world_size > 1 # --- Main Training Function --- def train_worker(local_rank, config): # Initialize distributed training if config.distributed: dist.init_process_group( backend='nccl', init_method='env://', rank=local_rank, world_size=config.world_size ) torch.cuda.set_device(local_rank) device = torch.device('cuda', local_rank) print(f"Worker {local_rank} initialized on device {device}") else: device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') local_rank = 0 # --- 1. Data Loading --- if local_rank == 0: print(f"Loading data from {config.train_data_path} and {config.val_data_path}...") train_data_arr = np.memmap(config.train_data_path, dtype=np.uint32, mode='r').reshape(-1, 3) val_data_arr = np.memmap(config.val_data_path, dtype=np.uint32, mode='r').reshape(-1, 3) vocab_size = int(max(train_data_arr[:, 2].max(), val_data_arr[:, 2].max())) + 1 if local_rank == 0: print(f"Inferred vocabulary size: {vocab_size}") print(f"Using {config.world_size} GPU(s) for training") train_dataset = PatientEventDataset(train_data_arr, config.block_length) val_dataset = PatientEventDataset(val_data_arr, config.block_length) # 计算每个GPU的批次大小 per_gpu_batch_size = config.batch_size // config.world_size # 优化数据加载器参数 if config.distributed: train_sampler = DistributedSampler(train_dataset, num_replicas=config.world_size, rank=local_rank, shuffle=True) val_sampler = DistributedSampler(val_dataset, num_replicas=config.world_size, rank=local_rank, shuffle=False) else: train_sampler = None val_sampler = None # 增加num_workers，使用persistent_workers减少进程创建开销 train_loader = DataLoader( train_dataset, batch_size=per_gpu_batch_size, # 使用每个GPU的批次大小 sampler=train_sampler, shuffle=(train_sampler is None), num_workers=8, # 增加worker数量 pin_memory=True, persistent_workers=True, # 保持worker进程 prefetch_factor=2 # 预取批次 ) val_loader = DataLoader( val_dataset, batch_size=per_gpu_batch_size, sampler=val_sampler, shuffle=False, num_workers=8, pin_memory=True, persistent_workers=True, prefetch_factor=2 ) # --- 2. Model, Optimizer, and Loss Initialization --- if local_rank == 0: print(f"Initializing model on {device}...") model = TimeAwareGPT2( vocab_size=vocab_size, n_embd=config.n_embd, n_layer=config.n_layer, n_head=config.n_head, pdrop=config.pdrop, token_pdrop=config.token_pdrop ).to(device) # 使用梯度累积来模拟更大的批次大小，减少通信频率 if config.distributed: # 使用find_unused_parameters=False来加速 model = DDP(model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=False) if local_rank == 0: if config.distributed: num_params = sum(p.numel() for p in model.module.parameters() if p.requires_grad) else: num_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f"Model initialized with {num_params/1e6:.2f}M trainable parameters.") print(f"Per GPU batch size: {per_gpu_batch_size}") loss_fn = CombinedLoss(config.ignored_token_ids) optimizer = Adam(model.parameters(), lr=config.lr_initial) # --- 3. Training Loop --- best_val_loss = float('inf') patience_counter = 0 if local_rank == 0: train_losses_ce, train_losses_surv, train_losses_total = [], [], [] val_losses_ce, val_losses_surv, val_losses_total = [], [], [] if local_rank == 0: print("Starting training...") for epoch in range(config.max_epoch): if config.distributed: train_sampler.set_epoch(epoch) # --- Learning Rate Scheduling --- if epoch < config.warmup_epochs: lr = config.lr_initial else: progress = (epoch - config.warmup_epochs) / (config.max_epoch - config.warmup_epochs) lr = config.lr_final + 0.5 * (config.lr_initial - config.lr_final) * (1 + math.cos(math.pi * progress)) for param_group in optimizer.param_groups: param_group['lr'] = lr # --- Training Phase --- model.train() train_loss_ce_acc, train_loss_surv_acc = 0.0, 0.0 train_steps = 0 # 只在rank 0显示进度条 if local_rank == 0: pbar = tqdm.tqdm(train_loader, desc=f"Epoch {epoch+1}/{config.max_epoch} [Train]") else: pbar = train_loader batch_start_time = time.time() for batch_idx, (event_seq, time_seq) in enumerate(pbar): event_seq, time_seq = event_seq.to(device, non_blocking=True), time_seq.to(device, non_blocking=True) # Prepare inputs and targets input_events = event_seq[:, :-1] input_times = time_seq[:, :-1] target_events = event_seq[:, 1:] target_wait_times = (time_seq[:, 1:] - time_seq[:, :-1]).float() # Forward pass logits = model(input_events, input_times) loss_ce, loss_survival = loss_fn(logits, target_events, target_wait_times) loss = loss_ce + loss_survival # Backward pass and optimization optimizer.zero_grad() loss.backward() # 梯度同步在DDP中自动处理 optimizer.step() # 异步记录损失，避免同步阻塞 train_loss_ce_acc += loss_ce.item() train_loss_surv_acc += loss_survival.item() train_steps += 1 if local_rank == 0 and batch_idx % 10 == 0: # 每10个批次更新一次 batch_time = time.time() - batch_start_time pbar.set_postfix({ 'loss_ce': f'{loss_ce.item():.4f}', 'loss_surv': f'{loss_survival.item():.4f}', 'lr': f'{lr:.2e}', 'batch_time': f'{batch_time:.3f}s' }) batch_start_time = time.time() # 只在epoch结束时同步一次损失，减少通信 if config.distributed: # 使用all_reduce同步损失 train_loss_ce_tensor = torch.tensor([train_loss_ce_acc], device=device) train_loss_surv_tensor = torch.tensor([train_loss_surv_acc], device=device) train_steps_tensor = torch.tensor([train_steps], device=device) dist.all_reduce(train_loss_ce_tensor) dist.all_reduce(train_loss_surv_tensor) dist.all_reduce(train_steps_tensor) avg_train_loss_ce = (train_loss_ce_tensor.item() / train_steps_tensor.item()) avg_train_loss_surv = (train_loss_surv_tensor.item() / train_steps_tensor.item()) else: avg_train_loss_ce = train_loss_ce_acc / train_steps avg_train_loss_surv = train_loss_surv_acc / train_steps # --- Validation Phase --- model.eval() val_loss_ce_acc, val_loss_surv_acc = 0.0, 0.0 val_steps = 0 with torch.no_grad(): for event_seq, time_seq in val_loader: event_seq, time_seq = event_seq.to(device, non_blocking=True), time_seq.to(device, non_blocking=True) input_events = event_seq[:, :-1] input_times = time_seq[:, :-1] target_events = event_seq[:, 1:] target_wait_times = (time_seq[:, 1:] - time_seq[:, :-1]).float() logits = model(input_events, input_times) loss_ce, loss_survival = loss_fn(logits, target_events, target_wait_times) val_loss_ce_acc += loss_ce.item() val_loss_surv_acc += loss_survival.item() val_steps += 1 # 同步验证损失 if config.distributed: val_loss_ce_tensor = torch.tensor([val_loss_ce_acc], device=device) val_loss_surv_tensor = torch.tensor([val_loss_surv_acc], device=device) val_steps_tensor = torch.tensor([val_steps], device=device) dist.all_reduce(val_loss_ce_tensor) dist.all_reduce(val_loss_surv_tensor) dist.all_reduce(val_steps_tensor) avg_val_loss_ce = (val_loss_ce_tensor.item() / val_steps_tensor.item()) avg_val_loss_surv = (val_loss_surv_tensor.item() / val_steps_tensor.item()) else: avg_val_loss_ce = val_loss_ce_acc / val_steps avg_val_loss_surv = val_loss_surv_acc / val_steps total_val_loss = avg_val_loss_ce + avg_val_loss_surv # 只在rank 0进行打印和保存 if local_rank == 0: train_losses_ce.append(avg_train_loss_ce) train_losses_surv.append(avg_train_loss_surv) train_losses_total.append(avg_train_loss_ce + avg_train_loss_surv) val_losses_ce.append(avg_val_loss_ce) val_losses_surv.append(avg_val_loss_surv) val_losses_total.append(total_val_loss) print(f"Epoch {epoch+1} Summary: \n" f" Train Loss: {avg_train_loss_ce + avg_train_loss_surv:.4f} (CE: {avg_train_loss_ce:.4f}, Surv: {avg_train_loss_surv:.4f})\n" f" Val Loss: {total_val_loss:.4f} (CE: {avg_val_loss_ce:.4f}, Surv: {avg_val_loss_surv:.4f})\n" f" Learning Rate: {lr:.6f}") # Early stopping check if total_val_loss < best_val_loss: best_val_loss = total_val_loss patience_counter = 0 print(f"Validation loss improved to {best_val_loss:.4f}. Saving checkpoint...") if config.distributed: torch.save(model.module.state_dict(), 'best_model_checkpoint.pt') else: torch.save(model.state_dict(), 'best_model_checkpoint.pt') else: if epoch >= config.warmup_epochs: patience_counter += 1 print(f"Validation loss did not improve. Patience: {patience_counter}/{config.early_stopping_patience}") if patience_counter >= config.early_stopping_patience: print("\nEarly stopping triggered due to no improvement in validation loss.") if config.distributed: stop_signal = torch.tensor(1, device=device) dist.broadcast(stop_signal, 0) break else: # 非rank 0进程检查停止信号 if config.distributed: stop_signal = torch.tensor(0, device=device) dist.broadcast(stop_signal, 0) if stop_signal.item() == 1: break # 清理和保存 if local_rank == 0 and best_val_loss != float('inf'): print(f"\nTraining finished. Loading best model from checkpoint with validation loss {best_val_loss:.4f}.") if config.distributed: model.module.load_state_dict(torch.load('best_model_checkpoint.pt')) torch.save(model.module.state_dict(), 'best_model.pt') else: model.load_state_dict(torch.load('best_model_checkpoint.pt')) torch.save(model.state_dict(), 'best_model.pt') print("Final best model saved to best_model.pt") if config.distributed: dist.destroy_process_group() def main(): config = TrainConfig() # 设置环境变量优化 os.environ['CUDA_LAUNCH_BLOCKING'] = '0' # 减少同步 os.environ['NCCL_DEBUG'] = 'WARN' # 减少NCCL日志 os.environ['NCCL_SOCKET_IFNAME'] = '^lo,docker' # 选择正确的网络接口 if config.distributed: print(f"Starting distributed training with {config.world_size} GPUs") mp.spawn( train_worker, args=(config,), nprocs=config.world_size, join=True ) else: print("Starting single GPU training") train_worker(0, config) if __name__ == '__main__': main()