Cosmos3T2VWNScheduler,
"Plugin": cosmos3_wn,
"源码": scheduler_wn_cosmos3.py,
"并行轴": ["tp", "cfg"],
"支持路径": ["T2V", "I2V", "T2AV", "I2AV"],
"采样器": "UniPC",
}}
/>
# 概述
Cosmos3 是面向 Physical AI 的 omnimodal world model。它可以从 text、image、video、audio 和 action trajectory 等组合输入中生成 video、image、audio、action 等输出，用在世界生成、未来预测、动作推理和 embodied policy learning 这类任务上。
本页讲的是 Cosmos3 的多卡生成路径，也就是 `cosmos3_wn` plugin。它覆盖 T2V/I2V/T2AV/I2AV：视频 latent 和可选 sound latent 在同一个 denoise loop 里推进，最后由 VAE/AVAE decode 成可保存的媒体。
PhyAI 目前在这条路径上支持两类并行：transformer forward 沿 `tp` 轴做 tensor parallel；当 `cfg=2` 且 `guidance_scale > 1` 时，cond / uncond 两个 CFG branch 会被分到两个 TP group 上并行运行。视频 VAE decode 也会按 rank 做空间 tile 切分，并用 halo overlap 合并边界。
# 并行拓扑
下面用 `TP=4`、`CFG=2`、`world_size=8` 作为例子。8 张卡被拆成两个 CFG group：rank 0-3 是 cond branch 的一个 TP group，rank 4-7 是 uncond branch 的另一个 TP group。每个 denoise step 结束前，两个 branch 的 velocity 会沿 `cfg` 轴 all-gather 到所有 rank，再在本地完成 CFG combine。
`cfg=2` 只在 `guidance_scale > 1` 时有收益。否则 scheduler 会打印 warning，因为 uncond 分支仍然会被分配 rank 和计算，但 CFG 实际没有打开。Cosmos3 只有 cond 和 uncond 两个 CFG 分支，示例脚本也把 `--cfg` 限制为 `1` 或 `2`。
# Decode 与输出
`Cosmos3T2VWNScheduler.step()` 只负责生成 latent。`Cosmos3WNEntry.step()` 会在 scheduler 返回后调用 decode：
| 请求类型 | Scheduler 返回 | Entry 返回 |
| --------- | ------------------------------------------------ | ---------------------------------------------------------- |
| T2V/I2V | `video` latent，shape `[1, C, t, h, w]` | pixels，shape `[B, 3, T, H, W]`，范围 `[0, 1]` |
| T2AV/I2AV | `{"video": video_latent, "sound": sound_latent}` | `{"video": pixels, "sound": waveform, "sample_rate": int}` |
视频 decode 使用 `Cosmos3VAERunner.decode()`。当并行规模大于 1 时，runner 会走 WAN VAE 的并行 decode，把空间 tile 分配给不同 rank，并在 halo 处理后合并结果。音频 decode 使用 `Cosmos3SoundVAERunner.decode()`。
VAE 的 8 卡切分示意如下，`cfg` 作为 outer axis，`tp` 作为 inner axis：
# 运行路径