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

[checkpointio] support unsharded checkpointIO for hybrid parallel (#4774)

Browse Source 
* support unsharded saving/loading for model

* support optimizer unsharded saving

* update doc

* support unsharded loading for optimizer

* small fix
This commit is contained in:
Baizhou Zhang
2023-09-26 10:58:03 +08:00
committed by GitHub
parent a2db75546d
commit 64a08b2dc3
4 changed files with 197 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tests/test_checkpoint_io/test_hybrid_parallel_plugin_checkpoint_io.py
View File

@@ -20,9 +20,8 @@ from colossalai.testing import (
from tests.kit.model_zoo import model_zoo
# TODO (Baizhou): Add test cases for shard=False
@clear_cache_before_run()
@parameterize("shard", [True])
@parameterize("shard", [True, False])
@parameterize("model_name", ["transformers_gpt"])
@parameterize("size_per_shard", [32])
@parameterize(