Function get_static_torch_model

colossalai/zero/gemini/utils.py:64–107

Get a static torch.nn.Module model from the given GeminiDDP module. You should notice that the original GeminiDDP model is not modified. Thus, you can use the original model in further training. But you should not use the returned torch model to train, this can cause unexpected errors.

(
    zero_ddp_model, device=torch.device("cpu"), dtype=torch.float32, only_rank_0=True
)

def get_static_torch_model(
    zero_ddp_model, device=torch.device("cpu"), dtype=torch.float32, only_rank_0=True
) -> torch.nn.Module:
    """Get a static torch.nn.Module model from the given GeminiDDP module.
    You should notice that the original GeminiDDP model is not modified.
    Thus, you can use the original model in further training.
    But you should not use the returned torch model to train, this can cause unexpected errors.

    Args:
        zero_ddp_model (GeminiDDP): a zero ddp model
        device (torch.device): the device of the final torch model
        dtype (torch.dtype): the dtype of the final torch model
        only_rank_0 (bool): if True, only rank0 has the converted torch model

    Returns:
        torch.nn.Module: a static torch model used for saving checkpoints or numeric checks
    """
    from colossalai.zero.gemini.gemini_ddp import GeminiDDP

    assert isinstance(zero_ddp_model, GeminiDDP)

    state_dict = zero_ddp_model.state_dict(only_rank_0=only_rank_0)
    colo_model = zero_ddp_model.module
    torch_model = _get_shallow_copy_model(colo_model)

    if not only_rank_0 or dist.get_rank() == 0:
        for (name, colo_module), (_, torch_module) in zip(
            _get_dfs_module_list(colo_model), _get_dfs_module_list(torch_model)
        ):
            # clean the parameter list of the new torch module
            torch_module._parameters = OrderedDict()
            for sufix_param_name, param in colo_module.named_parameters(recurse=False):
                # get the full name of the parameter
                full_param_name = name + ("." if name else "") + sufix_param_name
                assert (
                    full_param_name in state_dict
                ), f"Can not find parameter `{full_param_name}` in the GeminiDDP module"
                state_param = state_dict[full_param_name]
                torch_param = torch.nn.Parameter(state_param.data.to(device=device, dtype=dtype))

                setattr(torch_module, sufix_param_name, torch_param)
    dist.barrier()

    return torch_model