ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

LTX-2 多卡序列并行实战指南:一个 Denoising Step 如何用 all2all 内核摊到多张 GPU

LTX-2 多卡序列并行实战指南:一个 Denoising Step 如何用 all2all 内核摊到多张 GPU LTX-2 多卡序列并行实战指南一个 Denoising Step 如何用 all2all 内核摊到多张 GPU【免费下载链接】LTX-2Official Python inference and LoRA trainer package for the LTX-2 audio–video generative model.项目地址: https://gitcode.com/GitHub_Trending/lt/LTX-2LTX-2 的序列并行SP把视频 token 维均匀切到多张 GPU用自定义 all2all 内核在卡间交换注意力头让每个 denoising step 以更低延迟完成且结果与单卡推理数值等价读完本文你能讲清 SP 为什么忠实并知道如何把AttentionManager与SequenceParallelBuilder接进自己的多卡 runner。 多卡推理选哪个 stage序列并行的场景与选型单卡跑 LTX-2 时一个 denoising step 的耗时随视频 token 数线性增长——分辨率和帧数一上来单步延迟就是整条管线的瓶颈。SP 的解法很直接每张卡只算自己负责的那一段序列注意力部分靠卡间换头补齐。先钉死大前提MGPU 是延迟工具不是显存工具multigpu/README.md L5-L25。transformer 的工作副本在每张 GPU 上都是完整副本SP 额外做的是把激活内存分摊到各 rank它救不了模型放不进单卡那要靠 FP8 量化与权重 offload。四种多卡策略横向对比multigpu/README.md L36-L45 的 Capabilities 表策略切分维度定位适用场景序列并行 SPtoken 序列维忠实与单卡数值等价训练分布内分辨率、要求结果对齐单卡Tiled Data Parallelism每卡一个空间 tile面向训练分布外分辨率仅 upscale全分辨率放大阶段分布式解码器latent tile 并行解码加速 VAE decodeVAE 阶段分布式 GemmaAcceleratedevice_map切分文本编码器并行Gemma 阶段默认走 SP 的有两处sequence-parallel.md L20-L21stage 1ti2vid_two_stages_mgpu的低分辨率阶段shared stagedistilled_mgpu的共享阶段——同一个 SP 包裹同时覆盖 half-res 与 full-res 两次调用。而 two-stage 管线的 stage 2全分辨率官方选的是 TDP因为它面向训练分布外分辨率tiled-data-parallel.md。结论一句话分辨率在训练分布内、要求与单卡结果一致选 SP做分布外分辨率 upscale选 TDP。⚙️ 一个 denoising step 怎么被拆开SP 机制拆解每个 step 的执行主体是 SequenceParallelModelWrapper 的forwardsequence_parallel.py L290-L318。一个 step 的流转pad → 切片 → 换头 → 还原Pad 对齐pad_modality_for_uniform_sharding把视频 seq 维补齐到world_size的整数倍。若原本没有 attention mask构造一个 key-only padding maskshape(1, 1, T_padded)、沿 batch 和 query 广播——O(T)内存不必物化稠密的(B, T, T)矩阵若用户传了(B, T, T)mask则扩展 pad 行/列且pad 的 query 行可 attend 所有有效 key防止全 mask 行在 softmax 里产生 NaN输出反正会被切掉。本 rank 切片tile_modality_for_rank把 latent / timesteps / positions 切到本 rank 的 slice。切分依据compute_sequence_partition要求total_tokens % world_size 0否则直接ValueError。均匀切分有个隐藏收益all2all 自定义算子的 fake-impl 能从输入 shape 符号化推导输出 shapex.shape[1] * world_size/x.shape[1] // world_size而不是依赖 Python int 参数。运行模型视频自注意力attn1与 video→audio 交叉注意力已被替换成走 all2all 内核的版本本 rank 只在自己的 token 切片上跑前向。all-gather 还原gather_output_tokens先把本地输出 pad 到最大长度做均匀all_gather再按各 rank 真实 token 数裁回并拼接最后去掉第 1 步加的 pad 行恢复调用方传入的原始长度。每步的通信量管理在 step 边界完成set_seqlen_all2all把各 rank 的 token 数下发到 C 运行时随后一次torch.distributed.barrier保证所有 rank 对齐后才进模型sequence_parallel.py L310-L311。all2all 内核IPC 直写、SM 轮转与 barrier内核是ltx_kernels.All2All用CUDA-IPC peer buffer在各 rank 的 GPU 之间直接交换 token。all2all_heads.cu 头部注释把算法讲得很透L10-L30直写每个 GPU 通过 IPC 把数据直接写到目标 GPU 的内存 buffer没有中间拷贝接近峰值内存带宽。SM 轮转分配SM 按SM i 写 rank (i % world_size)分给目标 rank专门处理 SM 数不可整除的情况——文档举的例子132 个 SM、8 张卡时rank 0–3 各分到 17 个 SMrank 4–7 各 16 个每个 SM 组负责搬运分配给目标 rank 的全部 token。同步协议数据搬完后各 SM 原子递增目标 rank 的 barrier 计数器SM 0 等齐所有 rank 的信号后重置计数器供下一轮使用allgather.cu 里是带超时周期timeout_cycles的轮转变体。一个通俗类比把每张 GPU 想成邮局的分拣员手里攥着一批包裹token每个包裹上标着收件分区head 归属的 rank。all2all 干的事就是按收件分区把包裹重新分拣直写相当于分拣员直接把包裹投进对方邮局的格口IPC 直写不经过中转场SM 轮转相当于 132 名分拣员轮值 8 个邮局前 4 个邮局排 17 人、后 4 个排 16 人barrier 就是分拣完各邮局举手确认全齐了才开下一班车。分拣完后每个邮局恰好收到寄给自己的全部包裹——也就是本 rank 负责的那部分 head 上的所有 token。Python 侧用torch.library.custom_op注册了ltx_kernels::send_recv_heads与ltx_kernels::gather_headsall_to_all.py L30-L50好处是torch.compile含modereduce-overhead的 CUDA Graph 捕获能无 graph break 地 trace 过去。world_size作为 int 常量进入 traced graphDynamo 的 guard 只按 GPU 数量键控编译缓存——同一个图绝不会在另一个 GPU 数量下被重放而每个 step 都会变的 per-rank token 数留在 C 运行时状态set_rank_tokens里不经过算子all_to_all.py L21-L27。另外copy_outFalse时返回的是 IPC buffer 的零拷贝视图因为 IPC buffer 由cudaMalloc分配、不在静态图池内在 cudagraph_trees 下也安全L37-L39实例销毁由weakref.finalize自动释放 CUDA/IPC 资源L122-L128。数值等价为什么成立全局注意力没丢all2all 只是把 token×head 的二维数据做一次分布变换。交换之后每个 rank 持有所有 token × 本 rank 的那部分 head本地算注意力token 之间的交互是完整的——没有任何一对 token 被跳过。往返逐字节精确all2all 内核只搬运字节gather(send(x)) x逐字节成立sequence-parallel.md L14-L18。唯一差异是浮点归约顺序不同硬件分布下加法顺序变化会带来极小的浮点差但这属于并行计算常态不属于改变了模型。SP 忠实因为它只改变数据的分布方式从不改变模型行为。需求若是更低延迟 与单卡完全一致的结果SP 就是唯一正确的选择。 SP 接入指南AttentionManager、Builder 与 runner 三处改动两个核心 API 长什么样from ltx_core.multigpu.transformer.attention import AttentionManager attn_mgr AttentionManager( max_tokens: int, # 视频总 token 数上界超出报错 num_heads: int, # transformer.num_attention_heads head_dim: int, # transformer.attention_head_dim tensor_dtype: torch.dtype, group: dist.ProcessGroup, # self.groups.transformer_group copy_out_: bool False, )AttentionManager持有 all2all buffer每 rank 大小为ceil(max_tokens / world_size)个 tokensequence-parallel.md L57-L59。几个构造细节值得记num_heads必须能被world_size整除_All2AllRedistribute.redistribute里显式校验并抛ValueErrorattention.py L112-L113内部创建 4 个All2All实例q / k / v / headscopy_out_True时 k、v 与 q 共用实例attention.py L33-L58ltx_kernels是惰性 import的保证 multigpu 模块在未装内核的机器如 CPU CI上仍可导入attention.py L21-L24all2all_timeout_seconds默认 10.0s对应内核configs.cuh的DEFAULT_BARRIER_TIMEOUT_SECONDSattention.py L6-L8。setter 注释给出应对策略torch.compile首次前向时某个 rank 重编译可能让内核启动晚于稳态超时、触发 barrier所以首跑前临时调大、之后复位attention.py L75-L84。from ltx_pipelines.multigpu.sp_builder import SequenceParallelBuilder SequenceParallelBuilder( inner: ModelBuilderProtocol, # 该 stage 的单 GPU transformer builder attn_mgr: AttentionManager, registry: Registry, tracker: TransformerWeightTracker, )它是包裹型 builder只接受SingleGPUModelBuilder否则抛TypeErrorsp_builder.py L35-L36。构造时注入 registry 与 LoRA 加载设备再追加 SP 的 module-opsL37-L42build()时经TransformerWeightTracker构建模型外层包一层SequenceParallelModelWrapper返回L58-L62。由于是包裹而不是重建它继承 inner builder 的 checkpoint 路径、量化、编译与 LoRA 配置只叠加并行。module-ops 如何替换两个注意力槽位create_video_self_attention_module_ops在LTXModel上匹配每个BasicAVTransformerBlock改两处视频自注意力attn1attention_function换成All2AllAttentionmasked_attention_function换成MaskedAll2AllAttention。两个槽位都换是因为Attention.forward可能收到非 None 的video.self_attention_mask运行时按实际调用走对应分支sequence_parallel.py L246-L256。video→audio 交叉注意力video_to_audio_attn换成AudioAll2AllAttention/MaskedAudioAll2AllAttention。目前没有调用方传 maskmasked 分支是死代码但槽位照样预装——将来有人加 maskSP 管道已经就位不会静默绕过 All2Allsequence_parallel.py L257-L270。两类洗牌的差异值得分清attention.py视频自注意力Q/K/V 全走send_recv_heads本地只算local_heads heads // world_size个 head再gather_heads换回L137-L151音频交叉注意力Q 在本地按 rank 切片、不做跨 rank 洗牌音频序列足够短可以复制只有 K/V 走send_recv_heads输出沿 head 维all_gather_into_tensor收集L172-L222。runner.setup() 里动哪三处# runner.setup() 内按 stage 执行 model_cfg pipeline.stage_1._transformer_builder.model_config().get(transformer, {}) attn_mgr AttentionManager( max_tokens32768, num_headsmodel_cfg[num_attention_heads], head_dimmodel_cfg[attention_head_dim], tensor_dtypepipeline.dtype, groupself.groups.transformer_group, ) pipeline.stage_1._transformer_builder SequenceParallelBuilder( innerpipeline.stage_1._transformer_builder, attn_mgrattn_mgr, registryregistry, trackertracker, )三处改动对应三件事读配置从model_config()拿num_attention_heads/attention_head_dim不要手写魔数建管理器AttentionManager绑定self.groups.transformer_group——该组由 NCCLGroups 用dist.new_group(ranksall_ranks)创建换 builder用SequenceParallelBuilder替换该 stage 的._transformer_builder。registry是进程内所有 builder 共享的一个ModelRegistrycheckpoint 每进程只从磁盘加载一次tracker是对应transformer_group的TransformerWeightTracker。这正是 pipeline-setup.md 描述的 MGPU 模式MGPU 管线 单 GPU 管线 替换各 block 的 builder改完这一行并行就叠加上了。️ max_tokens 怎么估、排障看哪些报错max_tokens 上界的三个参考量max_tokens必须覆盖最大的那个 step三个官方数字可直接抄sequence-parallel.md L97-L100场景分辨率/帧数视频 token 数源码出处stage 1512×768×1216144ti2vid_two_stages_mgpu.py L45distilled full-res1024×1536×121~24576distilled_mgpu.py L43两者默认上界—32768_DEFAULT_SP_MAX_TOKENSL46超限时的报错原文sequence_parallel.py L305-L309Total video token count (...) exceeds attention_manager max_tokens (...). Use a smaller resolution or fewer frames.解法二选一按提示降分辨率/减帧数或显式调大sp_max_tokens——两个 MGPU runner 的setup()都接受该参数ti2vid_two_stages_mgpu.py L64、distilled_mgpu.py L61。另外两个常见报错heads (...) must be divisible by world_size (...)说明你的卡数不整除注意力头数attention.py L112-L113SequenceParallelBuilder wraps a SingleGPUModelBuilder, got ...说明你拿去包裹的不是单卡 buildersp_builder.py L35-L36。运行前提与不适用清单SP 依赖 NCCL 与 CUDA-IPC peer buffer硬前提四条multigpu/README.md L27-L34仅LinuxNCCL 与 CUDA-IPC 均不支持 macOS/Windows单节点 ≥2 张支持 P2PNVLink/PCIe的 CUDA GPU不支持多节点带 CUDA 的 PyTorchltx-kernels已构建uv sync --group kernels需要 CUDA toolkit/nvcc 与 gcc 或 clang。SP builder 的 docstring 也明确写着 Requiresltx-kernelsto be installedsp_builder.py L1-L4。不适用的场景训练分布外分辨率的 upscale用 TDP、把放不下的模型塞进单卡用量化 offload。SP 的边界分布内分辨率 忠实性要求两个条件都满足才轮到它。 跑通 SP两个官方 runner 的关键 CLI 参数端到端体验 SP 只需两个 runnermultigpu/README.md L55-L58ltx_pipelines.ti2vid_two_stages_mgpuSP 用于 stage 1 TDP 用于 stage 2 Accelerate Gemma 分布式 VAEltx_pipelines.distilled_mgpuSP 覆盖 shared stage同一次包裹吃掉 half-res 与 full-res 两次调用 Accelerate Gemma 分布式 VAE。执行架构MGPUController每 GPU 拉起一个 worker 进程以 SPMD lockstep 跑同一份 runner 代码multigpu/README.md L47-L53。checkpoint 路径与单卡 CLI 相同的 monolith XOR split 约定要么传 fat checkpoint Gemma 目录要么传 split 组件 flags不能混用L60-L63。# two-stage 管线跑在全部可见 GPU 上monolith 布局 python -m ltx_pipelines.ti2vid_two_stages_mgpu \ --checkpoint-path path/to/checkpoint.safetensors \ --distilled-lora path/to/distilled_lora.safetensors 1.0 \ --spatial-upsampler-path path/to/upsampler.safetensors \ --gemma-root path/to/gemma \ --prompt A beautiful sunset over the ocean \ --output-path output.mp4关键参数读法--distilled-lora后跟强度1.0--spatial-upsampler-path是 stage 2 upscale 用的模型split 布局下则换成--transformer-path/--video-vae-path/--audio-vae-path等组件 flagsmultigpu/README.md L75-L85。distilled runner 入口是python -m ltx_pipelines.distilled_mgpu参数族与上面一致prompt、seed、height/width、num_frames、frame_rate、num_inference_steps、output-path。 要点速记SP token 维均匀切分 all2all 换头 all-gather 还原注意力保持全局往返逐字节精确唯一变化是浮点归约顺序——这就是它数值等价的根源sequence-parallel.md L14-L18。MGPU 是延迟工具不是显存工具每张卡都是完整模型副本SP 只额外分摊激活内存multigpu/README.md L5-L25。接入只有三处改动建AttentionManager、用SequenceParallelBuilder包裹单卡 builder、在runner.setup()里换掉._transformer_buildersp_builder.py。max_tokens 按最大 step 估官方参考 6144 / ~24576默认 32768超限报错 Use a smaller resolution or fewer frames.可降规格或调大sp_max_tokens。跑 SP 的硬门槛Linux、单节点 ≥2 张 P2P 卡、uv sync --group kernels构建内核、num_heads % world_size 0all_to_all.py、attention.py L112-L113。【免费下载链接】LTX-2Official Python inference and LoRA trainer package for the LTX-2 audio–video generative model.项目地址: https://gitcode.com/GitHub_Trending/lt/LTX-2创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表