Pembongkaran aktivasi

penting

Di SMP v2.2.0, fungsionalitas pembongkaran aktivasi perpustakaan SMP tidak berfungsi. Gunakan pembongkaran PyTorch aktivasi asli sebagai gantinya.

Biasanya, pass maju menghitung aktivasi di setiap lapisan dan menyimpannya dalam memori GPU sampai pass mundur untuk lapisan yang sesuai selesai. Membongkar tensor ini ke memori CPU setelah forward pass dan mengambilnya kembali ke GPU saat dibutuhkan untuk backward pass lapisan dapat menghemat penggunaan memori GPU yang substanal. PyTorchmendukung aktivasi pembongkaran, tetapi implementasinya menyebabkan GPU menjadi idle sementara aktivasi diambil kembali dari CPU selama backward pass. Hal ini menyebabkan penurunan kinerja yang besar saat menggunakan pembongkaran aktivasi.

SMP v2 meningkatkan pembongkaran aktivasi ini, dan melakukan pra-pengambilan aktivasi sebelumnya sebelum aktivasi diperlukan agar GPU mulai meneruskan aktivasi tersebut secara mundur. Fitur pra-pengambilan ini membantu kemajuan pelatihan berjalan lebih efisien tanpa GPU idle, yang menghasilkan menawarkan manfaat dari penggunaan memori yang lebih rendah tanpa penurunan kinerja.

Anda dapat menyimpan PyTorch modul asli untuk pembongkaran aktivasi dalam skrip pelatihan Anda. Berikut ini adalah contoh struktur penerapan fitur pembongkaran aktivasi SMP di skrip Anda. Perhatikan bahwa pembongkaran aktivasi hanya berlaku jika digunakan bersama denganPos pemeriksaan aktivasi. Untuk mempelajari lebih lanjut tentang alat PyTorch pos pemeriksaan asli untuk pembongkaran aktivasi, lihat juga checkpoint_wrapper.py di PyTorch GitHubrepositori dan Activation Checkpointing di PyTorch blog Scaling Multi-modal Foundation Models in with Distributed. TorchMultimodal PyTorch

Untuk menerapkan fitur pembongkaran aktivasi SMP pada pos pemeriksaan PyTorch aktivasi, tambahkan sm_activation_offloading dan activation_loading_horizon parameter ke kamus konfigurasi SMP selama. Langkah 2: Luncurkan pekerjaan pelatihan

Cuplikan kode berikut menunjukkan cara menambahkan modul inisialisasi SMP torch.sagemaker.init() ke skrip pelatihan Anda dan mengatur kamus konfigurasi SMP dalam format JSON untuk peluncur pekerjaan pelatihan sambil mengikuti proses dua langkah yang diperkenalkan. Memulai dengan SageMaker pustaka paralelisme model v2 Anda tidak perlu membuat perubahan apa pun pada PyTorch model atau konfigurasi PyTorch FSDP Anda. Untuk informasi tentang parameter sm_activation_offloading dan activation_loading_horizon, lihat Parameter konfigurasi fitur inti SMP v2.

Konfigurasi SMP

{
    "activation_loading_horizon": 2,
    "sm_activation_offloading": True
}

Dalam naskah pelatihan

catatan

Saat mengaktifkan fitur pembongkaran aktivasi SMP, pastikan Anda juga menggunakan PyTorch offload_wrapper fungsi tersebut dan menerapkannya ke modul root. Fitur pembongkaran aktivasi SMP menggunakan modul root untuk menentukan kapan pass maju dilakukan untuk memulai pra-pengambilan.

import torch.sagemaker as tsm
tsm.init()

# Native PyTorch module for activation offloading
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    apply_activation_checkpointing, 
    offload_wrapper,
)

model = FSDP(...)

# Activation offloading requires activation checkpointing.
apply_activation_checkpointing(
    model,
    check_fn=checkpoint_transformer_layers_policy,
)

model = offload_wrapper(model)