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

[Tensor] activation is an attr of ColoTensor (#897)

Browse Source
This commit is contained in:
Jiarui Fang
2022-04-28 14:43:22 +08:00
committed by GitHub
parent e76f76c08b
commit 676f191532
4 changed files with 51 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -94,7 +94,10 @@ class ColoInitContext(InsertPostInitMethodToModuleSubClasses):
save_torch_payload = True if not self._lazy_memory_allocate else False
for name, param in name_list:
delattr(module, name)
setattr(module, name,
ColoTensor.init_from_torch_tensor(tensor=param.to(self._device), save_payload=save_torch_payload))
setattr(
module, name,
ColoTensor.init_from_torch_tensor(tensor=param.to(self._device),
save_payload=save_torch_payload,
is_model_data=True))
ColoModulize(module)