Merge 3cc6d0d074 into 561979c7e3

.gitignore

@@ -1,3 +1,4 @@
 /model/__pycache__
 /dataset
 /out
+/MiniMind2

README.md

@@ -235,6 +235,18 @@ git clone https://github.com/jingyaogong/minimind.git
 pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
 ```
 
+因为需要的库较多，建议使用虚拟环境来安装：
+```bash
+python -m venv .venv
+```
+
+再下载库：
+
+```bash
+pip install -r requirements.txt
+```
+
+
 ### 2.下载模型
 到项目根目录
 ```bash
@@ -243,11 +255,21 @@ git clone https://huggingface.co/jingyaogong/MiniMind2
 
 ### （可选）命令行问答
 
+如果在第1步中安装了虚拟环境，需要在使用命令行激活虚拟环境：
+```bash
+.venv/Scripts/activate
+```
+
 ```bash
 # load=0: load from pytorch model, load=1: load from transformers-hf model
 python eval_model.py --load 1 --model_mode 2
 ```
 
+或者使用
+```bash
+python model_chat.py
+```
+
 ### （可选）启动WebUI
 
 ```bash
@@ -451,7 +473,7 @@ quality（当然也还不算high，提升数据质量无止尽）。
 导出文件为`sft_512.jsonl`(~7.5GB)。
 
 [Magpie-SFT数据集](https://www.modelscope.cn/organization/Magpie-Align)
-收集了~1M条来自Qwen2/2.5的高质量对话，我将这部分数据进一步清洗，把总长度`<2048`的部分导出为`sft_2048.jsonl`(~9GB)。
+收集了大约1M条来自Qwen2/2.5的高质量对话，我将这部分数据进一步清洗，把总长度`<2048`的部分导出为`sft_2048.jsonl`(~9GB)。
 长度`<1024`的部分导出为`sft_1024.jsonl`(~5.5GB)，用大模型对话数据直接进行sft就属于“黑盒蒸馏”的范畴。
 
 进一步清洗前两步sft的数据（只保留中文字符占比高的内容），筛选长度`<512`的对话，得到`sft_mini_512.jsonl`(~1.2GB)。

model/model_minimind.py

@@ -290,7 +290,6 @@ class MOEFeedForward(nn.Module):
     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])

model_chat.py

@@ -0,0 +1,85 @@
+print("开始加载模型")
+
+import argparse
+import random
+import numpy as np
+from transformers import AutoTokenizer, TextStreamer, LlamaForCausalLM
+import torch
+
+def init_model():
+    transformers_model_path = './MiniMind2'
+    tokenizer = AutoTokenizer.from_pretrained(transformers_model_path)
+    model = LlamaForCausalLM.from_pretrained(transformers_model_path, trust_remote_code=True)
+    print(f'模型参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad)}')
+    return model.eval().to('cuda' if torch.cuda.is_available() else 'cpu'), tokenizer
+
+# 设置可复现的随机种子
+def setup_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+parser = argparse.ArgumentParser(description="Chat with MiniMind")
+parser.add_argument('--temp', default=0.85, type=float)
+parser.add_argument('--top_p', default=0.85, type=float)
+
+parser.add_argument('--max_len', default=8192, type=int)
+args = parser.parse_args()
+model, tokenizer = init_model()
+
+streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+print("模型加载完毕")
+
+messages = []
+
+random_seed = random.randint(0, 2048)
+
+setup_seed(random_seed)
+
+print("使用的随机种子为：", random_seed)
+
+for prompt in iter(lambda: input('👶: '), ''):
+    if prompt == "exit":
+        break
+
+    elif print == "clear history":
+        messages = []
+        print("成功清除对话历史")
+        continue
+
+    messages.append({"role": "user", "content": prompt})
+
+    new_prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    inputs = tokenizer(
+        new_prompt,
+        return_tensors="pt",
+        truncation=True
+    ).to('cuda' if torch.cuda.is_available() else 'cpu')
+
+    print('🤖️: ', end='')
+    generated_ids = model.generate(
+        inputs["input_ids"],
+        max_new_tokens=args.max_len,
+        num_return_sequences=1,
+        do_sample=True,
+        attention_mask=inputs["attention_mask"],
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+        streamer=streamer,
+        top_p=args.top_p,
+        temperature=args.temp
+    )
+
+    response = tokenizer.decode(generated_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+    messages.append({"role": "assistant", "content": response})
+    print('\n')

trainer/train_lora.py

@@ -1,6 +1,8 @@
 import os
 import sys
 
+from sympy import true
+
 __package__ = "trainer"
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
@@ -13,10 +15,11 @@ from torch import optim, nn
 import torch.distributed as dist
 from contextlib import nullcontext
 from torch.utils.data import DataLoader, DistributedSampler
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer
 from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
 from dataset.lm_dataset import SFTDataset
 from model.model_lora import load_lora, save_lora, apply_lora
+import safesentors.torch
 
 warnings.filterwarnings('ignore')
 
@@ -91,8 +94,15 @@ def train_epoch(epoch, wandb):
 
 
 def init_model(lm_config):
-    tokenizer = AutoTokenizer.from_pretrained('../model/')
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    model_path = os.path.join(current_dir, '..', 'model')
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
     model = MiniMindForCausalLM(lm_config)
+    if args.minimind2:
+        model_data_path = os.path.join(current_dir, '..', 'MiniMind2')
+        state_dict = safetensors.torch.load_file("model.safetensors", device=args.device)
+        model.load_state_dict(state_dict, strict= True)
+        return model.to(args.device), tokenizer
     moe_path = '_moe' if lm_config.use_moe else ''
     ckp = f'{args.save_dir}/full_sft_{lm_config.hidden_size}{moe_path}.pth'
     state_dict = torch.load(ckp, map_location=args.device)
@@ -105,9 +115,7 @@ def init_distributed_mode():
     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)
 
@@ -136,8 +144,16 @@ if __name__ == "__main__":
     parser.add_argument('--use_moe', default=False, type=bool)
     parser.add_argument("--data_path", type=str, default="../dataset/lora_medical.jsonl")
     parser.add_argument("--lora_name", type=str, default="lora_medical", help="根据任务保存成lora_(英文/医学/心理...)")
+    parser.add_argument("--minimind2", type=bool, default=true, help="是否使用从huggingface下载下来的MiniMind2模型")
     args = parser.parse_args()
 
+    if args.minimind2 == true:
+        args.hidden_size = 768
+        args.num_hidden_layers=16
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        args.data_path = os.path.join(current_dir, "../dataset/lora_medical.jsonl")
+        
+
     lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
                                use_moe=args.use_moe)
     args.save_dir = os.path.join(args.out_dir)
@@ -201,7 +217,7 @@ if __name__ == "__main__":
         sampler=train_sampler
     )
 
-    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
+    scaler = torch.cuda.amp.GradScaler("cuda", enabled=(args.dtype in ['float16', 'bfloat16']))
     iter_per_epoch = len(train_loader)
 
     for epoch in range(args.epochs):