diff --git a/README.md b/README.md index 1036a95894..f5c139e975 100644 --- a/README.md +++ b/README.md @@ -151,6 +151,7 @@ Every model is written from scratch to maximize performance and remove layers of | QwQ | 32B | Alibaba Group | [Qwen Team 2025](https://qwenlm.github.io/blog/qwq-32b/) | | QwQ-Preview | 32B | Alibaba Group | [Qwen Team 2024](https://qwenlm.github.io/blog/qwq-32b-preview/) | | Qwen3 | 0.6B, 1.7B, 4B, 8B, 14B, 32B | Alibaba Group | [Qwen Team 2025](https://arxiv.org/abs/2505.09388/) | +| Qwen3 MoE | 30B, 235B | Alibaba Group | [Qwen Team 2025](https://arxiv.org/abs/2505.09388/) | | R1 Distill Llama | 8B, 70B | DeepSeek AI | [DeepSeek AI 2025](https://github.com/deepseek-ai/DeepSeek-R1/blob/main/DeepSeek_R1.pdf) | | SmolLM2 | 135M, 360M, 1.7B | Hugging Face | [Hugging Face 2024](https://github.com/huggingface/smollm) | | Salamandra | 2B, 7B | Barcelona Supercomputing Centre | [BSC-LTC 2024](https://github.com/BSC-LTC/salamandra) | diff --git a/litgpt/adapter_v2.py b/litgpt/adapter_v2.py index 363d6eb2b1..fb4d12ae08 100644 --- a/litgpt/adapter_v2.py +++ b/litgpt/adapter_v2.py @@ -158,11 +158,12 @@ def _load_from_state_dict(self, state_dict: Dict, prefix: str, *args: Any, **kwa class LLaMAMLP(litgpt.model.LLaMAMLP): - def __init__(self, config: Config) -> None: + def __init__(self, config: Config, intermediate_size: Optional[int] = None) -> None: nn.Module.__init__(self) - self.fc_1 = AdapterV2Linear(config.n_embd, config.intermediate_size, bias=config.bias) - self.fc_2 = AdapterV2Linear(config.n_embd, config.intermediate_size, bias=config.bias) - self.proj = AdapterV2Linear(config.intermediate_size, config.n_embd, bias=config.bias) + self.intermediate_size = intermediate_size or config.intermediate_size + self.fc_1 = AdapterV2Linear(config.n_embd, self.intermediate_size, bias=config.bias) + self.fc_2 = AdapterV2Linear(config.n_embd, self.intermediate_size, bias=config.bias) + self.proj = AdapterV2Linear(self.intermediate_size, config.n_embd, bias=config.bias) self.config = config def _load_from_state_dict(self, state_dict: Dict, prefix: str, *args: Any, **kwargs: Any) -> None: @@ -191,7 +192,9 @@ class LLaMAMoE(litgpt.model.LLaMAMoE): def __init__(self, config: Config) -> None: nn.Module.__init__(self) self.gate = AdapterV2Linear(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 _load_from_state_dict(self, state_dict: Dict, prefix: str, *args: Any, **kwargs: Any) -> None: diff --git a/litgpt/config.py b/litgpt/config.py index 3dd5ba2e07..123101130a 100644 --- a/litgpt/config.py +++ b/litgpt/config.py @@ -2583,6 +2583,9 @@ def norm_class(self) -> Type: copy["hf_config"]["name"] = c["hf_config"]["name"].format(kind) configs.append(copy) +########## +# QwQ +########## qwq = [ # https://huggingface.co/Qwen/QwQ-32B/blob/main/config.json dict( @@ -2630,6 +2633,9 @@ def norm_class(self) -> Type: configs.extend(qwq) +########## +# Qwen3 +########## qwen_3 = [ # https://huggingface.co/Qwen/Qwen3-0.6B/blob/main/config.json dict( @@ -2771,6 +2777,85 @@ def norm_class(self) -> Type: ] configs.extend(qwen_3_32b) +qwen_3_moe = [ + # https://huggingface.co/Qwen/Qwen3-30B-A3B/blob/main/config.json + dict( + name="Qwen3-30B-A3B", + hf_config=dict(org="Qwen", name="Qwen3-30B-A3B"), + block_size=40960, + head_size=128, + vocab_size=151643, + padded_vocab_size=151936, + n_layer=48, + n_head=32, + n_embd=2048, + n_query_groups=4, + rotary_percentage=1.0, + parallel_residual=False, + bias=False, + norm_class_name="RMSNorm", + mlp_class_name="LLaMAMoE", + intermediate_size=6144, + moe_intermediate_size=768, + norm_eps=1e-6, + rope_base=1000000, + norm_qk=True, + n_expert=128, + n_expert_per_token=8, + ), + # https://huggingface.co/Qwen/Qwen3-30B-A3B-Base/blob/main/config.json + dict( + name="Qwen3-30B-A3B-Base", + hf_config=dict(org="Qwen", name="Qwen3-30B-A3B-Base"), + block_size=40960, + head_size=128, + vocab_size=151643, + padded_vocab_size=151936, + n_layer=48, + n_head=32, + n_embd=2048, + n_query_groups=4, + rotary_percentage=1.0, + parallel_residual=False, + bias=False, + norm_class_name="RMSNorm", + mlp_class_name="LLaMAMoE", + intermediate_size=6144, + moe_intermediate_size=768, + norm_eps=1e-6, + rope_base=1000000, + norm_qk=True, + n_expert=128, + n_expert_per_token=8, + ), + # https://huggingface.co/Qwen/Qwen3-235B-A22B/blob/main/config.json + dict( + name="Qwen3-235B-A22B", + hf_config=dict(org="Qwen", name="Qwen3-235B-A22B"), + block_size=40960, + head_size=128, + vocab_size=151643, + padded_vocab_size=151936, + n_layer=94, + n_head=64, + n_embd=4096, + n_query_groups=4, + rotary_percentage=1.0, + parallel_residual=False, + bias=False, + norm_class_name="RMSNorm", + mlp_class_name="LLaMAMoE", + intermediate_size=12288, + moe_intermediate_size=1536, + norm_eps=1e-6, + rope_base=1000000, + norm_qk=True, + n_expert=128, + n_expert_per_token=8, + ), +] +configs.extend(qwen_3_moe) + ############# # Salamandra diff --git a/litgpt/lora.py b/litgpt/lora.py index 0277412c67..6739b5b040 100644 --- a/litgpt/lora.py +++ b/litgpt/lora.py @@ -609,11 +609,12 @@ def _load_from_state_dict(self, state_dict: Dict, prefix: str, *args: Any, **kwa class LLaMAMLP(litgpt.model.LLaMAMLP): - def __init__(self, config: Config) -> None: + def __init__(self, config: Config, intermediate_size: Optional[int] = None) -> None: nn.Module.__init__(self) - self.fc_1 = create_lora_linear(config, config.n_embd, config.intermediate_size) - self.fc_2 = create_lora_linear(config, config.n_embd, config.intermediate_size) - self.proj = create_lora_linear(config, config.intermediate_size, config.n_embd) + self.intermediate_size = intermediate_size or config.intermediate_size + self.fc_1 = create_lora_linear(config, config.n_embd, self.intermediate_size) + self.fc_2 = create_lora_linear(config, config.n_embd, self.intermediate_size) + self.proj = create_lora_linear(config, self.intermediate_size, config.n_embd) self.config = config def _load_from_state_dict(self, state_dict: Dict, prefix: str, *args: Any, **kwargs: Any) -> None: @@ -642,7 +643,9 @@ class LLaMAMoE(litgpt.model.LLaMAMoE): def __init__(self, config: Config) -> None: nn.Module.__init__(self) self.gate = create_lora_linear(config, 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 _load_from_state_dict(self, state_dict: Dict, prefix: str, *args: Any, **kwargs: Any) -> None: diff --git a/litgpt/scripts/convert_hf_checkpoint.py b/litgpt/scripts/convert_hf_checkpoint.py index bdcc9b0588..bff61d290a 100644 --- a/litgpt/scripts/convert_hf_checkpoint.py +++ b/litgpt/scripts/convert_hf_checkpoint.py @@ -651,12 +651,28 @@ def copy_weights_qwen_3( "model.layers.{}.self_attn.k_norm.weight": "transformer.h.{}.attn.norm_k.weight", "model.layers.{}.self_attn.o_proj.weight": "transformer.h.{}.attn.proj.weight", "model.layers.{}.post_attention_layernorm.weight": "transformer.h.{}.norm_2.weight", - "model.layers.{}.mlp.gate_proj.weight": "transformer.h.{}.mlp.fc_1.weight", - "model.layers.{}.mlp.up_proj.weight": "transformer.h.{}.mlp.fc_2.weight", - "model.layers.{}.mlp.down_proj.weight": "transformer.h.{}.mlp.proj.weight", "model.norm.weight": "transformer.ln_f.weight", "lm_head.weight": "lm_head.weight", } + if config.mlp_class_name == "LLaMAMoE": + weight_map.update( + { + "model.layers.{}.mlp.experts.{}.gate_proj.weight": "transformer.h.{}.mlp.experts.{}.fc_1.weight", + "model.layers.{}.mlp.experts.{}.up_proj.weight": "transformer.h.{}.mlp.experts.{}.fc_2.weight", + "model.layers.{}.mlp.experts.{}.down_proj.weight": "transformer.h.{}.mlp.experts.{}.proj.weight", + "model.layers.{}.mlp.gate.weight": "transformer.h.{}.mlp.gate.weight", + } + ) + elif config.mlp_class_name == "LLaMAMLP": + weight_map.update( + { + "model.layers.{}.mlp.gate_proj.weight": "transformer.h.{}.mlp.fc_1.weight", + "model.layers.{}.mlp.up_proj.weight": "transformer.h.{}.mlp.fc_2.weight", + "model.layers.{}.mlp.down_proj.weight": "transformer.h.{}.mlp.proj.weight", + } + ) + else: + raise NotImplementedError if progress_per_file is not None: progress_per_file = progress_per_file / max(1, len(hf_weights) + len(qkv_weights)) diff --git a/litgpt/scripts/convert_lit_checkpoint.py b/litgpt/scripts/convert_lit_checkpoint.py index 3718e6e671..d7ce885c55 100644 --- a/litgpt/scripts/convert_lit_checkpoint.py +++ b/litgpt/scripts/convert_lit_checkpoint.py @@ -465,12 +465,28 @@ def copy_weights_qwen_3( "transformer.h.{}.attn.proj.weight": "model.layers.{}.self_attn.o_proj.weight", "transformer.h.{}.attn.norm_q.weight": "model.layers.{}.self_attn.q_norm.weight", "transformer.h.{}.attn.norm_k.weight": "model.layers.{}.self_attn.k_norm.weight", - "transformer.h.{}.mlp.fc_1.weight": "model.layers.{}.mlp.gate_proj.weight", - "transformer.h.{}.mlp.fc_2.weight": "model.layers.{}.mlp.up_proj.weight", - "transformer.h.{}.mlp.proj.weight": "model.layers.{}.mlp.down_proj.weight", "transformer.ln_f.weight": "model.norm.weight", "lm_head.weight": "lm_head.weight", } + if config.mlp_class_name == "LLaMAMoE": + weight_map.update( + { + "transformer.h.{}.mlp.gate.weight": "model.layers.{}.mlp.gate.weight", + "transformer.h.{}.mlp.experts.{}.fc_1.weight": "model.layers.{}.mlp.experts.{}.gate_proj.weight", + "transformer.h.{}.mlp.experts.{}.fc_2.weight": "model.layers.{}.mlp.experts.{}.up_proj.weight", + "transformer.h.{}.mlp.experts.{}.proj.weight": "model.layers.{}.mlp.experts.{}.down_proj.weight", + } + ) + elif config.mlp_class_name == "LLaMAMLP": + weight_map.update( + { + "transformer.h.{}.mlp.fc_1.weight": "model.layers.{}.mlp.gate_proj.weight", + "transformer.h.{}.mlp.fc_2.weight": "model.layers.{}.mlp.up_proj.weight", + "transformer.h.{}.mlp.proj.weight": "model.layers.{}.mlp.down_proj.weight", + } + ) + else: + raise NotImplementedError for from_name, param in lit_weights.items(): if from_name == "lm_head.weight" and untie_weights: diff --git a/tests/test_model.py b/tests/test_model.py index ab5dec50b2..da931c7297 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -32,6 +32,7 @@ from transformers.models.olmo2 import Olmo2Config, Olmo2ForCausalLM from transformers.models.qwen2 import Qwen2Config, Qwen2ForCausalLM from transformers.models.qwen3 import Qwen3Config, Qwen3ForCausalLM +from transformers.models.qwen3_moe import Qwen3MoeConfig, Qwen3MoeForCausalLM import litgpt.config as config_module from litgpt import GPT, Config @@ -1139,6 +1140,72 @@ def test_against_original_qwen_3(model_name, device, dtype): torch.testing.assert_close(ours_y, theirs_y) +@torch.inference_mode() +@pytest.mark.parametrize("model_name", ["Qwen3-30B-A3B", "Qwen3-235B-A22B"]) +@pytest.mark.parametrize( + ("device", "dtype"), + [ + (torch.device("cpu"), torch.float32), + pytest.param( + torch.device("cuda"), + torch.float16, + marks=[ + # the reference does softmax upscaled to fp32 during attention. additionally, the final layernorm input + # is slightly different + pytest.mark.xfail(raises=AssertionError, strict=False), + _RunIf(min_cuda_gpus=1), + ], + ), + ], +) +def test_against_original_qwen_3_moe(model_name, device, dtype): + torch.set_default_dtype(dtype) + + T = 20 + ours_config = Config.from_name( + model_name, + block_size=T, + n_layer=2, + n_head=16, + n_embd=32, + intermediate_size=86, + moe_intermediate_size=20, + n_expert=4, + n_expert_per_token=2, + ) + theirs_config = Qwen3MoeConfig( + vocab_size=ours_config.padded_vocab_size, + hidden_size=ours_config.n_embd, + head_dim=ours_config.head_size, + num_attention_heads=ours_config.n_head, + num_hidden_layers=ours_config.n_layer, + intermediate_size=ours_config.intermediate_size, + moe_intermediate_size=ours_config.moe_intermediate_size, + max_position_embeddings=ours_config.block_size, + rms_norm_eps=ours_config.norm_eps, + num_key_value_heads=ours_config.n_query_groups, + rope_theta=ours_config.rope_base, + tie_word_embeddings=False, + num_experts=ours_config.n_expert, + num_experts_per_tok=ours_config.n_expert_per_token, + norm_topk_prob=True, + ) + + theirs_model = Qwen3MoeForCausalLM(theirs_config).to(device) + theirs_state_dict = theirs_model.state_dict() + state_dict = {} + copy_weights_qwen_3(ours_config, {}, state_dict, theirs_state_dict) + ours_model = GPT(ours_config).to(device) + ours_model.load_state_dict(state_dict) + + # test end to end + x = torch.randint(low=0, high=ours_config.padded_vocab_size, size=(T,), device=device).unsqueeze(0) + assert x.size(1) == T + ours_y = ours_model(x) + theirs_y = theirs_model(x)["logits"].to(dtype) # HF converts logits to float + torch.testing.assert_close(ours_y, theirs_y) + + @torch.inference_mode() @pytest.mark.parametrize("model_name", ("salamandra-2b", "salamandra-7b")) @pytest.mark.parametrize( diff --git a/tutorials/download_model_weights.md b/tutorials/download_model_weights.md index 244a137948..080d4f8e06 100644 --- a/tutorials/download_model_weights.md +++ b/tutorials/download_model_weights.md @@ -49,6 +49,7 @@ LitGPT supports a variety of LLM architectures with publicly available weights. | QwQ | 32B | Alibaba Group | [Qwen Team 2025](https://qwenlm.github.io/blog/qwq-32b/) | | QwQ-Preview | 32B | Alibaba Group | [Qwen Team 2024](https://qwenlm.github.io/blog/qwq-32b-preview/) | | Qwen3 | 0.6B, 1.7B, 4B, 8B, 14B, 32B | Alibaba Group | [Qwen Team 2025](https://arxiv.org/abs/2505.09388/) | +| Qwen3 MoE | 30B, 235B | Alibaba Group | [Qwen Team 2025](https://arxiv.org/abs/2505.09388/) | | R1 Distll Llama | 8B, 70B | DeepSeek AI | [DeepSeek AI 2025](https://github.com/deepseek-ai/DeepSeek-R1/blob/main/DeepSeek_R1.pdf) | | RedPajama-INCITE | 3B, 7B | Together | [Together 2023](https://together.ai/blog/redpajama-models-v1) | | SmolLM2 | 135M, 360M, 1.7B | Hugging Face | [Hugging Face 2024](https://github.com/huggingface/smollm) |