Guía introductoria al aprendizaje profundo en 2023 (23) - ChatGLM2

Guía introductoria al aprendizaje profundo en 2023 (23) - ChatGLM2

En la sección "Ejecución de modelos grandes en su computadora", presentamos el modelo ChatGLM, que era uno de los mejores modelos chinos grandes en ese momento. Ahora, se ha actualizado a la segunda generación, ChatGLM2.

En ese momento, nuestras reservas técnicas no eran suficientes, solo podíamos dejarlo correr y no nos atrevíamos a explicar sus principios y códigos.

Ahora, tras el bombardeo indiscriminado de LLaMA 2 y el código de Baichuan, todo el mundo se ha adaptado al ritmo de mirar el código. Ahora es el momento de ver los principios y el código de ChatGLM2.

Ejecutar ChatGLM2

En primer lugar, todavía ejecutamos el código de ChatGLM2. En máquinas con más de 13 GB de memoria de video, ChatGLM2 puede funcionar sin problemas. Por ejemplo, estoy corriendo en un 15G T4.

El primer paso es instalar las bibliotecas correspondientes:

pip install protobuf
pip install transformers==4.30.2
pip install cpm_kernels
pip install torch>=2.0
pip install gradio
pip install mdtex2html
pip install sentencepiece
pip install accelerate
pip install sse-starlette
pip install streamlit>=1.24.0

En el segundo paso, puede usar la interfaz estándar de Transformers para llamar a ChatGLM2:

from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True, device='cuda')
model = model.eval()
response, history = model.chat(tokenizer, "生成scala语言的快速排序", history=[])
print(response)

La salida es la siguiente:

def quickSort(arr: Int[]): Int[] = {
    
    
  val pivot = arr(arr.length / 2)
  val left = 0
  val right = arr.length - 1
  while (left <= right) {
    
    
    while (arr(left) < pivot) {
    
    
      left = left + 1
    }
    arr(left) = pivot
    while (arr(right) > pivot) {
    
    
      right = right - 1
    }
    arr(right) = pivot
    left = left + 1
    right = right - 1
  }
  return arr
}

Si se ejecuta en una tarjeta gráfica con una memoria de video más pequeña, podemos usar el resultado de la cuantificación de 4 bits:

from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4",trust_remote_code=True).cuda()
model = model.eval()
response, history = model.chat(tokenizer, "生成Kotlin语言编写的快速排序", history=[])
print(response)

Esto es lo que obtuve funcionando en el 3060:

fun quickSort(arr: IntArray): IntArray {
    
        
    val left = 0                            
    val right = arr.size - 1                
    val quicksortFactor = arr.size / 2      
                                            
    while (left < right) {
    
                      
        quicksortFactor--.let {
    
                 
            let x = left                    
            let y = right                   
            let temp = arr[x]               
                                            
            if (temp < arr[y]) {
    
                
                x++                         
            } else {
    
                            
                y--                         
            }                               
                                            
            if (x == y) {
    
                       
                break                       
            }                               
                                            
            quicksortFactor++.let {
    
             
                arr[x] = arr[y]             
                    arr[y] = temp           
            }                               
        }                                   
    }                                       
                                            
    return arr                              
}                                           

Análisis de código CUDA cuantificado

He dicho que se han implementado muchos códigos de atención de múltiples cabezas, y hablaré de eso más adelante. En esta sección hablamos de algo que no hemos cubierto antes, cuantificar el código CUDA utilizado. La parte de LLaMA 2 no se menciona porque no tiene una parte de cuantificación y el código central CUDA de Baichuan aún no se ha abierto. Entonces, hablemos primero con el código de GLM.

Echemos un vistazo al Makefile de la parte central de CUDA:

NVCC=nvcc
OPTIONS=-gencode arch=compute_61,code=sm_61 \
		-gencode arch=compute_62,code=sm_62 \
		-gencode arch=compute_70,code=sm_70 \
		-gencode arch=compute_72,code=sm_72 \
		-gencode arch=compute_75,code=sm_75 \
		-gencode arch=compute_80,code=sm_80 \
		-gencode arch=compute_86,code=sm_86

TARGETS=$(patsubst %.cu, %.fatbin, $(wildcard *.cu))

all: $(TARGETS)

%.fatbin: %.cu
	$(NVCC) -fatbin $^ $(OPTIONS) -o $@

.PHONY : clean, copy
clean:
	rm $(TARGETS)

copy:
	cp $(TARGETS) ../kernels/

Podemos ver que el código aquí admite múltiples arquitecturas CUDA, incluidas 6.1, 6.2, 7.0, 7.2, 7.5, 8.0 y 8.6. La arquitectura aquí se refiere a la arquitectura de la GPU. Por ejemplo, la arquitectura del RTX 3090 es 8.6 y la arquitectura del RTX 3060 es 8.0.

• 6.1 y 6.2 corresponden a la arquitectura Pascal, como P100, GTX 1060
• 7.0 es la arquitectura Volta, como V100
• 7.5 es la arquitectura Turing, como RTX 2080, T4
• 8.0 y 8.6 son arquitectura Ampere, como A100, RTX 3090

Aunque ya es compatible con muchas arquitecturas, las arquitecturas más antiguas como Maxwell y Kepler se han ido con el viento.

Para admitir tantas arquitecturas, es necesario introducir un nuevo punto de conocimiento: fatbin.

El archivo .fatbin es un archivo de formato binario CUDA Fat. Este es un formato de archivo binario especial utilizado por la plataforma CUDA de NVIDIA. El archivo fatbin contiene código para varias arquitecturas de GPU y capacidades informáticas, y puede ejecutarse en muchos tipos diferentes de procesadores.

En la programación CUDA, el código GPU (comúnmente llamado kernel) a menudo se almacena en el código host de manera similar al ensamblaje en línea. Sin embargo, este enfoque tiene algunas dificultades en las aplicaciones prácticas, principalmente porque las diferentes arquitecturas y dispositivos de GPU pueden requerir diferentes versiones de código de GPU. CUDA Fat Binary resuelve este problema al contener múltiples versiones del código GPU, cada una optimizada para una arquitectura GPU específica.

Cuando se ejecuta un programa CUDA, el sistema de tiempo de ejecución CUDA examina el dispositivo en el que se está ejecutando y selecciona la versión del código GPU que es más apropiada para ese dispositivo del archivo binario pesado. De esta forma, se puede usar un archivo .fatbin para permitir que el mismo programa CUDA se ejecute en GPU con diferente potencia informática. No es necesario compilar por separado para diferentes GPU.

Echemos un vistazo a la implementación de la cuantificación de 8 bits. Este código es exactamente un ejemplo de cómo escribir el código CUDA más simple:

template<typename T>
__device__ void
int8WeightExtractionDevice(const int8_t* weight,
                                const T* scale_list,
                                T* output,
                                const int n,
                                const int k)
{
    
    
    for(int i = blockIdx.x * k + threadIdx.x; i < blockIdx.x * k + k; i += blockDim.x){
    
    
        output[i] = T(weight[i]) * scale_list[blockIdx.x];
    }
}

En el hilo de la GPU:

• Calcule el índice de peso leído por el subproceso actual: blockIdx.x representa la identificación del bloque, _k representa cada bloque procesa k valores y threadIdx.x representa la identificación del subproceso
• Lea el valor int8 en el índice actual de la matriz de peso
• Escálelo: multiplique el valor del peso por el factor de escala correspondiente a la identificación del bloque en scale_list
• El resultado se envía al índice correspondiente de la matriz de salida.

Finalmente, el cálculo general de copiado y escalado desde la matriz de peso a la matriz de salida se completa en paralelo a través del subproceso blockDim.x.

Si olvida el contenido de la sección CUDA, repasemos los conceptos de blockIdx, threadIdx y blockDim:

• blockIdx: CUDA organiza los hilos en bloques, y cada bloque tiene una identificación llamada blockIdx. Puede haber múltiples bloques, los diferentes bloques se distinguen por blockIdx
• threadIdx: Hay varios hilos en cada bloque, los diferentes hilos en el mismo bloque se distinguen por threadIdx. Los identificadores de subprocesos empiezan a contar desde 0
• blockDim: indica el número de hilos contenidos en cada bloque

Especifique al iniciar la función del kernel, por ejemplo, ejecute <<<32, 128>>> al llamar a la función del kernel, lo que significa que hay 32 bloques y hay 128 subprocesos en cada bloque.

Dentro del cuerpo de la función, hay un bucle for, tenga en cuenta que este bucle no es en serie como en una CPU de un solo núcleo, ¡sino que se ejecuta en cada subproceso de CUDA!
El valor inicial de la variable de ciclo i es blockIdx.x * k + threadIdx.x, que es un patrón común para asignar diferentes partes de datos a diferentes subprocesos CUDA. En cada ciclo, i aumenta blockDim.x, lo que significa que el intervalo de datos procesado por cada subproceso es del tamaño de un bloque.

En el bucle for, la función multiplica los pesos por los factores de escala correspondientes y almacena el resultado en la matriz de salida. Aquí, el tipo de peso se convierte a T y luego se multiplica por el elemento scale_list correspondiente. Tenga en cuenta el uso de scale_list[blockIdx.x], lo que significa que todos los subprocesos en el mismo bloque usan el mismo factor de escala.

El núcleo de CUDA está encapsulado en la función de host:

extern "C" __global__ void int8WeightExtractionHalf(const int8_t* weight,
                                const half* scale_list,
                                half* output,
                                const int n,
                                const int k){
    
    
                                    int8WeightExtractionDevice<half>(weight, scale_list, output, n, k);
                                }

extern "C" __global__ void int8WeightExtractionFloat(const int8_t* weight,
                                const float* scale_list,
                                float* output,
                                const int n,
                                const int k){
    
    
                                    int8WeightExtractionDevice<float>(weight, scale_list, output, n, k);
                                }

Echemos un vistazo a cómo llamar a esta función en Python:

def extract_weight_to_half(weight: torch.Tensor, scale_list: torch.Tensor, source_bit_width: int):
    if source_bit_width == 8:
        func = kernels.int8WeightExtractionHalf
    elif source_bit_width == 4:
        func = kernels.int4WeightExtractionHalf
    else:
        assert False, "Unsupported bit-width"

    with torch.cuda.device(weight.device):
        n, m = weight.size(0), weight.size(1)
        out = torch.empty(n, m * (8 // source_bit_width), dtype=torch.half, device="cuda")
        stream = torch.cuda.current_stream()

        gridDim = (n, 1, 1)
        blockDim = (min(round_up(m, 32), 1024), 1, 1)

        func(
            gridDim,
            blockDim,
            0,
            stream,
            [
                ctypes.c_void_p(weight.data_ptr()),
                ctypes.c_void_p(scale_list.data_ptr()),
                ctypes.c_void_p(out.data_ptr()),
                ctypes.c_int32(n),
                ctypes.c_int32(m),
            ],
        )
        return out

Podemos ver que si es para cuantización de 8 bits, se llamará a la función kernels.int8WeightExtractionHalf, que corresponde a la función que escribimos anteriormente.

Expliquemos cómo dividir el grado de paralelismo.
gridDim Esta variable representa el tamaño de cuadrícula de la función kernel de CUDA. En este código, se establece en (n, 1, 1), donde n es el tamaño de la primera dimensión del tensor de peso. Esto significa que hay n bloques en la grilla, y cada bloque es responsable de procesar una fila del tensor de peso.

blockDim es un triplete que especifica las dimensiones de cada bloque de subprocesos. En este código, blockDim se establece en (min(round_up(m, 32), 1024), 1, 1). Esto significa que el número de subprocesos en cada bloque de subprocesos es min(round_up(m, 32), 1024), que es el múltiplo de m (el tamaño de la segunda dimensión del tensor de peso) elevado a los 32 más cercanos, pero el máximo no es más de 1024. Esto se debe a la limitación de la arquitectura CUDA, la cantidad de subprocesos por bloque de subprocesos no puede exceder los 1024.

Hablemos de nuevo de la compresión de 4 bits:

__device__ void
int4WeightCompressionDevice(const int8_t* input,
                                int8_t* output,
                                const int n,
                                const int k)
{
    
    
    for(int i = blockIdx.x * k + threadIdx.x; i < blockIdx.x * k + k; i += blockDim.x){
    
    
        output[i] = (input[i * 2] << 4) | (input[i * 2 + 1] & 0b00001111);
    }
}

Dentro de int4WeightCompressionDevice, para cada hilo, calcula el índice i del elemento que debe procesar. Luego comprime los dos elementos en los índices i * 2 e i * 2 + 1 de la matriz de entrada en un solo elemento. El método de compresión consiste en desplazar el primer elemento a la izquierda en 4 bits y luego realizar una operación OR bit a bit con el segundo elemento. Finalmente, almacene el resultado en el índice i en la matriz de salida.

Aunque puede ser altamente paralelizado, de hecho, el código escrito en la GPU no es muy diferente al de la CPU, y no hay necesidad de aprender nuevas oraciones.

De la misma manera, veamos cómo convertir pesos comprimidos de 4 bits a pesos comprimidos de 8 bits:

template<typename T>
__device__ void
int4WeightExtractionDevice(const int8_t* weight,
                                const T* scale_list,
                                T* output,
                                const int n,
                                const int k)
{
    
    
    for(int i = blockIdx.x * k + threadIdx.x; i < blockIdx.x * k + k; i += blockDim.x){
    
    
        int8_t original = weight[i];
        int8_t high = original >> 4;
        int8_t low = original << 4; low = low >> 4;
        output[i * 2] = T(high) * scale_list[blockIdx.x];
        output[i * 2 + 1] = T(low) * scale_list[blockIdx.x];
    }
}

Con el conocimiento anterior, no hay necesidad de una explicación adicional aquí, ¿verdad?

Realización de la capa de cuantificación

Finalmente, veamos cómo se usa la cuantificación en las redes neuronales:

解释下面的代码：
import torch

from kernels import extract_weight_to_half


class W8A16Linear(torch.autograd.Function):
    @staticmethod
    def forward(ctx, inp: torch.Tensor, quant_w: torch.Tensor, scale_w: torch.Tensor, weight_bit_width):
        ctx.inp_shape = inp.size()
        ctx.weight_shape = quant_w.size()
        ctx.weight_bit_width = weight_bit_width
        out_features = quant_w.size(0)
        inp = inp.contiguous().view(-1, inp.size(-1))
        weight = extract_weight_to_half(quant_w, scale_w, weight_bit_width)
        output = inp.mm(weight.t())
        ctx.save_for_backward(inp, quant_w, scale_w)
        return output.view(*(ctx.inp_shape[:-1] + (out_features,)))

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        inp, quant_w, scale_w = ctx.saved_tensors
        weight = extract_weight_to_half(quant_w, scale_w, ctx.weight_bit_width)
        grad_output = grad_output.contiguous().view(-1, weight.size(0))
        grad_input = grad_output.mm(weight)
        grad_weight = grad_output.t().mm(inp)
        return grad_input.view(ctx.inp_shape), grad_weight.view(ctx.weight_shape), None

El método directo acepta cuatro parámetros: inp es un tensor de entrada, quant_w es un tensor de peso cuantificado, scale_w es un tensor de escala de peso y weight_bit_width es el ancho de bit del peso. Este método primero guarda la forma del tensor de entrada y el tensor de peso, luego convierte el tensor de entrada en continuo y remodelado. A continuación, extrae pesos de precisión media de los pesos cuantificados y la escala de peso mediante la función extract_weight_to_half de la que acabamos de hablar. Luego calcula la salida mediante la multiplicación de matrices y cambia el tamaño del resultado a la forma correcta. Finalmente, guarda la entrada, los pesos cuantificados y las escalas de peso para su uso durante la retropropagación.

El método de retroceso acepta un argumento: grad_output es un tensor de salida de gradiente. Este método primero toma la entrada guardada, los pesos cuantificados y las escalas de peso del contexto y extrae pesos de precisión media de ellos. Luego, convierte la salida del degradado en continuo y remodelado. A continuación, calcula los gradientes de entrada y ponderación mediante la multiplicación de matrices. Finalmente, devuelve los gradientes de entrada reformados y los gradientes de peso.

Este código implementa una capa lineal personalizada que utiliza pesos de precisión media para el cálculo y es compatible con el mecanismo de diferenciación automática de PyTorch.

Esto aún no ha terminado. Para lograr una mayor paralelización, la capa de cuantificación se puede encapsular aún más:

import torch
from torch.nn.parameter import Parameter

from SwissArmyTransformer.mpu import copy_to_model_parallel_region
from SwissArmyTransformer.mpu import gather_from_model_parallel_region
from SwissArmyTransformer.mpu import reduce_from_model_parallel_region
from SwissArmyTransformer.mpu import scatter_to_model_parallel_region
from SwissArmyTransformer.mpu import ColumnParallelLinear, RowParallelLinear

from .functional import W8A16Linear
from kernels import compress_int4_weight


class QuantizedColumnParallelLinear(ColumnParallelLinear):
    def __init__(self, weight_bit_width: int, weight=None, *args, **kwargs):
        super(QuantizedColumnParallelLinear, self).__init__(*args, **kwargs)
        self.weight_bit_width = weight_bit_width

        shape = self.weight.shape
        del self.weight

        if weight is None:
            self.weight = torch.empty(
                shape[0], shape[1] * weight_bit_width // 8, dtype=torch.int8, device=kwargs["device"]
            )
            self.weight_scale = torch.empty(shape[0], dtype=kwargs["params_dtype"], device=kwargs["device"])
        else:
            self.weight_scale = (weight.abs().max(dim=-1).values / ((2 ** (weight_bit_width - 1)) - 1)).half()
            self.weight = torch.round(weight / self.weight_scale[:, None]).to(torch.int8)
            if weight_bit_width == 4:
                self.weight = compress_int4_weight(self.weight)

        self.weight = Parameter(self.weight.to(kwargs["device"]), requires_grad=False)
        self.weight_scale = Parameter(self.weight_scale.to(kwargs["device"]), requires_grad=False)

    def forward(self, input_):
        # Set up backprop all-reduce.
        input_parallel = copy_to_model_parallel_region(input_)
        # Matrix multiply.
        output_parallel = W8A16Linear.apply(input_parallel, self.weight, self.weight_scale, self.weight_bit_width)
        if self.bias is not None:
            output_parallel = output_parallel + self.bias
        if self.gather_output:
            # All-gather across the partitions.
            output = gather_from_model_parallel_region(output_parallel)
        else:
            output = output_parallel
        return output

Este código define una clase denominada QuantizedColumnParallelLinear que hereda de la clase ColumnParallelLinear. Esta clase implementa una capa lineal paralela a columnas cuantificada.

__init__El método acepta varios parámetros, donde weight_bit_width es el ancho de bits del peso y el peso es un tensor de peso opcional. Este método primero llama al constructor de la clase principal y luego guarda el ancho de bits del peso. A continuación, obtiene la forma del peso y elimina el atributo de peso. Si no se proporcionan pesos, crea un tensor de pesos vacíos y un tensor de escala de pesos vacíos. De lo contrario, calcula la escala de ponderación en función de las ponderaciones proporcionadas y cuantifica las ponderaciones en números enteros. Si el ancho de bits es 4, el peso se comprime mediante la función compress_int4_weight. Finalmente, convierta los pesos y las escalas de peso a parámetros de PyTorch y guárdelos.

El método directo acepta un argumento: input_ es un tensor de entrada. Este método primero copia la entrada a la región paralela del modelo usando la función copy_to_model_parallel_region. Luego, use la función W8A16Linear.apply para calcular la salida. Si está sesgado, el sesgo se agrega a la salida. Si es necesario recopilar la salida, utilice la función de recopilación_de_modelo_paralelo_región para recopilar la salida. De lo contrario, devuelva la salida directamente.

Este código implementa una capa lineal paralela de columna cuantificada que se puede calcular en paralelo en varias GPU.

resumen

En esta sección, aprovechamos para hablar sobre la función de ChatGLM2 y, de paso, daremos una introducción completa a los métodos de cuantificación que se utilizarán desde CUDA hasta el paralelismo multi-GPU.
Si tiene alguna función que pueda acelerarse con el código del dispositivo CUDA, ¡no dude en implementarla! La aceleración de los algoritmos no se limita a cómo usar los marcos de trabajo de otras personas y las funciones existentes.