add context parallel group init to mp init #1174
Conversation
facebook-github-bot
added
the
CLA Signed
amylittleyang
force-pushed
the
add_context_parallel_init
branch
from
289e0c1
to
a8e7d06
Compare
| pipeline_length: int = 1, | ||
| *, | ||
| model_parallel_backend: Optional[str] = None, | ||
| pipeline_backend: Optional[str] = None, | ||
| ddp_backend: Optional[str] = None | ||
| ddp_backend: Optional[str] = None, | ||
| cp_backend: Optional[str] = None, |
There was a problem hiding this comment.
can we put cp_backend between model_parallel_backend and pipeline_backend to meet the ordering?
| global _CP_PREV_RANK | ||
| global _CP_NEXT_RANK | ||
| global _CP_ZERO_RANK |
There was a problem hiding this comment.
wondering why we add these?
all these can be derived with CP GROUP RANKS (in fact, PP does that for similar functionalities)
we can always utilize @functools.lru_cache if performance is concern?
There was a problem hiding this comment.
Sounds good let me move this part into fbcode. It's probably for convenience in the original PR where they put all comms related changes in this file.
| def model_parallel_is_initialized() -> bool: | ||
| """Check if model and data parallel groups are initialized.""" | ||
| if _MODEL_PARALLEL_GROUP is None or _DATA_PARALLEL_GROUP is None or _PIPELINE_PARALLEL_GROUP is None: | ||
| return False | ||
| return True |
There was a problem hiding this comment.
defined twice here and below?
also, should we add context parallel group to condition?
| return True | ||
|
|
||
|
|
||
| def get_context_parallel_group(): |
There was a problem hiding this comment.
add return type annotation for all newly added functions?
| global _PIPELINE_PARALLEL_RANKS | ||
| _PIPELINE_PARALLEL_RANKS = None |
| if torch.distributed.is_available() and torch.distributed.is_initialized(): | ||
| return torch.distributed.get_world_size(group=get_context_parallel_group()) | ||
| else: | ||
| return 0 |
There was a problem hiding this comment.
is 0 right default value? usually world size returns 1 by default?
There was a problem hiding this comment.
We should return torch.distributed.get_world_size(...) directly here. get_context_parallel_group() already checks the CP pg is initialized.
amylittleyang
force-pushed
the
add_context_parallel_init
branch
2 times, most recently
from
9fb8ed9
to
78471c8
Compare
| def get_context_parallel_rank() -> int: | ||
| """Return my rank for the context parallel group.""" | ||
| return torch.distributed.get_rank(group=get_context_parallel_group()) |
There was a problem hiding this comment.
looks like line 161-165 is the same function - maybe remove this?
amylittleyang
force-pushed
the
add_context_parallel_init
branch
from
78471c8
to
66d75ce
Compare
What does this PR do?
follow up from D55538929.
update model_parallel/initialize.py to also initialize context parallel groups.
original change made in this PR: https://github.com/fairinternal/llm_inference/pull/333/files
Before submitting
PR review
Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.