Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

llama : support Mamba Selective State Space Models #5328

Merged
merged 43 commits into from
Mar 8, 2024
Merged
Changes from 1 commit
Commits
Show all changes
43 commits
Select commit Hold shift + click to select a range
8cd0a28
mamba : begin working on support for Mamba SSM
compilade Jan 26, 2024
5a69a26
mamba : begin figuring out how to (ab)use the kv cache for Mamba
compilade Jan 27, 2024
f680364
mamba : recurrent inference almost works, but incoherent
compilade Jan 28, 2024
54d3e48
mamba : recurrent inference WORKS!!!
compilade Jan 28, 2024
74eea85
convert : optionally use d_conv and d_state from config.json for Mamba
compilade Jan 29, 2024
9e77061
mamba : refactor recurrent conv, resulting in 20% perf increase
compilade Jan 29, 2024
3f7233b
ggml : parallelize ggml_exp
compilade Jan 29, 2024
e9cc45e
mamba : simplify the conv step with a self-overlapping view
compilade Jan 31, 2024
81b57bb
mamba : fix self-overlapping view depth stride
compilade Jan 31, 2024
ffc116f
mamba : handle batches of more than 1 token
compilade Feb 1, 2024
78a853b
ggml : in ggml_ssm_scan, merge multiple rows in the same vec operation
compilade Feb 2, 2024
5816ae6
mamba : very basic quantization support
compilade Feb 2, 2024
a3f4a1c
mamba : fuse more steps of the SSM scan in the ggml_ssm_scan operator
compilade Feb 3, 2024
9f55809
convert : for Mamba, also consider the "MambaLMHeadModel" arch name
compilade Feb 4, 2024
cd0f33f
mamba : fix vocab size problems with official models
compilade Feb 4, 2024
de92f15
ggml : remove ggml_exp and ggml_soft_plus
compilade Feb 4, 2024
766db75
mamba : remove some useless comments
compilade Feb 4, 2024
c52fb3c
convert : fix flake8 linter errors
compilade Feb 5, 2024
6ff34da
mamba : apply suggestions from code review
compilade Feb 5, 2024
8a43ffc
mamba : multiple sequences, but one at a time
compilade Feb 14, 2024
e73eaa7
mamba : in comments, properly refer to KV cells instead of slots
compilade Feb 14, 2024
de50c54
mamba : reduce memory usage of ggml_ssm_scan
compilade Feb 18, 2024
9473ec2
mamba : simultaneous sequence processing
compilade Feb 19, 2024
3dcf798
mamba : support llama_kv_cache_seq_cp copy chains
compilade Feb 25, 2024
34e2fca
mamba : make the server and parallel examples work with whole sequences
compilade Feb 25, 2024
79d636c
mamba : dedicate an input tensor for state copy indices
compilade Feb 25, 2024
8f605cf
mamba : adapt perplexity, batched, and batched-bench examples
compilade Feb 27, 2024
206e8ee
mamba : stop abusing attention metadata
compilade Feb 28, 2024
1af1000
mamba : more correctly update the "used" field of the KV cache
compilade Mar 2, 2024
d52dd50
ggml : in ggml_ssm_scan, use a threshold for soft_plus
compilade Mar 3, 2024
b83fbc9
convert : for Mamba, fallback to internal NeoX tokenizer
compilade Mar 3, 2024
eefb794
mamba : support state saving and restoring
compilade Mar 3, 2024
2a99d1b
ggml : implicitly pass src tensors through dst for Mamba-related ops
compilade Mar 4, 2024
93fd4b8
mamba : clarify some comments
compilade Mar 4, 2024
5544f52
Merge branch 'master' into support-mamba-ssm
compilade Mar 5, 2024
916b586
Merge branch 'master' into support-mamba-ssm
compilade Mar 7, 2024
7cd5a1f
server : fix cache_tokens not getting correctly resized
compilade Mar 7, 2024
d8024a4
convert-hf : support new metadata keys for Mamba
compilade Mar 8, 2024
17e4d6c
mamba : rename metadata to be more similar to transformers library
compilade Mar 8, 2024
1c8ea55
mamba : add missing spaces
compilade Mar 8, 2024
d0d32dc
convert-hf : omit output.weight when identical with token_embd.weight
compilade Mar 8, 2024
3e5685f
readme : add Mamba to supported models, and add recent API changes
compilade Mar 8, 2024
39579d3
mamba : move state_seq and state_mask views outside layer loop
compilade Mar 8, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
mamba : move state_seq and state_mask views outside layer loop
A few tensors were also missing `struct` in front of `ggml_tensor`.
  • Loading branch information
compilade committed Mar 8, 2024
commit 39579d3ceb3e5a47b18e82988cbbc9ea3c348f5b
16 changes: 9 additions & 7 deletions llama.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -5540,9 +5540,11 @@ struct llm_build_context {
struct ggml_cgraph * build_s_copy() {
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, LLAMA_MAX_NODES, false);

GGML_ASSERT(kv_self.recurrent);

for (int il = 0; il < n_layer; ++il) {
ggml_tensor * conv_states = ggml_reshape_2d(ctx0, kv_self.k_l[il], hparams.n_embd_k_s(), kv_self.size);
ggml_tensor * ssm_states = ggml_reshape_2d(ctx0, kv_self.v_l[il], hparams.n_embd_v_s(), kv_self.size);
struct ggml_tensor * conv_states = ggml_reshape_2d(ctx0, kv_self.k_l[il], hparams.n_embd_k_s(), kv_self.size);
struct ggml_tensor * ssm_states = ggml_reshape_2d(ctx0, kv_self.v_l[il], hparams.n_embd_v_s(), kv_self.size);

conv_states = ggml_get_rows(ctx0, conv_states, lctx.inp_s_copy);
ssm_states = ggml_get_rows(ctx0, ssm_states, lctx.inp_s_copy);
Expand Down Expand Up @@ -8171,14 +8173,16 @@ struct llm_build_context {
inpL = llm_build_inp_embd(ctx0, hparams, batch, model.tok_embd, lctx.inp_tokens, lctx.inp_embd, cb);
cb(inpL, "inp_embd", -1);

struct ggml_tensor * state_mask = ggml_view_2d(ctx0, lctx.inp_s_mask, 1, n_kv, lctx.inp_s_mask->nb[0], 0);
struct ggml_tensor * state_seq = ggml_view_2d(ctx0, lctx.inp_s_seq, n_kv, n_tokens, n_kv*ggml_element_size(lctx.inp_s_seq), 0);

for (int il = 0; il < n_layer; ++il) {
// (ab)using the KV cache to store the states
ggml_tensor * conv_states = ggml_reshape_2d(ctx0, kv_self.k_l[il], hparams.n_embd_k_s(), kv_self.size);
ggml_tensor * ssm_states = ggml_reshape_2d(ctx0, kv_self.v_l[il], hparams.n_embd_v_s(), kv_self.size);
struct ggml_tensor * conv_states = ggml_reshape_2d(ctx0, kv_self.k_l[il], hparams.n_embd_k_s(), kv_self.size);
struct ggml_tensor * ssm_states = ggml_reshape_2d(ctx0, kv_self.v_l[il], hparams.n_embd_v_s(), kv_self.size);

// clear states of sequences which are starting at the beginning of this batch
{
ggml_tensor * state_mask = ggml_view_2d(ctx0, lctx.inp_s_mask, 1, n_kv, lctx.inp_s_mask->nb[0], 0);
conv_states = ggml_mul(ctx0,
ggml_view_2d(ctx0, conv_states, conv_states->ne[0], n_kv, conv_states->nb[1], kv_head*conv_states->nb[1]),
state_mask);
Expand All @@ -8203,8 +8207,6 @@ struct llm_build_context {
struct ggml_tensor * x = ggml_view_2d(ctx0, xz, d_inner, xz->ne[1], xz->nb[1], 0);
struct ggml_tensor * z = ggml_view_2d(ctx0, xz, d_inner, xz->ne[1], xz->nb[1], ggml_element_size(xz)*d_inner);

struct ggml_tensor * state_seq = ggml_view_2d(ctx0, lctx.inp_s_seq, n_kv, n_tokens, n_kv*ggml_element_size(lctx.inp_s_seq), 0);

// conv
{
// Custom operator which is needed only to ease simultaneous sequence processing.
Expand Down
Loading