Entrenamiento eficiente en memoria

🤗 PEFT hace que el ajuste fino sea eficiente en parámetros, pero no automáticamente eficiente en memoria. Esta recopilación reúne consejos para reducir la memoria de entrenamiento y enlaces a las guías detalladas.

Considera siempre las opciones básicas de elegir un modelo base más pequeño, un tamaño de lote más pequeño o una longitud de secuencia más corta para reducir el uso de memoria.

Resumen de la memoria de entrenamiento

Desglosemos la distribución de la memoria de entrenamiento para poder razonar sobre posibles contramedidas. Utilizaremos como ejemplo un modelo de lenguaje grande entrenado con el optimizador Adam. Al realizar un ajuste fino completo, tendremos las siguientes posiciones ocupando memoria:

1. parámetros del modelo base: el consumo de memoria depende en gran medida del dtype elegido. Cuantos menos bits por parámetro (16 para float16), menos memoria ocupará. Un modelo de 1B en float16 (16 bits/2 bytes por parámetro) ocupará aproximadamente 1e9 × 2 byte = 1.863 GiB de memoria.
2. todos los parámetros del modelo base entrenables ×3 para los gradientes (1×) y los estados del optimizador Adam (2×), por lo tanto, 5.59 GB de memoria
3. memoria para las activaciones intermedias entre capas, estas son difíciles de predecir pero dependen principalmente del dtype de cálculo utilizado y de la longitud de la secuencia / tamaño del lote.

Un modelo base más pequeño o el uso de un dtype de cómputo más pequeño reducirá todos los puntos, mientras que el uso de secuencias más cortas o lotes más pequeños afecta principalmente a la memoria de gradientes y activaciones. El empleo de métodos PEFT reducirá el número de parámetros entrenables y, por lo tanto, reducirá significativamente tanto los gradientes como el estado del optimizador, ahorrando mucha memoria.

Elegir el método adecuado

No todos los métodos de PEFT se construyen de la misma manera y algunas formulaciones son más fáciles de implementar de forma eficiente en memoria. Si tienes un presupuesto de memoria limitado, tiene sentido consultar el PEFT method comparison suite y filtrar por el uso máximo de memoria del acelerador. El uso medio de memoria del acelerador puede ser bastante similar entre métodos, pero no todos los métodos escalan de la misma manera con las activaciones y la longitud de la secuencia; algunos métodos son más propensos a picos de memoria que otros.

Considere usar tokens entrenables al apuntar a capas grandes, como las cabezas de modelado de lenguaje o las capas de incrustación, para ajustar finamente tokens específicos.

Cuantización

La cuantización es una de las mejores formas de reducir el consumo de memoria del modelo base y, dependiendo de la cuantización empleada, también reducirá la memoria de activaciones. Dado que los métodos PEFT solo ocuparán una pequeña parte del número total de parámetros, PEFT utiliza de forma predeterminada una precisión mayor que la del modelo base. Esto también puede tener el efecto de que los adaptadores mitiguen parte de la pérdida de calidad incurrida por los métodos de cuantización. Lea la guía de cuantización de PEFT.

Compilación

Los modelos que entrenamos están compuestos por operaciones como multiplicaciones de matrices, sumas y asignaciones, donde cada operación produce un nuevo resultado y, a continuación, necesita ocupar memoria. Si esos resultados intermedios no son necesarios, podemos fusionar estas operaciones y ahorrar memoria. Esta es solo una de las muchas optimizaciones que torch.compile puede realizar por ti, así que consulta la guía de PEFT para torch.compile.

Gradient Checkpointing

Puede intercambiar memoria por computación guardando solo cada n-ésimo gradiente entre capas y calculando el resto sobre la marcha. Consulte la documentación de gradient checkpointing de Transformers para obtener más información.

Cuando no se utilizan Diffusers ni Transformers, es posible que sea necesario implementar la propia lógica de verificación de gradientes, dependiendo del marco de entrenamiento que se esté utilizando.

Pérdida NLL por bloques

El uso de NLLLoss es muy común al entrenar modelos de lenguaje (o tareas de clasificación). Se asigna una matriz de tamaño batch × sequence × vocabulary. Con secuencias o vocabularios particularmente largos, esto puede resultar costoso rápidamente.

Al usar TRL puedes utilizar la integración del kernel Liger o usar Chunked NLLLoss. Este último divide la secuencia en fragmentos de tamaño 256 para mantener el consumo máximo de memoria constante.

En caso de que el tamaño de bloque predeterminado no sea óptimo para su configuración, consulte el PR original de TRL para obtener más información sobre cómo ajustar el tamaño de bloque.