Recuperación Mejorada de Generación con Huggingface Transformers y Ray

Recuperación Mejorada de Generación con Huggingface Transformers y Ray Improved Generation Retrieval with Huggingface Transformers and Ray.

Una publicación de blog invitada de Amog Kamsetty del equipo de Anyscale

Huggingface Transformers recientemente agregó el modelo Retrieval Augmented Generation (RAG), una nueva arquitectura de procesamiento del lenguaje natural (NLP) que aprovecha documentos externos (como Wikipedia) para aumentar su conocimiento y lograr resultados de vanguardia en tareas intensivas en conocimiento. En esta publicación de blog, presentamos la integración de Ray, una biblioteca para construir aplicaciones escalables, en el mecanismo de recuperación de documentos contextuales de RAG. Esto acelera las llamadas de recuperación en un 2x y mejora la escalabilidad del ajuste fino distribuido de RAG.

¿Qué es Retrieval Augmented Generation (RAG)?

Una visión general de RAG. El modelo recupera documentos contextuales de un conjunto de datos externo como parte de su ejecución. Estos documentos contextuales se utilizan junto con la entrada original para producir una salida. El GIF se tomó de la publicación de blog original de Facebook.

Recientemente, Huggingface se asoció con Facebook AI para introducir el modelo RAG como parte de su biblioteca Transformers.

• Hugging Face Reads, Feb. 2021 – Transformers de largo alcance
• Comprendiendo la Atención Esparsa por Bloques de BigBird
• Entrenamiento distribuido Entrena BART/T5 para resumir utilizando 🤗 Transformers y Amazon SageMaker

RAG funciona igual que cualquier otro modelo seq2seq. Sin embargo, RAG tiene un componente intermedio que recupera documentos contextuales de una base de conocimientos externa (como un corpus de texto de Wikipedia). Estos documentos se utilizan luego junto con la secuencia de entrada y se pasan al generador subyacente seq2seq.

Este paso de recuperación de información permite que RAG utilice múltiples fuentes de conocimiento: aquellas que están integradas en los parámetros del modelo y la información contenida en los pasajes contextuales, lo que le permite superar a otros modelos de vanguardia en tareas como la respuesta a preguntas. ¡Puedes probarlo por ti mismo utilizando esta demostración proporcionada por Huggingface!

Escalando el ajuste fino

Esta recuperación de documentos contextuales es crucial para los resultados de vanguardia de RAG, pero introduce una capa adicional de complejidad. Al escalar el proceso de entrenamiento mediante una rutina de entrenamiento de datos paralelos, una implementación ingenua de la búsqueda de documentos puede convertirse en un cuello de botella para el entrenamiento. Además, el índice de documentos utilizado en el componente de recuperación a menudo es bastante grande, lo que hace impracticable que cada trabajador de entrenamiento cargue su propia copia replicada del índice.

La implementación anterior del ajuste fino de RAG aprovechó el paquete de comunicación torch.distributed para la parte de recuperación de documentos. Sin embargo, esta implementación a veces resultó inflexible y limitada en escalabilidad.

En cambio, se requiere una implementación independiente del marco y más flexible para la programación concurrente ad hoc. Ray encaja perfectamente en este sentido. Ray es una biblioteca simple pero poderosa de Python para programación distribuida y paralela de propósito general. Al usar Ray para la recuperación de documentos distribuida, logramos una aceleración de 2x por llamada de recuperación en comparación con torch.distributed, y una mejor escalabilidad en general del ajuste fino.

Ray para la recuperación de documentos

Recuperación de documentos con la implementación de torch.distributed

La principal desventaja de la implementación de torch.distributed para la recuperación de documentos fue que se aferraba al mismo grupo de procesos utilizado para el entrenamiento y solo el trabajador de entrenamiento de rango 0 cargaba el índice en memoria.

Como resultado, esta implementación tenía algunas limitaciones:

1. Cuello de botella de sincronización: El trabajador de rango 0 debía recibir las entradas de todos los trabajadores, realizar la consulta del índice y luego enviar los resultados a los demás trabajadores. Esto limitaba el rendimiento con múltiples trabajadores de entrenamiento.
2. Específico de PyTorch: El grupo de procesos de recuperación de documentos debía aferrarse al grupo de procesos existente utilizado para el entrenamiento, lo que significa que también se debía utilizar PyTorch para el entrenamiento.

Recuperación de documentos con la implementación de Ray

Para superar estas limitaciones, introdujimos una nueva implementación de recuperación distribuida basada en Ray. Con las abstracciones de actores persistentes de Ray, se utilizan múltiples procesos separados de los procesos de entrenamiento para cargar el índice y manejar las consultas de recuperación. Con múltiples actores de Ray, la recuperación ya no es un cuello de botella y PyTorch ya no es un requisito para RAG.

Y como se puede ver a continuación, el uso de la implementación basada en Ray lleva a un mejor rendimiento de recuperación para el ajuste fino con múltiples GPU. Los siguientes resultados muestran los segundos por llamada de recuperación y podemos ver que a medida que aumentamos el número de GPU en las que entrenamos, el uso de Ray tiene un rendimiento comparativamente mejor que torch.distributed. Además, si aumentamos el número de procesos de Ray que realizan la recuperación, también obtenemos un mejor rendimiento con más trabajadores de entrenamiento, ya que un solo proceso de recuperación ya no es un cuello de botella.

	2 GPU	3 GPU	4 GPU
torch.distributed	2.12 seg/recuperación	2.62 seg/recuperación	3.438 seg/recuperación
Ray 2 procesos de recuperación	1.49 seg/recuperación	1.539 seg/recuperación	2.029 seg/recuperación
Ray 4 procesos de recuperación	1.145 seg/recuperación	1.484 seg/recuperación	1.66 seg/recuperación

Una comparación de rendimiento de diferentes implementaciones de recuperación. Para cada implementación de recuperación de documentos, ejecutamos 500 pasos de entrenamiento con un tamaño de lote por GPU de 8, y medimos el tiempo que tarda en recuperar los documentos contextuales para cada lote en el trabajador de entrenamiento de rango 0. Como muestran los resultados, el uso de múltiples procesos de recuperación mejora el rendimiento, especialmente a medida que escalamos el entrenamiento a múltiples GPU.

¿Cómo lo uso?

Huggingface proporciona un script de ajuste fino basado en PyTorch Lightning, y lo extendimos para agregar la implementación de recuperación de Ray como opción.

Para probarlo, primero instala los requisitos necesarios

pip install ray
pip install transformers
pip install -r transformers/examples/research_projects/rag/requirements.txt

Luego, puedes especificar las rutas de tus datos y otras configuraciones y ejecutar ¡finetune-rag-ray.sh!

# Script de ejemplo para ajustar fino RAG usando Ray para recuperación distribuida.

# Agrega el directorio padre al camino de Python para acceder a lightning_base.py
export PYTHONPATH="../":"${PYTHONPATH}"

# Inicia un clúster Ray de un solo nodo.
ray start --head

# Un ejemplo de ajuste fino, debes especificar data_dir, output_dir y model_name_or_path
# ejecuta ./examples/rag/finetune_rag_ray.sh --help para ver todas las opciones posibles

python examples/rag/finetune_rag.py \
    --data_dir $DATA_DIR \
    --output_dir $OUTPUT_DIR \
    --model_name_or_path $MODEL_NAME_OR_PATH \
    --model_type rag_sequence \
    --fp16 \
    --gpus 8 \
    --profile \
    --do_train \
    --do_predict \
    --n_val -1 \
    --train_batch_size 8 \
    --eval_batch_size 1 \
    --max_source_length 128 \
    --max_target_length 25 \
    --val_max_target_length 25 \
    --test_max_target_length 25 \
    --label_smoothing 0.1 \
    --dropout 0.1 \
    --attention_dropout 0.1 \
    --weight_decay 0.001 \
    --adam_epsilon 1e-08 \
    --max_grad_norm 0.1 \
    --lr_scheduler polynomial \
    --learning_rate 3e-05 \
    --num_train_epochs 100 \
    --warmup_steps 500 \
    --gradient_accumulation_steps 1 \
    --distributed_retriever ray \
    --num_retrieval_workers 4

# Detén el clúster Ray.
ray stop

¿Qué sigue?

Usando RAG con Huggingface transformers y la implementación de recuperación de Ray para un ajuste fino distribuido más rápido, puedes aprovechar RAG para la generación basada en recuperación en tus propias tareas intensivas en conocimiento.

Además, la sintonización de hiperparámetros es otro aspecto del ajuste fino de transformers y puede tener un gran impacto en la precisión. Para una sintonización de hiperparámetros escalable y fácil, echa un vistazo a la biblioteca Ray Tune. Al utilizar la integración de Ray Tune con PyTorch Lightning, o la integración incorporada con Huggingface transformers, puedes ejecutar experimentos para encontrar los hiperparámetros perfectos para tu modelo RAG.

Y por último, ¡mantente atento a una posible implementación de RAG en Tensorflow en Huggingface!

Si planeas probar la integración RAG+Ray, no dudes en compartir tus experiencias en el foro de Ray Discourse o unirte a la comunidad de Ray Slack para continuar la discusión, ¡nos encantaría saber de ti!

También publicado en https://medium.com/distributed-computing-with-ray/retrieval-augmented-generation-with-huggingface-transformers-and-ray-b09b56161b1e

We will continue to update Zepes; if you have any questions or suggestions, please contact us!