Las GPU han sido un producto escaso desde que los modelos grandes se convirtieron en una tendencia candente. Las reservas de muchas empresas no son necesariamente suficientes, y mucho menos de los desarrolladores individuales. ¿Hay alguna forma de utilizar la potencia informática para entrenar modelos de manera más eficiente?
En un blog reciente, Sebastian Raschka presentó el método de "acumulación de gradiente", que puede usar un modelo de entrenamiento de tamaño de lote más grande cuando la memoria de la GPU es limitada, sin pasar por las limitaciones del hardware.
Antes de esto, Sebastian Raschka también compartió un artículo sobre el uso de estrategias de entrenamiento multi-GPU para acelerar el ajuste fino del modelo de lenguaje a gran escala, incluidos mecanismos como el modelo o la fragmentación del tensor, que distribuyen los pesos y cálculos del modelo en diferentes dispositivos para resolver problemas de GPU. límite de memoria.
Ajuste fino del modelo BLOOM para la clasificación
Supongamos que estamos interesados en adoptar modelos de lenguaje grande preentrenados recientemente para tareas posteriores como la clasificación de texto. Entonces, podríamos optar por usar la alternativa de código abierto a GPT-3, el modelo BLOOM, específicamente la versión BLOOM con "solo" 560 millones de parámetros, que debería caber en la RAM de una GPU tradicional sin ningún problema (Google Colab gratis versión tiene una GPU con 15 Gb de RAM).
Una vez que comience, es probable que se encuentre con un problema: la memoria aumentará rápidamente durante el entrenamiento o el ajuste fino. La única forma de entrenar este modelo es hacer que el tamaño del lote sea 1 (tamaño del lote = 1).
El código para ajustar BLOOM para la tarea de clasificación de destino utilizando un tamaño de lote de 1 (tamaño de lote = 1) se muestra a continuación. También puede descargar el código completo desde la página del proyecto de GitHub:
https://github.com/rasbt/gradient-accumulation-blog/blob/main/src/1_batchsize-1.py
Puede copiar y pegar este código directamente en Google Colab, pero también debe arrastrar y soltar el archivo local_dataset_utilities.py adjunto en la misma carpeta desde la que importó algunas utilidades del conjunto de datos.
# pip install torch lightning matplotlib pandas torchmetrics watermark transformers datasets -Uimport osimport os.path as opimport timefrom datasets import load_datasetfrom lightning import Fabricimport torchfrom torch.utils.data import DataLoaderimport torchmetricsfrom transformers import AutoTokenizerfrom transformers import AutoModelForSequenceClassificationfrom watermark import watermarkfrom local_dataset_utilities import download_dataset, load_dataset_into_to_dataframe, partition_datasetfrom local_dataset_utilities import IMDBDatasetdef tokenize_text (batch):return tokenizer (batch ["text"], truncation=True, padding=True, max_length=1024)def train (num_epochs, model, optimizer, train_loader, val_loader, fabric):for epoch in range (num_epochs):train_acc = torchmetrics.Accuracy (task="multiclass", num_classes=2).to (fabric.device)for batch_idx, batch in enumerate (train_loader):model.train ()### FORWARD AND BACK PROPoutputs = model (batch ["input_ids"],attention_mask=batch ["attention_mask"],labels=batch ["label"])fabric.backward (outputs ["loss"])### UPDATE MODEL PARAMETERSoptimizer.step ()optimizer.zero_grad ()### LOGGINGif not batch_idx % 300:print (f"Epoch: {epoch+1:04d}/{num_epochs:04d}"f"| Batch {batch_idx:04d}/{len (train_loader):04d}"f"| Loss: {outputs ['loss']:.4f}")model.eval ()with torch.no_grad ():predicted_labels = torch.argmax (outputs ["logits"], 1)train_acc.update (predicted_labels, batch ["label"])### MORE LOGGINGmodel.eval ()with torch.no_grad ():val_acc = torchmetrics.Accuracy (task="multiclass", num_classes=2).to (fabric.device)for batch in val_loader:outputs = model (batch ["input_ids"],attention_mask=batch ["attention_mask"],labels=batch ["label"])predicted_labels = torch.argmax (outputs ["logits"], 1)val_acc.update (predicted_labels, batch ["label"])print (f"Epoch: {epoch+1:04d}/{num_epochs:04d}"f"| Train acc.: {train_acc.compute ()*100:.2f}%"f"| Val acc.: {val_acc.compute ()*100:.2f}%")train_acc.reset (), val_acc.reset ()if __name__ == "__main__":print (watermark (packages="torch,lightning,transformers", python=True))print ("Torch CUDA available?", torch.cuda.is_available ())device = "cuda" if torch.cuda.is_available () else "cpu"torch.manual_seed (123)# torch.use_deterministic_algorithms (True)############################# 1 Loading the Dataset##########################download_dataset ()df = load_dataset_into_to_dataframe ()if not (op.exists ("train.csv") and op.exists ("val.csv") and op.exists ("test.csv")):partition_dataset (df)imdb_dataset = load_dataset ("csv",data_files={"train": "train.csv","validation": "val.csv","test": "test.csv",},)############################################ 2 Tokenization and Numericalization#########################################tokenizer = AutoTokenizer.from_pretrained ("bigscience/bloom-560m", max_length=1024)print ("Tokenizer input max length:", tokenizer.model_max_length, flush=True)print ("Tokenizer vocabulary size:", tokenizer.vocab_size, flush=True)print ("Tokenizing ...", flush=True)imdb_tokenized = imdb_dataset.map (tokenize_text, batched=True, batch_size=None)del imdb_datasetimdb_tokenized.set_format ("torch", columns=["input_ids", "attention_mask", "label"])os.environ ["TOKENIZERS_PARALLELISM"] = "false"############################################ 3 Set Up DataLoaders#########################################train_dataset = IMDBDataset (imdb_tokenized, partition_key="train")val_dataset = IMDBDataset (imdb_tokenized, partition_key="validation")test_dataset = IMDBDataset (imdb_tokenized, partition_key="test")train_loader = DataLoader (dataset=train_dataset,batch_size=1,shuffle=True,num_workers=4,drop_last=True,)val_loader = DataLoader (dataset=val_dataset,batch_size=1,num_workers=4,drop_last=True,)test_loader = DataLoader (dataset=test_dataset,batch_size=1,num_workers=2,drop_last=True,)############################################ 4 Initializing the Model#########################################fabric = Fabric (accelerator="cuda", devices=1, precision="16-mixed")fabric.launch ()model = AutoModelForSequenceClassification.from_pretrained ("bigscience/bloom-560m", num_labels=2)optimizer = torch.optim.Adam (model.parameters (), lr=5e-5)model, optimizer = fabric.setup (model, optimizer)train_loader, val_loader, test_loader = fabric.setup_dataloaders (train_loader, val_loader, test_loader)############################################ 5 Finetuning#########################################start = time.time ()train (num_epochs=1,model=model,optimizer=optimizer,train_loader=train_loader,val_loader=val_loader,fabric=fabric,)end = time.time ()elapsed = end-startprint (f"Time elapsed {elapsed/60:.2f} min")with torch.no_grad ():model.eval ()test_acc = torchmetrics.Accuracy (task="multiclass", num_classes=2).to (fabric.device)for batch in test_loader:outputs = model (batch ["input_ids"],attention_mask=batch ["attention_mask"],labels=batch ["label"])predicted_labels = torch.argmax (outputs ["logits"], 1)test_acc.update (predicted_labels, batch ["label"])print (f"Test accuracy {test_acc.compute ()*100:.2f}%")
El autor usó Lightning Fabric porque permite a los desarrolladores cambiar de manera flexible la cantidad de GPU y la estrategia de entrenamiento de múltiples GPU al ejecutar este código en hardware diferente. También permite habilitar el entrenamiento de precisión mixta simplemente ajustando la bandera de precisión. En este caso, el entrenamiento de precisión mixta puede triplicar la velocidad de entrenamiento y reducir los requisitos de memoria en aproximadamente un 25 %.
El código principal que se muestra arriba se ejecuta en la función principal (el contexto de if __name__ == "__main__"). Incluso si solo se usa una sola GPU, se recomienda usar el entorno de tiempo de ejecución de PyTorch para realizar el entrenamiento de múltiples GPU. Luego, las siguientes tres secciones de código envueltas en if __name__ == "__main__" son responsables de la carga de datos:
#1 Carga el conjunto de datos
#2 tokenización y digitalización
#3 Configurar el cargador de datos
La Sección 4 está en Inicializar el modelo, y luego en la Sección 5, Ajuste fino, se llama a la función de tren, y aquí es donde las cosas comienzan a ponerse interesantes. En la función entrenar (...), se implementa un bucle PyTorch estándar. Una versión anotada del ciclo de entrenamiento principal se ve así:
El problema con un tamaño de lote de 1 (Tamaño de lote = 1) es que las actualizaciones de gradiente pueden volverse muy complicadas y difíciles, como se ve con la pérdida de entrenamiento fluctuante y el rendimiento deficiente del conjunto de prueba al entrenar el modelo a continuación:
...torch : 2.0.0lightning : 2.0.0transformers: 4.27.2Torch CUDA available? True...Epoch: 0001/0001 | Batch 23700/35000 | Loss: 0.0969Epoch: 0001/0001 | Batch 24000/35000 | Loss: 1.9902Epoch: 0001/0001 | Batch 24300/35000 | Loss: 0.0395Epoch: 0001/0001 | Batch 24600/35000 | Loss: 0.2546Epoch: 0001/0001 | Batch 24900/35000 | Loss: 0.1128Epoch: 0001/0001 | Batch 25200/35000 | Loss: 0.2661Epoch: 0001/0001 | Batch 25500/35000 | Loss: 0.0044Epoch: 0001/0001 | Batch 25800/35000 | Loss: 0.0067Epoch: 0001/0001 | Batch 26100/35000 | Loss: 0.0468Epoch: 0001/0001 | Batch 26400/35000 | Loss: 1.7139Epoch: 0001/0001 | Batch 26700/35000 | Loss: 0.9570Epoch: 0001/0001 | Batch 27000/35000 | Loss: 0.1857Epoch: 0001/0001 | Batch 27300/35000 | Loss: 0.0090Epoch: 0001/0001 | Batch 27600/35000 | Loss: 0.9790Epoch: 0001/0001 | Batch 27900/35000 | Loss: 0.0503Epoch: 0001/0001 | Batch 28200/35000 | Loss: 0.2625Epoch: 0001/0001 | Batch 28500/35000 | Loss: 0.1010Epoch: 0001/0001 | Batch 28800/35000 | Loss: 0.0035Epoch: 0001/0001 | Batch 29100/35000 | Loss: 0.0009Epoch: 0001/0001 | Batch 29400/35000 | Loss: 0.0234Epoch: 0001/0001 | Batch 29700/35000 | Loss: 0.8394Epoch: 0001/0001 | Batch 30000/35000 | Loss: 0.9497Epoch: 0001/0001 | Batch 30300/35000 | Loss: 0.1437Epoch: 0001/0001 | Batch 30600/35000 | Loss: 0.1317Epoch: 0001/0001 | Batch 30900/35000 | Loss: 0.0112Epoch: 0001/0001 | Batch 31200/35000 | Loss: 0.0073Epoch: 0001/0001 | Batch 31500/35000 | Loss: 0.7393Epoch: 0001/0001 | Batch 31800/35000 | Loss: 0.0512Epoch: 0001/0001 | Batch 32100/35000 | Loss: 0.1337Epoch: 0001/0001 | Batch 32400/35000 | Loss: 1.1875Epoch: 0001/0001 | Batch 32700/35000 | Loss: 0.2727Epoch: 0001/0001 | Batch 33000/35000 | Loss: 0.1545Epoch: 0001/0001 | Batch 33300/35000 | Loss: 0.0022Epoch: 0001/0001 | Batch 33600/35000 | Loss: 0.2681Epoch: 0001/0001 | Batch 33900/35000 | Loss: 0.2467Epoch: 0001/0001 | Batch 34200/35000 | Loss: 0.0620Epoch: 0001/0001 | Batch 34500/35000 | Loss: 2.5039Epoch: 0001/0001 | Batch 34800/35000 | Loss: 0.0131Epoch: 0001/0001 | Train acc.: 75.11% | Val acc.: 78.62%Time elapsed 69.97 minTest accuracy 78.53%
Dado que no hay muchas GPU disponibles para la fragmentación de tensores, ¿qué se puede hacer para entrenar modelos con tamaños de lote más grandes?
Una de esas soluciones es la acumulación de gradientes, que modifica el ciclo de entrenamiento antes mencionado.
¿Qué es la acumulación de gradiente?
La acumulación de gradientes es una forma de aumentar virtualmente el tamaño del lote durante el entrenamiento, lo cual es útil cuando la memoria GPU disponible es insuficiente para contener el tamaño del lote deseado. En la acumulación de gradientes, los gradientes se calculan para lotes más pequeños y se acumulan (generalmente sumados o promediados) en varias iteraciones, en lugar de actualizar los pesos del modelo después de cada lote. Una vez que el gradiente acumulativo alcanza el tamaño de lote "virtual" de destino, los pesos del modelo se actualizan utilizando el gradiente acumulativo.
Vea el ciclo de entrenamiento de PyTorch actualizado a continuación:
Si Accumulation_steps se establece en 2, entonces zero_grad() y Optimizer.step() solo se llamarán cada segundo. Por lo tanto, ejecutar el ciclo de entrenamiento modificado con Accumulation_steps=2 tiene el mismo efecto que duplicar el tamaño del lote.
Por ejemplo, si desea utilizar un tamaño de lote de 256, pero solo puede colocar un tamaño de lote de 64 en la memoria de GPU, puede realizar una acumulación de gradiente en cuatro lotes de tamaño 64. (Cuando se han procesado los cuatro lotes, los gradientes acumulativos equivalen a un solo tamaño de lote de 256). Esto simula efectivamente tamaños de lote más grandes sin requerir una memoria de GPU más grande o dividir tensores entre diferentes dispositivos.
Si bien la acumulación de gradientes puede ayudarnos a entrenar modelos con tamaños de lote más grandes, no reduce el cálculo total requerido. De hecho, a veces puede hacer que el proceso de entrenamiento sea un poco más lento porque las actualizaciones de peso se realizan con menos frecuencia. Sin embargo, nos ayuda a sortear la limitación de que los tamaños de lote muy pequeños provocan actualizaciones frecuentes y caóticas.
Por ejemplo, ahora ejecutemos el código anterior con un tamaño de lote de 1, que requiere 16 pasos de acumulación para simular un tamaño de lote igual a 16.
La salida es la siguiente:
...torch : 2.0.0lightning : 2.0.0transformers: 4.27.2Torch CUDA available? True...Epoch: 0001/0001 | Batch 23700/35000 | Loss: 0.0168Epoch: 0001/0001 | Batch 24000/35000 | Loss: 0.0006Epoch: 0001/0001 | Batch 24300/35000 | Loss: 0.0152Epoch: 0001/0001 | Batch 24600/35000 | Loss: 0.0003Epoch: 0001/0001 | Batch 24900/35000 | Loss: 0.0623Epoch: 0001/0001 | Batch 25200/35000 | Loss: 0.0010Epoch: 0001/0001 | Batch 25500/35000 | Loss: 0.0001Epoch: 0001/0001 | Batch 25800/35000 | Loss: 0.0047Epoch: 0001/0001 | Batch 26100/35000 | Loss: 0.0004Epoch: 0001/0001 | Batch 26400/35000 | Loss: 0.1016Epoch: 0001/0001 | Batch 26700/35000 | Loss: 0.0021Epoch: 0001/0001 | Batch 27000/35000 | Loss: 0.0015Epoch: 0001/0001 | Batch 27300/35000 | Loss: 0.0008Epoch: 0001/0001 | Batch 27600/35000 | Loss: 0.0060Epoch: 0001/0001 | Batch 27900/35000 | Loss: 0.0001Epoch: 0001/0001 | Batch 28200/35000 | Loss: 0.0426Epoch: 0001/0001 | Batch 28500/35000 | Loss: 0.0012Epoch: 0001/0001 | Batch 28800/35000 | Loss: 0.0025Epoch: 0001/0001 | Batch 29100/35000 | Loss: 0.0025Epoch: 0001/0001 | Batch 29400/35000 | Loss: 0.0000Epoch: 0001/0001 | Batch 29700/35000 | Loss: 0.0495Epoch: 0001/0001 | Batch 30000/35000 | Loss: 0.0164Epoch: 0001/0001 | Batch 30300/35000 | Loss: 0.0067Epoch: 0001/0001 | Batch 30600/35000 | Loss: 0.0037Epoch: 0001/0001 | Batch 30900/35000 | Loss: 0.0005Epoch: 0001/0001 | Batch 31200/35000 | Loss: 0.0013Epoch: 0001/0001 | Batch 31500/35000 | Loss: 0.0112Epoch: 0001/0001 | Batch 31800/35000 | Loss: 0.0053Epoch: 0001/0001 | Batch 32100/35000 | Loss: 0.0012Epoch: 0001/0001 | Batch 32400/35000 | Loss: 0.1365Epoch: 0001/0001 | Batch 32700/35000 | Loss: 0.0210Epoch: 0001/0001 | Batch 33000/35000 | Loss: 0.0374Epoch: 0001/0001 | Batch 33300/35000 | Loss: 0.0007Epoch: 0001/0001 | Batch 33600/35000 | Loss: 0.0341Epoch: 0001/0001 | Batch 33900/35000 | Loss: 0.0259Epoch: 0001/0001 | Batch 34200/35000 | Loss: 0.0005Epoch: 0001/0001 | Batch 34500/35000 | Loss: 0.4792Epoch: 0001/0001 | Batch 34800/35000 | Loss: 0.0003Epoch: 0001/0001 | Train acc.: 78.67% | Val acc.: 87.28%Time elapsed 51.37 minTest accuracy 87.37%
De acuerdo con los resultados anteriores, la fluctuación de la pérdida es menor que antes. Además, el rendimiento del equipo de prueba mejoró en un 10 %. Dado que el conjunto de entrenamiento se itera solo una vez, cada ejemplo de entrenamiento solo se encontrará una vez. Los modelos de entrenamiento para múltiples épocas pueden mejorar aún más el rendimiento predictivo.
También puede notar que este código también se ejecuta más rápido que el código de tamaño de lote 1 utilizado anteriormente. Si aumentamos el tamaño del lote virtual a 8 mediante la acumulación de gradientes, seguirá habiendo el mismo número de pases hacia adelante. Sin embargo, dado que el modelo se actualiza solo una vez cada ocho épocas, habrá menos pasos hacia atrás, lo que permitirá una iteración más rápida sobre las muestras dentro de una época (número de rondas de entrenamiento).
en conclusión
La acumulación de gradientes es una técnica que simula tamaños de lotes más grandes mediante la acumulación de múltiples gradientes de lotes pequeños antes de realizar actualizaciones de peso. Esta técnica ayuda cuando la memoria disponible es limitada y el tamaño del lote que cabe en la memoria es pequeño.
Pero primero, piense en un escenario en el que pueda ejecutar con un tamaño de lote, lo que significa que la memoria disponible es lo suficientemente grande para contener el tamaño de lote deseado. En ese caso, la acumulación de gradiente puede no ser necesaria. De hecho, puede ser más eficiente ejecutar con un tamaño de lote más grande, ya que permite más paralelismo y reduce la cantidad de actualizaciones de peso necesarias para entrenar el modelo.
En resumen, la acumulación de gradientes es una técnica práctica que se puede utilizar para reducir el impacto del ruido del tamaño de los mini lotes en la precisión de las actualizaciones de gradientes. Esta es, con mucho, una técnica simple y efectiva que nos permite eludir las limitaciones de hardware.
PD: ¿Se puede hacer que esto funcione más rápido?
ningún problema. Se puede hacer que se ejecute aún más rápido usando torch.compile introducido en PyTorch 2.0. Solo necesita agregar algo de model = torch.compile, como se muestra en la imagen a continuación:
El script completo está disponible en GitHub.
En este caso, torch.compile elimina otros diez minutos de tiempo de entrenamiento sin afectar el rendimiento del modelado:
poch: 0001/0001 | Batch 26400/35000 | Loss: 0.0320Epoch: 0001/0001 | Batch 26700/35000 | Loss: 0.0010Epoch: 0001/0001 | Batch 27000/35000 | Loss: 0.0006Epoch: 0001/0001 | Batch 27300/35000 | Loss: 0.0015Epoch: 0001/0001 | Batch 27600/35000 | Loss: 0.0157Epoch: 0001/0001 | Batch 27900/35000 | Loss: 0.0015Epoch: 0001/0001 | Batch 28200/35000 | Loss: 0.0540Epoch: 0001/0001 | Batch 28500/35000 | Loss: 0.0035Epoch: 0001/0001 | Batch 28800/35000 | Loss: 0.0016Epoch: 0001/0001 | Batch 29100/35000 | Loss: 0.0015Epoch: 0001/0001 | Batch 29400/35000 | Loss: 0.0008Epoch: 0001/0001 | Batch 29700/35000 | Loss: 0.0877Epoch: 0001/0001 | Batch 30000/35000 | Loss: 0.0232Epoch: 0001/0001 | Batch 30300/35000 | Loss: 0.0014Epoch: 0001/0001 | Batch 30600/35000 | Loss: 0.0032Epoch: 0001/0001 | Batch 30900/35000 | Loss: 0.0004Epoch: 0001/0001 | Batch 31200/35000 | Loss: 0.0062Epoch: 0001/0001 | Batch 31500/35000 | Loss: 0.0032Epoch: 0001/0001 | Batch 31800/35000 | Loss: 0.0066Epoch: 0001/0001 | Batch 32100/35000 | Loss: 0.0017Epoch: 0001/0001 | Batch 32400/35000 | Loss: 0.1485Epoch: 0001/0001 | Batch 32700/35000 | Loss: 0.0324Epoch: 0001/0001 | Batch 33000/35000 | Loss: 0.0155Epoch: 0001/0001 | Batch 33300/35000 | Loss: 0.0007Epoch: 0001/0001 | Batch 33600/35000 | Loss: 0.0049Epoch: 0001/0001 | Batch 33900/35000 | Loss: 0.1170Epoch: 0001/0001 | Batch 34200/35000 | Loss: 0.0002Epoch: 0001/0001 | Batch 34500/35000 | Loss: 0.4201Epoch: 0001/0001 | Batch 34800/35000 | Loss: 0.0018Epoch: 0001/0001 | Train acc.: 78.39% | Val acc.: 86.84%Time elapsed 43.33 minTest accuracy 87.91%
Tenga en cuenta que el ligero aumento en la precisión en comparación con antes probablemente se deba a la aleatoriedad.
¿Qué software es? http://143ai.com