Diretório de artigos
A biblioteca do Hugging Face transformerssuporta o método de instanciação do modelo automático (AutoModel) para carregar e usar automaticamente modelos como GPT e ChatGLM. Os parâmetros do AutoModel.from_pretrained()método device_mappodem realizar inferência de vários cartões em uma única máquina.
análise de parâmetro device_map
device_mapé AutoModel.from_pretrained()um parâmetro importante no método, que é usado para especificar em qual dispositivo de computação específico cada componente do modelo deve ser carregado para obter alocação e utilização eficazes de recursos . Este parâmetro é particularmente útil ao realizar treinamento de modelo paralelo ou distribuído.
device_mapExistem diversas opções de parâmetros auto, balanced, balanced_low_0, sequential, como segue:
“auto” 和 “balanced”: dividirá o modelo uniformemente em todas as GPUs. Principalmente é possível descobrir estratégias de alocação mais eficientes. A funcionalidade do parâmetro "balanceado" permanece estável. (Disponível sob demanda)“balanced_low_0”: O modelo será dividido igualmente em outras GPUs, exceto a primeira GPU, e ocupará menos recursos na primeira GPU. Esta opção é adequada para necessidades que requerem operações adicionais na primeira GPU, como a necessidade de executar a função de geração (processo iterativo) na primeira GPU. (recomendado)“sequential”: Aloque os fragmentos do modelo na ordem da GPU, começando da GPU 0 até a última GPU (então a última GPU geralmente não está ocupada e a diferença entre - "balanced_low_0" é a primeira ou a última e o preenchimento não balanceado), mas no uso real, a GPU 0 irá estourar diretamente a memória de vídeo. (não recomendado)
exemplo de código device_map="auto"
Aqui nosso ambiente é uma única máquina com duas placas gráficas. Use device_map="auto"para carregar ChatGLM-6Bo modelo e observar o uso da placa gráfica.
O código de exemplo é o seguinte:
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1'
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
# 加载模型
model_path = "./model/chatglm2-6b"
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModel.from_pretrained(model_path, trust_remote_code=True, device_map="auto")
text = '什么是机器学习?'
inputs = tokenizer(text, return_tensors="pt")
print(inputs)
outputs = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
-
A situação de ocupação da placa gráfica antes da execução do programa é a seguinte:
Você pode ver que a própria placa gráfica nº 0 está ocupada por outros programas com cerca de 13G de memória de vídeo. -
device_map="auto"A situação de ocupação da placa gráfica após o uso é a seguinte:
Depois de usar a política automática, as placas gráficas 0 e 1 ocupam cerca de 6~7G a mais de memória de vídeo, respectivamente.
Configuração manual
- Configurar como cartão único
device = "cuda:1"
model = AutoModel.from_pretrained(model_path, trust_remote_code=True, device_map=device)
- Configurando para múltiplas placas
Suponha que você queira que algumas partes do modelo estejam na primeira placa gráfica e outras partes na segunda placa gráfica. Você precisa saber o nome da camada do modelo ou alocá-lo razoavelmente de acordo com o tamanho do componente do modelo. modelo. No entanto, o nome da camada específica precisa ser determinado com base no modelo real. Aqui está um exemplo conceitual:
device = {
"transformer.h.0": "cuda:0", # 第一部分放在GPU 0
"transformer.h.1": "cuda:1", # 第二部分放在GPU 1
# ... 根据模型结构继续分配
}
model = AutoModelForCausalLM.from_pretrained(model_dir, device_map=device)