是否支持多卡推理 about vary HOT 2 OPEN

################指定卡
os.environ["CUDA_VISIBLE_DEVICES"]="0，1"
##############导入accelerate
from accelerate import init_empty_weights, infer_auto_device_map, load_checkpoint_and_dispatch
###############替换下列代码：
#model = varyQwenForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True, device_map='auto', trust_remote_code=True)
#################替换为
config = varyConfig.from_pretrained(model_name, trust_remote_code=True)
with init_empty_weights():
model = varyQwenForCausalLM._from_config(config, torch_dtype=torch.float16)
no_split_modules = model._no_split_modules
print(f"no_split_modules: {no_split_modules}", flush=True)
map_list = {0:"20GB", 1:"20GB"}
device_map = infer_auto_device_map(model, max_memory=map_list, no_split_module_classes=no_split_modules)

######### 通过device map加载模型
model = load_checkpoint_and_dispatch(model, checkpoint=model_dir, device_map=device_map)

就可以了

from vary.

哇！感谢大佬！

from vary.