github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-03 10:06:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

[gemini] init genimi individual directory (#754)

Browse Source
This commit is contained in:
Jiarui Fang
2022-04-14 16:40:26 +08:00
committed by GitHub
parent dcca614eee
commit 10ef8afdd2
8 changed files with 14 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tests/test_zero/test_stateful_tensor_mgr.py
View File

@@ -6,7 +6,7 @@ from colossalai.utils.cuda import get_current_device
from colossalai.utils.memory_tracer import MemStatsCollector
from colossalai.utils.memory_tracer.model_data_memtracer import GLOBAL_MODEL_DATA_TRACER
from colossalai.utils.memory import colo_set_process_memory_fraction
from colossalai.zero.utils import StatefulTensorMgr
from colossalai.gemini import StatefulTensorMgr
from colossalai.zero.sharded_param.sharded_param import ShardedParamV2
from colossalai.zero.sharded_param.tensorful_state import TensorState
from colossalai.utils import free_port
@@ -14,7 +14,9 @@ from colossalai.testing import rerun_on_exception
from torch.nn.parameter import Parameter
from typing import List
from functools import partial
from colossalai.zero.utils.tensor_placement_policy import AutoTensorPlacementPolicy
from colossalai.gemini import StatefulTensorMgr
from colossalai.gemini.tensor_placement_policy import AutoTensorPlacementPolicy
class Net(torch.nn.Module):