Pelatihan pada TPU

Pelatihan pada TPU dapat sedikit berbeda dari pelatihan pada multi-GPU, bahkan dengan Accelerate. Panduan ini bertujuan untuk menunjukkan di mana Anda harus berhati-hati dan mengapa, serta praktik terbaik secara umum.

Pelatihan di Notebook

Perhatian utama saat melatih pada TPU berasal dari notebook_launcher(). Seperti yang disebutkan dalam tutorial notebook, Anda perlu menyusun ulang kode pelatihan Anda menjadi fungsi yang dapat diteruskan ke fungsi notebook_launcher() dan berhati-hatilah agar tidak mendeklarasikan tensor apa pun di GPU.

Meskipun di TPU bagian terakhir itu tidak terlalu penting, bagian kritis yang perlu dipahami adalah bahwa saat Anda menjalankan kode dari notebook, Anda melakukannya melalui proses yang disebut forking. Saat meluncurkan dari command-line, Anda melakukan spawning, di mana proses python tidak sedang berjalan dan Anda spawn proses baru. Karena notebook Jupyter Anda sudah memanfaatkan proses python, Anda perlu fork proses baru darinya untuk meluncurkan kode Anda.

Hal ini menjadi penting terkait deklarasi model Anda. Pada proses TPU yang di-fork, disarankan untuk menginstansiasi model Anda sekali dan meneruskannya ke fungsi pelatihan Anda. Ini berbeda dengan pelatihan pada GPU, di mana Anda membuat n model yang gradiennya disinkronkan dan di-backpropagate pada momen tertentu. Sebaliknya, satu instans model dibagikan di antara semua node dan diteruskan bolak-balik. Hal ini sangat penting terutama saat melatih pada TPU berresource rendah, seperti yang disediakan di kernel Kaggle atau di Google Colaboratory.

Berikut adalah contoh fungsi pelatihan yang diteruskan ke notebook_launcher() jika pelatihan dilakukan pada CPU atau GPU:

Snippet kode ini didasarkan pada yang ada di notebook simple_nlp_example yang dapat ditemukan di sini dengan sedikit modifikasi demi kesederhanaan

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)

notebook_launcher akan secara default menggunakan 8 proses jika Accelerate telah dikonfigurasi untuk TPU

Jika Anda menggunakan contoh ini dan mendeklarasikan model di dalam loop pelatihan, maka pada sistem dengan sumber daya terbatas Anda berpotensi melihat error seperti:

ProcessExitedException: process 0 terminated with signal SIGSEGV

Error ini sangat kriptik, tetapi penjelasan dasarnya adalah RAM sistem Anda habis. Anda dapat sepenuhnya menghindari hal ini dengan mengonfigurasi ulang fungsi pelatihan agar menerima satu argumen model, dan mendeklarasikannya di sel di luar:

# 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"]
      )
  ...

Dan akhirnya memanggil fungsi pelatihan dengan:

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

Workaround di atas hanya diperlukan saat meluncurkan instance TPU dari Jupyter Notebook di server dengan sumber daya rendah seperti Google Colaboratory atau Kaggle. Jika menggunakan skrip atau meluncurkan di server yang jauh lebih bertenaga, deklarasi model sebelumnya tidak diperlukan.

Presisi Campuran dan Variabel Global

Seperti yang disebutkan dalam bagian presisi campuran dari tutorial migrasi, Accelerate mendukung fp16 dan bf16, dan keduanya dapat digunakan pada TPU. Meskipun demikian, idealnya bf16 yang digunakan karena sangat efisien.

Ada dua “lapisan” saat menggunakan bf16 dan Accelerate pada TPU, yaitu pada level dasar dan pada level operasi.

Pada level dasar, ini diaktifkan saat meneruskan mixed_precision="bf16" ke Accelerator, seperti:

accelerator = Accelerator(mixed_precision="bf16")

Secara default, ini akan mengonversi torch.float dan torch.double menjadi bfloat16 pada TPU. Konfigurasi spesifik yang diatur adalah variabel lingkungan XLA_USE_BF16 yang diset ke 1.

Ada konfigurasi lain yang dapat Anda lakukan, yaitu mengatur variabel lingkungan XLA_DOWNCAST_BF16. Jika diatur ke 1, maka torch.float adalah bfloat16 dan torch.double adalah float32.

Hal ini dilakukan pada objek Accelerator saat meneruskan downcast_bf16=True:

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

Menggunakan downcasting alih-alih bf16 di mana-mana sangat berguna ketika Anda mencoba menghitung metrik, mencatat nilai, dan lainnya, di mana tensor bf16 mentah tidak dapat digunakan.

Waktu Pelatihan pada TPU

Saat Anda menjalankan skrip, Anda mungkin akan memperhatikan bahwa pelatihan terasa sangat lambat pada awalnya. Hal ini terjadi karena TPU terlebih dahulu memproses beberapa batch data untuk melihat berapa banyak memori yang perlu dialokasikan sebelum akhirnya memanfaatkan alokasi memori yang telah dikonfigurasi ini dengan sangat efisien.

Jika Anda menyadari bahwa kode evaluasi Anda untuk menghitung metrik model Anda membutuhkan waktu lebih lama karena ukuran batch yang lebih besar digunakan, disarankan untuk mempertahankan ukuran batch yang sama dengan data pelatihan jika prosesnya terlalu lambat. Jika tidak, memori akan dialokasikan ulang ke ukuran batch baru ini setelah beberapa iterasi pertama.

Fakta bahwa memori telah dialokasikan tidak berarti memori tersebut akan digunakan atau bahwa ukuran batch akan meningkat saat kembali ke dataloader pelatihan Anda.