我相信PyTorch内部使用的环境变量是CUDA_VISIBILE_DEVICES,这是否能达到相同的效果？

哦，我忘了添加一个细节。CUDA_VISIBLE_DEVICES 通常可以正常工作，但是如果某些库初始化了 torch(例如 import torch 或 import transformers),那么它将无法正常工作。实际上，上述代码片段应该修改为：

import torch

def child_process(rank):
    os.environ["LOCAL_RANK"] = str(rank)
    from vllm import LLM

有趣的是，您是在说一旦导入torch,就会分配cuda设备吗？

是的，更准确地说，进行一些GPU操作会产生效果。
以下是重现的最小脚本：
案例A(运行良好)

import os
from multiprocessing import Pool, set_start_method

set_start_method("spawn", force=True)
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

import torch

def child(rank):
    os.environ["CUDA_VISIBLE_DEVICES"] = str(rank)
    from vllm import LLM

    llm = LLM("/nfs_data_storage/llama2_hf/Llama-2-7b-chat-hf")

def main():
    p = Pool(2)
    p.map(child, range(2))

if __name__ == "__main__":
    main()

输出A(预期)

[0] NVIDIA A100-SXM4-80GB | 29°C,   0 % | 13535 / 81920 MB |
[1] NVIDIA A100-SXM4-80GB | 31°C,   0 % | 13535 / 81920 MB |

案例B(不起作用)

import os
from multiprocessing import Pool, set_start_method

set_start_method("spawn", force=True)
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

import torch

torch.tensor(1).to("cuda:0")

def child(rank):
    os.environ["CUDA_VISIBLE_DEVICES"] = str(rank)
    from vllm import LLM

    llm = LLM("/nfs_data_storage/llama2_hf/Llama-2-7b-chat-hf")

def main():
    p = Pool(2)
    p.map(child, range(2))

if __name__ == "__main__":
    main()

输出B(未预期)

[0] NVIDIA A100-SXM4-80GB | 29°C,   0 % | 27440 / 81920 MB |
[1] NVIDIA A100-SXM4-80GB | 30°C,   0 % |     7 / 81920 MB |

以上是我使用LOCAL RANK的解决方案在第二个案例中也有效，但仅适用于vllm<=0.2.6

我也遇到过这个限制，希望在同一个脚本中运行多个模型，但它们位于不同的GPU上。理想情况下，应该有一种方法可以为每个 LLM(...) 初始化指定模型所在的GPU。设置 CUDA_VISIBILE_DEVICES 对于单个模型来说非常有效(可以在导入 torch 之前通过设置 os.environ["CUDA_VISIBLE_DEVICES"] 来实现编程方式)但对于同一进程中的多个模型不起作用。
这里有一些关于修复的讨论，但需要将其实施到发布版本中。