[Feature] Xiaomi MiMo-V2-Flash day0 support (#15207)

Co-authored-by: 谢学扬 <xiexueyang@xiaomi.com>
Co-authored-by: tz <tangzhen3@xiaomi.com>
Co-authored-by: 李家乐 <lijiale10@xiaomi.com>
Co-authored-by: 张晨 <zhangchen50@xiaomi.com>
Co-authored-by: Shaohui Liu <liushaohui3@xiaomi.com>
Co-authored-by: 王晨 <wangchen77@xiaomi.com>
Co-authored-by: jiangzihan <jiangzihan@xiaomi.com>
Co-authored-by: xiexueyang <xyxie_wangyi@163.com>
Co-authored-by: Linghao Zhang <zhanglinghao@xiaomi.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>
Co-authored-by: Liangsheng Yin <lsyincs@gmail.com>
Co-authored-by: JoyFuture <35593546+JoyFuture@users.noreply.github.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
Co-authored-by: Qiaolin Yu <liin1211@outlook.com>
Co-authored-by: root <root@bj9-ml-g8h20e-k8s-slave106-20251106.alicn.idc.xiaomi.com>
This commit is contained in:
Yingchun Lai
2025-12-19 11:40:07 +08:00
committed by GitHub
co-authored by 谢学扬 tz 李家乐 张晨 Shaohui Liu 王晨 jiangzihan xiexueyang Linghao Zhang ispobock Liangsheng Yin JoyFuture Liangsheng Yin Qiaolin Yu root
parent a0985dd5e5
commit 160a06cab2
38 changed files with 5396 additions and 169 deletions
+17
View File
@@ -498,6 +498,23 @@ class DefaultModelLoader(BaseModelLoader):
else:
weights_iterator = pt_weights_iterator(hf_weights_files)
if self.load_config.draft_model_idx is not None:
import re
pattern = r"model.mtp.layers.(\d+)."
filtered_weights = []
for name, tensor in weights_iterator:
group = re.match(pattern, name)
if group is not None:
idx = int(group.group(1))
if idx != self.load_config.draft_model_idx:
continue
new_name = name.replace(group.group(), "model.mtp.layers.0.")
else:
new_name = name
filtered_weights.append((source.prefix + new_name, tensor))
return tuple(filtered_weights)
# Apply the prefix.
return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator)