【代码解读】开源模型 minimind之pretrain

minimind原模型地址: https://github.com/jingyaogong/minimind

本文解读下开源模型minimind的预训练代码 train_pretrain.py，解释以代码注释的形式添加

1. 参数配置代码

parser = argparse.ArgumentParser(description="MiniMind Pretraining")
parser.add_argument("--out_dir", type=str, default="out") # 模型输出目录
# 若要以最快速度实现zero则epochs设置为1轮；否则应当利用有限的数据训练2~6个epochs。
parser.add_argument("--epochs", type=int, default=1)  # 遍历训练集
parser.add_argument("--batch_size", type=int, default=32) # 每次输入模型样本数量
parser.add_argument("--learning_rate", type=float, default=5e-4) # 学习率
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu") 
# CUDA是由NVIDIA开发的用于并行计算的平台（通用计算构建的运算平台，是建立GPU显卡驱动层之上，必须与之相匹配)
parser.add_argument("--dtype", type=str, default="bfloat16") # 数据类型
parser.add_argument("--use_wandb", action="store_true") # wandb日志可视化工具
parser.add_argument("--wandb_project", type=str, default="MiniMind-Pretrain")
parser.add_argument("--num_workers", type=int, default=1) # 控制数据加载和预处理的并行工作进程数量
parser.add_argument("--ddp", action="store_true") # 分布式的数据并行
parser.add_argument("--accumulation_steps", type=int, default=8) # 显存原因模型非常大时batch_size只能很小，更新效果差。梯度累计是指多算几个batch后，使用这些batch平均梯度更新模型
parser.add_argument("--grad_clip", type=float, default=1.0) # 解决梯度消失/爆炸的问题，当预更新的梯度小于阈值时，将预更新的梯度设为阈值
parser.add_argument("--warmup_iters", type=int, default=0) # 学习率预热阶段是为了在训练初期逐渐增加学习率，以帮助模型稳定地更新参数，避免因初始学习率过大而导致的振荡和不稳定。
parser.add_argument("--log_interval", type=int, default=100) #所以每隔10个batch会输出日志
parser.add_argument("--save_interval", type=int, default=100) #模型检查点（checkpoint）的保存间隔
parser.add_argument('--local_rank', type=int, default=-1)#当前进程的本地编号，帮助在分布式环境中区分不同的进程 https://blog.csdn.net/a8039974/article/details/145790920
parser.add_argument('--dim', default=512, type=int) #维度
parser.add_argument('--n_layers', default=8, type=int) #层数
parser.add_argument('--max_seq_len', default=512, type=int) #最大长度
parser.add_argument('--use_moe', default=False, type=bool) #moe
parser.add_argument("--data_path", type=str, default="./dataset/pretrain_hq.jsonl")
args = parser.parse_args()

2.数据加载代码

train_ds = PretrainDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len)train_sampler = DistributedSampler(train_ds) if ddp else Nonetrain_loader = DataLoader(train_ds,batch_size=args.batch_size,pin_memory=True,drop_last=False,shuffle=False,num_workers=args.num_workers,sampler=train_sampler
)

2.1. 分词器

tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')

3.模型代码

3.1. 模型初始化

# 模型配置，继承transformers.PretrainedConfig类，初始化模型参数配置
lm_config = LMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, use_moe=args.use_moe)# 初始化模型和分词器实例
model, tokenizer = init_model(lm_config)## init_model代码
def init_model(lm_config):tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')model = MiniMindLM(lm_config).to(args.device)Logger(f'LLM总参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')return model, tokenizer

3.2. 模型参数配置

dim: int = 512,       # token向量维度
n_layers: int = 8,    # 解码器transformer 层数（即Encoder和Decoder中的block个数），就多少层有多少个张量对
n_heads: int = 8,     # 多头注意力的头数
n_kv_heads: int = 2,   # ？用于键值对（Key-Value）注意力机制的头数
vocab_size: int = 6400, # 词汇大小
hidden_dim: int = None, # 隐藏层维度
multiple_of: int = 64,  # 隐藏层维度的倍数,隐藏层如果未指定会根据 dim 和 multiple_of 自动计算
norm_eps: float = 1e-5, # ？ 控制归一化数据稳定性
max_seq_len: int = 8192, # 最长序列长度，超出截断
rope_theta: int = 1e6,  # ？ 旋转位置编码
dropout: float = 0.0, # 在训练过程中随机丢弃（设置为0）神经网络中的一些单元及其连接，从而减少模型对特定特征的依赖，防止过拟合
flash_attn: bool = True, #是否加速注意力计算 FlashAttention的核心原理是将输入QKV分块，并保证每个块能够在一级缓存上完成注意力操作并将结果更新回HBM，从而降低对高带宽内存（HBM）的读写操作
####################################################
# Here are the specific configurations of MOE
# When use_moe is false, the following is invalid
####################################################
use_moe: bool = False,
####################################################
num_experts_per_tok: int = 2,   # 每个 token 的专家数量
n_routed_experts: int = 4,      # 总的专家数量
n_shared_experts: bool = True,  # 是否使用共享专家
scoring_func: str = 'softmax',  # 专家的评分函数
aux_loss_alpha: float = 0.1,    # 辅助损失的权重,用于平衡主损失和辅助损失
seq_aux: bool = True,           # 是否在序列级别上计算辅助损失
norm_topk_prob: bool = True,    # 是否对 top-k 概率进行标准化

3.3. 模型结构

3.3.1 主结构代码

# 继承transformer.PreTrainedModel类
model = MiniMindLM(lm_config).to(args.device)#模型初始化
class MiniMindLM(PreTrainedModel):config_class = LMConfigdef __init__(self, params: LMConfig = None):self.params = params or LMConfig()super().__init__(self.params)self.vocab_size, self.n_layers = params.vocab_size, params.n_layersself.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)self.dropout = nn.Dropout(params.dropout)self.layers = nn.ModuleList([MiniMindBlock(l, params) for l in range(self.n_layers)])self.norm = RMSNorm(params.dim, eps=params.norm_eps)self.output = nn.Linear(params.dim, params.vocab_size, bias=False)self.tok_embeddings.weight = self.output.weightself.register_buffer("pos_cis",precompute_pos_cis(dim=params.dim // params.n_heads, theta=params.rope_theta), # 旋转位置编码persistent=False)self.OUT = CausalLMOutputWithPast() #transformer库一个类，用于处理自回归语言模型的输出。它是一个有序的字典def forward(self,input_ids: Optional[torch.Tensor] = None,past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,use_cache: bool = False,**args):past_key_values = past_key_values or [None] * len(self.layers)start_pos = args.get('start_pos', 0)h = self.dropout(self.tok_embeddings(input_ids)) # 对应图1里 Input Embedding层， input_ids是tokenizer分词处理后的输出pos_cis = self.pos_cis[start_pos:start_pos + input_ids.size(1)]past_kvs = []for l, layer in enumerate(self.layers):h, past_kv = layer( # 调用了MiniMindBlock类，是Trnsformer Layer层的主要内容h, pos_cis, # 除第一层输入为Input Embedding层输出，其余层为上层MiniMindBlock的输出past_key_value=past_key_values[l],use_cache=use_cache)past_kvs.append(past_kv)logits = self.output(self.norm(h)) # 对应图1里的Trnsformer Layer结束后先RMSNorm 然后Linear得到回归结果aux_loss = sum(l.feed_forward.aux_loss for l in self.layers if isinstance(l.feed_forward, MOEFeedForward))self.OUT.__setitem__('logits', logits)self.OUT.__setitem__('aux_loss', aux_loss) # MOE的辅助损失self.OUT.__setitem__('past_key_values', past_kvs)return self.OUT

3.3.2. Transformer Layer代码

class MiniMindBlock(nn.Module):'''Trnsformer Layer层，每层包含一个attention(包含attention_norm）：对应图1中的GQA,调用了Attention类feed_forward（包含ffn_norm）：对应图1中的FFN，调用了FeedForward类'''def __init__(self, layer_id: int, config: LMConfig):super().__init__()self.n_heads = config.n_headsself.dim = config.dimself.head_dim = config.dim // config.n_headsself.attention = Attention(config)self.layer_id = layer_idself.attention_norm = RMSNorm(config.dim, eps=config.norm_eps) # self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)self.feed_forward = FeedForward(config) if not config.use_moe else MOEFeedForward(config)def forward(self, x, pos_cis, past_key_value=None, use_cache=False):h_attn, past_kv = self.attention(self.attention_norm(x),pos_cis,past_key_value=past_key_value,use_cache=use_cache)h = x + h_attnout = h + self.feed_forward(self.ffn_norm(h))return out, past_kv

3.3.3. GQA代码（Attenton类）

lass Attention(nn.Module):def __init__(self, args: LMConfig):super().__init__()self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_headsassert args.n_heads % self.n_kv_heads == 0self.n_local_heads = args.n_headsself.n_local_kv_heads = self.n_kv_headsself.n_rep = self.n_local_heads // self.n_local_kv_headsself.head_dim = args.dim // args.n_headsself.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)self.attn_dropout = nn.Dropout(args.dropout)self.resid_dropout = nn.Dropout(args.dropout)self.dropout = args.dropoutself.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn# print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))mask = torch.triu(mask, diagonal=1)self.register_buffer("mask", mask, persistent=False)def forward(self,x: torch.Tensor,pos_cis: torch.Tensor,past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,use_cache=False):bsz, seq_len, _ = x.shapexq, xk, xv = self.wq(x), self.wk(x), self.wv(x)xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)xk = xk.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)xq, xk = apply_rotary_emb(xq, xk, pos_cis) #  对应图1中RoPE,旋转位置编码计算,用于将旋转位置嵌入应用到输入张量x上# kv_cache实现if past_key_value is not None:xk = torch.cat([past_key_value[0], xk], dim=1)xv = torch.cat([past_key_value[1], xv], dim=1)past_kv = (xk, xv) if use_cache else Nonexq, xk, xv = (xq.transpose(1, 2),repeat_kv(xk, self.n_rep).transpose(1, 2),repeat_kv(xv, self.n_rep).transpose(1, 2))if self.flash and seq_len != 1:dropout_p = self.dropout if self.training else 0.0output = F.scaled_dot_product_attention(xq, xk, xv,attn_mask=None,dropout_p=dropout_p,is_causal=True)else:scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim) scores += self.mask[:, :, :seq_len, :seq_len] # 加mask掩码scores = F.softmax(scores.float(), dim=-1).type_as(xq)scores = self.attn_dropout(scores)output = scores @ xvoutput = output.transpose(1, 2).reshape(bsz, seq_len, -1)output = self.resid_dropout(self.wo(output))return output, past_kv# 函数 apply_rotary_emb 旋转位置编码
def apply_rotary_emb(xq, xk, pos_cis):def unite_shape(pos_cis, x):ndim = x.ndimassert 0 <= 1 < ndimassert pos_cis.shape == (x.shape[1], x.shape[-1])shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]return pos_cis.view(*shape)xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))pos_cis = unite_shape(pos_cis, xq_)xq_out = torch.view_as_real(xq_ * pos_cis).flatten(3)xk_out = torch.view_as_real(xk_ * pos_cis).flatten(3)return xq_out.type_as(xq), xk_out.type_as(xk)

3.3.4. FFN代码（FeedForward类）

class FeedForward(nn.Module):def __init__(self, config: LMConfig):super().__init__()if config.hidden_dim is None:hidden_dim = 4 * config.dimhidden_dim = int(2 * hidden_dim / 3)config.hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)self.w1 = nn.Linear(config.dim, config.hidden_dim, bias=False)self.w2 = nn.Linear(config.hidden_dim, config.dim, bias=False)self.w3 = nn.Linear(config.dim, config.hidden_dim, bias=False)self.dropout = nn.Dropout(config.dropout)def forward(self, x):return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))

3.3.5. RMSNorm类

class RMSNorm(torch.nn.Module):'''RMSNorm:替代传统的LayerNorm，通过均方根进行归一化，计算公式为：weight * (x / sqrt(mean(x²) + eps))。减少计算量，提升训练稳定性。'''def __init__(self, dim: int, eps: float):super().__init__()self.eps = epsself.weight = nn.Parameter(torch.ones(dim))def forward(self, x):return self.weight * (x.float() * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x)

4.模型训练

     最后输出模型存储目录： out/pretrain_${dim}.pth

# 迭代优化器
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)iter_per_epoch = len(train_loader)
for epoch in range(args.epochs):train_epoch(epoch, wandb)