În trecut, salvarea unui checkpoint de antrenare pe un cluster distribuit necesita adunarea tuturor ponderilor de pe toate GPU-urile pe un singur nod ‘Rank 0’ (prin `torch.save`), proces care bloca tot antrenamentul timp de 20-30 de minute și declanșa deseori erori de depășire de memorie CPU.
Distributed Checkpointing (DCP) a revoluționat salvarea modelelor distribuite. Fiecare proces GPU își scrie propria bucată de tensori fragmentați în paralel direct pe sistemul de fișiere partajat (sau stocare obiect), fără a comunica cu celelalte plăci.
O proprietate remarcabilă a DCP este Resharding on Load: puteți salva un model pe un cluster cu 16 GPU-uri și îl puteți încărca direct pe un cluster cu 4 sau 8 GPU-uri, PyTorch redistribuind automat feliile de memorie fără nicio intervenție manuală.
Sintaxă, opțiuni și concepte cheie
| Opțiune / Concept | Descriere practică |
|---|---|
dcp.save(state_dict, storage_writer) |
Operația de scriere distribuită paralelă pe disc. |
dcp.load(state_dict, storage_reader) |
Încărcarea automată cu adaptare la topologia GPU curentă. |
FileSystemWriter / FileSystemReader |
Backend-ul optimizat de scriere pe directoare partajate NFS / Lustre. |
Asynchronous Checkpointing |
Salvarea tensorilor în fundal pe fire secundare permițând continuarea imediată a antrenamentului. |
Exemple practice și configurații
Exemplul 1: Salvarea și încărcarea distribuită a unui model FSDP cu DCP
Sintaxa modernă introdusă în PyTorch 2.1+.
import torch.distributed.checkpoint as dcp
# Obținem starea modelului și a optimizatorului fără consolidare pe Rank 0:
stare = {"model": fsdp_model.state_dict(), "optimizator": optimizator.state_dict()}
# Salvare paralelă pe disc:
dcp.save(stare, checkpoint_id="/storage/checkpoints/pas_5000")
print(f"Rank {dist.get_rank()}: bucată de checkpoint salvată cu succes!")
# La reluare (chiar și cu un număr diferit de GPU-uri):
stare_goala = {"model": fsdp_model.state_dict(), "optimizator": optimizator.state_dict()}
dcp.load(stare_goala, checkpoint_id="/storage/checkpoints/pas_5000")
fsdp_model.load_state_dict(stare_goala["model"])Rezultat & Impact
Salvare finalizată în câteva secunde fără întreruperea fluxului masiv de calcul.Greșeli frecvente și bune practici
- Modificarea manuală a numelor cheilor din dicționarul de stare între salvare și încărcare, rupând maparea automată de resharding.
- Scrierea checkpoint-urilor pe sisteme de fișiere locale în loc de volume partajate între toate nodurile din cluster.