Entrenamiento en TPU

Un TPU (Tensor Processing Unit) es un tipo de hardware diseñado específicamente para entrenar modelos de manera eficiente. Accelerate admite el entrenamiento en TPU, pero hay algunas cosas que debes tener en cuenta, en particular la compilación de grafos. Este tutorial aborda brevemente la compilación y, para más detalles, consulta la guía Training on TPUs with Accelerate.

Compilación

Un TPU crea un grafo de todas las operaciones del paso de entrenamiento, como el paso hacia adelante, el paso hacia atrás y el paso del optimizador. Por eso, el primer paso de entrenamiento siempre tarda un poco, ya que construir y compilar este grafo lleva tiempo. Pero una vez que la compilación está completa, se almacena en caché y todos los pasos posteriores son mucho más rápidos.

La clave es evitar recompilar el código, ya que de lo contrario el entrenamiento será extremadamente lento. Esto significa que todas tus operaciones deben ser exactamente las mismas:

• todos los tensores en sus lotes deben tener la misma longitud (por ejemplo, sin relleno dinámico para tareas de PLN)
• su código debe ser estático (por ejemplo, sin capas con bucles for que tengan longitudes diferentes según la entrada, como una LSTM)

Vinculación de pesos

Un diseño común de modelos de lenguaje consiste en vincular los pesos de las capas de incrustación y softmax. Sin embargo, al mover el modelo a un TPU (ya sea manualmente o pasándolo al método prepare()), se rompe la vinculación de pesos y será necesario volver a vincularlos.

Para añadir un comportamiento especial (como el atado de pesos) en su script para TPU, primero establezca distributed_type en DistributedType.TPU. Después, puede usar el método tie_weights para atar los pesos.

if accelerator.distributed_type == DistributedType.TPU:
    model.tie_weights()