Apex's distributed dataparallael setting is not working as intended

username: jaewoo_song
cooked: <p>Hi, I’m currently working on encoder-decoder transformer model pre-training and trying to train it with apex’s distributed dataparallel settings.<br>
As I know, it executes multi-processing and each process is assigned to one GPU.<br>
And as I searched, this distributed training certainly has an advantage that it reduces memory imbalance problem than other methods do.</p>
<p>But in my training, GPU’s memory allocations are different and also GPU-utils are so low, which indicates that each GPU does not work fully.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/e/ce6f6b83018db613b177eb4253fc80e14ebdf87a.png" data-download-href="https://discuss.pytorch.org/uploads/default/ce6f6b83018db613b177eb4253fc80e14ebdf87a" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/c/e/ce6f6b83018db613b177eb4253fc80e14ebdf87a.png" alt="image" data-base62-sha1="tsd4EEDg5cGTVEedaUas8pbG7JU" width="366" height="499" data-dominant-color="191919"><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">image</span><span class="informations">637×868 22.9 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>
<ol>
<li>Is there any reason why each GPU’s memory is used differently? Is there I misunderstood?</li>
<li>Why is GPU’s utilization so low? Is there anything I missed?</li>
<li>One more thing, before I tried with full data, I tested the smaller set with same batch sizes. In that setting, the training was completed with no problems. But with full data, I had to decrease the batch size quite a lot since it raised CUDA out-of-memory error. I don’t know why this happened because I think at least the size of one batch put into each GPU is same in each iteration.</li>
</ol>
<p>Here is the training code only with essential parts.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch.distributed as dist
from apex.parallel import DistributedDataParallel as DDP
from apex import amp
from torch.utils.data.distributed import DistributedSampler


torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True

torch.cuda.set_device(args.local_rank)
device = torch.device(f'cuda:{args.local_rank}')
dist.init_process_group(backend='nccl', init_method='env://', rank=args.local_rank, world_size=args.nproc_per_node)

data = PretrainDataset(path, data)
ppd = PretrainPadCollate(pad_id)
sampler = DistributedSampler(data)
data_loader = DataLoader(
    data,
    collate_fn=ppd.pad_collate, 
    batch_size=args.batch, 
    shuffle=True if sampler is None else False, 
    pin_memory=True, 
    num_workers=4,
    sampler=sampler
)

model = ModelClass().to(device)
opt = Adafactor(model.parameters(), lr=args.lr, relative_step=False)
scheduler = lr_scheduler.ExponentialLR(opt, gamma=1/math.sqrt(args.warmup_steps))

model, opt = amp.initialize(model, opt, opt_level='O1')
model = DDP(model, delay_allreduce=True)

# Train iterations
train_step = 1
for batch in data_loader:
    opt.zero_grad()
    outputs = model(batch)
    
    loss = outputs[0]
    with amp.scale_loss(loss, opt) as scaled_loss:
        scaled_loss.backward()

    opt.step()

    if train_step &gt; args.warmup_steps:
        scheduler.gamma = 1/math.sqrt(train_step)
        scheduler.step()
</code></pre>
<p>Thank you.</p>

username: ptrblck
cooked: <p>We recommend to use the native DDP implementation instead of <code>apex/DDP</code>, as it’ll be further developed and improved. Could you switch and rerun your code?<br>
Also, use <code>torch.cuda.amp</code> instead of <code>apex/amp</code> for the same reason.</p>