mirror of https://github.com/hpcaitech/ColossalAI
aibig-modeldata-parallelismdeep-learningdistributed-computingfoundation-modelsheterogeneous-traininghpcinferencelarge-scalemodel-parallelismpipeline-parallelism
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
74 lines
2.6 KiB
74 lines
2.6 KiB
import torch |
|
import colossalai |
|
import pytest |
|
import torch.multiprocessing as mp |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
from colossalai.utils.cuda import get_current_device |
|
from colossalai.utils.memory import colo_device_memory_capacity, colo_set_process_memory_fraction |
|
from colossalai.zero.init_ctx import ZeroInitContext |
|
from colossalai.zero.sharded_model import ShardedModelV2 |
|
from colossalai.zero.shard_utils import BucketTensorShardStrategy |
|
from colossalai.utils import free_port |
|
from colossalai.testing import rerun_if_address_is_in_use |
|
from functools import partial |
|
|
|
|
|
class MyTestModel(torch.nn.Module): |
|
|
|
def __init__(self) -> None: |
|
super().__init__() |
|
self.proj1 = nn.Linear(512, 512) |
|
self.weight = nn.Parameter(torch.randn(1024, 512)) |
|
self.proj2 = nn.Linear(1024, 512) |
|
|
|
def forward(self, x): |
|
x = self.proj1(x) |
|
x = F.linear(x, self.weight) |
|
x = self.proj2(x) |
|
|
|
return x |
|
|
|
|
|
def run_mem_collector_testing(): |
|
cuda_capacity = colo_device_memory_capacity(get_current_device()) |
|
fraction = (50 * 1024**2) / cuda_capacity |
|
# limit max memory to 50MB |
|
colo_set_process_memory_fraction(fraction) |
|
shard_strategy = BucketTensorShardStrategy() |
|
with ZeroInitContext(target_device=get_current_device(), shard_strategy=shard_strategy, shard_param=True): |
|
model = MyTestModel() |
|
|
|
model = ShardedModelV2(module=model, |
|
shard_strategy=shard_strategy, |
|
reduce_scatter_bucket_size_mb=1, |
|
tensor_placement_policy='auto') |
|
|
|
data = torch.randn(2, 512, device=get_current_device()) |
|
|
|
output = model(data) |
|
loss = torch.mean(output) |
|
model.backward(loss) |
|
|
|
cuda_model_data_list = model._memstats_collector.model_data_list('cuda') |
|
assert cuda_model_data_list == [1311744, 1836032, 1836032, 1311744, 1836032, 1836032] |
|
|
|
cuda_non_model_data_list = model._memstats_collector.non_model_data_list('cuda') |
|
assert cuda_non_model_data_list[0] > cuda_non_model_data_list[1] |
|
assert cuda_non_model_data_list[-2] > cuda_non_model_data_list[-1] |
|
|
|
|
|
def run_dist(rank, world_size, port): |
|
colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') |
|
run_mem_collector_testing() |
|
|
|
|
|
@pytest.mark.dist |
|
@rerun_if_address_is_in_use() |
|
def test_mem_collector(world_size=2): |
|
run_func = partial(run_dist, world_size=world_size, port=free_port()) |
|
mp.spawn(run_func, nprocs=world_size) |
|
|
|
|
|
if __name__ == '__main__': |
|
test_mem_collector()
|
|
|