请问大佬，taichi可以同时使用多张显卡训练吗？怎么加速大规模矩阵的计算结果从显卡取回呢？

综合讨论区
#1

目前服务器上有两张3090，程序把一张3090的内存跑满了。请问有什么办法使用多显卡，或者是可以开双进程指定在两张显卡上各自运行吗？
此外，计算结果是一个300,000x1的向量，用field+ti.template()在程序中传递，请问有什么办法可以加快向量从显卡取回cpu吗，目前需要80s左右。

#2

你可以用CUDA_VISIBLE_DEVICES来指定：Global Settings | Taichi Docs

1 个赞
#3

谢谢大佬！

#5

抱歉打扰，根据文档，环境变量我设置如下：
export CUDA_VISIBLE_DEVICES=[(UUIDa), (UUIDb)]
export TI_VISIBLE_DEVICE=[(UUIDa), (UUIDb)]
结果还是只能在UUIDa上运行。
请问这样设置完，是否应该在后端上调整。

#6

设置为ti.vulkan之后，nvidia-smi显示为C+G运行该程序。

#7

欸我们暂时还不支持用GPU多卡跑Taichi，你可能需要开两个process，分别在UDA_VISIBLE_DEVICES=UUIDa和UDA_VISIBLE_DEVICES=UUIDb上面跑

1 个赞
#8

感谢大佬！我查一查怎么开两个process。

#9

抱歉打扰啦，没有查到如何在开两个process的情况下，如何分别设置环境变量，能否请您赐教？

#10

Taichi现在没有原生支持单机多卡，所以如果数据在一张卡跑不下，你可能需要自己做一下数据分割，然后分别跑两段程序：

import os
import subprocess

# Prepare the base environment variables
base_env = os.environ.copy()

# Define the environment variables for the first process
env1 = base_env.copy()
env1['VAR_NAME1'] = 'value1'

# Define the environment variables for the second process
env2 = base_env.copy()
env2['VAR_NAME2'] = 'value2'

# Start the first Python process with its environment variables
process1 = subprocess.Popen(['python', 'script1.py'], env=env1)

# Start the second Python process with its environment variables
process2 = subprocess.Popen(['python', 'script2.py'], env=env2)

# Wait for both processes to complete
process1.wait()
process2.wait()
1 个赞
#11

成功实现！感谢大佬！

1 个赞