Quantization awareness training multi-gpu suport?

robotcator123: <p>Does pytorch support multi-GPU in quantization awareness training?<br>
In this script <a href="https://github.com/pytorch/vision/blob/master/references/classification/train_quantization.py#L73" rel="nofollow noopener">https://github.com/pytorch/vision/blob/master/references/classification/train_quantization.py#L73</a>, it seems that it has the logic of multi-GPU.</p>

Mazhar_Shaikh: <p>Hi <a class="mention" href="/u/robotcator123">@robotcator123</a>,<br>
Multi gpu training is orthogonal to quantization aware training. Code written with Pytorch’s quantization aware training modules will work whether you are using a single gpu or using Data parallel on multiple gpus. Hope this helps!</p>

robotcator123: <p>Hi, <a class="mention" href="/u/mazhar_shaikh">@Mazhar_Shaikh</a>，</p>
<p>Actually, I train the mobilenet using this command in cluster.<br>
<code>python -m torch.distributed.launch --nproc_per_node=8 --use_env train.py --data-path=./imagenet_1k</code> , it seems that the code works fine.</p>
<p>But I change the training script to<br>
<code>python -m torch.distributed.launch --nproc_per_node=8 --use_env train_quantization.py --data-path=./imagenet_1k</code></p>
<p>it will raise error like this after print lost of unknown data:<br>
`Namespace(backend=‘qnnpack’, batch_size=32, cache_dataset=False, data_path=’~/test/imagenet_1k’, device=‘cuda’, dist_backend=‘nccl’, dist_url=‘env://’, distributed=True, epochs=90, eval_batch_size=128, gpu=0, lr=0.0001, lr_gamma=0.1, lr_step_size=30, model=‘mobilenet_v2’, momentum=0.9, num_batch_norm_update_epochs=3, num_calibration_batches=32, num_observer_update_epochs=4, output_dir=’.’, post_training_quantize=False, print_freq=10, rank=0, resume=’’, start_epoch=0, test_only=False, weight_decay=0.0001, workers=16, world_size=8)<br>
Loading data<br>
Loading data<br>
Loading training data<br>
Took 0.27007627487182617<br>
Loading validation data<br>
Creating data loaders<br>
Creating model mobilenet_v2<br>
Traceback (most recent call last):<br>
File “train_quantization.py”, line 258, in <br>
main(args)<br>
File “train_quantization.py”, line 77, in main<br>
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])<br>
File “xxx/.conda/envs/pytorch1.3/lib/python3.6/site-packages/torch/nn/parallel/distributed.py”, line 298, in <strong>init</strong><br>
self.broadcast_bucket_size)<br>
File “xxx/.conda/envs/pytorch1.3/lib/python3.6/site-packages/torch/nn/parallel/distributed.py”, line 480, in _distributed_broadcast_coalesced<br>
dist._broadcast_coalesced(self.process_group, tensors, buffer_size)<br>
TypeError: _broadcast_coalesced(): incompatible function arguments. The following argument types are supported:<br>
1. (process_group: torch.distributed.ProcessGroup, tensors: List[at::Tensor], buffer_size: int) -&gt; None</p>
<p>Invoked with: &lt;torch.distributed.ProcessGroupNCCL object at 0x7f943f78dd18&gt;, [tensor([[[[ 1.3185e-02, -4.3213e-03,  1.4823e-02],<br>
…<br>
…<br>
…<br>
subprocess.CalledProcessError: Command ‘[’/xxxx/pytorch1.3/bin/python’, ‘-u’, ‘train_quantization.py’, ‘–data-path=./imagenet_1k’]’ returned non-zero exit status 1.`<br>
sorry for hiding some personal information.</p>

robotcator123: <p>It seems that broadcast can not support None tensor, is there anybody know this problem?</p>

dskhudia: <p>Hi <a class="mention" href="/u/robotcator123">@robotcator123</a>, If you believe broadcast of None doesn’t work as expected, please open an issue against PyTorch with a minimal reproducible example.</p>

robotcator123: <p>Hi, <a class="mention" href="/u/dskhudia">@dskhudia</a>, Thank you for your response. I can give some reproduction steps. For more detailed producible example, maybe I will do it at weekends.</p>
<p>1: Download the imagenet1k dataset.<br>
2: pip install torchvision==0.5.0, this will upgrade the torch into 1.4.0.<br>
3: Use the script with the commands<br>
<code>python -m torch.distributed.launch --nproc_per_node=8 --use_env train_quant.py --data-path=./imagenet_1k </code></p>
<p>The train_quant.py script is borrowed from torchvision reference code.</p>
<details>
<summary>
Summary</summary>
<pre><code class="lang-auto">from __future__ import print_function
import datetime
import os
import time
import sys
import copy

import torch
import torch.utils.data
from torch import nn
import torchvision
import torch.quantization
import train_utils as utils
from train import train_one_epoch, evaluate, load_data

def main(args):
    if args.output_dir:
        utils.mkdir(args.output_dir)

    utils.init_distributed_mode(args)

    print(args)

    if args.post_training_quantize and args.distributed:
        raise RuntimeError("Post training quantization example should not be performed "
                           "on distributed mode")

    # Set backend engine to ensure that quantized model runs on the correct kernels
    if args.backend not in torch.backends.quantized.supported_engines:
        raise RuntimeError("Quantized backend not supported: " + str(args.backend))
    torch.backends.quantized.engine = args.backend

    device = torch.device(args.device)
    torch.backends.cudnn.benchmark = True

    # Data loading code
    print("Loading data")
    train_dir = os.path.join(args.data_path, 'train')
    val_dir = os.path.join(args.data_path, 'val')

    dataset, dataset_test, train_sampler, test_sampler = load_data(train_dir, val_dir,
                                                                   args.cache_dataset, args.distributed)
    data_loader = torch.utils.data.DataLoader(
        dataset, batch_size=args.batch_size,
        sampler=train_sampler, num_workers=args.workers, pin_memory=True)

    data_loader_test = torch.utils.data.DataLoader(
        dataset_test, batch_size=args.eval_batch_size,
        sampler=test_sampler, num_workers=args.workers, pin_memory=True)

    print("Creating model", args.model)
    # when training quantized models, we always start from a pre-trained fp32 reference model
    model = torchvision.models.quantization.__dict__[args.model](pretrained=True, quantize=args.test_only)
    model.to(device)

    if not (args.test_only or args.post_training_quantize):
        model.fuse_model()
        model.qconfig = torch.quantization.get_default_qat_qconfig(args.backend)
        torch.quantization.prepare_qat(model, inplace=True)

        optimizer = torch.optim.SGD(
            model.parameters(), lr=args.lr, momentum=args.momentum,
            weight_decay=args.weight_decay)

        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                                       step_size=args.lr_step_size,
                                                       gamma=args.lr_gamma)

    criterion = nn.CrossEntropyLoss()
    model_without_ddp = model
    if args.distributed:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
        model_without_ddp = model.module
        print (model.module)

    model.apply(torch.quantization.enable_observer)
    model.apply(torch.quantization.enable_fake_quant)
    start_time = time.time()
    for epoch in range(args.start_epoch, args.epochs):
        if args.distributed:
            train_sampler.set_epoch(epoch)
        print('Starting training for epoch', epoch)
        train_one_epoch(model, criterion, optimizer, data_loader, device, epoch,
                        args.print_freq)
        lr_scheduler.step()
        with torch.no_grad():
            if epoch &gt;= args.num_observer_update_epochs:
                print('Disabling observer for subseq epochs, epoch = ', epoch)
                model.apply(torch.quantization.disable_observer)
            if epoch &gt;= args.num_batch_norm_update_epochs:
                print('Freezing BN for subseq epochs, epoch = ', epoch)
                model.apply(torch.nn.intrinsic.qat.freeze_bn_stats)
            print('Evaluate QAT model')

            evaluate(model, criterion, data_loader_test, device=device)
            quantized_eval_model = copy.deepcopy(model)
            quantized_eval_model.eval()
            quantized_eval_model.to(torch.device('cpu'))
            torch.quantization.convert(quantized_eval_model, inplace=True)

            print('Evaluate Quantized model')
            evaluate(quantized_eval_model, criterion, data_loader_test,
                     device=torch.device('cpu'))

        model.train()

        print('Saving models after epoch ', epoch)

    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    print('Training time {}'.format(total_time_str))

def parse_args():                                                                                                                                                                                   
    import argparse
    parser = argparse.ArgumentParser(description='PyTorch Classification Training')

    parser.add_argument('--data-path',
                        default='/datasets01/imagenet_full_size/061417/',
                        help='dataset')
    parser.add_argument('--model',
                        default='mobilenet_v2',
                        help='model')
    parser.add_argument('--backend',
                        default='qnnpack',
                        help='fbgemm or qnnpack')
    parser.add_argument('--device',
                        default='cuda',
                        help='device')

    parser.add_argument('-b', '--batch-size', default=32, type=int,
                        help='batch size for calibration/training')
    parser.add_argument('--eval-batch-size', default=128, type=int,
                        help='batch size for evaluation')
    parser.add_argument('--epochs', default=90, type=int, metavar='N',
                        help='number of total epochs to run')
    parser.add_argument('--num-observer-update-epochs',
                        default=4, type=int, metavar='N',
                        help='number of total epochs to update observers')
    parser.add_argument('--num-batch-norm-update-epochs', default=3,
                        type=int, metavar='N',
                        help='number of total epochs to update batch norm stats')
    parser.add_argument('--num-calibration-batches',
                        default=32, type=int, metavar='N',
                        help='number of batches of training set for \
                              observer calibration ')

    parser.add_argument('-j', '--workers', default=16, type=int, metavar='N',
                        help='number of data loading workers (default: 16)')
    parser.add_argument('--lr',
                        default=0.0001, type=float,
                        help='initial learning rate')
    parser.add_argument('--momentum',
                        default=0.9, type=float, metavar='M',
                        help='momentum')
    parser.add_argument('-j', '--workers', default=16, type=int, metavar='N',
                        help='number of data loading workers (default: 16)')
    parser.add_argument('--lr',
                        default=0.0001, type=float,
                        help='initial learning rate')
    parser.add_argument('--momentum',
                        default=0.9, type=float, metavar='M',
                        help='momentum')
    parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
                        metavar='W', help='weight decay (default: 1e-4)',
                        dest='weight_decay')
    parser.add_argument('--lr-step-size', default=30, type=int,
                        help='decrease lr every step-size epochs')
    parser.add_argument('--lr-gamma', default=0.1, type=float,
                        help='decrease lr by a factor of lr-gamma')
    parser.add_argument('--print-freq', default=10, type=int,
                        help='print frequency')
    parser.add_argument('--output-dir', default='.', help='path where to save')
    parser.add_argument('--resume', default='', help='resume from checkpoint')
    parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                        help='start epoch')

    parser.add_argument(
        "--cache-dataset",
        dest="cache_dataset",
        help="Cache the datasets for quicker initialization. \
             It also serializes the transforms",
        action="store_true",
    )
    parser.add_argument(
        "--test-only",
        dest="test_only",
        help="Only test the model",
        action="store_true",
    )
    parser.add_argument(
        "--post-training-quantize",
        dest="post_training_quantize",
        help="Post training quantize the model",
        action="store_true",
    )
    # distributed training parameters
    parser.add_argument('--world-size', default=1, type=int,
                        help='number of distributed processes')
    parser.add_argument('--dist-url',
                        default='env://',
                        help='url used to set up distributed training')

    args = parser.parse_args()

    return args

if __name__ == "__main__":
    args = parse_args()
    main(args)
</code></pre>
</details>

raghuramank100: <p>The torchvision reference script (train_quantization.py) has not been tested for multiGPU support yet. Recently, we landed fixes to the code that should solve this issue:<br>
<aside class="onebox githubpullrequest">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/pull/33626" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Pull Request">
    <svg width="60" height="60" class="github-icon" viewbox="0 0 12 16" aria-hidden="true"><path d="M11 11.28V5c-.03-.78-.34-1.47-.94-2.06C9.46 2.35 8.78 2.03 8 2H7V0L4 3l3 3V4h1c.27.02.48.11.69.31.21.2.3.42.31.69v6.28A1.993 1.993 0 0 0 10 15a1.993 1.993 0 0 0 1-3.72zm-1 2.92c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zM4 3c0-1.11-.89-2-2-2a1.993 1.993 0 0 0-1 3.72v6.56A1.993 1.993 0 0 0 2 15a1.993 1.993 0 0 0 1-3.72V4.72c.59-.34 1-.98 1-1.72zm-.8 10c0 .66-.55 1.2-1.2 1.2-.65 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/pull/33626" target="_blank" rel="nofollow noopener">[quant] Regsiter fake_quant and observer attributes as buffers</a>
    </h4>

    <div class="branches">
      <code>pytorch:gh/supriyar/56/base</code> ← <code>pytorch:gh/supriyar/56/head</code>
    </div>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-02-21" data-time="19:16:06" data-timezone="UTC">07:16PM - 21 Feb 20 UTC</span>
      </div>

      <div class="user">
        <a href="https://github.com/supriyar" target="_blank" rel="nofollow noopener">
          <img alt="supriyar" src="https://avatars0.githubusercontent.com/u/2657489?v=4" class="onebox-avatar-inline" width="20" height="20">
          supriyar
        </a>
      </div>

      <div class="lines" title="5 commits changed 2 files with 18 additions and 13 deletions">
        <a href="https://github.com/pytorch/pytorch/pull/33626/files" target="_blank" rel="nofollow noopener">
          <span class="added">+18</span>
          <span class="removed">-13</span>
        </a>
      </div>
    </div>

  </div>
</div>
  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>Note that syncBN is not yet supported for quantization aware training.</p>

AliceLeeHX: <p>I have a try with the changed files, but it remains a bug.</p>
<pre><code class="lang-auto">    dist._broadcast_coalesced(self.process_group, tensors, buffer_size)
RuntimeError: Tensors must be CUDA and dense
</code></pre>

slippers_HUANG: <p>Hit the same issue by using pytorch 1.5.0.</p>

raghuramank100: <p>Creating a github issue to track this problem.<br>
<aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/issues/37270" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/37270" target="_blank" rel="nofollow noopener">Broadcasting does not work for Quantization aware training with multiple GPUs</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-04-25" data-time="00:44:07" data-timezone="UTC">12:44AM - 25 Apr 20 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/raghuramank100" target="_blank" rel="nofollow noopener">
          <img alt="raghuramank100" src="https://avatars1.githubusercontent.com/u/48835916?v=4" class="onebox-avatar-inline" width="20" height="20">
          raghuramank100
        </a>
      </div>
    </div>
  </div>
</div>

<div class="github-row">
  <p class="github-content">Repro code and error info are at:
https://discuss.pytorch.org/t/quantization-awareness-training-multi-gpu-suport/66106
Snippet of error at:
Traceback (most recent call last):
File “train_quantization.py”, line 258, in
main(args)
File “train_quantization.py”, line 77,...</p>
</div>

<div class="labels">
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">quantization</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">triaged</span>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>

robotcator123: <p>you can try to set the default value of scale and zero point, because it can not broadcast none tensor.</p>

robotcator123: <p>Great to see this merge request.</p>

slippers_HUANG: <p>I have tried to set default values of scale, zero point, quant_min and quant_max, and I could see the same error:<br>
“dist._broadcast_coalesced(self.process_group, tensors, buffer_size)<br>
RuntimeError: Tensors must be CUDA and dense”.</p>

robotcator123: <p>what’s the default value you set?</p>

slippers_HUANG: <p>scale = torch.FloatTensor([1])</p>
<p>zero_point = torch.FloatTensor([0])<br>
min_val = torch.FloatTensor([0])</p>
<p>max_val = torch.FloatTensor([255])</p>
<p>robotcator via PyTorch Forums &lt;<a href="mailto:noreply@discuss.pytorch.org">noreply@discuss.pytorch.org</a>&gt; 於 2020年5月20日 週三 上午11:05寫道：</p>

raghuramank100: <p>We have now added multi-GPU support for Quantization aware training in the nightly build, let us know if you see any issues</p>

slippers_HUANG: <p>Get the information “There is still work to do on verifying that BN is working correctly in<br>
QAT + DDP, but saving that for a separate PR.” from <a href="https://github.com/pytorch/vision/pull/2230" rel="nofollow noopener">https://github.com/pytorch/vision/pull/2230</a>.<br>
Could you provide the PR for tracking? Thanks.</p>

Vasiliy_Kuznetsov: <p>the PR to make BN work correctly with QAT+DDP is here: <a href="https://github.com/pytorch/pytorch/pull/38478" rel="nofollow noopener">https://github.com/pytorch/pytorch/pull/38478</a> .  This enables SyncBatchNorm to be swapped in to a fused QAT Conv-BN.  I will update the issue.  There were also a couple of bug fixes landed, such as <a href="https://github.com/pytorch/pytorch/pull/38368" rel="nofollow noopener">https://github.com/pytorch/pytorch/pull/38368</a></p>