def model_gathering(model: torch.nn.Module, model_shape_dict: dict, pinned_state_dict: dict) -> None:
    """
    Gather the model parameters from multiple GPUs.

    Args:
        model (torch.nn.Module): The model to gather.
        model_shape_dict (dict): The shape of the model parameters.
        device (torch.device): The device to gather the model to.
    """
    global_rank = dist.get_rank()
    global_size = dist.get_world_size()
    params = set()
    for name, param in model.named_parameters():
        params.add(name)
        all_params = [torch.empty_like(param.data) for _ in range(global_size)]
        dist.all_gather(all_params, param.data, group=dist.group.WORLD)
        if int(global_rank) == 0:
            all_params = torch.cat(all_params)
            gathered_param = remove_padding(all_params, model_shape_dict[name]).view(model_shape_dict[name])
            pinned_state_dict[name].copy_(gathered_param)
    if int(global_rank) == 0:
        for k, v in model.state_dict(keep_vars=True).items():
            if k not in params:
                pinned_state_dict[k].copy_(v)

    dist.barrier()

    Gather the model parameters from multiple GPUs.

    Args:
        model (torch.nn.Module): The model to gather.
        model_shape_dict (dict): The shape of the model parameters.
        device (torch.device): The device to gather the model to.
    
opensora/utils/ckpt.py