def METHOD_NAME(
    module: "flow.nn.Module",
    *,
    broadcast_buffers: bool = True,
    broadcast_parameters: bool = True,
    bucket_size: int = 10,
    use_bucket: bool = True,
):
    assert all(x.dtype == flow.float32 for x in module.parameters())
    if use_bucket and parse_boolean_from_env("ONEFLOW_DISABLE_VIEW", False):
        warnings.warn(
            "because the environment variable 'ONEFLOW_DISABLE_VIEW' is set to true, so the view mechanism is disabled, and we will set use_bucket=False"
        )
        use_bucket = False
    world_size = flow.env.get_world_size()
    if broadcast_parameters:
        with flow.no_grad():
            for x in module.parameters():
                requires_grad = x.requires_grad
                flow._C.comm_broadcast(x, inplace=True)
                # TODO: fix the bug that x's requires_grad is discarded
                # after flow._C.comm_broadcast
                x.requires_grad_(requires_grad)
    if use_bucket:
        all_grad_size = sum([x.numel() for x in module.parameters()])
        if all_grad_size > 0:
            device = list(module.parameters())[0].device
            assert all(x.device == device for x in module.parameters())
        reversed_param_list = list(
            reversed(
                list([param for param in module.parameters() if param.requires_grad])
            )
        )
        module._bucket_index = {
            x: i // bucket_size for i, x in enumerate(reversed_param_list)
        }
        module._buckets = [
            reversed_param_list[i : i + bucket_size]
            for i in range(0, len(reversed_param_list), bucket_size)
        ]
        module._params_group = ContiguousParamsGroup(module._buckets)
        module._bucket_tensors = module._params_group.grouped_parameters_grad
    ddp_state_for_reversed_params = OrderedDict(
        reversed([(x, [False, False]) for x in module.parameters() if x.requires_grad])
    )
    module._ddp_state_for_reversed_params = ddp_state_for_reversed_params
    # The gradient shoule be averaged by all the nodes, so besides allreduce,
    # a division by world_size is required.
    # Use x * (1 / world_size) instead of x / world_size for two reasons:
    # 1. multiplication is faster than division
    # 2. An inplace operation is needed here (for allreduce grouping)
    #    But we do not have inplace division in oneflow.
    mul_factor = 1 / world_size
    def inplace_mul_and_return_none(x):
        x.mul_(mul_factor)
        return None
    for param in module.parameters():
        if param.requires_grad:
            param._register_post_grad_accumulation_hook(inplace_mul_and_return_none)
            param._register_post_grad_accumulation_hook(
                allreduce_fn(module, param, use_bucket)
            )
    def post_forward_hook(module, input, output):
        ddp_state_for_reversed_params = module._ddp_state_for_reversed_params
        for state in ddp_state_for_reversed_params.values():
            state[0], state[1] = False, False
        output = ArgsTree(output).map_leaf(
            lambda x: flow._C.select_top_n(
                convert_to_tensor_tuple([x, *ddp_state_for_reversed_params.keys()]),
                n=1,
            )[0]
        )
        buffers = list(module.buffers())
        if len(buffers) > 0:
            flow._C.stream_touch(buffers)
        return output
    module.register_forward_hook(post_forward_hook)
    if broadcast_buffers:
        def pre_forward_hook(module, input):
            with flow.no_grad():
                buffers = list(module.buffers())
                flow._C.comm_broadcast(buffers, inplace=True)
        module.register_forward_pre_hook(pre_forward_hook)
    module._is_ddp_module = True
    return module