FSDP2

Fully Sharded Data Parallel (FSDP2) membagi model, gradien, dan status optimizer di antara GPU. Sebelum komputasi, setiap GPU mengumpulkan set parameter lengkap dari semua shard, lalu membebaskannya setelahnya. Sharding memungkinkan Anda melatih model yang lebih besar dari memori GPU tunggal, dengan biaya komunikasi yang lebih tinggi daripada DDP. Gunakan FSDP ketika model atau status optimizer Anda tidak muat di satu GPU.

                      ┌─────────────────┐
                      │  training data  │
                      └────────┬────────┘
            ┌──────────────────┼──────────────────┐
            │ shard 0          │ shard 1          │ shard 2
            ▼                  ▼                  ▼
     ┌─────────────┐    ┌─────────────┐    ┌─────────────┐
     │  param      │    │  param      │    │  param      │
     │  shard 0    │    │  shard 1    │    │  shard 2    │
     │  GPU 0      │    │  GPU 1      │    │  GPU 2      │
     └──────┬──────┘    └──────┬──────┘    └──────┬──────┘
            │                  │                  │
            └──────── all-gather (params) ────────┘
                               │
                    full params on each GPU
                               │
            ┌──────────────────┼──────────────────┐
            ▼                  ▼                  ▼
         forward             forward             forward
            │                  │                  │
            └───── reduce-scatter (grads) ────────┘
                               │
            ┌──────────────────┼──────────────────┐
            ▼                  ▼                  ▼
     grad shard 0       grad shard 1       grad shard 2
     optim shard 0      optim shard 1      optim shard 2
        step               step               step

Strategi sharding

FSDP2 mengontrol sharding dengan ~TrainingArguments.fsdp_config. Atur fsdp=True untuk mengaktifkan FSDP, dan atur reshard_after_forward dalam konfigurasi FSDP untuk memilih kompromi antara memori dan throughput.

auto_wrap_policy mengontrol bagaimana modul dibungkus menjadi unit FSDP. Nilainya secara default adalah "TRANSFORMER_BASED_WRAP", yang membungkus lapisan transformer model. Tanpa pembungkusan ("NO_WRAP"), seluruh model menjadi satu unit FSDP dan Anda kehilangan manfaat memori dari sharding.

Konfigurasi FSDP

Field-field ini mengontrol cara FSDP2 membungkus, memecah, dan memuat model. reshard_after_forward dan auto_wrap_policy dibahas dalam Strategi sharding.

• cpu_offload memindahkan parameter dan gradien ke CPU saat tidak digunakan untuk menghemat memori GPU.

• transformer_layer_cls_to_wrap mendefinisikan lapisan transformer yang akan dibungkus menjadi unit FSDP ketika auto_wrap_policy adalah "TRANSFORMER_BASED_WRAP". Setiap unit mengelola operasi gather dan scatter-nya sendiri. Hanya parameter dari unit saat ini yang dikumpulkan selama forward pass. Parameter dari unit-unit sebelumnya dilepaskan untuk menghemat memori.

  Membungkus hanya model tingkat atas tidak menghasilkan penghematan memori GPU. Membungkus setiap lapisan Linear individu membuat komunikasi antar-unit sangat mahal. Biarkan bidang ini kosong dan FSDP akan membaca nilai dari definisi model.

• min_num_params mengatur jumlah parameter minimum per modul untuk pembungkusan berbasis ukuran. Parameter ini hanya digunakan ketika auto_wrap_policy adalah "SIZE_BASED_WRAP".

• state_dict_type mengontrol format checkpoint. Nilainya secara default adalah "FULL_STATE_DICT" untuk satu checkpoint yang kompatibel dengan Transformers. Gunakan "SHARDED_STATE_DICT" untuk satu file checkpoint per rank, yang lebih cepat untuk model besar. Checkpoint sharded hanya dapat dimuat kembali ke FSDP, jadi simpan "FULL_STATE_DICT" untuk checkpoint final yang ingin Anda bagikan atau muat di luar FSDP.

• cpu_ram_efficient_loading memuat checkpoint dari disk hanya pada rank 0. GPU lain menginisialisasi model kosong dan menerima bobot melalui broadcast, sehingga menghindari beberapa proses memuat model besar ke RAM CPU.

• activation_checkpointing menghitung ulang aktivasi selama backward pass alih-alih menyimpannya. Gunakan ini sebagai pengganti gradient checkpointing di TrainingArguments. Mengaktifkan keduanya akan memicu error.

Konfigurasi pelatihan FSDP dengan menggunakan file konfigurasi Accelerate atau file konfigurasi FSDP yang diteruskan ke fsdp_config.

accelerate launch train.py

{
  "version": 2,
  "reshard_after_forward": true,
  "cpu_offload": false,
  "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
  "transformer_layer_cls_to_wrap": ["LlamaDecoderLayer"],
  "state_dict_type": "FULL_STATE_DICT",
  "cpu_ram_efficient_loading": true,
  "activation_checkpointing": true
}

from transformers import TrainingArguments

TrainingArguments(
    ...,
    fsdp=True,
    fsdp_config="path/to/fsdp.json",
)

Untuk model mixture-of-experts, fsdp_size dapat dikombinasikan dengan tp_size dan ep_size untuk membagi para ahli di satu dimensi mesh dan semua yang lainnya di dimensi lainnya. Lihat expert parallelism.

Evaluasi dengan paralelisme konteks

Context parallelism (CP) dengan FSDP2 diatur melalui cp_size pada ParallelismConfig Accelerate. Saat peluncuran, berikan --use_parallelism_config dan --parallelism_config_cp_size. parallelism_config yang tidak kosong dalam YAML Accelerate juga mengaktifkan use_parallelism_config. --parallelism_config_cp_size yang berdiri sendiri akan diabaikan dan tetap mempertahankan cp_size=1.

Di bawah FSDP2, CP membagi (shard) sekuens di antara rank sehingga setiap GPU memegang potongan dari konteks (allgather atau alltoall). Hal ini berbeda dari Ulysses sequence parallelism, yang berjalan melalui DeepSpeed ALST dan menggunakan sp_size pada konfigurasi yang sama. Pilih CP atau Ulysses SP. Jangan mengatur kedua cp_size dan sp_size di atas 1.

evaluate() menggunakan konteks paralel konteks yang sama dengan train() melalui prediction_step(). Setiap rank menjalankan shard sekuensnya selama konteks tersebut aktif. eval_loss adalah loss yang benar (tidak dikalikan dengan cp_size). Early stopping dan load_best_model_at_end menggunakan metrik tersebut.

Logits tetap terpecah (sharded) sepanjang urutan, dan tidak ada yang merakitnya kembali menjadi prediksi berpanjang penuh. Prediksi dan compute_metrics di atas logits tidak didukung di bawah CP.

Gunakan ulang peluncuran FSDP2 + cp_size yang sudah ada saat Anda membutuhkan eval_loss.

Jika Anda membuat subclass prediction_step() dan menghitung token secara mandiri, hitung num_items_in_batch sebelum memasuki konteks parallel-context. Selama konteks tersebut aktif, buffer sekuens di-shard secara in-place di antara rank. Menghitung setelahnya akan menghasilkan jumlah batch yang kurang. Trainer default sudah menghitung terlebih dahulu, dan Anda hanya perlu melakukan ini jika Anda meng-override prediction_step.

Langkah selanjutnya

• Lihat DDP untuk pelatihan paralel data ketika model Anda muat di satu GPU.
• Lihat DeepSpeed untuk optimasi ZeRO dan offloading NVMe.
• Untuk FSDP pada TPU dengan PyTorch/XLA, atur xla, xla_fsdp_settings, dan xla_fsdp_grad_ckpt di ~TrainingArguments.fsdp_config.
• Baca bab FSDP dari The Ultra-Scale Playbook untuk informasi lebih lanjut tentang cara kerja FSDP.