Modificare uno script di addestramento PyTorch

Se utilizzate una tecnica di ottimizzazione che si basa su norme globali di gradiente, ad esempio la norma del gradiente dell'intero modello, come alcune varianti di LAMB optimizer o di global gradient clipping, dovete raccogliere tutte le norme sulle partizioni del modello per verificarne la correttezza. A tale scopo è possibile utilizzare i tipi di dati di base per la comunicazione della libreria.

Tutti gli argomenti torch.Tensor relativi ai metodi forward di nn.Modules nel modello devono essere utilizzati nel calcolo dell'output del modulo. In altre parole, la libreria non supporta il caso in cui esista un argomento torch.Tensor per un modulo da cui l'output del modulo non dipende.

L'argomento della chiamata smp.DistributedModel.backward() deve dipendere da tutti gli output del modello. In altre parole, non può esserci un output della chiamata smp.DistributedModel.forward che non venga utilizzato nel calcolo del tensore utilizzato per la chiamata smp.DistributedModel.backward.

Se nel codice sono presenti chiamate torch.cuda.synchronize(), potrebbe essere necessario chiamare torch.cuda.set_device(smp.local_rank()) immediatamente prima della chiamata di sincronizzazione. Altrimenti potrebbero essere creati contesti non necessari nel dispositivo 0, che consumerà inutilmente memoriaCUDA.

Poiché la libreria colloca nn.Modules su dispositivi diversi, i moduli nel modello non devono dipendere da alcuno stato globale modificato all'interno smp.step. È consentito qualsiasi stato che rimanga fisso durante l'addestramento o che venga modificato all'esterno di smp.step in modo visibile a tutti i processi.

Non è necessario spostare il modello in GPU (ad esempio, utilizzandolomodel.to(device)) quando si utilizza la libreria. Se provate a spostare il modello GPU prima del partizionamento del modello (prima della prima smp.step chiamata), la chiamata move viene ignorata. La libreria sposta automaticamente la parte del modello assegnata a un rango al suo. GPU Una volta iniziato l'addestramento con la libreria, non spostate il modello CPU e non usatelo, poiché non avrà i parametri corretti per i moduli non assegnati alla partizione contenuta dal processo. Se desideri riaddestrare un modello o usarlo per l'inferenza senza la libreria dopo che è stato addestrato utilizzando la libreria di parallelismo dei modelli, il modo consigliato è salvare il modello completo usando il nostro checkpoint API e caricarlo di nuovo su un normale Module. PyTorch

Se disponi di un elenco di moduli tale per cui l'output di un modulo ne alimenta un altro, la sostituzione di tale elenco con nn.Sequential può migliorare significativamente le prestazioni.

L'aggiornamento del peso (optimizer.step()) deve avvenire all'esterno di smp.step perché è allora che l'intero passaggio indietro è terminato e i gradienti sono pronti. Quando si utilizza un modello ibrido con parallelismo di modelli e dati, a questo punto è garantito anche il completamento dei gradienti AllReduce .

Quando usi la libreria in combinazione con il parallelismo dei dati, assicurati che il numero di batch su tutti i ranghi paralleli dei dati sia lo stesso in modo da AllReduce non bloccare l'attesa di un rank che non partecipa al passaggio.

Se avvii un processo di addestramento utilizzando un tipo di istanza ml.p4d (ad esempio ml.p4d.24xlarge), è necessario impostare la variabile del caricatore di dati num_workers=0. Ad esempio, puoi definire il tuo DataLoader come di seguito indicato:

dataloader = torch.utils.data.DataLoader(
            data,
            batch_size=batch_size,
            num_workers=0,
            pin_memory=True,
            drop_last=True,
            shuffle=shuffle,
        )

Gli input per smp.step devono essere gli input del modello generati da DataLoader. Ciò è necessario perché smp.step suddivide internamente i tensori di input lungo la dimensione del batch e li inserisce in pipeline. Di conseguenza, passare DataLoader alla funzione smp.step per generare gli input del modello all'interno non funziona.

Ad esempio, se definisci un DataLoader come di seguito indicato:

train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True)

Dovresti accedere agli input del modello generati da train_loader e passarli a una funzione decorata smp.step. Non passare train_loader direttamente alla funzione smp.step.

def train(model, device, train_loader, optimizer):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        ...
        _, loss_mb = train_step(model, data, target)
        ...

@smp.step
def train_step(model, data, target):
    ...
    return output, loss

I tensori di input da smp.step devono essere spostati sul dispositivo corrente in uso .to()API, operazione che deve avvenire dopo la chiamata. torch.cuda.set_device(local_rank())

Ad esempio, puoi definire la funzione train nel modo seguente. Questa funzione aggiunge data e target al dispositivo corrente utilizzando .to() API prima di utilizzare quei tensori di input da chiamare. train_step

def train(model, device, train_loader, optimizer):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # smdistributed: Move input tensors to the GPU ID used by the current process,
        # based on the set_device call.
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        # Return value, loss_mb is a StepOutput object
        _, loss_mb = train_step(model, data, target)

        # smdistributed: Average the loss across microbatches.
        loss = loss_mb.reduce_mean()

        optimizer.step()

I tensori di input per questa funzione decorata smp.set sono stati spostati nel dispositivo attuale nella funzione train sopra riportata. Non è necessario spostare il modello nel dispositivo corrente. La libreria sposta automaticamente la parte del modello assegnata a un rango al suo. GPU

@smp.step
def train_step(model, data, target):
    output = model(data)
    loss = F.nll_loss(output, target, reduction="mean")
    model.backward(loss)
    return output, loss

Se utilizzi il parallelismo dei dati con quello nativo PyTorch DDP, il modulo torch.nn.parallel.DistributedDataParallelwrapper non è supportato dalla libreria. La libreria gestisce internamente l'integrazione con PyTorch DDP, inclusi i parametri broadcast e gradient. AllReduce Quando si utilizza la libreria, i buffer dei moduli vengono trasmessi una sola volta all'inizio dell'addestramento. Se il tuo modello ha buffer di moduli che devono essere sincronizzati tra gruppi paralleli di dati in ogni fase, puoi farlo tramite torch.distributedAPI, utilizzando il gruppo di processi che può essere ottenuto tramite. smp.get_dp_process_group()

Per l'addestramento di precisione misto, il modulo apex.amp non è supportato. Per utilizzare la libreria con precisione mista automatica si consiglia di utilizzare torch.cuda.amp, utilizzando tuttavia smp.amp.GradScaler anziché l'implementazione in torch.

torch.jit.ScriptModules e ScriptFunctions non sono supportati da smp.DistributedModel.

apex: FusedLayerNorm, FusedAdam, FusedLAMB e FusedNovoGrad da apex sono supportati. È possibile utilizzare le relative implementazioni della libreria tramite smp.optimizers e invece. smp.nn APIs