Adding some basic unit tests

.gitlab-ci.yml

 image: gitlab-master.nvidia.com/dl/dgx/pytorch:21.12-py3-devel
 
+[runners.docker]
+    gpus = "all"
+
 test:
   tags:
     - docker

tests/test_parallel_state.py

@@ -29,7 +29,7 @@ def test_initialize_model_parallel():
     with pytest.raises(RuntimeError):
         assert(ps.initialize_model_parallel(pipeline_model_parallel_size=world_size, tensor_model_parallel_size=world_size))
     with pytest.raises(RuntimeError):
-        assert(ps.initialize_model_parallel(virtual_pipeline_model_parallel_size=2))
+        assert(ps.initialize_model_parallel(virtual_pipeline_model_parallel_size=world_size))
     ps.initialize_model_parallel()
 
     assert(ps.model_parallel_is_initialized())
@@ -112,12 +112,6 @@ def test_get_tensor_model_parallel_src_rank():
     assert(ps.get_tensor_model_parallel_src_rank() == ((rank // world_size) * world_size))
     ps.destroy_model_parallel()
 
-def test_global_memory_buffer():
-    ps._GLOBAL_MEMORY_BUFFER = None
-    ps._set_global_memory_buffer()
-    assert(ps.get_global_memory_buffer() is not None)
-
-
 """
 
 def test_get_virtual_pipeline_model_parallel_world_size():