[BUG] qgZ doesn't work for odd number of nodes

Describe the bug

If we have odd number of nodes (e.g. 3 nodes with 2 gpus each), 2-stage all-to-all runs into size error when reducing final_output

Error Message:

  File "/home/jobuser/DeepSpeed/deepspeed/runtime/comm/coalesced_collectives.py", line 145, in all_to_all_quant_reduce
    output_lst[idx] = (sum(list(final_output.chunk(num_nodes))) / num_nodes).view(-1).narrow(0, 0, parition_size)
RuntimeError: The size of tensor a (171) must match the size of tensor b (170) at non-singleton dimension 0

To Reproduce
Steps to reproduce the behavior:

1. train.py

# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import os
import json
import argparse
import torch
import deepspeed
from torch.utils.data.distributed import DistributedSampler
import deepspeed.comm as dist


class SimpleModel(torch.nn.Module):

    def __init__(self, hidden_dim, empty_grad=False):
        super(SimpleModel, self).__init__()
        self.linear = torch.nn.Linear(hidden_dim, hidden_dim, bias=True)
        self.linear = torch.nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.cross_entropy_loss = torch.nn.CrossEntropyLoss()

    def forward(self, x, y):
        hidden = x
        hidden1 = self.linear(hidden)
        hidden2 = self.linear(hidden1)
        return self.cross_entropy_loss(hidden2, y)

h


def get_data_loader(model, total_samples, hidden_dim, device):
    batch_size = model.train_micro_batch_size_per_gpu()
    train_data = torch.randn(total_samples, hidden_dim, device=device, dtype=torch.half)
    train_label = torch.empty(total_samples, dtype=torch.long, device=device).random_(hidden_dim)
    train_dataset = torch.utils.data.TensorDataset(train_data, train_label)
    sampler = DistributedSampler(train_dataset)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, sampler=sampler)
    return train_loader

def print0(msg):
    if dist.get_rank() == 0:
        print(msg, flush=True)


rank = int(os.environ['RANK'])
print('seed:', 2222 + rank)
torch.random.manual_seed(2222 + rank)

# Argument Parser
parser = argparse.ArgumentParser(description='Dummy Training')
parser = deepspeed.add_config_arguments(parser)

cmd_args = parser.parse_args()


hidden_dim = 32

model = SimpleModel(hidden_dim, empty_grad=False)

model, _, _, _ = deepspeed.initialize(args=cmd_args,
                                      model=model,
                                      model_parameters=model.parameters(),
                                      dist_init_required=True)


def print_params(tag, model):
    if dist.get_rank() == 0:
        for n, p in model.named_parameters():
            print0("{} {}:{}".format(tag, n, p))


data_loader = get_data_loader(model=model, total_samples=256*32, hidden_dim=hidden_dim, device=model.device)
#print_params('pre-train', model)

for n, batch in enumerate(data_loader):
    loss = model(batch[0], batch[1])
    # if dist.get_rank() == 0:
    print("[{}] LOSS: {}".format(dist.get_rank(), loss.item()))
    model.backward(loss)
    model.step()
    #print_params('step={}'.format(n), model)
    #if n == 5: break

2. ds config

{
    "train_batch_size": 252,
    "steps_per_print": 1,
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 0.00015
        }
    },
    "fp16": {
        "enabled": true,
        "initial_scale_power": 8
    },
    "zero_optimization": {
        "stage": 3,
        "reduce_bucket_size": 20,
        "zero_hpz_partition_size": 1,
        "reduce_scatter": true,
        "zero_quantized_weights": false,
        "zero_quantized_gradients": true
    }
}

3. command

torchrun --nnodes 3 --nproc-per-node 2 --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" --rdzv_id=1234 --rdzv_backend=c10d  src/eztrain.py --deepspeed_config src/ez.json

System info (please complete the following information):

• GPU count and types: 3 machines with 2 A100 each

Analysis

1. quantization kernel is able to take any tensor size

Randomly set num_nodes, local_world_size, and dims of the tensor, then run quantize. There are no cuda errors.

def test_swizzle_quant():
    quantizer_module = op_builder.QuantizerBuilder().load()
    

    for _ in range(100):    
        num_nodes = random.randint(1, 10)
        local_world_size = random.randint(1, 10)
        global_world_size = local_world_size * num_nodes

        dim_0 = random.randint(20, 40)
        dim_1 = random.randint(20, 40)

        tensor = torch.ones((dim_0, dim_1), dtype=torch.float16, device=get_accelerator().current_device_name())
        
        intra_quant_group = max(tensor.shape[0], tensor.shape[1], global_world_size)
        inter_quant_group = intra_quant_group // local_world_size
        
        intra_quant_int4, intra_q_scales = quantizer_module.swizzle_quant(tensor, intra_quant_group, 4,
                                                                            quantizer_module.Symmetric, 1, num_nodes,
                                                                            local_world_size)

        local_output = torch.empty_like(intra_quant_int4)
        scale_output = torch.empty_like(intra_q_scales)

        global_input_tensor, global_scales = quantizer_module.quantized_reduction(
            local_output, scale_output, intra_quant_group, inter_quant_group, 4, quantizer_module.Symmetric,
            local_world_size)
        
        global_output = torch.empty_like(global_input_tensor)
        global_scale_output = torch.empty_like(global_scales)

        final_output = quantizer_module.dequantize(global_output, global_scale_output, global_scale_output.numel(),
                                                    4, quantizer_module.Symmetric)

$ pytest /home/jobuser/DeepSpeed/tests/unit/runtime/comm/test_coalesced_collectives.py::test_swizzle_quant -s
============================================================================================= slowest durations =============================================================================================
1.67s call     unit/runtime/comm/test_coalesced_collectives.py::test_swizzle_quant

(2 durations < 1s hidden.  Use -vv to show these durations.)
======================================================================================= 1 passed, 5 warnings in 6.11s =======================================================================================

5. all_to_all_single will error out if tensor_size % world_size != 0 with gloo backend
   Run 3 process, and all-to-all size 256 tensor

import torch
import torch.distributed as dist
import os

def run(rank, world_size):
    """ Distributed function performing all-to-all operation. """
    # Initialize the process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    dim = 256

    # Create a tensor for this rank
    tensor = torch.arange(dim, dtype=torch.int8)
    print(f"Rank {rank} has tensor {tensor.size()}")

    rank = dist.get_rank()
    world_size = dist.get_world_size()

    # Prepare a list of tensors to gather results
    gather_list = torch.empty(dim, dtype=torch.int8)

    # All-to-all communication

    dist.all_to_all_single(gather_list, tensor)

    # Print gathered results in the first rank

    print("[rank {}] All-to-all gathered tensors: {}".format(rank, gather_list.size()))

    # Clean up
    dist.destroy_process_group()

def main():
    world_size = 3
    # Set the MASTER_ADDR and MASTER_PORT environment variables
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12345'
    
    # Spawn the processes
    torch.multiprocessing.spawn(run,
                                args=(world_size,),
                                nprocs=world_size,
                                join=True)

if __name__ == "__main__":
    main()

$ python src/all_to_all_gloo.py 
Rank 2 has tensor torch.Size([256])
Rank 1 has tensor torch.Size([256])
Rank 0 has tensor torch.Size([256])
Traceback (most recent call last):
  File "/home/jobuser/src/all_to_all_gloo.py", line 49, in <module>
    main()
  File "/home/jobuser/src/all_to_all_gloo.py", line 43, in main
    torch.multiprocessing.spawn(run,
  File "/home/jobuser/.local/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 239, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "/home/jobuser/.local/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 197, in start_processes
    while not context.join():
  File "/home/jobuser/.local/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 160, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException: 

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/home/jobuser/.local/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/home/jobuser/src/all_to_all_gloo.py", line 27, in run
    dist.all_to_all_single(gather_list, tensor)
  File "/home/jobuser/.local/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 1436, in wrapper
    return func(*args, **kwargs)
  File "/home/jobuser/.local/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 3151, in all_to_all_single
    work.wait()
RuntimeError: [enforce fail at ../third_party/gloo/gloo/alltoall.cc:29] in->size % context->size == 0. 

6. all_to_all_single will not error out if tensor_size % world_size != 0 with nccl backend. I suspect something unexpected might happen silently in memory.

Run 2 processes, and all-to-all size 5 tensor. it should have unbalanced error but it doesn't.

import torch
import torch.distributed as dist
import os

def init_process():
    """Initialize the distributed environment."""
    dist.init_process_group(backend='nccl')

def all_to_all_example():
    """Perform a dummy all-to-all operation."""
    rank = dist.get_rank()
    size = dist.get_world_size()
    tensor_size = 5

    gpu_id = rank % torch.cuda.device_count()
    device = torch.device(f'cuda:{gpu_id}')

    # Create a tensor filled with the rank number for demonstration
    send_tensor = torch.arange(tensor_size, dtype=torch.int8).to(device)
    
    # Prepare a tensor for receiving data from other processes
    recv_tensor = torch.empty(tensor_size, dtype=torch.int8).to(device)
    
    # Perform the all-to-all operation
    dist.all_to_all_single(recv_tensor, send_tensor)

    print(f"Process {rank} received tensor: {recv_tensor}")

def main():
    init_process()
    all_to_all_example()
    dist.destroy_process_group()

if __name__ == "__main__":
    main()

$ torchrun --nnodes 1 --nproc-per-node 2 --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" --rdzv_id=1234 --rdzv_backend=c10d  src/all_to_all_nccl.py 
master_addr is only used for static rdzv_backend and when rdzv_endpoint is not specified.
WARNING:torch.distributed.run:
*****************************************
Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
*****************************************
Process 0 received tensor: tensor([ 0,  1,  0,  1, 29], device='cuda:0', dtype=torch.int8)
Process 1 received tensor: tensor([  2,   3,   2,   3, -51], device='cuda:1', dtype=torch.int8)

7. The tensor should be padded to be divisible by global_world_size * 2 according to the following calculation in all_to_all_quant_reduce.

   a. The initial tensor is of size (dim_1, dim_2, ... dim_n), and numel is A (=dim_1*dim_2...*dim_n).

   b. After swizzle_quant, intra_quant_int4 size is (A // 2) if A % 2 == 0. the tensor is quantized from fp16/bf16 to int4. However, intra_quant_int4 is actually represented by int8, which means every two int4 tensors is grouped into one and stored in int8 format. Note that if A % 2 != 0, the quantization can still process, but the size differ by cases. I still cannot find the underlying rule)

   c. At all_to_all_single(local_output, intra_quant_int4, group=groups[f'local_{intra_idx}']), we should assert that intra_quant_int4 % local_world_size == 0, which means (A // 2) % local_world_size == 0.

   d. At quantized_reduction, intra_quant_int4 is chunked to local_world_size pieces and reduce them together.

   e. global_input_tensor is of size A // (2 * local_world_size) after reduction

   f. At all_to_all_single(global_output, global_input_tensor, group=groups[f'global_{inter_idx}']), we should assert that global_input_tensor % n_nodes == 0, which means ( A // (2 * local_world_size) ) % n_nodes == 0

Start in deepspeed/runtime/comm/coalesced_collectives.py at all_to_all_quant_reduce, especially the final_output reduction around the reported line. Reproduce with the three-node, two-GPU command and review tests/unit/runtime/comm/test_coalesced_collectives.py::test_swizzle_quant. Done means the odd-node two-stage all-to-all no longer raises the reported tensor-size error.