Skip to content

[doc] Add missing docstrings to zero3 chunk search utilities #3572

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.

Sign up for GitHub

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

Merged
1SAA merged 1 commit into from
Apr 17, 2023
Merged

[doc] Add missing docstrings to zero3 chunk search utilities #3572

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 31 additions & 3 deletions colossalai/zero/gemini/chunk/search_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,8 +11,13 @@


def _filter_exlarge_params(model: nn.Module, size_dict: Dict[int, List[int]]) -> None:
"""
"""_filter_exlarge_params

Filter those parameters whose size is too large (more than 3x standard deviations) from others.

Args:
model (nn.Module): the model.
size_dict (Dict[int, List[int]]): the size dict of parameters.
"""
agg_size_list = []
for key in size_dict:
Expand All @@ -33,7 +38,16 @@ def _filter_exlarge_params(model: nn.Module, size_dict: Dict[int, List[int]]) ->


def _get_unused_byte(size_list: List[int], chunk_size: int) -> int:
"""Get unused byte for a certain chunk size.
"""_get_unused_byte

Get unused byte for a certain chunk size.

Args:
size_list (List[int]): the size list of parameters.
chunk_size (int): the chunk size.

Returns:
int: the unused byte.
"""
acc = 0
left = 0
Expand All @@ -45,7 +59,18 @@ def _get_unused_byte(size_list: List[int], chunk_size: int) -> int:
return left + acc


def _tensor_numel(local_param: ColoParameter, strict_ddp_flag: bool):
def _tensor_numel(local_param: ColoParameter, strict_ddp_flag: bool) -> int:
"""_tensor_numel

Get the number of elements of a tensor.

Args:
local_param (ColoParameter): The local parameter.
strict_ddp_flag (bool): whether to enable the strict ddp mode.

Returns:
int: the number of elements.
"""
if strict_ddp_flag and type(local_param) is ColoParameter:
return local_param.numel_global()
else:
Expand All @@ -61,6 +86,7 @@ def classify_params_by_dp_degree(param_order: OrderedParamGenerator,

Args:
param_order (OrderedParamGenerator): the order of param be visied
strict_ddp_flag (bool, optional): whether to enable the strict ddp mode. Defaults to False.

Returns:
Dict[int, List[ColoParameter]]: a dict contains the classification results.
Expand Down Expand Up @@ -96,6 +122,8 @@ def search_chunk_configuration(
memstas: Optional[MemStats] = None) -> Tuple[Dict, int, int]:
"""search_chunk_configuration

Search the chunk configuration for a model.

Args:
model (nn.Module): torch module
search_range_mb (float): searching range in mega byte.
Expand Down