diff --git a/colossalai/utils/memory.py b/colossalai/utils/memory.py index 799e68f4d..12c23c3a0 100644 --- a/colossalai/utils/memory.py +++ b/colossalai/utils/memory.py @@ -8,6 +8,7 @@ from colossalai.utils import get_current_device from colossalai.core import global_context as gpc from colossalai.context.parallel_mode import ParallelMode from colossalai.logging import get_dist_logger +from packaging import version _GLOBAL_CUDA_MEM_FRACTION = 1.0 @@ -106,7 +107,8 @@ def colo_device_memory_capacity(device: torch.device) -> int: assert isinstance(device, torch.device) if device.type == 'cpu': mem_info = _get_cpu_memory_info() - return mem_info.info.total / gpc.get_world_size(ParallelMode.DATA) + # In the context of 1-CPU-N-GPU, the memory capacity of the current process is 1/N overall CPU memory. + return mem_info.total / gpc.num_processes_on_current_node if device.type == 'cuda': return torch.cuda.get_device_properties(get_current_device()).total_memory * _GLOBAL_CUDA_MEM_FRACTION @@ -123,8 +125,9 @@ def colo_device_memory_used(device: torch.device) -> int: """ if device.type == 'cpu': mem_info = _get_cpu_memory_info() - # FIXME(jiaruifang) we need get how many processes are using the CPU memory. - ret = mem_info.used / gpc.get_world_size(ParallelMode.DATA) + # In the context of 1-CPU-N-GPU, the memory usage of the current process is 1/N CPU memory used. + # Each process consumes the same amount of memory. + ret = mem_info.used / gpc.num_processes_on_current_node return ret elif device.type == 'cuda': ret: int = torch.cuda.memory_allocated(device) @@ -142,6 +145,10 @@ def colo_set_process_memory_fraction(ratio: float) -> None: Args: ratio (float): a ratio between 0. ~ 1. """ + if version.parse(torch.__version__) < version.parse('1.8'): + logger = get_dist_logger('colo_set_process_memory_fraction') + logger.warning('colo_set_process_memory_fraction failed because torch version is less than 1.8') + return global _GLOBAL_CUDA_MEM_FRACTION _GLOBAL_CUDA_MEM_FRACTION = ratio torch.cuda.set_per_process_memory_fraction(_GLOBAL_CUDA_MEM_FRACTION, get_current_device()) diff --git a/colossalai/utils/memory_tracer/async_memtracer.py b/colossalai/utils/memory_tracer/async_memtracer.py index 3be66cd9c..4442a9e51 100644 --- a/colossalai/utils/memory_tracer/async_memtracer.py +++ b/colossalai/utils/memory_tracer/async_memtracer.py @@ -31,6 +31,7 @@ class AsyncMemoryMonitor: async_mem_monitor.finish() async_mem_monitor.save('log.pkl') + Args: power (int, optional): the power of time interva. Defaults to 10. diff --git a/tests/test_moe/test_moe_zero_init.py b/tests/test_moe/test_moe_zero_init.py index 5a5bee27b..7c5f30b14 100644 --- a/tests/test_moe/test_moe_zero_init.py +++ b/tests/test_moe/test_moe_zero_init.py @@ -16,7 +16,7 @@ from colossalai.zero.shard_utils import (BucketTensorShardStrategy, TensorShardS from colossalai.testing import rerun_on_exception from colossalai.utils import get_current_device -from tests.test_zero_data_parallel.common import CONFIG +from tests.test_zero.common import CONFIG class MoeModel(CheckpointModule): diff --git a/tests/test_moe/test_moe_zero_model.py b/tests/test_moe/test_moe_zero_model.py index a2268ea5c..af78c032c 100644 --- a/tests/test_moe/test_moe_zero_model.py +++ b/tests/test_moe/test_moe_zero_model.py @@ -16,7 +16,7 @@ from colossalai.engine.gradient_handler import MoeGradientHandler from colossalai.context import MOE_CONTEXT from colossalai.testing import assert_equal_in_group -from tests.test_zero_data_parallel.common import CONFIG, check_grads_padding, run_fwd_bwd +from tests.test_zero.common import CONFIG, check_grads_padding, run_fwd_bwd from tests.test_moe.test_moe_zero_init import MoeModel diff --git a/tests/test_moe/test_moe_zero_optim.py b/tests/test_moe/test_moe_zero_optim.py index 8fbf655f8..1ae25997b 100644 --- a/tests/test_moe/test_moe_zero_optim.py +++ b/tests/test_moe/test_moe_zero_optim.py @@ -20,7 +20,7 @@ from colossalai.engine.gradient_handler import MoeGradientHandler from colossalai.context import MOE_CONTEXT from colossalai.testing import assert_equal_in_group -from tests.test_zero_data_parallel.common import CONFIG, check_sharded_model_params +from tests.test_zero.common import CONFIG, check_sharded_model_params from tests.test_moe.test_moe_zero_init import MoeModel diff --git a/tests/test_utils/test_memory.py b/tests/test_utils/test_memory.py new file mode 100644 index 000000000..a6ff7c9a1 --- /dev/null +++ b/tests/test_utils/test_memory.py @@ -0,0 +1,32 @@ +import pytest + +import colossalai +from colossalai.utils.cuda import get_current_device +from colossalai.utils.memory import colo_set_process_memory_fraction, colo_device_memory_capacity +from colossalai.utils import free_port + +from functools import partial +import torch.multiprocessing as mp + + +def _run_colo_set_process_memory_fraction_and_colo_device_memory_capacity(): + frac1 = colo_device_memory_capacity(get_current_device()) + colo_set_process_memory_fraction(0.5) + frac2 = colo_device_memory_capacity(get_current_device()) + assert frac2 * 2 == frac1 + + +def run_dist(rank, world_size, port): + colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + _run_colo_set_process_memory_fraction_and_colo_device_memory_capacity() + + +@pytest.mark.dist +@pytest.mark.parametrize("world_size", [4, 5]) +def test_memory_utils(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_memory_utils(world_size=2) diff --git a/tests/test_zero_data_parallel/common.py b/tests/test_zero/common.py similarity index 100% rename from tests/test_zero_data_parallel/common.py rename to tests/test_zero/common.py diff --git a/tests/test_zero_data_parallel/test_found_inf.py b/tests/test_zero/test_found_inf.py similarity index 95% rename from tests/test_zero_data_parallel/test_found_inf.py rename to tests/test_zero/test_found_inf.py index 6bdd667e4..fc02bb67d 100644 --- a/tests/test_zero_data_parallel/test_found_inf.py +++ b/tests/test_zero/test_found_inf.py @@ -14,7 +14,7 @@ from colossalai.zero.sharded_model import ShardedModelV2 from colossalai.zero.sharded_optim import ShardedOptimizerV2 from colossalai.zero.sharded_optim._utils import has_inf_or_nan from tests.components_to_test.registry import non_distributed_component_funcs -from tests.test_zero_data_parallel.test_sharded_optim_v2 import _run_step +from tests.test_zero.test_sharded_optim_v2 import _run_step from common import CONFIG diff --git a/tests/test_zero_data_parallel/test_init_context.py b/tests/test_zero/test_init_context.py similarity index 100% rename from tests/test_zero_data_parallel/test_init_context.py rename to tests/test_zero/test_init_context.py diff --git a/tests/test_zero_data_parallel/test_shard_model_v2.py b/tests/test_zero/test_shard_model_v2.py similarity index 100% rename from tests/test_zero_data_parallel/test_shard_model_v2.py rename to tests/test_zero/test_shard_model_v2.py diff --git a/tests/test_zero_data_parallel/test_shard_param.py b/tests/test_zero/test_shard_param.py similarity index 98% rename from tests/test_zero_data_parallel/test_shard_param.py rename to tests/test_zero/test_shard_param.py index 79780fa51..e88a64c35 100644 --- a/tests/test_zero_data_parallel/test_shard_param.py +++ b/tests/test_zero/test_shard_param.py @@ -11,7 +11,7 @@ from colossalai.zero.shard_utils import (BucketTensorShardStrategy, TensorShardS from colossalai.zero.sharded_param import ShardedTensor from colossalai.zero.sharded_param.sharded_param import ShardedParamV2 from colossalai.testing import rerun_on_exception -from tests.test_zero_data_parallel.common import CONFIG, allclose +from tests.test_zero.common import CONFIG, allclose from colossalai.zero.sharded_param.tensorful_state import StatefulTensor diff --git a/tests/test_zero_data_parallel/test_sharded_optim_v2.py b/tests/test_zero/test_sharded_optim_v2.py similarity index 100% rename from tests/test_zero_data_parallel/test_sharded_optim_v2.py rename to tests/test_zero/test_sharded_optim_v2.py diff --git a/tests/test_zero_data_parallel/test_sharded_optim_with_sync_bn.py b/tests/test_zero/test_sharded_optim_with_sync_bn.py similarity index 100% rename from tests/test_zero_data_parallel/test_sharded_optim_with_sync_bn.py rename to tests/test_zero/test_sharded_optim_with_sync_bn.py diff --git a/tests/test_zero_data_parallel/test_state_dict.py b/tests/test_zero/test_state_dict.py similarity index 100% rename from tests/test_zero_data_parallel/test_state_dict.py rename to tests/test_zero/test_state_dict.py diff --git a/tests/test_zero_data_parallel/test_stateful_tensor_mgr.py b/tests/test_zero/test_stateful_tensor_mgr.py similarity index 100% rename from tests/test_zero_data_parallel/test_stateful_tensor_mgr.py rename to tests/test_zero/test_stateful_tensor_mgr.py diff --git a/tests/test_utils/test_tensor_move.py b/tests/test_zero/test_tensor_utils.py similarity index 86% rename from tests/test_utils/test_tensor_move.py rename to tests/test_zero/test_tensor_utils.py index cf1677d4b..d3b50d2dc 100644 --- a/tests/test_utils/test_tensor_move.py +++ b/tests/test_zero/test_tensor_utils.py @@ -5,7 +5,6 @@ from colossalai.utils.cuda import get_current_device from colossalai.zero.sharded_param import (StatefulTensor, colo_tensor_mem_usage, colo_model_data_tensor_move, colo_model_data_tensor_move_inline, colo_model_data_move_to_cpu, colo_model_tensor_clone) -from colossalai.utils.memory import colo_set_process_memory_fraction, colo_device_memory_capacity from colossalai.utils import free_port import torch @@ -32,13 +31,6 @@ def _run_colo_tensor_mem_usage(): assert g1 * 4 == g2 -def _run_colo_set_process_memory_fraction_and_colo_device_memory_capacity(): - frac1 = colo_device_memory_capacity(get_current_device()) - colo_set_process_memory_fraction(0.5) - frac2 = colo_device_memory_capacity(get_current_device()) - assert frac2 * 2 == frac1 - - def _run_colo_model_data_tensor_move_inline(): for t in [StatefulTensor(torch.randn(2, 3)), torch.randn(2, 3)]: colo_model_data_tensor_move_inline(t, get_current_device()) @@ -82,20 +74,20 @@ def _run_colo_model_tensor_clone(): def run_dist(rank, world_size, port): colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') - _run_colo_set_process_memory_fraction_and_colo_device_memory_capacity() + + _run_colo_tensor_mem_usage() _run_colo_model_data_tensor_move_inline() _run_colo_model_data_tensor_move() - _run_colo_tensor_mem_usage() _run_colo_model_data_move_to_cpu() _run_colo_model_tensor_clone() @pytest.mark.dist @pytest.mark.parametrize("world_size", [4, 5]) -def test_tensor_move(world_size): +def test_zero_tensor_utils(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_tensor_move(4) + test_zero_tensor_utils(world_size=2) diff --git a/tests/test_zero_data_parallel/test_zero_engine.py b/tests/test_zero/test_zero_engine.py similarity index 100% rename from tests/test_zero_data_parallel/test_zero_engine.py rename to tests/test_zero/test_zero_engine.py