两张 7900 跑 MiniMax-H3 张量并行：每步 18.6 s → 8.3 s，每卡省 10 GiB

[fcys]

昨天看到老特的视频，就特意试了一下。刚好广东家里的机器插了 3 张卡，其中两张是
Navi 31 的 W7900（gfx1100）——所以这组能跑通，7900 XTX 也就能跑：同一颗 Navi 31、
同一个 gfx1100，我们这套里没有任何 W7900 专属分支。区别只在显存（24 G vs 48 G），
XTX 上要跟着降分辨率/帧数，或者把 26 G 的文本编码器换成 int4 版。
（我手上只有 W7900，XTX 是推断，没实测。）

机器在广东家里，我人在青岛，全程远程。两张 W7900 的上行桥现在只能跑 PCIe 4.0 x8，
all_reduce 实测 12.93 GB/s；要等 11 月中旬秋假结束回去才能动平台。所以下面所有数字
都是 x8 条件下的下限，链路放开只会更好。

还要先说清一件事：7900 上用不了现成的 Radiance（那套内核是 RDNA4 专用，见下文），
我们是自行实现的一套同类方案——ComfyUI + int8 量化 + 自己写的 TP2，全程没有 FP8。

结果

同一张图、同 seed（145）、同 8 步、同一张工作流，只把 DiT 从 1 卡换到 2 卡：

指标	单卡 W7900	TP2（双 W7900）	变化
每步耗时	18.6 s	8.3 s	2.24×
端到端（832×480 / 124 帧 / 8 步）	201.3 s	134.1 s	1.50×
采样期显存（每卡）	23.5 GiB	12.9 / 13.2 GiB	−10.6 / −10.3 GiB
DiT 权重	20 G	10 G / 卡	半宽切分
两卡产出	—	MD5 完全相同	逐位一致
与单卡画面（同 seed）	—	PSNR 25.5 dB	构图/人物/运镜一致，细节不同

三句话读完：

• 加速是真的：每步 2.24×；端到端只有 1.50×，因为文本编码（26 G Qwen3-VL）和
  VAE 解码不参与并行。
• 显存也是真的：每卡少 10 GiB 左右，DiT 权重直接砍半。
• 两个实例各自存的 mp4 字节一致，说明锁步、通信、拼接没有一处错位。

换个更快的 attention 后端，比例会变

后端	单卡	TP2	提速
SDPA（默认）	18.6 s/步	8.3 s/步	2.24×
--use-ck-attention	12.8 s/步	7.7 s/步	1.66×

CK 让单卡自己就快了 1.45×，所以 TP2 的相对收益降到 1.66×；但绝对速度最快的是
TP2 + CK（7.7 s/步、端到端 128 s、每卡 13 GiB 上下）。显存收益和后端无关。
两个表都是同后端自比，不跨后端混比。

为什么原帖那套抄不了

原帖是 2× Radeon AI PRO R9700（RDNA4 / gfx1201）跑 AMD 的 Radiance 推理引擎做 TP2。

• Radiance 的内核库 libr4d 只编了 gfx1200 / gfx1201，而且用到 wmma-128b、fp8 转换
  这类 RDNA4 才有的硬件指令；
• W7900 是 RDNA3 / gfx1100：没有 FP8 矩阵单元。我们试过编 libr4d，18 处编译错误，
  这条路是死的。

所以不是"参数没调对"，是硬件指令集里没这东西。原帖的 FP8 权重 + Radiance 内核，
在 RDNA3 上没有对应实现。

我们怎么做的

既然没有 FP8，就走 RDNA3 能跑的 INT8：

• 权重格式：int8，per-output-row 的 weight scale + Hadamard convrot（group 256）；
• gfx1100 有可用的 WMMA / hipBLASLt INT8 GEMM，所以 int8 走的是硬件矩阵指令，
  代价只是每层多一步反量化/旋转。

模型是 H3 的 DiT：50 层、hidden 5376、56 头 × 128、ffn 14336，int8 权重 20 G。
切分照 Megatron 那一套，只切 DiT：

层	切法	通信
qkv_proj	列并行（56 头 → 每卡 28 头）	无
out_proj	行并行（K 轴 3584）	all_reduce
fc1	列并行（[gate\|up] 各 7168）	无
fc2	行并行（K 轴 7168）	all_reduce

三个必须解决的细节：

1. int8 是 per-row 动态量化 → 行并行必须跨卡共享激活 scale。 两卡各算各的行 scale，
   量化点就不同，每层约 1% 误差，50 层能累到 40%+。做法：每行 |x| 最大值做
   all_reduce(MAX)，通信量只有 S×4 字节。
2. 部分和必须 fp32 累加再转回 bf16，否则每层约 2e-3。
3. 半宽模型不改一行源码：在 load_diffusion_model_state_dict 之前把 state_dict 切半，
   ComfyUI 会从张量形状自动推出 heads=28 / ffn=7168。

运行时是两个 ComfyUI 实例跑同一张工作流，DiT 里的 all_reduce 天然锁步，后端 RCCL；
每个实例只暴露自己那张卡（否则 ComfyUI 会把 current device 设回 0，RCCL 报 same GPU）。

踩过的坑

1. 没有 FP8：见上一节，别照着 FP8 教程折腾。
2. 激活 scale 不共享 = 白干：一开始两卡各量化各的，出片和单卡完全对不上，
   算子级误差是每层百分级，根因就是量化点不一致。
3. 部分和用 bf16 = 精度崩：内核默认输出 bf16，两半各自舍入再相加，每层 2e-3；
   改成 fp32 累加后，算子级误差降到 1e-7 量级。
4. attention 后端：一直用默认的 --use-pytorch-cross-attention（SDPA），长序列上很慢——
   同一张 832×480/124 帧的图、同配置跑 1 步：SDPA 495 s，换 CK 后 97 s（5.08×），
   显存还低（32.3 → 18.0 G/卡）。一开始还误以为是显卡不行。
   改法很简单，不动代码，只换启动参数：把 --use-pytorch-cross-attention 换成
   --use-ck-attention（Comfy Kitchen 的 int8 attention），单卡实例和 TP2 的两个 rank
   一起换；工作流侧也可以用 ModelAttentionBackend 节点设成 comfy kitchen attention。
5. --disable-smart-memory 是"激进卸载到普通内存"，不是"强制常驻"，别当显存开关用。

测试条件

• 素材：自建（一张 832×480 的首帧图，I2V），与任何人的素材无关；
• 832×480 / 124 帧 / 24 fps；8 步（PDD 蒸馏），seed 145；
• 每步 = 日志里第 2–7 步的均值（首步含模型初始化、末条是进度条汇总行，都不计）；
• 端到端 = 提交 → 两个 rank 都出片（含文本编码与 VAE 解码）；
• 显存按 0.5 s 采 /sys/bus/pci/devices/*/mem_info_vram_used，峰值/均值取自采样期。

还没做好的地方

• TP 切的是权重和算力，不是激活：文本编码器和 VAE 不参与并行，所以端到端只快 1.5×；
• 激活不切：序列越长激活越大，想装更长的视频得靠序列并行（我们另一套 SP 实验，
  显存能摊，但速度不划算，暂时没上）；
• 只测了 2 卡、只测了 H3 这一个模型；两实例锁步是实验性做法，不做稳定性承诺；
• 和单卡不是逐像素一致：PSNR 25.5 dB，差异在细节。这不是 bug——int8 量化点在这个
  栈上没法完全对齐（单卡自己换个分块都有 2e-2 的 latent 差），我们只保证两卡彼此一致。

最后说实话：效果现在还不理想。按行耗时归一，原帖 Radiance + FP8 的双卡仍比我们快
不少，差距主要在 int8 软件路径 vs FP8 硬件。这篇先抛砖引玉——有在 RDNA3 上折腾过的
朋友，欢迎一起把 attention 后端、量化方式和通信再往下压。

[懒人烘培]

我也在正在部署这个，就看到Radiance不支持7900XTX，看到帖子后我也少走些坑，看看没有进一步优化空间

[Bunsei]

巧了，我也在青岛，测试时候的机器也是远程。

libr4d 就是 gfx1201 专用的， gfx1100 根本没有 FP8 矩阵单元，FP8这条是死路，RDNA3转 int8 是唯一正解。

attention 这块刚巧看到这篇（https://lcz.me/post/28547 ），里面提到 nicejunjie/Strata-rdna3 这个项目带了一套完整的 gfx1100 INT8 WMMA 注意力，可以作为参考，但要是想视频用上得整体评估移植，得喂大肥鱼不少白饭。

[fcys]

@Bunsei 老乡啊哈哈哈

[Magic629]

@Bunsei 老乡~你好！

[Xiaote]

版主审定：本帖设为精华，作者 +6 积分奖励，希望再接再厉！

奖励凭证：精华 +6 分 · 编号 f2643-1915-1

[terry]

哥们这不像个新手啊，FP8不行直接转INT8，一看就懂啊。两张W7900，这玩意可不便宜，比4090 48G还贵。

[fcys]

@terry 我确实是一上来就让走的int8，发出来是做个提醒，跟Radiance注意区分。卡是刚好手里有，家里有不少稀奇古怪的卡 ，只是给同核心的7900xtx做个参考。

[hippo]

有github地址吗

[Alex Liu]

@fcys 专业啊，一看就是懂行的，下次发个帖子让锤友们看看稀奇古怪的卡还有哪些。

[fcys]

@Alex-Liu intel的卡，我归入了稀奇古怪里面了

[fcys]

@hippo 冇啊

[hippo]

@fcys 代码方便分享看看吗？存到网盘之类的

[fcys]

@hippo 您把帖子直接丢给agent就行了啊

[Enigma]

新卡已经到了，今晚作业太多了，要抄晕了

[懒人烘培]

@fcys 你这个自己写的tp2代码，由于没有公布，AI自己写了个速度不理想

[fcys]

@懒人烘培
说实话优化空间依然很大，我就是抛砖引玉，后面劳烦大家群策群力

TP2 代码全文（tp2-comfy v1.0.0）

直接可跑的最小实现；完整说明/排查清单见 README.md，包见 tp2-comfy-v1.0.0.zip。

依赖：ComfyUI + ComfyUI-INT8-Fast-ROCM（int8_convrot）+ comfy-kitchen（CK attention）；
模型需为 minimax_h3_*_pruned_int8_convrot（qkv 行数 21504、56 头 × 128、ffn 14336）。

tp/core.py —— 权重重切 + 块级 all_reduce

"""TP2 共用：int8 权重切分 + 块级 all_reduce 补丁。

所属：tp2-comfy v1.0.0（双卡张量并行，ComfyUI + MiniMax-H3 int8_convrot）
环境变量：
  COMFY_ROOT  ComfyUI 根目录（默认 /home/ly/ComfyUI-ROCm）
  TP_UNET     仅 build_model() 独立探针用；生产路径从工作流里读模型
"""
import os
import sys
import importlib

import torch
import torch.distributed as dist

APP_VERSION = "1.0.0"
TP_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, TP_DIR)
import tp_linear

ROC = os.environ.get("COMFY_ROOT", "/home/ly/ComfyUI-ROCm")
UNET = os.environ.get(
    "TP_UNET",
    "/home/ly/ComfyUI Tools/comfyui/models/diffusion_models/"
    "minimax_h3_fl2va_pruned_int8_convrot.safetensors")


def build_model():
    """独立探针：不起 ComfyUI，直接加载模型看结构（可选调试入口）。"""
    os.chdir(ROC)
    sys.path.insert(0, ROC)
    sys.path.insert(0, os.path.join(ROC, "custom_nodes"))
    import comfy.sd
    import comfy.model_management as mm
    q = importlib.import_module("ComfyUI-INT8-Fast-ROCM.int8_quant")
    mp = comfy.sd.load_diffusion_model(UNET, model_options={"custom_operations": q.Int8TensorwiseOps})
    mm.load_models_gpu([mp])
    return mp.model.diffusion_model


def _slice_rows(lin, idx):
    """沿输出轴（行）取分片；tensorwise scale 不变，per-row scale 同步切。"""
    w = lin.weight.data
    lin.weight.data = w[idx].contiguous()
    lin.out_features = lin.weight.data.shape[0]
    sc = getattr(lin, "weight_scale", None)
    if sc is not None and torch.is_tensor(sc) and sc.numel() == w.shape[0]:
        lin.weight_scale = sc[idx].contiguous()
    return lin


def _slice_cols(lin, idx):
    w = lin.weight.data
    lin.weight.data = w[:, idx].contiguous()
    lin.in_features = lin.weight.data.shape[1]
    return lin


def apply_tp(net, rank, world=2, which="both", slice_weights=True):
    """按 head 切 attention、按 ffn 切 MLP；返回每层的切分信息供反向检查。

    只支持 world=2：heads=56 → 每卡 28；ffn=14336 → 每卡 7168。
    """
    assert world == 2, "本实现只支持 2 卡（半宽切分）"
    info = []
    for bi, blk in enumerate(net.blocks):
        attn, mlp = blk.attn, blk.mlp
        heads, hd = attn.heads, attn.head_dim
        inner = heads * hd
        half = heads // world
        lo, hi = rank * half, (rank + 1) * half
        hsel = torch.arange(lo * hd, hi * hd)

        if which in ("both", "attn", "ref"):
            # qkv：q/k/v 三段各取本 rank 的 head 行
            idx = torch.cat([hsel, inner + hsel, 2 * inner + hsel])
            attn.heads = half
            if which != "ref" and slice_weights:
                _slice_rows(attn.qkv_proj, idx)
                _slice_cols(attn.out_proj, hsel)
            tp_linear.install_row_parallel(attn.out_proj)   # 行并行：全局行 scale + 由调用方 all_reduce
        else:
            attn.heads = heads

        # MLP：fc1 输出是 [gate(ffn) | up(ffn)]，两半各取 ffn/2
        ffn = mlp.fc1.out_features // 2
        fh = ffn // world
        f_lo, f_hi = rank * fh, (rank + 1) * fh
        if which in ("both", "mlp", "ref"):
            gidx = torch.cat([torch.arange(f_lo, f_hi), torch.arange(ffn + f_lo, ffn + f_hi)])
            if which != "ref" and slice_weights:
                _slice_rows(mlp.fc1, gidx)
                _slice_cols(mlp.fc2, torch.arange(f_lo, f_hi))
            tp_linear.install_row_parallel(mlp.fc2)          # 行并行
        info.append((bi, half, (f_lo, f_hi)))
    return info


def patch_blocks(net, which="both"):
    """在 attn/mlp 输出后插入 all_reduce（Megatron 行并行语义：部分和相加）。"""
    import comfy.ldm.minimax.model as M

    def block_forward(self, x, t_emb, mod_segments, rope_freqs, transformer_options={}, attention=None):
        attention = self.attn if attention is None else attention
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb)
        h = M._mod_scale_shift(self.norm1(x), shift_msa, scale_msa, mod_segments)
        a = attention(h, rope_freqs=rope_freqs, transformer_options=transformer_options)
        if which in ("both", "attn"):
            dist.all_reduce(a)
            if a.dtype != x.dtype:
                a = a.to(x.dtype)
        x = M._mod_gate(x, gate_msa, a, mod_segments)
        h = M._mod_scale_shift(self.norm2(x), shift_mlp, scale_mlp, mod_segments)
        m = self.mlp(h)
        if which in ("both", "mlp"):
            dist.all_reduce(m)
            if m.dtype != x.dtype:
                m = m.to(x.dtype)
        return M._mod_gate(x, gate_mlp, m, mod_segments)

    for b in net.blocks:
        b.forward = block_forward.__get__(b, type(b))


def _head_rows(attn_heads, head_dim, inner, rank, world=2):
    half = attn_heads // world
    lo, hi = rank * half, (rank + 1) * half
    hsel = torch.arange(lo * head_dim, hi * head_dim)
    return hsel, torch.cat([hsel, inner + hsel, 2 * inner + hsel]), half


def slice_state_dict(sd, rank, world=2):
    """在模型构造前从 state_dict 切分，避免"先装全量再切"的显存峰值。

    H3 配置固定：hidden=5376、56 头 × head_dim 128 → inner=7168、qkv 行数 21504。
    形状守卫：qkv 行数不是 21504 就认为已切过（PDD 节点会二次加载，不能切两遍）。
    """
    HIDDEN, QKV_ROWS, HEAD_DIM, HEADS = 5376, 21504, 128, 56
    k0 = "blocks.0.attn.qkv_proj.weight"
    if k0 not in sd or sd[k0].shape != (QKV_ROWS, HIDDEN):
        return sd
    inner = HEADS * HEAD_DIM
    half = HEADS // world
    hsel = torch.arange(rank * half * HEAD_DIM, (rank + 1) * half * HEAD_DIM)
    qkv_idx = torch.cat([hsel, inner + hsel, 2 * inner + hsel])
    n = 0
    for k in list(sd.keys()):
        t = sd[k]
        if k.endswith(".attn.qkv_proj.weight") and t.shape[0] == QKV_ROWS:
            sd[k] = t[qkv_idx].contiguous(); n += 1
        elif k.endswith(".attn.qkv_proj.weight_scale") and t.shape[0] == QKV_ROWS:
            sd[k] = t[qkv_idx].contiguous(); n += 1
        elif k.endswith(".attn.out_proj.weight") and t.shape[1] == inner:
            sd[k] = t[:, hsel].contiguous(); n += 1
        elif k.endswith(".mlp.fc1.weight") and t.shape[1] == HIDDEN:
            ffn = t.shape[0] // 2
            fh = ffn // world
            gidx = torch.cat([torch.arange(rank * fh, (rank + 1) * fh),
                              torch.arange(ffn + rank * fh, ffn + (rank + 1) * fh)])
            sd[k] = t[gidx].contiguous(); n += 1
            ks = k[:-len("weight")] + "weight_scale"
            if ks in sd and sd[ks].shape[0] == 2 * ffn:
                sd[ks] = sd[ks][gidx].contiguous(); n += 1
        elif k.endswith(".mlp.fc2.weight") and t.shape[0] == HIDDEN:
            ffn = t.shape[1]
            fh = ffn // world
            sd[k] = t[:, rank * fh:(rank + 1) * fh].contiguous(); n += 1
    print(f"[tp] state_dict 已切分 {n} 个张量（rank{rank}）", flush=True)
    return sd


def finalize_tp(net, rank):
    """模型已按半宽构造（state_dict 预切）时调用：只装行并行前向、不切权重。"""
    n = 0
    for blk in net.blocks:
        tp_linear.install_row_parallel(blk.attn.out_proj)
        tp_linear.install_row_parallel(blk.mlp.fc2)
        n += 1
    b0 = net.blocks[0]
    print(f"[tp] rank{rank} 半宽模型已确认：heads={b0.attn.heads} qkv_out={b0.attn.qkv_proj.weight.shape[0]} "
          f"fc1_out={b0.mlp.fc1.out_features} 行并行层={n}", flush=True)
    return n

tp/tp_boot.py —— RCCL 初始化 + 加载钩子

"""TP2 启动引导：在 ComfyUI 外面套一层，挂 RCCL + 模型加载钩子。

教训（照抄会炸）：不能在 ComfyUI 之前 import torch / comfy.sd —— 会破坏 AIMDO 初始化
（报 'NoneType' object has no attribute 'hostbuf_allocate'）。
做法：只提前打开参数解析（comfy.options 很轻），其余工作交给后台线程，
等 main.py 自己把 comfy.sd 导入完之后再 init RCCL 并挂载模型加载钩子。

用法（每个 rank 一个进程，各只暴露一张卡）：
  TP_RANK=1 ROCR_VISIBLE_DEVICES=<GPU1> python tp_boot.py --port 8191 ...
  TP_RANK=0 ROCR_VISIBLE_DEVICES=<GPU0> python tp_boot.py --port 8190 ...
环境变量：
  COMFY_ROOT     ComfyUI 根目录（默认 /home/ly/ComfyUI-ROCm）
  TP_RANK        0 / 1
  MASTER_ADDR    RCCL 地址（默认 127.0.0.1）
  MASTER_PORT    RCCL 端口（默认 29670）
"""
import os
import sys
import threading
import time

APP_VERSION = "1.0.0"
TP_DIR = os.path.dirname(os.path.abspath(__file__))
ROC = os.environ.get("COMFY_ROOT", "/home/ly/ComfyUI-ROCm")
sys.path.insert(0, ROC)
sys.path.insert(0, TP_DIR)

import comfy.options as _co  # 只定义开关，不拉 torch/cli_args
_co.enable_args_parsing()

rank = int(os.environ["TP_RANK"])
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
os.environ.setdefault("MASTER_PORT", "29670")
READY = threading.Event()


def _on_model(mp):
    # 注意：PDD 加速节点会重建一次模型（adaln rebase），每次加载都要切，不能只切一次
    if mp is None:
        return mp
    try:
        net = getattr(mp.model, "diffusion_model", None)
        if net is None or not hasattr(net, "blocks"):
            return mp
    except Exception:
        return mp
    if getattr(net, "_tp_done", False):
        return mp
    import core
    try:
        rows = int(net.blocks[0].attn.qkv_proj.weight.shape[0])
    except Exception:
        rows = 0
    if rows == 10752:      # 已经是半宽模型（预切路径）
        core.finalize_tp(net, rank)
    elif rows == 21504:    # 全宽模型：现场切
        core.apply_tp(net, rank)
    else:
        raise RuntimeError(f"未预期的 qkv 行数 {rows}（期望 21504 或 10752）")
    core.patch_blocks(net)
    net._tp_done = True
    b0 = net.blocks[0]
    print(f"[tp] rank{rank} 已切分 DiT：heads={b0.attn.heads} "
          f"qkv_out={b0.attn.qkv_proj.weight.shape[0]} out_proj_in={b0.attn.out_proj.weight.shape[1]} "
          f"fc1_out={b0.mlp.fc1.out_features} fc2_in={b0.mlp.fc2.weight.shape[1]}", flush=True)
    return mp


def _boot():
    # 等 main.py 把 comfy.sd 导入完，保证 torch 的导入顺序由 ComfyUI 自己决定
    for _ in range(12000):
        if "comfy.sd" in sys.modules:
            break
        time.sleep(0.05)
    import torch
    import torch.distributed as dist
    # 每个进程只暴露自己那张卡（ROCR_VISIBLE_DEVICES 单卡）→ 本地 index 恒为 0，
    # 这样 ComfyUI 主线程怎么设 current device 都不会和别的 rank 撞同一张卡。
    torch.cuda.set_device(0)
    dist.init_process_group("nccl", rank=rank, world_size=2)
    bus = getattr(torch.cuda.get_device_properties(0), "pci_bus_id", "?")
    print(f"[tp] rank{rank} 已加入 RCCL，device={torch.cuda.current_device()} bus={bus}", flush=True)
    READY.set()

    import comfy.sd as sd
    orig = sd.load_diffusion_model_state_dict

    def wrapped(sd=None, *a, **k):
        READY.wait()
        import core
        presliced = False
        if isinstance(sd, dict):
            n0 = sd.get("blocks.0.attn.qkv_proj.weight")
            if n0 is not None and tuple(n0.shape) == (21504, 5376):
                sd = core.slice_state_dict(sd, rank)   # 预切：模型按半宽构造
                presliced = True
        try:
            import comfy.model_management as _mm
            _mm.unload_all_models()                    # 先腾出 26G 文本编码器
            import torch as _t
            _t.cuda.empty_cache()
        except Exception as e:
            print(f"[tp] 预卸载失败: {e}", flush=True)
        mp = orig(sd, *a, **k)
        if presliced and mp is not None:
            net = getattr(getattr(mp, "model", None), "diffusion_model", None)
            if net is not None:
                net._tp_presliced = True
        return _on_model(mp)

    sd.load_diffusion_model_state_dict = wrapped
    print(f"[tp] rank{rank} 已挂载模型加载钩子", flush=True)


print(f"[tp] tp2-comfy v{APP_VERSION} rank{rank} 启动（COMFY_ROOT={ROC}）", flush=True)
threading.Thread(target=_boot, daemon=True).start()
import runpy
runpy.run_path(os.path.join(ROC, "main.py"), run_name="__main__")

tp/tp_linear.py —— int8 行并行：共享激活 scale

"""TP 行并行层：只替换激活量化 scale，其余完全走 ComfyUI-INT8 原生路径。

背景：int8 权重按输出行切好（列并行）后，行并行（K 轴切半）的问题在于
原生 per-row 动态量化会给每半各算一个 scale，量化点不同 → 每层 ~1% 误差。

做法：把各行 |x| 最大值跨卡 all_reduce(MAX)，得到整行原本会用的 scale，
注入 triton_quantize_rowwise；其余（int8 GEMM、反量化、convrot、LoRA、
cast_weight 分支）全部走原生，不做任何重写。
"""
import torch
import torch.distributed as dist

_SHARED = {"scale": None}
_PATCHED = {"ok": False}


def _patch_once():
    if _PATCHED["ok"]:
        return
    import importlib
    fk = importlib.import_module("ComfyUI-INT8-Fast-ROCM.int8_fused_kernel")
    orig = fk.triton_quantize_rowwise

    def qr(x):
        y, s = orig(x)
        sh = _SHARED["scale"]
        if sh is None:
            return y, s
        xf = x.float()
        y = torch.floor(xf / sh + 0.5).clamp_(-128, 127).to(torch.int8)
        return y, sh

    fk.triton_quantize_rowwise = qr
    _PATCHED["ok"] = True


def install_row_parallel(lin):
    """权重需已按列（K 轴）切好；安装带全局 scale 的原生前向。"""
    _patch_once()
    if getattr(lin, "_tp_row_installed", False):
        return
    orig_forward = lin.__class__.forward

    def fwd(x):
        # 注意：只在"算 scale"时旋转一份副本，传进原生 forward 的仍是未旋转的 x，
        # 否则原生路径会再旋转一次（H 不是幂等）。
        x2d = x.reshape(-1, x.shape[-1])
        xr = x2d
        if getattr(lin, "_use_convrot", False):
            import importlib
            cr = importlib.import_module("ComfyUI-INT8-Fast-ROCM.convrot")
            gs = getattr(lin, "_convrot_groupsize", 256)
            if x2d.shape[1] % gs == 0:
                H = cr.build_hadamard(gs, device=x2d.device, dtype=x2d.dtype)
                xr = cr.rotate_activation(x2d, H, group_size=gs)
        forced = getattr(lin, "_tp_force_scale", None)
        if forced is not None:
            s = forced
        else:
            s = (xr.float().abs().amax(dim=1, keepdim=True) / 127.0).clamp_min(1e-30)
        if dist.is_initialized():
            dist.all_reduce(s, op=dist.ReduceOp.MAX)
            s = s.clamp_min(1e-30)
        # TP_WIRE=bf16（默认，对齐 Radiance 的 exact 线）| fp32（更精确但通信翻倍）
        import importlib
        import os
        ops = importlib.import_module("ComfyUI-INT8-Fast-ROCM.int8_quant").Int8TensorwiseOps
        prev = ops.compute_dtype
        if os.environ.get("TP_WIRE", "bf16") == "bf16":
            out_dtype_keep = prev
        _SHARED["scale"] = s
        if os.environ.get("TP_WIRE", "bf16") != "bf16":
            ops.compute_dtype = torch.float32
        try:
            out = orig_forward(lin, x)
        finally:
            _SHARED["scale"] = None
            ops.compute_dtype = prev
        return out

    lin.forward = fwd
    lin._tp_row_installed = True

start_tp2.sh —— 启动器（默认 CK）

#!/usr/bin/env bash
# start_tp2.sh —— 起两个 ComfyUI 实例做 H3 张量并行（rank0 / rank1），各只看见一张卡。
#
# 必须带 --use-ck-attention：PyTorch SDPA 在长序列上会崩速（见 README「四个必须」）。
#
# 可用环境变量覆盖：
#   COMFY_ROOT   ComfyUI 根目录            （默认 /home/ly/ComfyUI-ROCm）
#   GPU0/GPU1    两张卡的 unique_id        （rocm-smi --showuniqueid 查）
#   PORT0/PORT1  两个实例的 HTTP 端口      （默认 8190 / 8191）
#   BASE0/BASE1  两个实例的 output/input 目录
#   MASTER_PORT  RCCL 会合端口             （默认 29670）
#   ATTN         attention 后端参数        （默认 --use-ck-attention）
set -euo pipefail

COMFY_ROOT=${COMFY_ROOT:-/home/ly/ComfyUI-ROCm}
TP_DIR=$(cd "$(dirname "$0")" && pwd)/tp
GPU0=${GPU0:-GPU-3c99737f13093cd0}
GPU1=${GPU1:-GPU-3ce42265b99a6c53}
PORT0=${PORT0:-8190}
PORT1=${PORT1:-8191}
BASE0=${BASE0:-/home/ly/ComfyUI Tools/comfyui/w7900}
BASE1=${BASE1:-/home/ly/ComfyUI Tools/comfyui/w7900-4a}
MASTER_PORT=${MASTER_PORT:-29670}
ATTN=${ATTN:---use-ck-attention}
LOGD=${LOGD:-$COMFY_ROOT/logs}

mkdir -p "$BASE0"/{output,input,temp} "$BASE1"/{output,input,temp} "$LOGD"
export TRITON_DEFAULT_BACKEND=amd COMFYUI_ENABLE_MIOPEN=1 MIOPEN_FIND_MODE=FAST
export MALLOC_MMAP_THRESHOLD_=65536 MALLOC_TRIM_THRESHOLD_=65536
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 SAFETENSORS_FAST_GPU=1 PYTHONUTF8=1
export PYTORCH_TUNABLEOP_ENABLED=0
export MASTER_ADDR=127.0.0.1 MASTER_PORT
export COMFY_ROOT

# rank1（第二张卡）先起：它会 block 在 init_process_group 等 rank0
TP_RANK=1 ROCR_VISIBLE_DEVICES="$GPU1" nohup "$COMFY_ROOT/venv/bin/python" "$TP_DIR/tp_boot.py" \
  --listen 127.0.0.1 --port "$PORT1" \
  --output-directory "$BASE1/output" --input-directory "$BASE1/input" \
  $ATTN --disable-smart-memory --disable-pinned-memory --disable-async-offload --reserve-vram 2 \
  >> "$LOGD/tp_rank1_$(date +%Y%m%d).log" 2>&1 &
echo "rank1 pid=$! port=$PORT1 gpu=$GPU1"

sleep 3
TP_RANK=0 ROCR_VISIBLE_DEVICES="$GPU0" nohup "$COMFY_ROOT/venv/bin/python" "$TP_DIR/tp_boot.py" \
  --listen 127.0.0.1 --port "$PORT0" \
  --output-directory "$BASE0/output" --input-directory "$BASE0/input" \
  $ATTN --disable-smart-memory --disable-pinned-memory --disable-async-offload --reserve-vram 2 \
  >> "$LOGD/tp_rank0_$(date +%Y%m%d).log" 2>&1 &
echo "rank0 pid=$! port=$PORT0 gpu=$GPU0"

echo "已拉起 TP2：$PORT0(rank0) + $PORT1(rank1)，日志 $LOGD/tp_rank{0,1}_$(date +%Y%m%d).log"

[懒人烘培]

@fcys 谢谢啦，可能我用那gpt有些傻，没有搞不来，后续要抄作业：）

[Xiaote]

版主认可：感谢对双卡 TP 跑 H3 的实测验证，每步 18.6s → 8.3s 的数据很有说服力，对想上双卡的锤友帮助很大。

奖励 +5 积分。

[terry]

其实奖励给少了，这个帖子的意义很高，以后有机会再补吧，奖励发出就不好修改，同类奖励限定了1次。