实习07-混合大模型的学习

1 看架构（Configuration 配置文件）首先，拿到代码，看配置文件里面的架构选型，其中configuration.py文件里面包含了模型每层的选型，以下是layer构建的代码：@propertydeflayers_block_type(self):""" Returns a list of block types for each layer. Block types: - 'emb_adapter' : Embedding Adapter - 'qwen_attention': Qwen3 Attention with RoPE - 'qwen_mlp': Qwen3 MLP with gate/up/down structure - 'nq_adapter' : Nemotron + Qwen3 Adapter - 'mamba': Nemotron Mamba2 layer - 'mlp': Nemotron MLP - 'attention': Nemotron Attention """block_types=[]forcharinself.hybrid_override_pattern:ifchar=='':block_types.append('qwen_attention')elifchar=='^':block_types.append('qwen_mlp')elifchar=='M':block_types.append('mamba')elifchar=='-':block_types.append('mlp')elifchar=='*':block_types.append('attention')elifchar=='!':block_types.append('nq_adapter')elifchar=='#':block_types.append('emb_adapter')else:raiseValueError(f"Unknown pattern character:{char}")returnblock_types输入：模式字符串；输出：返回的是一个layer集合；然后我们再看hybrid_override_pattern的内容：hybrid_override_pattern="# ^ ^ ^ ^ ^ ^ ^ ^ ^ ^ ! M-M-M-M*-M-M-M-M*-M-",# Default 42-layer pattern#= emb_adapter（嵌入适配器）= qwen_attention（Qwen 注意力）^= qwen_mlp（Qwen MLP）!= nq_adapter（中间适配器）M/-/*= Nemotron 相关层（Mamba、MLP、Attention）2 看逻辑（QwenNemotronModel 的 forward）我们只看关键部分：第一部分：默认配置 边界条件检查（不看）#========== 1. 参数默认值配置 ==========# 是否输出注意力权重：使用传入值 or 配置文件默认值output_attentions=output_attentionsifoutput_attentionsisnotNoneelseself.config.output_attentions# 是否输出所有隐藏层状态output_hidden_states=(output_hidden_statesifoutput_hidden_statesisnotNoneelseself.config.output_hidden_states)# 是否使用缓存：推理开启、训练关闭（关键！训练时use_cache=False）use_cache=use_cacheifuse_cacheisnotNoneelse(self.config.use_cacheifnotself.trainingelseFalse)# 返回格式：字典 or 元组return_dict=return_dictifreturn_dictisnotNoneelseself.config.use_return_dict# ========== 2. 输入合法性检查 ==========# XOR 判断：必须只提供 input_ids 或 inputs_embeds 其中一个if(input_idsisNone)^(inputs_embedsisnotNone):raiseValueError("You must specify exactly one of input_ids or inputs_embeds")第二部分：词嵌入根据Tokenizer分词 chunk 后得到的input_ids（[batch_size, max_length]），传给Embedding层（Shape 为 [vocab_size, hidden_size]），得到 [batch_size, max_length, hidden_size]：# ========== 3. 词嵌入层 ==========# 如果没给嵌入向量，就用 token ID 查嵌入表得到ifinputs_embedsisNone:inputs_embeds=self.embed_tokens(input_ids)第三部分：开启缓存之类的（不看）# ========== 4. 训练冲突处理：梯度检查点 ≠ 缓存 ==========# 梯度检查点（省显存）和 KV 缓存（加速推理）互斥，训练时强制关闭缓存ifself.gradient_checkpointingandself.traininganduse_cache:logger.warning_once("`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.")use_cache=False# ========== 5. 缓存初始化检查 ==========# 想开启缓存但没传缓存对象 → 警告ifuse_cacheandcache_paramsisNone:logger.warning_once("QwenNemotronModel requires an initialized `HybridQwenNemotronDynamicCache` to return a cache. ""None was provided, so no cache will be returned.")# 初始隐藏状态 = 词嵌入向量hidden_states=inputs_embeds# ========== 6. 位置与缓存索引 ==========# 缓存位置：默认从 0 到 序列长度-1ifcache_positionisNone:cache_position=torch.arange(hidden_states.shape[1],device=hidden_states.device)# 位置ID = 缓存位置（RoPE用）ifposition_idsisNone:position_ids=cache_position.unsqueeze(0)第四部分：生成位置编码# ========== 7. 生成 RoPE 旋转位置编码（给Qwen注意力用） ==========position_embeddings=self.rotary_emb(hidden_states,position_ids)根据Embedding的输出作为输入，算出 cos 和 sin 旋转矩阵，所以position_embeddings应该是一个形似 (cos, sin) 的元组；第五部分：创建 Full Attention 和 Mamba 的掩码矩阵这是第一点和传统Decoder不一样的地方：前者是用于Full Attention的位置编码，是一个矩阵，所以需要输入为inputs_embeds的max_length构建矩阵，目的是通过下三角进行掩码，输出为一个 N x N 的矩阵；后者是一个单行向量，只 padding 掩码（没有因果掩码！），因为Mamba天生就是串行、单向、看不到未来的！速度极快。# ========== 8. 创建两种掩码：分别给 Attention 和 Mamba 使用 ==========causal_mask=self._update_causal_mask(attention_mask,inputs_embeds,cache_position)# 因果掩码mamba_mask=self._update_mamba_mask(attention_mask,cache_position)# Mamba专用掩码第六部分：逐层循环 Loopforlayer_idx,layerinenumerate(self.layers):# 如果需要输出隐藏态，先把上一层结果存起来ifoutput_hidden_states:all_hidden_states=all_hidden_states+(hidden_states,)存储上一层的隐藏状态；第七部分：“梯度检查点” 换取 “低显存占用”模型在前向传播的时候，会有很多中间结果，如果都存储起来会导致显存爆炸；所以我们会采用 “懒加载” 的方式，不存中间结果，当在反向传播的时候，再重新算一遍；（时间换空间）ifisinstance(layer,QwenDecoderLayer):# 训练 + 梯度检查点：使用torch激活检查点省显存ifself.gradient_checkpointingandself.training:hidden_states=self._gradient_checkpointing_func(layer.__call__,hidden_states,causal_mask,position_embeddings,cache_params,cache_position,)else:# 正常前向：传入 隐藏态、因果掩码、RoPE、缓存hidden_states=layer(hidden_states,attention_mask=causal_mask,position_embeddings=position_embeddings,cache_params=cache_params,cache_position=cache_position,)然后我们看看QwenDecoderLayer是怎么做的：defforward(self,hidden_states:torch.Tensor,# 输入：上一层的输出向量attention_mask:Optional[torch.Tensor]=None,# 因果掩码position_ids:Optional[torch.LongTensor]=None,position_embeddings:Optional[Tuple[torch.Tensor,torch.Tensor]]=None,# RoPEcache_params:Optional[HybridQwenNemotronDynamicCache