跳到正文
r/MachineLearning· /u/2dayyy6·· 4 小时前AI 评分52

循环 Transformer 能节省 HBM 带宽吗?在 B200 上剖析 Huginn 的权重复用

Do looped transformers save HBM traffic? Profiling weight reuse in Huginn on a B200 [P]

AI 导读

作者在 B200 上剖析循环 Transformer 架构 Huginn 的权重复用情况,测量每次迭代从内存读取的数据量和耗时。结果显示,8 个 matmul 共享一份 72 MiB 权重时比 8 份相同副本快 24%,且几乎不从 HBM 读取数据;权重增至 288 MiB、超过 L2 后这一差异消失。

正文
Do looped transformers save HBM traffic? Profiling weight reuse in Huginn on a B200 [P]

Do looped transformers, which reuse the same weights at every iteration, also save HBM traffic? In my new post I cover the looped architecture (Huginn), then model and measure its performance on a B200: how much data each iteration reads from memory, how long it takes, and why the shared weights usually don't stay on chip until the next iteration.
A few results:

- A chain of eight matmuls with one shared 72 MiB weight ran 24% faster than the same chain with eight identical copies, and read almost nothing from HBM. At 288 MiB, larger than L2, the difference disappeared.

- In a full Huginn Core (teaching implementation, random weights), eight iterations took about eight times as long as one and read about eight times as much data from HBM.

- Sharing saves storage, but by itself it doesn't stop the GPU from reading the weights from HBM again at every iteration.

👀Post: https://dinara-dl.github.io/posts/looped-transformer/

🤖Code and raw Nsight Compute data: https://github.com/dinara-dl/looped-transformer

I'd be curious whether anyone has tried keeping recurrent weights on chip across iterations, for example with a persistent kernel.

submitted by /u/2dayyy6
[link] [留言]

来源:r/MachineLearning · reddit.com