[MODEL] Add support for Zamba2 models (#13185)

Signed-off-by: Yury Tokpanov <yury@zyphra.com>
Signed-off-by: Quentin Anthony <qganthony@yahoo.com>
Co-authored-by: Quentin Anthony <qganthony@yahoo.com>
Co-authored-by: Tyler Michael Smith <tysmith@redhat.com>
Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
This commit is contained in:
yury-tokpanov
2025-03-18 08:56:21 -07:00
committed by GitHub
parent 8b793f7ec6
commit 452e8fd968
9 changed files with 1081 additions and 26 deletions

View File

@@ -245,7 +245,6 @@ class MambaMixer2(CustomOp):
assert num_heads % self.tp_size == 0, \
"Tensor parallel world size must divide num heads."
assert (n_groups % self.tp_size) == 0 or n_groups == 1, \
(
"If tensor parallel world size does not divide num_heads, "