Function preprocess_weights_tl1

utils/generate-dummy-bitnet-model.py:528–574

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

# based on t_mac.utils.preprocess_weights
def preprocess_weights_tl1(
    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)
    model_size = args.model_size

    KEMD = model_config[model_size]['hidden_size']
    # outer loop
    BMEMD = 256
    BYEMD = 256

    # inner loop (32row 32num/16index)
    bmEMD = 32
    byEMD = 8

    KGQA = model_config[model_size]['intermediate_size']

    BMGQA = 256
    BYGQA = 256

    bmGQA = 32
    byGQA = 8

    weight = np.reshape(weight, (weight_num // 2, 2))
    hi_weight = np.multiply(np.split(weight, 2, axis=1)[0], 3)
    lo_weight = np.split(weight, 2, axis=1)[1]

    weight = np.reshape((hi_weight + lo_weight), weight_num // 2)

    # row-major index
    weight = weight + 4
    weight = np.reshape(weight, (M, K // 2)).astype(np.uint8)

    if K == KEMD:
        weight = process_tl1(weight, BMEMD, BYEMD, bmEMD, byEMD, M, K)
    elif K == KGQA:
        weight = process_tl1(weight, BMGQA, BYGQA, bmGQA, byGQA, M, K)
    else:
        raise NotImplementedError

    return weight


def preprocess_two_weights_tl2(M, K, weight_num, BM, BY, bm, by, weight, final_weight):