def METHOD_NAME(model: nn.Module,
                       zero_stage: int = 1,
                       gemini_config: Optional[Dict] = None,
                       verbose: bool = False):
    """This wrapper function is used to wrap your training model for ZeRO DDP.
    Example:
        >>> with ColoInitContext():
        >>>     my_model = Bert()
        >>> my_optim = SGD(my_model.parameters(), lr = 1e-3)
        >>> zero_model = zero_model_wrapper(my_model, zero_stage=1)
        >>> zero_optim = zero_optim_wrapper(zero_model, my_optim)
    Args:
        model (nn.Module): The model used in ZeRO DDP.
        zero_stage (int, optional): The stage of ZeRO DDP. You can find more information in ZeRO's paper.
            https://arxiv.org/abs/1910.02054
        gemini_config (dict, optional): The configuration dictionary of `GeminiDDP`. `GeminiDDP` is enabled
            when the stage is set to 3. You can set the arguments of `GeminiDDP` in the gemini_config.
            Here is an example where we set the device of the model, the placement policy of Gemini, and the
            size of hidden dimension to help Gemini find out a unified chunk size.
            Example:
                >>> config_dict = dict(device=torch.cuda.current_device(), hidden_dim=1024, placement_policy='auto')
                >>> model = zero_model_wrapper(model, zero_stage=3, gemini_config=config_dict)
    """
    assert zero_stage in [1, 2, 3], "The stage of ZeRO should be 1, 2 or 3"
    if gemini_config is None:
        gemini_config = dict()
    if zero_stage in [1, 2]:
        wrapped_model = model
    else:
        wrapped_model = GeminiDDP(model, **gemini_config, verbose=verbose)
    setattr(wrapped_model, "_colo_zero_stage", zero_stage)
    return wrapped_model