| 61 | |
| 62 | |
| 63 | class LegacyInternTrainManager(InternTrainManager): |
| 64 | |
| 65 | def load_config(self, path, model_config=None): |
| 66 | from internlm.core.context import Config |
| 67 | if model_config is None: |
| 68 | model_config = torch.load(os.path.join(path, 'model_config.pt')) |
| 69 | elif isinstance(model_config, str) and model_config.endswith('.pt'): |
| 70 | model_config = torch.load(model_config) |
| 71 | elif isinstance(model_config, dict): |
| 72 | model_config = Config(model_config) |
| 73 | elif isinstance(model_config, str): |
| 74 | model_config = Config.from_file(model_config).model |
| 75 | else: |
| 76 | raise NotImplementedError( |
| 77 | 'model_config should be None, dict or filename.') |
| 78 | |
| 79 | return model_config |
| 80 | |
| 81 | def initialize_model(self): |
| 82 | from internlm.train.pipeline import initialize_model |
| 83 | model = initialize_model().model |
| 84 | |
| 85 | return model |
| 86 | |
| 87 | |
| 88 | @MODELS.register_module() |