| 30 | |
| 31 | |
| 32 | class CurrentInternTrainManager(InternTrainManager): |
| 33 | |
| 34 | def load_config(self, path, model_config=None): |
| 35 | if model_config is None: |
| 36 | from internlm.checkpoint.checkpoint_manager import try_load_config |
| 37 | model_config = try_load_config( |
| 38 | os.path.join(path, 'model_config.pt')) |
| 39 | elif isinstance(model_config, str) and model_config.endswith('.pt'): |
| 40 | from internlm.checkpoint.checkpoint_manager import try_load_config |
| 41 | model_config = try_load_config(model_config) |
| 42 | else: |
| 43 | from internlm.config import Config |
| 44 | if isinstance(model_config, dict): |
| 45 | model_config = Config(model_config) |
| 46 | elif isinstance(model_config, str): |
| 47 | model_config = Config.fromfile(model_config).model |
| 48 | else: |
| 49 | raise NotImplementedError( |
| 50 | 'model_config should be None, dict or filename.') |
| 51 | |
| 52 | return model_config |
| 53 | |
| 54 | def initialize_model(self): |
| 55 | from internlm.train.pipeline import (initialize_model, |
| 56 | initialize_parallel_communicator) |
| 57 | model = initialize_model().model |
| 58 | initialize_parallel_communicator(model) |
| 59 | |
| 60 | return model |
| 61 | |
| 62 | |
| 63 | class LegacyInternTrainManager(InternTrainManager): |