Repository: microsoft/DeepSpeed
File: deepspeed/runtime/zero/config.py (Python)

Code before:
    ZeRO3-Offload, ZeRO-Infinity and ZeRO-Inference.
    """

    max_live_parameters: int = Field(pp_int(1e9), ge=0, alias="stage3_max_live_parameters")
    """
    The maximum number of parameters resident per GPU before releasing. Smaller
    values use less memory, but perform more communication.
    """

    max_reuse_distance: int = Field(pp_int(1e9), ge=0, alias="stage3_max_reuse_distance")
    """
    Do not release a parameter if it will be reused within this threshold of
    parameters. Smaller values use less memory, but perform more communication.
    """

    gather_16bit_weights_on_model_save: bool = Field(False, alias="stage3_gather_16bit_weights_on_model_save")
    """
    Consolidate the weights before saving the model by ``save_16bit_model()``.
    Since the weights are partitioned across GPUs, they aren’t part of
    ``state_dict``, so this function automatically gathers the weights when
    this option is enabled and then saves the fp16 model weights.
    """

    coalesced_fetch_threshold: int = Field(pp_int(0), alias="stage3_coalesced_fetch_threshold")
    """
    The ratio of a module's number of parameters/required forward passes (layers)
    measures model granularity. Modules with values below this threshold will be
    treated as an integral unit (to avoid recursion) when fetching at stage3.
    This will reduce the host overhead and separated allgather overhead in fetching
    parameters introduced by hooks for fine-grained layers.
    """

    use_all_reduce_for_fetch_params: bool = Field(False, alias="stage3_use_all_reduce_for_fetch_params")
    """
    Use all_reduce op when fetching module parameters at stage3. This improves performance by reducing
    the overhead of concatenation and slicing on the host.
    """

    stage3_gather_fp16_weights_on_model_save: bool = Field(False,
                                                           json_schema_extra={
                                                               "deprecated": True,
                                                               "new_param": "gather_16bit_weights_on_model_save"
                                                           })
    """ Deprecated, please use ``gather_16bit_weights_on_model_save`` """

    ignore_unused_parameters: bool = True
    """
    Unused parameters in modules may be unexpected in static networks, but
    could be normal in dynamic networks. This controls whether or not training
    should terminate with an error message when unused parameters are detected.
    This is set to ``True`` by default, which means unused parameters are

Review comment (bug):
The description of this option is unclear to me. It seems to be related to a **new** concept of module granularity. I notice that the UT uses `fine-grained` vs `coarse-grained`, which I could understand, but those terms are not used in this description.  

Based on my understanding of this feature, I recommend rewriting the description to emphasize the following points: 

1. Granularity of a module is as computed recursively as ratio of `parameter_count`/ (1 + `descendant count`).
2. ZeRO3 classifies each module into `fine-grained` vs `coarse-grained` by comparing granularity value to `coalesced_fetch_threshold` value.
3.  `Fine-grained` modules are treated as integral units (along with their descendants) for parameter fetching purposes. 

Also, I recommend renaming `coalesced_fetch_threshold` to `module_granularity_threshold` to indicate that this how users can control the classification of `fine-grained` vs `coarse-grained`. 

Please let me know what I am getting wrong. Thanks! 


Code after:
    ZeRO3-Offload, ZeRO-Infinity and ZeRO-Inference.
    """

    max_live_parameters: int = Field(pp_int(1e9), ge=0, alias="stage3_max_live_parameters")
    """
    The maximum number of parameters resident per GPU before releasing. Smaller
    values use less memory, but perform more communication.
    """

    max_reuse_distance: int = Field(pp_int(1e9), ge=0, alias="stage3_max_reuse_distance")
    """
    Do not release a parameter if it will be reused within this threshold of
    parameters. Smaller values use less memory, but perform more communication.
    """

    gather_16bit_weights_on_model_save: bool = Field(False, alias="stage3_gather_16bit_weights_on_model_save")
    """
    Consolidate the weights before saving the model by ``save_16bit_model()``.
    Since the weights are partitioned across GPUs, they aren’t part of
    ``state_dict``, so this function automatically gathers the weights when
    this option is enabled and then saves the fp16 model weights.
    """

    module_granularity_threshold: int = Field(pp_int(0), alias="stage3_module_granularity_threshold")
    """
    The granularity of a module is determined by the ratio of "parameter_count / (1 + descendant count)".
    ZeRO3 classifies modules with a granularity below the threshold as fine-grained,
    which are treated as integral units during parameter fetching. This reduces host overhead
    and the separate allgather overhead introduced by hooks for fine-grained layers when fetching parameters.
    """

    use_all_reduce_for_fetch_params: bool = Field(False, alias="stage3_use_all_reduce_for_fetch_params")
    """
    Use all_reduce op when fetching module parameters at stage3. This improves performance by reducing
    the overhead of concatenation and slicing on the host.
    """

    stage3_gather_fp16_weights_on_model_save: bool = Field(False,
                                                           json_schema_extra={
                                                               "deprecated": True,
                                                               "new_param": "gather_16bit_weights_on_model_save"
                                                           })
    """ Deprecated, please use ``gather_16bit_weights_on_model_save`` """

    ignore_unused_parameters: bool = True
    """
    Unused parameters in modules may be unexpected in static networks, but
    could be normal in dynamic networks. This controls whether or not training
    should terminate with an error message when unused parameters are detected.
    This is set to ``True`` by default, which means unused parameters are
    ignored and training continues. Now is just used in stage 2.