github
/
minimind-v
mirror of https://github.com/jingyaogong/minimind-v.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[update] sync minimind

pull/107/head
jingyaogong 2026-01-15 21:36:28 +08:00
parent d2bdff3269
commit 326b24727e
7 changed files with 126 additions and 105 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
dataset/lm_dataset.py
View File

@ -23,13 +23,13 @@ class VLMDataset(Dataset):
self.max_length = max_length self.max_length = max_length
self.preprocess = preprocess self.preprocess = preprocess
self.image_token = image_special_token self.image_token = image_special_token
self.bos_id = tokenizer('<|im_start|>assistant', 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('<|im_end|>', add_special_tokens=False).input_ids self.eos_id = tokenizer(f'{tokenizer.eos_token}\n', add_special_tokens=False).input_ids
def __len__(self): def __len__(self):
return len(self.table) return len(self.table)
def _create_chat_prompt(self, conversations): def create_chat_prompt(self, conversations):
messages = [] messages = []
for i, turn in enumerate(conversations): for i, turn in enumerate(conversations):
role = 'user' if i % 2 == 0 else 'assistant' role = 'user' if i % 2 == 0 else 'assistant'
@ -40,8 +40,8 @@ class VLMDataset(Dataset):
add_generation_prompt=False add_generation_prompt=False
) )
def _generate_loss_mask(self, input_ids): def generate_labels(self, input_ids):
loss_mask = [0] * len(input_ids) labels = [-100] * len(input_ids)
i = 0 i = 0
while i < len(input_ids): while i < len(input_ids):
if input_ids[i:i + len(self.bos_id)] == self.bos_id: 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: if input_ids[end:end + len(self.eos_id)] == self.eos_id:
break break
end += 1 end += 1
for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)): for j in range(start, min(end + len(self.eos_id), self.max_length)):
loss_mask[j] = 1 labels[j] = input_ids[j]
i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids) i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
else: else:
i += 1 i += 1
return loss_mask return labels
def __getitem__(self, index: int): def __getitem__(self, index: int):
conversations = json.loads(self.table['conversations'][index].as_py()) conversations = json.loads(self.table['conversations'][index].as_py())
image_bytes = self.table['image_bytes'][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(prompt).input_ids[:self.max_length]
input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids)) input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids))
loss_mask = self._generate_loss_mask(input_ids) labels = self.generate_labels(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)
image = Image.open(io.BytesIO(image_bytes)) image = Image.open(io.BytesIO(image_bytes))
image_tensor = MiniMindVLM.image2tensor(image, self.preprocess).unsqueeze(0) 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数据读取和可视化 # 测试parquet数据读取和可视化
if __name__ == '__main__': if __name__ == '__main__':

17
eval_vlm.py
View File

@ -1,3 +1,4 @@
import time
import argparse import argparse
import os import os
import warnings import warnings
@ -5,7 +6,7 @@ import torch
from PIL import Image from PIL import Image
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
from model.model_vlm import MiniMindVLM, VLMConfig 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') warnings.filterwarnings('ignore')
def init_model(args): def init_model(args):
@ -22,8 +23,7 @@ def init_model(args):
else: else:
model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True) 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") model.vision_encoder, model.processor = MiniMindVLM.get_vision_model("./model/vision_model/clip-vit-base-patch16")
get_model_params(model, model.params)
print(f'VLM模型参数: {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.2f} M(illion)')
preprocess = model.processor preprocess = model.processor
return model.eval().to(args.device), tokenizer, preprocess 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('--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('--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('--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="运行设备") parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备")
args = parser.parse_args() args = parser.parse_args()
@ -59,15 +60,17 @@ def main():
inputs = tokenizer(inputs_text, return_tensors="pt", truncation=True).to(args.device) inputs = tokenizer(inputs_text, return_tensors="pt", truncation=True).to(args.device)
print(f'[图像]: {image_file}') print(f'[图像]: {image_file}')
print(f'👶: {prompt.replace('\n', '\\n')}') print(f'💬: {prompt.replace('\n', '\\n')}')
print('🤖️: ', end='') print('🤖: ', end='')
model.generate( st = time.time()
generated_ids = model.generate(
inputs=inputs["input_ids"], attention_mask=inputs["attention_mask"], inputs=inputs["input_ids"], attention_mask=inputs["attention_mask"],
max_new_tokens=args.max_new_tokens, do_sample=True, streamer=streamer, 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, 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 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__": if __name__ == "__main__":
main() main()

34
model/model_minimind.py
View File

@ -179,7 +179,7 @@ class Attention(nn.Module):
xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim) xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
cos, sin = position_embeddings 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实现 # kv_cache实现
if past_key_value is not None: if past_key_value is not None:
@ -193,14 +193,11 @@ class Attention(nn.Module):
repeat_kv(xv, self.n_rep).transpose(1, 2) 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) output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=True)
else: else:
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim) scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
scores = scores + torch.triu( scores[:, :, :, -seq_len:] += torch.triu(torch.full((seq_len, seq_len), float("-inf"), device=scores.device), diagonal=1)
torch.full((seq_len, seq_len), float("-inf"), device=scores.device),
diagonal=1
).unsqueeze(0).unsqueeze(0) # scores+mask
if attention_mask is not None: if attention_mask is not None:
extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
@ -284,7 +281,7 @@ class MoEGate(nn.Module):
fi = ce * self.n_routed_experts fi = ce * self.n_routed_experts
aux_loss = (Pi * fi).sum() * self.alpha aux_loss = (Pi * fi).sum() * self.alpha
else: else:
aux_loss = 0 aux_loss = scores.new_zeros(1).squeeze()
return topk_idx, topk_weight, aux_loss return topk_idx, topk_weight, aux_loss
@ -313,9 +310,11 @@ class MOEFeedForward(nn.Module):
flat_topk_idx = topk_idx.view(-1) flat_topk_idx = topk_idx.view(-1)
if self.training: if self.training:
x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0) 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): 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(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
y = y.view(*orig_shape) y = y.view(*orig_shape)
else: else:
@ -421,12 +420,7 @@ class MiniMindModel(nn.Module):
hidden_states = self.norm(hidden_states) hidden_states = self.norm(hidden_states)
aux_loss = sum( aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze())
layer.mlp.aux_loss
for layer in self.layers
if isinstance(layer.mlp, MOEFeedForward)
)
return hidden_states, presents, aux_loss return hidden_states, presents, aux_loss
@ -443,6 +437,7 @@ class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
def forward(self, def forward(self,
input_ids: Optional[torch.Tensor] = None, input_ids: Optional[torch.Tensor] = None,
attention_mask: 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, past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
use_cache: bool = False, use_cache: bool = False,
logits_to_keep: Union[int, torch.Tensor] = 0, 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 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, :]) 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 output.aux_loss = aux_loss
return output return output

16
model/model_vlm.py
View File

@ -113,6 +113,7 @@ class MiniMindVLM(MiniMindForCausalLM):
def forward(self, def forward(self,
input_ids: Optional[torch.Tensor] = None, input_ids: Optional[torch.Tensor] = None,
attention_mask: 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, past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
use_cache: bool = False, use_cache: bool = False,
logits_to_keep: Union[int, torch.Tensor] = 0, logits_to_keep: Union[int, torch.Tensor] = 0,
@ -155,13 +156,16 @@ class MiniMindVLM(MiniMindForCausalLM):
hidden_states = self.model.norm(hidden_states) hidden_states = self.model.norm(hidden_states)
aux_loss = sum( aux_loss = sum([l.mlp.aux_loss for l in self.model.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze())
layer.mlp.aux_loss
for layer in self.model.layers
if isinstance(layer.mlp, MOEFeedForward)
)
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep 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, :]) 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 output.aux_loss = aux_loss
return output return output

47
trainer/train_pretrain_vlm.py
View File

@ -22,26 +22,18 @@ warnings.filterwarnings('ignore')
def train_epoch(epoch, loader, iters, start_step=0, wandb=None): def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
loss_fct = nn.CrossEntropyLoss(reduction='none')
start_time = time.time() start_time = time.time()
for step, (X, Y, loss_mask, pixel_values) in enumerate(loader, start=start_step + 1): for step, (input_ids, labels, pixel_values) in enumerate(loader, start=start_step + 1):
X = X.to(args.device) input_ids = input_ids.to(args.device)
Y = Y.to(args.device) labels = labels.to(args.device)
loss_mask = loss_mask.to(args.device)
pixel_values = pixel_values.to(args.device) pixel_values = pixel_values.to(args.device)
lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate) lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
for param_group in optimizer.param_groups: for param_group in optimizer.param_groups:
param_group['lr'] = lr param_group['lr'] = lr
with autocast_ctx: with autocast_ctx:
res = model(X, pixel_values=pixel_values) res = model(input_ids, labels=labels, pixel_values=pixel_values)
loss = loss_fct( loss = res.loss + res.aux_loss
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
loss = loss / args.accumulation_steps loss = loss / args.accumulation_steps
scaler.scale(loss).backward() 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: if step % args.log_interval == 0 or step == iters - 1:
spend_time = time.time() - start_time spend_time = time.time() - start_time
current_loss = loss.item() * args.accumulation_steps 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'] current_lr = optimizer.param_groups[-1]['lr']
eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60 eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60
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')
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, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})
if wandb: wandb.log({"loss": current_loss, "lr": current_lr, "epoch_Time": eta_min})
if (step % args.save_interval == 0 or step == iters - 1) and is_main_process(): if (step % args.save_interval == 0 or step == iters - 1) and is_main_process():
model.eval() model.eval()
moe_suffix = '_moe' if vlm_config.use_moe else '' moe_suffix = '_moe' if vlm_config.use_moe else ''
ckp = f'{args.save_dir}/{args.save_weight}_{vlm_config.hidden_size}{moe_suffix}.pth' ckp = f'{args.save_dir}/{args.save_weight}_{vlm_config.hidden_size}{moe_suffix}.pth'
if isinstance(model, torch.nn.parallel.DistributedDataParallel): raw_model = model.module if isinstance(model, DistributedDataParallel) else model
state_dict = model.module.state_dict() raw_model = getattr(raw_model, '_orig_mod', raw_model)
else: state_dict = raw_model.state_dict()
state_dict = model.state_dict()
clean_state_dict = { clean_state_dict = {
key: value for key, value in state_dict.items() if not key.startswith('vision_encoder.') 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() model.train()
del state_dict, clean_state_dict 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__": 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_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('--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('--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("--use_wandb", action="store_true", help="是否使用wandb")
parser.add_argument("--wandb_project", type=str, default="MiniMind-V-Pretrain", help="wandb项目名") parser.add_argument("--wandb_project", type=str, default="MiniMind-V-Pretrain", help="wandb项目名")
args = parser.parse_args() 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) wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
# ========== 5. 定义模型、数据、优化器 ========== # ========== 5. 定义模型、数据、优化器 ==========
model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight, model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight, device=args.device, freeze_llm=bool(args.freeze_llm))
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, train_ds = VLMDataset(args.data_path, tokenizer, preprocess=preprocess,
image_special_token=vlm_config.image_special_token, image_special_token=vlm_config.image_special_token,
max_length=vlm_config.max_seq_len) max_length=vlm_config.max_seq_len)
@ -158,7 +152,7 @@ if __name__ == "__main__":
# ========== 7. DDP包模型 ========== # ========== 7. DDP包模型 ==========
if dist.is_initialized(): 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]) model = DistributedDataParallel(model, device_ids=[local_rank])
# ========== 8. 开始训练 ========== # ========== 8. 开始训练 ==========
@ -172,3 +166,6 @@ if __name__ == "__main__":
else: # 默认从头开始 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) 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) train_epoch(epoch, loader, len(loader), 0, wandb)
# ========== 9. 清理分布进程 ==========
if dist.is_initialized(): dist.destroy_process_group()

47
trainer/train_sft_vlm.py
View File

@ -22,26 +22,18 @@ warnings.filterwarnings('ignore')
def train_epoch(epoch, loader, iters, start_step=0, wandb=None): def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
loss_fct = nn.CrossEntropyLoss(reduction='none')
start_time = time.time() start_time = time.time()
for step, (X, Y, loss_mask, pixel_values) in enumerate(loader, start=start_step + 1): for step, (input_ids, labels, pixel_values) in enumerate(loader, start=start_step + 1):
X = X.to(args.device) input_ids = input_ids.to(args.device)
Y = Y.to(args.device) labels = labels.to(args.device)
loss_mask = loss_mask.to(args.device)
pixel_values = pixel_values.to(args.device) pixel_values = pixel_values.to(args.device)
lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate) lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
for param_group in optimizer.param_groups: for param_group in optimizer.param_groups:
param_group['lr'] = lr param_group['lr'] = lr
with autocast_ctx: with autocast_ctx:
res = model(X, pixel_values=pixel_values) res = model(input_ids, labels=labels, pixel_values=pixel_values)
loss = loss_fct( loss = res.loss + res.aux_loss
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
loss = loss / args.accumulation_steps loss = loss / args.accumulation_steps
scaler.scale(loss).backward() 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: if step % args.log_interval == 0 or step == iters - 1:
spend_time = time.time() - start_time spend_time = time.time() - start_time
current_loss = loss.item() * args.accumulation_steps 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'] current_lr = optimizer.param_groups[-1]['lr']
eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60 eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60
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')
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, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})
if wandb: wandb.log({"loss": current_loss, "lr": current_lr, "epoch_Time": eta_min})
if (step % args.save_interval == 0 or step == iters - 1) and is_main_process(): if (step % args.save_interval == 0 or step == iters - 1) and is_main_process():
model.eval() model.eval()
moe_suffix = '_moe' if vlm_config.use_moe else '' moe_suffix = '_moe' if vlm_config.use_moe else ''
ckp = f'{args.save_dir}/{args.save_weight}_{vlm_config.hidden_size}{moe_suffix}.pth' ckp = f'{args.save_dir}/{args.save_weight}_{vlm_config.hidden_size}{moe_suffix}.pth'
if isinstance(model, torch.nn.parallel.DistributedDataParallel): raw_model = model.module if isinstance(model, DistributedDataParallel) else model
state_dict = model.module.state_dict() raw_model = getattr(raw_model, '_orig_mod', raw_model)
else: state_dict = raw_model.state_dict()
state_dict = model.state_dict()
clean_state_dict = { clean_state_dict = {
key: value for key, value in state_dict.items() if not key.startswith('vision_encoder.') 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() model.train()
del state_dict, clean_state_dict 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__": 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("--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_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('--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("--use_wandb", action="store_true", help="是否使用wandb")
parser.add_argument("--wandb_project", type=str, default="MiniMind-V-SFT", help="wandb项目名") parser.add_argument("--wandb_project", type=str, default="MiniMind-V-SFT", help="wandb项目名")
args = parser.parse_args() 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) wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
# ========== 5. 定义模型、数据、优化器 ========== # ========== 5. 定义模型、数据、优化器 ==========
model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight, model, tokenizer, preprocess = init_vlm_model(vlm_config, from_weight=args.from_weight, device=args.device)
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, train_ds = VLMDataset(args.data_path, tokenizer, preprocess=preprocess,
image_special_token=vlm_config.image_special_token, image_special_token=vlm_config.image_special_token,
max_length=vlm_config.max_seq_len) max_length=vlm_config.max_seq_len)
@ -157,7 +151,7 @@ if __name__ == "__main__":
# ========== 7. DDP包模型 ========== # ========== 7. DDP包模型 ==========
if dist.is_initialized(): 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]) model = DistributedDataParallel(model, device_ids=[local_rank])
# ========== 8. 开始训练 ========== # ========== 8. 开始训练 ==========
@ -171,3 +165,6 @@ if __name__ == "__main__":
else: # 默认从头开始 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) 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) train_epoch(epoch, loader, len(loader), 0, wandb)
# ========== 9. 清理分布进程 ==========
if dist.is_initialized(): dist.destroy_process_group()

39
trainer/trainer_utils.py
View File

@ -2,16 +2,32 @@
训练工具函数集合 训练工具函数集合
""" """
import os import os
import sys
__package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import random import random
import math import math
import gc
import numpy as np import numpy as np
import torch import torch
import torch.distributed as dist import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import Sampler from torch.utils.data import Sampler
from transformers import AutoTokenizer from transformers import AutoTokenizer
from model.model_vlm import MiniMindVLM 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(): def is_main_process():
@ -24,7 +40,7 @@ def Logger(content):
def get_lr(current_step, total_steps, lr): 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(): 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(): for name, param in model.model.named_parameters():
if f'layers.{last_layer_idx}.' in name: 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} 百万') 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 preprocess = model.processor
return model.to(device), tokenizer, preprocess 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' resume_path = f'{save_dir}/{weight}_{vlm_config.hidden_size}{moe_path}_resume.pth'
if model is not None: if model is not None:
from torch.nn.parallel import DistributedDataParallel raw_model = model.module if isinstance(model, DistributedDataParallel) else model
state_dict = model.module.state_dict() if isinstance(model, DistributedDataParallel) else model.state_dict() raw_model = getattr(raw_model, '_orig_mod', raw_model)
state_dict = raw_model.state_dict()
# 移除vision_encoder参数（不需要保存，因为是预训练的） # 移除vision_encoder参数（不需要保存，因为是预训练的）
clean_state_dict = {k: v for k, v in state_dict.items() if not k.startswith('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' 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(): for key, value in kwargs.items():
if value is not None: if value is not None:
if hasattr(value, 'state_dict'): if hasattr(value, 'state_dict'):
if isinstance(value, DistributedDataParallel): raw_value = value.module if isinstance(value, DistributedDataParallel) else value
resume_data[key] = value.module.state_dict() raw_value = getattr(raw_value, '_orig_mod', raw_value)
else: resume_data[key] = raw_value.state_dict()
resume_data[key] = value.state_dict()
else: else:
resume_data[key] = value resume_data[key] = value