Training de Deep Learning en AWS Inferentia
Entrenamiento de Deep Learning en AWS Inferentia
Otro truco para ahorrar dinero en el entrenamiento de modelos de IA
El tema de esta publicación es el chip de IA de desarrollo propio de AWS, AWS Inferentia, más específicamente, la segunda generación de AWS Inferentia2. Esta es una continuación de nuestra publicación del año pasado sobre AWS Trainium y se une a una serie de publicaciones sobre el tema de los aceleradores de IA dedicados. A diferencia de los chips que hemos explorado en nuestras publicaciones anteriores de la serie, AWS Inferentia fue diseñado para la inferencia de modelos de IA y está dirigido específicamente a aplicaciones de inferencia de aprendizaje profundo. Sin embargo, el hecho de que AWS Inferentia2 y AWS Trainium compartan la misma arquitectura subyacente NeuronCore-v2 y la misma pila de software (el SDK Neuron de AWS), plantea la pregunta: ¿Se puede utilizar AWS Inferentia también para cargas de trabajo de entrenamiento de IA?
Es cierto que hay algunos elementos de las especificaciones de la familia de instancias Amazon EC2 Inf2 (que están alimentadas por aceleradores AWS Inferentia) que podrían hacerlas menos adecuadas para algunas cargas de trabajo de entrenamiento en comparación con la familia de instancias Amazon EC2 Trn1. Por ejemplo, aunque tanto Inf2 como Trn1 admiten una interconexión de dispositivo a dispositivo NeuronLink-v2 de alta capacidad de ancho de banda y baja latencia, los dispositivos Trainium están conectados en una topología de toroide 2D en lugar de una topología de anillo, lo que potencialmente puede afectar el rendimiento de los operadores de comunicación colectiva (consulte aquí para obtener más detalles). Sin embargo, algunas cargas de trabajo de entrenamiento pueden no requerir las características únicas de la arquitectura Trn1 y pueden funcionar igual de bien en las arquitecturas Inf1 e Inf2.
De hecho, la capacidad de entrenar en los aceleradores Trainium e Inferentia aumentaría considerablemente la variedad de instancias de entrenamiento a nuestra disposición y nuestra capacidad para ajustar la elección de la instancia de entrenamiento a las necesidades específicas de cada proyecto de aprendizaje profundo. En nuestra publicación reciente, Selección de instancias para el aprendizaje profundo, desarrollamos el valor de tener una amplia variedad de tipos de instancias diferentes para el entrenamiento de DL. Mientras que la familia Trn1 incluye solo dos tipos de instancias, habilitar el entrenamiento en Inf2 agregaría otros cuatro tipos de instancias adicionales. Incluir Inf1 en la mezcla agregaría otros cuatro más.
Nuestra intención en esta publicación es demostrar la oportunidad de entrenar en AWS Inferentia. Definiremos un modelo de visión simple y compararemos el rendimiento de entrenarlo en las familias de instancias Amazon EC2 Trn1 y Amazon EC2 Inf2. Muchas gracias a Ohad Klein y Yitzhak Levi por sus contribuciones a esta publicación.
Avisos legales
- Tenga en cuenta que, en el momento de escribir esto, hay algunas arquitecturas de modelos de DL que aún no son compatibles con el SDK Neuron. Por ejemplo, si bien se admite la inferencia de modelos CNN, el entrenamiento de modelos CNN aún no es compatible. La documentación del SDK incluye una matriz de soporte de modelos que detalla las características admitidas por arquitectura de modelo, marco de entrenamiento (por ejemplo, TensorFlow y PyTorch) y versión de arquitectura Neuron.
- Los experimentos que describiremos se ejecutaron en Amazon EC2 con la versión más reciente de la AMI de aprendizaje profundo para Neuron disponible en el momento de escribir esto, “Deep Learning AMI Neuron PyTorch 1.13 (Ubuntu 20.04) 20230720”, que incluye la versión 2.8 del SDK Neuron. Dado que el SDK Neuron sigue en desarrollo activo, es probable que los resultados comparativos que logramos cambien con el tiempo. Se recomienda encarecidamente que vuelva a evaluar los hallazgos de esta publicación con las versiones más actualizadas de las bibliotecas subyacentes.
- Nuestra intención en esta publicación es demostrar el potencial de entrenar en instancias alimentadas por AWS Inferentia. Por favor, no vea esta publicación como una recomendación para el uso de estas instancias o cualquiera de los otros productos que mencionemos. Hay muchas variables que influyen en cómo elegir un entorno de entrenamiento, que pueden variar considerablemente según los detalles de su proyecto. En particular, diferentes modelos pueden mostrar resultados de rendimiento y precio relativos completamente diferentes al ejecutarse en dos tipos de instancias diferentes.
Modelo de juguete
Similar a los experimentos que describimos en nuestra publicación anterior, definimos un modelo de clasificación simple respaldado por Vision Transformer (ViT) (usando el paquete Python timm versión 0.9.5) junto con un conjunto de datos generado aleatoriamente.
- Aprendizaje por refuerzo Programación dinámica y Monte Carlo –...
- Investigadores de Virginia Tech y Microsoft presentan el Algoritmo ...
- Las mejores herramientas de IA para estudiantes (septiembre de 2023)
from torch.utils.data import Datasetimport time, osimport torchimport torch_xla.core.xla_model as xmimport torch_xla.distributed.parallel_loader as plfrom timm.models.vision_transformer import VisionTransformer
# usar datos aleatoriosclass FakeDataset(Dataset): def __len__(self): return 1000000 def __getitem__(self, index): rand_image = torch.randn([3, 224, 224], dtype=torch.float32) label = torch.tensor(data=[index % 1000], dtype=torch.int64) return rand_image, labeldef train(batch_size=16, num_workers=4): # Inicializar el grupo de procesos XLA para torchrun import torch_xla.distributed.xla_backend torch.distributed.init_process_group('xla') # multiprocesamiento: asegurarse de que cada trabajador tenga los mismos pesos iniciales torch.manual_seed(0) dataset = FakeDataset() model = VisionTransformer() # cargar el modelo en el dispositivo XLA device = xm.xla_device() model = model.to(device) optimizer = torch.optim.Adam(model.parameters()) data_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=num_workers) data_loader = pl.MpDeviceLoader(data_loader, device) loss_function = torch.nn.CrossEntropyLoss() summ, tsumm = 0, 0 count = 0 for step, (inputs, target) in enumerate(data_loader, start=1): t0 = time.perf_counter() inputs = inputs.to(device) targets = torch.squeeze(target.to(device), -1) optimizer.zero_grad() outputs = model(inputs) loss = loss_function(outputs, targets) loss.backward() xm.optimizer_step(optimizer) batch_time = time.perf_counter() - t0 if idx > 10: # omitir los primeros pasos summ += batch_time count += 1 t0 = time.perf_counter() if idx > 500: break print(f'tiempo promedio por paso: {summ/count}')if __name__ == '__main__': os.environ['XLA_USE_BF16'] = '1' # establecer el número de trabajadores del cargador de datos según el número de vCPUs # por ejemplo, 4 para trn1, 2 para inf2.xlarge, 8 para inf2.12xlarge e inf2.48xlarge train(num_workers=4)# Comando de inicialización:# torchrun --nproc_per_node=2 train.pyResultados
En la tabla a continuación, comparamos la velocidad y el rendimiento en precio de varios tipos de instancias de Amazon EC2 Trn1 y Amazon EC2 Inf2.
Aunque está claro que los tipos de instancias con Trainium admiten un mejor rendimiento absoluto (es decir, velocidades de entrenamiento más rápidas), el entrenamiento en las instancias con Inferentia resultó en un rendimiento en precio ~39% mejor (para los tipos de instancias de dos núcleos) y mayor (para los tipos de instancias más grandes).
Nuevamente, advertimos que no se deben tomar decisiones de diseño basadas únicamente en estos resultados. Algunas arquitecturas de modelos pueden funcionar correctamente en instancias Trn1 pero pueden fallar en Inf2. Otros pueden tener éxito en ambos, pero mostrar resultados de rendimiento comparativos muy diferentes a los mostrados aquí.
Tenga en cuenta que hemos omitido el tiempo requerido para compilar el modelo de DL. Aunque esto solo se requiere la primera vez que se ejecuta el modelo, los tiempos de compilación pueden ser bastante altos (por ejemplo, más de diez minutos para nuestro modelo de juguete). Dos formas de reducir la sobrecarga de la compilación del modelo son la compilación en paralelo y la compilación sin conexión. Es importante asegurarse de que su script no incluya operaciones (o cambios en el gráfico) que desencadenen compilaciones frecuentes. Consulte la documentación del SDK de Neuron para obtener más detalles.
Resumen
Aunque se comercializa como un chip de inferencia de IA, parece que AWS Inferentia ofrece otra opción más para el entrenamiento de modelos de aprendizaje profundo. En nuestra publicación anterior sobre AWS Trainium, destacamos algunos de los desafíos que puede encontrar al adaptar sus modelos para entrenar en un nuevo ASIC de IA. La posibilidad de entrenar los mismos modelos en tipos de instancias con AWS Inferentia también podría aumentar la recompensa potencial de sus esfuerzos.
