def set_group_size(plugin_config: dict):
    """
    Set the group size for tensor parallelism and sequence parallelism.

    Args:
        plugin_config (dict): Plugin configuration.
    """
    tp_size = int(plugin_config.get("tp_size", 1))
    sp_size = int(plugin_config.get("sp_size", 1))
    if tp_size > 1:
        assert sp_size == 1
        plugin_config["tp_size"] = tp_size = min(tp_size, torch.cuda.device_count())
        log_message(f"Using TP with size {tp_size}")
    if sp_size > 1:
        assert tp_size == 1
        plugin_config["sp_size"] = sp_size = min(sp_size, torch.cuda.device_count())
        log_message(f"Using SP with size {sp_size}")

    Set the group size for tensor parallelism and sequence parallelism.

    Args:
        plugin_config (dict): Plugin configuration.
    
opensora/utils/cai.py