[DEV] Cherry-pick: M4 + Dist Checkpoint: Replace global parallel state with explicit group parameters - #2152
Conversation
skyw
left a comment
skyw
left a comment
There was a problem hiding this comment.
Looks like it is mainly apply the boilerplate to all necessary places. LGTM.
The reason as of why tp_group is passed outside metadata is not clear. would recommend to document it some where in the code for users to understand and also prevent further changes from breaking it unintentionally.
| self.num_local_experts = num_local_experts | ||
| self.local_experts = torch.nn.ModuleList() | ||
| self.ep_group = pg_collection.ep | ||
| self.tp_group = pg_collection.expt_tp |
There was a problem hiding this comment.
Q: Any reason to not use expt_tp_group to be consistent? Are there any code depend on the naming?
There was a problem hiding this comment.
we were trying to keep it consistent in sharded state dict. it sometimes goes to parent's sharded state dict or default sharded state dict. In these modules it will look for self.tp_group.
|
|
||
| clean_metadata = metadata.copy() | ||
| # Remove dp_cp_group as it's not serializable | ||
| clean_metadata.pop('dp_cp_group', None) |
There was a problem hiding this comment.
Q: Is it possible to have other process groups reach here?
There was a problem hiding this comment.
I don't think so. dp_cp_group should be the only thing.
Signed-off-by: dimapihtar <dpihtar@gmail.com>
Signed-off-by: dimapihtar <dpihtar@gmail.com>
Signed-off-by: dimapihtar <dpihtar@gmail.com>
Signed-off-by: dimapihtar <dpihtar@gmail.com>
Signed-off-by: dimapihtar <dpihtar@gmail.com>
Signed-off-by: dimapihtar <dpihtar@gmail.com>
Signed-off-by: dimapihtar <dpihtar@gmail.com>
Signed-off-by: dimapihtar <dpihtar@gmail.com>
Signed-off-by: dimapihtar <dpihtar@gmail.com>
|
/ok to test ddd3fca |
|
/ok to test 7d92f3c |
@yaoyu-33, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/ |
Signed-off-by: dimapihtar <dpihtar@gmail.com>
Signed-off-by: dimapihtar <dpihtar@gmail.com>
Signed-off-by: dimapihtar <dpihtar@gmail.com>
Signed-off-by: dimapihtar <dpihtar@gmail.com>
|
/ok to test fcb950a |
|
Regarding this MR, have we passed the functional test on dev/MR? |
|
/ok to test c67d12e |
|
unit and functional test both passed on original pr to main: #2053 |
|
BTW, what is the corresponding main PR? Could you update it in the desc? |
# Conflicts: # megatron/training/checkpointing.py
|
/ok to test e1da770 |
6 participants
This PR refactors the parallel group management to eliminate dependencies on global parallel_state.xxx APIs and instead use explicit group parameters (tp_group, pp_group, dp_cp_group) with fallbacks to existing global state when not provided.
Key Changes
1. Explicit Group Parameters
tp_group,pp_group,dp_cp_groupparameters to key functions in:megatron/training/checkpointing.pymegatron/training/utils.pymegatron/core/utils.pyOptional[torch.distributed.ProcessGroup]parameters withNonedefaultsNone, code falls back to existingmpu.get_xxx_group()APIs for backward compatibility2. Enhanced Metadata Handling
_build_sharded_state_dict_metadata()to includedp_cp_groupin metadatadp_cp_groupnow consistently sourced from metadata across checkpoint operations3. Improved Group Sourcing Strategy
module.tp_groupandmodule.pp_groupget_pg_size()andget_pg_rank()utilities for group introspection4. Function Signature Updates
Key functions updated with explicit group parameters:
save_checkpoint()load_checkpoint()get_rng_state()_build_sharded_state_dict_metadata()