File tree Expand file tree Collapse file tree 1 file changed +4
-2
lines changed
scripts/checkpoint_converters Expand file tree Collapse file tree 1 file changed +4
-2
lines changed Original file line number Diff line number Diff line change 1616import re
1717from argparse import ArgumentParser
1818from collections import defaultdict
19+ from datetime import timedelta
1920
21+ import megatron .core .parallel_state as ps
2022import torch
2123import torch .distributed as dist
2224from megatron .core .dist_checkpointing .serialization import load_plain_tensors
2527from megatron .core .transformer .spec_utils import import_module
2628from megatron .training .arguments import core_transformer_config_from_args
2729from omegaconf .omegaconf import OmegaConf
28- from datetime import timedelta
29- import megatron .core .parallel_state as ps
3030from torch ._C ._distributed_c10d import PrefixStore
3131from torch .distributed import rendezvous
32+
3233from nemo .collections .nlp .models .language_modeling .megatron_mamba_model import MegatronMambaModel
3334from nemo .collections .nlp .parts .megatron_trainer_builder import MegatronLMPPTrainerBuilder
3435from nemo .collections .nlp .parts .utils_funcs import torch_dtype_from_precision
@@ -95,6 +96,7 @@ def get_args():
9596 args = parser .parse_args ()
9697 return args
9798
99+
98100try :
99101
100102 class Utils :
You can’t perform that action at this time.
0 commit comments