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

Background mlflow model registration #1078

Merged
merged 19 commits into from
Apr 3, 2024
Merged

Conversation

irenedea
Copy link
Contributor

@irenedea irenedea commented Mar 30, 2024

Spawn a new process that handles the mlflow model registration (which uploads the model to UC). This way, global rank 0 isn't stuck on with a long operation with variable speeds, which could lead to nccl timeouts. This also allows other operations to happen concurrently such as composer checkpoint saving, etc.

At the end of training, we wait for all child processes to finish and clean up temporary directories.

Manual Test

test-hf-mlflow-background-register-ib6LgH
test-hf-mlflow-background-register-llama70b-Lhda8z

@irenedea irenedea marked this pull request as ready for review March 31, 2024 18:59
@irenedea irenedea requested a review from dakinggg March 31, 2024 19:17
llmfoundry/callbacks/hf_checkpointer.py Outdated Show resolved Hide resolved
llmfoundry/callbacks/hf_checkpointer.py Outdated Show resolved Hide resolved
llmfoundry/callbacks/hf_checkpointer.py Show resolved Hide resolved
Copy link
Collaborator

@dakinggg dakinggg left a comment

Choose a reason for hiding this comment

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

LGTM!

llmfoundry/callbacks/hf_checkpointer.py Outdated Show resolved Hide resolved
irenedea and others added 2 commits April 2, 2024 14:49
Co-authored-by: Daniel King <43149077+dakinggg@users.noreply.github.com>
@irenedea irenedea enabled auto-merge (squash) April 2, 2024 21:49
@irenedea irenedea merged commit 23c3173 into mosaicml:main Apr 3, 2024
9 checks passed
dakinggg pushed a commit to dakinggg/llm-foundry that referenced this pull request Apr 4, 2024
KuuCi pushed a commit that referenced this pull request Apr 18, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants