diff --git a/dataset/lm_dataset.py b/dataset/lm_dataset.py index f22360d..795e715 100644 --- a/dataset/lm_dataset.py +++ b/dataset/lm_dataset.py @@ -39,7 +39,16 @@ class PretrainDataset(Dataset): super().__init__() self.tokenizer = tokenizer self.max_length = max_length - self.samples = load_dataset('json', data_files=data_path, split='train') + # data_path 支持三种形式： + # 1. 单文件路径 str → 单 jsonl + # 2. 逗号分隔多路径 str → 从 CLI 传入的多文件混合（"a.jsonl,b.jsonl"） + # 3. list[str] → API 调用直接传列表 + # 混合场景：1B pretrain 用 SkyPile 25B（raw web）+ minimind pretrain 1.76B（curated Q&A） + if isinstance(data_path, str) and ',' in data_path: + data_files = [p.strip() for p in data_path.split(',') if p.strip()] + else: + data_files = data_path + self.samples = load_dataset('json', data_files=data_files, split='train') def __len__(self): return len(self.samples) diff --git a/model/model_minimind.py b/model/model_minimind.py index 0236639..11f6000 100755 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -205,6 +205,13 @@ class MiniMindModel(nn.Module): freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.head_dim, end=config.max_position_embeddings, rope_base=config.rope_theta, rope_scaling=config.rope_scaling) self.register_buffer("freqs_cos", freqs_cos, persistent=False) self.register_buffer("freqs_sin", freqs_sin, persistent=False) + # gradient checkpointing 标志：True 时 forward 不存 activation， + # backward 时重算——省 activation 显存但慢 ~30%。1B + seq=1024 在 80GB 无需， + # 但 seq=2048 或 8B+ 模型必须。由 trainer 通过 set_grad_checkpointing() 打开。 + self.grad_checkpointing = False + + def set_grad_checkpointing(self, enable: bool = True): + self.grad_checkpointing = enable def warmup_rope(self, device=None): """Meta-device init 下 buffer 可能为空，manual 调用重新 precompute。正常 code path 不需要。""" @@ -226,14 +233,25 @@ class MiniMindModel(nn.Module): # 兼容，在 model.to(device) 后手动调 warmup_rope() 即可，不要放 forward hot path。) position_embeddings = (self.freqs_cos[start_pos:start_pos + seq_length], self.freqs_sin[start_pos:start_pos + seq_length]) presents = [] + # gradient checkpointing 与 use_cache=True 不兼容（重算时无法访问 KV cache）， + # 训练场景一般 use_cache=False，此处兜底 guard 提示。 + use_ckpt = self.grad_checkpointing and self.training and not use_cache for layer, past_key_value in zip(self.layers, past_key_values): - hidden_states, present = layer( - hidden_states, - position_embeddings, - past_key_value=past_key_value, - use_cache=use_cache, - attention_mask=attention_mask - ) + if use_ckpt: + # use_reentrant=False 是 torch>=2.1 推荐（旧 reentrant 版本对 DDP/compile 有问题） + hidden_states, present = torch.utils.checkpoint.checkpoint( + layer, hidden_states, position_embeddings, + past_key_value, use_cache, attention_mask, + use_reentrant=False, + ) + else: + hidden_states, present = layer( + hidden_states, + position_embeddings, + past_key_value=past_key_value, + use_cache=use_cache, + attention_mask=attention_mask + ) presents.append(present) hidden_states = self.norm(hidden_states) aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze()) diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index e9348a9..43384db 100644 --- a/trainer/train_full_sft.py +++ b/trainer/train_full_sft.py @@ -75,18 +75,23 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): Logger(f'[{time.strftime("%H:%M:%S")}] Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), it/s: {window_its:.2f}, MFU: {mfu:.1f}%, 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, "it/s": window_its, "MFU": mfu}) - if (step % args.save_interval == 0 or step == iters) and is_main_process(): - model.eval() - moe_suffix = '_moe' if lm_config.use_moe else '' - ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth' - raw_model = model.module if isinstance(model, DistributedDataParallel) else model - raw_model = getattr(raw_model, '_orig_mod', raw_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', scaler=scaler) - model.train() - del state_dict + if step % args.save_interval == 0 or step == iters: + if is_main_process(): + model.eval() + moe_suffix = '_moe' if lm_config.use_moe else '' + ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth' + raw_model = model.module if isinstance(model, DistributedDataParallel) else model + raw_model = getattr(raw_model, '_orig_mod', raw_model) + state_dict = raw_model.state_dict() + ckp_tmp = ckp + '.tmp' + torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp_tmp) + os.replace(ckp_tmp, ckp) + lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, + epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler) + model.train() + del state_dict + if dist.is_initialized(): + dist.barrier() del input_ids, labels, res, loss diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index 9f2a033..cf654b6 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -6,6 +6,7 @@ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) import datasets # noqa: F401 # Windows pyarrow/torch DLL conflict workaround (issue #771) import argparse +import math import time import warnings import torch @@ -25,7 +26,32 @@ from trainer.trainer_utils import ( warnings.filterwarnings('ignore') -def train_epoch(epoch, loader, iters, start_step=0, wandb=None): +@torch.no_grad() +def eval_holdout(model, eval_loader, max_batches, autocast_ctx, device): + """ + rank-0 only 跑 holdout PPL；调用方需保证只 rank-0 进入。 + 返回 (avg_loss, ppl)。max_batches 控制 eval 时长（1B 时 20-50 batch 足够，~1s）。 + """ + was_training = model.training + model.eval() + total_loss, total_tokens = 0.0, 0 + for i, (input_ids, labels) in enumerate(eval_loader): + if i >= max_batches: break + input_ids = input_ids.to(device, non_blocking=True) + labels = labels.to(device, non_blocking=True) + with autocast_ctx: + res = model(input_ids, labels=labels) + # loss 已经是 mean over non-pad tokens；乘 n_tokens 便于跨 batch weighted avg + n_tokens = (labels != -100).sum().item() + if n_tokens > 0: + total_loss += res.loss.item() * n_tokens + total_tokens += n_tokens + if was_training: model.train() + avg_loss = total_loss / max(total_tokens, 1) + return avg_loss, math.exp(avg_loss) + + +def train_epoch(epoch, loader, iters, start_step=0, wandb=None, eval_loader=None): start_time = time.time() last_log_time, last_log_step = start_time, start_step last_step = start_step @@ -40,7 +66,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): input_ids = input_ids.to(args.device, non_blocking=True) labels = labels.to(args.device, non_blocking=True) last_step = step - lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate) + lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate, args.warmup_steps) for param_group in optimizer.param_groups: param_group['lr'] = lr @@ -79,17 +105,34 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): Logger(f'[{time.strftime("%H:%M:%S")}] Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), it/s: {window_its:.2f}, MFU: {mfu:.1f}%, 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.eval() - moe_suffix = '_moe' if lm_config.use_moe else '' - ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth' - raw_model = model.module if isinstance(model, DistributedDataParallel) else model - raw_model = getattr(raw_model, '_orig_mod', raw_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, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints') - model.train() - del state_dict + # ---- eval loop（可选） ---- + if args.eval_interval > 0 and step % args.eval_interval == 0 and eval_loader is not None: + if is_main_process(): + val_loss, val_ppl = eval_holdout(model, eval_loader, args.eval_max_batches, autocast_ctx, args.device) + Logger(f'[eval] step={step} val_loss={val_loss:.4f} val_ppl={val_ppl:.2f}') + if wandb: wandb.log({"val_loss": val_loss, "val_ppl": val_ppl, "step": step}) + if dist.is_initialized(): dist.barrier() + + if step % args.save_interval == 0 or step == iters: + if is_main_process(): + model.eval() + moe_suffix = '_moe' if lm_config.use_moe else '' + ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth' + raw_model = model.module if isinstance(model, DistributedDataParallel) else model + raw_model = getattr(raw_model, '_orig_mod', raw_model) + state_dict = raw_model.state_dict() + # atomic: tmp + rename 避免读者读到半写文件（当前的 torch.save 直接写不 atomic， + # 极端情况 kill -9 会留 corrupt ckpt；lm_checkpoint 内部已 atomic，此处也补上） + ckp_tmp = ckp + '.tmp' + torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp_tmp) + os.replace(ckp_tmp, ckp) + lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints') + model.train() + del state_dict + # 所有 rank 在此汇合，让非 rank-0 等 rank-0 写完再进下一 step + # （防止 NCCL 隐式同步下的 timing race，也确保 ckpt 文件对 resume 一致） + if dist.is_initialized(): + dist.barrier() del input_ids, labels, res, loss @@ -127,6 +170,17 @@ if __name__ == "__main__": parser.add_argument("--use_compile", default=1, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是，A100 上推荐默认 1）") parser.add_argument("--gpu_peak_tflops", type=float, default=0, help="GPU bf16 理论峰值 TFLOPS，用于 MFU 计算；默认 0 表示按 GPU 名自动检测（A100=312）") + parser.add_argument("--warmup_steps", type=int, default=0, + help="linear warmup 步数。1B from-scratch 建议 500-2000；64M 用现有 pretrain ckpt 微调设 0 即可") + parser.add_argument("--grad_checkpointing", type=int, default=0, choices=[0, 1], + help="激活值 gradient checkpointing：省 30-50%% activation 显存，慢 ~30%%。" + "1B + seq=1024 在 80GB 用不上；seq=2048 或 8B+ 或大 batch 时需要打开") + parser.add_argument("--eval_interval", type=int, default=0, + help="每 N step 跑一次 holdout eval（0=关闭）。1B 建议 500-1000") + parser.add_argument("--eval_data_path", type=str, default=None, + help="holdout jsonl 路径；不设则关闭 eval") + parser.add_argument("--eval_max_batches", type=int, default=20, + help="每次 eval 最多跑几个 batch（20*batch_size 通常够 ~1s 内出数）") args = parser.parse_args() # ========== 1. 初始化环境和随机种子 ========== @@ -155,6 +209,9 @@ if __name__ == "__main__": # ========== 5. 定义模型、数据、优化器 ========== model, tokenizer = init_model(lm_config, args.from_weight, device=args.device) + if args.grad_checkpointing == 1: + model.model.set_grad_checkpointing(True) + Logger('gradient checkpointing enabled (activation 省 30-50%%, 慢 ~30%%)') 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')) @@ -189,6 +246,16 @@ if __name__ == "__main__": if dist.is_initialized(): model = DistributedDataParallel(model, device_ids=[local_rank]) + # ========== 7b. eval loader（可选） ========== + eval_loader = None + if args.eval_interval > 0 and args.eval_data_path and is_main_process(): + # rank-0 only：只 main rank 需要 eval loader（其他 rank barrier 等待） + eval_ds = PretrainDataset(args.eval_data_path, tokenizer, max_length=args.max_seq_len) + eval_loader = DataLoader(eval_ds, batch_size=args.batch_size, + shuffle=False, num_workers=2, pin_memory=True) + Logger(f'[eval] loader ready: {len(eval_ds)} samples from {args.eval_data_path}, ' + f'eval_interval={args.eval_interval}, max_batches={args.eval_max_batches}') + # ========== 8. 开始训练 ========== for epoch in range(start_epoch, args.epochs): train_sampler and train_sampler.set_epoch(epoch) @@ -199,9 +266,9 @@ if __name__ == "__main__": pin_memory=True, persistent_workers=args.num_workers > 0) 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) + train_epoch(epoch, loader, len(loader) + skip, start_step, wandb, eval_loader) else: - train_epoch(epoch, loader, len(loader), 0, wandb) + train_epoch(epoch, loader, len(loader), 0, wandb, eval_loader) # ========== 9. 清理分布进程 ========== if dist.is_initialized(): diff --git a/trainer/train_qat.py b/trainer/train_qat.py index 91e5592..da10819 100644 --- a/trainer/train_qat.py +++ b/trainer/train_qat.py @@ -88,20 +88,25 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): "epoch_time": eta_min, "it_per_s": window_its, "mfu_pct": mfu, }) - if (step % args.save_interval == 0 or step == iters) and is_main_process(): - model.eval() - moe_suffix = '_moe' if lm_config.use_moe else '' - ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth' - raw_model = model.module if isinstance(model, DistributedDataParallel) else model - raw_model = getattr(raw_model, '_orig_mod', raw_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', scaler=scaler, - ) - model.train() - del state_dict + if step % args.save_interval == 0 or step == iters: + if is_main_process(): + model.eval() + moe_suffix = '_moe' if lm_config.use_moe else '' + ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth' + raw_model = model.module if isinstance(model, DistributedDataParallel) else model + raw_model = getattr(raw_model, '_orig_mod', raw_model) + state_dict = raw_model.state_dict() + ckp_tmp = ckp + '.tmp' + torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp_tmp) + os.replace(ckp_tmp, ckp) + lm_checkpoint( + lm_config, weight=args.save_weight, model=model, optimizer=optimizer, + epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler, + ) + model.train() + del state_dict + if dist.is_initialized(): + dist.barrier() del input_ids, labels, res, loss diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index 66608a1..efeb71c 100644 --- a/trainer/trainer_utils.py +++ b/trainer/trainer_utils.py @@ -37,8 +37,19 @@ def Logger(content): print(content) -def get_lr(current_step, total_steps, lr): - return lr*(0.1 + 0.45*(1 + math.cos(math.pi * current_step / total_steps))) +def get_lr(current_step, total_steps, lr, warmup_steps=0): + """ + Linear warmup + cosine decay。 + - warmup_steps=0（默认）: 纯 cosine，向后兼容 + - warmup_steps>0: 前 warmup_steps 步从 0.1*lr 线性升到 lr，之后 cosine 到 0.1*lr + + 1B from-scratch 必须 warmup（无 warmup 时前 500 步 loss 发散概率 >50%， + 因为随机初始化的权重在 full lr 下第一步梯度过大打乱结构）。 + """ + if warmup_steps > 0 and current_step < warmup_steps: + return lr * (0.1 + 0.9 * current_step / warmup_steps) + progress = (current_step - warmup_steps) / max(total_steps - warmup_steps, 1) + return lr * (0.1 + 0.45 * (1 + math.cos(math.pi * progress))) # A100/H100/3090 等常见卡的 bf16 理论峰值 TFLOPS。MoE 等价 FLOPs 是 active params 不是 total， @@ -115,6 +126,14 @@ def setup_seed(seed: int, deterministic: bool = False): torch.set_float32_matmul_precision('high') def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoch=0, step=0, wandb=None, save_dir='../checkpoints', **kwargs): + """ + save mode (model != None): 只在 rank-0 写盘，非 main rank 早返回。 + load mode (model == None): 所有 rank 读同一份 ckpt，无冲突。 + 调用方在 save 后应加 dist.barrier()，让非-0 rank 等待 rank-0 写完再进入下一 step。 + """ + # DDP 保护：save 必须由单一 rank 执行；否则 4 rank 同写一个 file 会 corrupt。 + if model is not None and not is_main_process(): + return None os.makedirs(save_dir, exist_ok=True) moe_path = '_moe' if lm_config.use_moe else '' ckp_path = f'{save_dir}/{weight}_{lm_config.hidden_size}{moe_path}.pth'