Optimization: Convert Hidden States to Bfloat16
- Ngôn ngữ chính
- C++
- Star
- 7k
- Fork
- 660
- Merge trung bình
- 20 giờ 43 phút
- Pull request đã merge (30 ngày)
- 33
Mô tả
In line with #543 and #164 I have worked a bit with changing the precision of the `pre_att_rms_out` from float to BF16. First, as also suggested, I made only a tiny code change. The patch would be:
```
diff --git a/gemma/activations.h b/gemma/activations.h
index 86345e2..a994e98 100644
--- a/gemma/activations.h
+++ b/gemma/activations.h
@@ -41,7 +41,7 @@ struct Activations {
RowVectorBatch logits;
// Attention
- RowVectorBatch pre_att_rms_out;
+ RowVectorBatch pre_att_rms_out;
RowVectorBatch att; // attention vector
RowVectorBatch att_out; // attention output
// Accumulation of attention outputs over heads
@@ -88,7 +88,7 @@ struct Activations {
logits = RowVectorBatch(Extents2D(batch_size, vocab_size));
}
- pre_att_rms_out = RowVectorBatch(Extents2D(batch_size, model_dim));
+ pre_att_rms_out = RowVectorBatch(Extents2D(batch_size, model_dim));
att = RowVectorBatch(
Extents2D(batch_size, heads * weights_config.seq_len));
att_out = RowVectorBatch(Extents2D(batch_size, heads * qkv_dim));
diff --git a/gemma/gemma-inl.h b/gemma/gemma-inl.h
index ccb34f0..e7eece4 100644
--- a/gemma/gemma-inl.h
+++ b/gemma/gemma-inl.h
@@ -293,7 +293,7 @@ class GemmaAttention {
// Proceed row by row because there will be wraparound.
for (size_t interleaved_idx = 0; interleaved_idx < num_interleaved;
++interleaved_idx) {
- const float* x = activations_.pre_att_rms_out.Batch(interleaved_idx);
+ const hwy::bfloat16_t* x = activations_.pre_att_rms_out.Batch(interleaved_idx);
const size_t query_idx = interleaved_idx % num_queries_;
const size_t batch_idx = interleaved_idx / num_queries_;
KVCache& kv_cache = kv_caches_[query_idx];
```
I wonder if this change was too small?
My (unscientific) measurements showed no speed gain in doing so. I have tried it on an x86 laptop that doesn't support BF16 and on an ARM chip that does (Cortex X3, inside a tablet). I am a bit surprised that the latency hasn't improved.
I understand the matmul is done with bf16 types. The decompression of B seems to take up a major chunk of the runtime, whilst "decompression" (depending on the type of pre_att_rms_out) of A is unnoticeable. The profiler says:
```
MM.NT.DecB : 140318080 x 26 = 190.491139
MM.NT_K.DecB : 12238848 x 121 = 77.107746
MM.NT : 107076 x 10869 = 60.612900
Gen.input : 2 x 485175727 = 50.539138
MM.NT_K : 21248 x 15169 = 16.787086
...
MM.DecompressA : 128324 x 3 = 0.022761
```
Which makes me wonder if B should be loaded in BF16 instead of decompressing it? If the precision of pre_att_rms_out predominantly affects the decompression time and the RMSNorm computation time, maybe it's unsurprising that the latency effects are not of first order.
Hướng dẫn đóng góp
Hướng nghiên cứu
Bắt đầu với gemma/activations.h và gemma/gemma-inl.h, nơi thay đổi BF16 được đề xuất được thể hiện. Tái hiện các phép đo độ trễ trên các nền tảng x86 và ARM, đồng thời so sánh các mục profiler cho việc giải nén và phép nhân ma trận. Được xem là hoàn tất khi xác định được liệu việc chuyển đổi các trạng thái ẩn hoặc thay đổi cách tải B có tạo ra cải thiện đáng kể có thể đo được hay không.
Do mô hình lập chỉ mục viết ra từ nội dung của issue.
Đánh giá
- Công nghệ
- cpp
- Lĩnh vực
- machine-learning, performance
- Loại issue
- Tái cấu trúc
- Độ khó
- 4/5
- Thời gian dự kiến
- 3-5 ngày
- Mức độ hoạt động
- Đình trệ
- Độ rõ ràng
- Khá rõ ràng
- Mức phù hợp với người mới
- 35/100