LTX-2 多 GPU 推理控制器(MGPUController)完全指南:持久化工作舰队、SPMD 调度与流式结果回收
LTX-2 多 GPU 推理控制器MGPUController完全指南持久化工作舰队、SPMD 调度与流式结果回收【免费下载链接】LTX-2Official Python inference and LoRA trainer package for the LTX-2 audio–video generative model.项目地址: https://gitcode.com/GitHub_Trending/lt/LTX-2导读本文面向希望在单机多卡如 H100、B200上降低 LTX-2 单次生成延迟的开发者系统讲解ltx-pipelines中多 GPUMGPU推理基础设施的控制面——MGPUController与MGPURunner。读完本文你将掌握如何构建一个按 GPU 数量常驻的持久化工作进程舰队、如何以 SPMD 锁步方式提交并流式回收任务结果、如何正确传递张量与处理各类错误以及两个官方 runner两阶段与蒸馏管线是如何把控制器接入标准 CLI 的真实范式。本文以 controller.md 为骨架并结合 controller.py、runner.py、fleet.py 等源码逐层印证。背景速览LTX-2 的 MGPU 模块是降低延迟的工具而非扩充显存的工具——transformer 的可变工作副本在每张 GPU 上都是完整副本无法让装不下的模型变得装得下其目标是靠序列并行SP、分块数据并行TDP、分布式解码与分布式 Gemma 把每一步去噪和 VAE 解码分摊到多卡上见 多 GPU 总览。本文只讲驱动这一切的控制中枢并行技术的细节请移步 sequence-parallel.md、tiled-data-parallel.md 与 pipeline-setup.md。MGPUController一次常驻、一次一任务的 GPU 舰队MGPUController的本质是一个持久的、同一时刻只运行一个任务one-job-at-a-time的 GPU 舰队。它按 GPU 数量为每个 GPU 派生一个 worker 进程运行用户定义的runnerMGPURunner的子类并以 SPMD 锁步lockstep方式执行任务再把结果流式送回调用方。其构造函数签名如下controller.pyMGPUController( runner_cls: type[MGPURunner], *, num_gpus: int | None None, # GPUs 0..num_gpus-1默认所有可见 GPU devices: Sequence[int] | None None, # 放置到指定的物理 GPU 上例如 [2, 3] logs_specs: LogsSpecs | None None, )两个选址参数的含义与约束num_gpus使用0..num_gpus-1号 GPU为None时取torch.cuda.device_count()即所有可见 GPU源码中start()内会解析这一默认值见 controller.py。devices把舰队放到指定的物理 GPU 子集上rank r 运行在cuda:devices[r]。这样两个控制器可以在一台机器上、各占一组互不重叠的 GPU 并行运行。num_gpus与devices互斥。源码中两者同时传入会直接抛ValueError(Pass either num_gpus or devices, not both.)devices还要求非空且元素唯一controller.py。一个容易被忽略的实现细节worker 进程中所有 GPU 仍然可见每个 rank 只是通过torch.cuda.set_device(...)绑定到自己的物理 GPU 上而非设置CUDA_VISIBLE_DEVICES——源码注释解释了原因spawn 引导过程可能已触碰 CUDA、冻结了设备列表晚到的环境变量覆盖并不可靠fleet.py。生命周期start / stream / shutdown / is_alive方法 / 属性作用start(*, timeout30min, **setup_kwargs)派生舰队在每个 rank 上运行setup(**setup_kwargs)阻塞直到所有 rank 上报就绪。timeout约束 NCCL 初始化 CUDA 初始化 setup()的总时长它必须超过最慢的模型加载时间。stream(*, timeoutNone, **kwargs) - Stream派发一个任务并立即返回。迭代返回的Stream即可收集结果。shutdown(*, graceful_timeout60.0)拆除舰队同时也是强制终止开关——可以从其他线程安全调用用于回收一个无法 drain 的任务。is_alive属性舰队在线且未被毒化poisoned时为True是一个不需要 try/except 的健康检查。关于start()的timeout源码里有更精确的语义值得展开_await_ready用time.monotonic()计算截止时间循环里先fleet.poll()检查是否有 worker 进程退出能捕获崩溃但poll()只能看到进程退出、看不到卡死再用截止时间兜底捕获那些活着但永远不上报就绪的卡死 worker抛出带缺失 rank 列表的TimeoutError提示卡在init_process_group / create_local_nccl_groups / setup()应查看 worker 日志。因此想要无限等待就传一个很大的 timeoutcontroller.py。另外start()内部若在等待就绪期间出现任何异常会自动调用shutdown()拆掉半启动的舰队调用方无需再清理。stream()的timeout单位为秒与start()不同它是按任务计时的从派发时刻起算worker 死亡或超时并不会在stream()返回时立刻暴露而是在你下一次迭代 Stream 时才以异常形式浮出水面controller.py。shutdown()的完整行为是先向任务队列放入None哨兵relay 会把哨兵广播给所有 rank随后drain(graceful_timeout)等待 worker 自行退出超时未退则terminate()SIGTERM → SIGKILL永不抛错、可重复调用。任务进行中调用shutdown()会先等完graceful_timeout再强杀传0可跳过等待controller.py、fleet.py。MGPURunner每个 pipeline 实现一次的抽象基类MGPURunner是每个 pipeline 需要继承实现的抽象基类。控制器把 runner 的子类按值by value派发给每个 worker——这意味着定义在__main__或测试模块里的 runner 同样受支持然后每个 worker 构建一个实例、注入 NCCL 组、调用一次setup()之后按任务调用该实例。class MyRunner(MGPURunner): torch.inference_mode() def setup(self, *, checkpoint_path: str, ...) - None: # 构建 pipeline并把各 block 的 builder 换成 MGPU builder见 pipeline-setup.md ... torch.inference_mode() def __call__(self, *, prompt: str, ...) - Iterator[...]: video, audio self._pipeline(...) yield output_path # __call__ 必须是生成器即使只 yield 一次也要用 yield核心契约runner.pysetup()与__call__()在每个 rank上都会执行。self.groups提供各组件对应的NCCLGroupsgemma_group、transformer_group、vae_group。这些组由 nccl_groups.py 中的create_local_nccl_groups()在 worker 内集体创建——dist.new_group是集合操作因此所有 rank 都要参与、且每个组都包含全体 rank。框架不会自动套用推理模式inference mode需要你显式给setup/__call__加上torch.inference_mode()装饰器。__call__必须是生成器至少yield一次每个yield都会作为独立元素被转发给 Stream且不会跨 rank 聚合。若 runner 不是生成器next()会抛TypeError并被视为致命错误fleet.py。结果是通过yield流出的return值罕见路径随终止信号StopIteration.value按 rank 收集——通常都是None因为结果都是 yield 出来的。想读取它用result yield from stream。runner 如何按值漂洋过海_RunnerShipperfleet.py中的_RunnerShipper是支撑定义在__main__也能跑的关键。普通的按引用 pickle 在全新 spawn 的 worker 中无法按模块名导入__main__里的类会抛ModuleNotFoundError。_RunnerShipper改用cloudpickle.dumps把 runner 类按值序列化并强制register_pickle_by_value其__reduce__指向全局可导入的cloudpickle.loads因此只有 runner 本身按值跨进程控制器自己的队列载荷仍走标准 picklerfleet.py。完整使用范式from ltx_pipelines.multigpu import MGPUController controller MGPUController(MyRunner, num_gpus8) controller.start(model_pathsModelPaths.from_monolith(..., ...)) # setup kwargs stream controller.stream(prompta cat, seed42) try: for item in stream: # 每个元素对应一次 yield按到达顺序返回不会跨 rank 聚合 show(item) finally: stream.drain() # 即使提前退出也要释放控制器 controller.shutdown()这个try: ... finally: stream.drain()模式是官方文档反复强调的推荐写法controller 源码 docstring 与Stream.drain的注释中均有。drain()会从派发线程耗尽 Stream、丢弃未消费的 yield并把可恢复的Symmetric/AsymmetricRunnerError吞掉纯清理语义但 worker 死亡 / 超时 / 舰队失步这类致命错误仍然会浮出因为调用方必须知情controller.py。官方 runner 的接线示范两个随仓库提供的 runner 各自带 CLI它们的__main__就是控制器最标准的接线模板ti2vid_two_stages_mgpu.pyTI2VidTwoStagesRunnerSP 阶段 1 TDP 阶段 2 Accelerate Gemma 分布式 VAEdistilled_mgpu.pyDistilledRunnerSP 共享阶段 Accelerate Gemma 分布式 VAE以TI2VidTwoStagesRunner为例其__main__的关键结构是vae_queue torch.multiprocessing.get_context(spawn).SimpleQueue() controller MGPUController(TI2VidTwoStagesRunner) controller.start( model_pathsargs.model_paths, prompt_enhancer_gemma_rootargs.prompt_enhancer_gemma_root, spatial_upsampler_pathargs.spatial_upsampler_path, vae_queuevae_queue, distilled_lora_pathargs.distilled_lora[0].path, compilation_configargs.compile, diffvae_optimizationargs.diffvae_optimization, ) try: for _ in controller.stream( output_pathargs.output_path, promptargs.prompt, seedargs.seed, heightargs.height, widthargs.width, num_framesargs.num_frames, frame_rateargs.frame_rate, video_guider_paramsMultiModalGuiderParams(...), audio_guider_paramsMultiModalGuiderParams(...), imagesargs.images, enhance_promptargs.enhance_prompt, hdrargs.hdr, generated_keyframesargs.num_generated_keyframes, ): pass # 驱动任务跑完runner 把视频文件作为副作用写出 finally: controller.shutdown()注意两处细节vae_queue是一个spawn 上下文的SimpleQueue用于把分布式 VAE 的解码结果传回 driver rank 组装主循环中pass只是驱动任务到结束真正的产出是 runner 内部encode_video(...)写出的文件driver rank 在__call__末尾yield output_path非 driver rankyield None。命令行的完整选项monolith 整体权重 vs 拆分包两种 checkpoint 布局见 多 GPU 总览。张量透明输入与输出的零 pickle 通道张量在 MGPU 中是透明的这是整个框架最实用也最容易被误用的特性输入把张量作为顶层 kwargs传递stream(latentt, steps30)。relayrank 0会把它通过共享内存 / CUDA IPC 从控制器送到 relay 的 GPU再由 relay 用NCCL 广播到每个 rank——__call__收到的张量已经在本地 GPU 上。如果一个张量嵌套在 list/dict 参数里则不会走 NCCL 广播而是退回到更慢的pickle 路径。输出从任意 rankyield张量包括嵌套在 dict 里都会通过结果队列按共享内存 / CUDA IPC返回——无论嵌套多深都不做 pickle。其余所有数据都必须是可 pickle 且小的。源码印证了这套机制的实现fleet.py_replace_tensors_by_placeholders只处理顶层Tensor 值把它们替换成轻量的_TensorPlaceholder仅携带idx / shape / dtype可 pickle张量本体单独拎出relay 端announce_job先dist.broadcast_object_list广播骨架含占位符再按同一顺序对每个张量执行dist.broadcastNCCL 要求 CUDA contiguous因此先t.to(device).contiguous()非 relay 端await_job按占位符的idx顺序预分配torch.empty(shape, dtype, device)接收广播再_fill_tensors_into_placeholders还原输出方向yield 的_JobResult只走结果队列张量随队列的共享内存 / CUDA IPC 传输全程无二次 pickle。空闲舰队如何不触发 NCCL watchdogBroadcastCoordinator一个易被忽视但很重要的细节任务与任务之间非 relay 的 worker 如果直接阻塞在下一个 broadcast 集合操作上空闲期过长会触发NCCL watchdog 超时。解决方案是_broadcast.py中的BroadcastCoordinator它利用进程组的分布式 store 里一个递增计数器PrefixStore(ltx_pipeline_broadcast/, ...)做信令——worker 以 10ms 间隔轮询计数器变化relay 在每次广播前把计数器加 1 写回 store只有当 relay 发出有任务要发的信号后worker 才真正进入 NCCL 集合broadcast.py。这样空闲的舰队保持廉价轮询、绝不会卡死在 pending broadcast 上。契约与限制控制器的行为建立在一组严格的契约之上违反任何一条都会以特定方式失败仅限单机MASTER_ADDRlocalhost实际为127.0.0.1、RANK LOCAL_RANK、每 GPU 一个 rank不支持多节点。spawn 时由_find_free_port()动态挑选空闲端口写入各 worker 环境变量fleet.py。一次一个任务没有任务队列、没有流水线。必须把Stream消费到底才能发起下一个stream()。放弃任务不会被自动清理在下一次stream()时抛ControllerBusyError直到stream.drain()或shutdown()被调用。推荐try: ... finally: stream.drain()。ControllerBusyError携带在途任务的job_id便于日志定位controller.py。SPMD 锁步各 rank 的 yield逐个转发按结果队列顺序不会聚合只有各 rank 的终止信号terminal会被收集全部到齐后流结束。线程所有权baton-lock任何线程都可以调用stream()但每个任务归属于派发它的那个线程——只有该线程可以迭代或drain()它的StreamStream.__next__中通过threading.get_ident()强制校验跨线程迭代直接抛RuntimeError。MGPUController内部只有一个threading.Lock且仅守护在途任务的检查并置位临界区在并发的stream()调用者中一个成功进入、其余抛ControllerBusyError。锁绝不在派发 / 迭代 / 收集期间持有controller.py。此外还有一个失步desync检测每个_Job携带单调递增的job_id任何 rank 回传的job_id与派发的不一致即判定舰队失步并毒化控制器controller.py。错误处理哪些可恢复、哪些必须重建情况结果Runner 抛出非预期异常致命——已失步的 NCCL 集合无法挽回。控制器被毒化必须新建一个控制器。Runner 在每个 rank 上一致地抛出RunnerError或ValueError会被自动转换可恢复。迭代Stream时会重新抛出SymmetricRunnerError舰队存活——修正输入后重试即可。注意要在任何集合操作之外抛出例如在第一个广播之前校验广播 kwargs 时。部分 rank 抛RunnerError、其余正常结束AsymmetricRunnerError——会被显眼地浮出因为它潜在地预示挂起风险但不会终止舰队。Worker 死亡 / 超过单任务timeout在下一次迭代Stream时浮出控制器被毒化。补充几个源码级细节让这张表更可操作自动转换worker 循环里_run_job对ValueError捕获并合成RunnerError(str(err))fleet.py——这正是官方 runner 注释里pipeline 对非法输入抛ValueError跨 rank 对称控制器将其转成可恢复的RunnerError其余任何异常都是致命的依据见 ti2vid_two_stages_mgpu.py。分类逻辑_classify_terminal把每个 rank 的终止信号干净的StopIteration或RunnerError按 rank 排序后分类——全部RunnerError→SymmetricRunnerError混合 →AsymmetricRunnerError并打 error 级日志全部干净 → 返回各 rank 的返回值列表通常全是Nonecontroller.py。drain()的语义它通过self迭代触发__next__因此同样受派发线程约束它会吞掉Symmetric/AsymmetricRunnerError清理场景不值得浮出但 worker 死亡 / 超时 / 失步的致命错误仍然抛出。毒化是一次性的_fatal_error一旦设置后续所有stream()调用都会以RuntimeError(MGPUController is dead; create a new one.)拒绝is_alive因此成为无需 try/except 的健康探针。公共 API 一览控制器的公共导出面在 multigpu/init.py 中统一声明from ltx_pipelines.multigpu import ...可拿到MGPUController——控制器start / stream / drain / shutdownMGPURunner——每个 pipeline 实现的基类setup 生成器__call__Stream——controller.stream(...)返回的句柄迭代它逐个接收各 rank 的 yieldRunnerError——runner 为可恢复的对称失败抛出的异常SymmetricRunnerError/AsymmetricRunnerError——调用侧从任务结果重抛的异常ControllerBusyError——上一个任务尚未收集时再次stream()抛出的异常NCCLGroups——注入 runnersetup的按组件 NCCL 进程组容器组合使用控制器如何驱动完整的 MGPU pipeline控制器本身不关心并行策略它只负责舰队 调度 流回收。真正把单 GPU pipeline 变成多 GPU 的是 runner 的setup()构建标准 pipeline 后把每个 block 的 builder 换成 MGPU builderSequenceParallelBuilder、TiledDataParallelBuilder、AccelerateGemmaBuilder、DistributedDecoderBuilder并共享一个ModelRegistry与一个TransformerWeightTracker后者负责 LoRA 热切换工作副本全量驻留、干净权重按md5(key) % world_size分片存储。完整模式见 pipeline-setup.md两个官方 runner 的setup()约 50–160 行就是可直接对照的参考实现ti2vid_two_stages_mgpu.py、distilled_mgpu.py。前置条件与运行环境运行 MGPU 推理前请确认详见 多 GPU 总览LinuxNCCL 与 CUDA-IPC peer buffers 仅支持 LinuxmacOS/Windows 不可用单节点≥2 块 CUDA GPU具备 P2P 访问NVLink/PCIe带 CUDA 的 PyTorchltx-kernels已构建SP 的 all2all 内核是硬性依赖可经uv sync --group kernels构建需要 CUDA toolkit / nvcc 与 C 编译器。由于 MGPU 面向延迟而非显存若目标是让放不下的模型跑起来应优先参考 FP8 量化与权重 offloading见 optimization.md。【免费下载链接】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),仅供参考
上一篇/下一篇内容由系统自动关联
返回资讯列表 →