github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-15 22:19:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

[refactor] zero directory (#724)

Browse Source
This commit is contained in:
Jiarui Fang
2022-04-11 23:13:02 +08:00
committed by GitHub
parent 20ab1f5520
commit 4d90a7b513
19 changed files with 196 additions and 223 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.zero.shard_utils.tensor_utils import colo_model_data_tensor_move, colo_model_data_tensor_move_inline
from colossalai.zero.sharded_param.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

7
tests/test_utils/test_tensor_move.py
View File

@@ -1,11 +1,12 @@
import pytest
import colossalai
from colossalai.utils.cuda import get_current_device
from colossalai.zero.shard_utils.tensor_utils import 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.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
from colossalai.zero.sharded_param.tensorful_state import StatefulTensor
import colossalai
import torch