bigscience-workshop/petals

cpufeature fails to be installed on some non-x86 CPUs

Open

#279 建立於 2023年3月1日

在 GitHub 查看
 (6 留言) (0 反應) (0 負責人)Python (424 fork)batch import
1daybuggood first issuehelp wanted

倉庫指標

Star
 (8,248 star)
PR 合併指標
 (30 天內沒有已合併 PR)

描述

We use the cpufeature module to detect if a CPU supports AVX512, however, cpufeature fails to be installed on some non-x86 CPUs (e.g., ARMs).

We need to remove this dependency and use code like this to detect if the CPU supports AVX512:

def is_fast_bfloat16_supported() -> bool:
    if os.name != "posix" or platform.processor() != "x86_64":
        return False

    try:
        eax = _run_x86_assembler(
            b"\xB8\x07\x00\x00\x00"  # mov eax, 7
            b"\xB9\x01\x00\x00\x00"  # mov ecx, 1
            b"\x0f\xa2"  # cpuid
            b"\xC3"  # ret
        )
        avx512_bf16 = bool(eax & (1 << 5))

        return amx_bf16 or avx512_bf16
    except:
        logger.warning("Failed to check for the CPU's native bfloat16 support:", exc_info=True)
        return False


def _run_x86_assembler(machine_code: bytes) -> int:
    memory = mmap.mmap(
        -1,
        len(machine_code),
        flags=mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS,
        prot=mmap.PROT_WRITE | mmap.PROT_READ | mmap.PROT_EXEC,
    )
    try:
        memory.write(machine_code)
        address = ctypes.addressof(ctypes.c_int.from_buffer(memory))
        functype = ctypes.CFUNCTYPE(ctypes.c_uint32)
        func = functype(address)
        return func()
    finally:
        memory.close()

Note that the code above checks for the presence of AVX512_BF16 instruction set, but actually pure AVX512 is enough. This needs to be changed before integrating this code, see https://en.wikichip.org/wiki/x86/avx-512 for the relevant instruction params.

Before merging, it's necessary to test on (a) x86-64 CPUs without AVX512, (b) x86-64 CPUs with AVX512, (c) ARM CPUs (e.g. available on Oracle free-tier cloud or other places).

貢獻者指南