Entrenamiento en TPUs

El entrenamiento en TPUs puede diferir ligeramente del entrenamiento en múltiples GPUs, incluso con Accelerate. Esta guía tiene como objetivo mostrarle dónde debe tener cuidado y por qué, así como las mejores prácticas en general.

Entrenamiento en un Notebook

El principal punto de atención al entrenar en TPUs proviene de notebook_launcher(). Como se mencionó en el tutorial de notebook, es necesario reestructurar el código de entrenamiento en una función que pueda pasarse a la función notebook_launcher() y tener cuidado de no declarar tensores en la GPU.

Aunque en un TPU esa parte no es tan importante, un aspecto crítico a comprender es que al ejecutar código desde un notebook lo haces a través de un proceso llamado forking. Al lanzar desde la línea de comandos, se realiza spawning, donde no hay un proceso de python en ejecución y se spawnea un nuevo proceso. Dado que tu notebook de Jupyter ya está utilizando un proceso de python, necesitas forkar un nuevo proceso a partir de él para lanzar tu código.

Esto cobra importancia al declarar el modelo. En los procesos TPU bifurcados, se recomienda instanciar el modelo una sola vez y pasar esta instancia a la función de entrenamiento. Esto difiere del entrenamiento en GPUs, donde se crean n modelos que sincronizan y retropropagan sus gradientes en ciertos momentos. En su lugar, una instancia del modelo se comparte entre todos los nodos y se pasa de un lado a otro. Esto es especialmente importante al entrenar en TPUs de recursos limitados, como los proporcionados en los kernels de Kaggle o en Google Colaboratory.

A continuación se muestra un ejemplo de una función de entrenamiento pasada a notebook_launcher() si el entrenamiento se realiza en CPUs o GPUs:

Este fragmento de código se basa en el del cuaderno simple_nlp_example que se encuentra aquí, con ligeras modificaciones por simplicidad

def training_function():
    # Initialize accelerator
    accelerator = Accelerator()
    model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
    train_dataloader, eval_dataloader = create_dataloaders(
        train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
    )

    # Instantiate optimizer
    optimizer = AdamW(params=model.parameters(), lr=hyperparameters["learning_rate"])

    # Prepare everything
    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
    # prepare method.
    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader
    )

    num_epochs = hyperparameters["num_epochs"]
    # Now we train the model
    for epoch in range(num_epochs):
        model.train()
        for step, batch in enumerate(train_dataloader):
            outputs = model(**batch)
            loss = outputs.loss
            accelerator.backward(loss)

            optimizer.step()
            optimizer.zero_grad()

from accelerate import notebook_launcher

notebook_launcher(training_function)

El notebook_launcher usará 8 procesos de forma predeterminada si Accelerate se ha configurado para un TPU

Si usas este ejemplo y declaras el modelo dentro del bucle de entrenamiento, en un sistema con recursos limitados es posible que veas un error como:

ProcessExitedException: process 0 terminated with signal SIGSEGV

Este error es extremadamente críptico, pero la explicación básica es que se agotó la RAM del sistema. Puedes evitarlo por completo reconfigurando la función de entrenamiento para que acepte un único argumento model y lo declares en una celda externa:

# In another Jupyter cell
model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)

+ def training_function(model):
      # Initialize accelerator
      accelerator = Accelerator()
-     model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)
      train_dataloader, eval_dataloader = create_dataloaders(
          train_batch_size=hyperparameters["train_batch_size"], eval_batch_size=hyperparameters["eval_batch_size"]
      )
  ...

Y por último, llamando a la función de entrenamiento con:

  from accelerate import notebook_launcher
- notebook_launcher(training_function)
+ notebook_launcher(training_function, (model,))

El workaround anterior solo es necesario al iniciar una instancia de TPU desde un Jupyter Notebook en un servidor con recursos limitados, como Google Colaboratory o Kaggle. Si se utiliza un script o se inicia en un servidor mucho más potente, no es necesario declarar el modelo con antelación.

Precisión mixta y variables globales

Como se mencionó en la sección de precisión mixta del tutorial de migración, Accelerate admite fp16 y bf16, y ambos pueden usarse en TPUs. Dicho esto, lo ideal es utilizar bf16, ya que es extremadamente eficiente.

Hay dos “capas” al usar bf16 y Accelerate en TPU, a nivel base y a nivel de operación.

A nivel básico, esto se habilita al pasar mixed_precision="bf16" a Accelerator, por ejemplo:

accelerator = Accelerator(mixed_precision="bf16")

Por defecto, esto convertirá torch.float y torch.double a bfloat16 en TPU. La configuración específica que se establece es que la variable de entorno XLA_USE_BF16 se fija en 1.

Existe una configuración adicional que puedes realizar: establecer la variable de entorno XLA_DOWNCAST_BF16. Si se establece en 1, entonces torch.float es bfloat16 y torch.double es float32.

Esto se realiza en el objeto Accelerator al pasar downcast_bf16=True:

accelerator = Accelerator(mixed_precision="bf16", downcast_bf16=True)

Usar downcasting en lugar de bf16 en todas partes es útil cuando se intenta calcular métricas, registrar valores y más, donde los tensores bf16 crudos serían inutilizables.

Tiempos de entrenamiento en TPUs

Al ejecutar su script, es posible que note que el entrenamiento parece excepcionalmente lento al principio. Esto se debe a que las TPU primero procesan algunos lotes de datos para determinar cuánta memoria asignar antes de utilizar de manera extremadamente eficiente esta asignación de memoria configurada.

Si notas que tu código de evaluación para calcular las métricas de tu modelo tarda más debido al uso de un tamaño de lote mayor, se recomienda mantener el mismo tamaño de lote que los datos de entrenamiento si es demasiado lento. De lo contrario, la memoria se reasignará a este nuevo tamaño de lote después de las primeras iteraciones.

El hecho de que la memoria esté asignada no significa que se vaya a utilizar ni que el tamaño del lote aumente al volver a su dataloader de entrenamiento.