bigscience-workshop/petals

cpufeature fails to be installed on some non-x86 CPUs

Open

#279 opened on Mar 1, 2023

View on GitHub
 (6 comments) (0 reactions) (0 assignees)Python (424 forks)batch import
1daybuggood first issuehelp wanted

Repository metrics

Stars
 (8,248 stars)
PR merge metrics
 (No merged PRs in 30d)

Description

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

Contributor guide