使用DeepSpeed分布式训练时，如何设置单设备最大GPU显存以运行长序列Llama？

一、设置单设备最大GPU显存使用量 #

DeepSpeed支持通过配置文件精准限制每个GPU的显存使用上限，避免显存溢出。在你的deepspeed_config.json中添加max_memory字段即可：

{
  "train_batch_size": 32,  // 根据任务需求调整
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 2e-5
    }
  },
  "max_memory": {
    "gpu": "42G",  // 给46G显存的卡预留4G左右系统余量
    "cpu": "100G"  // 可选，设置CPU内存上限，防止数据溢出到磁盘
  },
  // 其他DeepSpeed配置项（如ZeRO、混合精度等）
}

• 该配置会自动应用到所有可见GPU设备；
• 多卡场景下，若需针对不同GPU设置不同上限，可将"gpu"的值改为数组，比如["42G", "40G"]。

二、适配Llama长序列（8000-10000长度） #

Llama默认2048的序列长度无法满足需求，需从位置编码和模型配置两方面修改：

方法1：RoPE线性缩放（推荐） #

Llama采用旋转位置编码（RoPE），可通过配置动态扩展序列长度，无需重新训练位置嵌入：

from transformers import LlamaConfig, LlamaForCausalLM

# 加载原模型配置并修改
config = LlamaConfig.from_pretrained("meta-llama/Llama-2-7b-hf")
config.max_position_embeddings = 10000  # 设置目标序列长度
config.rope_scaling = {
    "type": "linear",
    "factor": 10000 / 2048  # 计算缩放因子
}

# 用修改后的配置加载模型
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", config=config)

方法2：位置嵌入插值扩展 #

如果不使用RoPE缩放，可直接对原位置嵌入进行插值拉伸：

import torch
from transformers import LlamaForCausalLM

model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
target_seq_len = 10000

# 插值扩展位置嵌入权重
old_pos_emb = model.model.embed_positions.weight
new_pos_emb = torch.nn.functional.interpolate(
    old_pos_emb.unsqueeze(0),
    size=target_seq_len,
    mode='linear',
    align_corners=False
).squeeze(0)

# 替换原嵌入层并更新配置
model.model.embed_positions = torch.nn.Embedding(target_seq_len, old_pos_emb.size(1))
model.model.embed_positions.weight.data = new_pos_emb
model.config.max_position_embeddings = target_seq_len

三、显存优化额外建议 #

结合你的46G显存，配合以下DeepSpeed配置可进一步降低显存占用：

1. 启用ZeRO优化：在deepspeed_config.json中添加ZeRO阶段配置（推荐stage 2）：

{
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  }
}

2. 开启混合精度：启用FP16或BF16减少显存占用：

{
  "fp16": {
    "enabled": true
  }
}

内容的提问来源于stack exchange，提问作者hjc