将Million MoE的思想加入

.gitignore

@@ -1,3 +1,5 @@
 /model/__pycache__
 /dataset
-/out
+/out
+wandb/
+**/*.log

eval_model.py

@@ -16,7 +16,7 @@ def init_model(args):
     tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
     if args.load == 0:
         moe_path = '_moe' if args.use_moe else ''
-        modes = {0: 'pretrain', 1: 'full_sft', 2: 'rlhf', 3: 'reason'}
+        modes = {0: 'pretrain', 1: 'full_sft', 2: 'rlhf', 3: 'reason', 4: 'grpo'}
         ckp = f'./{args.out_dir}/{modes[args.model_mode]}_{args.dim}{moe_path}.pth'
 
         model = MiniMindLM(LMConfig(
@@ -123,7 +123,7 @@ def main():
     parser.add_argument('--stream', default=True, type=bool)
     parser.add_argument('--load', default=0, type=int, help="0: 原生torch权重，1: transformers加载")
     parser.add_argument('--model_mode', default=1, type=int,
-                        help="0: 预训练模型，1: SFT-Chat模型，2: RLHF-Chat模型，3: Reason模型")
+                        help="0: 预训练模型，1: SFT-Chat模型，2: RLHF-Chat模型，3: Reason模型，4: RLAIF-Chat模型")
     args = parser.parse_args()
 
     model, tokenizer = init_model(args)
@@ -143,7 +143,7 @@ def main():
             messages,
             tokenize=False,
             add_generation_prompt=True
-        )[-args.max_seq_len + 1:] if args.model_mode != 0 else (tokenizer.bos_token + prompt)
+        )[-args.max_seq_len - 1:] if args.model_mode != 0 else (tokenizer.bos_token + prompt)
 
         answer = new_prompt
         with torch.no_grad():

model/dataset.py

@@ -196,5 +196,50 @@ class DPODataset(Dataset):
         return loss_mask
 
 
+class RLAIFDataset(Dataset):
+    def __init__(self, jsonl_path, tokenizer, max_length=1024):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.samples = self.load_data(jsonl_path)
+        self.bos_id = tokenizer('<s>assistant', add_special_tokens=False).input_ids
+        self.eos_id = tokenizer('</s>', add_special_tokens=False).input_ids
+
+    def __len__(self):
+        return len(self.samples)
+
+    def load_data(self, path):
+        samples = []
+        with open(path, 'r', encoding='utf-8') as f:
+            for line_num, line in enumerate(f, 1):
+                data = json.loads(line.strip())
+                samples.append(data)
+        return samples
+
+    def _create_chat_prompt(self, conversations):
+        """构建符合ChatML格式的对话"""
+        messages = []
+        answer = ''
+        for i, turn in enumerate(conversations):
+            role = 'user' if i % 2 == 0 else 'assistant'
+            messages.append({"role": role, "content": turn['content']})
+            answer = turn['content']
+        return self.tokenizer.apply_chat_template(
+            messages[:-1],
+            tokenize=False,
+            add_generation_prompt=True
+        ), answer
+
+    def __getitem__(self, index):
+        sample = self.samples[index]
+        # 构建对话提示
+        prompt, answer = self._create_chat_prompt(sample['conversations'])
+
+        return {
+            'prompt': prompt,
+            'answer': answer
+        }
+
+
 if __name__ == "__main__":
     pass

model/model.py

@@ -12,7 +12,7 @@ from torch import nn
 from transformers import PreTrainedModel
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
-
+# RMSNorm 类定义了一个用于归一化输入张量的模块。
 class RMSNorm(torch.nn.Module):
     def __init__(self, dim: int, eps: float = 1e-6):
         super().__init__()
@@ -25,7 +25,7 @@ class RMSNorm(torch.nn.Module):
     def forward(self, x):
         return self.weight * self._norm(x.float()).type_as(x)
 
-
+# precompute_pos_cis 函数用于预计算位置编码。
 def precompute_pos_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
     freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
     t = torch.arange(end, device=freqs.device)  # type: ignore
@@ -33,7 +33,7 @@ def precompute_pos_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
     pos_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
     return pos_cis
 
-
+# apply_rotary_emb 函数用于应用旋转位置编码。
 def apply_rotary_emb(xq, xk, pos_cis):
     def unite_shape(pos_cis, x):
         ndim = x.ndim
@@ -49,7 +49,7 @@ def apply_rotary_emb(xq, xk, pos_cis):
     xk_out = torch.view_as_real(xk_ * pos_cis).flatten(3)
     return xq_out.type_as(xq), xk_out.type_as(xk)
 
-
+# repeat_kv 函数用于重复键值对。
 def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
     """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
     bs, slen, n_kv_heads, head_dim = x.shape
@@ -88,13 +88,15 @@ class Attention(nn.Module):
                 x: torch.Tensor,
                 pos_cis: torch.Tensor,
                 past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-                use_cache=False):
-        bsz, seq_len, _ = x.shape
-        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
-        xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
-        xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
-        xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
+                use_cache=False,
+                db_value=None):
+        bsz, seq_len, _ = x.shape #bsz: 批量大小, seq_len: 序列长度, _: 隐藏维度
+        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x) #将输入张量x分别通过线性层wq, wk, wv进行变换，得到查询、键和值。
+        xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim) #将变换后的张量xq重塑为形状为(bsz, seq_len, n_local_heads, head_dim)的形状。
+        xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim) #将变换后的张量xk重塑为形状为(bsz, seq_len, n_local_kv_heads, head_dim)的形状。
+        xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim) #将变换后的张量xv重塑为形状为(bsz, seq_len, n_local_kv_heads, head_dim)的形状。
 
+        # 应用旋转位置编码
         xq, xk = apply_rotary_emb(xq, xk, pos_cis)
         # kv_cache实现
         if past_key_value is not None:
@@ -102,11 +104,40 @@ class Attention(nn.Module):
             xv = torch.cat([past_key_value[1], xv], dim=1)
         past_kv = (xk, xv) if use_cache else None
 
+        # 重复键值对
         xq, xk, xv = (
             xq.transpose(1, 2),
             repeat_kv(xk, self.n_rep).transpose(1, 2),
             repeat_kv(xv, self.n_rep).transpose(1, 2)
         )
+        
+        # 如果提供了db_value，根据头的数量调整它的形状并与xv合并
+        if db_value is not None:
+            # 确保db_value的形状与xv兼容，假设db_value形状为[B, N, H, D]
+            if db_value.ndim == 4:  # [B, N, H, D]
+                db_value = db_value.transpose(1, 2)  # -> [B, H, N, D]
+                
+                # 检查是否需要调整D维度
+                if db_value.shape[-1] != xv.shape[-1]:
+                    # 如果db_value的维度与xv不同，可以添加一个投影层
+                    # 或者在这里使用简单的调整方法
+                    # 这里我们简单地通过均值池化或重复来调整维度
+                    if db_value.shape[-1] > xv.shape[-1]:
+                        # 降维
+                        factor = db_value.shape[-1] // xv.shape[-1]
+                        db_value = db_value.view(bsz, self.n_local_heads, seq_len, factor, xv.shape[-1])
+                        db_value = db_value.mean(dim=3)
+                    else:
+                        # 升维
+                        factor = xv.shape[-1] // db_value.shape[-1]
+                        db_value = db_value.unsqueeze(-1).repeat(1, 1, 1, 1, factor)
+                        db_value = db_value.view(bsz, self.n_local_heads, seq_len, xv.shape[-1])
+                
+                # 将db_value与xv相加或融合
+                # 这里我们简单地将它们相加，但你也可以使用其他融合方法
+                xv = xv + db_value
+        
+        # 使用Flash Attention
         if self.flash and seq_len != 1:
             dropout_p = self.dropout if self.training else 0.0
             output = F.scaled_dot_product_attention(
@@ -221,7 +252,6 @@ class MOEFeedForward(nn.Module):
         x = x.view(-1, x.shape[-1])
         flat_topk_idx = topk_idx.view(-1)
         if self.training:
-            # 训练模式下，重复输入数据
             x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
             y = torch.empty_like(x, dtype=torch.float16)
             for i, expert in enumerate(self.experts):
@@ -229,7 +259,6 @@ class MOEFeedForward(nn.Module):
             y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
             y = y.view(*orig_shape)
         else:
-            # 推理模式下，只选择最优专家
             y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
         if self.config.n_shared_experts is not None:
             y = y + self.shared_experts(identity)
@@ -242,9 +271,10 @@ class MOEFeedForward(nn.Module):
         idxs = flat_expert_indices.argsort()
         tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
         token_idxs = idxs // self.config.num_experts_per_tok
-        # 例如当tokens_per_expert=[6, 15, 20, 26, 33, 38, 46, 52]
-        # 当token_idxs=[3, 7, 19, 21, 24, 25,  4,  5,  6, 10, 11, 12...]
-        # 意味着当token_idxs[:6] -> [3,  7, 19, 21, 24, 25,  4]位置的token都由专家0处理，token_idxs[6:15]位置的token都由专家1处理......
+        # 当tokens_per_expert = [6, 15, 20, 26]，tokens_per_expert.shape[0]即为专家数量（此时为4）
+        # 且token_idxs = [3, 7, 19, 21, 24, 25,  4,  5,  6, 10, 11, 12...] 时
+        # 意味token_idxs[:6] -> [3, 7, 19, 21, 24, 25]这6个位置属于专家0处理的token（每个token有可能被多个专家处理，这取决于num_experts_per_tok）
+        # 接下来9个位置token_idxs[6:15] -> [4,  5,  6, 10, 11, 12...]属于专家1处理的token...依此类推
         for i, end_idx in enumerate(tokens_per_expert):
             start_idx = 0 if i == 0 else tokens_per_expert[i - 1]
             if start_idx == end_idx:
@@ -254,14 +284,13 @@ class MOEFeedForward(nn.Module):
             expert_tokens = x[exp_token_idx]
             expert_out = expert(expert_tokens).to(expert_cache.dtype)
             expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
-            # 使用 scatter_add_ 进行 sum 操作
             expert_cache.scatter_add_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out)
 
         return expert_cache
 
 
 class MiniMindBlock(nn.Module):
-    def __init__(self, layer_id: int, config: LMConfig):
+    def __init__(self, layer_id: int, config: LMConfig, weight_down_embed=None):
         super().__init__()
         self.n_heads = config.n_heads
         self.dim = config.dim
@@ -272,13 +301,86 @@ class MiniMindBlock(nn.Module):
         self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)
         self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
         self.feed_forward = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
+        
+        # Product Key 相关参数
+        self.weight_down_embed = weight_down_embed
+        # 假设num_experts是已定义的总专家数量的平方根
+        self.num_keys = int(math.sqrt(self.weight_down_embed.num_embeddings)) if weight_down_embed is not None else 0
+        
+        # 查询生成的参数
+        self.dim_key = config.dim // 2  # 一般用特征维度的一半
+        
+        # 创建查询生成模块
+        if weight_down_embed is not None:
+            self.to_queries = nn.Sequential(
+                nn.Linear(config.dim, self.dim_key * self.n_heads * 2, bias=False),
+                nn.Unflatten(2, (2, self.n_heads, self.dim_key))  # 替代Rearrange
+            )
+            
+            # 存储Product Keys
+            self.keys = nn.Parameter(torch.randn(self.n_heads, self.num_keys, 2, self.dim_key) * 0.02)
+            
+            # 超参数
+            self.product_key_topk = min(16, self.num_keys)  # 确保不超过num_keys
+            self.num_experts_per_head_topk = 1  # 最终每个头选取的专家数
 
     def forward(self, x, pos_cis, past_key_value=None, use_cache=False):
+        db_value = None
+        
+        # 如果有weight_down_embed，使用Product Key机制
+        if self.weight_down_embed is not None:
+            # 1. 生成queries
+            queries = self.to_queries(x)  # [b, n, 2, h, d]
+            queries = queries.permute(2, 0, 1, 3, 4)  # [2, b, n, h, d]
+            
+            # 2. 计算queries与keys的相似度
+            sim = torch.einsum('p b n h d, h k p d -> p b n h k', queries, self.keys)
+            
+            # 3. 在两个子空间分别做top-k
+            scores_and_indices = [sim[p].topk(self.product_key_topk, dim=-1) for p in range(2)]
+            scores_x, scores_y = scores_and_indices[0][0], scores_and_indices[1][0]
+            indices_x, indices_y = scores_and_indices[0][1], scores_and_indices[1][1]
+            
+            # 4. 组合两个子空间的分数和索引
+            all_scores = scores_x.unsqueeze(-1) + scores_y.unsqueeze(-2)
+            all_scores = all_scores.view(*all_scores.shape[:-2], -1)
+            
+            all_indices = (indices_x.unsqueeze(-1) * self.num_keys) + indices_y.unsqueeze(-2)
+            all_indices = all_indices.view(*all_indices.shape[:-2], -1)
+            
+            # 5. 最终top-k选择
+            scores, pk_indices = all_scores.topk(self.num_experts_per_head_topk, dim=-1)
+            indices = all_indices.gather(-1, pk_indices)
+            
+            # 6. 从embedding中获取专家值
+            # [b, n, h, k] -> [b, n, h, k, 1] -> [b, n, h, k, d]
+            indices_expanded = indices.unsqueeze(-1).expand(-1, -1, -1, -1, self.weight_down_embed.embedding_dim)
+            
+            # 将索引从3D展平为1D以便gather操作
+            batch_size, seq_len = x.shape[0], x.shape[1]
+            flat_indices = indices.view(-1)
+            
+            # 从embedding中获取值
+            db_values = self.weight_down_embed(flat_indices)
+            
+            # 重塑回原始形状
+            db_value = db_values.view(batch_size, seq_len, self.n_heads, 
+                                    self.num_experts_per_head_topk, -1)
+            
+            # 使用分数加权
+            db_value = db_value * F.relu(scores.unsqueeze(-1))
+            
+            # 合并多个专家的输出（如果每个头有多个专家）
+            if self.num_experts_per_head_topk > 1:
+                db_value = db_value.sum(dim=3)  # [b, n, h, d]
+        
+        # 注意力计算
         h_attn, past_kv = self.attention(
             self.attention_norm(x),
             pos_cis,
             past_key_value=past_key_value,
-            use_cache=use_cache
+            use_cache=use_cache,
+            db_value=db_value
         )
         h = x + h_attn
         out = h + self.feed_forward(self.ffn_norm(h))
@@ -294,7 +396,20 @@ class MiniMindLM(PreTrainedModel):
         self.vocab_size, self.n_layers = params.vocab_size, params.n_layers
         self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)
         self.dropout = nn.Dropout(params.dropout)
-        self.layers = nn.ModuleList([MiniMindBlock(l, params) for l in range(self.n_layers)])
+        
+        # 修改专家数量和知识维度，确保能开方
+        self.num_experts = 1000 * 1000  # 1M专家，确保是完全平方数
+        # 将knowledge_dim设置为与head_dim相同，以便在attention中直接使用
+        self.head_dim = params.dim // params.n_heads
+        self.knowledge_dim = self.head_dim
+        
+        # 定义weight_down_embed，用于存储专家知识
+        self.weight_down_embed = nn.Embedding(self.num_experts, self.knowledge_dim)
+        # 初始化embedding权重
+        nn.init.normal_(self.weight_down_embed.weight, std=0.02)
+        
+        # 将self.weight_down_embed传递给每个MiniMindBlock
+        self.layers = nn.ModuleList([MiniMindBlock(l, params, self.weight_down_embed) for l in range(self.n_layers)])
         self.norm = RMSNorm(params.dim, eps=params.norm_eps)
         self.output = nn.Linear(params.dim, params.vocab_size, bias=False)
         self.tok_embeddings.weight = self.output.weight
@@ -325,6 +440,7 @@ class MiniMindLM(PreTrainedModel):
         slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
         logits = self.output(self.norm(h)[:, slice_indices, :])
         aux_loss = sum(l.feed_forward.aux_loss for l in self.layers if isinstance(l.feed_forward, MOEFeedForward))
+        self.OUT.__setitem__('last_hidden_state', h)
         self.OUT.__setitem__('logits', logits)
         self.OUT.__setitem__('aux_loss', aux_loss)
         self.OUT.__setitem__('past_key_values', past_kvs)
@@ -356,9 +472,7 @@ class MiniMindLM(PreTrainedModel):
             for seq in generated
         ]
         output = torch.cat(generated, dim=0)
-        res = output.view(input_ids.size(0), num_return_sequences, -1)
-        res = res.squeeze(0) if input_ids.size(0) == 1 else res
-        res = res.squeeze(1) if num_return_sequences == 1 else res
+        res = output.view(input_ids.size(0) * num_return_sequences, -1)
         return res
 
     def _stream(self, input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args):

scripts/chat_openai_api.py

@@ -7,13 +7,13 @@ client = OpenAI(
 stream = True
 conversation_history_origin = []
 conversation_history = conversation_history_origin.copy()
+history_messages_num = 2  # 设置为偶数（Q+A），为0则每次不携带历史对话进行独立QA
 while True:
-    conversation_history = conversation_history_origin.copy()
     query = input('[Q]: ')
     conversation_history.append({"role": "user", "content": query})
     response = client.chat.completions.create(
         model="minimind",
-        messages=conversation_history,
+        messages=conversation_history[-history_messages_num:],
         stream=stream
     )
     if not stream:

scripts/serve_openai_api.py

@@ -55,7 +55,7 @@ class ChatRequest(BaseModel):
     model: str
     messages: list
     temperature: float = 0.7
-    top_p: int = 0.92
+    top_p: float = 0.92
     max_tokens: int = 8192
     stream: bool = False
 

train_full_sft.py

@@ -1,4 +1,6 @@
 import os
+# 设置环境变量
+os.environ["WANDB_MODE"] = "offline"  # 或者使用 "dryrun"
 import platform
 import argparse
 import time
@@ -19,51 +21,57 @@ from model.model import MiniMindLM
 from model.LMConfig import LMConfig
 from model.dataset import SFTDataset
 
+
 warnings.filterwarnings('ignore')
 
-
+# 日志记录函数，用于打印训练信息。
 def Logger(content):
     if not ddp or dist.get_rank() == 0:
         print(content)
 
-
+# 学习率计算函数，用于计算当前学习率。
 def get_lr(current_step, total_steps, lr):
     return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
 
-
+# 训练一个epoch的函数，用于训练模型。
 def train_epoch(epoch, wandb):
-    loss_fct = nn.CrossEntropyLoss(reduction='none')
+    loss_fct = nn.CrossEntropyLoss(reduction='none') #交叉熵损失函数，用于计算损失。
     start_time = time.time()
     for step, (X, Y, loss_mask) in enumerate(train_loader):
+        # 将数据移动到指定设备。
         X = X.to(args.device)
         Y = Y.to(args.device)
         loss_mask = loss_mask.to(args.device)
+        # 计算当前学习率。
         lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)
+        # 更新学习率。
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
         with ctx:
-            res = model(X)
+            res = model(X) #获取输出
             loss = loss_fct(
                 res.logits.view(-1, res.logits.size(-1)),
                 Y.view(-1)
-            ).view(Y.size())
-
+            ).view(Y.size()) #计算损失
+            
+            # 计算损失
             loss = (loss * loss_mask).sum() / loss_mask.sum()
             loss += res.aux_loss
             loss = loss / args.accumulation_steps
 
-        scaler.scale(loss).backward()
+        scaler.scale(loss).backward() #用于处理混合精度训练。它的作用是自动缩放损失值，以防止在使用低精度（如 FP16）计算时出现数值不稳定的问题。
 
         if (step + 1) % args.accumulation_steps == 0:
-            scaler.unscale_(optimizer)
-            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            scaler.unscale_(optimizer) #PyTorch 自动混合精度(AMP)训练的一部分。它"反缩放"之前为防止在混合精度训练中出现下溢而缩放的梯度。
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) #应用梯度裁剪以防止梯度爆炸。它会缩放梯度，使其范数不超过args.grad_clip。
 
-            scaler.step(optimizer)
-            scaler.update()
+            scaler.step(optimizer) #使用优化器更新模型权重，但由缩放器控制以适应混合精度训练。
+            scaler.update() #根据本次迭代是否有梯度溢出来更新下一次迭代的缩放因子。
 
-            optimizer.zero_grad(set_to_none=True)
+            optimizer.zero_grad(set_to_none=True) #清空梯度。
 
+        # 如果达到日志记录间隔，则记录日志。
         if step % args.log_interval == 0:
             spend_time = time.time() - start_time
             Logger(
@@ -94,7 +102,7 @@ def train_epoch(epoch, wandb):
             torch.save(state_dict, ckp)
             model.train()
 
-
+# 初始化模型函数，用于初始化模型。
 def init_model(lm_config):
     tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
     model = MiniMindLM(lm_config)
@@ -106,7 +114,7 @@ def init_model(lm_config):
     model = model.to(args.device)
     return model, tokenizer
 
-
+# 初始化分布式模式函数，用于初始化分布式模式。
 def init_distributed_mode():
     if not ddp: return
     global ddp_local_rank, DEVICE
@@ -122,12 +130,12 @@ def init_distributed_mode():
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind Full SFT")
     parser.add_argument("--out_dir", type=str, default="out")
-    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--epochs", type=int, default=3)
     parser.add_argument("--batch_size", type=int, default=32)
     parser.add_argument("--learning_rate", type=float, default=5e-5)
     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("--use_wandb", action="store_true")
+    parser.add_argument("--use_wandb", default=True, action="store_true")
     parser.add_argument("--wandb_project", type=str, default="MiniMind-Full-SFT")
     parser.add_argument("--num_workers", type=int, default=1)
     parser.add_argument("--ddp", action="store_true")
@@ -137,11 +145,11 @@ if __name__ == "__main__":
     parser.add_argument("--log_interval", type=int, default=100)
     parser.add_argument("--save_interval", type=int, default=100)
     parser.add_argument('--local_rank', type=int, default=-1)
-    parser.add_argument('--dim', default=512, type=int)
-    parser.add_argument('--n_layers', default=8, type=int)
-    parser.add_argument('--max_seq_len', default=512, type=int)
+    parser.add_argument('--dim', default=1024, type=int) #模型维度，用于控制模型的大小。
+    parser.add_argument('--n_layers', default=24, type=int) #层数，用于控制模型层数。
+    parser.add_argument('--max_seq_len', default=1024, type=int) #最大序列长度，用于控制输入序列的最大长度。
     parser.add_argument('--use_moe', default=False, type=bool)
-    parser.add_argument("--data_path", type=str, default="./dataset/sft_mini_512.jsonl")
+    parser.add_argument("--data_path", type=str, default="./dataset/sft_1024.jsonl")
 
     args = parser.parse_args()
 
@@ -161,6 +169,7 @@ if __name__ == "__main__":
     torch.manual_seed(base_seed)
     torch.cuda.manual_seed(base_seed)
 
+    # 如果使用分布式模式，则初始化分布式模式。
     if ddp:
         init_distributed_mode()
         args.device = torch.device(DEVICE)
@@ -169,6 +178,7 @@ if __name__ == "__main__":
         # 同时设置 CUDA 的随机种子
         torch.cuda.manual_seed(base_seed + rank)
 
+    # 如果使用WandB，则初始化WandB。
     if args.use_wandb and (not ddp or ddp_local_rank == 0):
         import wandb
 
@@ -176,8 +186,10 @@ if __name__ == "__main__":
     else:
         wandb = None
 
+    # 初始化模型。
     model, tokenizer = init_model(lm_config)
 
+    # 初始化数据集。
     train_ds = SFTDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)
     train_sampler = DistributedSampler(train_ds) if ddp else None
     train_loader = DataLoader(
@@ -190,8 +202,8 @@ if __name__ == "__main__":
         sampler=train_sampler
     )
 
-    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
-    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
+    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))  #创建一个梯度缩放器(GradScaler)，用于混合精度训练。当模型使用半精度格式(float16或bfloat16)训练时启用，它帮助防止梯度下溢并提高训练效率。
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) # 创建AdamW优化器实例，负责更新模型参数。它接收模型的所有参数和指定的学习率作为输入。AdamW是Adam优化器的变体，增加了权重衰减的正则化。
 
     if ddp:
         model._ddp_params_and_buffers_to_ignore = {"pos_cis"}

train_pretrain.py

@@ -1,4 +1,6 @@
 import os
+# 设置环境变量
+os.environ["WANDB_MODE"] = "offline"  # 或者使用 "dryrun"
 import platform
 import argparse
 import time
@@ -23,47 +25,55 @@ warnings.filterwarnings('ignore')
 
 
 def Logger(content):
+    # 如果没有使用ddp或者ddp的主设备，那么就打印
     if not ddp or dist.get_rank() == 0:
         print(content)
 
 
 def get_lr(current_step, total_steps, lr):
+    # 更新学习率
+    # \text{get\_lr}(c, t, l) = \frac{l}{10} + 0.5 \cdot l \cdot \left(1 + \cos\left(\frac{\pi \cdot c}{t}\right)\right)
     return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
 
 
 def train_epoch(epoch, wandb):
-    loss_fct = nn.CrossEntropyLoss(reduction='none')
+    loss_fct = nn.CrossEntropyLoss(reduction='none') #交叉熵损失（Cross-Entropy Loss）；当 reduction='none' 时，nn.CrossEntropyLoss 不会对损失进行任何汇总操作，而是返回每个样本的单独损失值。
     start_time = time.time()
     for step, (X, Y, loss_mask) in enumerate(train_loader):
+        # 将数据加载到设备上
         X = X.to(args.device)
         Y = Y.to(args.device)
         loss_mask = loss_mask.to(args.device)
 
+        # 更新学习率
         lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
         with ctx:
-            res = model(X)
+            res = model(X) #获取输出
             loss = loss_fct(
                 res.logits.view(-1, res.logits.size(-1)),
                 Y.view(-1)
-            ).view(Y.size())
-            loss = (loss * loss_mask).sum() / loss_mask.sum()
-            loss += res.aux_loss
+            ).view(Y.size())#计算损失
+            loss = (loss * loss_mask).sum() / loss_mask.sum() #计算总的loss
+            # 为了批次堆叠进行的处理，真正的batch size为num gpu*batch size per gpu*accumulation steps
+            loss += res.aux_loss 
             loss = loss / args.accumulation_steps
 
-        scaler.scale(loss).backward()
+        scaler.scale(loss).backward() #用于处理混合精度训练。它的作用是自动缩放损失值，以防止在使用低精度（如 FP16）计算时出现数值不稳定的问题。
 
+        # 如果达到堆叠数目就进行处理
         if (step + 1) % args.accumulation_steps == 0:
-            scaler.unscale_(optimizer)
-            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            scaler.unscale_(optimizer) #PyTorch 自动混合精度(AMP)训练的一部分。它"反缩放"之前为防止在混合精度训练中出现下溢而缩放的梯度。
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) #应用梯度裁剪以防止梯度爆炸。它会缩放梯度，使其范数不超过args.grad_clip。
 
-            scaler.step(optimizer)
-            scaler.update()
+            scaler.step(optimizer) #使用优化器更新模型权重，但由缩放器控制以适应混合精度训练。
+            scaler.update() #根据本次迭代是否有梯度溢出来更新下一次迭代的缩放因子。
 
-            optimizer.zero_grad(set_to_none=True)
+            optimizer.zero_grad(set_to_none=True) #为下一次迭代清零所有梯度。set_to_none=True参数通过将梯度设置为None而不是零来提高内存效率。
 
+        # 打印日志
         if step % args.log_interval == 0:
             spend_time = time.time() - start_time
             Logger(
@@ -81,37 +91,41 @@ def train_epoch(epoch, wandb):
                            "lr": optimizer.param_groups[-1]['lr'],
                            "epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
 
+        # 保存模型
         if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
             model.eval()
             moe_path = '_moe' if lm_config.use_moe else ''
             ckp = f'{args.save_dir}/pretrain_{lm_config.dim}{moe_path}.pth'
 
             if isinstance(model, torch.nn.parallel.DistributedDataParallel):
-                state_dict = model.module.state_dict()
+                state_dict = model.module.state_dict() #获取模型参数
             else:
-                state_dict = model.state_dict()
+                state_dict = model.state_dict() #获取模型参数
 
-            torch.save(state_dict, ckp)
+            torch.save(state_dict, ckp) #只保存参数
             model.train()
 
 
 def init_model(lm_config):
+    # 加载tokenizer
     tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
+    # 加载模型
     model = MiniMindLM(lm_config).to(args.device)
+    # 打印模型参数
     Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
     return model, tokenizer
 
 
 def init_distributed_mode():
-    if not ddp: return
-    global ddp_local_rank, DEVICE
+    if not ddp: return #如果没有启用分布式数据并行(DDP)，直接返回，不执行任何操作。
+    global ddp_local_rank, DEVICE #声明这两个变量为全局变量，以便在函数外部也能访问它们。
 
-    dist.init_process_group(backend="nccl")
-    ddp_rank = int(os.environ["RANK"])
-    ddp_local_rank = int(os.environ["LOCAL_RANK"])
-    ddp_world_size = int(os.environ["WORLD_SIZE"])
-    DEVICE = f"cuda:{ddp_local_rank}"
-    torch.cuda.set_device(DEVICE)
+    dist.init_process_group(backend="nccl") #初始化分布式进程组，使用NCCL后端（NVIDIA Collective Communications Library），这是NVIDIA GPU之间通信的优化库。
+    ddp_rank = int(os.environ["RANK"]) #从环境变量获取当前进程的全局编号。
+    ddp_local_rank = int(os.environ["LOCAL_RANK"]) #从环境变量获取当前进程的本地编号。
+    ddp_world_size = int(os.environ["WORLD_SIZE"]) #从环境变量获取当前进程组中的进程总数。
+    DEVICE = f"cuda:{ddp_local_rank}" #根据本地编号选择GPU设备。
+    torch.cuda.set_device(DEVICE) #设置当前进程的GPU设备。
 
 
 # torchrun --nproc_per_node 2 1-pretrain.py
@@ -119,34 +133,35 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind Pretraining")
     parser.add_argument("--out_dir", type=str, default="out")
     # 若要以最快速度实现zero则epochs设置为1轮；否则应当利用有限的数据训练2~6个epochs。
-    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--epochs", type=int, default=3)
     parser.add_argument("--batch_size", type=int, default=32)
     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("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu") #如果GPU可用，则使用GPU，否则使用CPU。
     parser.add_argument("--dtype", type=str, default="bfloat16")
-    parser.add_argument("--use_wandb", action="store_true")
+    parser.add_argument("--use_wandb", default=True, action="store_true")
     parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain")
-    parser.add_argument("--num_workers", type=int, default=1)
+    parser.add_argument("--num_workers", type=int, default=8)
     parser.add_argument("--ddp", action="store_true")
-    parser.add_argument("--accumulation_steps", type=int, default=8)
-    parser.add_argument("--grad_clip", type=float, default=1.0)
-    parser.add_argument("--warmup_iters", type=int, default=0)
-    parser.add_argument("--log_interval", type=int, default=100)
-    parser.add_argument("--save_interval", type=int, default=100)
-    parser.add_argument('--local_rank', type=int, default=-1)
-    parser.add_argument('--dim', default=512, type=int)
-    parser.add_argument('--n_layers', default=8, type=int)
-    parser.add_argument('--max_seq_len', default=512, type=int)
-    parser.add_argument('--use_moe', default=False, type=bool)
-    parser.add_argument("--data_path", type=str, default="./dataset/pretrain_hq.jsonl")
+    parser.add_argument("--accumulation_steps", type=int, default=8) #梯度累积步数，用于控制梯度更新频率。
+    parser.add_argument("--grad_clip", type=float, default=1.0) #梯度裁剪阈值，用于防止梯度爆炸。
+    parser.add_argument("--warmup_iters", type=int, default=0) #预热迭代次数，用于控制学习率预热过程。
+    parser.add_argument("--log_interval", type=int, default=100) #日志打印间隔，用于控制日志打印的频率。
+    parser.add_argument("--save_interval", type=int, default=100) #模型保存间隔，用于控制模型保存的频率。
+    parser.add_argument('--local_rank', type=int, default=-1) #本地进程编号，用于分布式训练。
+    parser.add_argument('--dim', default=1024, type=int) #模型维度，用于控制模型的大小。
+    parser.add_argument('--n_layers', default=24, type=int) #层数，用于控制模型层数。
+    parser.add_argument('--max_seq_len', default=1024, type=int) #最大序列长度，用于控制输入序列的最大长度。
+    parser.add_argument('--use_moe', default=False, type=bool) #是否使用MOE，用于控制是否使用MOE。
+    parser.add_argument("--data_path", type=str, default="./dataset/pretrain_hq.jsonl") #数据路径，用于控制数据集的路径。
     args = parser.parse_args()
 
-    lm_config = LMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, use_moe=args.use_moe)
-    args.save_dir = os.path.join(args.out_dir)
-    os.makedirs(args.save_dir, exist_ok=True)
-    os.makedirs(args.out_dir, exist_ok=True)
-    tokens_per_iter = args.batch_size * lm_config.max_seq_len
-    device_type = "cuda" if "cuda" in args.device else "cpu"
+    lm_config = LMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, use_moe=args.use_moe) #创建LMConfig对象，用于控制模型配置。
+    args.save_dir = os.path.join(args.out_dir) #创建保存目录。
+    os.makedirs(args.save_dir, exist_ok=True) #创建保存目录。
+    os.makedirs(args.out_dir, exist_ok=True) #创建输出目录。
+    tokens_per_iter = args.batch_size * lm_config.max_seq_len #计算每个迭代步骤的token数量。
+    print(f"tokens_per_iter: {tokens_per_iter}")
+    device_type = "cuda" if "cuda" in args.device else "cpu" #确定设备类型。
 
     args.wandb_run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"