From ed006a254f53ff70d26bab4493bf3cc88c7b345e Mon Sep 17 00:00:00 2001 From: baihaochen <1291308544@qq.com> Date: Fri, 22 May 2026 00:27:06 +0800 Subject: [PATCH] [feat] add DeepSpeed pretrain (ZeRO-2, BF16, auto LR, gradient accumulation) --- trainer/ds_config_pretrain.json | 37 +++++ trainer/train_pretrain_deepspeed.py | 216 ++++++++++++++++++++++++++++ 2 files changed, 253 insertions(+) create mode 100644 trainer/ds_config_pretrain.json create mode 100644 trainer/train_pretrain_deepspeed.py diff --git a/trainer/ds_config_pretrain.json b/trainer/ds_config_pretrain.json new file mode 100644 index 0000000..d03d56b --- /dev/null +++ b/trainer/ds_config_pretrain.json @@ -0,0 +1,37 @@ +{ + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + "gradient_accumulation_steps": "auto", + "gradient_clipping": 1.0, + "zero_optimization": { + "stage": 2, + "allgather_partitions": true, + "allgather_bucket_size": 5e6, + "reduce_scatter": true, + "reduce_bucket_size": 5e6, + "overlap_comm": true, + "contiguous_gradients": true + }, + "bf16": { + "enabled": true + }, + "optimizer": { + "type": "AdamW", + "params": { + "lr": "auto", + "betas": [0.9, 0.95], + "eps": 1e-8, + "weight_decay": 0.1 + } + }, + "scheduler": { + "type": "WarmupDecayLR", + "params": { + "warmup_min_lr": 0, + "warmup_max_lr": "auto", + "warmup_num_steps": "auto", + "total_num_steps": "auto" + } + }, + "wall_clock_breakdown": false +} diff --git a/trainer/train_pretrain_deepspeed.py b/trainer/train_pretrain_deepspeed.py new file mode 100644 index 0000000..a87688d --- /dev/null +++ b/trainer/train_pretrain_deepspeed.py @@ -0,0 +1,216 @@ +""" +MiniMind Pretrain — DeepSpeed 版本 +使用 DeepSpeed ZeRO-2 替代 DDP，自动管理混合精度、梯度累积、梯度裁剪和学习率调度 +启动方式: + deepspeed --num_gpus=6 train_pretrain_deepspeed.py --use_wandb + torchrun --nproc_per_node=6 train_pretrain_deepspeed.py --use_wandb +""" +import os, sys + +__package__ = "trainer" +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +import argparse, time, json, warnings +import torch, torch.distributed as dist, deepspeed +from torch.utils.data import DataLoader, DistributedSampler +from model.model_minimind import MiniMindConfig +from model.model_minimind_mla import MiniMindMLAConfig +from dataset.lm_dataset import PretrainDataset +from trainer.trainer_utils import ( + Logger, is_main_process, setup_seed, init_model, SkipBatchSampler, get_model_suffix, + init_distributed_mode +) + +warnings.filterwarnings('ignore') + +# DeepSpeed 默认配置路径 +DEFAULT_DS_CONFIG = os.path.join(os.path.dirname(__file__), "ds_config_pretrain.json") + + +def _get_ds_config(args, total_steps): + """加载 DeepSpeed 配置，并填入训练总步数""" + ds_path = getattr(args, 'deepspeed_config', None) or DEFAULT_DS_CONFIG + with open(ds_path) as f: + cfg = json.load(f) + + # 自动推算 batch size / accumulation + if cfg.get("train_micro_batch_size_per_gpu") == "auto": + cfg["train_micro_batch_size_per_gpu"] = args.batch_size + if cfg.get("gradient_accumulation_steps") == "auto": + cfg["gradient_accumulation_steps"] = args.accumulation_steps + if cfg.get("train_batch_size") == "auto": + world = dist.get_world_size() if dist.is_initialized() else 1 + cfg["train_batch_size"] = args.batch_size * args.accumulation_steps * world + + # 填入学习率和总步数 + cfg["optimizer"]["params"]["lr"] = args.learning_rate + sched = cfg.get("scheduler", {}) + if sched.get("type") == "WarmupDecayLR": + p = sched.setdefault("params", {}) + if p.get("warmup_max_lr") == "auto": + p["warmup_max_lr"] = args.learning_rate + if p.get("warmup_num_steps") == "auto": + p["warmup_num_steps"] = max(1, total_steps // 20) + if p.get("total_num_steps") == "auto": + p["total_num_steps"] = total_steps + return cfg + + +def train_epoch(epoch, loader, iters, start_step=0, wandb=None): + start_time = time.time() + last_step = start_step + for step, (input_ids, labels) in enumerate(loader, start=start_step + 1): + input_ids = input_ids.to(args.device) + labels = labels.to(args.device) + last_step = step + + # DeepSpeed 自动管理混合精度，不需要 autocast_ctx + res = model_engine(input_ids, labels=labels) + loss = res.loss + res.aux_loss + # DeepSpeed 自动处理梯度累积和 loss scaling + model_engine.backward(loss) + # step：梯度裁剪 → 更新参数 → lr scheduler step → zero_grad + model_engine.step() + + if step % args.log_interval == 0 or step == iters: + spend_time = time.time() - start_time + current_loss = loss.item() + current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0 + current_logits_loss = current_loss - current_aux_loss + current_lr = model_engine.get_lr()[0] if model_engine.get_lr() else args.learning_rate + eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60 + Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min') + if wandb: + wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min}) + + if (step % args.save_interval == 0 or step == iters) and is_main_process(): + model_engine.eval() + moe_suffix = get_model_suffix(lm_config) + ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth' + # 保存权重（轻量） + raw_model = model_engine.module + raw_model = getattr(raw_model, '_orig_mod', raw_model) + torch.save({k: v.half().cpu() for k, v in raw_model.state_dict().items()}, ckp) + # 保存完整训练状态（含 optimizer / lr scheduler） + ds_tag = f'{args.save_weight}_{lm_config.hidden_size}{moe_suffix}' + model_engine.save_checkpoint('../checkpoints', tag=ds_tag) + Logger(f'DeepSpeed checkpoint saved: ../checkpoints/{ds_tag}') + model_engine.train() + del input_ids, labels, res, loss + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="MiniMind Pretraining (DeepSpeed)") + parser.add_argument("--save_dir", type=str, default="../out") + parser.add_argument('--save_weight', default='pretrain', type=str) + parser.add_argument("--epochs", type=int, default=2) + parser.add_argument("--batch_size", type=int, default=128) + parser.add_argument("--learning_rate", type=float, default=5e-4) + parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu") + parser.add_argument("--dtype", type=str, default="bfloat16") + parser.add_argument("--num_workers", type=int, default=8) + parser.add_argument("--accumulation_steps", type=int, default=8) + parser.add_argument("--grad_clip", type=float, default=1.0) + parser.add_argument("--log_interval", type=int, default=100) + parser.add_argument("--save_interval", type=int, default=1000) + parser.add_argument('--hidden_size', default=768, type=int) + parser.add_argument('--num_hidden_layers', default=8, type=int) + parser.add_argument('--max_seq_len', default=340, type=int) + parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1]) + parser.add_argument('--use_mla', default=0, type=int, choices=[0, 1]) + parser.add_argument('--kv_lora_rank', default=128, type=int) + parser.add_argument("--data_path", type=str, default="../dataset/pretrain_t2t_mini.jsonl") + parser.add_argument('--from_weight', default='none', type=str) + parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1]) + parser.add_argument("--use_wandb", action="store_true") + parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain-DS") + parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1]) + parser.add_argument("--deepspeed_config", type=str, default=DEFAULT_DS_CONFIG, help="DeepSpeed 配置 JSON") + parser.add_argument("--local_rank", type=int, default=-1, help="DeepSpeed 自动注入") + args = parser.parse_args() + + # ========== 1. 初始化环境 ========== + local_rank = init_distributed_mode() + if dist.is_initialized(): + args.device = f"cuda:{local_rank}" + setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) + + # ========== 2. 模型配置 ========== + os.makedirs(args.save_dir, exist_ok=True) + if args.use_mla: + lm_config = MiniMindMLAConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, + use_moe=bool(args.use_moe), kv_lora_rank=args.kv_lora_rank) + else: + lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, + use_moe=bool(args.use_moe)) + + # ========== 3. 初始化模型和 tokenizer ========== + model, tokenizer = init_model(lm_config, args.from_weight, device=args.device) + + # ========== 4. 数据集 ========== + train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len) + train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None + + # ========== 5. 算总步数，填充 DeepSpeed 配置 ========== + world_size = dist.get_world_size() if dist.is_initialized() else 1 + steps_per_epoch = len(train_ds) // (args.batch_size * args.accumulation_steps * world_size) + total_steps = steps_per_epoch * args.epochs + ds_config = _get_ds_config(args, total_steps) + + # ========== 6. DeepSpeed 初始化 ========== + model_engine, optimizer, _, _ = deepspeed.initialize( + model=model, + config=ds_config, + model_parameters=model.parameters(), + ) + + # ========== 7. wandb ========== + wandb = None + ckp_data = None + if args.use_wandb and is_main_process(): + import swanlab as wandb + wandb_run_name = f"MiniMind-Pretrain-DS-Epoch-{args.epochs}-BS-{args.batch_size}-LR-{args.learning_rate}" + wandb.init(project=args.wandb_project, name=wandb_run_name) + + # ========== 8. 从 DeepSpeed checkpoint 恢复 ========== + start_epoch, start_step = 0, 0 + ds_tag = f'{args.save_weight}_{lm_config.hidden_size}{get_model_suffix(lm_config)}' + if args.from_resume == 1: + try: + _, client_state = model_engine.load_checkpoint('../checkpoints', tag=ds_tag) + start_epoch = client_state.get('epoch', 0) + start_step = client_state.get('step', 0) + Logger(f'从 DeepSpeed checkpoint 恢复: epoch={start_epoch}, step={start_step}') + except Exception as e: + Logger(f'未找到 DeepSpeed checkpoint: {e}，从头开始训练') + + # 预热：DeepSpeed 在 bf16 下不需要 loss scaling + Logger(f'DeepSpeed 配置: world_size={world_size}, ' + f'micro_batch={ds_config["train_micro_batch_size_per_gpu"]}, ' + f'accumulation={ds_config["gradient_accumulation_steps"]}, ' + f'total_steps={total_steps}') + + # ========== 9. torch.compile（可选）========== + if args.use_compile == 1: + Logger('[WARN] torch.compile 与 DeepSpeed 可能会冲突，建议关闭') + # 如果坚持要用，只编译内部 module + if hasattr(model_engine, 'module'): + model_engine.module = torch.compile(model_engine.module) + + # ========== 10. 训练 ========== + for epoch in range(start_epoch, args.epochs): + train_sampler and train_sampler.set_epoch(epoch) + 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) + + 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) + else: + train_epoch(epoch, loader, len(loader), 0, wandb) + + if dist.is_initialized(): + dist.destroy_process_group()