(self,
path: str,
model_config: Optional[str] = None,
parallel_config: Optional[str] = None,
model_type: str = 'INTERNLM2',
model_dtype: Optional[str] = None,
ckpt_type: Optional[str] = None)
| 206 | self.end_str = end_str |
| 207 | |
| 208 | def _load_model(self, |
| 209 | path: str, |
| 210 | model_config: Optional[str] = None, |
| 211 | parallel_config: Optional[str] = None, |
| 212 | model_type: str = 'INTERNLM2', |
| 213 | model_dtype: Optional[str] = None, |
| 214 | ckpt_type: Optional[str] = None): |
| 215 | # funcs |
| 216 | from internlm.checkpoint.load_funcs import (LOAD_FUNC_DICT, |
| 217 | merge_pp_within_tp) |
| 218 | from internlm.core.context import global_context as gpc |
| 219 | from internlm.initialize.launch import launch |
| 220 | from internlm.utils.storage_manager import (get_storage_manager, |
| 221 | init_storage_manager) |
| 222 | |
| 223 | # config |
| 224 | model_config = self.manager.load_config(path, model_config) |
| 225 | model_config['parallel_output'] = False |
| 226 | model_config['dtype'] = self._convert_dtype(model_config['dtype'], |
| 227 | model_dtype=model_dtype) |
| 228 | |
| 229 | world_size = int(os.getenv('WORLD_SIZE', '1')) |
| 230 | tp_size = world_size # TODO |
| 231 | self.logger.info(f'world size: {world_size} tp: {tp_size}') |
| 232 | if parallel_config is None: |
| 233 | parallel_config = dict(zero1=dict(size=1, fsdp=False), |
| 234 | pipeline=dict(size=1), |
| 235 | tensor=dict(size=tp_size, mode='mtp'), |
| 236 | sequence_parallel=False) |
| 237 | config = dict(model=model_config, |
| 238 | parallel=parallel_config, |
| 239 | data=dict(use_packed_dataset=False), |
| 240 | model_type=model_type, |
| 241 | use_cuda_flash_attn=model_config.get( |
| 242 | 'use_flash_attn', True)) |
| 243 | launch( |
| 244 | config=config, |
| 245 | seed=42, |
| 246 | local_rank=int(os.getenv('RANK', '0')), |
| 247 | rank=int(os.getenv('LOCAL_RANK', '0')), |
| 248 | world_size=int(os.getenv('WORLD_SIZE', '1')), |
| 249 | host=os.getenv('MASTER_ADDR', '127.0.0.1'), |
| 250 | port=int(os.getenv('MASTER_PORT', random.randint(12000, 32000))), |
| 251 | ) |
| 252 | self.logger.info(f'Config: {gpc.config}') |
| 253 | |
| 254 | self.model = self.manager.initialize_model() |
| 255 | |
| 256 | # load state dict |
| 257 | try: |
| 258 | get_storage_manager() |
| 259 | except AssertionError: |
| 260 | init_storage_manager(False, None, None) |
| 261 | get_storage_manager() |
| 262 | if ckpt_type is None or ckpt_type == 'internevo': |
| 263 | state_dict = merge_pp_within_tp(path, del_model_prefix=True) |
| 264 | load_info = self.model.load_state_dict(state_dict, strict=False) |
| 265 | self.logger.info(load_info) |
no test coverage detected