alibaba / alibaba/ChatLearn

[Feature]Use Megatron-core dist_checkpointing to load checkpoint with different parallel strategies

Open
#169 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
450
Forks
40
PR merge metrics
No merged PRs in 30d

Description

**Is your feature request related to a problem? Please describe.**
Currently ChatLearn use a tool to convert checkpoint when different parallel strategies are detected, https://github.com/alibaba/ChatLearn/blob/main/chatlearn/utils/megatron_utils.py#L164

The online conversion has been addressed in Megatron core dist_checkpointing. https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/dist_checkpointing.html

**Describe the solution you'd like**
Use Megatron-core dist checkpointing to save and load checkpoint.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with chatlearn/utils/megatron_utils.py around line 164 and read the Megatron-core dist_checkpointing documentation linked in the issue. Trace the current checkpoint conversion path and determine how save and load should use distributed checkpointing across parallel strategies. Done means checkpoints can be saved and loaded without the existing online conversion step.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems, machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.