[feat] add DeepSpeed pretrain (ZeRO-2, BF16, auto LR, gradient accumulation)
baihaochen
parent
579cc79cd8
commit
ed006a254f
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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()
|
||||
Loading…
Reference in New Issue