Repository: microsoft/DeepSpeed
File: deepspeed/utils/zero_to_fp32.py (Python)

Code before:
        print(f"added {len(buffers)} buffers")

    if not exclude_frozen_parameters:
        _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)

    _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)

    # recover shared parameters
    for pair in zero_model_states[0].shared_params:
        if pair[1] in state_dict:
            state_dict[pair[0]] = state_dict[pair[1]]

    return state_dict


def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None, exclude_frozen_parameters=False, lazy_merge=False):
    """
    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
    via a model hub.

    Args:
        - ``checkpoint_dir``: path to the desired checkpoint folder
        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
        - ``exclude_frozen_parameters``: exclude frozen parameters
        - ``lazy_merge``: a more memory-efficient feature

    Returns:
        - pytorch ``state_dict``

    Note: this approach may not work if your application doesn't have sufficient free CPU memory and
    you may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
    the checkpoint.

    A typical usage might be ::

        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
        # do the training and checkpoint saving
        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
        model = model.cpu() # move to cpu
        model.load_state_dict(state_dict)
        # submit to model hub or save the model to share with others

    In this example the ``model`` will no longer be usable in the deepspeed context of the same
    application. i.e. you will need to re-initialize the deepspeed engine, since
    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.

    If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.

    """
    if tag is None:

Review comment (performance):
Provide brief description of why more memory-efficient, and perhaps mention important usage concepts like pseudo tensors and `contiguous()`.
 
Also, please add a unit test.

Code after:
            state_dict[pair[0]] = state_dict[pair[1]]

    return state_dict


def to_torch_tensor(state_dict, return_empty_tensor=False):
    """
    Convert state_dict of GatheredTensor to torch tensor
    """
    converted_tensors = {}
    for name, tensor in state_dict.items():
        tensor_id = id(tensor)
        if tensor_id in converted_tensors:
            shared_tensor = state_dict[converted_tensors[tensor_id]]
            state_dict[name] = shared_tensor
        else:
            converted_tensors[tensor_id] = name
            if return_empty_tensor:
                state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
            else:
                state_dict[name] = tensor.contiguous()
    return state_dict


def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
                                             tag=None,
                                             exclude_frozen_parameters=False,
                                             lazy_mode=False):
    """
    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
    via a model hub.

    Args:
        - ``checkpoint_dir``: path to the desired checkpoint folder
        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
        - ``exclude_frozen_parameters``: exclude frozen parameters
        - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
          Convert the pesduo tensor to torch tensor by ``.contiguous()``

    Returns:
        - pytorch ``state_dict``

    A typical usage might be ::

        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
        # do the training and checkpoint saving
        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
        model = model.cpu() # move to cpu
        model.load_state_dict(state_dict)
        # submit to model hub or save the model to share with others