Ajuste Llama 2 usando QLoRA e impleméntelo en Amazon SageMaker con AWS Inferentia2

En esta publicación, mostramos el ajuste de un modelo Llama 2 utilizando un método de ajuste fino con eficiencia de parámetros (PEFT) e implementamos el modelo ajustado en AWS Inferentia2. Utilizamos el kit de desarrollo de software (SDK) AWS Neuron para acceder al dispositivo AWS Inferentia2 y beneficiarnos de su alto rendimiento. Luego utilizamos un contenedor de inferencia de modelos grande impulsado por ai/” target=”_blank” rel=”noopener”>Biblioteca Java profunda (DJLServing) como nuestra solución de servicio de modelos.

Descripción general de la solución

Ajuste eficiente de Llama2 usando QLoRa

La familia Llama 2 de modelos de lenguaje grande (LLM) es una colección de modelos de texto generativo previamente entrenados y ajustados que varían en escala de 7 mil millones a 70 mil millones de parámetros. Llama 2 fue preentrenada con 2 billones de tokens de datos de fuentes disponibles públicamente. Los clientes de AWS a veces optan por ajustar los modelos de Llama 2 utilizando los propios datos de los clientes para lograr un mejor rendimiento en las tareas posteriores. Sin embargo, debido a la gran cantidad de parámetros del modelo Llama 2, un ajuste completo podría resultar prohibitivamente costoso y consumir mucho tiempo. El enfoque de ajuste fino eficiente de parámetros (PEFT) puede abordar este problema ajustando solo una pequeña cantidad de parámetros adicionales del modelo mientras congela la mayoría de los parámetros del modelo previamente entrenado. Para obtener más información sobre PEFT, puede leer esto correo. En esta publicación utilizamos QLoRa para poner a punto un modelo Llama 2 7B.

Implemente un modelo ajustado en Inf2 utilizando Amazon SageMaker

AWS Inferentia2 es un acelerador de aprendizaje automático (ML) especialmente diseñado para cargas de trabajo de inferencia y ofrece un alto rendimiento con un costo hasta un 40 % menor para cargas de trabajo de IA generativa y LLM en comparación con otras instancias optimizadas de inferencia en AWS. En esta publicación, utilizamos la instancia Inf2 de Amazon Elastic Compute Cloud (Amazon EC2), que presenta AWS Inferentia2, los aceleradores Inferentia2 de segunda generación, cada uno de los cuales contiene dos NeuronCores-v2. Cada NeuronCore-v2 es una unidad informática heterogénea e independiente, con cuatro motores principales: motores tensoriales, vectoriales, escalares y GPSIMD. Incluye una memoria SRAM administrada por software en chip para maximizar la localidad de los datos. Dado que se han publicado varios blogs sobre Inf2, el lector puede consultar esta publicación y nuestra documentación para obtener más información sobre Inf2.

Para implementar modelos en Inf2, necesitamos AWS Neuron SDK como capa de software que se ejecuta sobre el hardware de Inf2. AWS Neuron es el SDK que se utiliza para ejecutar cargas de trabajo de aprendizaje profundo en instancias basadas en AWS Inferentia y AWS Trainium. Permite el ciclo de vida de desarrollo de ML de un extremo a otro para crear nuevos modelos, entrenarlos y optimizarlos e implementarlos para producción. AWS Neuron incluye un aprendizaje profundo compilador, tiempo de ejecucióny herramientas que se integran de forma nativa con marcos populares como TensorFlow y PyTorch. En este blog vamos a utilizar transformers-neuronx, que forma parte del SDK de AWS Neuron para flujos de trabajo de inferencia de decodificadores de transformadores. Él apoya una gama de modelos populares, incluido Llama 2.

Para implementar modelos en Amazon SageMaker, normalmente utilizamos un contenedor que contiene las bibliotecas necesarias, como Neuron SDK y transformers-neuronx así como el componente de servicio de modelos. Amazon SageMaker mantiene contenedores de aprendizaje profundo (DLC) con bibliotecas populares de código abierto para alojar modelos grandes. En esta publicación utilizamos el Contenedor de inferencia de modelo grande para Neuron. Este contenedor tiene todo lo que necesitas para implementar tu modelo Llama 2 en Inf2. Para obtener recursos para comenzar con LMI en Amazon SageMaker, consulte muchas de nuestras publicaciones existentes (blog 1, blog 2, blog 3) sobre este tema. En resumen, puede ejecutar el contenedor sin escribir ningún código adicional. Puedes usar el controlador predeterminado para una experiencia de usuario perfecta y pase uno de los nombres de modelo admitidos y cualquier parámetro configurable de tiempo de carga. Esto compila y ofrece un LLM en una instancia de Inf2. Por ejemplo, para implementar OpenAssistant/llama2-13b-orca-8k-3319puede proporcionar la siguiente configuración (como serving.properties archivo). En serving.propertiesespecificamos el tipo de modelo como llama2-13b-orca-8k-3319, el tamaño del lote es 4, el grado paralelo del tensor es 2, y eso es todo. Para obtener la lista completa de parámetros configurables, consulte Todas las opciones de configuración de DJL.

# Engine to use: MXNet, PyTorch, TensorFlow, ONNX, PaddlePaddle, DeepSpeed, etc.
engine = Python 
# default handler for model serving
option.entryPoint = djl_python.transformers_neuronx
# The Hugging Face ID of a model or the s3 url of the model artifacts. 
option.model_id = meta-llama/Llama-2-7b-chat-hf
#the dynamic batch size, default is 1.
option.batch_size=4
# This option specifies number of tensor parallel partitions performed on the model.
option.tensor_parallel_degree=2
# The input sequence length
option.n_positions=512
#Enable iteration level batching using one of "auto", "scheduler", "lmi-dist"
option.rolling_batch=auto
# The data type to which you plan to cast the model default
option.dtype=fp16
# worker load model timeout
option.model_loading_timeout=1500

Alternativamente, puede escribir su propio archivo de controlador de modelo como se muestra en este ejemplopero eso requiere implementar los métodos de carga e inferencia del modelo para que sirvan como puente entre las API de DJLServing.

Requisitos previos

La siguiente lista describe los requisitos previos para implementar el modelo descrito en esta publicación de blog. Puede implementar desde la Consola de administración de AWS o utilizando la última versión de la interfaz de línea de comandos de AWS (AWS CLI).

Tutorial

En la siguiente sección, analizaremos el código en dos partes:

1. Ajuste un modelo Llama2-7b y cargue los artefactos del modelo en una ubicación de depósito de Amazon S3 especificada.
2. Implemente el modelo en Inferentia2 utilizando el contenedor de servicio DJL alojado en Amazon SageMaker.

Los ejemplos de código completos con instrucciones se pueden encontrar en este GitHub repositorio.

Parte 1: ajustar un modelo Llama2-7b usando PEFT

Vamos a utilizar el método recientemente introducido en el artículo. QLoRA: Ajuste del adaptador de bajo rango con reconocimiento de cuantificación para la generación de lenguaje por Tim Dettmers et al. QLoRA es una nueva técnica para reducir la huella de memoria de modelos de lenguaje grandes durante el ajuste fino, sin sacrificar el rendimiento.

Nota: El ajuste fino del modelo llama2-7b que se muestra a continuación se probó en una computadora portátil Amazon SageMaker Studio con un kernel optimizado para GPU Python 2.0 utilizando un ml.g5.2xgrande tipo de instancia. Como práctica recomendada, recomendamos utilizar un entorno de desarrollo integrado (IDE) de Amazon SageMaker Studio lanzado en su propia nube privada virtual de Amazon (Amazon VPC). Esto le permite controlar, monitorear e inspeccionar el tráfico de red dentro y fuera de su VPC utilizando capacidades de seguridad y redes estándar de AWS. Para obtener más información, consulte Protección de la conectividad de Amazon SageMaker Studio mediante una VPC privada.

Cuantizar el modelo base

Primero cargamos un modelo cuantificado con cuantificación de 4 bits usando Transformadores de cara abrazada biblioteca de la siguiente manera:

# The base pretrained model for fine-tuning
model_name = "NousResearch/Llama-2-7b-chat-hf"

# The instruction dataset to use
dataset_name = "mlabonne/guanaco-llama2-1k"

#Activate 4-bit precision base model loading
use_4bit = True
bnb_4bit_compute_dtype = "float16"
bnb_4bit_quant_type = "nf4"
use_nested_quant = False

compute_dtype = getattr(torch, bnb_4bit_compute_dtype)

bnb_config = BitsAndBytesConfig(
load_in_4bit=use_4bit,
bnb_4bit_quant_type=bnb_4bit_quant_type,
bnb_4bit_compute_dtype=compute_dtype,
bnb_4bit_use_double_quant=use_nested_quant,
)

# Load base model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map=device_map
)
model.config.pretraining_tp = 1

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

Cargar conjunto de datos de entrenamiento

A continuación, cargamos el conjunto de datos para alimentar el modelo y realizar el paso de ajuste que se muestra a continuación:

# Load dataset (you can process it here)
dataset = load_dataset(dataset_name, split="train")

Adjuntar una capa adaptadora

Aquí adjuntamos una pequeña capa adaptadora entrenable, configurada como LoraConfig definido en la cara de abrazo peft biblioteca.

# include linear layers to apply LoRA to.
modules = find_all_linear_names(model)

## Setting up LoRA configuration
lora_r = 64

# Alpha parameter for LoRA scaling
lora_alpha = 16

# Dropout probability for LoRA layers
lora_dropout = 0.1

peft_config = LoraConfig(
lora_alpha=lora_alpha,
lora_dropout=lora_dropout,
r=lora_r,
bias="none",
task_type="CAUSAL_LM",
target_modules=modules)

Entrenar un modelo

Usando la configuración LoRA que se muestra arriba, ajustaremos el modelo Llama2 junto con los hiperparámetros. A continuación se muestra un fragmento de código para entrenar el modelo:

# Set training parameters
training_arguments = TrainingArguments(...)

trainer = SFTTrainer(
model=model,
train_dataset=dataset,
peft_config=peft_config, # LoRA config
dataset_text_field="text",
max_seq_length=max_seq_length,
tokenizer=tokenizer,
args=training_arguments,
packing=packing,
)

# Train model
trainer.train()

# Save trained model
trainer.model.save_pretrained(new_model)

Fusionar peso del modelo

El modelo ajustado ejecutado anteriormente creó un nuevo modelo que contiene los pesos del adaptador LoRA entrenados. En el siguiente fragmento de código, fusionaremos el adaptador con el modelo base para poder usar el modelo ajustado para la inferencia.

# Reload model in FP16 and merge it with LoRA weights
base_model = AutoModelForCausalLM.from_pretrained(
model_name,
low_cpu_mem_usage=True,
return_dict=True,
torch_dtype=torch.float16,
device_map=device_map,
)
model = PeftModel.from_pretrained(base_model, new_model)
model = model.merge_and_unload()

save_dir = "merged_model"
model.save_pretrained(save_dir, safe_serialization=True, max_shard_size="2GB")

# Reload tokenizer to save it
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"
tokenizer.save_pretrained(save_dir)

Cargar el peso del modelo en Amazon S3

En el paso final de la parte 1, guardaremos los pesos del modelo combinado en una ubicación especificada de Amazon S3. El peso del modelo lo utilizará un contenedor de servicio de modelos en Amazon SageMaker para alojar el modelo mediante una instancia de Inferentia2.

model_data_s3_location = "s3://<bucket_name>/<prefix>/"
!cd {save_dir} && aws s3 cp —recursive . {model_data_s3_location}

Parte 2: Alojar el modelo QLoRA para inferencia con AWS Inf2 utilizando SageMaker LMI Container

En esta sección, recorreremos los pasos para implementar un modelo ajustado de QLoRA en un entorno de alojamiento de Amazon SageMaker. Usaremos un sirviendo DJL contenedor de SageMaker contenido descargableque se integra con el transformadores-neuronx biblioteca para alojar este modelo. La configuración facilita la carga de modelos en aceleradores AWS Inferentia2, paraleliza el modelo en múltiples NeuronCores y permite el servicio a través de puntos finales HTTP.

Preparar artefactos modelo

DJL admite muchas bibliotecas de optimización de aprendizaje profundo, incluidas DeepSpeed, FasterTransformer y más. Para configuraciones específicas del modelo, proporcionamos un serving.properties con parámetros clave, como tensor_parallel_degree y model_id para definir las opciones de carga del modelo. El model_id podría ser un ID de modelo de Hugging Face o una ruta de Amazon S3 donde se almacenan los pesos del modelo. En nuestro ejemplo, proporcionamos la ubicación de Amazon S3 de nuestro modelo ajustado. El siguiente fragmento de código muestra las propiedades utilizadas para la publicación del modelo:

%%writefile serving.properties
engine=Python
option.entryPoint=djl_python.transformers_neuronx
option.model_id=<model data s3 location>
option.batch_size=4
option.neuron_optimize_level=2
option.tensor_parallel_degree=8
option.n_positions=512
option.rolling_batch=auto
option.dtype=fp16
option.model_loading_timeout=1500

Por favor consulte esto documentación para obtener más información sobre las opciones configurables disponibles a través de serving.properties. Tenga en cuenta que utilizamos option.n_position=512 en este blog para una compilación más rápida de AWS Neuron. Si desea probar una longitud de token de entrada mayor, recomendamos al lector precompilar el modelo con anticipación (consulte Modelo de precompilación de AOT en EC2). De lo contrario, podría encontrarse con un error de tiempo de espera si el tiempo de compilación es demasiado.

Después de la serving.properties archivo está definido, empaquetaremos el archivo en un tar.gz formato, como sigue:

%%sh
mkdir mymodel
mv serving.properties mymodel/
tar czvf mymodel.tar.gz mymodel/
rm -rf mymodel

Luego, cargaremos el archivo tar.gz en una ubicación de depósito de Amazon S3:

s3_code_prefix = "large-model-lmi/code"
bucket = sess.default_bucket()  # bucket to house artifacts
code_artifact = sess.upload_data("mymodel.tar.gz", bucket, s3_code_prefix)
print(f"S3 Code or Model tar ball uploaded to --- > {code_artifact}")

Crear un punto final de modelo de Amazon SageMaker

Para utilizar una instancia de Inf2 para el servicio, utilizamos un contenedor LMI de Amazon SageMaker con soporte DJL neuronX. Consulte esta publicación para obtener más información sobre el uso de un contenedor DJL NeuronX para inferencia. El siguiente código muestra cómo implementar un modelo utilizando el SDK de Python de Amazon SageMaker:

# Retrieves the DJL-neuronx docker image URI
image_uri = image_uris.retrieve(
framework="djl-neuronx",
region=sess.boto_session.region_name,
version="0.24.0"
)

# Define inf2 instance type to use for serving
instance_type = "ml.inf2.48xlarge"

endpoint_name = sagemaker.utils.name_from_base("lmi-model")

# Deploy the model for inference
model.deploy(initial_instance_count=1,
instance_type=instance_type,
container_startup_health_check_timeout=1500,
volume_size=256,
endpoint_name=endpoint_name)

# our requests and responses will be in json format so we specify the serializer and the deserializer
predictor = sagemaker.Predictor(
endpoint_name=endpoint_name,
sagemaker_session=sess,
serializer=serializers.JSONSerializer(),
)

Punto final del modelo de prueba

Una vez que el modelo se haya implementado correctamente, podemos validar el punto final enviando una solicitud de muestra al predictor:

prompt="What is machine learning?"
input_data = f"<s>(INST) <<SYS>>\nAs a data scientist\n<</SYS>>\n{prompt} (/INST)"

response = predictor.predict(
{"inputs": input_data, "parameters": {"max_new_tokens":300, "do_sample":"True"}}
)

print(json.loads(response)('generated_text'))

El resultado de muestra se muestra a continuación:

En el contexto del análisis de datos, el aprendizaje automático (ML) se refiere a una técnica estadística capaz de extraer poder predictivo de un conjunto de datos con una complejidad y precisión cada vez mayores al reducir iterativamente el alcance de una estadística.

Machine Learning no es una técnica estadística nueva, sino más bien una combinación de técnicas existentes. Además, no ha sido diseñado para usarse con un conjunto de datos específico o para producir un resultado específico. Más bien, fue diseñado para ser lo suficientemente flexible como para adaptarse a cualquier conjunto de datos y hacer predicciones sobre cualquier resultado.

Limpiar

Si decide que ya no desea mantener el punto final de SageMaker en ejecución, puede eliminarlo mediante AWS SDK para Python (boto3), AWS CLI o Amazon SageMaker Console. Además, también puede cerrar los recursos de Amazon SageMaker Studio que ya no sean necesarios.

Conclusión

En esta publicación, le mostramos cómo ajustar un modelo Llama2-7b usando un adaptador LoRA con cuantificación de 4 bits usando una sola instancia de GPU. Luego implementamos el modelo en una instancia de Inf2 alojada en Amazon SageMaker utilizando un contenedor de servicio DJL. Finalmente, validamos el punto final del modelo de Amazon SageMaker con una predicción de generación de texto utilizando el SDK de Python de SageMaker. Continúe y pruébelo, nos encantaría escuchar sus comentarios. Esté atento a las actualizaciones sobre más capacidades y nuevas innovaciones con AWS Inferentia.

Para obtener más ejemplos sobre AWS Neuron, consulte muestras-de-neuronas-aws.

Sobre los autores

wei teh es arquitecto senior de soluciones especializado en IA/ML en AWS. Le apasiona ayudar a los clientes a avanzar en su recorrido por AWS, centrándose en los servicios de Amazon Machine Learning y las soluciones basadas en aprendizaje automático. Fuera del trabajo, disfruta de actividades al aire libre como acampar, pescar y hacer caminatas con su familia.

Qingwey li es especialista en aprendizaje automático en Amazon Web Services. Recibió su doctorado. en Investigación de Operaciones después de arruinar la cuenta de becas de investigación de su asesor y no entregar el Premio Nobel que prometió. Actualmente, ayuda a los clientes de la industria de seguros y servicios financieros a crear soluciones de aprendizaje automático en AWS. En su tiempo libre le gusta leer y enseñar.