[PROPOSAL]: Addition of a SIMD function wrapper like jax.pmap

[BigBalloon8]

Proposal:

see issue #3651

Motivation

Colossal AI provides lots of ways to easily parallelize the training and deployment of a model (using data parallelism, pipeline parallelism, etc.) but training large models is not the only thing in ML and data science that could benefit from parallelization.

jax.pmap provides an easy way to parallelize your code across multiple devices. torch2.0 introduced torch.vmap which added single-device vectorization for SIMD (single input multiple data) functions like jax.vmap but nothing like jax.pmap.

Design

I propose cmap (colossal ai map).
cmap map would act as a Python wrapper around your function. It would have all the same args and kwargs as torch.vmap with the addition of a few more such as the output destination and the process group

1. cmap takes your input function along with standard vmap kwargs see torch.vmap
2. cmap then creates a wrapper function that splits the input tensors along the given dim/dims and allocates each chunk to its own device
3. then the input function is vmapped across each device's chunk in parallel.
4. the output of the vmaped chunk is either gathered on single devices within the process group or all-gathered to all the devices
5. finally the outputs are concatenated along a given dim
6. then that wrapper function is returned as the wrapped cmap function

Possible sample code

import torch
from colossalai import cmap

x = torch.randn(64, 128)
y = torch.randn(64, 128)
batch_dot = cmap(torch.dot)
batch_dot(x,y)
#---------------------------------------------

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(128, 128)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.linear(x)
        x = self.relu(x)
        return x

model = Model()
batched_model = cmap(model)
input = torch.randn(64, 128)
batched_model(input)
#---------------------------------------------

@cmap(in_dims=(0,None), dst=0)
def batch_dot(x, y):
    return torch.dot(x, y)
a = torch.randn(64, 128)
b = torch.randn(128)
out = batch_dot(a, b)  # out = None on all process except process 0
#---------------------------------------------

N = 5
f = lambda x: x ** 2
x = torch.randn(N, requires_grad=True)
y = f(x)
I_N = torch.eye(N)

# vectorized gradient computation
def get_vjp(v):
    return torch.autograd.grad(y, x, v, retain_graph=True) 
jacobian = cmap(get_vjp)(I_N)

Possible cmap code (pseudo-code)

def cmap(func: Callable,
                in_dims: Union[int, Tuple[int]] = 0,
                out_dims: Union[int, Tuple[int]] = 0,
                group = None,
                dst: int = -1 ):
    def wrapper(*args, **kwargs):
        # takes all args and returns args for current device
        args = split_input(args, in_dims)

        func_out = torch.vmap(func, in_dims, out_dims)(*args, **kwrags)

        if dst == -1:
            out = allgather(func_out, group=group)
        
        else:
            out = gather(func_out, dst=dst, group=group)
    
        return out
    return wrapper

Self-Service

• I'd be willing to do some initial work on this proposal myself.

[binmakeswell]

Hi @mayfieldmobster Thanks so much for your contribution! We will review it later. Thanks.