import os import sys import json import time from time import gmtime, strftime import torch.distributed as dist import torch import torch.nn.functional as F from torch.nn.parallel import DistributedDataParallel as DDP import comet_ml # Ensure project root is in path sys.path.append('your_work_path') from config import Config from dolphindb_klines import DolphinDBStreamingDataset from model.kronos import KronosTokenizer # Import shared utilities from utils.training_utils import ( setup_ddp, cleanup_ddp, set_seed, get_model_size, format_time, ) def create_dataloaders(config: dict, rank: int): """ 创建ddbdataloader实例，并返回训练和验证的dataloader。 """ print(f"[Rank {rank}] Creating distributed dataloaders...") train_loader = DolphinDBStreamingDataset('train', config).ddb_loader val_loader = DolphinDBStreamingDataset('val', config).ddb_loader print(f"[Rank {rank}] Dataloaders created.") return train_loader, val_loader def train_model(model, device, config, save_dir, logger, rank, world_size): """ 主训练循环，包含训练和验证逻辑，以及日志记录和模型保存。 """ start_time = time.time() #steps_per根据自己的数据量来进行估计，或者设置为一个较大的数值，确保每个epoch都能完整地迭代训练数据。 steps_per = 6600 if rank == 0: effective_bs = config['batch_size'] * world_size * config['accumulation_steps'] print(f"[Rank {rank}] BATCHSIZE (per GPU): {config['batch_size']}") print(f"[Rank {rank}] Effective total batch size: {effective_bs}") train_loader, val_loader = create_dataloaders(config, rank) optimizer = torch.optim.AdamW( model.parameters(), lr=config['tokenizer_learning_rate'], weight_decay=config['adam_weight_decay'] ) scheduler = torch.optim.lr_scheduler.OneCycleLR( optimizer=optimizer, max_lr=config['tokenizer_learning_rate'], steps_per_epoch=steps_per, epochs=config['epochs'], pct_start=0.03, div_factor=10 ) best_val_loss = float('inf') dt_result = {} batch_idx_global_train = 0 for epoch_idx in range(config['epochs']): epoch_start_time = time.time() model.train() for i, (ori_batch_x, _) in enumerate(train_loader): ori_batch_x = torch.nan_to_num(ori_batch_x, nan=0.0).float() lookback = config['lookback_window'] past = ori_batch_x[:, :lookback, :] mean = past.mean(dim=1, keepdim=True) std = past.std(dim=1, keepdim=True) + 1e-5 ori_batch_x = (ori_batch_x - mean) / std ori_batch_x = torch.clamp(ori_batch_x, -config['clip'], config['clip']) ori_batch_x = ori_batch_x.to(device, non_blocking=True) # --- Gradient Accumulation Loop --- current_batch_total_loss = 0.0 for j in range(config['accumulation_steps']): start_idx = j * (ori_batch_x.shape[0] // config['accumulation_steps']) end_idx = (j + 1) * (ori_batch_x.shape[0] // config['accumulation_steps']) batch_x = ori_batch_x[start_idx:end_idx] # Forward pass zs, bsq_loss, _, _ = model(batch_x) z_pre, z = zs # Loss calculation recon_loss_pre = F.mse_loss(z_pre, batch_x) recon_loss_all = F.mse_loss(z, batch_x) recon_loss = recon_loss_pre + recon_loss_all loss = (recon_loss + bsq_loss) / 2 # Assuming w_1=w_2=1 loss_scaled = loss / config['accumulation_steps'] current_batch_total_loss += loss.item() loss_scaled.backward() # --- Optimizer Step after Accumulation --- torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=2.0) optimizer.step() scheduler.step() optimizer.zero_grad() # --- Logging (Master Process Only) --- if rank == 0 and (batch_idx_global_train + 1) % config['log_interval'] == 0: avg_loss = current_batch_total_loss / config['accumulation_steps'] print( f"[Rank {rank}, Epoch {epoch_idx + 1}/{config['epochs']}, Step {i + 1}/{steps_per}] " f"LR {optimizer.param_groups[0]['lr']:.6f}, Loss: {avg_loss:.4f}" ) if rank == 0 and logger: avg_loss = current_batch_total_loss / config['accumulation_steps'] logger.log_metric('train_tokenizer_loss_batch', avg_loss, step=batch_idx_global_train) logger.log_metric(f'train_vqvae_vq_loss_each_batch', bsq_loss.item(), step=batch_idx_global_train) logger.log_metric(f'train_recon_loss_pre_each_batch', recon_loss_pre.item(), step=batch_idx_global_train) logger.log_metric(f'train_recon_loss_each_batch', recon_loss_all.item(), step=batch_idx_global_train) logger.log_metric('tokenizer_learning_rate', optimizer.param_groups[0]["lr"], step=batch_idx_global_train) batch_idx_global_train += 1 # --- Validation Loop --- model.eval() tot_val_loss_sum_rank = 0.0 val_sample_count_rank = 0 with torch.no_grad(): for ori_batch_x, _ in val_loader: ori_batch_x = torch.nan_to_num(ori_batch_x, nan=0.0).float() lookback = config['lookback_window'] past = ori_batch_x[:, :lookback, :] mean = past.mean(dim=1, keepdim=True) std = past.std(dim=1, keepdim=True) + 1e-5 ori_batch_x = (ori_batch_x - mean) / std ori_batch_x = torch.clamp(ori_batch_x, -config['clip'], config['clip']) ori_batch_x = ori_batch_x.to(device, non_blocking=True) zs, _, _, _ = model(ori_batch_x) _, z = zs val_loss_item = F.mse_loss(z, ori_batch_x) tot_val_loss_sum_rank += val_loss_item.item() * ori_batch_x.size(0) val_sample_count_rank += ori_batch_x.size(0) # Reduce validation losses from all processes val_loss_sum_tensor = torch.tensor(tot_val_loss_sum_rank, device=device) val_count_tensor = torch.tensor(val_sample_count_rank, device=device) dist.all_reduce(val_loss_sum_tensor, op=dist.ReduceOp.SUM) dist.all_reduce(val_count_tensor, op=dist.ReduceOp.SUM) avg_val_loss = val_loss_sum_tensor.item() / val_count_tensor.item() if val_count_tensor.item() > 0 else 0 # --- End of Epoch Summary & Checkpointing (Master Process Only) --- if rank == 0: print(f"\n--- Epoch {epoch_idx + 1}/{config['epochs']} Summary ---") print(f"Validation Loss: {avg_val_loss:.4f}") print(f"Time This Epoch: {format_time(time.time() - epoch_start_time)}") print(f"Total Time Elapsed: {format_time(time.time() - start_time)}\n") if logger: logger.log_metric('val_tokenizer_loss_epoch', avg_val_loss, epoch=epoch_idx) if avg_val_loss < best_val_loss: best_val_loss = avg_val_loss save_path = f"{save_dir}/checkpoints/best_model" model.module.save_pretrained(save_path) print(f"Best model saved to {save_path} (Val Loss: {best_val_loss:.4f})") if logger: logger.log_model("best_model", save_path) dist.barrier() # Ensure all processes finish the epoch before starting the next one. dt_result['best_val_loss'] = best_val_loss return model, dt_result def main(config: dict): """ Main function to orchestrate the DDP training process. """ rank, world_size, local_rank = setup_ddp() device = torch.device(f"cuda:{local_rank}") set_seed(config['seed'], rank) save_dir = os.path.join(config['save_path'], config['tokenizer_save_folder_name']) print(save_dir) # 日志输出（仅主进程），如果要使用的话需要按照项目教程配置comet_ml的相关参数 comet_logger, master_summary = None, {} if rank == 0: os.makedirs(os.path.join(save_dir, 'checkpoints'), exist_ok=True) master_summary = { 'start_time': strftime("%Y-%m-%dT%H-%M-%S", gmtime()), 'save_directory': save_dir, 'world_size': world_size, } if config['use_comet']: comet_logger = comet_ml.Experiment( api_key=config['comet_config']['api_key'], project_name=config['comet_config']['project_name'], workspace=config['comet_config']['workspace'], ) comet_logger.add_tag(config['comet_tag']) comet_logger.set_name(config['comet_name']) comet_logger.log_parameters(config) print("Comet Logger Initialized.") dist.barrier() # Ensure save directory is created before proceeding # Model Initialization model = KronosTokenizer.from_pretrained(config['pretrained_tokenizer_path']) model.to(device) model = DDP(model, device_ids=[local_rank], find_unused_parameters=False) if rank == 0: print(f"Model Size: {get_model_size(model.module)}") # Start Training _, dt_result = train_model( model, device, config, save_dir, comet_logger, rank, world_size ) # Finalize and save summary (master process only) if rank == 0: master_summary['final_result'] = dt_result with open(os.path.join(save_dir, 'summary.json'), 'w') as f: json.dump(master_summary, f, indent=4) print('Training finished. Summary file saved.') if comet_logger: comet_logger.end() cleanup_ddp() if __name__ == '__main__': # Usage: torchrun --standalone --nproc_per_node=NUM_GPUS train_tokenizer.py if "WORLD_SIZE" not in os.environ: raise RuntimeError("This script must be launched with `torchrun`.") config_instance = Config() main(config_instance.__dict__)