Slow backward prop with Apex

Hi there,

I am using Apex with DALI to boost up resnet50 training on imagenet. My env is as below:
pytorch: 1.4.0
cuda: 10.1
apex: 0.1
dali: 0.20.0

I am use a single V100 GPU. With apex, the speed is almost double, but still far from over 1000 images/second, merely at 250. I have checked with DALI guys and we found that the backward prop is suspiciously slow, 10x less than forward prop. I have printed out the time for each individual steps in training loop, with --opt-level to O1

forward time: 0.017492055892944336
loss time: 0.0001506805419921875
update loss and acc time: 0.05937790870666504
backward prop time: 0.16233062744140625

Using fp32 by --opt-level to O0, the time is as below:

forward time: 0.012087583541870117
loss time: 0.00022459030151367188
update loss and acc time: 0.09990501403808594
backward prop time: 0.06609177589416504

It seems, time for backward prop with apex is almost 3x slowers than w/o it.

The code for training is as below:

for i, data in enumerate(train_loader):
        input = data[0]["data"].cuda(non_blocking=True)
        target = data[0]["label"].squeeze().cuda(non_blocking=True).long()
        train_loader_len = int(math.ceil(train_loader._size / args.batch_size))

        adjust_learning_rate(optimizer, epoch, i, train_loader_len)

        if args.prof:
            if i > 10:
                break
        # measure data loading time
        data_time.update(time.time() - end)

        temp_time = time.time()
        # compute output
        output = model(input)
        print("\tforward time: {}".format(time.time() - temp_time))
        temp_time = time.time()
        loss = criterion(output, target)
        print("\tloss time: {}".format(time.time() - temp_time))
        temp_time = time.time()

        # measure accuracy and record loss
        prec1, prec5 = accuracy(output.data, target, topk=(1, 5))

        if args.distributed:
            reduced_loss = reduce_tensor(loss.data)
            prec1 = reduce_tensor(prec1)
            prec5 = reduce_tensor(prec5)
        else:
            reduced_loss = loss.data

        losses.update(to_python_float(reduced_loss), input.size(0))
        top1.update(to_python_float(prec1), input.size(0))
        top5.update(to_python_float(prec5), input.size(0))
        print("\tupdate loss and acc time: {}".format(time.time() - temp_time))
        temp_time = time.time()

        # compute gradient and do SGD step
        optimizer.zero_grad()
        if args.fp16:
            with amp.scale_loss(loss, optimizer) as scaled_loss:
                scaled_loss.backward()
        else:
            loss.backward()
        optimizer.step()

        print("\tbackward prop time: {}".format(time.time() - temp_time))
        temp_time = time.time()

        torch.cuda.synchronize()
        # measure elapsed time
        batch_time.update(time.time() - end)

        end = time.time()

        if args.local_rank == 0 and args.print_freq > 0 and i % args.print_freq == 0 and i > 1:
            print("Epoch: [{0}][{1}/{2}]\t"
                  "Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t"
                  "Speed {3:.3f} ({4:.3f})\t"
                  "Data {data_time.val:.3f} ({data_time.avg:.3f})\t"
                  "Loss {loss.val:.4f} ({loss.avg:.4f})\t"
                  "Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t"
                  "Prec@5 {top5.val:.3f} ({top5.avg:.3f})".format(
                   epoch, i, train_loader_len,
                   args.total_batch_size / batch_time.val,
                   args.total_batch_size / batch_time.avg,
                   batch_time=batch_time,
                   data_time=data_time, loss=losses, top1=top1, top5=top5))

        if i > args.nb_iteration > 0:
            quit()

        print("\trecording time: {}".format(time.time() - temp_time))

    return batch_time.avg

Any tips would be appreciated!

Best wishes,
Chiang

Start by reproducing the training loop shown in the issue with the listed PyTorch, CUDA, Apex, and DALI versions, then compare the O1 and O0 timing measurements around loss, backward, and optimizer steps. Profile the backward pass to identify the source of the slowdown and document a reproducible explanation or confirmed fix.