github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-23 18:39:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

[doc] added documentation to chunk and chunk manager (#1094)

Browse Source 
* [doc] added documentation to chunk and chunk manager

* polish code

* polish code

* polish code
This commit is contained in:
Frank Lee
2022-06-10 15:33:06 +08:00
committed by GitHub
parent 1f894e033f
commit cb18922c47
3 changed files with 215 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -120,7 +120,7 @@ class ColoDDPV2(ColoDDP):
def _setup_grads_ptr(self):
for p in self.module.parameters():
if self.chunk_manager.get_chunk(p).is_free or not p.requires_grad:
if self.chunk_manager.get_chunk(p).is_empty or not p.requires_grad:
p.grad = None
else:
p.grad = p.data
@@ -154,7 +154,7 @@ class ColoDDPV2(ColoDDP):
chunk = self.chunk_manager.get_chunk(p)
reduced = self.chunk_manager.reduce_chunk(chunk)
self.chunk_manager.release_chunk(chunk)
if reduced and not chunk.is_free:
if reduced and not chunk.is_empty:
self.overflow_counter += chunk.has_inf_or_nan
self.chunk_manager.move_chunk(chunk, self.grads_device[p])
return empty_grad