Los chips de inteligencia artificial de AWS ofrecen alto rendimiento y bajo costo para los modelos Llama 3.1 en AWS

Hoy estamos emocionados de anunciar AWS Trainium y Inferencia de AWS Soporte para el ajuste fino y la inferencia de los modelos Llama 3.1. La familia Llama 3.1 de modelos de lenguajes grandes multilingües (LLM) es una colección de modelos generativos preentrenados y ajustados por instrucción en tamaños 8B, 70B y 405B. Publicación anteriorCubrimos cómo implementar modelos de Llama 3 en instancias basadas en AWS Trainium e Inferentia en Inicio rápido de Amazon SageMakerEn esta publicación, describimos cómo comenzar a ajustar e implementar la familia de modelos Llama 3.1 en chips de IA de AWS para aprovechar sus beneficios en cuanto a precio y rendimiento.

Descripción general de los modelos Llama 3.1

La familia de modelos generativos multilingües Llama 3.1 es una colección de modelos generativos preentrenados y ajustados a las instrucciones en tamaños 8B, 70B y 405B (texto entrante/texto saliente y código saliente). Todos los modelos admiten contextos extensos (128k) y están optimizados para la inferencia con soporte para atención de consultas agrupadas (GQA).

Los modelos optimizados de instrucciones Llama 3.1 (8B, 70B, 405B) están optimizados para casos de uso de diálogo multilingüe y superan a muchos de los modelos de chat disponibles públicamente en los puntos de referencia comunes de la industria. Se han entrenado para generar llamadas a herramientas para algunas herramientas específicas para capacidades como búsqueda, generación de imágenes, ejecución de código y razonamiento matemático. Además, admiten el uso de herramientas de disparo cero.

Llama 3.1 405B es el modelo de aprendizaje a largo plazo más grande del mundo disponible públicamente según Meta. El modelo establece un nuevo estándar para la inteligencia artificial (IA) y es ideal para aplicaciones de nivel empresarial e investigación y desarrollo. Es ideal para tareas como la generación de datos sintéticos, donde los resultados del modelo se pueden utilizar para mejorar modelos Llama más pequeños después de un ajuste fino, y destilaciones de modelos para transferir conocimiento a modelos más pequeños desde el modelo 405B. Este modelo se destaca en conocimientos generales, generación de textos extensos, traducción multilingüe, traducción automática, codificación, matemáticas, uso de herramientas, comprensión contextual mejorada y razonamiento y toma de decisiones avanzados.

En términos arquitectónicos, el núcleo LLM para Llama 3 y Llama 3.1 tiene la misma arquitectura densa. Son modelos de lenguaje autorregresivos que utilizan una arquitectura de transformador optimizada. Las versiones optimizadas utilizan un ajuste fino supervisado (SFT) y aprendizaje de refuerzo con retroalimentación humana (RLHF) para alinearse con las preferencias humanas de ayuda y seguridad.

El Guía de uso responsable de Meta Puede ayudarle a implementar ajustes adicionales que puedan ser necesarios para personalizar y optimizar los modelos con mitigaciones de seguridad adecuadas.

Trainium impulsa Llama 3.1 en Amazon Bedrock y Amazon SageMaker

La forma más rápida de Comience a utilizar Llama 3.1 en AWS es a través de Amazon Bedrock, que funciona con nuestra infraestructura de IA especialmente diseñada, que incluye AWS Trainium. A través de su API completamente administrada, La roca madre del Amazonas ofrece los beneficios de nuestra infraestructura de IA especialmente diseñada y simplifica el acceso a estos modelos potentes para que usted pueda concentrarse en crear aplicaciones de IA diferenciadas.

Si necesita un mayor control sobre los recursos subyacentes, puede Ajuste e implemente modelos de Llama 3.1 con SageMakerPróximamente se incorporará soporte de Trainium para Llama 3.1 en SageMaker JumpStart.

AWS Trainium y AWS Inferentia2 permiten un alto rendimiento y un bajo coste para los modelos Llama 3.1

Si desea crear sus propias canalizaciones de ML para entrenamiento e inferencia para lograr una mayor flexibilidad y control, puede comenzar con Llama 3.1 en chips de IA de AWS usando Nube de computación elástica de Amazon (Amazon EC2) Instancias Trn1 e Inf2. Veamos cómo puedes comenzar a utilizar los nuevos modelos Llama 3.1 8/70B en Trainium usando SDK de AWS Neuron.

Ajuste de Llama 3.1 en Trainium

Para comenzar a ajustar Llama 3.1 8B o Llama 3.1 70B, puede utilizar el NeuronX distribuido Biblioteca NeuronX Distributed ofrece implementaciones de algunas de las técnicas de inferencia y entrenamiento distribuido más populares. Para comenzar a realizar ajustes, puede utilizar los siguientes ejemplos:

Ambas muestras están construidas sobre Clúster paralelo de AWS para administrar la infraestructura del clúster Trainium y Slurm para la administración de la carga de trabajo. El siguiente es el comando Slurm de ejemplo para iniciar el entrenamiento para Llama3.1 70B:

sbatch --exclusive \
--nodes 32 \
--cpus-per-task 128 \
--wrap="srun bash $(pwd)/run_llama3_70B_tp_pp.sh"

Dentro del script Slurm, lanzamos un proceso de entrenamiento distribuido en nuestro clúster. En los scripts de ejecución, cargamos los pesos y la configuración preentrenados que proporciona Meta, y lanzamos el proceso de entrenamiento:

torchrun $DISTRIBUTED_ARGS run_llama_nxd.py \
    —train_batch_size $BS \
    —use_meta_device_init 1 \
    —training_dir $DATA_PATH \
    —training_config $SCRIPT_DIR/${MODEL_SIZE}config_llama${LLAMA_VERSION} \
    —max_steps $max_steps \
    —seq_len $SEQ_LEN \
    —pipeline_parallel_size $PP_DEGREE \
    —tensor_parallel_size $TP_DEGREE \
    —num_microbatches $NUM_MICROBATCHES \
    —lr 0.000015 \
    —min_lr 1e-06 \
    —beta1 0.9 \
    —beta2 0.95 \
    —weight_decay 0.1 \
    —warmup_steps 2000 \
    —constant_steps 0 \
    —use_zero1_optimizer 1 \
    —use_selective_checkpoint 1 \
    —use_flash_attention 1 \
    —qkv_linear 1 \
    —kv_replicator 4 \
    —pretrained_weight 1 \
    —save_load_xser 1 \
    —checkpoint_dir "/shared/llama${LLAMA_VERSION}${MODEL_SIZE}/" \
    —checkpoint_freq $checkpoint_freq \
    —num_kept_checkpoint -1 \
    —loading_step -1 \
    —tb_dir $tb_dir |& tee $LOG_PATH/log
exit ${PIPESTATUS[0]}

Implementar Llama 3.1 en Trainium

Cuando su modelo esté listo para implementarse, puede hacerlo actualizando el ID del modelo en el código de muestra anterior de Llama 3 8B Neuron:

model_id = "meta-llama/Meta-Llama-3.1-8B"
neuron_model = LlamaForSampling.from_pretrained(model_id, neuron_config=neuron_config, batch_size=1, tp_degree=24, amp='bf16', n_positions=4096)
neuron_model.to_neuron()

Puede utilizar el mismo código de inferencia de muestra:

tokenizer = AutoTokenizer.from_pretrained(model_id)
prompt = "Hello, I'm a language model and I like to"
input_ids = tokenizer.encode(prompt, return_tensors="pt")

# run inference with top-k sampling
with torch.inference_mode():
    start = time.time()
    generated_sequences = neuron_model.sample(input_ids, sequence_length=2048, top_k=50)
    elapsed = time.time() - start

generated_sequences = [tokenizer.decode(seq) for seq in generated_sequences]
print(f'generated sequences {generated_sequences} in {elapsed} seconds')

Para obtener detalles paso a paso, consulte los nuevos ejemplos de Llama 3.1:

También puede utilizar la biblioteca Optimum Neuron de Hugging Face para implementar rápidamente modelos directamente desde SageMaker a través del centro de modelos de Hugging Face. Desde el centro de tarjetas de modelos de Llama 3.1, elija Desplegarentonces Creador de Sagey finalmente AWS Inferentia y TrainiumCopie el código de ejemplo en un cuaderno de SageMaker y luego seleccione Correr.

import json
import sagemaker
import boto3
from sagemaker.huggingface import HuggingFaceModel, get_huggingface_llm_image_uri

try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client("iam")
    role = iam.get_role(RoleName="sagemaker_execution_role")["Role"]["Arn"]

# Hub Model configuration. https://huggingface.co/models
hub = {
    "HF_MODEL_ID": "meta-llama/Meta-Llama-3.1-8B",
    "HF_NUM_CORES": "2",
    "HF_AUTO_CAST_TYPE": "fp16",
    "MAX_BATCH_SIZE": "8",
    "MAX_INPUT_LENGTH": "3686",
    "MAX_TOTAL_TOKENS": "4096",
    "HF_TOKEN": "<REPLACE WITH YOUR TOKEN>",
}

assert hub["HF_TOKEN"] != "<REPLACE WITH YOUR TOKEN>", "Please replace '<REPLACE WITH YOUR TOKEN>' with your Hugging Face Hub API token"


# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
    image_uri=get_huggingface_llm_image_uri("huggingface-neuronx", version="0.0.23"),
    env=hub,
    role=role,
)

# deploy model to SageMaker Inference
predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type="ml.inf2.xlarge",
    container_startup_health_check_timeout=1800,
    volume_size=512,
)

# send request
predictor.predict(
    {
        "inputs": "What is is the capital of France?",
        "parameters": {
            "do_sample": True,
            "max_new_tokens": 128,
            "temperature": 0.7,
            "top_k": 50,
            "top_p": 0.95,
        }
    }
)

Además, si desea utilizar vLLM para implementar los modelos, puede consultar la Guía de dosificación continua para crear el entorno. Después de crear el entorno, puede usar vLLM para implementar modelos Llama 3.1 8/70B en AWS Trainium o Inferentia. A continuación, se muestra un ejemplo para implementar Llama 3.1 8B:

from vllm import LLM, SamplingParams
# Sample prompts.
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
# Create a sampling params object.
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
# Create an LLM.
llm = LLM(
    model="meta-llama/Meta-Llama-3.1-8B",
    max_num_seqs=8,
    # The max_model_len and block_size arguments are required to be same as max sequence length,
    # when targeting neuron device. Currently, this is a known limitation in continuous batching
    # support in transformers-neuronx.
    max_model_len=128,
    block_size=128,
    # The device can be automatically detected when AWS Neuron SDK is installed.
    # The device argument can be either unspecified for automated detection, or explicitly assigned.
    device="neuron",
    tensor_parallel_size=8)
# Generate texts from the prompts. The output is a list of RequestOutput objects
# that contain the prompt, generated text, and other information.
outputs = llm.generate(prompts, sampling_params)
# Print the outputs.
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

Conclusión

AWS Trainium e Inferentia ofrecen un alto rendimiento y un bajo costo para ajustar e implementar los modelos de Llama 3.1. Estamos ansiosos por ver cómo usará estos modelos poderosos y nuestra infraestructura de IA diseñada específicamente para crear aplicaciones de IA diferenciadas. Para obtener más información sobre cómo comenzar a usar los chips de IA de AWS, consulte Muestras de modelos y tutoriales en la documentación de AWS Neuron.

Sobre los autores

John Gray es arquitecto de soluciones sénior en Annapurna Labs, AWS, con sede en Seattle. En esta función, John trabaja con clientes en sus casos de uso de inteligencia artificial y aprendizaje automático, diseña soluciones para resolver de manera rentable sus problemas comerciales y los ayuda a construir un prototipo escalable utilizando chips de inteligencia artificial de AWS.

Panigrahi rosado Trabaja con clientes para crear soluciones basadas en ML para resolver problemas comerciales estratégicos en AWS. En su función actual, trabaja en la optimización del entrenamiento y la inferencia de modelos de IA generativos en chips de IA de AWS.

Kamran KhanDirector de Desarrollo Comercial de AWS Inferentina/Trianium en AWS. Tiene más de una década de experiencia ayudando a los clientes a implementar y optimizar cargas de trabajo de inferencia y entrenamiento de aprendizaje profundo mediante AWS Inferentia y AWS Trainium.

Shruti Koparkar es gerente sénior de marketing de productos en AWS. Ayuda a los clientes a explorar, evaluar y adoptar la infraestructura informática acelerada de Amazon EC2 para sus necesidades de aprendizaje automático.