Fix reload_model_params failure when loading MoE models with explicit state_dict

megatron/core/optimizer/distrib_optimizer.py

@@ -2528,7 +2528,7 @@ def _build_model_param_to_state_dict_param_map(self, state_dict):
             for name, model_param in model_chunk.named_parameters():
                 while name.startswith("module."):
                     name = name[len("module.") :]
-                matched_keys = [k for k in names_in_state_dict if name in k]
+                matched_keys = [k for k in names_in_state_dict if k.endswith(name)]
                 assert (
                     len(matched_keys) == 1
                 ), f"Parameter {name} has {len(matched_keys)} matches in state dict"