mixed precision training(fp16) always give error.

Hi

I try to use apex lib for my own model, and it is impossible to use mixed precision training.
Whenever I set opt_level to O1/O2, program does not work.

First, I attach my environment system for better understanding.

I used environment collection script from pytorch github.

Collecting environment information...
PyTorch version: 1.2.0a0+f6aac41
Is debug build: No
CUDA used to build PyTorch: 10.1.168

OS: Ubuntu 18.04.2 LTS
GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0
CMake version: version 3.14.0

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 10.1.168
GPU models and configuration: 
GPU 0: Tesla P40
GPU 1: Tesla P40
GPU 2: Tesla P40

Nvidia driver version: 418.39
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.1

Versions of relevant libraries:
[pip] msgpack-numpy==0.4.3.2
[pip] numpy==1.16.4
[pip] torch==1.2.0a0+f6aac41
[pip] torchtext==0.4.0
[pip] torchvision==0.3.0a0
[conda] magma-cuda100             2.1.0                         5    local
[conda] mkl                       2019.1                      144  
[conda] mkl-include               2019.1                      144  
[conda] nomkl                     3.0                           0  
[conda] torch                     1.2.0a0+f6aac41          pypi_0    pypi
[conda] torchtext                 0.4.0                    pypi_0    pypi
[conda] torchvision               0.3.0a0                  pypi_0    pypi

I always end up with 3 different type errors when I set opt_level to O1/O2.

1. Never Ending Gradient overflow.

2019-11-20 18:08:56,128 STDOUT Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 32768.0
2019-11-20 18:08:56,128 STDOUT Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 32768.0
2019-11-20 18:08:56,128 STDOUT Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 32768.0
2019-11-20 18:09:10,507 STDOUT Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 16384.0
2019-11-20 18:09:10,508 STDOUT Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 16384.0
2019-11-20 18:09:10,508 STDOUT Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 16384.0
...
2019-11-21 09:51:36,521 STDOUT Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4096.0
2019-11-21 10:27:59,618 STDOUT Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4096.0
2019-11-21 10:27:59,622 STDOUT Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4096.0
2019-11-21 10:27:59,622 STDOUT Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4096.0
2019-11-21 11:03:45,922 STDOUT Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4096.0
2019-11-21 11:03:45,922 STDOUT Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4096.0
2019-11-21 11:03:45,922 STDOUT Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4096.0

2. CUDNN_STATUS_INTERNAL_ERROR

Traceback (most recent call last):
  File "train_parallel.py", line 450, in <module>
    main()
  File "train_parallel.py", line 184, in main
    train(config, train_loader, model, criterion, optimizer, epoch)
  File "train_parallel.py", line 308, in train
    query_rep = model(query)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 525, in __call__
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 385, in forward
    output = self.module(*inputs[0], **kwargs[0])
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 525, in __call__
    result = self.forward(*input, **kwargs)
  File "/root/c3_workspace/model/model.py", line 45, in forward
    x = self.convs(x)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 525, in __call__
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/container.py", line 92, in forward
    input = module(input)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 525, in __call__
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/container.py", line 92, in forward
    input = module(input)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 525, in __call__
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/conv.py", line 345, in forward
    self.padding, self.dilation, self.groups)
  File "/opt/conda/lib/python3.6/site-packages/apex/amp/wrap.py", line 28, in wrapper
    return orig_fn(*new_args, **kwargs)
RuntimeError: cuDNN error: CUDNN_STATUS_INTERNAL_ERROR

3. CUDA out of memory with plenty of spaces

Traceback (most recent call last):
  File "train_parallel.py", line 450, in <module>
    main()
  File "train_parallel.py", line 184, in main
    train(config, train_loader, model, criterion, optimizer, epoch)
  File "train_parallel.py", line 308, in train
    query_rep = model(query)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 525, in __call__
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 385, in forward
    output = self.module(*inputs[0], **kwargs[0])
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 525, in __call__
    result = self.forward(*input, **kwargs)
  File "/root/c3_workspace/model/model.py", line 45, in forward
    x = self.convs(x)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 525, in __call__
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/container.py", line 92, in forward
    input = module(input)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 525, in __call__
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/container.py", line 92, in forward
    input = module(input)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 525, in __call__
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/conv.py", line 345, in forward
    self.padding, self.dilation, self.groups)
  File "/opt/conda/lib/python3.6/site-packages/apex/amp/wrap.py", line 28, in wrapper
    return orig_fn(*new_args, **kwargs)
RuntimeError: CUDA out of memory. Tried to allocate 160.00 MiB (GPU 0; 23.88 GiB total capacity; 2.94 GiB already allocated; 20.36 GiB free; 24.86 MiB cached)

I tried several settings and always ending up meeting this 3 different errors that mentioned above.

1. Data prefetch On / Off
2. Different batch size (size down if cuda mem error happened)
3. Different number of P40 (2~4)

My code for model and training scripts are written below.

class SNRM(nn.Module):
    def __init__(self, config, weights=None):
        super(SNRM, self).__init__()
        self.ngram = config['network']['ngram']
        # use same embeddings for query and document
        if weights is not None:
            # This code works at 1.1 version
            self.embed = nn.Embedding.from_pretrained(weights, freeze=False, padding_idx=PAD)
            # self.embed = nn.Embedding.from_pretrained(weights, padding_idx=PAD)
            self.embed_dim = self.embed.embedding_dim
        else:
            self.embed_dim = config['network']['d_embed']
            self.embed = nn.Embedding(config['n_vocab'], self.embed_dim, padding_idx=PAD)

        # use same cnn for query and document
        self.filter_nums = [self.embed_dim] + config['network']['d_hiddens']
        conv_list = []
        for i, (f1, f2) in enumerate(pairwise(self.filter_nums)):
            conv_list.append(
                nn.Sequential(
                    nn.Conv2d(f1, f2, (1, self.ngram if i == 0 else 1),
                              padding=(0, int(self.ngram / 2)) if i == 0 else (0, 0)),
                    nn.ReLU(),
                    nn.Dropout(config['network']['dropout'])
                )
            )
        self.convs = nn.Sequential(*conv_list)

    def forward(self, x):
        x = self.embed(x)  # [batch, max_len, D]
        x = x.transpose(1, 2).unsqueeze(2)  # [B x D x 1 x max_len]
        x = self.convs(x)
        x = x.squeeze(2).transpose(1, 2)  # [B x max_len x Out]

        # average pooling
        x = torch.mean(x, dim=1)
        return x

When I used data prefteching,

def main():
    config = parse()

    if config['local_rank'] == 0:
        print("CUDNN VERSION: {}".format(torch.backends.cudnn.version()))
        print("NUMBER of GPUs: {}".format(torch.cuda.device_count()))

    cudnn.benchmark = True
    best_score = 0

    if config['deterministic']:
        cudnn.benchmark = False
        cudnn.deterministic = True
        torch.manual_seed(config['local_rank'])
        torch.set_printoptions(precision=10)

    # FOR DISTRIBUTED:  If we are running under torch.distributed.launch,
    # the 'WORLD_SIZE' environment variable will also be set automatically.
    config['distributed'] = False

    if 'WORLD_SIZE' in os.environ:
        config['distributed'] = int(os.environ['WORLD_SIZE']) > 1

    print("\n\nlocal_rank (GPU NUMBER): ", config['local_rank'], "\n\n")

    config['gpu'] = 0
    config['world_size'] = 1

    if config['distributed']:
        config['gpu'] = config['local_rank']
        # FOR DISTRIBUTED:  Set the device according to local_rank.
        torch.cuda.set_device(config['gpu'])

        # FOR DISTRIBUTED:  Initialize the backend.  torch.distributed.launch will provide
        # environment variables, and requires that you use init_method=`env://`.
        torch.distributed.init_process_group(backend='nccl',
                                             init_method='env://')
        config['world_size'] = torch.distributed.get_world_size()

    assert torch.backends.cudnn.enabled, "Amp requires cudnn backend to be enabled."

    dictionary, num, weights = load_w2v(config)
    weights = torch.FloatTensor(weights)

    if config['local_rank'] == 0:
        print('Loaded and imported w2v weights (%d) words' % num)

    model = SNRM(config, weights)
    model = model.cuda()

    # Scale learning rate based on global batch size
    # follow https://github.com/NVIDIA/apex/blob/master/examples/imagenet/main_amp.py code
    lr = config['optimizer']['lr']
    batch_size = config['learning']['batch_size']
    world_size = config['world_size']
    config['optimizer']['lr'] = lr * float(batch_size * world_size) / 256.

    if config['local_rank'] == 0:
        print('learning rate : ', config['optimizer']['lr'])

    parameters = filter(lambda p: p.requires_grad, model.parameters())
    optimizer = getattr(torch.optim, config['optimizer']['optim'])(parameters, config['optimizer']['lr'])
    # scheduler = ReduceLROnPlateau(optimizer, 'min')
    model, optimizer = amp.initialize(model, optimizer,
                                      opt_level=config['opt_level'])

    if config['local_rank'] == 0:
        print('amp initialized')

    if config['distributed']:
        model = torch.nn.parallel.DistributedDataParallel(model,
                                                          device_ids=[config['gpu']],
                                                          output_device=config['gpu'],
                                                          find_unused_parameters=True)
    if config['local_rank'] == 0:
        print('DDP(model) Finished')

    criterion = nn.MarginRankingLoss(margin=1.0).cuda()

    if config['local_rank'] == 0:
        print('criterion Finished')

    collection = Collections()
    collection.load_query(config['query_train_file'])

    if config['local_rank'] == 0:
        print('Start Loading Training Set')

    assert 'doc_mapper' in config.keys()

    doc_mapper = read_pickle(config['doc_mapper'])

    new_mapper = {k: os.path.join('enclo_docs/', v) for k, v in doc_mapper.items()}
    del doc_mapper

    train_dataset = DistributedRetrievalDataset(collection, new_mapper, dictionary, config, torch.LongTensor)

    if config['local_rank'] == 0:
        print('Finish Loading Training Set')

    collection.clear_query()
    collection.load_query(config['query_valid_file'])

    valid_dataset = DistributedRetrievalDataset(collection, new_mapper, dictionary, config, torch.LongTensor)
    if config['local_rank'] == 0:
        print('Finish Loading Valid Set')

    train_sampler = None
    valid_sampler = None
    if config['distributed']:
        train_sampler = DistributedSampler(train_dataset)
        valid_sampler = DistributedSampler(valid_dataset)

    train_loader = DataLoader(dataset=train_dataset,
                              batch_size=config['learning']['batch_size'],
                              shuffle=(train_sampler is None),
                              sampler=train_sampler,
                              pin_memory=True,
                              num_workers=0)
    valid_loader = DataLoader(dataset=valid_dataset,
                              batch_size=config['learning']['batch_size'],
                              shuffle=False,
                              sampler=valid_sampler,
                              pin_memory=True,
                              num_workers=0)

    if config['local_rank'] == 0:
        print('Total # of training instance: %d' % len(train_dataset))
        print('Total # of validation instance: %d' % len(valid_dataset))

    print('local_rank: ', config['local_rank'], '# of training instance: %d' % len(train_loader))
    print('local_rank: ', config['local_rank'], '# of validation instance: %d' % len(valid_loader))

    for epoch in range(config['learning']['epochs']):
        if config['distributed']:
            train_sampler.set_epoch(epoch)

        train(config, train_loader, model, criterion, optimizer, epoch)

        if (epoch + 1) % config['valid'] == 0:
            score = validate(config, valid_loader, model, criterion)
            if config['local_rank'] == 0:
                print("At", epoch + 1, "score", score)
                is_best = score > best_score
                best_score = max(score, best_score)

                save_checkpoint({
                    'epoch': epoch + 1,
                    'state_dict': model.state_dict(),
                    'best_score': best_score,
                    'optimizer': optimizer.state_dict()
                }, is_best, 'data_out/checkpoint_ep{0:003d}.pth.tar'.format(epoch + 1))


def save_checkpoint(state, is_best, filename):
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, 'data_out/model_best.pth.tar')


def load_w2v(config):
    dictionary = Dictionary()
    dictionary.load(config['dictionary'])
    config['n_vocab'] = len(dictionary)
    print('Loaded dictionary')
    w2v = gensim.models.KeyedVectors.load_word2vec_format(config['w2v_file'], binary=True)
    emb_dim = w2v.vector_size
    weights, num = [], 0
    for token in dictionary.idx2token:
        if token in w2v:
            weights.append(w2v[token])
            num = num + 1
        else:
            weights.append(np.random.uniform(0, 1, emb_dim))
    return dictionary, num, weights


class data_prefetcher():
    def __init__(self, loader):
        self.loader = iter(loader)
        self.stream = torch.cuda.Stream()

        self.preload()

    def preload(self):
        try:
            self.next_q, self.next_d1, self.next_d2, self.next_s = next(self.loader)
            self.next_s = self.next_s.type(torch.float)
        except StopIteration:
            self.next_q = None
            self.next_d1 = None
            self.next_d2 = None
            self.next_s = None
            return
        # if record_stream() doesn't work, another option is to make sure device inputs are created
        # on the main stream.
        # self.next_input_gpu = torch.empty_like(self.next_input, device='cuda')
        # self.next_target_gpu = torch.empty_like(self.next_target, device='cuda')
        # Need to make sure the memory allocated for next_* is not still in use by the main stream
        # at the time we start copying to next_*:
        # self.stream.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(self.stream):
            self.next_q = self.next_q.cuda(non_blocking=True)
            self.next_d1 = self.next_d1.cuda(non_blocking=True)
            self.next_d2 = self.next_d1.cuda(non_blocking=True)
            self.next_s = self.next_s.cuda(non_blocking=True)

            # more code for the alternative if record_stream() doesn't work:
            # copy_ will record the use of the pinned source tensor in this side stream.
            # self.next_input_gpu.copy_(self.next_input, non_blocking=True)
            # self.next_target_gpu.copy_(self.next_target, non_blocking=True)
            # self.next_input = self.next_input_gpu
            # self.next_target = self.next_target_gpu

            # With Amp, it isn't necessary to manually convert data to half.
            # if args.fp16:
            #     self.next_input = self.next_input.half()
            # else:

    def next(self):
        torch.cuda.current_stream().wait_stream(self.stream)
        query = self.next_q
        doc1 = self.next_d1
        doc2 = self.next_d2
        score = self.next_s

        if query is not None:
            query.record_stream(torch.cuda.current_stream())
        if doc1 is not None:
            doc1.record_stream(torch.cuda.current_stream())
        if doc2 is not None:
            doc2.record_stream(torch.cuda.current_stream())
        if score is not None:
            score.record_stream(torch.cuda.current_stream())

        self.preload()
        return query, doc1, doc2, score


def train(config, train_loader, model, criterion, optimizer, epoch):
    print("GPU: ", config['local_rank'], "EPOCH :", epoch, "training starts")

    losses = AverageMeter()
    norms = AverageMeter()

    model.train()

    prefetcher = data_prefetcher(train_loader)
    query, doc_a, doc_b, score = prefetcher.next()
    # if config['local_rank'] == 0:
    #     print("[Data]", query, doc_a, doc_b)

    i = 0
    while query is not None:
        i += 1

        optimizer.zero_grad()

        # if config['local_rank'] == 0 and i < 10:
        #     print(i, 'query', query, 'doc', doc_a)

        query_rep = model(query)
        doc_a_rep = model(doc_a)
        doc_b_rep = model(doc_b)

        # if config['local_rank'] == 0 and i < 10:
        #     print(i, 'query_rep', query_rep, 'doc_rep', doc_a_rep)

        logits_doc_a = torch.mul(query_rep, doc_a_rep).sum(1)
        logits_doc_b = torch.mul(query_rep, doc_b_rep).sum(1)

        # if config['local_rank'] == 0 and i < 10:
        #     print(i, 'logita', logits_doc_a, 'logitb', logits_doc_b)

        loss = criterion(logits_doc_a, logits_doc_b, score)
        l1_norm = torch.norm(torch.cat([query_rep, doc_a_rep, doc_b_rep], dim=1), p=1)

        # if config['local_rank'] == 0 and i < 10:
        #     print(i, 'loss', loss, 'norm', l1_norm)

        loss += config['optimizer']['weight_decay'] * l1_norm

        with amp.scale_loss(loss, optimizer) as scaled_loss:
            scaled_loss.backward()

        optimizer.step()
        # if config['local_rank'] == 0 and i < 10:
        #     for param in model.parameters():
        #         print("[model parameter] ", param.grad.float().sum())

        if i % config['print_step'] == 0:

            if config['distributed']:
                reduced_loss = reduce_tensor(config, loss.data)
                reduced_norm = reduce_tensor(config, l1_norm.data)
            else:
                reduced_loss = loss.data
                reduced_norm = l1_norm.data

            if config['local_rank'] == 0:
                print("reduced_loss", to_python_float(reduced_loss))
            losses.update(to_python_float(reduced_loss), query.size(0))
            norms.update(to_python_float(reduced_norm), query.size(0))
            # print("[before sync]GPU: ", config['local_rank'], "EPOCH :", epoch,
            #      "LOSS", loss, "NORM", l1_norm)

            torch.cuda.synchronize()

            if config['local_rank'] == 0:
                print('Epoch: [{0}], step: [{1}]\t'
                      'Loss {l.val:.10f} ({l.avg:.4f})\t'
                      'L1Norm {norm.val:.10f} ({norm.avg:.4f})'.format(
                    epoch, i, l=losses, norm=norms))

        query, doc_a, doc_b, score = prefetcher.next()


    if config['local_rank'] == 0:
        print('Epoch: [{0}]\t'
              'Loss {l.val:.10f} ({l.avg:.4f})\t'
              'L1Norm {norm.val:.10f} ({norm.avg:.4f})'.format(
            epoch, l=losses, norm=norms))


def validate(config, val_loader, model, criterion):
    scores = AverageMeter()
    query_sparsities = AverageMeter()
    doc_sparsities = AverageMeter()

    model.eval()

    prefetcher = data_prefetcher(val_loader)
    query, doc_a, doc_b, score = prefetcher.next()
    i = 0
    while query is not None:
        i += 1
        query_rep = model(query)
        doc_a_rep = model(doc_a)
        doc_b_rep = model(doc_b)
        logits_doc_a = torch.mul(query_rep, doc_a_rep).sum(1)
        logits_doc_b = torch.mul(query_rep, doc_b_rep).sum(1)

        query_sparsity = sparsity_ratio(query_rep)
        doc_a_sparsity = sparsity_ratio(doc_a_rep)
        doc_b_sparsity = sparsity_ratio(doc_b_rep)

        total_score = ((logits_doc_a - logits_doc_b) * score).sum()

        if config['distributed']:
            reduced_score = reduce_tensor(config, total_score.data)
            reduced_query_sparsity = reduce_tensor(config, query_sparsity.data)
            reduced_doc_a_sparsity = reduce_tensor(config, doc_a_sparsity.data)
            reduced_doc_b_sparsity = reduce_tensor(config, doc_b_sparsity.data)
        else:
            reduced_score = score.data
            reduced_query_sparsity = query_sparsity.data
            reduced_doc_a_sparsity = doc_a_sparsity.data
            reduced_doc_b_sparsity = doc_b_sparsity.data

        scores.update(to_python_float(reduced_score))
        query_sparsities.update(to_python_float(reduced_query_sparsity.mean()), query.size(0))
        doc_sparsities.update(to_python_float(reduced_doc_a_sparsity.mean()), query.size(0))
        doc_sparsities.update(to_python_float(reduced_doc_b_sparsity.mean()), query.size(0))

        query, doc_a, doc_b, score = prefetcher.next()

    if config['local_rank'] == 0:
        print('Score {score.val:.10f} ({score.avg:.4f})\t'
              'Query Sparsity {query_s.val:.4f} ({query_s.avg:.4f})\t'
              'Doc Sparsity {doc_s.val:.4f} ({doc_s.avg:.4f})'.format(
            score=scores, query_s=query_sparsities, doc_s=doc_sparsities))

    return scores.avg


def reduce_tensor(config, tensor):
    rt = tensor.clone()
    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
    rt /= config['world_size']
    return rt


class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

When I do not use data prefetching,

import os

import gensim
import numpy as np
import torch
import torch.nn as nn
import yaml
import pprint
import argparse
import shutil

import torch.distributed as dist
from torch.utils.data import DataLoader, RandomSampler
from torch.utils.data.distributed import DistributedSampler
import torch.backends.cudnn as cudnn

from model.dataset import DistributedRetrievalDataset
from model.dictionary import Collections, Dictionary
from model.model import SNRM
from model.util import sparsity_ratio, read_pickle, get_ordinal

from apex.fp16_utils import *
from apex import amp, optimizers
from apex.multi_tensor_apply import multi_tensor_applier


# python -m torch.distributed.launch --nproc_per_node={$n_gpu} train_parallel.py

def parse():
    # from https://github.com/NVIDIA/apex/blob/master/examples/simple/distributed/distributed_data_parallel.py
    parser = argparse.ArgumentParser()
    # FOR DISTRIBUTED:  Parse for the local_rank argument, which will be supplied
    # automatically by torch.distributed.launch.
    parser.add_argument("--local_rank", default=0, type=int)
    parser.add_argument("-p", "--parameter_yaml", default='parameters/snrm-parallel-1st.yaml', type=str)

    args = parser.parse_args()

    with open(args.parameter_yaml, 'r') as f:
        config = yaml.load(f, Loader=yaml.FullLoader)

    config['local_rank'] = args.local_rank
    if config['local_rank'] == 0:
        pprint.pprint(config)
    return config


def main():
    config = parse()

    if config['local_rank'] == 0:
        print("CUDNN VERSION: {}".format(torch.backends.cudnn.version()))
        print("NUMBER of GPUs: {}".format(torch.cuda.device_count()))

    cudnn.benchmark = True
    best_score = 0

    if config['deterministic']:
        cudnn.benchmark = False
        cudnn.deterministic = True
        torch.manual_seed(config['local_rank'])
        torch.set_printoptions(precision=10)

    # FOR DISTRIBUTED:  If we are running under torch.distributed.launch,
    # the 'WORLD_SIZE' environment variable will also be set automatically.
    config['distributed'] = False

    if 'WORLD_SIZE' in os.environ:
        config['distributed'] = int(os.environ['WORLD_SIZE']) > 1

    print("\n\nlocal_rank (GPU NUMBER): ", config['local_rank'], "\n\n")

    config['gpu'] = 0
    config['world_size'] = 1

    if config['distributed']:
        config['gpu'] = config['local_rank']
        # FOR DISTRIBUTED:  Set the device according to local_rank.
        torch.cuda.set_device(config['gpu'])

        # FOR DISTRIBUTED:  Initialize the backend.  torch.distributed.launch will provide
        # environment variables, and requires that you use init_method=`env://`.
        torch.distributed.init_process_group(backend='nccl',
                                             init_method='env://')
        config['world_size'] = torch.distributed.get_world_size()

    assert torch.backends.cudnn.enabled, "Amp requires cudnn backend to be enabled."

    dictionary, num, weights = load_w2v(config)
    weights = torch.FloatTensor(weights)

    if config['local_rank'] == 0:
        print('Loaded and imported w2v weights (%d) words' % num)

    model = SNRM(config, weights)
    model = model.cuda()

    # Scale learning rate based on global batch size
    # follow https://github.com/NVIDIA/apex/blob/master/examples/imagenet/main_amp.py code
    lr = config['optimizer']['lr']
    batch_size = config['learning']['batch_size']
    world_size = config['world_size']
    config['optimizer']['lr'] = lr * float(batch_size * world_size) / 256.

    if config['local_rank'] == 0:
        print('learning rate : ', config['optimizer']['lr'])

    optimizer = torch.optim.SGD(model.parameters(), config['optimizer']['lr'])
    # scheduler = ReduceLROnPlateau(optimizer, 'min')
    model, optimizer = amp.initialize(model, optimizer,
                                      opt_level=config['opt_level'])

    if config['local_rank'] == 0:
        print('amp initialized')

    if config['distributed']:
        model = torch.nn.parallel.DistributedDataParallel(model,
                                                          device_ids=[config['gpu']],
                                                          output_device=config['gpu'])

    if config['local_rank'] == 0:
        print('DDP(model) Finished')

    criterion = nn.MarginRankingLoss(margin=1.0).cuda()

    if config['local_rank'] == 0:
        print('criterion Finished')

    collection = Collections()
    collection.load_query(config['query_train_file'])

    if config['local_rank'] == 0:
        print('Start Loading Training Set')

    assert 'doc_mapper' in config.keys()

    doc_mapper = read_pickle(config['doc_mapper'])

    new_mapper = {k: os.path.join('enclo_docs/', v) for k, v in doc_mapper.items()}
    del doc_mapper

    train_dataset = DistributedRetrievalDataset(collection, new_mapper, dictionary, config, torch.LongTensor)

    if config['local_rank'] == 0:
        print('Finish Loading Training Set')

    collection.clear_query()
    collection.load_query(config['query_valid_file'])

    valid_dataset = DistributedRetrievalDataset(collection, new_mapper, dictionary, config, torch.LongTensor)
    if config['local_rank'] == 0:
        print('Finish Loading Valid Set')

    train_sampler = None
    valid_sampler = None
    if config['distributed']:
        train_sampler = DistributedSampler(train_dataset)
        valid_sampler = DistributedSampler(valid_dataset)

    train_loader = DataLoader(dataset=train_dataset,
                              batch_size=config['learning']['batch_size'],
                              shuffle=(train_sampler is None),
                              sampler=train_sampler,
                              pin_memory=True,
                              num_workers=0)
    valid_loader = DataLoader(dataset=valid_dataset,
                              batch_size=config['learning']['batch_size'],
                              shuffle=False,
                              sampler=valid_sampler,
                              pin_memory=True,
                              num_workers=0)

    if config['local_rank'] == 0:
        print('Total # of training instance: %d' % len(train_dataset))
        print('Total # of validation instance: %d' % len(valid_dataset))

    print('local_rank: ', config['local_rank'], '# of training instance: %d' % len(train_loader))
    print('local_rank: ', config['local_rank'], '# of validation instance: %d' % len(valid_loader))

    for epoch in range(config['learning']['epochs']):
        model.train()

        model.train()
        train_loss = 0
        l1_norms = []
        step = 0
        for query, doc_a, doc_b, score in train_loader:
            query = query.cuda()
            doc_a = doc_a.cuda()
            doc_b = doc_b.cuda()
            score = score.type(torch.float).cuda()

            query_rep = model(query)
            doc_a_rep = model(doc_a)
            doc_b_rep = model(doc_b)

            logits_doc_a = torch.mul(query_rep, doc_a_rep).sum(1)
            logits_doc_b = torch.mul(query_rep, doc_b_rep).sum(1)

            loss = criterion(logits_doc_a, logits_doc_b, score)

            optimizer.zero_grad()

            l1_norm = torch.norm(torch.cat([query_rep, doc_a_rep, doc_b_rep], dim=1), p=1)
            loss += config['optimizer']['weight_decay'] * l1_norm

            with amp.scale_loss(loss, optimizer) as scaled_loss:
                scaled_loss.backward()

            l1_norms.append(l1_norm.item())

            optimizer.step()

            train_loss += loss.item()
            step += 1
            if step % config['print_step'] == 0 and config['local_rank'] == 0:
                print(f'In {get_ordinal(epoch + 1)} epoch, step: {step},'
                      f'train loss(avg): {train_loss / step}, L1-norm: {np.mean(l1_norms)}')

        print(f"In {get_ordinal(epoch+1)} epoch, GPU:{config['local_rank']}'s "
              f"train loss(all/avg) : {train_loss}, {train_loss/step}, L1-norm : {np.mean(l1_norms)} ")

        if (epoch + 1) % config['print'] == 0:
            model.eval()
            with torch.no_grad():
                sparsity_query, sparsity_doc = [], []
                total_score = 0

                for query, doc_a, doc_b, score in valid_loader:
                    # ex) query / doc_a / doc_b / -1
                    # if doc_a > doc_b : score = 1 else score = -1
                    # higher score, higher quality
                    query = query.cuda()  # B x max_q
                    doc_a = doc_a.cuda() # B x max_d
                    doc_b = doc_b.cuda()  # B x max_d
                    score = score.type(torch.float)
                    score = score.cuda()  # [B]

                    query_rep = model(query)
                    doc_a_rep = model(doc_a)
                    doc_b_rep = model(doc_b)

                    sparsity_query.extend(sparsity_ratio(query_rep).cpu().numpy())
                    sparsity_doc.extend(sparsity_ratio(doc_a_rep).cpu().numpy())
                    sparsity_doc.extend(sparsity_ratio(doc_b_rep).cpu().numpy())

                    logits_doc_a = torch.mul(query_rep, doc_a_rep).sum(1)  # single array []
                    logits_doc_b = torch.mul(query_rep, doc_b_rep).sum(1)
                    total_score += ((logits_doc_a - logits_doc_b) * score).sum().item()

                print(f"At GPU: {config['local_rank']}, score : {total_score}, "
                      f"query sparsity: {np.mean(sparsity_query)}, doc sparsity: {np.mean(sparsity_doc)}")

        if (epoch + 1) % config['save_point'] == 0 and config['local_rank'] == 0:
            save_point_name = config['model_name'].split('.')[0] + '_' + str(epoch + 1) + '.pkl'
            torch.save(model.state_dict(), save_point_name)
            # save file name {$model_name}_{$epoch}.pkl


def load_w2v(config):
    dictionary = Dictionary()
    dictionary.load(config['dictionary'])
    config['n_vocab'] = len(dictionary)
    print('Loaded dictionary')
    w2v = gensim.models.KeyedVectors.load_word2vec_format(config['w2v_file'], binary=True)
    emb_dim = w2v.vector_size
    weights, num = [], 0
    for token in dictionary.idx2token:
        if token in w2v:
            weights.append(w2v[token])
            num = num + 1
        else:
            weights.append(np.random.uniform(0, 1, emb_dim))
    return dictionary, num, weights


def reduce_tensor(config, tensor):
    rt = tensor.clone()
    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
    rt /= config['world_size']
    return rt


if __name__ == '__main__':
    main()

I can not find the reason for the bug.
Is there any part of my code that does not work with current apex lib?
Especially, for cuda memory error or CUDNN error, there's no way to find which part of the code makes that error.
Is it because of lack of NV-link in my P40 cluster environment??

And also, even when I tried to use fp32 option, sometime training script gives cuda mememory error (even if there are plenty of free mem).
Does DDP option need a lot of gpu memeory?

For example. when I used single gpu machine setting, batch size 64 works.
However, when I used 3 gpu, batch size lower than 48 only works with the same parameter except batch size.
For the 4 gpu, its getting worse, even batch size=32 gives error.
Should I reduce the batch size to be inversely proportional to number of gpus?

Start with the training entry points in train_parallel.py and the failing model path in model/model.py, then inspect the apex/amp/wrap.py frame and the amp.initialize call for the O1/O2 runs. Reproduce the reported gradient overflow, cuDNN, and CUDA errors with the supplied environment and training configuration; done means the mixed-precision run completes without these failures.