diff --git a/dataset/lm_dataset.py b/dataset/lm_dataset.py index 3deeae0..dc81aec 100644 --- a/dataset/lm_dataset.py +++ b/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.images_path = images_path - + self.table = pq.read_table(parquet_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) - - 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 + return len(self.table) 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] - image_paths = sample['image'] - prompt = self._create_chat_prompt(sample['conversations']) + conversations = json.loads(self.table['conversations'][index].as_py()) + image_bytes = self.table['image_bytes'][index].as_py() + + 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 = [] - for image_name in image_paths.split(','): - 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) + image = Image.open(io.BytesIO(image_bytes)) + image_tensor = MiniMindVLM.image2tensor(image, self.preprocess).unsqueeze(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)}条') diff --git a/trainer/train_pretrain_vlm.py b/trainer/train_pretrain_vlm.py index 8a10f34..bb2743f 100644 --- a/trainer/train_pretrain_vlm.py +++ b/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("--images_path", type=str, default="../dataset/pretrain_images", help="训练图像路径") + parser.add_argument("--data_path", type=str, default="../dataset/pretrain_data.parquet", 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 diff --git a/trainer/train_sft_vlm.py b/trainer/train_sft_vlm.py index 1cec71e..0df6f1f 100644 --- a/trainer/train_sft_vlm.py +++ b/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("--images_path", type=str, default="../dataset/sft_images", help="训练图像路径") + parser.add_argument("--data_path", type=str, default="../dataset/sft_data.parquet", 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 diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index 4fc0ceb..83bdf7e 100644 --- a/trainer/trainer_utils.py +++ b/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 - # 默认全参训练时的可选配置（已注释） - # # 只解冻注意力机制中的投影层参数 - # 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 + # pretrain阶段：解冻最后1层 + last_layer_idx = vlm_config.num_hidden_layers - 1 + for name, param in model.model.named_parameters(): + if f'layers.{last_layer_idx}.' in name: + 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):