Fully Sharded Data Parallel

Untuk mempercepat pelatihan model berukuran besar pada ukuran batch yang lebih besar, kita dapat menggunakan model data paralel yang sepenuhnya di-shard. Paradigma data paralel jenis ini memungkinkan muat lebih banyak data dan model yang lebih besar dengan cara meng-shard state optimizer, gradien, dan parameter. Untuk membaca lebih lanjut tentang hal ini dan manfaatnya, lihat blog Fully Sharded Data Parallel. Kami telah mengintegrasikan fitur pelatihan Fully Sharded Data Parallel (FSDP) terbaru dari PyTorch. Yang perlu Anda lakukan hanyalah mengaktifkannya melalui config.

Cara kerjanya secara bawaan

Pada mesin Anda, cukup jalankan:

accelerate config

dan menjawab pertanyaan yang diajukan. Hal ini akan menghasilkan file konfigurasi yang akan digunakan secara otomatis untuk mengatur opsi default dengan benar saat melakukan

accelerate launch my_script.py --args_to_my_script

Sebagai contoh, berikut adalah cara menjalankan examples/nlp_example.py (dari root repo) dengan FSDP diaktifkan:

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_forward_prefetch: false
  fsdp_cpu_ram_efficient_loading: true
  fsdp_offload_params: false
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_transformer_layer_cls_to_wrap: BertLayer
  fsdp_use_orig_params: true
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

accelerate launch examples/nlp_example.py

Saat ini, Accelerate mendukung konfigurasi berikut melalui CLI:

fsdp_sharding_strategy: [1] FULL_SHARD (memecah optimizer states, gradients, dan parameters), [2] SHARD_GRAD_OP (memecah optimizer states dan gradients), [3] NO_SHARD (DDP), [4] HYBRID_SHARD (memecah optimizer states, gradients, dan parameters di dalam setiap node, sementara setiap node memiliki salinan penuh), [5] HYBRID_SHARD_ZERO2 (memecah optimizer states dan gradients di dalam setiap node, sementara setiap node memiliki salinan penuh). Untuk informasi lebih lanjut, silakan merujuk ke dokumentasi resmi PyTorch.

fsdp_offload_params : Menentukan apakah parameter dan gradien akan di-offload ke CPU

fsdp_auto_wrap_policy: [1] TRANSFORMER_BASED_WRAP, [2] SIZE_BASED_WRAP, [3] NO_WRAP

fsdp_transformer_layer_cls_to_wrap: Hanya berlaku untuk Transformers. Saat menggunakan fsdp_auto_wrap_policy=TRANSFORMER_BASED_WRAP, pengguna dapat memberikan string yang dipisahkan koma berisi nama kelas layer transformer (peka huruf besar/kecil) yang akan dibungkus, misalnya BertLayer, GPTJBlock, T5Block, BertLayer,BertEmbeddings,BertSelfOutput. Hal ini penting karena submodul yang berbagi bobot (misalnya layer embedding) tidak boleh berakhir di unit FSDP yang berbeda. Dengan kebijakan ini, pembungkusan terjadi untuk setiap blok yang berisi Multi-Head Attention diikuti beberapa layer MLP. Layer yang tersisa, termasuk embedding yang dibagikan, dibungkus dengan nyaman di unit FSDP terluar yang sama. Oleh karena itu, gunakan ini untuk model berbasis transformer. Anda dapat menggunakan model._no_split_modules untuk model Transformer dengan menjawab yes untuk Do you want to use the model's_no_split_modulesto wrap. It will try to usemodel._no_split_modules` when possible.

fsdp_min_num_params: jumlah parameter minimum saat menggunakan fsdp_auto_wrap_policy=SIZE_BASED_WRAP.

fsdp_backward_prefetch_policy: [1] BACKWARD_PRE, [2] BACKWARD_POST, [3] NO_PREFETCH

fsdp_forward_prefetch: jika True, maka FSDP secara eksplisit melakukan prefetch untuk all-gather berikutnya saat dieksekusi dalam forward pass. Sebaiknya hanya digunakan untuk model dengan graf statis karena prefetching mengikuti urutan eksekusi dari iterasi pertama. Artinya, jika urutan sub-modul berubah secara dinamis selama eksekusi model, jangan aktifkan fitur ini.

fsdp_state_dict_type: [1] FULL_STATE_DICT, [2] LOCAL_STATE_DICT, [3] SHARDED_STATE_DICT

fsdp_use_orig_params: Jika True, memungkinkan requires_grad yang tidak seragam saat inisialisasi, yang berarti mendukung parameter beku dan parameter yang dapat dilatih secara berselang-seling. Pengaturan ini berguna dalam kasus seperti fine-tuning yang efisien parameter seperti yang dibahas dalam postingan ini. Opsi ini juga memungkinkan adanya beberapa grup parameter optimizer. Opsi ini harus True saat membuat optimizer sebelum menyiapkan/membungkus model dengan FSDP.

fsdp_cpu_ram_efficient_loading: Hanya berlaku untuk model Transformers. Jika True, hanya proses pertama yang memuat checkpoint model pretrained, sementara semua proses lainnya memiliki bobot kosong. Pengaturan ini harus di设为 False jika Anda mengalami kesalahan saat memuat model Transformers pretrained melalui metode from_pretrained. Ketika pengaturan ini True, fsdp_sync_module_states juga harus True, jika tidak semua proses kecuali proses utama akan memiliki bobot acak yang menyebabkan perilaku tak terduga selama pelatihan. Agar ini berfungsi, pastikan grup proses terdistribusi diinisialisasi sebelum memanggil metode from_pretrained Transformers. Saat menggunakan API Trainer, grup proses terdistribusi diinisialisasi saat Anda membuat instance dari kelas TrainingArguments.

fsdp_sync_module_states: Jika True, setiap unit FSDP yang dibungkus secara individual akan menyiarkan parameter modul dari rank 0.

Untuk kontrol tambahan yang lebih halus, Anda dapat menentukan parameter FSDP lainnya melalui FullyShardedDataParallelPlugin. Saat membuat objek FullyShardedDataParallelPlugin, berikan parameter yang bukan bagian dari konfigurasi accelerate atau jika Anda ingin menimpanya. Parameter FSDP akan dipilih berdasarkan file konfigurasi accelerate atau argumen perintah peluncuran, dan parameter lain yang Anda berikan langsung melalui objek FullyShardedDataParallelPlugin akan mengatur/menimpa parameter tersebut.

Berikut adalah contoh:

from accelerate import FullyShardedDataParallelPlugin
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig

fsdp_plugin = FullyShardedDataParallelPlugin(
    state_dict_config=FullStateDictConfig(offload_to_cpu=False, rank0_only=False),
    optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=False, rank0_only=False),
)

accelerator = Accelerator(fsdp_plugin=fsdp_plugin)

Menyimpan dan memuat

Cara baru yang direkomendasikan untuk checkpointing saat menggunakan model FSDP adalah dengan menggunakan SHARDED_STATE_DICT sebagai StateDictType saat mengatur konfigurasi accelerate. Berikut adalah potongan kode untuk menyimpan menggunakan utilitas save_state dari accelerate.

accelerator.save_state("ckpt")

Periksa folder checkpoint untuk melihat model dan optimizer sebagai shard per proses:

ls ckpt
# optimizer_0  pytorch_model_0  random_states_0.pkl  random_states_1.pkl  scheduler.bin

cd ckpt

ls optimizer_0
# __0_0.distcp  __1_0.distcp

ls pytorch_model_0
# __0_0.distcp  __1_0.distcp

Untuk memuatnya kembali guna melanjutkan pelatihan, gunakan utilitas load_state dari accelerate

accelerator.load_state("ckpt")

Saat menggunakan save_pretrained dari transformers, berikan state_dict=accelerator.get_state_dict(model) untuk menyimpan state dict model. Berikut adalah contohnya:

  unwrapped_model.save_pretrained(
      args.output_dir,
      is_main_process=accelerator.is_main_process,
      save_function=accelerator.save,
+     state_dict=accelerator.get_state_dict(model),
)

State Dict

accelerator.get_state_dict akan memanggil implementasi model.state_dict yang mendasarinya menggunakan context manager FullStateDictConfig(offload_to_cpu=True, rank0_only=True) untuk mendapatkan state dict hanya untuk rank 0, dan hasilnya akan di-offload ke CPU.

Anda kemudian dapat meneruskan state ke metode save_pretrained. Ada beberapa mode untuk StateDictType dan FullStateDictConfig yang dapat Anda gunakan untuk mengontrol perilaku state_dict. Untuk informasi lebih lanjut, lihat dokumentasi PyTorch.

Jika Anda memilih untuk menggunakan StateDictType.SHARDED_STATE_DICT, bobot model saat Accelerator.save_state akan dipecah menjadi n file untuk setiap sub-pecahan pada model. Untuk menggabungkannya kembali menjadi satu kamus tunggal agar dapat dimuat kembali ke model setelah pelatihan, Anda dapat menggunakan utilitas merge_weights:

from accelerate.utils import merge_fsdp_weights

# Our weights are saved usually in a `pytorch_model_fsdp_{model_number}` folder
merge_fsdp_weights("pytorch_model_fsdp_0", "output_path", safe_serialization=True)

Output akhir kemudian akan disimpan ke model.safetensors atau pytorch_model.bin (jika safe_serialization=False diteruskan).

Hal ini juga dapat dipanggil menggunakan CLI:

accelerate merge-weights pytorch_model_fsdp_0/ output_path

Pemetaan antara strategi sharding FSDP dan Tahapan ZeRO DeepSpeed

• FULL_SHARD dipetakan ke ZeRO Stage-3 DeepSpeed. Memecah (sharding) status optimizer, gradien, dan parameter.
• SHARD_GRAD_OP dipetakan ke ZeRO Stage-2 DeepSpeed. Memecah (sharding) status optimizer dan gradien.
• NO_SHARD dipetakan ke ZeRO Stage-0. Tidak ada sharding, di mana setiap GPU memiliki salinan penuh dari model, state optimizer, dan gradien.
• HYBRID_SHARD dipetakan ke ZeRO++ Stage-3 di mana zero_hpz_partition_size=<num_gpus_per_node>. Di sini, hal ini akan memecah optimizer states, gradients, dan parameters di dalam setiap node, sementara setiap node memiliki salinan penuh.

Beberapa hal yang perlu diperhatikan

• Jika terdapat beberapa model, berikan optimizer ke panggilan prepare dalam urutan yang sama dengan model yang sesuai, selain itu accelerator.save_state() dan accelerator.load_state() akan menghasilkan perilaku yang salah/tidak terduga.
• Fitur ini tidak kompatibel dengan --predict_with_generate dalam skrip run_translation.py dari pustaka Transformers.

Untuk kontrol yang lebih besar, pengguna dapat memanfaatkan FullyShardedDataParallelPlugin. Setelah membuat instance dari kelas ini, pengguna dapat meneruskannya ke instansiasi kelas Accelerator. Untuk informasi lebih lanjut tentang opsi-opsi ini, silakan merujuk ke kode FullyShardedDataParallel PyTorch.

Bagi yang tertarik dengan kesamaan dan perbedaan antara FSDP dan DeepSpeed, silakan lihat panduan konsep di sini!