From 326b24727e9c1d1374a091e25f11ddb35eb060e1 Mon Sep 17 00:00:00 2001 From: jingyaogong Date: Thu, 15 Jan 2026 21:36:28 +0800 Subject: [PATCH] [update] sync minimind --- dataset/lm_dataset.py | 31 ++++++++++++----------- eval_vlm.py | 17 +++++++------ model/model_minimind.py | 34 +++++++++++++------------ model/model_vlm.py | 16 +++++++----- trainer/train_pretrain_vlm.py | 47 ++++++++++++++++------------------- trainer/train_sft_vlm.py | 47 ++++++++++++++++------------------- trainer/trainer_utils.py | 39 +++++++++++++++++++++-------- 7 files changed, 126 insertions(+), 105 deletions(-) diff --git a/dataset/lm_dataset.py b/dataset/lm_dataset.py index e2def8f..911de85 100644 --- a/dataset/lm_dataset.py +++ b/dataset/lm_dataset.py @@ -23,13 +23,13 @@ class VLMDataset(Dataset): self.max_length = max_length self.preprocess = preprocess self.image_token = image_special_token - self.bos_id = tokenizer('<|im_start|>assistant', add_special_tokens=False).input_ids - self.eos_id = tokenizer('<|im_end|>', add_special_tokens=False).input_ids + self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant\n', add_special_tokens=False).input_ids + self.eos_id = tokenizer(f'{tokenizer.eos_token}\n', add_special_tokens=False).input_ids def __len__(self): return len(self.table) - def _create_chat_prompt(self, conversations): + def create_chat_prompt(self, conversations): messages = [] for i, turn in enumerate(conversations): role = 'user' if i % 2 == 0 else 'assistant' @@ -40,8 +40,8 @@ class VLMDataset(Dataset): add_generation_prompt=False ) - def _generate_loss_mask(self, input_ids): - loss_mask = [0] * len(input_ids) + def generate_labels(self, input_ids): + labels = [-100] * len(input_ids) i = 0 while i < len(input_ids): if input_ids[i:i + len(self.bos_id)] == self.bos_id: @@ -51,30 +51,31 @@ class VLMDataset(Dataset): if input_ids[end:end + len(self.eos_id)] == self.eos_id: break end += 1 - for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)): - loss_mask[j] = 1 + for j in range(start, min(end + len(self.eos_id), self.max_length)): + labels[j] = input_ids[j] i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids) else: i += 1 - return loss_mask + return labels def __getitem__(self, index: int): conversations = json.loads(self.table['conversations'][index].as_py()) image_bytes = self.table['image_bytes'][index].as_py() - prompt = self._create_chat_prompt(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) - - X = torch.tensor(input_ids[:-1], dtype=torch.long) - Y = torch.tensor(input_ids[1:], dtype=torch.long) - loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long) + labels = self.generate_labels(input_ids) image = Image.open(io.BytesIO(image_bytes)) image_tensor = MiniMindVLM.image2tensor(image, self.preprocess).unsqueeze(0) + # # === 调试打印 === + # print(f"\n--- Sample {index} ---") + # for i, (x, y) in enumerate(zip(input_ids[:-1], labels[1:])): + # print(f"{i:3d}: X={self.tokenizer.decode([x])!r:16s} ---> Y={self.tokenizer.decode([input_ids[i+1]])!r:16s} label={y}") + # # ================ - return X, Y, loss_mask, image_tensor + return torch.tensor(input_ids, dtype=torch.long), torch.tensor(labels, dtype=torch.long), image_tensor # 测试parquet数据读取和可视化 if __name__ == '__main__': diff --git a/eval_vlm.py b/eval_vlm.py index 93357ce..5f3963f 100644 --- a/eval_vlm.py +++ b/eval_vlm.py @@ -1,3 +1,4 @@ +import time import argparse import os import warnings @@ -5,7 +6,7 @@ import torch from PIL import Image from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer from model.model_vlm import MiniMindVLM, VLMConfig -from trainer.trainer_utils import setup_seed +from trainer.trainer_utils import setup_seed, get_model_params warnings.filterwarnings('ignore') def init_model(args): @@ -22,8 +23,7 @@ def init_model(args): else: model = AutoModelForCausalLM.from_pretrained(args.load_from, 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:.2f} M(illion)') + get_model_params(model, model.params) preprocess = model.processor return model.eval().to(args.device), tokenizer, preprocess @@ -40,6 +40,7 @@ def main(): parser.add_argument('--temperature', default=0.65, type=float, help="生成温度，控制随机性（0-1，越大越随机）") parser.add_argument('--top_p', default=0.85, type=float, help="nucleus采样阈值（0-1）") parser.add_argument('--image_dir', default='./dataset/eval_images/', type=str, help="测试图像目录") + parser.add_argument('--show_speed', default=1, type=int, help="显示decode速度（tokens/s）") parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备") args = parser.parse_args() @@ -59,15 +60,17 @@ def main(): inputs = tokenizer(inputs_text, return_tensors="pt", truncation=True).to(args.device) print(f'[图像]: {image_file}') - print(f'👶: {prompt.replace('\n', '\\n')}') - print('🤖️: ', end='') - model.generate( + print(f'💬: {prompt.replace('\n', '\\n')}') + print('🤖: ', end='') + st = time.time() + generated_ids = model.generate( inputs=inputs["input_ids"], attention_mask=inputs["attention_mask"], max_new_tokens=args.max_new_tokens, do_sample=True, streamer=streamer, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id, top_p=args.top_p, temperature=args.temperature, pixel_values=pixel_values ) - print('\n\n') + gen_tokens = len(generated_ids[0]) - len(inputs["input_ids"][0]) + print(f'\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s\n\n') if args.show_speed else print('\n\n') if __name__ == "__main__": main() diff --git a/model/model_minimind.py b/model/model_minimind.py index ad62a68..b3910a8 100644 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -179,7 +179,7 @@ class Attention(nn.Module): 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]) + xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin) # kv_cache实现 if past_key_value is not None: @@ -193,14 +193,11 @@ class Attention(nn.Module): repeat_kv(xv, self.n_rep).transpose(1, 2) ) - if self.flash and seq_len > 1 and (attention_mask is None or torch.all(attention_mask == 1)): + if self.flash and (seq_len > 1) and (past_key_value is None) and (attention_mask is None or torch.all(attention_mask == 1)): output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, 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 + scores[:, :, :, -seq_len:] += torch.triu(torch.full((seq_len, seq_len), float("-inf"), device=scores.device), diagonal=1) if attention_mask is not None: extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) @@ -284,7 +281,7 @@ class MoEGate(nn.Module): fi = ce * self.n_routed_experts aux_loss = (Pi * fi).sum() * self.alpha else: - aux_loss = 0 + aux_loss = scores.new_zeros(1).squeeze() return topk_idx, topk_weight, aux_loss @@ -313,9 +310,11 @@ class MOEFeedForward(nn.Module): 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) + y = torch.empty_like(x, dtype=x.dtype) for i, expert in enumerate(self.experts): - y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(y.dtype) # 确保类型一致 + expert_out = expert(x[flat_topk_idx == i]) + if expert_out.shape[0] > 0: y[flat_topk_idx == i] = expert_out.to(y.dtype) + else: y[flat_topk_idx == i] = expert_out.to(y.dtype) + 0 * sum(p.sum() for p in expert.parameters()) y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1) y = y.view(*orig_shape) else: @@ -421,12 +420,7 @@ class MiniMindModel(nn.Module): hidden_states = self.norm(hidden_states) - aux_loss = sum( - layer.mlp.aux_loss - for layer in self.layers - if isinstance(layer.mlp, MOEFeedForward) - ) - + aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze()) return hidden_states, presents, aux_loss @@ -443,6 +437,7 @@ class MiniMindForCausalLM(PreTrainedModel, GenerationMixin): def forward(self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, + labels: 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, @@ -456,6 +451,13 @@ class MiniMindForCausalLM(PreTrainedModel, GenerationMixin): ) 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, :]) - output = CausalLMOutputWithPast(logits=logits, past_key_values=past_key_values, hidden_states=hidden_states) + + loss = None + if labels is not None: + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100) + + output = CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=past_key_values, hidden_states=hidden_states) output.aux_loss = aux_loss return output diff --git a/model/model_vlm.py b/model/model_vlm.py index f17e93e..3babdb8 100644 --- a/model/model_vlm.py +++ b/model/model_vlm.py @@ -113,6 +113,7 @@ class MiniMindVLM(MiniMindForCausalLM): def forward(self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, + labels: 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, @@ -155,13 +156,16 @@ class MiniMindVLM(MiniMindForCausalLM): hidden_states = self.model.norm(hidden_states) - aux_loss = sum( - layer.mlp.aux_loss - for layer in self.model.layers - if isinstance(layer.mlp, MOEFeedForward) - ) + aux_loss = sum([l.mlp.aux_loss for l in self.model.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze()) 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, :]) - output = CausalLMOutputWithPast(logits=logits, past_key_values=presents, hidden_states=hidden_states) + + loss = None + if labels is not None: + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100) + + output = CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=presents, hidden_states=hidden_states) output.aux_loss = aux_loss return output diff --git a/trainer/train_pretrain_vlm.py b/trainer/train_pretrain_vlm.py index bb2743f..8d6256a 100644 --- a/trainer/train_pretrain_vlm.py +++ b/trainer/train_pretrain_vlm.py @@ -22,26 +22,18 @@ warnings.filterwarnings('ignore') def train_epoch(epoch, loader, iters, start_step=0, wandb=None): - loss_fct = nn.CrossEntropyLoss(reduction='none') start_time = time.time() - for step, (X, Y, loss_mask, pixel_values) in enumerate(loader, start=start_step + 1): - X = X.to(args.device) - Y = Y.to(args.device) - loss_mask = loss_mask.to(args.device) + for step, (input_ids, labels, pixel_values) in enumerate(loader, start=start_step + 1): + input_ids = input_ids.to(args.device) + labels = labels.to(args.device) pixel_values = pixel_values.to(args.device) lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate) for param_group in optimizer.param_groups: param_group['lr'] = lr with autocast_ctx: - res = model(X, pixel_values=pixel_values) - loss = loss_fct( - res.logits.view(-1, res.logits.size(-1)), - Y.view(-1) - ).view(Y.size()) - - loss = (loss * loss_mask).sum() / loss_mask.sum() - loss += res.aux_loss + res = model(input_ids, labels=labels, pixel_values=pixel_values) + loss = res.loss + res.aux_loss loss = loss / args.accumulation_steps scaler.scale(loss).backward() @@ -58,21 +50,20 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): if step % args.log_interval == 0 or step == iters - 1: spend_time = time.time() - start_time current_loss = loss.item() * args.accumulation_steps + current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0 + current_logits_loss = current_loss - current_aux_loss current_lr = optimizer.param_groups[-1]['lr'] eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60 - - Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}) loss:{current_loss:.6f} lr:{current_lr:.12f} epoch_Time:{eta_min}min:') - - if wandb: wandb.log({"loss": current_loss, "lr": current_lr, "epoch_Time": eta_min}) + Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min') + if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min}) if (step % args.save_interval == 0 or step == iters - 1) and is_main_process(): model.eval() moe_suffix = '_moe' if vlm_config.use_moe else '' ckp = f'{args.save_dir}/{args.save_weight}_{vlm_config.hidden_size}{moe_suffix}.pth' - if isinstance(model, torch.nn.parallel.DistributedDataParallel): - state_dict = model.module.state_dict() - else: - state_dict = model.state_dict() + raw_model = model.module if isinstance(model, DistributedDataParallel) else model + raw_model = getattr(raw_model, '_orig_mod', raw_model) + state_dict = raw_model.state_dict() clean_state_dict = { key: value for key, value in state_dict.items() if not key.startswith('vision_encoder.') } @@ -83,7 +74,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): model.train() del state_dict, clean_state_dict - del X, Y, loss_mask, pixel_values, res, loss + del input_ids, labels, pixel_values, res, loss if __name__ == "__main__": @@ -108,6 +99,7 @@ if __name__ == "__main__": 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）") + parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）") parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb") parser.add_argument("--wandb_project", type=str, default="MiniMind-V-Pretrain", help="wandb项目名") args = parser.parse_args() @@ -138,8 +130,10 @@ if __name__ == "__main__": wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume) # ========== 5. 定义模型、数据、优化器 ========== - model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight, - device=args.device, freeze_llm=bool(args.freeze_llm)) + model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight, device=args.device, freeze_llm=bool(args.freeze_llm)) + if args.use_compile == 1: + model = torch.compile(model) + Logger('torch.compile enabled') train_ds = VLMDataset(args.data_path, tokenizer, preprocess=preprocess, image_special_token=vlm_config.image_special_token, max_length=vlm_config.max_seq_len) @@ -158,7 +152,7 @@ if __name__ == "__main__": # ========== 7. DDP包模型 ========== if dist.is_initialized(): - model._ddp_params_and_buffers_to_ignore = {"pos_cis"} + model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"} model = DistributedDataParallel(model, device_ids=[local_rank]) # ========== 8. 开始训练 ========== @@ -172,3 +166,6 @@ if __name__ == "__main__": else: # 默认从头开始 loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=(train_sampler is None), sampler=train_sampler, num_workers=args.num_workers, pin_memory=True) train_epoch(epoch, loader, len(loader), 0, wandb) + + # ========== 9. 清理分布进程 ========== + if dist.is_initialized(): dist.destroy_process_group() diff --git a/trainer/train_sft_vlm.py b/trainer/train_sft_vlm.py index 0df6f1f..38285df 100644 --- a/trainer/train_sft_vlm.py +++ b/trainer/train_sft_vlm.py @@ -22,26 +22,18 @@ warnings.filterwarnings('ignore') def train_epoch(epoch, loader, iters, start_step=0, wandb=None): - loss_fct = nn.CrossEntropyLoss(reduction='none') start_time = time.time() - for step, (X, Y, loss_mask, pixel_values) in enumerate(loader, start=start_step + 1): - X = X.to(args.device) - Y = Y.to(args.device) - loss_mask = loss_mask.to(args.device) + for step, (input_ids, labels, pixel_values) in enumerate(loader, start=start_step + 1): + input_ids = input_ids.to(args.device) + labels = labels.to(args.device) pixel_values = pixel_values.to(args.device) lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate) for param_group in optimizer.param_groups: param_group['lr'] = lr with autocast_ctx: - res = model(X, pixel_values=pixel_values) - loss = loss_fct( - res.logits.view(-1, res.logits.size(-1)), - Y.view(-1) - ).view(Y.size()) - - loss = (loss * loss_mask).sum() / loss_mask.sum() - loss += res.aux_loss + res = model(input_ids, labels=labels, pixel_values=pixel_values) + loss = res.loss + res.aux_loss loss = loss / args.accumulation_steps scaler.scale(loss).backward() @@ -58,21 +50,20 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): if step % args.log_interval == 0 or step == iters - 1: spend_time = time.time() - start_time current_loss = loss.item() * args.accumulation_steps + current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0 + current_logits_loss = current_loss - current_aux_loss current_lr = optimizer.param_groups[-1]['lr'] eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60 - - Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}) loss:{current_loss:.6f} lr:{current_lr:.12f} epoch_Time:{eta_min}min:') - - if wandb: wandb.log({"loss": current_loss, "lr": current_lr, "epoch_Time": eta_min}) + Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min') + if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min}) if (step % args.save_interval == 0 or step == iters - 1) and is_main_process(): model.eval() moe_suffix = '_moe' if vlm_config.use_moe else '' ckp = f'{args.save_dir}/{args.save_weight}_{vlm_config.hidden_size}{moe_suffix}.pth' - if isinstance(model, torch.nn.parallel.DistributedDataParallel): - state_dict = model.module.state_dict() - else: - state_dict = model.state_dict() + raw_model = model.module if isinstance(model, DistributedDataParallel) else model + raw_model = getattr(raw_model, '_orig_mod', raw_model) + state_dict = raw_model.state_dict() clean_state_dict = { key: value for key, value in state_dict.items() if not key.startswith('vision_encoder.') } @@ -83,7 +74,7 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None): model.train() del state_dict, clean_state_dict - del X, Y, loss_mask, pixel_values, res, loss + del input_ids, labels, pixel_values, res, loss if __name__ == "__main__": @@ -107,6 +98,7 @@ if __name__ == "__main__": 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_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）") parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb") parser.add_argument("--wandb_project", type=str, default="MiniMind-V-SFT", help="wandb项目名") args = parser.parse_args() @@ -137,8 +129,10 @@ if __name__ == "__main__": wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume) # ========== 5. 定义模型、数据、优化器 ========== - model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight, - device=args.device) + model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight, device=args.device) + if args.use_compile == 1: + model = torch.compile(model) + Logger('torch.compile enabled') train_ds = VLMDataset(args.data_path, tokenizer, preprocess=preprocess, image_special_token=vlm_config.image_special_token, max_length=vlm_config.max_seq_len) @@ -157,7 +151,7 @@ if __name__ == "__main__": # ========== 7. DDP包模型 ========== if dist.is_initialized(): - model._ddp_params_and_buffers_to_ignore = {"pos_cis"} + model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"} model = DistributedDataParallel(model, device_ids=[local_rank]) # ========== 8. 开始训练 ========== @@ -171,3 +165,6 @@ if __name__ == "__main__": else: # 默认从头开始 loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=(train_sampler is None), sampler=train_sampler, num_workers=args.num_workers, pin_memory=True) train_epoch(epoch, loader, len(loader), 0, wandb) + + # ========== 9. 清理分布进程 ========== + if dist.is_initialized(): dist.destroy_process_group() diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index 83bdf7e..10d443e 100644 --- a/trainer/trainer_utils.py +++ b/trainer/trainer_utils.py @@ -2,16 +2,32 @@ 训练工具函数集合 """ import os +import sys +__package__ = "trainer" +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) import random import math -import gc import numpy as np import torch import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel from torch.utils.data import Sampler from transformers import AutoTokenizer from model.model_vlm import MiniMindVLM - + + +def get_model_params(model, config, ignore_patterns=['vision_encoder']): + def should_count(n): return not any(p in n for p in ignore_patterns) + total = sum(p.numel() for n, p in model.named_parameters() if should_count(n)) / 1e6 + n_routed = getattr(config, 'n_routed_experts', getattr(config, 'num_experts', 0)) + n_active = getattr(config, 'num_experts_per_tok', 0) + n_shared = getattr(config, 'n_shared_experts', 0) + expert = sum(p.numel() for n, p in model.named_parameters() if 'mlp.experts.0.' in n and should_count(n)) / 1e6 + shared_expert = sum(p.numel() for n, p in model.named_parameters() if 'mlp.shared_experts.0.' in n and should_count(n)) / 1e6 + base = total - (expert * n_routed) - (shared_expert * n_shared) + active = base + (expert * n_active) + (shared_expert * n_shared) + if active < total: Logger(f'Model Params: {total:.2f}M-A{active:.2f}M') + else: Logger(f'Model Params: {total:.2f}M') def is_main_process(): @@ -24,7 +40,7 @@ def Logger(content): def get_lr(current_step, total_steps, lr): - return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps)) + return lr*(0.1 + 0.45*(1 + math.cos(math.pi * current_step / total_steps))) def init_distributed_mode(): @@ -70,8 +86,9 @@ def init_vlm_model(vlm_config, from_weight='pretrain_vlm', tokenizer_path='../mo 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} 百万') + + get_model_params(model, vlm_config) + Logger(f'Trainable Params: {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f}M') preprocess = model.processor return model.to(device), tokenizer, preprocess @@ -83,8 +100,9 @@ def vlm_checkpoint(vlm_config, weight='pretrain_vlm', model=None, optimizer=None resume_path = f'{save_dir}/{weight}_{vlm_config.hidden_size}{moe_path}_resume.pth' if model is not None: - from torch.nn.parallel import DistributedDataParallel - state_dict = model.module.state_dict() if isinstance(model, DistributedDataParallel) else model.state_dict() + raw_model = model.module if isinstance(model, DistributedDataParallel) else model + raw_model = getattr(raw_model, '_orig_mod', raw_model) + state_dict = raw_model.state_dict() # 移除vision_encoder参数（不需要保存，因为是预训练的） clean_state_dict = {k: v for k, v in state_dict.items() if not k.startswith('vision_encoder.')} ckp_tmp = ckp_path + '.tmp' @@ -110,10 +128,9 @@ def vlm_checkpoint(vlm_config, weight='pretrain_vlm', model=None, optimizer=None for key, value in kwargs.items(): if value is not None: if hasattr(value, 'state_dict'): - if isinstance(value, DistributedDataParallel): - resume_data[key] = value.module.state_dict() - else: - resume_data[key] = value.state_dict() + raw_value = value.module if isinstance(value, DistributedDataParallel) else value + raw_value = getattr(raw_value, '_orig_mod', raw_value) + resume_data[key] = raw_value.state_dict() else: resume_data[key] = value