Skip to content

[checkpointio]support distributed checkpoint io for model saving. #6181

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

Open
wants to merge 8 commits into
base: feature/dist-ckp-io
Choose a base branch
from
Open

[checkpointio]support distributed checkpoint io for model saving. #6181

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
370e545
support distribute checkpoint io
flybird11111 Jan 16, 2025
c34ba4e
fix
flybird11111 Jan 17, 2025
e3f9de3
Modify the design
flybird11111 Jan 20, 2025
a28fdde
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jan 20, 2025
f388bbe
fix
flybird11111 Jan 20, 2025
c5b0882
Remove duplicates
flybird11111 Jan 21, 2025
8e6902c
fix
flybird11111 Jan 22, 2025
7e880a1
fix async io
flybird11111 Feb 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions colossalai/booster/plugin/hybrid_parallel_plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,9 @@ def __init__(
self.require_grad_sync = True
self.overlap_allgather = overlap_allgather
self.use_fp8 = use_fp8
self.param_origin_shape = {}
for name, param in module.named_parameters():
self.param_origin_shape[name] = param.shape

shardformer = ShardFormer(shard_config)
if custom_policy is not None:
Expand Down
Loading
Loading