灵活组网

灵活组网#

LoongForge 框架的核心特性之一是原生支持灵活组网。 通过简单的 YAML 配置文件，用户可以自由组合视觉/音频编码器、模态对齐投影层和语言模型基座，快速构建定制化的多模态大模型。 只需编辑 ../configs/models 下的模型配置 YAML 文件，即可在不修改任何代码的情况下重组或切换整个模型结构。

例如，不同的视觉编码器（Qwen2.5-ViT、Qwen3-ViT、LLaVA-OV-1.5-ViT、InternViT 等）可以与不同的 LLM 基座（LLaMA 系列、DeepSeek 系列、Qwen 系列 等）拼接。 这种配置驱动组网极大地降低了探索和定制多模态架构的成本，实现零代码模型构建。

1. 核心组件#

OmniEncoderModel#

所有编码器组件的抽象基类。 它将图像、视频、音频或其他模态数据转换为 LLM 可理解的嵌入表示。

  • 抽象化：封装视觉/音频编码器和投影层，并统一管理文本嵌入。

  • 关键实现：

    • 多模态兼容性：通过简单插入新分支即可添加新模态。

# loongforge/models/omni_models/omni_encoder_model.py
class OmniEncoderModel(torch.nn.Module):
    def __init__(self, config, ...):
        # 文本模态
        self.text_encoder = LanguageModelEmbedding(...)

        # 图像模态
        if hasattr(config, "image_encoder"):
            self.image_encoder = AutoModel.from_config(config.image_encoder, ...)

        # 视频模态（可轻松扩展）
        if hasattr(config, "video_encoder"):
            self.video_encoder = AutoModel.from_config(config.video_encoder, ...)

  • 异构张量并行：使用 hook 机制在异构设备间实现编码器张量并行。

# 自动注册 hook，在前向前后切换并行状态
self.image_encoder.register_forward_pre_hook(
    make_encoder_forward_pre_hook("image_encoder")
)
self.image_encoder.register_forward_hook(
    make_encoder_forward_hook("text_decoder")  # 切回解码器状态
)

OmniCombinationModel#

多模态组合的核心组件。 定义数据何时以及如何在模态组件之间流动；不包含实际计算逻辑。

  • 逻辑解耦：通过外部配置动态决定是否加载编码器或基座模型，实现组件级解耦。

# loongforge/models/omni_models/omni_combination_model.py
class OmniCombinationModel(BaseMegatronModule):
    def __init__(self, config, ...):
        # 1. 动态初始化编码器模型
        if config.image_encoder is not None and add_encoder:
            self.encoder_model = OmniEncoderModel(config, ...)

        # 2. 动态初始化 LLM
        if config.foundation is not None and add_decoder:
            self.foundation_model = AutoModel.from_config(config.foundation, ...)

OmniModelProvider#

解析全局参数，处理分布式初始化（如流水线并行拆分），最终实例化 OmniCombinationModel。

关键能力：

  • 分布式流水线并行：检测当前流水线并行阶段并按需加载编码器/基座模型，实现跨 GPU 流水线部署。

# loongforge/models/omni_models/omni_model_provider.py
def omni_model_provider(...):
    # 自动检测当前 rank 是否为第一个 PP 阶段；决定是否加载编码器
    # 这对于将编码器和解码器放置在不同 GPU 上至关重要
    if args.encoder_pipeline_model_parallel_size in [0, None]:
        add_encoder = mpu.is_pipeline_first_stage()

    # 使用环境感知标志构建模型
    return OmniCombinationModel(..., add_encoder=add_encoder, ...)

  • 参数桥接与适配：获取全局训练参数和模型配置，将训练关键超参数（language_vocab_size、language_max_sequence_length 等）注入模型初始化，保证模型配置与训练环境的一致性。

args = get_args()
model_config = get_model_config()

model = OmniCombinationModel(
    model_config,
    language_vocab_size=args.padded_vocab_size,
    language_max_sequence_length=args.max_position_embeddings,
    fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
    ...
)

2. 使用方法#

我们采用基于 Hydra 的配置系统，通过 defaults 列表组合组件。 以 configs/models/internvl2.5/internvl2_5_8b.yaml 为例：

defaults:
  # 1. 选择视觉编码器（ViT）
  - ../../models/image_encoder@model.image_encoder: intern_vit_0.3b

  # 2. 选择图像投影层
  - ../../models/image_projector@model.image_projector: intern_mlp_adapter

  # 3. 选择 LLM 基座
  - ../../models/internlm2.5@model.foundation: internlm2_5_8b

  - _self_

model:
  model_type: intern_vl
  # ... 其他全局参数

假设你想要一个结合 InternViT 与 Qwen2.5-7B 的全新 VLM：

  1. 定义组件配置：确保 configs/models/ 包含 Qwen2.5 配置，且 configs/models/image_encoder/ 包含 InternViT 配置。

  2. 创建组合配置：创建一个新的 YAML，只需在 defaults 中列出所需组件。

defaults:
  # 1. 视觉编码器
  # 'intern_vit_0.3b' 指向 configs/models/image_encoder/intern_vit_0.3b.yaml
  - ../../models/image_encoder@model.image_encoder: intern_vit_0.3b

  # 2. 图像投影层
  - ../../models/image_projector@model.image_projector: intern_mlp_adapter

  # 3. LLM 基座
  # 'internlm2_5_8b' 指向 configs/models/internlm2.5/internlm2_5_8b.yaml
  - ../../models/internlm2.5@model.foundation: internlm2_5_8b

  - _self_

model:
  # 多模态模型类型
  model_type: intern_vl

  # 基座模型的损失函数
  loss_func: ${loss_func:loss_func_internvl}

  # 显式设置投影层输出维度以匹配 LLM hidden_size
  image_projector:
    hidden_size: 4096
    ffn_hidden_size: 4096

  # LLM 特定设置
  foundation:
    rotary_emb_func: "DynamicRotaryEmbedding"
    rotary_base: 1000000

    # Megatron 层规范，支持 Transformer-Engine 加速
    model_spec: ["loongforge.models.foundation.internlm.internlm_layer_spec",
                 "get_internlm_layer_with_te_spec"]

当前的 configs/models 目录已体现这种即插即用的组件库：

  • image_encoder/：视觉编码器配置

  • image_projector/：投影层配置

  • llama3/、qwen2/、deepseek2/ 等：LLM 基座配置

通过这种方式，构建多模态大模型就像拼搭 LEGO 积木一样简单。