Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[feature] A new ZeRO implementation #1644

Merged
merged 17 commits into from
Oct 9, 2022
Merged

[feature] A new ZeRO implementation #1644

Show file tree
Hide file tree
Changes from 13 commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
41121cb
[zero] use ChunkV2 for zero training
1SAA Aug 12, 2022
d5f83d7
[zero] fix cuda, cpu memory leak
1SAA Aug 15, 2022
04082a4
[zero] add zero optimizer
1SAA Aug 16, 2022
7f31381
[zero] add unit test for ZeROOptimV2
1SAA Aug 16, 2022
c01c62e
[zero] make fp32_param in pin_memory
1SAA Aug 17, 2022
9f5195b
[zero] add state_dict, load_state_dict, unit test for ZeRODDP
1SAA Sep 6, 2022
5feb621
[hotfix] fix the init device for each chunk
1SAA Sep 8, 2022
0f60d90
[hotfix] fix error for gathered chunks
1SAA Sep 8, 2022
81f33bb
[zero] add support for different paramter groups
1SAA Sep 15, 2022
b7fd676
[zero] add state_dict for zero optimizer
1SAA Sep 20, 2022
d36ff1c
[polish] polish zero code and add docstrings
1SAA Sep 21, 2022
662e38d
[refactor] refactor zero and its test files
1SAA Sep 21, 2022
dac2388
[polish] fix all unit tests with zero
1SAA Sep 21, 2022
5d114a4
rename test_zero_optim.py
1SAA Oct 8, 2022
40f4189
fix ZeroOptimizer unit test
1SAA Oct 8, 2022
54ebbdc
add comments
1SAA Oct 8, 2022
2d73ac8
move pair init to ZeroDDP
1SAA Oct 8, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
8 changes: 2 additions & 6 deletions colossalai/gemini/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,6 @@
from .chunk import TensorInfo, Chunk, TensorState
from .chunk_mgr import ChunkManager
from .chunk import TensorInfo, TensorState
from .stateful_tensor_mgr import StatefulTensorMgr
from .tensor_placement_policy import TensorPlacementPolicyFactory
from .gemini_mgr import GeminiManager

__all__ = [
'StatefulTensorMgr', 'TensorPlacementPolicyFactory', 'GeminiManager', 'ChunkManager', 'TensorInfo', 'Chunk',
'TensorState'
]
__all__ = ['StatefulTensorMgr', 'TensorPlacementPolicyFactory', 'GeminiManager', 'TensorInfo', 'TensorState']
316 changes: 0 additions & 316 deletions colossalai/gemini/chunk.py
View file
Open in desktop

This file was deleted.

3 changes: 3 additions & 0 deletions colossalai/gemini/chunk/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
from .chunk import TensorState, TensorInfo, ChunkFullError, Chunk
from .manager import ChunkManager
from .search_utils import clasify_params, search_chunk_configuration