Recientemente, nuestros amigos de Unsloth han compartido un problema relacionado con la acumulación de gradientes que está afectando al entrenador (Trainer) de transformadores. El informe inicial proviene de @bnjmn_marie (¡bien hecho!).
La acumulación de gradientes debería ser matemáticamente equivalente al entrenamiento con un batch completo; sin embargo, no se lograron pérdidas coherentes entre las ejecuciones de entrenamiento cuando se activaba y desactivaba esta opción.
¿Qué es la Acumulación de Gradientes?
La acumulación de gradientes permite procesar más datos utilizando una menor cantidad de memoria. Esto se logra acumulando las actualizaciones de gradientes durante varias iteraciones antes de realizar una actualización del modelo. Sin embargo, si no se implementa correctamente, puede resultar en pérdidas inexactas durante el entrenamiento.
¿De dónde proviene este problema?
En el código de modelado de cada modelo, transformers
ofrece funciones de pérdida «por defecto» que son las más comúnmente utilizadas para la tarea del modelo. Estas funciones se determinan según el tipo de aprendizaje: respuesta a preguntas, clasificación de tokens, LM causal, LM enmascarado, entre otros.
Esta función de pérdida por defecto no estaba pensada para ser personalizable; solo se calcula cuando se pasan labels
y input_ids
como entradas al modelo, de modo que el usuario no tiene que calcularla. Aunque útil, esta función tiene limitaciones por diseño: para cualquier variación, se espera que no se pasen etiquetas directamente y que los usuarios obtengan los logits del modelo para calcular la pérdida externamente.
No obstante, el Trainer de transformadores, así como muchos otros, aprovechan en gran medida estos métodos por la sencillez que ofrecen. Esto puede resultar en un arma de doble filo; ofrecer una API simple que cambia según el uso no es ideal, y nosotros mismos hemos sido sorprendidos.
Para ser precisos, la pérdida correcta para la acumulación de gradientes en tareas a nivel de token, como el entrenamiento de LM causal, debe calcularse sumando la pérdida total a través de todos los batches en un paso de acumulación de gradientes y dividiéndola por el número total de tokens que no son de relleno en esos batches. Esto no es lo mismo que el promedio de las pérdidas por cada batch.
La solución es bastante sencilla, como se muestra a continuación:
def ForCausalLMLoss(logits, labels, vocab_size, **kwargs):
logits = logits.float()
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
shift_logits = shift_logits.view(-1, vocab_size)
shift_labels = shift_labels.view(-1)
shift_labels = shift_labels.to(shift_logits.device)
num_items = kwargs.pop("num_items", None)
+ loss = nn.functional.cross_entropy(shift_logits, shift_labels, ignore_index=-100, reduction="sum")
+ loss = loss / num_items
- loss = nn.functional.cross_entropy(shift_logits, shift_labels, ignore_index=-100)
return loss
¿Cómo lo estamos solucionando?
Para abordar este problema, estamos cambiando la forma en que funcionan nuestros modelos y entrenamientos de dos maneras:
- Si los usuarios están utilizando las funciones de pérdida «por defecto», tomaremos automáticamente en cuenta los cambios necesarios al usar la acumulación de gradientes, asegurando que la pérdida adecuada se informe y utilice, corrigiendo así el problema principal.
- Para garantizar que futuros problemas en el cálculo de pérdidas no bloqueen a los usuarios, expondremos una API que permita a los usuarios pasar sus propias funciones de pérdida directamente al
Trainer
, facilitándoles sus propias soluciones hasta que resolvamos cualquier problema interno y realicemos un nuevo lanzamiento de transformers.
Todos los modelos que heredan de PreTrainedModel
ahora cuentan con una propiedad loss_function
, que se determina por:
- el
config.loss_type
: esto garantiza que cualquier persona pueda usar su pérdida personalizada. Puedes hacer esto modificando elLOSS_MAPPING
:
def my_super_loss(logits, labels):
return loss = nn.functional.cross_entropy(logits, labels, ignore_index=-100)
LOSS_MAPPING["my_loss_type"] = my_super_loss
Estamos trabajando para implementar el primer cambio en los modelos más populares de esta PR: https://github.com/huggingface/transformers/pull/34191#pullrequestreview-2372725010. Después de esto, lanzaremos un llamado a contribuciones para ayudar a propagar esto al resto de modelos de manera que la mayoría de ellos estén soportados para el próximo lanzamiento.
Asimismo, estamos trabajando activamente para implementar el segundo cambio en esta PR: https://github.com/huggingface/transformers/pull/34198, lo que permitirá a los usuarios utilizar su propia función de pérdida y hacer uso del número de muestras vistas por batch para ayudar a calcular su pérdida (y se realizará el cálculo correcto de pérdida durante la acumulación de gradientes a medida que se soporten más modelos de esta modificación previa).
Impacto Esperado
Con estos cambios, esperamos que los modelos que heredan de PreTrainedModel
se beneficien de una nueva propiedad llamada loss_function
, que permitirá a los usuarios personalizar aún más sus entrenamientos.
Invitamos a la comunidad a mantenernos al tanto de cualquier inconveniente a través de nuestros reportes, ya que cada comentario contribuye a mejorar la plataforma y adaptarla a las diversas necesidades de sus usuarios. La rapidez con la que hemos abordado este problema – en menos de 24 horas – es un reflejo de nuestro compromiso de ofrecer una herramienta robusta y eficiente para el entrenamiento de modelos de inteligencia artificial.
Conclusión
Estamos emocionados por estos avances y el impacto positivo que tendrán en la eficiencia y precisión de los entrenamientos de AI. Por favor, instalen desde main
para beneficiarse de estas correcciones:
pip install git+https://github.com/huggingface/transformers
No duden en visitarnos y aportar sus comentarios en nuestro repositorio de problemas en GitHub: Aquí. Su colaboración es crucial para que sigamos mejorando.
¡El equipo de Transformers!