Skip to content

Qwen3 MoE Preliminary: add intermediate_size argument to MLP modules #2046

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 4 commits into
base: main
Choose a base branch
from
Open

Qwen3 MoE Preliminary: add intermediate_size argument to MLP modules #2046

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a3f2ca0
initial
May 17, 2025
c2b96be
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 17, 2025
73d247e
Merge branch 'main' into preliminary-qwen3_moe
Borda May 22, 2025
2c888ed
Merge branch 'main' into preliminary-qwen3_moe
Borda May 23, 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
1 change: 1 addition & 0 deletions litgpt/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@ class Config:
rope_adjustments: Optional[dict] = None
# Transformer block (MLP)
intermediate_size: Optional[int] = None
moe_intermediate_size: Optional[int] = None
bias: bool = True
mlp_class_name: Literal["GptNeoxMLP", "LLaMAMLP", "GemmaMLP", "LLaMAMoE"] = "GptNeoxMLP"
gelu_approximate: str = "none"
Expand Down
20 changes: 12 additions & 8 deletions litgpt/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -516,10 +516,11 @@ def _load_from_state_dict(self, state_dict: dict, prefix: str, *args: Any, **kwa


class GptNeoxMLP(nn.Module):
def __init__(self, config: Config) -> None:
def __init__(self, config: Config, intermediate_size: Optional[int] = None) -> None:
super().__init__()
self.fc = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias)
self.proj = nn.Linear(config.intermediate_size, config.n_embd, bias=config.bias)
self.intermediate_size = intermediate_size or config.intermediate_size
self.fc = nn.Linear(config.n_embd, self.intermediate_size, bias=config.bias)
self.proj = nn.Linear(self.intermediate_size, config.n_embd, bias=config.bias)
self.config = config

def forward(self, x: torch.Tensor) -> torch.Tensor:
Expand All @@ -529,11 +530,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:


class LLaMAMLP(nn.Module):
def __init__(self, config: Config) -> None:
def __init__(self, config: Config, intermediate_size: Optional[int] = None) -> None:
super().__init__()
self.fc_1 = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias)
self.fc_2 = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias)
self.proj = nn.Linear(config.intermediate_size, config.n_embd, bias=config.bias)
self.intermediate_size = intermediate_size or config.intermediate_size
self.fc_1 = nn.Linear(config.n_embd, self.intermediate_size, bias=config.bias)
self.fc_2 = nn.Linear(config.n_embd, self.intermediate_size, bias=config.bias)
self.proj = nn.Linear(self.intermediate_size, config.n_embd, bias=config.bias)
self.config = config

def forward(self, x: torch.Tensor) -> torch.Tensor:
Expand All @@ -555,7 +557,9 @@ class LLaMAMoE(nn.Module):
def __init__(self, config: Config) -> None:
super().__init__()
self.gate = nn.Linear(config.n_embd, config.n_expert, bias=False)
self.experts = nn.ModuleList(LLaMAMLP(config) for _ in range(config.n_expert))
self.experts = nn.ModuleList(
LLaMAMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.n_expert)
)
self.config = config

def forward(self, x: torch.Tensor) -> torch.Tensor:
Expand Down
Loading