[ADD] support Distributed Data Parallel

[jinwonkim93]

Title

Colossal AI-based Distributed Data Parallel with oslo interface

Description

The purpose of this implementation is to enable DDP in Oslo, with the reducer method being identical to that of Colossal AI, but adapted to fit Oslo's interface. To enhance user experience, we replaced model.backward() with loss.backward() and added model.zero_grad() temporarily to the code. Any feedback is welcome :)

If you don't use model.zero_grad() there will be unexpected error.

test_data_parallel.py

import os
import torch.multiprocessing as mp

import torch
from torch.nn.parallel import DistributedDataParallel as DDP
from torch import nn
from torch import optim
import torch.distributed as dist

from oslo.torch.distributed.parallel_context import ParallelContext


def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12345"
    os.environ["RANK"] = str(rank)
    os.environ["LOCAL_RANK"] = str(rank)
    os.environ["WORLD_SIZE"] = str(world_size)
    os.environ["LOCAL_WORLD_SIZE"] = str(world_size)


def cleanup():
    dist.destroy_process_group()


class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def train(rank, world_size):
    print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)
    parallel_context = ParallelContext.from_torch(data_parallel_size=world_size)

    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.zeros(20, 10).to(rank))
    labels = torch.zeros(20, 5).to(rank)
    loss_fn(outputs, labels).backward()
    optimizer.step()
    print(outputs)
    cleanup()


def main(world_size):
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)


if __name__ == "__main__":
    main(2)

test_oslo_data_parallel.py

import os
import torch.multiprocessing as mp

import torch
from torch import nn
from torch import optim
import torch.distributed as dist

import oslo
from oslo.torch.distributed.parallel_context import ParallelContext
from oslo.torch.nn.parallel.data_parallel import DistributedDataParallel as DDP


def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12345"
    os.environ["RANK"] = str(rank)
    os.environ["LOCAL_RANK"] = str(rank)
    os.environ["WORLD_SIZE"] = str(world_size)
    os.environ["LOCAL_WORLD_SIZE"] = str(world_size)


def cleanup():
    dist.destroy_process_group()


class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def train(rank, world_size):
    print(f"Running oslo DDP example on rank {rank}.")
    setup(rank, world_size)
    parallel_context = ParallelContext.from_torch(data_parallel_size=world_size)

    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, parallel_context)

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
    oslo.ready(ddp_model, parallel_context)
    optimizer.zero_grad()
    outputs = ddp_model(torch.zeros(20, 10).to(rank))
    labels = torch.zeros(20, 5).to(rank)
    loss = loss_fn(outputs, labels)
    ddp_model.backward(loss)
    optimizer.step()
    print(outputs)
    cleanup()


def main(world_size):
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)


if __name__ == "__main__":
    main(2)

pytorch DDP

Oslo DDP

By checking the model's parameters, oslo DDP is working as expected.

After Cleaning

Oslo DDP
oslo-ddp-time.log

Torch DDP
torch-ddp-time.log

Linked Issues

• resolved #00

[hyunwoongko]

Could you run precommit run --all-files.

[KKIEEK]

I think it would be better to merge _DistirbutedDataParallelWrapper class into _DistributedDataParallel.

[hyunwoongko]

please resolve conflict! @jinwonkim93

[jinwonkim93]

please resolve conflict! @jinwonkim93

Completed. One question is there a reason for "ALL" rather than "all"?

[hyunwoongko]

no. I prefer __ALL__, but we don't use both of them because we think import is enough.
Is there any file which contains __all__?

[jinwonkim93]

no. I prefer __ALL__, but we don't use both of them because we think import is enough. Is there any file which contains __all__?

oslo/oslo/torch/nn/parallel/data_parallel/zero/__init__.py

Line 5 in dcad48e

__all__ = ["ZeroRedundancyOptimizer"]

oslo/oslo/torch/nn/parallel/data_parallel/zero/sharded_optim/__init__.py

Line 5 in dcad48e

__all__ = ["ZeroRedundancyOptimizer"]

oslo/oslo/torch/utils/__init__.py

Line 5 in dcad48e

__all__ = ["get_free_port", "set_seed"]

[hyunwoongko]

I didn't add them. all the code added by DP new members. so it's okay to change them to uppercase.

[jinwonkim93]

I didn't add them. all the code added by DP new members. so it's okay to change them to uppercase.

Okay. I think it is ready to merged. what do you think?

[hyunwoongko]

@jinwonkim93 looks good to me.