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

[Tensor] init ColoParameter (#914)

Browse Source
This commit is contained in:
Jiarui Fang
2022-05-06 12:57:14 +08:00
committed by GitHub
parent 193d629311
commit ab95ec9aea
6 changed files with 77 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
colossalai/utils/model/colo_init_context.py
View File

@@ -1,6 +1,6 @@
from .utils import InsertPostInitMethodToModuleSubClasses
import torch
from colossalai.tensor import ColoTensor
from colossalai.tensor import ColoTensor, ColoParameter
import types
from torch import nn
@@ -100,10 +100,7 @@ class ColoInitContext(InsertPostInitMethodToModuleSubClasses):
tensor_detached = param.to(self._device).detach()
tensor_detached.requires_grad = requires_grad
setattr(
module, name,
ColoTensor.init_from_torch_tensor(tensor=tensor_detached,
save_payload=save_torch_payload,
is_model_data=True))
setattr(module, name,
ColoParameter.init_from_torch_tensor(tensor=tensor_detached, save_payload=save_torch_payload))
ColoModulize(module)