diff --git a/eval_llm.py b/eval_llm.py index 9359ca8..54a6dbc 100755 --- a/eval_llm.py +++ b/eval_llm.py @@ -6,7 +6,7 @@ import torch from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from model.model_lora import * -from trainer.trainer_utils import setup_seed, get_model_params +from trainer.trainer_utils import setup_seed, get_model_params, get_default_device warnings.filterwarnings('ignore') def init_model(args): @@ -45,7 +45,7 @@ def main(): parser.add_argument('--open_thinking', default=0, type=int, help="是否开启自适应思考（0=否，1=是）") parser.add_argument('--historys', default=0, type=int, help="携带历史对话轮数（需为偶数，0表示不携带历史）") parser.add_argument('--show_speed', default=1, type=int, help="显示decode速度（tokens/s）") - parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备") + parser.add_argument('--device', default=get_default_device(), type=str, help="运行设备 (cuda:0/mps/cpu)") args = parser.parse_args() prompts = [ @@ -91,4 +91,4 @@ def main(): print(f'\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s\n\n') if args.show_speed else print('\n\n') if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/trainer/train_agent.py b/trainer/train_agent.py index 445101c..be7b9ea 100644 --- a/trainer/train_agent.py +++ b/trainer/train_agent.py @@ -16,7 +16,6 @@ import warnings import torch import torch.nn.functional as F import torch.distributed as dist -from contextlib import nullcontext from torch import optim from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, DistributedSampler @@ -24,7 +23,11 @@ from torch.optim.lr_scheduler import CosineAnnealingLR from transformers import AutoTokenizer from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import AgentRLDataset -from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel +from trainer.trainer_utils import (Logger, is_main_process, lm_checkpoint, + init_distributed_mode, setup_seed, + SkipBatchSampler, init_model, LMForRewardModel, + get_default_device, get_autocast_context, + get_grad_scaler, should_pin_memory) from trainer.rollout_engine import create_rollout_engine, compute_per_token_logps warnings.filterwarnings('ignore') @@ -330,11 +333,13 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model policy_loss = (((per_token_loss * completion_mask).sum(dim=1)[valid_rows] / token_counts[valid_rows].clamp(min=1)).mean() if valid_rows.any() else per_token_loss.sum() * 0.0) loss = (policy_loss + aux_loss) / args.accumulation_steps - loss.backward() + scaler.scale(loss).backward() if step % args.accumulation_steps == 0: - if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) - optimizer.step(); scheduler.step(); optimizer.zero_grad() + if args.grad_clip > 0: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) + scaler.step(optimizer); scaler.update(); scheduler.step(); optimizer.zero_grad() if step % args.log_interval == 0 or step == iters: pl = loss.item() * args.accumulation_steps @@ -357,7 +362,8 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model state_dict = raw_model.state_dict() torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp) lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, - epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scheduler=scheduler) + epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', + scheduler=scheduler, scaler=scaler) model.train() del state_dict @@ -367,8 +373,10 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model del completions, rewards, grouped_rewards, mean_r, std_r, advantages, completion_mask if last_step > start_step and last_step % args.accumulation_steps != 0: - if args.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) - optimizer.step(); scheduler.step(); optimizer.zero_grad() + if args.grad_clip > 0: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) + scaler.step(optimizer); scaler.update(); scheduler.step(); optimizer.zero_grad() if __name__ == "__main__": @@ -378,8 +386,8 @@ if __name__ == "__main__": parser.add_argument("--epochs", type=int, default=1, help="训练轮数") parser.add_argument("--batch_size", type=int, default=2, help="批次大小") parser.add_argument("--learning_rate", type=float, default=3e-7, help="学习率") - parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") - parser.add_argument("--dtype", type=str, default="bfloat16", help="数据类型 bfloat16/float16") + parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)") + parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="训练数据类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") @@ -412,7 +420,7 @@ if __name__ == "__main__": parser.add_argument("--sglang_shared_path", type=str, default="./sglang_ckpt_agent", help="SGLang共享存储路径") args = parser.parse_args() - local_rank = init_distributed_mode() + local_rank = init_distributed_mode(args.device) if dist.is_initialized(): args.device = f"cuda:{local_rank}" setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) @@ -421,9 +429,8 @@ if __name__ == "__main__": max_seq_len=args.max_seq_len + args.max_gen_len, use_moe=bool(args.use_moe)) ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume == 1 else None - device_type = "cuda" if "cuda" in args.device else "cpu" - dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 - autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) + autocast_ctx = get_autocast_context(args.device, args.dtype) + scaler = get_grad_scaler(args.device, args.dtype) wandb = None if args.use_wandb and is_main_process(): @@ -464,6 +471,7 @@ if __name__ == "__main__": model.load_state_dict(ckp_data['model']) optimizer.load_state_dict(ckp_data['optimizer']) scheduler.load_state_dict(ckp_data['scheduler']) + if 'scaler' in ckp_data: scaler.load_state_dict(ckp_data['scaler']) start_epoch = ckp_data['epoch'] start_step = ckp_data.get('step', 0) @@ -480,7 +488,8 @@ if __name__ == "__main__": setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist() skip = start_step if (epoch == start_epoch and start_step > 0) else 0 batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) - loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True, collate_fn=collate_fn) + loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, + pin_memory=should_pin_memory(args.device), collate_fn=collate_fn) if skip > 0: Logger(f'Epoch [{epoch+1}/{args.epochs}]: skip {start_step} steps') rl_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang")) diff --git a/trainer/train_distillation.py b/trainer/train_distillation.py index 085b2a7..4ebeecc 100644 --- a/trainer/train_distillation.py +++ b/trainer/train_distillation.py @@ -11,13 +11,16 @@ import warnings import torch import torch.nn.functional as F import torch.distributed as dist -from contextlib import nullcontext from torch import optim from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, DistributedSampler from model.model_minimind import MiniMindConfig from dataset.lm_dataset import SFTDataset -from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler +from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint, + init_distributed_mode, setup_seed, init_model, + SkipBatchSampler, get_default_device, + get_autocast_context, get_grad_scaler, + should_pin_memory) warnings.filterwarnings('ignore') @@ -151,8 +154,8 @@ if __name__ == "__main__": parser.add_argument("--epochs", type=int, default=6, help="训练轮数") parser.add_argument("--batch_size", type=int, default=32, help="batch size") parser.add_argument("--learning_rate", type=float, default=5e-6, help="初始学习率") - parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") - parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") + parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)") + parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") @@ -177,7 +180,7 @@ if __name__ == "__main__": args = parser.parse_args() # ========== 1. 初始化环境和随机种子 ========== - local_rank = init_distributed_mode() + local_rank = init_distributed_mode(args.device) if dist.is_initialized(): args.device = f"cuda:{local_rank}" setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) @@ -188,9 +191,8 @@ if __name__ == "__main__": ckp_data = lm_checkpoint(lm_config_student, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None # ========== 3. 设置混合精度 ========== - device_type = "cuda" if "cuda" in args.device else "cpu" - dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 - autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) + autocast_ctx = get_autocast_context(args.device, args.dtype) + scaler = get_grad_scaler(args.device, args.dtype) # ========== 4. 配wandb ========== wandb = None @@ -210,7 +212,6 @@ if __name__ == "__main__": Logger(f'教师模型总参数量：{sum(p.numel() for p in teacher_model.parameters()) / 1e6:.3f} M') train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None - scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) # ========== 6. 从ckp恢复状态 ========== @@ -235,7 +236,8 @@ if __name__ == "__main__": setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist() skip = start_step if (epoch == start_epoch and start_step > 0) else 0 batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) - loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True) + loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, + pin_memory=should_pin_memory(args.device)) if skip > 0: Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') train_epoch(epoch, loader, len(loader) + skip, teacher_model, lm_config_student, start_step, wandb, args.alpha, args.temperature) @@ -245,4 +247,4 @@ if __name__ == "__main__": # ========== 9. 清理分布进程 ========== if dist.is_initialized(): dist.barrier() - dist.destroy_process_group() \ No newline at end of file + dist.destroy_process_group() diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index 479eb85..9bf58cf 100644 --- a/trainer/train_dpo.py +++ b/trainer/train_dpo.py @@ -11,13 +11,16 @@ import warnings import torch import torch.nn.functional as F import torch.distributed as dist -from contextlib import nullcontext from torch import optim from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, DistributedSampler from model.model_minimind import MiniMindConfig from dataset.lm_dataset import DPODataset -from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler +from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint, + init_distributed_mode, setup_seed, init_model, + SkipBatchSampler, get_default_device, + get_autocast_context, get_grad_scaler, + should_pin_memory) warnings.filterwarnings('ignore') @@ -135,8 +138,8 @@ if __name__ == "__main__": parser.add_argument("--epochs", type=int, default=1, help="训练轮数") parser.add_argument("--batch_size", type=int, default=4, help="batch size") parser.add_argument("--learning_rate", type=float, default=4e-8, help="初始学习率（建议<=5e-8避免遗忘）") - parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") - parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") + parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)") + parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") @@ -156,7 +159,7 @@ if __name__ == "__main__": args = parser.parse_args() # ========== 1. 初始化环境和随机种子 ========== - local_rank = init_distributed_mode() + local_rank = init_distributed_mode(args.device) if dist.is_initialized(): args.device = f"cuda:{local_rank}" setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) @@ -166,9 +169,8 @@ if __name__ == "__main__": ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None # ========== 3. 设置混合精度 ========== - device_type = "cuda" if "cuda" in args.device else "cpu" - dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 - autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) + autocast_ctx = get_autocast_context(args.device, args.dtype) + scaler = get_grad_scaler(args.device, args.dtype) # ========== 4. 配wandb ========== wandb = None @@ -190,7 +192,6 @@ if __name__ == "__main__": train_ds = DPODataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None - scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) # ========== 6. 从ckp恢复状态 ========== @@ -215,7 +216,8 @@ if __name__ == "__main__": setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist() skip = start_step if (epoch == start_epoch and start_step > 0) else 0 batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) - loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True) + loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, + pin_memory=should_pin_memory(args.device)) if skip > 0: Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') train_epoch(epoch, loader, len(loader) + skip, ref_model, lm_config, start_step, wandb, args.beta) @@ -225,4 +227,4 @@ if __name__ == "__main__": # ========== 9. 清理分布进程 ========== if dist.is_initialized(): dist.barrier() - dist.destroy_process_group() \ No newline at end of file + dist.destroy_process_group() diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index c22eb80..5aa7bb7 100644 --- a/trainer/train_full_sft.py +++ b/trainer/train_full_sft.py @@ -10,13 +10,16 @@ import time import warnings import torch import torch.distributed as dist -from contextlib import nullcontext from torch import optim, nn from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, DistributedSampler from model.model_minimind import MiniMindConfig from dataset.lm_dataset import SFTDataset -from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler +from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint, + init_distributed_mode, setup_seed, init_model, + SkipBatchSampler, get_default_device, + get_autocast_context, get_grad_scaler, + should_pin_memory) warnings.filterwarnings('ignore') @@ -88,8 +91,8 @@ if __name__ == "__main__": parser.add_argument("--epochs", type=int, default=2, help="训练轮数") parser.add_argument("--batch_size", type=int, default=16, help="batch size") parser.add_argument("--learning_rate", type=float, default=1e-5, help="初始学习率") - parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") - parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") + parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)") + parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") @@ -109,7 +112,7 @@ if __name__ == "__main__": args = parser.parse_args() # ========== 1. 初始化环境和随机种子 ========== - local_rank = init_distributed_mode() + local_rank = init_distributed_mode(args.device) if dist.is_initialized(): args.device = f"cuda:{local_rank}" setup_seed(args.seed + (dist.get_rank() if dist.is_initialized() else 0)) @@ -119,9 +122,8 @@ if __name__ == "__main__": ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None # ========== 3. 设置混合精度 ========== - device_type = "cuda" if "cuda" in args.device else "cpu" - dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 - autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) + autocast_ctx = get_autocast_context(args.device, args.dtype) + scaler = get_grad_scaler(args.device, args.dtype) # ========== 4. 配wandb ========== wandb = None @@ -136,7 +138,6 @@ if __name__ == "__main__": model, tokenizer = init_model(lm_config, args.from_weight, device=args.device) train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None - scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) # ========== 6. 从ckp恢复状态 ========== @@ -161,7 +162,8 @@ if __name__ == "__main__": setup_seed(args.seed + epoch); indices = torch.randperm(len(train_ds)).tolist() skip = start_step if (epoch == start_epoch and start_step > 0) else 0 batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) - loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True) + loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, + pin_memory=should_pin_memory(args.device)) if skip > 0: Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') train_epoch(epoch, loader, len(loader) + skip, start_step, wandb) @@ -171,4 +173,4 @@ if __name__ == "__main__": # ========== 9. 清理分布进程 ========== if dist.is_initialized(): dist.barrier() - dist.destroy_process_group() \ No newline at end of file + dist.destroy_process_group() diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py index 200adb6..7033e76 100755 --- a/trainer/train_grpo.py +++ b/trainer/train_grpo.py @@ -14,7 +14,6 @@ import torch import torch.nn.functional as F import torch.distributed as dist from transformers import AutoTokenizer -from contextlib import nullcontext from torch import optim from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, DistributedSampler @@ -22,7 +21,11 @@ from torch.optim.lr_scheduler import CosineAnnealingLR from transformers import AutoModel from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import RLAIFDataset -from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel +from trainer.trainer_utils import (Logger, is_main_process, lm_checkpoint, + init_distributed_mode, setup_seed, + SkipBatchSampler, init_model, LMForRewardModel, + get_default_device, get_autocast_context, + get_grad_scaler, should_pin_memory) from trainer.rollout_engine import create_rollout_engine warnings.filterwarnings('ignore') @@ -142,12 +145,14 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod per_token_loss = -(torch.min(per_token_loss1, per_token_loss2) - args.beta * per_token_kl) policy_loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1).clamp(min=1)).mean() loss = (policy_loss + aux_loss) / args.accumulation_steps # scalar - loss.backward() + scaler.scale(loss).backward() if step % args.accumulation_steps == 0: if args.grad_clip > 0: + scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) - optimizer.step() + scaler.step(optimizer) + scaler.update() scheduler.step() optimizer.zero_grad() @@ -186,7 +191,8 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod state_dict = raw_model.state_dict() torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp) lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, - epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scheduler=scheduler) + epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', + scheduler=scheduler, scaler=scaler) model.train() del state_dict @@ -197,8 +203,10 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod if step > start_step and step % args.accumulation_steps != 0: if args.grad_clip > 0: + scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) - optimizer.step() + scaler.step(optimizer) + scaler.update() scheduler.step() optimizer.zero_grad() @@ -210,8 +218,8 @@ if __name__ == "__main__": parser.add_argument("--epochs", type=int, default=1, help="训练轮数") parser.add_argument("--batch_size", type=int, default=2, help="batch size") parser.add_argument("--learning_rate", type=float, default=3e-7, help="初始学习率") - parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") - parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") + parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)") + parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") @@ -244,7 +252,7 @@ if __name__ == "__main__": args = parser.parse_args() # ========== 1. 初始化环境和随机种子 ========== - local_rank = init_distributed_mode() + local_rank = init_distributed_mode(args.device) if dist.is_initialized(): args.device = f"cuda:{local_rank}" setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) @@ -255,9 +263,8 @@ if __name__ == "__main__": ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None # ========== 3. 设置混合精度 ========== - device_type = "cuda" if "cuda" in args.device else "cpu" - dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 - autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) + autocast_ctx = get_autocast_context(args.device, args.dtype) + scaler = get_grad_scaler(args.device, args.dtype) # ========== 4. 配wandb ========== wandb = None @@ -303,6 +310,7 @@ if __name__ == "__main__": model.load_state_dict(ckp_data['model']) optimizer.load_state_dict(ckp_data['optimizer']) scheduler.load_state_dict(ckp_data['scheduler']) + if 'scaler' in ckp_data: scaler.load_state_dict(ckp_data['scaler']) start_epoch = ckp_data['epoch'] start_step = ckp_data.get('step', 0) @@ -321,7 +329,8 @@ if __name__ == "__main__": setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist() skip = start_step if (epoch == start_epoch and start_step > 0) else 0 batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) - loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True) + loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, + pin_memory=should_pin_memory(args.device)) if skip > 0: Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') grpo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang")) @@ -331,4 +340,4 @@ if __name__ == "__main__": # ========== 9. 清理分布进程 ========== if dist.is_initialized(): dist.barrier() - dist.destroy_process_group() \ No newline at end of file + dist.destroy_process_group() diff --git a/trainer/train_lora.py b/trainer/train_lora.py index 19a1499..455fd0f 100644 --- a/trainer/train_lora.py +++ b/trainer/train_lora.py @@ -10,14 +10,17 @@ import time import warnings import torch import torch.distributed as dist -from contextlib import nullcontext from torch import optim, nn from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, DistributedSampler from model.model_minimind import MiniMindConfig from dataset.lm_dataset import SFTDataset from model.model_lora import save_lora, apply_lora -from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler +from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint, + init_distributed_mode, setup_seed, init_model, + SkipBatchSampler, get_default_device, + get_autocast_context, get_grad_scaler, + should_pin_memory) warnings.filterwarnings('ignore') @@ -82,8 +85,8 @@ if __name__ == "__main__": parser.add_argument("--epochs", type=int, default=10, help="训练轮数") parser.add_argument("--batch_size", type=int, default=32, help="batch size") parser.add_argument("--learning_rate", type=float, default=1e-4, help="初始学习率") - parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") - parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") + parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)") + parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") @@ -102,7 +105,7 @@ if __name__ == "__main__": args = parser.parse_args() # ========== 1. 初始化环境和随机种子 ========== - local_rank = init_distributed_mode() + local_rank = init_distributed_mode(args.device) if dist.is_initialized(): args.device = f"cuda:{local_rank}" setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) @@ -112,9 +115,8 @@ if __name__ == "__main__": ckp_data = lm_checkpoint(lm_config, weight=args.lora_name, save_dir='../checkpoints') if args.from_resume==1 else None # ========== 3. 设置混合精度 ========== - device_type = "cuda" if "cuda" in args.device else "cpu" - dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 - autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) + autocast_ctx = get_autocast_context(args.device, args.dtype) + scaler = get_grad_scaler(args.device, args.dtype) # ========== 4. 配wandb ========== wandb = None @@ -148,7 +150,6 @@ if __name__ == "__main__": # ========== 6. 定义数据和优化器 ========== train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None - scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) optimizer = optim.AdamW(lora_params, lr=args.learning_rate) # ========== 7. 从ckp恢复状态 ========== @@ -173,7 +174,8 @@ if __name__ == "__main__": setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist() skip = start_step if (epoch == start_epoch and start_step > 0) else 0 batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) - loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True) + loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, + pin_memory=should_pin_memory(args.device)) if skip > 0: Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') train_epoch(epoch, loader, len(loader) + skip, lora_params, start_step, wandb) @@ -183,4 +185,4 @@ if __name__ == "__main__": # ========== 10. 清理分布进程 ========== if dist.is_initialized(): dist.barrier() - dist.destroy_process_group() \ No newline at end of file + dist.destroy_process_group() diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py index 0cf844d..2aab528 100644 --- a/trainer/train_ppo.py +++ b/trainer/train_ppo.py @@ -12,7 +12,6 @@ import warnings import torch import torch.distributed as dist import torch.nn.functional as F -from contextlib import nullcontext from torch import optim, nn from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, DistributedSampler @@ -20,7 +19,11 @@ from torch.nn.utils import clip_grad_norm_ from torch.optim.lr_scheduler import CosineAnnealingLR from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import RLAIFDataset -from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel +from trainer.trainer_utils import (Logger, is_main_process, lm_checkpoint, + init_distributed_mode, setup_seed, + SkipBatchSampler, init_model, LMForRewardModel, + get_default_device, get_autocast_context, + get_grad_scaler, should_pin_memory) from trainer.rollout_engine import create_rollout_engine warnings.filterwarnings('ignore') @@ -224,7 +227,7 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched else: loss = (policy_loss + args.vf_coef * value_loss + aux_loss) / args.accumulation_steps - loss.backward() + scaler.scale(loss).backward() policy_loss_sum += policy_loss.item() value_loss_sum += value_loss.item() @@ -237,20 +240,26 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched grad_accum_step += 1 if grad_accum_step % args.accumulation_steps == 0: + scaler.unscale_(actor_optimizer) + scaler.unscale_(critic_optimizer) clip_grad_norm_(actor_model.parameters(), args.grad_clip) clip_grad_norm_(critic_model.parameters(), args.grad_clip) - actor_optimizer.step() - critic_optimizer.step() + scaler.step(actor_optimizer) + scaler.step(critic_optimizer) + scaler.update() actor_scheduler.step() critic_scheduler.step() actor_optimizer.zero_grad() critic_optimizer.zero_grad() if grad_accum_step % args.accumulation_steps != 0: + scaler.unscale_(actor_optimizer) + scaler.unscale_(critic_optimizer) clip_grad_norm_(actor_model.parameters(), args.grad_clip) clip_grad_norm_(critic_model.parameters(), args.grad_clip) - actor_optimizer.step() - critic_optimizer.step() + scaler.step(actor_optimizer) + scaler.step(critic_optimizer) + scaler.update() actor_scheduler.step() critic_scheduler.step() actor_optimizer.zero_grad() @@ -297,7 +306,8 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched lm_checkpoint(lm_config, weight=args.save_weight, model=actor_model, optimizer=actor_optimizer, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scheduler=actor_scheduler, critic_model=critic_model, - critic_optimizer=critic_optimizer, critic_scheduler=critic_scheduler) + critic_optimizer=critic_optimizer, critic_scheduler=critic_scheduler, + scaler=scaler) actor_model.train() del actor_state @@ -314,8 +324,8 @@ if __name__ == "__main__": parser.add_argument("--batch_size", type=int, default=2, help="batch size") parser.add_argument("--learning_rate", type=float, default=3e-7, help="Actor学习率") parser.add_argument("--critic_learning_rate", type=float, default=5e-7, help="Critic学习率") - parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") - parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") + parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)") + parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") @@ -353,7 +363,7 @@ if __name__ == "__main__": args = parser.parse_args() # ========== 1. 初始化环境和随机种子 ========== - local_rank = init_distributed_mode() + local_rank = init_distributed_mode(args.device) if dist.is_initialized(): args.device = f"cuda:{local_rank}" setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) @@ -363,9 +373,8 @@ if __name__ == "__main__": ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None # ========== 3. 设置混合精度 ========== - device_type = "cuda" if "cuda" in args.device else "cpu" - dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 - autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) + autocast_ctx = get_autocast_context(args.device, args.dtype) + scaler = get_grad_scaler(args.device, args.dtype) # ========== 4. 配wandb ========== wandb = None @@ -419,6 +428,7 @@ if __name__ == "__main__": critic_optimizer.load_state_dict(ckp_data['critic_optimizer']) actor_scheduler.load_state_dict(ckp_data['scheduler']) critic_scheduler.load_state_dict(ckp_data['critic_scheduler']) + if 'scaler' in ckp_data: scaler.load_state_dict(ckp_data['scaler']) start_epoch = ckp_data['epoch'] start_step = ckp_data.get('step', 0) @@ -438,7 +448,8 @@ if __name__ == "__main__": setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist() skip = start_step if (epoch == start_epoch and start_step > 0) else 0 batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) - loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True) + loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, + pin_memory=should_pin_memory(args.device)) if skip > 0: Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') ppo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, actor_scheduler, critic_scheduler, reward_model, start_step, wandb) @@ -448,4 +459,4 @@ if __name__ == "__main__": # ========== 9. 清理分布进程 ========== if dist.is_initialized(): dist.barrier() - dist.destroy_process_group() \ No newline at end of file + dist.destroy_process_group() diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index 83a0382..62bca57 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -10,13 +10,16 @@ import time import warnings import torch import torch.distributed as dist -from contextlib import nullcontext from torch import optim, nn from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, DistributedSampler from model.model_minimind import MiniMindConfig from dataset.lm_dataset import PretrainDataset -from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, init_model, SkipBatchSampler +from trainer.trainer_utils import (get_lr, Logger, is_main_process, lm_checkpoint, + init_distributed_mode, setup_seed, init_model, + SkipBatchSampler, get_default_device, + get_autocast_context, get_grad_scaler, + should_pin_memory) warnings.filterwarnings('ignore') @@ -87,8 +90,8 @@ if __name__ == "__main__": parser.add_argument("--epochs", type=int, default=2, help="训练轮数") parser.add_argument("--batch_size", type=int, default=32, help="batch size") parser.add_argument("--learning_rate", type=float, default=5e-4, help="初始学习率") - parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备") - parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") + parser.add_argument("--device", type=str, default=get_default_device(), help="训练设备 (cuda:0/mps/cpu)") + parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16"], help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数") parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值") @@ -108,7 +111,7 @@ if __name__ == "__main__": args = parser.parse_args() # ========== 1. 初始化环境和随机种子 ========== - local_rank = init_distributed_mode() + local_rank = init_distributed_mode(args.device) if dist.is_initialized(): args.device = f"cuda:{local_rank}" setup_seed(args.seed + (dist.get_rank() if dist.is_initialized() else 0)) @@ -118,9 +121,8 @@ if __name__ == "__main__": ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None # ========== 3. 设置混合精度 ========== - device_type = "cuda" if "cuda" in args.device else "cpu" - dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 - autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) + autocast_ctx = get_autocast_context(args.device, args.dtype) + scaler = get_grad_scaler(args.device, args.dtype) # ========== 4. 配wandb ========== wandb = None @@ -135,7 +137,6 @@ if __name__ == "__main__": model, tokenizer = init_model(lm_config, args.from_weight, device=args.device) train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len) train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None - scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16')) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) # ========== 6. 从ckp恢复状态 ========== @@ -160,7 +161,8 @@ if __name__ == "__main__": setup_seed(args.seed + epoch); indices = torch.randperm(len(train_ds)).tolist() skip = start_step if (epoch == start_epoch and start_step > 0) else 0 batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) - loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True) + loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, + pin_memory=should_pin_memory(args.device)) if skip > 0: Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') train_epoch(epoch, loader, len(loader) + skip, start_step, wandb) @@ -170,4 +172,4 @@ if __name__ == "__main__": # ========== 9. 清理分布进程 ========== if dist.is_initialized(): dist.barrier() - dist.destroy_process_group() \ No newline at end of file + dist.destroy_process_group()