llmcompressor.utils.fsdp.context

def fix_fsdp_module_name(name: str) -> str:
    """
    Remove FSDP wrapper prefixes from a module name.
    Accounts for scenario where FSDP_WRAPPED_MODULE is
    at the end of the name, as well as in the middle.

    :param name: name to strip
    :return: stripped name
    """
    if FullyShardedDataParallel is None:
        return name

    return name.replace(FSDP_WRAPPED_MODULE + ".", "").replace(
        "." + FSDP_WRAPPED_MODULE, ""
    )

def main_process_first_context():
    """
    Creates a context manager where the main process runs the block before all other
    processes. Returns a nullcontext when called from a single process application.
    """
    if Accelerator is None:
        return nullcontext()

    return Accelerator().main_process_first()