bigscience-workshop/petals

cpufeature fails to be installed on some non-x86 CPUs

Open

#279 geöffnet am 1. März 2023

Auf GitHub ansehen
 (6 Kommentare) (0 Reaktionen) (0 zugewiesene Personen)Python (424 Forks)batch import
1daybuggood first issuehelp wanted

Repository-Metriken

Stars
 (8.248 Stars)
PR-Merge-Metriken
 (Keine gemergten PRs in 30 T)

Beschreibung

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