AI說書 - 從0開始 - 14

我想要一天分享一點「LLM從底層堆疊的技術」，並且每篇文章長度控制在三分鐘以內，讓大家不會壓力太大，但是又能夠每天成長一點。

已經在AI說書 - 從0開始 - 12以及AI說書 - 從0開始 - 13中見識到TPU的威力了，現在我們把參數放大到真實大型語言模型的規模，看看運算時間的等級。

程式配置為：

import tensorflow as tf
import numpy as np
import time

n = 32768
d = 12288

input_seq = tf.random.normal((n, d), dtype = tf.float32)

start_time = time.time()
_ = tf.matmul(input_seq, input_seq, transpose_b = True)

attention_time = time.time() - start_time
print(f"Self-attention computation time: {attention_time} seconds")

對此，驚人的執行結果如下：

如果想看看一些TPU的資訊，可以執行以下程式：

import os
from tensorflow.python.profiler import profiler_client

tpu_profile_service_address = os.environ['COLAB_TPU_ADDR'].replace('8470', '8466')
print(profiler_client.monitor(tpu_profile_service_address, 100, 2))

對此，執行結果如下：

可以看到TPU幾乎沒被使用到。