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

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

背景

本文方案

什么是 SegMoE？

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

关于名称

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

推理

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

效果

免责声明及正在进行的工作

代码分析

SparseMoeBlock

# 受 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  # 返回最终的隐藏状态

SegMoEPipeline

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

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

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

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

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

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

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

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

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

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

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

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

算法实现思路

定义专家和对应的 positive_prompt/negative_prompt

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

将多个 expert 转换为一个 MoE 模型

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  # 返回包含所有层门控向量的字典

总结

总结