github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-14 05:33:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

[refactor] memory utils (#577)

Browse Source
This commit is contained in:
Jiarui Fang
2022-04-01 09:22:33 +08:00
committed by GitHub
parent 104cbbb313
commit e956d93ac2
15 changed files with 261 additions and 202 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_utils/test_commons.py
View File

@@ -1,4 +1,4 @@
from colossalai.utils.memory_utils.utils import colo_model_data_tensor_move, colo_model_data_tensor_move_inline
from colossalai.zero.shard_utils.tensor_utils import colo_model_data_tensor_move, colo_model_data_tensor_move_inline
from colossalai.utils import free_port
from colossalai.testing import rerun_on_exception
from colossalai.zero.sharded_param import ShardedTensor

2
tests/test_utils/test_tensor_move.py
View File

@@ -1,7 +1,7 @@
import pytest
from colossalai.utils.cuda import get_current_device
from colossalai.utils.memory_utils.utils import colo_model_data_tensor_move, colo_model_data_tensor_move_inline
from colossalai.zero.shard_utils.tensor_utils import colo_model_data_tensor_move, colo_model_data_tensor_move_inline
from colossalai.utils import free_port
from colossalai.zero.sharded_param import ShardedTensor
import colossalai