[Modelo grande] Biblioteca de transformadores dispositivo_map de inferencia de múltiples tarjetas de una sola máquina

La biblioteca de Hugging Face transformersadmite el método de creación de instancias del modelo automático (AutoModel) para cargar y utilizar automáticamente modelos como GPT y ChatGLM. Los parámetros del AutoModel.from_pretrained()método device_mappueden realizar inferencias de múltiples tarjetas en una sola máquina.

análisis del parámetro device_map

device_mapes AutoModel.from_pretrained()un parámetro importante en el método, que se utiliza para especificar en qué dispositivo informático específico se debe cargar cada componente del modelo para lograr una asignación y utilización efectiva de los recursos . Este parámetro es particularmente útil cuando se realiza un entrenamiento distribuido o paralelo de modelos.

device_mapHay varias opciones para los parámetros auto, balanced, balanced_low_0, sequential, como se muestra a continuación:

  • “auto” 和 “balanced”: dividirá el modelo de manera uniforme en todas las GPU. Principalmente es posible descubrir estrategias de asignación más eficientes. La funcionalidad del parámetro "equilibrado" permanece estable. (Disponible bajo demanda)
  • “balanced_low_0”: El modelo se dividirá equitativamente en otras GPU excepto la primera GPU y ocupará menos recursos en la primera GPU. Esta opción es adecuada para necesidades que requieren operaciones adicionales en la primera GPU, como la necesidad de ejecutar la función de generación (proceso iterativo) en la primera GPU. (recomendado)
  • “sequential”: Asigne fragmentos de modelo en el orden de la GPU, desde GPU 0 hasta la última GPU (luego, la última GPU a menudo no está ocupada y la diferencia entre - "balanced_low_0" es la primera o la última, y ​​el llenado desequilibrado), pero en uso real, la GPU 0 explotará directamente la memoria de video. (no recomendado)

ejemplo de código device_map="auto"

Aquí nuestro entorno es una sola máquina con dos tarjetas gráficas. Se utiliza device_map="auto"para cargar ChatGLM-6Bel modelo y observar el uso de la tarjeta gráfica.

El código de muestra es el siguiente:

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))
  • La situación de ocupación de la tarjeta gráfica antes de ejecutar el programa es la siguiente:
    Insertar descripción de la imagen aquí
    Puede ver que la tarjeta gráfica No. 0 está ocupada por otros programas con aproximadamente 13 G de memoria de video.

  • device_map="auto"La situación de ocupación de la tarjeta gráfica después de su uso es la siguiente:
    Insertar descripción de la imagen aquí
    después de usar la política automática, las tarjetas gráficas 0 y 1 ocupan alrededor de 6~7G más de memoria de video respectivamente.

Configuración manual

  • Configurar como tarjeta única
device = "cuda:1"
model = AutoModel.from_pretrained(model_path, trust_remote_code=True, device_map=device)
  • Configuración para tarjetas múltiples
    Supongamos que desea que algunas partes del modelo estén en la primera tarjeta gráfica y otras partes en la segunda tarjeta gráfica. Necesita conocer el nombre de la capa del modelo o asignarlo razonablemente de acuerdo con el tamaño del componente. modelo. Sin embargo, el nombre de la capa específica debe determinarse en función del modelo real. A continuación se muestra un ejemplo conceptual:
device = {
    
    
    "transformer.h.0": "cuda:0",  # 第一部分放在GPU 0
    "transformer.h.1": "cuda:1",  # 第二部分放在GPU 1
    # ... 根据模型结构继续分配
}
model = AutoModelForCausalLM.from_pretrained(model_dir, device_map=device)

Referencias

Supongo que te gusta

Origin blog.csdn.net/u012856866/article/details/140498484
Recomendado
Clasificación