Function preprocess_weights_tl2

utils/generate-dummy-bitnet-model.py:699–773

(
    w: np.ndarray,
    bits = 2,
    g    = 4,
)

def preprocess_weights_tl2(
    w: np.ndarray,
    bits = 2,
    g    = 4,
) -> Tuple[np.ndarray, np.ndarray]:
    M, K = w.shape
    weight = w
    weight = np.where(np.abs(weight) < 1e-6, 0, weight).astype(np.float32)
    weight = np.sign(weight)
    weight_num = np.prod(weight.shape)

    # for three num 6 bit ->

    # outer loop
    KEMD = 1536
    BMEMD = 256
    BYEMD = 96

    KGQA = 4096
    BMGQA = 128
    BYGQA = 96

    # inner loop (32row 32num/16index)
    bm3 = 32
    by3 = 6

    if K == KEMD:
        BM3 = BMEMD
        BY3 = BYEMD
    elif K == KGQA:
        BM3 = BMGQA
        BY3 = BYGQA
    else:
        raise NotImplementedError

    BM2 = BM3
    BY2 = 32
    # inner loop (32row 32num/16index)
    bm2 = 32
    by2 = 4

    if (weight.shape[1] % BY3 != 0):
        slice_k_idx = weight.shape[1] - weight.shape[1] % BY3
        slice_weights = np.split(weight, [slice_k_idx], axis=1)
        three_weight = slice_weights[0]
        two_weight = slice_weights[1]
    else:
        three_weight = weight

    final_weight = []

    preprocess_three_weights_tl2(three_weight.shape[0],
                         three_weight.shape[1],
                         three_weight.shape[0] * three_weight.shape[1],
                         BM3,
                         BY3,
                         bm3,
                         by3,