github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-17 15:36:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

[NFC] polish colossalai/nn/parallel/data_parallel.py code style (#1570)

Browse Source 
Co-authored-by: JThh <jiatong.han@u.nus.edu>
This commit is contained in:
Jiatong Han 2022-09-08 16:33:14 +08:00 committed by Frank Lee
parent f566c9b98d
commit 3263cdf57f
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/nn/parallel/data_parallel.py
View File

@ -350,7 +350,7 @@ class ZeroDDP(ColoDDP):
for tensor in chunk.get_tensors(): for tensor in chunk.get_tensors():
rec_p = torch.empty([0]) rec_p = torch.empty([0])
if record_flag: if record_flag:
rec_p = tensor.cpu() # move the whole tensor to CPU mem rec_p = tensor.cpu() # move the whole tensor to CPU mem
assert tensor not in param_to_save_data assert tensor not in param_to_save_data
param_to_save_data[tensor] = rec_p param_to_save_data[tensor] = rec_p
# release the actual memory of the chunk # release the actual memory of the chunk
@ -406,7 +406,7 @@ class ZeroDDP(ColoDDP):
state_dict = state_dict.copy() state_dict = state_dict.copy()
if metadata is not None: if metadata is not None:
# mypy isn't aware that "_metadata" exists in state_dict # mypy isn't aware that "_metadata" exists in state_dict
state_dict._metadata = metadata # type: ignore[attr-defined] state_dict._metadata = metadata # type: ignore[attr-defined]
prefix = '' prefix = ''
local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {}) local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})