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.

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

Optimizer load gathered state and record delta feature are supported now #184

Merged
merged 3 commits into from
Feb 23, 2024

Conversation

MayDomine
Copy link
Collaborator

AdamOffloadOptimizer new feature

Description

为AdamOffload提供两个新的新接口

  1. 位于AdamOffloadOptimizer的__init__(self)里传递record_delta=True,将记录adam每个step更新量的统计信息。

  2. AdamOffloadOptimizer的state_dict(self)里传递gather=True,将对所有rank的opt进行gather得到完整的grad ckpt(注意该功能在大模型上可能会使cpu内存炸掉),并且该完整ckpt可以使用optimizer.load_state_dict(torch.load(filename))直接加载(无论在单卡还是多卡),即多卡可以同时加载一个完整的grad ckpt,因为在写文件时会添加一个标识意味着该文件为完整ckpt,从而load每个rank内再切出属于自己的那一份grad

Type of Change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • This change requires a documentation update

Checklist

  • I have read the CONTRIBUTING document.
  • My code follows the code style of this project.
  • My change requires a change to the documentation.
  • I have updated the documentation accordingly.
  • I have added tests to cover my changes.
  • All new and existing tests passed.

Additional Information

Any additional information, configuration, or data that might be necessary for the review.

@MayDomine MayDomine requested a review from Achazwl February 22, 2024 11:03
@MayDomine MayDomine self-assigned this Feb 22, 2024
@MayDomine MayDomine merged commit 0def29f into dev Feb 23, 2024
6 checks passed
@MayDomine MayDomine deleted the load_opt branch February 26, 2024 05:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant