这是报错的日志： Launch params (512, 1, 1) are larger than launch bounds (256) for kernel _Z25multi_tensor_apply_kernelI18TensorListMetadataILi4EE11AdamFunctorIfiEJffffff10adamMode_tfEEvlPViT_T0_DpT1_ please add __launch_bounds__ to kernel define or use --gpu-max-threads-per-block recompile program ! Launch params (512, 1, 1) are larger than launch bounds (256) for kernel _Z25multi_tensor_apply_kernelI18TensorListMetadataILi4EE11AdamFunctorIfiEJffffff10adamMode_tfEEvlPViT_T0_DpT1_ please add __launch_bounds__ to kernel define or use --gpu-max-threads-per-block recompile program ! Launch params (512, 1, 1) are larger than launch bounds (256) for kernel _Z25multi_tensor_apply_kernelI18TensorListMetadataILi4EE11AdamFunctorIfiEJffffff10adamMode_tfEEvlPViT_T0_DpT1_ please add __launch_bounds__ to kernel define or use --gpu-max-threads-per-block recompile program ! Launch params (512, 1, 1) are larger than launch bounds (256) for kernel _Z25multi_tensor_apply_kernelI18TensorListMetadataILi4EE11AdamFunctorIfiEJffffff10adamMode_tfEEvlPViT_T0_DpT1_ please add __launch_bounds__ to kernel define or use --gpu-max-threads-per-block recompile program ! Launch params (512, 1, 1) are larger than launch bounds (256) for kernel _Z25multi_tensor_apply_kernelI18TensorListMetadataILi4EE11AdamFunctorIfiEJffffff10adamMode_tfEEvlPViT_T0_DpT1_ please add __launch_bounds__ to kernel define or use --gpu-max-threads-per-block recompile program ! Launch params (512, 1, 1) are larger than launch bounds (256) for kernel _Z25multi_tensor_apply_kernelI18TensorListMetadataILi4EE11AdamFunctorIfiEJffffff10adamMode_tfEEvlPViT_T0_DpT1_ please add __launch_bounds__ to kernel define or use --gpu-max-threads-per-block recompile program ! Launch params (512, 1, 1) are larger than launch bounds (256) for kernel _Z25multi_tensor_apply_kernelI18TensorListMetadataILi4EE11AdamFunctorIfiEJffffff10adamMode_tfEEvlPViT_T0_DpT1_ please add __launch_bounds__ to kernel define or use --gpu-max-threads-per-block recompile program ! Launch params (512, 1, 1) are larger than launch bounds (256) for kernel _Z25multi_tensor_apply_kernelI18TensorListMetadataILi4EE11AdamFunctorIfiEJffffff10adamMode_tfEEvlPViT_T0_DpT1_ please add __launch_bounds__ to kernel define or use --gpu-max-threads-per-block recompile program ! 训练代码如下： import os import csv import time import torch import warnings import deepspeed import argparse from torch.nn import functional as F from server_module import StarForge from server_dataloader import pretraining_loader os.environ['MAX JOBS'] = '4' os.environ['NCCL_DEBUG'] = 'INFO' os.environ['NCCL_TIMEOUT'] = '3600' os.environ['NCCL_P2P_DISABLE'] = '1' os.environ['CUDA_LAUNCH_BLOCKING'] = '0' os.environ['HSA_FORCE_FINE_GRAIN_PCIE'] = '1' warnings.filterwarnings('ignore') ds_config_k100_64g = { 'train_batch_size': 256, 'train_micro_batch_size_per_gpu': 4, 'gradient_accumulation_steps': 8, 'zero_optimization': { 'stage': 2, 'allgather_partitions': True, 'allgather_bucket_size': 2e8, 'overlap_comm': True, 'reduce_scatter': True, 'reduce_bucket_size': 2e8, 'contiguous_gradients': True # 'offload_optimizer': { # 'device': None, # 'pin_memory': True # }, # 'contiguous_gradients': True, # 'overlap_comm': True }, 'bf16': { 'enabled': True, 'loss_scale': 0, 'loss_scale_window': 1000, 'initial_scale_power': 16, 'hysteresis': 2, 'min_loss_scale': 1 }, 'gradient_clipping': 1.0, 'optimizer': { 'type': 'AdamW', 'params': { 'lr': 2e-4, 'betas': [0.9, 0.95], 'eps': 1e-8, 'weight_decay': 0.1 } }, 'scheduler': { 'type': 'WarmupDecayLR', 'params': { 'total_num_steps': 1000000, 'warmup_min_lr': 0, 'warmup_max_lr': 1e-4, 'warmup_num_steps': 2000 } }, 'activation_checkpointing': { 'partition_activations': True, 'contiguous_memory_optimization': False, 'cpu_checkpointing': False }, 'steps_per_print': 2000, 'wall_clock_breakdown': False } def parse_args(): parser = argparse.ArgumentParser(description='LLM Training') parser.add_argument('--local_rank', type=int, default=-1, help='DeepSpeed local rank') parser.add_argument('--save_dir', type=str, default='./checkpoints', help='保存检查点的目录') parser.add_argument('--save_interval', type=int, default=10000, help='多少个 Global Steps 保存一次检查') parser.add_argument('--print_interval', type=int, default=100, help='多少个 Global Steps 打印一次日志') parser.add_argument('--resume_path', type=str, default=None, help='恢复训练的checkpoint路径') parser = deepspeed.add_config_arguments(parser) return parser.parse_args() def main(): ``` # 1. Init args = parse_args() deepspeed.init_distributed() if args.save_dir and not os.path.exists(args.save_dir): if torch.distributed.get_rank() == 0: os.makedirs(args.save_dir, exist_ok=True) # 2. DataLoader grad_acc_steps = ds_config_k100_64g['gradient_accumulation_steps'] train_loader, val_loader = pretraining_loader( dataset_path='./server_data_arrow', batch_size=ds_config_k100_64g['train_micro_batch_size_per_gpu'], val_size=10, num_workers=2 ) # 3. Model & DeepSpeed Config deepspeed.checkpointing.configure(None, deepspeed_config=ds_config_k100_64g) model = StarForge() model_engine, optimizer, _, scheduler = deepspeed.initialize( args=args, model=model, config=ds_config_k100_64g, model_parameters=model.parameters() ) # 4. Load Checkpoint start_epoch = 0 micro_step = 0 load_path = args.resume_path if args.resume_path else args.save_dir try: load_path, client_state = model_engine.load_checkpoint(load_path) if load_path is not None: print(f'[Rank {model_engine.local_rank}] Checkpoint loaded from {load_path}') micro_step = client_state.get('micro_step', 0) start_epoch = client_state.get('epoch', 0) else: print(f'[Rank {model_engine.local_rank}] No checkpoint found, starting from scratch.') except Exception as e: print(f'Loading checkpoint failed: {e}') # 5. Logging Init if model_engine.local_rank == 0: log_file = 'training_log_simple.csv' if micro_step == 0: with open(log_file, 'w', newline='') as f: writer = csv.writer(f) writer.writerow(['epoch', 'micro_step', 'global_step', 'loss', 'lr', 'throughput_samples_per_sec']) print(f'Start Training from Micro Step {micro_step} (Epoch {start_epoch})...') model_engine.train() # 6. Training Loop start_time = time.time() total_samples = 0 for epoch in range(start_epoch, 1): if hasattr(train_loader.sampler, 'set_epoch'): train_loader.sampler.set_epoch(epoch) # micro batch size for batch in train_loader: try: micro_step += 1 # --- A. Data Move --- input_ids = batch['input_ids'].to(model_engine.device) attention_mask = batch['attention_mask'].to(model_engine.device) labels = batch['labels'].to(model_engine.device) # 处理 Label Padding (-100) labels = torch.where(attention_mask == 1, labels, torch.tensor(-100, device=labels.device)) # 统计样本数用于计算吞吐量 current_batch_size = input_ids.size(0) total_samples += current_batch_size * torch.distributed.get_world_size() # --- B. Forward --- outputs = model_engine(input_ids) # --- C. Loss --- loss = F.cross_entropy( outputs.view(-1, outputs.size(-1)), labels.view(-1), ignore_index=-100 ) # --- D. Backward & Step (Automatic Gradient Clipping here) --- model_engine.backward(loss) model_engine.step() # --- E. Logging (Simplified) --- # 只有当完成一次梯度累积（即完成一次参数更新）时，才算一个 Global Step # 我们只在 Global Step 更新时打印，或者按照设定的 interval is_gradient_update = (micro_step % grad_acc_steps == 0) if is_gradient_update: global_step = micro_step // grad_acc_steps if model_engine.local_rank == 0 and (global_step % args.print_interval == 0): elapsed = time.time() - start_time throughput = total_samples / elapsed current_lr = optimizer.param_groups[0]['lr'] # 简单的控制台打印 print(f'[Epoch {epoch}][Micro {micro_step}][Global {global_step}] Loss: {loss.item():.4f} | LR: {current_lr:.2e} | Spd: {throughput:.1f} samples/s') # CSV 记录 with open('training_log_simple.csv', 'a', newline='') as f: writer = csv.writer(f) writer.writerow([epoch, micro_step, global_step, loss.item(), current_lr, throughput]) # 重置计时器 start_time = time.time() total_samples = 0 # --- F. Checkpointing --- # 按 Global Step 的倍数保存（例如每 1000 次参数更新保存一次） # 注意：保存操作比较耗时，确保 interval 不要太小 if args.save_interval > 0 and is_gradient_update: global_step = micro_step // grad_acc_steps if global_step % args.save_interval == 0: client_state = {'micro_step': micro_step, 'epoch': epoch} save_tag = f'global_step{global_step}' model_engine.save_checkpoint(args.save_dir, tag=save_tag, client_state=client_state) if model_engine.local_rank == 0: print(f' -> Checkpoint saved: {save_tag}') except RuntimeError as e: if 'out of memory' in str(e): print(f'OOM Error on rank {model_engine.local_rank}') torch.cuda.empty_cache() else: raise e # End of Epoch Save client_state = {'micro_step': micro_step, 'epoch': epoch + 1} model_engine.save_checkpoint(args.save_dir, tag=f'epoch_{epoch}', client_state=client_state) print('Training finished.') ``` if __name__ == '__main__': main()