llama : support Jamba hybrid Transformer-Mamba models #7531
Conversation
This will be necessary to support Jamba (and other recurrent models mixed with Attention). Doesn't compile yet, and finding a slot isn't yet done correctly for recurrent states.
* llama : begin work on support for variable GQA This will also be useful for Jamba if we consider the Mamba layers to have 0 KV heads. * llama : gracefully fail when not finding hybrid slot
compilade
added
enhancement
llama.cpp
Outdated
| switch (hparams.n_layer) { | ||
| // TODO: Jamba layers are a bit heterogenous, so naming this is hard. | ||
| case 12: // 900M 8x???M | ||
| case 32: // 51B 16x?B | ||
| default: model.type = e_model::MODEL_UNKNOWN; |
There was a problem hiding this comment.
I'm not sure what model size type(s) I should give to Jamba models.
github-actions
bot
added
the
ggml
|
📈 llama.cpp server for bench-server-baseline on Standard_NC4as_T4_v3 for Expand details for performance related PR only
More---
config:
xyChart:
titleFontSize: 12
width: 900
height: 600
themeVariables:
xyChart:
titleColor: "#000000"
---
xychart-beta
title "llama.cpp bench-server-baseline on Standard_NC4as_T4_v3
duration=10m 557 iterations"
y-axis "llamacpp:prompt_tokens_seconds"
x-axis "llamacpp:prompt_tokens_seconds" 1717475210 --> 1717475834
line [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 306.61, 306.61, 306.61, 306.61, 306.61, 572.5, 572.5, 572.5, 572.5, 572.5, 579.51, 579.51, 579.51, 579.51, 579.51, 601.73, 601.73, 601.73, 601.73, 601.73, 638.34, 638.34, 638.34, 638.34, 638.34, 702.62, 702.62, 702.62, 702.62, 702.62, 704.56, 704.56, 704.56, 704.56, 704.56, 718.91, 718.91, 718.91, 718.91, 718.91, 723.54, 723.54, 723.54, 723.54, 723.54, 739.59, 739.59, 739.59, 739.59, 739.59, 771.46, 771.46, 771.46, 771.46, 771.46, 802.48, 802.48, 802.48, 802.48, 802.48, 815.12, 815.12, 815.12, 815.12, 815.12, 804.65, 804.65, 804.65, 804.65, 804.65, 797.38, 797.38, 797.38, 797.38, 797.38, 800.86, 800.86, 800.86, 800.86, 800.86, 805.61, 805.61, 805.61, 805.61, 805.61, 803.64, 803.64, 803.64, 803.64, 803.64, 824.04, 824.04, 824.04, 824.04, 824.04, 823.3, 823.3, 823.3, 823.3, 823.3, 830.32, 830.32, 830.32, 830.32, 830.32, 832.47, 832.47, 832.47, 832.47, 832.47, 846.38, 846.38, 846.38, 846.38, 846.38, 842.07, 842.07, 842.07, 842.07, 842.07, 844.76, 844.76, 844.76, 844.76, 844.76, 861.96, 861.96, 861.96, 861.96, 861.96, 855.54, 855.54, 855.54, 855.54, 855.54, 854.58, 854.58, 854.58, 854.58, 854.58, 856.84, 856.84, 856.84, 856.84, 856.84, 860.17, 860.17, 860.17, 860.17, 860.17, 858.21, 858.21, 858.21, 858.21, 858.21, 861.33, 861.33, 861.33, 861.33, 861.33, 871.29, 871.29, 871.29, 871.29, 871.29, 847.29, 847.29, 847.29, 847.29, 847.29, 832.73, 832.73, 832.73, 832.73, 832.73, 831.59, 831.59, 831.59, 831.59, 831.59, 831.76, 831.76, 831.76, 831.76, 831.76, 835.52, 835.52, 835.52, 835.52, 835.52, 836.15, 836.15, 836.15, 836.15, 836.15, 836.37, 836.37, 836.37, 836.37, 836.37, 817.57, 817.57, 817.57, 817.57, 817.57, 820.16, 820.16, 820.16, 820.16, 820.16, 820.49, 820.49, 820.49, 820.49, 820.49, 820.0, 820.0, 820.0, 820.0, 820.0, 817.08, 817.08, 817.08, 817.08, 817.08, 820.83, 820.83, 820.83, 820.83, 820.83, 823.82, 823.82, 823.82, 823.82, 823.82, 823.03, 823.03, 823.03, 823.03, 823.03, 827.7, 827.7, 827.7, 827.7, 827.7, 826.96, 826.96, 826.96, 826.96, 826.96, 833.12, 833.12, 833.12, 833.12, 833.12, 832.75, 832.75, 832.75, 832.75, 832.75, 832.65, 832.65, 832.65, 832.65, 832.65, 826.23, 826.23, 826.23, 826.23, 826.23, 827.38, 827.38, 827.38, 827.38, 827.38, 827.43, 827.43, 827.43, 827.43, 827.43, 827.46, 827.46, 827.46, 827.46, 827.46, 825.87, 825.87, 825.87, 825.87, 825.87, 828.84, 828.84, 828.84, 828.84, 828.84, 829.05, 829.05, 829.05, 829.05, 829.05, 829.15, 829.15, 829.15]
More---
config:
xyChart:
titleFontSize: 12
width: 900
height: 600
themeVariables:
xyChart:
titleColor: "#000000"
---
xychart-beta
title "llama.cpp bench-server-baseline on Standard_NC4as_T4_v3
duration=10m 557 iterations"
y-axis "llamacpp:predicted_tokens_seconds"
x-axis "llamacpp:predicted_tokens_seconds" 1717475210 --> 1717475834
line [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 42.1, 42.1, 42.1, 42.1, 42.1, 30.42, 30.42, 30.42, 30.42, 30.42, 28.2, 28.2, 28.2, 28.2, 28.2, 28.69, 28.69, 28.69, 28.69, 28.69, 29.63, 29.63, 29.63, 29.63, 29.63, 30.55, 30.55, 30.55, 30.55, 30.55, 32.02, 32.02, 32.02, 32.02, 32.02, 32.76, 32.76, 32.76, 32.76, 32.76, 33.41, 33.41, 33.41, 33.41, 33.41, 33.56, 33.56, 33.56, 33.56, 33.56, 34.05, 34.05, 34.05, 34.05, 34.05, 33.99, 33.99, 33.99, 33.99, 33.99, 33.35, 33.35, 33.35, 33.35, 33.35, 33.38, 33.38, 33.38, 33.38, 33.38, 32.25, 32.25, 32.25, 32.25, 32.25, 31.71, 31.71, 31.71, 31.71, 31.71, 30.36, 30.36, 30.36, 30.36, 30.36, 30.81, 30.81, 30.81, 30.81, 30.81, 30.82, 30.82, 30.82, 30.82, 30.82, 30.39, 30.39, 30.39, 30.39, 30.39, 30.41, 30.41, 30.41, 30.41, 30.41, 30.5, 30.5, 30.5, 30.5, 30.5, 30.85, 30.85, 30.85, 30.85, 30.85, 30.97, 30.97, 30.97, 30.97, 30.97, 31.24, 31.24, 31.24, 31.24, 31.24, 31.45, 31.45, 31.45, 31.45, 31.45, 31.23, 31.23, 31.23, 31.23, 31.23, 31.18, 31.18, 31.18, 31.18, 31.18, 31.36, 31.36, 31.36, 31.36, 31.36, 31.43, 31.43, 31.43, 31.43, 31.43, 31.63, 31.63, 31.63, 31.63, 31.63, 31.71, 31.71, 31.71, 31.71, 31.71, 31.78, 31.78, 31.78, 31.78, 31.78, 31.61, 31.61, 31.61, 31.61, 31.61, 31.48, 31.48, 31.48, 31.48, 31.48, 31.35, 31.35, 31.35, 31.35, 31.35, 31.43, 31.43, 31.43, 31.43, 31.43, 31.54, 31.54, 31.54, 31.54, 31.54, 31.71, 31.71, 31.71, 31.71, 31.71, 31.79, 31.79, 31.79, 31.79, 31.79, 31.85, 31.85, 31.85, 31.85, 31.85, 31.71, 31.71, 31.71, 31.71, 31.71, 31.42, 31.42, 31.42, 31.42, 31.42, 31.06, 31.06, 31.06, 31.06, 31.06, 29.65, 29.65, 29.65, 29.65, 29.65, 29.37, 29.37, 29.37, 29.37, 29.37, 29.37, 29.37, 29.37, 29.37, 29.37, 29.4, 29.4, 29.4, 29.4, 29.4, 29.46, 29.46, 29.46, 29.46, 29.46, 29.58, 29.58, 29.58, 29.58, 29.58, 29.61, 29.61, 29.61, 29.61, 29.61, 29.57, 29.57, 29.57, 29.57, 29.57, 29.58, 29.58, 29.58, 29.58, 29.58, 29.45, 29.45, 29.45, 29.45, 29.45, 29.55, 29.55, 29.55, 29.55, 29.55, 29.69, 29.69, 29.69, 29.69, 29.69, 29.83, 29.83, 29.83, 29.83, 29.83, 29.9, 29.9, 29.9, 29.9, 29.9, 29.96, 29.96, 29.96, 29.96, 29.96, 29.97, 29.97, 29.97, 29.97, 29.97, 30.03, 30.03, 30.03]
Details
More---
config:
xyChart:
titleFontSize: 12
width: 900
height: 600
themeVariables:
xyChart:
titleColor: "#000000"
---
xychart-beta
title "llama.cpp bench-server-baseline on Standard_NC4as_T4_v3
duration=10m 557 iterations"
y-axis "llamacpp:kv_cache_usage_ratio"
x-axis "llamacpp:kv_cache_usage_ratio" 1717475210 --> 1717475834
line [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.14, 0.14, 0.14, 0.14, 0.14, 0.37, 0.37, 0.37, 0.37, 0.37, 0.25, 0.25, 0.25, 0.25, 0.25, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.14, 0.14, 0.14, 0.14, 0.14, 0.15, 0.15, 0.15, 0.15, 0.15, 0.14, 0.14, 0.14, 0.14, 0.14, 0.12, 0.12, 0.12, 0.12, 0.12, 0.16, 0.16, 0.16, 0.16, 0.16, 0.25, 0.25, 0.25, 0.25, 0.25, 0.26, 0.26, 0.26, 0.26, 0.26, 0.12, 0.12, 0.12, 0.12, 0.12, 0.36, 0.36, 0.36, 0.36, 0.36, 0.36, 0.36, 0.36, 0.36, 0.36, 0.41, 0.41, 0.41, 0.41, 0.41, 0.14, 0.14, 0.14, 0.14, 0.14, 0.15, 0.15, 0.15, 0.15, 0.15, 0.23, 0.23, 0.23, 0.23, 0.23, 0.2, 0.2, 0.2, 0.2, 0.2, 0.19, 0.19, 0.19, 0.19, 0.19, 0.16, 0.16, 0.16, 0.16, 0.16, 0.19, 0.19, 0.19, 0.19, 0.19, 0.11, 0.11, 0.11, 0.11, 0.11, 0.13, 0.13, 0.13, 0.13, 0.13, 0.32, 0.32, 0.32, 0.32, 0.32, 0.21, 0.21, 0.21, 0.21, 0.21, 0.1, 0.1, 0.1, 0.1, 0.1, 0.12, 0.12, 0.12, 0.12, 0.12, 0.13, 0.13, 0.13, 0.13, 0.13, 0.13, 0.13, 0.13, 0.13, 0.13, 0.23, 0.23, 0.23, 0.23, 0.23, 0.21, 0.21, 0.21, 0.21, 0.21, 0.28, 0.28, 0.28, 0.28, 0.28, 0.3, 0.3, 0.3, 0.3, 0.3, 0.2, 0.2, 0.2, 0.2, 0.2, 0.09, 0.09, 0.09, 0.09, 0.09, 0.14, 0.14, 0.14, 0.14, 0.14, 0.13, 0.13, 0.13, 0.13, 0.13, 0.16, 0.16, 0.16, 0.16, 0.16, 0.45, 0.45, 0.45, 0.45, 0.45, 0.56, 0.56, 0.56, 0.56, 0.56, 0.56, 0.56, 0.56, 0.56, 0.56, 0.64, 0.64, 0.64, 0.64, 0.64, 0.36, 0.36, 0.36, 0.36, 0.36, 0.21, 0.21, 0.21, 0.21, 0.21, 0.2, 0.2, 0.2, 0.2, 0.2, 0.17, 0.17, 0.17, 0.17, 0.17, 0.13, 0.13, 0.13, 0.13, 0.13, 0.11, 0.11, 0.11, 0.11, 0.11, 0.29, 0.29, 0.29, 0.29, 0.29, 0.27, 0.27, 0.27, 0.27, 0.27, 0.24, 0.24, 0.24, 0.24, 0.24, 0.19, 0.19, 0.19, 0.19, 0.19, 0.13, 0.13, 0.13, 0.13, 0.13, 0.12, 0.12, 0.12, 0.12, 0.12, 0.17, 0.17, 0.17, 0.17, 0.17, 0.15, 0.15, 0.15, 0.15, 0.15, 0.22, 0.22, 0.22, 0.22, 0.22, 0.12, 0.12, 0.12, 0.12, 0.12, 0.16, 0.16, 0.16]
More---
config:
xyChart:
titleFontSize: 12
width: 900
height: 600
themeVariables:
xyChart:
titleColor: "#000000"
---
xychart-beta
title "llama.cpp bench-server-baseline on Standard_NC4as_T4_v3
duration=10m 557 iterations"
y-axis "llamacpp:requests_processing"
x-axis "llamacpp:requests_processing" 1717475210 --> 1717475834
line [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 3.0, 3.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0, 4.0, 4.0, 7.0, 7.0, 7.0, 7.0, 7.0, 4.0, 4.0, 4.0, 4.0, 4.0, 3.0, 3.0, 3.0, 3.0, 3.0, 5.0, 5.0, 5.0, 5.0, 5.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 5.0, 5.0, 5.0, 5.0, 5.0, 4.0, 4.0, 4.0, 4.0, 4.0, 3.0, 3.0, 3.0, 3.0, 3.0, 6.0, 6.0, 6.0, 6.0, 6.0, 4.0, 4.0, 4.0, 4.0, 4.0, 8.0, 8.0, 8.0, 8.0, 8.0, 6.0, 6.0, 6.0, 6.0, 6.0, 7.0, 7.0, 7.0, 7.0, 7.0, 5.0, 5.0, 5.0, 5.0, 5.0, 4.0, 4.0, 4.0, 4.0, 4.0, 8.0, 8.0, 8.0, 8.0, 8.0, 4.0, 4.0, 4.0, 4.0, 4.0, 6.0, 6.0, 6.0, 6.0, 6.0, 6.0, 6.0, 6.0, 6.0, 6.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 2.0, 2.0, 4.0, 4.0, 4.0, 4.0, 4.0, 5.0, 5.0, 5.0, 5.0, 5.0, 6.0, 6.0, 6.0, 6.0, 6.0, 2.0, 2.0, 2.0, 2.0, 2.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 4.0, 4.0, 4.0, 4.0, 4.0, 2.0, 2.0, 2.0, 2.0, 2.0, 4.0, 4.0, 4.0, 4.0, 4.0, 1.0, 1.0, 1.0, 1.0, 1.0, 5.0, 5.0, 5.0, 5.0, 5.0, 2.0, 2.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 3.0, 3.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 8.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 4.0, 7.0, 7.0, 7.0, 7.0, 7.0, 6.0, 6.0, 6.0, 6.0, 6.0, 2.0, 2.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 3.0, 3.0, 2.0, 2.0, 2.0, 2.0, 2.0, 4.0, 4.0, 4.0, 4.0, 4.0, 6.0, 6.0, 6.0, 6.0, 6.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0, 2.0, 2.0, 2.0, 2.0, 2.0, 5.0, 5.0, 5.0, 5.0, 5.0, 8.0, 8.0, 8.0, 8.0, 8.0, 5.0, 5.0, 5.0, 5.0, 5.0, 6.0, 6.0, 6.0, 6.0, 6.0, 3.0, 3.0, 3.0]
|
|
Great job! Works for me too, it's very fast. There were some warnings during compilation, but nothing major.
|
|
Amazing work! |
ggml.c
Outdated
| if (n_rs > 1) { | ||
| // multiple sequences means it's hard to know when it's the first time a state is read, | ||
| // so copy them all over to the destination, just to be sure. | ||
| for (int i3 = 0; i3 < n_kv; ++i3) { | ||
| for (int i3 = 0; i3 < n_rs; ++i3) { |
There was a problem hiding this comment.
I'm looking at adding the missing Metal kernels for SSM_CONV and SSM_SCAN. I'm wondering if this part of the kernels where we copy src0 -> dst could be extracted outside of the operation via ggml_cpy + ggml_view or ggml_acc? Would simplify the implementation
Also, I still haven't understood the details of the computation, but if we find a way to express these ops via existing ops all together (e.g. using ggml_conv, ggml_mul_mat, ...), it would be preferred to do so, in order to reduce the amount of kernels that we have to write.
There was a problem hiding this comment.
I'm wondering if this part of the kernels where we copy
src0->dstcould be extracted outside of the operation viaggml_cpy+ggml_vieworggml_acc? Would simplify the implementation
Yes, this is definitely possible. I'll find a way to extract the copies outside.
if we find a way to express these ops via existing ops all together (e.g. using ggml_conv, ggml_mul_mat, ...), it would be preferred to do so, in order to reduce the amount of kernels that we have to write.
For SSM_SCAN, I think there's a way to fully express it in terms of other ops, though it will use much more memory because of the big intermediate tensors, and new operators like SOFT_PLUS and EXP would be needed instead. But different lengths of simultaneous sequences might make a custom operator still necessary. I'll think about ways to make it simpler, especially since other recurrent architectures (like RWKV) will also need to work on multiple sequences per batch.
For simplifying SSM_CONV, I don't think ggml_conv supports working on independent 1D rolling windows with varying sequence lengths.
When working on a single sequence, though, it's quite simple to do the equivalent of ggml_ssm_conv with a self-overlapping view, as I did in my original implementation which I described in more detail in #5328 (comment):
Setting nb[2] to the element size makes the view self-overlapping.
But this would create too many nodes in the compute graph when done with multiple sequences (unless they're always all the same length in which case the 4th dimension could be used), so a custom operator is necessary.
There was a problem hiding this comment.
One idea that we might consider is to unfuse the n_rs dimension from the SSM ops and make them work per 1 recurrent state. Then, during inference and right before the SSM operations, we split the batch into same-sequence chunks and SSM them individually. After that we concat back the results into the full hidden state for the batch
The main goal would be to simplify the SSM operators, and potentially express them as other existing ops if possible. But additionally, I'm considering a similar processing mode for the standard transformer KV cache in which we don't rely on a "unified" buffer for all the sequences, but instead each sequence has it's own separate KV cache buffer. In that mode, we would do a similar same-sequence batch split before the attention. The main purpose of supporting this mode would be to achieve reproducible results during parallel decoding (currently, decoding the same sequence in parallel can yield slightly different results due to the unified KV cache).
Just throwing some thoughts that I have so far - will continue looking at the PR in the next days
Edit: I was writing this comment before I saw you posted - will take a look tomorrow
There was a problem hiding this comment.
One idea that we might consider is to unfuse the
n_rsdimension from the SSM ops and make them work per 1 recurrent state. Then, during inference and right before the SSM operations, we split the batch into same-sequence chunks and SSM them individually. After that we concat back the results into the full hidden state for the batch
Yes, this would be doable, but would make the number of compute graph nodes scale with the number of sequences. (EDIT: if it's split when making ubatches, then the number of compute graph nodes can stay constant)
Another way would be to make all sequences have the same number of new tokens in a ubatch, to allow using another dimension instead of having to loop when building the compute graphs. This would still allow batching multiple sequences with recurrent models, but without the need for new custom operators for each architecture, and still with a constant number of compute graph nodes.
The recurrent steps are simpler for ubatches with sequence lengths of 1, but prompt processing performance would be much slower than with a per-recurrent-architecture operator for longer sequences. Still thinking about ways to generalize this while keeping good performance.
But additionally, I'm considering a similar processing mode for the standard transformer KV cache in which we don't rely on a "unified" buffer for all the sequences, but instead each sequence has it's own separate KV cache buffer. In that mode, we would do a similar same-sequence batch split before the attention.
For the transformer KV cache, if there's logic to make all sequences within a ubatch to have the same number of new tokens, I think a mode to split batches sequence-wise will be simpler and could re-use much of the same code.
I also think there's a way to keep the unified KV cache (one buffer) and chunk it to make each sequence have their own independent contiguous reserved cells. Batching sequences together might still be possible though, if the KQ mask gets another dimension (the number of sequences in the ubatch, and the number of new tokens per sequence instead of the batch size) so that these equal-sized "chunks" get processed independently in parallel. But this might not work (because the newly-calculated KV cells have to be copied in a bunch of not-regularly-spaced places), unless... unless maybe with some kind of ggml_set_rows? Not sure about the transposed V cache, though.
A sequence-wise processing mode is likely simpler, although it's not really parallel processing then (the model weights are all read at each ubatch).
There was a problem hiding this comment.
if it's split when making ubatches, then the number of compute graph nodes can stay constant
No, it has to be split only for the attention so that the rest of the ops are still batched. Otherwise we will sacrifice a lot of performance
Another way would be to make all sequences have the same number of new tokens in a ubatch, to allow using another dimension instead of having to loop when building the compute graphs. This would still allow batching multiple sequences with recurrent models, but without the need for new custom operators for each architecture, and still with a constant number of compute graph nodes.
For the transformer KV cache, if there's logic to make all sequences within a ubatch to have the same number of new tokens, I think a mode to split batches sequence-wise will be simpler and could re-use much of the same code.
Not sure how that would work. Adding dummy tokens sounds too much overhead (at least
in the case of the regular transformer). Any other ideas?
A sequence-wise processing mode is likely simpler, although it's not really parallel processing then (the model weights are all read at each ubatch).
From a broad PoV, if we have an implementation that works with a single-sequence and any batch size, then to extend it to multi-sequence batches we can split the batch into same-sequence tokens right before the attention and merge it back after the attention. Each split will do what we already do for the single-sequence solution, using separate cache for each sequence. I didn't consider the number of nodes until you noted - so that might be a problem indeed.
I'm currently working on a big refactor of how Mamba (and Jamba) works to make all sequences of a sub-batch be of the same length (initially only for models with recurrent states), and to make recurrent state slots contiguous, with the goal of simplifying the SSM operations (and removing GGML_OP_SSM_CONV), so that GPU support will be much easier to implement after that.
Looking forward to this!
There was a problem hiding this comment.
No, it has to be split only for the attention so that the rest of the ops are still batched. Otherwise we will sacrifice a lot of performance
It will sacrifice some performance, but only in the cases where a batch contains an unequal number of tokens for each affected sequence. So this should not affect large prompt processing or parallel text generation, if both are not done in the same batch.
Not sure how that would work. Adding dummy tokens sounds too much overhead (at least
in the case of the regular transformer). Any other ideas?
This is not about adding dummy tokens, but about making the number of new tokens in each ubatch the same per sequence. I think the overhead will be minmal, though there is still some.
Let me illustrate.
Let's say there's a batch with new tokens for 4 sequences of length 16, 7, 1, 1, respectively.
0: ################
1: #######
2: #
3: #
Splitting that into equal-length sequences would make 3 ubatches, like so:
0: #
1: #
2: #
3: #
0: ######
1: ######
0: #########
Each of these shapes are nice and rectangular, which is good for recurrent architectures because their operations can be more easily batched across sequences this way.
But I'm not yet sure if it would also benefit Transformers, which is why I'm thinking of initially only enabling the equal-length splitting for recurrent (or hybrid) model architectures.
From a broad PoV, if we have an implementation that works with a single-sequence and any batch size, then to extend it to multi-sequence batches we can split the batch into same-sequence tokens right before the attention and merge it back after the attention. Each split will do what we already do for the single-sequence solution, using separate cache for each sequence. I didn't consider the number of nodes until you noted - so that might be a problem indeed.
Doing this with a constant number of graph nodes is pretty much what using same-length sequences (as illustrated above) allows, because the split into same-sequence tokens can then simply become another tensor dimension.
There was a problem hiding this comment.
Aha, got it. Good idea. I'm also not sure if this can help Transformers, but it's something to think about 👍
|
Progress? |
|
@theogbob , you may tag author, @compilade, which progress? ) |
|
Progress @compilade ? |
|
@compilade Thanks for all the work here! I've also been working through a very similar architecture for bamba independently. I suspect your implementation of It looks like this branch is pretty out of date with the latest refactors in the codebase. I have a version of my branch that I got working against the rebased tip of your We just released an updated V2 of |
@gabe-l-hart Another approach like using per-layer cache types would need considerable additional refactoring which would conflict even more with #12799 (although it might simplify #13194). (Now that I write this out, you're making me realize that all the kv-cache needs for hybrid models is per-type (e.g. self-attention and recurrent) top-level metadata (the cells) and some data buffers (of which there seem to always be up to 2 per layer (k and v, or r and s), since no layer ever has both Attention and recurrent states (at least this seems true for the hybrid models I've seen so far)). That is pretty much what is implemented here with the zero-sized layers, but this hints towards possible future simplifications (which will be doable after resolving conflicts from #12799 here).)
I too would be interested in consolidating with your work, or at least making it easier for you to get Bamba supported. How would you prefer this to happen? Note that I will update the
Yes, this branch is not very up to date, but it's fixable. The main reason close to no progress was being made here was because I don't find it particularly fun to resolve thousands of lines of conflicts. Or at least I need to dedicate a good chunk of time to that so that I don't get lost half-way (since the conflict resolutions of this size mean mostly re-thinking the approach and porting it to the new structures). So this PR might have been neglected for a while because the moments where I had enough time and the moments where I wanted to fix this and/or reply1 to "progress?" comments did not align. But I am in a period where I'm starting to have more spare time, and so I could dedicate a day (or more) to resolve the conflicts here and in #9126 (but I suspect it's going to take more than a day).
That's awesome!
When branches drift that much, merging is usually simpler to handle than rebasing and still leaves a trail of tested versions, and allows resolving conflicts once (per merge) instead of at every commit which change conflicting parts. But I see what you mean, and I'd love to get help with the conflict resolutions, but it's unfortunately something which almost has to be done in one go (because
Right, the Mamba2 branch (in #9126) modifies a bit how the SSM operator works (to minimize useless copies), and that will need to be adapted to the CUDA version of the operator which was added in #10558. Footnotes
|
|
Thank you for the detailed response! It's really helpful. I 100% hear you on the giant merge conflicts, and I agree at this stage merging is better than rebasing. I spent yesterday trying to resolve It looks like the biggest change since I last synced is around moving to more abstract interfaces for things. In particular, it looks like all caching has moved behind the As always, thanks for the outstanding work here, 3-hour comments included! |
|
Ok, I found my merge bugs in https://github.com/gabe-l-hart/llama.cpp/tree/BambaAbstractMemory and I'm now able to run a lightweight As a separate question, this probably isn't the right place to centralize this discussion. Would it be best to create a central issue to discuss the convergence of Details# Download lightweight mamba2 model
huggingface-cli download AntonV/mamba2-370m-hf --local-dir ~/models/mamba2-370m-hf
# Convert to GGUF
python convert_hf_to_gguf.py ~/models/mamba2-370m-hf/
# Run a sample query
./build/bin/llama-cli -m ~/models/mamba2-370m-hf/mamba2-370M-hf-F16.gguf -p "Hello world" -ngl 0 --temp 0 -n 20 |
@gabe-l-hart Amazing! I've also merged from latest It's very helpful to compare both merges to compare the approaches1 (and sometimes notice when changes are missing). It does reduce the stress of a bad merge. Thank you! (although it seems like Multi-sequence inference is broken, though (that's also true on $ ./build/bin/llama-parallel -m ~/models/mamba2-370m-hf/mamba2-370M-hf-F16.gguf -np 5 -ns 8 --temp 0 --repeat-penalty 1.1Part of the problem is caused by an early
Yes, I think that should be more appropriate. It's true that technically Mamba2 isn't directly related to Jamba. If I create the issue, I will tag you and refer to the relevant PRs and issues. Footnotes
|
|
@compilade Great to hear that you got the merge working, and not at all surprised that I missed some nuance beyond basic single-sequence generation. I'll look to pick up your changes on my branch.
I've never used I did also start taking a whack at the hybrid cache based on the new layers of abstraction in The trickiest part seems to be the intermixing of |
|
It looks like the work of hoisting the cache abstraction is almost all done in #12799! I'll move to build off of that branch. |
compilade
force-pushed
the
compilade/refactor-kv-cache
branch
from
aa4039d to
2bcaf64
Compare
There was a problem hiding this comment.
I've updated this to use the hybrid cache implementation from #13979. As discussed a while ago, state checkpoints for recurrent state rollbacks will be implemented in a separate PR (which I did not begin).
This seems to still work; I've tested a previous conversion of https://huggingface.co/pszemraj/jamba-900M-v0.13-KIx2, and a new conversion, and both work. I will test the official Jamba-Mini-1.6 in the next days.
I've also shortened the main description of the PR, since the scope has reduced a lot over time (most of the changes were split into other pull requests which were merged in the past year).
| def get_vocab_base_pre(self, tokenizer) -> str: | ||
| del tokenizer # unused | ||
|
|
||
| return "gpt-2" |
There was a problem hiding this comment.
This pre-tokenizer override is pretty much only used by https://huggingface.co/pszemraj/jamba-900M-v0.13-KIx2.
The official Jamba models and finetunes use a sentencepiece tokenizer.model.
src/llama-model.cpp
Outdated
| const auto * mctx_hyb = static_cast<const llama_memory_hybrid_context *>(mctx); | ||
|
|
||
| auto * inp_rs = build_rs_inp(mctx_hyb->get_recr()); | ||
|
|
||
| auto * inp_attn = build_attn_inp_kv_unified(mctx_hyb->get_attn()); |
There was a problem hiding this comment.
build_rs and build_attn now use inp->mctx instead of casting this->mctx again, and so this removes the need for build_inp_mem_hybrid and hybrid-specific overloads of build_rs and build_attn, because the bare build_rs_inp and build_attn_inp_kv_unified can now be used directly instead by giving them the correct mctx to use (defaults to casting from this->mctx, and so existing model graph builders don't need to change).
This will also potentially make it easier to support hybrid models with sliding-window attention.
There was a problem hiding this comment.
This is a good change. Thanks.
There was a problem hiding this comment.
Hm, however this will not be compatible with the llama_graph_result_i::update(mctx) mechanism from #14482 for reusing compute graphs. The idea is when we can reuse the old graph, to call res->update(mctx) with the new memory context in order to make all inputs from the previous graph result point to that new memory context.
In this case, we will call update([hybrid_context]) and the rs input wouldn't know that it needs to get the recurrent sub-context.
There was a problem hiding this comment.
@gabe-l-hart Heads up I am still contemplating this change. If we can't figure a reasonable alternative for the mechanism in #14482 we will have to go back to using a separate hybrid memory input as on master.
There was a problem hiding this comment.
Makes sense. I think we could probably make the virtual inheritance thing work even with the separate hybrid memory input, so it probably still makes sense to put GR4 after Jamba.
There was a problem hiding this comment.
If I'm reading the changes in #14482 correctly, the issue you're describing would come up here when llm_graph_result::update gets called and in turn calls update(mctx) on each of the input types. In this case, the mctx would be a hybrid context, but the input would be expecting one of the child context types. I recall trying to add cast operators to the hybrid context so that it could be cast to the child types, but I think that caused problems because it didn't inherit from them. It seems like it might make sense for the hybrid cache to have some mechanism for acting like either of the child types. I'll play around a little and see if I can make a dummy version.
There was a problem hiding this comment.
I the simplest solution (but one that would violate encapsulation pretty badly) would be to use dynamic_cast on mctx in the implementation of update for the input types that might need to consume different context types.
void use_a(const Interface* i) {
// const TypA * a = static_cast<const TypA *>(i);
const TypA * a = dynamic_cast<const TypA *>(i);
if (!a) {
a = static_cast<const TypAB *>(i)->get_a();
}
std::cout << "----" << std::endl;
std::cout << "Using A" << std::endl;
a->doit();
a->just_a();
}There was a problem hiding this comment.
(not suggesting this as a final solution, just brainstorming)
There was a problem hiding this comment.
Another not-that-good option would be to handle this with multiple inheritance. We could create an intermediate base class derived from llama_memory_context_i that holds the common members between the non-hybrid cache types (ubatches, status, possibly i_*), then use virtual inheritance for kv_cache_unified and recurrent from there. This would allow hybrid to inherit from both (rather than "hasa" both). It would also avoid the ubatch vector copying that's currently happening in both hybrid and iswa.
There was a problem hiding this comment.
If we can't figure a reasonable alternative for the mechanism in #14482 we will have to go back to using a separate hybrid memory input as on
master.
@ggerganov In that case, I think the hybrid graph input could be re-added, but it should contain sub-cache graph inputs, without adding them to res->inputs. This would still allow the mechanism in #14482 to be used (since the hybrid input expects an hybrid mctx, and can update its sub-inputs), and it would also continue to allow passing the sub-cache graph inputs to build_rs and build_attn (to avoid unnecessary overloads).
I've implemented this in 20f8e43.
| @@ -303,33 +303,6 @@ class llm_graph_input_attn_cross : public llm_graph_input_i { | |||
| const llama_cross * cross = nullptr; | |||
| }; | |||
|
|
|||
| class llm_graph_input_mem_hybrid : public llm_graph_input_i { | |||
There was a problem hiding this comment.
Nice, I really like removing these duplicate inputs in favor of allowing the caller to explicitly fetch the hybrid caches' children and pass them to the creation methods.
But this time it contains the sub-cache graph inputs. This *should* make it easier to handle updating the inputs when caching the graph (eventually).
| static std::unique_ptr<llm_graph_input_attn_kv_unified> build_attn_inp_kv_unified_impl( | ||
| ggml_context * ctx0, | ||
| const llama_ubatch & ubatch, | ||
| const llama_hparams & hparams, | ||
| const llama_cparams & cparams, | ||
| const llama_kv_cache_unified_context * mctx_cur) { |
There was a problem hiding this comment.
(from changes in 20f8e43)
I'm not entirely sure where would be the most convenient place for shared graph input building functions, so I've used static functions with an _impl suffix.
This is used both in llm_graph_context::build_attn_inp_kv_unified() and in llm_graph_context::build_inp_mem_hybrid().
I guess they could be added to llm_graph_context as private methods, or somewhere else.
This adds support for Jamba (fixes #6372). (https://arxiv.org/abs/2403.19887)
(this has been open for a while, and this description was very different originally (much broader scope), feel free to look at the edit history)
New features
llama.cppState checkpoints for recurrent modelsWorks best whenn_parallelis at least 3 or 4 times the number of actual usersAllows backtracking tokens from the end of the last generation without having to reprocess the whole contextVery useful with theserverexample when trimming the stop string{model}.attention.head_count_kvcan now also be an array ofint32_t, one value per layer0kv heads are considered recurrent layers (Mamba, in the case of Jamba).Internal changes
build_mamba_layerfunctions to a shared parent class between bothllm_build_mambaandllm_build_jamba.llm_graph_context::build_inp_mem_hybridllm_graph_input_mem_hybridllm_graph_input_rsandllm_graph_input_attn_kv_unified, and causes unnecessary duplication and overloads ofbuild_rsandbuild_attn.Future ideas
--parallelto a big value while not unnecessarily limiting the context size of the clients of theserverif there aren't many. (related to Parallelization / Batching Explanation #4130 (reply in thread))Testing
convert-hf-to-gguf.py)main)serverwith backtrackingparallelExample output of
jamba-900M-v0.13-KIx2(click to expand)