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

[Tensor] get named parameters for model using ColoTensors (#874)

Browse Source
This commit is contained in:
Jiarui Fang
2022-04-26 14:08:01 +08:00
committed by GitHub
parent 2883040286
commit e43f83aa5c
3 changed files with 59 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tests/test_tensor/test_net_tp.py
View File

@@ -7,6 +7,7 @@ from colossalai.testing import parameterize, rerun_if_address_is_in_use
from colossalai.utils.cuda import get_current_device
from colossalai.utils import free_port
from colossalai.utils import ColoInitContext
from colossalai.tensor import named_params_with_colotensor
from functools import partial
@@ -19,6 +20,8 @@ def run_simple_net():
with ColoInitContext(device=get_current_device()):
model = model_builder(checkpoint=True)
for param in named_params_with_colotensor(model):
print(param)
# we set the Specs for weight of each linear.
# model.proj1.weight.set_spec('1Drow')
# model.proj2.weight.set_spec('1Drow')