Huggingface Distributed Fine-tuning with Accelerate

Nidhi_Rustagi: <p>I am trying to Finetune BertForSequenceClassification on custom dataset with Dsitributed training using Accelerate. Runtime environment is AWS Sagemaker notebook with 4 GPUs and 24 GB RAM.<br>
Executed my code earlier for one step to ensure batch size does not cause memory issues and that I get train/val errors.<br>
But notebook launcher now fails to execute the code, the cell simply completes execution without executing a single step of training.</p>
<pre><code class="lang-auto">def training_loop():
    model_path = "models/category_classifier_bert.pt"    
    accelerator = Accelerator(mixed_precision="fp16")
    device = accelerator.device
    
    model = BertForSequenceClassification.from_pretrained("bert-base-uncased",num_labels=NUM_LABELS)
    print("Loading model ")
    learning_rate = 1e-5
    optimizer = AdamW(model.parameters(), lr=learning_rate)
    print("Loading optimizer")
    train_dl, val_dl = get_dataloaders()
    total_steps = int((len(train_dl) * EPOCHS)/4*BATCH_SIZE)
    print("total_steps : ",total_steps)
    # Set up the learning rate scheduler
    scheduler = get_linear_schedule_with_warmup(optimizer,
                                                num_warmup_steps=100, # Default value
                                                num_training_steps=total_steps)
    print("Loading model on accelerators .... ")
    model, optimizer, train_dl, val_dl, scheduler = accelerator.prepare(model, optimizer, train_dl, val_dl, scheduler)
    #val_dl = accelerator.prepare(val_dl)
    
    metric = Accuracy(task="multiclass", num_classes=NUM_LABELS).to(device)
    # MulticlassAccuracy(num_classes=NUM_LABELS).to(device)
    print("Training ... ")
    for epoch in range(EPOCHS):
        model.train()
        for step,batch in enumerate(train_dl):
            input_ids = batch["input_ids"].to(device)
            targets = batch["label"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            optimizer.zero_grad()

            output = model(input_ids, attention_mask)
            loss = CrossEntropyLoss(output.logits, targets)

            accelerator.backward(loss)
            optimizer.step()
            
            if step%5000 == 0:
                print(step, loss)               
                accelerator.wait_for_everyone()
                unwrapped_model = accelerator.unwrap_model(model)
                accelerator.save(unwrapped_model.state_dict(), model_path)
            
from accelerate import notebook_launcher
notebook_launcher(training_loop, num_processes=4)

</code></pre>
<p>The cell completes execution with out performing a single step of training. Are there any cache or buffer issues that I needto take care of ? From the output, it appears that NOT all the GPUs are identified and hence the cell execution terminates before completing execution.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/f/2fccbc14029b14686b707706bca4173ebc1697c1.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/2fccbc14029b14686b707706bca4173ebc1697c1" title="output"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/f/2fccbc14029b14686b707706bca4173ebc1697c1_2_568x499.jpeg" alt="output" data-base62-sha1="6OR59ALcQ087SrU8mJQoOygufpD" width="568" height="499" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/f/2fccbc14029b14686b707706bca4173ebc1697c1_2_568x499.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/2/f/2fccbc14029b14686b707706bca4173ebc1697c1_2_852x748.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/f/2fccbc14029b14686b707706bca4173ebc1697c1.jpeg 2x" data-dominant-color="E9D3D4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">output</span><span class="informations">1110×976 428 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>