FSDP2

Fully Sharded Data Parallel (FSDP2), modeli, gradyanları ve optimizör durumlarını GPU'lar arasında parçalar. Hesaplama öncesinde her GPU, tüm parçalardan tam bir parametre seti toplar, ardından bunları serbest bırakır. Parçalama, tek bir GPU'nun belleğinden daha büyük modelleri eğitmenize olanak tanır; bunun bedeli ise DDP'ye kıyasla daha fazla iletişimdir. Modeliniz veya optimizör durumlarınız tek bir GPU'ya sığmıyorsa FSDP kullanın.

                      ┌─────────────────┐
                      │  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

Sharding stratejileri

FSDP2, şardingi ~TrainingArguments.fsdp_config ile kontrol eder. FSDP'yi etkinleştirmek için fsdp=True olarak ayarlayın ve bellek ile verimlilik arasındaki dengeyi seçmek için FSDP yapılandırmasında reshard_after_forward değerini belirleyin.

auto_wrap_policy, modüllerin FSDP birimlerine nasıl sarılacağını kontrol eder. Varsayılan değeri "TRANSFORMER_BASED_WRAP"'dir ve modelin transformer katmanlarını sarar. Sarılma yapılmazsa ("NO_WRAP"), tüm model tek bir FSDP birimi olur ve sharding'in sağladığı bellek avantajını kaybedersiniz.

FSDP'yi yapılandır

Bu alanlar, FSDP2'nin modelin nasıl sarıldığını, bölündüğünü ve yüklendiğini kontrol eder. reshard_after_forward ve auto_wrap_policy, Bölme stratejileri bölümünde ele alınır.

• cpu_offload, GPU belleğini tasarruflu kullanmak için parametreleri ve gradyanları kullanılmadıkları zaman CPU'ya taşır.

• transformer_layer_cls_to_wrap, auto_wrap_policy "TRANSFORMER_BASED_WRAP" olduğunda, bir FSDP birimine sarılacak transformer katmanını tanımlar. Her birim kendi toplama (gather) ve dağıtma (scatter) işlemlerini yönetir. İleri geçiş sırasında yalnızca geçerli birimin parametreleri toplanır. Bellek tasarrufu sağlamak için önceki birimlerin parametreleri serbest bırakılır.

  Sadece üst düzey modeli sarmak GPU bellek tasarrufu sağlamaz. Her bir Linear katmanını ayrı ayrı sarmak birimler arası iletişimi çok pahalı hale getirir. Bu alanı boş bırakın ve FSDP değeri model tanımından okusun.

• min_num_params, boyuta dayalı sarmalama için modül başına minimum parametre sayısını belirler. Yalnızca auto_wrap_policy "SIZE_BASED_WRAP" olduğunda kullanılır.

• state_dict_type, checkpoint biçimini kontrol eder. Varsayılan olarak, tek bir Transformers uyumlu checkpoint için "FULL_STATE_DICT" kullanılır. Büyük modeller için daha hızlı olan, rank başına bir checkpoint dosyası içeren "SHARDED_STATE_DICT" seçeneği kullanılabilir. Parçalı checkpoint'ler yalnızca FSDP içine geri yüklenebildiğinden, paylaşmak veya FSDP dışında yüklemek istediğiniz son checkpoint için bir "FULL_STATE_DICT" kaydedin.

• cpu_ram_efficient_loading, kontrol noktasını yalnızca 0. sıradan diskten yükler. Diğer GPU'lar boş bir model başlatır ve ağırlıkları yayın yoluyla alır; bu sayede büyük bir modelin CPU RAM'ine birden fazla süreç tarafından yüklenmesi önlenir.

• activation_checkpointing, geriye yayılım sırasında aktivasyonları saklamak yerine yeniden hesaplar. TrainingArguments içinde gradient checkpointing yerine bunu kullanın. İkisini birden ayarlamak bir hata verir.

FSDP eğitimini, bir Accelerate yapılandırma dosyası veya fsdp_config parametresine geçirilen bir FSDP yapılandırma dosyası ile yapılandırın.

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",
)

Mixture-of-experts modelleri için, fsdp_size, uzmanları bir mesh boyutunda, diğer her şeyi ise diğer boyutta parçalamak üzere tp_size ve ep_size ile birleştirilebilir. expert parallelism bölümüne bakın.

Bağlam paralelliği ile değerlendirme

FSDP2 ile bağlam paralelliği (CP), Accelerate ParallelismConfig içindeki cp_size parametresiyle ayarlanır. Başlatırken hem --use_parallelism_config hem de --parallelism_config_cp_size parametrelerini iletin. Accelerate YAML dosyasında boş olmayan bir parallelism_config alanı, use_parallelism_config seçeneğini de etkinleştirir. Yalnızca --parallelism_config_cp_size verilirse bu parametre yok sayılır ve cp_size=1 değeri korunur.

FSDP2 altında CP, diziyi ranklar arasında paylaştırır; böylece her GPU bağlamın bir dilimini tutar (allgather veya alltoall). Bu, DeepSpeed ALST üzerinden çalışan ve aynı yapılandırmada sp_size kullanan Ulysses sequence parallelism'den farklıdır. CP veya Ulysses SP'den birini seçin. cp_size ve sp_size'nin ikisini birden 1'in üzerinde ayarlamayın.

evaluate(), prediction_step() üzerinden train() ile aynı bağlam paralel bağlamını kullanır. Her rank, bu bağlam etkinken kendi dizi dilimini çalıştırır. eval_loss doğru kayıp değeridir (cp_size ile çarpılmamıştır). Erken durdurma ve load_best_model_at_end bu metriği kullanır.

Logit'ler, diz boyu boyunca parçalı (sharded) kalır ve bunları tam uzunlukta tahminlere dönüştüren bir yeniden birleştirme işlemi yoktur. CP altında logit'ler üzerinden tahminler ve compute_metrics desteklenmez.

eval_loss gerektiğinde mevcut FSDP2 + cp_size başlatma komutunuzu yeniden kullanın.

prediction_step() alt sınıfını oluşturup jetonları kendiniz sayıyorsanız, num_items_in_batch değerini context-parallel bağlamına girmeden önce sayın. Bu bağlam etkinken, dizi tamponları yerinde olarak rank'lar arasında paylaştırılır. Sonradan sayım yaparsanız batch eksik sayılır. Varsayılan Trainer zaten önce sayım yapar ve yalnızca prediction_step işlevini geçersiz kılıyorsanız buna ihtiyacınız vardır.

Sonraki adımlar

• Modeliniz tek bir GPU'ya sığıyorsa veri paralel eğitim için DDP'ye bakın.
• ZeRO optimizasyonu ve NVMe dışa aktarma için DeepSpeed'e bakın.
• PyTorch/XLA ile TPU'lar üzerinde FSDP için ~TrainingArguments.fsdp_config içinde xla, xla_fsdp_settings ve xla_fsdp_grad_ckpt değerlerini ayarlayın.
• FSDP'nin nasıl çalıştığına dair daha fazla bilgi için The Ultra-Scale Playbook'taki FSDP bölümünü okuyun.