在多 GPU (PyTriton) 环境中忽略明确的设备分配

我在一个Python进程(通过PyTriton)中运行多个句子转换器模型. 我的目标是在不同的物理GPU(如cuda:0,cuda:1)上协调不同的模型.

尽管通过一个明确的设备参数(例如:device='cuda:1'),并试图使用film.cuda.set device()上下文管理器,但我观察到在cuda:0(默认设备)上有大量的内存分配,该设备会随模型大小而缩放,建议模型重量或其初始化阶段在被移动前在错误的GPU上进行.

环境:

句式转换器版本: 5.5.1

火炬版本: 2.1.20.

OS: Ubuntu 22.04 (英语).

预期行为 : 所有模型权重和内部缓冲器应完全分配给设备参数中提供的目标GPU.

实际行为: 模型在cuda:0上一致分配内存. 我用nvidia-smi验证了这个 cuda:0上的内存足迹与模型重量的大小相匹配,证实这不仅仅是标准"Global CUDA Context"的间接费用,而是真正的中转问题.

我已经试过了: 为了孤立这个问题,我尝试了以下几点:

• 人工模块注射:初始化模型. Transformer和模型. 将它们分开移动到 . to(device) , 然后再交给判决转换器 。

• 硬件锁定:在初始化阶段使用火炬.cuda. set device(gpu id)上下文管理器.

• RAM-First Loading:先用设备='cpu'来装入模型,再在第一个推论前调用.to(device).

• model kwargs:通过 AutoModel. from pretrained(在绕过句子 Transfer 构造器时)注入设备 映射 。

这些都无法阻止库达0的初入/分配重量.