În paralelizarea tradițională de tip DDP (Distributed Data Parallel), fiecare placă video reține o copie completă a modelului. Când modelul depășește 10-20 de miliarde de parametri, acesta nu mai încape în memoria unui singur GPU, făcând DDP inutilizabil.
FSDP (Fully Sharded Data Parallel) implementează algoritmul Zero Redundancy Optimizer (ZeRO-3) nativ în nucleul PyTorch. În loc să dubleze datele, FSDP fragmentează (shards) cei trei mari consumatori de memorie pe întregul cluster de acceleratoare: stările optimizatorului AdamW (75% din memorie), gradienții și parametrii modelului.
În timpul trecerii directe (forward pass) și a celei inverse (backward pass), FSDP adună dinamic ponderile necesare pentru fiecare strat neuronal prin comunicații ultra-rapide `all-gather`, calculează operația și eliberează imediat memoria, permițând antrenarea modelelor de 70B+ pe clustere accesibile fără librării terțe greoaie.
Sintaxă, opțiuni și concepte cheie
| Opțiune / Concept | Descriere practică |
|---|---|
FullyShardedDataParallel |
Modulul wrapper din PyTorch ce încapsulează straturile rețelei neuronale. |
ShardingStrategy.FULL_SHARD |
Strategia optimă ce fragmentează parametrii, gradienții și optimizatorul (ZeRO-3). |
auto_wrap_policy |
Regula automată de împachetare pe straturi (de exemplu pe fiecare bloc `TransformerBlock`). |
backward_prefetch |
Prefetching-ul ponderilor stratului anterior pe durata calculului curent pentru a ascunde latența de rețea. |
CPU Offload |
Opțiunea de a muta temporar stările optimizatorului în RAM-ul sistemului pe configurații cu GPU-uri mici. |
Exemple practice și configurații
Exemplul 1: Împachetarea unui model de limbaj cu FSDP și politică de transformator în PyTorch
Configurarea completă pentru antrenare multi-GPU distribuită.
import torch
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
import functools
# Inițializăm procesul distribuit (NCCL pentru GPU-uri Nvidia)
dist.init_process_group(backend="nccl")
torch.cuda.set_device(dist.get_rank())
model = ModelTransformerMare().to(dist.get_rank())
politica_wrap = functools.partial(size_based_auto_wrap_policy, min_num_params=1e6)
# Împachetăm modelul cu FSDP:
fsdp_model = FSDP(
model,
auto_wrap_policy=politica_wrap,
sharding_strategy=ShardingStrategy.FULL_SHARD,
device_id=torch.cuda.current_device()
)
optimizator = torch.optim.AdamW(fsdp_model.parameters(), lr=1e-5)Rezultat & Impact
Model scalat pe 8 GPU-uri cu zero duplicare redundantă a memoriei.Greșeli frecvente și bune practici
- Omiterea configurării `auto_wrap_policy`, ceea ce face ca FSDP să împacheteze întregul model ca un singur bloc gigant, pierzând orice economie de memorie.
- Utilizarea FSDP pe rețele fără suport pentru comunicație de mare viteză (cum ar fi lipsa NVLink între GPU-uri locale), unde traficul `all-gather` devine un blocaj masiv.