mcarilli:

First try it without amp, which gives an idea how strongly dataloader bound you are overall. For example, if switching to dummy data gives a 2X speedup right away, you know the dataloader is a big bottleneck.
Then try it with amp, and see if you observe a speedup relative to the above dummy data+no amp case, which gives an idea if Amp is working.


Here are the results of the test. I tested the native speed of my model, then the speed with dummy data and then dummy data and Amp:

testing ddp
testing: model=mlm batch_size=4 num_workers=0 mixed_pres=False pin_memory=False use_dummy_data=False
Training complete in: 0:03:47.287354
model=mlm batch_size=4 num_workers=0 mixed_pres=False pin_memory=False use_dummy_data=True
Training complete in: 0:02:40.539612
testing: model=mlm batch_size=4 num_workers=0 mixed_pres=True pin_memory=False use_dummy_data=True
Training complete in: 0:02:41.635868

so while my dataloader seems to slow down the training quite a bit I still see no speedup using Amp.
I very lazily updated my training loop to:
    for epoch in range(epochs):
        if not use_dummy_data:
            for i, inputs in enumerate(train_loader):
                for k, v in inputs.items():
                    if isinstance(v, torch.Tensor):
                        inputs[k] = v.cuda(rank, non_blocking=True)
                optimizer.zero_grad()
                if mixed_pres:
                    with torch.cuda.amp.autocast():
                        outputs = model(**inputs)
                        loss = outputs[0]
                        loss = loss.mean()
                    # Backward and optimize
                    scaler.scale(loss).backward()
                    scaler.step(optimizer)
                    scaler.update()
                else:
                    outputs = model(**inputs)
                    loss = outputs[0]
                    loss = loss.mean()
                    loss.backward()
                    optimizer.step()
                if i >= n_steps:
                    break
        else:
            inputs = next(enumerate(train_loader))[1]
            inputs.to(f"cuda:{rank}")
            for i in range(n_steps):
                optimizer.zero_grad()
                if mixed_pres:
                    with torch.cuda.amp.autocast():
                        outputs = model(**inputs)
                        loss = outputs[0]
                        loss = loss.mean()
                    # Backward and optimize
                    scaler.scale(loss).backward()
                    scaler.step(optimizer)
                    scaler.update()
                else:
                    outputs = model(**inputs)
                    loss = outputs[0]
                    loss = loss.mean()
                    loss.backward()
                    optimizer.step()