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

feat(python): support checkpointing for moe #242

Merged
merged 25 commits into from
Oct 15, 2021
Merged

feat(python): support checkpointing for moe #242

merged 25 commits into from
Oct 15, 2021

Conversation

liuhatry
Copy link
Member

No description provided.

github-actions[bot]
github-actions bot reviewed Sep 30, 2021
View reviewed changes
Copy link
Contributor

@github-actions github-actions bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

blackfmt

bagua/torch_api/checkpointing.py|148|
bagua/torch_api/checkpointing.py|153|
bagua/torch_api/checkpointing.py|158|
bagua/torch_api/checkpointing.py|165|
bagua/torch_api/checkpointing.py|187|
bagua/torch_api/checkpointing.py|191|
bagua/torch_api/checkpointing.py|195|
bagua/torch_api/checkpointing.py|198|
bagua/torch_api/checkpointing.py|201|
bagua/torch_api/checkpointing.py|206|
bagua/torch_api/checkpointing.py|211|
bagua/torch_api/checkpointing.py|213|
bagua/torch_api/checkpointing.py|216|
bagua/torch_api/checkpointing.py|219|
bagua/torch_api/checkpointing.py|223|
bagua/torch_api/checkpointing.py|225|

bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
@liuhatry @github-actions
Apply suggestions from code review
76a7dd0 
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
github-actions[bot]
github-actions bot reviewed Sep 30, 2021
View reviewed changes
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
bagua/torch_api/checkpointing.py Outdated Show resolved Hide resolved
liuhatry and others added 4 commits September 30, 2021 09:01
@liuhatry @github-actions
Apply suggestions from code review
fc7151e 
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
@liuhatry
code style
7df56ed
@liuhatry
update
d8b50c3
@liuhatry
upadte
a8dbc00
@liuhatry liuhatry marked this pull request as draft September 30, 2021 08:24
liuhatry and others added 4 commits September 30, 2021 16:58
@liuhatry
update moe
2f86953
@liuhatry @github-actions
Update bagua/torch_api/checkpoint/checkpointing.py
6bd0197 
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
@liuhatry
udpate
401dea2
@liuhatry
Merge branch 'moe_checkpoint' of https://github.com/BaguaSys/bagua in…
3b9a1eb 
…to moe_checkpoint
@liuhatry liuhatry marked this pull request as ready for review September 30, 2021 09:58
NOBLES5E
NOBLES5E requested changes Oct 8, 2021
View reviewed changes
Copy link
Contributor

@NOBLES5E NOBLES5E left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

see comments

liuhatry and others added 2 commits October 8, 2021 15:45
@liuhatry
code format
46bdc15
@liuhatry @github-actions
Update bagua/torch_api/checkpoint/checkpointing.py
c13f765 
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
@liuhatry liuhatry changed the title feat(python): support checkpointing for moe feat: support checkpointing for moe Oct 8, 2021
@liuhatry liuhatry changed the title feat: support checkpointing for moe feat(python): support checkpointing for moe Oct 8, 2021
NOBLES5E
NOBLES5E reviewed Oct 9, 2021
View reviewed changes
bagua/torch_api/checkpoint/checkpointing.py Outdated


def _save_moe_checkpoint(
iteration, checkpoints_path, num_experts, model, optimizer=None, lr_scheduler=None
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lack type annotation

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

still lack type annotations on _xxx functions

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

NOBLES5E
NOBLES5E reviewed Oct 9, 2021
View reviewed changes
bagua/torch_api/checkpoint/checkpointing.py Outdated


def save_checkpoint(
iteration, checkpoints_path, model, optimizer=None, lr_scheduler=None
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lack type annotation, similar for other functions

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

@liuhatry @github-actions
Update bagua/torch_api/checkpoint/checkpointing.py
39bf7aa 
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
NOBLES5E
NOBLES5E reviewed Oct 15, 2021
View reviewed changes
bagua/torch_api/checkpoint/checkpointing.py Outdated
def load_checkpoint(
checkpoints_path: str,
model: BaguaModule,
optimizer: torch.optim.Optimizer = None,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Optional[xxx]

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

NOBLES5E
NOBLES5E reviewed Oct 15, 2021
View reviewed changes
bagua/torch_api/checkpoint/checkpointing.py Outdated
checkpoints_path: str,
model: BaguaModule,
optimizer: torch.optim.Optimizer = None,
lr_scheduler: torch.optim.lr_scheduler._LRScheduler = None,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Optional[xxx]

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

NOBLES5E
NOBLES5E reviewed Oct 15, 2021
View reviewed changes
bagua/torch_api/checkpoint/checkpointing.py Outdated
iteration: int,
checkpoints_path: str,
model: BaguaModule,
optimizer: torch.optim.Optimizer = None,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Optional[xxx]

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

NOBLES5E
NOBLES5E reviewed Oct 15, 2021
View reviewed changes
bagua/torch_api/checkpoint/checkpointing.py Outdated
def _get_moe_state_dict(
full_state_dict: Dict[str, torch.Tensor],
num_local_experts: int,
expp_rank: int,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what is expp short for?

@NOBLES5E NOBLES5E merged commit 5dcd77b into master Oct 15, 2021
@NOBLES5E NOBLES5E deleted the moe_checkpoint branch October 15, 2021 08:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@NOBLES5E NOBLES5E NOBLES5E requested changes

@github-actions GitHub Actions GitHub Actions left review comments

Assignees
No one assigned
Labels
PR: merged
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@liuhatry @NOBLES5E