未验证 提交 cc3a7c9c 编写于 作者: L Lev Kurilenko 提交者: GitHub

Fix Meta Tensor checkpoint load for BLOOM models (#3885)

This PR fixes Meta Tensor checkpoint loading for BLOOM models where the SD keys start with transformer..
上级 d6f62217
......@@ -29,10 +29,14 @@ def load_model_with_checkpoint(r_module,
error_msgs = []
def prefix_check():
# if keys start with 'model.', don't skip level 0 prefix
# if keys start with 'model.' or 'transformer.', don't skip level 0 prefix
for key in sd[0].keys():
# OPT models
if re.match("^model[.]", key):
return False
# BLOOM models
if re.match("^transformer[.]", key):
return False
return True
skip_level_0_prefix = prefix_check() and container.policy.use_load_prefix
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册