llmcompressor.utils.fsdp.context

函数

• fix_fsdp_module_name –

  从模块名称中移除 FSDP 包装器前缀。

• main_process_first_context –

  创建一个上下文管理器，其中主进程在所有其他进程之前运行该块。

fix_fsdp_module_name

fix_fsdp_module_name(name: str) -> str

从模块名称中移除 FSDP 包装器前缀。考虑到 FSDP_WRAPPED_MODULE 位于名称的末尾以及中间的情况。

参数

• 名称

  (str) –

  需要剥离的名称

返回

• str –

  剥离后的名称

llmcompressor/utils/fsdp/context.py 中的源代码

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, ""
    )

main_process_first_context

main_process_first_context()

创建一个上下文管理器，其中主进程在所有其他进程之前运行该块。当从单进程应用程序调用时，返回一个 nullcontext。

llmcompressor/utils/fsdp/context.py 中的源代码

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()