混元大模型：腾讯开源的目前最大的MoE模型（论文代码详解）

1.简介

在人工智能领域，大型语言模型（LLMs）的发展日新月异，它们在自然语言处理（NLP）、计算机视觉（CV）、语音识别和AI4Science等多个领域展现出了卓越的性能。最近，腾讯的Hunyuan团队推出了一款名为Hunyuan-Large的开源模型，这是一款基于Transformer的Mixture of Experts（MoE）模型，拥有惊人的3890亿参数和52亿激活参数，能够处理高达256K的token。这篇文章详细介绍了Hunyuan-Large的设计、性能和开源细节，它在多个基准测试中表现出色，包括语言理解和生成、逻辑推理、数学问题解决、编程、长文本处理和聚合任务等。

关键特点：

• 大规模合成数据：Hunyuan-Large在预训练阶段使用了比以往文献中更大的合成数据集，这有助于模型学习更丰富的表示，并更好地泛化到未见过的数据。
• 混合专家路由策略：模型采用了共享专家和专门专家的混合路由策略，以及创新的回收路由方法，以提高训练效率和模型性能。
• 关键值缓存压缩技术：通过关键值（KV）缓存压缩技术，Hunyuan-Large显著降低了内存压力，提高了推理效率。
• 专家特定学习率策略：模型为不同专家设置了特定的学习率，优化了训练过程。

性能表现： Hunyuan-Large在多个英文和中文的基准测试中展现了卓越的性能，与同类参数规模的最佳密集模型和MoE模型相比，它在常识理解、问答、数学推理、编程和聚合任务等任务中取得了最佳的整体性能。

开源贡献： 腾讯不仅通过其AI聊天机器人Yuanbao为用户提供服务，还通过开源Hunyuan-Large模型的代码和检查点，促进了技术传播和应用发展。这一举措有望激发未来的研究创新和实际应用，对社会产生积极影响。

这篇文章为我们揭开了Hunyuan-Large的神秘面纱，展示了腾讯在大型语言模型领域的最新进展和对开源社区的贡献。随着Hunyuan-Large的开源，我们期待它将如何推动未来AI技术的发展和应用。

官方网站：腾讯混元

权重地址：https://huggingface.co/tencent/Tencent-Hunyuan-Large

代码地址：https://github.com/Tencent/Tencent-Hunyuan-Large

论文地址：https://arxiv.org/abs/2411.02265

API接口：腾讯混元大模型_大语言模型_自然语言大模型- 腾讯云

2.论文

Pre-Training

数据

• 步骤1：指令生成。为了确保指令的多样性，我们使用高质量、知识丰富的数据源，如网页、基于Web的问答数据、代码库、书籍和其他资源作为种子。这些种子与不同的指令生成提示相配合，使我们能够生成覆盖不同领域的各种指令，这些指令具有不同的期望指令风格和复杂度。
• 步骤2：指令进化。为进一步改善这些初步指示的质素，我们会按以下三项指引加以改善：（a）提高指示的清晰度和资料的丰富程度。(b)通过自我指导增强扩展低资源域指令。(c)改进指令以增加其难度级别。这些不断发展的高质量和挑战性指令使我们的模型能够更有效地从合成数据中获益，从而跨越原始的能力界限。
• 步骤3：生成答案。我们利用几个专门的模型来为上述进化的指令生成信息丰富和准确的答案。这些模型大小不一，并且是精心设计的专用模型，用于合成针对各个领域中的指令的专家级响应。
• 步骤4：答案过滤。为了过滤合成的解释-响应对，作者采用了一个批判模型，并进行自我一致性检查，其中我们生成多个答案，以执行任务，如客观问答任务的自我一致性过滤，确保可靠性和准确性。这个过程使我们能够有效地删除任何低质量或不一致的数据，确保在预训练中利用高质量的文本。

模型结构

模型参数：

共享专家和混合专家

混元大模型采用了一种混合路由策略，既包括一个共享专家来处理所有token，也结合了传统的top-k路由策略，利用多个MOE（Mixture of Experts）专家处理专业领域知识。

在混元模型中，特别设定了1位共享专家，负责捕捉所有token共有的知识。同时，还有16位专业专家，他们负责动态学习特定领域的知识。通常情况下，每个token会激活得分最高的前1位专业专家。

这里我们以单个token的选择为例，画了一张草图

循环路由

传统的top-k路由通常与MoE中专家的最大负载有关，其中过载专家的token在训练期间被丢弃。容量因子越大，丢弃的令牌越少，但训练效率越低。过多的token丢弃会导致关键信息的丢失，进而影响训练的稳定性。

作者为原始top-k路由过程中丢弃的令牌开发了一种新的循环路由策略，该技术需要将过载专家的token额外随机分配给没有超过其容量的其他专家。该方法力求在优化训练效率的同时保留重要信息，从而确保模型训练的整体有效性和效率。

专家特定学习率量表

以前的工作已经阐明了LLM中亚当式优化器的最佳学习率和批量大小B之间更合适的联系：

表示AdamW的学习率。表示训练速度和数据效率之间的权衡点，在Kaplan等人中指出。

但这对MoE模型不合适，因为不同的专家在训练的数据方面不平衡。因此作者对其做了更改。具体来说，对于专业专家，有效的批次大小应该大致除以专家的数量，从而使他们的最佳学习率表示为。共享专家和专业专家之间的学习率比例比例为，在作者的设置中约为0.31。

因此，在配置混元大模型的学习率时，作者为共享专家分配了最优的，并根据这个比例刻意地降低了专业专家的学习率。

预训练

MoE尺寸法则

通常，密集模型的训练计算预算使用来估计，其中N表示参数的数量，D表示训练令牌。然而，对于具有较长序列的MOE模型（例如，8K、32K和256K），计算预算公式因注意力复杂度和稀疏激活而异。

经过细致的计算，作者确定了MoE模型的精确计算预算C，其中公式中的N表示激活参数的数量：

为了获得精确的估计，作者采用了临界批量大小，它优化了时间和计算效率之间的权衡，而用于计算最小计算预算 的公式，它考虑了批量大小B对计算预算的影响。公式如下：

这里，C是在给定模型规模和数据规模下的计算预算，B是批量大小，而 是在给定模型规模 LL 下的临界批量大小，它优化了训练速度和数据效率之间的权衡。

这个公式是基于Kaplan et al. (2020)和Li et al. (2024a)的研究，它帮助研究者确定在给定模型规模和数据规模下，最小计算预算 如何随着批量大小 B的变化而变化。通过这种分析，研究者可以更好地理解批量大小对计算预算的影响，并据此做出合理的模型设计和训练决策。

总的来说，公式3是一个关于计算预算和批量大小之间关系的数学表达式，它为Hunyuan-Large模型的设计和优化提供了重要的参考依据。通过这种分析，研究者可以更精确地确定在特定批量大小下模型的最优计算预算，以实现成本效益最大化。

图3是一个关于模型规模和计算预算之间关系的可视化表示，它为Hunyuan-Large模型的设计和优化提供了重要的参考依据。

Learning Rate Scheduling 学习率调度

学习率时间表被描绘成三个连续的阶段：初始预热阶段，随后是一个长期的阶段逐渐衰减，并在一个简洁的退火阶段达到高潮。渐进衰减的扩展阶段的优点是它善于平衡解空间的探索与向最优解的收敛。通过在初始预训练阶段保持较高的学习速率，模型能够有效地在解空间的不同区域中导航，从而避免过早收敛到次优局部最小值。随着训练的进行，学习速率的递增减少确保了系统地收敛到更优的解。

在最后的5%的预训练记号中，作者引入了简短的退火阶段，其中学习速率被降低到其峰值的十分之一。这种方法有助于模型细致地微调其参数，从而实现上级程度的泛化，从而提高其整体性能。此外，在此阶段，我们优先使用最高质量的数据集，这在增强模型在退火阶段的性能方面起着关键作用。

图4是一个关于训练数据量和计算预算之间关系的可视化表示，它为Hunyuan-Large模型的训练和优化提供了重要的参考依据。

长上下文的预训练

在退火阶段之后，浑源-Large将在更长的序列（高达256 K个令牌）上进行训练，以启用其长上下文能力。

具体地，长上下文预训练阶段包含两个阶段（即，随着32 K → 256 K逐渐增加令牌长度）。我们采用RoPE用于构建位置嵌入，并在256 K预训练阶段将RoPE基本频率扩展到10亿。

对于数据，我们完全依赖于从书籍和代码中获得的自然长上下文数据（占语料库的近25%），并将其与正常长度的预训练数据（近75%）混合，以形成我们的长上下文预训练语料库

我们还发现，它不需要太多的培训LLM获得长上下文的能力。在32 K和256 K阶段的每一个阶段中，我们都使用了大约100亿个令牌的长上下文预训练语料库。每个阶段的长上下文预训练都可以获得令人满意的长上下文能力。

Post-Training

这一阶段包括一个监督微调（SFT）阶段和一个从人的反馈中强化学习（RLHF）阶段。

SFT数据

整个SFT数据量超过100万。

1. 指令提取。作者开发了一个专门针对数学、逻辑推理和基于知识的问答等领域的指令抽取模型，其主要目标是从公开的可用数据源(如网页、百科全书等)中有效地提取适合于指令调整的数据。提取的数据包括指令和相应的参考答案。
2. 指令泛化。具体地说，作者设计并训练了一个指令概括系统，该系统能够在逐步增加目标指令的难度和复杂性的同时对其进行概括。这个系统的中心配方在于通过合成简单和复杂指令之间的大量映射来训练模型。此外，我们构建了一个结构良好的指令分类法及其相应的分类模型，旨在分析和平衡各种指令类型在SFT数据中的分布。
3. 指令平衡。通过指令提取和泛化过程，作者积累了超过1000万条指令。然而，许多生成的指令具有非常相似的语义，指令类型分布自然是不平衡的。为了提高指令复杂度，同时保持均衡的指令分布，作者为每条指令附加标签。这些标签包含多个尺寸。通过仔细标记这些标签，作者可以更准确地理解和分析指令集的特性。通过在SFT过程中保证不同类型指令的数量充足和均衡分布，可以有效缓解特定指令类型的过拟合或欠拟合问题，从而提高模型的泛化能力和对不同应用场景的适应性
4. 数据质量控制。SFT数据的质量是上级性能的基础。作者主要通过以下三种方法来保证我们的SFT数据的高质量。
   1. 基于规则的过滤。作者发现了SFT数据中的一些常见问题，如数据截断错误、重复、乱码和格式错误。因此，作者开发了一套基于规则的数据过滤策略，以防止上述指令提取和生成模型产生不良输出。
   2. 基于模型的过滤。为了从大量合成的指令数据中自动提取高质量的SFT数据，作者基于混元系列的70B密度模型训练了一个批评性模型。该模型为每个指令样本分配一个四级质量分数，评估生成的答复的准确性、相关性、完整性、有用性和清晰度等方面，以及其他可能的数据质量问题。
   3. 基于人工的过滤。在模型训练之前，通过基于规则和基于模型的方法过滤的SFT数据进一步经过人工注释，确保答案符合期望的任务特定响应模式，并避免引入额外的低质量问题。

在SFT中，作者根据高质量数据（超过100万）对预训练模型进行微调，共3个epoch。

DPO强化学习

作者采用线下和线上培训相结合的单阶段培训策略，我们利用预编译的偏好数据集来增强可控性，同时利用当前的策略模型来为每个提示生成多个响应，并使用我们的奖励模型来选择最喜欢和最不偏好的响应。

训练

关于MoE的训练方法，论文中没有讲，但这里我想应该做个简单的介绍。

• 预训练阶段：MoE模型通常有一个预训练阶段，这个阶段中，模型会使用大量的数据进行训练，以获得基础的语言或任务处理能力。
• 任务特定训练：在预训练之后，MoE模型可能会进入一个任务特定的训练阶段。MoE模型的不同专家可能会在同一个数据集上进行训练，但是通过门控网络（gating network）的控制，每个专家可能会专注于数据的不同子集。这种机制允许模型在处理特定输入时只激活相关的专家，从而提高效率和性能。

参考：https://www.zhihu.com/question/634845272

3.代码

混元大模型的代码其实和其他MoE模型差不多，如果读者已经对MoE的代码，相当熟悉了，可以直接跳过；但如果读者还没有了解过MoE的代码实现，不妨看看混元大模型的代码，个人认为是结构比较清晰的，非常适合上手。

代码地址：GitHub - Tencent/Tencent-Hunyuan-Large

权重地址：https://huggingface.co/tencent/Tencent-Hunyuan-Large/tree/main

由于代码权重实在太大了，作者无法进行运行，故代码详解部分只做代码模型架构的解释。

所有模型架构的位置在models/modeling_hunyuan.py下。接下来，我们将从外到内逐步讲解混元大模型的结构。

HunYuanForCausalLM

通过__init__()函数可以知道，model就是Hunyuan模型，我们在下文再详细介绍；lm_head就是将模型输出结果映射到词表长度的线性层，用于预测概率。

def __init__(self, config: HunYuanConfig):
    super().__init__(config)
    self.model = HunYuanModel(config)
    self.vocab_size = config.vocab_size
    self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

接下来是forward()函数，这里就是运行模型并预测每个词的概率，以及训练时的损失函数计算，如果你对大模型有所了解，这部分和其他模型基本上没什么差别。

class HunYuanForCausalLM(HunYuanPreTrainedModel):
    
    def forward(...) -> Union[Tuple, CausalLMOutputWithPast]:
        
        # 1.配置信息
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # 2.模型输出 decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
        outputs = self.model(       # HunYuanModel
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        # 3.线性层预测概率  （分多卡和单卡版）
        hidden_states = outputs[0]
        if self.config.pretraining_tp > 1:      # 如果多卡并行
            lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)    # 按pretraining_tp分割成多个部分。
            logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]    # 对每个部分进行线性变换，得到多个 logits。
            logits = torch.cat(logits, dim=-1)  # 在最后一个维度上拼接起来
        else:       # 如果单卡
            logits = self.lm_head(hidden_states)
        logits = logits.float()

        # 4.计算损失
        loss = None
        if labels is not None:      # 自回归损失函数计算
            # Shift so that tokens < n predict n  将 logits 的最后一个维度去掉一位，labels 则去掉第一个维度的第一位，使 logits 的每个位置预测下一个位置的 label。
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens  展平 tokens
            loss_fct = CrossEntropyLoss()
            shift_logits = shift_logits.view(-1, self.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            # Enable model parallelism
            shift_labels = shift_labels.to(shift_logits.device)     # 移动到与 shift_logits 相同的设备上。
            loss = loss_fct(shift_logits, shift_labels)     # 计算损失。

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        # 5.返回结果
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

HunYuanModel

通过__init__()函数，我们可以知道模型的核心层由诸多HunYuanDecoderLayer组成，同时代码中还设置了attention的计算方法、归一化层的方法等。

class HunYuanModel(HunYuanPreTrainedModel):
    def __init__(self, config: HunYuanConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList(    # 核心层
            [HunYuanDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]        
        )
        self._use_sdpa = config._attn_implementation == "sdpa"
        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
        self.norm = HunYuanRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        self.cla = config.use_cla
        self.cla_share_factor = config.cla_share_factor

        self.gradient_checkpointing = False
        # Initialize weights and apply final processing
        self.post_init()

其forward()函数如下，整体逻辑是先生成注意力掩码、位置信息等信息，然后调用每一层的decoder_layer进行处理，并保存中间结果，如隐藏状态和注意力权重。最后输出。

class HunYuanModel(HunYuanPreTrainedModel):
    @add_start_docstrings_to_model_forward(HUNYUAN_INPUTS_DOCSTRING)
    def forward(...) -> Union[Tuple, BaseModelOutputWithPast]:
        # 初始化部分省略

        past_key_values_length = 0
        if use_cache:
            use_legacy_cache = not isinstance(past_key_values, Cache)
            if use_legacy_cache:        # 是否使用kv-cache
                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
            past_key_values_length = past_key_values.get_usable_length(seq_length)

        if position_ids is None:        # 嵌入位置信息
            device = input_ids.device if input_ids is not None else inputs_embeds.device
            position_ids = torch.arange(
                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
            )
            position_ids = position_ids.unsqueeze(0)

        if inputs_embeds is None:       # embedding
            inputs_embeds = self.embed_tokens(input_ids)
        
        # Fix lora with gradient checkpointing training
        if self.training and inputs_embeds.is_leaf:
            inputs_embeds.requires_grad = True

        if self._use_flash_attention_2:     # 选择注意力掩码
            # 2d mask is passed through the layers
            attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
        elif self._use_sdpa and not output_attentions:
            # output_attentions=True can not be supported when using SDPA, and we fall back on
            # the manual implementation that requires a 4D causal mask in all cases.
            attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
                attention_mask,
                (batch_size, seq_length),
                inputs_embeds,
                past_key_values_length,
            )
        else:
            # 4d mask is passed through the layers
            attention_mask = _prepare_4d_causal_attention_mask(
                attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
            )

        # embed positions
        hidden_states = inputs_embeds

        # decoder layers
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        next_decoder_cache = None

        prev_kv_states = None
        for layer_idx, decoder_layer in enumerate(self.layers):     # 多层解码器进行计算
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    decoder_layer.__call__,
                    hidden_states,
                    attention_mask,
                    position_ids,
                    past_key_values,
                    output_attentions,
                    use_cache,
                    prev_kv_states,
                )
            else:
                layer_outputs = decoder_layer(
                    hidden_states,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    past_key_value=past_key_values,
                    output_attentions=output_attentions,
                    use_cache=use_cache,
                    kv_states=prev_kv_states
                )
                
            # 保存中间结果，如隐藏状态和注意力权重。
            hidden_states = layer_outputs[0]

            if use_cache:       # KV-cache
                next_decoder_cache = layer_outputs[2 if output_attentions else 1]

            if output_attentions:       # 注意力权重。
                all_self_attns += (layer_outputs[1],)

            kv_states = layer_outputs[-1]       # kv

            if self.cla and layer_idx % self.cla_share_factor == 0:
                prev_kv_states = kv_states

        hidden_states = self.norm(hidden_states)

        # 输出处理add hidden states from the last decoder layer
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        next_cache = None
        if use_cache:
            next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
        if not return_dict:
            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=next_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
        )

其父类里面定义了如何初始化

class HunYuanPreTrainedModel(PreTrainedModel):
    ...

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):   # 使用正态分布初始化权重，并将偏置项设置为零。
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):  # 使用正态分布初始化权重，并将指定的填充索引位置的权重设置为零。
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()

DecoderLayer

从DecoderLayer的__init__()函数，我们可以看出

1. 混元模型的注意力有不同的模式，在同一文件下定义了三个类，可以自行选择，包括HunYuanAttention（原始注意力，即Attention Is All You Need）、HunYuanFlashAttention2、HunYuanSdpaAttention
2. mlp层可以根据num_experts的值选择HunYuanMoE或者HunYuanMLP
3. 归一化为HunYuanRMSNorm

class HunYuanDecoderLayer(nn.Module):
    def __init__(self, config: HunYuanConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.layer_idx = layer_idx

        self.self_attn = HUNYUAN_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)

        if config.num_experts > 1:
            self.mlp = HunYuanMoE(config, layer_idx=layer_idx)
        else:
            self.mlp = HunYuanMLP(config, layer_idx=layer_idx, is_shared_mlp=False)
        self.input_layernorm = HunYuanRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = HunYuanRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

其forward方法如下，这里和普通的decoder其实没有什么差别，熟悉的可以直接跳过。

class HunYuanDecoderLayer(nn.Module):
    def forward(...) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
        ...

        residual = hidden_states

        hidden_states = self.input_layernorm(hidden_states)     # 输入归一化

        # Self Attention
        hidden_states, self_attn_weights, present_key_value, kv_states = self.self_attn(        # 自注意力机制
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
            kv_states=kv_states,
            **kwargs,
        )
        hidden_states = residual + hidden_states        # 残差连接

        # Fully Connected
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)        # 后注意力归一化
        hidden_states = self.mlp(hidden_states)     # 全连接层
        hidden_states = residual + hidden_states    # 残差连接

        outputs = (hidden_states,)

        if output_attentions:   # 是否输出注意力权重
            outputs += (self_attn_weights,)

        if use_cache:       # 是否输出缓存
            outputs += (present_key_value,)

        outputs += (kv_states,)     # 是否返回键值状态

        return outputs

HunYuanAttention

混元大模型实现的注意力方法有三种，这里我们只看原始版本的。

初始化方法如下：

• 主要内容是qkv矩阵的初始化，具体来说，如果是cross-attention，就不会初始化kv矩阵,如果是self-attention，就会初始化kv矩阵
• 以及是否使用归一化层

class HunYuanAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(self, config: HunYuanConfig, layer_idx: Optional[int] = None):
        ...

        # layer_idx 从 0 开始
        self.attention_type = 'cross' if config.use_cla and layer_idx % config.cla_share_factor != 0 else 'self'
       
        ...

        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
        if self.attention_type == 'self':
            self.k_proj = nn.Linear(
                self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias
            )
            self.v_proj = nn.Linear(
                self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias
            )
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)
        if self.use_qk_norm:
            self.query_layernorm = HunYuanRMSNorm(self.head_dim, eps=config.rms_norm_eps)
            self.key_layernorm = HunYuanRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self._init_rope()

在forward层中涉及了不少多GPU并行运算的代码，但主体内容和普通attention其实差不多，部分不重要的代码被我删去了。

class HunYuanAttention(nn.Module):
    def forward(...) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        ...

        bsz, q_len, _ = hidden_states.size()

        if self.config.pretraining_tp > 1:      # 多GPU并行训练
            query_slices = self.q_proj.weight.split(        # 切分 q_proj.weight
                (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0
            )
            query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]    # 计算 query_states。
            query_states = torch.cat(query_states, dim=-1)      # 将不同设备的结果合并。

            if self.attention_type == "cross" and kv_states is not None and isinstance(kv_states, tuple):   # 交叉注意力且 kv_states 不为空，直接使用 kv_states 作为键值状态。
                orig_key_states, orig_value_states = kv_states
                key_states, value_states = kv_states
            else:       # 否则，切分 k_proj.weight 和 v_proj.weight 并计算 key_states 和 value_states。
                key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp
                key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
                value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)

                key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]
                key_states = torch.cat(key_states, dim=-1)

                value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]
                value_states = torch.cat(value_states, dim=-1)
                orig_key_states, orig_value_states = key_states, value_states

        else:       # 单GPU
            query_states = self.q_proj(hidden_states)       # q
            if self.attention_type == "cross" and kv_states is not None and isinstance(kv_states, tuple):
                orig_key_states, orig_value_states = kv_states
                key_states, value_states = kv_states
            else:
                key_states = self.k_proj(hidden_states)
                value_states = self.v_proj(hidden_states)
                orig_key_states, orig_value_states = key_states, value_states

        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

        ...

        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)        # 计算 cos 和 sin
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)   # 应用旋转位置编码（RoPE）

        if self.use_qk_norm:    # 对 Query-Key 层归一化
            query_states = self.query_layernorm(query_states)
            key_states = self.key_layernorm(key_states)

        if past_key_value is not None:      # KV-cache  如果有过去的键值对，更新键值状态。
            cache_kwargs = {"sin": sin, "cos": cos}  # Specific to RoPE models
            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)

        key_states = repeat_kv(key_states, self.num_key_value_groups)       # 重复键值状态
        value_states = repeat_kv(value_states, self.num_key_value_groups)

        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)    # 计算注意力权重

        if attention_mask is not None:      # 如果有注意力掩码，应用注意力掩码。
            attn_weights = attn_weights + attention_mask    # 添加注意力掩码

        
        attn_output = torch.matmul(attn_weights, value_states)      # 计算注意力输出 a*v

        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
            raise ValueError(
                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
                f" {attn_output.size()}"
            )

        attn_output = attn_output.transpose(1, 2).contiguous()

        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)

        if self.config.pretraining_tp > 1:  # 如果 pretraining_tp 大于1，切分输出权重并计算最终输出。
            attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)
            o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)
            attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])
        else:
            attn_output = self.o_proj(attn_output)  # 单GPU直接计算最终输出

        if not output_attentions:
            attn_weights = None

        return attn_output, attn_weights, past_key_value, (orig_key_states, orig_value_states)

其中的主要代码如下

RoPE

定义的原版RoPE类

class HunYuanRotaryEmbedding(nn.Module):        # Rope旋转位置编码
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        super().__init__()

        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))    # 计算逆频率 inv_freq
        inv_freq = inv_freq.bfloat16()
        self.register_buffer("inv_freq", inv_freq, persistent=False)    # 注册缓冲区：使用 self.register_buffer 方法将计算结果注册为模型的缓冲区变量，并且设置 persistent=False 表示这个缓冲区不会被保存在模型的状态字典中。

        # Build here to make `torch.jit.trace` work.
        self._set_cos_sin_cache(    # 调用 _set_cos_sin_cache 方法预计算缓存
            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
        )

    def _set_cos_sin_cache(self, seq_len, device, dtype):
        self.max_seq_len_cached = seq_len
        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32)   # 根据输入序列长度 seq_len，计算时间步长 t

        freqs = torch.outer(t, self.inv_freq)   # 计算两个张量的外积
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        emb = torch.cat((freqs, freqs), dim=-1).float()     # 计算频率 freqs 并拼接成 emb
        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)   # 计算并缓存 cos 和 sin 值。
        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)

    def forward(self, x, seq_len=None):
        # x: [bs, num_attention_heads, seq_len, head_size]
        if seq_len > self.max_seq_len_cached:   # 如果超过，更新缓存
            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)

        return (        # 返回相应的 cos 和 sin 缓存值。
            self.cos_cached[:seq_len].to(dtype=x.dtype),
            self.sin_cached[:seq_len].to(dtype=x.dtype),
        )

初始化rope的部分在HunYuanAttention下面的_init_rope()，这里省略，我们直接来看rope的核心代码。

def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):    # 应用Rope
   
    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

def rotate_half(x):     # Rope里面旋转的那部分
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2:]
    return torch.cat((-x2, x1), dim=-1)

repeat_kv

输入：hidden_states:形状为 (batch, num_key_value_heads, seqlen, head_dim) 的四维张量。

输出：最后，将张量的形状调整为 (batch, num_attention_heads, seqlen, head_dim)，其中 num_attention_heads = num_key_value_heads * n_rep。

def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:     # 用于在特定维度上重复张量。
    """
    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
    """
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)   # 使用 expand 方法在新的维度上扩展张量，使其形状变为 (batch, num_key_value_heads, n_rep, seqlen, head_dim)。
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)    # 将张量的形状调整为 (batch, num_attention_heads, seqlen, head_dim)，其中 num_attention_heads = num_key_value_heads * n_rep。

重复张量的意义

1. 多头注意力机制：
   1. 在多头注意力机制中，通常会有多个注意力头（num_attention_heads），每个头负责捕捉不同的特征。num_key_value_heads 是实际用于计算 key 和 value 的头的数量。
   2. 通过 repeat_kv() 函数，可以将 num_key_value_heads 扩展到 num_attention_heads，使得每个注意力头都能访问到相同的 key 和 value 信息。
2. 提高计算效率：
   1. 在某些情况下，为了减少计算量，可能会使用较少的 num_key_value_heads 来计算 key 和 value，然后通过重复这些 key 和 value 来匹配 num_attention_heads。
   2. 这样可以在不显著增加计算资源的情况下，提高模型的性能和灵活性。
3. 适应不同的模型架构：
   1. 不同的模型架构可能有不同的 num_attention_heads 和 num_key_value_heads 配置。repeat_kv() 函数提供了一种灵活的方式来适配这些不同的配置

HunYuanMOE

HunYuanMLP

从代码可以看出，这个MLP层的实质由门控层和两层线性层组成。其运算过程如下：

其中，x是输入，out是输出，up和down分别为两层线性层，gate为门控层，act为激活函数。

class HunYuanMLP(nn.Module):
    def __init__(self, config: HunYuanConfig, layer_idx=None, is_shared_mlp=False):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.hidden_size = config.hidden_size
        if is_shared_mlp:
            self.intermediate_size = config.intermediate_size * config.num_shared_expert
        else:
            self.intermediate_size = config.intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)    # 门控层
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        if self.config.pretraining_tp > 1:      # 多卡
            slice = self.intermediate_size // self.config.pretraining_tp    # 权重矩阵按照指定的切片大小进行切分。
            gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)
            up_proj_slices = self.up_proj.weight.split(slice, dim=0)
            down_proj_slices = self.down_proj.weight.split(slice, dim=1)

            gate_proj = torch.cat(      # 对输入 x 进行多次线性变换，并将结果在最后一个维度上拼接起来
                [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1
            )
            up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)

            intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)
            down_proj = [
                F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp)
            ]
            down_proj = sum(down_proj)
        else:       # 单卡
            down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

        return down_proj

HunYuanMoE

这个层实现了混合专家的具体功能，具体来说：

1. 如果有共享的MLP层，输入数据先经过共享MLP，得到hidden_states_mlp
2. 通过门控层(HunYuanTopKGate)计算每个专家的负载
3. 拆分输入数据，然后将不同数据交给不同的专家进行处理，得到expert_outputs
   1. 根据dispatch_mask，使用torch.einsum分配给不同的专家
   2. 使用沿第一个维度（通常是批次维度）切分成self.num_experts个块，每个块对应一个专家的处理输入。
   3. 遍历每个块和专家，计算输出
4. 将不同专家的输出拼起来，并恢复原来的形状，得到combined_output
5. 如果有共享的MLP层，将专家输出结果和共享MLP的输出拼起来
6. 返回结果。

class HunYuanMoE(nn.Module):
    def __init__(self, config: HunYuanConfig, layer_idx: Optional[int] = None):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.moe_topk = config.moe_topk
        self.num_experts = config.num_experts
        if config.use_mixed_mlp_moe:    # 创建共享的 MLP 层
            self.shared_mlp = HunYuanMLP(config, layer_idx=layer_idx, is_shared_mlp=True)
        self.gate = HunYuanTopKGate(config, layer_idx=layer_idx)    # 门控机制
        self.experts = nn.ModuleList(   # 创建专家池
            [HunYuanMLP(config, layer_idx=layer_idx, is_shared_mlp=False) for _ in range(config.num_experts)]
        )

    def forward(self, hidden_states):
        bsz, seq_len, hidden_size = hidden_states.shape

        if self.config.use_mixed_mlp_moe:
            hidden_states_mlp = self.shared_mlp(hidden_states)      # 对输入进行共享MLP处理

        l_moe, combine_weights, dispatch_mask, exp_counts = self.gate(hidden_states)    # 门控机制计算每个专家的分配权重和掩码。

        reshaped_input = hidden_states.reshape(-1, hidden_size)

        # 根据掩码将输入分派给不同的专家。
        dispatched_input = torch.einsum("sec,sm->ecm", dispatch_mask.type_as(hidden_states), reshaped_input)    # [s,num_expect,容量],[s,hidden_size]->[num_expect,容量,hidden_size]
        chunks = dispatched_input.chunk(self.num_experts, dim=0)    # 按专家数量切分，每个专家处理一个输入部分。
        expert_outputs = []

        for chunk, expert in zip(chunks, self.experts):     # 每个专家处理分派给它的输入部分。
            expert_outputs.append(expert(chunk))

        # 将专家的输出重新组合成原始输入的形状。
        expert_output = torch.cat(expert_outputs, dim=0)
        combined_output = torch.einsum("sec,ecm->sm", combine_weights.type_as(hidden_states), expert_output)    # [s,num_expect,容量],[num_expect,容量,hidden_size]->[s,hidden_size]
        combined_output = combined_output.reshape(bsz, seq_len, hidden_size)

        if self.config.use_mixed_mlp_moe:       # 混合MLP输出与组合输出相加
            output = hidden_states_mlp + combined_output
        else:       # 如果没有启用混合MLP模式，直接返回组合后的专家输出。
            output = combined_output

        return output

HunYuanTopKGate

这个门控机制负责决定每个输入token应该被分配给哪些专家进行处理。简而言之，这个类的作用是根据模型配置和输入的隐藏状态，来计算每个token分配给不同专家的概率，并据此进行路由。

class HunYuanTopKGate(nn.Module):
    def __init__(self, config: HunYuanConfig, layer_idx: Optional[int] = None):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.moe_topk = config.moe_topk
        self.drop_tokens = config.moe_drop_tokens
        self.min_capacity = 8
        self.random_routing_dropped_token = config.moe_random_routing_dropped_token
        self.wg = nn.Linear(config.hidden_size, config.num_experts, bias=False, dtype=torch.float32)    # 用于将hidden_size映射到专家数量

    def forward(self, hidden_states):
        bsz, seq_len, hidden_size = hidden_states.shape
        hidden_states = hidden_states.reshape(-1, hidden_size)
        if self.wg.weight.dtype == torch.float32:
            hidden_states = hidden_states.float()
        logits = self.wg(hidden_states)     # 通过线性层 self.wg 计算门控逻辑 logits
        if self.moe_topk == 1:
            gate_output = top1gating(logits, random_routing_dropped_token=self.random_routing_dropped_token)
        else:
            gate_output = topkgating(logits, self.moe_topk)     # 路由机制

        return gate_output

路由机制 topkgating()

这个函数是整个路由机制的核心。这种机制允许每个token被分配给概率最高的前k个专家，提供了一种更加灵活的路由策略，可以在模型的负载均衡和处理能力之间取得平衡。

def topkgating(logits: Tensor, topk: int):
    logits = logits.float()     # 线性层hidden_size -> num_expects的结果，[s,m=num_expects]
    gates = F.softmax(logits, dim=1)    # 计算每个token对每个专家的路由概率。
    expert_capacity = topk * gates.shape[0]
    num_experts = int(gates.shape[1])
    # Top-k router probability and corresponding expert indices for each token.
    # Shape: [tokens_per_group, num_selected_experts].
    expert_gate, expert_index = torch.topk(gates, topk)     # 使用 torch.topk 确定每个token的Top-K专家及其对应的路由概率。
    expert_mask = F.one_hot(expert_index, num_experts)      # 使用 F.one_hot 生成专家掩码
    # For a given token, determine if it was routed to a given expert.
    # Shape: [tokens_per_group, num_experts]
    expert_mask_aux = expert_mask.max(dim=-2)[0]
    tokens_per_group_and_expert = torch.mean(expert_mask_aux.float(), dim=-2)   # 计算每个专家的负载
    router_prob_per_group_and_expert = torch.mean(gates.float(), dim=-2)
    l_aux = num_experts**2 * torch.mean(tokens_per_group_and_expert * router_prob_per_group_and_expert)     # 生成辅助损失 l_aux

    # 计算每个token的优先级
    gates_s = torch.clamp(      # 计算专家门控权重的加权和，并对其进行裁剪，确保结果不会小于一个极小值。
        torch.matmul(expert_mask.float(), gates.unsqueeze(-1)).sum(dim=1), min=torch.finfo(gates.dtype).eps
    )
    router_probs = gates / gates_s      # 计算路由概率
    # Make num_selected_experts the leading axis to ensure that top-1 choices have priority over top-2 choices, which have priority over top-3 choices, etc.
    expert_index = torch.transpose(expert_index, 0, 1)
    expert_index = expert_index.reshape(-1)     # Shape: [num_selected_experts * tokens_per_group]

    # Create mask out of indices.
    expert_mask = F.one_hot(expert_index, num_experts).to(torch.int32)     # 计算每个专家被选择的次数 Shape: [tokens_per_group * num_selected_experts, num_experts].
    exp_counts = torch.sum(expert_mask, dim=0).detach()     # 计算每个专家被选择的次数

    # 计算每个令牌在目标专家中的优先级。Experts have a fixed capacity that we cannot exceed. A token's priority within the expert's buffer is given by the masked, cumulative capacity of its target expert.
    # Shape: [tokens_per_group * num_selected_experts, num_experts].
    token_priority = torch.cumsum(expert_mask, dim=0) * expert_mask - 1
    # Shape: [num_selected_experts, tokens_per_group, num_experts].
    token_priority = token_priority.reshape((topk, -1, num_experts))
    # Shape: [tokens_per_group, num_selected_experts, num_experts].
    token_priority = torch.transpose(token_priority, 0, 1)
    # For each token, across all selected experts, select the only non-negative (unmasked) priority. Now, for group G routing to expert E, token T has non-negative priority (i.e. token_priority[G,T,E] >= 0) if and only if E is its targeted expert.
    # Shape: [tokens_per_group, num_experts].
    token_priority = torch.max(token_priority, dim=1)[0]        # 在所有选择的专家中，选择唯一的非负优先级。

    # 生成有效的调度掩码
    # Token T can only be routed to expert E if its priority is positive and less than the expert capacity. One-hot matrix will ignore indices outside the range [0, expert_capacity).
    # Shape: [tokens_per_group, num_experts, expert_capacity].
    valid_mask = torch.logical_and(token_priority >= 0, token_priority < expert_capacity)   # 布尔张量 valid_mask，标记了token_priority中在[0, expert_capacity)范围内的元素。
    token_priority = torch.masked_fill(token_priority, ~valid_mask, 0)      # 将 token_priority 中不在有效范围内的值设为 0。
    dispatch_mask = F.one_hot(token_priority, expert_capacity).to(torch.bool)   # 将 token_priority 转换为 one-hot 编码，生成 dispatch_mask
    valid_mask = valid_mask.unsqueeze(-1).expand(-1, -1, expert_capacity)
    dispatch_mask = torch.masked_fill(dispatch_mask, ~valid_mask, 0)    # 将 dispatch_mask 中对应 valid_mask 为 False 的位置设为 0。

    # The combine array will be used for combining expert outputs, scaled by the router probabilities. Shape: [num_groups, tokens_per_group, num_experts, expert_capacity].
    combine_weights = torch.einsum("...te,...tec->...tec", router_probs, dispatch_mask)
    exp_counts_capacity = torch.sum(dispatch_mask)      # 计算调度掩码的期望容量
    exp_capacity_rate = exp_counts_capacity / (logits.shape[0]*topk)    # 计算容量利用率
    return [l_aux, exp_capacity_rate], combine_weights, dispatch_mask, exp_counts

代码的关键返回值有：

• dispatch_mask 是一个布尔掩码，用于指示每个token应该被发送到哪些专家。
• combine_weights 是用于合并专家输出的权重

dispatch_mask

具体来说，dispatch_mask 是通过以下步骤得到的：

• 首先，计算每个token的优先级，这个优先级是基于token被分配到专家的顺序。
• 通过使用one-hot编码，将每个token的专家索引转换为一个掩码，指示该token是否被分配给特定的专家。
• 接下来，生成一个有效掩码，确保每个token的优先级在有效范围内（即优先级非负且小于专家的容量）。
• 然后，将这个有效掩码应用到优先级上，填充无效的优先级值为0。最后，将优先级转换为one-hot编码形式，生成最终的dispatch_mask，这个掩码的形状通常为 (tokens_per_group, num_experts, expert_capacity)，指示每个token在每个专家的缓存中应该占据的位置。

这个过程确保了每个token只被路由到其优先级有效且不超过专家容量的专家，从而实现了高效的负载分配和输出合并。

combine_weights

具体来说，combine_weights 的计算过程是这样的：

1. 首先，通过softmax函数将logits（即每个token分配给每个专家的原始分数）转换成概率分布，这些概率表示每个token被分配到每个专家的可能性。

2. 接着，计算每个token分配给所有专家的总概率（gates_s），这个总概率是通过对softmax概率进行求和得到的。

3. 然后，对softmax概率进行归一化处理，得到每个token分配给每个专家的归一化概率（router_probs）。这是通过将softmax概率除以每个token分配给所有专家的总概率来实现的，目的是确保每个token分配给所有专家的归一化概率之和为1。

4. 接下来，使用torch.topk函数从归一化概率中选择每个token的前topk个最可能的专家，并将这些选择转换为one-hot编码形式的掩码（dispatch_mask），这个掩码指示每个token应该被发送到哪些专家。

5. 最后，通过torch.einsum函数结合归一化的门控信号（router_probs）和调度掩码（dispatch_mask）来计算combine_weights。这个操作实际上是将每个token的归一化概率与它应该被发送到的每个专家的调度掩码相乘，得到一个四维数组，其中包含了每个token对于每个专家输出的贡献权重。

combine_weights 的形状通常是 (num_groups, tokens_per_group, num_experts, expert_capacity)，其中 num_groups 是批次中的组数，tokens_per_group 是每个组中的token数，num_experts 是专家的数量，expert_capacity 是每个专家的容量。这个权重数组在后续步骤中用于将各个专家的输出按照它们对每个token的贡献权重合并起来，形成最终的模型输出。这样，每个专家只对其被分配的token贡献输出，而没有被分配的token则不包含在该专家的输出中。

top1gating()

top1gating 函数的目的是实现基于 logits 的 Top-1 门控机制，选择每个 token 最优的专家（expert），并计算辅助损失 l_aux 和专家容量利用率 exp_capacity_rate。该函数还处理随机路由丢弃的 token，确保每个专家的负载均衡。

def top1gating(logits: Tensor, random_routing_dropped_token: bool = False):
    """Implements Top1Gating on logits."""
    # everything is in fp32 in this function
    logits = logits.float()
    gates = F.softmax(logits, dim=1)    # 计算门控概率
    capacity = gates.shape[0]

    # Create a mask for 1st's expert per token
    # noisy gating
    indices1_s = torch.argmax(gates, dim=1)     # 选择最佳专家
    num_experts = int(gates.shape[1])
    mask1 = F.one_hot(indices1_s, num_classes=num_experts)

    # gating decisions
    # exp_counts = torch.sum(mask1, dim=0).detach().to('cpu')
    exp_counts = torch.sum(mask1, dim=0).detach()

    # Compute l_aux     计算辅助损失
    me = torch.mean(gates, dim=0)
    ce = torch.mean(mask1.float(), dim=0)
    l_aux = torch.sum(me * ce) * num_experts
    mask1_rand = mask1

    top_idx = torch.topk(mask1_rand, k=capacity, dim=0)[1]

    new_mask1 = mask1 * torch.zeros_like(mask1).scatter_(0, top_idx, 1)
    mask1 = new_mask1
    mask1_bk = mask1
    if random_routing_dropped_token:        # 随机路由丢弃  处理丢弃的token，重新分配给未满的专家
        not_full = capacity - new_mask1.sum(dim=0)   # 计算每个专家剩余的容量
        sorted_notfull, indices_notfull = torch.sort(not_full, descending=True)    # 对未满的专家进行排序
        sorted_notfull = sorted_notfull.to(torch.int64)
        not_full_experts_ids = torch.repeat_interleave(indices_notfull, sorted_notfull)    # 重复未满专家的索引
        shuffle_not_full_ids = torch.randperm(not_full_experts_ids.shape[0])    # 随机打乱未满专家的索引
        # 重新分配被丢弃的token
        not_full_experts_ids = not_full_experts_ids[shuffle_not_full_ids]    #  计算 new_mask1 中每个token被分配到的专家索引
        indices1_s_after_drop = torch.argmax(new_mask1, dim=1)
        # get drop idx
        drop_mask = 1 - new_mask1.sum(dim=1)    # 标识出那些没有被分配给任何专家的token（即被丢弃的token）
        drop_mask = drop_mask.bool()
        drop_idx = drop_mask.nonzero().view(-1)    # 获取被丢弃token的索引
        drop_num = drop_mask.sum().to(torch.int64)    # 计算被丢弃token的数量
        indices1_s_after_drop.scatter_(0, drop_idx, not_full_experts_ids[:drop_num])    # 将随机选择的未满专家索引分配给被丢弃的token
        nodrop_mask1 = F.one_hot(indices1_s_after_drop, num_classes=num_experts)    #  将更新后的专家索引转换为one-hot编码形式。
        mask1 = nodrop_mask1    # 每个token都被分配给了至少一个专家。

    # Compute locations in capacity buffer      计算位置索引
    locations1 = torch.cumsum(mask1, dim=0) - 1

    # Store the capacity location for each token
    locations1_s = torch.sum(locations1 * mask1, dim=1)

    # Normalize gate probabilities  归一化门控概率
    mask1_float = mask1.float()
    gates = gates * mask1_float

    # 归一化门控概率
    locations1_sc = F.one_hot(locations1_s, num_classes=capacity).float()   # one hot to float
    combine_weights = torch.einsum("se,sc->sec", gates, locations1_sc)

    dispatch_mask = combine_weights.bool()

    exp_counts_capacity = torch.sum(mask1_bk)
    exp_capacity_rate = exp_counts_capacity / (logits.shape[0])
    return [l_aux, exp_capacity_rate], combine_weights, dispatch_mask, exp_counts

总的来说，top1gating 中的随机路由丢弃是一种处理策略，用于确保在多专家模型（MoE）中，没有被任何专家处理的token（即被丢弃的token）能够被重新分配给那些尚未达到容量上限的专家。这个过程可以总结为以下两句话：

1. 识别未满专家：首先，通过计算每个专家的容量与已分配token的数量差，找出那些尚未达到容量上限的专家。

2. 重新分配被丢弃token：然后，将那些未被分配给任何专家的token随机分配给上述未满的专家，以此来优化专家的负载均衡，并确保所有token都能得到处理。

其他

归一化 HunYuanRMSNorm

这里和T5模型的归一化层类似，通过方差和权重实现

计算公式如下：

1. 计算均方根（RMS）： 。其中，是输入向量x的第i个元素，d 是特征维度。

2. 归一化输入向量：

3. 应用缩放和偏移参数： 。其中，γ 是可学习的缩放参数（与输入维度相同），β 是可学习的偏移参数（与输入维度相同），⊙ 表示元素级的乘法操作。

class HunYuanRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        """
        HunYuanRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)      # 方差
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)   # 应用权重，实现归一化
        return self.weight * hidden_states.to(input_dtype)

HunYuanForSequenceClassification

这个类用于实现大模型进行分类或回归，主要过程包括：

1. 输入处理：接收多种类型的输入参数，如 input_ids、attention_mask 等。
2. 模型调用：调用内部的 model 方法进行前向传播，获取隐藏状态。
3. 分类/回归：根据配置和标签计算损失，支持回归、单标签分类和多标签分类。
4. 输出格式化：根据 return_dict 参数决定返回格式，可以是元组或字典

class HunYuanForSequenceClassification(HunYuanPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.model = HunYuanModel(config)
        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)

        ...

    @add_start_docstrings_to_model_forward(HUNYUAN_INPUTS_DOCSTRING)
    def forward(...) -> Union[Tuple, SequenceClassifierOutputWithPast]:
        ...

        transformer_outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = transformer_outputs[0]
        logits = self.score(hidden_states)      # 线性层分类/回归

        ...

        if self.config.pad_token_id is None:
            sequence_lengths = -1
        else:
            if input_ids is not None:
                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
                    logits.device
                )
            else:
                sequence_lengths = -1

        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]    # 提取每个序列的最后一个有效（非填充）元素的输出，以便用于序列级别的决策

        loss = None
        if labels is not None:      # 根据配置和标签计算损失，支持回归、单标签分类和多标签分类。
            labels = labels.to(logits.device)
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                loss_fct = MSELoss()
                if self.num_labels == 1:
                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
                else:
                    loss = loss_fct(pooled_logits, labels)
            elif self.config.problem_type == "single_label_classification":
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
            elif self.config.problem_type == "multi_label_classification":
                loss_fct = BCEWithLogitsLoss()
                loss = loss_fct(pooled_logits, labels)
        if not return_dict:     # 输出格式化
            output = (pooled_logits,) + transformer_outputs[1:]
            return ((loss,) + output) if loss is not None else output

        return SequenceClassifierOutputWithPast(
            loss=loss,
            logits=pooled_logits,
            past_key_values=transformer_outputs.past_key_values,
            hidden_states=transformer_outputs.hidden_states,
            attentions=transformer_outputs.attentions,
        )

其中的核心部分如下：

pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] 

在自回归模型中，如Transformer或其变体，每个元素（例如，单词或字符）通常是基于之前元素的预测结果来生成的。这意味着在序列中，除了最后一个元素之外，其他每个元素都是基于前面的上下文计算得到的，而最后一个元素则是模型基于整个序列的上下文生成的预测值。因此，这里这样做的目的是提取每个序列的最后一个有效（非填充）元素的输出，用于分类或回归。

这段代码的功能是：

1. 确定序列长度：首先，代码通过检查 input_ids 中的填充标记（pad_token_id）来确定每个序列的实际长度（即非填充部分的长度）。这是通过找到每个序列中填充标记第一次出现的位置并将其转换为序列长度来实现的。

2. 提取池化输出：然后，使用确定的序列长度，从 logits 中提取每个序列的最后一个非填充元素的值。这个值通常被视为整个序列的“池化”表示，可以用于序列级别的任务，如分类或回归。

4.总结

随着Hunyuan-Large的开源，我们不仅见证了腾讯在AI领域的又一重要里程碑，也看到了开源精神在推动科技进步中的力量。Hunyuan-Large的发布，不仅为研究人员和开发者提供了一个强大的工具，更为全球的AI社区提供了一个共同探索和创新的平台。通过开源，腾讯邀请全世界的研究人员共同参与到模型的优化和应用开发中，这无疑将加速AI技术的演进，拓宽其应用的边界。

在未来，我们期待看到Hunyuan-Large在各种实际应用中的潜力得到充分释放，无论是在提升用户体验、优化业务流程，还是在解决复杂的科学问题上，它都有可能发挥关键作用。同时，我们也期待开源社区能够围绕Hunyuan-Large产生更多创新的火花，共同推动人工智能向更加智能、高效、普惠的方向发展。

亲爱的读者，如果您对Hunyuan-Large模型的探索和开源精神感到兴奋，或者对AI技术的未来充满期待，那么请不要犹豫，给我们一个赞👍，让更多人看到这篇文章。您的每一次点赞都是对我们最大的鼓励和支持！同时，别忘了关注我们的账号，及时获取最新的AI技术动态和深度解析。我们承诺，将不断为您带来高质量的内容，让您在AI的海洋中乘风破浪。

最后，如果您觉得这篇文章有价值，不妨收藏起来，方便日后回顾和分享。您的每一次收藏，都是对知识的积累和对智慧的传承。

感谢您的陪伴和支持，让我们在AI的道路上一起前行，探索未知，创造可能！