Come funziona il Parallelismo tensoriale

Il parallelismo tensoriale avviene a livello di nn.Modules; partiziona moduli specifici nel modello tra classificazioni tensoriali parallele. Ciò si aggiunge alla partizione esistente del set di moduli utilizzati nel parallelismo delle pipeline.

Quando un modulo viene partizionato tramite il parallelismo tensoriale, la sua propagazione in avanti e all'indietro viene distribuita. La libreria gestisce la comunicazione necessaria tra i dispositivi per implementare l'esecuzione distribuita di questi moduli. I moduli sono partizionati su più classificazioni di dati parallele. Contrariamente alla distribuzione tradizionale dei carichi di lavoro, ogni classificazione di parallelismo dei dati non ha la replica completa del modello quando viene utilizzato il parallelismo tensoriale della libreria. Invece, ogni classificazione di parallelismo dei dati può avere solo una partizione dei moduli distribuiti, oltre alla totalità dei moduli che non sono distribuiti.

Esempio: considera il parallelismo tensoriale tra le classificazioni di dati parallele, dove il grado di parallelismo dei dati è 4 e il grado di parallelismo tensoriale è 2. Supponiamo di avere un gruppo parallelo di dati che contiene il seguente albero di moduli, dopo aver partizionato il set di moduli.

A
├── B
|   ├── E
|   ├── F
├── C
└── D
    ├── G
    └── H

Supponiamo che il parallelismo tensoriale sia supportato per i moduli B, G e H. Un possibile risultato della partizione parallela tensoriale di questo modello potrebbe essere:

dp_rank 0 (tensor parallel rank 0): A, B:0, C, D, G:0, H
dp_rank 1 (tensor parallel rank 1): A, B:1, C, D, G:1, H
dp_rank 2 (tensor parallel rank 0): A, B:0, C, D, G:0, H
dp_rank 3 (tensor parallel rank 1): A, B:1, C, D, G:1, H

Ogni riga rappresenta l'insieme di moduli memorizzati in dp_rank, e la notazione X:y rappresenta la frazione y del modulo X. Tieni presente quanto segue:

1. Il partizionamento avviene tra sottoinsiemi di classificazioni di dati paralleli, che chiamiamo TP_GROUP, non l'intero DP_GROUP, in modo che la partizione esatta del modello venga replicata su dp_rank 0 e dp_rank 2 e analogamente su dp_rank 1 e dp_rank 3.

2. I moduli E e F non fanno più parte del modello, poiché il loro modulo principale B è partizionato e qualsiasi esecuzione che normalmente fa parte di E e F avviene all'interno del modulo (partizionato) B.

3. Anche se H è supportato per il parallelismo tensoriale, in questo esempio non è partizionato, il che evidenzia che il partizionamento di un modulo dipende dall'input dell'utente. Il fatto che un modulo sia supportato per il parallelismo tensoriale non significa necessariamente che sia partizionato.

Come la libreria adatta il PyTorch nn.Linear parallelismo tensoriale al modulo

Quando il parallelismo tensoriale viene eseguito su classificazioni di dati parallele, un sottoinsieme dei parametri, dei gradienti e degli stati dell'ottimizzatore viene partizionato tra i dispositivi paralleli tensoriali per i moduli partizionati. Per il resto dei moduli, i dispositivi paralleli tensoriali funzionano in modo normale in parallelo ai dati. Per eseguire il modulo partizionato, un dispositivo raccoglie innanzitutto le parti necessarie di tutti i campioni di dati su dispositivi peer nello stesso gruppo di parallelismo tensoriale. Il dispositivo esegue quindi la frazione locale del modulo su tutti questi campioni di dati, seguita da un altro ciclo di sincronizzazione che combina le parti dell'output per ciascun campione di dati e restituisce i campioni di dati combinati al punto GPUs da cui ha avuto origine il campione di dati. La figura seguente mostra un esempio di questo processo su un modulo partizionato nn.Linear.

La prima figura mostra un modello piccolo con un modulo nn.Linear grande con parallelismo dei dati sulle due classificazioni di parallelismo tensoriale. Il modulo nn.Linear viene replicato nelle due classificazioni parallele.

La seconda figura mostra il parallelismo tensoriale applicato su un modello più grande durante la suddivisione del modulo nn.Linear. Ciascun tp_rank contiene metà del modulo lineare e la totalità del resto delle operazioni. Durante il funzionamento del modulo lineare, ogni tp_rank raccoglie la metà pertinente di tutti i campioni di dati e la trasmette attraverso la propria metà del modulo nn.Linear. Il risultato deve essere ridotto (con la somma come operazione di riduzione) in modo che ogni classificazione abbia l'output lineare finale dei propri campioni di dati. Il resto del modello viene eseguito nel tipico modo parallelo dei dati.