Guía introductoria al aprendizaje profundo en 2023 (19) - Análisis del código fuente de LLaMA 2

Guía introductoria al aprendizaje profundo en 2023 (19) - Análisis del código fuente de LLaMA 2

En la sección anterior, aprendimos cómo usar las API de finalización y chat de LLaMA 2. En esta sección echamos un vistazo al código fuente de LLaMA 2.

Función de finalización text_completion análisis de código fuente

En la sección anterior, hablamos sobre el método de programación de LLaMA 2. Vamos a repasarlo:

    generator = Llama.build(
        ckpt_dir=ckpt_dir,
        tokenizer_path=tokenizer_path,
        max_seq_len=max_seq_len,
        max_batch_size=max_batch_size,
    )

    prompts = [
        "上下五千年，英雄万万千。黄沙百战穿金甲，不破楼兰终不还",
    ]
    results = generator.text_completion(
        prompts,
        max_gen_len=max_gen_len,
        temperature=temperature,
        top_p=top_p,
    )

Primero echemos un vistazo a lo que significan los parámetros de la función text_completion.El prototipo de esta función es:

def text_completion(
self,
prompts: List[str],
temperature: float = 0.6,
top_p: float = 0.9,
max_gen_len: Optional[int] = None,
logprobs: bool = False,
echo: bool = False,
) -> List[CompletionPrediction]:

Echemos un vistazo al significado de estos parámetros:

• solicitudes: esta es una lista de cadenas, cada una de las cuales es una solicitud para generar texto.
• temperatura (por defecto 0.6): Este es un parámetro que controla la aleatoriedad del texto generado. Cuanto mayor sea el valor de la temperatura, más aleatorio será el texto generado; cuanto menor sea el valor de la temperatura, más probable es que el texto generado tenga la salida más probable.
• top_p (el valor predeterminado es 0.9): este es un parámetro que controla la diversidad del texto generado. Comienza desde la palabra con la probabilidad más alta y se acumula hasta la palabra con la probabilidad total que excede top_p, y luego selecciona aleatoriamente una palabra de estas palabras como palabras generadas. Este método también se conoce como muestreo de núcleo o muestreo top-p.
• max_gen_len: parámetro opcional, que indica la longitud máxima del texto generado. Si no se especifica, se utilizará la longitud de secuencia máxima en los parámetros del modelo menos 1.
• logprobs (el valor predeterminado es falso): si es verdadero, la probabilidad de registro de cada palabra generada se incluirá en el resultado devuelto.
• echo (falso por defecto): este es un parámetro que controla si se incluyen indicaciones de entrada en el texto generado.

Después de entender los parámetros, veamos la implementación completa de text_completion:

    def text_completion(
        self,
        prompts: List[str],
        temperature: float = 0.6,
        top_p: float = 0.9,
        max_gen_len: Optional[int] = None,
        logprobs: bool = False,
        echo: bool = False,
    ) -> List[CompletionPrediction]:
        if max_gen_len is None:
            max_gen_len = self.model.params.max_seq_len - 1
        prompt_tokens = [self.tokenizer.encode(x, bos=True, eos=False) for x in prompts]
        generation_tokens, generation_logprobs = self.generate(
            prompt_tokens=prompt_tokens,
            max_gen_len=max_gen_len,
            temperature=temperature,
            top_p=top_p,
            logprobs=logprobs,
            echo=echo,
        )
        if logprobs:
            return [
                {
    
    
                    "generation": self.tokenizer.decode(t),
                    "tokens": [self.tokenizer.decode(x) for x in t],
                    "logprobs": logprobs_i,
                }
                for t, logprobs_i in zip(generation_tokens, generation_logprobs)
            ]
        return [{
    
    "generation": self.tokenizer.decode(t)} for t in generation_tokens]

En resumen, hay tres pasos. Este text_completion es en realidad la función contenedora de generar:

• Codificación: llamada tokenizer.encode
• Generar: llamada generar
• Decodificación: llamada tokenizer.decode

Participio

import os
from logging import getLogger
from typing import List

from sentencepiece import SentencePieceProcessor


logger = getLogger()


class Tokenizer:
    def __init__(self, model_path: str):
        # reload tokenizer
        assert os.path.isfile(model_path), model_path
        self.sp_model = SentencePieceProcessor(model_file=model_path)
        logger.info(f"Reloaded SentencePiece model from {
      
      model_path}")

        # BOS / EOS token IDs
        self.n_words: int = self.sp_model.vocab_size()
        self.bos_id: int = self.sp_model.bos_id()
        self.eos_id: int = self.sp_model.eos_id()
        self.pad_id: int = self.sp_model.pad_id()
        logger.info(
            f"#words: {
      
      self.n_words} - BOS ID: {
      
      self.bos_id} - EOS ID: {
      
      self.eos_id}"
        )
        assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()

    def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
        assert type(s) is str
        t = self.sp_model.encode(s)
        if bos:
            t = [self.bos_id] + t
        if eos:
            t = t + [self.eos_id]
        return t

    def decode(self, t: List[int]) -> str:
        return self.sp_model.decode(t)

La primera es usar el componente de segmentación de palabras SentencePieceProcessor. SentencePieceProcessor es un componente de la biblioteca SentencePiece, que implementa las funciones de subpalabra (subword) tokenize y detokenize.

Sus funciones principales incluyen:

• Tokenize el texto en subpalabras. El enfoque basado en datos utilizado por SentencePiece puede aprender el vocabulario del texto y tokenizar el texto en unidades de subpalabras.
• Combine la subpalabra detokenize en el texto original. La secuencia de subpalabras tokenizadas se puede volver a fusionar en el texto original.
• Proporciona gestión de vocabulario. Se puede obtener información como el vocabulario de subpalabras de tokenize.
• Admite tokenizar y detokenizar texto en varios idiomas.
• Proporciona una implementación eficiente. La capa inferior está realizada por C++, que puede procesar rápidamente texto a gran escala.
• Proporcione una variedad de opciones de modelos, como BPE, unigram, etc.
• Admite modelos de sublematización de entrenamiento personalizados.

Bien, volvamos al código en sí. Este código implementa una clase Tokenizer basada en SentencePiece, que puede tokenizar y detokenizar texto.

Lógica principal:

• Cargue el archivo de modelo de SentencePiece model_path en la inicialización.
• Obtenga el tamaño de vocabulario n_words del modelo y la identificación del token especial (bos_id, eos_id, pad_id).
• El método de codificación puede tokenizar el literal de cadena s en una lista de id. Puede elegir agregar bos_id al principio y eos_id al final.
• El método de decodificación puede decodificar una lista de identificadores a una cadena literal.

Esto construye una clase Tokenizer que encapsula el tokenize/detokenize de SentencePiece. Se puede cargar un modelo personalizado de SentencePiece, y luego el texto se puede sublematizar fácilmente.

De esta forma, el modelo SentencePiece entrenado se puede reutilizar para proporcionar funciones confiables de tokenización y destokenización para tareas posteriores de NLP.

Finalmente, hablemos de algunos símbolos especiales bos_id, eos_id y pad_id:

• bos_id: El id del Principio de la Oración. Se utiliza para indicar el inicio de una secuencia.
• eos_id: El id del final de la oración. Se utiliza para indicar el final de una secuencia.
• pad_id: El id del relleno. Cuando es necesario alinear varias longitudes de secuencia, se puede usar pad_id para rellenar detrás de secuencias más cortas.

Función de chat chat_completion

Antes de ingresar a la función de generación, veamos cómo se implementa chat_completion.

    def chat_completion(
        self,
        dialogs: List[Dialog],
        temperature: float = 0.6,
        top_p: float = 0.9,
        max_gen_len: Optional[int] = None,
        logprobs: bool = False,
    ) -> List[ChatPrediction]:
        if max_gen_len is None:
            max_gen_len = self.model.params.max_seq_len - 1
        prompt_tokens = []
        for dialog in dialogs:
            if dialog[0]["role"] != "system":
                dialog = [
                    {
    
    
                        "role": "system",
                        "content": DEFAULT_SYSTEM_PROMPT,
                    }
                ] + dialog
            dialog = [
                {
    
    
                    "role": dialog[1]["role"],
                    "content": B_SYS
                    + dialog[0]["content"]
                    + E_SYS
                    + dialog[1]["content"],
                }
            ] + dialog[2:]
            assert all([msg["role"] == "user" for msg in dialog[::2]]) and all(
                [msg["role"] == "assistant" for msg in dialog[1::2]]
            ), (
                "model only supports 'system', 'user' and 'assistant' roles, "
                "starting with 'system', then 'user' and alternating (u/a/u/a/u...)"
            )
            dialog_tokens: List[int] = sum(
                [
                    self.tokenizer.encode(
                        f"{
      
      B_INST} {
      
      (prompt['content']).strip()} {
      
      E_INST} {
      
      (answer['content']).strip()} ",
                        bos=True,
                        eos=True,
                    )
                    for prompt, answer in zip(
                        dialog[::2],
                        dialog[1::2],
                    )
                ],
                [],
            )
            assert (
                dialog[-1]["role"] == "user"
            ), f"Last message must be from user, got {
      
      dialog[-1]['role']}"
            dialog_tokens += self.tokenizer.encode(
                f"{
      
      B_INST} {
      
      (dialog[-1]['content']).strip()} {
      
      E_INST}",
                bos=True,
                eos=False,
            )
            prompt_tokens.append(dialog_tokens)

        generation_tokens, generation_logprobs = self.generate(
            prompt_tokens=prompt_tokens,
            max_gen_len=max_gen_len,
            temperature=temperature,
            top_p=top_p,
            logprobs=logprobs,
        )
        if logprobs:
            return [
                {
    
    
                    "generation": {
    
    
                        "role": "assistant",
                        "content": self.tokenizer.decode(t),
                    },
                    "tokens": [self.tokenizer.decode(x) for x in t],
                    "logprobs": logprobs_i,
                }
                for t, logprobs_i in zip(generation_tokens, generation_logprobs)
            ]
        return [
            {
    
    "generation": {
    
    "role": "assistant", "content": self.tokenizer.decode(t)}}
            for t in generation_tokens
        ]

Veamos primero los parámetros:

• diálogos: una lista de diálogos, donde cada diálogo es una lista de diccionarios que representan un diálogo.
• temperatura: un flotador que representa la temperatura utilizada al generar el texto. El valor predeterminado es 0,6.
• top_p: un flotador que indica el muestreo top-p utilizado al generar el texto. El valor predeterminado es 0,9.
• max_gen_len: un entero opcional que representa la longitud máxima del texto generado. Si no se especifica, se utiliza la longitud de secuencia máxima en los parámetros del modelo menos uno.
• logprobs: un booleano que indica si devolver las probabilidades de registro del texto generado. El valor predeterminado es Falso.
  La función devuelve una lista de ChatPredictions, donde cada elemento es un diccionario que contiene la respuesta generada y la información relacionada.

La función primero verifica si max_gen_len es Ninguno y, de ser así, lo establece en la longitud de secuencia máxima en los parámetros del modelo menos uno. Luego, para cada conversación, la función hace lo siguiente:

• Agregue un aviso del sistema predeterminado al comienzo de la conversación si la función del primer mensaje no es "sistema".
• Combine el primer y el segundo mensaje en uno y actualice la conversación.
• Verifique que los roles de los mensajes en la conversación sean los esperados (es decir, comience con "sistema" y luego alterne entre "usuario" y "asistente").
• Para cada conjunto de indicaciones y respuestas adyacentes (es decir, cada dos mensajes), use un tokenizador para codificarlos y concatenar los tokens codificados.
• Comprueba si el último mensaje es de un usuario.
• Codifique el último mensaje y agregue el token codificado a la lista de tokens.

A continuación, la función llama al método de generación para generar una respuesta y devuelve el resultado correspondiente según el valor del parámetro logprobs. Si logprobs es True, devuelve una lista de diccionarios que contienen respuestas generadas, tokens y probabilidades de registro; de lo contrario, devuelve una lista de diccionarios que contienen solo respuestas generadas. Todas estas respuestas generadas tienen la función de "asistente" y se decodifican mediante un tokenizador.

En términos generales, solo agrega procesamiento de lógica comercial para roles de diálogo, y el núcleo es llamar a la función de generación.

Muestreo de temperatura y p superior

Antes de entrar en la explicación de la función de generación, hablemos de un pequeño punto de conocimiento, que es el papel de la temperatura. Veamos el siguiente código:

            if temperature > 0:
                probs = torch.softmax(logits[:, -1] / temperature, dim=-1)
                next_token = sample_top_p(probs, top_p)
            else:
                next_token = torch.argmax(logits[:, -1], dim=-1)

La temperatura es un hiperparámetro utilizado para controlar la diversidad del texto generado. Cuando la temperatura es más alta, la distribución de probabilidad es más plana, por lo que los marcadores muestreados son más diversos. Cuando la temperatura es baja, la distribución de probabilidad es más nítida, por lo que el marcador muestreado se inclina más hacia el que tiene la probabilidad más alta. Cuando la temperatura sea igual a 0, seleccione directamente la marca con mayor probabilidad.

Entonces, ¿cómo se implementa sample_top_p? Escribí la explicación en el comentario del código:

def sample_top_p(probs, p):
    # 这行代码将输入的概率 probs 按照降序排序。probs_sort 是排序后的概率，probs_idx 是对应的索引。
    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
    # 这行代码计算 probs_sort 的累积和。累积和是从第一个元素开始，依次将序列中的每个元素与前面所有元素的和相加得到的。
    probs_sum = torch.cumsum(probs_sort, dim=-1)
    # 这行代码生成一个布尔掩码，用于指示哪些累积和减去当前概率的值大于 p。这用于确定哪些概率应该被设为0，以保证被抽样的概率和不超过 p。
    mask = probs_sum - probs_sort > p
    # 这行代码使用上述生成的掩码，将那些使累积和减去当前概率的值大于 p 的 probs_sort 中的元素设为0。
    probs_sort[mask] = 0.0
    # 这行代码将 probs_sort 中的每个元素除以它们的和，以便重新归一化概率分布。
    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
    # 这行代码从归一化的 probs_sort 中抽取一个样本。torch.multinomial 是PyTorch中的多项式分布抽样函数，它根据每个元素的权重抽取样本。
    next_token = torch.multinomial(probs_sort, num_samples=1)
    # 这行代码使用 torch.gather 函数从 probs_idx 中收集对应 next_token 的索引，这样就能得到原始概率 probs 中对应的索引。
    next_token = torch.gather(probs_idx, -1, next_token)
    return next_token

En general, sample_top_p conserva la distribución aproximada ordenada por probabilidad, pero filtra el ruido de baja probabilidad de la parte de cola larga. Luego toma muestras de la distribución renormalizada, que conserva la calidad al tiempo que agrega la aleatoriedad adecuada.

generar función

Bueno, finalmente comenzamos a explorar la función de generación de núcleo:

    @torch.inference_mode()
    def generate(
        self,
        prompt_tokens: List[List[int]],
        max_gen_len: int,
        temperature: float = 0.6,
        top_p: float = 0.9,
        logprobs: bool = False,
        echo: bool = False,
    ) -> Tuple[List[List[int]], Optional[List[List[float]]]]:

El primero son los parámetros de esta función, con los que ya estamos familiarizados. Incluye tokens de solicitud de entrada (prompt_tokens), longitud máxima de generación (max_gen_len), parámetro de temperatura (temperatura, que afecta la aleatoriedad del texto generado), top_p (utilizado para determinar el umbral de probabilidad del conjunto de tokens retenido durante el proceso de muestreo, también conocido como " muestreo de núcleo"), si devolver la probabilidad logarítmica de cada token (logprobs) y si devolver el indicador de entrada (eco).

        params = self.model.params
        bsz = len(prompt_tokens)
        assert bsz <= params.max_batch_size, (bsz, params.max_batch_size)

        min_prompt_len = min(len(t) for t in prompt_tokens)
        max_prompt_len = max(len(t) for t in prompt_tokens)
        assert max_prompt_len <= params.max_seq_len
        total_len = min(params.max_seq_len, max_gen_len + max_prompt_len)

        pad_id = self.tokenizer.pad_id
        tokens = torch.full((bsz, total_len), pad_id, dtype=torch.long, device="cuda")
        for k, t in enumerate(prompt_tokens):
            tokens[k, : len(t)] = torch.tensor(t, dtype=torch.long, device="cuda")
        if logprobs:
            token_logprobs = torch.zeros_like(tokens, dtype=torch.float)

        prev_pos = 0
        eos_reached = torch.tensor([False] * bsz, device="cuda")
        input_text_mask = tokens != pad_id

A continuación, inicialice un tensor de tokens con una longitud de total_len según los prompt_tokens proporcionados y complete el pad_id del modelo. Luego, copie el contenido de prompt_tokens en la ubicación correspondiente del tensor de tokens.

        for cur_pos in range(min_prompt_len, total_len):
            logits = self.model.forward(tokens[:, prev_pos:cur_pos], prev_pos)
            if logprobs:
                token_logprobs[:, prev_pos + 1 : cur_pos + 1] = -F.cross_entropy(
                    input=logits.transpose(1, 2),
                    target=tokens[:, prev_pos + 1 : cur_pos + 1],
                    reduction="none",
                    ignore_index=pad_id,
                )
            if temperature > 0:
                probs = torch.softmax(logits[:, -1] / temperature, dim=-1)
                next_token = sample_top_p(probs, top_p)
            else:
                next_token = torch.argmax(logits[:, -1], dim=-1)

            next_token = next_token.reshape(-1)
            # only replace token if prompt has already been generated
            next_token = torch.where(
                input_text_mask[:, cur_pos], tokens[:, cur_pos], next_token
            )
            tokens[:, cur_pos] = next_token
            eos_reached |= (~input_text_mask[:, cur_pos]) & (
                next_token == self.tokenizer.eos_id
            )
            prev_pos = cur_pos
            if all(eos_reached):
                break

Luego, para cada posición en el tensor de tokens, calcule los logits del siguiente token y genere el siguiente token en función de esos logits. Si el parámetro logprobs es verdadero, se calcula la probabilidad de registro de cada token. Si la temperatura es mayor que 0, los logits se escalan usando la función softmax y el parámetro de temperatura, y luego se usa el muestreo top-p para generar el siguiente token. De lo contrario, seleccione directamente el token con los logits más grandes. El token recién generado reemplazará la posición correspondiente en el tensor de tokens.

Si el token generado es un token final (eos_id), actualice el token eos_reached. Si todas las secuencias ya han generado marcadores finales, detenga la generación.

        if logprobs:
            token_logprobs = token_logprobs.tolist()
        out_tokens, out_logprobs = [], []
        for i, toks in enumerate(tokens.tolist()):
            # cut to max gen len
            start = 0 if echo else len(prompt_tokens[i])
            toks = toks[start : len(prompt_tokens[i]) + max_gen_len]
            probs = None
            if logprobs:
                probs = token_logprobs[i][start : len(prompt_tokens[i]) + max_gen_len]
            # cut to eos tok if any
            if self.tokenizer.eos_id in toks:
                eos_idx = toks.index(self.tokenizer.eos_id)
                toks = toks[:eos_idx]
                probs = probs[:eos_idx] if logprobs else None
            out_tokens.append(toks)
            out_logprobs.append(probs)
        return (out_tokens, out_logprobs if logprobs else None)

Finalmente, token_logprobs se convierte en una lista si el parámetro logprobs es verdadero. Luego, para cada fila en el tensor de tokens (es decir, cada secuencia generada), si el parámetro de eco es falso, se elimina la parte de la pista. Luego, si hay una etiqueta final, se elimina la parte posterior a la etiqueta final. Finalmente, se devuelven los tokens generados y las probabilidades de registro (si el parámetro logprobs es verdadero).

Esta función devuelve una tupla, cuyo primer elemento es una lista que contiene cada secuencia de token generada. El segundo elemento es una lista que contiene cada secuencia de probabilidad de registro generada (si el parámetro logprobs es verdadero).

construir constructor

Finalmente, hablemos de la parte de construir Llama:

    @staticmethod
    def build(
        ckpt_dir: str,
        tokenizer_path: str,
        max_seq_len: int,
        max_batch_size: int,
        model_parallel_size: Optional[int] = None,
    ) -> "Llama":
        if not torch.distributed.is_initialized():
            torch.distributed.init_process_group("nccl")
        if not model_parallel_is_initialized():
            if model_parallel_size is None:
                model_parallel_size = int(os.environ.get("WORLD_SIZE", 1))
            initialize_model_parallel(model_parallel_size)

        local_rank = int(os.environ.get("LOCAL_RANK", 0))
        torch.cuda.set_device(local_rank)

        # seed must be the same in all processes
        torch.manual_seed(1)

        if local_rank > 0:
            sys.stdout = open(os.devnull, "w")

        start_time = time.time()
        checkpoints = sorted(Path(ckpt_dir).glob("*.pth"))
        assert len(checkpoints) > 0, f"no checkpoint files found in {
      
      ckpt_dir}"
        assert model_parallel_size == len(
            checkpoints
        ), f"Loading a checkpoint for MP={
      
      len(checkpoints)} but world size is {
      
      model_parallel_size}"
        ckpt_path = checkpoints[get_model_parallel_rank()]
        checkpoint = torch.load(ckpt_path, map_location="cpu")
        with open(Path(ckpt_dir) / "params.json", "r") as f:
            params = json.loads(f.read())

        model_args: ModelArgs = ModelArgs(
            max_seq_len=max_seq_len,
            max_batch_size=max_batch_size,
            **params,
        )
        tokenizer = Tokenizer(model_path=tokenizer_path)
        model_args.vocab_size = tokenizer.n_words
        torch.set_default_tensor_type(torch.cuda.HalfTensor)
        model = Transformer(model_args)
        model.load_state_dict(checkpoint, strict=False)
        print(f"Loaded in {
      
      time.time() - start_time:.2f} seconds")

        return Llama(model, tokenizer)

Aunque es una sección tan grande, en realidad es un trabajo de inicialización.

• Configuración distribuida: en primer lugar, este código comprueba si el entorno distribuido de PyTorch se ha inicializado y, en caso contrario, lo inicializa. Luego, verifique si el entorno paralelo del modelo se ha inicializado; de lo contrario, obtenga el valor de la variable de entorno WORLD_SIZE como el tamaño paralelo del modelo e inicialícelo.

• Configuración del dispositivo: obtenga el valor de la variable de entorno LOCAL_RANK como el rango local y configure el dispositivo actual como la GPU correspondiente al rango.

• Configuración de semilla aleatoria: para asegurarse de que todos los procesos generen el mismo número aleatorio, establezca la semilla aleatoria en 1.

• Configuración de salida estándar: si el rango local es mayor que 0, redirija la salida estándar a un dispositivo vacío, es decir, no muestre ninguna salida.

• Cargar modelo de punto de control: busque todos los archivos de punto de control en el directorio de puntos de control, ordenados por nombre de archivo. Luego, se selecciona un archivo de punto de control basado en la clasificación paralela del modelo y se carga el punto de control. Luego, cargue los parámetros del modelo.

• Cree un modelo y tokenizador: cree un objeto de parámetro de modelo utilizando los parámetros del modelo cargado y el max_seq_len y max_batch_size proporcionados. Luego, cargue el tokenizador y establezca el tamaño del vocabulario del parámetro del modelo en el tamaño del vocabulario del tokenizador. Luego, establezca el tipo de tensor predeterminado en punto flotante de precisión media (para ahorrar memoria y recursos informáticos). Luego, construya el modelo de Transformador y cargue el punto de control del modelo.

• Finalmente, construya un objeto Llama que contenga el modelo cargado y el tokenizador, y devuelva ese objeto.

resumen

En esta sección, hemos aprendido el código fuente de LLaMA 2, incluida la implementación de la función de finalización text_completion y la función de chat chat_completion, y el principio de su implementación real de la función de generación. También aprendimos el principio de la temperatura y el muestreo de p superior.

Puede ser un poco difícil para los estudiantes que no han obtenido una generación de aprendizaje profundo.

El código de LLaMA sigue siendo parte del modelo, del que hablaremos en la siguiente sección. De lo contrario, hay demasiados puntos de conocimiento y todos son propensos a la hipoxia cerebral :)