Skip to content

Add support for Mamba-MoE #129

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.

Sign up for GitHub

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

Open
wants to merge 10 commits into
base: main
Choose a base branch
from
Open

Add support for Mamba-MoE #129

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
0a43edd
enable mamba moe
lchu6 Jan 23, 2025
1845392
add 30b mamba moe and 120b mamba moe
lchu6 Jan 23, 2025
c6cca86
lint
lchu6 Jan 23, 2025
7d26b45
remove distribution report
lchu6 Jan 24, 2025
0bbfabe
add both losses to report
lchu6 Jan 24, 2025
de631af
update 30b and 120b config for mamba moe
lchu6 Jan 24, 2025
c5271d2
add 236b config
lchu6 Jan 28, 2025
0ed114e
lint
lchu6 Jan 28, 2025
7a29b60
make meta device init
lchu6 Jan 31, 2025
eae33c3
make meta device init
lchu6 Jan 31, 2025
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
75 changes: 75 additions & 0 deletions fms_fsdp/utils/config_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,81 @@ def get_model_config(model_variant):
"pad_vocab_size_multiple": 16,
"tie_embeddings": False,
}
elif model_variant == "mamba_9.8b_moe":
model_config = {
"d_model": 4096,
"d_intermediate": 448,
"n_layer": 32,
"vocab_size": 128256,
"ssm_cfg": {"layer": "Mamba2"},
"attn_layer_idx": [9, 18, 27],
"attn_cfg": {
"causal": True,
"d_conv": 0,
"head_dim": 128,
"num_heads": 32,
"num_heads_kv": 8,
"out_proj_bias": False,
"qkv_proj_bias": False,
"rotary_emb_dim": 64,
},
"mlp_cfg": {"n_expert": 32, "load_balancing_loss": True, "top_k": 8},
"rms_norm": True,
"residual_in_fp32": True,
"fused_add_norm": True,
"pad_vocab_size_multiple": 16,
"tie_embeddings": False,
}
elif model_variant == "mamba_30b_moe":
model_config = {
"d_model": 6144,
"d_intermediate": 336,
"n_layer": 48,
"vocab_size": 128256,
"ssm_cfg": {"layer": "Mamba2"},
"attn_layer_idx": [9, 18, 27, 36, 45],
"attn_cfg": {
"causal": True,
"d_conv": 0,
"head_dim": 128,
"num_heads": 48,
"num_heads_kv": 8,
"out_proj_bias": False,
"qkv_proj_bias": False,
"rotary_emb_dim": 64,
},
"mlp_cfg": {"n_expert": 64, "load_balancing_loss": True, "top_k": 8},
"rms_norm": True,
"residual_in_fp32": True,
"fused_add_norm": True,
"pad_vocab_size_multiple": 16,
"tie_embeddings": False,
}
elif model_variant == "mamba_120b_moe":
model_config = {
"d_model": 8192,
"d_intermediate": 112,
"n_layer": 108,
"vocab_size": 128256,
"ssm_cfg": {"layer": "Mamba2"},
"attn_layer_idx": [9, 18, 27, 36, 45, 54, 63, 72, 81, 90, 99],
"attn_cfg": {
"causal": True,
"d_conv": 0,
"head_dim": 128,
"num_heads": 64,
"num_heads_kv": 8,
"out_proj_bias": False,
"qkv_proj_bias": False,
"rotary_emb_dim": 64,
},
"mlp_cfg": {"n_expert": 256, "load_balancing_loss": True, "top_k": 16},
"rms_norm": True,
"residual_in_fp32": True,
"fused_add_norm": True,
"pad_vocab_size_multiple": 16,
"tie_embeddings": False,
}
else:
raise ValueError(f"model variant {model_variant} not supported.")

Expand Down
16 changes: 14 additions & 2 deletions fms_fsdp/utils/train_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
from dataclasses import asdict
from functools import partial

from transformers.models.granitemoe.modeling_granitemoe import load_balancing_loss_func
Copy link

@fabianlim fabianlim Jan 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we should copy in this code later on



try:
import packaging.version
Expand Down Expand Up @@ -86,9 +88,19 @@ def train(

optimizer.zero_grad()
output = model(input)
output = output.logits if hasattr(output, "logits") else output
logits = output.logits if hasattr(output, "logits") else output
ce_loss = torch.nn.CrossEntropyLoss()
loss = ce_loss(output.view(-1, output.size(-1)), label.view(-1).long())
loss = ce_loss(logits.view(-1, logits.size(-1)), label.view(-1).long())
if "moe" in cfg.model_variant:
aux_outputs = output.aux_outputs
if aux_outputs is not None:
top_k = model.config.mlp_cfg.get("top_k", 2)
aux_loss = load_balancing_loss_func(
aux_outputs,
num_experts=model.config.mlp_cfg["n_expert"],
top_k=top_k,
)
loss += 0.2 * aux_loss

loss.backward()
ddp_stats[1] += model.clip_grad_norm_(cfg.grad_clip_thresh).item()
Expand Down
Loading