🤗 Transformers 中 Cosmos3 Edge Reasoner 的架构解析与多模态推理实战指南

本指南以官方模型文档 docs/source/en/model_doc/cosmos3_edge.md 为核心骨架，并结合仓库内真实源码与测试进行纵深讲解。读完本文，你将理解 Cosmos3 Edge 的 Reasoner 在 Transformers 中的集成范围与整体架构、掌握通过 AutoModelForImageTextToText/AutoProcessor 进行图文与视频理解推理的完整用法，并能读懂其三个配置类、打包式视觉处理与多模态 RoPE 位置编码的设计细节。

Cosmos3 Edge 是 NVIDIA Cosmos3 系列的多模态推理模型，Transformers 于 2026-07-16 将其接入仓库。需要特别强调它的集成边界：Transformers 只集成 Reasoner 塔，checkpoint 中的扩散式 Generator、VAE、scheduler 等生成组件仍由 Diffusers 负责。换言之，在本仓库中，你拿到的是一个"看得懂世界、能输出文字推理"的前端语言模型，而不是完整的图像/视频生成链路。

一、模型总览：Reasoner 塔的组成

根据模型文档介绍，Cosmos3 Edge 的 Reasoner 使用一个稠密（dense）、与 Llama 兼容的语言塔，由 28 个 decoder block 组成，每个 block 内含标准 self-attention 与 MLP。它的视觉侧采用 SigLIP2 视觉编码器，具备如下特性：

• 接受打包（packed）的可变分辨率 patch 输入；
• 使用**序列边界（sequence boundaries）**把不同图像、视频帧在视觉注意力阶段彼此隔离；
• 将 patch 按空间 2×2 分组后投影进语言模型；
• 图像与视频输入统一使用多模态旋转位置编码（M-RoPE）；
• 视频提示（prompt）会被展开为每个采样帧对应一个带时间戳的视觉 span。

这些描述都能在源码中找到精确对应实现，下文逐条展开。

二、架构源码级解读：从像素到语言 Token

完整的模型实现位于 modeling_cosmos3_edge.py（该文件由 modular_cosmos3_edge.py 自动生成，贡献者需修改 modular 源文件）。

2.1 语言塔：Llama 兼容的稠密 decoder

语言塔在 Cosmos3EdgeTextModel 中定义（modeling_cosmos3_edge.py），结构上非常接近 Llama 系模型：

• embed_tokens 词嵌入 → 28 个 Cosmos3EdgeTextDecoderLayer → 末端 RMSNorm；
• 每个 decoder layer 为标准的 Pre-Norm 残差结构：input_layernorm → self-attention → 残差 → post_attention_layernorm → MLP → 残差；
• 注意力为 GQA（Grouped Query Attention）：num_attention_heads=16、num_key_value_heads=8；
• 归一化使用 Cosmos3EdgeTextRMSNorm，eps=1e-5；
• MLP 激活函数为 relu2（ReLU 平方，代码注释称其为 "ReLU-squared MLP"），且 MLP 无 bias（mlp_bias=False）。

类级设计上，Cosmos3EdgeTextDecoderLayer 继承自 GradientCheckpointingLayer，支持梯度检查点；_no_split_modules 包含文本层与视觉层，便于设备切分与 FSDP。

2.2 视觉塔：打包式 SigLIP2

视觉部分在 Cosmos3EdgeVisionModel 中实现，结构与 SigLIP2 对齐：

• Cosmos3EdgeVisionEmbeddings：由于输入已经是"patchify 之后"的扁平 patch 序列，这里用 nn.Linear 做 patch embedding，再叠加位置编码；
• 位置编码采用可学习的方形参考网格（num_patches=256，即 16×16），在 resize_positional_embeddings 中为每个打包的图像/视频帧独立地做 bilinear 插值到该帧的 (H, W) 网格，再按 2×2 block-major 顺序重排（modeling_cosmos3_edge.py）；
• 27 层 Cosmos3EdgeVisionEncoderLayer，其视觉注意力不是因果的（is_causal=False）；
• 视觉注意力通过 cu_seqlens（累计 patch 计数边界）把每张图、每个视频帧限制在自己的序列内做 full attention：当请求 Flash Attention 时，直接把 cu_seq_lens_q/k 与 max_seqlen 传给底层实现；使用 eager 时则按 cu_seqlens 把序列 torch.split 成多段分别算再拼接（modeling_cosmos3_edge.py）。这就是文档所说"序列边界让图像与视频帧在视觉注意力中保持独立"的实现机制。

2.3 视觉-语言投影：Patch Merger

Cosmos3EdgePatchMerger 把视觉输出映射进语言空间（modeling_cosmos3_edge.py）：

• 先按 spatial_merge_size**2（即 2×2=4）把相邻 patch 分组拼接；
• 经 LayerNorm 后进入两段 MLP：hidden_size×4 (4608) → projector_hidden_size (11520) → text hidden_size (2048)，中间使用 GELU 激活。

文档所描述的"把 patch 按空间 2×2 分组后投影进语言模型"，正是该模块的职责。

2.4 图像与视频的 M-RoPE 位置编码

文档强调"图像和视频输入使用多模态 rotary position IDs"。对应的实现是 Cosmos3EdgeTextRotaryEmbedding，一种 interleaved M-RoPE（modeling_cosmos3_edge.py）：

• compute_default_rope_parameters 按 mrope_section=[24, 20, 20] 把维度分成三份，分别给 temporal、height、width 三个轴生成 inverse-frequency 矩阵（高度与宽度轴通过 indices % 3 交错分布）；
• forward 中 position_ids 形状为 (3, batch, seq)，三行分别代表 T/H/W 坐标；
• get_rope_index 依据 mm_token_type_ids 把输入序列切分成文本(0)/图像(1)/视频(2) 段：文本 token 用普通一维递增位置，视觉 token 由 get_vision_position_ids 生成三维网格位置；
• 一个细节：video_grid_thw 中每行代表一个源视频，而 Edge 的模板为每个采样帧输出一个视觉 span，因此 get_rope_index 开头会先把每个视频的 temporal 网格 repeat_interleave 展开为逐帧的单帧网格（modeling_cosmos3_edge.py）。

2.5 检查点加载细节：丢弃 Generator 权重

由于一个 checkpoint 内同时包含 Reasoner 与扩散 Generator 权重，加载时通过 _keys_to_ignore_on_load_unexpected = _COSMOS3_EDGE_DROPPED_GENERATOR_KEYS（正则列表，覆盖 action_modality_embed、time_embedder、proj_in/out、norm_moe_gen、mlp_moe_gen、add_q/k/v_proj 等）静默忽略 Generator 相关键，避免误报 unexpected key（modeling_cosmos3_edge.py）。这正是"Transformers 仅集成 Reasoner 塔"在权重层面的落地方式。

三、快速上手：图像问答推理

模型文档给出如下最小可运行示例，可直接保存为脚本执行：

from transformers import AutoModelForImageTextToText, AutoProcessor

model_id = "nvidia/Cosmos3-Edge"
model = AutoModelForImageTextToText.from_pretrained(model_id, device_map="auto")
processor = AutoProcessor.from_pretrained(model_id)

messages = [
    {
        "role": "user",
        "content": [
            {
                "type": "image",
                "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg",
            },
            {"type": "text", "text": "Describe this image."},
        ],
    }
]

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

generated_ids = model.generate(**inputs, max_new_tokens=128)
generated_ids = [output_ids[len(input_ids) :] for input_ids, output_ids in zip(inputs.input_ids, generated_ids)]
print(processor.batch_decode(generated_ids, skip_special_tokens=True))

几个值得展开的实操要点：

1. 模型与处理器都通过 Auto API 加载。AutoModelForImageTextToText 会路由到 Cosmos3EdgeForConditionalGeneration，AutoProcessor 则组合出 Cosmos3EdgeProcessor。
2. apply_chat_template 一站式完成多模态 prompt 构造。传入 messages 后，处理器会把 type: "image" 项中的 <|image_pad|> 占位符替换成与图片 token 数等量的 pad token，再 tokenize 并附加生成 prompt。对应逻辑在 processing_cosmos3_edge.py 的 replace_image_token：每个被投影的 2×2 patch 组产生一个 token，即 num_image_tokens = prod(grid_thw) // merge_size**2。
3. 生成后需裁剪输入部分。model.generate(**inputs) 返回的 generated_ids 包含输入前缀，因此用 output_ids[len(input_ids):] 切片去掉前缀，再用 processor.batch_decode（或 processor.post_process_image_text_to_text）解码，skip_special_tokens=True 以过滤 <|image_pad|>、<|vision_start|> 等特殊 token。
4. 处理器输出中除 input_ids、attention_mask 外还包含视觉输入：pixel_values、image_grid_thw，并额外返回 mm_token_type_ids（见 processing_cosmos3_edge.py 的 model_input_names）。多模态输入时缺少 mm_token_type_ids 会直接抛错，因为 M-RoPE 依赖它区分文本/图像/视频段。

3.1 视频输入：带时间戳的逐帧 span

如果 messages 中放的是视频内容项，流程与图像一致，但 prompt 组织方式不同。replace_video_token 会把一个视频展开成逐帧的、带时间戳的视觉段（processing_cosmos3_edge.py）：

<0.0 seconds><|vision_start|><|video_pad|>…<|vision_end|><1.0 seconds><|vision_start|>…

• 每帧 token 数为 grid_h × grid_w // merge_size**2；
• 时间戳由 _calculate_timestamps 依据 frames_indices 与视频 fps 计算（取每个时间 patch 中心帧对应时刻）；
• 若无法推断 fps，处理器会告警并默认 fps=24，因此建议显式传入准确的 video_metadata。

相应地，视频帧的采样策略由 Cosmos3EdgeVideoProcessor 承担（见第五节）。

四、预处理管线：smart_resize + patchify

Cosmos3 Edge 采用"打包式可变分辨率"视觉输入，核心处理逻辑在 image_processing_cosmos3_edge.py（另有基于 PIL 的 image_processing_pil_cosmos3_edge.py 与视频版 video_processing_cosmos3_edge.py）。

4.1 smart_resize：动态分辨率缩放

smart_resize 不把输入强行缩放到固定尺寸，而是按 factor = patch_size × merge_size 将高宽规整到 factor 的整数倍，同时把总像素约束在 [min_pixels, max_pixels] 区间：

• 图像处理器默认 size={"shortest_edge": 256*256, "longest_edge": 4096*4096}，patch_size=16、merge_size=2；
• 若缩放后长宽比超过 200 会抛错；宽或高不足 factor 时先等比放大；
• 视觉编码器侧 patch_size 为 16（与部分 Qwen 系 14 不同）。

4.2 patchify：block-major 2×2 展平

处理器的 patchify 与 Qwen2-VL 类模型风格一致：把图像切成 patch_size×patch_size 的 patch，并按 2×2 block-major 顺序展平成 (seq_len, patch_dim) 的扁平张量，便于把不同尺寸的图 concat 到一个 batch：

• 图像输出 pixel_values + image_grid_thw（每行 [1, grid_h, grid_w]）；
• 视频输出 pixel_values_videos + video_grid_thw（每行 [grid_t, grid_h, grid_w]），且仅支持 temporal_patch_size=1（初始化时若不满足会直接抛错）；
• 若帧数不能整除 temporal_patch_size，会 pad 最后一帧补齐。

4.3 视频帧采样策略

Cosmos3EdgeVideoProcessor.sample_frames 默认按 2 帧/秒（fps=2） 均匀采样，并将帧数夹在 min_frames=4 与 max_frames=768 之间（video_processing_cosmos3_edge.py）。可通过 num_frames（指定帧数）或 fps（指定帧率）二选一覆盖默认行为，二者互斥。

五、三个配置类与默认超参数

配置类定义于 configuration_cosmos3_edge.py，采用 Cosmos3EdgeConfig 持有 text_config、vision_config 两个子配置（与 Qwen2-VL 的多模态配置风格一致），均支持以 dict 或子配置实例传入，并带 @strict 架构校验。

5.1 Cosmos3EdgeTextConfig

参数	默认值	说明
vocab_size	131072	词表大小
hidden_size	2048	隐藏层维度
intermediate_size	9216	MLP 中间维度
num_hidden_layers	28	decoder 层数
num_attention_heads	16	注意力头数
num_key_value_heads	8	KV 头数（GQA）
hidden_act	"relu2"	激活函数
max_position_embeddings	131072	最大位置长度
initializer_range	0.02	初始化范围
rms_norm_eps	1e-5	RMSNorm epsilon
use_cache	True	是否使用 KV cache
pad_token_id / bos_token_id / eos_token_id	0 / 1 / 11	特殊 token id
attention_dropout	0.0	注意力 dropout
head_dim	128	每头维度
rope_parameters	{"rope_type": "default", "rope_theta": 1e8, "mrope_section": [24, 20, 20]}	M-RoPE 参数，mrope_section 三轴和必须等于 head_dim/2 = 64

架构校验 validate_architecture 会强制：hidden_size 能被 num_attention_heads 整除、rope_type 必须为 "default"、mrope_section 三段之和必须等于 head_dim/2（configuration_cosmos3_edge.py）。另外该配置还声明了张量并行（base_model_tp_plan）与流水线并行（base_model_pp_plan）计划，q/k/v/up 投影按 colwise、o/down 投影按 rowwise 切分。

5.2 Cosmos3EdgeVisionConfig

参数	默认值	说明
hidden_size	1152	视觉隐藏维度
intermediate_size	4304	视觉 MLP 中间维度
num_hidden_layers	27	视觉编码器层数
num_attention_heads	16	注意力头数
num_channels	3	输入通道数
patch_size	16	patch 大小
hidden_act	"gelu_pytorch_tanh"	视觉 MLP 激活
layer_norm_eps	1e-6	LayerNorm epsilon
attention_dropout	0.0	dropout
num_patches	256	可学习参考位置网格的 patch 数（16×16）
spatial_merge_size	2	2×2 空间合并

5.3 Cosmos3EdgeConfig（顶层）

参数	默认值	说明
text_config	默认文本子配置	可传 dict
vision_config	默认视觉子配置	可传 dict
projector_hidden_size	11520	视觉→语言投影 MLP 的中间维度
image_token_id	19	图像占位 token
video_token_id	18	视频占位 token
vision_start_token_id	20	`<
vision_end_token_id	21	`<
tie_word_embeddings	False	是否绑定词嵌入权重

六、前向推理与生成时的特殊处理

6.1 多模态前向流程

Cosmos3EdgeModel.forward 的拼接逻辑与 Qwen2-VL 的思路一致（源码注释也明言两者共享 multimodal API，但 Edge 使用自己的 packed vision 与 Llama 派生文本组件）：

1. inputs_embeds = embed_tokens(input_ids)；
2. 若传 pixel_values：get_image_features 产出视觉特征（vision tower → PatchMerger → 按 prod(grid_thw)//4 切分回各图），再经 get_placeholder_mask 校验占位 token 数与特征数一致，并用 masked_scatter 将特征写入图像占位位置；
3. 视频同理（get_video_features 复用与图像完全相同的视觉塔与投影路径）；
4. compute_3d_position_ids 计算 M-RoPE 所需的三维位置 id，然后交给语言塔做因果解码。

6.2 生成（decode）阶段的注意事项

• Cosmos3EdgeForConditionalGeneration 重写了 _prepare_position_ids_for_generation：多模态输入下用 get_rope_index 计算 3D 位置并缓存 rope_deltas；继续从已有 KV cache 生成时，用 text_positions + rope_deltas 快速恢复正确的位置（modeling_cosmos3_edge.py）；
• 纯文本输入时，自动退化为三维展开的普通位置 id（rope_deltas 置零）；
• 该模型还实现了 _expand_inputs_for_generation，正确处理 beam search 下图像/视频相关张量（pixel_values、image_grid_thw 等按图像/视频粒度 repeat，而非按帧粒度，避免 video_grid_thw 一行对一源视频、占位符却逐帧出现导致的错位）；
• 支持 Flash Attention 与 SDPA（_supports_flash_attn/_supports_sdpa 均为 True）；由于打包式可变分辨率视觉输入依赖逐网格的 Python 级 reshape，模型不支持全图编译（_can_compile_fullgraph=False）。

七、测试与进一步阅读

想验证上述理解或作为二次开发的参照，可阅读以下测试：

• tests/models/cosmos3_edge/test_modeling_cosmos3_edge.py：含纯文本 Cosmos3EdgeTextModelTester 与走共享 VLM 测试套件的 Cosmos3EdgeVisionText2TextModelTester。注意其中微缩配置把 mrope_section 设成 [2,1,1]、head_dim=8，验证了"三段和等于 head_dim/2"的约束；同时 create_pixel_values 注释印证"Edge 消费扁平化空间 patch，2×2 patch 网格合并为一个语言 token"；
• tests/models/cosmos3_edge/test_processing_cosmos3_edge.py、tests/models/cosmos3_edge/test_image_processing_cosmos3_edge.py、tests/models/cosmos3_edge/test_video_processing_cosmos3_edge.py：覆盖 chat template、时间戳计算、图像/视频 patchify 与采样。

相关源码清单（均以仓库根目录为起点）：

• 模型与权重加载：src/transformers/models/cosmos3_edge/modeling_cosmos3_edge.py
• 配置：src/transformers/models/cosmos3_edge/configuration_cosmos3_edge.py
• 处理器与模板替换：src/transformers/models/cosmos3_edge/processing_cosmos3_edge.py
• 图像/视频预处理：src/transformers/models/cosmos3_edge/image_processing_cosmos3_edge.py、src/transformers/models/cosmos3_edge/video_processing_cosmos3_edge.py
• Modular 源文件（改动入口）：src/transformers/models/cosmos3_edge/modular_cosmos3_edge.py

结语

Cosmos3 Edge 在 Transformers 中的集成是一个典型的"多模态理解 + 与生态分工"案例：Reasoner 用 28 层 Llama 兼容稠密语言塔负责推理与生成文本，SigLIP2 视觉塔通过打包式可变分辨率输入与 2×2 block-major patch 组织高效编码图像与视频，interleaved M-RoPE 则让多张图、多段视频在同一序列里共享三维坐标体系并保持帧间独立。理解这些源码级机制后，无论是调用 nvidia/Cosmos3-Edge 做图文/视频推理，还是在其上做微调或二次开发，你都能准确判断配置改动的边界与影响。