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

Bamba architecture #10810

Draft
wants to merge 43 commits into
base: master
Choose a base branch
from
Draft
Changes from 1 commit
Commits
Show all changes
43 commits
Select commit Hold shift + click to select a range
1f0fea7
llama : initial Mamba-2 support
compilade Aug 1, 2024
dceff23
ggml : SIMD ggml_ssm_scan for Mamba-2
compilade Aug 19, 2024
2bfe9de
llama : support running Mamba-Codestral-7B-v0.1
compilade Aug 19, 2024
aff9692
llama : fix Mamba-2 conv state saving
compilade Aug 21, 2024
e04910d
llama : remove unused variable
compilade Aug 22, 2024
fa358e7
llama : add missing break
compilade Aug 22, 2024
38913dc
convert_hf : prefer SentencePiece tokenizer for Mamba-2 when present
compilade Aug 22, 2024
0e601ca
Merge branch 'master' into compilade/mamba2
compilade Sep 18, 2024
273e7a4
llama : avoid redundant state copy for Mamba 1 and 2
compilade Sep 30, 2024
7d6cb36
Merge branch 'master' into compilade/mamba2
compilade Oct 1, 2024
2c77d79
metal : attempt to adapt SSM_SCAN for Mamba-2
compilade Oct 2, 2024
87b97d0
metal : fix SSM_SCAN pipeline scope
compilade Oct 2, 2024
03d0e6e
metal : use log and exp instead of log1pf and expf in SSM_SCAN
compilade Oct 2, 2024
7a351ab
metal : remove unused arguments for SSM_SCAN
compilade Oct 2, 2024
8b15bc6
metal : add back n_seqs to SSM_SCAN args
compilade Oct 2, 2024
5b8ec2b
metal : fix SSM_SCAN state head offset
compilade Oct 2, 2024
62b09b3
metal : fix wrong number of tokens per sequence in SSM_SCAN
compilade Oct 3, 2024
038d958
Merge branch 'master' into compilade/mamba2
compilade Oct 12, 2024
805512a
ggml : remove unused fast broadcast path in GGML_MUL
compilade Oct 12, 2024
7d16e1b
Merge branch 'master' into compilade/mamba2
compilade Nov 1, 2024
3bc7103
ggml : avoid multiply by D in GGML_OP_SSM_SCAN
compilade Nov 4, 2024
8d8f065
Merge branch 'master' into compilade/mamba2
compilade Nov 4, 2024
b4e9c59
convert : fix flake8 lint
compilade Nov 4, 2024
1ee6c48
Merge branch 'master' into compilade/mamba2
compilade Nov 25, 2024
9a68f75
feat(jamba): First pass at GGUF conversion for Jamba models
gabe-l-hart Nov 26, 2024
246dfdb
feat(jamba): Add jamba architecture to llama.cpp enums
gabe-l-hart Nov 26, 2024
e3525e9
feat(convert): Full pass at hparam conversion
gabe-l-hart Dec 2, 2024
fd98682
fix(bamba conv): Jamba -> Bamba
gabe-l-hart Dec 3, 2024
1c1e008
fix(bamba): Jamba->Bamba in llama.cpp
gabe-l-hart Dec 3, 2024
e0af809
feat(bamba): hparam parsing in llama.cpp
gabe-l-hart Dec 3, 2024
fd3bb30
fix(bamba conv): Fizes in tensor name and hparam conversion for llama…
gabe-l-hart Dec 4, 2024
3ee0ae3
feat(bamba): Full tensor parsing for bamba
gabe-l-hart Dec 4, 2024
dfe8d3d
fix(bamba conv): Remove chunk size and consolidate head count w/ time…
gabe-l-hart Dec 5, 2024
41fc019
fix(bamba): Remove ssm_head_count and ssm_chunk_size in llama.cpp
gabe-l-hart Dec 5, 2024
e7b1abb
feat(bamba): Partially complete work on constructing the forward graph
gabe-l-hart Dec 5, 2024
f2478bc
fix: Get n_head_kv per-layer in build_bamba
gabe-l-hart Dec 9, 2024
d3a34e0
fix: per-layer recurrent embd_[kv]_s
gabe-l-hart Dec 9, 2024
92653d0
WIP: Partial work towards separate hybrid cache
gabe-l-hart Dec 9, 2024
44bf431
fix: Only allocate kv cache tensors for the appropriate layers in hyb…
gabe-l-hart Dec 10, 2024
4543ed5
feat: Update the logic in llama_decode_internal for kv_hybrid cache
gabe-l-hart Dec 10, 2024
204e78f
fix: A number of places where hybrid needs to be handled
gabe-l-hart Dec 10, 2024
97e6ba8
fix: Remove outdated TODO in convrsion script
gabe-l-hart Dec 12, 2024
b83e9a6
fix: Remove unused LLM_KV_ATTENTION_LAYER_COUNT
gabe-l-hart Dec 12, 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
Next Next commit
fix: per-layer recurrent embd_[kv]_s
For hybrid models, this value should be 0 for the non-recurrent layers

Branch: BambaArchitecture

Signed-off-by: Gabe Goodhart <ghart@us.ibm.com>
gabe-l-hart committed Dec 12, 2024
commit d3a34e0282579ba08d773bb7760f4a6a2060cb8a
49 changes: 28 additions & 21 deletions src/llama.cpp
Original file line number Diff line number Diff line change
@@ -2485,7 +2485,7 @@ struct llama_hparams {
uint32_t ssm_head_dim = 0;

// for hybrid state space models
std::array<bool, LLAMA_MAX_LAYERS> ssm_layer_arr;
std::array<bool, LLAMA_MAX_LAYERS> recurrent_layer_arr;

float f_clamp_kqv = 0.0f;
float f_max_alibi_bias = 0.0f;
@@ -2544,7 +2544,7 @@ struct llama_hparams {
if (this->ssm_dt_b_c_rms != other.ssm_dt_b_c_rms) return true;
if (this->ssm_head_dim != other.ssm_head_dim) return true;

if (this->ssm_layer_arr != other.ssm_layer_arr) return true;
if (this->recurrent_layer_arr != other.recurrent_layer_arr) return true;

if (this->rescale_every_n_layers != other.rescale_every_n_layers) return true;
if (this->time_mix_extra_dim != other.time_mix_extra_dim) return true;
@@ -2616,30 +2616,34 @@ struct llama_hparams {
return n_embd_head_v * n_head_kv;
}

uint32_t n_embd_k_s() const { // dimension of the rolling state embeddings
uint32_t n_embd_k_s(uint32_t il = 0) const { // dimension of the rolling state embeddings
// corresponds to Mamba's conv_states size or RWKV's token_shift states size
if (!recurrent_layer(il)) {
return 0;
}
if (wkv_head_size != 0) {
// for RWKV models
return 2 * n_embd;
} else {
// TODO: maybe support other convolution strides than 1
// NOTE: since the first column of the conv_state is shifted out each time, it's not actually needed
return (ssm_d_conv > 0 ? ssm_d_conv - 1 : 0) * (ssm_d_inner + 2*ssm_n_group*ssm_d_state);
}
// TODO: maybe support other convolution strides than 1
// NOTE: since the first column of the conv_state is shifted out each time, it's not actually needed
return (ssm_d_conv > 0 ? ssm_d_conv - 1 : 0) * (ssm_d_inner + 2*ssm_n_group*ssm_d_state);
}

uint32_t n_embd_v_s() const { // dimension of the recurrent state embeddings
uint32_t n_embd_v_s(uint32_t il = 0) const { // dimension of the recurrent state embeddings
if (!recurrent_layer(il)) {
return 0;
}
if (wkv_head_size != 0) {
// corresponds to RWKV's wkv_states size
return n_embd * wkv_head_size;
} else {
// corresponds to Mamba's ssm_states size
return ssm_d_state * ssm_d_inner;
}
// corresponds to Mamba's ssm_states size
return ssm_d_state * ssm_d_inner;
}

bool ssm_layer(uint32_t il) const {
return ssm_layer_arr[il];
bool recurrent_layer(uint32_t il) const {
return recurrent_layer_arr[il];
}
};

@@ -3555,8 +3559,8 @@ static bool llama_kv_cache_init(
cache.v_l.reserve(n_layer);

for (int i = 0; i < (int) n_layer; i++) {
const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(i) + hparams.n_embd_k_s();
const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(i) + hparams.n_embd_v_s();
const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(i) + hparams.n_embd_k_s(i);
const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(i) + hparams.n_embd_v_s(i);

ggml_backend_buffer_type_t buft;
if (offload) {
@@ -5509,7 +5513,10 @@ static void llm_load_hparams(
std::fill(hparams.n_head_arr.begin(), hparams.n_head_arr.end(), 0);
std::fill(hparams.n_head_kv_arr.begin(), hparams.n_head_kv_arr.end(), 0);
std::fill(hparams.n_ff_arr.begin(), hparams.n_ff_arr.end(), 0);
std::fill(hparams.ssm_layer_arr.begin(), hparams.ssm_layer_arr.end(), false);
std::fill(
hparams.recurrent_layer_arr.begin(),
hparams.recurrent_layer_arr.end(),
llama_model_is_recurrent(&model));

ml.get_key_or_arr(LLM_KV_FEED_FORWARD_LENGTH, hparams.n_ff_arr, hparams.n_layer);
ml.get_key_or_arr(LLM_KV_ATTENTION_HEAD_COUNT, hparams.n_head_arr, hparams.n_layer);
@@ -5999,12 +6006,12 @@ static void llm_load_hparams(
std::fill(hparams.n_head_kv_arr.begin(), hparams.n_head_kv_arr.end(), 0);

// Attention params
std::fill(hparams.ssm_layer_arr.begin(), hparams.ssm_layer_arr.end(), true);
std::fill(hparams.recurrent_layer_arr.begin(), hparams.recurrent_layer_arr.end(), true);
std::vector<uint32_t> attn_layer_indices;
ml.get_arr(LLM_KV_ATTENTION_LAYER_INDICES, attn_layer_indices);
for (const auto attn_idx : attn_layer_indices) {
GGML_ASSERT(attn_idx < hparams.n_layer);
hparams.ssm_layer_arr[attn_idx] = false;
hparams.recurrent_layer_arr[attn_idx] = false;
// Correctly set n_head and n_head_kv for attention layers
hparams.n_head_arr[attn_idx] = n_head_attn;
hparams.n_head_kv_arr[attn_idx] = n_head_kv_attn;
@@ -7162,7 +7169,7 @@ static void llm_load_print_meta(llama_model_loader & ml, llama_model & model) {
}

if (model.arch == LLM_ARCH_BAMBA) {
LLAMA_LOG_INFO("%s: ssm_layer_arr = %s\n", __func__, print_f([&](uint32_t il) { return uint32_t(hparams.ssm_layer(il)); }, hparams.n_layer).c_str());
LLAMA_LOG_INFO("%s: recurrent_layer_arr = %s\n", __func__, print_f([&](uint32_t il) { return uint32_t(hparams.recurrent_layer(il)); }, hparams.n_layer).c_str());
}
}

@@ -8769,7 +8776,7 @@ static bool llm_load_tensors(
// norm
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);

if (hparams.ssm_layer(i)) {
if (hparams.recurrent_layer(i)) {
// ssm layers
layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, d_in_proj}, 0);
layer.ssm_in_b = create_tensor(tn(LLM_TENSOR_SSM_IN, "bias", i), {n_embd, d_in_proj}, llama_model_loader::TENSOR_NOT_REQUIRED);
@@ -14677,7 +14684,7 @@ struct llm_build_context {
LLM_NORM_RMS, cb, il);
cb(cur, "attn_norm", il);

if (hparams.ssm_layer(il)) {
if (hparams.recurrent_layer(il)) {
// ssm layer
cur = llm_build_mamba2(ctx0, lctx, ubatch, gf, cur, state_copy,
rs_zero, kv_head, n_kv, cb, il);