Modificare uno script PyTorch di addestramento

Se utilizzi una tecnica di ottimizzazione che si basa su norme riguardanti i gradienti globali, ad esempio una norma sui gradienti di tutto il modello, come alcune varianti dell'ottimizzatore LAMB o il ritaglio del gradiente globale, è necessario raccogliere tutte le norme delle 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 CUDA non necessari nel dispositivo 0, che consumerebbe inutilmente memoria.

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 nella GPU (ad esempio, utilizzando model.to(device)) quando si utilizza la libreria. Se provi a spostare il modello nella GPU prima che il modello sia partizionato (prima della prima chiamata smp.step), la chiamata di spostamento viene ignorata. La libreria sposta automaticamente la parte del modello assegnato a una classificazione nella sua GPU. Una volta iniziato l'addestramento con la libreria, non spostare il modello nella CPU e usalo, poiché non avrà i parametri corretti per i moduli non assegnati alla partizione tenuta dal processo. Se desideri riqualificare un modello o utilizzarlo per l'inferenza senza la libreria dopo che è stato addestrato utilizzando la libreria di parallelismo dei modelli, il modo consigliato è salvare il modello completo utilizzando la nostra API di checkpointing 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 per smp.step devono essere spostati sul dispositivo attuale utilizzando l'API .to(), 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 attuale utilizzando l'API .to() prima di utilizzare i tensori di input per 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 assegnato a una classificazione nella sua 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 si utilizza il parallelismo dei dati con il PyTorch DDP nativo, 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 proprio modello ha buffer di moduli che devono essere sincronizzati tra gruppi paralleli di dati in ogni fase, è possibile farlo tramite l'API torch.distributed, 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. Puoi invece utilizzare le relative implementazioni della libreria tramite le API smp.optimizers e smp.nn.