单RTX4090运行Merlyn模型失败，如何用双GPU部署？

可以用两块RTX 4090运行该模型，进程被杀的核心原因是单卡半精度模式下，模型权重加推理上下文的显存占用仍超过单卡24GB上限，通过Accelerate或DeepSpeed实现模型并行，将模型权重拆分到两块GPU上，即可解决显存溢出问题。

方法一：使用Accelerate实现模型并行 #

Accelerate是轻量化的并行工具，配置和代码修改都很简洁：

1. 先通过命令行初始化Accelerate配置：

accelerate config

根据提示选择「多GPU」「模型并行」「半精度(fp16)」等选项，完成配置文件生成。

2. 修改推理代码如下：

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from accelerate import Accelerator

# 初始化Accelerator
accelerator = Accelerator()
model_path = "MerlynMind/merlyn-education-corpus-qa"

# 加载tokenizer和模型（直接指定半精度，Accelerate自动处理并行）
tokenizer = AutoTokenizer.from_pretrained(model_path, fast_tokenizer=True)
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16)

# 准备模型和tokenizer
model, tokenizer = accelerator.prepare(model, tokenizer)

# 推理示例
prompt = "请输入你的问题"
inputs = tokenizer(prompt, return_tensors="pt").to(accelerator.device)
outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

方法二：使用DeepSpeed实现显存优化（支持ZeRO） #

DeepSpeed适合更极致的显存优化，比如ZeRO（零冗余优化器），能进一步降低单卡显存占用：

1. 先安装DeepSpeed：

pip install deepspeed

2. 创建DeepSpeed配置文件ds_config.json，示例ZeRO-2配置（兼顾并行效率和显存节省）：

{
  "train_batch_size": 1,
  "train_micro_batch_size_per_gpu": 1,
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 2e-5
    }
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "none"
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 200000000,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 200000000,
    "contiguous_gradients": true
  },
  "fp16": {
    "enabled": true
  }
}

3. 修改推理代码：

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import deepspeed

model_path = "MerlynMind/merlyn-education-corpus-qa"
tokenizer = AutoTokenizer.from_pretrained(model_path, fast_tokenizer=True)
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16)

# 初始化DeepSpeed加载模型
model, _, _, _ = deepspeed.initialize(model=model, config_params="ds_config.json")

# 推理示例
prompt = "请输入你的问题"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

注意事项 #

• 确保PyTorch、transformers、accelerate/deepspeed版本兼容，建议使用PyTorch 2.x+版本以获得更好的并行性能。
• 如果推理时上下文过长导致显存仍紧张，可在加载模型时添加gradient_checkpointing=True参数，进一步降低显存占用，但会略微增加推理时间。
• 两种方法都不需要手动调用model.to(device)或model.half()（代码中已通过torch_dtype=torch.float16指定半精度，并行框架自动处理设备分配）。

内容的提问来源于stack exchange，提问作者José Adrián Pardo Pérez