修改了一些bug和漏洞

README.md

@@ -473,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

@@ -10,7 +10,7 @@ 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) / 1e6:.2f}M(illion)')
+    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
 
 # 设置可复现的随机种子
@@ -24,10 +24,10 @@ def setup_seed(seed):
     torch.backends.cudnn.benchmark = False
 
 parser = argparse.ArgumentParser(description="Chat with MiniMind")
-parser.add_argument('--temperature', default=0.85, type=float)
+parser.add_argument('--temp', default=0.85, type=float)
 parser.add_argument('--top_p', default=0.85, type=float)
 
-parser.add_argument('--max_seq_len', default=8192, type=int)
+parser.add_argument('--max_len', default=8192, type=int)
 args = parser.parse_args()
 model, tokenizer = init_model()
 
@@ -36,6 +36,13 @@ 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
@@ -44,7 +51,6 @@ for prompt in iter(lambda: input('👶: '), ''):
         messages = []
         print("成功清除对话历史")
         continue
-    setup_seed(random.randint(0, 2048))
 
     messages.append({"role": "user", "content": prompt})
 
@@ -63,7 +69,7 @@ for prompt in iter(lambda: input('👶: '), ''):
     print('🤖️: ', end='')
     generated_ids = model.generate(
         inputs["input_ids"],
-        max_new_tokens=args.max_seq_len,
+        max_new_tokens=args.max_len,
         num_return_sequences=1,
         do_sample=True,
         attention_mask=inputs["attention_mask"],
@@ -71,7 +77,7 @@ for prompt in iter(lambda: input('👶: '), ''):
         eos_token_id=tokenizer.eos_token_id,
         streamer=streamer,
         top_p=args.top_p,
-        temperature=args.temperature
+        temperature=args.temp
     )
 
     response = tokenizer.decode(generated_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)

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,7 +15,7 @@ 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
@@ -91,8 +93,14 @@ 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')
+        model.from_pretrained(model_data_path)
+        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 +113,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 +142,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 +215,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):