diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/baichuan.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/baichuan.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8447afbe02eff845453026206143d7073e573191 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/baichuan.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/chatglm.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/chatglm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1b94ac198e7b22d6bde6e7d9f32077f6046aa7ee Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/chatglm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/commandr.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/commandr.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a577b09b12210dddb22d1ff2c945cd739cfc8438 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/commandr.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/dbrx.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/dbrx.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f88b393f69da96a4c679b9e4adcdcc185a97d915 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/dbrx.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/deepseek.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/deepseek.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c8c0194881a0cd37fb017ed80c0a1ba9ed007418 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/deepseek.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/deepseek_janus_pro.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/deepseek_janus_pro.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2a21ed5300ecc94514821c0fa59cdf29b14ac568 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/deepseek_janus_pro.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/deepseek_nextn.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/deepseek_nextn.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ed94c0ece8cf2467d24167a98ce0b422fceb8633 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/deepseek_nextn.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/deepseek_vl2.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/deepseek_vl2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..940abffac1a49b5a3c86196966f5e60cb59d9c12 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/deepseek_vl2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/exaone.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/exaone.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8d2f8dc78d2ef559990f25ca00e48341a002c120 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/exaone.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gemma.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gemma.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..37d19702282b5dcac5c0e851ccf1f1c9008e7f4f Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gemma.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gemma2.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gemma2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f61b586f3568ccf648b593cb71fca5eb9725c31b Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gemma2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gemma2_reward.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gemma2_reward.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..073622e9d088a98cfcaecc7a65af95c0135fc5ac Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gemma2_reward.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gemma3_causal.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gemma3_causal.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6450a17249cc9f4910a8abd9121ae705f6e72dc7 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gemma3_causal.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gemma3_mm.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gemma3_mm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb1e89e5a35dddd12aaaac0f0fabc5c372176721 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gemma3_mm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gpt2.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gpt2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fd3e73018b41c6dc2345f6dd72175190b31dad2e Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gpt2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gpt_bigcode.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gpt_bigcode.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cca3388dba9d97d2e9375053d555ef834201817d Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/gpt_bigcode.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/granite.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/granite.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..59e5308c7580aaee78de7a9a2a907c21779a0c6a Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/granite.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/grok.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/grok.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0e988b5164abf5df4f06ab84f07f0e9dbc509e45 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/grok.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/internlm2.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/internlm2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..68a6c3d9335b8a1a08a4ce2fdf3b50a05c8f332d Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/internlm2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/internlm2_reward.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/internlm2_reward.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..363d5fe695ae39326af792f2da42dc389caabe0b Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/internlm2_reward.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f3d7255376177b3653bcc96f6b8e7b6b27ee8d11 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama4.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama4.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6b3f3e3e601813f8b814905d39831c4e9d429cef Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama4.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama_classification.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..162f8ab42c217b43dde1a4b503ab46310e5a68ce Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama_eagle.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama_eagle.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2a12564c34b350292417d21d102a00398c326c89 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama_eagle.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama_eagle3.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama_eagle3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ed2bb00649afca010f2babae37db63f4c8e609a4 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama_eagle3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama_reward.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama_reward.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ae34639bc4e7106e5c21b53ed0ed4278085787f1 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llama_reward.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llava.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llava.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd094b7bd90789e9a449162610995672954f7aa2 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llava.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llavavid.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llavavid.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..42513d5ced9ce87a66305a2a3141e2d1e85695ef Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/llavavid.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/minicpm.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/minicpm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..028e811b1f7c45dc040da0997a47fd346530fb43 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/minicpm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/minicpmv.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/minicpmv.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d0ae68b9f97081f048fddcdc38f86a6d0aa128e4 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/minicpmv.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/mistral.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/mistral.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f34397ddf05a115f04abf1f8d1734deef8ffd3ca Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/mistral.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/mixtral.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/mixtral.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..472e7a5fd4e023faec1203b8b13d8b9a2ef9f4cd Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/mixtral.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/mixtral_quant.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/mixtral_quant.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c5b18b7984a309e11facffcfc05cc17cd7e9abbb Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/mixtral_quant.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/mllama4.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/mllama4.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a0032dda794bab6879a3f45ef6679e16622041c1 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/mllama4.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/olmo.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/olmo.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a04e34520294aeb9b2f72afcc2bff109d7c7c93f Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/olmo.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/olmo2.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/olmo2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3a1cc8266270bc25db21076c52b3c6eb461b6ad6 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/olmo2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/olmoe.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/olmoe.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6bdb39dc9bfaa15d5bee04d8aaf99465baab869e Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/olmoe.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/phi3_small.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/phi3_small.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8b6c83e234fb0168cd8c564bba17bcbce9c92bfd Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/phi3_small.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..33783d1f32f9f3762abb2f08ebf27bad94a22806 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2_classification.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2_classification.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ba4d1c71fd2e58989e41ac987f8188a7d8b8067d Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2_classification.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2_eagle.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2_eagle.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57ae4ed9d0561cf2fd0c0ce15cf588ba6dc6ee21 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2_eagle.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2_moe.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2_moe.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..70ce35a4d02469184fdd3f9e530a439c0c9e1c8f Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2_moe.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2_rm.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2_rm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e31bd0316224a06cb95683af8b2bbe345ce1b8b Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2_rm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2_vl.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2_vl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f55f99a93d9db4776f8d61d7588bcea5c8f4c8b Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen2_vl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen3.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4a6150cf05a298352451d49bb015b5e9fd1bf686 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen3_moe.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen3_moe.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0c28a2e1f2c3c32cce15cff8f3be1c7cd82f1aa2 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/qwen3_moe.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/registry.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/registry.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f2ca706822a1c107fc92e22083e26bd6fd1f5991 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/registry.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/roberta.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/roberta.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a32aefbe010cb128deeb0ce60686075d23e018b4 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/roberta.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/stablelm.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/stablelm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ddadb953f36cedd0f604521e150a937b99ba6ccd Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/stablelm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/torch_native_llama.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/torch_native_llama.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eeb1f38b586fe12c47695e7555f6098b6024d9cc Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/torch_native_llama.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/xverse.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/xverse.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e0f12b314c00aef0c6a09030c3c445a8a5c5eb24 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/xverse.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/xverse_moe.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/xverse_moe.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..52fab55c796b0132c1e265e18ba0f8a6426f97e7 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/xverse_moe.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/__pycache__/yivl.cpython-312.pyc b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/yivl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..72c2e44b2560e1342fd1d4883c62c9d0ca4c73e3 Binary files /dev/null and b/lib/python3.12/site-packages/sglang/srt/models/__pycache__/yivl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sglang/srt/models/bert.py b/lib/python3.12/site-packages/sglang/srt/models/bert.py new file mode 100644 index 0000000000000000000000000000000000000000..46d2e7265fd9ab2244452f15dbf53d39e3990f26 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/bert.py @@ -0,0 +1,398 @@ +# SPDX-License-Identifier: Apache-2.0 +from typing import Any, Dict, Iterable, Optional, Set, Tuple + +import torch +from torch import nn + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.activation import get_act_fn +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.pooler import EmbeddingPoolerOutput, Pooler, PoolingType +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import AttentionType, RadixAttention +from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader + +BertConfig = None + + +class BertEmbedding(nn.Module): + + def __init__(self, config: BertConfig): + + super().__init__() + self.size = config.hidden_size + self.word_embeddings = VocabParallelEmbedding( + config.vocab_size, config.hidden_size + ) + self.position_embeddings = VocabParallelEmbedding( + config.max_position_embeddings, config.hidden_size + ) + self.token_type_embeddings = VocabParallelEmbedding( + config.type_vocab_size, config.hidden_size + ) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.position_ids = nn.Parameter( + torch.empty((1, config.max_position_embeddings)), + ) + + self.position_embedding_type = config.position_embedding_type + if self.position_embedding_type != "absolute": + raise ValueError( + "Only 'absolute' position_embedding_type" + " is supported" + ) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + input_shape = input_ids.size() + + # Input embeddings. + inputs_embeds = self.word_embeddings(input_ids) + + # Position embeddings. + position_embeddings = self.position_embeddings(position_ids) + + token_type_ids = torch.zeros( + input_shape, dtype=torch.long, device=inputs_embeds.device + ) + + token_type_embeddings = self.token_type_embeddings(token_type_ids) + + embeddings = inputs_embeds + token_type_embeddings + position_embeddings + embeddings = self.LayerNorm(embeddings) + return embeddings + + +class BertEncoder(nn.Module): + + def __init__( + self, + config: BertConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.quant_config = quant_config + self.layer = nn.ModuleList( + [ + BertLayer( + config=config, + layer_id=layer_idx, + quant_config=quant_config, + prefix=f"{prefix}.layer.{layer_idx}", + ) + for layer_idx in range(config.num_hidden_layers) + ] + ) + + def forward( + self, hidden_states: torch.Tensor, forward_batch: ForwardBatch + ) -> torch.Tensor: + for layer in self.layer: + hidden_states = layer(hidden_states, forward_batch) + return hidden_states + + +class BertLayer(nn.Module): + + def __init__( + self, + config: BertConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + + self.attention = BertAttention( + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + layer_id=layer_id, + layer_norm_eps=config.layer_norm_eps, + quant_config=quant_config, + prefix=f"{prefix}.attention", + ) + + self.intermediate = BertIntermediate( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.intermediate", + ) + + self.output = BertOutput( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + layer_norm_eps=config.layer_norm_eps, + quant_config=quant_config, + prefix=f"{prefix}.output", + ) + + def forward(self, hidden_states: torch.Tensor, forward_batch: ForwardBatch): + attn_output = self.attention(hidden_states, forward_batch) + intermediate_output = self.intermediate(attn_output) + output = self.output(intermediate_output, attn_output) + return output + + +class BertAttention(nn.Module): + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + layer_norm_eps: float, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + + self.self_attn = BertSelfAttention( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=f"{prefix}.output", + ) + + self.output = BertSelfOutput( + hidden_size=hidden_size, + layer_norm_eps=layer_norm_eps, + quant_config=quant_config, + prefix=f"{prefix}.output", + ) + + def forward( + self, hidden_states: torch.Tensor, forward_batch: ForwardBatch + ) -> torch.Tensor: + self_output = self.self_attn(hidden_states, forward_batch) + return self.output(self_output, hidden_states) + + +class BertSelfAttention(nn.Module): + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + + self.total_num_heads = num_attention_heads + assert self.total_num_heads % tp_size == 0 + + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = self.total_num_heads + self.head_dim = self.hidden_size // self.total_num_heads + assert self.head_dim * self.total_num_heads == self.hidden_size + + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.qkv_proj = QKVParallelLinear( + hidden_size=self.hidden_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.attn = RadixAttention( + num_heads=self.num_heads, + head_dim=self.head_dim, + scaling=self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + prefix=f"{prefix}.attn", + attn_type=AttentionType.ENCODER_ONLY, + ) + + def forward( + self, hidden_states: torch.Tensor, forward_batch: ForwardBatch + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + output = self.attn(q, k, v, forward_batch) + return output + + +class BertSelfOutput(nn.Module): + + def __init__( + self, + hidden_size: int, + layer_norm_eps: float, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.dense = RowParallelLinear( + input_size=hidden_size, + output_size=hidden_size, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.dense", + ) + self.LayerNorm = nn.LayerNorm(hidden_size, eps=layer_norm_eps) + + def forward( + self, hidden_states: torch.Tensor, input_tensor: torch.Tensor + ) -> torch.Tensor: + hidden_states, _ = self.dense(hidden_states) + hidden_states = self.LayerNorm(hidden_states + input_tensor) + return hidden_states + + +class BertIntermediate(nn.Module): + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.dense = ColumnParallelLinear( + input_size=hidden_size, + output_size=intermediate_size, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.dense", + ) + self.intermediate_act_fn = get_act_fn(hidden_act) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, _ = self.dense(hidden_states) + hidden_states = self.intermediate_act_fn(hidden_states) + return hidden_states + + +class BertOutput(nn.Module): + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + layer_norm_eps: float, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + + self.dense = RowParallelLinear( + input_size=intermediate_size, + output_size=hidden_size, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.dense", + ) + + self.LayerNorm = nn.LayerNorm(hidden_size, eps=layer_norm_eps) + + def forward( + self, hidden_states: torch.Tensor, input_tensor: torch.Tensor + ) -> torch.Tensor: + hidden_states, _ = self.dense(hidden_states) + hidden_states = self.LayerNorm(hidden_states + input_tensor) + return hidden_states + + +class BertModel(nn.Module): + + def __init__( + self, + *, + config: BertConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.embeddings = BertEmbedding(config) + self.encoder = BertEncoder( + config=config, quant_config=quant_config, prefix=f"encoder" + ) + self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) + # self.pooler = BertPooler(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + get_embedding: bool = False, + ) -> torch.Tensor: + assert get_embedding == True + # Your tokenized IDs + + hidden_states = self.embeddings( + input_ids=input_ids, + position_ids=positions, + ) + + hidden_states = self.encoder(hidden_states, forward_batch=forward_batch) + return self.pooler(hidden_states, forward_batch) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "query", "q"), + ("qkv_proj", "key", "k"), + ("qkv_proj", "value", "v"), + ] + + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + name = name.replace("self", "self_attn") + if "pooler" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +class Contriever(BertModel): + pass + + +EntryClass = [BertModel, Contriever] diff --git a/lib/python3.12/site-packages/sglang/srt/models/chatglm.py b/lib/python3.12/site-packages/sglang/srt/models/chatglm.py new file mode 100644 index 0000000000000000000000000000000000000000..9cf585a02bde6f7f4682c8c6abdc1e861572c156 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/chatglm.py @@ -0,0 +1,427 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from +# https://github.com/THUDM/ChatGLM2-6B +"""Inference-only ChatGLM model compatible with THUDM weights.""" + +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn +from torch.nn import LayerNorm + +from sglang.srt.configs import ChatGLMConfig +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix + +LoraConfig = None + + +class GLMAttention(nn.Module): + def __init__( + self, + config, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.hidden_size = config.hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.multi_query_attention = config.multi_query_attention + self.total_num_kv_heads = ( + config.multi_query_group_num + if config.multi_query_attention + else config.num_attention_heads + ) + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = config.hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + self.query_key_value = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=config.add_bias_linear or config.add_qkv_bias, + quant_config=quant_config, + prefix=add_prefix("query_key_value", prefix), + ) + self.dense = RowParallelLinear( + self.total_num_heads * self.head_dim, + config.hidden_size, + bias=config.add_bias_linear, + quant_config=quant_config, + prefix=add_prefix("dense", prefix), + ) + + # https://huggingface.co/THUDM/chatglm3-6b-32k/blob/e210410255278dd9d74463cf396ba559c0ef801c/modeling_chatglm.py#L141 + rope_ratio = getattr(config, "rope_ratio", 1.0) + max_positions = getattr(config, "seq_length", 8192) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim // 2, + max_position=max_positions, + base=10000 * rope_ratio, + is_neox_style=False, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.query_key_value(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(position_ids, q, k) + context_layer = self.attn( + q, + k, + v, + forward_batch, + ) + attn_output, _ = self.dense(context_layer) + return attn_output + + +class GLMMLP(nn.Module): + """MLP. + + MLP will take the input with h hidden state, project it to 4*h + hidden dimension, perform nonlinear transformation, and project the + state back into h hidden dimension. + """ + + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + + self.add_bias = config.add_bias_linear + + # Project to 4h. + self.dense_h_to_4h = MergedColumnParallelLinear( + config.hidden_size, + [config.ffn_hidden_size] * 2, + bias=config.add_bias_linear, + quant_config=quant_config, + prefix=add_prefix("dense_h_to_4h", prefix), + ) + + self.activation_func = SiluAndMul() + + # Project back to h. + self.dense_4h_to_h = RowParallelLinear( + config.ffn_hidden_size, + config.hidden_size, + bias=config.add_bias_linear, + quant_config=quant_config, + prefix=add_prefix("dense_4h_to_h", prefix), + ) + + def forward(self, hidden_states): + # [s, b, 4hp] + intermediate_parallel, _ = self.dense_h_to_4h(hidden_states) + intermediate_parallel = self.activation_func(intermediate_parallel) + # [s, b, h] + output, _ = self.dense_4h_to_h(intermediate_parallel) + return output + + +class GLMBlock(nn.Module): + """A single transformer layer. + + Transformer layer takes input with size [s, b, h] and returns an + output of the same size. + """ + + def __init__( + self, + config, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.apply_residual_connection_post_layernorm = ( + config.apply_residual_connection_post_layernorm + ) + + self.fp32_residual_connection = config.fp32_residual_connection + + layer_norm_func = RMSNorm if config.rmsnorm else LayerNorm + # Layernorm on the input data. + self.input_layernorm = layer_norm_func( + config.hidden_size, eps=config.layernorm_epsilon + ) + + # Self attention. + self.self_attention = GLMAttention( + config, layer_id, quant_config, prefix=add_prefix("self_attention", prefix) + ) + self.hidden_dropout = config.hidden_dropout + + # Layernorm on the attention output + self.post_attention_layernorm = layer_norm_func( + config.hidden_size, eps=config.layernorm_epsilon + ) + + # MLP + self.mlp = GLMMLP(config, quant_config, prefix=add_prefix("mlp", prefix)) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + # hidden_states: [num_tokens, h] + # Layer norm at the beginning of the transformer layer. + layernorm_output = self.input_layernorm(hidden_states) + # Self attention. + attention_output = self.self_attention( + hidden_states=layernorm_output, + position_ids=position_ids, + forward_batch=forward_batch, + ) + + # Residual connection. + if self.apply_residual_connection_post_layernorm: + residual = layernorm_output + else: + residual = hidden_states + + layernorm_input = residual + attention_output + + # Layer norm post the self attention. + layernorm_output = self.post_attention_layernorm(layernorm_input) + + # Second residual connection. + if self.apply_residual_connection_post_layernorm: + residual = layernorm_output + else: + residual = layernorm_input + + output = self.mlp(layernorm_output) + residual + + return output + + +class GLMTransformer(nn.Module): + """Transformer class.""" + + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.post_layer_norm = config.post_layer_norm + + # Number of layers. + self.num_layers = config.num_layers + + # Transformer layers. + self.layers = nn.ModuleList( + [ + GLMBlock( + config, + i, + quant_config, + prefix=add_prefix(f"layers.{i}", prefix), + ) + for i in range(self.num_layers) + ] + ) + + if self.post_layer_norm: + layer_norm_func = RMSNorm if config.rmsnorm else LayerNorm + # Final layer norm before output. + self.final_layernorm = layer_norm_func( + config.hidden_size, eps=config.layernorm_epsilon + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + for i in range(self.num_layers): + layer = self.layers[i] + hidden_states = layer( + hidden_states=hidden_states, + position_ids=position_ids, + forward_batch=forward_batch, + ) + # Final layer norm. + if self.post_layer_norm: + hidden_states = self.final_layernorm(hidden_states) + + return hidden_states + + +class ChatGLMM(nn.Module): + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + + self.embedding = VocabParallelEmbedding( + config.padded_vocab_size, + config.hidden_size, + prefix=add_prefix("embedding", prefix), + ) + + self.num_layers = config.num_layers + self.multi_query_group_num = config.multi_query_group_num + self.kv_channels = config.kv_channels + self.encoder = GLMTransformer( + config, quant_config, add_prefix("encoder", prefix) + ) + + self.output_layer = ParallelLMHead( + config.padded_vocab_size, + config.hidden_size, + prefix=add_prefix("output_layer", prefix), + ) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + inputs_embeds = self.embedding(input_ids) + + # Run encoder. + hidden_states = self.encoder( + hidden_states=inputs_embeds, + position_ids=position_ids, + forward_batch=forward_batch, + ) + return hidden_states + + +class ChatGLMForCausalLM(nn.Module): + packed_modules_mapping = { + "query_key_value": ["query_key_value"], + "dense_h_to_4h": ["dense_h_to_4h"], + } + # LoRA specific attributes + supported_lora_modules = [ + "query_key_value", + "dense", + "dense_h_to_4h", + "dense_4h_to_h", + ] + embedding_modules = {} + embedding_padding_modules = [] + + def __init__( + self, + config: ChatGLMConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config: ChatGLMConfig = config + self.quant_config = quant_config + self.max_position_embeddings = getattr(config, "max_sequence_length", 8192) + self.transformer = ChatGLMM( + config, quant_config, prefix=add_prefix("transformer", prefix) + ) + self.lm_head = self.transformer.output_layer + self.logits_processor = LogitsProcessor(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + hidden_states = self.transformer(input_ids, positions, forward_batch) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + params_dict = dict(self.named_parameters(remove_duplicate=False)) + for name, loaded_weight in weights: + if "rotary_pos_emb.inv_freq" in name: + continue + if "word_embeddings" in name: + name = name.replace(".word_embeddings", "") + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +class ChatGLMModel(ChatGLMForCausalLM): + pass + + +EntryClass = [ChatGLMModel] diff --git a/lib/python3.12/site-packages/sglang/srt/models/clip.py b/lib/python3.12/site-packages/sglang/srt/models/clip.py new file mode 100644 index 0000000000000000000000000000000000000000..eecb9208def870940b0151b30128c99d68521a9e --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/clip.py @@ -0,0 +1,568 @@ +# Adapted from +# https://github.com/huggingface/transformers/blob/af9b2eaa54c150741f298d6db939af6328e1dc38/src/transformers/models/clip/modeling_clip.py + +from functools import partial +from typing import Iterable, List, Optional, Tuple, Type, Union + +import torch +import torch.nn as nn +from transformers import CLIPConfig, CLIPTextConfig, CLIPVisionConfig +from transformers.modeling_attn_mask_utils import _create_4d_causal_attention_mask + +from sglang.srt.layers.activation import QuickGELU +from sglang.srt.layers.attention.vision import VisionAttention +from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear +from sglang.srt.layers.pooler import EmbeddingPoolerOutput, Pooler, PoolingType +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.managers.schedule_batch import MultimodalInputs +from sglang.srt.model_executor.model_runner import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix, flatten_nested_list + + +class CLIPVisionEmbeddings(nn.Module): + + def __init__(self, config: CLIPVisionConfig): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.image_size = config.image_size + self.patch_size = config.patch_size + assert self.image_size % self.patch_size == 0 + + self.class_embedding = nn.Parameter(torch.randn(self.embed_dim)) + + self.patch_embedding = nn.Conv2d( + in_channels=config.num_channels, + out_channels=self.embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + bias=False, + ) + + self.num_patches = (self.image_size // self.patch_size) ** 2 + self.num_positions = self.num_patches + 1 + self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) + self.register_buffer( + "position_ids", + torch.arange(self.num_positions).expand((1, -1)), + persistent=False, + ) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + batch_size = pixel_values.shape[0] + target_dtype = self.patch_embedding.weight.dtype + patch_embeds = self.patch_embedding( + pixel_values.to(dtype=target_dtype) + ) # shape = [*, width, grid, grid] + patch_embeds = patch_embeds.flatten(2).transpose(1, 2) + + class_embeds = self.class_embedding.expand(batch_size, 1, -1) + embeddings = torch.cat([class_embeds, patch_embeds], dim=1) + embeddings = embeddings + self.position_embedding(self.position_ids) + + return embeddings + + +class CLIPTextEmbeddings(nn.Module): + def __init__(self, config: CLIPTextConfig): + super().__init__() + embed_dim = config.hidden_size + + self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) + self.position_embedding = nn.Embedding( + config.max_position_embeddings, embed_dim + ) + + # position_ids (1, len position emb) is contiguous in memory and exported when serialized + self.register_buffer( + "position_ids", + torch.arange(config.max_position_embeddings).expand((1, -1)), + persistent=False, + ) + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + ) -> torch.Tensor: + seq_length = ( + input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2] + ) + + if position_ids is None: + position_ids = self.position_ids[:, :seq_length] + + if inputs_embeds is None: + inputs_embeds = self.token_embedding(input_ids) + + position_embeddings = self.position_embedding(position_ids) + embeddings = inputs_embeds + position_embeddings + + return embeddings + + +class CLIPMLP(nn.Module): + + def __init__( + self, + config, + act_layer: Type[nn.Module] = QuickGELU, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.fc1 = ColumnParallelLinear( + config.hidden_size, + config.intermediate_size, + quant_config=quant_config, + prefix=add_prefix("fc1", prefix), + ) + self.act = act_layer() + self.fc2 = RowParallelLinear( + config.intermediate_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("fc2", prefix), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x_parallel, _ = self.fc1(x) + x_parallel = self.act(x_parallel) + x, _ = self.fc2(x_parallel) + return x + + +class CLIPEncoderLayer(nn.Module): + + def __init__( + self, + config: CLIPVisionConfig, + act_layer: Type[nn.Module] = QuickGELU, + norm_layer: Type[nn.Module] = None, + attn_implementation: Optional[str] = "sdpa", + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + if norm_layer is None: + norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps) + self.layer_norm1 = norm_layer(config.hidden_size) + self.layer_norm2 = norm_layer(config.hidden_size) + if attn_implementation == "sdpa": + use_context_forward = False + softmax_in_single_precision = False + elif attn_implementation == "flash_attention_2": + softmax_in_single_precision = False + use_context_forward = True + elif attn_implementation == "eager": + softmax_in_single_precision = True + use_context_forward = False + self.self_attn = VisionAttention( + embed_dim=config.hidden_size, + num_heads=config.num_attention_heads, + projection_size=config.hidden_size, + use_qkv_parallel=True, + use_context_forward=use_context_forward, + softmax_in_single_precision=softmax_in_single_precision, + flatten_batch=True, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + self.mlp = CLIPMLP( + config, + act_layer=act_layer, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + causal_attention_mask: torch.Tensor, + ) -> torch.Tensor: + + residual = hidden_states + hidden_states = self.layer_norm1(hidden_states) + # CLIP text model uses both `causal_attention_mask` and `attention_mask` + if attention_mask is not None and causal_attention_mask is not None: + attn_mask = attention_mask + causal_attention_mask + elif causal_attention_mask is not None: + attn_mask = causal_attention_mask + else: + attn_mask = attention_mask + hidden_states = self.self_attn( + hidden_states, + attention_mask=attn_mask, + # causal_attention_mask=causal_attention_mask, + ) + + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +class CLIPEncoder(nn.Module): + """ + Transformer encoder consisting of `config.num_hidden_layers` self + attention layers. Each layer is a [`CLIPEncoderLayer`]. + + Args: + config: CLIPConfig + """ + + def __init__( + self, + config: CLIPVisionConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + + num_hidden_layers = config.num_hidden_layers + norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps) + self.layers = nn.ModuleList( + [ + CLIPEncoderLayer( + config=config, + norm_layer=norm_layer, + attn_implementation="sdpa", + quant_config=quant_config, + prefix=add_prefix(f"layers.{layer_idx}", prefix), + ) + for layer_idx in range(num_hidden_layers) + ] + ) + + def forward( + self, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor = None, + causal_attention_mask: torch.Tensor = None, + return_all_hidden_states: bool = False, + ) -> Union[torch.Tensor, list[torch.Tensor]]: + hidden_states_pool = [inputs_embeds] + hidden_states = inputs_embeds + + for encoder_layer in self.layers: + hidden_states = encoder_layer( + hidden_states, attention_mask, causal_attention_mask + ) + if return_all_hidden_states: + hidden_states_pool.append(hidden_states) + if return_all_hidden_states: + return hidden_states_pool + return hidden_states + + +class CLIPTextTransformer(nn.Module): + def __init__( + self, + config: CLIPTextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + embed_dim = config.hidden_size + self.embeddings = CLIPTextEmbeddings(config) + self.encoder = CLIPEncoder( + config=config, + quant_config=quant_config, + prefix=add_prefix("encoder", prefix), + ) + self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + @property + def device(self) -> torch.device: + return self.encoder.layers[0].layer_norm1.weight.device + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + ): + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + hidden_states = self.embeddings(input_ids, position_ids) + causal_attention_mask = _create_4d_causal_attention_mask( + input_ids.shape, hidden_states.dtype, device=hidden_states.device + ) + encoder_outputs = self.encoder( + hidden_states, attention_mask, causal_attention_mask + ) + last_hidden_state = self.final_layer_norm(encoder_outputs) + return last_hidden_state + + +class CLIPTextModel(nn.Module): + def __init__( + self, + config: CLIPTextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.text_model = CLIPTextTransformer( + config=config, + quant_config=quant_config, + prefix=add_prefix("text_model", prefix), + ) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + ): + return self.text_model(input_ids, position_ids) + + +class CLIPVisionTransformer(nn.Module): + + def __init__( + self, + config: CLIPVisionConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + embed_dim = config.hidden_size + + self.embeddings = CLIPVisionEmbeddings(config) + + # NOTE: This typo of "layrnorm" is not fixed on purpose to match + # the original transformers code and name of the model weights. + self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + self.encoder = CLIPEncoder( + config=config, + quant_config=quant_config, + prefix=add_prefix("encoder", prefix), + ) + + num_hidden_layers = config.num_hidden_layers + if len(self.encoder.layers) > config.num_hidden_layers: + raise ValueError( + f"The original encoder only has {num_hidden_layers} " + f"layers, but you requested {len(self.encoder.layers)} layers." + ) + + self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + @property + def device(self) -> torch.device: + return self.encoder.layers[0].layer_norm1.weight.device + + def forward( + self, + pixel_values: torch.Tensor, + ) -> torch.Tensor: + hidden_states = self.embeddings(pixel_values.to(self.device)) + hidden_states = self.pre_layrnorm(hidden_states) + + return_all_hidden_states = False + + last_hidden_state = self.encoder( + inputs_embeds=hidden_states, + return_all_hidden_states=return_all_hidden_states, + ) + + last_hidden_state = self.post_layernorm(last_hidden_state) + + return last_hidden_state + + +class CLIPVisionModel(nn.Module): + def __init__( + self, + config: CLIPVisionConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.vision_model = CLIPVisionTransformer( + config, quant_config, prefix=add_prefix("vision_model", prefix) + ) + + def forward(self, pixel_values: torch.Tensor): + return self.vision_model(pixel_values) + + +class CLIPModel(nn.Module): + def __init__( + self, + config: CLIPConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + if not isinstance(config.text_config, CLIPTextConfig): + raise TypeError( + "config.text_config is expected to be of type CLIPTextConfig but is of type" + f" {type(config.text_config)}." + ) + + if not isinstance(config.vision_config, CLIPVisionConfig): + raise TypeError( + "config.vision_config is expected to be of type CLIPVisionConfig but is of type" + f" {type(config.vision_config)}." + ) + + text_config = config.text_config + vision_config = config.vision_config + + self.projection_dim = config.projection_dim + self.text_embed_dim = text_config.hidden_size + self.vision_embed_dim = vision_config.hidden_size + self.visual_projection = nn.Linear( + self.vision_embed_dim, self.projection_dim, bias=False + ) + self.text_projection = nn.Linear( + self.text_embed_dim, self.projection_dim, bias=False + ) + self.logit_scale = nn.Parameter( + torch.tensor(self.config.logit_scale_init_value) + ) + + text_model = CLIPTextModel( + text_config, quant_config, prefix=add_prefix("text_model", prefix) + ) + vision_model = CLIPVisionModel( + vision_config, quant_config, prefix=add_prefix("vision_model", prefix) + ) + self.text_model = text_model.text_model + self.vision_model = vision_model.vision_model + self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) + monkey_patch_weight_loader() + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + get_embedding: bool = True, + ): + assert get_embedding, "CLIPEmbeddingModel is only used for embedding" + mm_inputs = [] + if forward_batch.mm_inputs is not None: + mm_inputs = forward_batch.mm_inputs + pixel_values_list = [ + item.pixel_values + for item in flatten_nested_list( + [mm_input.mm_items for mm_input in mm_inputs if mm_input is not None] + ) + ] + if len(pixel_values_list) != 0: + pixel_values = torch.concat(pixel_values_list) + vision_outputs = self.vision_model(pixel_values) + pooled_output = vision_outputs[:, 0, :] + image_embeds = self.visual_projection(pooled_output) + image_embeds = nn.functional.normalize(image_embeds, p=2, dim=1) + return EmbeddingPoolerOutput(embeddings=image_embeds) + + else: + text_outputs = self.text_model(input_ids, position_ids=positions) + pooled_output = self.pooler(text_outputs[0], forward_batch) + return EmbeddingPoolerOutput( + embeddings=self.text_projection(pooled_output.embeddings) + ) + + def pad_input_ids(self, input_ids: List[int], image_inputs: MultimodalInputs): + # Clip embeddings models handle text/image separately, so we don't need to pad input ids + return input_ids + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ] + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if "position_ids" in name: + continue + if "out_proj" in name: + name = name.replace("out_proj", "proj") + for param_name, shard_name, shard_id in stacked_params_mapping: + if shard_name not in name: + continue + name = name.replace(shard_name, param_name) + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +# monkey patch weight loader to remove open_clip file +def monkey_patch_weight_loader(): + import glob + import os + + from sglang.srt.model_loader.loader import DefaultModelLoader + from sglang.srt.model_loader.weight_utils import ( + download_weights_from_hf, + filter_files_not_needed_for_inference, + ) + + def prepare_weights( + self, model_name_or_path: str, revision: Optional[str], fall_back_to_pt: bool + ) -> Tuple[str, List[str], bool]: + model_name_or_path = ( + self._maybe_download_from_modelscope(model_name_or_path, revision) + or model_name_or_path + ) + + is_local = os.path.isdir(model_name_or_path) + use_safetensors = False + allow_patterns = ["*.bin"] + + if not is_local: + hf_folder = download_weights_from_hf( + model_name_or_path, + self.load_config.download_dir, + allow_patterns, + revision, + ignore_patterns=self.load_config.ignore_patterns, + ) + else: + hf_folder = model_name_or_path + + hf_weights_files: List[str] = [] + for pattern in allow_patterns: + hf_weights_files += glob.glob(os.path.join(hf_folder, pattern)) + + hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files) + + # remove open_clip file + hf_weights_files = [ + file for file in hf_weights_files if "open_clip" not in file + ] + + if len(hf_weights_files) == 0: + raise RuntimeError( + f"Cannot find any model weights with `{model_name_or_path}`" + ) + + return hf_folder, hf_weights_files, use_safetensors + + setattr(DefaultModelLoader, "_prepare_weights", prepare_weights) + + +EntryClass = CLIPModel diff --git a/lib/python3.12/site-packages/sglang/srt/models/commandr.py b/lib/python3.12/site-packages/sglang/srt/models/commandr.py new file mode 100644 index 0000000000000000000000000000000000000000..ebbf8ed64029b771e9cffbe4a74030f165d889bc --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/commandr.py @@ -0,0 +1,418 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# Copyright 2024 Cohere and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from +# https://github.com/vllm-project/vllm/blob/c7f2cf2b7f67bce5842fedfdba508440fe257375/vllm/model_executor/models/commandr.py#L1 + +# This file is based on the LLama model definition file in transformers +"""PyTorch Cohere model.""" + +from typing import Iterable, Optional, Tuple + +import torch +import torch.utils.checkpoint +from torch import nn +from torch.nn.parameter import Parameter +from transformers import PretrainedConfig + +from sglang.srt.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from sglang.srt.utils import add_prefix, get_compiler_backend, set_weight_attrs + + +@torch.compile(backend=get_compiler_backend()) +def layer_norm_func(hidden_states, weight, variance_epsilon): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + mean = hidden_states.mean(-1, keepdim=True) + variance = (hidden_states - mean).pow(2).mean(-1, keepdim=True) + hidden_states = (hidden_states - mean) * torch.rsqrt(variance + variance_epsilon) + hidden_states = weight.to(torch.float32) * hidden_states + return hidden_states.to(input_dtype) + + +class LayerNorm(nn.Module): + def __init__(self, param_shape=None, eps=1e-5): + super().__init__() + self.weight = nn.Parameter(torch.ones(param_shape)) + self.variance_epsilon = eps + set_weight_attrs(self.weight, {"weight_loader": self.weight_loader}) + + def forward(self, hidden_states, residuals=None): + hidden_states = layer_norm_func( + hidden_states, self.weight, self.variance_epsilon + ) + return hidden_states, residuals + + def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): + tp_rank = get_tensor_model_parallel_rank() + shard_dim = 0 if param.dim() != 1 else None + param_data = param.data + if shard_dim is not None: + shard_size = param_data.shape[shard_dim] + start_idx = tp_rank * shard_size + loaded_weight = loaded_weight.narrow(shard_dim, start_idx, shard_size) + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + +# Copied from transformers.models.llama.modeling_llama.LlamaMLP Llama->Cohere +class CohereMLP(nn.Module): + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_up_proj = MergedColumnParallelLinear( + self.hidden_size, + [self.intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + self.intermediate_size, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class CohereAttention(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + tp_size = get_tensor_model_parallel_world_size() + self.config = config + self.attention_dropout = config.attention_dropout + self.hidden_size = config.hidden_size + self.total_num_heads = config.num_attention_heads + self.num_heads = self.total_num_heads // tp_size + self.head_dim = self.hidden_size // self.total_num_heads + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.max_position_embeddings = getattr( + config, "model_max_length", None + ) or getattr(config, "max_position_embeddings", 8192) + self.rope_theta = config.rope_theta + self.rope_scaling = getattr(config, "rope_scaling", None) + self.use_qk_norm = getattr(config, "use_qk_norm", False) + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=self.max_position_embeddings, + base=self.rope_theta, + rope_scaling=self.rope_scaling, + is_neox_style=False, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + if self.use_qk_norm: + self.q_norm = LayerNorm( + param_shape=(self.num_heads, self.head_dim), eps=config.layer_norm_eps + ) + self.k_norm = LayerNorm( + param_shape=(self.num_kv_heads, self.head_dim), + eps=config.layer_norm_eps, + ) + + def _apply_qk_norm(self, q, k): + q = q.view(*q.shape[:-1], -1, self.head_dim) + k = k.view(*k.shape[:-1], -1, self.head_dim) + q, _ = self.q_norm(q) + k, _ = self.k_norm(k) + q = q.view(*q.shape[:-2], -1) + k = k.view(*k.shape[:-2], -1) + return q, k + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + if self.use_qk_norm: + q, k = self._apply_qk_norm(q, k) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class CohereDecoderLayer(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = CohereAttention( + config, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + + self.mlp = CohereMLP( + config, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + self.input_layernorm = LayerNorm( + param_shape=(config.hidden_size), eps=config.layer_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Self Attention + residual = hidden_states + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states_attention = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + hidden_states_mlp = self.mlp(hidden_states) + # Add everything together + hidden_states = residual + hidden_states_attention + hidden_states_mlp + + return hidden_states, residual + + +class CohereModel(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.vocab_size = config.vocab_size + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, config.hidden_size + ) + self.layers = nn.ModuleList( + [ + CohereDecoderLayer( + config, + i, + quant_config=quant_config, + prefix=add_prefix(f"layers.{i}", prefix), + ) + for i in range(config.num_hidden_layers) + ] + ) + self.norm = LayerNorm( + param_shape=(config.hidden_size), eps=config.layer_norm_eps + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + hidden_states = self.embed_tokens(input_ids) + residual = None + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + forward_batch, + residual, + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class CohereForCausalLM(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.logits_processor = LogitsProcessor(config) + self.model = CohereModel( + config, quant_config, prefix=add_prefix("model", prefix) + ) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + hidden_states = self.model( + input_ids, + positions, + forward_batch, + ) + return self.logits_processor( + input_ids, hidden_states, self.model.embed_tokens, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.named_parameters()) + loaded_params = set() + for name, loaded_weight in weights: + for param_name, shard_name, shard_id in stacked_params_mapping: + if shard_name not in name: + continue + name = name.replace(shard_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # lm_head is not used in vllm as it is tied with embed_token. + # To prevent errors, skip loading lm_head.weight. + if "lm_head.weight" in name: + continue + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Remapping the name of FP8 kv-scale. + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + +class Cohere2ForCausalLM(CohereForCausalLM): + pass + + +EntryClass = [CohereForCausalLM, Cohere2ForCausalLM] diff --git a/lib/python3.12/site-packages/sglang/srt/models/dbrx.py b/lib/python3.12/site-packages/sglang/srt/models/dbrx.py new file mode 100644 index 0000000000000000000000000000000000000000..15cef015c6dd683ff4b9131a7421435ff1ca021c --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/dbrx.py @@ -0,0 +1,454 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from: +# https://github.com/vllm-project/vllm/blob/c7f2cf2b7f67bce5842fedfdba508440fe257375/vllm/model_executor/models/dbrx.py#L1 + +from typing import Iterable, Optional, Tuple + +import torch +import torch.nn as nn + +from sglang.srt.configs import DbrxConfig +from sglang.srt.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from sglang.srt.layers.linear import ( + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.fused_moe_triton import fused_moe +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + DEFAULT_VOCAB_PADDING_SIZE, + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from sglang.srt.utils import add_prefix, set_weight_attrs + + +class DbrxRouter(nn.Module): + """A Router implementation for DBRX that returns logits for each expert + per token. + """ + + def __init__( + self, + config: DbrxConfig, + params_dtype: Optional[torch.dtype] = None, + prefix: str = "", + ): + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + self.num_total_experts = config.ffn_config.moe_num_experts + self.d_model = config.d_model + self.layer = ReplicatedLinear( + self.d_model, + self.num_total_experts, + bias=False, + params_dtype=params_dtype, + quant_config=None, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + router_logits, _ = self.layer(hidden_states) + return router_logits + + +class DbrxExperts(nn.Module): + """A tensor-parallel MoE implementation for DBRX. + + Each expert's weights are sharded across all ranks and a fused MoE + kernel is used for the forward pass, and finally we reduce the outputs + across ranks. + """ + + def __init__( + self, + config: DbrxConfig, + quant_config: Optional[QuantizationConfig] = None, + params_dtype: Optional[torch.dtype] = None, + prefix: str = "", + ): + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + self.num_total_experts = config.ffn_config.moe_num_experts + self.top_k = config.ffn_config.moe_top_k + self.d_model = config.d_model + self.intermediate_size = config.ffn_config.ffn_hidden_size // self.tp_size + + if params_dtype is None: + params_dtype = torch.get_default_dtype() + self.params_dtype = params_dtype + + self.router = DbrxRouter(config, self.params_dtype) + self.ws = nn.Parameter( + torch.empty( + self.num_total_experts, + 2 * self.intermediate_size, + self.d_model, + device="cuda", + dtype=self.params_dtype, + ) + ) + self.w2s = nn.Parameter( + torch.empty( + self.num_total_experts, + self.d_model, + self.intermediate_size, + device="cuda", + dtype=self.params_dtype, + ) + ) + + set_weight_attrs( + self.ws, + { + "weight_loader": self.weight_loader, + }, + ) + set_weight_attrs( + self.w2s, + { + "weight_loader": self.weight_loader, + }, + ) + + def weight_loader( + self, param: nn.Parameter, loaded_weight: torch.Tensor, weight_name: str + ): + tp_rank = get_tensor_model_parallel_rank() + param_data = param.data + shard_size = self.intermediate_size + shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size) + # DBRX uses GLU for each experts. + # GLU has 3 linear layers: w1, v1 and w2. + if weight_name.endswith("w1"): + loaded_weight = torch.reshape( + loaded_weight, + [-1, self.intermediate_size * self.tp_size, self.d_model], + ) + param_data[:, 0:shard_size, :] = loaded_weight[:, shard, :] + if weight_name.endswith("v1"): + loaded_weight = torch.reshape( + loaded_weight, + [-1, self.intermediate_size * self.tp_size, self.d_model], + ) + param_data[:, shard_size : 2 * shard_size, :] = loaded_weight[:, shard, :] + if weight_name.endswith("w2"): + loaded_weight = torch.reshape( + loaded_weight, + [-1, self.intermediate_size * self.tp_size, self.d_model], + ).transpose(1, 2) + param_data[:] = loaded_weight[:, :, shard] + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_tokens, hidden_size = hidden_states.shape + hidden_states = hidden_states.view(-1, self.d_model) + # router_logits: (num_tokens, n_experts) + router_logits = self.router(hidden_states) + final_hidden_states = fused_moe( + hidden_states, + self.ws, + self.w2s, + router_logits, + self.top_k, + renormalize=True, + inplace=True, + ) + + if self.tp_size > 1: + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + + return final_hidden_states.view(num_tokens, hidden_size) + + +class DbrxAttention(nn.Module): + def __init__( + self, + config: DbrxConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.d_model = config.d_model + self.total_num_heads = config.n_heads + self.head_dim = self.d_model // self.total_num_heads + self.total_num_kv_heads = config.attn_config.kv_n_heads + self.clip_qkv = config.attn_config.clip_qkv + self.rope_theta = config.attn_config.rope_theta + self.max_position = config.max_seq_len + + # pylint: disable=invalid-name + self.Wqkv = QKVParallelLinear( + self.d_model, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=add_prefix("Wqkv", prefix), + ) + self.out_proj = RowParallelLinear( + self.d_model, + self.d_model, + bias=False, + quant_config=quant_config, + prefix=add_prefix("out_proj", prefix), + ) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=self.max_position, + base=int(self.rope_theta), + is_neox_style=True, + ) + + tp_world_size = get_tensor_model_parallel_world_size() + self.tp_size = tp_world_size + assert self.total_num_heads % tp_world_size == 0 + self.num_heads = self.total_num_heads // tp_world_size + if self.total_num_kv_heads >= tp_world_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_world_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_world_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_world_size) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.Wqkv(hidden_states) + if self.clip_qkv is not None: + qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(position_ids, q, k) + attn_output = self.attn(q, k, v, forward_batch) + hidden_states, _ = self.out_proj(attn_output) + return hidden_states + + +class DbrxFusedNormAttention(nn.Module): + def __init__( + self, + config: DbrxConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.d_model = config.d_model + self.attn = DbrxAttention( + config, + layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + self.norm_1 = nn.LayerNorm(self.d_model) + self.norm_2 = nn.LayerNorm(self.d_model) + + def forward( + self, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.norm_1(hidden_states) + x = self.attn( + position_ids=position_ids, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + hidden_states = residual + x + residual = hidden_states + hidden_states = self.norm_2(hidden_states) + return hidden_states, residual + + +class DbrxBlock(nn.Module): + def __init__( + self, + config: DbrxConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.norm_attn_norm = DbrxFusedNormAttention( + config, + layer_id, + quant_config=quant_config, + prefix=add_prefix("norm_attn_norm", prefix), + ) + self.ffn = DbrxExperts(config, quant_config=quant_config) + + def forward( + self, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + hidden_states, residual = self.norm_attn_norm( + position_ids=position_ids, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + hidden_states = self.ffn(hidden_states) + hidden_states = hidden_states + residual + return hidden_states + + +class DbrxModel(nn.Module): + def __init__( + self, + config: DbrxConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.wte = VocabParallelEmbedding( + config.vocab_size, + config.d_model, + ) + self.blocks = nn.ModuleList( + [ + DbrxBlock( + config, + i, + quant_config=quant_config, + prefix=add_prefix(f"blocks.{i}", prefix), + ) + for i in range(config.n_layers) + ] + ) + self.norm_f = nn.LayerNorm(config.d_model, eps=1e-5) + for module in self.modules(): + if hasattr(module, "bias") and isinstance(module.bias, nn.Parameter): + # Remove the bias term in Linear and LayerNorm. + module.register_parameter("bias", None) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.wte(input_ids) + else: + hidden_states = input_embeds + for i in range(len(self.blocks)): + block = self.blocks[i] + hidden_states = block(position_ids, hidden_states, forward_batch) + hidden_states = self.norm_f(hidden_states) + return hidden_states + + +class DbrxForCausalLM(nn.Module): + def __init__( + self, + config: DbrxConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.quant_config = quant_config + self.unpadded_vocab_size = config.vocab_size + self.transformer = DbrxModel( + config, quant_config=quant_config, prefix=add_prefix("transformer", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.d_model, + org_num_embeddings=config.vocab_size, + padding_size=DEFAULT_VOCAB_PADDING_SIZE, + prefix=add_prefix("lm_head", prefix), + ) + self.logits_processor = LogitsProcessor(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + hidden_states = self.transformer(input_ids, positions, forward_batch) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + expert_params_mapping = [ + ( + "ws" if weight_name in ["w1", "v1"] else "w2s", + f"experts.mlp.{weight_name}", + ) + for weight_name in ["w1", "v1", "w2"] + ] + params_dict = dict(self.named_parameters(remove_duplicate=False)) + for name, loaded_weight in weights: + for param_name, weight_name in expert_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, weight_name) + break + else: + # Remapping the name of FP8 kv-scale. + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +EntryClass = DbrxForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/deepseek.py b/lib/python3.12/site-packages/sglang/srt/models/deepseek.py new file mode 100644 index 0000000000000000000000000000000000000000..95bfe001a2da811e2a4b54ccd229553a395384e1 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/deepseek.py @@ -0,0 +1,485 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from: +# https://github.com/vllm-project/vllm/blob/14f91fe67c2342f2fe859dc6a5c40810df0e1c61/vllm/model_executor/models/deepseek.py +"""Inference-only Deepseek model.""" + +from typing import Any, Dict, Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.fused_moe_triton import fused_moe +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix + + +class DeepseekMLP(nn.Module): + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + reduce_results: bool = True, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=add_prefix("down_proj", prefix), + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class DeepseekMoE(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.rank = get_tensor_model_parallel_rank() + self.tp_size = get_tensor_model_parallel_world_size() + self.n_routed_experts = config.n_routed_experts + self.top_k = config.num_experts_per_tok + if self.tp_size > self.n_routed_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {self.n_routed_experts}." + ) + + self.experts = nn.ModuleList( + [ + DeepseekMLP( + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix(f"{idx}.experts", prefix), + ) + for idx in range(self.n_routed_experts) + ] + ) + self.pack_params() + + self.gate = ReplicatedLinear( + config.hidden_size, + self.n_routed_experts, + bias=False, + quant_config=None, + prefix=add_prefix("gate", prefix), + ) + + if config.n_shared_experts is not None: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + self.shared_experts = DeepseekMLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix("shared_experts", prefix), + ) + + def pack_params(self): + w1 = [] + w2 = [] + for expert in self.experts: + w1.append(expert.gate_up_proj.weight) + w2.append(expert.down_proj.weight) + self.w1 = torch._utils._flatten_dense_tensors(w1) + w1s = torch._utils._unflatten_dense_tensors(self.w1, w1) + for data, param in zip(w1s, w1): + param.data = data + self.w1 = self.w1.view(len(w1), *w1s[0].shape) + + self.w2 = torch._utils._flatten_dense_tensors(w2) + w2s = torch._utils._unflatten_dense_tensors(self.w2, w2) + for data, param in zip(w2s, w2): + param.data = data + + self.w2 = self.w2.view(len(w2), *w2s[0].shape) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + if self.config.n_shared_experts is not None: + shared_output = self.shared_experts(hidden_states) + # router_logits: (num_tokens, n_experts) + router_logits, _ = self.gate(hidden_states) + final_hidden_states = fused_moe.fused_moe( + hidden_states, + self.w1, + self.w2, + router_logits, + self.top_k, + renormalize=self.config.norm_topk_prob, + inplace=True, + ) + + if self.config.n_shared_experts is not None: + final_hidden_states = final_hidden_states + shared_output + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + + return final_hidden_states.view(num_tokens, hidden_dim) + + +class DeepseekAttention(nn.Module): + + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + max_position_embeddings: int = 8192, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class DeepseekDecoderLayer(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + self.self_attn = DeepseekAttention( + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + layer_id=layer_id, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + if ( + config.n_routed_experts is not None + and layer_id >= config.first_k_dense_replace + and layer_id % config.moe_layer_freq == 0 + ): + self.mlp = DeepseekMoE( + config=config, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + else: + self.mlp = DeepseekMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> torch.Tensor: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class DeepseekModel(nn.Module): + + fall_back_to_pt_during_load = False + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) + self.layers = nn.ModuleList( + [ + DeepseekDecoderLayer( + config, + layer_id, + quant_config=quant_config, + prefix=add_prefix(f"layers.{layer_id}", prefix), + ) + for layer_id in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + + residual = None + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, hidden_states, forward_batch, residual + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class DeepseekForCausalLM(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.model = DeepseekModel( + config, quant_config, prefix=add_prefix("model", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + self.logits_processor = LogitsProcessor(config) + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.embed_tokens + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Skip experts that are not assigned to this worker. + if ( + "mlp.experts." in name or "mlp.shared_experts." in name + ) and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Skip experts that are not assigned to this worker. + if ( + "mlp.experts." in name or "mlp.shared_experts." in name + ) and name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +EntryClass = DeepseekForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/deepseek_janus_pro.py b/lib/python3.12/site-packages/sglang/srt/models/deepseek_janus_pro.py new file mode 100644 index 0000000000000000000000000000000000000000..04dcf787e2f08ca067c10065ec5ba5fb4b0a3f72 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/deepseek_janus_pro.py @@ -0,0 +1,2065 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Copied and Adapted from: +# https://github.com/deepseek-ai/Janus + + +import collections +import math +import os +from dataclasses import field +from enum import Enum +from functools import partial +from itertools import repeat +from typing import ( + Callable, + Final, + Iterable, + Literal, + Optional, + Sequence, + Set, + Tuple, + Type, + Union, +) + +import torch +import torch.nn.functional as F +from einops import rearrange +from torch import Tensor, _assert, nn +from torch.nn.init import trunc_normal_ +from transformers import AutoModel, PreTrainedModel + +from sglang.srt.configs.janus_pro import * +from sglang.srt.layers.attention.vision import VisionAttention +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization import QuantizationConfig +from sglang.srt.managers.mm_utils import ( + MultiModalityDataPaddingPatternTokenPairs, + general_mm_embed_routine, +) +from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.llama import LlamaForCausalLM +from sglang.utils import logger + +################################################################################# +# VQ Model Configs # +################################################################################# + + +# Copied from: +# https://github.com/deepseek-ai/Janus/tree/main/janus/models/vq_model.py +@dataclass +class ModelArgs: + codebook_size: int = 16384 + codebook_embed_dim: int = 8 + codebook_l2_norm: bool = True + codebook_show_usage: bool = True + commit_loss_beta: float = 0.25 + entropy_loss_ratio: float = 0.0 + + encoder_ch_mult: List[int] = field(default_factory=lambda: [1, 1, 2, 2, 4]) + decoder_ch_mult: List[int] = field(default_factory=lambda: [1, 1, 2, 2, 4]) + z_channels: int = 256 + dropout_p: float = 0.0 + + +def named_apply( + fn: Callable, + module: nn.Module, + name="", + depth_first: bool = True, + include_root: bool = False, +) -> nn.Module: + if not depth_first and include_root: + fn(module=module, name=name) + for child_name, child_module in module.named_children(): + child_name = ".".join((name, child_name)) if name else child_name + named_apply( + fn=fn, + module=child_module, + name=child_name, + depth_first=depth_first, + include_root=True, + ) + if depth_first and include_root: + fn(module=module, name=name) + return module + + +def VQ_16(**kwargs): + return VQModel( + ModelArgs( + encoder_ch_mult=[1, 1, 2, 2, 4], decoder_ch_mult=[1, 1, 2, 2, 4], **kwargs + ) + ) + + +VQ_models = {"VQ-16": VQ_16} + +import collections.abc + + +# From PyTorch internals +def _ntuple(n): + def parse(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + return tuple(x) + return tuple(repeat(x, n)) + + return parse + + +def _trunc_normal_(tensor, mean, std, a, b): + # Cut & paste from PyTorch official master until it's in a few official releases - RW + # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf + def norm_cdf(x): + # Computes standard normal cumulative distribution function + return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0 + + if (mean < a - 2 * std) or (mean > b + 2 * std): + logger.warn( + "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. " + "The distribution of values may be incorrect.", + stacklevel=2, + ) + + # Values are generated by using a truncated uniform distribution and + # then using the inverse CDF for the normal distribution. + # Get upper and lower cdf values + l = norm_cdf((a - mean) / std) + u = norm_cdf((b - mean) / std) + + # Uniformly fill tensor with values from [l, u], then translate to + # [2l-1, 2u-1]. + tensor.uniform_(2 * l - 1, 2 * u - 1) + + # Use inverse cdf transform for normal distribution to get truncated + # standard normal + if tensor.dtype in [torch.float16, torch.bfloat16]: + # The `erfinv_` op is not (yet?) defined in float16+cpu, bfloat16+gpu + og_dtype = tensor.dtype + tensor = tensor.to(torch.float32) + tensor.erfinv_() + tensor = tensor.to(og_dtype) + else: + tensor.erfinv_() + + # Transform to proper mean, std + tensor.mul_(std * math.sqrt(2.0)) + tensor.add_(mean) + + # Clamp to ensure it's in the proper range + if tensor.dtype == torch.float16: + # The `clamp_` op is not (yet?) defined in float16+cpu + tensor = tensor.to(torch.float32) + tensor.clamp_(min=a, max=b) + else: + tensor.clamp_(min=a, max=b) + + +def trunc_normal_tf_( + tensor: torch.Tensor, + mean: float = 0.0, + std: float = 1.0, + a: float = -2.0, + b: float = 2.0, +): + """Fills the input Tensor with values drawn from a truncated + normal distribution. The values are effectively drawn from the + normal distribution :math:`\\mathcal{N}(\text{mean}, \text{std}^2)` + with values outside :math:`[a, b]` redrawn until they are within + the bounds. The method used for generating the random values works + best when :math:`a \\leq \text{mean} \\leq b`. + NOTE: this 'tf' variant behaves closer to Tensorflow / JAX impl where the + bounds [a, b] are applied when sampling the normal distribution with mean=0, std=1.0 + and the result is subsquently scaled and shifted by the mean and std args. + Args: + tensor: an n-dimensional `torch.Tensor` + mean: the mean of the normal distribution + std: the standard deviation of the normal distribution + a: the minimum cutoff value + b: the maximum cutoff value + """ + with torch.no_grad(): + _trunc_normal_(tensor, 0, 1.0, a, b) + tensor.mul_(std).add_(mean) + + +to_2tuple = _ntuple(2) + + +class Format(str, Enum): + NCHW = "NCHW" + NHWC = "NHWC" + NCL = "NCL" + NLC = "NLC" + + +def nchw_to(x: torch.Tensor, fmt: Format): + if fmt == Format.NHWC: + x = x.permute(0, 2, 3, 1) + elif fmt == Format.NLC: + x = x.flatten(2).transpose(1, 2) + elif fmt == Format.NCL: + x = x.flatten(2) + return x + + +def resample_patch_embed( + patch_embed, + new_size: List[int], + interpolation: str = "bicubic", + antialias: bool = True, + verbose: bool = False, +): + """Resample the weights of the patch embedding kernel to target resolution. + We resample the patch embedding kernel by approximately inverting the effect + of patch resizing. + + Code based on: + https://github.com/google-research/big_vision/blob/b00544b81f8694488d5f36295aeb7972f3755ffe/big_vision/models/proj/flexi/vit.py + + With this resizing, we can for example load a B/8 filter into a B/16 model + and, on 2x larger input image, the result will match. + + Args: + patch_embed: original parameter to be resized. + new_size (tuple(int, int): target shape (height, width)-only. + interpolation (str): interpolation for resize + antialias (bool): use anti-aliasing filter in resize + verbose (bool): log operation + Returns: + Resized patch embedding kernel. + """ + import numpy as np + + try: + from torch import vmap + except ImportError: + from torch.func import vmap + + assert len(patch_embed.shape) == 4, "Four dimensions expected" + assert len(new_size) == 2, "New shape should only be hw" + old_size = patch_embed.shape[-2:] + if tuple(old_size) == tuple(new_size): + return patch_embed + + if verbose: + logger.info( + f"Resize patch embedding {patch_embed.shape} to {new_size}, w/ {interpolation} interpolation." + ) + + def resize(x_np, _new_size): + x_tf = torch.Tensor(x_np)[None, None, ...] + x_upsampled = F.interpolate( + x_tf, size=_new_size, mode=interpolation, antialias=antialias + )[0, 0, ...].numpy() + return x_upsampled + + def get_resize_mat(_old_size, _new_size): + mat = [] + for i in range(np.prod(_old_size)): + basis_vec = np.zeros(_old_size) + basis_vec[np.unravel_index(i, _old_size)] = 1.0 + mat.append(resize(basis_vec, _new_size).reshape(-1)) + return np.stack(mat).T + + resize_mat = get_resize_mat(old_size, new_size) + resize_mat_pinv = torch.tensor( + np.linalg.pinv(resize_mat.T), device=patch_embed.device + ) + + def resample_kernel(kernel): + resampled_kernel = resize_mat_pinv @ kernel.reshape(-1) + return resampled_kernel.reshape(new_size) + + v_resample_kernel = vmap(vmap(resample_kernel, 0, 0), 1, 1) + orig_dtype = patch_embed.dtype + patch_embed = patch_embed.float() + patch_embed = v_resample_kernel(patch_embed) + patch_embed = patch_embed.to(orig_dtype) + return patch_embed + + +# Copied from: +# https://github.com/deepseek-ai/Janus/tree/main/janus/models/siglip_vit.py +class PatchEmbed(nn.Module): + """2D Image to Patch Embedding""" + + output_fmt: Format + dynamic_img_pad: torch.jit.Final[bool] + + def __init__( + self, + img_size: Optional[int] = 224, + patch_size: int = 16, + in_chans: int = 3, + embed_dim: int = 768, + norm_layer: Optional[Callable] = None, + flatten: bool = True, + output_fmt: Optional[str] = None, + bias: bool = True, + strict_img_size: bool = True, + dynamic_img_pad: bool = False, + ): + super().__init__() + self.patch_size = tuple(to_2tuple(patch_size)) + self.img_size, self.grid_size, self.num_patches = self._init_img_size(img_size) + + if output_fmt is not None: + self.flatten = False + self.output_fmt = Format(output_fmt) + else: + # flatten spatial dim and transpose to channels last, kept for bwd compat + self.flatten = flatten + self.output_fmt = Format.NCHW + self.strict_img_size = strict_img_size + self.dynamic_img_pad = dynamic_img_pad + + self.proj = nn.Conv2d( + in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias + ) + self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() + + def _init_img_size(self, img_size: Union[int, Tuple[int, int]]): + assert self.patch_size + if img_size is None: + return None, None, None + img_size = to_2tuple(img_size) + grid_size = tuple([s // p for s, p in zip(img_size, self.patch_size)]) + num_patches = grid_size[0] * grid_size[1] + return img_size, grid_size, num_patches + + def set_input_size( + self, + img_size: Optional[Union[int, Tuple[int, int]]] = None, + patch_size: Optional[Union[int, Tuple[int, int]]] = None, + ): + new_patch_size = None + if patch_size is not None: + new_patch_size = to_2tuple(patch_size) + if new_patch_size is not None and new_patch_size != self.patch_size: + with torch.no_grad(): + new_proj = nn.Conv2d( + self.proj.in_channels, + self.proj.out_channels, + kernel_size=new_patch_size, + stride=new_patch_size, + bias=self.proj.bias is not None, + ) + new_proj.weight.copy_( + resample_patch_embed(self.proj.weight, new_patch_size, verbose=True) + ) + if self.proj.bias is not None: + new_proj.bias.copy_(self.proj.bias) + self.proj = new_proj + self.patch_size = new_patch_size + img_size = img_size or self.img_size + if img_size != self.img_size or new_patch_size is not None: + self.img_size, self.grid_size, self.num_patches = self._init_img_size( + img_size + ) + + def feat_ratio(self, as_scalar=True) -> Union[Tuple[int, int], int]: + if as_scalar: + return max(self.patch_size) + else: + return self.patch_size + + def dynamic_feat_size(self, img_size: Tuple[int, int]) -> Tuple[int, int]: + """Get grid (feature) size for given image size taking account of dynamic padding. + NOTE: must be torchscript compatible so using fixed tuple indexing + """ + if self.dynamic_img_pad: + return math.ceil(img_size[0] / self.patch_size[0]), math.ceil( + img_size[1] / self.patch_size[1] + ) + else: + return img_size[0] // self.patch_size[0], img_size[1] // self.patch_size[1] + + def forward(self, x): + B, C, H, W = x.shape + if self.img_size is not None: + if self.strict_img_size: + _assert( + H == self.img_size[0], + f"Input height ({H}) doesn't match model ({self.img_size[0]}).", + ) + _assert( + W == self.img_size[1], + f"Input width ({W}) doesn't match model ({self.img_size[1]}).", + ) + elif not self.dynamic_img_pad: + _assert( + H % self.patch_size[0] == 0, + f"Input height ({H}) should be divisible by patch size ({self.patch_size[0]}).", + ) + _assert( + W % self.patch_size[1] == 0, + f"Input width ({W}) should be divisible by patch size ({self.patch_size[1]}).", + ) + if self.dynamic_img_pad: + pad_h = (self.patch_size[0] - H % self.patch_size[0]) % self.patch_size[0] + pad_w = (self.patch_size[1] - W % self.patch_size[1]) % self.patch_size[1] + x = F.pad(x, (0, pad_w, 0, pad_h)) + x = self.proj(x) + if self.flatten: + x = x.flatten(2).transpose(1, 2) # NCHW -> NLC + elif self.output_fmt != Format.NCHW: + x = nchw_to(x, self.output_fmt) + x = self.norm(x) + return x + + +class Mlp(nn.Module): + """MLP as used in Vision Transformer, MLP-Mixer and related networks + + NOTE: When use_conv=True, expects 2D NCHW tensors, otherwise N*C expected. + """ + + def __init__( + self, + in_features, + hidden_features=None, + out_features=None, + act_layer=nn.GELU, + norm_layer=None, + bias=True, + drop=0.0, + use_conv=False, + ): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + bias = to_2tuple(bias) + drop_probs = to_2tuple(drop) + linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear + + self.fc1 = linear_layer(in_features, hidden_features, bias=bias[0]) + self.act = act_layer() + self.drop1 = nn.Dropout(drop_probs[0]) + self.norm = ( + norm_layer(hidden_features) if norm_layer is not None else nn.Identity() + ) + self.fc2 = linear_layer(hidden_features, out_features, bias=bias[1]) + self.drop2 = nn.Dropout(drop_probs[1]) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop1(x) + x = self.norm(x) + x = self.fc2(x) + x = self.drop2(x) + return x + + +def drop_path( + x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True +): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + + This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, + the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... + See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for + changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use + 'survival rate' as the argument. + + """ + if drop_prob == 0.0 or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * ( + x.ndim - 1 + ) # work with diff dim tensors, not just 2D ConvNets + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0 and scale_by_keep: + random_tensor.div_(keep_prob) + return x * random_tensor + + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + self.scale_by_keep = scale_by_keep + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training, self.scale_by_keep) + + def extra_repr(self): + return f"drop_prob={round(self.drop_prob, 3):0.3f}" + + +class VisionTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = False, + qk_norm: bool = False, + proj_drop: float = 0.0, + attn_drop: float = 0.0, + init_values: Optional[float] = None, + drop_path: float = 0.0, + act_layer: nn.Module = nn.GELU, + norm_layer: nn.Module = nn.LayerNorm, + mlp_layer: nn.Module = Mlp, + ) -> None: + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = VisionAttention( + embed_dim=dim, + num_heads=num_heads, + projection_size=dim, + use_qkv_parallel=True, + use_context_forward=False, + softmax_in_single_precision=False, + dropout=attn_drop, + ) + + self.ls1 = ( + LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + ) + self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + self.norm2 = norm_layer(dim) + self.mlp = mlp_layer( + in_features=dim, + hidden_features=int(dim * mlp_ratio), + act_layer=act_layer, + drop=proj_drop, + ) + self.ls2 = ( + LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + ) + self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x)))) + x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x)))) + return x + + +LayerType = Union[str, Callable, Type[torch.nn.Module]] + + +class PatchDropout(nn.Module): + """ + https://arxiv.org/abs/2212.00794 and https://arxiv.org/pdf/2208.07220 + """ + + return_indices: torch.jit.Final[bool] + + def __init__( + self, + prob: float = 0.5, + num_prefix_tokens: int = 1, + ordered: bool = False, + return_indices: bool = False, + ): + super().__init__() + assert 0 <= prob < 1.0 + self.prob = prob + self.num_prefix_tokens = ( + num_prefix_tokens # exclude CLS token (or other prefix tokens) + ) + self.ordered = ordered + self.return_indices = return_indices + + def forward( + self, x + ) -> Union[torch.Tensor, Tuple[torch.Tensor, Optional[torch.Tensor]]]: + if not self.training or self.prob == 0.0: + if self.return_indices: + return x, None + return x + + if self.num_prefix_tokens: + prefix_tokens, x = ( + x[:, : self.num_prefix_tokens], + x[:, self.num_prefix_tokens :], + ) + else: + prefix_tokens = None + + B = x.shape[0] + L = x.shape[1] + num_keep = max(1, int(L * (1.0 - self.prob))) + keep_indices = torch.argsort(torch.randn(B, L, device=x.device), dim=-1)[ + :, :num_keep + ] + if self.ordered: + # NOTE does not need to maintain patch order in typical transformer use, + # but possibly useful for debug / visualization + keep_indices = keep_indices.sort(dim=-1)[0] + x = x.gather(1, keep_indices.unsqueeze(-1).expand((-1, -1) + x.shape[2:])) + + if prefix_tokens is not None: + x = torch.cat((prefix_tokens, x), dim=1) + + if self.return_indices: + return x, keep_indices + return x + + +def resample_abs_pos_embed( + posemb: torch.Tensor, + new_size: List[int], + old_size: Optional[List[int]] = None, + num_prefix_tokens: int = 1, + interpolation: str = "bicubic", + antialias: bool = True, + verbose: bool = False, +): + # sort out sizes, assume square if old size not provided + num_pos_tokens = posemb.shape[1] + num_new_tokens = new_size[0] * new_size[1] + num_prefix_tokens + if num_new_tokens == num_pos_tokens and new_size[0] == new_size[1]: + return posemb + + if old_size is None: + hw = int(math.sqrt(num_pos_tokens - num_prefix_tokens)) + old_size = hw, hw + + if num_prefix_tokens: + posemb_prefix, posemb = ( + posemb[:, :num_prefix_tokens], + posemb[:, num_prefix_tokens:], + ) + else: + posemb_prefix, posemb = None, posemb + + # do the interpolation + embed_dim = posemb.shape[-1] + orig_dtype = posemb.dtype + posemb = posemb.float() # interpolate needs float32 + posemb = posemb.reshape(1, old_size[0], old_size[1], -1).permute(0, 3, 1, 2) + posemb = F.interpolate( + posemb, size=new_size, mode=interpolation, antialias=antialias + ) + posemb = posemb.permute(0, 2, 3, 1).reshape(1, -1, embed_dim) + posemb = posemb.to(orig_dtype) + + # add back extra (class, etc) prefix tokens + if posemb_prefix is not None: + posemb = torch.cat([posemb_prefix, posemb], dim=1) + + if not torch.jit.is_scripting() and verbose: + logger.info(f"Resized position embedding: {old_size} to {new_size}.") + + return posemb + + +def init_weights(self): + if self.pos_embed is not None: + trunc_normal_(self.pos_embed, std=self.pos_embed.shape[1] ** -0.5) + trunc_normal_(self.latent, std=self.latent_dim**-0.5) + + +def init_weights_vit_timm(module: nn.Module, name: str = "") -> None: + """ViT weight initialization, original timm impl (for reproducibility)""" + if isinstance(module, nn.Linear): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif hasattr(module, "init_weights"): + module.init_weights() + + +class VisionTransformer(nn.Module): + """Vision Transformer + + A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale` + - https://arxiv.org/abs/2010.11929 + """ + + dynamic_img_size: Final[bool] + + def __init__( + self, + img_size: Union[int, Tuple[int, int]] = 224, + patch_size: Union[int, Tuple[int, int]] = 16, + in_chans: int = 3, + num_classes: int = 1000, + global_pool: Literal["", "avg", "token", "map"] = "token", + embed_dim: int = 768, + depth: int = 12, + num_heads: int = 12, + mlp_ratio: float = 4.0, + qkv_bias: bool = True, + qk_norm: bool = False, + init_values: Optional[float] = None, + class_token: bool = True, + no_embed_class: bool = False, + reg_tokens: int = 0, + pre_norm: bool = False, + fc_norm: Optional[bool] = None, + dynamic_img_size: bool = False, + dynamic_img_pad: bool = False, + drop_rate: float = 0.0, + pos_drop_rate: float = 0.0, + patch_drop_rate: float = 0.0, + proj_drop_rate: float = 0.0, + attn_drop_rate: float = 0.0, + drop_path_rate: float = 0.0, + weight_init: Literal["skip", "jax", "jax_nlhb", "moco", ""] = "", + embed_layer: Callable = PatchEmbed, + _norm_layer: Optional[LayerType] = None, + _act_layer: Optional[LayerType] = None, + block_fn: Type[nn.Module] = VisionTransformerBlock, + mlp_layer: Type[nn.Module] = Mlp, + ignore_head: bool = False, + ) -> None: + """ + Args: + img_size: Input image size. + patch_size: Patch size. + in_chans: Number of image input channels. + num_classes: Mumber of classes for classification head. + global_pool: Type of global pooling for final sequence (default: 'token'). + embed_dim: Transformer embedding dimension. + depth: Depth of transformer. + num_heads: Number of attention heads. + mlp_ratio: Ratio of mlp hidden dim to embedding dim. + qkv_bias: Enable bias for qkv projections if True. + init_values: Layer-scale init values (layer-scale enabled if not None). + class_token: Use class token. + no_embed_class: Don't include position embeddings for class (or reg) tokens. + reg_tokens: Number of register tokens. + fc_norm: Pre head norm after pool (instead of before), if None, enabled when global_pool == 'avg'. + drop_rate: Head dropout rate. + pos_drop_rate: Position embedding dropout rate. + attn_drop_rate: Attention dropout rate. + drop_path_rate: Stochastic depth rate. + weight_init: Weight initialization scheme. + embed_layer: Patch embedding layer. + _norm_layer: Normalization layer. + _act_layer: MLP activation layer. + block_fn: Transformer block layer. + """ + super().__init__() + assert global_pool in ("", "avg", "token", "map") + assert class_token or global_pool != "token" + use_fc_norm = global_pool == "avg" if fc_norm is None else fc_norm + # norm_layer = get_norm_layer(norm_layer) or partial(nn.LayerNorm, eps=1e-6) + # act_layer = get_act_layer(act_layer) or nn.GELU + norm_layer = partial(nn.LayerNorm, eps=1e-6) + act_layer = nn.GELU + + self.num_classes = num_classes + self.global_pool = global_pool + self.num_features = self.embed_dim = ( + embed_dim # num_features for consistency with other models + ) + self.num_prefix_tokens = 1 if class_token else 0 + self.num_prefix_tokens += reg_tokens + self.num_reg_tokens = reg_tokens + self.has_class_token = class_token + self.no_embed_class = ( + no_embed_class # don't embed prefix positions (includes reg) + ) + self.dynamic_img_size = dynamic_img_size + self.grad_checkpointing = False + self.ignore_head = ignore_head + + embed_args = {} + if dynamic_img_size: + # flatten deferred until after pos embed + embed_args.update(dict(strict_img_size=False, output_fmt="NHWC")) + self.patch_embed = embed_layer( + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + bias=not pre_norm, # disable bias if pre-norm is used (e.g. CLIP) + dynamic_img_pad=dynamic_img_pad, + **embed_args, + ) + num_patches = self.patch_embed.num_patches + + self.cls_token = ( + nn.Parameter(torch.zeros(1, 1, embed_dim)) if class_token else None + ) + self.reg_token = ( + nn.Parameter(torch.zeros(1, reg_tokens, embed_dim)) if reg_tokens else None + ) + embed_len = ( + num_patches if no_embed_class else num_patches + self.num_prefix_tokens + ) + self.pos_embed = nn.Parameter(torch.randn(1, embed_len, embed_dim) * 0.02) + self.pos_drop = nn.Dropout(p=pos_drop_rate) + if patch_drop_rate > 0: + self.patch_drop = PatchDropout( + patch_drop_rate, + num_prefix_tokens=self.num_prefix_tokens, + ) + else: + self.patch_drop = nn.Identity() + self.norm_pre = norm_layer(embed_dim) if pre_norm else nn.Identity() + + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, depth) + ] # stochastic depth decay rule + self.blocks = nn.Sequential( + *[ + block_fn( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_norm=qk_norm, + init_values=init_values, + proj_drop=proj_drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[i], + norm_layer=norm_layer, + act_layer=act_layer, + mlp_layer=mlp_layer, + ) + for i in range(depth) + ] + ) + self.norm = norm_layer(embed_dim) if not use_fc_norm else nn.Identity() + + # Classifier Head + if global_pool == "map": + AttentionPoolLatent.init_weights = init_weights + self.attn_pool = AttentionPoolLatent( + self.embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + norm_layer=norm_layer, + ) + else: + self.attn_pool = None + self.fc_norm = norm_layer(embed_dim) if use_fc_norm else nn.Identity() + self.head_drop = nn.Dropout(drop_rate) + self.head = ( + nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() + ) + + if weight_init != "skip": + self.init_weights(weight_init) + + def init_weights(self, mode: Literal["jax", "jax_nlhb", "moco", ""] = "") -> None: + assert mode in ("jax", "jax_nlhb", "moco", "") + # head_bias = -math.log(self.num_classes) if "nlhb" in mode else 0.0 + trunc_normal_(self.pos_embed, std=0.02) + if self.cls_token is not None: + nn.init.normal_(self.cls_token, std=1e-6) + named_apply(init_weights_vit_timm, self) + + @torch.jit.ignore + def no_weight_decay(self) -> Set: + return {"pos_embed", "cls_token", "dist_token"} + + @torch.jit.ignore + def group_matcher(self, coarse: bool = False) -> Dict: + return dict( + stem=r"^cls_token|pos_embed|patch_embed", # stem and embed + blocks=[(r"^blocks\.(\d+)", None), (r"^norm", (99999,))], + ) + + @torch.jit.ignore + def get_classifier(self) -> nn.Module: + return self.head + + def reset_classifier(self, num_classes: int, global_pool=None) -> None: + self.num_classes = num_classes + if global_pool is not None: + assert global_pool in ("", "avg", "token", "map") + if global_pool == "map" and self.attn_pool is None: + assert ( + False + ), "Cannot currently add attention pooling in reset_classifier()." + elif global_pool != "map " and self.attn_pool is not None: + self.attn_pool = None # remove attention pooling + self.global_pool = global_pool + self.head = ( + nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity() + ) + + def _pos_embed(self, x: torch.Tensor) -> torch.Tensor: + if self.dynamic_img_size: + B, H, W, C = x.shape + pos_embed = resample_abs_pos_embed( + self.pos_embed, + [H, W], + num_prefix_tokens=0 if self.no_embed_class else self.num_prefix_tokens, + ) + x = x.view(B, -1, C) + else: + pos_embed = self.pos_embed + + to_cat = [] + if self.cls_token is not None: + to_cat.append(self.cls_token.expand(x.shape[0], -1, -1)) + if self.reg_token is not None: + to_cat.append(self.reg_token.expand(x.shape[0], -1, -1)) + + if self.no_embed_class: + # deit-3, updated JAX (big vision) + # position embedding does not overlap with class token, add then concat + x = x + pos_embed + if to_cat: + x = torch.cat(to_cat + [x], dim=1) + else: + # original timm, JAX, and deit vit impl + # pos_embed has entry for class token, concat then add + if to_cat: + x = torch.cat(to_cat + [x], dim=1) + x = x + pos_embed + + return self.pos_drop(x) + + def _intermediate_layers( + self, + x: torch.Tensor, + n: Union[int, Sequence] = 1, + ) -> List[torch.Tensor]: + outputs, num_blocks = [], len(self.blocks) + take_indices = set( + range(num_blocks - n, num_blocks) if isinstance(n, int) else n + ) + + # forward pass + x = self.patch_embed(x) + x = self._pos_embed(x) + x = self.patch_drop(x) + x = self.norm_pre(x) + for i, blk in enumerate(self.blocks): + x = blk(x) + if i in take_indices: + outputs.append(x) + + return outputs + + def forward_features(self, x: torch.Tensor) -> torch.Tensor: + x = self.patch_embed(x) + x = self._pos_embed(x) + x = self.patch_drop(x) + x = self.norm_pre(x) + x = self.blocks(x) + x = self.norm(x) + return x + + def forward_head(self, x: torch.Tensor, pre_logits: bool = False) -> torch.Tensor: + if self.attn_pool is not None: + x = self.attn_pool(x) + elif self.global_pool == "avg": + x = x[:, self.num_prefix_tokens :].mean(dim=1) + elif self.global_pool: + x = x[:, 0] # class token + x = self.fc_norm(x) + x = self.head_drop(x) + return x if pre_logits else self.head(x) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.forward_features(x) + if not self.ignore_head: + x = self.forward_head(x) + return x + + +def model_name_to_cls(cls_name): + if "MlpProjector" in cls_name: + cls = MlpProjector + + elif "CLIPVisionTower" in cls_name: + cls = CLIPVisionTower + + elif "VQ" in cls_name: + + cls = VQ_models[cls_name] + elif "vision_head" in cls_name: + cls = vision_head + else: + raise ValueError(f"class_name {cls_name} is invalid.") + + return cls + + +class vision_head(torch.nn.Module): + def __init__(self, params): + super().__init__() + self.output_mlp_projector = torch.nn.Linear( + params["n_embed"], params["image_token_embed"] + ) + self.vision_activation = torch.nn.GELU() + self.vision_head = torch.nn.Linear( + params["image_token_embed"], params["image_token_size"] + ) + + def forward(self, x): + x = self.output_mlp_projector(x) + x = self.vision_activation(x) + x = self.vision_head(x) + return x + + +SigLIP_MODEL_CONFIG = { + "siglip_so400m_patch14_384": { + "image_size": 336, + "patch_size": 14, + "width": 1152, + "layers": 27, + "heads": 16, + "mlp_ratio": 3.7362, + "global_pool": "map", + "use_checkpoint": False, + }, + "siglip_so400m_patch14_224": { + "image_size": 224, + "patch_size": 14, + "width": 1152, + "layers": 27, + "heads": 16, + "mlp_ratio": 3.7362, + "global_pool": "map", + "use_checkpoint": False, + }, + "siglip_large_patch16_384": { + "image_size": 384, + "patch_size": 16, + "width": 1024, + "layers": 24, + "heads": 16, + "mlp_ratio": 4, + "global_pool": "map", + "use_checkpoint": False, + }, +} + + +def create_siglip_vit( + model_name: str = "siglip_so400m_patch14_384", + image_size: int = 384, + select_layer: int = -1, + ckpt_path: str = "", + **kwargs, +): + assert ( + model_name in SigLIP_MODEL_CONFIG.keys() + ), f"model name should be in {SigLIP_MODEL_CONFIG.keys()}" + + vision_cfg = SigLIPVisionCfg(**SigLIP_MODEL_CONFIG[model_name]) + + if select_layer <= 0: + layers = min(vision_cfg.layers, vision_cfg.layers + select_layer + 1) + else: + layers = min(vision_cfg.layers, select_layer) + + model = VisionTransformer( + img_size=image_size, + patch_size=vision_cfg.patch_size, + embed_dim=vision_cfg.width, + depth=layers, + num_heads=vision_cfg.heads, + mlp_ratio=vision_cfg.mlp_ratio, + class_token=vision_cfg.class_token, + global_pool=vision_cfg.global_pool, + ignore_head=kwargs.get("ignore_head", True), + weight_init=kwargs.get("weight_init", "skip"), + num_classes=0, + ) + + if ckpt_path: + state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=True) + + incompatible_keys = model.load_state_dict(state_dict, strict=False) + print( + f"SigLIP-ViT restores from {ckpt_path},\n" + f"\tincompatible_keys:', {incompatible_keys}." + ) + + return model + + +class Normalize(torch.nn.Module): + """Normalize a tensor image with mean and standard deviation. + This transform does not support PIL Image. + Given mean: ``(mean[1],...,mean[n])`` and std: ``(std[1],..,std[n])`` for ``n`` + channels, this transform will normalize each channel of the input + ``torch.*Tensor`` i.e., + ``output[channel] = (input[channel] - mean[channel]) / std[channel]`` + + .. note:: + This transform acts out of place, i.e., it does not mutate the input tensor. + + Args: + mean (sequence): Sequence of means for each channel. + std (sequence): Sequence of standard deviations for each channel. + inplace(bool,optional): Bool to make this operation in-place. + + """ + + def __init__(self, mean, std, inplace=False): + super().__init__() + # _log_api_usage_once(self) + self.mean = mean + self.std = std + self.inplace = inplace + + def forward(self, tensor: Tensor) -> Tensor: + """ + Args: + tensor (Tensor): Tensor image to be normalized. + + Returns: + Tensor: Normalized Tensor image. + """ + return F.normalize(tensor, self.mean, self.std, self.inplace) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(mean={self.mean}, std={self.std})" + + +class CLIPVisionTower(nn.Module): + def __init__( + self, + model_name: str = "siglip_large_patch16_384", + image_size: Union[Tuple[int, int], int] = 336, + select_feature: str = "patch", + select_layer: int = -2, + select_layers: list = None, + ckpt_path: str = "", + pixel_mean: Optional[List[float]] = None, + pixel_std: Optional[List[float]] = None, + **kwargs, + ): + super().__init__() + + self.model_name = model_name + self.select_feature = select_feature + self.select_layer = select_layer + self.select_layers = select_layers + + vision_tower_params = { + "model_name": model_name, + "image_size": image_size, + "ckpt_path": ckpt_path, + "select_layer": select_layer, + } + vision_tower_params.update(kwargs) + self.vision_tower, self.forward_kwargs = self.build_vision_tower( + vision_tower_params + ) + + if pixel_mean is not None and pixel_std is not None: + image_norm = Normalize(mean=pixel_mean, std=pixel_std) + else: + image_norm = None + + self.image_norm = image_norm + + @property + def device(self) -> torch.device: + return next(self.vision_tower.parameters()).device + + @property + def dtype(self): + return next(self.vision_tower.parameters()).dtype + + def build_vision_tower(self, vision_tower_params): + if self.model_name.startswith("siglip"): + self.select_feature = "same" + vision_tower = create_siglip_vit(**vision_tower_params) + forward_kwargs = dict() + + elif self.model_name.startswith("sam"): + # vision_tower = create_sam_vit(**vision_tower_params) + forward_kwargs = dict() + + else: # huggingface + from transformers import CLIPVisionModel + + vision_tower = CLIPVisionModel.from_pretrained(**vision_tower_params) + forward_kwargs = dict(output_hidden_states=True) + + return vision_tower, forward_kwargs + + def feature_select(self, image_forward_outs): + if isinstance(image_forward_outs, torch.Tensor): + # the output has been the self.select_layer"s features + image_features = image_forward_outs + else: + image_features = image_forward_outs.hidden_states[self.select_layer] + + if self.select_feature == "patch": + # if the output has cls_token + image_features = image_features[:, 1:] + elif self.select_feature == "cls_patch": + image_features = image_features + elif self.select_feature == "same": + image_features = image_features + + else: + raise ValueError(f"Unexpected select feature: {self.select_feature}") + return image_features + + def forward(self, images): + """ + + Args: + images (torch.Tensor): [b, 3, H, W] + + Returns: + image_features (torch.Tensor): [b, n_patch, d] + """ + + if self.image_norm is not None: + images = self.image_norm(images) + + image_forward_outs = self.vision_tower(images, **self.forward_kwargs) + image_features = self.feature_select(image_forward_outs) + return image_features + + +class MlpProjector(nn.Module): + def __init__(self, cfg): + super().__init__() + + self.cfg = cfg + + if cfg["projector_type"] == "identity": + modules = nn.Identity() + + elif cfg["projector_type"] == "linear": + modules = nn.Linear(cfg["input_dim"], cfg["n_embed"]) + + elif cfg["projector_type"] == "mlp_gelu": + mlp_depth = cfg.get("depth", 1) + modules = [nn.Linear(cfg["input_dim"], cfg["n_embed"])] + for _ in range(1, mlp_depth): + modules.append(nn.GELU()) + modules.append(nn.Linear(cfg["n_embed"], cfg["n_embed"])) + modules = nn.Sequential(*modules) + + elif cfg["projector_type"] == "low_high_hybrid_split_mlp_gelu": + mlp_depth = cfg.get("depth", 1) + self.high_up_proj = nn.Linear(cfg["input_dim"], cfg["n_embed"] // 2) + self.low_up_proj = nn.Linear(cfg["input_dim"], cfg["n_embed"] // 2) + + modules = [] + for _ in range(1, mlp_depth): + modules.append(nn.GELU()) + modules.append(nn.Linear(cfg["n_embed"], cfg["n_embed"])) + modules = nn.Sequential(*modules) + + else: + raise ValueError(f"Unknown projector type: {cfg['projector_type']}") + + self.layers = modules + + def forward( + self, x_or_tuple: Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor] + ): + """ + + Args: + x_or_tuple (Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]: if it is a tuple of torch.Tensor, + then it comes from the hybrid vision encoder, and x = high_res_x, low_res_x); + otherwise it is the feature from the single vision encoder. + + Returns: + x (torch.Tensor): [b, s, c] + """ + + if isinstance(x_or_tuple, tuple): + # self.cfg.projector_type == "low_high_hybrid_split_mlp_gelu": + high_x, low_x = x_or_tuple + high_x = self.high_up_proj(high_x) + low_x = self.low_up_proj(low_x) + x = torch.cat([high_x, low_x], dim=-1) + else: + x = x_or_tuple + + return self.layers(x) + + +class LayerScale(nn.Module): + def __init__( + self, + dim: int, + init_values: float = 1e-5, + inplace: bool = False, + ) -> None: + super().__init__() + self.inplace = inplace + self.gamma = nn.Parameter(init_values * torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.mul_(self.gamma) if self.inplace else x * self.gamma + + +# use torch.scaled_dot_product_attention where possible +_HAS_FUSED_ATTN = hasattr(torch.nn.functional, "scaled_dot_product_attention") +if "TIMM_FUSED_ATTN" in os.environ: + _USE_FUSED_ATTN = int(os.environ["TIMM_FUSED_ATTN"]) +else: + _USE_FUSED_ATTN = ( + 1 # 0 == off, 1 == on (for tested use), 2 == on (for experimental use) + ) + +# Set to True if exporting a model with Same padding via ONNX +_EXPORTABLE = False + + +def use_fused_attn(experimental: bool = False) -> bool: + # NOTE: ONNX export cannot handle F.scaled_dot_product_attention as of pytorch 2.0 + if not _HAS_FUSED_ATTN or _EXPORTABLE: + return False + if experimental: + return _USE_FUSED_ATTN > 1 + return _USE_FUSED_ATTN > 0 + + +class AttentionPoolLatent(nn.Module): + """Attention pooling w/ latent query""" + + fused_attn: torch.jit.Final[bool] + + def __init__( + self, + in_features: int, + out_features: int = None, + embed_dim: int = None, + num_heads: int = 8, + feat_size: Optional[int] = None, + mlp_ratio: float = 4.0, + qkv_bias: bool = True, + qk_norm: bool = False, + latent_len: int = 1, + latent_dim: int = None, + pos_embed: str = "", + pool_type: str = "token", + norm_layer: Optional[nn.Module] = None, + drop: float = 0.0, + ): + super().__init__() + embed_dim = embed_dim or in_features + out_features = out_features or in_features + assert embed_dim % num_heads == 0 + self.num_heads = num_heads + self.head_dim = embed_dim // num_heads + self.feat_size = feat_size + self.scale = self.head_dim**-0.5 + self.pool = pool_type + self.fused_attn = use_fused_attn() + + if pos_embed == "abs": + assert feat_size is not None + self.pos_embed = nn.Parameter(torch.zeros(feat_size, in_features)) + else: + self.pos_embed = None + + self.latent_dim = latent_dim or embed_dim + self.latent_len = latent_len + self.latent = nn.Parameter(torch.zeros(1, self.latent_len, embed_dim)) + + self.q = nn.Linear(embed_dim, embed_dim, bias=qkv_bias) + self.kv = nn.Linear(embed_dim, embed_dim * 2, bias=qkv_bias) + self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() + self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() + self.proj = nn.Linear(embed_dim, embed_dim) + self.proj_drop = nn.Dropout(drop) + + self.norm = ( + norm_layer(out_features) if norm_layer is not None else nn.Identity() + ) + self.mlp = Mlp(embed_dim, int(embed_dim * mlp_ratio)) + + self.init_weights() + + def init_weights(self): + if self.pos_embed is not None: + trunc_normal_tf_(self.pos_embed, std=self.pos_embed.shape[1] ** -0.5) + trunc_normal_tf_(self.latent, std=self.latent_dim**-0.5) + + def forward(self, x): + B, N, C = x.shape + + if self.pos_embed is not None: + # FIXME interpolate + x = x + self.pos_embed.unsqueeze(0).to(x.dtype) + + q_latent = self.latent.expand(B, -1, -1) + q = ( + self.q(q_latent) + .reshape(B, self.latent_len, self.num_heads, self.head_dim) + .transpose(1, 2) + ) + + kv = ( + self.kv(x) + .reshape(B, N, 2, self.num_heads, self.head_dim) + .permute(2, 0, 3, 1, 4) + ) + k, v = kv.unbind(0) + + q, k = self.q_norm(q), self.k_norm(k) + + if self.fused_attn: + x = F.scaled_dot_product_attention(q, k, v) + else: + q = q * self.scale + attn = q @ k.transpose(-2, -1) + attn = attn.softmax(dim=-1) + x = attn @ v + x = x.transpose(1, 2).reshape(B, self.latent_len, C) + x = self.proj(x) + x = self.proj_drop(x) + + x = x + self.mlp(self.norm(x)) + + # optional pool if latent seq_len > 1 and pooled output is desired + if self.pool == "token": + x = x[:, 0] + elif self.pool == "avg": + x = x.mean(1) + + +class Encoder(nn.Module): + def __init__( + self, + in_channels=3, + ch=128, + ch_mult=(1, 1, 2, 2, 4), + num_res_blocks=2, + norm_type="group", + dropout=0.0, + resamp_with_conv=True, + z_channels=256, + ): + super().__init__() + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.conv_in = nn.Conv2d(in_channels, ch, kernel_size=3, stride=1, padding=1) + + # downsampling + in_ch_mult = (1,) + tuple(ch_mult) + self.conv_blocks = nn.ModuleList() + for i_level in range(self.num_resolutions): + conv_block = nn.Module() + # res & attn + res_block = nn.ModuleList() + attn_block = nn.ModuleList() + block_in = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + for _ in range(self.num_res_blocks): + res_block.append( + ResnetBlock( + block_in, block_out, dropout=dropout, norm_type=norm_type + ) + ) + block_in = block_out + if i_level == self.num_resolutions - 1: + attn_block.append(AttnBlock(block_in, norm_type)) + conv_block.res = res_block + conv_block.attn = attn_block + # downsample + if i_level != self.num_resolutions - 1: + conv_block.downsample = Downsample(block_in, resamp_with_conv) + self.conv_blocks.append(conv_block) + + # middle + self.mid = nn.ModuleList() + self.mid.append( + ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type) + ) + self.mid.append(AttnBlock(block_in, norm_type=norm_type)) + self.mid.append( + ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type) + ) + + # end + self.norm_out = Normalize(block_in, norm_type) + self.conv_out = nn.Conv2d( + block_in, z_channels, kernel_size=3, stride=1, padding=1 + ) + + def forward(self, x): + h = self.conv_in(x) + # downsampling + for i_level, block in enumerate(self.conv_blocks): + for i_block in range(self.num_res_blocks): + h = block.res[i_block](h) + if len(block.attn) > 0: + h = block.attn[i_block](h) + if i_level != self.num_resolutions - 1: + h = block.downsample(h) + + # middle + for mid_block in self.mid: + h = mid_block(h) + + # end + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h + + +class Decoder(nn.Module): + def __init__( + self, + z_channels=256, + ch=128, + ch_mult=(1, 1, 2, 2, 4), + num_res_blocks=2, + norm_type="group", + dropout=0.0, + resamp_with_conv=True, + out_channels=3, + ): + super().__init__() + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + + block_in = ch * ch_mult[self.num_resolutions - 1] + # z to block_in + self.conv_in = nn.Conv2d( + z_channels, block_in, kernel_size=3, stride=1, padding=1 + ) + + # middle + self.mid = nn.ModuleList() + self.mid.append( + ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type) + ) + self.mid.append(AttnBlock(block_in, norm_type=norm_type)) + self.mid.append( + ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type) + ) + + # upsampling + self.conv_blocks = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + conv_block = nn.Module() + # res & attn + res_block = nn.ModuleList() + attn_block = nn.ModuleList() + block_out = ch * ch_mult[i_level] + for _ in range(self.num_res_blocks + 1): + res_block.append( + ResnetBlock( + block_in, block_out, dropout=dropout, norm_type=norm_type + ) + ) + block_in = block_out + if i_level == self.num_resolutions - 1: + attn_block.append(AttnBlock(block_in, norm_type)) + conv_block.res = res_block + conv_block.attn = attn_block + # downsample + if i_level != 0: + conv_block.upsample = Upsample(block_in, resamp_with_conv) + self.conv_blocks.append(conv_block) + + # end + self.norm_out = Normalize(block_in, norm_type) + self.conv_out = nn.Conv2d( + block_in, out_channels, kernel_size=3, stride=1, padding=1 + ) + + @property + def last_layer(self): + return self.conv_out.weight + + def forward(self, z): + # z to block_in + h = self.conv_in(z) + + # middle + for mid_block in self.mid: + h = mid_block(h) + + # upsampling + for i_level, block in enumerate(self.conv_blocks): + for i_block in range(self.num_res_blocks + 1): + h = block.res[i_block](h) + if len(block.attn) > 0: + h = block.attn[i_block](h) + if i_level != self.num_resolutions - 1: + h = block.upsample(h) + + # end + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h + + +class VectorQuantizer(nn.Module): + def __init__(self, n_e, e_dim, beta, entropy_loss_ratio, l2_norm, show_usage): + super().__init__() + self.n_e = n_e + self.e_dim = e_dim + self.beta = beta + self.entropy_loss_ratio = entropy_loss_ratio + self.l2_norm = l2_norm + self.show_usage = show_usage + + self.embedding = nn.Embedding(self.n_e, self.e_dim) + self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) + if self.l2_norm: + self.embedding.weight.data = F.normalize( + self.embedding.weight.data, p=2, dim=-1 + ) + if self.show_usage: + # self.register_buffer("codebook_used", nn.Parameter(torch.zeros(65536))) + self.codebook_used = nn.Parameter(torch.zeros(65536)) + + def forward(self, z): + # reshape z -> (batch, height, width, channel) and flatten + z = torch.einsum("b c h w -> b h w c", z).contiguous() + z_flattened = z.view(-1, self.e_dim) + # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z + + if self.l2_norm: + z = F.normalize(z, p=2, dim=-1) + z_flattened = F.normalize(z_flattened, p=2, dim=-1) + embedding = F.normalize(self.embedding.weight, p=2, dim=-1) + else: + embedding = self.embedding.weight + + d = ( + torch.sum(z_flattened**2, dim=1, keepdim=True) + + torch.sum(embedding**2, dim=1) + - 2 + * torch.einsum( + "bd,dn->bn", z_flattened, torch.einsum("n d -> d n", embedding) + ) + ) + + min_encoding_indices = torch.argmin(d, dim=1) + z_q = embedding[min_encoding_indices].view(z.shape) + perplexity = None + min_encodings = None + vq_loss = None + commit_loss = None + entropy_loss = None + + # compute loss for embedding + if self.training: + vq_loss = torch.mean((z_q - z.detach()) ** 2) + commit_loss = self.beta * torch.mean((z_q.detach() - z) ** 2) + entropy_loss = self.entropy_loss_ratio * compute_entropy_loss(-d) + + # preserve gradients + z_q = z + (z_q - z).detach() + + # reshape back to match original input shape + z_q = torch.einsum("b h w c -> b c h w", z_q) + + return ( + z_q, + (vq_loss, commit_loss, entropy_loss), + (perplexity, min_encodings, min_encoding_indices), + ) + + def get_codebook_entry(self, indices, shape=None, channel_first=True): + # shape = (batch, channel, height, width) if channel_first else (batch, height, width, channel) + if self.l2_norm: + embedding = F.normalize(self.embedding.weight, p=2, dim=-1) + else: + embedding = self.embedding.weight + z_q = embedding[indices] # (b*h*w, c) + + if shape is not None: + if channel_first: + z_q = z_q.reshape(shape[0], shape[2], shape[3], shape[1]) + # reshape back to match original input shape + z_q = z_q.permute(0, 3, 1, 2).contiguous() + else: + z_q = z_q.view(shape) + return z_q + + +class ResnetBlock(nn.Module): + def __init__( + self, + in_channels, + out_channels=None, + conv_shortcut=False, + dropout=0.0, + norm_type="group", + ): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + + self.norm1 = Normalize(in_channels, norm_type) + self.conv1 = nn.Conv2d( + in_channels, out_channels, kernel_size=3, stride=1, padding=1 + ) + self.norm2 = Normalize(out_channels, norm_type) + self.dropout = nn.Dropout(dropout) + self.conv2 = nn.Conv2d( + out_channels, out_channels, kernel_size=3, stride=1, padding=1 + ) + + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = nn.Conv2d( + in_channels, out_channels, kernel_size=3, stride=1, padding=1 + ) + else: + self.nin_shortcut = nn.Conv2d( + in_channels, out_channels, kernel_size=1, stride=1, padding=0 + ) + + def forward(self, x): + h = x + h = self.norm1(h) + h = nonlinearity(h) + h = self.conv1(h) + h = self.norm2(h) + h = nonlinearity(h) + h = self.dropout(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + x = self.conv_shortcut(x) + else: + x = self.nin_shortcut(x) + return x + h + + +class AttnBlock(nn.Module): + def __init__(self, in_channels, norm_type="group"): + super().__init__() + self.norm = Normalize(in_channels, norm_type) + self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.proj_out = nn.Conv2d( + in_channels, in_channels, kernel_size=1, stride=1, padding=0 + ) + + def forward(self, x): + h_ = x + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + # compute attention + b, c, h, w = q.shape + q = q.reshape(b, c, h * w) + q = q.permute(0, 2, 1) # b,hw,c + k = k.reshape(b, c, h * w) # b,c,hw + w_ = torch.bmm(q, k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j] + w_ = w_ * (int(c) ** (-0.5)) + w_ = F.softmax(w_, dim=2) + + # attend to values + v = v.reshape(b, c, h * w) + w_ = w_.permute(0, 2, 1) # b,hw,hw (first hw of k, second of q) + h_ = torch.bmm(v, w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j] + h_ = h_.reshape(b, c, h, w) + + h_ = self.proj_out(h_) + + return x + h_ + + +def nonlinearity(x): + # swish + return x * torch.sigmoid(x) + + +def Normalize(in_channels, norm_type="group"): + assert norm_type in ["group", "batch"] + if norm_type == "group": + return nn.GroupNorm( + num_groups=32, num_channels=in_channels, eps=1e-6, affine=True + ) + elif norm_type == "batch": + return nn.SyncBatchNorm(in_channels) + + +class Upsample(nn.Module): + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + self.conv = nn.Conv2d( + in_channels, in_channels, kernel_size=3, stride=1, padding=1 + ) + + def forward(self, x): + if x.dtype != torch.float32: + x = F.interpolate(x.to(torch.float), scale_factor=2.0, mode="nearest").to( + torch.bfloat16 + ) + else: + x = F.interpolate(x, scale_factor=2.0, mode="nearest") + + if self.with_conv: + x = self.conv(x) + return x + + +class Downsample(nn.Module): + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + # no asymmetric padding in torch conv, must do it ourselves + self.conv = nn.Conv2d( + in_channels, in_channels, kernel_size=3, stride=2, padding=0 + ) + + def forward(self, x): + if self.with_conv: + pad = (0, 1, 0, 1) + x = F.pad(x, pad, mode="constant", value=0) + x = self.conv(x) + else: + x = F.avg_pool2d(x, kernel_size=2, stride=2) + return x + + +def compute_entropy_loss(affinity, loss_type="softmax", temperature=0.01): + flat_affinity = affinity.reshape(-1, affinity.shape[-1]) + flat_affinity /= temperature + probs = F.softmax(flat_affinity, dim=-1) + log_probs = F.log_softmax(flat_affinity + 1e-5, dim=-1) + if loss_type == "softmax": + target_probs = probs + else: + raise ValueError("Entropy loss {} not supported".format(loss_type)) + avg_probs = torch.mean(target_probs, dim=0) + avg_entropy = -torch.sum(avg_probs * torch.log(avg_probs + 1e-5)) + sample_entropy = -torch.mean(torch.sum(target_probs * log_probs, dim=-1)) + loss = sample_entropy - avg_entropy + return loss + + +class VQModel(nn.Module): + def __init__(self, config: ModelArgs): + super().__init__() + self.config = config + self.encoder = Encoder( + ch_mult=config.encoder_ch_mult, + z_channels=config.z_channels, + dropout=config.dropout_p, + ) + self.decoder = Decoder( + ch_mult=config.decoder_ch_mult, + z_channels=config.z_channels, + dropout=config.dropout_p, + ) + + self.quantize = VectorQuantizer( + config.codebook_size, + config.codebook_embed_dim, + config.commit_loss_beta, + config.entropy_loss_ratio, + config.codebook_l2_norm, + config.codebook_show_usage, + ) + self.quant_conv = nn.Conv2d(config.z_channels, config.codebook_embed_dim, 1) + self.post_quant_conv = nn.Conv2d( + config.codebook_embed_dim, config.z_channels, 1 + ) + + def encode(self, x): + h = self.encoder(x) + h = self.quant_conv(h) + quant, emb_loss, info = self.quantize(h) + return quant, emb_loss, info + + def decode(self, quant): + quant = self.post_quant_conv(quant) + dec = self.decoder(quant) + return dec + + def decode_code(self, code_b, shape=None, channel_first=True): + quant_b = self.quantize.get_codebook_entry(code_b, shape, channel_first) + dec = self.decode(quant_b) + return dec + + def forward(self, input): + quant, diff, _ = self.encode(input) + dec = self.decode(quant) + return dec, diff + + +class MultiModalityPreTrainedModel(PreTrainedModel): + config_class = MultiModalityConfig + base_model_prefix = "multi_modality" + _no_split_modules = [] + _skip_keys_device_placement = "past_key_values" + + +# Copied and adapted from: +# https://github.com/deepseek-ai/Janus/tree/main/janus/models/modeling_vlm.py +class MultiModalityCausalLM(MultiModalityPreTrainedModel): + + def __init__( + self, + config: MultiModalityConfig, + quant_config: Optional[QuantizationConfig] = None, + ): + super().__init__(config) + + vision_config = config.vision_config + vision_cls = model_name_to_cls(vision_config.cls) + self.vision_model = vision_cls(**vision_config.params) + + aligner_config = config.aligner_config + aligner_cls = model_name_to_cls(aligner_config.cls) + self.aligner = aligner_cls(aligner_config.params) + + gen_vision_config = config.gen_vision_config + gen_vision_cls = model_name_to_cls(gen_vision_config.cls) + self.gen_vision_model = gen_vision_cls() + + gen_aligner_config = config.gen_aligner_config + gen_aligner_cls = model_name_to_cls(gen_aligner_config.cls) + self.gen_aligner = gen_aligner_cls(gen_aligner_config.params) + + gen_head_config = config.gen_head_config + gen_head_cls = model_name_to_cls(gen_head_config.cls) + self.gen_head = gen_head_cls(gen_head_config.params) + + self.gen_embed = torch.nn.Embedding( + gen_vision_config.params["image_token_size"], + gen_vision_config.params["n_embed"], + ) + + language_config = config.language_config + self.language_model = LlamaForCausalLM( + language_config, quant_config=quant_config + ) + self.logits_processor = LogitsProcessor(config) + + def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + pixel_values = torch.concat([item.pixel_values for item in items], dim=0) + bs, n = pixel_values.shape[0:2] + pixel_values = pixel_values.to( + device=self.vision_model.device, dtype=self.vision_model.dtype + ) + images = rearrange(pixel_values, "b n c h w -> (b n) c h w") + + # [b x n, T2, D] + images_embeds = self.aligner(self.vision_model(images)) + + # [b x n, T2, D] -> [b, n x T2, D] + images_embeds = rearrange(images_embeds, "(b n) t d -> b (n t) d", b=bs, n=n) + + return images_embeds + + def get_input_embeddings(self) -> nn.Embedding: + return self.language_model.get_input_embeddings() + + @torch.no_grad() + def forward( + self, + input_ids: torch.LongTensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + get_embedding: bool = False, + ) -> torch.Tensor: + hidden_states = general_mm_embed_routine( + input_ids=input_ids, + forward_batch=forward_batch, + image_data_embedding_func=self.get_image_feature, + language_model=self.language_model, + positions=positions, + ) + + return hidden_states + + def prepare_gen_img_embeds(self, image_ids: torch.LongTensor): + return self.gen_aligner(self.gen_embed(image_ids)) + + def pad_input_ids(self, input_ids: List[int], image_inputs: MultimodalInputs): + im_start_id = image_inputs.im_start_id + im_end_id = image_inputs.im_end_id + media_token_pairs = [(im_start_id, im_end_id)] + + helper = MultiModalityDataPaddingPatternTokenPairs(media_token_pairs) + + return helper.pad_input_tokens(input_ids, image_inputs) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq~" in name or "projector" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + if name.startswith("model.vision_tower") and name not in params_dict: + continue + + # skip generation sub model + if "gen" in name: + continue + + # adapt to VisionAttention + name = name.replace(r"self_attn.out_proj", r"self_attn.proj") + if "vision_model.vision_tower" in name: + name = name.replace("attn.qkv", "attn.qkv_proj") + + for param_name, weight_name, shard_id in stacked_params_mapping: + # replace the name and load with customized loader + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + + # # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", None) + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +AutoModel.register(config_class=MultiModalityConfig, model_class=MultiModalityCausalLM) +EntryClass = [MultiModalityCausalLM] diff --git a/lib/python3.12/site-packages/sglang/srt/models/deepseek_nextn.py b/lib/python3.12/site-packages/sglang/srt/models/deepseek_nextn.py new file mode 100644 index 0000000000000000000000000000000000000000..4849177af7ebe13835024ee8d501c267bc8dafb2 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/deepseek_nextn.py @@ -0,0 +1,439 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""Inference-only DeepSeek NextN Speculative Decoding.""" +import logging +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ReplicatedLinear +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.ep_moe.layer import EPMoE +from sglang.srt.layers.moe.fused_moe_triton import FusedMoE +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.quantization.fp8_utils import ( + block_quant_to_tensor_quant, + normalize_e4m3fn_to_e4m3fnuz, +) +from sglang.srt.layers.quantization.int8_utils import ( + block_dequant as int8_block_dequant, +) +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.deepseek_v2 import DeepseekV2DecoderLayer, DeepseekV3ForCausalLM +from sglang.srt.utils import BumpAllocator, add_prefix, is_cuda, is_hip + +_is_hip = is_hip() +_is_cuda = is_cuda() + +if _is_cuda: + from sgl_kernel import awq_dequantize +else: + from vllm._custom_ops import awq_dequantize + + +logger = logging.getLogger(__name__) + + +class DeepseekModelNextN(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + enable_tp=not global_server_args_dict["enable_dp_attention"], + prefix=add_prefix("embed_tokens", prefix), + ) + + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) + + self.decoder = DeepseekV2DecoderLayer( + config, + 0, + quant_config=quant_config, + is_nextn=True, + prefix=add_prefix("decoder", prefix), + ) + + self.shared_head = nn.Module() + self.shared_head.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + zero_allocator = BumpAllocator( + buffer_size=2, + dtype=torch.float32, + device=( + input_embeds.device if input_embeds is not None else input_ids.device + ), + ) + + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + + hidden_states = self.eh_proj( + torch.cat( + ( + self.enorm(hidden_states), + self.hnorm(forward_batch.spec_info.hidden_states), + ), + dim=-1, + ) + ) + + residual = None + hidden_states, residual = self.decoder( + positions, hidden_states, forward_batch, residual, zero_allocator + ) + + if not forward_batch.forward_mode.is_idle(): + hidden_states, _ = self.shared_head.norm(hidden_states, residual) + return hidden_states + + +class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM): + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + nn.Module.__init__(self) + self.config = config + self.tp_size = get_tensor_model_parallel_world_size() + self.quant_config = quant_config + self.determine_n_share_experts_fusion("DeepseekV3ForCausalLMNextN") + + self.model = DeepseekModelNextN( + config, quant_config, prefix=add_prefix("model", prefix) + ) + + if global_server_args_dict["enable_dp_attention"]: + self.lm_head = ReplicatedLinear( + config.hidden_size, + config.vocab_size, + bias=False, + prefix=add_prefix("model.shared_head.head", prefix), + ) + self.logits_processor = LogitsProcessor(config, skip_all_gather=True) + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("model.shared_head.head", prefix), + ) + self.logits_processor = LogitsProcessor(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + hidden_states = self.model(input_ids, positions, forward_batch) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + if hasattr(self.config, "num_nextn_predict_layers"): + num_nextn_layers = self.config.num_nextn_predict_layers + assert num_nextn_layers == 1, "Only 1 nextn layer is supportted" + assert num_nextn_layers == self.config.num_hidden_layers + else: + raise ValueError("num_nextn_predict_layers is not in the config") + + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + if self.n_share_experts_fusion > 0: + logger.info( + f"Cloning {self.n_share_experts_fusion} " + "replicas of the shared expert into MoE for DeepseekV3ForCausalLMNextN" + ) + weights_list = list(weights) + weights_dict = dict(weights_list) + if self.quant_config is None or self.quant_config.get_name() == "w8a8_int8": + suffix_list = [ + "down_proj.weight", + "down_proj.weight_scale", + "gate_proj.weight", + "gate_proj.weight_scale", + "up_proj.weight", + "up_proj.weight_scale", + ] + else: + suffix_list = [ + "down_proj.weight", + "down_proj.weight_scale_inv", + "gate_proj.weight", + "gate_proj.weight_scale_inv", + "up_proj.weight", + "up_proj.weight_scale_inv", + ] + names_to_remove = [] + for suffix in suffix_list: + shared_expert_weight_name = ( + f"model.layers.0.mlp.shared_experts.{suffix}" + ) + for num_repeat in range(self.n_share_experts_fusion): + weights_list.append( + ( + f"model.layers.0." + f"mlp.experts." + f"{self.config.n_routed_experts + num_repeat}" + f".{suffix}", + weights_dict[shared_expert_weight_name], + ) + ) + names_to_remove += [shared_expert_weight_name] + weights = [w for w in weights_list if w[0] not in names_to_remove] + + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + MoEImpl = EPMoE if global_server_args_dict["enable_ep_moe"] else FusedMoE + expert_params_mapping = MoEImpl.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts + self.n_share_experts_fusion, + ) + + # Fuse q_a_proj and kv_a_proj_with_mqa along output dimension when q_lora_rank is not None + fuse_qkv_a_proj = hasattr(self.config, "q_lora_rank") and ( + self.config.q_lora_rank is not None + ) + cached_a_proj = {} if fuse_qkv_a_proj else None + + nextn_layer_prefix = "model.layers.0" + nextn_spec_weight_names = [ + "shared_head.norm", + "eh_proj", + "enorm", + "hnorm", + ] + + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if not name.startswith(nextn_layer_prefix): + continue + + # Use shared head and embed weights from target model + if "shared_head.head" in name or "embed_tokens" in name: + continue + + is_decoder = True + # For nextn specific weights + for weight_name in nextn_spec_weight_names: + if weight_name in name: + name = name.replace(nextn_layer_prefix, "model") + is_decoder = False + break + # For decoder layer weights + if is_decoder: + name = name.replace(nextn_layer_prefix, "model.decoder") + + if "rotary_emb.inv_freq" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip non-stacked layers and experts (experts handled below). + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since we handle the experts below in expert_params_mapping, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below in expert_params_mapping + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if ("mlp.experts." in name) and name not in params_dict: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + # Handle fused_qkv_a_proj + if fuse_qkv_a_proj and ( + "q_a_proj" in name or "kv_a_proj_with_mqa" in name + ): + cached_a_proj[name] = loaded_weight + q_a_proj_name = ( + name + if "q_a_proj" in name + else name.replace("kv_a_proj_with_mqa", "q_a_proj") + ) + kv_a_proj_name = ( + name + if "kv_a_proj_with_mqa" in name + else name.replace("q_a_proj", "kv_a_proj_with_mqa") + ) + + # When both q_a_proj and kv_a_proj_with_mqa has been cached, load the fused weight to parameter + if ( + q_a_proj_name in cached_a_proj + and kv_a_proj_name in cached_a_proj + ): + + q_a_proj_weight = cached_a_proj[q_a_proj_name] + kv_a_proj_weight = cached_a_proj[kv_a_proj_name] + fused_weight = torch.cat( + [q_a_proj_weight, kv_a_proj_weight], dim=0 + ) + + param_name = name.replace( + "q_a_proj", "fused_qkv_a_proj_with_mqa" + ) + param = params_dict[param_name] + + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, fused_weight) + cached_a_proj.pop(q_a_proj_name) + cached_a_proj.pop(kv_a_proj_name) + else: + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + self_attn = self.model.decoder.self_attn + if hasattr(self_attn.kv_b_proj, "qweight"): + # AWQ compatible + if _is_cuda: + w = awq_dequantize( + self_attn.kv_b_proj.qweight, + self_attn.kv_b_proj.scales, + self_attn.kv_b_proj.qzeros, + ).T + else: + w = awq_dequantize( + self_attn.kv_b_proj.qweight, + self_attn.kv_b_proj.scales, + self_attn.kv_b_proj.qzeros, + 0, + 0, + 0, + ).T + else: + w = self_attn.kv_b_proj.weight + # NOTE(HandH1998): Since `bmm_fp8` only supports per-tensor scale, we have to requantize `self_attn.kv_b_proj`. + # This may affect the accuracy of fp8 model. + if hasattr(self.quant_config, "weight_block_size") and w.dtype in ( + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + ): + weight_block_size = self.quant_config.weight_block_size + if weight_block_size is not None: + assert hasattr(self_attn.kv_b_proj, "weight_scale_inv") + if _is_hip: + weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=w, + weight_scale=self_attn.kv_b_proj.weight_scale_inv, + input_scale=None, + ) + else: + weight = w + weight_scale = self_attn.kv_b_proj.weight_scale_inv + + w, scale = block_quant_to_tensor_quant( + weight, weight_scale, weight_block_size + ) + self_attn.w_scale = scale + if w.dtype == torch.int8: + if hasattr(self.quant_config, "weight_block_size"): + # block-wise int8 need it + weight_block_size = self.quant_config.weight_block_size + if weight_block_size is not None: + assert hasattr(self_attn.kv_b_proj, "weight_scale_inv") + weight = w + weight_scale = self_attn.kv_b_proj.weight_scale_inv + w = int8_block_dequant(weight, weight_scale, weight_block_size).to( + torch.bfloat16 + ) + else: + # channel-wise int8 need it + assert hasattr(self_attn.kv_b_proj, "weight_scale") + w = w.to(torch.bfloat16) * self_attn.kv_b_proj.weight_scale.to( + torch.bfloat16 + ) + w_kc, w_vc = w.unflatten( + 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim) + ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1) + self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2) + self_attn.w_vc = w_vc.contiguous().transpose(1, 2) + if hasattr(self_attn.kv_b_proj, "weight_scale") and self_attn.w_scale is None: + self_attn.w_scale = self_attn.kv_b_proj.weight_scale + if _is_hip: + self_attn.w_scale *= 2.0 + + +EntryClass = [DeepseekV3ForCausalLMNextN] diff --git a/lib/python3.12/site-packages/sglang/srt/models/deepseek_v2.py b/lib/python3.12/site-packages/sglang/srt/models/deepseek_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..c24c098c02966f79fca0b5cdd74ac83050e0fad8 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/deepseek_v2.py @@ -0,0 +1,1807 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from: +# https://github.com/vllm-project/vllm/blob/fb6af8bc086328ca6659e72d11ffd4309ce4de22/vllm/model_executor/models/deepseek_v2.py +"""Inference-only DeepseekV2 model.""" + +import logging +import os +from dataclasses import dataclass +from enum import Enum, IntEnum, auto +from typing import Any, Dict, Iterable, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn +from tqdm import tqdm +from transformers import PretrainedConfig + +from sglang.srt.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + parallel_state, + tensor_model_parallel_all_reduce, +) +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.dp_attention import ( + dp_gather_partial, + dp_scatter, + get_attention_dp_size, + get_attention_tp_rank, + get_attention_tp_size, + tp_all_gather, + tp_reduce_scatter, +) +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE, EPMoE +from sglang.srt.layers.moe.ep_moe.token_dispatcher import DeepEPDispatcher +from sglang.srt.layers.moe.fused_moe_triton import FusedMoE +from sglang.srt.layers.moe.topk import select_experts +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.quantization.deep_gemm import _ENABLE_JIT_DEEPGEMM +from sglang.srt.layers.quantization.fp8_kernel import ( + per_tensor_quant_mla_deep_gemm_masked_fp8, + per_tensor_quant_mla_fp8, +) +from sglang.srt.layers.quantization.fp8_utils import ( + block_quant_to_tensor_quant, + channel_quant_to_tensor_quant, + normalize_e4m3fn_to_e4m3fnuz, +) +from sglang.srt.layers.quantization.int8_utils import ( + block_dequant as int8_block_dequant, +) +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope, get_rope_wrapper +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.managers.expert_distribution import ExpertDistributionRecorder +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import ( + BumpAllocator, + DeepEPMode, + add_prefix, + get_bool_env_var, + get_int_env_var, + is_cuda, + is_hip, +) + +_is_hip = is_hip() +_is_cuda = is_cuda() + +if _is_cuda: + from sgl_kernel import awq_dequantize, bmm_fp8, merge_state_v2 + + from sglang.srt.layers.quantization.deep_gemm import ( + grouped_gemm_nt_f8f8bf16_masked as deep_gemm_grouped_gemm_nt_f8f8bf16_masked, + ) +else: + from vllm._custom_ops import awq_dequantize + +if _is_hip: + from sglang.srt.layers.attention.triton_ops.rocm_mla_decode_rope import ( + decode_attention_fwd_grouped_rope, + ) + +expert_distribution_recorder = ExpertDistributionRecorder() + +logger = logging.getLogger(__name__) + + +class AttnForwardMethod(IntEnum): + # Use multi-head attention + MHA = auto() + + # Use absorbed multi-latent attention + MLA = auto() + + # Use multi-head attention, but with KV cache chunked. + # This method can avoid OOM when prefix lengths are long. + MHA_CHUNKED_KV = auto() + + +class DeepseekV2MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + reduce_results: bool = True, + prefix: str = "", + tp_rank: Optional[int] = None, + tp_size: Optional[int] = None, + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=add_prefix("down_proj", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x, forward_mode: Optional[ForwardMode] = None): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class MoEGate(nn.Module): + def __init__( + self, + config, + prefix: str = "", + ): + super().__init__() + self.weight = nn.Parameter( + torch.empty((config.n_routed_experts, config.hidden_size)) + ) + if config.topk_method == "noaux_tc": + self.e_score_correction_bias = nn.Parameter( + torch.empty((config.n_routed_experts)) + ) + else: + self.e_score_correction_bias = None + + def forward(self, hidden_states): + logits = F.linear(hidden_states, self.weight, None) + return logits + + +class DeepseekV2MoE(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + self.routed_scaling_factor = config.routed_scaling_factor + self.n_shared_experts = config.n_shared_experts + self.n_share_experts_fusion = global_server_args_dict["n_share_experts_fusion"] + + if self.tp_size > config.n_routed_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {config.n_routed_experts}." + ) + + if config.hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {config.hidden_act}. " + "Only silu is supported for now." + ) + + self.gate = MoEGate(config=config, prefix=add_prefix("gate", prefix)) + + MoEImpl = ( + DeepEPMoE + if global_server_args_dict["enable_deepep_moe"] + else (EPMoE if global_server_args_dict["enable_ep_moe"] else FusedMoE) + ) + + self.experts = MoEImpl( + num_experts=config.n_routed_experts + self.n_share_experts_fusion, + top_k=config.num_experts_per_tok + min(self.n_share_experts_fusion, 1), + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + renormalize=config.norm_topk_prob, + quant_config=quant_config, + use_grouped_topk=True, + num_expert_group=config.n_group, + topk_group=config.topk_group, + correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + prefix=add_prefix("experts", prefix), + **( + dict(deepep_mode=DeepEPMode[global_server_args_dict["deepep_mode"]]) + if global_server_args_dict["enable_deepep_moe"] + else {} + ), + ) + + if config.n_shared_experts is not None and self.n_share_experts_fusion == 0: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + # disable tp for shared experts when enable deepep moe + if not global_server_args_dict["enable_deepep_moe"]: + self.shared_experts = DeepseekV2MLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix("shared_experts", prefix), + ) + else: + self.shared_experts = DeepseekV2MLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix("shared_experts", prefix), + tp_rank=0, + tp_size=1, + ) + + if global_server_args_dict["enable_deepep_moe"]: + # TODO: we will support tp < ep in the future + self.ep_size = get_tensor_model_parallel_world_size() + self.num_experts = config.n_routed_experts + self.top_k = config.num_experts_per_tok + self.renormalize = config.norm_topk_prob + self.topk_group = config.topk_group + self.num_expert_group = config.n_group + self.correction_bias = ( + self.gate.e_score_correction_bias.data + if self.gate.e_score_correction_bias is not None + else None + ) + + self.deepep_dispatcher = DeepEPDispatcher( + group=parallel_state.get_tp_group().device_group, + router_topk=self.top_k, + permute_fusion=True, + num_experts=config.n_routed_experts, + num_local_experts=config.n_routed_experts // self.tp_size, + hidden_size=config.hidden_size, + params_dtype=config.torch_dtype, + deepep_mode=DeepEPMode[global_server_args_dict["deepep_mode"]], + async_finish=True, # TODO + return_recv_hook=True, + ) + + def forward( + self, hidden_states: torch.Tensor, forward_mode: Optional[ForwardMode] = None + ) -> torch.Tensor: + if not global_server_args_dict["enable_deepep_moe"]: + return self.forward_normal(hidden_states) + else: + return self.forward_deepep(hidden_states, forward_mode) + + def forward_normal(self, hidden_states: torch.Tensor) -> torch.Tensor: + shared_output = self._forward_shared_experts(hidden_states) + # router_logits: (num_tokens, n_experts) + router_logits = self.gate(hidden_states) + final_hidden_states = ( + self.experts(hidden_states=hidden_states, router_logits=router_logits) + * self.routed_scaling_factor + ) + if shared_output is not None: + final_hidden_states = final_hidden_states + shared_output + if self.tp_size > 1: + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + return final_hidden_states + + def forward_deepep( + self, hidden_states: torch.Tensor, forward_mode: ForwardMode + ) -> torch.Tensor: + shared_output = None + if ( + forward_mode is not None + and not forward_mode.is_idle() + and hidden_states.shape[0] > 0 + ): + # router_logits: (num_tokens, n_experts) + router_logits = self.gate(hidden_states) + shared_output = self._forward_shared_experts(hidden_states) + topk_weights, topk_idx = select_experts( + hidden_states=hidden_states, + router_logits=router_logits, + top_k=self.top_k, + use_grouped_topk=True, + renormalize=self.renormalize, + topk_group=self.topk_group, + num_expert_group=self.num_expert_group, + correction_bias=self.correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + ) + else: + topk_idx = torch.full( + (0, self.top_k), -1, dtype=torch.int, device=hidden_states.device + ) + topk_weights = torch.empty( + (0, self.top_k), dtype=torch.float32, device=hidden_states.device + ) + if self.ep_size > 1: + # TODO(ch-wan): allow users to set num_max_dispatch_tokens_per_rank value + ( + hidden_states, + topk_idx, + topk_weights, + reorder_topk_ids, + seg_indptr, + masked_m, + expected_m, + ) = self.deepep_dispatcher.dispatch( + hidden_states, + topk_idx, + topk_weights, + forward_mode=forward_mode, + ) + final_hidden_states = self.experts( + hidden_states=hidden_states, + reorder_topk_ids=reorder_topk_ids, + seg_indptr=seg_indptr, + masked_m=masked_m, + expected_m=expected_m, + forward_mode=forward_mode, + ) + if self.ep_size > 1: + final_hidden_states = self.deepep_dispatcher.combine( + final_hidden_states, + topk_idx, + topk_weights, + forward_mode, + ) + final_hidden_states *= self.routed_scaling_factor + + if shared_output is not None: + final_hidden_states = final_hidden_states + shared_output + + return final_hidden_states + + def _forward_shared_experts(self, hidden_states): + if self.n_share_experts_fusion == 0: + return self.shared_experts(hidden_states) + else: + return None + + +def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: + import math + + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +class DeepseekV2AttentionMLA(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + hidden_size: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int, + kv_lora_rank: int, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + max_position_embeddings: int = 8192, + quant_config: Optional[QuantizationConfig] = None, + reduce_results: bool = True, + layer_id: int = None, + prefix: str = "", + ) -> None: + super().__init__() + self.layer_id = layer_id + self.hidden_size = hidden_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.dp_size = get_attention_dp_size() + attn_tp_rank = get_attention_tp_rank() + attn_tp_size = get_attention_tp_size() + + self.num_heads = num_heads + assert num_heads % attn_tp_size == 0 + self.num_local_heads = num_heads // attn_tp_size + self.scaling = self.qk_head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + # For tensor parallel attention + if self.q_lora_rank is not None: + self.fused_qkv_a_proj_with_mqa = ReplicatedLinear( + self.hidden_size, + self.q_lora_rank + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("fused_qkv_a_proj_with_mqa", prefix), + ) + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) + self.q_b_proj = ColumnParallelLinear( + q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("q_b_proj", prefix), + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + ) + else: + self.q_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("q_proj", prefix), + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + ) + self.kv_a_proj_with_mqa = ReplicatedLinear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("kv_a_proj_with_mqa", prefix), + ) + + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + quant_config=quant_config, + prefix=add_prefix("kv_b_proj", prefix), + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + ) + # O projection. + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=add_prefix("o_proj", prefix), + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + ) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) + + if rope_scaling: + rope_scaling["rope_type"] = "deepseek_yarn" + + self.rotary_emb = get_rope( + qk_rope_head_dim, + rotary_dim=qk_rope_head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + is_neox_style=False, + ) + + if rope_scaling: + mscale_all_dim = rope_scaling.get("mscale_all_dim", False) + scaling_factor = rope_scaling["factor"] + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) + self.scaling = self.scaling * mscale * mscale + else: + self.rotary_emb.forward = self.rotary_emb.forward_native + + self.attn_mqa = RadixAttention( + self.num_local_heads, + self.kv_lora_rank + self.qk_rope_head_dim, + self.scaling, + num_kv_heads=1, + layer_id=layer_id, + v_head_dim=self.kv_lora_rank, + quant_config=quant_config, + prefix=add_prefix("attn_mqa", prefix), + ) + + self.attn_mha = RadixAttention( + self.num_local_heads, + self.qk_nope_head_dim + self.qk_rope_head_dim, + self.scaling, + num_kv_heads=self.num_local_heads, + layer_id=layer_id, + v_head_dim=self.v_head_dim, + quant_config=quant_config, + prefix=add_prefix("attn_mha", prefix), + ) + + self.w_kc = None + self.w_vc = None + self.w_scale = None + + self.w_scale_k = None + self.w_scale_v = None + self.use_deep_gemm_bmm = False + + self.flashinfer_mla_disable_ragged = global_server_args_dict[ + "flashinfer_mla_disable_ragged" + ] + self.disable_chunked_prefix_cache = global_server_args_dict[ + "disable_chunked_prefix_cache" + ] + self.attention_backend = global_server_args_dict["attention_backend"] + self.rocm_fused_decode_mla = get_bool_env_var( + "SGLANG_ROCM_FUSED_DECODE_MLA", "false" + ) + + # TODO: Design a finer way to determine the threshold + self.chunked_prefix_cache_threshold = get_int_env_var( + "SGL_CHUNKED_PREFIX_CACHE_THRESHOLD", 8192 + ) + + def dispatch_attn_forward_method( + self, forward_batch: ForwardBatch + ) -> AttnForwardMethod: + if self.attention_backend == "flashinfer": + # Flashinfer MLA: Do not absorb when enabling ragged prefill + if ( + not self.flashinfer_mla_disable_ragged + and forward_batch.forward_mode.is_extend() + and not forward_batch.forward_mode.is_target_verify() + and not forward_batch.forward_mode.is_draft_extend() + and sum(forward_batch.extend_prefix_lens_cpu) == 0 + ): + return AttnForwardMethod.MHA + else: + return AttnForwardMethod.MLA + elif self.attention_backend == "fa3": + # Flash Attention: Use MHA with chunked KV cache when prefilling on long sequences. + if forward_batch.extend_prefix_lens_cpu is not None: + sum_extend_prefix_lens = sum(forward_batch.extend_prefix_lens_cpu) + if ( + forward_batch.forward_mode.is_extend() + and not self.disable_chunked_prefix_cache + and not forward_batch.forward_mode.is_target_verify() + and not forward_batch.forward_mode.is_draft_extend() + and ( + sum_extend_prefix_lens >= self.chunked_prefix_cache_threshold + or sum_extend_prefix_lens == 0 + ) + ): + return AttnForwardMethod.MHA_CHUNKED_KV + else: + return AttnForwardMethod.MLA + else: + # Triton: Use normal computation for prefill and use weight absorption for extend/decode + if ( + forward_batch.forward_mode.is_extend() + and not forward_batch.forward_mode.is_target_verify() + and not forward_batch.forward_mode.is_draft_extend() + and sum(forward_batch.extend_prefix_lens_cpu) == 0 + ): + return AttnForwardMethod.MHA + else: + return AttnForwardMethod.MLA + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + ) -> torch.Tensor: + if hidden_states.shape[0] == 0: + assert ( + not self.o_proj.reduce_results + ), "short-circuiting allreduce will lead to hangs" + return hidden_states + + attn_forward_method = self.dispatch_attn_forward_method(forward_batch) + + if attn_forward_method == AttnForwardMethod.MHA: + return self.forward_normal(positions, hidden_states, forward_batch) + elif attn_forward_method == AttnForwardMethod.MHA_CHUNKED_KV: + return self.forward_normal_chunked_kv( + positions, hidden_states, forward_batch + ) + else: + if _is_hip: + if ( + self.rocm_fused_decode_mla + and forward_batch.forward_mode.is_decode() + ): + return self.forward_absorb_fused_mla_rope( + positions, hidden_states, forward_batch + ) + else: + return self.forward_absorb( + positions, hidden_states, forward_batch, zero_allocator + ) + else: + return self.forward_absorb( + positions, hidden_states, forward_batch, zero_allocator + ) + + def forward_normal( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + if self.q_lora_rank is not None: + q, latent_cache = self.fused_qkv_a_proj_with_mqa(hidden_states)[0].split( + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1 + ) + q = self.q_a_layernorm(q) + q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) + else: + q = self.q_proj(hidden_states)[0].view( + -1, self.num_local_heads, self.qk_head_dim + ) + latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0] + + _, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + latent_cache = latent_cache.unsqueeze(1) + kv_a = self.kv_a_layernorm(kv_a.contiguous()) + kv = self.kv_b_proj(kv_a)[0] + kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim) + k_nope = kv[..., : self.qk_nope_head_dim] + v = kv[..., self.qk_nope_head_dim :] + k_pe = latent_cache[:, :, self.kv_lora_rank :] + q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) + q[..., self.qk_nope_head_dim :] = q_pe + k = torch.empty_like(q) + k[..., : self.qk_nope_head_dim] = k_nope + k[..., self.qk_nope_head_dim :] = k_pe + + latent_cache[:, :, : self.kv_lora_rank] = kv_a.unsqueeze(1) + latent_cache[:, :, self.kv_lora_rank :] = k_pe + + # Save latent cache + forward_batch.token_to_kv_pool.set_kv_buffer( + self.attn_mha, forward_batch.out_cache_loc, latent_cache, None + ) + attn_output = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False) + attn_output = attn_output.reshape(-1, self.num_local_heads * self.v_head_dim) + output, _ = self.o_proj(attn_output) + return output + + def forward_absorb( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + ) -> torch.Tensor: + if self.q_lora_rank is not None: + q, latent_cache = self.fused_qkv_a_proj_with_mqa(hidden_states)[0].split( + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1 + ) + q = self.q_a_layernorm(q) + q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) + else: + q = self.q_proj(hidden_states)[0].view( + -1, self.num_local_heads, self.qk_head_dim + ) + latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0] + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + if self.use_deep_gemm_bmm: + q_nope_val, q_nope_scale, masked_m, expected_m, aligned_m = ( + per_tensor_quant_mla_deep_gemm_masked_fp8( + q_nope.transpose(0, 1), dtype=torch.float8_e4m3fn + ) + ) + q_nope_out = q_nope.new_empty( + (self.num_local_heads, aligned_m, self.kv_lora_rank) + ) + deep_gemm_grouped_gemm_nt_f8f8bf16_masked( + (q_nope_val, q_nope_scale), + (self.w_kc, self.w_scale_k), + q_nope_out, + masked_m, + expected_m, + ) + q_nope_out = q_nope_out[:, :expected_m, :] + elif self.w_kc.dtype == torch.float8_e4m3fnuz: + # TODO(kernel): add bmm_fp8 for torch.float8_e4m3fnuz + q_nope_out = torch.bmm( + q_nope.to(torch.bfloat16).transpose(0, 1), + self.w_kc.to(torch.bfloat16) * self.w_scale, + ) + elif self.w_kc.dtype == torch.float8_e4m3fn: + q_nope_val, q_nope_scale = per_tensor_quant_mla_fp8( + q_nope.transpose(0, 1), + zero_allocator.allocate(1), + ) + q_nope_out = bmm_fp8( + q_nope_val, self.w_kc, q_nope_scale, self.w_scale, torch.bfloat16 + ) + else: + q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc) + + q_nope_out = q_nope_out.transpose(0, 1) + + k_nope = latent_cache[..., : self.kv_lora_rank] + k_nope = self.kv_a_layernorm(k_nope).unsqueeze(1) + k_pe = latent_cache[..., self.kv_lora_rank :].unsqueeze(1) + + q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) + + if self.attention_backend == "fa3": + attn_output = self.attn_mqa( + q_nope_out, k_nope, k_nope, forward_batch, q_rope=q_pe, k_rope=k_pe + ) + else: + q = torch.cat([q_nope_out, q_pe], dim=-1) + k = torch.cat([k_nope, k_pe], dim=-1) + attn_output = self.attn_mqa(q, k, k_nope, forward_batch) + attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank) + + if self.use_deep_gemm_bmm: + attn_output_val, attn_output_scale, masked_m, expected_m, aligned_m = ( + per_tensor_quant_mla_deep_gemm_masked_fp8( + attn_output.transpose(0, 1), dtype=torch.float8_e4m3fn + ) + ) + attn_bmm_output = attn_output.new_empty( + (self.num_local_heads, aligned_m, self.v_head_dim) + ) + deep_gemm_grouped_gemm_nt_f8f8bf16_masked( + (attn_output_val, attn_output_scale), + (self.w_vc, self.w_scale_v), + attn_bmm_output, + masked_m, + expected_m, + ) + attn_bmm_output = attn_bmm_output[:, :expected_m, :] + elif self.w_vc.dtype == torch.float8_e4m3fnuz: + # TODO(kernel): add bmm_fp8 for torch.float8_e4m3fnuz + attn_bmm_output = torch.bmm( + attn_output.to(torch.bfloat16).transpose(0, 1), + self.w_vc.to(torch.bfloat16) * self.w_scale, + ) + elif self.w_vc.dtype == torch.float8_e4m3fn: + attn_output_val, attn_output_scale = per_tensor_quant_mla_fp8( + attn_output.transpose(0, 1), + zero_allocator.allocate(1), + ) + attn_bmm_output = bmm_fp8( + attn_output_val, + self.w_vc, + attn_output_scale, + self.w_scale, + torch.bfloat16, + ) + else: + attn_bmm_output = torch.bmm(attn_output.transpose(0, 1), self.w_vc) + attn_output = attn_bmm_output.transpose(0, 1).flatten(1, 2) + output, _ = self.o_proj(attn_output) + + return output + + def forward_absorb_fused_mla_rope( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + ) -> torch.Tensor: + enable_rope_fusion = ( + os.getenv("SGLANG_FUSED_MLA_ENABLE_ROPE_FUSION", "1") == "1" + ) + q_len = hidden_states.shape[0] + q_input = hidden_states.new_empty( + q_len, self.num_local_heads, self.kv_lora_rank + self.qk_rope_head_dim + ) + if self.q_lora_rank is not None: + q, latent_cache = self.fused_qkv_a_proj_with_mqa(hidden_states)[0].split( + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1 + ) + q = self.q_a_layernorm(q) + q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) + else: + q = self.q_proj(hidden_states)[0].view( + -1, self.num_local_heads, self.qk_head_dim + ) + latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0] + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + if self.w_kc.dtype == torch.float8_e4m3fnuz: + # TODO(kernel): add bmm_fp8 for torch.float8_e4m3fnuz + q_nope_out = torch.bmm( + q_nope.to(torch.bfloat16).transpose(0, 1), + self.w_kc.to(torch.bfloat16) * self.w_scale, + ) + elif self.w_kc.dtype == torch.float8_e4m3fn: + q_nope_val, q_nope_scale = per_tensor_quant_mla_fp8( + q_nope.transpose(0, 1), + zero_allocator.allocate(1), + dtype=torch.float8_e4m3fn, + ) + q_nope_out = bmm_fp8( + q_nope_val, self.w_kc, q_nope_scale, self.w_scale, torch.bfloat16 + ) + else: + q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc) + q_input[..., : self.kv_lora_rank] = q_nope_out.transpose(0, 1) + v_input = latent_cache[..., : self.kv_lora_rank] + v_input = self.kv_a_layernorm(v_input.contiguous()).unsqueeze(1) + k_input = latent_cache.unsqueeze(1) + k_input[..., : self.kv_lora_rank] = v_input + + if not enable_rope_fusion: + k_pe = k_input[..., self.kv_lora_rank :] + q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) + q_input[..., self.kv_lora_rank :] = q_pe + k_input[..., self.kv_lora_rank :] = k_pe + k_pe_output = None + else: + k_pe_output = torch.empty_like(k_input[..., self.kv_lora_rank :]) + + q_input[..., self.kv_lora_rank :] = q_pe + + # attn_output = self.attn_mqa(q_input, k_input, v_input, forward_batch) + # Use Fused ROPE with use_rope=OFF. + attn_output = torch.empty( + (q_len, self.num_local_heads, self.kv_lora_rank), + dtype=q.dtype, + device=q.device, + ) + attn_logits, _, kv_indptr, kv_indices, _, _, _ = ( + forward_batch.attn_backend.forward_metadata + ) + cos_sin_cache = self.rotary_emb.cos_sin_cache + num_kv_split = forward_batch.attn_backend.num_kv_splits + sm_scale = self.attn_mqa.scaling + if attn_logits is None: + attn_logits = torch.empty( + ( + forward_batch.batch_size, + self.num_local_heads, + num_kv_split, + self.kv_lora_rank + 1, + ), + dtype=torch.float32, + device=q.device, + ) + + # save current latent cache. + forward_batch.token_to_kv_pool.set_kv_buffer( + self.attn_mqa, forward_batch.out_cache_loc, k_input, None + ) + key_cache_buf = forward_batch.token_to_kv_pool.get_key_buffer( + self.attn_mqa.layer_id + ) + val_cache_buf = key_cache_buf[..., : self.kv_lora_rank] + + decode_attention_fwd_grouped_rope( + q_input, + key_cache_buf, + val_cache_buf, + attn_output, + kv_indptr, + kv_indices, + k_pe_output, + self.kv_lora_rank, + self.rotary_emb.rotary_dim, + cos_sin_cache, + positions, + attn_logits, + num_kv_split, + sm_scale, + logit_cap=self.attn_mqa.logit_cap, + use_rope=enable_rope_fusion, + is_neox_style=self.rotary_emb.is_neox_style, + ) + + if enable_rope_fusion: + k_input[..., self.kv_lora_rank :] = k_pe_output + forward_batch.token_to_kv_pool.set_kv_buffer( + self.attn_mqa, forward_batch.out_cache_loc, k_input, None + ) + + attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank) + + if self.w_vc.dtype == torch.float8_e4m3fnuz: + # TODO(kernel): add bmm_fp8 for torch.float8_e4m3fnuz + attn_bmm_output = torch.bmm( + attn_output.to(torch.bfloat16).transpose(0, 1), + self.w_vc.to(torch.bfloat16) * self.w_scale, + ) + elif self.w_vc.dtype == torch.float8_e4m3fn: + attn_output_val, attn_output_scale = per_tensor_quant_mla_fp8( + attn_output.transpose(0, 1), + zero_allocator.allocate(1), + dtype=torch.float8_e4m3fn, + ) + attn_bmm_output = bmm_fp8( + attn_output_val, + self.w_vc, + attn_output_scale, + self.w_scale, + torch.bfloat16, + ) + else: + attn_bmm_output = torch.bmm(attn_output.transpose(0, 1), self.w_vc) + attn_output = attn_bmm_output.transpose(0, 1).flatten(1, 2) + output, _ = self.o_proj(attn_output) + + return output + + def _chunked_prefix_attn_mha( + self, + q: torch.Tensor, + accum_output: torch.Tensor, + accum_lse: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + + assert forward_batch.num_prefix_chunks is not None + for i in range(forward_batch.num_prefix_chunks): + forward_batch.set_prefix_chunk_idx(i) + + # Fetch latent cache from memory pool with precomputed chunked kv indices + latent_cache_buf = forward_batch.token_to_kv_pool.get_key_buffer( + self.attn_mha.layer_id + ) + latent_cache = latent_cache_buf[ + forward_batch.prefix_chunk_kv_indices[i] + ].contiguous() + + kv_a_normed, k_pe = latent_cache.split( + [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + kv_a_normed = kv_a_normed.squeeze(1).contiguous() + kv = self.kv_b_proj(kv_a_normed)[0] + kv = kv.view( + -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim + ) + v = kv[..., self.qk_nope_head_dim :] + k_nope = kv[..., : self.qk_nope_head_dim] + + k = torch.empty( + ( + k_nope.shape[0], + self.num_local_heads, + self.qk_nope_head_dim + self.qk_rope_head_dim, + ), + dtype=v.dtype, + device=v.device, + ) + k[..., : self.qk_nope_head_dim] = k_nope + k[..., self.qk_nope_head_dim :] = k_pe + + output, lse = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False) + lse = torch.transpose(lse, 0, 1).contiguous() + tmp_output = torch.empty_like(accum_output) + tmp_lse = torch.empty_like(accum_lse) + merge_state_v2(output, lse, accum_output, accum_lse, tmp_output, tmp_lse) + accum_output, accum_lse = tmp_output, tmp_lse + + return accum_output + + def forward_normal_chunked_kv( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + # In normal mha, the k and v tensors will become overly large when the prefix length is long. + # To avoid this, we split the kv cache into chunks and process them one after another. + # Since mha is compute friendly, the for loop induced here will not introduce significant overhead. + # The top comments in https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/mla/common.py + # will be helpful for understanding the purpose of this function. + + # First do normal mha forward to get output for extended part + if self.q_lora_rank is not None: + q, latent_cache = self.fused_qkv_a_proj_with_mqa(hidden_states)[0].split( + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1 + ) + q = self.q_a_layernorm(q) + q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) + else: + q = self.q_proj(hidden_states)[0].view( + -1, self.num_local_heads, self.qk_head_dim + ) + latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0] + _, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + latent_cache = latent_cache.unsqueeze(1) + kv_a = self.kv_a_layernorm(kv_a.contiguous()) + kv = self.kv_b_proj(kv_a)[0] + kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim) + k_nope = kv[..., : self.qk_nope_head_dim] + v = kv[..., self.qk_nope_head_dim :] + k_pe = latent_cache[:, :, self.kv_lora_rank :] + + q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) + q[..., self.qk_nope_head_dim :] = q_pe + k = torch.empty_like(q) + k[..., : self.qk_nope_head_dim] = k_nope + k[..., self.qk_nope_head_dim :] = k_pe + + latent_cache[:, :, : self.kv_lora_rank] = kv_a.unsqueeze(1) + latent_cache[:, :, self.kv_lora_rank :] = k_pe + + # Save latent cache + forward_batch.token_to_kv_pool.set_kv_buffer( + self.attn_mha, forward_batch.out_cache_loc, latent_cache, None + ) + + # Do mha for extended part without prefix + forward_batch.set_attn_attend_prefix_cache(False) + attn_output, lse = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False) + lse = torch.transpose(lse, 0, 1).contiguous() + + # Do mha attention with chunked prefix cache if there are any sequence with prefix + if any(forward_batch.extend_prefix_lens_cpu): + # Only initialize the info once + if forward_batch.num_prefix_chunks is None: + forward_batch.prepare_chunked_prefix_cache_info(q.device) + + forward_batch.set_attn_attend_prefix_cache(True) + attn_output = self._chunked_prefix_attn_mha( + q=q, + accum_output=attn_output, + accum_lse=lse, + forward_batch=forward_batch, + ) + + attn_output = attn_output.reshape(-1, self.num_local_heads * self.v_head_dim) + output, _ = self.o_proj(attn_output) + return output + + +class _FFNInputMode(Enum): + # The MLP sublayer requires 1/tp_size tokens as input + SCATTERED = auto() + # The MLP sublayer requires all tokens as input + FULL = auto() + + +@dataclass +class _DecoderLayerInfo: + is_sparse: bool + ffn_input_mode: _FFNInputMode + + +class DeepseekV2DecoderLayer(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + is_nextn: bool = False, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + self.enable_dp_attention = global_server_args_dict["enable_dp_attention"] + self.layer_id = layer_id + self.dp_size = get_attention_dp_size() + self.attn_tp_size = get_attention_tp_size() + self.attn_tp_rank = get_attention_tp_rank() + self.self_attn = DeepseekV2AttentionMLA( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=( + config.q_lora_rank if hasattr(config, "q_lora_rank") else None + ), + kv_lora_rank=config.kv_lora_rank, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + layer_id=layer_id, + reduce_results=False, + prefix=add_prefix("self_attn", prefix), + ) + + self.info = self._compute_info(config, layer_id=layer_id, is_nextn=is_nextn) + previous_layer_info = self._compute_info( + config, layer_id=layer_id - 1, is_nextn=False + ) + + if self.info.is_sparse: + self.mlp = DeepseekV2MoE( + config=config, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + else: + if self._enable_moe_dense_fully_dp(): + mlp_tp_rank, mlp_tp_size = 0, 1 + else: + mlp_tp_rank, mlp_tp_size = None, None + self.mlp = DeepseekV2MLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + tp_rank=mlp_tp_rank, + tp_size=mlp_tp_size, + ) + + self.input_is_scattered = ( + previous_layer_info.ffn_input_mode == _FFNInputMode.SCATTERED + ) + self.is_last_layer = self.layer_id == config.num_hidden_layers - 1 + + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + @staticmethod + def _enable_moe_dense_fully_dp(): + return global_server_args_dict["moe_dense_tp_size"] == 1 + + @staticmethod + def _compute_info(config: PretrainedConfig, layer_id: int, is_nextn: bool): + is_sparse = is_nextn or ( + config.n_routed_experts is not None + and layer_id >= config.first_k_dense_replace + and layer_id % config.moe_layer_freq == 0 + ) + ffn_input_mode = ( + _FFNInputMode.SCATTERED + if (global_server_args_dict["enable_deepep_moe"] and is_sparse) + or (DeepseekV2DecoderLayer._enable_moe_dense_fully_dp() and not is_sparse) + else _FFNInputMode.FULL + ) + return _DecoderLayerInfo(is_sparse=is_sparse, ffn_input_mode=ffn_input_mode) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + zero_allocator: BumpAllocator, + ) -> torch.Tensor: + if self.info.ffn_input_mode == _FFNInputMode.SCATTERED: + return self.forward_ffn_with_scattered_input( + positions, hidden_states, forward_batch, residual, zero_allocator + ) + elif self.info.ffn_input_mode == _FFNInputMode.FULL: + return self.forward_ffn_with_full_input( + positions, hidden_states, forward_batch, residual, zero_allocator + ) + else: + raise NotImplementedError + + def forward_ffn_with_full_input( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + zero_allocator: BumpAllocator, + ) -> torch.Tensor: + + if hidden_states.shape[0] == 0: + residual = hidden_states + else: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + assert not ( + self.attn_tp_size != 1 and self.input_is_scattered + ), "moe_layer_freq > 1 is not supported when attn_tp_size > 1" + + # Self Attention + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + zero_allocator=zero_allocator, + ) + + # Gather + if get_tensor_model_parallel_world_size() > 1: + # all gather and all reduce + if self.dp_size != 1: + if self.attn_tp_rank == 0: + hidden_states += residual + hidden_states, local_hidden_states = ( + forward_batch.gathered_buffer, + hidden_states, + ) + dp_gather_partial(hidden_states, local_hidden_states, forward_batch) + dp_scatter(residual, hidden_states, forward_batch) + hidden_states = self.post_attention_layernorm(hidden_states) + else: + hidden_states = tensor_model_parallel_all_reduce(hidden_states) + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual + ) + else: + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual + ) + + # Fully Connected + hidden_states = self.mlp(hidden_states) + + # TODO(ch-wan): ues reduce-scatter in MLP to avoid this scatter + # Scatter + if self.dp_size != 1: + # important: forward batch.gathered_buffer is used both after scatter and after gather. + # be careful about this! + hidden_states, global_hidden_states = ( + forward_batch.gathered_buffer[: forward_batch.input_ids.shape[0]], + hidden_states, + ) + dp_scatter(hidden_states, global_hidden_states, forward_batch) + + return hidden_states, residual + + def forward_ffn_with_scattered_input( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + zero_allocator: BumpAllocator, + ) -> torch.Tensor: + + if hidden_states.shape[0] == 0: + residual = hidden_states + else: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + if self.attn_tp_size != 1 and self.input_is_scattered: + hidden_states, local_hidden_states = ( + forward_batch.gathered_buffer[: forward_batch.input_ids.shape[0]], + hidden_states, + ) + tp_all_gather( + list(hidden_states.tensor_split(self.attn_tp_size)), local_hidden_states + ) + + # Self Attention + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + zero_allocator=zero_allocator, + ) + + if self.attn_tp_size != 1: + if self.input_is_scattered: + tensor_list = list(hidden_states.tensor_split(self.attn_tp_size)) + hidden_states = tensor_list[self.attn_tp_rank] + tp_reduce_scatter(hidden_states, tensor_list) + if hidden_states.shape[0] != 0: + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual + ) + else: + if self.attn_tp_rank == 0: + hidden_states += residual + tensor_list = list(hidden_states.tensor_split(self.attn_tp_size)) + hidden_states = tensor_list[self.attn_tp_rank] + tp_reduce_scatter(hidden_states, tensor_list) + residual = hidden_states + if hidden_states.shape[0] != 0: + hidden_states = self.post_attention_layernorm(hidden_states) + else: + if hidden_states.shape[0] != 0: + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual + ) + + if not ( + self._enable_moe_dense_fully_dp() + and (not self.info.is_sparse) + and hidden_states.shape[0] == 0 + ): + hidden_states = self.mlp(hidden_states, forward_batch.forward_mode) + + if self.is_last_layer and self.attn_tp_size != 1: + hidden_states += residual + residual = None + hidden_states, local_hidden_states = ( + forward_batch.gathered_buffer[: forward_batch.input_ids.shape[0]], + hidden_states, + ) + tp_all_gather( + list(hidden_states.tensor_split(self.attn_tp_size)), local_hidden_states + ) + + return hidden_states, residual + + +class DeepseekV2Model(nn.Module): + fall_back_to_pt_during_load = False + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.padding_id = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + enable_tp=not global_server_args_dict["enable_dp_attention"], + ) + self.layers = nn.ModuleList( + [ + DeepseekV2DecoderLayer( + config, + layer_id, + quant_config=quant_config, + prefix=add_prefix(f"layers.{layer_id}", prefix), + ) + for layer_id in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.dp_size = get_attention_dp_size() + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + zero_allocator = BumpAllocator( + # TODO for two-batch-overlap, we need a larger buffer size + buffer_size=len(self.layers) * 2, + dtype=torch.float32, + device=( + input_embeds.device if input_embeds is not None else input_ids.device + ), + ) + + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + + residual = None + for i in range(len(self.layers)): + expert_distribution_recorder.set_current_layer(i) + layer = self.layers[i] + hidden_states, residual = layer( + positions, hidden_states, forward_batch, residual, zero_allocator + ) + if not forward_batch.forward_mode.is_idle(): + if residual is None: + hidden_states = self.norm(hidden_states) + else: + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class DeepseekV2ForCausalLM(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.tp_size = get_tensor_model_parallel_world_size() + self.quant_config = quant_config + self.determine_n_share_experts_fusion() + self.model = DeepseekV2Model( + config, quant_config, prefix=add_prefix("model", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + self.logits_processor = LogitsProcessor(config) + self.dp_size = get_attention_dp_size() + + def determine_n_share_experts_fusion( + self, architecture: str = "DeepseekV3ForCausalLM" + ): + self.n_share_experts_fusion = global_server_args_dict["n_share_experts_fusion"] + if self.n_share_experts_fusion > 0: + # Only Deepseek V3/R1 can use shared experts fusion optimization now. + if ( + self.config.architectures[0] != architecture + or self.config.n_routed_experts != 256 + ): + self.n_share_experts_fusion = 0 + global_server_args_dict["n_share_experts_fusion"] = 0 + logger.info( + "Only Deepseek V3/R1 can use shared experts fusion optimization. Shared experts fusion optimization is disabled." + ) + else: + assert ( + self.n_share_experts_fusion == self.tp_size + ), f"Shared experts fusion optimization is enabled in DeepSeek V3/R1, set it to {self.tp_size} can get best optimized performace." + elif self.n_share_experts_fusion == 0: + if ( + torch.cuda.get_device_capability("cuda") >= (9, 0) + and self.config.architectures[0] == architecture + and self.config.n_routed_experts == 256 + and (not global_server_args_dict["enable_deepep_moe"]) + ): + self.n_share_experts_fusion = self.tp_size + global_server_args_dict["n_share_experts_fusion"] = self.tp_size + logger.info( + "Deepseek V3/R1 with fp8 can use shared experts fusion optimization when SM version >=90. Shared experts fusion optimization is enabled." + ) + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.embed_tokens + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def post_load_weights(self): + + # Perform post-processing after loading weights + for layer_id in range(self.config.num_hidden_layers): + self_attn = self.model.layers[layer_id].self_attn + if hasattr(self_attn.kv_b_proj, "qweight"): + # AWQ compatible + if _is_cuda: + w = awq_dequantize( + self_attn.kv_b_proj.qweight, + self_attn.kv_b_proj.scales, + self_attn.kv_b_proj.qzeros, + ).T + else: + w = awq_dequantize( + self_attn.kv_b_proj.qweight, + self_attn.kv_b_proj.scales, + self_attn.kv_b_proj.qzeros, + 0, + 0, + 0, + ).T + else: + w = self_attn.kv_b_proj.weight + # NOTE(HandH1998): Since `bmm_fp8` only supports per-tensor scale, we have to requantize `self_attn.kv_b_proj`. + # This may affect the accuracy of fp8 model. + # Fix deepseek v3 blockwise bmm by using deep_gemm + use_deep_gemm_bmm = False + model_dtype = torch.get_default_dtype() + + if w.dtype in ( + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + ): + if hasattr(self.quant_config, "weight_block_size"): + weight_block_size = self.quant_config.weight_block_size + if weight_block_size is not None: + assert hasattr(self_attn.kv_b_proj, "weight_scale_inv") + if _is_hip: + weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=w, + weight_scale=self_attn.kv_b_proj.weight_scale_inv, + input_scale=None, + ) + else: + weight = w + weight_scale = self_attn.kv_b_proj.weight_scale_inv + + if ( + _is_cuda + and _ENABLE_JIT_DEEPGEMM + and weight_block_size[0] == 128 + and weight_block_size[1] == 128 + and model_dtype == torch.bfloat16 + ): + block_scale = weight_scale + use_deep_gemm_bmm = True + else: + w, scale = block_quant_to_tensor_quant( + weight, weight_scale, weight_block_size + ) + self_attn.w_scale = scale + else: + weight = w + weight_scale = self_attn.kv_b_proj.weight_scale + w, scale = channel_quant_to_tensor_quant(weight, weight_scale) + self_attn.w_scale = scale + + if w.dtype == torch.int8: + if hasattr(self.quant_config, "weight_block_size"): + # block-wise int8 need it + weight_block_size = self.quant_config.weight_block_size + if weight_block_size is not None: + assert hasattr(self_attn.kv_b_proj, "weight_scale_inv") + weight = w + weight_scale = self_attn.kv_b_proj.weight_scale_inv + w = int8_block_dequant( + weight, weight_scale, weight_block_size + ).to(torch.bfloat16) + else: + # channel-wise int8 need it + w = w.to(torch.bfloat16) * self_attn.kv_b_proj.weight_scale.to( + torch.bfloat16 + ) + + w_kc, w_vc = w.unflatten( + 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim) + ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1) + if not use_deep_gemm_bmm: + self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2) + self_attn.w_vc = w_vc.contiguous().transpose(1, 2) + if ( + hasattr(self_attn.kv_b_proj, "weight_scale") + and self_attn.w_scale is None + ): + self_attn.w_scale = self_attn.kv_b_proj.weight_scale + if _is_hip: + self_attn.w_scale *= 2.0 + else: + num_tiles_k = self_attn.qk_nope_head_dim // weight_block_size[1] + num_tiles_n = self_attn.v_head_dim // weight_block_size[0] + ws_kc, ws_vc = block_scale.unflatten( + 0, (-1, (num_tiles_k + num_tiles_n)) + ).split([num_tiles_k, num_tiles_n], dim=1) + self_attn.w_scale_k = ws_kc.transpose(1, 2).contiguous() + self_attn.w_scale_v = ws_vc.contiguous() + self_attn.w_kc = w_kc.transpose(1, 2).contiguous() + self_attn.w_vc = w_vc.contiguous() + self_attn.use_deep_gemm_bmm = True + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + if self.n_share_experts_fusion > 0: + weights_list = list(weights) + weights_dict = dict(weights_list) + if self.quant_config is None or self.quant_config.get_name() == "w8a8_int8": + suffix_list = [ + "down_proj.weight", + "down_proj.weight_scale", + "gate_proj.weight", + "gate_proj.weight_scale", + "up_proj.weight", + "up_proj.weight_scale", + ] + else: + suffix_list = [ + "down_proj.weight", + "down_proj.weight_scale_inv", + "gate_proj.weight", + "gate_proj.weight_scale_inv", + "up_proj.weight", + "up_proj.weight_scale_inv", + ] + names_to_remove = [] + for moe_layer in tqdm( + range( + self.config.first_k_dense_replace, + self.config.num_hidden_layers, + self.config.moe_layer_freq, + ), + desc=f"Cloning {self.n_share_experts_fusion} " + "replicas of the shared expert into MoE", + ): + for suffix in suffix_list: + shared_expert_weight_name = ( + f"model.layers.{moe_layer}.mlp.shared_experts.{suffix}" + ) + for num_repeat in range(self.n_share_experts_fusion): + weights_list.append( + ( + f"model.layers.{moe_layer}." + f"mlp.experts." + f"{self.config.n_routed_experts + num_repeat}" + f".{suffix}", + weights_dict[shared_expert_weight_name], + ) + ) + names_to_remove += [shared_expert_weight_name] + weights = [w for w in weights_list if w[0] not in names_to_remove] + + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + MoEImpl = ( + DeepEPMoE + if global_server_args_dict["enable_deepep_moe"] + else (EPMoE if global_server_args_dict["enable_ep_moe"] else FusedMoE) + ) + expert_params_mapping = MoEImpl.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts + self.n_share_experts_fusion, + ) + + # Fuse q_a_proj and kv_a_proj_with_mqa along output dimension when q_lora_rank is not None + fuse_qkv_a_proj = hasattr(self.config, "q_lora_rank") and ( + self.config.q_lora_rank is not None + ) + cached_a_proj = {} if fuse_qkv_a_proj else None + + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + # TODO(HandH1998): Modify it when nextn is supported. + if hasattr(self.config, "num_nextn_predict_layers"): + num_nextn_layers = self.config.num_nextn_predict_layers + if num_nextn_layers > 0 and name.startswith("model.layers"): + name_list = name.split(".") + if ( + len(name_list) >= 3 + and int(name_list[2]) >= self.config.num_hidden_layers + ): + continue + if "rotary_emb.inv_freq" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip non-stacked layers and experts (experts handled below). + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since we handle the experts below in expert_params_mapping, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below in expert_params_mapping + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if ("mlp.experts." in name) and name not in params_dict: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if fuse_qkv_a_proj and ( + "q_a_proj" in name or "kv_a_proj_with_mqa" in name + ): + cached_a_proj[name] = loaded_weight + q_a_proj_name = ( + name + if "q_a_proj" in name + else name.replace("kv_a_proj_with_mqa", "q_a_proj") + ) + kv_a_proj_name = ( + name + if "kv_a_proj_with_mqa" in name + else name.replace("q_a_proj", "kv_a_proj_with_mqa") + ) + + # When both q_a_proj and kv_a_proj_with_mqa has been cached, load the fused weight to parameter + if ( + q_a_proj_name in cached_a_proj + and kv_a_proj_name in cached_a_proj + ): + + q_a_proj_weight = cached_a_proj[q_a_proj_name] + kv_a_proj_weight = cached_a_proj[kv_a_proj_name] + fused_weight = torch.cat( + [q_a_proj_weight, kv_a_proj_weight], dim=0 + ) + + param_name = name.replace( + "q_a_proj", "fused_qkv_a_proj_with_mqa" + ) + param = params_dict[param_name] + + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, fused_weight) + cached_a_proj.pop(q_a_proj_name) + cached_a_proj.pop(kv_a_proj_name) + else: + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + self.post_load_weights() + + def get_embed_and_head(self): + return self.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + del self.model.embed_tokens.weight + del self.lm_head.weight + self.model.embed_tokens.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + +class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM): + pass + + +EntryClass = [DeepseekV2ForCausalLM, DeepseekV3ForCausalLM] diff --git a/lib/python3.12/site-packages/sglang/srt/models/exaone.py b/lib/python3.12/site-packages/sglang/srt/models/exaone.py new file mode 100644 index 0000000000000000000000000000000000000000..430c1d58b8f35be1b31dcaadf2d4fc5693a52027 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/exaone.py @@ -0,0 +1,372 @@ +# Copyright 2024 The LGcns AI Engineering Team +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from llama2.py +"""Inference-only Exaone model compatible with HuggingFace weights.""" + +from typing import Any, Dict, Iterable, Optional, Tuple + +import torch +from torch import nn + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix + + +class ExaoneGatedMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.c_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("c_proj", prefix), + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.c_proj(x) + return x + + +class ExaoneAttention(nn.Module): + def __init__( + self, + config, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + rope_theta: float = 500000, + rope_scaling: Optional[Dict[str, Any]] = None, + rope_is_neox_style: bool = True, + max_position_embeddings: int = 4096, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + # MistralConfig has an optional head_dim introduced by Mistral-Nemo + self.head_dim = getattr( + config, "head_dim", self.hidden_size // self.total_num_heads + ) + self.rotary_dim = int( + self.head_dim * getattr(config, "partial_rotary_factor", 1) + ) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.out_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("out_proj", prefix), + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.rotary_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + is_neox_style=rope_is_neox_style, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.out_proj(attn_output) + return output + + +class ExaoneDecoderLayer(nn.Module): + def __init__( + self, + config, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + rope_theta = getattr(config, "rope_theta", 500000) + rope_scaling = getattr(config, "rope_scaling", None) + if rope_scaling is not None and getattr( + config, "original_max_position_embeddings", None + ): + rope_scaling["original_max_position_embeddings"] = ( + config.original_max_position_embeddings + ) + rope_is_neox_style = getattr(config, "rope_is_neox_style", True) + max_position_embeddings = getattr(config, "max_position_embeddings", 4096) + self.self_attn = ExaoneAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + layer_id=layer_id, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + rope_is_neox_style=rope_is_neox_style, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + self.mlp = ExaoneGatedMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.activation_function, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + rms_norm_eps = config.layer_norm_epsilon + self.ln_1 = RMSNorm(config.hidden_size, eps=rms_norm_eps) + self.ln_2 = RMSNorm(config.hidden_size, eps=rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.ln_1(hidden_states) + else: + hidden_states, residual = self.ln_1(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + # Fully Connected + hidden_states, residual = self.ln_2(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class ExaoneModel(nn.Module): + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.wte = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) + self.h = nn.ModuleList( + [ + ExaoneDecoderLayer( + config, + i, + quant_config=quant_config, + prefix=add_prefix(f"h.{i}", prefix), + ) + for i in range(config.num_hidden_layers) + ] + ) + rms_norm_eps = config.layer_norm_epsilon + self.ln_f = RMSNorm(config.hidden_size, eps=rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.wte(input_ids) + else: + hidden_states = input_embeds + residual = None + for i in range(len(self.h)): + layer = self.h[i] + hidden_states, residual = layer( + positions, + hidden_states, + forward_batch, + residual, + ) + hidden_states, _ = self.ln_f(hidden_states, residual) + return hidden_states + + +class ExaoneForCausalLM(nn.Module): + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.transformer = ExaoneModel( + config, quant_config=quant_config, prefix=add_prefix("transformer", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, config.hidden_size, prefix=add_prefix("lm_head", prefix) + ) + self.logits_processor = LogitsProcessor(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> LogitsProcessorOutput: + hidden_states = self.transformer( + input_ids, positions, forward_batch, input_embeds + ) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "c_fc_0", 0), + ("gate_up_proj", "c_fc_1", 1), + ] + params_dict = dict(self.named_parameters()) + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name or "projector" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + if name.startswith("model.vision_tower") and name not in params_dict: + continue + + name = name.replace("attn.attention", "self_attn") + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +EntryClass = ExaoneForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/gemma.py b/lib/python3.12/site-packages/sglang/srt/models/gemma.py new file mode 100644 index 0000000000000000000000000000000000000000..d8074487cb67a61d6d7718f09c65872e6b891e5d --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/gemma.py @@ -0,0 +1,362 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from: +# https://github.com/vllm-project/vllm/blob/c7f2cf2b7f67bce5842fedfdba508440fe257375/vllm/model_executor/models/gemma.py#L1 +"""Inference-only Gemma model compatible with HuggingFace weights.""" + +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.activation import GeluAndMul +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix + + +class GemmaMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + self.act_fn = GeluAndMul("none") + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class GemmaAttention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + layer_id: int = 0, + max_position_embeddings: int = 8192, + rope_theta: float = 10000, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=self.rope_theta, + is_neox_style=True, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class GemmaDecoderLayer(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = GemmaAttention( + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + head_dim=config.head_dim, + layer_id=layer_id, + max_position_embeddings=config.max_position_embeddings, + rope_theta=config.rope_theta, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + self.mlp = GemmaMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class GemmaModel(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) + self.layers = nn.ModuleList( + [ + GemmaDecoderLayer( + config, + i, + quant_config=quant_config, + prefix=add_prefix(f"layers.{i}", prefix), + ) + for i in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + + # Normalize the embedding by sqrt(hidden_size) + hidden_states *= self.config.hidden_size**0.5 + + residual = None + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + forward_batch, + residual, + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class GemmaForCausalLM(nn.Module): + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": [ + "gate_proj", + "up_proj", + ], + } + + # LoRA specific attributes + supported_lora_modules = [ + "qkv_proj", + "o_proj", + "gate_up_proj", + "down_proj", + ] + # Gemma does not apply LoRA to the embedding layer. + embedding_modules = {} + embedding_padding_modules = [] + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.model = GemmaModel( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) + ) + self.logits_processor = LogitsProcessor(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + return self.logits_processor( + input_ids, hidden_states, self.model.embed_tokens, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.named_parameters()) + loaded_params = set() + for name, loaded_weight in weights: + for param_name, shard_name, shard_id in stacked_params_mapping: + if shard_name not in name: + continue + name = name.replace(shard_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # lm_head is not used in vllm as it is tied with embed_token. + # To prevent errors, skip loading lm_head.weight. + if "lm_head.weight" in name: + continue + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # GemmaRMSNorm is different from Llama's in that it multiplies + # (1 + weight) to the output, instead of just weight. + if "norm.weight" in name: + loaded_weight += 1.0 + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + +EntryClass = GemmaForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/gemma2.py b/lib/python3.12/site-packages/sglang/srt/models/gemma2.py new file mode 100644 index 0000000000000000000000000000000000000000..9056b0b0c67b501997182f50082c0a2fe1a6ad1b --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/gemma2.py @@ -0,0 +1,462 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from: +# https://github.com/vllm-project/vllm/blob/56b325e977435af744f8b3dca7af0ca209663558/vllm/model_executor/models/gemma2.py + +from typing import Iterable, Optional, Set, Tuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.activation import GeluAndMul +from sglang.srt.layers.layernorm import GemmaRMSNorm +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from sglang.srt.utils import add_prefix, make_layers + + +# Aligned with HF's implementation, using sliding window inclusive with the last token +# SGLang assumes exclusive +def get_attention_sliding_window_size(config): + return config.sliding_window - 1 + + +class Gemma2MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + hidden_activation: str, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + if not (hidden_act == hidden_activation == "gelu_pytorch_tanh"): + raise ValueError( + "Gemma2 uses `gelu_pytorch_tanh` as the hidden activation " + "function. Please set `hidden_act` and `hidden_activation` to " + "`gelu_pytorch_tanh`." + ) + self.act_fn = GeluAndMul() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class Gemma2Attention(nn.Module): + def __init__( + self, + layer_id: int, + config: PretrainedConfig, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + max_position_embeddings: int, + rope_theta: float, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.layer_id = layer_id + self.config = config + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = config.query_pre_attn_scalar**-0.5 + self.rope_theta = rope_theta + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=config.attention_bias, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=config.attention_bias, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=self.rope_theta, + is_neox_style=True, + ) + + use_sliding_window = layer_id % 2 == 0 and hasattr(config, "sliding_window") + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + logit_cap=self.config.attn_logit_softcapping, + sliding_window_size=( + get_attention_sliding_window_size(config) + if use_sliding_window + else None + ), + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class Gemma2DecoderLayer(nn.Module): + def __init__( + self, + layer_id: int, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Gemma2Attention( + layer_id=layer_id, + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + head_dim=config.head_dim, + max_position_embeddings=config.max_position_embeddings, + rope_theta=config.rope_theta, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + self.hidden_size = config.hidden_size + self.mlp = Gemma2MLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + hidden_activation=config.hidden_activation, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = GemmaRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.pre_feedforward_layernorm = GemmaRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_feedforward_layernorm = GemmaRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + hidden_states = self.post_attention_layernorm(hidden_states) + + hidden_states, residual = self.pre_feedforward_layernorm( + hidden_states, residual + ) + hidden_states = self.mlp(hidden_states) + hidden_states = self.post_feedforward_layernorm(hidden_states) + return hidden_states, residual + + +class Gemma2Model(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) + self.layers = make_layers( + config.num_hidden_layers, + lambda idx, prefix: Gemma2DecoderLayer( + layer_id=idx, + config=config, + quant_config=quant_config, + ), + prefix=add_prefix("layers", prefix), + ) + self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + # Normalize the embedding by sqrt(hidden_size) + # The normalizer's data type should be downcasted to the model's + # data type such as bfloat16, not float32. + # See https://github.com/huggingface/transformers/pull/29402 + normalizer = self.config.hidden_size**0.5 + self.register_buffer("normalizer", torch.tensor(normalizer)) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + normalizer = torch.tensor(self.config.hidden_size**0.5, dtype=torch.float16) + hidden_states *= normalizer + + residual = None + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + forward_batch, + residual, + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class Gemma2ForCausalLM(nn.Module): + # BitandBytes specific attributes + default_bitsandbytes_target_modules = [ + ".gate_proj.", + ".down_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ".o_proj.", + ] + bitsandbytes_stacked_params_mapping = { + # shard_name, weight_name, index + "q_proj": ("qkv_proj", 0), + "k_proj": ("qkv_proj", 1), + "v_proj": ("qkv_proj", 2), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": [ + "gate_proj", + "up_proj", + ], + } + + # LoRA specific attributes + supported_lora_modules = [ + "qkv_proj", + "o_proj", + "gate_up_proj", + "down_proj", + ] + # Gemma does not apply LoRA to the embedding layer. + embedding_modules = {} + embedding_padding_modules = [] + supports_lora = True + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.model = Gemma2Model( + config, quant_config, prefix=add_prefix("model", prefix) + ) + self.logits_processor = LogitsProcessor(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + return self.logits_processor( + input_ids, hidden_states, self.model.embed_tokens, forward_batch + ) + + def get_hidden_dim(self, module_name): + # return input_dim, output_dim + if module_name in ["q_proj", "qkv_proj"]: + return ( + self.config.hidden_size, + self.config.head_dim * self.config.num_attention_heads, + ) + elif module_name in ["o_proj"]: + return ( + self.config.head_dim * self.config.num_attention_heads, + self.config.hidden_size, + ) + elif module_name in ["kv_proj"]: + return ( + self.config.hidden_size, + self.config.head_dim * self.config.num_key_value_heads, + ) + elif module_name == "gate_up_proj": + return self.config.hidden_size, self.config.intermediate_size + elif module_name == "down_proj": + return self.config.intermediate_size, self.config.hidden_size + else: + raise NotImplementedError() + + def get_module_name(self, name): + params_mapping = { + "q_proj": "qkv_proj", + "k_proj": "qkv_proj", + "v_proj": "qkv_proj", + "gate_proj": "gate_up_proj", + "up_proj": "gate_up_proj", + } + return params_mapping.get(name, name) + + def get_attention_sliding_window_size(self): + return get_attention_sliding_window_size(self.config) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.named_parameters()) + loaded_params: Set[str] = set() + for name, loaded_weight in weights: + for param_name, shard_name, shard_id in stacked_params_mapping: + if shard_name not in name: + continue + name = name.replace(shard_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # lm_head is not used in vllm as it is tied with embed_token. + # To prevent errors, skip loading lm_head.weight. + if "lm_head.weight" in name: + continue + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Remapping the name of FP8 kv-scale. + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + +EntryClass = Gemma2ForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/gemma2_reward.py b/lib/python3.12/site-packages/sglang/srt/models/gemma2_reward.py new file mode 100644 index 0000000000000000000000000000000000000000..03bea4d1085568c09293f9aed2ff3fae614c1530 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/gemma2_reward.py @@ -0,0 +1,70 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import Gemma2Config + +from sglang.srt.layers.pooler import EmbeddingPoolerOutput, Pooler, PoolingType +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.models.gemma2 import Gemma2ForCausalLM, Gemma2Model +from sglang.srt.utils import add_prefix + + +class Gemma2ForSequenceClassification(nn.Module): + def __init__( + self, + config: Gemma2Config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.num_labels = config.num_labels + self.model = Gemma2Model( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) + ) + self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) + self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=False) + + self.eos_token_id = config.eos_token_id + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + get_embedding: bool = True, + ) -> EmbeddingPoolerOutput: + assert ( + get_embedding + ), "Gemma2ForSequenceClassification is only used for embedding" + + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + last_token_hidden = self.pooler(hidden_states, forward_batch).embeddings + scores = self.score(last_token_hidden) + + return EmbeddingPoolerOutput(scores) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + Gemma2ForCausalLM.load_weights(self, weights) + + +EntryClass = [Gemma2ForSequenceClassification] diff --git a/lib/python3.12/site-packages/sglang/srt/models/gemma3_causal.py b/lib/python3.12/site-packages/sglang/srt/models/gemma3_causal.py new file mode 100644 index 0000000000000000000000000000000000000000..511d9c7e82b5658c2a9d2415658ad9cf5cfd5310 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/gemma3_causal.py @@ -0,0 +1,695 @@ +# Copyright 2025 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +import copy +from typing import Iterable, Optional, Set, Tuple + +import einops +import torch +import torch.nn.functional as F +from torch import nn +from transformers import ( + ROPE_INIT_FUNCTIONS, + AutoModel, + Gemma3TextConfig, + PretrainedConfig, + PreTrainedModel, +) + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.activation import GeluAndMul +from sglang.srt.layers.layernorm import Gemma3RMSNorm +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import apply_rotary_pos_emb +from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from sglang.srt.utils import add_prefix, make_layers + + +# Aligned with HF's implementation, using sliding window inclusive with the last token +# SGLang assumes exclusive +def get_attention_sliding_window_size(config): + return config.sliding_window - 1 + + +# Adapted from: +# https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/gemma3.py +def extract_layer_index(prefix: str) -> int: + """Extract the layer index from a prefix string.""" + parts = prefix.split(".") + for part in parts: + if part.startswith("layers."): + layer_str = part.split(".")[-1] + try: + return int(layer_str) + except ValueError: + continue + return -1 + + +class Gemma3MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_activation: str, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + if hidden_activation != "gelu_pytorch_tanh": + raise ValueError( + "Gemma3 uses `gelu_pytorch_tanh` as the hidden activation " + "function. Please set `hidden_activation` to " + "`gelu_pytorch_tanh`." + ) + self.act_fn = GeluAndMul() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class Gemma3Attention(nn.Module): + def __init__( + self, + layer_id: int, + config: Gemma3TextConfig, + max_position_embeddings: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.layer_id = layer_id + self.config = config + tp_size = get_tensor_model_parallel_world_size() + + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = config.num_key_value_heads + + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + + hidden_size = config.hidden_size + + head_dim = getattr( + config, "head_dim", hidden_size // config.num_attention_heads + ) + self.head_dim = head_dim + + self.q_size = self.num_heads * self.head_dim + + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = config.query_pre_attn_scalar**-0.5 + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=config.attention_bias, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=config.attention_bias, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + + # Determine if layer uses sliding window based on pattern + self.is_sliding = bool((layer_id + 1) % config.sliding_window_pattern) + + # Initialize the rotary embedding. + if self.is_sliding: + # Local attention. Override the values in config.json. + self.rope_theta = config.rope_local_base_freq + self.rope_scaling = {"rope_type": "default"} + # FIXME(mick): idk why vllm does this + # self.sliding_window = config.interleaved_sliding_window + self.sliding_window = get_attention_sliding_window_size(config) + else: + # Global attention. Use the values in config.json. + self.rope_theta = config.rope_theta + self.rope_scaling = config.rope_scaling + self.sliding_window = None + + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + logit_cap=0.0, + # Module must also define `get_attention_sliding_window_size` to correctly initialize + # attention backend in `ForwardBatch`. + sliding_window_size=self.sliding_window, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + # Gemma3 adds normalization for q and k + self.q_norm = Gemma3RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps) + self.k_norm = Gemma3RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps) + + def naive_attn_with_masks( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + q = q.view(-1, self.num_heads, self.head_dim) + # Expand the key and value to handle GQA. + num_queries_per_kv = self.num_heads // self.num_kv_heads + k = k.view(-1, self.num_kv_heads, self.head_dim) + k = k.repeat_interleave(num_queries_per_kv, dim=-2) + v = v.view(-1, self.num_kv_heads, self.head_dim) + v = v.repeat_interleave(num_queries_per_kv, dim=-2) + + if self.is_sliding: + attn_masks = kwargs["local_attn_masks"] + else: + attn_masks = kwargs["global_attn_masks"] + + seq_lens = kwargs["seq_lens"] + start_idx = 0 + for seq_len, attn_mask in zip(seq_lens, attn_masks): + end_idx = start_idx + seq_len + query = q[start_idx:end_idx].unsqueeze(0) + key = k[start_idx:end_idx].unsqueeze(0) + value = v[start_idx:end_idx].unsqueeze(0) + + # Transpose. + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + + output = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask, + self.scaling, + ) + output = output.transpose(1, 2).flatten(-2, -1) + out[start_idx:end_idx] = output + start_idx = end_idx + return out + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + forward_batch: ForwardBatch, + **kwargs, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + # [s, h * head_dim] + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + # [s, h, head_dim] + q = q.unflatten(-1, (self.num_heads, self.head_dim)) + # -> [h, s, head_dim] + q = q.transpose(0, 1).unsqueeze(0) + q = self.q_norm(q) + k = k.unflatten(-1, (self.num_kv_heads, self.head_dim)) + # -> [h, s, head_dim] + k = k.transpose(0, 1).unsqueeze(0) + k = self.k_norm(k) + + # q, k = self.rotary_emb(positions, q, k) + cos, sin = position_embeddings + q, k = apply_rotary_pos_emb(q, k, cos, sin) + + # [b, h, s, head_dim] -> [b, s, h, head_dim] + q = q.permute(0, 2, 1, 3) + k = k.permute(0, 2, 1, 3) + + attn_output = self.attn(q, k, v, forward_batch=forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class Gemma3DecoderLayer(nn.Module): + def __init__( + self, + layer_id: int, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Gemma3Attention( + layer_id=layer_id, + config=config, + max_position_embeddings=config.max_position_embeddings, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + self.hidden_size = config.hidden_size + self.mlp = Gemma3MLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_activation=config.hidden_activation, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + self.input_layernorm = Gemma3RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = Gemma3RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.pre_feedforward_layernorm = Gemma3RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_feedforward_layernorm = Gemma3RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.is_sliding = self.self_attn.is_sliding + self.layer_id = layer_id + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + position_embeddings_global: torch.Tensor, + position_embeddings_local: torch.Tensor, + forward_batch: ForwardBatch, + **kwargs, + ) -> tuple[ + torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]] + ]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + # apply global RoPE to non-sliding layer only + if self.self_attn.is_sliding: + position_embeddings = position_embeddings_local + else: + position_embeddings = position_embeddings_global + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + position_embeddings=position_embeddings, + forward_batch=forward_batch, + **kwargs, + ) + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.pre_feedforward_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = self.post_feedforward_layernorm(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + return outputs + + +class Gemma3RotaryEmbedding(nn.Module): + def __init__(self, config: Gemma3TextConfig, device=None): + super().__init__() + # BC: "rope_type" was originally "type" + if hasattr(config, "rope_scaling") and config.rope_scaling is not None: + self.rope_type = config.rope_scaling.get( + "rope_type", config.rope_scaling.get("type") + ) + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = self.inv_freq + + def _dynamic_frequency_update(self, position_ids, device): + """ + dynamic RoPE layers should recompute `inv_freq` in the following situations: + 1 - growing beyond the cached sequence length (allow scaling) + 2 - the current sequence length is in the original scale (avoid losing precision with small sequences) + """ + seq_len = torch.max(position_ids) + 1 + if seq_len > self.max_seq_len_cached: # growth + inv_freq, self.attention_scaling = self.rope_init_fn( + self.config, device, seq_len=seq_len + ) + self.register_buffer( + "inv_freq", inv_freq, persistent=False + ) # TODO joao: may break with compilation + self.max_seq_len_cached = seq_len + + if ( + seq_len < self.original_max_seq_len + and self.max_seq_len_cached > self.original_max_seq_len + ): # reset + # This .to() is needed if the model has been moved to a device after being initialized (because + # the buffer is automatically moved, but not the original copy) + self.original_inv_freq = self.original_inv_freq.to(device) + self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) + self.max_seq_len_cached = self.original_max_seq_len + + @torch.no_grad() + def forward(self, x, position_ids): + if "dynamic" in self.rope_type: + self._dynamic_frequency_update(position_ids, device=x.device) + + # Core RoPE block + inv_freq_expanded = ( + self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + ) + position_ids_expanded = position_ids[:, None, :].float() + # Force float32 (see https://github.com/huggingface/transformers/pull/29285) + device_type = x.device.type + device_type = ( + device_type + if isinstance(device_type, str) and device_type != "mps" + else "cpu" + ) + with torch.autocast(device_type=device_type, enabled=False): + freqs = ( + inv_freq_expanded.float().to(x.device) @ position_ids_expanded.float() + ).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + + # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention + cos = cos * self.attention_scaling + sin = sin * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class Gemma3TextScaledWordEmbedding(nn.Embedding): + """ + This module overrides nn.Embeddings' forward by multiplying with embeddings scale. + """ + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + padding_idx: int, + embed_scale: Optional[float] = 1.0, + ): + super().__init__(num_embeddings, embedding_dim, padding_idx) + self.embed_scale = embed_scale + + def forward(self, input_ids: torch.Tensor): + return super().forward(input_ids) * self.embed_scale + + +class Gemma3TextModel(PreTrainedModel): + def __init__( + self, + config: Gemma3TextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__(config=config) + self.config = config + self.quant_config = quant_config + + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + # Gemma3 downcasts the below to float16, causing sqrt(3072)=55.4256 to become 55.5. See https://github.com/huggingface/transformers/pull/29402 + self.embed_tokens = Gemma3TextScaledWordEmbedding( + config.vocab_size, + config.hidden_size, + self.padding_idx, + embed_scale=self.config.hidden_size**0.5, + ) + + self.norm = Gemma3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Gemma3RotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # when we want to create a local RoPE layer. Config defaults should hold values for global RoPE + config = copy.deepcopy(config) + config.rope_theta = config.rope_local_base_freq + config.rope_scaling = {"rope_type": "default"} + self.rotary_emb_local = Gemma3RotaryEmbedding(config=config) + + self.layers = make_layers( + config.num_hidden_layers, + lambda idx, prefix: Gemma3DecoderLayer( + layer_id=idx, + config=config, + quant_config=quant_config, + prefix=prefix, + ), + prefix=add_prefix("layers", prefix), + ) + self.norm = Gemma3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_init() + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + **kwargs, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + + if positions.dim() == 1: + positions = einops.rearrange(positions, "s -> 1 s") + + position_embeddings_global = self.rotary_emb(hidden_states, positions) + position_embeddings_local = self.rotary_emb_local(hidden_states, positions) + for layer in self.layers: + layer_outputs = layer( + positions=positions, + position_embeddings_global=position_embeddings_global, + position_embeddings_local=position_embeddings_local, + hidden_states=hidden_states, + forward_batch=forward_batch, + **kwargs, + ) + hidden_states = layer_outputs[0] + + hidden_states = self.norm(hidden_states) + + return hidden_states + + +class Gemma3ForCausalLM(PreTrainedModel): + config_class = Gemma3TextConfig + + _tied_weights_keys = ["lm_head.weight"] + _tp_plan = {"lm_head": "colwise_rep"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + config_class = Gemma3TextConfig + base_model_prefix = "language_model" + + # BitandBytes specific attributes + default_bitsandbytes_target_modules = [ + ".gate_proj.", + ".down_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ".o_proj.", + ] + bitsandbytes_stacked_params_mapping = { + # shard_name, weight_name, index + "q_proj": ("qkv_proj", 0), + "k_proj": ("qkv_proj", 1), + "v_proj": ("qkv_proj", 2), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": [ + "gate_proj", + "up_proj", + ], + } + + # LoRA specific attributes + supported_lora_modules = [ + "qkv_proj", + "o_proj", + "gate_up_proj", + "down_proj", + ] + # Gemma does not apply LoRA to the embedding layer. + embedding_modules = {} + embedding_padding_modules = [] + supports_lora = True + + def __init__( + self, + config: Gemma3TextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__(config=config) + self.config = config + self.quant_config = quant_config + self.model = Gemma3TextModel( + config, quant_config, prefix=add_prefix("model", prefix) + ) + self.logits_processor = LogitsProcessor(config) + + if self.config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + self.post_init() + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.embed_tokens + + def get_attention_sliding_window_size(self): + return get_attention_sliding_window_size(self.config) + + def dtype(self) -> torch.dtype: + return next(self.parameters()).dtype + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + **kwargs, + ) -> LogitsProcessor: + hidden_states = self.model( + input_ids, positions, forward_batch, input_embeds, **kwargs + ) + + return self.logits_processor( + input_ids, hidden_states, self.model.embed_tokens, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.named_parameters()) + loaded_params: Set[str] = set() + for name, loaded_weight in weights: + for param_name, shard_name, shard_id in stacked_params_mapping: + # if param_name in name: + # print(f"{param_name} is already in {name}") + if shard_name not in name: + continue + name = name.replace(shard_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # lm_head is not used in vllm as it is tied with embed_token. + # To prevent errors, skip loading lm_head.weight. + if "lm_head.weight" in name: + continue + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Remapping the name of FP8 kv-scale. + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + # unloaded_params = params_dict.keys() - loaded_params + # if unloaded_params: + # logger.warning( + # "Some weights are not initialized from checkpoints: %s", unloaded_params + # ) + return loaded_params + + +EntryClass = Gemma3ForCausalLM +AutoModel.register(Gemma3TextConfig, Gemma3ForCausalLM, exist_ok=True) diff --git a/lib/python3.12/site-packages/sglang/srt/models/gemma3_mm.py b/lib/python3.12/site-packages/sglang/srt/models/gemma3_mm.py new file mode 100644 index 0000000000000000000000000000000000000000..92a690b2db0569f42d5b6dfd508a516e1b90d510 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/gemma3_mm.py @@ -0,0 +1,402 @@ +# Copyright 2025 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from: +# https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/gemma3_mm.py + +import logging +from functools import lru_cache +from typing import Dict, Iterable, List, Optional, Set, Tuple, TypedDict + +import torch +from torch import nn +from transformers import AutoModel, Gemma3Config, PreTrainedModel + +from sglang.srt.hf_transformers_utils import get_processor +from sglang.srt.layers.layernorm import Gemma3RMSNorm +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.managers.mm_utils import ( + MultiModalityDataPaddingPatternTokenPairs, + general_mm_embed_routine, +) +from sglang.srt.managers.schedule_batch import ( + MultimodalDataItem, + MultimodalInputs, + flatten_nested_list, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from sglang.srt.models.gemma3_causal import Gemma3ForCausalLM +from sglang.srt.utils import add_prefix + +logger = logging.getLogger(__name__) + +cached_get_processor = lru_cache(get_processor) + + +class Gemma3ImagePixelInputs(TypedDict): + pixel_values: torch.Tensor + """Shape: `(batch_size * num_images, num_channels, height, width)`""" + + +class Gemma3MultiModalProjector(nn.Module): + """Projector for Gemma3 multimodal.""" + + def __init__(self, config: Gemma3Config): + super().__init__() + + self.mm_input_projection_weight = nn.Parameter( + torch.zeros( + config.vision_config.hidden_size, config.text_config.hidden_size + ) + ) + + self.mm_soft_emb_norm = Gemma3RMSNorm( + config.vision_config.hidden_size, eps=config.vision_config.layer_norm_eps + ) + + self.patches_per_image = int( + config.vision_config.image_size // config.vision_config.patch_size + ) + self.tokens_per_side = int(config.mm_tokens_per_image**0.5) + self.kernel_size = self.patches_per_image // self.tokens_per_side + self.avg_pool = nn.AvgPool2d( + kernel_size=self.kernel_size, stride=self.kernel_size + ) + + def forward(self, vision_outputs: torch.Tensor) -> torch.Tensor: + batch_size, seq_length, hidden_size = vision_outputs.shape + + # Reshape for pooling + reshaped_vision_outputs = vision_outputs.transpose(1, 2) + reshaped_vision_outputs = reshaped_vision_outputs.reshape( + batch_size, hidden_size, self.patches_per_image, self.patches_per_image + ) + reshaped_vision_outputs = reshaped_vision_outputs.contiguous() + + # Apply pooling + pooled_vision_outputs = self.avg_pool(reshaped_vision_outputs) + pooled_vision_outputs = pooled_vision_outputs.flatten(2) + pooled_vision_outputs = pooled_vision_outputs.transpose(1, 2) + + # Apply normalization + normed_vision_outputs = self.mm_soft_emb_norm(pooled_vision_outputs) + + # Project to text embedding space + projected_vision_outputs = torch.matmul( + normed_vision_outputs, self.mm_input_projection_weight + ) + + return projected_vision_outputs.type_as(vision_outputs) + + +class Gemma3ForConditionalGeneration(PreTrainedModel): + config_class = Gemma3Config + """Gemma3 multimodal model for conditional generation.""" + + # BitandBytes specific attributes + default_bitsandbytes_target_modules = [ + ".gate_proj.", + ".down_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ".o_proj.", + ] + bitsandbytes_stacked_params_mapping = { + # shard_name, weight_name, index + "q_proj": ("qkv_proj", 0), + "k_proj": ("qkv_proj", 1), + "v_proj": ("qkv_proj", 2), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": [ + "gate_proj", + "up_proj", + ], + } + + # LoRA specific attributes + supported_lora_modules = [ + "qkv_proj", + "o_proj", + "gate_up_proj", + "down_proj", + ] + # Gemma does not apply LoRA to the embedding layer. + embedding_modules = {} + embedding_padding_modules = [] + supports_lora = True + + def __init__( + self, + config: Gemma3Config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__(config=config) + self.config = config + self.quant_config = quant_config + # Vision components + # TODO: replace with vision attention + # self.vision_tower = SiglipVisionModel( + # config.vision_config, + # quant_config, + # prefix=add_prefix("vision_tower", prefix), + # ) + self.vision_tower = AutoModel.from_config(config=config.vision_config) + self.multi_modal_projector = Gemma3MultiModalProjector(config) + self.vocab_size = config.text_config.vocab_size + + # Text model + self.language_model = Gemma3ForCausalLM( + config.text_config, quant_config, prefix=add_prefix("model", prefix) + ) + if self.language_model.logits_processor.logit_scale: + logit_scale = getattr(config, "logit_scale", 1.0) + self.language_model.logits_processor.logit_scale *= logit_scale + self.post_init() + + def pad_input_ids( + self, input_ids: List[int], image_inputs: MultimodalInputs + ) -> List[int]: + """Pad input IDs with image tokens.""" + # Get special token IDs + im_start_id: int = image_inputs.im_start_id + im_end_id: int = image_inputs.im_end_id + + media_token_pairs = [(im_start_id, im_end_id)] + pattern = MultiModalityDataPaddingPatternTokenPairs(media_token_pairs) + ids = pattern.pad_input_tokens(input_ids, image_inputs) + return ids + + def prepare_attn_masks( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + mask_dtype: torch.dtype, + **kwargs, + ) -> Dict: + """Prepare attention masks for multimodal inputs.""" + kwargs["has_images"] = True + + # Distinguish sequences by position id 0 + start_indices = (positions == 0).cpu().nonzero() + num_seqs = len(start_indices) + seq_lens = [] + + for i in range(num_seqs): + start_idx = start_indices[i].item() + if i < num_seqs - 1: + end_idx = start_indices[i + 1].item() + else: + end_idx = len(input_ids) + seq_lens.append(end_idx - start_idx) + + kwargs["seq_lens"] = seq_lens + + # Create attention masks + global_attn_masks = [] + local_attn_masks = [] + sliding_window = self.config.text_config.interleaved_sliding_window + + start_idx = 0 + for seq_len in seq_lens: + end_idx = start_idx + seq_len + input_token_ids = input_ids[start_idx:end_idx] + start_idx = end_idx + + # Create global causal mask + global_attn_mask = torch.empty( + 1, + 1, + seq_len, + seq_len, + dtype=mask_dtype, + device=input_ids.device, + ) + global_attn_mask.fill_(float("-inf")) + global_attn_mask = global_attn_mask.triu(diagonal=1) + + # Consider bidirectional attention between image tokens + img_mask = torch.zeros_like(global_attn_mask) + img_pos = input_token_ids == self.config.image_token_index + img_mask[:, :, :, img_pos] += 1 + img_mask[:, :, img_pos, :] += 1 + global_attn_mask = torch.where(img_mask == 2, 0, global_attn_mask) + global_attn_masks.append(global_attn_mask) + + # Create local causal mask with sliding window + local_attn_mask = torch.ones_like(global_attn_mask) + local_attn_mask = torch.tril(local_attn_mask, diagonal=-sliding_window) + local_attn_mask = torch.where( + local_attn_mask == 0, global_attn_mask, float("-inf") + ) + local_attn_masks.append(local_attn_mask) + + kwargs["global_attn_masks"] = global_attn_masks + kwargs["local_attn_masks"] = local_attn_masks + return kwargs + + def get_input_embeddings(self) -> nn.Embedding: + return self.language_model.get_input_embeddings() + + def get_attention_sliding_window_size(self): + """ + This value is used to initialize attention backends in `ForwardBatch`. + """ + return self.language_model.get_attention_sliding_window_size() + + def get_image_feature(self, items: List[MultimodalDataItem]): + """ + Projects the last hidden state from the vision model into language model space. + + Returns: + image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). + """ + pixel_values = torch.stack( + flatten_nested_list([item.pixel_values for item in items]), dim=0 + ) + pixel_values = pixel_values.to("cuda") + pixel_values = pixel_values.to(dtype=self.language_model.dtype()) + + vision_outputs = self.vision_tower(pixel_values=pixel_values).last_hidden_state + image_features = self.multi_modal_projector(vision_outputs) + return image_features + + @torch.no_grad() + def forward( + self, + input_ids: torch.LongTensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + **kwargs: object, + ) -> LogitsProcessor: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.text_config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.text_config.vocab_size]`. + + logits_to_keep (`int` or `torch.Tensor`, *optional*): + If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all + `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that + token can save memory, which becomes pretty significant for long sequences or large vocabulary size. + If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension. + This is useful when using packed tensor format (single dimension for batch and sequence length). + + Returns: + + Example: + + ```python + >>> from PIL import Image + >>> import requests + >>> from transformers import AutoProcessor, Gemma3ForConditionalGeneration + + >>> model = Gemma3ForConditionalGeneration.from_pretrained("google/Gemma3-test-224px-hf") + >>> processor = AutoProcessor.from_pretrained("google/Gemma3-test-224px-hf") + + >>> prompt = "answer en Where is the cow standing?" + >>> url = "https://huggingface.co/gv-hf/Gemma3-test-224px-hf/resolve/main/cow_beach_1.png" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> inputs = processor(images=image, text=prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(**inputs, max_length=30) + >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "answer en Where is the cow standing?\nbeach" + ```""" + + # Important: position_ids in Gemma3 are 1-indexed + # This really does cost me sometime + positions += 1 + + # Replace image id with PAD if the image token if OOV, to avoid index-errors + if input_ids is not None and self.config.image_token_index >= self.vocab_size: + special_image_mask = input_ids == self.config.image_token_index + llm_input_ids = input_ids.clone() + llm_input_ids[special_image_mask] = 0 + else: + llm_input_ids = input_ids + + hs = general_mm_embed_routine( + input_ids=llm_input_ids, + forward_batch=forward_batch, + language_model=self.language_model, + image_data_embedding_func=self.get_image_feature, + positions=positions, + ) + + return hs + + def tie_weights(self): + return self.language_model.tie_weights() + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + """Load weights for the model.""" + params_dict = dict(self.named_parameters()) + loaded_params: Set[str] = set() + + for name, loaded_weight in weights: + if "language_model" in name: + # Gemma3ForCausalLM.load_weights(self, [(name.replace("language_model.", ""), loaded_weight)]) + causal_loaded_params = Gemma3ForCausalLM.load_weights( + self, [(name, loaded_weight)] + ) + loaded_params.update(causal_loaded_params) + continue + else: + # Skip lm_head.weight as it's tied with embed_tokens + if "lm_head.weight" in name: + continue + + # Skip loading extra bias for GPTQ models + if name.endswith(".bias") and name not in params_dict: + continue + + # Remapping the name of FP8 kv-scale + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + unloaded_params = params_dict.keys() - loaded_params + if unloaded_params: + pass + # raise RuntimeError( + # f"Some weights are not initialized from checkpoints: {unloaded_params}") + return loaded_params + + +EntryClass = Gemma3ForConditionalGeneration + +AutoModel.register(Gemma3Config, Gemma3ForConditionalGeneration, exist_ok=True) diff --git a/lib/python3.12/site-packages/sglang/srt/models/gpt2.py b/lib/python3.12/site-packages/sglang/srt/models/gpt2.py new file mode 100644 index 0000000000000000000000000000000000000000..1ec33406f47d4f3a8aec357145efc39e1059a33c --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/gpt2.py @@ -0,0 +1,288 @@ +# coding=utf-8 +# Adapted from +# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/gpt2/modeling_gpt2.py +# Copyright 2023 The vLLM team. +# Copyright 2018 The OpenAI Team Authors and HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Inference-only GPT-2 model compatible with HuggingFace weights.""" +from typing import Iterable, Optional, Tuple, Type + +import torch +from torch import nn +from transformers import GPT2Config + +from sglang.srt.distributed.parallel_state import get_tensor_model_parallel_world_size +from sglang.srt.layers.activation import NewGELU +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix + + +class GPT2Attention(nn.Module): + + def __init__( + self, + layer_id: int, + config: GPT2Config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.hidden_size = config.hidden_size + total_num_heads = config.num_attention_heads + tensor_model_parallel_world_size = get_tensor_model_parallel_world_size() + assert total_num_heads % tensor_model_parallel_world_size == 0 + self.num_heads = total_num_heads // tensor_model_parallel_world_size + self.head_dim = self.hidden_size // total_num_heads + self.scale = self.head_dim**-0.5 + + self.c_attn = QKVParallelLinear( + self.hidden_size, + self.head_dim, + total_num_heads, + bias=True, + quant_config=quant_config, + prefix=add_prefix("c_attn", prefix), + ) + self.c_proj = RowParallelLinear( + self.hidden_size, + self.hidden_size, + bias=True, + quant_config=quant_config, + prefix=add_prefix("c_proj", prefix), + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + scaling=self.scale, + num_kv_heads=total_num_heads, + layer_id=layer_id, + quant_config=quant_config, + ) + + def forward( + self, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.c_attn(hidden_states) + q, k, v = qkv.chunk(chunks=3, dim=-1) + attn_output = self.attn(q, k, v, forward_batch) + attn_output, _ = self.c_proj(attn_output) + return attn_output + + +class GPT2MLP(nn.Module): + + def __init__( + self, + intermediate_size: int, + config: GPT2Config, + act_layer: Type[nn.Module] = NewGELU, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + hidden_size = config.hidden_size + self.c_fc = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=True, + quant_config=quant_config, + prefix=add_prefix("c_fc", prefix), + ) + self.c_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=True, + quant_config=quant_config, + prefix=add_prefix("c_proj", prefix), + ) + self.act = act_layer() + + def forward( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + hidden_states, _ = self.c_fc(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states, _ = self.c_proj(hidden_states) + return hidden_states + + +class GPT2Block(nn.Module): + + def __init__( + self, + layer_id: int, + config: GPT2Config, + act_layer: Type[nn.Module] = NewGELU, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + hidden_size = config.hidden_size + inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size + + self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon) + self.attn = GPT2Attention( + layer_id, config, quant_config, prefix=add_prefix("attn", prefix) + ) + self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon) + self.mlp = GPT2MLP( + inner_dim, + config, + act_layer=act_layer, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + + def forward( + self, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.ln_1(hidden_states) + attn_output = self.attn( + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + # residual connection + hidden_states = attn_output + residual + + residual = hidden_states + hidden_states = self.ln_2(hidden_states) + feed_forward_hidden_states = self.mlp(hidden_states) + # residual connection + hidden_states = residual + feed_forward_hidden_states + return hidden_states + + +class GPT2Model(nn.Module): + + def __init__( + self, + config: GPT2Config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + assert not config.add_cross_attention + assert not config.scale_attn_by_inverse_layer_idx + assert not config.reorder_and_upcast_attn + self.embed_dim = config.hidden_size + self.wte = VocabParallelEmbedding(config.vocab_size, self.embed_dim) + self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim) + self.h = nn.ModuleList( + [ + GPT2Block( + i, + config, + quant_config=quant_config, + prefix=add_prefix(f"h.{i}", prefix), + ) + for i in range(config.num_hidden_layers) + ] + ) + + self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + inputs_embeds = self.wte(input_ids) + position_embeds = self.wpe(position_ids) + hidden_states = inputs_embeds + position_embeds + + for i in range(len(self.h)): + layer = self.h[i] + hidden_states = layer(hidden_states, forward_batch) + + hidden_states = self.ln_f(hidden_states) + return hidden_states + + +class GPT2LMHeadModel(nn.Module): + + def __init__( + self, + config: GPT2Config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.quant_config = quant_config + self.transformer = GPT2Model( + config, quant_config, prefix=add_prefix("transformer", prefix) + ) + self.lm_head = self.transformer.wte + + self.logits_processor = LogitsProcessor(config) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + hidden_states = self.transformer(input_ids, positions, forward_batch) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + params_dict = dict(self.named_parameters(remove_duplicate=False)) + for name, loaded_weight in weights: + if "lm_head.weight" in name: + # GPT-2 ties the weights of the embedding layer and the final + # linear layer. + continue + if ".attn.bias" in name or ".attn.masked_bias" in name: + # Skip attention mask. + # NOTE: "c_attn.bias" should not be skipped. + continue + if not name.startswith("transformer."): + name = "transformer." + name + + param = params_dict[name] + # The HF's GPT-2 implementation uses Conv1D instead of Linear. + # Because of this, we need to transpose the weights. + # Note(zhuohan): the logic below might break quantized models. + for conv1d_weight_name in ["c_attn", "c_proj", "c_fc"]: + if conv1d_weight_name not in name: + continue + if not name.endswith(".weight"): + continue + loaded_weight = loaded_weight.t() + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +EntryClass = GPT2LMHeadModel diff --git a/lib/python3.12/site-packages/sglang/srt/models/gpt_bigcode.py b/lib/python3.12/site-packages/sglang/srt/models/gpt_bigcode.py new file mode 100644 index 0000000000000000000000000000000000000000..f49a5d304198233adc1077f36dc0446d4880b5eb --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/gpt_bigcode.py @@ -0,0 +1,304 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from: +# https://github.com/vllm-project/vllm/blob/07eb6f19f3b0ee9f7adf6eb689607028aa40bfd5/vllm/model_executor/models/gpt_bigcode.py +"""Inference-only GPTBigCode model compatible with HuggingFace weights.""" + +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import GPTBigCodeConfig + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.activation import get_act_fn +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix + + +class GPTBigCodeAttention(nn.Module): + + def __init__( + self, + layer_id: int, + config: GPTBigCodeConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.hidden_size = config.hidden_size + total_num_heads = config.num_attention_heads + self.tensor_model_parallel_world_size = get_tensor_model_parallel_world_size() + assert total_num_heads % self.tensor_model_parallel_world_size == 0 + self.num_heads = total_num_heads // self.tensor_model_parallel_world_size + self.head_dim = self.hidden_size // total_num_heads + self.scale = self.head_dim**-0.5 + + self.multi_query = config.multi_query + if self.multi_query: + total_num_kv_heads = 1 + self.num_kv_heads = 1 + else: + total_num_kv_heads = total_num_heads + self.num_kv_heads = self.num_heads + self.kv_dim = self.head_dim * self.num_kv_heads + self.c_attn = QKVParallelLinear( + self.hidden_size, + self.head_dim, + total_num_heads, + total_num_kv_heads, + bias=True, + quant_config=quant_config, + prefix=add_prefix("c_attn", prefix), + ) + + self.c_proj = RowParallelLinear( + self.hidden_size, + self.hidden_size, + bias=True, + quant_config=quant_config, + prefix=add_prefix("c_proj", prefix), + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + scaling=self.scale, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.c_attn(hidden_states) + q, k, v = qkv.split( + [ + self.hidden_size // self.tensor_model_parallel_world_size, + self.kv_dim, + self.kv_dim, + ], + dim=-1, + ) + attn_output = self.attn(q, k, v, forward_batch) + attn_output, _ = self.c_proj(attn_output) + return attn_output + + +class GPTBigMLP(nn.Module): + + def __init__( + self, + intermediate_size: int, + config: GPTBigCodeConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + hidden_size = config.hidden_size + self.c_fc = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=True, + quant_config=quant_config, + prefix=add_prefix("c_fc", prefix), + ) + self.c_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=True, + quant_config=quant_config, + prefix=add_prefix("c_proj", prefix), + ) + self.act = get_act_fn( + config.activation_function, quant_config, intermediate_size + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, _ = self.c_fc(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states, _ = self.c_proj(hidden_states) + return hidden_states + + +class GPTBigCodeBlock(nn.Module): + + def __init__( + self, + layer_id: int, + config: GPTBigCodeConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + hidden_size = config.hidden_size + inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size + + self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon) + self.attn = GPTBigCodeAttention( + layer_id, config, quant_config, prefix=add_prefix("attn", prefix) + ) + self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon) + self.mlp = GPTBigMLP( + inner_dim, config, quant_config, prefix=add_prefix("mlp", prefix) + ) + + def forward( + self, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.ln_1(hidden_states) + attn_output = self.attn( + hidden_states=hidden_states, forward_batch=forward_batch + ) + # residual connection + hidden_states = attn_output + residual + + residual = hidden_states + hidden_states = self.ln_2(hidden_states) + feed_forward_hidden_states = self.mlp(hidden_states) + # residual connection + hidden_states = residual + feed_forward_hidden_states + return hidden_states + + +class GPTBigCodeModel(nn.Module): + + def __init__( + self, + config: GPTBigCodeConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + assert not config.add_cross_attention + + self.embed_dim = config.hidden_size + lora_vocab = 0 + self.vocab_size = config.vocab_size + lora_vocab + self.wte = VocabParallelEmbedding( + self.vocab_size, + self.embed_dim, + org_num_embeddings=config.vocab_size, + prefix=add_prefix("wte", prefix), + ) + self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim) + self.h = nn.ModuleList( + [ + GPTBigCodeBlock( + i, config, quant_config, prefix=add_prefix(f"h.{i}", prefix) + ) + for i in range(config.num_hidden_layers) + ] + ) + self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + inputs_embeds = self.wte(input_ids) + position_embeds = self.wpe(position_ids) + hidden_states = inputs_embeds + position_embeds + + for i in range(len(self.h)): + layer = self.h[i] + hidden_states = layer(hidden_states, forward_batch) + + hidden_states = self.ln_f(hidden_states) + return hidden_states + + +class GPTBigCodeForCausalLM(nn.Module): + packed_modules_mapping = {"c_attn": ["c_attn"]} + + supported_lora_modules = ["c_fc", "c_proj", "wte", "c_attn"] + + embedding_modules = { + "wte": "input_embeddings", + "lm_head": "output_embeddings", + } + + embedding_padding_modules = [] + + def __init__( + self, + config: GPTBigCodeConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + + self.config = config + + self.quant_config = quant_config + self.transformer = GPTBigCodeModel( + config, quant_config, prefix=add_prefix("transformer", prefix) + ) + self.lm_head = self.transformer.wte + self.unpadded_vocab_size = config.vocab_size + self.logits_processor = LogitsProcessor(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + hidden_states = self.transformer(input_ids, positions, forward_batch) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + params_dict = dict(self.named_parameters(remove_duplicate=False)) + for name, loaded_weight in weights: + if "lm_head.weight" in name: + continue + if ".attn.bias" in name: + # Skip attention mask. + # NOTE: "c_attn.bias" should not be skipped. + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + # TODO (@robertgshaw2-neuralmagic): move to fp8 linear method + if "c_attn.input_scale" in name or "c_attn.weight_scale" in name: + weight_loader(param, loaded_weight, "q") + weight_loader(param, loaded_weight, "k") + weight_loader(param, loaded_weight, "v") + else: + weight_loader(param, loaded_weight) + + +EntryClass = GPTBigCodeForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/granite.py b/lib/python3.12/site-packages/sglang/srt/models/granite.py new file mode 100644 index 0000000000000000000000000000000000000000..26fccc48d14311b2e6b627fea80b1ac2a38157ab --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/granite.py @@ -0,0 +1,530 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from +# https://github.com/vllm-project/vllm/blob/c7f2cf2b7f67bce5842fedfdba508440fe257375/vllm/model_executor/models/llama.py#L1 +"""Inference-only Granite model compatible with HuggingFace weights.""" + +import logging +from typing import Any, Dict, Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import GraniteConfig + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput +from sglang.srt.layers.pooler import Pooler, PoolingType +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix +from sglang.utils import get_exception_traceback + +logger = logging.getLogger(__name__) + + +class GraniteMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class GraniteAttention(nn.Module): + def __init__( + self, + config: GraniteConfig, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + rope_is_neox_style: bool = True, + max_position_embeddings: int = 8192, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + # MistralConfig has an optional head_dim introduced by Mistral-Nemo + self.head_dim = getattr( + config, "head_dim", self.hidden_size // self.total_num_heads + ) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = config.attention_multiplier + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + is_neox_style=rope_is_neox_style, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class GraniteDecoderLayer(nn.Module): + def __init__( + self, + config: GraniteConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.residual_multiplier = config.residual_multiplier + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + if rope_scaling is not None and getattr( + config, "original_max_position_embeddings", None + ): + rope_scaling["original_max_position_embeddings"] = ( + config.original_max_position_embeddings + ) + rope_is_neox_style = getattr(config, "rope_is_neox_style", True) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + self.self_attn = GraniteAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + layer_id=layer_id, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + rope_is_neox_style=rope_is_neox_style, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + self.mlp = GraniteMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = ( + self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + * self.residual_multiplier + ) # multiplier for Maximal Update Parameterization + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) * self.residual_multiplier + return hidden_states, residual + + +class GraniteModel(nn.Module): + def __init__( + self, + config: GraniteConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, config.hidden_size + ) + self.layers = nn.ModuleList( + [ + GraniteDecoderLayer( + config, + i, + quant_config=quant_config, + prefix=add_prefix(f"layers.{i}", prefix), + ) + for i in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + hidden_states *= self.config.embedding_multiplier + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + forward_batch, + residual, + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class GraniteForCausalLM(nn.Module): + def __init__( + self, + config: GraniteConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.model = GraniteModel( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) + ) + # If tie_word_embeddings == True, then input and output embeddings are + # the same tensor. Enforce during object creation so that weights will + # load correctly even if the LM head weights don't have a separate entry + # in the state dict. + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + if self.config.tie_word_embeddings: + self.lm_head.tie_weights(self.model.embed_tokens) + + # Granite logit scaling factors are applied via division, but + # LogitsProcessor expects a multiplicative factor. + if hasattr(config, "logits_scaling"): + logit_scale = 1.0 / config.logits_scaling + else: + logit_scale = None + self.logits_processor = LogitsProcessor(config, logit_scale=logit_scale) + self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) + self.stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + get_embedding: bool = False, + ) -> LogitsProcessorOutput: + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + if not get_embedding: + logits_processor_output: LogitsProcessorOutput = self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + return logits_processor_output + else: + return self.pooler(hidden_states, forward_batch) + + def get_hidden_dim(self, module_name): + # return input_dim, output_dim + if module_name in ["q_proj", "o_proj", "qkv_proj"]: + return self.config.hidden_size, self.config.hidden_size + elif module_name in ["kv_proj"]: + return self.config.hidden_size, self.config.hidden_size // ( + self.config.num_attention_heads // self.config.num_key_value_heads + ) + elif module_name == "gate_up_proj": + return self.config.hidden_size, self.config.intermediate_size + elif module_name == "down_proj": + return self.config.intermediate_size, self.config.hidden_size + else: + raise NotImplementedError() + + def get_module_name(self, name): + params_mapping = { + "q_proj": "qkv_proj", + "k_proj": "qkv_proj", + "v_proj": "qkv_proj", + "gate_proj": "gate_up_proj", + "up_proj": "gate_up_proj", + } + return params_mapping.get(name, name) + + def get_module_name_from_weight_name(self, name): + for param_name, weight_name, shard_id, num_shard in self.stacked_params_mapping: + if weight_name in name: + return ( + name.replace(weight_name, param_name)[: -len(".weight")], + num_shard, + ) + return name[: -len(".weight")], 1 + + def get_num_params(self): + params_dict = dict(self.named_parameters()) + return len(params_dict) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + + params_dict = dict(self.named_parameters()) + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name or "projector" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + if name.startswith("model.vision_tower") and name not in params_dict: + continue + if "lm_head.weight" in name and self.config.tie_word_embeddings: + # Input and output embeddings are tied, so the output embeddings + # may not be present in the checkpoint. We assume that the input + # embeddings are always present in the checkpoint. + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # This block only runs if the preceding for loop doesn't find + # a match for `name` in `stacked_params_mapping`. + + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Skip loading kv_scale from ckpts towards new design. + if name.endswith(".kv_scale") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + def get_weights_by_name( + self, name: str, truncate_size: int = 100, tp_size: int = 1 + ) -> Optional[torch.Tensor]: + """Get the weights of the parameter by its name. Similar to `get_parameter` in Hugging Face. + + Only used for unit test with an unoptimized performance. + For optimized performance, please use torch.save and torch.load. + """ + try: + if name == "lm_head.weight" and self.config.tie_word_embeddings: + logger.info( + "word embedding is tied for this model, return embed_tokens.weight as lm_head.weight." + ) + return ( + self.model.embed_tokens.weight.cpu() + .to(torch.float32) + .numpy() + .tolist()[:truncate_size] + ) + + mapped_name = name + mapped_shard_id = None + for param_name, weight_name, shard_id in self.stacked_params_mapping: + if weight_name in name: + mapped_name = name.replace(weight_name, param_name) + mapped_shard_id = shard_id + break + params_dict = dict(self.named_parameters()) + param = params_dict[mapped_name] + if mapped_shard_id is not None: + if mapped_shard_id in ["q", "k", "v"]: + num_heads = self.config.num_attention_heads // tp_size + num_kv_heads = self.config.num_key_value_heads // tp_size + head_dim = ( + self.config.hidden_size // self.config.num_attention_heads + ) + if mapped_shard_id == "q": + offset = 0 + size = num_heads * head_dim + elif mapped_shard_id == "k": + offset = num_heads * head_dim + size = num_kv_heads * head_dim + elif mapped_shard_id == "v": + offset = (num_heads + num_kv_heads) * head_dim + size = num_kv_heads * head_dim + weight = param.data.narrow(0, offset, size) + elif mapped_shard_id in [0, 1]: + intermediate_size = self.config.intermediate_size + slice_size = intermediate_size // tp_size + if mapped_shard_id == 0: # gate_proj + offset = 0 + size = slice_size + elif mapped_shard_id == 1: # up_proj + offset = slice_size + size = slice_size + + weight = param.data.narrow(0, offset, size) + else: + weight = param.data + else: + weight = param.data + if tp_size > 1 and ("o_proj" in name or "down_proj" in name): + gathered_weights = [torch.zeros_like(weight) for _ in range(tp_size)] + torch.distributed.all_gather(gathered_weights, weight) + weight = torch.cat(gathered_weights, dim=1) + return weight.cpu().to(torch.float32).numpy().tolist()[:truncate_size] + + except Exception: + logger.error( + f"Error getting weights by name {name} in GraniteForCausalLM: {get_exception_traceback()}" + ) + return None + + +EntryClass = [GraniteForCausalLM] diff --git a/lib/python3.12/site-packages/sglang/srt/models/internlm2.py b/lib/python3.12/site-packages/sglang/srt/models/internlm2.py new file mode 100644 index 0000000000000000000000000000000000000000..28f51cef42378f0c5a411af04aa0085dc5e8aea4 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/internlm2.py @@ -0,0 +1,354 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from https://raw.githubusercontent.com/vllm-project/vllm/7f62077af5159c625fe3ad1c812e6c1a2b93ba3b/vllm/model_executor/models/internlm2.py + +from typing import Any, Dict, Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix + + +class InternLM2MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.w2 = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("w2", prefix), + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.w2(x) + return x + + +class InternLM2Attention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + max_position_embeddings: int = 8192, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.wqkv = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=add_prefix("wqkv", prefix), + ) + self.wo = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("wo", prefix), + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + self.num_kv_heads, + layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.wqkv(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.wo(attn_output) + return output + + +class InternLMDecoderLayer(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + self.attention = InternLM2Attention( + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attention", prefix), + ) + self.feed_forward = InternLM2MLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("feed_forward", prefix), + ) + self.attention_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.ffn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.attention_norm(hidden_states) + else: + hidden_states, residual = self.attention_norm(hidden_states, residual) + hidden_states = self.attention( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + # Fully Connected + hidden_states, residual = self.ffn_norm(hidden_states, residual) + hidden_states = self.feed_forward(hidden_states) + return hidden_states, residual + + +class InternLM2Model(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.tok_embeddings = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=add_prefix("tok_embeddings", prefix), + ) + self.layers = nn.ModuleList( + [ + InternLMDecoderLayer( + config, i, quant_config, prefix=add_prefix(f"layers.{i}", prefix) + ) + for i in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.tok_embeddings(input_ids) + else: + hidden_states = input_embeds + residual = None + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + forward_batch, + residual, + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class InternLM2ForCausalLM(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.model = InternLM2Model( + config, quant_config, prefix=add_prefix("model", prefix) + ) + self.output = ParallelLMHead( + config.vocab_size, config.hidden_size, prefix=add_prefix("output", prefix) + ) + self.logits_processor = LogitsProcessor(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + return self.logits_processor( + input_ids, hidden_states, self.output, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "w1", 0), + ("gate_up_proj", "w3", 1), + ] + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + if "wqkv" in name: + config = self.config + kv_groups = config.num_attention_heads // config.num_key_value_heads + head_dim = config.hidden_size // config.num_attention_heads + loaded_weight = loaded_weight.view( + -1, 2 + kv_groups, head_dim, loaded_weight.shape[-1] + ) + wq, wk, wv = torch.split(loaded_weight, [kv_groups, 1, 1], dim=1) + wq = wq.reshape(-1, wq.shape[-1]) + wk = wk.reshape(-1, wk.shape[-1]) + wv = wv.reshape(-1, wv.shape[-1]) + weight_loader = param.weight_loader + weight_loader(param, wq, "q") + weight_loader(param, wk, "k") + weight_loader(param, wv, "v") + else: + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + +EntryClass = InternLM2ForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/internlm2_reward.py b/lib/python3.12/site-packages/sglang/srt/models/internlm2_reward.py new file mode 100644 index 0000000000000000000000000000000000000000..68be8d001e71ff6b887f708b4f8fd8a1c2bc570a --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/internlm2_reward.py @@ -0,0 +1,64 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.layers.pooler import EmbeddingPoolerOutput, Pooler, PoolingType +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.models.internlm2 import InternLM2ForCausalLM, InternLM2Model +from sglang.srt.utils import add_prefix + + +class InternLM2ForRewardModel(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.vocab_size = config.vocab_size + self.model = InternLM2Model( + config, quant_config, prefix=add_prefix("model", prefix) + ) + self.v_head = nn.Linear(config.hidden_size, 1, bias=False) + self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=False) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + get_embedding: bool = True, + ) -> EmbeddingPoolerOutput: + assert get_embedding, "InternLM2ForRewardModel is only used for embedding" + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + last_token_hidden = self.pooler(hidden_states, forward_batch).embeddings + scores = self.v_head(last_token_hidden) + return EmbeddingPoolerOutput(scores) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + return InternLM2ForCausalLM.load_weights(self, weights) + + +EntryClass = InternLM2ForRewardModel diff --git a/lib/python3.12/site-packages/sglang/srt/models/llama4.py b/lib/python3.12/site-packages/sglang/srt/models/llama4.py new file mode 100644 index 0000000000000000000000000000000000000000..95edfa40e3faa62450c4ccf192c8696fdbcc9d7a --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/llama4.py @@ -0,0 +1,486 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from +# https://github.com/vllm-project/vllm/blob/v0.8.3/vllm/model_executor/models/llama4.py +"""Inference-only LLaMA model compatible with HuggingFace weights.""" + +import logging +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +from torch import nn +from transformers import Llama4TextConfig + +from sglang.srt.distributed import ( + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from sglang.srt.layers.dp_attention import ( + dp_gather_partial, + dp_scatter, + get_attention_dp_size, + get_attention_tp_rank, + get_attention_tp_size, +) +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.moe.fused_moe_triton import FusedMoE +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.models.llama import LlamaForCausalLM, LlamaMLP +from sglang.srt.utils import add_prefix, fast_topk, get_compiler_backend, make_layers + +logger = logging.getLogger(__name__) + + +class Llama4MoE(nn.Module): + + @torch.compile(dynamic=True, backend=get_compiler_backend()) + @staticmethod + def custom_routing_function( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + ) -> Tuple[torch.Tensor, torch.Tensor]: + router_scores_aK, router_indices_aK = fast_topk(gating_output, topk, dim=-1) + router_scores_aK = torch.sigmoid(router_scores_aK.float()).to( + hidden_states.dtype + ) + return ( + router_scores_aK.view(-1).reshape(router_scores_aK.shape), + router_indices_aK.to(torch.int32), + ) + + def __init__( + self, + config: Llama4TextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + self.top_k = config.num_experts_per_tok + + intermediate_size_moe = config.intermediate_size + self.router = ReplicatedLinear( + config.hidden_size, + config.num_local_experts, + bias=False, + quant_config=None, + prefix=add_prefix("router", prefix), + ) + + self.experts = FusedMoE( + num_experts=config.num_local_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + custom_routing_function=Llama4MoE.custom_routing_function, + intermediate_size=intermediate_size_moe, + reduce_results=False, + renormalize=False, + quant_config=quant_config, + apply_router_weight_on_input=True, + prefix=add_prefix("experts", prefix), + ) + + self.shared_expert = LlamaMLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size_moe, + hidden_act="silu", + quant_config=quant_config, + prefix=add_prefix("shared_expert", prefix), + reduce_results=False, # We need to do scatter before reduce + ) + + def forward(self, hidden_states): + # router_scores: [num_tokens, num_experts] + router_logits, _ = self.router(hidden_states) + shared_out = self.shared_expert(hidden_states) + routed_out = self.experts( + hidden_states=hidden_states, + router_logits=router_logits, + ) + out_aD = routed_out + shared_out + + if self.tp_size > 1: + out_aD = tensor_model_parallel_all_reduce(out_aD) + + return out_aD + + +class Llama4Attention(nn.Module): + + def __init__( + self, + config: Llama4TextConfig, + layer_id: int, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + max_position_embeddings: int = 8192, + quant_config: Optional[QuantizationConfig] = None, + bias: bool = False, + bias_o_proj: bool = False, + prefix: str = "", + ) -> None: + super().__init__() + self.layer_id = layer_id + self.hidden_size = hidden_size + self.use_rope = int((layer_id + 1) % 4 != 0) + self.use_qk_norm = config.use_qk_norm and self.use_rope + + self.dp_size = get_attention_dp_size() + attn_tp_rank = get_attention_tp_rank() + attn_tp_size = get_attention_tp_size() + + self.total_num_heads = num_heads + assert self.total_num_heads % attn_tp_size == 0 + self.num_heads = self.total_num_heads // attn_tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= attn_tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % attn_tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert attn_tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // attn_tp_size) + self.head_dim = config.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.attn_temperature_tuning = config.attn_temperature_tuning + self.floor_scale = config.floor_scale + self.attn_scale = config.attn_scale + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.n_rep = self.num_heads // self.num_kv_heads + self.qk_norm = ( + RMSNorm( + hidden_size=self.head_dim, + eps=config.rms_norm_eps, + ) + if self.use_qk_norm + else None + ) + self.qkv_proj = QKVParallelLinear( + hidden_size=hidden_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=bias, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=hidden_size, + bias=bias_o_proj, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + reduce_results=False, + ) + is_neox_style = True + is_gguf = quant_config and quant_config.get_name() == "gguf" + if is_gguf and config.model_type in ["llama", "llama4"]: + is_neox_style = False + + self.rotary_emb = ( + get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=int(rope_theta), + rope_scaling=rope_scaling if rope_scaling != "default" else None, + is_neox_style=is_neox_style, + ) + if self.use_rope + else None + ) + + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + prefix=add_prefix("attn", prefix), + use_irope=self.use_rope, + ) + + def _get_attn_scale(self, positions: torch.Tensor) -> torch.Tensor: + floor = torch.floor((positions + 1.0) / self.floor_scale) + attn_scale = torch.log(floor + 1.0) * self.attn_scale + 1.0 + return attn_scale.unsqueeze(-1) + + @torch.compile(dynamic=True, backend=get_compiler_backend()) + def _mul_attn_scale(self, positions, q): + attn_scale = self._get_attn_scale(positions) + return (q * attn_scale).to(q.dtype) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + + qk, v = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1) + + if self.rotary_emb is not None: + q_view, k_view = qk.split([self.q_size, self.kv_size], dim=-1) + q_out_unused, k_out_unused = self.rotary_emb(positions, q_view, k_view) + del q_view, k_view, q_out_unused, k_out_unused + + if self.qk_norm is not None: + # TODO there are still 2 redundant direct_copy_kernel_cuda for this `reshape` and (in attn backend) q.contiguous(), maybe we can fuse them later + qk = qk.reshape(-1, self.head_dim).contiguous().bfloat16() + qk = self.qk_norm(qk).to(torch.bfloat16) + qk = qk.reshape(-1, self.q_size + self.kv_size) + + q, k = qk.split([self.q_size, self.kv_size], dim=-1) + + # We are applying temperature tuning (https://arxiv.org/abs/2501.19399) to NoPE layers, where + # the inference-time temperature tuning function is customized to not affect short context + # while working at very long context + # https://arxiv.org/abs/2501.19399 + if self.attn_temperature_tuning and not self.use_rope: + q = self._mul_attn_scale(positions=positions, q=q) + + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class Llama4DecoderLayer(nn.Module): + def __init__( + self, + config: Llama4TextConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.layer_id = layer_id + self.hidden_size = config.hidden_size + rope_theta = config.rope_theta + rope_scaling = config.rope_scaling + max_position_embeddings = config.max_position_embeddings + self.dp_size = get_attention_dp_size() + self.attn_tp_size = get_attention_tp_size() + self.attn_tp_rank = get_attention_tp_rank() + + self.self_attn = Llama4Attention( + config=config, + layer_id=layer_id, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + bias=False, + bias_o_proj=False, + prefix=add_prefix("self_attn", prefix), + ) + is_moe_layer = (layer_id + 1) % config.interleave_moe_layer_step == 0 + if is_moe_layer: + self.feed_forward = Llama4MoE( + config=config, + quant_config=quant_config, + prefix=add_prefix("feed_forward", prefix), + ) + else: + self.feed_forward = LlamaMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size_mlp, + hidden_act="silu", + quant_config=quant_config, + prefix=add_prefix("feed_forward", prefix), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + if hidden_states.shape[0] == 0: + residual = hidden_states + else: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + # Gather + if get_tensor_model_parallel_world_size() > 1: + # all gather and all reduce + if self.dp_size != 1: + if self.attn_tp_rank == 0: + hidden_states += residual + hidden_states, local_hidden_states = ( + forward_batch.gathered_buffer, + hidden_states, + ) + dp_gather_partial(hidden_states, local_hidden_states, forward_batch) + dp_scatter(residual, hidden_states, forward_batch) + hidden_states = self.post_attention_layernorm(hidden_states) + else: + hidden_states = tensor_model_parallel_all_reduce(hidden_states) + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual + ) + else: + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual + ) + + # Fully Connected + hidden_states = self.feed_forward(hidden_states) + + # TODO(ch-wan): ues reduce-scatter in MLP to avoid this scatter + # Scatter + if self.dp_size != 1: + # important: forward batch.gathered_buffer is used both after scatter and after gather. + # be careful about this! + hidden_states, global_hidden_states = ( + forward_batch.gathered_buffer[: forward_batch.input_ids.shape[0]], + hidden_states, + ) + dp_scatter(hidden_states, global_hidden_states, forward_batch) + + return hidden_states, residual + + +class Llama4Model(nn.Module): + def __init__( + self, + config: Llama4TextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("embed_tokens", prefix), + enable_tp=not global_server_args_dict["enable_dp_attention"], + ) + self.layers = make_layers( + config.num_hidden_layers, + lambda idx, prefix: Llama4DecoderLayer( + config=config, layer_id=idx, quant_config=quant_config, prefix=prefix + ), + prefix=add_prefix("layers", prefix), + ) + + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.layers_to_capture = [] + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[torch.Tensor]]]: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + aux_hidden_states = [] + for i in range(len(self.layers)): + if i in self.layers_to_capture: + aux_hidden_states.append(hidden_states + residual) + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + forward_batch, + residual, + ) + if not forward_batch.forward_mode.is_idle(): + hidden_states, _ = self.norm(hidden_states, residual) + + if len(aux_hidden_states) == 0: + return hidden_states + + return hidden_states, aux_hidden_states + + +class Llama4ForCausalLM(LlamaForCausalLM): + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + def __init__( + self, + config: Llama4TextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__(config, quant_config, prefix) + + def get_input_embeddings(self): + return self.model.embed_tokens + + def _init_model( + self, + config: Llama4TextConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + return Llama4Model(config, quant_config=quant_config, prefix=prefix) + + +EntryClass = [Llama4ForCausalLM] diff --git a/lib/python3.12/site-packages/sglang/srt/models/llama_eagle.py b/lib/python3.12/site-packages/sglang/srt/models/llama_eagle.py new file mode 100644 index 0000000000000000000000000000000000000000..b04d334bd3dfe044f1448e3a8f7795397847295e --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/llama_eagle.py @@ -0,0 +1,146 @@ +""" +Copyright 2023-2024 SGLang Team +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from sglang.srt.utils import add_prefix + +# Adapted from +# https://github.com/SafeAILab/EAGLE/blob/main/eagle/model/cnets.py +"""Inference-only LLaMA-EAGLE model compatible with HuggingFace weights.""" + +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import LlamaConfig + +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.models.llama import LlamaDecoderLayer, LlamaForCausalLM + + +class LlamaDecoderLayer(LlamaDecoderLayer): + def __init__( + self, + config: LlamaConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__(config, layer_id, quant_config, prefix) + + # Skip the input_layernorm + # https://github.com/SafeAILab/EAGLE/blob/35c78f6cdc19a73e05cf5c330b4c358dad970c6a/eagle/model/cnets.py#L427 + if layer_id == 0: + del self.input_layernorm + setattr(self, "input_layernorm", lambda x: x) + + +class LlamaModel(nn.Module): + def __init__( + self, + config: LlamaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.vocab_size = config.vocab_size + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=add_prefix("embed_tokens", prefix), + ) + self.layers = nn.ModuleList( + [ + LlamaDecoderLayer( + config, + i, + quant_config=quant_config, + prefix=add_prefix(f"layers.{i}", prefix), + ) + for i in range(config.num_hidden_layers) + ] + ) + self.fc = torch.nn.Linear(config.hidden_size * 2, config.hidden_size) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + + hidden_states = self.fc( + torch.cat((hidden_states, forward_batch.spec_info.hidden_states), dim=-1) + ) + + residual = None + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + forward_batch, + residual, + ) + return hidden_states + residual + + +class LlamaForCausalLMEagle(LlamaForCausalLM): + def __init__( + self, + config: LlamaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + nn.Module.__init__(self) + self.config = config + self.quant_config = quant_config + self.model = LlamaModel( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) + ) + # Llama 3.2 1B Instruct set tie_word_embeddings to True + # Llama 3.1 8B Instruct set tie_word_embeddings to False + if self.config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + getattr(config, "hot_vocab_size", config.vocab_size), + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + + self.logits_processor = LogitsProcessor(config) + self.capture_aux_hidden_states = False + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + for name, loaded_weight in weights: + if "lm_head" not in name: + name = "model." + name + super().load_weights([(name, loaded_weight)]) + + +EntryClass = [LlamaForCausalLMEagle] diff --git a/lib/python3.12/site-packages/sglang/srt/models/llama_eagle3.py b/lib/python3.12/site-packages/sglang/srt/models/llama_eagle3.py new file mode 100644 index 0000000000000000000000000000000000000000..137a6da56e828628c1468b3fd757349f3b0f1c2e --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/llama_eagle3.py @@ -0,0 +1,196 @@ +""" +Copyright 2023-2024 SGLang Team +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from sglang.srt.utils import add_prefix + +# Adapted from +# https://github.com/SafeAILab/EAGLE/blob/main/eagle/model/cnets.py +"""Inference-only LLaMA-EAGLE model compatible with HuggingFace weights.""" + +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import LlamaConfig + +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import QKVParallelLinear, RowParallelLinear +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.models.llama import LlamaAttention, LlamaDecoderLayer, LlamaForCausalLM + + +class LlamaDecoderLayer(LlamaDecoderLayer): + def __init__( + self, + config: LlamaConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__(config, layer_id, quant_config, prefix) + + # override qkv + self.self_attn.qkv_proj = QKVParallelLinear( + 2 * self.hidden_size, + self.self_attn.head_dim, + self.self_attn.total_num_heads, + self.self_attn.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + + self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + embeds: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + + residual = hidden_states + embeds = self.input_layernorm(embeds) + hidden_states = self.hidden_norm(hidden_states) + + hidden_states = torch.cat([embeds, hidden_states], dim=-1) + # Self Attention + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + + # Fully Connected + hidden_states = self.mlp(hidden_states) + + return hidden_states, residual + + +class LlamaModel(nn.Module): + def __init__( + self, + config: LlamaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.vocab_size = config.vocab_size + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=add_prefix("embed_tokens", prefix), + ) + self.midlayer = LlamaDecoderLayer(config, 0, quant_config, prefix) + if hasattr(config, "target_hidden_size"): + self.fc = torch.nn.Linear(config.target_hidden_size * 3, config.hidden_size) + else: + self.fc = torch.nn.Linear(config.hidden_size * 3, config.hidden_size) + + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + embeds = self.embed_tokens(input_ids) + else: + embeds = input_embeds + + hidden_states = forward_batch.spec_info.hidden_states + if hidden_states.shape[-1] != embeds.shape[-1]: + hidden_states = self.fc(hidden_states) + + residual = None + hidden_states, residual = self.midlayer( + positions, + embeds, + hidden_states, + forward_batch, + residual, + ) + + hidden_states_to_logits, hidden_states_to_aux = self.norm( + hidden_states, residual + ) + + # For draft decode, we capture the hidden state before norm + return hidden_states_to_logits, [hidden_states_to_aux] + + +class LlamaForCausalLMEagle3(LlamaForCausalLM): + def __init__( + self, + config: LlamaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + nn.Module.__init__(self) + self.config = config + self.quant_config = quant_config + + if self.config.num_hidden_layers != 1: + raise ValueError("EAGLE3 currently only supports 1 layer") + + self.model = LlamaModel( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) + ) + # Llama 3.2 1B Instruct set tie_word_embeddings to True + # Llama 3.1 8B Instruct set tie_word_embeddings to False + if self.config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.draft_vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + + self.logits_processor = LogitsProcessor(config) + self.capture_aux_hidden_states = True + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + for name, loaded_weight in weights: + if "d2t" in name: + # d2t stores diffs between draft id and target id + self.hot_token_id = loaded_weight + torch.arange(loaded_weight.shape[0]) + + if "d2t" not in name and "t2d" not in name and "lm_head" not in name: + new_name = f"model.{name}" + super().load_weights([(new_name, loaded_weight)]) + elif "lm_head" in name: + super().load_weights([(name, loaded_weight)]) + + def get_hot_token_id(self): + return self.hot_token_id + + +EntryClass = [LlamaForCausalLMEagle3] diff --git a/lib/python3.12/site-packages/sglang/srt/models/llama_embedding.py b/lib/python3.12/site-packages/sglang/srt/models/llama_embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..ba448f7fcd3bf615330b461c9ae922dfe771395c --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/llama_embedding.py @@ -0,0 +1,87 @@ +from typing import Iterable, Tuple + +import torch +from torch import nn +from transformers import LlamaConfig + +from sglang.srt.layers.pooler import EmbeddingPoolerOutput, Pooler, PoolingType +from sglang.srt.model_executor.model_runner import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.llama import LlamaModel +from sglang.srt.utils import add_prefix + + +class LlamaEmbeddingModel(nn.Module): + def __init__( + self, + config: LlamaConfig, + quant_config=None, + prefix: str = "", + ) -> None: + super().__init__() + self.model = LlamaModel( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) + ) + self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + get_embedding: bool = True, + ) -> EmbeddingPoolerOutput: + assert ( + get_embedding + ), "LlamaEmbeddingModel / MistralModel is only used for embedding" + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + return self.pooler(hidden_states, forward_batch) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.model.named_parameters()) + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name or "projector" in name: + return + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + return + if name.startswith("model.vision_tower") and name not in params_dict: + return + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + return + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +class MistralModel(LlamaEmbeddingModel): + pass + + +EntryClass = [LlamaEmbeddingModel, MistralModel] diff --git a/lib/python3.12/site-packages/sglang/srt/models/llava.py b/lib/python3.12/site-packages/sglang/srt/models/llava.py new file mode 100644 index 0000000000000000000000000000000000000000..f7beda233c10fc0966a8721d514310827e2d8b4e --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/llava.py @@ -0,0 +1,586 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Inference-only LLaVa model compatible with HuggingFace weights.""" + +import math +import re +from typing import Iterable, List, Optional, Tuple + +import numpy as np +import torch +from torch import nn +from transformers import ( + CLIPVisionConfig, + CLIPVisionModel, + LlavaConfig, + MistralConfig, + Qwen2Config, + SiglipVisionModel, +) +from transformers.models.llava.modeling_llava import LlavaMultiModalProjector + +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.managers.schedule_batch import Modality, MultimodalInputs +from sglang.srt.mm_utils import ( + get_anyres_image_grid_shape, + unpad_image, + unpad_image_shape, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.llama import LlamaForCausalLM +from sglang.srt.models.mistral import MistralForCausalLM +from sglang.srt.models.qwen2 import Qwen2ForCausalLM +from sglang.srt.utils import add_prefix, flatten_nested_list + + +class LlavaBaseForCausalLM(nn.Module): + def pad_input_ids(self, input_ids: List[int], image_inputs: MultimodalInputs): + image_sizes = flatten_nested_list( + [item.image_sizes for item in image_inputs.mm_items] + ) + + pad_values = [item.pad_value for item in image_inputs.mm_items] + + # hardcode for spatial_unpad + anyres + if any( + item.modality == Modality.MULTI_IMAGES or item.modality == Modality.VIDEO + for item in image_inputs.mm_items + ): + image_aspect_ratio = "pad" + else: + image_aspect_ratio = "anyres" + offset_list = [] + image_inputs.image_pad_len = [] + for image_idx, image_s in enumerate(image_sizes): + if len(image_sizes) > 16: + # 2x2 pooling with stride 2 + new_image_feature_len = ( + math.ceil(self.image_size / self.patch_size / 2) ** 2 + ) + else: + new_image_feature_len = self.image_feature_len # multi-image + + height = width = self.num_patches_per_side + if "anyres" in image_aspect_ratio: + num_patch_width, num_patch_height = get_anyres_image_grid_shape( + image_s, + self.image_grid_pinpoints, + self.vision_tower.config.image_size, + ) + h = num_patch_height * height + w = num_patch_width * width + new_h, new_w = unpad_image_shape(h, w, image_s) + + if "anyres_max" in self.config.image_aspect_ratio: + matched_anyres_max_num_patches = re.match( + r"anyres_max_(\d+)", self.config.image_aspect_ratio + ) + if matched_anyres_max_num_patches: + max_num_patches = int(matched_anyres_max_num_patches.group(1)) + # times = math.sqrt(h * w / (max_num_patches * unit**2)) + times = math.sqrt( + new_h * new_w / (max_num_patches * self.image_feature_len) + ) + if times > 1.1: + new_h = int(new_h // times) + new_w = int(new_w // times) + new_image_feature_len += new_h * (new_w + 1) + + try: + offset = input_ids.index(self.config.image_token_index) + except ValueError: + offset = 0 + # old_len + pad_len - 1, because we need to remove image_token_id + input_ids = ( + input_ids[:offset] + + [pad_values[image_idx % len(pad_values)]] * new_image_feature_len + + input_ids[offset + 1 :] + ) + offset_list.append(offset) + image_inputs.image_pad_len.append(new_image_feature_len) + + image_inputs.image_offsets = offset_list + return input_ids + + def encode_images(self, pixel_values: torch.Tensor) -> torch.Tensor: + image_outputs = self.vision_tower(pixel_values, output_hidden_states=True) + # NOTE: This is not memory efficient. (output_hidden_states=True) will save all the hidden stated. + + selected_image_feature = image_outputs.hidden_states[self.vision_feature_layer] + if self.vision_feature_select_strategy in ["default", "patch"]: + selected_image_feature = selected_image_feature[:, 1:] + elif self.vision_feature_select_strategy == "full": + selected_image_feature = selected_image_feature + else: + raise ValueError( + f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}" + ) + image_features = self.multi_modal_projector(selected_image_feature) + + return image_features + + @torch.no_grad() + def forward( + self, + input_ids: torch.LongTensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + image_inputs = forward_batch.mm_inputs + + if forward_batch.forward_mode.is_extend(): + # Clamp input ids. This is because the input_ids for the image tokens are + # filled with the hash values of the image for the prefix matching in the radix attention. + # There values are useless because their embeddings will be replaced by vision embeddings anyway. + input_ids.clamp_(min=0, max=self.config.vocab_size - 1) + + # Embed text inputs + input_embeds = self.language_model.model.embed_tokens(input_ids) + + # Got List[List[str]] extend it to List[str] + # The length of the List should be equal to batch size + modalities_list = [] + max_image_offset = [] + for im in image_inputs: + if im: + modalities_list.extend([item.modality for item in im.mm_items]) + if im and im.image_offsets: + max_image_offset.append( + np.max(np.array(im.image_offsets) + np.array(im.image_pad_len)) + ) + else: + max_image_offset.append(-1) + + start_positions = positions[forward_batch.extend_start_loc].cpu().numpy() + need_vision = start_positions <= np.array(max_image_offset) + + if need_vision.any(): + bs = forward_batch.batch_size + pixel_values = flatten_nested_list( + [ + [item.pixel_values for item in image_inputs[i].mm_items] + for i in range(bs) + if need_vision[i] + ] + ) + image_sizes = [ + flatten_nested_list( + [item.image_sizes for item in image_inputs[i].mm_items] + ) + for i in range(bs) + if need_vision[i] + ] + + ########## Encode Image ######## + + if pixel_values[0].ndim == 4: + # llava-hd: BS, num_patch, C=3, H=336, W=336, num_patch obtained from process_images + np.concatenate(pixel_values, axis=0) + # ndim=4 + concat_images = torch.tensor( + np.concatenate(pixel_values, axis=0), + device=self.vision_tower.device, + ) + image_features = self.encode_images(concat_images) + split_sizes = [image.shape[0] for image in pixel_values] + image_features = torch.split(image_features, split_sizes, dim=0) + # hd image_features: BS, num_patch, 576, 4096 + else: + # normal pixel: BS, C=3, H=336, W=336 + pixel_values = torch.tensor( + np.array(pixel_values), device=self.vision_tower.device + ) + image_features = self.encode_images(pixel_values) + # image_features: BS, 576, 4096 + + if self.mm_patch_merge_type.startswith("spatial"): + new_image_features = [] + height = width = self.num_patches_per_side + for image_idx, image_feature in enumerate(image_features): + if modalities_list[image_idx] == Modality.IMAGE: + image_aspect_ratio = ( + self.config.image_aspect_ratio + ) # single image + elif ( + modalities_list[image_idx] == Modality.MULTI_IMAGES + or modalities_list[image_idx] == Modality.VIDEO + ): + image_aspect_ratio = "pad" # multi image + # image_aspect_ratio = ( + # "anyres" if len(image_sizes[image_idx]) == 1 else "pad" + # ) + if ( + image_feature.shape[0] > 1 + and "anyres" in image_aspect_ratio + and modalities_list[image_idx] == Modality.IMAGE + ): + base_image_feature = image_feature[0] + image_feature = image_feature[1:] + assert height * width == base_image_feature.shape[0] + + if "anyres_max" in image_aspect_ratio: + matched_anyres_max_num_patches = re.match( + r"anyres_max_(\d+)", image_aspect_ratio + ) + if matched_anyres_max_num_patches: + max_num_patches = int( + matched_anyres_max_num_patches.group(1) + ) + + if ( + image_aspect_ratio == "anyres" + or "anyres_max" in image_aspect_ratio + ): + vision_tower_image_size = self.image_size + try: + num_patch_width, num_patch_height = ( + get_anyres_image_grid_shape( + image_sizes[image_idx][0], + self.config.image_grid_pinpoints, + vision_tower_image_size, + ) + ) + except Exception as e: + print(f"Error: {e}") + num_patch_width, num_patch_height = 2, 2 + image_feature = image_feature.view( + num_patch_height, num_patch_width, height, width, -1 + ) + else: + image_feature = image_feature.view( + 2, 2, height, width, -1 + ) + + # ( + # num_patch_width, + # num_patch_height, + # ) = get_anyres_image_grid_shape( + # image_sizes[image_idx][0], + # self.image_grid_pinpoints, + # self.vision_tower.config.image_size, + # ) + + # image_feature = image_feature.view( + # num_patch_height, num_patch_width, height, width, -1 + # ) + + if "unpad" in self.mm_patch_merge_type: + unit = image_feature.shape[2] + image_feature = image_feature.permute( + 4, 0, 2, 1, 3 + ).contiguous() + image_feature = image_feature.flatten(1, 2).flatten( + 2, 3 + ) + image_feature = unpad_image( + image_feature, image_sizes[image_idx][0] + ) + if ( + "anyres_max" in image_aspect_ratio + and matched_anyres_max_num_patches + ): + c, h, w = image_feature.shape + times = math.sqrt( + h * w / (max_num_patches * unit**2) + ) + if times > 1.1: + image_feature = image_feature[None] + image_feature = nn.functional.interpolate( + image_feature, + [int(h // times), int(w // times)], + mode="bilinear", + )[0] + image_feature = torch.cat( + ( + image_feature, + self.language_model.model.image_newline[ + :, None, None + ].expand(*image_feature.shape[:-1], 1), + ), + dim=-1, + ) + image_feature = image_feature.flatten(1, 2).transpose( + 0, 1 + ) + else: + image_feature = image_feature.permute( + 0, 2, 1, 3, 4 + ).contiguous() + image_feature = image_feature.flatten(0, 3) + image_feature = torch.cat( + (base_image_feature, image_feature), dim=0 + ) + image_feature = image_feature.unsqueeze(0) + else: + if modalities_list[image_idx] == Modality.VIDEO: # video + # 2x2 pooling + num_of_frames = image_feature.shape[0] + image_feature = image_feature.view( + num_of_frames, height, width, -1 + ) + image_feature = image_feature.permute( + 0, 3, 1, 2 + ).contiguous() # N, C, H, W + height, weight = image_feature.shape[2:] + scaled_shape = [ + math.ceil(height / 2), + math.ceil(weight / 2), + ] + image_feature = nn.functional.interpolate( + image_feature, size=scaled_shape, mode="bilinear" + ) + image_feature = ( + image_feature.flatten(2) + .transpose(1, 2) + .contiguous() + ) # N, C, H*W + if "unpad" in self.mm_patch_merge_type: + image_feature = torch.cat( + ( + image_feature, + # Expand to (bs, 1, hidden_dim) and concat at the end of the image tokens + self.language_model.model.image_newline[ + None, None + ].expand( + image_feature.shape[0], + 1, + image_feature.shape[-1], + ), + ), + dim=1, + ) + + new_image_features.append(image_feature) + image_features = new_image_features + + # Fill in the placeholder for the image + extend_start_loc_cpu = forward_batch.extend_start_loc.cpu().numpy() + extend_seq_lens = forward_batch.extend_seq_lens.cpu().numpy() + prefix_lens_cpu = forward_batch.extend_prefix_lens_cpu + pt = 0 + for i in range(bs): + if not need_vision[i]: + continue + + start_idx = extend_start_loc_cpu[i] + seq_len = extend_seq_lens[i] + prefix_len = prefix_lens_cpu[i] + + # Multiple images + for image_idx, image_offset in enumerate( + image_inputs[i].image_offsets + ): + if ( + image_offset + image_inputs[i].image_pad_len[image_idx] + <= prefix_len + ): + continue + if image_offset >= prefix_len + seq_len: + break + + tmp_image_feature = image_features[pt][image_idx] + pad_len = tmp_image_feature.shape[0] + + input_offset = image_offset - prefix_len + left_idx = start_idx + input_offset + right_idx = left_idx + pad_len + assert right_idx > start_idx + if input_offset < 0: + left_idx = start_idx + tmp_image_feature = tmp_image_feature[-input_offset:] + if right_idx > start_idx + seq_len: + tmp_image_feature = tmp_image_feature[ + : start_idx + seq_len - right_idx + ] + right_idx = start_idx + seq_len + try: + input_embeds[left_idx:right_idx] = tmp_image_feature + except RuntimeError as e: + print(f"RuntimeError in image encoding: {e}") + print(f"{input_embeds.shape=}, {tmp_image_feature.shape=}") + print( + f"{start_idx=}, {image_offset=}, {prefix_len=}, {pad_len=}" + ) + pt += 1 + + return self.language_model( + input_ids, positions, forward_batch, input_embeds=input_embeds + ) + elif forward_batch.forward_mode.is_decode(): + return self.language_model(input_ids, positions, forward_batch) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + # Load clip vision model by cfg['mm_vision_tower']: + # huggingface_name or path_of_clip_relative_to_llava_model_dir + # We put the initialization here instead of __init__ to allow it being reused by other subclasses. + vision_path = self.config.mm_vision_tower + if "clip" in vision_path: + self.vision_tower = CLIPVisionModel.from_pretrained( + vision_path, torch_dtype=torch.float16 + ).cuda() + elif "siglip" in vision_path: + self.vision_tower = SiglipVisionModel.from_pretrained( + vision_path, torch_dtype=torch.float16 + ).cuda() + # Siglip needs all feature tokens + self.config.mm_vision_select_feature = "full" + self.vision_tower.eval() + + self.vision_feature_layer = self.config.mm_vision_select_layer + self.vision_feature_select_strategy = self.config.mm_vision_select_feature + self.image_size = self.vision_tower.config.image_size + self.patch_size = self.vision_tower.config.patch_size + + self.mm_patch_merge_type = getattr(self.config, "mm_patch_merge_type", "flat") + self.image_aspect_ratio = getattr(self.config, "image_aspect_ratio", "square") + self.image_grid_pinpoints = getattr(self.config, "image_grid_pinpoints", None) + + self.image_feature_len = int((self.image_size // self.patch_size) ** 2) + if ( + self.vision_feature_select_strategy == "patch" + or self.vision_feature_select_strategy == "full" + ): + pass + elif self.vision_feature_select_strategy == "cls_patch": + self.image_feature_len += 1 + else: + raise ValueError(f"Unexpected select feature: {self.select_feature}") + + # load mm_projector + projector_weights = { + "model.mm_projector.0": "multi_modal_projector.linear_1", + "model.mm_projector.2": "multi_modal_projector.linear_2", + "model.vision_tower.vision_tower": "vision_tower", + # Update the vision tower weights if we find them in the checkpoint (it may be finetuned). + "model.image_newline": "language_model.model.image_newline", + } + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if "projector" in name or "vision_tower" in name or "image_newline" in name: + for weight_name, param_name in projector_weights.items(): + if weight_name in name: + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + else: + self.language_model.load_weights([(name, loaded_weight)]) + + @property + def num_patches_per_side(self): + return self.image_size // self.patch_size + + +class LlavaLlamaForCausalLM(LlavaBaseForCausalLM): + def __init__( + self, + config: LlavaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + self.vision_tower = None + self.config.vision_config.hidden_size = config.mm_hidden_size + self.config.text_config.hidden_size = config.hidden_size + + self.multi_modal_projector = LlavaMultiModalProjector(config) + self.language_model = LlamaForCausalLM( + config, + quant_config=quant_config, + prefix=add_prefix("language_model", prefix), + ) + if "unpad" in getattr(config, "mm_patch_merge_type", ""): + self.language_model.model.image_newline = nn.Parameter( + torch.empty(config.text_config.hidden_size, dtype=torch.float16) + ) + + +class LlavaQwenForCausalLM(LlavaBaseForCausalLM): + def __init__( + self, + config: LlavaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + self.vision_tower = None + + if getattr(self.config, "vision_config", None) is None: + self.config.vision_config = CLIPVisionConfig(self.config.mm_vision_tower) + if getattr(self.config, "text_config", None) is None: + self.config.text_config = Qwen2Config(self.config._name_or_path) + + self.config.vision_config.hidden_size = config.mm_hidden_size + self.config.text_config.hidden_size = config.hidden_size + + if getattr(self.config, "projector_hidden_act", None) is None: + self.config.projector_hidden_act = "gelu" + if getattr(self.config, "image_token_index", None) is None: + self.config.image_token_index = 151646 + + self.multi_modal_projector = LlavaMultiModalProjector(config) + self.language_model = Qwen2ForCausalLM( + config, + quant_config=quant_config, + prefix=add_prefix("language_model", prefix), + ) + if "unpad" in getattr(config, "mm_patch_merge_type", ""): + self.language_model.model.image_newline = nn.Parameter( + torch.empty(config.text_config.hidden_size, dtype=torch.float16) + ) + + +class LlavaMistralForCausalLM(LlavaBaseForCausalLM): + def __init__( + self, + config: LlavaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + self.vision_tower = None + + if getattr(self.config, "vision_config", None) is None: + self.config.vision_config = CLIPVisionConfig(self.config.mm_vision_tower) + if getattr(self.config, "text_config", None) is None: + self.config.text_config = MistralConfig(self.config._name_or_path) + + self.config.vision_config.hidden_size = config.mm_hidden_size + self.config.text_config.hidden_size = config.hidden_size + + if getattr(self.config, "projector_hidden_act", None) is None: + self.config.projector_hidden_act = "gelu" + if getattr(self.config, "image_token_index", None) is None: + self.config.image_token_index = 32000 + + self.multi_modal_projector = LlavaMultiModalProjector(config) + self.language_model = MistralForCausalLM( + config, + quant_config=quant_config, + prefix=add_prefix("language_model", prefix), + ) + if "unpad" in getattr(config, "mm_patch_merge_type", ""): + self.language_model.model.image_newline = nn.Parameter( + torch.empty(config.text_config.hidden_size, dtype=torch.float16) + ) + + +EntryClass = [LlavaLlamaForCausalLM, LlavaQwenForCausalLM, LlavaMistralForCausalLM] diff --git a/lib/python3.12/site-packages/sglang/srt/models/llavavid.py b/lib/python3.12/site-packages/sglang/srt/models/llavavid.py new file mode 100644 index 0000000000000000000000000000000000000000..fe6ec90554941dc1bf419bb5100545a4481db51a --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/llavavid.py @@ -0,0 +1,283 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Inference-only LLaVa video model compatible with HuggingFace weights.""" + +from typing import Iterable, List, Optional, Tuple + +import numpy as np +import torch +from torch import nn +from transformers import CLIPVisionModel, LlavaConfig +from transformers.models.llava.modeling_llava import LlavaMultiModalProjector + +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.managers.schedule_batch import MultimodalInputs, flatten_nested_list +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.llama import LlamaForCausalLM +from sglang.srt.utils import add_prefix + + +class LlavaVidForCausalLM(nn.Module): + def __init__( + self, + config: LlavaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.vision_tower = None + self.config.vision_config.hidden_size = config.mm_hidden_size + self.config.text_config.hidden_size = config.hidden_size + self.multi_modal_projector = LlavaMultiModalProjector(config) + self.mm_spatial_pool_stride = getattr(self.config, "mm_spatial_pool_stride", 2) + self.resampler = nn.AvgPool2d( + kernel_size=self.mm_spatial_pool_stride, stride=self.mm_spatial_pool_stride + ) + self.language_model = LlamaForCausalLM( + config, + quant_config=quant_config, + prefix=add_prefix("language_model", prefix), + ) + self.num_frames = getattr(self.config, "num_frames", 16) + if "unpad" in getattr(config, "mm_patch_merge_type", ""): + self.language_model.model.image_newline = nn.Parameter( + torch.empty(config.text_config.hidden_size, dtype=torch.float16) + ) + + def pad_input_ids(self, input_ids: List[int], image_inputs: MultimodalInputs): + pad_values = [item.pad_value for item in image_inputs.mm_items] + new_image_feature_len = self.image_feature_len + + pad_ids = pad_values * ( + (new_image_feature_len + len(pad_values)) // len(pad_values) + ) + offset = input_ids.index(self.config.image_token_index) + # old_len + pad_len - 1, because we need to remove image_token_id + new_input_ids = ( + input_ids[:offset] + + pad_ids[:new_image_feature_len] + + input_ids[offset + 1 :] + ) + image_inputs.image_offsets = [offset] + return new_input_ids + + def encode_images(self, pixel_values: torch.Tensor) -> torch.Tensor: + image_outputs = self.vision_tower(pixel_values, output_hidden_states=True) + # NOTE: This is not memory efficient. (output_hidden_states=True) will save all the hidden stated. + + selected_image_feature = image_outputs.hidden_states[self.vision_feature_layer] + if self.vision_feature_select_strategy in ["default", "patch"]: + selected_image_feature = selected_image_feature[:, 1:] + elif self.vision_feature_select_strategy == "full": + selected_image_feature = selected_image_feature + else: + raise ValueError( + f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}" + ) + + height = width = self.num_patches_per_side + num_of_frames = selected_image_feature.shape[0] + selected_image_feature = selected_image_feature.view( + num_of_frames, height, width, -1 + ) + selected_image_feature = selected_image_feature.permute(0, 3, 1, 2).contiguous() + selected_image_feature = ( + self.resampler(selected_image_feature) + .flatten(2) + .transpose(1, 2) + .contiguous() + ) + + image_features = self.multi_modal_projector(selected_image_feature) + + return image_features + + @torch.no_grad() + def forward( + self, + input_ids: torch.LongTensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + image_inputs = forward_batch.mm_inputs + if forward_batch.forward_mode.is_extend(): + bs = forward_batch.batch_size + + # Clamp input ids. See llava.py for more details + input_ids = input_ids.clamp_(min=0, max=self.config.vocab_size - 1) + + # Embed text inputs + input_embeds = self.language_model.model.embed_tokens(input_ids) + + # Whether the requests need vision inputs + max_image_offset = [] + for im in image_inputs: + if im and im.image_offsets: + max_image_offset.append(max(im.image_offsets)) + else: + max_image_offset.append(-1) + start_positions = positions[forward_batch.extend_start_loc].cpu().numpy() + need_vision = start_positions <= np.array(max_image_offset) + + if need_vision.any(): + pixel_values = flatten_nested_list( + [ + [item.pixel_values for item in image_inputs[i].mm_items] + for i in range(bs) + if need_vision[i] + ] + ) + image_offsets = [ + flatten_nested_list( + [item.image_offsets for item in image_inputs[i].mm_items] + ) + for i in range(bs) + if need_vision[i] + ] + + ########## Encode Image ######## + + if pixel_values[0].ndim == 4: + # llava-hd: BS, num_patch, C=3, H=336, W=336, num_patch obtained from process_images + np.concatenate(pixel_values, axis=0) + # ndim=4 + concat_images = torch.tensor( + np.concatenate(pixel_values, axis=0), + device=self.vision_tower.device, + ) + # image_features = self.encode_images(concat_images) + # split_sizes = [image.shape[0] for image in pixel_values] + # image_features = torch.split(image_features, split_sizes, dim=0) + image_features = self.encode_images( + concat_images + ) # , prompts)#, image_counts, long_video=long_video) + split_sizes = [image.shape[0] for image in pixel_values] + image_features = torch.split(image_features, split_sizes, dim=0) + + # hd image_features: BS, num_patch, 576, 4096 + else: + # normal pixel: BS, C=3, H=336, W=336 + pixel_values = torch.tensor( + np.array(pixel_values), device=self.vision_tower.device + ) + image_features = self.encode_images(pixel_values) + # image_features: BS, 576, 4096 + + new_image_features = [] + for image_idx, image_feature in enumerate(image_features): + new_image_features.append(image_feature.flatten(0, 1)) + image_features = new_image_features + + # Fill in the placeholder for the image + extend_start_loc_cpu = forward_batch.extend_start_loc.cpu().numpy() + prefix_lens_cpu = forward_batch.extend_prefix_lens_cpu + pt = 0 + for i in range(bs): + if not need_vision[i]: + continue + + start_idx = extend_start_loc_cpu[i] + prefix_len = prefix_lens_cpu[i] + + # Multiple images + for image_offset in image_offsets[i]: + if image_offset < prefix_len: + continue + + tmp_image_feature = image_features[pt] + pad_len = tmp_image_feature.shape[0] + + left_idx = start_idx + (image_offset - prefix_len) + right_idx = start_idx + (image_offset - prefix_len) + pad_len + try: + input_embeds[left_idx:right_idx] = tmp_image_feature + except RuntimeError as e: + print(f"RuntimeError in image encoding: {e}") + print(f"{input_embeds.shape=}, {tmp_image_feature.shape=}") + print( + f"{start_idx=}, {image_offset=}, {prefix_len=}, {pad_len=}" + ) + pt += 1 + + return self.language_model( + input_ids, positions, forward_batch, input_embeds=input_embeds + ) + elif forward_batch.forward_mode.is_decode(): + return self.language_model(input_ids, positions, forward_batch) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + # Load clip vision model by cfg['mm_vision_tower']: + # huggingface_name or path_of_clip_relative_to_llava_model_dir + # We put the initialization here instead of __init__ to allow it being reused by other subclasses. + vision_path = self.config.mm_vision_tower + self.vision_tower = CLIPVisionModel.from_pretrained( + vision_path, torch_dtype=torch.float16 + ).cuda() + self.vision_tower.eval() + + self.vision_feature_layer = self.config.mm_vision_select_layer + self.vision_feature_select_strategy = self.config.mm_vision_select_feature + self.image_size = self.vision_tower.config.image_size + self.patch_size = self.vision_tower.config.patch_size + + self.mm_patch_merge_type = getattr(self.config, "mm_patch_merge_type", "flat") + self.image_aspect_ratio = getattr(self.config, "image_aspect_ratio", "square") + self.image_grid_pinpoints = getattr(self.config, "image_grid_pinpoints", None) + + print(f"target_frames: {self.num_frames}") + self.image_feature_len = self.num_frames * int( + (self.image_size / self.patch_size / self.mm_spatial_pool_stride) ** 2 + ) + if self.vision_feature_select_strategy == "patch": + pass + elif self.vision_feature_select_strategy == "cls_patch": + self.image_feature_len += 1 + else: + raise ValueError(f"Unexpected select feature: {self.select_feature}") + + # load mm_projector + projector_weights = { + "model.mm_projector.0": "multi_modal_projector.linear_1", + "model.mm_projector.2": "multi_modal_projector.linear_2", + "model.vision_resampler.mm_projector.0": "multi_modal_projector.linear_1", + "model.vision_resampler.mm_projector.2": "multi_modal_projector.linear_2", + "model.vision_tower.vision_tower": "vision_tower", + # Update the vision tower weights if we find them in the checkpoint (it may be finetuned). + "model.image_newline": "language_model.model.image_newline", + } + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + # FIXME: why projector weights read two times? + if "projector" in name or "vision_tower" in name or "image_newline" in name: + for weight_name, param_name in projector_weights.items(): + if weight_name in name: + name = name.replace(weight_name, param_name) + if name in params_dict: + param = params_dict[name] + else: + print(f"Warning: {name} not found in the model") + continue + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + else: + self.language_model.load_weights([(name, loaded_weight)]) + + @property + def num_patches_per_side(self): + return self.image_size // self.patch_size + + +EntryClass = LlavaVidForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/minicpm.py b/lib/python3.12/site-packages/sglang/srt/models/minicpm.py new file mode 100644 index 0000000000000000000000000000000000000000..2df170f3894cc8125cfdde63b74cc281243b4422 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/minicpm.py @@ -0,0 +1,401 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Inference-only MiniCPM model compatible with HuggingFace weights.""" + +import math +from typing import Any, Dict, Iterable, Optional, Tuple + +import torch +from torch import nn + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix + + +class MiniCPMMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class MiniCPMAttention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + max_position_embeddings: int = 8192, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + ) + # set rope as fp32 instead of bf16 + self.rotary_emb.cos_sin_cache = self.rotary_emb._compute_cos_sin_cache() + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + orig_dtype = q.dtype + q, k = q.float(), k.float() + q, k = self.rotary_emb(positions, q, k) + q, k = q.to(orig_dtype), k.to(orig_dtype) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class MiniCPMDecoderLayer(nn.Module): + def __init__( + self, + config, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + self.self_attn = MiniCPMAttention( + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + layer_id=layer_id, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + self.mlp = MiniCPMMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Self Attention + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + hidden_states = residual + hidden_states * ( + self.config.scale_depth / math.sqrt(self.config.num_hidden_layers) + ) + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states * ( + self.config.scale_depth / math.sqrt(self.config.num_hidden_layers) + ) + + return hidden_states, None + + +class MiniCPMModel(nn.Module): + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + prefix=add_prefix("embed_tokens", prefix), + ) + self.layers = nn.ModuleList( + [ + MiniCPMDecoderLayer( + config, + i, + quant_config=quant_config, + prefix=add_prefix(f"layers.{i}", prefix), + ) + for i in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) * self.config.scale_emb + else: + hidden_states = input_embeds + residual = None + + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + forward_batch, + residual, + ) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class MiniCPMForCausalLM(nn.Module): + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + + self.num_experts = getattr(self.config, "num_experts", 0) + self.quant_config = quant_config + self.model = MiniCPMModel( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) + ) + # self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size) + if not self.config.tie_word_embeddings: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + prefix=add_prefix("lm_head", prefix), + ) + + self.scale_width = self.config.hidden_size / self.config.dim_model_base + + self.logits_processor = LogitsProcessor(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is not None: + input_embeds = input_embeds * self.config.scale_emb + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + hidden_states = hidden_states / self.scale_width + if self.config.tie_word_embeddings: + lm_head = self.model.embed_tokens + else: + lm_head = self.lm_head + return self.logits_processor(input_ids, hidden_states, lm_head, forward_batch) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + expert_params_mapping = [ + # (param_name, weight_name, expert_id) + ( + "ws" if weight_name in ["w1", "w3"] else "w2s", + f"experts.{expert_id}.{weight_name}.weight", + expert_id, + ) + for expert_id in range(self.num_experts) + for weight_name in ["w1", "w2", "w3"] + ] + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + if self.config.tie_word_embeddings and "lm_head.weight" in name: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for param_name, weight_name, expert_id in expert_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, loaded_weight, weight_name, expert_id=expert_id + ) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + +EntryClass = MiniCPMForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/minicpm3.py b/lib/python3.12/site-packages/sglang/srt/models/minicpm3.py new file mode 100644 index 0000000000000000000000000000000000000000..1156c3e470dadad52990eb38a4528b1f8bc12695 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/minicpm3.py @@ -0,0 +1,518 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Inference-only MiniCPM3 model compatible with HuggingFace weights.""" + +import math +from typing import Any, Dict, Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix, is_cuda + +if is_cuda(): + from sgl_kernel import bmm_fp8 + + +class MiniCPM3MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +def input_to_float8(x, dtype=torch.float8_e4m3fn): + finfo = torch.finfo(dtype) + min_val, max_val = x.aminmax() + amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12) + scale = finfo.max / amax + x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max) + return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal() + + +class MiniCPM3AttentionMLA(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + hidden_size: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int, + kv_lora_rank: int, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + max_position_embeddings: int = 8192, + quant_config: Optional[QuantizationConfig] = None, + layer_id=None, + prefix: str = "", + ) -> None: + super().__init__() + self.layer_id = layer_id + self.hidden_size = hidden_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.num_heads = num_heads + tp_size = get_tensor_model_parallel_world_size() + assert num_heads % tp_size == 0 + self.num_local_heads = num_heads // tp_size + self.scaling = self.qk_head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + if self.q_lora_rank is not None: + self.q_a_proj = ReplicatedLinear( + self.hidden_size, + self.q_lora_rank, + bias=False, + quant_config=quant_config, + prefix=add_prefix("q_a_proj", prefix), + ) + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) + self.q_b_proj = ColumnParallelLinear( + q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("q_b_proj", prefix), + ) + else: + self.q_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("q_proj", prefix), + ) + + self.kv_a_proj_with_mqa = ReplicatedLinear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("kv_a_proj_with_mqa", prefix), + ) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + quant_config=quant_config, + prefix=add_prefix("kv_b_proj", prefix), + ) + # O projection. + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + self.rotary_emb = get_rope( + qk_rope_head_dim, + rotary_dim=qk_rope_head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + ) + + self.attn = RadixAttention( + self.num_local_heads, + self.kv_lora_rank + self.qk_rope_head_dim, + self.scaling, + num_kv_heads=1, + layer_id=layer_id, + v_head_dim=self.kv_lora_rank, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + self.w_kc = None + self.w_vc = None + self.w_scale = None + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + q_len = hidden_states.shape[0] + q_input = hidden_states.new_empty( + q_len, self.num_local_heads, self.kv_lora_rank + self.qk_rope_head_dim + ) + if self.q_lora_rank is not None: + q = self.q_a_proj(hidden_states)[0] + q = self.q_a_layernorm(q) + q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) + else: + q = self.q_proj(hidden_states)[0].view( + -1, self.num_local_heads, self.qk_head_dim + ) + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + if self.w_kc.dtype == torch.float8_e4m3fn: + q_nope_val, q_nope_scale = input_to_float8( + q_nope.transpose(0, 1), torch.float8_e4m3fn + ) + q_nope_out = bmm_fp8( + q_nope_val, self.w_kc, q_nope_scale, self.w_scale, torch.bfloat16 + ) + else: + q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc) + q_input[..., : self.kv_lora_rank] = q_nope_out.transpose(0, 1) + + latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0] + v_input = latent_cache[..., : self.kv_lora_rank] + v_input = self.kv_a_layernorm(v_input.contiguous()).unsqueeze(1) + k_input = latent_cache.unsqueeze(1) + k_input[..., : self.kv_lora_rank] = v_input + k_pe = k_input[..., self.kv_lora_rank :] + + original_shapes = [q_pe.shape, k_pe.shape] + q_pe, k_pe = self.rotary_emb( + positions, q_pe.reshape(q_pe.shape[0], -1), k_pe.reshape(k_pe.shape[0], -1) + ) + q_pe, k_pe = q_pe.view(original_shapes[0]), k_pe.view(original_shapes[1]) + q_input[..., self.kv_lora_rank :] = q_pe + k_input[..., self.kv_lora_rank :] = k_pe + + attn_output = self.attn(q_input, k_input, v_input, forward_batch) + attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank) + + if self.w_vc.dtype == torch.float8_e4m3fn: + attn_output_val, attn_output_scale = input_to_float8( + attn_output.transpose(0, 1), torch.float8_e4m3fn + ) + attn_bmm_output = bmm_fp8( + attn_output_val, + self.w_vc, + attn_output_scale, + self.w_scale, + torch.bfloat16, + ) + else: + attn_bmm_output = torch.bmm(attn_output.transpose(0, 1), self.w_vc) + attn_output = attn_bmm_output.transpose(0, 1).flatten(1, 2) + output, _ = self.o_proj(attn_output) + + return output + + +class MiniCPM3DecoderLayer(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + self.self_attn = MiniCPM3AttentionMLA( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=self.hidden_size // config.num_attention_heads, + q_lora_rank=( + config.q_lora_rank if hasattr(config, "q_lora_rank") else None + ), + kv_lora_rank=config.kv_lora_rank, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + layer_id=layer_id, + prefix=add_prefix("self_attn", prefix), + ) + + self.mlp = MiniCPM3MLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Self Attention + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + hidden_states = residual + hidden_states * ( + self.config.scale_depth / math.sqrt(self.config.num_hidden_layers) + ) + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states * ( + self.config.scale_depth / math.sqrt(self.config.num_hidden_layers) + ) + + return hidden_states, None + + +class MiniCPM3Model(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + prefix=add_prefix("embed_tokens", prefix), + ) + self.layers = nn.ModuleList( + [ + MiniCPM3DecoderLayer( + config, + i, + quant_config=quant_config, + prefix=add_prefix(f"layers.{i}", prefix), + ) + for i in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) * self.config.scale_emb + else: + hidden_states = input_embeds + residual = None + + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + forward_batch, + residual, + ) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class MiniCPM3ForCausalLM(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + + self.num_experts = getattr(self.config, "num_experts", 0) + self.quant_config = quant_config + self.model = MiniCPM3Model( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) + ) + # self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size) + if not self.config.tie_word_embeddings: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + prefix=add_prefix("lm_head", prefix), + ) + + self.scale_width = self.config.hidden_size / self.config.dim_model_base + + self.logits_processor = LogitsProcessor(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is not None: + input_embeds = input_embeds * self.config.scale_emb + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + hidden_states = hidden_states / self.scale_width + if self.config.tie_word_embeddings: + lm_head = self.model.embed_tokens + else: + lm_head = self.lm_head + return self.logits_processor(input_ids, hidden_states, lm_head, forward_batch) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + expert_params_mapping = [ + # (param_name, weight_name, expert_id) + ( + "ws" if weight_name in ["w1", "w3"] else "w2s", + f"experts.{expert_id}.{weight_name}.weight", + expert_id, + ) + for expert_id in range(self.num_experts) + for weight_name in ["w1", "w2", "w3"] + ] + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + if self.config.tie_word_embeddings and "lm_head.weight" in name: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for param_name, weight_name, expert_id in expert_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, loaded_weight, weight_name, expert_id=expert_id + ) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + for layer_id in range(self.config.num_hidden_layers): + self_attn = self.model.layers[layer_id].self_attn + w_kc, w_vc = self_attn.kv_b_proj.weight.unflatten( + 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim) + ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1) + self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2) + self_attn.w_vc = w_vc.contiguous().transpose(1, 2) + if hasattr(self_attn.kv_b_proj, "weight_scale"): + self_attn.w_scale = self_attn.kv_b_proj.weight_scale + del self_attn.kv_b_proj + + +EntryClass = MiniCPM3ForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/minicpmo.py b/lib/python3.12/site-packages/sglang/srt/models/minicpmo.py new file mode 100644 index 0000000000000000000000000000000000000000..75df1dfae57ab34b07a9bfaecafd3f514a078204 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/minicpmo.py @@ -0,0 +1,1926 @@ +# Copied and adapted from: https://huggingface.co/openbmb/MiniCPM-o-2_6/blob/main/modeling_minicpmo.py + +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Inference-only MiniCPM-o model compatible with HuggingFace weights.""" + +import math +from dataclasses import dataclass +from typing import Any, Iterable, List, Literal, Optional, Tuple, Union + +import numpy as np +import torch +import torch.nn.functional as F +import torch.nn.utils.parametrize as P +import torch.types +from torch import nn +from torch.nn.utils import parametrizations +from tqdm import tqdm +from transformers import LlamaConfig, LlamaModel, PretrainedConfig, PreTrainedModel +from transformers.activations import ACT2FN +from transformers.cache_utils import DynamicCache, EncoderDecoderCache +from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput +from transformers.models.whisper.modeling_whisper import ( + WHISPER_ATTENTION_CLASSES, + WhisperConfig, + WhisperEncoder, +) + +from sglang.srt.layers.quantization import QuantizationConfig +from sglang.srt.managers.mm_utils import ( + MultiModalityDataPaddingPatternTokenPairs, + general_mm_embed_routine, +) +from sglang.srt.managers.schedule_batch import ( + Modality, + MultimodalDataItem, + MultimodalInputs, + flatten_nested_list, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.utils import set_default_torch_dtype +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.minicpmv import ( + Idefics2VisionTransformer, + MiniCPMBaseModel, + Resampler2_5, +) +from sglang.srt.models.qwen2 import Qwen2ForCausalLM +from sglang.srt.utils import logger + +try: + from transformers import LogitsWarper + from vector_quantize_pytorch import GroupedResidualFSQ + from vocos import Vocos + from vocos.pretrained import instantiate_class + + _tts_deps = True +except: + LogitsWarper = None + _tts_deps = False + + +def apply_spk_emb( + input_ids: torch.Tensor = None, + spk_emb: torch.Tensor = None, + input_embeds: torch.Tensor = None, + spk_emb_token_id: int = 0, + num_spk_embs: int = 1, +): + """ + Replace consecutive `num_spk_embs` speaker embedding placeholders in input_embeds with pre-prepared speaker embeddings. This is an in-place replacement, no new tensor is created, so no value is returned. + + Args: + input_ids (torch.Tensor): Input ID tensor, shape [batch_size, seq_len_max] + spk_emb (torch.Tensor): Speaker embedding tensor, shape [batch_size, num_spk_emb, hidden_dim] + input_embeds (torch.Tensor): Input embedding tensor, shape [batch_size, seq_len_max, hidden_dim] + spk_emb_token_id (int): ID of the speaker embedding token + num_spk_embs (int): Number of speaker embeddings + + Returns: + None + """ + + batch_size = input_ids.shape[0] + + for idx in range(batch_size): + input_ids_ = input_ids[idx] # [seq_len_max] + spk_emb_ = spk_emb[idx] # [num_spk_emb] + mask_ = input_ids_ == spk_emb_token_id # [batch_size, seq_len_max] + nonzero_position_idx = mask_.nonzero(as_tuple=False) # [num_spk_emb, 1] + assert nonzero_position_idx.shape[0] == num_spk_embs + begin_idx = nonzero_position_idx.min() + end_idx = nonzero_position_idx.max() + input_embeds[idx, begin_idx : end_idx + 1, :] = spk_emb_ + + return + + +@dataclass +class ConditionalChatTTSGenerationOutput(ModelOutput): + """ + Output class for ConditionalChatTTS generation. + + Args: + new_ids (torch.LongTensor): Newly generated audio code sequence, shape (batch_size, sequence_length, num_vq). + audio_input_ids (torch.LongTensor): Updated input IDs including condition and generated audio codes, shape (batch_size, full_sequence_length, num_vq). + past_key_values (Tuple[Tuple[torch.FloatTensor]]): Tuple containing pre-computed keys and values used for attention mechanism. Each element has shape (batch_size, num_heads, sequence_length, embed_size_per_head). + finished (bool): Boolean indicating whether generation is complete. + + """ + + new_ids: torch.LongTensor = None + audio_input_ids: torch.LongTensor = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + finished: bool = None + + +def make_streaming_chunk_mask_generation( + inputs_embeds: torch.Tensor, + past_seen_tokens: int, + streaming_tts_text_mask: torch.Tensor, + streaming_reserved_length: int = 300, + streaming_audio_chunk_size: int = 50, + streaming_text_chunk_size: int = 10, + num_spk_emb: int = 1, + use_spk_emb: bool = True, +) -> torch.Tensor: + """ + In streaming audio generation, determine which `text` positions the TTS model can attend to when generating each chunk of `audio` tokens. + + This function creates a mask that allows the model to attend to a specific chunk of text + tokens when generating each chunk of audio tokens, enabling streaming TTS generation. + + Args: + inputs_embeds (torch.Tensor): Input embeddings tensor. + past_seen_tokens (int): Number of tokens already seen by the model. + streaming_tts_text_mask (torch.Tensor): Mask for the text tokens. + streaming_reserved_length (int, optional): Number of reserved tokens for streaming. Defaults to 300. + streaming_text_chunk_size (int, optional): Size of each text chunk. Defaults to 7. + + Returns: + torch.Tensor: Causal mask for streaming TTS generation, shape is [batch_size=1, 1, seq_len=1, past_seen_tokens+1] + + Raises: + AssertionError: If the batch size is not 1 (only supports batch size of 1 for inference). + """ + assert inputs_embeds.shape[0] == 1 + + dtype = inputs_embeds.dtype + device = inputs_embeds.device + min_dtype = torch.finfo(dtype).min + + # Add `1` to the past seen tokens to account for new `tokens` during `generate` + causal_mask = torch.full( + (1, past_seen_tokens + inputs_embeds.shape[1]), + fill_value=0, + dtype=dtype, + device=device, + ) + + # Calculate the start of invisible text tokens + invisible_text_tokens_start = ( + min( + math.ceil( + (past_seen_tokens - streaming_reserved_length) + / streaming_audio_chunk_size + ) + * streaming_text_chunk_size, + streaming_reserved_length, + ) + + 1 + + num_spk_emb * use_spk_emb + ) # Add 1 for [Stts] and N for [spk_emb] tokens if `use_spk_emb` is True + + invisible_text_tokens_end = ( + streaming_reserved_length + 1 + num_spk_emb * use_spk_emb + 1 + ) # Add 1 for [Ptts] (aka `audio_bos_token_id`) + + # Set invisible text tokens to min_dtype (effectively -inf) + causal_mask[0, invisible_text_tokens_start:invisible_text_tokens_end] = min_dtype + + # Mask padding positions in the text mask + causal_mask[ + 0, 0 : 1 + num_spk_emb * use_spk_emb + streaming_reserved_length + 1 + ].masked_fill_(streaming_tts_text_mask == 0, min_dtype) + + # Add extra dimensions for batch and heads + causal_mask = causal_mask.unsqueeze(0).unsqueeze(0) + + return causal_mask + + +# Borrowed from `https://github.com/2noise/ChatTTS/blob/main/ChatTTS/model/dvae.py` +class ConvNeXtBlock(nn.Module): + def __init__( + self, + dim: int, + intermediate_dim: int, + kernel: int, + dilation: int, + layer_scale_init_value: float = 1e-6, + ): + # ConvNeXt Block copied from Vocos. + super().__init__() + self.dwconv = nn.Conv1d( + dim, + dim, + kernel_size=kernel, + padding=dilation * (kernel // 2), + dilation=dilation, + groups=dim, + ) + + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, intermediate_dim) + self.act = nn.GELU() + self.pwconv2 = nn.Linear(intermediate_dim, dim) + self.coef = ( + nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True) + if layer_scale_init_value > 0 + else None + ) + + def forward(self, x: torch.Tensor, cond=None) -> torch.Tensor: + residual = x + + y = self.dwconv(x) + y.transpose_(1, 2) # (B, C, T) -> (B, T, C) + x = self.norm(y) + del y + y = self.pwconv1(x) + del x + x = self.act(y) + del y + y = self.pwconv2(x) + del x + if self.coef is not None: + y *= self.coef + y.transpose_(1, 2) # (B, T, C) -> (B, C, T) + + x = y + residual + del y + + return x + + +# Borrowed from `https://github.com/2noise/ChatTTS/blob/main/ChatTTS/model/dvae.py` +class DVAEDecoder(nn.Module): + def __init__( + self, + idim: int, + odim: int, + n_layer=12, + bn_dim=64, + hidden=256, + kernel=7, + dilation=2, + up=False, + ): + super().__init__() + self.up = up + self.conv_in = nn.Sequential( + nn.Conv1d(idim, bn_dim, 3, 1, 1), + nn.GELU(), + nn.Conv1d(bn_dim, hidden, 3, 1, 1), + ) + self.decoder_block = nn.ModuleList( + [ + ConvNeXtBlock( + hidden, + hidden * 4, + kernel, + dilation, + ) + for _ in range(n_layer) + ] + ) + self.conv_out = nn.Conv1d(hidden, odim, kernel_size=1, bias=False) + + def forward(self, x: torch.Tensor, conditioning=None) -> torch.Tensor: + # B, C, T + y = self.conv_in(x) + del x + for f in self.decoder_block: + y = f(y, conditioning) + + x = self.conv_out(y) + del y + return x + + +# Borrowed from `https://github.com/2noise/ChatTTS/blob/main/ChatTTS/model/dvae.py` +class GFSQ(nn.Module): + def __init__( + self, + dim: int, + levels: List[int], + G: int, + R: int, + eps=1e-5, + transpose=True, + ): + super(GFSQ, self).__init__() + self.quantizer = GroupedResidualFSQ( + dim=dim, + levels=list(levels), + num_quantizers=R, + groups=G, + ) + self.n_ind = math.prod(levels) + self.eps = eps + self.transpose = transpose + self.G = G + self.R = R + + def _embed(self, x: torch.Tensor): + if self.transpose: + x = x.transpose(1, 2) + x = x.view(x.size(0), x.size(1), self.G, self.R).permute(2, 0, 1, 3) + feat = self.quantizer.get_output_from_indices(x) + return feat.transpose_(1, 2) if self.transpose else feat + + def __call__(self, x: torch.Tensor) -> torch.Tensor: + return super().__call__(x) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.transpose: + x.transpose_(1, 2) + _, ind = self.quantizer(x) + ind = ind.permute(1, 2, 0, 3).contiguous() + ind = ind.view(ind.size(0), ind.size(1), -1) + return ind.transpose_(1, 2) if self.transpose else ind + + +# Borrowed from `https://github.com/2noise/ChatTTS/blob/main/ChatTTS/model/dvae.py` +class DVAE(nn.Module): + def __init__( + self, + ): + super().__init__() + + coef = torch.rand(100) + self.coef = nn.Parameter(coef.unsqueeze(0).unsqueeze_(2)) + + self.downsample_conv = nn.Sequential( + nn.Conv1d(100, 512, 3, 1, 1), + nn.GELU(), + nn.Conv1d(512, 512, 4, 2, 1), + nn.GELU(), + ) + + self.encoder = DVAEDecoder( + idim=512, + odim=1024, + hidden=256, + n_layer=12, + bn_dim=128, + ) + + self.decoder = DVAEDecoder( + idim=512, + odim=512, + hidden=256, + n_layer=12, + bn_dim=128, + ) + + self.out_conv = nn.Conv1d(512, 100, 3, 1, 1, bias=False) + + self.vq_layer = GFSQ( + dim=1024, + levels=(5, 5, 5, 5), + G=2, + R=2, + ) + + @torch.inference_mode() + def forward( + self, inp: torch.Tensor, mode: Literal["encode", "decode"] = "decode" + ) -> torch.Tensor: + if mode == "encode" and hasattr(self, "encoder") and self.vq_layer is not None: + mel = inp.clone() + x: torch.Tensor = self.downsample_conv( + torch.div(mel, self.coef.view(100, 1).expand(mel.shape), out=mel), + ).unsqueeze_(0) + del mel + x = self.encoder(x) + ind = self.vq_layer(x) + del x + return ind + + if self.vq_layer is not None: + vq_feats = self.vq_layer._embed(inp) + else: + vq_feats = inp + + vq_feats = ( + vq_feats.view( + (vq_feats.size(0), 2, vq_feats.size(1) // 2, vq_feats.size(2)), + ) + .permute(0, 2, 3, 1) + .flatten(2) + ) + + dec_out = self.out_conv( + self.decoder( + x=vq_feats, + ), + ) + + del vq_feats + + return torch.mul(dec_out, self.coef, out=dec_out) + + +# Borrowed from `https://github.com/2noise/ChatTTS/blob/main/ChatTTS/model/processors.py` +class CustomRepetitionPenaltyLogitsProcessorRepeat: + def __init__(self, penalty: float, max_input_ids: int, past_window: int): + if not isinstance(penalty, float) or not (penalty > 0): + raise ValueError( + f"`penalty` has to be a strictly positive float, but is {penalty}" + ) + + self.penalty = penalty + self.max_input_ids = max_input_ids + self.past_window = past_window + + def __call__( + self, input_ids: torch.LongTensor, scores: torch.FloatTensor + ) -> torch.FloatTensor: + if input_ids.size(1) > self.past_window: + input_ids = input_ids.narrow(1, -self.past_window, self.past_window) + freq = F.one_hot(input_ids, scores.size(1)).sum(1) + if freq.size(0) > self.max_input_ids: + freq.narrow( + 0, self.max_input_ids, freq.size(0) - self.max_input_ids + ).zero_() + alpha = torch.pow(self.penalty, freq) + scores = scores.contiguous() + inp = scores.multiply(alpha) + oth = scores.divide(alpha) + con = scores < 0 + out = torch.where(con, inp, oth) + del inp, oth, scores, con, alpha + return out + + +class ConditionalChatTTS(PreTrainedModel): + """A conditional text-to-speech model that can generate speech from text with speaker conditioning. + + This model extends PreTrainedModel to provide text-to-speech capabilities with: + - LLM hidden state conditioning + - Streaming generation + + The model uses a transformer architecture with LLM hidden states and can operate in both + streaming and non-streaming modes for flexible deployment. + + The model process sequence in the following format: + | text bos token | LLM embedding projected to tts embedding space | text tokens (fixed length, reserved for future tokens) | audio bos token | audio tokens (audio token length is not fixed)| audio eos token | + + The format is designed to support LLM-conditioned streaming audio generation. + + Usage: + To support streaming generation, two global variables should be maintained outside of the model. + 1. `audio_input_ids`: stores *discrete* audio codes. It is a tensor with shape [1, sequence length+1, num_vq]. + 2. `past_key_values`: stores the KV cache for both text tokens and audio codes. It is a list of tuples, each tuple contains two tensors with shape [1, num_attention_heads, sequence length, hidden_size // num_attention_heads] + + where `num_vq` is the number of audio codebooks, in default setting, it is `4`. + + 1. Create an empty `past_key_values` with + ```python + initial_kv_cache_length = 1 + model.num_spk_embs + model.streaming_text_reserved_len # where `1` denotes the `bos` token + dtype = model.emb_text.weight.dtype + device = model.emb_text.weight.device + past_key_values = [ + ( + torch.zeros(1, model.config.num_attention_heads, initial_kv_cache_length, model.config.hidden_size // model.config.num_attention_heads, dtype=dtype, device=device), + torch.zeros(1, model.config.num_attention_heads, initial_kv_cache_length, model.config.hidden_size // model.config.num_attention_heads, dtype=dtype, device=device) + ) + for _ in range(model.config.num_hidden_layers) + ] + + 2. At the same time, create an empty `audio_input_ids` with shape [1, sequence length, num_vq], `num_vq` denotes multiple layer audio codebooks. But here we also include text tokens in the sequence, but they will be zeros, and will not be used, just a placeholder. + + ```python + initial_audio_input_ids_length = 1 + model.num_spk_embs + model.streaming_text_reserved_len + 1 + # [bos token, speaker embeddings, text tokens, audio bos token] + audio_input_ids = torch.zeros(batch_size=1, initial_audio_input_ids_length, model.num_vq) + ``` + + 2. Prefill some text tokens to TTS model (for example, 10 tokens) using `prefill_text` method. + + ```python + outputs = llm.generate(**kwargs) + llm_tokens = some_function_to_extract_llm_tokens(outputs) + lm_spk_emb_last_hidden_states = some_function_to_extract_lm_spk_emb_last_hidden_states(outputs) + tts_text_input_ids = tts_tokenizer.encode(llm_tokenizer.decode(llm_tokens)) + # here assume we are prefilling text token 0 to text token 9 (included), totally 10 tokens. + begin = 0 + end = 9+1 + position_ids = torch.arange(begin, end, dtype=torch.long, device=device) + + past_key_values = model.prefill_text( + input_ids=tts_text_input_ids, + position_ids=position_ids, + past_key_values=past_key_values, + lm_spk_emb_last_hidden_states=lm_spk_emb_last_hidden_states, + ) + ``` + + 3. Make a `streaming_tts_text_mask` to denote which position contains valid text tokens, similar to `attention_mask` in standard causal attention. + + ```python + streaming_tts_text_mask = torch.zeros(model.streaming_reserved_length) + streaming_tts_text_mask[0:end] = 1 # denotes these post + ``` + + 3. Generate audio codes using `generate` method. + + ```python + outputs = model.generate( + input_ids=audio_input_ids, + past_key_values=past_key_values, + streaming_tts_text_mask=streaming_tts_text_mask, + max_new_token=50, + ) + + # update past_key_values and input_ids + past_key_values = outputs.past_key_values + audio_input_ids = outputs.input_ids + ``` + + The `past_key_values` is extended by `max_new_token=50`, and `audio_input_ids` is also extended by `max_new_token=50` after `generate` calling. + + 4. Notice that after prefilling `10` text tokens, the model can generate up to `50` audio tokens, if you want to generate more audio tokens, you need to prefill next `10` text tokens. And it is okay to only generate `25` audio tokens for faster initial response. + + 5. Repeat steps `2,3,4` as needed in your streaming audio generation cases, but ensure usage complies with the following guidelines discussed above. + """ + + config_class = PretrainedConfig + _no_split_modules = [] + + def __init__(self, config: PretrainedConfig): + super().__init__(config) + + self.use_speaker_embedding = config.use_speaker_embedding + self.use_llm_hidden_state = config.use_llm_hidden_state + self.num_spk_embs = config.num_spk_embs + self.spk_emb_token_id = config.spk_emb_token_id + + self.use_text = config.use_text + self.streaming = config.streaming + self.streaming_text_chunk_size = config.streaming_text_chunk_size + self.streaming_audio_chunk_size = config.streaming_audio_chunk_size + self.streaming_text_reserved_len = config.streaming_text_reserved_len + self.audio_bos_token_id = config.audio_bos_token_id + self.num_mel_bins = config.num_mel_bins + self.num_vq = config.num_vq + self.num_audio_tokens = config.num_audio_tokens + + self.top_p = config.top_p + self.top_k = config.top_k + self.repetition_penalty = config.repetition_penalty + + if self.config.use_mlp: + self.projector = MultiModalProjector(config.llm_dim, config.hidden_size) + else: + self.projector = nn.Linear(config.llm_dim, config.hidden_size, bias=False) + self.emb_code = nn.ModuleList( + [ + nn.Embedding(config.num_audio_tokens, config.hidden_size) + for _ in range(config.num_vq) + ] + ) + self.emb_text = nn.Embedding(config.num_text_tokens, config.hidden_size) + self.head_code = nn.ModuleList( + [ + parametrizations.weight_norm( + nn.Linear(config.hidden_size, config.num_audio_tokens, bias=False), + name="weight", + ) + for _ in range(config.num_vq) + ] + ) + + dvae = DVAE() + self.dvae = dvae + + model_config = LlamaConfig( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + num_attention_heads=config.num_attention_heads, + num_hidden_layers=config.num_hidden_layers, + max_position_embeddings=config.max_position_embeddings, + attn_implementation=config.attn_implementation, + ) + + model = LlamaModel(model_config) + self.model = model + + @torch.inference_mode() + def merge_inputs_embeds( + self, + input_ids: torch.Tensor, + lm_spk_emb_last_hidden_states: Optional[torch.Tensor] = None, + ): + """Merge `input_ids` and `lm_spk_emb_last_hidden_states` to `inputs_embeds`. + + Args: + input_ids (torch.Tensor): Input token IDs. + lm_spk_emb_last_hidden_states (Optional[torch.Tensor], optional): Last hidden states of speaker embeddings from the language model. Defaults to None. + + Raises: + NotImplementedError: If speaker embedding is not used and language model hidden states are not implemented. + + Returns: + torch.Tensor: Prepared input embeddings for the model. + """ + assert input_ids.shape[0] == 1 + + # Embed input_ids to input_embeds + inputs_embeds = self.emb_text(input_ids) + + # Inject speaker embedding to input_embeds if it exists + if self.use_speaker_embedding: + spk_emb_mask = input_ids == self.spk_emb_token_id + if spk_emb_mask.any(): + assert lm_spk_emb_last_hidden_states is not None + # Project spk emb to tts hidden size first, [batch_size, num_spk_emb, llm_dim] -> [batch_size, num_spk_emb, self.hidden_size] + lm_spk_emb_last_hidden_states = lm_spk_emb_last_hidden_states.to( + self.projector.linear1.weight.dtype + ) + projected_spk_emb = self.projector(lm_spk_emb_last_hidden_states) + projected_spk_emb = F.normalize(projected_spk_emb, p=2, dim=-1) + apply_spk_emb( + input_ids=input_ids, + spk_emb=projected_spk_emb, + input_embeds=inputs_embeds, + spk_emb_token_id=self.spk_emb_token_id, + num_spk_embs=self.num_spk_embs, + ) + else: + raise NotImplementedError + + return inputs_embeds + + @torch.inference_mode() + def prefill_text( + self, + input_ids: torch.Tensor, + position_ids: torch.LongTensor, + past_key_values: List[Tuple[torch.Tensor, torch.Tensor]], + lm_spk_emb_last_hidden_states: Optional[torch.Tensor] = None, + ): + """Prefill a chunk of new text tokens in streaming setting. + Specifically speaking, update `past_key_values` using new text tokens, then the model will read the new text tokens. + + Args: + input_ids (Tensor): Tensor of shape [batch_size, seq_len] + position_ids (LongTensor): Tensor of shape [batch_size, seq_len] + past_key_values (List[Tuple[Tensor]]): KV Cache of all layers, each layer is a tuple (Tensor, Tensor) denoting keys and values. Each tensor is of seq_len = `self.streaming_text_reserved_len`. `past_key_values` will be updated. + lm_spk_emb_last_hidden_states (Tensor, optional): Tensor of shape [batch_size, num_spk_emb, llm_dim]. Defaults to None. + + Note that all `batch_size` should be `1`. + """ + assert input_ids.shape[0] == 1 + assert past_key_values is not None + + # Merge text and LLM embeddings + inputs_embeds = self.merge_inputs_embeds( + input_ids=input_ids, + lm_spk_emb_last_hidden_states=lm_spk_emb_last_hidden_states, + ) + + # Clone KV Cache + past_key_values_for_prefill = [] + for i in range(len(past_key_values)): + past_key_values_for_prefill.append( + ( + past_key_values[i][0][:, :, : position_ids[:, 0], :].clone(), + past_key_values[i][1][:, :, : position_ids[:, 0], :].clone(), + ) + ) + + # ModelMiniCPMVBaseModel + outputs_prefill: BaseModelOutputWithPast = self.model( + attention_mask=None, # because for text, it is standard causal attention mask, do nothing + position_ids=position_ids, # position_ids denotes the position of new text tokens in the sequence + past_key_values=past_key_values_for_prefill, # `past_key_values` will be updated by the model + inputs_embeds=inputs_embeds, # contains text and language model embedding + use_cache=True, + output_attentions=False, + cache_position=position_ids, # which new positions will use this cache, basically the same as position_ids + ) + + # Get model updated KV Cache + past_key_values_for_prefill_updated = outputs_prefill.past_key_values + + # Update generated KV Cache to input `past_key_values` + for layer_idx in range(len(past_key_values)): + # Update keys + past_key_values[layer_idx][0][ + :, :, position_ids[:, 0] : position_ids[:, -1] + 1, : + ] = past_key_values_for_prefill_updated[layer_idx][0][ + :, :, position_ids[:, 0] : position_ids[:, -1] + 1 + ].clone() + # Update values + past_key_values[layer_idx][1][ + :, :, position_ids[:, 0] : position_ids[:, -1] + 1, : + ] = past_key_values_for_prefill_updated[layer_idx][1][ + :, :, position_ids[:, 0] : position_ids[:, -1] + 1 + ].clone() + + # TODO: del past_key_values_for_prefill_updated recursively + # TODO: del outputs_prefill recursively + + return past_key_values + + @torch.inference_mode() + def prefill_audio_ids( + self, + input_ids: torch.Tensor, + past_key_values: List[Tuple[torch.Tensor, torch.Tensor]], + streaming_tts_text_mask=None, + add_audio_bos: bool = True, + ): + """Prefill a chunk of audio ids to the model. Used in sliding-window long audio generation. + Specifically, prefill many audio ids (typically from last window) to the model in the new window. + + Args: + input_ids (torch.Tensor): (1, seq_len, num_vq) Audio input token ids. + past_key_values (List[Tuple[torch.Tensor, torch.Tensor]]): Past key values for attention mechanism. + """ + assert input_ids.shape[0] == 1 + assert past_key_values is not None + + code_emb = [self.emb_code[i](input_ids[:, :, i]) for i in range(self.num_vq)] + inputs_embeds = torch.stack(code_emb, 3).sum(3) # [1,seq_len,768] + input_len = input_ids.shape[1] + + if add_audio_bos: + narrowed_input_ids = torch.tensor( + [[self.audio_bos_token_id]], dtype=torch.long, device=self.device + ) + bos_inputs_embeds = self.emb_text(narrowed_input_ids) + inputs_embeds = torch.cat([bos_inputs_embeds, inputs_embeds], dim=1) + input_len += 1 + + past_key_values_length = past_key_values[0][0].shape[2] + position_ids = torch.arange( + past_key_values_length, + past_key_values_length + input_len, + dtype=torch.long, + device=self.device, + ).unsqueeze(0) + + cache_position = position_ids.clone() + causal_mask = make_streaming_chunk_mask_generation( + inputs_embeds=inputs_embeds, + past_seen_tokens=past_key_values[0][0].shape[2], + streaming_tts_text_mask=streaming_tts_text_mask, + streaming_reserved_length=self.streaming_text_reserved_len, + streaming_text_chunk_size=self.streaming_text_chunk_size, + ) # [1, 1, 1, past_key_values_length + input_len] + + # Model forward + outputs: BaseModelOutputWithPast = self.model( + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=True, + output_attentions=False, + cache_position=cache_position, + ) + past_key_values = outputs.past_key_values + return past_key_values + + @torch.inference_mode() + def generate( + self, + input_ids: torch.Tensor, + past_key_values: List[Tuple[torch.Tensor, torch.Tensor]], + temperature: torch.Tensor, + eos_token: Union[int, torch.Tensor], + streaming_tts_text_mask=None, + force_no_stop=False, + min_new_token=10, + max_new_token=50, + logits_warpers: List[LogitsWarper] = [], + logits_processors: List[CustomRepetitionPenaltyLogitsProcessorRepeat] = [], + show_tqdm=False, + ): + """Generate audio codes in streaming setting or non-streaming setting. + Specifically speaking, generate audio codes when not all text tokens are prefilled. + + Always pass a valid `past_key_values` to the method. The method does not do `prefill` by itself. It relies on `prefill_text` method to provide valid `past_key_values`. Please refer to docstring of this class for more details. + + In this method, we borrowed a lot of codes from `https://github.com/2noise/ChatTTS/blob/main/ChatTTS/model/gpt.py`. + + Args: + input_ids (torch.Tensor): Input token ids. + past_key_values (List[Tuple[torch.Tensor, torch.Tensor]]): Past key values for attention mechanism. + temperature (torch.Tensor): Temperature for sampling. + eos_token (Union[int, torch.Tensor]): End of sequence token. + streaming_tts_text_mask (Optional[torch.Tensor], optional): Mask for streaming TTS text. Defaults to None. + max_new_token (int, optional): Maximum number of new tokens to generate. Defaults to 50. + logits_warpers (List[LogitsWarper], optional): List of logits warpers. Defaults to []. + logits_processors (List[CustomRepetitionPenaltyLogitsProcessorRepeat], optional): List of logits processors. Defaults to []. + show_tqdm (bool, optional): Whether to show progress bar. Defaults to True. + + Returns: + GenerationOutputs: Generation outputs. + """ + + # We only support batch size `1` for now + assert input_ids.shape[0] == 1 + assert past_key_values is not None + + # fix: this should not be `input_ids.shape[1]` + # start_idx = input_ids.shape[1] + start_idx = ( + 1 + + self.num_spk_embs * self.use_speaker_embedding + + self.streaming_text_reserved_len + + 1 + ) + + finish = torch.zeros(input_ids.shape[0], device=input_ids.device).bool() + + temperature = ( + temperature.unsqueeze(0) + .expand(input_ids.shape[0], -1) + .contiguous() + .view(-1, 1) + ) + + progress = input_ids.shape[1] + + # Pre-allocate input_ids, shape is [batch_size=1, max_possible_seq_len, self.num_vqs] + input_ids_buf = torch.zeros( + input_ids.shape[0], # batch_size + progress + + max_new_token, # max_possible_seq_len = input_ids.shape[1] + max_new_token + input_ids.shape[2], # self.num_vqs + dtype=input_ids.dtype, + device=input_ids.device, + ) + + # Copy existing `input_ids` to `input_ids_buf` + input_ids_buf.narrow(1, 0, progress).copy_(input_ids) + + del input_ids + input_ids = input_ids_buf.narrow(1, 0, progress) + + pbar: Optional[tqdm] = None + if show_tqdm: + pbar = tqdm( + total=max_new_token, + desc="code", + bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}(max) [{elapsed}, {rate_fmt}{postfix}]", + ) + + condition_length = ( + 1 + + self.num_spk_embs * self.use_speaker_embedding + + self.streaming_text_reserved_len + + 1 + ) + + for i in range(max_new_token): + # Prepare generation inputs + audio_bos = False + + # If this is the first audio token, the case is SPECIAL + if progress == condition_length: + audio_bos = True + + assert progress == ( + past_key_values[0][0].shape[2] + 1 + ) # If you are using according to the guidelines, this should be passed. + + if audio_bos: + # Generate the first token, activate the model with `self.audio_bos_token_id`, the model will predict + # a new audio token. This is a special case because without the `audio bos token`, it is impossible + # to generate the first audio token in our streaming setting. + narrowed_input_ids = torch.tensor( + [[self.audio_bos_token_id]], dtype=torch.long, device=self.device + ) + inputs_embeds = self.emb_text(narrowed_input_ids) + del narrowed_input_ids + else: + # Generate the following audio tokens, it is applicable to all other cases, including second and the + # following calling of `generate`. + narrowed_input_ids = input_ids.narrow( + dim=1, start=input_ids.shape[1] - 1, length=1 + ) + code_emb = [ + self.emb_code[i](narrowed_input_ids[:, :, i]) + for i in range(self.num_vq) + ] + inputs_embeds = torch.stack(code_emb, 3).sum(3) + + position_ids = torch.tensor( + [past_key_values[0][0].shape[2]], dtype=torch.long, device=self.device + ).unsqueeze(0) + + cache_position = position_ids.clone() + + # Make causal mask + causal_mask = make_streaming_chunk_mask_generation( + inputs_embeds=inputs_embeds, + past_seen_tokens=past_key_values[0][0].shape[2], + streaming_tts_text_mask=streaming_tts_text_mask, + streaming_reserved_length=self.streaming_text_reserved_len, + streaming_text_chunk_size=self.streaming_text_chunk_size, + ) + + # Model forward + outputs: BaseModelOutputWithPast = self.model( + attention_mask=causal_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=True, + output_attentions=False, + cache_position=cache_position, + ) + + del position_ids + del inputs_embeds + del cache_position + del causal_mask + + hidden_states = outputs.last_hidden_state + past_key_values = outputs.past_key_values + + with P.cached(): + logits = torch.empty( + hidden_states.size(0), + hidden_states.size(1), + self.num_audio_tokens, + self.num_vq, + dtype=torch.float, + device=self.device, + ) + for num_vq_iter in range(self.num_vq): + x: torch.Tensor = self.head_code[num_vq_iter](hidden_states) + logits[..., num_vq_iter] = x + del x + + del hidden_states + + # logits = logits[:, -1].float() + logits = logits.narrow(1, -1, 1).squeeze_(1).float() + + # logits = rearrange(logits, "b c n -> (b n) c") + logits = logits.permute(0, 2, 1) + logits = logits.reshape(-1, logits.size(2)) + # logits_token = rearrange(input_ids[:, start_idx:], "b c n -> (b n) c") + input_ids_sliced = input_ids.narrow( + 1, + start_idx, + input_ids.size(1) - start_idx, + ).permute(0, 2, 1) + logits_token = input_ids_sliced.reshape( + input_ids_sliced.size(0) * input_ids_sliced.size(1), + -1, + ).to(self.device) + del input_ids_sliced + + logits /= temperature + + if not audio_bos: + for logitsProcessors in logits_processors: + logits = logitsProcessors(logits_token, logits) + if not audio_bos: + for logitsWarpers in logits_warpers: + logits = logitsWarpers(logits_token, logits) + + del logits_token + + if i < min_new_token: + logits[:, eos_token] = -torch.inf + + if force_no_stop: + logits[:, eos_token] = -torch.inf + + scores = F.softmax(logits, dim=-1) + + del logits + idx_next = torch.multinomial(scores, num_samples=1) # .to(finish.device) + + del scores + + # idx_next = rearrange(idx_next, "(b n) 1 -> b n", n=self.num_vq) + idx_next = idx_next.view(-1, self.num_vq) + finish_or = idx_next.eq(eos_token).any(1) + finish.logical_or_(finish_or) + + del finish_or + # Store new `token` into `input_ids_buf` + input_ids_buf.narrow(1, progress, 1).copy_(idx_next.unsqueeze_(1)) + + if i == 0 and finish.any(): + # raise Exception + break + + del idx_next + progress += 1 + input_ids = input_ids_buf.narrow(1, 0, progress) + + if finish.all(): + break + + if pbar is not None: + pbar.update(1) + + if pbar is not None: + pbar.close() + + if not finish.all(): + if show_tqdm: + logger.info(f"incomplete result. hit max_new_token: {max_new_token}") + + del input_ids_buf + + if finish.all(): + # the last may contains eos token + genrated_input_ids = input_ids[:, condition_length:-1, :] + else: + # there is no eos token + genrated_input_ids = input_ids[:, condition_length:, :] + + return ConditionalChatTTSGenerationOutput( + new_ids=genrated_input_ids, + audio_input_ids=input_ids, # for update purpose + past_key_values=past_key_values, # for update purpose + finished=finish.all(), + ) + + @torch.inference_mode() + def decode_to_mel_specs( + self, + result_list: List[torch.Tensor], + ): + """Decode discrete audio codes to mel spectrograms. + + Borrowed from `https://github.com/2noise/ChatTTS/blob/main/ChatTTS/core.py` + + Args: + result_list (List[torch.Tensor]): Audio codes output from `generate`. + + Returns: + torch.Tensor: Mel spectrograms. + """ + + decoder = self.dvae + max_x_len = -1 + if len(result_list) == 0: + return np.array([], dtype=np.float32) + for result in result_list: + if result.size(0) > max_x_len: + max_x_len = result.size(0) + batch_result = torch.zeros( + (len(result_list), result_list[0].size(1), max_x_len), + dtype=result_list[0].dtype, + device=result_list[0].device, + ) + for i in range(len(result_list)): + src = result_list[i] + batch_result[i].narrow(1, 0, src.size(0)).copy_(src.permute(1, 0)) + del src + + mel_specs = decoder(batch_result) + del batch_result + return mel_specs + + +# Copied from transformers.models.whisper.modeling_whisper.WhisperEncoderLayer and add use_cache for streaming inference +class MiniCPMWhisperEncoderLayer(nn.Module): + def __init__(self, config: WhisperConfig, layer_idx: int = None): + super().__init__() + self.embed_dim = config.d_model + self.self_attn = WHISPER_ATTENTION_CLASSES[config._attn_implementation]( + embed_dim=self.embed_dim, + num_heads=config.encoder_attention_heads, + dropout=config.attention_dropout, + config=config, + layer_idx=layer_idx, + ) + self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) + self.dropout = config.dropout + self.activation_fn = ACT2FN[config.activation_function] + self.activation_dropout = config.activation_dropout + self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim) + self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim) + self.final_layer_norm = nn.LayerNorm(self.embed_dim) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + layer_head_mask: torch.Tensor, + output_attentions: bool = False, + past_key_values: Optional[EncoderDecoderCache] = None, + use_cache: Optional[bool] = False, + ) -> torch.Tensor: + r""" + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, embed_dim)`): + Hidden states to be fed into the encoder layer. + attention_mask (`torch.FloatTensor` of shape `(batch_size, 1, tgt_len, src_len)`): + Attention mask where padding elements are indicated by large negative values. + layer_head_mask (`torch.FloatTensor` of shape `(encoder_attention_heads,)`): + Mask to nullify selected heads of the attention modules. + output_attentions (`bool`, *optional*): + Whether or not to return the attention weights. + past_key_values (`EncoderDecoderCache`, *optional*): + Past key-value pairs used for incremental decoding. + use_cache (`bool`, *optional*): + Whether or not to return updated `past_key_values` for caching. + + Returns: + A tuple of shape `(hidden_states, optional(attn_weights), optional(past_key_values))`. + """ + residual = hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + hidden_states, attn_weights, past_key_values = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + layer_head_mask=layer_head_mask, + output_attentions=output_attentions, + past_key_value=past_key_values, + ) + hidden_states = nn.functional.dropout( + hidden_states, p=self.dropout, training=False + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = nn.functional.dropout( + hidden_states, p=self.activation_dropout, training=False + ) + hidden_states = self.fc2(hidden_states) + hidden_states = nn.functional.dropout( + hidden_states, p=self.dropout, training=False + ) + hidden_states = residual + hidden_states + + if hidden_states.dtype == torch.float16 and ( + torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any() + ): + clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + hidden_states = torch.clamp( + hidden_states, min=-clamp_value, max=clamp_value + ) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attn_weights,) + + if use_cache: + outputs += (past_key_values,) + + return outputs + + +# Copied from from transformers.models.whisper.modeling_whisper.WhisperEncoder and add use_cache for streaming inference +class MiniCPMWhisperEncoder(WhisperEncoder): + + def __init__(self, config: WhisperConfig): + super().__init__(config) + self.layers = nn.ModuleList( + [ + MiniCPMWhisperEncoderLayer(config, layer_idx=i) + for i in range(config.encoder_layers) + ] + ) + + def forward( + self, + input_features, + attention_mask=None, + head_mask=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + past_key_values: Optional[EncoderDecoderCache] = None, + use_cache: Optional[bool] = None, + ): + r""" + Forward pass of the Whisper encoder. + + Args: + input_features (`torch.FloatTensor` of shape `(batch_size, feature_size, sequence_length)`): + Float values of log-mel features extracted from the raw audio waveform. Typically generated + by a feature extractor (e.g., `WhisperFeatureExtractor`) that processes `.flac` or `.wav` + files into padded 2D mel spectrogram frames. These features are projected via convolution layers + (`conv1` and `conv2`) and then transformed into embeddings for the encoder. + + attention_mask (`torch.Tensor`, *optional*): + Not used by Whisper for masking `input_features`, but included for API compatibility with + other models. If provided, it is simply ignored within the model. By default, Whisper + effectively ignores silence in the input log-mel spectrogram. + + head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*): + Mask to nullify selected attention heads. The elements should be either 1 or 0, where: + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked** (i.e., the attention head is dropped). + + output_attentions (`bool`, *optional*): + Whether or not to return the attention tensors of all encoder layers. If set to `True`, the + returned tuple (or `BaseModelOutputWithPast`) will contain an additional element with + attention weights for each encoder layer. + + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. If set to `True`, the returned + tuple (or `BaseModelOutputWithPast`) will contain a tuple of hidden states, including the + initial embedding output as well as the outputs of each layer. + + return_dict (`bool`, *optional*): + Whether or not to return a `BaseModelOutputWithPast` (a subclass of `ModelOutput`) instead + of a plain tuple. If set to `True`, the output will be a `BaseModelOutputWithPast` object, + otherwise it will be a tuple. + + past_key_values (`EncoderDecoderCache`, *optional*): + When using caching for faster inference, this is an object that stores the key-value pairs + for attention states. If provided, the model will append new states to the existing cache + and return the updated cache. This speeds up sequential decoding or chunked inference. + + - If `past_key_values` is `None`, no past states are used or returned. + - If `past_key_values` is not `None` and `use_cache=True`, the model will use the provided + cache and return the updated cache (as `next_encoder_cache`). + + use_cache (`bool`, *optional*): + Whether or not the model should use caching (`past_key_values`) to speed up processing + during inference. When set to `True`, the model will: + - Inspect and use `past_key_values` if provided. + - Return updated `past_key_values` (under the name `next_encoder_cache` in + `BaseModelOutputWithPast`). + + Returns: + `BaseModelOutputWithPast` or `tuple` (depending on `return_dict`): + If `return_dict=True`, a `BaseModelOutputWithPast` is returned, which contains: + - **last_hidden_state** (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + The output of the final encoder layer. + - **hidden_states** (`tuple(torch.FloatTensor)`, *optional*, returned if `output_hidden_states=True`): + Hidden states of the model at each layer (including the initial projection). + - **attentions** (`tuple(torch.FloatTensor)`, *optional*, returned if `output_attentions=True`): + Attention weights from each encoder layer. + - **past_key_values** (an object of type `EncoderDecoderCache` or `None`, *optional*): + Updated cache of key-value pairs if `use_cache=True`. + + If `return_dict=False`, a tuple is returned, where the format is: + `(last_hidden_state, hidden_states, attentions)`, with `hidden_states` and `attentions` + only present if their respective `output_*` arguments are set to `True`. + + """ + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + # Ignore copy + input_features = input_features.to( + dtype=self.conv1.weight.dtype, device=self.conv1.weight.device + ) + + inputs_embeds = nn.functional.gelu(self.conv1(input_features)) + inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds)) + + inputs_embeds = inputs_embeds.permute(0, 2, 1) + + embed_pos = self.embed_positions.weight + past_key_values_length = 0 + if use_cache: + if past_key_values is None: + past_key_values = EncoderDecoderCache(DynamicCache(), DynamicCache()) + elif isinstance(past_key_values, list): + past_key_values = EncoderDecoderCache( + DynamicCache.from_legacy_cache(past_key_values), DynamicCache() + ) + elif isinstance(past_key_values, DynamicCache): + past_key_values = EncoderDecoderCache(past_key_values, DynamicCache()) + else: + pass + past_key_values_length = ( + past_key_values.self_attention_cache.get_usable_length( + inputs_embeds.shape[1] + ) + ) + if inputs_embeds.shape[1] + past_key_values_length > embed_pos.shape[0]: + logger.warning( + "seems the audio is longer than 30s. repeating the last part of the audio" + ) + embed_pos_front = embed_pos[past_key_values_length:, :] + embed_pos = torch.cat( + ( + embed_pos_front, + torch.repeat_interleave( + embed_pos[-1, :].unsqueeze(0), + inputs_embeds.shape[1] + - embed_pos.shape[0] + + past_key_values_length, + dim=0, + ), + ) + ) + else: + embed_pos = embed_pos[ + past_key_values_length : inputs_embeds.shape[1] + + past_key_values_length, + :, + ] + else: + embed_pos = embed_pos[: inputs_embeds.shape[1], :] + + hidden_states = inputs_embeds + embed_pos + hidden_states = nn.functional.dropout( + hidden_states, p=self.dropout, training=False + ) + + encoder_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + # check if head_mask has a correct number of layers specified if desired + if head_mask is not None: + assert head_mask.size()[0] == ( + len(self.layers) + ), f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}." + + for idx, encoder_layer in enumerate(self.layers): + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) + to_drop = False + + # Ignore copy + if to_drop: + layer_outputs = (None, None) + else: + layer_outputs = encoder_layer( + hidden_states, + attention_mask, + layer_head_mask=(head_mask[idx] if head_mask is not None else None), + output_attentions=output_attentions, + past_key_values=past_key_values, + use_cache=use_cache, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_encoder_cache = layer_outputs[2 if output_attentions else 1] + else: + next_encoder_cache = None + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + hidden_states = self.layer_norm(hidden_states) + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [hidden_states, encoder_states, all_attentions] + if v is not None + ) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + hidden_states=encoder_states, + attentions=all_attentions, + past_key_values=next_encoder_cache, + ) + + +class MultiModalProjector(nn.Module): + def __init__(self, in_dim, out_dim): + super().__init__() + self.linear1 = nn.Linear(in_features=in_dim, out_features=out_dim, bias=True) + self.relu = nn.ReLU() + self.linear2 = nn.Linear(in_features=out_dim, out_features=out_dim, bias=True) + + def forward(self, audio_features): + hidden_states = self.relu(self.linear1(audio_features)) + hidden_states = self.linear2(hidden_states) + return hidden_states + + +class MiniCPMO(MiniCPMBaseModel): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + ) -> None: + super().__init__(config=config, quant_config=quant_config) + + self.llm = self.init_llm(config=config, quant_config=quant_config) + + self.embed_dim = self.llm.config.hidden_size + + # init vision module + if self.config.init_vision: + # print("vision-understanding enabled") + self.vpm = self.init_vision_module(config=config, quant_config=quant_config) + self.vision_dim = self.vpm.embed_dim + self.resampler = self.init_resampler(self.embed_dim, self.vision_dim) + + # init audio module + self.config.init_audio = True + if self.config.init_audio: + # print("audio-understanding enabled") + self.apm = self.init_audio_module() + audio_output_dim = int(self.apm.config.encoder_ffn_dim // 4) + self.audio_avg_pooler = nn.AvgPool1d( + self.config.audio_pool_step, stride=self.config.audio_pool_step + ) + self.audio_projection_layer = MultiModalProjector( + in_dim=audio_output_dim, out_dim=self.embed_dim + ) + self.audio_encoder_layer = -1 + + # init tts module + self.config.init_tts = False + logger.info("TTS is disabled for now") + if self.config.init_tts: + # print("tts enabled") + assert ( + _tts_deps + ), "please make sure vector_quantize_pytorch and vocos are installed." + self.tts = self.init_tts_module() + + def init_tts_module(self): + model = ConditionalChatTTS(self.config.tts_config) + return model + + def init_audio_module(self): + model = MiniCPMWhisperEncoder(self.config.audio_config) + return model + + def init_llm( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> nn.Module: + return Qwen2ForCausalLM(config=config, quant_config=quant_config, prefix=prefix) + + def init_vision_module( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig], + prefix: str = "", + ): + if self.config._attn_implementation == "flash_attention_2": + self.config.vision_config._attn_implementation = "flash_attention_2" + else: + self.config.vision_config._attn_implementation = "eager" + model = Idefics2VisionTransformer( + config=config.vision_config, quant_config=quant_config, prefix=prefix + ) + if self.config.drop_vision_last_layer: + model.encoder.layers = model.encoder.layers[:-1] + + setattr(model, "embed_dim", model.embeddings.embed_dim) + setattr(model, "patch_size", model.embeddings.patch_size) + + return model + + def init_resampler( + self, + embed_dim: int, + vision_dim: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> nn.Module: + with set_default_torch_dtype(torch.float16): + # The resampler in 2.6 remains consistent with the one in 2.5. + resampler = Resampler2_5( + num_queries=self.config.query_num, + embed_dim=embed_dim, + num_heads=embed_dim // 128, + kv_dim=vision_dim, + quant_config=quant_config, + prefix=prefix, + ) + + return resampler.to(device="cuda", dtype=torch.get_default_dtype()) + + def pad_input_ids(self, input_ids: List[int], mm_input: MultimodalInputs): + # Get all special token IDs + im_start_id: int = mm_input.im_start_id + im_end_id: int = mm_input.im_end_id + slice_start_id: int = mm_input.slice_start_id + slice_end_id: int = mm_input.slice_end_id + + media_token_pairs = [ + (im_start_id, im_end_id), + (slice_start_id, slice_end_id), + (mm_input.audio_start_id, mm_input.audio_end_id), + ] + pattern = MultiModalityDataPaddingPatternTokenPairs(media_token_pairs) + + return pattern.pad_input_tokens(input_ids, mm_input) + + def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor): + """ + Computes the output length of the convolutional layers and the output length of the audio encoder + """ + input_lengths_after_cnn = (input_lengths - 1) // 2 + 1 + input_lengths_after_pooling = ( + input_lengths_after_cnn - self.config.audio_pool_step + ) // self.config.audio_pool_step + 1 + input_lengths_after_pooling = input_lengths_after_pooling.to(dtype=torch.int32) + + return input_lengths_after_cnn, input_lengths_after_pooling + + def get_audio_embedding_streaming(self, items: List[MultimodalDataItem]): + r""" + Extract audio embeddings in a streaming manner using cached key-value pairs. + + This method processes incoming audio features incrementally and stores/updates `past_key_values` + for faster inference on subsequent audio frames. It only supports batch_size=1 and is intended + for streaming scenarios. + + Returns: + List[List[torch.Tensor]]: audio embeddings + """ + wavforms = flatten_nested_list( + [item.audio_features for item in items if item.audio_features] + ) + # list, [[x1, x2], [y1], [z1]] + audio_feature_lens_raw = flatten_nested_list( + [item.audio_feature_lens for item in items if item.audio_feature_lens] + ) + + # exist audio + if len(wavforms) > 0: + audio_feature_lens = torch.hstack(audio_feature_lens_raw) + batch_size, _, max_mel_seq_len = wavforms.shape + assert batch_size == 1 + max_seq_len = (max_mel_seq_len - 1) // 2 + 1 + + if self.audio_past_key_values is not None: + cache_length = self.audio_past_key_values[0][0].shape[2] + apm_max_len = self.apm.embed_positions.weight.shape[0] + if cache_length + max_seq_len >= apm_max_len: + logger.warning( + f"audio_past_key_values length {cache_length + max_seq_len} exceed {apm_max_len}, reset." + ) + self.audio_past_key_values = None + + audio_outputs = self.apm( + wavforms, past_key_values=self.audio_past_key_values, use_cache=True + ) + audio_states = ( + audio_outputs.last_hidden_state + ) # [:, :audio_feat_lengths, :] + self.audio_past_key_values = audio_outputs.past_key_values + + audio_embeds = self.audio_projection_layer(audio_states) + + audio_embeds = audio_embeds.transpose(1, 2) + audio_embeds = self.audio_avg_pooler(audio_embeds) + audio_embeds = audio_embeds.transpose(1, 2) + + _, feature_lens_after_pooling = self._get_feat_extract_output_lengths( + audio_feature_lens + ) + + num_audio_tokens = feature_lens_after_pooling + + final_audio_embeds = [] + idx = 0 + for i in range(len(audio_feature_lens_raw)): + target_audio_embeds = [] + for _ in range(len(audio_feature_lens_raw[i])): + target_audio_embeds.append( + audio_embeds[idx, : num_audio_tokens[idx], :] + ) + idx += 1 + final_audio_embeds.append(target_audio_embeds) + return final_audio_embeds + else: + return [] + + def subsequent_chunk_mask( + self, + size: int, + chunk_size: int, + num_left_chunks: int = -1, + device: torch.device = torch.device("cpu"), + num_lookhead: int = 0, + ) -> torch.Tensor: + """Create mask for subsequent steps (size, size) with chunk size, + this is for streaming encoder + + Args: + size (int): size of mask + chunk_size (int): size of chunk + num_left_chunks (int): number of left chunks + <0: use full chunk + >=0: use num_left_chunks + device (torch.device): "cpu" or "cuda" or torch.Tensor.device + + Returns: + torch.Tensor: mask + + """ + ret = torch.zeros(size, size, device=device, dtype=torch.bool) + for i in range(size): + if num_left_chunks < 0: + start = 0 + else: + start = max((i // chunk_size - num_left_chunks) * chunk_size, 0) + ending = min((i // chunk_size + 1) * chunk_size + num_lookhead, size) + ret[i, start:ending] = True + return ret + + def get_audio_embedding(self, items: List[MultimodalDataItem], chunk_length=-1): + r""" + Extract full audio embeddings with optional chunk-based attention. + + This method computes embeddings for all audio frames at once, either using full attention (when + `chunk_length` is -1) or chunk-based attention (when `chunk_length` is a positive number). It does + not use key-value caching and is suitable for non-streaming inference. + + Args: + chunk_length (int, optional): Determines whether to use full attention (-1) or chunk-based + attention (>0) during embedding computation. + + Returns: + List[List[torch.Tensor]]: audio embeddings + """ + # (bs, 80, frames) or [], multi audios need filled in advance + wavforms = flatten_nested_list( + [item.audio_features for item in items if item.audio_features] + ) + # list, [[x1, x2], [y1], [z1]] + audio_feature_lens_raw = flatten_nested_list( + [item.audio_feature_lens for item in items if item.audio_feature_lens] + ) + + final_audio_embeds = [] + + assert isinstance(wavforms, list) + assert isinstance(wavforms[0], torch.Tensor) + # exist audio + for wavform in wavforms: + if len(wavform) > 0: + audio_feature_lens = torch.hstack(audio_feature_lens_raw) + batch_size, _, max_mel_seq_len = wavform.shape + max_seq_len = (max_mel_seq_len - 1) // 2 + 1 + + # Create a sequence tensor of shape (batch_size, max_seq_len) + seq_range = ( + torch.arange( + 0, + max_seq_len, + dtype=audio_feature_lens.dtype, + device=audio_feature_lens.device, + ) + .unsqueeze(0) + .expand(batch_size, max_seq_len) + ) + lengths_expand = audio_feature_lens.unsqueeze(1).expand( + batch_size, max_seq_len + ) + # Create mask + padding_mask = seq_range >= lengths_expand # 1 for padded values + + audio_attention_mask_ = padding_mask.view( + batch_size, 1, 1, max_seq_len + ).expand(batch_size, 1, max_seq_len, max_seq_len) + audio_attention_mask = audio_attention_mask_.to( + dtype=self.apm.conv1.weight.dtype, + device=self.apm.conv1.weight.device, + ) + + if chunk_length > 0: + chunk_num_frame = int(chunk_length * 50) + chunk_mask = self.subsequent_chunk_mask( + size=max_seq_len, + chunk_size=chunk_num_frame, + num_left_chunks=-1, + device=audio_attention_mask_.device, + ) + audio_attention_mask_ = torch.logical_or( + audio_attention_mask_, torch.logical_not(chunk_mask) + ) + + audio_attention_mask[audio_attention_mask_] = float("-inf") + audio_states = self.apm( + wavform, + output_hidden_states=True, + attention_mask=audio_attention_mask, + ).hidden_states[self.audio_encoder_layer] + audio_embeds = self.audio_projection_layer(audio_states) + + audio_embeds = audio_embeds.transpose(1, 2) + audio_embeds = self.audio_avg_pooler(audio_embeds) + audio_embeds = audio_embeds.transpose(1, 2) + + _, feature_lens_after_pooling = self._get_feat_extract_output_lengths( + audio_feature_lens + ) + + num_audio_tokens = feature_lens_after_pooling + + idx = 0 + for i in range(len(audio_feature_lens_raw)): + target_audio_embeds = [] + for _ in range(len(audio_feature_lens_raw[i])): + target_audio_embeds.append( + audio_embeds[idx, : num_audio_tokens[idx], :] + ) + idx += 1 + final_audio_embeds.append(target_audio_embeds) + return final_audio_embeds + + def get_audio_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + embedding = self.get_omni_embedding( + items=items, + chunk_length=self.config.audio_chunk_length, + stream_input=False, + ) + return embedding + + def get_omni_embedding( + self, + items: List[MultimodalDataItem], + chunk_length=-1, + stream_input=False, + ): + """ + Args: + chunk_length: whisper use full attention or chunk attention + stream_input: use streaming audio embedding + Returns: + final embeddings with audio feature + """ + + if stream_input: + audio_embeddings = self.get_audio_embedding_streaming(items) + else: + audio_embeddings = self.get_audio_embedding(items, chunk_length) + bs = len(audio_embeddings) + # batch size + audio_embs = torch.cat(flatten_nested_list(audio_embeddings), dim=0) + + return audio_embs + + def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + # list of tensors + pixel_values = flatten_nested_list([item.pixel_values for item in items]) + tgt_sizes = torch.stack( + flatten_nested_list([item.tgt_size for item in items]), dim=0 + ) + assert len(pixel_values) == tgt_sizes.shape[0] + + device = self.vpm.embeddings.position_embedding.weight.device + dtype = self.vpm.embeddings.position_embedding.weight.dtype + all_pixel_values_lst = [ + i.flatten(end_dim=1).permute(1, 0) for i in pixel_values + ] + + max_patches = (tgt_sizes[:, 0] * tgt_sizes[:, 1]).max().item() + assert isinstance(max_patches, int) + all_pixel_values = torch.nn.utils.rnn.pad_sequence( + all_pixel_values_lst, batch_first=True, padding_value=0.0 + ) + + B, L, _ = all_pixel_values.shape + all_pixel_values = all_pixel_values.permute(0, 2, 1).reshape(B, 3, -1, L) + patch_attn_mask = torch.zeros( + (B, 1, max_patches), dtype=torch.bool, device=device + ) + + tgt_sizes_tensor = tgt_sizes.clone().to(device=patch_attn_mask.device) + mask_shapes = tgt_sizes_tensor[:, 0] * tgt_sizes_tensor[:, 1] + patch_attn_mask[:, 0, :] = torch.arange( + patch_attn_mask.size(2), device=patch_attn_mask.device + ).unsqueeze(0) < mask_shapes.unsqueeze(1) + + vision_embedding = self.vpm( + all_pixel_values.type(dtype), + patch_attention_mask=patch_attn_mask, + tgt_sizes=tgt_sizes, + ) + return self.resampler(vision_embedding, tgt_sizes) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + **kwargs: Any, + ) -> torch.Tensor: + + mm_input = forward_batch.merge_mm_inputs() + placeholder_token_ids = ( + ([mm_input.im_token_id] + [item.pad_value for item in mm_input.mm_items]) + if forward_batch.contains_mm_inputs() + else [] + ) + hidden_states = general_mm_embed_routine( + input_ids=input_ids, + forward_batch=forward_batch, + language_model=self.llm, + image_data_embedding_func=self.get_image_feature, + audio_data_embedding_func=self.get_audio_feature, + placeholder_tokens={ + Modality.IMAGE: placeholder_token_ids, + Modality.AUDIO: placeholder_token_ids, + }, + positions=positions, + ) + return hidden_states + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + + if "rotary_emb.inv_freq~" in name or "projector" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + + # For weight_norm parametrization, handle both old and new formats + if self.config.init_tts and "tts" in name: + # Handle loading from older checkpoints with weight_g/weight_v format + if ".weight_g" in name or ".weight_v" in name: + name = name.replace( + ".weight_g", ".parametrizations.weight.original0" + ) + name = name.replace( + ".weight_v", ".parametrizations.weight.original1" + ) + elif ".weight" in name and name not in params_dict: + param_name = name.replace( + ".weight", ".parametrizations.weight.original0" + ) + if param_name in params_dict: + name = param_name + + # adapt to VisionAttention + if "vpm" in name: + name = name.replace(r"self_attn.out_proj", r"self_attn.proj") + + if not self.config.init_tts and "tts" in name: + continue + if not self.config.init_audio and ("apm" in name or "audio" in name): + continue + if not self.config.init_vision and "vpm" in name: + continue + + if ( + "sampler" in name + or "apm" in name + or ("tts" in name and "self_attn" in name) + or ("tts.model.layers" in name and ".mlp" in name) + ): + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + # replace the name and load with customized loader + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +EntryClass = [MiniCPMO] diff --git a/lib/python3.12/site-packages/sglang/srt/models/mistral.py b/lib/python3.12/site-packages/sglang/srt/models/mistral.py new file mode 100644 index 0000000000000000000000000000000000000000..70d2ca3e04c0a23da2da1816bcced0a5b702e083 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/mistral.py @@ -0,0 +1,23 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Inference-only Mistral model.""" + +from sglang.srt.models.llama import LlamaForCausalLM + + +class MistralForCausalLM(LlamaForCausalLM): + pass + + +EntryClass = MistralForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/mixtral.py b/lib/python3.12/site-packages/sglang/srt/models/mixtral.py new file mode 100644 index 0000000000000000000000000000000000000000..8d5a03f0adb4eb1e0ce595fefecb7cc367e8a552 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/mixtral.py @@ -0,0 +1,408 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from +# https://github.com/vllm-project/vllm/blob/c7f2cf2b7f67bce5842fedfdba508440fe257375/vllm/model_executor/models/mixtral.py#L1 +"""Inference-only Mixtral model.""" + +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import MixtralConfig + +from sglang.srt.distributed import ( + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.ep_moe.layer import EPMoE +from sglang.srt.layers.moe.fused_moe_triton import FusedMoE +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.managers.schedule_batch import global_server_args_dict +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix + + +class MixtralMoE(nn.Module): + """A tensor-parallel MoE implementation for Mixtral that shards each expert + across all ranks. + + Each expert's weights are sharded across all ranks and a fused MoE + kernel is used for the forward pass, and finally we reduce the outputs + across ranks. + """ + + def __init__( + self, + num_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + params_dtype: Optional[torch.dtype] = None, + quant_config: Optional[QuantizationConfig] = None, + tp_size: Optional[int] = None, + prefix: str = "", + ): + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + self.hidden_size = hidden_size + + # Gate always runs at half / full precision for now. + self.gate = ReplicatedLinear( + hidden_size, + num_experts, + bias=False, + params_dtype=params_dtype, + quant_config=None, + prefix=add_prefix("gate", prefix), + ) + MoEImpl = EPMoE if global_server_args_dict["enable_ep_moe"] else FusedMoE + self.experts = MoEImpl( + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + params_dtype=params_dtype, + renormalize=True, + quant_config=quant_config, + tp_size=tp_size, + prefix=add_prefix("experts", prefix), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # NOTE: hidden_states can have either 1D or 2D shape. + orig_shape = hidden_states.shape + hidden_states = hidden_states.view(-1, self.hidden_size) + # router_logits: (num_tokens, n_experts) + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts(hidden_states, router_logits) + if self.tp_size > 1: + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + return final_hidden_states.view(orig_shape) + + +class MixtralAttention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + max_position: int = 4096 * 32, + rope_theta: float = 10000, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position, + base=int(self.rope_theta), + is_neox_style=True, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class MixtralDecoderLayer(nn.Module): + def __init__( + self, + config: MixtralConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + # Requires transformers > 4.32.0 + rope_theta = getattr(config, "rope_theta", 10000) + self.self_attn = MixtralAttention( + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + max_position=config.max_position_embeddings, + num_kv_heads=config.num_key_value_heads, + layer_id=layer_id, + rope_theta=rope_theta, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + self.block_sparse_moe = MixtralMoE( + num_experts=config.num_local_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + quant_config=quant_config, + prefix=add_prefix("block_sparse_moe", prefix), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> torch.Tensor: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.block_sparse_moe(hidden_states) + return hidden_states, residual + + +class MixtralModel(nn.Module): + def __init__( + self, + config: MixtralConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=add_prefix("embed_tokens", prefix), + ) + self.layers = nn.ModuleList( + [ + MixtralDecoderLayer( + config, + i, + quant_config=quant_config, + prefix=add_prefix(f"layers.{i}", prefix), + ) + for i in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, hidden_states, forward_batch, residual + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class MixtralForCausalLM(nn.Module): + + def __init__( + self, + config: MixtralConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.model = MixtralModel( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, config.hidden_size, prefix=add_prefix("lm_head", prefix) + ) + self.logits_processor = LogitsProcessor(config) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ] + + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + MoEImpl = EPMoE if global_server_args_dict["enable_ep_moe"] else FusedMoE + expert_params_mapping = MoEImpl.make_expert_params_mapping( + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.num_local_experts, + ) + + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if ( + name.endswith(".bias") or name.endswith("_bias") + ) and name not in params_dict: + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + + if ( + name.endswith(".bias") or name.endswith("_bias") + ) and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + break + else: + # Skip loading extra bias for GPTQ models. + if ( + name.endswith(".bias") or name.endswith("_bias") + ) and name not in params_dict: + continue + # Skip loading kv_scale from ckpts towards new design. + if name.endswith(".kv_scale") and name not in params_dict: + continue + if name is None: + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + +EntryClass = MixtralForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/mixtral_quant.py b/lib/python3.12/site-packages/sglang/srt/models/mixtral_quant.py new file mode 100644 index 0000000000000000000000000000000000000000..cfb23cdf3e8cf69f729a636d61bb00d8302f3ad3 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/mixtral_quant.py @@ -0,0 +1,426 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from +# https://github.com/vllm-project/vllm/blob/c7f2cf2b7f67bce5842fedfdba508440fe257375/vllm/model_executor/models/mixtral_quant.py#L1 +"""Inference-only Mixtral model.""" + +from typing import Iterable, Optional, Tuple + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn +from transformers import MixtralConfig + +from sglang.srt.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix + + +class MixtralMLP(nn.Module): + def __init__( + self, + num_experts: int, + hidden_size: int, + intermediate_size: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.num_experts = num_experts + self.ffn_dim = intermediate_size + self.hidden_dim = hidden_size + + self.w1 = ReplicatedLinear( + self.hidden_dim, + self.ffn_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("w1", prefix), + ) + self.w2 = ReplicatedLinear( + self.ffn_dim, + self.hidden_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("w2", prefix), + ) + self.w3 = ReplicatedLinear( + self.hidden_dim, + self.ffn_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("w3", prefix), + ) + + # TODO: Use vllm's SiluAndMul + self.act_fn = nn.SiLU() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + w1_out, _ = self.w1(hidden_states) + w1_out = self.act_fn(w1_out) + w3_out, _ = self.w3(hidden_states) + current_hidden_states = w1_out * w3_out + current_hidden_states, _ = self.w2(current_hidden_states) + return current_hidden_states + + +class MixtralMoE(nn.Module): + def __init__( + self, + config: MixtralConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.rank = get_tensor_model_parallel_rank() + self.tp_size = get_tensor_model_parallel_world_size() + self.num_total_experts = config.num_local_experts + self.top_k = config.num_experts_per_tok + if self.tp_size > self.num_total_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {self.num_total_experts}." + ) + # Split experts equally between ranks + self.expert_indicies = np.array_split( + range(self.num_total_experts), self.tp_size + )[self.rank].tolist() + if not self.expert_indicies: + raise ValueError(f"Rank {self.rank} has no experts assigned to it.") + + self.experts = nn.ModuleList( + [ + ( + MixtralMLP( + self.num_total_experts, + config.hidden_size, + config.intermediate_size, + quant_config=quant_config, + prefix=add_prefix(f"experts.{idx}", prefix), + ) + if idx in self.expert_indicies + else None + ) + for idx in range(self.num_total_experts) + ] + ) + self.gate = ReplicatedLinear( + config.hidden_size, + self.num_total_experts, + bias=False, + quant_config=None, + prefix=add_prefix("gate", prefix), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + router_logits, _ = self.gate(hidden_states) + + routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk( + routing_weights, self.top_k, dim=-1 + ) + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + + final_hidden_states = None + for expert_idx in self.expert_indicies: + expert_layer = self.experts[expert_idx] + expert_mask = selected_experts == expert_idx + expert_weights = (routing_weights * expert_mask).sum(dim=-1, keepdim=True) + + current_hidden_states = expert_layer(hidden_states).mul_(expert_weights) + if final_hidden_states is None: + final_hidden_states = current_hidden_states + else: + final_hidden_states.add_(current_hidden_states) + + return tensor_model_parallel_all_reduce(final_hidden_states) + + +class MixtralAttention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + max_position: int = 4096 * 32, + rope_theta: float = 10000, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position, + base=int(self.rope_theta), + is_neox_style=True, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class MixtralDecoderLayer(nn.Module): + def __init__( + self, + config: MixtralConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + # Requires transformers > 4.32.0 + rope_theta = getattr(config, "rope_theta", 10000) + self.self_attn = MixtralAttention( + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + max_position=config.max_position_embeddings, + num_kv_heads=config.num_key_value_heads, + layer_id=layer_id, + rope_theta=rope_theta, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + self.block_sparse_moe = MixtralMoE( + config=config, + quant_config=quant_config, + prefix=add_prefix("block_sparse_moe", prefix), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> torch.Tensor: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.block_sparse_moe(hidden_states) + return hidden_states, residual + + +class MixtralModel(nn.Module): + def __init__( + self, + config: MixtralConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=add_prefix("embed_tokens", prefix), + ) + self.layers = nn.ModuleList( + [ + MixtralDecoderLayer( + config, + i, + quant_config=quant_config, + prefix=add_prefix(f"layers.{i}", prefix), + ) + for i in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, hidden_states, forward_batch, residual + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class QuantMixtralForCausalLM(nn.Module): + def __init__( + self, + config: MixtralConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.model = MixtralModel( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, config.hidden_size, prefix=add_prefix("lm_head", prefix) + ) + self.logits_processor = LogitsProcessor(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ] + + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Skip experts that are not assigned to this worker. + if "block_sparse_moe.experts." in name and name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +EntryClass = QuantMixtralForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/mllama.py b/lib/python3.12/site-packages/sglang/srt/models/mllama.py new file mode 100644 index 0000000000000000000000000000000000000000..aea1fdf716273459545285844c9645311aaf6da1 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/mllama.py @@ -0,0 +1,1064 @@ +# Adapted from: +# https://github.com/vllm-project/vllm/blob/7193774b1ff8603ad5bf4598e5efba0d9a39b436/vllm/model_executor/models/mllama.py +"""PyTorch Mllama model.""" +import math +from typing import Iterable, List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +import transformers.models.mllama.configuration_mllama as config_mllama +from torch import nn +from transformers.modeling_outputs import BaseModelOutput, CausalLMOutputWithPast +from transformers.models.mllama.modeling_mllama import ( + _prepare_aspect_ratio_attention_mask, +) + +import sglang.srt.distributed.parallel_state as ps +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.activation import get_act_fn +from sglang.srt.layers.attention.vision import VisionAttention +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.vocab_parallel_embedding import ( + DEFAULT_VOCAB_PADDING_SIZE, + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.managers.schedule_batch import MultimodalInputs +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.llama import LlamaDecoderLayer, LlamaMLP +from sglang.srt.utils import add_prefix + + +class ColumnParallelConv2dPatch(torch.nn.Module): + """Conv2D Patching layer with model parallelism. + Column parallel over unfolded input. + Arguments: + in_channels: Input channels. + out_channels: Output channels. + kernel_size: Size of convolution kernel. + stride (default 1): Stride for convolution. + bias (default False): Use bias in Conv2d. + Input: (bsz, in_channels, width, height) + Output: (bsz, num_tokens, out_channels) + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: Union[int, Tuple[int, int]], + stride: Union[int, Tuple[int, int]], + bias: bool = False, + ) -> None: + super().__init__() + if isinstance(kernel_size, int): + kernel_size = (kernel_size, kernel_size) + self._unfold = torch.nn.Unfold(kernel_size=kernel_size, stride=stride) + self._linear = ColumnParallelLinear( + in_channels * kernel_size[0] * kernel_size[1], + out_channels, + bias=bias, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self._unfold(x) + x = x.permute(0, 2, 1) + x, _ = self._linear(x) + return x + + +class MllamaPrecomputedAspectRatioEmbedding(nn.Module): + + def __init__(self, config: config_mllama.MllamaVisionConfig, is_gated: bool = True): + super().__init__() + self.max_num_tiles = config.max_num_tiles + self.hidden_size = config.hidden_size + self.max_aspect_ratio_id = config.max_aspect_ratio_id + self.is_gated = is_gated + + self.embedding = nn.Embedding( + self.max_aspect_ratio_id + 1, self.max_num_tiles * self.hidden_size + ) + if is_gated: + self.gate = nn.Parameter(torch.zeros(1)) + + def forward( + self, hidden_state: torch.Tensor, aspect_ratio_ids: torch.Tensor + ) -> torch.Tensor: + embeddings = self.embedding(aspect_ratio_ids) + embeddings = embeddings.reshape(-1, self.max_num_tiles, 1, self.hidden_size) + + if self.is_gated: + embeddings = embeddings * self.gate.tanh() + + hidden_state = hidden_state + embeddings + return hidden_state + + +class MllamaPrecomputedPositionEmbedding(nn.Module): + def __init__(self, config: config_mllama.MllamaVisionConfig): + super().__init__() + self.max_num_tiles = config.max_num_tiles + self.max_aspect_ratio_id = config.max_aspect_ratio_id + self.num_patches = (config.image_size // config.patch_size) ** 2 + 1 + self.hidden_size = config.hidden_size + self.scale = config.hidden_size**-0.5 + + self.gate = nn.Parameter(torch.zeros(1)) + + # position embedding + position_embedding = torch.randn(self.num_patches, self.hidden_size) + self.embedding = nn.Parameter(self.scale * position_embedding) + + # tile position embedding + self.tile_embedding = nn.Embedding( + self.max_aspect_ratio_id + 1, + self.max_num_tiles * self.num_patches * self.hidden_size, + ) + + def forward( + self, hidden_state: torch.Tensor, aspect_ratio_ids: torch.Tensor + ) -> torch.Tensor: + # position embeddings + gated_position_embedding = (1 - self.gate.tanh()) * self.embedding + hidden_state = hidden_state + gated_position_embedding.view( + 1, 1, self.num_patches, self.hidden_size + ) + + # precomputed tile position embeddings + tile_position_embedding = self.tile_embedding(aspect_ratio_ids) + batch_size = hidden_state.shape[0] + tile_position_embedding = tile_position_embedding.reshape( + batch_size, self.max_num_tiles, self.num_patches, self.hidden_size + ) + gated_tile_position_embedding = self.gate.tanh() * tile_position_embedding + hidden_state = hidden_state + gated_tile_position_embedding + + return hidden_state + + +class MllamaVisionMLP(nn.Module): + def __init__( + self, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.activation_fn = get_act_fn(config.hidden_act) + self.fc1 = ColumnParallelLinear( + config.hidden_size, + config.intermediate_size, + bias=True, + quant_config=quant_config, + prefix=add_prefix("fc1", prefix), + ) + self.fc2 = RowParallelLinear( + config.intermediate_size, + config.hidden_size, + bias=True, + quant_config=quant_config, + prefix=add_prefix("fc2", prefix), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, _ = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states, _ = self.fc2(hidden_states) + + return hidden_states + + +class MllamaVisionEncoderLayer(nn.Module): + def __init__( + self, + config: config_mllama.MllamaVisionConfig, + quant_config: Optional[QuantizationConfig] = None, + is_gated: bool = False, + prefix: str = "", + ): + super().__init__() + + self.hidden_size = config.hidden_size + self.num_attention_heads = config.attention_heads + self.is_gated = is_gated + self.intermediate_size = config.intermediate_size + + self.self_attn = VisionAttention( + self.hidden_size, + self.num_attention_heads, + self.hidden_size, + use_qkv_parallel=True, + quant_config=quant_config, + dropout=0.0, + use_context_forward=False, + softmax_in_single_precision=False, + flatten_batch=False, + prefix=add_prefix("self_attn", prefix), + ) + self.mlp = MllamaVisionMLP( + config, quant_config, prefix=add_prefix("mlp", prefix) + ) + + self.input_layernorm = nn.LayerNorm(self.hidden_size, eps=config.norm_eps) + self.post_attention_layernorm = nn.LayerNorm( + self.hidden_size, eps=config.norm_eps + ) + + # there used to be an if else here, no code path + if is_gated: + self.gate_attn = nn.Parameter(torch.ones(1) * math.pi / 4) + self.gate_ffn = nn.Parameter(torch.ones(1) * math.pi / 4) + + def forward( + self, + hidden_state: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + ): + # Self Attention + residual = hidden_state + hidden_state = self.input_layernorm(hidden_state) + hidden_state = self.self_attn(hidden_state, attention_mask=attention_mask) + gate_attn = 1 if not self.is_gated else self.gate_attn.tanh() + hidden_state = residual + gate_attn * hidden_state + + # Feed forward + residual = hidden_state + hidden_state = self.post_attention_layernorm(hidden_state) + hidden_state = self.mlp(hidden_state) + gate_ffn = 1 if not self.is_gated else self.gate_ffn.tanh() + hidden_state = residual + gate_ffn * hidden_state + + return hidden_state + + +class MllamaVisionEncoder(nn.Module): + def __init__( + self, + config: config_mllama.MllamaVisionConfig, + quant_config: Optional[QuantizationConfig] = None, + num_layers=32, + is_gated=False, + output_hidden_states=None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.layers = nn.ModuleList( + [ + MllamaVisionEncoderLayer( + config, + quant_config, + is_gated, + prefix=add_prefix(f"layers.{i}", prefix), + ) + for i in range(num_layers) + ] + ) + self.output_hidden_states = output_hidden_states or [] + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + ) -> Union[Tuple, BaseModelOutput]: + encoder_states = () + + for i, encoder_layer in enumerate(self.layers): + if i in self.output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + hidden_states = encoder_layer( + hidden_states, + attention_mask, + ) + + if len(self.layers) - 1 in self.output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + return hidden_states, encoder_states + + +class MllamaVisionModel(nn.Module): + def __init__( + self, + config: config_mllama.MllamaVisionConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.image_size = config.image_size + self.patch_size = config.patch_size + self.max_num_tiles = config.max_num_tiles + self.hidden_size = config.hidden_size + self.in_channels = config.num_channels + self.intermediate_layers_indices = config.intermediate_layers_indices + + self.num_patches = (self.image_size // self.patch_size) ** 2 + 1 + self.scale = config.hidden_size**-0.5 + + self.patch_embedding = ColumnParallelConv2dPatch( + in_channels=config.num_channels, + out_channels=self.hidden_size, + kernel_size=self.patch_size, + stride=self.patch_size, + bias=False, + ) + + self.class_embedding = nn.Parameter(self.scale * torch.randn(self.hidden_size)) + self.gated_positional_embedding = MllamaPrecomputedPositionEmbedding(config) + + self.pre_tile_positional_embedding = MllamaPrecomputedAspectRatioEmbedding( + config, is_gated=True + ) + self.post_tile_positional_embedding = MllamaPrecomputedAspectRatioEmbedding( + config, is_gated=True + ) + + # layer norms + self.layernorm_pre = nn.LayerNorm(self.hidden_size) + self.layernorm_post = nn.LayerNorm(self.hidden_size) + + # encoders + self.transformer = MllamaVisionEncoder( + config, + quant_config, + config.num_hidden_layers, + is_gated=False, + output_hidden_states=config.intermediate_layers_indices, + prefix=add_prefix("transformer", prefix), + ) + self.global_transformer = MllamaVisionEncoder( + config, + quant_config, + config.num_global_layers, + is_gated=True, + prefix=add_prefix("global_transformer", prefix), + ) + + def apply_class_embedding(self, hidden_state: torch.Tensor) -> torch.Tensor: + batch_size, _, hidden_size = hidden_state.shape + class_embedding = self.class_embedding.expand(batch_size, 1, hidden_size) + hidden_state = torch.cat([class_embedding, hidden_state], dim=1) + return hidden_state + + def forward( + self, + pixel_values: torch.Tensor, + aspect_ratio_ids: torch.Tensor, + aspect_ratio_mask: torch.Tensor, + ) -> torch.Tensor: + batch_size, num_concurrent_media, num_tiles, num_channels, height, width = ( + pixel_values.shape + ) + + pixel_values = pixel_values.reshape( + batch_size * num_concurrent_media * num_tiles, num_channels, height, width + ) + aspect_ratio_ids = aspect_ratio_ids.reshape( + batch_size * num_concurrent_media, -1 + ) + + # patch embedding + patch_embeds = self.patch_embedding( + pixel_values.to(self.layernorm_pre.weight.dtype) + ) + hidden_state = patch_embeds + hidden_state = ps.get_tp_group().all_gather(hidden_state) + + # tile embeddings + _, num_patches, dim = hidden_state.shape + hidden_state = hidden_state.reshape( + batch_size * num_concurrent_media, num_tiles, -1, dim + ) + hidden_state = self.pre_tile_positional_embedding( + hidden_state, aspect_ratio_ids + ) + + # apply cls token + hidden_state = hidden_state.reshape( + batch_size * num_concurrent_media * num_tiles, num_patches, dim + ) + hidden_state = self.apply_class_embedding(hidden_state) + num_patches += 1 + + # apply position embeddings + hidden_state = hidden_state.reshape( + batch_size * num_concurrent_media, num_tiles, num_patches, dim + ) + hidden_state = self.gated_positional_embedding(hidden_state, aspect_ratio_ids) + + # apply encoder + hidden_state = self.layernorm_pre(hidden_state) + + # Compute the number of tokens to pad + num_padding_patches = (8 - (hidden_state.shape[-2] % 8)) % 8 + # Compute padding tuple for pad function + padding = ( + 0, + 0, + 0, + num_padding_patches, + ) # (pad_left, pad_right, pad_left for dim -2, pad_right for dim -2) + # Pad the tensor + hidden_state = F.pad(hidden_state, padding, mode="constant", value=0) + slice_index = -num_padding_patches if num_padding_patches > 0 else None + + attention_mask = aspect_ratio_mask.reshape( + batch_size * num_concurrent_media, -1 + ) + attention_mask = _prepare_aspect_ratio_attention_mask( + aspect_ratio_mask=attention_mask, + num_patches=self.num_patches, + target_length=hidden_state.shape[2], + dtype=self.layernorm_pre.weight.dtype, + ) + + hidden_state = hidden_state.view(batch_size * num_concurrent_media, -1, dim) + output = self.transformer( + hidden_state, + attention_mask=attention_mask, + ) + hidden_state, intermediate_hidden_states = output[0], output[1] + intermediate_hidden_states = torch.stack(intermediate_hidden_states, dim=-1) + + # apply global encoder + hidden_state = self.layernorm_post(hidden_state) + hidden_state = hidden_state.reshape( + batch_size * num_concurrent_media, + num_tiles, + num_patches + num_padding_patches, + dim, + ) + hidden_state = self.post_tile_positional_embedding( + hidden_state, aspect_ratio_ids + ) + hidden_state = hidden_state.reshape( + batch_size * num_concurrent_media, + num_tiles * (num_patches + num_padding_patches), + dim, + ) + hidden_state = self.global_transformer( + hidden_state, attention_mask=attention_mask + )[0] + hidden_state = hidden_state.reshape( + batch_size * num_concurrent_media, + num_tiles, + num_patches + num_padding_patches, + dim, + ) + hidden_state = hidden_state[:, :, :slice_index] + + # adding intermediate layer outputs + hidden_state = hidden_state.reshape( + batch_size, num_concurrent_media, num_tiles, num_patches, dim + ) + intermediate_hidden_states = intermediate_hidden_states.reshape( + batch_size * num_concurrent_media, + num_tiles, + num_patches + num_padding_patches, + -1, + ) + intermediate_hidden_states = intermediate_hidden_states[:, :, :slice_index] + intermediate_hidden_states = intermediate_hidden_states.reshape( + batch_size, num_concurrent_media, num_tiles, num_patches, -1 + ) + hidden_state = torch.cat([hidden_state, intermediate_hidden_states], dim=-1) + return hidden_state + + +class MllamaTextRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +class MllamaTextCrossAttention(nn.Module): + def __init__( + self, + config: Optional[config_mllama.MllamaTextConfig] = None, + layer_id: Optional[int] = None, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.model_parallel_size = get_tensor_model_parallel_world_size() + self.num_heads = self.config.num_attention_heads + self.num_local_heads = self.num_heads // self.model_parallel_size + self.num_key_value_heads = self.config.num_key_value_heads + self.num_local_key_value_heads = ( + self.num_key_value_heads // self.model_parallel_size + ) + self.dropout = config.dropout + self.hidden_size = config.hidden_size + self.head_dim = config.hidden_size // self.num_heads + self.layer_id = layer_id + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.q_local_size = self.num_local_heads * self.head_dim + self.kv_local_size = self.num_local_key_value_heads * self.head_dim + + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.num_heads, + self.num_key_value_heads, + bias=False, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.num_heads * self.head_dim, + self.hidden_size, + bias=False, + input_is_parallel=True, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + # vllm.model_executor.layers.layernorm.RMSNorm has precision issue, + # use huggingface's instead + self.q_norm = MllamaTextRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = MllamaTextRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.scaling = self.head_dim**-0.5 + + self.attn = RadixAttention( + self.num_local_heads, + self.head_dim, + self.scaling, + self.num_local_key_value_heads, + layer_id=layer_id, + is_cross_attention=True, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor], + cross_attention_states: Optional[torch.Tensor], + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv_dec, _ = self.qkv_proj(hidden_states) + q, _, _ = qkv_dec.split( + [self.q_local_size, self.kv_local_size, self.kv_local_size], dim=-1 + ) + if cross_attention_states is None: + k = None + v = None + else: + qkv_enc, _ = self.qkv_proj(cross_attention_states) + _, k, v = qkv_enc.split( + [self.q_local_size, self.kv_local_size, self.kv_local_size], dim=-1 + ) + k = k.view(-1, self.num_local_key_value_heads, self.head_dim) + v = v.view(-1, self.num_local_key_value_heads, self.head_dim) + k = self.k_norm(k) + q = q.view(-1, self.num_local_heads, self.head_dim) + q = self.q_norm(q) + + output = self.attn(q, k, v, forward_batch) + out, _ = self.o_proj(output) + return out + + +class MllamaCrossAttentionDecoderLayer(torch.nn.Module): + """Cross-attention transformer block with tanh-gated attention + and feedforward.""" + + def __init__( + self, + config: config_mllama.MllamaTextConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig], + prefix: str = "", + ) -> None: + super().__init__() + self.layer_id = layer_id + self.cross_attn = MllamaTextCrossAttention( + config=config, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("cross_attn", prefix), + ) + + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.cross_attn_attn_gate = torch.nn.Parameter(torch.zeros(1)) + + self.mlp = LlamaMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.cross_attn_mlp_gate = torch.nn.Parameter(torch.zeros(1)) + + def forward( + self, + hidden_states: torch.Tensor, + cross_attention_states: torch.Tensor, + cross_attention_mask: torch.Tensor, + full_text_row_masked_out_mask: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + hidden_states = self.cross_attn( + hidden_states=hidden_states, + attention_mask=cross_attention_mask, + cross_attention_states=cross_attention_states, + forward_batch=forward_batch, + ) + hidden_states = full_text_row_masked_out_mask * hidden_states + hidden_states = residual + self.cross_attn_attn_gate.tanh() * hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = full_text_row_masked_out_mask * hidden_states + hidden_states = residual + self.cross_attn_mlp_gate.tanh() * hidden_states + return hidden_states + + +class MllamaTextModel(nn.Module): + config_class = config_mllama.MllamaTextConfig + base_model_prefix = "model" + + def __init__( + self, + config: config_mllama.MllamaTextConfig, + quant_config: Optional[QuantizationConfig], + prefix: str = "", + ): + super().__init__() + self.padding_id = config.pad_token_id + self.vocab_size = config.vocab_size + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size + 8, + config.hidden_size, + prefix=add_prefix("embed_tokens", prefix), + ) + self.cross_attention_layers = config.cross_attention_layers + + layers = [] + for layer_id in range(config.num_hidden_layers): + if layer_id in self.cross_attention_layers: + layers.append( + MllamaCrossAttentionDecoderLayer( + config, + layer_id, + quant_config=quant_config, + prefix=add_prefix(f"layers.{layer_id}", prefix), + ) + ) + else: + # TODO: force LlamaDecoderLayer to config.attention_bias=False + layers.append( + LlamaDecoderLayer( + config, + quant_config=quant_config, + layer_id=layer_id, + prefix=add_prefix(f"layers.{layer_id}", prefix), + ) + ) + + self.layers = nn.ModuleList(layers) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.LongTensor, + positions: Optional[torch.LongTensor], + cross_attention_states: Optional[torch.LongTensor], + cross_attention_mask: Optional[torch.LongTensor], + full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor, torch.Tensor]], + forward_batch: ForwardBatch, + skip_cross_attention: bool, + ) -> torch.Tensor: + inputs_embeds = self.embed_tokens(input_ids) + hidden_states = inputs_embeds + + for _, decoder_layer in enumerate(self.layers): + if isinstance(decoder_layer, MllamaCrossAttentionDecoderLayer): + if not skip_cross_attention: + hidden_states = decoder_layer( + hidden_states=hidden_states, + cross_attention_states=cross_attention_states, + cross_attention_mask=cross_attention_mask, + full_text_row_masked_out_mask=full_text_row_masked_out_mask, + forward_batch=forward_batch, + ) + elif isinstance(decoder_layer, LlamaDecoderLayer): + hidden_states, residual = decoder_layer( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + residual=None, + ) + hidden_states = hidden_states + residual + else: + raise ValueError(f"Unknown decoder layer type {type(decoder_layer)}") + hidden_states = self.norm(hidden_states) + return hidden_states + + +class MllamaForCausalLM(nn.Module): + config_class = config_mllama.MllamaTextConfig + base_model_prefix = "language_model" + _no_split_modules = [ + "MllamaCrossAttentionDecoderLayer", + "MllamaSelfAttentionDecoderLayer", + ] + + def __init__( + self, + config: config_mllama.MllamaTextConfig, + quant_config: Optional[QuantizationConfig], + prefix: str = "", + ): + super().__init__() + self.vocab_size = config.vocab_size + self.model = MllamaTextModel( + config, quant_config, prefix=add_prefix("model", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + padding_size=DEFAULT_VOCAB_PADDING_SIZE, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + + def forward( + self, + input_ids: torch.LongTensor, + positions: Optional[torch.LongTensor], + cross_attention_states: Optional[torch.LongTensor], + cross_attention_mask: Optional[torch.LongTensor], + full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor, torch.Tensor]], + forward_batch: ForwardBatch, + skip_cross_attention: bool, + ) -> torch.Tensor: + hidden_states = self.model( + input_ids=input_ids, + positions=positions, + cross_attention_states=cross_attention_states, + cross_attention_mask=cross_attention_mask, + full_text_row_masked_out_mask=full_text_row_masked_out_mask, + forward_batch=forward_batch, + skip_cross_attention=skip_cross_attention, + ) + return hidden_states + + +class MllamaForConditionalGeneration(nn.Module): + # BitandBytes specific attributes + default_bitsandbytes_target_modules = [ + ".gate_proj.", + ".down_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ".o_proj.", + ] + # in TP, these weights are partitioned along the column dimension (dim=-1) + column_parallel_weights_modules = [".down_proj.", ".o_proj."] + bitsandbytes_stacked_params_mapping = { + # shard_name, weight_name, index + "q_proj": ("qkv_proj", 0), + "k_proj": ("qkv_proj", 1), + "v_proj": ("qkv_proj", 2), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def __init__( + self, + config: config_mllama.MllamaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.quant_config = quant_config + self.vocab_size = config.text_config.vocab_size + self.hidden_size = config.text_config.hidden_size + self.max_num_tiles = config.vision_config.max_num_tiles + self.vision_output_dim = config.vision_config.vision_output_dim + self.pad_token_id = ( + config.pad_token_id if config.pad_token_id is not None else -1 + ) + self.image_size = config.vision_config.image_size + + self.vision_model = MllamaVisionModel( + config.vision_config, + quant_config=quant_config, + prefix=add_prefix("vision_model", prefix), + ) + self.language_model = MllamaForCausalLM( + config.text_config, + quant_config=quant_config, + prefix=add_prefix("language_model", prefix), + ) + self.multi_modal_projector = ReplicatedLinear( + config.vision_config.vision_output_dim, + config.text_config.hidden_size, + bias=True, + quant_config=quant_config, + prefix="multi_modal_projector", + ) + self.logits_processor = LogitsProcessor(config.text_config) + self.capture_mode = False + + def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): + pixel_values = torch.cat( + [item.pixel_values for item in mm_inputs.mm_items], dim=0 + ) + pad_values = [item.pad_value for item in mm_inputs.mm_items] + + num_concurrent_media, num_tiles = pixel_values.shape[1:3] + num_patches = self.vision_model.num_patches + image_len = num_concurrent_media * num_tiles * num_patches + mm_inputs.num_image_tokens = image_len + + pad_ids = pad_values * ((image_len + len(pad_values)) // len(pad_values)) + + return pad_ids[:image_len] + input_ids + + def _batch_image_inputs(self, forward_batch: ForwardBatch): + if forward_batch.forward_mode.is_decode() or all(forward_batch.encoder_cached): + return None, None, None, None + + # pixel_values: shape (bs, num_image, num_tiles, 3, image_res, image_res) + max_num_images = max_num_tiles = bs = 0 + for i, mm_input in enumerate(forward_batch.mm_inputs): + + if not forward_batch.encoder_cached[i] and mm_input is not None: + pixel_values = torch.cat( + [item.pixel_values for item in mm_input.mm_items], dim=0 + ) + # max_num_images = max(max_num_images, sum(1 if item.is_image() else 0 for item in mm_input.items)) + max_num_images = max(max_num_images, pixel_values.shape[1]) + + max_num_tiles = max(max_num_tiles, pixel_values.shape[2]) + bs += 1 + + if max_num_images * max_num_tiles * bs == 0: + return None, None, None, None + + with forward_batch.out_cache_loc.device: + batched_images = torch.zeros( + bs, + max_num_images, + max_num_tiles, + 3, + self.image_size, + self.image_size, + dtype=torch.float32, + ) + batched_ar_ids = torch.ones( + bs, max_num_images, dtype=torch.int64, device="cuda" + ) + batched_ar_mask = torch.zeros( + bs, max_num_images, max_num_tiles, dtype=torch.int64 + ) + i = 0 + encoder_lens_need = [] + + for k, mm_input in enumerate(forward_batch.mm_inputs): + if forward_batch.encoder_cached[k] or mm_input is None: + continue + + encoder_lens_need.append(forward_batch.encoder_lens[k]) + pixel_values = torch.cat( + [item.pixel_values for item in mm_input.mm_items], dim=0 + ) + for j in range(pixel_values.shape[1]): + img = pixel_values[0, j] + num_tiles = img.shape[0] + batched_images[i, j, :num_tiles] = img + batched_ar_ids[i, j] = mm_input.mm_items[0].aspect_ratio_id[0, j] + + batched_ar_mask[i, j, :num_tiles] = mm_input.mm_items[ + 0 + ].aspect_ratio_mask[0, j] + i += 1 + + return batched_images, batched_ar_ids, batched_ar_mask, encoder_lens_need + + def flat_encoder_result( + self, cross_attention_states: torch.Tensor, encoder_lens_need: List[int] + ): + # NOTE: not all encoders need computation, some are cached + head_dim = cross_attention_states.shape[-1] + total_encoder_len = sum(encoder_lens_need) + cross_attention_states_flat = torch.zeros( + total_encoder_len, + head_dim, + device=cross_attention_states.device, + dtype=cross_attention_states.dtype, + ) + + i = start_pos = 0 + for encoder_len in encoder_lens_need: + if encoder_len == 0: + continue + end_pos = start_pos + encoder_len + cross_attention_states_flat[start_pos:end_pos] = cross_attention_states[i][ + :encoder_len + ] + i += 1 + start_pos += encoder_len + + return cross_attention_states_flat + + def get_full_text_row_masked_out_mask(self, forward_batch: ForwardBatch): + if forward_batch.forward_mode.is_decode(): + full_text_row_masked_out_mask = forward_batch.encoder_lens != 0 + else: + full_text_row_masked_out_mask = torch.ones( + forward_batch.extend_seq_lens.sum(), dtype=torch.bool + ) + start_pos = 0 + + for seq_len, encoder_len in zip( + forward_batch.seq_lens.tolist(), forward_batch.encoder_lens_cpu + ): + if encoder_len == 0: + full_text_row_masked_out_mask[start_pos : start_pos + seq_len] = ( + False + ) + start_pos += encoder_len + + full_text_row_masked_out_mask = full_text_row_masked_out_mask.to( + forward_batch.seq_lens.device + ) + + return full_text_row_masked_out_mask.reshape(-1, 1) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> Union[Tuple, CausalLMOutputWithPast]: + batched_images, batched_ar_ids, batched_ar_mask, encoder_lens_need = ( + self._batch_image_inputs(forward_batch) + ) + + # TODO: support multi-image by this mask + cross_attention_mask = None + cross_attention_states = None + + if self.capture_mode: + # NOTE: when doing cuda graph capture, we do not want to skip cross attention + # Make is a constant value to avoid cuda graph capture issue + skip_cross_attention = False + else: + # NOTE: we do not need image_inputs when prefill + assert len(forward_batch.encoder_lens) == len(forward_batch.seq_lens) + assert len(forward_batch.encoder_lens_cpu) == len(forward_batch.seq_lens) + skip_cross_attention = forward_batch.encoder_lens.max() == 0 + + if not skip_cross_attention: + full_text_row_masked_out_mask = self.get_full_text_row_masked_out_mask( + forward_batch + ) + else: + full_text_row_masked_out_mask = None + + if batched_images is not None: + # NOTE: llama's reference implementation runs vision model on CPU + cross_attention_states = self.vision_model( + batched_images, batched_ar_ids, batched_ar_mask + ) + cross_attention_states, _ = self.multi_modal_projector( + cross_attention_states + ) + + bs, _, _, _, image_token_dim = cross_attention_states.shape + cross_attention_states = cross_attention_states.view( + bs, -1, image_token_dim + ) + + cross_attention_states = self.flat_encoder_result( + cross_attention_states, encoder_lens_need + ) + + hidden_states = self.language_model( + input_ids=input_ids, + positions=positions, + cross_attention_states=cross_attention_states, + cross_attention_mask=cross_attention_mask, + full_text_row_masked_out_mask=full_text_row_masked_out_mask, + forward_batch=forward_batch, + skip_cross_attention=skip_cross_attention, + ) + return self.logits_processor( + input_ids, hidden_states, self.language_model.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + params_dict = dict(self.named_parameters()) + updated_params = set() + for name, loaded_weight in weights: + if "patch_embedding.weight" in name: + name = name.replace( + "patch_embedding.weight", "patch_embedding._linear.weight" + ) + loaded_weight = loaded_weight.view(loaded_weight.shape[0], -1) + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + updated_params.add(name) + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + if "vision_model" in name: + # adapt to VisionAttention + name = name.replace("self_attn.o_proj", "self_attn.proj") + param = params_dict.pop(name) + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +EntryClass = MllamaForConditionalGeneration diff --git a/lib/python3.12/site-packages/sglang/srt/models/mllama4.py b/lib/python3.12/site-packages/sglang/srt/models/mllama4.py new file mode 100644 index 0000000000000000000000000000000000000000..21b6c51d6a0c1a890103d07495c153a79b9baff1 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/mllama4.py @@ -0,0 +1,227 @@ +from collections.abc import Iterable +from typing import List, Optional, Set, Tuple + +import torch +from torch import nn +from transformers import Llama4Config, Llama4VisionModel +from transformers.models.llama4.modeling_llama4 import Llama4MultiModalProjector + +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.fused_moe_triton import FusedMoE +from sglang.srt.layers.quantization import QuantizationConfig +from sglang.srt.managers.mm_utils import ( + MultiModalityDataPaddingPatternMultimodalTokens, + general_mm_embed_routine, +) +from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix + + +class Llama4ForConditionalGeneration(nn.Module): + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + def __init__( + self, + config: Llama4Config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.quant_config = quant_config + + self.vision_model = Llama4VisionModel(config.vision_config) + self.multi_modal_projector = Llama4MultiModalProjector(config) + + # Initialize the language model + from sglang.srt.models.llama4 import Llama4ForCausalLM + + self.language_model = Llama4ForCausalLM( + config.text_config, + quant_config=quant_config, + prefix=add_prefix("language_model", prefix), + ) + + self.logits_processor = LogitsProcessor(config.text_config) + + def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): + # Get all special token IDs + im_token_id: int = mm_inputs.im_token_id + + pattern = MultiModalityDataPaddingPatternMultimodalTokens([im_token_id]) + return pattern.pad_input_tokens(input_ids, mm_inputs) + + def get_image_feature( + self, + items: List[MultimodalDataItem], + ) -> torch.Tensor: + pixel_values = ( + torch.concat([item.pixel_values for item in items]) + .to(next(self.vision_model.parameters()).device) + .type(next(self.vision_model.parameters()).dtype) + ) + + image_outputs = self.vision_model(pixel_values, output_hidden_states=False) + image_features = image_outputs.last_hidden_state + vision_flat = image_features.view(-1, image_features.size(-1)) + projected_vision_flat = self.multi_modal_projector(vision_flat) + return projected_vision_flat + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + **kwargs: object, + ) -> torch.Tensor: + + hs = general_mm_embed_routine( + input_ids=input_ids, + forward_batch=forward_batch, + language_model=self.language_model, + image_data_embedding_func=self.get_image_feature, + positions=positions, + ) + + return hs + + def permute_qk_weight_for_rotary( + self, + name: str, + loaded_weight: torch.Tensor, + ) -> Tuple[str, torch.Tensor]: + + def permute(w: torch.Tensor, n_heads: int): + attn_in = self.language_model.config.head_dim * n_heads + attn_out = self.language_model.config.hidden_size + + return ( + w.view(n_heads, attn_in // n_heads // 2, 2, attn_out) + .transpose(1, 2) + .reshape(attn_in, attn_out) + ) + + modules = name.split(".") + + # rotary embeds should be sliced + if ("wk" in modules or "k_proj" in modules) and modules[-1] == "weight": + loaded_weight = permute( + loaded_weight, self.language_model.config.num_key_value_heads + ) + elif ("wq" in modules or "q_proj" in modules) and modules[-1] == "weight": + loaded_weight = permute( + loaded_weight, self.language_model.config.num_attention_heads + ) + + return name, loaded_weight + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]: + + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".self_attn.qkv_proj", ".self_attn.q_proj", "q"), + (".self_attn.qkv_proj", ".self_attn.k_proj", "k"), + (".self_attn.qkv_proj", ".self_attn.v_proj", "v"), + (".shared_expert.gate_up_proj", ".shared_expert.gate_proj", 0), + (".shared_expert.gate_up_proj", ".shared_expert.up_proj", 1), + (".feed_forward.gate_up_proj", ".feed_forward.gate_proj", 0), + (".feed_forward.gate_up_proj", ".feed_forward.up_proj", 1), + ] + + params_dict = dict(self.named_parameters()) + + num_experts = self.config.text_config.num_local_experts + + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=num_experts, + ) + + for name, loaded_weight in weights: + if not "vision" in name: + name, loaded_weight = self.permute_qk_weight_for_rotary( + name, loaded_weight + ) + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + + if "vision" in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + if ".experts" in name: + # NOTE: llama4 fp8 has different weight format for experts + if ( + "experts.gate_up_proj" not in name + and "experts.down_proj" not in name + ): + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + break + else: + if ".gate_up_proj" in name: + name_list = [ + name.replace( + ".experts.gate_up_proj", ".experts.w13_weight" + ) + ] * 2 + loaded_weight_list = loaded_weight.chunk(2, dim=-1) + shard_id_list = ["w1", "w3"] + else: + name_list = [ + name.replace(".experts.down_proj", ".experts.w2_weight") + ] + shard_id_list = ["w2"] + loaded_weight_list = [loaded_weight] + for name, loaded_weight, shard_id in zip( + name_list, loaded_weight_list, shard_id_list + ): + param = params_dict[name] + weight_loader = param.weight_loader + for expert_id in range(num_experts): + weight_loader( + param, + loaded_weight[expert_id].T, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + +EntryClass = Llama4ForConditionalGeneration diff --git a/lib/python3.12/site-packages/sglang/srt/models/olmo.py b/lib/python3.12/site-packages/sglang/srt/models/olmo.py new file mode 100644 index 0000000000000000000000000000000000000000..08239374daa7f2de872e4dc099b2138fa26fa333 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/olmo.py @@ -0,0 +1,377 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from +# https://github.com/vllm-project/vllm/blob/c7f2cf2b7f67bce5842fedfdba508440fe257375/vllm/model_executor/models/olmo.py#L1 +"""Inference-only OLMo model compatible with HuggingFace weights.""" +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import OlmoConfig + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix, make_layers + + +class OlmoAttention(nn.Module): + """ + This is the attention block where the output is computed as + ``Attention(LN(x))`` in ``MLP(LN(x + Attention(LN(x))))`` + (plus another skip connection). + """ + + def __init__( + self, + config: OlmoConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + tensor_model_parallel_world_size = get_tensor_model_parallel_world_size() + self.total_num_heads = config.num_attention_heads + + assert self.hidden_size % self.total_num_heads == 0 + assert self.total_num_heads % tensor_model_parallel_world_size == 0 + + self.num_heads = self.total_num_heads // tensor_model_parallel_world_size + self.head_dim = self.hidden_size // self.total_num_heads + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + self.clip_qkv = config.clip_qkv + + # Attention input projection. Projects x -> (q, k, v) + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + bias=config.attention_bias, + prefix=add_prefix("qkv_proj", prefix), + ) + + # Rotary embeddings. + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=self.max_position_embeddings, + base=self.rope_theta, + ) + self.scaling = self.head_dim**-0.5 + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + # Attention output projection. + self.o_proj = RowParallelLinear( + self.hidden_size, + self.hidden_size, + bias=config.attention_bias, + prefix=add_prefix("o_proj", prefix), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + if self.clip_qkv is not None: + qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv) + q, k, v = qkv.chunk(chunks=3, dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class OlmoMLP(nn.Module): + """ + This is the MLP block where the output is computed as + ``MLP(LN(x))`` in ``MLP(LN(x + Attention(LN(x))))`` + (plus another skip connection). + """ + + def __init__( + self, + config: OlmoConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + + # Feed-forward input projection. + self.gate_up_proj = MergedColumnParallelLinear( + self.hidden_size, + [self.intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + + # Activation function. + self.act_fn = SiluAndMul() + + # Feed-forward output projection. + self.down_proj = RowParallelLinear( + self.intermediate_size, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + + def forward( + self, + x: torch.Tensor, + ) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class OlmoDecoderLayer(nn.Module): + """ + This is a typical transformer block where the output is + computed as ``MLP(LN(x + Attention(LN(x))))`` + (plus another skip connection). + """ + + def __init__( + self, + config: OlmoConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + # Attention block. + self.self_attn = OlmoAttention( + config, + layer_id, + quant_config, + prefix=add_prefix("self_attn", prefix), + ) + + # MLP block. + self.mlp = OlmoMLP( + config, + quant_config, + prefix=add_prefix("mlp", prefix), + ) + + # LayerNorm + self.input_layernorm = nn.LayerNorm( + config.hidden_size, elementwise_affine=False, bias=False + ) + self.post_attention_layernorm = nn.LayerNorm( + config.hidden_size, elementwise_affine=False, bias=False + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]: + # Attention block. + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(positions, hidden_states, forward_batch) + hidden_states = hidden_states + residual + + # MLP block. + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +class OlmoModel(nn.Module): + + def __init__( + self, + config: OlmoConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=add_prefix("embed_tokens", prefix), + ) + self.layers = make_layers( + config.num_hidden_layers, + lambda idx, prefix: OlmoDecoderLayer( + layer_id=idx, + config=config, + quant_config=quant_config, + prefix=prefix, + ), + prefix=add_prefix("layers", prefix), + ) + self.norm = nn.LayerNorm( + config.hidden_size, elementwise_affine=False, bias=False + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + """ + :param input_ids: A tensor of shape `(batch_size, seq_len)`. + """ + # Get embeddings of input. + # shape: (batch_size, seq_len, d_model) + + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + + # Apply blocks one-by-one. + for layer_id, decoder_layer in enumerate(self.layers): + # shape: (batch_size, seq_len, d_model) + hidden_states = decoder_layer( + positions, + hidden_states, + forward_batch, + ) + + # Apply final layer norm. + # shape: (batch_size, seq_len or 1, d_model) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class OlmoForCausalLM(nn.Module): + """ + Extremely barebones HF model wrapper. + """ + + def __init__( + self, + config: OlmoConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.model = OlmoModel(config, quant_config, prefix=add_prefix("model", prefix)) + if config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.unpadded_vocab_size = config.vocab_size + self.lm_head = ParallelLMHead( + self.unpadded_vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + self.logits_processor = LogitsProcessor(config) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + hidden_states = self.model( + input_ids=input_ids, + positions=positions, + forward_batch=forward_batch, + input_embeds=input_embeds, + ) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.named_parameters(remove_duplicate=False)) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + if self.config.tie_word_embeddings and "lm_head.weight" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +EntryClass = OlmoForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/olmoe.py b/lib/python3.12/site-packages/sglang/srt/models/olmoe.py new file mode 100644 index 0000000000000000000000000000000000000000..612120fe939b85a224ed99f1b7857817abbb5be9 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/olmoe.py @@ -0,0 +1,432 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from: +# https://github.com/vllm-project/vllm/pull/7922 + +"""Inference-only OLMoE model compatible with HuggingFace weights.""" + +from typing import Any, Dict, Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.fused_moe_triton import FusedMoE +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix, make_layers, print_warning_once + + +class OlmoeMoE(nn.Module): + """A tensor-parallel MoE implementation for Olmoe that shards each expert + across all ranks. + + Each expert's weights are sharded across all ranks and a fused MoE + kernel is used for the forward pass, and finally we reduce the outputs + across ranks. + """ + + def __init__( + self, + num_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + params_dtype: Optional[torch.dtype] = None, + quant_config: Optional[QuantizationConfig] = None, + tp_size: Optional[int] = None, + prefix: str = "", + ): + super().__init__() + self.hidden_size = hidden_size + + # Gate always runs at half / full precision for now. + self.gate = ReplicatedLinear( + hidden_size, + num_experts, + bias=False, + quant_config=None, + prefix=add_prefix("gate", prefix), + ) + + self.experts = FusedMoE( + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + reduce_results=True, + renormalize=False, + quant_config=quant_config, + tp_size=tp_size, + prefix=add_prefix("experts", prefix), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # NOTE: hidden_states can have either 1D or 2D shape. + orig_shape = hidden_states.shape + hidden_states = hidden_states.view(-1, self.hidden_size) + # router_logits: (num_tokens, n_experts) + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=router_logits + ) + return final_hidden_states.view(orig_shape) + + +class OlmoeAttention(nn.Module): + + def __init__( + self, + layer_id: int, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + max_position_embeddings: int = 4096, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.q_norm = RMSNorm(hidden_size, eps=1e-5) + self.k_norm = RMSNorm(hidden_size, eps=1e-5) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + is_neox_style=True, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + layer_id=layer_id, + num_kv_heads=self.num_kv_heads, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.q_norm(q.contiguous()), self.k_norm(k.contiguous()) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class OlmoeDecoderLayer(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 4096) + + self.self_attn = OlmoeAttention( + layer_id, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + + self.mlp = OlmoeMoE( + num_experts=config.num_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=1e-5) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=1e-5) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> torch.Tensor: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class OlmoeModel(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=add_prefix("embed_tokens", prefix), + ) + self.layers = make_layers( + config.num_hidden_layers, + lambda idx, prefix: OlmoeDecoderLayer( + config=config, + quant_config=quant_config, + layer_id=idx, + prefix=prefix, + ), + prefix=add_prefix("layers", prefix), + ) + self.norm = RMSNorm(config.hidden_size, eps=1e-5) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, hidden_states, forward_batch, residual + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class OlmoeForCausalLM(nn.Module): + + fall_back_to_pt_during_load = False + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.model = OlmoeModel( + config, quant_config, prefix=add_prefix("model", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + self.logits_processor = LogitsProcessor(config) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + ) + + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip non-stacked layers and experts (experts handled below). + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since we handle the experts below in expert_params_mapping, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below in expert_params_mapping + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if "mlp.experts" in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Remapping the name of FP8 kv-scale. + if name.endswith("kv_scale"): + remapped_kv_scale_name = name.replace( + ".kv_scale", ".attn.kv_scale" + ) + if remapped_kv_scale_name not in params_dict: + print_warning_once( + "Found kv scale in the checkpoint " + f"(e.g. {name}), but not found the expected " + f"name in the model " + f"(e.g. {remapped_kv_scale_name}). " + "kv-scale is not loaded." + ) + continue + else: + name = remapped_kv_scale_name + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + +EntryClass = OlmoeForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/qwen2_5_vl.py b/lib/python3.12/site-packages/sglang/srt/models/qwen2_5_vl.py new file mode 100644 index 0000000000000000000000000000000000000000..17697a83554368f3779d56bf006743afe88f1060 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/qwen2_5_vl.py @@ -0,0 +1,609 @@ +# coding=utf-8 +# Adapted from +# https://github.com/huggingface/transformers/blob/19e6e80e10118f855137b90740936c0b11ac397f/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py +# Copyright 2024 The Qwen team. +# Copyright 2023 The vLLM team. +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Inference-only Qwen2-VL model compatible with HuggingFace weights.""" +import logging +from functools import lru_cache, partial +from typing import Iterable, List, Optional, Tuple, Type + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from transformers.activations import ACT2FN +from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm +from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import ( + Qwen2_5_VLConfig, + Qwen2_5_VLVisionConfig, +) +from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( + Qwen2_5_VisionPatchEmbed, + Qwen2_5_VisionRotaryEmbedding, +) + +from sglang.srt.hf_transformers_utils import get_processor +from sglang.srt.layers.attention.vision import VisionAttention +from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.pooler import Pooler, PoolingType +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead +from sglang.srt.managers.mm_utils import ( + MultiModalityDataPaddingPatternMultimodalTokens, + general_mm_embed_routine, +) +from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.qwen2 import Qwen2Model +from sglang.srt.models.qwen2_vl import Qwen2VLVideoInputs +from sglang.srt.utils import add_prefix + +logger = logging.getLogger(__name__) + + +class Qwen2_5_VLMLP(nn.Module): + + def __init__( + self, + in_features: int, + hidden_features: int = None, + bias: bool = True, + hidden_act="silu", + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.gate_proj = ColumnParallelLinear( + in_features, + hidden_features, + bias=bias, + quant_config=quant_config, + prefix=add_prefix("gate_proj", prefix), + ) + self.up_proj = ColumnParallelLinear( + in_features, + hidden_features, + bias=bias, + quant_config=quant_config, + prefix=add_prefix("up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + hidden_features, + in_features, + bias=bias, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + self.act = ACT2FN[hidden_act] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x_parallel_gate, _ = self.gate_proj(x) + x_parallel_gate = self.act(x_parallel_gate) + x_parallel_up, _ = self.up_proj(x) + x_parallel = x_parallel_gate * x_parallel_up + x, _ = self.down_proj(x_parallel) + return x + + +class Qwen2_5_VisionBlock(nn.Module): + + def __init__( + self, + dim: int, + intermediate_dim: int, + num_heads: int, + hidden_act="silu", + norm_layer: Type[nn.Module] = None, + attn_implementation: Optional[str] = "sdpa", + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + if norm_layer is None: + norm_layer = partial(nn.LayerNorm, eps=1e-6) + self.norm1 = Qwen2RMSNorm(dim, eps=1e-6) + self.norm2 = Qwen2RMSNorm(dim, eps=1e-6) + if attn_implementation == "sdpa": + use_context_forward = False + softmax_in_single_precision = False + flatten_batch = True + elif attn_implementation == "flash_attention_2": + softmax_in_single_precision = False + use_context_forward = True + flatten_batch = True + elif attn_implementation == "eager": + softmax_in_single_precision = True + use_context_forward = False + flatten_batch = True + + self.attn = VisionAttention( + embed_dim=dim, + num_heads=num_heads, + projection_size=dim, + use_qkv_parallel=True, + use_context_forward=use_context_forward, + softmax_in_single_precision=softmax_in_single_precision, + flatten_batch=flatten_batch, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + self.mlp = Qwen2_5_VLMLP( + dim, + intermediate_dim, + hidden_act=hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + + def forward( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + position_embeddings: torch.Tensor, + ) -> torch.Tensor: + hidden_states = self.norm1(x) + hidden_states = rearrange(hidden_states, "s b ... -> b s ...") + attn = self.attn( + hidden_states, + cu_seqlens=cu_seqlens, + position_embeddings=position_embeddings, + ) + attn = rearrange(attn, "b s ... -> s b ...") + x = x + attn + norm2 = self.norm2(x) + mlp = self.mlp(norm2) + x = x + mlp + return x + + +class Qwen2_5_VisionPatchMerger(nn.Module): + + def __init__( + self, + dim: int, + context_dim: int, + spatial_merge_size: int = 2, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = context_dim * (spatial_merge_size**2) + self.ln_q = Qwen2RMSNorm(context_dim, eps=1e-6) + self.mlp = nn.ModuleList( + [ + ColumnParallelLinear( + self.hidden_size, + self.hidden_size, + bias=True, + quant_config=quant_config, + prefix=add_prefix("mlp.0", prefix), + ), + nn.GELU(), + RowParallelLinear( + self.hidden_size, + dim, + bias=True, + quant_config=quant_config, + prefix=add_prefix("mlp.2", prefix), + ), + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.ln_q(x) + x = x.view(-1, self.hidden_size) + + mlp_fc1, mlp_act, mlp_fc2 = self.mlp + x_parallel, _ = mlp_fc1(x) + x_parallel = mlp_act(x_parallel) + out, _ = mlp_fc2(x_parallel) + return out + + +class Qwen2_5_VisionTransformer(nn.Module): + + def __init__( + self, + vision_config: Qwen2_5_VLVisionConfig, + norm_eps: float = 1e-6, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + patch_size: int = vision_config.patch_size + temporal_patch_size: int = vision_config.temporal_patch_size + spatial_merge_size: int = vision_config.spatial_merge_size + self.spatial_merge_size = spatial_merge_size + self.spatial_merge_unit: int = spatial_merge_size * spatial_merge_size + in_channels: int = vision_config.in_channels + hidden_size: int = vision_config.hidden_size + depth: int = vision_config.depth + num_heads: int = vision_config.num_heads + self.fullatt_block_indexes = vision_config.fullatt_block_indexes + self.window_size = vision_config.window_size + self.patch_size = vision_config.patch_size + mlp_hidden_size: int = vision_config.intermediate_size + self.patch_embed = Qwen2_5_VisionPatchEmbed( + patch_size=patch_size, + temporal_patch_size=temporal_patch_size, + in_channels=in_channels, + embed_dim=hidden_size, + ) + + norm_layer = partial(nn.LayerNorm, eps=norm_eps) + head_dim = hidden_size // num_heads + self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2) + self.blocks = nn.ModuleList( + [ + Qwen2_5_VisionBlock( + dim=hidden_size, + intermediate_dim=mlp_hidden_size, + num_heads=num_heads, + hidden_act=vision_config.hidden_act, + norm_layer=norm_layer, + attn_implementation="sdpa", + quant_config=quant_config, + prefix=add_prefix(f"blocks.{i}", prefix), + ) + for i in range(depth) + ] + ) + self.merger = Qwen2_5_VisionPatchMerger( + dim=vision_config.out_hidden_size, + context_dim=hidden_size, + spatial_merge_size=spatial_merge_size, + quant_config=quant_config, + prefix=add_prefix("merger", prefix), + ) + + def get_window_index(self, grid_thw): + cu_window_seqlens: list = [0] + window_index_id = 0 + vit_merger_window_size = ( + self.window_size // self.spatial_merge_size // self.patch_size + ) + window_index: list = [] + for grid_t, grid_h, grid_w in grid_thw: + llm_grid_h, llm_grid_w = ( + grid_h // self.spatial_merge_size, + grid_w // self.spatial_merge_size, + ) + index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape( + grid_t, llm_grid_h, llm_grid_w + ) + pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size + pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size + num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size + num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size + index_padded = F.pad(index, (0, pad_w, 0, pad_h), "constant", -100) + index_padded = index_padded.reshape( + grid_t, + num_windows_h, + vit_merger_window_size, + num_windows_w, + vit_merger_window_size, + ) + index_padded = index_padded.permute(0, 1, 3, 2, 4).reshape( + grid_t, + num_windows_h * num_windows_w, + vit_merger_window_size, + vit_merger_window_size, + ) + seqlens = (index_padded != -100).sum([2, 3]).reshape(-1) + index_padded = index_padded.reshape(-1) + index_new = index_padded[index_padded != -100] + window_index.append(index_new + window_index_id) + cu_seqlens_tmp = ( + seqlens.cumsum(0) * self.spatial_merge_unit + cu_window_seqlens[-1] + ) + cu_window_seqlens.extend(cu_seqlens_tmp.tolist()) + window_index_id += (grid_t * llm_grid_h * llm_grid_w).item() + window_index = torch.cat(window_index, dim=0) + return window_index, cu_window_seqlens + + @property + def dtype(self) -> torch.dtype: + return self.patch_embed.proj.weight.dtype + + @property + def device(self) -> torch.device: + return self.blocks[0].mlp.gate_proj.weight.device + + def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor: + pos_ids = [] + for i in range(grid_thw.size(0)): + t, h, w = grid_thw[i].tolist() + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) + + hpos_ids = hpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + hpos_ids = hpos_ids.permute(0, 2, 1, 3) + hpos_ids = hpos_ids.flatten() + + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) + wpos_ids = wpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + wpos_ids = wpos_ids.permute(0, 2, 1, 3) + wpos_ids = wpos_ids.flatten() + + pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) + pos_ids = torch.cat(pos_ids, dim=0) + max_grid_size = grid_thw[:, 1:].max() + rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) + rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) + return rotary_pos_emb + + def forward( + self, + x: torch.Tensor, + grid_thw: torch.Tensor, + ) -> torch.Tensor: + # patchify + x = x.to(device=self.device, dtype=self.dtype) + x = self.patch_embed(x) + + # compute position embedding + rotary_pos_emb = self.rot_pos_emb(grid_thw) + + window_index, cu_window_seqlens = self.get_window_index(grid_thw) + cu_window_seqlens = torch.tensor( + cu_window_seqlens, + device=x.device, + dtype=torch.int32, + ) + cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens) + + seq_len, _ = x.size() + + x = x.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + x = x[window_index, :, :] + x = x.reshape(seq_len, -1) + rotary_pos_emb = rotary_pos_emb.reshape( + seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1 + ) + rotary_pos_emb = rotary_pos_emb[window_index, :, :] + rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) + + # compute cu_seqlens + cu_seqlens = torch.cat( + [ + torch.tensor([0], device=grid_thw.device), + (grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).cumsum(dim=0), + ] + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), "constant", 0) + + # transformers + x = x.unsqueeze(1) + for layer_num, blk in enumerate(self.blocks): + if layer_num in self.fullatt_block_indexes: + cu_seqlens_now = cu_seqlens + else: + cu_seqlens_now = cu_window_seqlens + x = blk( + x, cu_seqlens=cu_seqlens_now, position_embeddings=position_embeddings + ) + + # adapter + x = self.merger(x) + + reverse_indices = torch.argsort(window_index) + x = x[reverse_indices, :] + + return x + + +cached_get_processor = lru_cache(get_processor) + + +class Qwen2_5_VLForConditionalGeneration(nn.Module): + # BitandBytes specific attributes + default_bitsandbytes_target_modules = [ + ".gate_proj.", + ".down_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ".o_proj.", + ] + bitsandbytes_stacked_params_mapping = { + # shard_name, weight_name, index + "q_proj": ("qkv_proj", 0), + "k_proj": ("qkv_proj", 1), + "v_proj": ("qkv_proj", 2), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def __init__( + self, + config: Qwen2_5_VLConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + self.visual = Qwen2_5_VisionTransformer( + config.vision_config, + norm_eps=getattr(config, "rms_norm_eps", 1e-6), + # NOTE: Qwen2_5-VL vision encoder currently supports BitsAndBytes 4-bit quantization. + # Other quantization methods (e.g., GPTQ, AWQ) are untested and may not be supported. + quant_config=quant_config, + prefix=add_prefix("visual", prefix), + ) + + self.model = Qwen2Model( + config, + quant_config, + prefix=add_prefix("model", prefix), + ) + + if config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + self.is_mrope_enabled = "mrope_section" in self.config.rope_scaling + + self.logits_processor = LogitsProcessor(config) + self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) + + def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): + # Get all special token IDs + im_token_id: int = mm_inputs.im_token_id + pattern = MultiModalityDataPaddingPatternMultimodalTokens([im_token_id]) + return pattern.pad_input_tokens(input_ids, mm_inputs) + + def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + # in qwen-vl, last dim is the same + pixel_values = torch.cat([item.pixel_values for item in items], dim=0).type( + self.visual.dtype + ) + image_grid_thws = torch.concat([item.image_grid_thws for item in items], dim=0) + assert pixel_values.dim() == 2, pixel_values.dim() + assert image_grid_thws.dim() == 2, image_grid_thws.dim() + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thws) + return image_embeds + + def _process_video_input(self, video_input: Qwen2VLVideoInputs) -> torch.Tensor: + pixel_values_videos = video_input["pixel_values_videos"].type(self.visual.dtype) + video_embeds = self.visual( + pixel_values_videos, grid_thw=video_input["video_grid_thw"] + ) + return video_embeds + + def get_input_embeddings(self): + return self.model.embed_tokens + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + get_embedding: bool = False, + ): + """Run forward pass for Qwen2_5-VL. + + Args: + input_ids: Flattened (concatenated) input_ids corresponding to a + batch. + positions: Flattened (concatenated) position ids corresponding to a + batch. + **NOTE**: If mrope is enabled (default setting for Qwen2-VL + opensource models), the shape will be `(3, seq_len)`, + otherwise it will be `(seq_len,). + (Use input_metadata.mrope_positions to replace it) + """ + if self.is_mrope_enabled: + positions = forward_batch.mrope_positions + + if not ( + forward_batch.forward_mode.is_decode() + or not forward_batch.contains_image_inputs() + ): + if self.is_mrope_enabled: + assert positions.ndim == 2 and positions.size(0) == 3, ( + "multimodal section rotary embedding requires " + f"(3, seq_len) positions, but got {positions.size()}" + ) + + hidden_states = general_mm_embed_routine( + input_ids=input_ids, + forward_batch=forward_batch, + language_model=self.model, + image_data_embedding_func=self.get_image_feature, + positions=positions, + ) + + if not get_embedding: + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + else: + return self.pooler(hidden_states, forward_batch) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + ("gate_up_proj", "up_proj", 1), + ("gate_up_proj", "gate_proj", 0), + ] + params_dict = dict(self.named_parameters(remove_duplicate=False)) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "visual" in name: + continue + name = name.replace(weight_name, param_name) + + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + if "visual" in name: + # adapt to VisionAttention + name = name.replace(r"attn.qkv.", r"attn.qkv_proj.") + + try: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + except KeyError: + print(params_dict.keys()) + raise + + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +EntryClass = [Qwen2_5_VLForConditionalGeneration] diff --git a/lib/python3.12/site-packages/sglang/srt/models/qwen2_classification.py b/lib/python3.12/site-packages/sglang/srt/models/qwen2_classification.py new file mode 100644 index 0000000000000000000000000000000000000000..366213e70b8f4d0d6e2b8c02e273ae765f4ff491 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/qwen2_classification.py @@ -0,0 +1,75 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import Qwen2Config + +from sglang.srt.layers.pooler import EmbeddingPoolerOutput, Pooler, PoolingType +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.models.qwen2 import Qwen2ForCausalLM, Qwen2Model +from sglang.srt.utils import add_prefix + + +class Qwen2ForSequenceClassification(nn.Module): + def __init__( + self, + config: Qwen2Config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.model = Qwen2Model( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) + ) + self.score = nn.Linear(config.hidden_size, config.num_labels) + self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=False) + + self.eos_token_id = config.eos_token_id + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + get_embedding: bool = True, + ) -> EmbeddingPoolerOutput: + assert ( + get_embedding + ), "Qwen2ForSequenceClassification is only used for embedding" + + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + logits = self.score(hidden_states) + pooled_logits = self.pooler(logits, forward_batch).embeddings + + return EmbeddingPoolerOutput(pooled_logits) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + # Filter out lm_head weights of Qwen2ForCausalLM + filtered_weights = [ + (name, w) for name, w in weights if not name.startswith("lm_head") + ] + return Qwen2ForCausalLM.load_weights(self, filtered_weights) + + +EntryClass = [ + Qwen2ForSequenceClassification, +] diff --git a/lib/python3.12/site-packages/sglang/srt/models/qwen2_rm.py b/lib/python3.12/site-packages/sglang/srt/models/qwen2_rm.py new file mode 100644 index 0000000000000000000000000000000000000000..f5ed9eae23a3ee05f06265a88724cf7ac5043136 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/qwen2_rm.py @@ -0,0 +1,78 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import Qwen2Config + +from sglang.srt.layers.pooler import EmbeddingPoolerOutput, Pooler, PoolingType +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.models.qwen2 import Qwen2ForCausalLM, Qwen2Model +from sglang.srt.utils import add_prefix + + +class Qwen2ForRewardModel(nn.Module): + def __init__( + self, + config: Qwen2Config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.num_labels = 1 + self.model = Qwen2Model( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) + ) + self.score = nn.Sequential( + nn.Linear(config.hidden_size, config.hidden_size), + nn.ReLU(), + nn.Linear(config.hidden_size, self.num_labels), + ) + self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=False) + + self.eos_token_id = config.eos_token_id + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + get_embedding: bool = True, + ) -> EmbeddingPoolerOutput: + assert get_embedding, "Qwen2ForRewardModel is only used for embedding" + + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + logits = self.score(hidden_states) + pooled_logits = self.pooler(logits, forward_batch).embeddings + + return EmbeddingPoolerOutput(pooled_logits) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + # Filter out lm_head weights of Qwen2ForCausalLM + filtered_weights = [ + (name, w) for name, w in weights if not name.startswith("lm_head") + ] + return Qwen2ForCausalLM.load_weights(self, filtered_weights) + + +EntryClass = [ + Qwen2ForRewardModel, +] diff --git a/lib/python3.12/site-packages/sglang/srt/models/qwen2_vl.py b/lib/python3.12/site-packages/sglang/srt/models/qwen2_vl.py new file mode 100644 index 0000000000000000000000000000000000000000..d29b775fc33376fd3ffb7488e035df7d43696682 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/qwen2_vl.py @@ -0,0 +1,613 @@ +# coding=utf-8 +# Adapted from +# https://github.com/huggingface/transformers/blob/19e6e80e10118f855137b90740936c0b11ac397f/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py +# Copyright 2024 The Qwen team. +# Copyright 2023 The vLLM team. +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Inference-only Qwen2-VL model compatible with HuggingFace weights.""" +import logging +from functools import lru_cache, partial +from typing import Iterable, List, Optional, Tuple, Type, TypedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from transformers import Qwen2VLConfig +from transformers.models.qwen2_vl.configuration_qwen2_vl import Qwen2VLVisionConfig + +from sglang.srt.hf_transformers_utils import get_processor +from sglang.srt.layers.activation import QuickGELU +from sglang.srt.layers.attention.vision import VisionAttention +from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.pooler import Pooler, PoolingType +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead +from sglang.srt.managers.mm_utils import ( + MultiModalityDataPaddingPatternMultimodalTokens, + general_mm_embed_routine, +) +from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.qwen2 import Qwen2Model +from sglang.srt.utils import add_prefix + +logger = logging.getLogger(__name__) + + +# === Vision Inputs === # + + +class Qwen2VLImageInputs(TypedDict): + pixel_values: torch.Tensor + """Shape: + `(num_patches, num_channels * patch_size * patch_size)` + """ + + image_grid_thw: torch.Tensor + """Shape: `(num_images, 3)` + + This should be in `(grid_t, grid_h, grid_w)` format. + """ + + +class Qwen2VLVideoInputs(TypedDict): + pixel_values_videos: torch.Tensor + """Shape: + `(num_patches, + num_channels * temporal_patch_size * patch_size * patch_size)` + """ + + video_grid_thw: torch.Tensor + """Shape: `(num_videos, 3)` + + This should be in `(grid_t, grid_h, grid_w)` format. + """ + + +# === Vision Encoder === # + + +class Qwen2VisionMLP(nn.Module): + + def __init__( + self, + in_features: int, + hidden_features: int = None, + act_layer: Type[nn.Module] = QuickGELU, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.fc1 = ColumnParallelLinear( + in_features, + hidden_features, + quant_config=quant_config, + prefix=add_prefix("fc1", prefix), + ) + self.act = act_layer() + self.fc2 = RowParallelLinear( + hidden_features, + in_features, + quant_config=quant_config, + prefix=add_prefix("fc2", prefix), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x_parallel, _ = self.fc1(x) + x_parallel = self.act(x_parallel) + x, _ = self.fc2(x_parallel) + return x + + +class Qwen2VisionBlock(nn.Module): + + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float, + act_layer: Type[nn.Module] = QuickGELU, + norm_layer: Type[nn.Module] = None, + attn_implementation: Optional[str] = "sdpa", + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + if norm_layer is None: + norm_layer = partial(nn.LayerNorm, eps=1e-6) + self.norm1 = norm_layer(dim) + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + if attn_implementation == "sdpa": + use_context_forward = False + softmax_in_single_precision = False + elif attn_implementation == "flash_attention_2": + softmax_in_single_precision = False + use_context_forward = True + elif attn_implementation == "eager": + softmax_in_single_precision = True + use_context_forward = False + + self.attn = VisionAttention( + embed_dim=dim, + num_heads=num_heads, + projection_size=dim, + use_qkv_parallel=True, + use_context_forward=use_context_forward, + softmax_in_single_precision=softmax_in_single_precision, + flatten_batch=True, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + self.mlp = Qwen2VisionMLP( + dim, + mlp_hidden_dim, + act_layer=act_layer, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + + def forward( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + position_embeddings: torch.Tensor, + ) -> torch.Tensor: + hidden_states = self.norm1(x) + hidden_states = rearrange(hidden_states, "s b ... -> b s ...") + attn = self.attn( + hidden_states, + cu_seqlens=cu_seqlens, + position_embeddings=position_embeddings, + ) + attn = rearrange(attn, "b s ... -> s b ...") + x = x + attn + x = x + self.mlp(self.norm2(x)) + return x + + +class Qwen2VisionPatchEmbed(nn.Module): + + def __init__( + self, + patch_size: int = 14, + temporal_patch_size: int = 2, + in_chans: int = 3, + embed_dim: int = 1152, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.temporal_patch_size = temporal_patch_size + self.embed_dim = embed_dim + + kernel_size = [temporal_patch_size, patch_size, patch_size] + self.proj = nn.Conv3d( + in_chans, embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=False + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + L, C = x.shape + x = x.view(L, -1, self.temporal_patch_size, self.patch_size, self.patch_size) + x = self.proj(x).view(L, self.embed_dim) + return x + + +class Qwen2VisionPatchMerger(nn.Module): + + def __init__( + self, + d_model: int, + context_dim: int, + norm_layer: Type[nn.Module] = None, + spatial_merge_size: int = 2, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = context_dim * (spatial_merge_size**2) + if norm_layer is None: + norm_layer = partial(nn.LayerNorm, eps=1e-6) + self.ln_q = norm_layer(context_dim) + self.mlp = nn.ModuleList( + [ + ColumnParallelLinear( + self.hidden_size, + self.hidden_size, + bias=True, + quant_config=quant_config, + prefix=add_prefix("mlp.0", prefix), + ), + nn.GELU(), + RowParallelLinear( + self.hidden_size, + d_model, + bias=True, + quant_config=quant_config, + prefix=add_prefix("mlp.2", prefix), + ), + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.ln_q(x) + x = x.view(-1, self.hidden_size) + + mlp_fc1, mlp_act, mlp_fc2 = self.mlp + x_parallel, _ = mlp_fc1(x) + x_parallel = mlp_act(x_parallel) + out, _ = mlp_fc2(x_parallel) + return out + + +class Qwen2VisionRotaryEmbedding(nn.Module): + + def __init__(self, dim: int, theta: float = 10000.0) -> None: + super().__init__() + self.dim = dim + self.theta = theta + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self._seq_len_cached = 0 + self._freqs_cached = None + + def update_freqs_cache(self, seqlen: int) -> None: + if seqlen > self._seq_len_cached: + seqlen *= 2 + self._seq_len_cached = seqlen + self.inv_freq = 1.0 / ( + self.theta + ** ( + torch.arange( + 0, self.dim, 2, dtype=torch.float, device=self.inv_freq.device + ) + / self.dim + ) + ) + seq = torch.arange( + seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype + ) + freqs = torch.outer(seq, self.inv_freq) + self._freqs_cached = freqs + + def forward(self, seqlen: int) -> torch.Tensor: + self.update_freqs_cache(seqlen) + return self._freqs_cached[:seqlen] + + +class Qwen2VisionTransformer(nn.Module): + + def __init__( + self, + vision_config: Qwen2VLVisionConfig, + norm_eps: float = 1e-6, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + patch_size: int = vision_config.patch_size + temporal_patch_size: int = vision_config.temporal_patch_size + spatial_merge_size: int = vision_config.spatial_merge_size + in_chans: int = vision_config.in_chans + hidden_size: int = vision_config.hidden_size + embed_dim: int = vision_config.embed_dim + depth: int = vision_config.depth + num_heads: int = vision_config.num_heads + mlp_ratio: float = vision_config.mlp_ratio + + self.spatial_merge_size = spatial_merge_size + + self.patch_embed = Qwen2VisionPatchEmbed( + patch_size=patch_size, + temporal_patch_size=temporal_patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + ) + + norm_layer = partial(nn.LayerNorm, eps=norm_eps) + head_dim = embed_dim // num_heads + self.rotary_pos_emb = Qwen2VisionRotaryEmbedding(head_dim // 2) + self.blocks = nn.ModuleList( + [ + Qwen2VisionBlock( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + norm_layer=norm_layer, + attn_implementation="sdpa", + quant_config=quant_config, + prefix=add_prefix(f"blocks.{i}", prefix), + ) + for i in range(depth) + ] + ) + self.merger = Qwen2VisionPatchMerger( + d_model=hidden_size, + context_dim=embed_dim, + norm_layer=norm_layer, + quant_config=quant_config, + prefix=add_prefix("merger", prefix), + ) + + @property + def dtype(self) -> torch.dtype: + return self.patch_embed.proj.weight.dtype + + @property + def device(self) -> torch.device: + return self.blocks[0].mlp.fc2.weight.device + + def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor: + pos_ids = [] + for i in range(grid_thw.size(0)): + t, h, w = grid_thw[i].tolist() + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) + hpos_ids = ( + hpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + .permute(0, 2, 1, 3) + .flatten() + ) + wpos_ids = ( + wpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + .permute(0, 2, 1, 3) + .flatten() + ) + pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) + pos_ids = torch.cat(pos_ids, dim=0) + max_grid_size = grid_thw[:, 1:].max() + rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) + rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) + return rotary_pos_emb + + def forward( + self, + x: torch.Tensor, + grid_thw: torch.Tensor, + ) -> torch.Tensor: + # patchify + x = x.to(device=self.device, dtype=self.dtype) + x = self.patch_embed(x) + + # compute position embedding + rotary_pos_emb = self.rot_pos_emb(grid_thw) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) + # compute cu_seqlens + cu_seqlens = torch.repeat_interleave( + grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] + ).cumsum(dim=0, dtype=torch.int32) + cu_seqlens = F.pad(cu_seqlens, (1, 0), "constant", 0) + + # transformers + x = x.unsqueeze(1) + for blk in self.blocks: + x = blk(x, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings) + + # adapter + x = self.merger(x) + return x + + +cached_get_processor = lru_cache(get_processor) + + +class Qwen2VLForConditionalGeneration(nn.Module): + # BitandBytes specific attributes + default_bitsandbytes_target_modules = [ + ".gate_proj.", + ".down_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ".o_proj.", + ] + bitsandbytes_stacked_params_mapping = { + # shard_name, weight_name, index + "q_proj": ("qkv_proj", 0), + "k_proj": ("qkv_proj", 1), + "v_proj": ("qkv_proj", 2), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def calculate_num_image_tokens(self, image_grid_thw: Tuple[int, int, int]): + processor = cached_get_processor(self.config._name_or_path) + grid_t, grid_h, grid_w = image_grid_thw + num_image_tokens = ( + grid_t + * grid_h + * grid_w + // processor.image_processor.merge_size + // processor.image_processor.merge_size + ) + return num_image_tokens + + def __init__( + self, + config: Qwen2VLConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + self.visual = Qwen2VisionTransformer( + config.vision_config, + norm_eps=getattr(config, "rms_norm_eps", 1e-6), + # NOTE: Qwen2-VL vision encoder currently supports BitsAndBytes 4-bit quantization. + # Other quantization methods (e.g., GPTQ, AWQ) are untested and may not be supported. + quant_config=quant_config, + prefix=add_prefix("visual", prefix), + ) + + self.model = Qwen2Model( + config, quant_config, prefix=add_prefix("model", prefix) + ) + + if config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + + self.is_mrope_enabled = "mrope_section" in self.config.rope_scaling + self.logits_processor = LogitsProcessor(config) + self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) + + def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): + # Get all special token IDs + im_token_id: int = mm_inputs.im_token_id + + pattern = MultiModalityDataPaddingPatternMultimodalTokens([im_token_id]) + return pattern.pad_input_tokens(input_ids, mm_inputs) + + def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: + # in qwen-vl, last dim is the same + pixel_values = torch.cat([item.pixel_values for item in items], dim=0).type( + self.visual.dtype + ) + image_grid_thws = torch.concat([item.image_grid_thws for item in items], dim=0) + assert pixel_values.dim() == 2, pixel_values.dim() + assert image_grid_thws.dim() == 2, image_grid_thws.dim() + image_embeds = self.visual(pixel_values, grid_thw=image_grid_thws) + return image_embeds + + def _process_video_input(self, video_input: Qwen2VLVideoInputs) -> torch.Tensor: + pixel_values_videos = video_input["pixel_values_videos"].type(self.visual.dtype) + video_embeds = self.visual( + pixel_values_videos, grid_thw=video_input["video_grid_thw"] + ) + return video_embeds + + def get_input_embeddings(self): + return self.model.embed_tokens + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + get_embedding: bool = False, + ): + """Run forward pass for Qwen2-VL. + + Args: + input_ids: Flattened (concatenated) input_ids corresponding to a + batch. + positions: Flattened (concatenated) position ids corresponding to a + batch. + **NOTE**: If mrope is enabled (default setting for Qwen2-VL + opensource models), the shape will be `(3, seq_len)`, + otherwise it will be `(seq_len,). + (Use input_metadata.mrope_positions to replace it) + """ + if self.is_mrope_enabled: + positions = forward_batch.mrope_positions + + if not ( + forward_batch.forward_mode.is_decode() + or not forward_batch.contains_image_inputs() + ): + if self.is_mrope_enabled: + assert positions.ndim == 2 and positions.size(0) == 3, ( + "multimodal section rotary embedding requires " + f"(3, seq_len) positions, but got {positions.size()}" + ) + hidden_states = general_mm_embed_routine( + input_ids=input_ids, + forward_batch=forward_batch, + language_model=self.model, + image_data_embedding_func=self.get_image_feature, + positions=positions, + ) + + if get_embedding: + return self.pooler(hidden_states, forward_batch) + else: + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "up_proj", 1), + ("gate_up_proj", "gate_proj", 0), + ] + params_dict = dict(self.named_parameters(remove_duplicate=False)) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if self.config.tie_word_embeddings and "lm_head.weight" in name: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + if "visual" in name: + # adapt to VisionAttention + name = name.replace(r"attn.qkv.", r"attn.qkv_proj.") + + try: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + except KeyError: + print(params_dict.keys()) + raise + + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +EntryClass = Qwen2VLForConditionalGeneration diff --git a/lib/python3.12/site-packages/sglang/srt/models/qwen3_moe.py b/lib/python3.12/site-packages/sglang/srt/models/qwen3_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..c8786e4e5fa686d67763e8c3433529df9075aa9f --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/qwen3_moe.py @@ -0,0 +1,423 @@ +# Adapted from qwen2_moe.py + +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + + +"""Inference-only Qwen3MoE model compatible with HuggingFace weights.""" + +from functools import partial +from typing import Any, Dict, Iterable, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn + +from sglang.srt.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + split_tensor_along_last_dim, + tensor_model_parallel_all_gather, + tensor_model_parallel_all_reduce, +) +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.fused_moe_triton import FusedMoE +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.qwen2_moe import Qwen2MoeMLP as Qwen3MoeMLP +from sglang.srt.models.qwen2_moe import Qwen2MoeModel +from sglang.srt.utils import add_prefix + +Qwen3MoeConfig = None + + +class Qwen3MoeSparseMoeBlock(nn.Module): + def __init__( + self, + config: Qwen3MoeConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + + if self.tp_size > config.num_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {config.num_experts}." + ) + + self.experts = FusedMoE( + num_experts=config.num_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + reduce_results=False, + renormalize=config.norm_topk_prob, + quant_config=quant_config, + prefix=add_prefix("experts", prefix), + ) + + self.gate = ReplicatedLinear( + config.hidden_size, + config.num_experts, + bias=False, + quant_config=None, + prefix=add_prefix("gate", prefix), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + + # router_logits: (num_tokens, n_experts) + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=router_logits + ) + if self.tp_size > 1: + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + + return final_hidden_states.view(num_tokens, hidden_dim) + + +class Qwen3MoeAttention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + max_position_embeddings: int = 8192, + head_dim: Optional[int] = None, + rms_norm_eps: float = 1e-06, + attention_bias: bool = False, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % self.tp_size == 0 + self.num_heads = self.total_num_heads // self.tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= self.tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % self.tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert self.tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // self.tp_size) + self.head_dim = head_dim or hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.tp_rank = get_tensor_model_parallel_rank() + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=attention_bias, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=attention_bias, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + prefix=add_prefix("attn", prefix), + ) + + self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + + def _apply_qk_norm( + self, q: torch.Tensor, k: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + q_by_head = q.reshape(-1, self.head_dim) + q_by_head = self.q_norm(q_by_head) + q = q_by_head.view(q.shape) + k_by_head = k.reshape(-1, self.head_dim) + k_by_head = self.k_norm(k_by_head) + k = k_by_head.view(k.shape) + return q, k + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self._apply_qk_norm(q, k) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class Qwen3MoeDecoderLayer(nn.Module): + def __init__( + self, + config: Qwen3MoeConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + rms_norm_eps = config.rms_norm_eps + attention_bias = config.attention_bias + self.self_attn = Qwen3MoeAttention( + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + layer_id=layer_id, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + head_dim=head_dim, + rms_norm_eps=rms_norm_eps, + attention_bias=attention_bias, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + + # Note: Qwen/Qwen2-57B-A14B-Instruct does not have + # `mlp_only_layers` in the config. + mlp_only_layers = ( + [] if not hasattr(config, "mlp_only_layers") else config.mlp_only_layers + ) + if (layer_id not in mlp_only_layers) and ( + config.num_experts > 0 and (layer_id + 1) % config.decoder_sparse_step == 0 + ): + self.mlp = Qwen3MoeSparseMoeBlock( + config=config, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + else: + self.mlp = Qwen3MoeMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> torch.Tensor: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class Qwen3MoeModel(Qwen2MoeModel): + def __init__( + self, + config: Qwen3MoeConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__( + config=config, + quant_config=quant_config, + prefix=prefix, + decoder_layer_type=Qwen3MoeDecoderLayer, + ) + + +class Qwen3MoeForCausalLM(nn.Module): + + fall_back_to_pt_during_load = False + + def __init__( + self, + config: Qwen3MoeConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.model = Qwen3MoeModel( + config, quant_config, prefix=add_prefix("model", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + self.logits_processor = LogitsProcessor(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + ) + + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip non-stacked layers and experts (experts handled below). + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since we handle the experts below in expert_params_mapping, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below in expert_params_mapping + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if "mlp.experts" in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + +EntryClass = Qwen3MoeForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/registry.py b/lib/python3.12/site-packages/sglang/srt/models/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..fc63bf12554722b8b7b41d08288becc208dc2a59 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/registry.py @@ -0,0 +1,99 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/models/registry.py + +import importlib +import logging +import pkgutil +from dataclasses import dataclass, field +from functools import lru_cache +from typing import AbstractSet, Dict, List, Optional, Tuple, Type, Union + +import torch.nn as nn + +logger = logging.getLogger(__name__) + + +@dataclass +class _ModelRegistry: + # Keyed by model_arch + models: Dict[str, Union[Type[nn.Module], str]] = field(default_factory=dict) + + def get_supported_archs(self) -> AbstractSet[str]: + return self.models.keys() + + def _raise_for_unsupported(self, architectures: List[str]): + all_supported_archs = self.get_supported_archs() + + if any(arch in all_supported_archs for arch in architectures): + raise ValueError( + f"Model architectures {architectures} failed " + "to be inspected. Please check the logs for more details." + ) + + raise ValueError( + f"Model architectures {architectures} are not supported for now. " + f"Supported architectures: {all_supported_archs}" + ) + + def _try_load_model_cls(self, model_arch: str) -> Optional[Type[nn.Module]]: + if model_arch not in self.models: + return None + + return self.models[model_arch] + + def _normalize_archs( + self, + architectures: Union[str, List[str]], + ) -> List[str]: + if isinstance(architectures, str): + architectures = [architectures] + if not architectures: + logger.warning("No model architectures are specified") + + return architectures + + def resolve_model_cls( + self, + architectures: Union[str, List[str]], + ) -> Tuple[Type[nn.Module], str]: + architectures = self._normalize_archs(architectures) + + for arch in architectures: + model_cls = self._try_load_model_cls(arch) + if model_cls is not None: + return (model_cls, arch) + + return self._raise_for_unsupported(architectures) + + +@lru_cache() +def import_model_classes(): + model_arch_name_to_cls = {} + package_name = "sglang.srt.models" + package = importlib.import_module(package_name) + for _, name, ispkg in pkgutil.iter_modules(package.__path__, package_name + "."): + if not ispkg: + try: + module = importlib.import_module(name) + except Exception as e: + logger.warning(f"Ignore import error when loading {name}. " f"{e}") + continue + if hasattr(module, "EntryClass"): + entry = module.EntryClass + if isinstance( + entry, list + ): # To support multiple model classes in one module + for tmp in entry: + assert ( + tmp.__name__ not in model_arch_name_to_cls + ), f"Duplicated model implementation for {tmp.__name__}" + model_arch_name_to_cls[tmp.__name__] = tmp + else: + assert ( + entry.__name__ not in model_arch_name_to_cls + ), f"Duplicated model implementation for {entry.__name__}" + model_arch_name_to_cls[entry.__name__] = entry + + return model_arch_name_to_cls + + +ModelRegistry = _ModelRegistry(import_model_classes()) diff --git a/lib/python3.12/site-packages/sglang/srt/models/roberta.py b/lib/python3.12/site-packages/sglang/srt/models/roberta.py new file mode 100644 index 0000000000000000000000000000000000000000..d9e8c2c7ae1ed581a00c4be9792597e9a0e8713c --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/roberta.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: Apache-2.0 + +import itertools +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn + +from sglang.srt.layers.pooler import Pooler, PoolingType +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.bert import BertEncoder + +RobertaConfig = None + + +class RobertaEmbedding(nn.Module): + + def __init__(self, config: RobertaConfig): + super().__init__() + self.size = config.hidden_size + self.word_embeddings = VocabParallelEmbedding( + config.vocab_size, config.hidden_size + ) + self.padding_idx = config.pad_token_id + self.position_embeddings = nn.Embedding( + config.max_position_embeddings, + config.hidden_size, + padding_idx=self.padding_idx, + ) + + self.token_type_embeddings = nn.Embedding( + config.type_vocab_size, config.hidden_size + ) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + self.position_ids = nn.Parameter( + torch.empty((1, config.max_position_embeddings)), + ) + + self.position_embedding_type = config.position_embedding_type + if self.position_embedding_type != "absolute": + raise ValueError( + "Only 'absolute' position_embedding_type" + " is supported" + ) + + def forward( + self, + input_ids: torch.Tensor, + seq_lens: torch.Tensor, + position_ids: torch.Tensor, + inputs_embeds=None, + token_type_ids: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + input_shape = input_ids.size() + inputs_embeds = self.word_embeddings(input_ids) + + # adpated from vllm: https://github.com/vllm-project/vllm/commit/4a18fd14ba4a349291c798a16bf62fa8a9af0b6b/vllm/model_executor/models/roberta.py + + pos_list = [] + token_list = [] + offset = 0 + for seq_len in seq_lens: + pos_list.append(position_ids[offset : offset + seq_len]) + token_list.append(input_ids[offset : offset + seq_len]) + offset += seq_len + + new_pos_list = [] + for positions, tokens in zip(pos_list, token_list): + # Verify assumption that incoming position are + # always a sequence from 0 to N. + expected_pos = torch.arange( + positions.size()[0], dtype=torch.long, device=inputs_embeds.device + ) + assert torch.equal(positions, expected_pos) + new_pos_list.append( + create_position_ids_from_input_ids(tokens, self.padding_idx) + ) + position_ids = torch.cat(new_pos_list) + + # Position embeddings. + position_embeddings = self.position_embeddings(position_ids) + if token_type_ids is None: + token_type_ids = torch.zeros( + input_shape, dtype=torch.long, device=inputs_embeds.device + ) + + token_type_embeddings = self.token_type_embeddings(token_type_ids) + embeddings = inputs_embeds + token_type_embeddings + position_embeddings + embeddings = self.LayerNorm(embeddings) + return embeddings + + +class XLMRobertaModel(nn.Module): + def __init__( + self, + *, + config: RobertaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + + self.config = config + self.embeddings = RobertaEmbedding(config) + self.encoder = BertEncoder(config=config, quant_config=quant_config, prefix="") + self.pooler = Pooler(pooling_type=PoolingType.CLS, normalize=True) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + get_embedding: bool = False, + ) -> torch.Tensor: + assert get_embedding == True + # Your tokenized IDs + + hidden_states = self.embeddings( + input_ids=input_ids, + position_ids=positions, + seq_lens=forward_batch.seq_lens, + ) + + hidden_states = self.encoder(hidden_states, forward_batch=forward_batch) + pooler_out = self.pooler(hidden_states, forward_batch) + return pooler_out + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "query", "q"), + ("qkv_proj", "key", "k"), + ("qkv_proj", "value", "v"), + ] + + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + name = name.replace("self", "self_attn") + if "pooler" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +# Adapted from transformers +def create_position_ids_from_input_ids( + input_ids, padding_idx, past_key_values_length=0 +): + mask = input_ids.ne(padding_idx).int() + incremental_indices = ( + torch.cumsum(mask, dim=0).type_as(mask) + past_key_values_length + ) * mask + return incremental_indices.long() + padding_idx + + +EntryClass = [XLMRobertaModel] diff --git a/lib/python3.12/site-packages/sglang/srt/models/torch_native_llama.py b/lib/python3.12/site-packages/sglang/srt/models/torch_native_llama.py new file mode 100644 index 0000000000000000000000000000000000000000..1c99c52a80840efca075115e11b0bb781ab97944 --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/torch_native_llama.py @@ -0,0 +1,524 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from +# https://github.com/vllm-project/vllm/blob/c7f2cf2b7f67bce5842fedfdba508440fe257375/vllm/model_executor/models/llama.py#L1 +""" +Inference-only LLaMA model compatible with HuggingFace weights. + +This model supports tensor parallelism (TP) using the PyTorch tensor parallel package. +Reference: https://pytorch.org/docs/stable/distributed.tensor.parallel.html + +Here is a quick example to enable TP: +```python +from sglang.srt.model_parallel import tensor_parallel + +device_mesh = torch.distributed.init_device_mesh("cuda", (tp_size,)) +tensor_parallel(model, device_mesh) +``` + +An end-to-end example can be found in `python/sglang/bench_one_batch.py`. +You can run it with the following command: +```bash +$ python3 -m sglang.bench_one_batch --correct \ + --model meta-llama/Meta-Llama-3-8B \ + --json-model-override-args '{"architectures": ["TorchNativeLlamaForCausalLM"]}' \ + --tensor-parallel-size 2 \ + --disable-cuda-graph +``` +We will eanble CUDA Graph support soon. +""" + +import types +from typing import Any, Dict, Iterable, Optional, Tuple + +import torch +from torch import nn +from torch.nn.parameter import Parameter +from transformers import LlamaConfig + +from sglang.srt.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.logits_processor import LogitsProcessor, LogitsProcessorOutput +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix + +tp_size = get_tensor_model_parallel_world_size() +tp_rank = get_tensor_model_parallel_rank() + + +def gate_up_proj_weight_loader( + self, + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: int, +): + # shard_id: (shard_offset, shard_size) + gate_up_offsets = {} + current_shard_offset = 0 + for i, output_size in enumerate(self.output_sizes): + # Everything shrinks by tp_size if TP enabled + output_size = output_size // tp_size + gate_up_offsets[i] = (current_shard_offset, output_size) + current_shard_offset += output_size + # Re-size the param to the size after TP + if current_shard_offset != param.shape[0]: + # The clone will free the original, full tensor + param.data = param.data.narrow(0, 0, current_shard_offset).clone() + + # Now load gate or up + assert loaded_shard_id < len(self.output_sizes) + param_data = param.data + shard_offset, shard_size = gate_up_offsets[loaded_shard_id] + param_data = param_data.narrow(0, shard_offset, shard_size) + loaded_weight = loaded_weight.narrow(0, tp_rank * shard_size, shard_size) + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + +class LlamaMLP(nn.Module): + _tp_plan = { + "gate_up_proj": "Colwise_Sharded", + "down_proj": "Rowwise", + } + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = torch.nn.Linear( + hidden_size, + intermediate_size * 2, + bias=False, + ) + self.gate_up_proj.output_sizes = [intermediate_size] * 2 + self.gate_up_proj.weight_loader = types.MethodType( + gate_up_proj_weight_loader, self.gate_up_proj + ) + self.gate_up_proj.weight.weight_loader = self.gate_up_proj.weight_loader + self.down_proj = torch.nn.Linear(intermediate_size, hidden_size, bias=False) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x = self.down_proj(x) + return x + + +def qkv_proj_weight_loader( + self, + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: str, +): + num_heads = self.num_heads // tp_size + num_kv_heads = self.num_kv_heads // tp_size + # shard_id: (shard_offset, shard_size) + qkv_offsets = { + "q": (0, num_heads * self.head_size), + "k": (num_heads * self.head_size, num_kv_heads * self.head_size), + "v": ( + (num_heads + num_kv_heads) * self.head_size, + num_kv_heads * self.head_size, + ), + } + total_size = qkv_offsets["v"][0] + qkv_offsets["v"][1] + # Re-size the param to the size after TP + if total_size != param.shape[0]: + # The clone will free the original, full tensor + param.data = param.data.narrow(0, 0, total_size).clone() + + # Now load q, k or v + shard_offset, shard_size = qkv_offsets[loaded_shard_id] + param_data = param.data + param_data = param_data.narrow(0, shard_offset, shard_size) + loaded_weight = loaded_weight.narrow(0, tp_rank * shard_size, shard_size) + assert param_data.shape == loaded_weight.shape + param_data.copy_(loaded_weight) + + +class LlamaAttention(nn.Module): + _tp_plan = { + "qkv_proj": "Colwise_Sharded", + "o_proj": "Rowwise", + } + + def __init__( + self, + config: LlamaConfig, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + rope_is_neox_style: bool = True, + max_position_embeddings: int = 8192, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + # MistralConfig has an optional head_dim introduced by Mistral-Nemo + self.head_dim = getattr( + config, "head_dim", self.hidden_size // self.total_num_heads + ) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = torch.nn.Linear( + hidden_size, + (self.total_num_heads + 2 * self.total_num_kv_heads) * self.head_dim, + bias=False, + ) + self.qkv_proj.head_size = self.head_dim + self.qkv_proj.num_heads = self.total_num_heads + self.qkv_proj.num_kv_heads = self.total_num_kv_heads + self.qkv_proj.weight_loader = types.MethodType( + qkv_proj_weight_loader, self.qkv_proj + ) + self.qkv_proj.weight.weight_loader = self.qkv_proj.weight_loader + self.qkv_proj.weight.output_dim = 0 + self.o_proj = torch.nn.Linear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + ) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + is_neox_style=rope_is_neox_style, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output = self.o_proj(attn_output) + return output + + +class LlamaDecoderLayer(nn.Module): + def __init__( + self, + config: LlamaConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + if rope_scaling is not None and getattr( + config, "original_max_position_embeddings", None + ): + rope_scaling["original_max_position_embeddings"] = ( + config.original_max_position_embeddings + ) + rope_is_neox_style = getattr(config, "rope_is_neox_style", True) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + self.self_attn = LlamaAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + layer_id=layer_id, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + rope_is_neox_style=rope_is_neox_style, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + self.mlp = LlamaMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class LlamaModel(nn.Module): + def __init__( + self, + config: LlamaConfig, + quant_config: Optional[QuantizationConfig] = None, + ) -> None: + super().__init__() + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) + self.layers = nn.ModuleList( + [ + LlamaDecoderLayer( + config, i, quant_config=quant_config, prefix=f"model.layers.{i}" + ) + for i in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + forward_batch, + residual, + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class TorchNativeLlamaForCausalLM(nn.Module): + def __init__( + self, + config: LlamaConfig, + quant_config: Optional[QuantizationConfig] = None, + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.supports_torch_tp = True + self.model = LlamaModel(config, quant_config=quant_config) + if self.config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size) + self.logits_processor = LogitsProcessor(config) + + # turning off autotune for fp8dq since it doesn't give speedup and + # increases compile time significantly + torch._inductor.config.max_autotune_gemm_backends = "ATEN" + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> LogitsProcessorOutput: + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def get_hidden_dim(self, module_name): + if module_name in ["q_proj", "o_proj", "qkv_proj"]: + return self.config.hidden_size, self.config.hidden_size + elif module_name in ["kv_proj"]: + return self.config.hidden_size, self.config.hidden_size // ( + self.config.num_attention_heads // self.config.num_key_value_heads + ) + elif module_name == "gate_up_proj": + return self.config.hidden_size, self.config.intermediate_size + elif module_name == "down_proj": + return self.config.intermediate_size, self.config.hidden_size + else: + raise NotImplementedError() + + def get_module_name(self, name): + params_mapping = { + "q_proj": "qkv_proj", + "k_proj": "qkv_proj", + "v_proj": "qkv_proj", + "gate_proj": "gate_up_proj", + "up_proj": "gate_up_proj", + } + return params_mapping.get(name, name) + + def get_module_name_from_weight_name(self, name): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id, num_shard) + ("qkv_proj", "q_proj", "q", 3), + ("qkv_proj", "k_proj", "k", 3), + ("qkv_proj", "v_proj", "v", 3), + ("gate_up_proj", "gate_proj", 0, 2), + ("gate_up_proj", "up_proj", 1, 2), + ] + for param_name, weight_name, shard_id, num_shard in stacked_params_mapping: + if weight_name in name: + return ( + name.replace(weight_name, param_name)[: -len(".weight")], + num_shard, + ) + return name[: -len(".weight")], 1 + + def get_num_params(self): + params_dict = dict(self.named_parameters()) + return len(params_dict) + + def load_weights_to_module( + self, + fqn: str, + weights: Iterable[Tuple[str, torch.Tensor]], + ): + """Load weights onto submodule pointed by path `fqn`.""" + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + module = self.get_submodule(fqn) + params_dict = dict(module.named_parameters(prefix=fqn, recurse=False)) + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name or "projector" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + if name.startswith("model.vision_tower") and name not in params_dict: + continue + if self.config.tie_word_embeddings and "lm_head.weight" in name: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") or name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") or name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + def load_weights( + self, + weights: Iterable[Tuple[str, torch.Tensor]], + ): + """Load weights onto the full model.""" + self.load_weights_to_module("", weights) + + +class TorchNativePhi3ForCausalLM(TorchNativeLlamaForCausalLM): + pass + + +EntryClass = [TorchNativeLlamaForCausalLM, TorchNativePhi3ForCausalLM] diff --git a/lib/python3.12/site-packages/sglang/srt/models/xverse.py b/lib/python3.12/site-packages/sglang/srt/models/xverse.py new file mode 100644 index 0000000000000000000000000000000000000000..f84755b03635aa2ce3c8a7c722605a26d5a8985e --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/xverse.py @@ -0,0 +1,382 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +# Adapted from +# https://github.com/vllm-project/vllm/blob/c7f2cf2b7f67bce5842fedfdba508440fe257375/vllm/model_executor/models/xverse.py#L1 +"""Inference-only XVERSE model compatible with HuggingFace weights.""" + +from typing import Any, Dict, Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import LlamaConfig + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.model_runner import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix + + +class XverseMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class XverseAttention(nn.Module): + def __init__( + self, + config: LlamaConfig, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + rope_is_neox_style: bool = True, + max_position_embeddings: int = 8192, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + # MistralConfig has an optional head_dim introduced by Mistral-Nemo + self.head_dim = getattr( + config, "head_dim", self.hidden_size // self.total_num_heads + ) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + is_neox_style=rope_is_neox_style, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class XverseDecoderLayer(nn.Module): + def __init__( + self, + config: LlamaConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + if rope_scaling is not None and getattr( + config, "original_max_position_embeddings", None + ): + rope_scaling["original_max_position_embeddings"] = ( + config.original_max_position_embeddings + ) + rope_is_neox_style = getattr(config, "rope_is_neox_style", True) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + num_kv_heads = getattr( + config, "num_key_value_heads", config.num_attention_heads + ) + self.self_attn = XverseAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=num_kv_heads, + layer_id=layer_id, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + rope_is_neox_style=rope_is_neox_style, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + self.mlp = XverseMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class XverseModel(nn.Module): + def __init__( + self, + config: LlamaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=add_prefix("embed_tokens", prefix), + ) + self.layers = nn.ModuleList( + [ + XverseDecoderLayer( + config, + i, + quant_config=quant_config, + prefix=add_prefix(f"layers.{i}", prefix), + ) + for i in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + forward_batch, + residual, + ) + # print(f"layer[{i}].hidden_states: {hidden_states}") + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class XverseForCausalLM(nn.Module): + def __init__( + self, + config: LlamaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.model = XverseModel( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, config.hidden_size, prefix=add_prefix("lm_head", prefix) + ) + self.logits_processor = LogitsProcessor(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights( + self, weights: Iterable[Tuple[str, torch.Tensor]], name=None, loaded_weight=None + ): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.named_parameters()) + + def load_weights_per_param(name, loaded_weight): + if "rotary_emb.inv_freq" in name or "projector" in name: + return + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + return + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if name.startswith("model.vision_tower") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + return + if name.startswith("model.vision_tower") and name not in params_dict: + return + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + if name is None or loaded_weight is None: + for name, loaded_weight in weights: + load_weights_per_param(name, loaded_weight) + else: + load_weights_per_param(name, loaded_weight) + + +EntryClass = XverseForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/xverse_moe.py b/lib/python3.12/site-packages/sglang/srt/models/xverse_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..0ea9ed95012a4fce66e420833425f3a8f810a16a --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/xverse_moe.py @@ -0,0 +1,472 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Inference-only XVERSE MoE model.""" + +from typing import Any, Dict, Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.fused_moe_triton import fused_moe +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.utils import add_prefix + + +class XverseMLP(nn.Module): + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + reduce_results: bool = True, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=add_prefix("down_proj", prefix), + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class XverseMoE(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ): + super().__init__() + self.config = config + self.rank = get_tensor_model_parallel_rank() + self.tp_size = get_tensor_model_parallel_world_size() + self.n_routed_experts = config.num_experts + self.top_k = config.moe_top_k + if self.tp_size > self.n_routed_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {self.n_routed_experts}." + ) + + self.experts = nn.ModuleList( + [ + XverseMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix(f"experts.{i}", prefix), + ) + for i in range(self.n_routed_experts) + ] + ) + self.pack_params() + + self.router = ReplicatedLinear( + config.hidden_size, + self.n_routed_experts, + bias=False, + quant_config=None, + prefix=add_prefix("router", prefix), + ) + + if config.num_shared_experts is not None: + intermediate_size = config.intermediate_size * config.num_shared_experts + self.shared_experts = XverseMLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix("shared_experts", prefix), + ) + + def pack_params(self): + w1 = [] + w2 = [] + for expert in self.experts: + w1.append(expert.gate_up_proj.weight) + w2.append(expert.down_proj.weight) + self.w1 = torch._utils._flatten_dense_tensors(w1) + w1s = torch._utils._unflatten_dense_tensors(self.w1, w1) + for data, param in zip(w1s, w1): + param.data = data + self.w1 = self.w1.view(len(w1), *w1s[0].shape) + + self.w2 = torch._utils._flatten_dense_tensors(w2) + w2s = torch._utils._unflatten_dense_tensors(self.w2, w2) + for data, param in zip(w2s, w2): + param.data = data + + self.w2 = self.w2.view(len(w2), *w2s[0].shape) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + if self.config.num_shared_experts is not None: + shared_output = self.shared_experts(hidden_states) + # router_logits: (num_tokens, n_experts) + router_logits, _ = self.router(hidden_states) + final_hidden_states = fused_moe( + hidden_states, + self.w1, + self.w2, + router_logits, + self.top_k, + renormalize=getattr(self.config, "norm_topk_prob", False), + inplace=True, + ) + + if self.config.num_shared_experts is not None: + final_hidden_states = final_hidden_states + shared_output + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + + return final_hidden_states.view(num_tokens, hidden_dim) + + +class XverseAttention(nn.Module): + + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + max_position_embeddings: int = 8192, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=False, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + ) + + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.o_proj(attn_output) + return output + + +class XverseDecoderLayer(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + num_key_value_heads = getattr( + config, "num_key_value_heads", config.num_attention_heads + ) + self.self_attn = XverseAttention( + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=num_key_value_heads, + layer_id=layer_id, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + if config.num_experts is not None: + self.mlp = XverseMoE( + config=config, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + else: + self.mlp = XverseMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> torch.Tensor: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class XverseModel(nn.Module): + + fall_back_to_pt_during_load = False + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=add_prefix("embed_tokens", prefix), + ) + self.layers = nn.ModuleList( + [ + XverseDecoderLayer( + config, + layer_id, + quant_config=quant_config, + prefix=add_prefix(f"layers.{layer_id}", prefix), + ) + for layer_id in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + hidden_states = self.embed_tokens(input_ids) + residual = None + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, hidden_states, forward_batch, residual + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class XverseMoeForCausalLM(nn.Module): + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.quant_config = quant_config + self.model = XverseModel( + config, quant_config, prefix=add_prefix("model", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + ) + self.logits_processor = LogitsProcessor(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + hidden_states = self.model(input_ids, positions, forward_batch) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.named_parameters()) + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Skip experts that are not assigned to this worker. + if ( + "mlp.experts." in name or "mlp.shared_experts." in name + ) and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Skip experts that are not assigned to this worker. + if ( + "mlp.experts." in name or "mlp.shared_experts." in name + ) and name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +EntryClass = XverseMoeForCausalLM diff --git a/lib/python3.12/site-packages/sglang/srt/models/yivl.py b/lib/python3.12/site-packages/sglang/srt/models/yivl.py new file mode 100644 index 0000000000000000000000000000000000000000..4c50b0d3c2b2bab1304493a2e0f70627676421ca --- /dev/null +++ b/lib/python3.12/site-packages/sglang/srt/models/yivl.py @@ -0,0 +1,115 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Inference-only Yi-VL model.""" + +from typing import Iterable, Optional, Tuple + +import torch +import torch.nn as nn +from transformers import CLIPVisionModel, LlavaConfig + +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.llava import LlavaLlamaForCausalLM + + +class YiVLForCausalLM(LlavaLlamaForCausalLM): + def __init__( + self, + config: LlavaConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__(config, quant_config, prefix=prefix) + + self.multi_modal_projector = YiVLMultiModalProjector(self.config) + self.vision_tower_subfolder = self.config.mm_vision_tower.replace( + "./", "" + ) # Everything after "./" + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + # We have to use the subfolder of the main model directory (e.g. 01-ai/Yi-VL-6B) + self.vision_tower = CLIPVisionModel.from_pretrained( + self.config._name_or_path, + torch_dtype=torch.float16, + subfolder=self.vision_tower_subfolder, + ).to("cuda") + + self.vision_tower.eval() + + self.vision_feature_layer = self.config.mm_vision_select_layer + self.vision_feature_select_strategy = self.config.mm_vision_select_feature + self.image_size = self.vision_tower.config.image_size + self.patch_size = self.vision_tower.config.patch_size + + self.mm_patch_merge_type = getattr(self.config, "mm_patch_merge_type", "flat") + self.image_aspect_ratio = getattr(self.config, "image_aspect_ratio", "square") + self.image_grid_pinpoints = getattr(self.config, "image_grid_pinpoints", None) + + self.image_feature_len = int((self.image_size / self.patch_size) ** 2) + if self.vision_feature_select_strategy == "patch": + pass + elif self.vision_feature_select_strategy == "cls_patch": + self.image_feature_len += 1 + else: + raise ValueError(f"Unexpected select feature: {self.select_feature}") + + # load mm_projector + # TODO: support TP? + projector_weights = { + "model.mm_projector.0": "multi_modal_projector.linear_1", + "model.mm_projector.1": "multi_modal_projector.ln_1", + "model.mm_projector.3": "multi_modal_projector.linear_2", + "model.mm_projector.4": "multi_modal_projector.ln_2", + "model.vision_tower.vision_tower": "vision_tower", # Update the vision tower weights if we find them in the checkpoint (it may be finetuned). + } + params_dict = dict(self.named_parameters()) + weights = list(weights) + for name, loaded_weight in weights: + if "projector" in name or "vision_tower" in name: + for weight_name, param_name in projector_weights.items(): + if weight_name in name: + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + # load language model + self.language_model.load_weights(weights) + + +class YiVLMultiModalProjector(nn.Module): + def __init__(self, config: LlavaConfig): + super().__init__() + + self.linear_1 = nn.Linear( + config.vision_config.hidden_size, config.text_config.hidden_size + ) + self.ln_1 = nn.LayerNorm(config.text_config.hidden_size) + self.act = nn.GELU() + self.linear_2 = nn.Linear( + config.text_config.hidden_size, config.text_config.hidden_size + ) + self.ln_2 = nn.LayerNorm(config.text_config.hidden_size) + + def forward(self, image_features): + hidden_states = self.linear_1(image_features) + hidden_states = self.ln_1(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.linear_2(hidden_states) + hidden_states = self.ln_2(hidden_states) + return hidden_states + + +EntryClass = YiVLForCausalLM