mirror of https://github.com/hpcaitech/ColossalAI
78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
|
import colossalai
|
||
|
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
import pytest
|
||
|
|
import torch.multiprocessing as mp
|
||
|
|
import torch.distributed as dist
|
||
|
|
from colossalai.testing import rerun_if_address_is_in_use
|
||
|
|
from functools import partial
|
||
|
|
from colossalai.fx import ColoTracer
|
||
|
|
from colossalai.utils.model.lazy_init_context import LazyInitContext
|
||
|
|
from colossalai.fx.passes.shard_1d_pass import transformer_mlp_pass
|
||
|
|
from colossalai.utils import free_port
|
||
|
|
from colossalai.tensor import ProcessGroup
|
||
|
|
|
||
|
|
|
||
|
|
class MLP(torch.nn.Module):
|
||
|
|
|
||
|
|
def __init__(self, dim: int):
|
||
|
|
super().__init__()
|
||
|
|
self.linear1 = torch.nn.Linear(dim, dim)
|
||
|
|
self.linear2 = torch.nn.Linear(dim, dim)
|
||
|
|
self.dropout = torch.nn.Dropout(0)
|
||
|
|
self.relu = torch.nn.ReLU()
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
x = self.linear1(x)
|
||
|
|
x = self.dropout(x)
|
||
|
|
x = self.relu(x)
|
||
|
|
x = self.linear2(x)
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
def run_workflow(world_size):
|
||
|
|
# initailization
|
||
|
|
with LazyInitContext() as ctx:
|
||
|
|
model = MLP(16)
|
||
|
|
|
||
|
|
# tracing
|
||
|
|
tracer = ColoTracer()
|
||
|
|
graph = tracer.trace(model)
|
||
|
|
gm = torch.fx.GraphModule(model, graph, model.__class__.__name__)
|
||
|
|
|
||
|
|
# annotate
|
||
|
|
annotated_gm = transformer_mlp_pass(gm, process_group=ProcessGroup())
|
||
|
|
annotated_gm.recompile()
|
||
|
|
|
||
|
|
# materialization and sharding
|
||
|
|
ctx.lazy_init_parameters(annotated_gm)
|
||
|
|
|
||
|
|
# # check sharding
|
||
|
|
assert list(model.linear1.weight.shape) == [16 // world_size, 16]
|
||
|
|
assert list(model.linear1.bias.shape) == [16 // world_size]
|
||
|
|
assert list(model.linear2.weight.shape) == [16, 16 // world_size]
|
||
|
|
|
||
|
|
# test forward to make sure that IR transform will produce the same results
|
||
|
|
# like how ColoTensor would do it normally
|
||
|
|
data = torch.rand(4, 16)
|
||
|
|
non_fx_out = model(data)
|
||
|
|
fx_out = annotated_gm(data)
|
||
|
|
assert torch.equal(non_fx_out, fx_out)
|
||
|
|
|
||
|
|
|
||
|
|
def run_dist(rank, world_size, port):
|
||
|
|
colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
|
||
|
|
run_workflow(world_size)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.dist
|
||
|
|
@pytest.mark.parametrize('world_size', [1, 2])
|
||
|
|
@rerun_if_address_is_in_use()
|
||
|
|
def test_complete_workflow(world_size):
|
||
|
|
run_func = partial(run_dist, world_size=world_size, port=free_port())
|
||
|
|
mp.spawn(run_func, nprocs=world_size)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
test_complete_workflow(2)
|