Fix new code import

Signed-off-by: Mikołaj Błaż <mblaz@nvidia.com>

nemo/utils/callbacks/dist_ckpt_io.py

@@ -32,10 +32,7 @@
     from megatron.core import dist_checkpointing
     from megatron.core.dist_checkpointing.dict_utils import extract_matching_values
     from megatron.core.dist_checkpointing.mapping import ShardedBase
-    from megatron.core.dist_checkpointing.serialization import (
-        get_default_load_sharded_strategy,
-        get_default_save_sharded_strategy,
-    )
+    from megatron.core.dist_checkpointing.serialization import get_default_save_sharded_strategy
     from megatron.core.dist_checkpointing.strategies import tensorstore
     from megatron.core.dist_checkpointing.strategies.async_utils import AsyncCallsQueue, AsyncRequest
     from megatron.core.dist_checkpointing.strategies.base import SaveShardedStrategy
@@ -46,6 +43,13 @@
     from megatron.core.dist_checkpointing.strategies.torch import TorchDistSaveShardedStrategy
     from megatron.core.parallel_state import get_data_parallel_group
 
+    try:
+        from megatron.core.dist_checkpointing.serialization import get_default_load_sharded_strategy
+    except ImportError:
+        from megatron.core.dist_checkpointing.serialization import (
+            _verify_checkpoint_and_load_strategy as get_default_load_sharded_strategy,
+        )
+
     HAVE_MEGATRON_CORE = True
 
 except (ImportError, ModuleNotFoundError) as e: