diff --git a/eval_llm.py b/eval_llm.py index 9359ca8..c7f679d 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): @@ -27,7 +27,9 @@ def init_model(args): else: model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True) get_model_params(model, model.config) - return model.half().eval().to(args.device), tokenizer + if 'cuda' in args.device: + model = model.half() + return model.eval().to(args.device), tokenizer def main(): parser = argparse.ArgumentParser(description="MiniMind模型推理与对话") @@ -45,7 +47,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="运行设备") args = parser.parse_args() prompts = [ diff --git a/trainer/train_agent.py b/trainer/train_agent.py index 261b00c..fc666b7 100644 --- a/trainer/train_agent.py +++ b/trainer/train_agent.py @@ -24,7 +24,7 @@ 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 from trainer.rollout_engine import create_rollout_engine, compute_per_token_logps warnings.filterwarnings('ignore') @@ -384,7 +384,7 @@ 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("--device", type=str, default=get_default_device(), help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="数据类型 bfloat16/float16") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") @@ -427,9 +427,14 @@ 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" + device_type = "cuda" if "cuda" in args.device else ("mps" if "mps" 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) + if device_type == "cuda": + autocast_ctx = torch.cuda.amp.autocast(dtype=dtype) + elif device_type == "mps": + autocast_ctx = torch.amp.autocast("mps", dtype=dtype) + else: + autocast_ctx = nullcontext() wandb = None if args.use_wandb and is_main_process(): @@ -486,7 +491,7 @@ 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=(device_type == "cuda"), 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..fb4d131 100644 --- a/trainer/train_distillation.py +++ b/trainer/train_distillation.py @@ -17,7 +17,7 @@ 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 warnings.filterwarnings('ignore') @@ -151,7 +151,7 @@ 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("--device", type=str, default=get_default_device(), help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") @@ -188,9 +188,14 @@ 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" + device_type = "cuda" if "cuda" in args.device else ("mps" if "mps" 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) + if device_type == "cuda": + autocast_ctx = torch.cuda.amp.autocast(dtype=dtype) + elif device_type == "mps": + autocast_ctx = torch.amp.autocast("mps", dtype=dtype) + else: + autocast_ctx = nullcontext() # ========== 4. 配wandb ========== wandb = None @@ -210,7 +215,7 @@ 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')) + scaler = torch.cuda.amp.GradScaler(enabled=(device_type == "cuda" and args.dtype == "float16")) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) # ========== 6. 从ckp恢复状态 ========== @@ -235,7 +240,7 @@ 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=(device_type == "cuda")) 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) diff --git a/trainer/train_dpo.py b/trainer/train_dpo.py index 479eb85..6e1c73d 100644 --- a/trainer/train_dpo.py +++ b/trainer/train_dpo.py @@ -17,7 +17,7 @@ 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 warnings.filterwarnings('ignore') @@ -135,7 +135,7 @@ 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("--device", type=str, default=get_default_device(), help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") @@ -166,9 +166,14 @@ 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" + device_type = "cuda" if "cuda" in args.device else ("mps" if "mps" 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) + if device_type == "cuda": + autocast_ctx = torch.cuda.amp.autocast(dtype=dtype) + elif device_type == "mps": + autocast_ctx = torch.amp.autocast("mps", dtype=dtype) + else: + autocast_ctx = nullcontext() # ========== 4. 配wandb ========== wandb = None @@ -190,7 +195,7 @@ 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')) + scaler = torch.cuda.amp.GradScaler(enabled=(device_type == "cuda" and args.dtype == "float16")) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) # ========== 6. 从ckp恢复状态 ========== @@ -215,7 +220,7 @@ 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=(device_type == "cuda")) 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) diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index c22eb80..01bdef7 100644 --- a/trainer/train_full_sft.py +++ b/trainer/train_full_sft.py @@ -16,7 +16,7 @@ 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 warnings.filterwarnings('ignore') @@ -88,7 +88,7 @@ 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("--device", type=str, default=get_default_device(), help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") @@ -119,9 +119,14 @@ 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" + device_type = "cuda" if "cuda" in args.device else ("mps" if "mps" 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) + if device_type == "cuda": + autocast_ctx = torch.cuda.amp.autocast(dtype=dtype) + elif device_type == "mps": + autocast_ctx = torch.amp.autocast("mps", dtype=dtype) + else: + autocast_ctx = nullcontext() # ========== 4. 配wandb ========== wandb = None @@ -136,7 +141,7 @@ 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')) + scaler = torch.cuda.amp.GradScaler(enabled=(device_type == "cuda" and args.dtype == "float16")) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) # ========== 6. 从ckp恢复状态 ========== @@ -161,7 +166,7 @@ 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=(device_type == "cuda")) 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) diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py index 76ea045..cd996c4 100755 --- a/trainer/train_grpo.py +++ b/trainer/train_grpo.py @@ -22,7 +22,7 @@ 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 from trainer.rollout_engine import create_rollout_engine warnings.filterwarnings('ignore') @@ -211,7 +211,7 @@ 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("--device", type=str, default=get_default_device(), help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") @@ -256,9 +256,14 @@ 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" + device_type = "cuda" if "cuda" in args.device else ("mps" if "mps" 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) + if device_type == "cuda": + autocast_ctx = torch.cuda.amp.autocast(dtype=dtype) + elif device_type == "mps": + autocast_ctx = torch.amp.autocast("mps", dtype=dtype) + else: + autocast_ctx = nullcontext() # ========== 4. 配wandb ========== wandb = None @@ -322,7 +327,7 @@ 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=(device_type == "cuda")) 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")) diff --git a/trainer/train_lora.py b/trainer/train_lora.py index 19a1499..43e1561 100644 --- a/trainer/train_lora.py +++ b/trainer/train_lora.py @@ -17,7 +17,7 @@ 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 warnings.filterwarnings('ignore') @@ -82,7 +82,7 @@ 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("--device", type=str, default=get_default_device(), help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") @@ -112,9 +112,14 @@ 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" + device_type = "cuda" if "cuda" in args.device else ("mps" if "mps" 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) + if device_type == "cuda": + autocast_ctx = torch.cuda.amp.autocast(dtype=dtype) + elif device_type == "mps": + autocast_ctx = torch.amp.autocast("mps", dtype=dtype) + else: + autocast_ctx = nullcontext() # ========== 4. 配wandb ========== wandb = None @@ -148,7 +153,7 @@ 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')) + scaler = torch.cuda.amp.GradScaler(enabled=(device_type == "cuda" and args.dtype == "float16")) optimizer = optim.AdamW(lora_params, lr=args.learning_rate) # ========== 7. 从ckp恢复状态 ========== @@ -173,7 +178,7 @@ 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=(device_type == "cuda")) 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) diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py index 07c7738..624ec7b 100644 --- a/trainer/train_ppo.py +++ b/trainer/train_ppo.py @@ -20,7 +20,7 @@ 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 from trainer.rollout_engine import create_rollout_engine warnings.filterwarnings('ignore') @@ -315,7 +315,7 @@ 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("--device", type=str, default=get_default_device(), help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数") @@ -364,9 +364,14 @@ 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" + device_type = "cuda" if "cuda" in args.device else ("mps" if "mps" 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) + if device_type == "cuda": + autocast_ctx = torch.cuda.amp.autocast(dtype=dtype) + elif device_type == "mps": + autocast_ctx = torch.amp.autocast("mps", dtype=dtype) + else: + autocast_ctx = nullcontext() # ========== 4. 配wandb ========== wandb = None @@ -439,7 +444,7 @@ 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=(device_type == "cuda")) 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) diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index 83a0382..e554950 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -16,7 +16,7 @@ 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 warnings.filterwarnings('ignore') @@ -87,7 +87,7 @@ 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("--device", type=str, default=get_default_device(), help="训练设备") parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型") parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数") parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数") @@ -118,9 +118,14 @@ 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" + device_type = "cuda" if "cuda" in args.device else ("mps" if "mps" 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) + if device_type == "cuda": + autocast_ctx = torch.cuda.amp.autocast(dtype=dtype) + elif device_type == "mps": + autocast_ctx = torch.amp.autocast("mps", dtype=dtype) + else: + autocast_ctx = nullcontext() # ========== 4. 配wandb ========== wandb = None @@ -135,7 +140,7 @@ 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')) + scaler = torch.cuda.amp.GradScaler(enabled=(device_type == "cuda" and args.dtype == "float16")) optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) # ========== 6. 从ckp恢复状态 ========== @@ -160,7 +165,7 @@ 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=(device_type == "cuda")) 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) diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index b6211d5..cdd67d7 100644 --- a/trainer/trainer_utils.py +++ b/trainer/trainer_utils.py @@ -15,6 +15,15 @@ from torch.utils.data import Sampler from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification from model.model_minimind import MiniMindForCausalLM +def get_default_device(): + """运行时自动选择最优设备,CUDA > MPS > CPU""" + if torch.cuda.is_available(): + return "cuda" + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return "mps" + return "cpu" + + def get_model_params(model, config): total = sum(p.numel() for p in model.parameters()) / 1e6 n_routed = getattr(config, 'n_routed_experts', getattr(config, 'num_experts', 0)) @@ -45,6 +54,9 @@ def init_distributed_mode(): if int(os.environ.get("RANK", -1)) == -1: return 0 # 非DDP模式 + if not torch.cuda.is_available(): + return 0 # MPS/CPU 不支持 NCCL DDP,静默跳过 + dist.init_process_group(backend="nccl") local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(local_rank) @@ -55,10 +67,11 @@ def setup_seed(seed: int): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoch=0, step=0, wandb=None, save_dir='../checkpoints', **kwargs): os.makedirs(save_dir, exist_ok=True) @@ -103,7 +116,10 @@ def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoc torch.save(resume_data, resume_tmp) os.replace(resume_tmp, resume_path) del state_dict, resume_data - torch.cuda.empty_cache() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + elif hasattr(torch, "mps") and torch.backends.mps.is_available(): + torch.mps.empty_cache() else: # 加载模式 if os.path.exists(resume_path): ckp_data = torch.load(resume_path, map_location='cpu') @@ -116,7 +132,9 @@ def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoc return None -def init_model(lm_config, from_weight='pretrain', tokenizer_path='../model', save_dir='../out', device='cuda'): +def init_model(lm_config, from_weight='pretrain', tokenizer_path='../model', save_dir='../out', device=None): + if device is None: + device = get_default_device() tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) model = MiniMindForCausalLM(lm_config) @@ -158,7 +176,9 @@ class SkipBatchSampler(Sampler): class LMForRewardModel: - def __init__(self, model_path, device="cuda", dtype=torch.float16): + def __init__(self, model_path, device=None, dtype=torch.float16): + if device is None: + device = get_default_device() self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) self.model = AutoModel.from_pretrained(model_path, torch_dtype=dtype, trust_remote_code=True) self.model = self.model.to(device).eval()