250427

README.md

@@ -64,7 +64,16 @@
 ### 👉**最近更新**
 
 <details close> 
-<summary> <b>2025-02-20 (newest 🎉)</b> </summary>
+<summary> <b>2025-04-27 (newest 🎉)</b> </summary>
+
+- 兼容性更新
+- 适配[「minimind仓库新特性」](https://github.com/jingyaogong/minimind/issues/370)
+- 规范化部分代码
+
+</details>
+
+<details close> 
+<summary> <b>2025-02-20</b> </summary>
 
 - MiniMind2-V伴随MiniMind2同步更新
 - 大幅减少所有冗余代码，规范代码格式

README_en.md

@@ -7,7 +7,6 @@
 
 <div align="center">
 
-![visitors](https://visitor-badge.laobi.icu/badge?page_id=jingyaogong/minimind-v)
 [![GitHub Repo stars](https://img.shields.io/github/stars/jingyaogong/minimind-v?style=social)](https://github.com/jingyaogong/minimind-v/stargazers)
 [![GitHub Code License](https://img.shields.io/github/license/jingyaogong/minimind-v?v=1)](LICENSE)
 [![GitHub last commit](https://img.shields.io/github/last-commit/jingyaogong/minimind-v)](https://github.com/jingyaogong/minimind-v/commits/master)
@@ -71,7 +70,17 @@ Is the training process difficult? Now, let's explore the answers and feel the j
 ### 👉**Recent Updates**
 
 <details close> 
-<summary> <b>2025-02-20 (newest 🎉)</b> </summary>
+<summary> <b>2025-04-27 (newest 🎉)</b> </summary>
+
+- Compatibility updates
+- Adapted to the new feature in the "minimind" repository
+- Standardized parts of the code
+
+</details>
+
+
+<details close> 
+<summary> <b>2025-02-20</b> </summary>
 
 - MiniMind2-V updated alongside MiniMind2
 - Significant reduction of all redundant code, standardized code format

dataset/lm_dataset.py

@@ -2,7 +2,7 @@ import json
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader
 import torch
-from .model_vlm import MiniMindVLM
+from model.model_vlm import MiniMindVLM
 import os
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -20,8 +20,8 @@ class VLMDataset(Dataset):
         self.max_length = max_length
         self.preprocess = preprocess
         self.image_token = image_special_token
-        self.bos_id = tokenizer('<s>assistant', add_special_tokens=False).input_ids
+        self.bos_id = tokenizer('<|im_start|>assistant', add_special_tokens=False).input_ids
-        self.eos_id = tokenizer('</s>', add_special_tokens=False).input_ids
+        self.eos_id = tokenizer('<|im_end|>', add_special_tokens=False).input_ids
 
     def __len__(self):
         return len(self.samples)

eval_vlm.py

@@ -5,12 +5,8 @@ import numpy as np
 import torch
 import warnings
 from PIL import Image
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
-from model.model_vlm import MiniMindVLM
+from model.model_vlm import MiniMindVLM, VLMConfig
-from model.VLMConfig import VLMConfig
-from transformers import logging as hf_logging
-
-hf_logging.set_verbosity_error()
 
 warnings.filterwarnings('ignore')
 
@@ -20,22 +16,23 @@ def count_parameters(model):
 
 
 def init_model(lm_config, device):
-    tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
+    tokenizer = AutoTokenizer.from_pretrained('./model')
     if args.load == 0:
         moe_path = '_moe' if args.use_moe else ''
         modes = {0: 'pretrain_vlm', 1: 'sft_vlm', 2: 'sft_vlm_multi'}
-        ckp = f'./{args.out_dir}/{modes[args.model_mode]}_{args.dim}{moe_path}.pth'
+        ckp = f'./{args.out_dir}/{modes[args.model_mode]}_{args.hidden_size}{moe_path}.pth'
-        model = MiniMindVLM(lm_config)
+        model = MiniMindVLM(lm_config, vision_model_path="./model/vision_model/clip-vit-base-patch16")
         state_dict = torch.load(ckp, map_location=device)
         model.load_state_dict({k: v for k, v in state_dict.items() if 'mask' not in k}, strict=False)
     else:
-        transformers_model_path = 'MiniMind2-V'
+        transformers_model_path = 'MiniMind2-Small-V'
         tokenizer = AutoTokenizer.from_pretrained(transformers_model_path)
         model = AutoModelForCausalLM.from_pretrained(transformers_model_path, trust_remote_code=True)
+        model.vision_encoder, model.processor = MiniMindVLM.get_vision_model("./model/vision_model/clip-vit-base-patch16")
 
     print(f'VLM参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
 
-    vision_model, preprocess = MiniMindVLM.get_vision_model()
+    vision_model, preprocess = model.vision_encoder, model.processor
     return model.eval().to(device), tokenizer, vision_model.eval().to(device), preprocess
 
 
@@ -56,62 +53,59 @@ if __name__ == "__main__":
     parser.add_argument('--temperature', default=0.65, type=float)
     parser.add_argument('--top_p', default=0.85, type=float)
     parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str)
-    # MiniMind2-Small (26M)：(dim=512, n_layers=8)
+    # MiniMind2-Small (26M)：(hidden_size=512, num_hidden_layers=8)
-    # MiniMind2 (104M)：(dim=768, n_layers=16)
+    # MiniMind2 (104M)：(hidden_size=768, num_hidden_layers=16)
-    parser.add_argument('--dim', default=512, type=int)
+    parser.add_argument('--hidden_size', default=512, type=int)
-    parser.add_argument('--n_layers', default=8, type=int)
+    parser.add_argument('--num_hidden_layers', default=8, type=int)
     parser.add_argument('--max_seq_len', default=8192, type=int)
     parser.add_argument('--use_moe', default=False, type=bool)
     # 默认单图推理，设置为2为多图推理
     parser.add_argument('--use_multi', default=1, type=int)
     parser.add_argument('--stream', default=True, type=bool)
-    parser.add_argument('--load', default=0, type=int, help="0: 原生torch权重，1: transformers加载")
+    parser.add_argument('--load', default=1, type=int, help="0: 原生torch权重，1: transformers加载")
     parser.add_argument('--model_mode', default=1, type=int,
                         help="0: Pretrain模型，1: SFT模型，2: SFT-多图模型 (beta拓展)")
     args = parser.parse_args()
 
-    lm_config = VLMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, use_moe=args.use_moe)
+    lm_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
-
+                          max_seq_len=args.max_seq_len, use_moe=args.use_moe)
     model, tokenizer, vision_model, preprocess = init_model(lm_config, args.device)
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
 
-    def chat_with_vlm(prompt, pixel_tensors, image_names):
+    def chat_with_vlm(prompt, pixel_values, image_names):
         messages = [{"role": "user", "content": prompt}]
-
         new_prompt = tokenizer.apply_chat_template(
             messages,
             tokenize=False,
             add_generation_prompt=True
         )[-args.max_seq_len + 1:]
 
+        inputs = tokenizer(
+            new_prompt,
+            return_tensors="pt",
+            truncation=True
+        ).to(args.device)
+
         print(f'[Image]: {image_names}')
-        with torch.no_grad():
+        print('🤖️: ', end='')
-            x = torch.tensor(tokenizer(new_prompt)['input_ids'], device=args.device).unsqueeze(0)
+        generated_ids = model.generate(
-            outputs = model.generate(
+            inputs["input_ids"],
-                x,
+            max_new_tokens=args.max_seq_len,
-                eos_token_id=tokenizer.eos_token_id,
+            num_return_sequences=1,
-                max_new_tokens=args.max_seq_len,
+            do_sample=True,
-                temperature=args.temperature,
+            attention_mask=inputs["attention_mask"],
-                top_p=args.top_p,
+            pad_token_id=tokenizer.pad_token_id,
-                stream=True,
+            eos_token_id=tokenizer.eos_token_id,
-                pad_token_id=tokenizer.pad_token_id,
+            streamer=streamer,
-                pixel_tensors=pixel_tensors
+            top_p=args.top_p,
-            )
+            temperature=args.temperature,
-            print('🤖️: ', end='')
+            pixel_values=pixel_values
-            try:
+        )
-                if not args.stream:
+
-                    print(tokenizer.decode(outputs.squeeze()[x.shape[1]:].tolist(), skip_special_tokens=True), end='')
+        response = tokenizer.decode(generated_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
-                else:
+        messages.append({"role": "assistant", "content": response})
-                    history_idx = 0
+        print('\n\n')
-                    for y in outputs:
-                        answer = tokenizer.decode(y[0].tolist(), skip_special_tokens=True)
-                        if (answer and answer[-1] == '<EFBFBD>') or not answer:
-                            continue
-                        print(answer[history_idx:], end='', flush=True)
-                        history_idx = len(answer)
-            except StopIteration:
-                print("No answer")
-            print('\n')
 
 
     # 单图推理：每1个图像单独推理

model/LMConfig.py

@@ -1,61 +0,0 @@
-from transformers import PretrainedConfig
-from typing import List
-
-
-class LMConfig(PretrainedConfig):
-    model_type = "minimind"
-
-    def __init__(
-            self,
-            dim: int = 512,
-            n_layers: int = 8,
-            n_heads: int = 8,
-            n_kv_heads: int = 2,
-            vocab_size: int = 6400,
-            hidden_dim: int = None,
-            multiple_of: int = 64,
-            norm_eps: float = 1e-5,
-            max_seq_len: int = 8192,
-            rope_theta: int = 1e6,
-            dropout: float = 0.0,
-            flash_attn: bool = True,
-            ####################################################
-            # Here are the specific configurations of MOE
-            # When use_moe is false, the following is invalid
-            ####################################################
-            use_moe: bool = False,
-            ####################################################
-            num_experts_per_tok: int = 2,
-            n_routed_experts: int = 4,
-            n_shared_experts: bool = True,
-            scoring_func: str = 'softmax',
-            aux_loss_alpha: float = 0.1,
-            seq_aux: bool = True,
-            norm_topk_prob: bool = True,
-            **kwargs,
-    ):
-        self.dim = dim
-        self.n_layers = n_layers
-        self.n_heads = n_heads
-        self.n_kv_heads = n_kv_heads
-        self.vocab_size = vocab_size
-        self.hidden_dim = hidden_dim
-        self.multiple_of = multiple_of
-        self.norm_eps = norm_eps
-        self.max_seq_len = max_seq_len
-        self.rope_theta = rope_theta
-        self.dropout = dropout
-        self.flash_attn = flash_attn
-        ####################################################
-        # Here are the specific configurations of MOE
-        # When use_moe is false, the following is invalid
-        ####################################################
-        self.use_moe = use_moe
-        self.num_experts_per_tok = num_experts_per_tok  # 每个token选择的专家数量
-        self.n_routed_experts = n_routed_experts  # 总的专家数量
-        self.n_shared_experts = n_shared_experts  # 共享专家
-        self.scoring_func = scoring_func  # 评分函数，默认为'softmax'
-        self.aux_loss_alpha = aux_loss_alpha  # 辅助损失的alpha参数
-        self.seq_aux = seq_aux  # 是否在序列级别上计算辅助损失
-        self.norm_topk_prob = norm_topk_prob  # 是否标准化top-k概率
-        super().__init__(**kwargs)

model/VLMConfig.py

@@ -1,16 +0,0 @@
-from .LMConfig import LMConfig
-from typing import List
-
-
-class VLMConfig(LMConfig):
-    model_type = "minimind-v"
-
-    def __init__(
-            self,
-            image_special_token: str = '@' * 196,
-            image_ids: List = [34] * 196,
-            **kwargs,
-    ):
-        self.image_special_token = image_special_token
-        self.image_ids = image_ids
-        super().__init__(**kwargs)

model/model.py

@@ -1,385 +0,0 @@
-import math
-import struct
-import inspect
-import time
-
-from .LMConfig import LMConfig
-from typing import Any, Optional, Tuple, List, Union
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torch import nn
-from transformers import PreTrainedModel
-from transformers.modeling_outputs import CausalLMOutputWithPast
-
-
-class RMSNorm(torch.nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
-
-    def _norm(self, x):
-        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
-
-    def forward(self, x):
-        return self.weight * self._norm(x.float()).type_as(x)
-
-
-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
-    freqs = torch.outer(t, freqs).float()  # type: ignore
-    pos_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
-    return pos_cis
-
-
-def apply_rotary_emb(xq, xk, pos_cis):
-    def unite_shape(pos_cis, x):
-        ndim = x.ndim
-        assert 0 <= 1 < ndim
-        assert pos_cis.shape == (x.shape[1], x.shape[-1])
-        shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
-        return pos_cis.view(*shape)
-
-    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
-    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
-    pos_cis = unite_shape(pos_cis, xq_)
-    xq_out = torch.view_as_real(xq_ * pos_cis).flatten(3)
-    xk_out = torch.view_as_real(xk_ * pos_cis).flatten(3)
-    return xq_out.type_as(xq), xk_out.type_as(xk)
-
-
-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
-    if n_rep == 1:
-        return x
-    return (
-        x[:, :, :, None, :]
-        .expand(bs, slen, n_kv_heads, n_rep, head_dim)
-        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
-    )
-
-
-class Attention(nn.Module):
-    def __init__(self, args: LMConfig):
-        super().__init__()
-        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
-        assert args.n_heads % self.n_kv_heads == 0
-        self.n_local_heads = args.n_heads
-        self.n_local_kv_heads = self.n_kv_heads
-        self.n_rep = self.n_local_heads // self.n_local_kv_heads
-        self.head_dim = args.dim // args.n_heads
-        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
-        self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
-        self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
-        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
-        self.attn_dropout = nn.Dropout(args.dropout)
-        self.resid_dropout = nn.Dropout(args.dropout)
-        self.dropout = args.dropout
-        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn
-        # print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
-        mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
-        mask = torch.triu(mask, diagonal=1)
-        self.register_buffer("mask", mask, persistent=False)
-
-    def forward(self,
-                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)
-
-        xq, xk = apply_rotary_emb(xq, xk, pos_cis)
-        # kv_cache实现
-        if past_key_value is not None:
-            xk = torch.cat([past_key_value[0], xk], dim=1)
-            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)
-        )
-        if self.flash and seq_len != 1:
-            dropout_p = self.dropout if self.training else 0.0
-            output = F.scaled_dot_product_attention(
-                xq, xk, xv,
-                attn_mask=None,
-                dropout_p=dropout_p,
-                is_causal=True
-            )
-        else:
-            scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
-            scores += self.mask[:, :, :seq_len, :seq_len]
-            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
-            scores = self.attn_dropout(scores)
-            output = scores @ xv
-
-        output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
-        output = self.resid_dropout(self.wo(output))
-        return output, past_kv
-
-
-class FeedForward(nn.Module):
-    def __init__(self, config: LMConfig):
-        super().__init__()
-        if config.hidden_dim is None:
-            hidden_dim = 4 * config.dim
-            hidden_dim = int(2 * hidden_dim / 3)
-            config.hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)
-        self.w1 = nn.Linear(config.dim, config.hidden_dim, bias=False)
-        self.w2 = nn.Linear(config.hidden_dim, config.dim, bias=False)
-        self.w3 = nn.Linear(config.dim, config.hidden_dim, bias=False)
-        self.dropout = nn.Dropout(config.dropout)
-
-    def forward(self, x):
-        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
-
-
-class MoEGate(nn.Module):
-    def __init__(self, config: LMConfig):
-        super().__init__()
-        self.config = config
-        self.top_k = config.num_experts_per_tok
-        self.n_routed_experts = config.n_routed_experts
-
-        self.scoring_func = config.scoring_func
-        self.alpha = config.aux_loss_alpha
-        self.seq_aux = config.seq_aux
-
-        self.norm_topk_prob = config.norm_topk_prob
-        self.gating_dim = config.dim
-        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
-        self.reset_parameters()
-
-    def reset_parameters(self) -> None:
-        import torch.nn.init as init
-        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
-
-    def forward(self, hidden_states):
-        bsz, seq_len, h = hidden_states.shape
-        hidden_states = hidden_states.view(-1, h)
-        logits = F.linear(hidden_states, self.weight, None)
-        if self.scoring_func == 'softmax':
-            scores = logits.softmax(dim=-1)
-        else:
-            raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')
-
-        topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
-
-        if self.top_k > 1 and self.norm_topk_prob:
-            denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
-            topk_weight = topk_weight / denominator
-
-        if self.training and self.alpha > 0.0:
-            scores_for_aux = scores
-            aux_topk = self.top_k
-            topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
-            if self.seq_aux:
-                scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
-                ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
-                ce.scatter_add_(1, topk_idx_for_aux_loss,
-                                torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_(
-                    seq_len * aux_topk / self.n_routed_experts)
-                aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha
-            else:
-                mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
-                ce = mask_ce.float().mean(0)
-                Pi = scores_for_aux.mean(0)
-                fi = ce * self.n_routed_experts
-                aux_loss = (Pi * fi).sum() * self.alpha
-        else:
-            aux_loss = 0
-        return topk_idx, topk_weight, aux_loss
-
-
-class MOEFeedForward(nn.Module):
-    def __init__(self, config: LMConfig):
-        super().__init__()
-        self.config = config
-        self.experts = nn.ModuleList([
-            FeedForward(config)
-            for _ in range(config.n_routed_experts)
-        ])
-        self.gate = MoEGate(config)
-        if config.n_shared_experts is not None:
-            self.shared_experts = FeedForward(config)
-
-    def forward(self, x):
-        identity = x
-        orig_shape = x.shape
-        bsz, seq_len, _ = x.shape
-        # 使用门控机制选择专家
-        topk_idx, topk_weight, aux_loss = self.gate(x)
-        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):
-                y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(y.dtype)  # 确保类型一致
-            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)
-        self.aux_loss = aux_loss
-        return y
-
-    @torch.no_grad()
-    def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
-        expert_cache = torch.zeros_like(x)
-        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]，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:
-                continue
-            expert = self.experts[i]
-            exp_token_idx = token_idxs[start_idx:end_idx]
-            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]])
-            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):
-        super().__init__()
-        self.n_heads = config.n_heads
-        self.dim = config.dim
-        self.head_dim = config.dim // config.n_heads
-        self.attention = Attention(config)
-
-        self.layer_id = layer_id
-        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)
-
-    def forward(self, x, pos_cis, past_key_value=None, use_cache=False):
-        h_attn, past_kv = self.attention(
-            self.attention_norm(x),
-            pos_cis,
-            past_key_value=past_key_value,
-            use_cache=use_cache
-        )
-        h = x + h_attn
-        out = h + self.feed_forward(self.ffn_norm(h))
-        return out, past_kv
-
-
-class MiniMindLM(PreTrainedModel):
-    config_class = LMConfig
-
-    def __init__(self, params: LMConfig = None):
-        self.params = params or LMConfig()
-        super().__init__(self.params)
-        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.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
-        self.register_buffer("pos_cis",
-                             precompute_pos_cis(dim=params.dim // params.n_heads, theta=params.rope_theta),
-                             persistent=False)
-        self.OUT = CausalLMOutputWithPast()
-
-    def forward(self,
-                input_ids: Optional[torch.Tensor] = None,
-                past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
-                use_cache: bool = False,
-                logits_to_keep: Union[int, torch.Tensor] = 0,
-                **args):
-        past_key_values = past_key_values or [None] * len(self.layers)
-        start_pos = args.get('start_pos', 0)
-        h = self.dropout(self.tok_embeddings(input_ids))
-        pos_cis = self.pos_cis[start_pos:start_pos + input_ids.size(1)]
-        past_kvs = []
-        for l, layer in enumerate(self.layers):
-            h, past_kv = layer(
-                h, pos_cis,
-                past_key_value=past_key_values[l],
-                use_cache=use_cache
-            )
-            past_kvs.append(past_kv)
-
-        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)
-        return self.OUT
-
-    @torch.inference_mode()
-    def generate(self, input_ids, eos_token_id=2, max_new_tokens=1024, temperature=0.75, top_p=0.90,
-                 stream=False, rp=1., use_cache=True, pad_token_id=0, num_return_sequences=1, **args):
-        # 流式生成
-        if stream:
-            return self._stream(input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)
-
-        # 直接生成
-        generated = []
-        for i in range(input_ids.size(0)):
-            non_pad = input_ids[i][input_ids[i] != pad_token_id].unsqueeze(0)
-            for _ in range(num_return_sequences):
-                out = self._stream(non_pad, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)
-                tokens_list = [tokens[:, -1:] for tokens in out]
-                gen = torch.cat(tokens_list, dim=-1) if tokens_list else non_pad
-                full_sequence = torch.cat([non_pad, gen], dim=-1)
-                generated.append(full_sequence)
-
-        max_length = max(seq.size(1) for seq in generated)
-        generated = [
-            torch.cat(
-                [seq, torch.full((1, max_length - seq.size(1)), pad_token_id, dtype=seq.dtype, device=seq.device)],
-                dim=-1)
-            for seq in generated
-        ]
-        output = torch.cat(generated, dim=0)
-        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):
-        start, first_seq, past_kvs = input_ids.shape[1], True, None
-        while input_ids.shape[1] < max_new_tokens - 1:
-            if first_seq or not use_cache:
-                out, first_seq = self(input_ids, past_key_values=past_kvs, use_cache=use_cache, **args), False
-            else:
-                out = self(input_ids[:, -1:], past_key_values=past_kvs, use_cache=use_cache,
-                           start_pos=input_ids.shape[1] - 1, **args)
-            logits, past_kvs = out.logits[:, -1, :], out.past_key_values
-            logits[:, list(set(input_ids.tolist()[0]))] /= rp
-            logits /= (temperature + 1e-9)
-            if top_p is not None and top_p < 1.0:
-                sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
-                sorted_probs = F.softmax(sorted_logits, dim=-1)
-                cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
-                sorted_indices_to_remove = cumulative_probs > top_p
-                sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
-                sorted_indices_to_remove[:, 0] = False
-                indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
-                logits[indices_to_remove] = -float('Inf')
-            input_ids_next = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
-            input_ids = torch.cat((input_ids, input_ids_next), dim=1)
-            yield input_ids[:, start:]
-            if input_ids_next.item() == eos_token_id:
-                break

model/model_minimind.py

@@ -0,0 +1,446 @@
+# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
+#                                             MiniMind Config
+# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
+
+from transformers import PretrainedConfig
+
+
+class MiniMindConfig(PretrainedConfig):
+    model_type = "minimind"
+
+    def __init__(
+            self,
+            dropout: float = 0.0,
+            bos_token_id: int = 1,
+            eos_token_id: int = 2,
+            hidden_act: str = 'silu',
+            hidden_size: int = 512,
+            intermediate_size: int = None,
+            max_position_embeddings: int = 32768,
+            num_attention_heads: int = 8,
+            num_hidden_layers: int = 8,
+            num_key_value_heads: int = 2,
+            vocab_size: int = 6400,
+            rms_norm_eps: float = 1e-05,
+            rope_theta: int = 1000000.0,
+            flash_attn: bool = True,
+            ####################################################
+            # Here are the specific configurations of MOE
+            # When use_moe is false, the following is invalid
+            ####################################################
+            use_moe: bool = False,
+            num_experts_per_tok: int = 2,
+            n_routed_experts: int = 4,
+            n_shared_experts: int = 1,
+            scoring_func: str = 'softmax',
+            aux_loss_alpha: float = 0.1,
+            seq_aux: bool = True,
+            norm_topk_prob: bool = True,
+            **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.dropout = dropout
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.hidden_act = hidden_act
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.max_position_embeddings = max_position_embeddings
+        self.num_attention_heads = num_attention_heads
+        self.num_hidden_layers = num_hidden_layers
+        self.num_key_value_heads = num_key_value_heads
+        self.vocab_size = vocab_size
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+        self.flash_attn = flash_attn
+        ####################################################
+        # Here are the specific configurations of MOE
+        # When use_moe is false, the following is invalid
+        ####################################################
+        self.use_moe = use_moe
+        self.num_experts_per_tok = num_experts_per_tok  # 每个token选择的专家数量
+        self.n_routed_experts = n_routed_experts  # 总的专家数量
+        self.n_shared_experts = n_shared_experts  # 共享专家
+        self.scoring_func = scoring_func  # 评分函数，默认为'softmax'
+        self.aux_loss_alpha = aux_loss_alpha  # 辅助损失的alpha参数
+        self.seq_aux = seq_aux  # 是否在序列级别上计算辅助损失
+        self.norm_topk_prob = norm_topk_prob  # 是否标准化top-k概率
+
+
+# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
+#                                             MiniMind Model
+# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
+
+import math
+import torch
+from torch import nn
+from transformers.activations import ACT2FN
+from typing import Optional, Tuple, List, Union
+import torch.nn.functional as F
+from transformers import PreTrainedModel, GenerationMixin, PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        return self.weight * self._norm(x.float()).type_as(x)
+
+
+def precompute_freqs_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)
+    freqs = torch.outer(t, freqs).float()
+    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1)
+    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1)
+    return freqs_cos, freqs_sin
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    def rotate_half(x):
+        return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
+
+    q_embed = (q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))
+    k_embed = (k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))
+    return q_embed, k_embed
+
+
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
+    bs, slen, num_key_value_heads, head_dim = x.shape
+    if n_rep == 1:
+        return x
+    return (
+        x[:, :, :, None, :]
+        .expand(bs, slen, num_key_value_heads, n_rep, head_dim)
+        .reshape(bs, slen, num_key_value_heads * n_rep, head_dim)
+    )
+
+
+class Attention(nn.Module):
+    def __init__(self, args: MiniMindConfig):
+        super().__init__()
+        self.num_key_value_heads = args.num_attention_heads if args.num_key_value_heads is None else args.num_key_value_heads
+        assert args.num_attention_heads % self.num_key_value_heads == 0
+        self.n_local_heads = args.num_attention_heads
+        self.n_local_kv_heads = self.num_key_value_heads
+        self.n_rep = self.n_local_heads // self.n_local_kv_heads
+        self.head_dim = args.hidden_size // args.num_attention_heads
+        self.q_proj = nn.Linear(args.hidden_size, args.num_attention_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(args.num_attention_heads * self.head_dim, args.hidden_size, bias=False)
+        self.attn_dropout = nn.Dropout(args.dropout)
+        self.resid_dropout = nn.Dropout(args.dropout)
+        self.dropout = args.dropout
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn
+        # print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
+
+    def forward(self,
+                x: torch.Tensor,
+                position_embeddings: Tuple[torch.Tensor, torch.Tensor],  # 修改为接收cos和sin
+                past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+                use_cache=False,
+                attention_mask: Optional[torch.Tensor] = None):
+        bsz, seq_len, _ = x.shape
+        xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(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)
+
+        cos, sin = position_embeddings
+        xq, xk = apply_rotary_pos_emb(xq, xk, cos[:seq_len], sin[:seq_len])
+
+        # kv_cache实现
+        if past_key_value is not None:
+            xk = torch.cat([past_key_value[0], xk], dim=1)
+            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)
+        )
+
+        if False and self.flash and seq_len != 1:
+            dropout_p = self.dropout if self.training else 0.0
+            attn_mask = None
+            if attention_mask is not None:
+                attn_mask = attention_mask.view(bsz, 1, 1, -1).expand(bsz, self.n_local_heads, seq_len, -1)
+                attn_mask = attn_mask.bool() if attention_mask is not None else None
+
+            output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=True)
+        else:
+            scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
+            scores = scores + torch.triu(
+                torch.full((seq_len, seq_len), float("-inf"), device=scores.device),
+                diagonal=1
+            ).unsqueeze(0).unsqueeze(0)  # scores+mask
+
+            if attention_mask is not None:
+                extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+                extended_attention_mask = (1.0 - extended_attention_mask) * -1e9
+                scores = scores + extended_attention_mask
+
+            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
+            scores = self.attn_dropout(scores)
+            output = scores @ xv
+
+        output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
+        output = self.resid_dropout(self.o_proj(output))
+        return output, past_kv
+
+
+class FeedForward(nn.Module):
+    def __init__(self, config: MiniMindConfig):
+        super().__init__()
+        if config.intermediate_size is None:
+            intermediate_size = int(config.hidden_size * 8 / 3)
+            config.intermediate_size = 64 * ((intermediate_size + 64 - 1) // 64)
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.dropout = nn.Dropout(config.dropout)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        return self.dropout(self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)))
+
+
+class MoEGate(nn.Module):
+    def __init__(self, config: MiniMindConfig):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.n_routed_experts = config.n_routed_experts
+
+        self.scoring_func = config.scoring_func
+        self.alpha = config.aux_loss_alpha
+        self.seq_aux = config.seq_aux
+
+        self.norm_topk_prob = config.norm_topk_prob
+        self.gating_dim = config.hidden_size
+        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        import torch.nn.init as init
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+
+    def forward(self, hidden_states):
+        bsz, seq_len, h = hidden_states.shape
+        hidden_states = hidden_states.view(-1, h)
+        logits = F.linear(hidden_states, self.weight, None)
+        if self.scoring_func == 'softmax':
+            scores = logits.softmax(dim=-1)
+        else:
+            raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')
+
+        topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
+
+        if self.top_k > 1 and self.norm_topk_prob:
+            denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
+            topk_weight = topk_weight / denominator
+
+        if self.training and self.alpha > 0.0:
+            scores_for_aux = scores
+            aux_topk = self.top_k
+            topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
+            if self.seq_aux:
+                scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
+                ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
+                ce.scatter_add_(1, topk_idx_for_aux_loss,
+                                torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_(
+                    seq_len * aux_topk / self.n_routed_experts)
+                aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha
+            else:
+                mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
+                ce = mask_ce.float().mean(0)
+                Pi = scores_for_aux.mean(0)
+                fi = ce * self.n_routed_experts
+                aux_loss = (Pi * fi).sum() * self.alpha
+        else:
+            aux_loss = 0
+        return topk_idx, topk_weight, aux_loss
+
+
+class MOEFeedForward(nn.Module):
+    def __init__(self, config: MiniMindConfig):
+        super().__init__()
+        self.config = config
+        self.experts = nn.ModuleList([
+            FeedForward(config)
+            for _ in range(config.n_routed_experts)
+        ])
+        self.gate = MoEGate(config)
+        if config.n_shared_experts > 0:
+            self.shared_experts = nn.ModuleList([
+                FeedForward(config)
+                for _ in range(config.n_shared_experts)
+            ])
+
+    def forward(self, x):
+        identity = x
+        orig_shape = x.shape
+        bsz, seq_len, _ = x.shape
+        # 使用门控机制选择专家
+        topk_idx, topk_weight, aux_loss = self.gate(x)
+        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):
+                y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(y.dtype)  # 确保类型一致
+            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 > 0:
+            for expert in self.shared_experts:
+                y = y + expert(identity)
+        self.aux_loss = aux_loss
+        return y
+
+    @torch.no_grad()
+    def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
+        expert_cache = torch.zeros_like(x)
+        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]，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:
+                continue
+            expert = self.experts[i]
+            exp_token_idx = token_idxs[start_idx:end_idx]
+            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]])
+            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: MiniMindConfig):
+        super().__init__()
+        self.num_attention_heads = config.num_attention_heads
+        self.hidden_size = config.hidden_size
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.self_attn = Attention(config)
+
+        self.layer_id = layer_id
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
+
+    def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
+        residual = hidden_states
+        hidden_states, present_key_value = self.self_attn(
+            self.input_layernorm(hidden_states), position_embeddings,
+            past_key_value, use_cache, attention_mask
+        )
+        hidden_states += residual
+        hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
+        return hidden_states, present_key_value
+
+
+class MiniMindModel(nn.Module):
+    def __init__(self, config: MiniMindConfig):
+        super().__init__()
+        self.config = config
+        self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.dropout = nn.Dropout(config.dropout)
+        self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.hidden_size // config.num_attention_heads,
+                                                    end=config.max_position_embeddings, theta=config.rope_theta)
+        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
+        self.register_buffer("freqs_sin", freqs_sin, persistent=False)
+
+    def forward(self,
+                input_ids: Optional[torch.Tensor] = None,
+                attention_mask: Optional[torch.Tensor] = None,
+                past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
+                use_cache: bool = False,
+                **kwargs):
+        batch_size, seq_length = input_ids.shape
+        past_key_values = past_key_values or [None] * len(self.layers)
+        start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
+
+        hidden_states = self.dropout(self.embed_tokens(input_ids))
+
+        position_embeddings = (
+            self.freqs_cos[start_pos:start_pos + seq_length],
+            self.freqs_sin[start_pos:start_pos + seq_length]
+        )
+
+        presents = []
+        for layer_idx, (layer, past_key_value) in enumerate(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
+            )
+            presents.append(present)
+
+        hidden_states = self.norm(hidden_states)
+
+        aux_loss = sum(
+            layer.mlp.aux_loss
+            for layer in self.layers
+            if isinstance(layer.mlp, MOEFeedForward)
+        )
+
+        return hidden_states, presents, aux_loss
+
+
+class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = MiniMindConfig
+
+    def __init__(self, config: MiniMindConfig = None):
+        self.config = config or MiniMindConfig()
+        super().__init__(self.config)
+        self.model = MiniMindModel(self.config)
+        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
+        self.model.embed_tokens.weight = self.lm_head.weight
+        self.OUT = CausalLMOutputWithPast()
+
+    def forward(self,
+                input_ids: Optional[torch.Tensor] = None,
+                attention_mask: Optional[torch.Tensor] = None,
+                past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
+                use_cache: bool = False,
+                logits_to_keep: Union[int, torch.Tensor] = 0,
+                **args):
+        h, past_kvs, aux_loss = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            **args
+        )
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(h[:, slice_indices, :])
+        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)
+        return self.OUT

model/model_vlm.py

@@ -1,21 +1,36 @@
-from .VLMConfig import VLMConfig
+import os
-from .model import *
+
+import torch
+import warnings
+from .model_minimind import *
 from typing import Optional, Tuple, List
 from torch import nn
-import warnings
 from transformers import CLIPProcessor, CLIPModel
-import torch
+from typing import List
 
 warnings.filterwarnings('ignore')
 
 
+class VLMConfig(MiniMindConfig):
+    model_type = "minimind-v"
+
+    def __init__(
+            self,
+            image_special_token: str = '@' * 196,
+            image_ids: List = [34] * 196,
+            **kwargs,
+    ):
+        self.image_special_token = image_special_token
+        self.image_ids = image_ids
+        super().__init__(**kwargs)
+
 class VisionProj(nn.Module):
-    def __init__(self, ve_dim=768, lm_dim=512):
+    def __init__(self, ve_hidden_size=768, hidden_size=512):
         super().__init__()
-        self.ve_dim = ve_dim
+        self.ve_hidden_size = ve_hidden_size
-        self.lm_dim = lm_dim
+        self.hidden_size = hidden_size
         self.vision_proj = nn.Sequential(
-            nn.Linear(self.ve_dim, self.lm_dim)
+            nn.Linear(self.ve_hidden_size, self.hidden_size)
         )
 
     def forward(self, image_encoders):
@@ -24,18 +39,22 @@ class VisionProj(nn.Module):
 
 
 # 继承自语言模型
-class MiniMindVLM(MiniMindLM):
+class MiniMindVLM(MiniMindForCausalLM):
     config_class = VLMConfig
 
-    def __init__(self, params: VLMConfig = None):
+    def __init__(self, params: VLMConfig = None, vision_model_path="./model/vision_model/clip-vit-base-patch16"):
         super().__init__(params)
         if not params: params = VLMConfig()
         self.params = params
-        self.vision_encoder, self.processor = self.__class__.get_vision_model()
+        self.vision_encoder, self.processor = self.__class__.get_vision_model(vision_model_path)
-        self.vision_proj = VisionProj(lm_dim=params.dim)
+        self.vision_proj = VisionProj(hidden_size=params.hidden_size)
 
     @staticmethod
-    def get_vision_model(model_path="./model/vision_model/clip-vit-base-patch16"):
+    def get_vision_model(model_path: str):
+        from transformers import logging as hf_logging
+        hf_logging.set_verbosity_error()
+        if not os.path.exists(model_path):
+            return None, None
         model = CLIPModel.from_pretrained(model_path)
         processor = CLIPProcessor.from_pretrained(model_path)
         # 冻结 vision_encoder 的所有参数
@@ -93,38 +112,57 @@ class MiniMindVLM(MiniMindLM):
 
     def forward(self,
                 input_ids: Optional[torch.Tensor] = None,
+                attention_mask: Optional[torch.Tensor] = None,
                 past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
                 use_cache: bool = False,
+                logits_to_keep: Union[int, torch.Tensor] = 0,
+                pixel_values: Optional[torch.FloatTensor] = None,
                 **args):
-        start_pos = args.get('start_pos', 0)
+        batch_size, seq_length = input_ids.shape
-        pixel_tensors = args.get('pixel_tensors', None)
+        past_key_values = past_key_values or [None] * len(self.model.layers)
-        h = self.tok_embeddings(input_ids)
+        start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
 
-        if pixel_tensors is not None and start_pos == 0:
+        hidden_states = self.model.dropout(self.model.embed_tokens(input_ids))
-            if len(pixel_tensors.shape) == 6:
+
-                pixel_tensors = pixel_tensors.squeeze(2)
+        if pixel_values is not None and start_pos == 0:
-            bs, num, c, im_h, im_w = pixel_tensors.shape
+            if len(pixel_values.shape) == 6:
+                pixel_values = pixel_values.squeeze(2)
+            bs, num, c, im_h, im_w = pixel_values.shape
             stack_dim = 1 if bs > 1 else 0
             vision_tensors = torch.stack([
-                MiniMindVLM.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.vision_encoder)
+                MiniMindVLM.get_image_embeddings(pixel_values[:, i, :, :, :], self.vision_encoder)
                 for i in range(num)
             ], dim=stack_dim)
-            h = self.count_vision_proj(tokens=input_ids, h=h, vision_tensors=vision_tensors, seqlen=input_ids.shape[1])
+            hidden_states = self.count_vision_proj(tokens=input_ids, h=hidden_states, vision_tensors=vision_tensors,
+                                                   seqlen=input_ids.shape[1])
 
-        pos_cis = self.pos_cis[start_pos:start_pos + input_ids.shape[1]]
+        position_embeddings = (
-        past_kvs = []
+            self.model.freqs_cos[start_pos:start_pos + seq_length],
-        for l, layer in enumerate(self.layers):
+            self.model.freqs_sin[start_pos:start_pos + seq_length]
-            h, past_kv = layer(
+        )
-                h, pos_cis,
+
-                past_key_value=past_key_values[l] if past_key_values else None,
+        presents = []
-                use_cache=use_cache
+        for layer_idx, (layer, past_key_value) in enumerate(zip(self.model.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
             )
-            past_kvs.append(past_kv)
+            presents.append(present)
 
-        logits = self.output(self.norm(h))
+        hidden_states = self.model.norm(hidden_states)
-        aux_loss = sum(l.feed_forward.aux_loss for l in self.layers if isinstance(l.feed_forward, MOEFeedForward))
 
+        aux_loss = sum(
+            layer.mlp.aux_loss
+            for layer in self.model.layers
+            if isinstance(layer.mlp, MOEFeedForward)
+        )
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        self.OUT.__setitem__('last_hidden_state', hidden_states)
         self.OUT.__setitem__('logits', logits)
         self.OUT.__setitem__('aux_loss', aux_loss)
-        self.OUT.__setitem__('past_key_values', past_kvs)
+        self.OUT.__setitem__('past_key_values', presents)
         return self.OUT

model/tokenizer_config.json

@@ -4,7 +4,7 @@
   "add_prefix_space": false,
   "added_tokens_decoder": {
     "0": {
-      "content": "<unk>",
+      "content": "<|endoftext|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -12,7 +12,7 @@
       "special": true
     },
     "1": {
-      "content": "<s>",
+      "content": "<|im_start|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -20,7 +20,7 @@
       "special": true
     },
     "2": {
-      "content": "</s>",
+      "content": "<|im_end|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -29,15 +29,15 @@
     }
   },
   "additional_special_tokens": [],
-  "bos_token": "<s>",
+  "bos_token": "<|im_start|>",
   "clean_up_tokenization_spaces": false,
-  "eos_token": "</s>",
+  "eos_token": "<|im_end|>",
   "legacy": true,
   "model_max_length": 32768,
-  "pad_token": "<unk>",
+  "pad_token": "<|endoftext|>",
   "sp_model_kwargs": {},
   "spaces_between_special_tokens": false,
   "tokenizer_class": "PreTrainedTokenizerFast",
-  "unk_token": "<unk>",
+  "unk_token": "<|endoftext|>",
-  "chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{{ '<s>system\\n' + system_message + '</s>\\n' }}{% else %}{{ '<s>system\\n你是 MiniMind，是一个有用的人工智能助手。</s>\\n' }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<s>user\\n' + content + '</s>\\n<s>assistant\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '</s>' + '\\n' }}{% endif %}{% endfor %}"
+  "chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{{ '<|im_start|>system\\n' + system_message + '<|im_end|>\\n' }}{% else %}{{ '<|im_start|>system\\nYou are a helpful assistant<|im_end|>\\n' }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<|im_start|>user\\n' + content + '<|im_end|>\\n<|im_start|>assistant\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '<|im_end|>' + '\\n' }}{% endif %}{% endfor %}"
 }

scripts/convert_vlm.py

@@ -0,0 +1,44 @@
+import os
+import sys
+
+__package__ = "scripts"
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+import torch
+import warnings
+from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaConfig, LlamaForCausalLM
+from model.model_vlm import MiniMindVLM, VLMConfig
+
+warnings.filterwarnings('ignore', category=UserWarning)
+
+
+def convert_torch2transformers_minimind(torch_path, transformers_path, dtype=torch.bfloat16):
+    VLMConfig.register_for_auto_class()
+    MiniMindVLM.register_for_auto_class("AutoModelForCausalLM")
+    lm_model = MiniMindVLM(lm_config, vision_model_path="../model/vision_model/clip-vit-base-patch16")
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    state_dict = torch.load(torch_path, map_location=device)
+    lm_model.load_state_dict(state_dict, strict=False)
+    lm_model = lm_model.to(dtype)  # 转换模型权重精度
+    model_params = sum(p.numel() for p in lm_model.parameters() if p.requires_grad)
+    print(f'模型参数: {model_params / 1e6} 百万 = {model_params / 1e9} B (Billion)')
+    del lm_model.vision_encoder
+    lm_model.save_pretrained(transformers_path, safe_serialization=False)
+    tokenizer = AutoTokenizer.from_pretrained('../model/')
+    tokenizer.save_pretrained(transformers_path)
+    print(f"模型已保存为 Transformers-MiniMind-V 格式: {transformers_path}")
+
+
+def convert_transformers2torch(transformers_path, torch_path):
+    model = AutoModelForCausalLM.from_pretrained(transformers_path, trust_remote_code=True)
+    torch.save(model.state_dict(), torch_path)
+    print(f"模型已保存为 PyTorch 格式: {torch_path}")
+
+
+if __name__ == '__main__':
+    lm_config = VLMConfig(hidden_size=768, num_hidden_layers=16, max_seq_len=8192, use_moe=False)
+
+    torch_path = f"../out/sft_vlm_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth"
+
+    transformers_path = '../MiniMind2-V'
+
+    convert_torch2transformers_minimind(torch_path, transformers_path)

scripts/web_demo_vlm.py

@@ -1,11 +1,17 @@
+import os
+import sys
+
+__package__ = "scripts"
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 import argparse
 import torch
 import warnings
 import gradio as gr
+from queue import Queue
+from threading import Thread
 from PIL import Image
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
-from model.model_vlm import MiniMindVLM
+from model.model_vlm import MiniMindVLM, VLMConfig
-from model.VLMConfig import VLMConfig
 from transformers import logging as hf_logging
 
 hf_logging.set_verbosity_error()
@@ -13,29 +19,42 @@ warnings.filterwarnings('ignore')
 
 
 def init_model(lm_config):
-    tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
+    tokenizer = AutoTokenizer.from_pretrained('../model')
 
     if args.load == 0:
         moe_path = '_moe' if lm_config.use_moe else ''
-        ckp = f'./out/sft_vlm_{lm_config.dim}{moe_path}.pth'
+        ckp = f'../out/sft_vlm_{lm_config.hidden_size}{moe_path}.pth'
-        model = MiniMindVLM(lm_config)
+        model = MiniMindVLM(lm_config, vision_model_path="../model/vision_model/clip-vit-base-patch16")
         state_dict = torch.load(ckp, map_location=args.device)
         model.load_state_dict({k: v for k, v in state_dict.items() if 'mask' not in k}, strict=False)
     else:
-        transformers_model_path = './MiniMind2-V'
+        transformers_model_path = '../MiniMind2-V'
         tokenizer = AutoTokenizer.from_pretrained(transformers_model_path)
         model = AutoModelForCausalLM.from_pretrained(transformers_model_path, trust_remote_code=True)
+        model.vision_encoder, model.processor = MiniMindVLM.get_vision_model("../model/vision_model/clip-vit-base-patch16")
 
     print(f'VLM参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
 
-    vision_model, preprocess = MiniMindVLM.get_vision_model()
+    vision_model, preprocess = model.vision_encoder, model.processor
     return model.eval().to(args.device), tokenizer, vision_model.to(args.device), preprocess
 
 
+class CustomStreamer(TextStreamer):
+    def __init__(self, tokenizer, queue):
+        super().__init__(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        self.queue = queue
+        self.tokenizer = tokenizer
+
+    def on_finalized_text(self, text: str, stream_end: bool = False):
+        self.queue.put(text)
+        if stream_end:
+            self.queue.put(None)
+
+
 def chat(prompt, current_image_path):
     global temperature, top_p
     image = Image.open(current_image_path).convert('RGB')
-    pixel_tensors = MiniMindVLM.image2tensor(image, preprocess).to(model.device).unsqueeze(0)
+    pixel_values = MiniMindVLM.image2tensor(image, preprocess).to(model.device).unsqueeze(0)
 
     prompt = f'{lm_config.image_special_token}\n{prompt}'
     messages = [{"role": "user", "content": prompt}]
@@ -45,32 +64,37 @@ def chat(prompt, current_image_path):
         tokenize=False,
         add_generation_prompt=True
     )[-args.max_seq_len + 1:]
+
     with torch.no_grad():
-        x = torch.tensor(tokenizer(new_prompt)['input_ids'], device=model.device).unsqueeze(0)
+        inputs = tokenizer(
-        outputs = model.generate(
+            new_prompt,
-            x,
+            return_tensors="pt",
-            eos_token_id=tokenizer.eos_token_id,
+            truncation=True
-            max_new_tokens=args.max_seq_len,
+        ).to(args.device)
-            temperature=temperature,
+        queue = Queue()
-            top_p=top_p,
+        streamer = CustomStreamer(tokenizer, queue)
-            stream=True,
+
-            pad_token_id=tokenizer.pad_token_id,
+        def _generate():
-            pixel_tensors=pixel_tensors
+            model.generate(
-        )
+                inputs.input_ids,
-        try:
+                max_new_tokens=args.max_seq_len,
-            if not args.stream:
+                do_sample=True,
-                print(tokenizer.decode(outputs.squeeze()[x.shape[1]:].tolist(), skip_special_tokens=True), end='')
+                temperature=temperature,
-            else:
+                top_p=top_p,
-                history_idx = 0
+                attention_mask=inputs.attention_mask,
-                for y in outputs:
+                pad_token_id=tokenizer.pad_token_id,
-                    answer = tokenizer.decode(y[0].tolist(), skip_special_tokens=True)
+                eos_token_id=tokenizer.eos_token_id,
-                    if (answer and answer[-1] == '<EFBFBD>') or not answer:
+                streamer=streamer,
-                        continue
+                pixel_values=pixel_values
-                    yield answer[history_idx:]
+            )
-                    history_idx = len(answer)
+
-        except StopIteration:
+        Thread(target=_generate).start()
-            print("No answer")
+
-        print('\n')
+        while True:
+            text = queue.get()
+            if text is None:
+                break
+            yield text
 
 
 def launch_gradio_server(server_name="0.0.0.0", server_port=7788):
@@ -163,14 +187,15 @@ if __name__ == '__main__':
     parser.add_argument('--temperature', default=0.65, type=float)
     parser.add_argument('--top_p', default=0.85, type=float)
     parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str)
-    parser.add_argument('--dim', default=768, type=int)
+    parser.add_argument('--hidden_size', default=512, type=int)
-    parser.add_argument('--n_layers', default=16, type=int)
+    parser.add_argument('--num_hidden_layers', default=8, type=int)
     parser.add_argument('--max_seq_len', default=8192, type=int)
     parser.add_argument('--use_moe', default=False, type=bool)
     parser.add_argument('--stream', default=True, type=bool)
-    parser.add_argument('--load', default=0, type=int, help="0: 原生torch权重，1: transformers加载")
+    parser.add_argument('--load', default=1, type=int, help="0: 原生torch权重，1: transformers加载")
     args = parser.parse_args()
 
-    lm_config = VLMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, use_moe=args.use_moe)
+    lm_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
+                          max_seq_len=args.max_seq_len, use_moe=args.use_moe)
     model, tokenizer, vision_model, preprocess = init_model(lm_config)
     launch_gradio_server(server_name="0.0.0.0", server_port=8888)

trainer/train_pretrain_vlm.py

@@ -1,26 +1,24 @@
-import os
-import platform
 import argparse
 import time
 import math
 import warnings
-import json
 
-import pandas as pd
+warnings.filterwarnings('ignore')
+import os
+import sys
 import torch
-import torch.nn.functional as F
 import torch.distributed as dist
-from contextlib import nullcontext
 
+__package__ = "trainer"
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from contextlib import nullcontext
 from torch import optim, nn
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
 from transformers import AutoTokenizer, AutoModel
-from model.model_vlm import MiniMindVLM
+from model.model_vlm import MiniMindVLM, VLMConfig
-from model.VLMConfig import VLMConfig
+from dataset.lm_dataset import VLMDataset
-from model.dataset import VLMDataset
-
-warnings.filterwarnings('ignore')
 
 
 def Logger(content):
@@ -35,17 +33,17 @@ def get_lr(current_step, total_steps, lr):
 def train_epoch(epoch, wandb):
     loss_fct = nn.CrossEntropyLoss(reduction='none')
     start_time = time.time()
-    for step, (X, Y, loss_mask, pixel_tensors) in enumerate(train_loader):
+    for step, (X, Y, loss_mask, pixel_values) in enumerate(train_loader):
         X = X.to(args.device)
         Y = Y.to(args.device)
         loss_mask = loss_mask.to(args.device)
-        pixel_tensors = pixel_tensors.to(args.device)
+        pixel_values = pixel_values.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, pixel_tensors=pixel_tensors)
+            res = model(X, pixel_values=pixel_values)
             loss = loss_fct(
                 res.logits.view(-1, res.logits.size(-1)),
                 Y.view(-1)
@@ -86,7 +84,7 @@ def train_epoch(epoch, wandb):
         if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
             model.eval()
             moe_path = '_moe' if model_config.use_moe else ''
-            ckp = f'{args.save_dir}/pretrain_vlm_{model_config.dim}{moe_path}.pth'
+            ckp = f'{args.save_dir}/pretrain_vlm_{model_config.hidden_size}{moe_path}.pth'
             if isinstance(model, torch.nn.parallel.DistributedDataParallel):
                 state_dict = model.module.state_dict()
             else:
@@ -94,16 +92,17 @@ def train_epoch(epoch, wandb):
             clean_state_dict = {
                 key: value for key, value in state_dict.items() if not key.startswith('vision_encoder.')
             }
+            clean_state_dict = {k: v.half() for k, v in clean_state_dict.items()}  # 半精度保存
             torch.save(clean_state_dict, ckp)
             model.train()
 
 
 def init_model(model_config: VLMConfig):
-    tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
+    tokenizer = AutoTokenizer.from_pretrained('../model', use_fast=True)
     moe_path = '_moe' if model_config.use_moe else ''
     # 加载纯语言模型权重
-    ckp = f'./out/lm_{model_config.dim}{moe_path}.pth'
+    ckp = f'{args.save_dir}/llm_{model_config.hidden_size}{moe_path}.pth'
-    model = MiniMindVLM(model_config)
+    model = MiniMindVLM(model_config, vision_model_path="../model/vision_model/clip-vit-base-patch16")
     state_dict = torch.load(ckp, map_location=args.device)
     model.load_state_dict(state_dict, strict=False)
 
@@ -111,16 +110,10 @@ def init_model(model_config: VLMConfig):
     for name, param in model.named_parameters():
         if 'vision_proj' not in name:
             param.requires_grad = False
-    # 可训练
-    if hasattr(model, "layers"):
-        last_two_layers = model.layers[-1:]
-        for layer in last_two_layers:
-            for param in layer.parameters():
-                param.requires_grad = True
 
     Logger(f'VLM可训练参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
 
-    _, preprocess = MiniMindVLM.get_vision_model()
+    _, preprocess = model.vision_encoder, model.processor
     return model.to(args.device), tokenizer, preprocess
 
 
@@ -138,7 +131,7 @@ def init_distributed_mode():
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind-V Pretrain")
-    parser.add_argument("--out_dir", type=str, default="out")
+    parser.add_argument("--out_dir", type=str, default="../out")
     parser.add_argument("--epochs", type=int, default=4)
     parser.add_argument("--batch_size", type=int, default=16)
     parser.add_argument("--learning_rate", type=float, default=4e-4)
@@ -147,8 +140,8 @@ if __name__ == "__main__":
     parser.add_argument("--use_wandb", default=False, action="store_true")
     parser.add_argument("--wandb_project", type=str, default="MiniMind-V")
     parser.add_argument("--num_workers", type=int, default=8)
-    parser.add_argument("--data_path", type=str, default="./dataset/pretrain_data.jsonl")
+    parser.add_argument("--data_path", type=str, default="../dataset/pretrain_data.jsonl")
-    parser.add_argument("--images_path", type=str, default="./dataset/pretrain_images")
+    parser.add_argument("--images_path", type=str, default="../dataset/pretrain_images")
     parser.add_argument("--ddp", action="store_true")
     parser.add_argument("--accumulation_steps", type=int, default=1)
     parser.add_argument("--grad_clip", type=float, default=1.0)
@@ -156,13 +149,14 @@ 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('--hidden_size', default=512, type=int)
-    parser.add_argument('--n_layers', default=8, type=int)
+    parser.add_argument('--num_hidden_layers', default=8, type=int)
     parser.add_argument('--max_seq_len', default=640, type=int)
     parser.add_argument('--use_moe', default=False, type=bool)
     args = parser.parse_args()
 
-    model_config = VLMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len)
+    model_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
+                             max_seq_len=args.max_seq_len)
     max_seq_len = model_config.max_seq_len
     args.save_dir = os.path.join(args.out_dir)
     os.makedirs(args.save_dir, exist_ok=True)

trainer/train_sft_vlm.py

@@ -1,24 +1,22 @@
-import os
-import platform
 import argparse
 import time
 import math
 import warnings
-import json
+import os
-
+import sys
-import pandas as pd
 import torch
-import torch.nn.functional as F
 import torch.distributed as dist
-from contextlib import nullcontext
 
+__package__ = "trainer"
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from contextlib import nullcontext
 from torch import optim, nn
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, DistributedSampler
 from transformers import AutoTokenizer, AutoModel
-from model.model_vlm import MiniMindVLM
+from model.model_vlm import MiniMindVLM, VLMConfig
-from model.VLMConfig import VLMConfig
+from dataset.lm_dataset import VLMDataset
-from model.dataset import *
 
 warnings.filterwarnings('ignore')
 
@@ -35,17 +33,17 @@ def get_lr(current_step, total_steps, lr):
 def train_epoch(epoch, wandb):
     loss_fct = nn.CrossEntropyLoss(reduction='none')
     start_time = time.time()
-    for step, (X, Y, loss_mask, pixel_tensors) in enumerate(train_loader):
+    for step, (X, Y, loss_mask, pixel_values) in enumerate(train_loader):
         X = X.to(args.device)
         Y = Y.to(args.device)
         loss_mask = loss_mask.to(args.device)
-        pixel_tensors = pixel_tensors.to(args.device)
+        pixel_values = pixel_values.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, pixel_tensors=pixel_tensors)
+            res = model(X, pixel_values=pixel_values)
             loss = loss_fct(
                 res.logits.view(-1, res.logits.size(-1)),
                 Y.view(-1)
@@ -86,7 +84,7 @@ def train_epoch(epoch, wandb):
         if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
             model.eval()
             moe_path = '_moe' if model_config.use_moe else ''
-            ckp = f'{args.save_dir}/sft_vlm_{model_config.dim}{moe_path}.pth'
+            ckp = f'{args.save_dir}/sft_vlm_{model_config.hidden_size}{moe_path}.pth'
             if isinstance(model, torch.nn.parallel.DistributedDataParallel):
                 state_dict = model.module.state_dict()
             else:
@@ -94,23 +92,29 @@ def train_epoch(epoch, wandb):
             clean_state_dict = {
                 key: value for key, value in state_dict.items() if not key.startswith('vision_encoder.')
             }
+            clean_state_dict = {k: v.half() for k, v in clean_state_dict.items()}  # 半精度保存
             torch.save(clean_state_dict, ckp)
             model.train()
 
 
 def init_model(model_config: VLMConfig):
-    tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
+    tokenizer = AutoTokenizer.from_pretrained('../model')
     moe_path = '_moe' if model_config.use_moe else ''
-    ckp = f'./out/pretrain_vlm_{model_config.dim}{moe_path}.pth'
+    ckp = f'{args.save_dir}/pretrain_vlm_{model_config.hidden_size}{moe_path}.pth'
 
-    model = MiniMindVLM(model_config)
+    model = MiniMindVLM(model_config, vision_model_path="../model/vision_model/clip-vit-base-patch16")
     state_dict = torch.load(ckp, map_location=args.device)
     model.load_state_dict(state_dict, strict=False)
-    model = model.to(args.device)
+
+    # 默认全参训练
+    # # 只解冻注意力机制中的投影层参数
+    # for name, param in model.model.named_parameters():
+    #     if any(proj in name for proj in ['q_proj', 'k_proj', 'v_proj', 'o_proj']):
+    #         param.requires_grad = True
 
     Logger(f'VLM可训练参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
 
-    _, preprocess = MiniMindVLM.get_vision_model()
+    _, preprocess = model.vision_encoder, model.processor
     return model.to(args.device), tokenizer, preprocess
 
 
@@ -128,17 +132,17 @@ def init_distributed_mode():
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="MiniMind-V Pretrain")
-    parser.add_argument("--out_dir", type=str, default="out")
+    parser.add_argument("--out_dir", type=str, default="../out")
-    parser.add_argument("--epochs", type=int, default=6)
+    parser.add_argument("--epochs", type=int, default=2)
-    parser.add_argument("--batch_size", type=int, default=8)
+    parser.add_argument("--batch_size", type=int, default=4)
     parser.add_argument("--learning_rate", type=float, default=1e-6)
     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", default=False, action="store_true")
     parser.add_argument("--wandb_project", type=str, default="MiniMind-V")
     parser.add_argument("--num_workers", type=int, default=8)
-    parser.add_argument("--data_path", type=str, default="./dataset/sft_data.jsonl")
+    parser.add_argument("--data_path", type=str, default="../dataset/sft_data.jsonl")
-    parser.add_argument("--images_path", type=str, default="./dataset/sft_images")
+    parser.add_argument("--images_path", type=str, default="../dataset/sft_images")
     parser.add_argument("--ddp", action="store_true")
     parser.add_argument("--accumulation_steps", type=int, default=1)
     parser.add_argument("--grad_clip", type=float, default=1.0)
@@ -146,13 +150,14 @@ if __name__ == "__main__":
     parser.add_argument("--log_interval", type=int, default=10)
     parser.add_argument("--save_interval", type=int, default=10)
     parser.add_argument('--local_rank', type=int, default=-1)
-    parser.add_argument('--dim', default=512, type=int)
+    parser.add_argument('--hidden_size', default=512, type=int)
-    parser.add_argument('--n_layers', default=8, type=int)
+    parser.add_argument('--num_hidden_layers', default=8, type=int)
     parser.add_argument('--max_seq_len', default=1536, type=int)
     parser.add_argument('--use_moe', default=False, type=bool)
     args = parser.parse_args()
 
-    model_config = VLMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len)
+    model_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
+                             max_seq_len=args.max_seq_len)
     max_seq_len = model_config.max_seq_len
     args.save_dir = os.path.join(args.out_dir)
     os.makedirs(args.save_dir, exist_ok=True)