fix(test): support relocated MCore schedule callables

Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>

Author: svcnemo-autobot

tests/unit_tests/training/test_gpt_step.py

@@ -760,7 +760,10 @@ def current_apply_expert_bias(_self, routing_map, padding_mask=None):
 
     def test_mcore_schedule_plan_routes_with_chunk_padding_mask(self, monkeypatch):
         """The pinned MCore EP-overlap callable must pass its chunk-local router mask."""
-        from megatron.core.models.common import fine_grained_callables
+        try:
+            from megatron.core.models.common import fine_grained_callables
+        except ImportError:
+            from megatron.core.models.gpt import fine_grained_callables
 
         observed = {}