devbox
人工智能uu
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

无需训练的 Diffusion MoE 方案 (SegMoE): Segmind Mixture of Diffusion Experts_免训练 moe

作者：人工智能uu | 2024-08-21 19:21:09

赞

踩

免训练 moe

Paper name

SegMoE: Segmind Mixture of Diffusion Experts

Paper Reading Note

Blog URL: https://blog.segmind.com/introducing-segmoe-segmind-mixture-of-diffusion-experts/

Project URL: https://huggingface.co/blog/segmoe

Code URL: https://github.com/segmind/segmoe

Intro Video URL: https://www.youtube.com/watch?v=6Q4BJOcvwGE

TL;DR

2024 年 Segmind 研发的全球首个用于 Stable Diffusion 的开源专家混合（Mixture of Experts，MoEs）框架。这是一种能够将多个稳定扩散模型动态组合在一起的框架，无需训练即可在短时间内创建更大的 MoE 模型。

Introduction

背景

  • 专家混合是一类稀疏深度学习模型，包含由线性层和门控层组成的稀疏线性层。这些门控层会将输入路由到少数几个线性层，从而节省计算资源并减少推理时间。这种架构使得模型在保持推理时间较低的同时，能够拥有大量的参数。
  • 受 Mixtral 模型的启发，创建一个类似于 mergekit 的框架，以结合 Stable Diffusion 模型，这将显著增加模型的知识量，同时保持接近单个 Stable Diffusion 模型的推理时间。

本文方案

  • 我们的框架 SegMoE 将 Stable Diffusion 1.5 和 Stable Diffusion XL 模型结合成一个专家混合风格的模型，提升了模型的提示依从性和多样性。未来，我们计划支持更多模型，并支持 SegMoE 模型的训练，这可能进一步提高模型的质量，并为文本生成图像提供一个新的 SOTA（最先进）模型。
什么是 SegMoE？

SegMoE 模型遵循与 Stable Diffusion 相同的架构。与 Mixtral 8x7b 类似，SegMoE 模型包含多个模型。这种方式是通过用稀疏 MoE 层替换部分前馈层来实现的。MoE 层包含一个路由器网络，用于选择哪个专家模型最有效地处理哪些 token。你可以使用 segmoe 包来创建你自己的 MoE 模型！整个过程只需几分钟。

关于名称

SegMoE 模型的命名为 SegMoE-AxB，其中 A 代表 MoE 组合在一起的专家模型数量，B 代表参与每个图像生成的专家数量。根据配置设置，模型的某些层（前馈块、注意力层或全部）会被复制，而其余参数则与 Stable Diffusion 模型相同。

推理

我们在 Hub 上发布了三个合并模型：

  • SegMoE 2x1 有两个专家模型。
  • SegMoE 4x2 有四个专家模型。
  • SegMoE SD 4x2 有四个 Stable Diffusion 1.5 专家模型。
效果
  • SegMoE 4x2
    在这里插入图片描述
免责声明及正在进行的工作
  • 速度较慢：如果每个 token 的专家数量大于 1，MoE 将在多个专家模型之间执行计算。这使得它比单个 SD 1.5 或 SDXL 模型更慢。
  • 高 VRAM 使用：MoEs 推理速度很快，但仍需要大量的 VRAM（因此需要昂贵的 GPU）。这使得它们在本地设置中使用具有挑战性，但它们非常适合多 GPU 部署。作为参考，SegMoE-4x2 在半精度下需要 24GB 的 VRAM。
  • SegMoE 已全面集成到 Hugging Face 生态系统中，并得到 diffusers 的支持。

代码分析

  • 主要分析代码：https://github.com/segmind/segmoe/blob/main/segmoe/main.py
SparseMoeBlock
  • SparseMoeBlock 的实现与 Mixtral 基本一致，稍微做了一定简化
# 受 transformers.models.mixtral.modeling_mixtral.MixtralSparseMoeBlock 的启发
class SparseMoeBlock(nn.Module):
    def __init__(self, config, experts):
        super().__init__()
        self.hidden_dim = config["hidden_size"]  # 隐藏层的维度
        self.num_experts = config["num_local_experts"]  # 可用专家的数量
        self.top_k = config["num_experts_per_tok"]  # 每个 token 选择的专家数量
        self.out_dim = config.get("out_dim", self.hidden_dim)  # 输出的维度，如果未指定则默认为 hidden_dim

        # 门控机制，通过线性层生成每个 token 的专家选择概率
        self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)
        self.experts = nn.ModuleList([deepcopy(exp) for exp in experts])  # 复制专家层

    def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
        batch_size, sequence_length, f_map_sz = hidden_states.shape  # 获取批量大小、序列长度和特征维度
        hidden_states = hidden_states.view(-1, f_map_sz)  # 展平批量和序列维度
        # 通过门控机制计算路由 logits
        router_logits = self.gate(hidden_states)
        _, selected_experts = torch.topk(
            router_logits.sum(dim=0, keepdim=True), self.top_k, dim=1
        )  # 选择 top_k 专家
        routing_weights = F.softmax(
            router_logits[:, selected_experts[0]], dim=1, dtype=torch.float
        )  # 计算选择的专家的路由权重

        # 将路由权重转换回输入的 dtype
        routing_weights = routing_weights.to(hidden_states.dtype)

        final_hidden_states = torch.zeros(
            (batch_size * sequence_length, self.out_dim),
            dtype=hidden_states.dtype,
            device=hidden_states.device,
        )  # 初始化最终的隐藏状态

        # 遍历所有选中的专家，并在每个专家上执行计算
        for i, expert_idx in enumerate(selected_experts[0].tolist()):
            expert_layer = self.experts[expert_idx]

            # 计算当前专家的隐藏状态
            current_hidden_states = routing_weights[:, i].view(
                batch_size * sequence_length, -1
            ) * expert_layer(hidden_states)

            # 由于 `index_add_` 仅支持使用 torch 张量进行索引，因此我们使用 `top_x` 张量。
            final_hidden_states = final_hidden_states + current_hidden_states

        final_hidden_states = final_hidden_states.reshape(
            batch_size, sequence_length, self.out_dim
        )  # 恢复最终隐藏状态的形状
        return final_hidden_states  # 返回最终的隐藏状态

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38
  • 39
  • 40
  • 41
  • 42
  • 43
  • 44
  • 45
  • 46
  • 47
  • 48
  • 49
  • 50
SegMoEPipeline
  • 定义了一个名为 SegMoEPipeline 的类，用于处理和管理混合扩散专家模型。

以下是对 SegMoEPipeline 中涉及的各个函数的定义和功能的解释，算法方面重点函数是 self.get_gate_params 的实现：

  1. remove_all_forward_hooks: 这个函数用于删除模型中的所有前向钩子。它会遍历模型的所有子模块，如果子模块有前向钩子，就会将它们删除。

  2. SparseMoeBlock: 这个类定义了一个稀疏的Mixture of Experts（MoE）块。它包含一个门控网络和多个专家网络。在前向传播过程中，它会根据输入的隐藏状态选择合适的专家网络进行计算，并将结果组合起来。

  3. getActivation: 这个函数用于获取模型中的激活值。它接受一个模型、输入和输出作为参数，并返回一个钩子函数。这个钩子函数会在模型的前向传播过程中被调用，并将激活值保存到一个字典中。

  4. SegMoEPipeline: 这个类定义了SegMoEPipeline，它是一个用于动态组合多个Stable Diffusion模型的管道。它接受一个配置文件或路径作为参数，并根据配置文件中的设置加载和组合模型。

  5. load_from_scratch: 这个函数用于从头开始加载SegMoEPipeline。它会根据配置文件中的设置加载基础模型和专家模型，并将它们组合起来。

  6. call: 这个函数用于调用SegMoEPipeline进行推理。它会将输入传递给基础模型，并返回推理结果。

  7. create_empty: 这个函数用于创建一个空的UNet2DConditionModel模型。它会根据配置文件中的设置初始化模型，并将模型中的层替换为稀疏的MoE层。

  8. save_pretrained: 这个函数用于将SegMoEPipeline保存到磁盘上。它会将模型的权重和配置保存到指定的路径下。

  9. cast_hook: 这个函数用于在模型的前向传播过程中获取隐藏状态。它会为模型中的每个层注册一个前向钩子，并在前向传播过程中将隐藏状态保存到一个字典中。

  10. get_hidden_states: 这个函数用于获取模型中的隐藏状态。它接受一个模型、正向提示和负向提示作为参数，并返回一个包含隐藏状态的字典。

  11. get_gate_params: 这个函数用于获取门控网络的参数。它接受一个专家模型列表、正向提示列表和负向提示列表作为参数，并返回一个包含门控网络参数的字典。

算法实现思路

定义专家和对应的 positive_prompt/negative_prompt
  • 定义 experts 的时候会给定该 expert 的 positive_prompt 和 negative_prompt，类似下面。一个完整示例见：https://github.com/segmind/segmoe/blob/main/segmoe_config_4x2.yaml
    • 这些 positive_prompt 和 negative_prompt 会用来计算 MoE router 的权重
base_model: Base Model Path, Model Card or CivitAI Download Link
num_experts: Number of experts to use
moe_layers: Type of Layers to Mix (can be "ff", "attn" or "all"). Defaults to "attn"
num_experts_per_tok: Number of Experts to use 
type: Type of the individual models (can be "sd" or "sdxl"). Defaults to "sdxl"
experts:
  - source_model: Expert 1 Path, Model Card or CivitAI Download Link
    positive_prompt: Positive Prompt for computing gate weights
    negative_prompt: Negative Prompt for computing gate weights
  - source_model: Expert 2 Path, Model Card or CivitAI Download Link
    positive_prompt: Positive Prompt for computing gate weights
    negative_prompt: Negative Prompt for computing gate weights
  - source_model: Expert 3 Path, Model Card or CivitAI Download Link
    positive_prompt: Positive Prompt for computing gate weights
    negative_prompt: Negative Prompt for computing gate weights
  - source_model: Expert 4 Path, Model Card or CivitAI Download Link
    positive_prompt: Positive Prompt for computing gate weights
    negative_prompt: Negative Prompt for computing gate weights

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
将多个 expert 转换为一个 MoE 模型
  • 把多个模型的 ffn 和 attn 中的 to_q/to_k/to_v 都会替换为 SparseMoeBlock，也即模型每一层都能自动选择不同的专家模型的参数
MoE router 参数初始化
  • 因为 router 是没有经过训练的，需要通过给定的模型 postive prompt 和 negative prompt 进行 MoE router 参数初始化。算法逻辑主要是从多个专家模型中提取隐藏状态，并对这些状态进行归一化处理和汇总。这种思路其实就是根据不同专家模型输出的 hidden_states 中参数整合后作为 MoE router 参数，这样如果是与 hidden_states 更相关的输入可以达到更高的 MoE router 激活值（其实是有很强的先验假设在这里）。
@torch.no_grad
def get_gate_params(
    self,
    experts,  # 专家模型列表
    positive, # 正样本列表，与专家模型一一对应
    negative, # 负样本列表，与专家模型一一对应
):
    gate_vects = {}  # 用于存储每一层的门控向量
    for i, expert in enumerate(tqdm.tqdm(experts, desc="Expert Prompts")):
        expert.to(self.device)  # 将专家模型移动到指定设备
        expert.unet.to(
            device=self.device,
            dtype=self.torch_dtype,
            memory_format=torch.channels_last,
        )  # 将专家模型的unet部分移动到指定设备，并设置数据类型和内存格式
        hidden_states = self.get_hidden_states(expert, positive[i], negative[i])  
        # 获取专家模型的隐藏状态
        del expert  # 删除专家模型以释放显存
        gc.collect()  # 手动进行垃圾回收
        torch.cuda.empty_cache()  # 清空CUDA缓存
        for h in hidden_states:
            if i == 0:
                gate_vects[h] = []  # 初始化门控向量列表
            hidden_states[h] /= (
                hidden_states[h].norm(p=2, dim=-1, keepdim=True).clamp(min=1e-8)
            )  # 对隐藏状态进行L2归一化
            gate_vects[h].append(hidden_states[h])  # 将归一化后的隐藏状态添加到门控向量列表中
    for h in hidden_states:
        gate_vects[h] = torch.stack(
            gate_vects[h], dim=0
        )  # 将门控向量堆叠成张量，形状为 (num_expert, num_layer, hidden_size)
        gate_vects[h].permute(1, 0)  # 对张量进行维度置换，形状为 (num_layer, num_expert)

    return gate_vects  # 返回包含所有层门控向量的字典

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34

总结

  • 无需训练的多模型 MoE 方案，代码完全开源。
    • MoE router 这里虽然是无需训练的，但是这里的初始化方法或许在从头训练的 MoE router 参数初始化上也能用上。
  • 这种无需训练的 MoE 方案目前 follow 的工作不多，可能还是有一定先天劣势，比如模型的生成质量上其实没有明显比单专家模型会更好，同时会占用更多的显存，速度也会变慢。
声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/人工智能uu/article/detail/1013027?site
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号