Initialize megatron_util environment for megatron_based model. If argument `megatron_cfg` is not specified, then the megatorn_cfg will be load from configuration.json file in the model_dir. Args: megatron_cfg (Dict, optional): Megatron Config will be send to megatron_util.
(megatron_cfg=None, model_dir=None, **kwargs)
| 34 | |
| 35 | |
| 36 | def init_megatron_util(megatron_cfg=None, model_dir=None, **kwargs): |
| 37 | """Initialize megatron_util environment for megatron_based model. |
| 38 | |
| 39 | If argument `megatron_cfg` is not specified, then the megatorn_cfg will be load |
| 40 | from configuration.json file in the model_dir. |
| 41 | |
| 42 | Args: |
| 43 | megatron_cfg (Dict, optional): Megatron Config will be send to megatron_util. |
| 44 | model_dir (str, optional): The model path for configuration. Defaults to None. |
| 45 | """ |
| 46 | from modelscope.utils.hub import read_config |
| 47 | from megatron_util import initialize_megatron |
| 48 | |
| 49 | assert not (megatron_cfg is None and model_dir is None), \ |
| 50 | 'cfg and model_dir cannot both be None when initializing megatron_util' |
| 51 | if megatron_cfg is None: |
| 52 | cfg = read_config(model_dir) |
| 53 | try: |
| 54 | megatron_cfg = cfg.megatron |
| 55 | except AttributeError: |
| 56 | try: |
| 57 | model_type = cfg.model.type |
| 58 | except AttributeError: |
| 59 | # Fit models without model type, such as mglm |
| 60 | model_type = cfg.pipeline.type |
| 61 | megatron_cfg = _DEFAULT_CFG_WITH_MODEL_TYPE[model_type] \ |
| 62 | if model_type in _DEFAULT_CFG_WITH_MODEL_TYPE else {} |
| 63 | megatron_cfg.update(kwargs) |
| 64 | initialize_megatron(megatron_cfg) |
| 65 | global _IS_MEGATRON_INITIALIZED |
| 66 | _IS_MEGATRON_INITIALIZED = True |
| 67 | |
| 68 | |
| 69 | def is_megatron_initialized() -> bool: |
searching dependent graphs…