[Bugfix] Fix GraniteMoeShared weight loading broken by #41184 (#47031)

Signed-off-by: <Michal Ganczarenko> <michal.ganczarenko@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
This commit is contained in:
Michał Ganczarenko
2026-07-01 22:39:12 +08:00
committed by GitHub
co-authored by Kunshang Ji
parent c638f9216a
commit 63fcce4de1
@@ -214,11 +214,11 @@ class GraniteMoeSharedModel(nn.Module):
for e in range(p.size(0)):
w1_name = n.replace(
".block_sparse_moe.input_linear.weight",
f".block_sparse_moe.experts.routed_experts.{e}.w1.weight",
f".block_sparse_moe.experts.{e}.w1.weight",
)
w3_name = n.replace(
".block_sparse_moe.input_linear.weight",
f".block_sparse_moe.experts.routed_experts.{e}.w3.weight",
f".block_sparse_moe.experts.{e}.w3.weight",
)
w1_param, w3_param = p[e].chunk(2, dim=0)
assert w1_name not in new_weights
@@ -229,7 +229,7 @@ class GraniteMoeSharedModel(nn.Module):
for e in range(p.size(0)):
w2_name = n.replace(
".block_sparse_moe.output_linear.weight",
f".block_sparse_moe.experts.routed_experts.{e}.w2.weight",
f".block_sparse_moe.experts.{e}.w2.weight",
)
w2_param = p[e]
assert w2_name not in new_weights