跳到正文
Hugging Face Blog·· 20 小时前精选AI 评分81

Meta 发布开源权重多模态 MoE 模型 Muse Spark Open,约 816B 参数、42B 激活

Welcome Muse Spark Open: open-weight frontier multimodal mixture-of-experts model from Meta

AI 导读

Meta 发布开源权重 MoE 多模态模型 Muse Spark Open,总参数约 816B、每 token 激活约 42B,支持文本、图像、音频和视频输入,上下文窗口最高 1M token,面向 agentic 编码训练。

推荐理由

官方博客给出模型规模、架构细节和从 llama.cpp 到 TRL 微调的完整部署路径,读者可据此评估落地成本。

正文 · AI 翻译

Muse Spark Open 是 Meta 推出的开放权重模型。这款混合专家模型将文本、图像、音频和视频整合到一个约 8160 亿参数的模型中,每个 token 约 420 亿激活参数,上下文窗口高达一百万 token。它生成文本,为开发者提供一个统一的接口来询问不同类型的输入。

继 Muse Glimmer 之后,Muse Spark Open 在更大规模上引入了混合专家架构,是迄今为止最强大的美国开放权重模型。

Meta 训练 Muse Spark Open 用于智能体编程:阅读代码库、编写和调试代码、调用工具。它还可以接收图像、音频和视频作为输入,因此同一个模型在处理你的代码时,可以读取屏幕截图、听取语音备忘录或观看屏幕录制。本文将展示如何运行它、如何微调它,以及我们用它构建了什么。所有模型请见此处。

整体能力与架构

Meta 将 Muse Spark Open 描述为一个针对编程优化的模型,其训练基于真实的开发者工作流,例如跨多个文件的重构和长时间的调试会话。Muse Spark Open 是一个以智能体为中心的模型,能够在长任务中进行规划、行动、观察和自我纠正,包括结合图像浏览网页。

特性 Muse Spark Open
模型规模 总参数约 8160 亿;每个 token 约 420 亿激活参数
专家路由 每个 MoE 层 256 个专家;每个 token 选择 8 个
输入 文本、图像、音频和视频
输出 文本
上下文窗口 高达 100 万 token
检查点格式 BF16 和 NVFP4

检查点仓库为 Muse-Spark-Open-816B-A42B 和 Muse-Spark-Open-816B-A42B-NVFP4。

model diagram

Muse Spark Open 是 Muse 系列模型的全尺寸模型。Meta 于 8 月发布的 Muse Glimmer 是同一系列的小型稠密模型,适合在本地运行。两个模型使用相同的分词器和词表。聊天模板遵循 Muse Glimmer 模板,并为新增的模态做了补充。在 Muse Spark Open 中,输入嵌入和输出层是两个相互独立的矩阵,每个约 16.6 亿参数。

混合专家

该语言模型有 62 层,隐藏维度为 8192。每层有 128 个查询头和 16 个键值头(分组查询注意力)。每个头有 64 个维度。注意力块包含一个输出门和一个输出归一化层。模型还对查询和键应用归一化(QK-norm)。每个块通过一个可学习的残差门将其输入与输出结合,而不是简单的相加。

Muse Glimmer 在每层使用一个稠密前馈网络。Muse Spark Open 用混合专家(MoE)块替换了这个稠密网络。每个 MoE 块有 256 个专家,路由器为每个 token 选择 8 个专家。每个专家是一个前馈网络,先将 8192 维投影到 4096 维再投影回 8192 维,然后对其输出应用自己的归一化。这些专家持有全部模型参数的约 97%。

每个 token 仅运行 8 个专家,因此模型计算每个 token 大约使用 42B 参数。全部 256 个专家都保留在内存中,因此部署需要 816B 参数模型的内存,但能以 42B 模型的快得多的速度运行。BF16 检查点为 safetensors 格式,大小约 1.5 TB。NVFP4 检查点以 4 位精度存储参数,精度损失可忽略不计,专为 NVIDIA Blackwell GPU 设计。在 transformers 中,BF16 检查点通过张量并行(TP=32)跨 4 个节点(每个节点 8 块 H200 GPU)运行,NVFP4 则在单个 B200 节点上运行。如果你想部署 Muse Spark Open,请查看下方的 llama.cpp、vLLM 和 SGLang 方案。

感知编码器

感知编码器将图像和视频帧转换为供语言模型使用的 token,与 Muse Glimmer 类似。该编码器拥有 1.8B 参数、50 层、隐藏维度为 1536、16 个注意力头。它将每张图像分割为 14×14 的图像块,并生成 4096 维的特征。位置信息来自可学习的绝对嵌入和 2D RoPE。每层的前馈网络使用 LayerNorm 和 GELU。每第四层和最后一层使用全注意力,其余层使用稀疏注意力。

image pipeline diagram

Muse Glimmer 30B 使用相同架构的编码器。该设计遵循 Meta 的 Perception Encoder 工作。视频输入通过与图像相同的编码器。

音频编码器

audio diagram

Muse Spark Open 通过其音频编码器接受音频输入。音频分词器将音频转换为 token,聊天模板用 <|audio|> token 标记音频位置。编码器每 80 ms 输出一个音频 token,即每秒 12.5 个 token。也就是说,如果你有一段 10 秒的音频,输入音频 token 的数量将是 125。 编码器本身是一个卷积金字塔,其中每个阶段由一个不同块大小的块嵌入器(patch embedder)和 2 个 1d 卷积块(卷积核大小为 7、步长为 1)组成。这些阶段按 [16, 6, 5, 4] 的块大小重复。在最后一个阶段结束时,时间帧被送入一个最终的块:RMSNorm → 线性投影 → RMSNorm。

长上下文

推理时上下文窗口为 1M token。配置将 max_position_embeddings 设为 1,048,576。图像、视频和音频 token 与文本共用同一个窗口。

分组查询注意力使这个规模的模型的 KV cache 保持较小。每个 token 会向缓存添加 2 × 16 头 × 64 维 × 62 层 = 126,976 个值。以 BF16 计算,每个 token 约为 250 KB,完整的 1M token 上下文约为 266 GB。

推理支持

Transformers

Muse Spark Open 的实现提供了用于从多模态输入生成文本的 MuseSparkForConditionalGeneration 和用于准备对话的 MuseSparkProcessor。处理器结合了分词器、图像和视频处理器、音频特征提取器以及聊天模板。

使用 DistributedConfig 来配置并行设置。tp_size 是密集层(dense layers)分片所跨的设备数量,ep_size 是专家分布所跨的设备数量。

你可以按如下方式初始化模型和处理器。

import os

from transformers import AutoModelForMultimodalLM, AutoTokenizer, DistributedConfig

distributed_config = DistributedConfig(
    tp_size=int(os.environ["WORLD_SIZE"]),
    ep_size=int(os.environ["WORLD_SIZE"]),
)

tokenizer = AutoTokenizer.from_pretrained("meta-models/Muse-Spark-Open-816B-A42B
")
model = AutoModelForMultimodalLM.from_pretrained(
    "meta-models/Muse-Spark-Open-816B-A42B",
    distributed_config=distributed_config,
)

messages = [
    {"role": "user", "content": "Explain in one sentence why the sky is usually blue."},
]
inputs = tokenizer.apply_chat_template(
    messages,
    reasoning_effort="medium",  # default: xhigh, options: low/medium/high/xhigh
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
).to(model.device).to(model.dtype)

input_len = inputs["input_ids"].shape[-1]

outputs = model.generate(**inputs, max_new_tokens=2000)

纯文本推理

准备好对话后,将消息传给 processor.apply_chat_template,并启用 add_generation_prompt=True、tokenize=True、return_dict=True 和 return_tensors="pt"。将处理器的输出传给 model.generate,并解码新生成的 token(不包含输入前缀)。处理器有一个推理力度调节项,将符号级别(low/medium/high/xhigh)映射为模型本身接受的整数(16、64、256、512)。请注意,Muse Spark Open 无法关闭推理,因此如果想以更少的 token 解决较简单的任务,请将模型推理保持在 “low”。你也可以在系统提示中加入类似 Reasoning strength: 1 的行来设置数值化的力度。

messages = [
    {"role": "system", "content": "You should only answer with a number."},
    {"role": "user", "content": "What is 17 * 23?"},
]

inputs = processor.apply_chat_template(
    messages,
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
    reasoning_effort="high", # options:low/medium/high/xhigh
).to(model.device).to(model.dtype)

output = model.generate(**inputs, max_new_tokens=2000)
generated_tokens = output[0][inputs["input_ids"].shape[1] :]
print(processor.decode(generated_tokens, skip_special_tokens=False))

图像与文本

Muse Spark Open 处理器会将每个图像标记展开为处理后的图像所需的补丁 token 数量,并由图像边界 token 包裹。这种展开使文本位置与传给模型的视觉特征保持对齐。该处理器还能处理单个提示中的多张图像。

image_url = (
    "https://huggingface.co/datasets/merve/vl-test-suite/"
    "resolve/main/pills.jpg"
)
messages = [
    {
        "role": "user",
        "content": [
            {
                "type": "image",
                "image": image_url,
            },
            {
                "type": "text",
                "text": "Do any of the components in this supplement interact?",
            },
        ],
    },
]

inputs = processor.apply_chat_template(
    messages,
    tokenize=True,
    add_generation_prompt=True,
    reasoning_effort="medium",
    return_dict=True,
    return_tensors="pt",
).to(model.device).to(model.dtype)
input_len = inputs["input_ids"].shape[-1]

outputs = model.generate(**inputs, max_new_tokens=2000)
response = processor.decode(outputs[0][input_len:], skip_special_tokens=False)

processor.parse_response(response)

视频推理

视频提示包含帧时间信息以及视觉内容。该处理器可以处理带有音频的视频(得益于音频编码器,这是相比 Muse Glimmer 的升级)。处理器会对视频帧和音频进行下采样以对齐各帧。默认情况下,视频以 2 FPS 采样,上限为 96 帧。你可以通过设置视频处理器的 fps 来修改此设置:

processor.video_processor.fps = 4.0

一般经验法则是:视频是按时间顺序排列的帧,下采样越多,在特定事件之间丢失的帧就越多,因此如果想要更准确的生成结果,就需要保持较高的 fps。不过这里存在权衡,因为你会需要更多内存,并增加预填充的延迟。帧与标注时间戳的文本输入交错排列(例如 “Time: 3.5s”)。要启用传入带音频的视频,请将 load_audio_from_video=True 传给 apply_chat_template。

在使用 return_metadata=True 准备输入时保留视频的元数据,以便处理器能分配正确的帧率。

import torch

messages = [
      {
          "role": "user",
          "content": [
{
                  "type": "text",
                  "text": (
                      "Describe what happens in this video. "
                      "Use timestamps and include relevant speech or sounds."
                  ),
              }
              {"type": "video", "url": "https://huggingface.co/datasets/merve/vlm_test_images/resolve/main/ramen.mp4"},
              ,
          ],
      }
  ]

  inputs = processor.apply_chat_template(
      messages,
      add_generation_prompt=True,
      tokenize=True,
      return_dict=True,
      return_tensors="pt",
      load_audio_from_video=True,
      return_metadata=True,
  ).to(model.device).to(model.dtype)

  with torch.inference_mode():
      output_ids = model.generate(
          **inputs,
          max_new_tokens=256,
          do_sample=False,
      )

  generated_ids = output_ids[:, inputs["input_ids"].shape[1] :]
  response = processor.batch_decode(
      generated_ids,
      skip_special_tokens=True,
  )[0]

  print(response)

音频推理

Muse Spark Open 的音频塔接收原始波形。处理器期望 24 kHz 的单声道音频,并将波形填充到 1,920 个样本的整数倍,相当于每 80 毫秒一个音频嵌入。以其他采样率直接提供给处理器的音频必须先进行重采样。你可以使用 transformers 的 load_audio 工具来完成。load_audio 使用 torchcodec、librosa 或 torchaudio 后端,因此请确保在运行推理之前已在本地安装其中之一。

不过,如果你只是想将 apply_chat_template 与消息一起使用,这一切都会在底层自动处理:

import torch
  from transformers import AutoProcessor, AutoModelForMultimodalLM


  model_id = "meta-models/Muse-Spark-Open-816B-A42B"
"

  processor = AutoProcessor.from_pretrained(model_id)

  model = AutoModelForMultimodalLM.from_pretrained(
      model_id,
        ).eval()

  messages = [
      {
          "role": "user",
          "content": [
              {
                  "type": "audio",
                  "url": "https://huggingface.co/buckets/huggingface/audio-samples/resolve/bbaudio_question_39.mp3",
              },
              {
                  "type": "text",
                  "text": "Answer the question.",
              },
          ],
      }
  ]

  # Downloads the audio, decodes it as mono, resamples it to 24 kHz, pads it, creates input_values_mask, and expands the audio placeholders.
  inputs = processor.apply_chat_template(
      messages,
      add_generation_prompt=True,
      tokenize=True,
      return_dict=True,
      return_tensors="pt",
      load_audio_backend="auto",
  ).to(model.device).to(model.dtype)

  with torch.inference_mode():
      output_ids = model.generate(
          **inputs,
          max_new_tokens=128,
          do_sample=False,
      )

  # Decode only newly generated tokens
  prompt_length = inputs["input_ids"].shape[1]
  generated_text = processor.batch_decode(
      output_ids[:, prompt_length:],
      skip_special_tokens=True,
  )

  print(generated_text[0])

处理器会在音频边界标记之间插入匹配数量的音频 token。它还会返回一个标识真实样本的掩码,使模型能够排除被填充的音频嵌入。

工具调用

Muse Spark Open 可以进行纯文本和多模态工具调用。你可以按如下方式让模型调用工具。


from transformers import AutoProcessor, AutoModelForMultimodalLM

MODEL_ID = "meta-models/Muse-Spark-Open-816B-A42B"

# Load model
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForMultimodalLM.from_pretrained(
    MODEL_ID,
    dtype="auto",
    device_map="auto"
)

tools = [
    {
        "type": "function",
        "function": {
            "name": "weather.get",
            "description": "Get the current weather for a city.",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string"},
                },
                "required": ["city"],
            },
        },
    }
]


messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "image": "https://huggingface.co/datasets/merve/vl-test-suite/resolve/main/SF.png"},
            {"type": "text", "text": "I'm going to the city in this picture. What clothes should I wear?"},
        ],
    },
]

inputs = processor.apply_chat_template(
    messages,
    tools=tools,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
    add_generation_prompt=True,
    reasoning_strength="high"
).to(model.device).to(model.dtype)

input_len = inputs["input_ids"].shape[-1]
outputs = model.generate(**inputs)
response = processor.parse_response(
      completion,
      prefix=inputs["input_ids"][0],
      tools=tools,
  )
print(response)

Muse Spark Open 使用基于频道的聊天模板(Harmony,由 gpt-oss 首次引入)。请注意,你不必了解其内部机制,但检查模型路由到的频道会很有用,以防得到低质量的输出。在下面的输入消息中:

<|begin_of_text|>
  <|start|>system<|message|>
  You are a helpful AI assistant.
  Knowledge cutoff: 2026-01-04.
  Current date: ...
  ...
  <|eot|>
  <|start|>user<|message|>...
  <|eot|>
  <|start|>assistant

<|start|>assistant 被有意留空,由模型通过 to= 选择接收者。模型有不同的接收者,例如 self 让模型进行推理,

to=self<|message|>the model reasons here<|eom|> # reasoning
<|start|>assistant to=user<|message|>final answer<|eot|> # final answer

self 是推理内容,commentary 是工具调用之前及之间对最终用户可见的评论,user 是最终回复。
工具调用的解析方式如下(假设你已将 tools 传递给 apply_chat_template:

<|start|>assistant to=weather.get<|message|>
  <atem:function_calls>
  <atem:invoke name="weather.get">
  <atem:parameter name="city">San Francisco</atem:parameter>
  </atem:invoke>
  </atem:function_calls>
  <|eot|>

Llama.cpp

Muse Spark 2-bit 运行在 Llama Server 上

Muse Spark Open 已集成到 llama.cpp 中,并包含在预构建的二进制文件里。你可以用以下命令安装:

curl -LsSf https://llama.app/install.sh | sh

你可以按如下方式启动服务器,并在模型 ID 后附加量化后缀:


llama serve -hf ggml-org/Muse-Spark-Open-816B-A42B-GGUF:Q2_K

包含量化版本的仓库在这里。服务器启动后,你可以访问 localhost:9931,通过内置的 WebUI 进行对话。在 WebUI 中,你可以添加音频和图片进行推理,并调整推理强度。
注意,Muse Spark Open 无法关闭推理,因此如果你希望在较简单的任务上消耗更少的 token,请将模型推理强度保持为“low”。
我们在 Mac M3 Ultra(512GB)上测试了量化版本。经 imatrix 校准的 2-bit 量化运行速度为 18 tokens/s,KV 缓存仍有充足空间(模型占用 304 GB)。

你也可以按如下方式查询服务器。欢迎查看API 参考,了解多模态、工具调用等更多内容。


## text-only inference
curl http://localhost:8080/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
        "messages": [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "Write a limerick about python exceptions"}
        ]
    }'

## image-text inference
curl -s http://127.0.0.1:8080/v1/chat/completions \
  -H "Content-Type: application/json" -H "Authorization: Bearer <your-key>" \
  -d '{"model":"muse-spark-open","messages":[{"role":"user","content":[
        {"type":"text","text":"What is in this image?"},
        {"type":"image_url","image_url":{"url":"data:image/png;base64,..."}}
      ]}],"temperature":0}'

## audio-text inference
curl -s http://127.0.0.1:8080/v1/chat/completions \
  -H "Content-Type: application/json" -H "Authorization: Bearer <your-key>" \
  -d '{"model":"muse-spark-open","messages":[{"role":"user","content":[
        {"type":"text","text":"Transcribe this audio."},
        {"type":"input_audio","input_audio":{"data":"data:audio/wav;base64,..."}}
      ]}],"temperature":0}'

llama.cpp 团队提供了多种不同的量化版本。我们在 Apple Silicon 芯片上测试了 Q2_K、Q4_K_S 和 Q_8。你可以通过在模型 ID 后添加量化 ID 后缀来运行每种量化版本,llama serve -hf ggml-org/Muse-Spark-Open-816B-A42B:Q2_K。

SGLang

Muse Spark Open 的 SGLang 集成支持文本、图像、音频和视频输入。NVFP4 参考配置可跨四块 GB300 GPU 运行。

在 SGLang 集成环境中,从 Hugging Face Hub 下载检查点并启动服务器:

export MODEL_ID=meta-models/Muse-Spark-Open-816B-A42B-NVFP4
export CKPT="$(hf download "$MODEL_ID")"

SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION=1 \
uv run sglang_meta serve \
  --model-path "$CKPT" \
  --served-model-name "$MODEL_ID" \
  --quantization modelopt_fp4 \
  --moe-runner-backend flashinfer_trtllm \
  --tp-size 4 \
  --ep-size 4 \
  --chat-template "$CKPT/muse_spark_1_2_oss.jinja" \
  --port 30000

该检查点包含聊天模板,我们会显式传入。启动器会配置 Muse 的推理和工具调用解析器。

通过 OpenAI 兼容 API 发送请求:

curl http://localhost:30000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "meta-models/Muse-Spark-Open-816B-A42B-NVFP4",
    "messages": [{
      "role": "user",
      "content": "Explain how a mixture-of-experts model works."
    }],
    "max_tokens": 2048
  }'

响应将 reasoning_content 中的推理轨迹与 content 中的最终答案分离开来。

vLLM

Muse Spark Open 的 vLLM 参考集成支持文本和图像输入,所用 NVFP4 检查点需要 Blackwell 系列 GPU。

安装参考插件及其依赖后,从 Hugging Face Hub 下载检查点并启动服务器:

export MODEL_ID=meta-models/Muse-Spark-Open-816B-A42B-NVFP4
export CKPT="$(hf download "$MODEL_ID")"

VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 \
VLLM_PLUGINS=musespark1x_models,ptq_nvfp4 \
vllm serve "$CKPT" \
  --served-model-name "$MODEL_ID" \
  --tokenizer-mode musespark1x \
  --config-format auto \
  --load-format internal_safetensor \
  --tensor-parallel-size 4 \
  --enable-expert-parallel \
  --max-model-len 131072 \
  --gpu-memory-utilization 0.92 \
  --reasoning-parser musespark1x \
  --tool-call-parser musespark1x \
  --enable-auto-tool-choice \
  --port 8000

此配置设置了 128K token 的服务上下文。该插件提供了模型实现、检查点加载器,以及推理和工具调用解析器。

使用相同的 Hub 模型 ID 发送请求:

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "meta-models/Muse-Spark-Open-816B-A42B-NVFP4",
    "messages": [{
      "role": "user",
      "content": "Explain how a mixture-of-experts model works."
    }],
    "max_tokens": 2048
  }'

MLX

Muse Spark Open 现已在 oMLX 和 mlx-lm 中可用。针对 oMLX 优化的 oQ 量化版本有 3-bit 和 4-bit,可直接使用。

我们在 M3 Ultra(512 GB)上测试了 oQ3e 量化版本。该模型占用约 344 GB,为 KV 缓存留出了空间,在 1024-token 提示词和 128 个生成 token 的条件下,预填充速度约为 225 tok/s,生成速度约为 21 tok/s。

通过 oMLX 界面下载模型后,你可以在内置聊天 UI 中与它对话。在聊天 UI 中,你可以添加图片,并实时调整推理强度(从 low 到 xhigh)。模型的推理内容显示为一个可折叠的思考块,与最终答案分离开来。

注意,推理无法关闭,默认(Auto)设置使用 xhigh。如果你希望在较简单的任务上使用更少的 token,请将其设为“low”。 omlx

该语言模型也在官方 mlx-lm 仓库的这个 PR 中可用。你可以用相同的优质 oQ3e 量化版本这样运行它:

mlx_lm.generate --model mlx-community/Muse-Spark-Open-816B-A42B-oQ3e \
  --prompt "Explain convolutional networks like I am 5" \
  --max-tokens 5000

这是一个很大的模型,如上所述,即使是 3-bit 量化也需要超过 256 GB 的内存。你也可以在多台计算机上运行分布式推理,详情请参阅文档。

通过 Inference Endpoints 部署 Muse Spark Open

你可以在 Inference Endpoints 上部署 Muse Spark Open,NVFP4 变体提供经过预验证的 SGLang 配置,可运行在 8 × NVIDIA RTX Pro 6000 98 GB 上(131072 上下文长度),或者如果你想要 1048576 的上下文长度,可以使用 8 x NVIDIA H200 141GB。

要部署,只需使用 Hugging Face CLI 运行一条命令:

hf endpoints catalog deploy --repo meta-models/Muse-Spark-Open-816B-A42B-NVFP4 --name muse-spark

或者打开 endpoints.huggingface.co/new/meta-models/Muse-Spark-Open-816B-A42B-NVFP4 并选择 Create Endpoint。

该端点提供 OpenAI 兼容的 API,同时支持位于 /v1/chat/completions 的 Chat Completions API 和位于 /v1/responses 的 Responses API。

curl "https://<your-endpoint>.endpoints.huggingface.cloud/v1/chat/completions" \
  -H "Authorization: Bearer $HF_TOKEN" \
  -H "Content-Type: application/json" \
  -d '{
    "model": "meta-models/Muse-Spark-Open-816B-A42B-NVFP4",
    "messages": [{
      "role": "user",
      "content": [
        {"type": "image_url", "image_url": {"url": "https://huggingface.co/datasets/merve/vl-test-suite/resolve/main/pills.jpg"}},
        {"type": "text", "text": "Do any of the components in this supplement interact?"}
      ]
    }],
    "chat_template_kwargs": {"enable_thinking": false}
  }'

任何 OpenAI 兼容的客户端都可以使用相同的 URL,并以你的 Hugging Face token 作为 API key。

使用 TRL 进行微调 Quentin Gallouédec Ben Burtenshaw

你可以使用 TRL SFTTrainer 微调整个 816B 检查点。本节展示了我们如何在 8 个节点(每个节点 8 块 H100 GPU)上用 TRL 微调整个 816B 检查点。该方案在注意力层上训练 LoRA 适配器,专家保持冻结。完整脚本位于 TRL examples 中。

首先,我们需要把模型切分到 64 个 GPU 设备上。鉴于模型规模很大,我们需要在多个维度上进行并行。我们使用 DistributedConfig 对张量、数据和专家进行并行化,它告诉 transformers 在加载模型时如何切分:

from transformers.distributed import DistributedConfig


distributed_config = DistributedConfig(fsdp_size=64, ep_size=8)

在上面的代码片段中,fsdp_size=64 使用 FSDP2 将模型切分到全部 64 块 GPU 上,ep_size=8 将每层的 256 个专家分成 8 组。Transformers 会把每个 token 发送到持有其专家的 GPU 上,因此每块 GPU 只在自己的 batch 切片上训练。

现在,我们可以为本次训练配置 SFTConfig。这里我们执行了一个简单的任务,即在开源数据集上适配模型的聊天模板,但你也可以将它应用到自己的用例中。


import torch
from datasets import load_dataset
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer

train_dataset = load_dataset("allenai/tulu-3-sft-mixture", split="train[:500000]")

training_args = SFTConfig(
    output_dir="muse-spark-lora-sft",
    model_init_kwargs={"dtype": torch.bfloat16, "distributed_config": distributed_config},
    per_device_train_batch_size=1,
    max_steps=500,
    max_length=2048,
    packing=True,
    save_only_model=True,
    logging_steps=1,
    report_to="trackio",
)

trainer = SFTTrainer(
    model="meta-models/Muse-Spark-Open-816B-A42B",
    args=training_args,
    train_dataset=train_dataset,
    peft_config=LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]),
)
trainer.train()

这里的核心部分是 packing=True,它将样本拼接成 2048-token 的序列且无需填充。块对角掩码让各样本保持独立。save_only_model=True 只保存适配器而不保存优化器状态,因为训练器无法恢复加载时就已切分的模型。LoraConfig 针对四个注意力投影,每层的 256 个专家保持冻结。

然后我们可以在每个节点上运行以下命令:

torchrun --nnodes 8 --nproc_per_node 8 --rdzv_backend c10d --rdzv_endpoint "$MASTER_IP:29700" sft_muse_spark.py

我们训练了 500 步,在 64 块 H100 上大约花了 2 小时,速度约为每秒 8,000 token。平均 token 准确率从 72% 提升到 84%,损失从 1.26 降至 0.57,表明微调效果良好。

Loss

Acc

演示

Muse Spark Open 是一个编程与智能体模型。我们想亲自验证一下,因此尝试了一系列多模态和智能体任务,这些任务会用到模型的图像和音频输入。以下所有内容你都可以直接运行或 fork。

3D 资产生成器

我们将 2-bit 量化的 Muse Spark Open(由 llama.cpp 提供服务)与 Pi 编程智能体相结合,在 Three.js 中生成并迭代改进 3D 资产。提示词可以在这里找到,由于模型是量化的,需要稍多一些引导。在 Mac Ultra 上的 llama.cpp 中我们获得了 18 TPS。

Pi demo

我们首先描述该资产,然后 Pi 使用模型为该资产生成 Three.js 代码,在浏览器中渲染结果,并从多个角度截取屏幕截图。随后它把这些图像发送回 llama server 上的同一模型以获取视觉反馈。Pi 利用该反馈来编辑代码,并重复这一过程。

基于 OpenEnv 和 Hugging Face Sandbox 的浏览器交互

我们在计算机使用任务上测试了 Spark,对其性能感到惊叹,因此我们构建了一个精简的演示,同时也能教你计算机使用智能体是如何工作的。你还可以根据自己的需求进一步定制。这既是一个教学示例,也可以通过几行代码运行。你可以在这里找到它。

Muse Spark NVFP4 浏览器智能体(结合 OpenEnv 与 HF Sandboxes)

Hugging Face Sandboxes 是构建在 Hugging Face Jobs 之上、便于执行代码的沙箱。该演示由三个组件组成:

  1. 通过 Inference Endpoints 服务的 Muse Spark Open,
  2. OpenEnv 使用 helium 和 selenium 封装网页驱动操作(点击、输入、滚动等),定义浏览器环境的 reset、step、observation、termination 接口,
  3. Chromium 网页浏览器

其中远程沙箱包含该环境和浏览器,如下图所示。

Diagram

你可以通过几个简单步骤轻松尝试,示例中有详细说明。启动带有 Muse Spark Open 的 Inference Endpoint(NVFP4 也运行良好!),将你的端点 URL 添加为环境变量,把你的 OpenEnv 服务器发布到 Spaces 并同样提供该地址,然后你就可以运行智能体循环了。 Muse Spark Open 只能看到之前的操作、系统提示词、URL 和浏览器截图来决定采取什么操作,仅此而已(没有 DOM、没有 HTML 代码,什么都没有),因此这是一个真正的多模态浏览器演示。

呼叫你的浏览器

最后,我们想把 Spark 的多模态智能体能力推向极限,因此我们搭建了这个 Space,将模型的音频、图像和智能体特性结合成一个浏览器智能体。我们给模型一个关于所需物品的一般性描述,它会与集成的浏览器交互,在 Amazon 上查找选项。

Muse Spark Open 接收语音输入,推理出符合描述的产品类型,在 Amazon 上搜索,并根据标准选出最合适的一款。

结语

Muse Spark Open 是 Muse 系列的全尺寸开放权重模型。可以从 BF16 或 NVFP4 检查点开始,或试用其中一个出色的 llama.cpp quants。如果你用 Muse Spark Open 构建了什么,欢迎分享,并在仓库的讨论区告诉我们。

来源:Hugging Face Blog · huggingface.co