跳转至

运行

chitu generate 执行单次生成,chitu serve 启动 EPE 服务。CLI 支持 zimage、 flux1、flux2-klein、qwen-image、wan、llada-image 和 hunyuan-image3。

单卡生成

chitu generate \
  --model zimage \
  --model-path /path/to/Z-Image \
  --prompt "a red cube on a white table" \
  --output outputs/zimage.png

查看某个模型的完整参数:

chitu generate --model zimage --help

所有 generate 入口都会打印 elapsed_s=...,并在产物旁写入同名 .json。 elapsed_s 与兼容字段 generate_seconds 都以秒记录本次生成耗时,包括文本编码、 去噪、VAE 解码和结果后处理,不含模型加载、图片保存或视频编码。CUDA 计时边界会 等待设备工作完成;这是单次请求耗时,首轮可能包含编译或初始化成本。 JSON 同时记录 prompt、seed、步数和并行参数;FlexCache 模型还记录缓存统计。

静态多卡生成

--cp-method agkv|ulysses 选择 CP 算法,--cp-backend fast|nccl 选择通信后端。 例如 agkv + fast 使用 Fast AGKV,ulysses + fast 使用 Fast Ulysses。 方法默认为 agkv;显式指定后端会覆盖对应的 transport 环境变量。 旧的 --attention-mode、--agkv-transport 和 --ulysses-transport 仍可使用, 但新旧参数冲突或 transport 与方法不匹配时会报错。

所有 rank 执行同一条命令,只有 leader 写出结果。默认 transport 是 Torch/NCCL:

torchrun --standalone --nproc-per-node=4 -m chitu_diffusion.cli \
  generate \
  --model flux1 \
  --model-path /path/to/FLUX.1-dev \
  --output outputs/flux1.png

Fast AGKV:

torchrun --standalone --nproc-per-node=4 -m chitu_diffusion.cli \
  generate \
  --model zimage \
  --model-path /path/to/Z-Image \
  --cp-method agkv --cp-backend fast \
  --no-cfg-parallel \
  --output outputs/zimage-fast-agkv.png

Fast Ulysses:

torchrun --standalone --nproc-per-node=4 -m chitu_diffusion.cli \
  generate \
  --model zimage \
  --model-path /path/to/Z-Image \
  --cp-method ulysses \
  --cp-backend fast \
  --no-cfg-parallel \
  --output outputs/zimage-fast-ulysses.png

以上两例关闭 CFG 分组,使用 CP4。它们选择 Fast 通信后端;Wan、Qwen-Image 和 Z-Image 同时启用 Attention Planner,完整 wrapper 的投影融合和通信重叠仍需显式 接入,见 Fast CP 启用方式。

Slurm 上生成 Qwen-Image,可使用仓库脚本申请单节点四卡:

脚本默认使用集群默认分区;需要指定分区时,在个人环境中设置 SRUN_PARTITION。

先完成 Fast CP 安装与验证,并将 QWEN_IMAGE_MODEL_PATH 设为已有的本地权重目录。统一环境脚本会从 .deps/fast-cp/ 配置扩展、CuTe 和 NVSHMEM 的加载路径:

source tools/install/fast_cp_env.sh
: "${QWEN_IMAGE_MODEL_PATH:?请设置为本地 Qwen-Image 权重目录}"
bash tools/cluster/srun_direct.sh 1 4 \
  -m chitu_diffusion.cli generate \
  --model qwen-image --model-path "$QWEN_IMAGE_MODEL_PATH" \
  --cp-method agkv --cp-backend fast --no-cfg-parallel \
  --height 1024 --width 1024 --steps 50 \
  --output outputs/qwen-fast-agkv.png

Hunyuan Image 3 的 TP/CFG/CP/EP 拓扑在启动时确定。以下使用 72 GiB 卡上验证的 TP2×CFG2×CP2×EP2,并以 VAEP8 解码:

torchrun --standalone --nproc-per-node=8 -m chitu_diffusion.cli \
  generate \
  --model hunyuan-image3 \
  --model-path /path/to/HunyuanImage-3 \
  --expert-parallel-degree 2 \
  --steps 50 \
  --output outputs/hunyuan_image3.png

只需满足 world = TP × CFG × CP、CFG ∈ {1, 2}、EP 整除 CFG × CP,即可改成 其他单机拓扑。未给出的 degree 由 world size 推导,例如下面等价于 TP2×CFG1×CP4×EP4:

torchrun --standalone --nproc-per-node=8 -m chitu_diffusion.cli \
  generate --model hunyuan-image3 --cfg-parallel-degree 1 ...

VAE 解码

Z-Image、FLUX.1、FLUX.2-klein、LLaDA-Image、Wan、Qwen-Image、Hunyuan Image 3 和 MiniMax-H3 的图像/视频静态并行解码统一使用逐层 行分片:卷积交换相邻 rank 的边界行,GroupNorm 通信全局统计量,attention 使用本地 Q 和 all-gather 后的全局 K/V(AGKV)。视频保留原有时序缓存。 该路径保留整图解码语义,浮点归约与卷积算法差异仍可能造成数值误差。

提供 VAE 开关的离线入口使用 --parallel-vae 开启,--no-parallel-vae 让 leader 整图解码;固定服务组使用 parallelism.vae 配置。 Z-Image 和 LLaDA-Image 默认关闭,其余上述模型默认开启。单 rank、分片过短 或不支持的 decoder 使用 leader 路径;不支持的结构会提示原因。 开启逐层并行前需关闭 Diffusers 自带 tiling,否则会报错。

--vae-parallel-halo N 保留兼容(非负),但上述逐层路径按各层卷积核自动 推导 halo,不使用该值。FLUX.2-klein 只提供启停开关,不需要手工 halo 参数。 sidecar 的 requested_parallel_vae、requested_vae_parallel_halo 保留请求值, vae_decode_mode(layerwise / leader)、parallel_vae 和 vae_parallel_degree 记录实际执行方式,逐层路径的 vae_parallel_halo 为 null。

本次实现和验证范围为静态并行,不包含动态 EPE。 Hunyuan 的全局 GroupNorm/attention 和 MiniMax 的全局 RoPE、特殊 token 都已纳入通信语义;MiniMax 串行参考也改为关闭空间 tiling,音频仍由 leader 解码。 统一算子范式见 VAEP 架构。 支持范围、误差和性能测量见 VAE 验证记录。

FlexCache

缓存参数直接附加到 generate:

chitu generate \
  --model wan \
  --model-path /path/to/Wan2.1-T2V-1.3B \
  --steps 50 \
  --cache-strategy magcache \
  --output outputs/wan-magcache.mp4

策略专属参数使用对应前缀。例如 MeanCache:

--cache-strategy meancache --meancache-fresh-steps 25

EPE 服务

配置中的 GPU 数必须等于 torchrun 进程数。四卡示例:

torchrun --standalone --nproc-per-node=4 -m chitu_diffusion.cli \
  serve --stage-config examples/stage-zimage.yaml

使用 Fast AGKV 的固定 CP4 服务,先修改配置中的模型路径:

torchrun --standalone --nproc-per-node=4 -m chitu_diffusion.cli \
  serve --stage-config examples/stage-zimage-fast-cp.yaml

通信方式在 parallelism.cp 下设置;示例关闭 CFG 分组并使用 static_cp 调度。 这与上面的 generate 命令启用范围一致,不会自动安装完整 wrapper。

八卡 static-CP MoE 服务示例:

torchrun --standalone --nproc-per-node=8 -m chitu_diffusion.cli \
  serve --stage-config examples/stage-hunyuan-image3.yaml

leader 提供 HTTP 服务,其他 rank 执行 worker loop。宿主已经管理进程时,可参考 examples/epe_embedded.py 使用 EmbeddedDiffusionRuntime。

配置里的 parallelism 分为公共的 cp、vae、scheduler 和模型专属的 model, 字段含义见 EPE 并行配置。

Slurm

集群入口会为每张 GPU 启动一个 task,并设置 RANK、LOCAL_RANK 和 WORLD_SIZE:

bash tools/cluster/srun_direct.sh 1 4 -m chitu_diffusion.cli \
  generate \
  --model zimage \
  --model-path /path/to/Z-Image \
  --output outputs/zimage.png

混合 GPU 节点可通过 SRUN_GPU_TYPE 指定 Slurm 注册的型号,例如:

SRUN_PARTITION=Star SRUN_GPU_TYPE=PRO6000 \
  bash tools/cluster/srun_direct.sh 1 2 -m chitu_diffusion.cli \
  generate --model zimage --model-path /path/to/Z-Image \
  --output outputs/zimage-pro6000.png

未设置 SRUN_GPU_TYPE 时保留 gpu:N 的调度方式;型号名称以集群的 GRES 配置为准。

仓库根目录 examples/ 还包含各模型 Python API、静态 CP、EPE embedded 和服务配置 示例。