[feat] support parquet data

dataset/lm_dataset.py

@@ -1,21 +1,24 @@
+import sys
+import os
+__package__ = "dataset"
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 import json
+import torch
+import io
 from PIL import Image
 from torch.utils.data import Dataset, DataLoader
-import torch
 from model.model_vlm import MiniMindVLM
-import os
+import pyarrow.parquet as pq
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 
 class VLMDataset(Dataset):
-    def __init__(self, jsonl_path, images_path, tokenizer, preprocess=None, max_length=512,
+    def __init__(self, parquet_path, tokenizer, preprocess=None, max_length=512,
                  image_special_token='@' * 196):
 
         super().__init__()
-        self.samples = self.load_data(jsonl_path)
+        self.table = pq.read_table(parquet_path)
-        self.images_path = images_path
-
         self.tokenizer = tokenizer
         self.max_length = max_length
         self.preprocess = preprocess
@@ -24,15 +27,7 @@ class VLMDataset(Dataset):
         self.eos_id = tokenizer('<|im_end|>', add_special_tokens=False).input_ids
 
     def __len__(self):
-        return len(self.samples)
+        return len(self.table)
-
-    def load_data(self, path):
-        samples = []
-        with open(path, 'r', encoding='utf-8') as f:
-            for line_num, line in enumerate(f, 1):
-                data = json.loads(line.strip())
-                samples.append(data)
-        return samples
 
     def _create_chat_prompt(self, conversations):
         messages = []
@@ -64,9 +59,10 @@ class VLMDataset(Dataset):
         return loss_mask
 
     def __getitem__(self, index: int):
-        sample = self.samples[index]
+        conversations = json.loads(self.table['conversations'][index].as_py())
-        image_paths = sample['image']
+        image_bytes = self.table['image_bytes'][index].as_py()
-        prompt = self._create_chat_prompt(sample['conversations'])
+        
+        prompt = self._create_chat_prompt(conversations)
         input_ids = self.tokenizer(prompt).input_ids[:self.max_length]
         input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids))
         loss_mask = self._generate_loss_mask(input_ids)
@@ -75,12 +71,16 @@ class VLMDataset(Dataset):
         Y = torch.tensor(input_ids[1:], dtype=torch.long)
         loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long)
 
-        image_tensors = []
+        image = Image.open(io.BytesIO(image_bytes))
-        for image_name in image_paths.split(','):
+        image_tensor = MiniMindVLM.image2tensor(image, self.preprocess).unsqueeze(0)
-            image_name = image_name.strip()
-            image = Image.open(f'{self.images_path}/{image_name}')
-            image_tensor = MiniMindVLM.image2tensor(image, self.preprocess)
-            image_tensors.append(image_tensor)
-        image_tensors = torch.stack(image_tensors, dim=0)
 
-        return X, Y, loss_mask, image_tensors
+        return X, Y, loss_mask, image_tensor
+
+if __name__ == '__main__':
+    import matplotlib.pyplot as plt; plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'SimHei']
+    for path in ['pretrain_data.parquet', 'sft_data.parquet']:
+        t = pq.read_table(path); fig, ax = plt.subplots(1, 5, figsize=(20, 4))
+        for i in range(5):
+            ax[i].imshow(Image.open(io.BytesIO(t['image_bytes'][i].as_py()))); ax[i].axis('off')
+            ax[i].set_title(json.loads(t['conversations'][i].as_py())[1]['content'][:30], fontsize=8)
+        out = path.replace('.parquet', '_preview.png'); plt.savefig(out); print(f'已保存{out}, 共{len(t)}条')

trainer/train_pretrain_vlm.py

@@ -84,7 +84,6 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
             del state_dict, clean_state_dict
 
         del X, Y, loss_mask, pixel_values, res, loss
-        torch.cuda.empty_cache()
 
 
 if __name__ == "__main__":
@@ -100,13 +99,12 @@ if __name__ == "__main__":
     parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
     parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
     parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
-    parser.add_argument("--save_interval", type=int, default=100, help="模型保存间隔")
+    parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
     parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--max_seq_len', default=640, type=int, help="训练的最大截断长度")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
-    parser.add_argument("--data_path", type=str, default="../dataset/pretrain_data.jsonl", help="训练数据路径")
+    parser.add_argument("--data_path", type=str, default="../dataset/pretrain_data.parquet", help="训练数据路径")
-    parser.add_argument("--images_path", type=str, default="../dataset/pretrain_images", help="训练图像路径")
     parser.add_argument('--from_weight', default='llm', type=str, help="基于哪个权重训练，为none则不基于任何权重训练")
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
     parser.add_argument('--freeze_llm', default=1, type=int, choices=[0, 1], help="是否冻结LLM参数（0=否，1=是，仅训练vision_proj）")
@@ -142,7 +140,7 @@ if __name__ == "__main__":
     # ========== 5. 定义模型、数据、优化器 ==========
     model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight, 
                                                    device=args.device, freeze_llm=bool(args.freeze_llm))
-    train_ds = VLMDataset(args.data_path, args.images_path, tokenizer, preprocess=preprocess,
+    train_ds = VLMDataset(args.data_path, tokenizer, preprocess=preprocess,
                           image_special_token=vlm_config.image_special_token,
                           max_length=vlm_config.max_seq_len)
     train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None

trainer/train_sft_vlm.py

@@ -84,7 +84,6 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
             del state_dict, clean_state_dict
 
         del X, Y, loss_mask, pixel_values, res, loss
-        torch.cuda.empty_cache()
 
 
 if __name__ == "__main__":
@@ -100,13 +99,12 @@ if __name__ == "__main__":
     parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
     parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
     parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
-    parser.add_argument("--save_interval", type=int, default=100, help="模型保存间隔")
+    parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
     parser.add_argument('--hidden_size', default=512, type=int, help="隐藏层维度")
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--max_seq_len', default=1536, type=int, help="训练的最大截断长度")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
-    parser.add_argument("--data_path", type=str, default="../dataset/sft_data.jsonl", help="训练数据路径")
+    parser.add_argument("--data_path", type=str, default="../dataset/sft_data.parquet", help="训练数据路径")
-    parser.add_argument("--images_path", type=str, default="../dataset/sft_images", help="训练图像路径")
     parser.add_argument('--from_weight', default='pretrain_vlm', type=str, help="基于哪个权重训练，为none则不基于任何权重训练")
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
     parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
@@ -141,7 +139,7 @@ if __name__ == "__main__":
     # ========== 5. 定义模型、数据、优化器 ==========
     model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight, 
                                                    device=args.device)
-    train_ds = VLMDataset(args.data_path, args.images_path, tokenizer, preprocess=preprocess,
+    train_ds = VLMDataset(args.data_path, tokenizer, preprocess=preprocess,
                           image_special_token=vlm_config.image_special_token,
                           max_length=vlm_config.max_seq_len)
     train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None

trainer/trainer_utils.py

@@ -65,11 +65,11 @@ def init_vlm_model(vlm_config, from_weight='pretrain_vlm', tokenizer_path='../mo
             if 'vision_proj' not in name:
                 param.requires_grad = False
     
-    # 默认全参训练时的可选配置（已注释）
+    # pretrain阶段：解冻最后1层
-    # # 只解冻注意力机制中的投影层参数
+    last_layer_idx = vlm_config.num_hidden_layers - 1
-    # for name, param in model.model.named_parameters():
+    for name, param in model.model.named_parameters():
-    #     if any(proj in name for proj in ['q_proj', 'k_proj', 'v_proj', 'o_proj']):
+        if f'layers.{last_layer_idx}.' in name:
-    #         param.requires_grad = True
+            param.requires_grad = True
     
     Logger(f'所加载VLM Model可训练参数：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
     preprocess = model.processor
@@ -121,7 +121,6 @@ def vlm_checkpoint(vlm_config, weight='pretrain_vlm', model=None, optimizer=None
         torch.save(resume_data, resume_tmp)
         os.replace(resume_tmp, resume_path)
         del state_dict, clean_state_dict, resume_data
-        gc.collect()
         torch.cuda.empty_cache()
     else:  # 加载模式
         if os.path.exists(resume_path):