It looks like this approach results in an error indicating that some ops used in all_gather are not implemented for quantized tensors, although we have built out prototype APIs for some quantized collective communication. See:

  

      github.com
  

  
    pytorch/pytorch/blob/master/test/distributed/algorithms/quantization/test_quantization.py


    # Owner(s): ["oncall: distributed"]

import torch
import os
import torch.cuda
import sys
import torch.distributed as dist
import torch.distributed.algorithms.quantization.quantization as quant
from torch.distributed.algorithms.quantization.quantization import DQuantType
from torch.testing._internal.common_distributed import (
    MultiProcessTestCase,
    init_multigpu_helper,
    requires_gloo,
    skip_if_rocm,
    skip_if_lt_x_gpu,
    requires_nccl,
)
from torch.testing._internal.common_utils import sandcastle_skip_if, run_tests, TEST_WITH_DEV_DBG_ASAN, NO_MULTIPROCESSING_SPAWN

torch.backends.cuda.matmul.allow_tf32 = False



  This file has been truncated. show original

  

  
    
    
  

  


  

      github.com
  

  
    pytorch/pytorch/blob/master/torch/distributed/algorithms/quantization/quantization.py


    import functools
import torch
import torch.distributed as dist


from enum import Enum


TORCH_HALF_MIN = torch.finfo(torch.float16).min
TORCH_HALF_MAX = torch.finfo(torch.float16).max

class DQuantType(Enum):
    """
    Different quantization methods for auto_quantize API are identified here.
    auto_quantize API currently supports fp16 and bfp16 methods.
    """
    FP16 = "fp16",
    BFP16 = "bfp16"

    def __str__(self) -> str:



  This file has been truncated. show original

  

  
    
    
  

  


Although please note that these APIs are prototype and may thus have some issues if you decide to give it a try. cc @wanchaol