DeepSpeed czy FSDP w PyTorch: jak radzić sobie z Out-of-Memory przy trenowaniu dużych modeli?

Out-of-Memory podczas treningu: problem pamięci, nie tylko „za małej karty”

Błąd CUDA Out of Memory pojawiający się podczas treningu dużego modelu nie musi oznaczać, że projekt wymaga natychmiast zakupu mocniejszych GPU. Często źródłem problemu jest nieoptymalna konfiguracja treningu, zbyt długi kontekst, niewłaściwie ustawiony mikrobatch albo replikowanie danych, które mogłyby zostać podzielone między procesy. DeepSpeed i PyTorch FSDP rozwiązują część tych problemów, ale sensowny wybór wymaga najpierw ustalenia, co dokładnie konsumuje pamięć. ^[3]

Najtrudniejsza sytuacja występuje wtedy, gdy model mieści się na karcie podczas inferencji, lecz trening kończy się OOM już po pierwszym lub kilku krokach. To normalne: uczenie wymaga przechowywania znacznie większej liczby tensorów niż samo generowanie odpowiedzi. W modelach transformerowych skala zużycia pamięci szczególnie szybko rośnie wraz z długością sekwencji, liczbą warstw i rozmiarem ukrytym modelu.

Co naprawdę zajmuje pamięć GPU podczas uczenia

Podczas treningu pamięć GPU jest obciążona przez cztery główne grupy danych: parametry modelu, gradienty, stany optymalizatora oraz aktywacje zapisane na potrzeby propagacji wstecznej. Parametry są wagami modelu. Gradienty określają kierunek ich aktualizacji. Stany optymalizatora przechowują dodatkowe informacje wykorzystywane do stabilnego uczenia, a aktywacje to tensory pośrednie z przebiegu forward pass, potrzebne później w obliczeniach gradientów.

Przy standardowym AdamW koszt pamięciowy może być zaskakująco wysoki. Dla parametrów trzymanych w formacie FP16 lub BF16 często nadal potrzebna jest kopia główna w FP32, a optymalizator przechowuje dwa dodatkowe stany momentów. W praktyce sam stan modelu, bez aktywacji, może zajmować wielokrotność rozmiaru wag widocznego w pliku checkpointu. Dlatego model zapisany na dysku jako kilkanaście gigabajtów nie jest automatycznie możliwy do trenowania na GPU o podobnej pojemności VRAM.

Aktywacje są z kolei silnie zależne od danych wejściowych. W uproszczeniu ich zużycie rośnie wraz z iloczynem liczby warstw, rozmiaru mikrobatcha i długości sekwencji. W architekturach z pełnym mechanizmem uwagi istotne znaczenie ma też koszt operacji attention, który dla klasycznych implementacji może rosnąć bardzo niekorzystnie wraz z długością kontekstu. Zwiększenie max_length z 2048 do 4096 tokenów może więc wywołać OOM mimo identycznego batch size i niezmienionego modelu.

Dlaczego błąd OOM bywa mylący

Komunikat o braku pamięci nie zawsze oznacza, że VRAM został w całości zajęty przez aktualnie potrzebne tensory. PyTorch używa własnego mechanizmu cache’owania alokacji CUDA. Dzięki temu kolejne operacje są szybsze, ale narzędzia systemowe mogą pokazywać pamięć jako zajętą, choć część z niej jest zarezerwowana przez alokator i może zostać użyta ponownie. Różnicę można obserwować przez torch.cuda.memory_allocated() oraz torch.cuda.memory_reserved().

Jeżeli wartość memory_reserved() jest wyraźnie większa niż memory_allocated(), problemem może być fragmentacja lub sposób rozdzielania bloków pamięci. Nie jest to jednak argument za automatycznym wywoływaniem torch.cuda.empty_cache() po każdym kroku. Ta funkcja zwalnia nieużywany cache do sterownika CUDA, ale nie redukuje pamięci wymaganej przez model, gradienty czy aktywacje. Nadmierne korzystanie z niej zwykle spowalnia trening.

W środowisku współdzielonym trzeba też sprawdzić, czy na GPU nie działają inne procesy: sesje notebooków, stare procesy Pythona, serwery inferencyjne albo procesy po nieudanym treningu. Narzędzie nvidia-smi pokazuje ogólny obraz, natomiast szczegółowa analiza powinna uwzględniać pamięć każdego procesu w treningu rozproszonym. Nierówny podział danych, inne długości sekwencji na poszczególnych rankach lub asynchroniczne ładowanie modelu mogą sprawić, że OOM wystąpi tylko na jednej karcie.

Pomiar przed zmianą architektury treningu

Zanim pojawi się DeepSpeed ZeRO lub sharding parametrów w FSDP, dobrze jest uchwycić moment skoku zużycia VRAM. Pomaga prosty pomiar po załadowaniu modelu, po przejściu forward, po obliczeniu straty i po backward. Jeżeli pamięć rośnie przede wszystkim po forward pass, głównym podejrzanym są aktywacje i długość sekwencji. Jeżeli model nie mieści się już podczas inicjalizacji optymalizatora, dominują parametry oraz jego stany.

def report_memory(label):
    allocated = torch.cuda.memory_allocated() / 1024**3
    reserved = torch.cuda.memory_reserved() / 1024**3
    peak = torch.cuda.max_memory_allocated() / 1024**3
    print(f"{label}: allocated={allocated:.2f} GB, "
          f"reserved={reserved:.2f} GB, peak={peak:.2f} GB")

torch.cuda.reset_peak_memory_stats()
report_memory("po załadowaniu modelu")

outputs = model(**batch)
report_memory("po forward")

loss = outputs.loss
loss.backward()
report_memory("po backward")

Takie dane ułatwiają decyzję. Jeśli pojedynczy mikrobatch nie mieści się z powodu aktywacji, samo przejście z DDP na pełne shardowanie może nie wystarczyć. Jeśli zaś problemem są stany modelu i optymalizatora, FSDP lub DeepSpeed ZeRO-3 mogą przynieść bardzo dużą ulgę pamięciową.

Pierwsza linia obrony przed OOM w czystym PyTorch

Sharding nie powinien maskować podstawowych błędów konfiguracji. Najprostsze optymalizacje często obniżają użycie pamięci na tyle, że nie trzeba komplikować infrastruktury. Co ważne, nie konkurują one z DeepSpeed ani FSDP. W dojrzałej konfiguracji treningu zwykle działają razem.

Mixed precision: FP16, BF16 i realna stabilność

Mixed precision ogranicza zużycie pamięci przez część tensorów oraz przyspiesza obliczenia na GPU wspierających odpowiednie jednostki sprzętowe. Format FP16 zmniejsza rozmiar wielu danych o połowę względem FP32, ale ma wąski zakres reprezentowanych wartości. Z tego powodu trening w FP16 często wymaga mechanizmu loss scaling, który ogranicza ryzyko zanikania gradientów.

BF16 ma taki sam rozmiar jak FP16, lecz zakres wykładnika zbliżony do FP32. W praktyce mixed precision BF16 jest zwykle wygodniejsza dla nowoczesnych modeli językowych, ponieważ rzadziej wymaga ręcznego reagowania na problemy numeryczne. Nie każda karta GPU obsługuje BF16 równie dobrze, dlatego przed wyborem dtype należy sprawdzić możliwości konkretnej generacji sprzętu.

Przykładowe użycie automatycznej precyzji w PyTorch może wyglądać następująco:

scaler = torch.amp.GradScaler("cuda", enabled=use_fp16)

for batch in dataloader:
    optimizer.zero_grad(set_to_none=True)

    with torch.autocast(
        device_type="cuda",
        dtype=torch.bfloat16 if use_bf16 else torch.float16,
        enabled=use_bf16 or use_fp16
    ):
        outputs = model(**batch)
        loss = outputs.loss / gradient_accumulation_steps

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Przy BF16 scaler zwykle nie jest potrzebny. W rzeczywistym kodzie trzeba również zadbać, by dane wejściowe, warstwy niestandardowe i wybrane operacje były zgodne z automatyczną precyzją. Niektóre fragmenty obliczeń, zwłaszcza redukcje numerycznie wrażliwe, mogą nadal działać w FP32.

Gradient accumulation i checkpointing aktywacji

Gradient accumulation pozwala zachować większy efektywny batch size bez umieszczania całego batcha w pamięci naraz. Zamiast wykonać jeden krok dla przykładowo 32 sekwencji, można wykonać osiem mikrobatchy po cztery sekwencje, sumując gradienty i aktualizując parametry dopiero po zakończeniu cyklu. Pamięć aktywacji zależy wtedy od mikrobatcha, a nie od całego batcha globalnego.

Trzeba przy tym prawidłowo przeskalować stratę oraz uwzględnić wpływ akumulacji na harmonogram uczenia, logowanie i synchronizację gradientów w treningu rozproszonym. Gradient accumulation nie zmniejsza natomiast pamięci parametrów, gradientów i stanów AdamW. Gdy te elementy są dominującym kosztem, potrzebne staje się sharding.

Gradient checkpointing, nazywany też activation checkpointingiem, redukuje pamięć aktywacji przez celowe nieprzechowywanie części tensorów z forward pass. Podczas backward pass wybrane fragmenty modelu są obliczane ponownie. To świadoma wymiana pamięci na czas obliczeń. Dla głębokich transformerów bywa jedną z najbardziej skutecznych metod ograniczania OOM, zwłaszcza przy długim kontekście.

W bibliotekach Hugging Face aktywacja tej funkcji jest często prosta:

model.gradient_checkpointing_enable()
model.config.use_cache = False

Wyłączenie use_cache ma znaczenie, ponieważ cache kluczy i wartości służący generowaniu autoregresyjnemu nie jest potrzebny podczas standardowego treningu, a potrafi niepotrzebnie zwiększać zużycie VRAM. W czystym PyTorch checkpointing można wdrażać przez torch.utils.checkpoint, najlepiej na granicach logicznych bloków modelu.

Ustawienia danych oraz optymalizatora, które często są pomijane

Padding do jednego wysokiego max_length jest kosztowny, gdy rzeczywiste przykłady mają bardzo różną długość. Dynamic padding dostosowuje długość tensora do najdłuższego przykładu w danym batchu. Jeszcze lepszy efekt może dać grupowanie przykładów o podobnej długości. Ogranicza to liczbę tokenów paddingu, a więc również pamięć aktywacji i czas obliczeń.

Warto przeanalizować, czy wszystkie sekwencje muszą mieć maksymalną długość kontekstu. Dla części zadań redukcja kontekstu, filtrowanie bardzo długich rekordów lub dzielenie dokumentów na sensowne fragmenty daje większą poprawę niż kolejna warstwa optymalizacji infrastruktury. To nie jest wyłącznie kwestia techniczna: skrócenie wejścia może zmienić jakość zadania, dlatego decyzja powinna wynikać z charakteru danych, a nie tylko z dostępnej pamięci GPU.

Adam i AdamW są popularne, ale pamięciożerne. Optymalizatory 8-bitowe mogą obniżyć koszt przechowywania stanów, jednak wprowadzają dodatkową zależność i wymagają testów stabilności, szczególnie przy nietypowych modelach. Przy fine-tuningu metodami PEFT, takimi jak LoRA lub QLoRA, liczba trenowanych parametrów jest znacznie mniejsza, co ogranicza rozmiar gradientów i stanów optymalizatora. Nie usuwa to jednak problemu aktywacji oraz pamięci potrzebnej do przechowywania bazowego modelu. ^[2]

FSDP w PyTorch: shardowanie blisko natywnego ekosystemu

PyTorch Fully Sharded Data Parallel, czyli FSDP, jest mechanizmem treningu rozproszonego zaprojektowanym dla modeli, których pełny stan nie powinien być replikowany na każdym GPU. Stanowi naturalną ewolucję względem klasycznego DistributedDataParallel. Dla zespołów już pracujących bezpośrednio z PyTorch jest to często atrakcyjna droga, ponieważ zmniejsza liczbę zewnętrznych warstw konfiguracji. ^[4]

Jak działa Fully Sharded Data Parallel

W klasycznym DDP każdy proces przechowuje pełną kopię parametrów modelu, pełne gradienty i zazwyczaj pełne stany optymalizatora. DDP dobrze skaluje obliczenia, ale nie rozwiązuje problemu modelu zbyt dużego dla pojedynczego GPU. FSDP dzieli te dane między procesy. Każdy rank przechowuje jedynie fragment pełnego stanu, a nie pełną replikę.

Przed obliczeniem danego modułu FSDP wykonuje komunikację all-gather, aby czasowo zebrać parametry potrzebne do obliczeń. Po propagacji wstecznej wykorzystuje reduce-scatter, który scala gradienty i rozdziela ich fragmenty między procesy. Dzięki temu pamięć trwała potrzebna na parametry, gradienty i stany optymalizatora jest znacząco mniejsza niż w DDP.

Konfiguracja FSDP: wrapping ma znaczenie

W praktyce FSDP nie powinien zwykle owijać całego modelu jako jednego, monolitycznego modułu. Taka konfiguracja może działać, ale ogranicza możliwość nakładania komunikacji na obliczenia i prowadzi do dużych, chwilowych skoków pamięci podczas zbierania parametrów. Znacznie lepszym podejściem jest wrapping logicznych bloków modelu, na przykład pojedynczych warstw transformera.

Dla modelu zbudowanego z powtarzalnych bloków można zastosować politykę automatycznego owijania. W przypadku modeli Hugging Face trzeba wskazać klasę warstwy odpowiadającą blokowi architektury, przykładowo LlamaDecoderLayer, GPTNeoXLayer albo BertLayer. Nazwa zależy od konkretnej implementacji modelu.

Przy planowaniu kolejnych kroków pomocny może być tekst: Frameworki do zarządzania flotą urządzeń IoT w dużych sieciach rozproszonych.

import torch
import torch.distributed as dist
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    ShardingStrategy,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from functools import partial

# Przykład: należy podstawić właściwą klasę bloku dla użytej architektury.
from transformers.models.llama.modeling_llama import LlamaDecoderLayer

dist.init_process_group("nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

auto_wrap_policy = partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={LlamaDecoderLayer},
)

mixed_precision_policy = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.bfloat16,
    buffer_dtype=torch.bfloat16,
)

model = FSDP(
    model,
    auto_wrap_policy=auto_wrap_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    mixed_precision=mixed_precision_policy,
    device_id=torch.cuda.current_device(),
    use_orig_params=True,
    sync_module_states=True,
)

use_orig_params=True jest często wygodne przy współczesnych modelach i narzędziach fine-tuningowych, ponieważ zachowuje bardziej intuicyjny dostęp do oryginalnych parametrów. Ma to znaczenie między innymi dla PEFT, LoRA, niektórych schedulerów oraz kodu, który oczekuje standardowego układu named_parameters(). Nie oznacza to jednak, że każde połączenie FSDP i adapterów zadziała bez dodatkowych testów. Szczególnie ostrożnie należy traktować własne moduły, parametry współdzielone i nietypowe mechanizmy zamrażania warstw.

Opcja sync_module_states=True synchronizuje stan modelu między rankami podczas inicjalizacji. Jest istotna, gdy każdy proces sam ładuje model lub gdy model został utworzony tylko na jednym ranku i trzeba zapewnić spójny start pozostałych procesów. W typowym uruchomieniu wieloprocesowym należy też poprawnie ustawić urządzenie CUDA przed konstrukcją modelu i przed inicjalizacją FSDP.

Strategie shardowania i ich konsekwencje

Najbardziej agresywną, a zarazem najczęściej rozważaną strategią jest FULL_SHARD. Odpowiada ona idei pełnego shardowania parametrów, gradientów i stanów optymalizatora. Zapewnia największą oszczędność pamięci trwałej na pojedynczej karcie, ale wymaga częstej komunikacji podczas wykonywania modelu.

SHARD_GRAD_OP pozostawia parametry w większym stopniu dostępne po forward pass, a sharduje przede wszystkim gradienty i stan optymalizatora. Może być rozsądnym kompromisem, gdy model prawie mieści się w pamięci, ale pełne shardowanie powoduje zbyt duży narzut komunikacyjny. Nie daje jednak tak dużej redukcji VRAM jak FULL_SHARD. ^[1]

W większych klastrach ważna staje się także topologia sprzętu. Jeśli serwer ma osiem GPU połączonych szybkim NVLink, a wiele serwerów komunikuje się przez wolniejszą sieć, można rozważyć HYBRID_SHARD. W takim układzie sharding odbywa się wewnątrz grupy GPU w jednym węźle, natomiast między węzłami utrzymywana jest replikacja. Ogranicza to koszt komunikacji przez sieć międzyserwerową, ale zwiększa pamięć potrzebną na każdej maszynie.

Nie istnieje jedna strategia optymalna dla wszystkich przypadków. Model, który działa szybko na ośmiu GPU w jednym serwerze, może zachowywać się zupełnie inaczej po rozproszeniu na kilka maszyn. Wtedy ograniczeniem przestaje być wyłącznie VRAM, a staje się przepustowość i opóźnienie sieci.

FSDP a momenty szczytowego użycia pamięci

FSDP redukuje pamięć trwałą, ale nie usuwa wszystkich pików zużycia VRAM. Parametry potrzebne do wykonania konkretnego bloku są chwilowo zbierane przez all-gather. Jeżeli pojedynczy owinięty moduł jest bardzo duży, jego pełna postać może nadal nie zmieścić się na GPU. To jedna z najważniejszych przyczyn sytuacji, w której konfiguracja teoretycznie używa pełnego shardowania, a mimo to kończy się OOM.

Zbyt gruby wrapping oznacza duże jednorazowe zbiory parametrów. Zbyt drobny wrapping oznacza natomiast większą liczbę operacji komunikacyjnych, większy narzut sterowania i potencjalnie słabsze wykorzystanie GPU. Granica na poziomie jednego bloku transformera jest często dobrym punktem startowym, ale przy bardzo dużych warstwach MoE, nietypowych modułach wizji albo rozbudowanych blokach multimodalnych może wymagać dalszego podziału.

W konfiguracjach FSDP warto analizować nie tylko końcową wartość pamięci po backward pass, lecz także maksimum podczas pojedynczego kroku. W szczególności należy sprawdzić, czy OOM pojawia się podczas forward, backward, aktualizacji optymalizatora czy zapisu checkpointu. Każdy z tych momentów może wskazywać na inny problem.

Checkpointy FSDP: pełny model nie zawsze mieści się na jednym ranku

Zapisywanie checkpointów jest jednym z miejsc, w których dobrze działający trening rozproszony potrafi nagle zużyć zbyt dużo pamięci CPU lub GPU. Jeśli każdy rank spróbuje odtworzyć pełny state_dict modelu, korzyść z shardowania może zostać chwilowo utracona.

FSDP obsługuje checkpointy shardowane, w których każdy proces zapisuje fragment stanu. Taki format jest praktyczny przy wznawianiu treningu na podobnej liczbie GPU i pozwala uniknąć centralizowania całego modelu w jednym procesie. Pełny checkpoint jest z kolei wygodniejszy do późniejszego inferowania, publikowania modelu lub ładowania go w środowisku bez FSDP.

Warto od początku zdecydować, który wariant jest potrzebny operacyjnie. Checkpoint do awaryjnego wznowienia treningu nie musi być tym samym artefaktem co checkpoint przeznaczony do wdrożenia. W projektach produkcyjnych często zapisuje się shardowany stan treningowy okresowo, a pełny model znacznie rzadziej, na przykład po zakończeniu etapu uczenia.

Przy zapisie pełnego modelu trzeba również kontrolować, który rank wykonuje operację wejścia-wyjścia. W przeciwnym razie kilka procesów może próbować jednocześnie zapisać ten sam plik albo niepotrzebnie gromadzić pełen stan w pamięci.

Gdzie FSDP wymaga najwię მეტი cierpliwości

FSDP jest blisko PyTorch, ale nie oznacza to, że znika złożoność treningu rozproszonego. Kod musi poprawnie inicjalizować procesy, przypisywać ranki do urządzeń, synchronizować dane i obsługiwać checkpointy. W przypadku Hugging Face Trainer część integracji jest dostępna z poziomu argumentów uruchomieniowych, jednak zaawansowane scenariusze nadal wymagają zrozumienia, co dzieje się pod spodem.

Problematyczne bywają parametry współdzielone między modułami, warstwy używane warunkowo oraz modele, których nie wszystkie gałęzie są wykonywane w każdym kroku. Trzeba też przetestować generowanie tekstu i ewaluację, ponieważ tryb inferencyjny może wymagać innego sposobu zbierania parametrów niż trening.

Warto traktować FSDP jako narzędzie do systematycznego skalowania, a nie jako pojedynczą flagę naprawiającą każdy OOM. Jeśli podstawowym problemem są aktywacje dla bardzo długiego kontekstu, nadal kluczowe będą długość sekwencji, checkpointing aktywacji i wielkość mikrobatcha. Jeśli problemem są parametry oraz stany optymalizatora, FSDP daje znacznie bardziej bezpośredni efekt.

DeepSpeed i ZeRO: szeroki zestaw narzędzi do treningu na granicy pamięci

DeepSpeed jest ekosystemem treningowym rozwijanym wokół idei wydajnego uczenia dużych modeli. Jego najbardziej rozpoznawalnym elementem jest ZeRO, czyli Zero Redundancy Optimizer. Podobnie jak FSDP, ZeRO eliminuje zbędne replikowanie stanu między procesami, ale DeepSpeed oferuje również rozbudowane mechanizmy offloadu, konfiguracje hybrydowe, optymalizacje komunikacji oraz narzędzia przeznaczone dla bardzo dużych modeli językowych.

W praktyce DeepSpeed jest często wybierany wtedy, gdy projekt przekracza wygodny zakres pojedynczego serwera GPU albo gdy potrzebne są agresywne techniki oszczędzania pamięci kosztem większej złożoności konfiguracji i możliwego spadku szybkości.

ZeRO-1, ZeRO-2 i ZeRO-3

Poziomy ZeRO różnią się tym, które elementy stanu treningowego są dzielone między procesy.

• ZeRO-1 sharduje stany optymalizatora. Parametry i gradienty pozostają replikowane.
• ZeRO-2 sharduje stany optymalizatora oraz gradienty, ale parametry nadal są replikowane na każdym GPU.
• ZeRO-3 sharduje również parametry modelu. Jest najbliższy pełnemu shardowaniu stosowanemu przez FSDP i daje największą oszczędność pamięci.

ZeRO-1 bywa dobrym rozwiązaniem, gdy głównym kosztem są stany AdamW, a model nadal mieści się na pojedynczej karcie. ZeRO-2 redukuje dodatkowo koszt gradientów. ZeRO-3 ma największy potencjał dla modeli zbyt dużych dla jednego GPU, ale wprowadza najwięcej komunikacji i wymaga staranniejszego strojenia.

Wybór poziomu nie powinien wynikać wyłącznie z zasady „wyższy jest lepszy”. Jeżeli ZeRO-2 rozwiązuje problem pamięci, przejście do ZeRO-3 może niepotrzebnie zwiększyć koszt komunikacji. Najbardziej agresywna konfiguracja ma sens wtedy, gdy bez niej trening po prostu nie mieści się w dostępnej pamięci.

Minimalna konfiguracja DeepSpeed

DeepSpeed najczęściej konfiguruje się przez plik JSON. Poniższy przykład przedstawia konfigurację ZeRO-2 z BF16. Jest to punkt startowy, a nie uniwersalny profil dla każdego modelu.

{
  "train_micro_batch_size_per_gpu": 2,
  "gradient_accumulation_steps": 8,
  "bf16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_scatter": true
  },
  "gradient_clipping": 1.0,
  "steps_per_print": 100
}

Dla ZeRO-3 konfiguracja może zostać rozszerzona o parametry odnoszące się do komunikacji i rozmiaru fragmentów zbieranych w pamięci:

{
  "train_micro_batch_size_per_gpu": 1,
  "gradient_accumulation_steps": 16,
  "bf16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_scatter": true,
    "allgather_partitions": true,
    "allgather_bucket_size": 200000000,
    "reduce_bucket_size": 200000000
  }
}

Wartości allgather_bucket_size i reduce_bucket_size nie są magicznymi liczbami. Większe bufory mogą poprawić efektywność komunikacji, ale same zajmują pamięć. Przy kartach z ograniczonym VRAM zbyt duże buckety mogą wywołać OOM, mimo że sam model teoretycznie powinien się zmieścić. Strojenie należy zaczynać od bezpiecznych wartości i mierzyć zarówno szczytowe użycie pamięci, jak i czas kroku treningowego.

Źródła

• [1] PyTorch – FullyShardedDataParallel — PyTorch 2.14 documentation ↑
• [2] DeepSpeed – DeepSpeed documentation — ZeRO and ZeRO-Infinity ↑
• [3] PyTorch Tutorials – FSDP Advanced Tutorial ↑
• [4] PyTorch – PyTorch Distributed Overview ↑