[inference] streaming Linear 1D Row inference (#1874)

tests/test_fx/test_complete_workflow.py

@@ -32,7 +32,7 @@ class MLP(torch.nn.Module):
         return x
 
 
-def run_workflow(world_size):
+def run_workflow(world_size, dev):
     # initailization
     with LazyInitContext() as ctx:
         model = MLP(16)
@@ -46,7 +46,7 @@ def run_workflow(world_size):
     gm = torch.fx.GraphModule(model, graph, model.__class__.__name__)
 
     # annotate
-    annotated_gm = transformer_mlp_pass(gm, process_group=ProcessGroup())
+    annotated_gm = transformer_mlp_pass(gm, process_group=ProcessGroup(tp_degree=world_size))
     annotated_gm.recompile()
 
     # materialization and sharding
@@ -61,22 +61,25 @@ def run_workflow(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)
+    data = torch.rand(4, 16, device=dev)
     non_fx_out = model(data)
     fx_out = annotated_gm(data)
     assert torch.equal(non_fx_out, fx_out), f'{non_fx_out} vs {fx_out}'
 
 
-def run_dist(rank, world_size, port):
+def run_dist(rank, world_size, dev, port):
     colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
-    run_workflow(world_size)
+    run_workflow(world_size, dev)
 
 
 @pytest.mark.dist
 @pytest.mark.parametrize('world_size', [1, 2])
+@pytest.mark.parametrize('dev', ['cuda', 'cpu'])
 @rerun_if_address_is_in_use()
-def test_complete_workflow(world_size):
-    run_func = partial(run_dist, world_size=world_size, port=free_port())
+def test_complete_workflow(world_size, dev):
+    if dev == 'cpu' and world_size > 1:
+        return
+    run_func = partial(run_dist, world_size=world_size, dev=dev, port=free_port())
     mp.spawn(run_func, nprocs=world_size)
 
 

tests/test_layers/test_1d/test_1d.py

@@ -1,46 +1,49 @@
-#!/usr/bin/env python
-# -*- encoding: utf-8 -*-
-
-from functools import partial
-
-import pytest
-import torch
-import torch.multiprocessing as mp
-from colossalai.core import global_context as gpc
-from colossalai.logging import disable_existing_loggers
-from colossalai.initialize import launch
-from colossalai.utils import free_port
-from colossalai.testing import rerun_if_address_is_in_use
-from checks_1d.check_layer_1d import *
-
-CONFIG = dict(parallel=dict(pipeline=dict(size=1), tensor=dict(size=4, mode='1d')),)
-
-
-def check_layer(rank, world_size, port):
-    disable_existing_loggers()
-    launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
-
-    check_linear_col()
-    check_linear_row()
-    check_embed()
-    check_vocab_parallel_embed()
-    check_classifier_no_given_weight()
-    check_vocab_parallel_classifier_no_given_weight()
-    check_classifier_given_embed_weight()
-    check_vocab_parallel_classifier_given_embed_weight()
-    check_vocab_parallel_loss()
-
-    gpc.destroy()
-    torch.cuda.empty_cache()
-
-
-@pytest.mark.dist
-@rerun_if_address_is_in_use()
-def test_1d():
-    world_size = 4
-    run_func = partial(check_layer, world_size=world_size, port=free_port())
-    mp.spawn(run_func, nprocs=world_size)
-
-
-if __name__ == '__main__':
-    test_1d()
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+from functools import partial
+
+import pytest
+import torch
+import torch.multiprocessing as mp
+from checks_1d.check_layer_1d import *
+
+from colossalai.core import global_context as gpc
+from colossalai.initialize import launch
+from colossalai.logging import disable_existing_loggers
+from colossalai.testing import rerun_if_address_is_in_use
+from colossalai.utils import free_port
+
+CONFIG = dict(parallel=dict(pipeline=dict(size=1), tensor=dict(size=4, mode='1d')),)
+
+
+def check_layer(rank, world_size, port):
+    disable_existing_loggers()
+    launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+
+    check_linear_col()
+    check_linear_row()
+    check_embed()
+    check_vocab_parallel_embed()
+    check_classifier_no_given_weight()
+    check_vocab_parallel_classifier_no_given_weight()
+    check_classifier_given_embed_weight()
+    check_vocab_parallel_classifier_given_embed_weight()
+    check_vocab_parallel_loss()
+
+    check_linear_row_stream_inference()
+
+    gpc.destroy()
+    torch.cuda.empty_cache()
+
+
+@pytest.mark.dist
+@rerun_if_address_is_in_use()
+def test_1d():
+    world_size = 4
+    run_func = partial(check_layer, world_size=world_size, port=free_port())
+    mp.spawn(run_func, nprocs=world_size)
+
+
+if __name__ == '__main__':
+    test_1d()