219 lines
11 KiB
Python
219 lines
11 KiB
Python
"""
|
||
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
|
||
)
|
||
|
||
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_deepspeed', 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.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. 初始化环境（DeepSpeed 全权负责）==========
|
||
# deepspeed.initialize() 内部会自动 init_process_group + set_device
|
||
local_rank = int(os.environ.get("LOCAL_RANK", 0))
|
||
if not dist.is_initialized():
|
||
dist.init_process_group(backend="nccl")
|
||
torch.cuda.set_device(local_rank)
|
||
args.device = f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu"
|
||
setup_seed(42 + local_rank)
|
||
|
||
# ========== 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()
|