distributed training.

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)

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,7 +233,18 @@ 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):
+            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,

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():
+        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()
-            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
+                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
 

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():
+        # ---- 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()
-            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
+                # 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():

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():
+        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()
-            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
+                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
 

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'