github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-04 02:26:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Gemini] remove GLOBAL_MODEL_DATA_TRACER (#2091)

Browse Source
This commit is contained in:
Jiarui Fang
2022-12-06 22:30:16 +08:00
committed by GitHub
parent 28e55c2530
commit 1fca5d79ea
4 changed files with 6 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
tests/test_zero/test_init_context.py
View File

@@ -3,23 +3,22 @@
from functools import partial
import colossalai
import pytest
import torch
import torch.multiprocessing as mp
from common import CONFIG
import colossalai
from colossalai.gemini.memory_tracer.utils import colo_model_mem_usage
from colossalai.logging import get_dist_logger
from colossalai.testing import parameterize, rerun_if_address_is_in_use
from colossalai.utils import free_port
from colossalai.utils.cuda import get_current_device
from colossalai.gemini.memory_tracer.model_data_memtracer import \
colo_model_mem_usage
from colossalai.utils.memory import colo_device_memory_used
from colossalai.zero.init_ctx import ZeroInitContext
from colossalai.zero.shard_utils import (BucketTensorShardStrategy, TensorShardStrategy)
from colossalai.zero.shard_utils import BucketTensorShardStrategy, TensorShardStrategy
from tests.components_to_test.registry import non_distributed_component_funcs
from common import CONFIG
@parameterize("init_device_type", ['cpu', 'cuda'])
@parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy])