def model_sharding(model: torch.nn.Module, device: torch.device = None):
    """
    Sharding the model parameters across multiple GPUs.

    Args:
        model (torch.nn.Module): The model to shard.
        device (torch.device): The device to shard the model to.
    """
    global_rank = dist.get_rank()
    world_size = dist.get_world_size()
    for _, param in model.named_parameters():
        if device is None:
            device = param.device
        padding_size = (world_size - param.numel() % world_size) % world_size
        if padding_size > 0:
            padding_param = torch.nn.functional.pad(param.data.view(-1), [0, padding_size])
        else:
            padding_param = param.data.view(-1)
        splited_params = padding_param.split(padding_param.numel() // world_size)
        splited_params = splited_params[global_rank]
        param.data = splited_params.to(device)

    Sharding the model parameters across multiple GPUs.

    Args:
        model (torch.nn.Module): The model to shard.
        device (torch.device): The device to shard the model to.
    
opensora/utils/ckpt.py