(rank, model, path, lora=False, strict=True, logger=None)
| 106 | |
| 107 | @staticmethod |
| 108 | def load_model(rank, model, path, lora=False, strict=True, logger=None): |
| 109 | if os.path.isdir(path): |
| 110 | # dcp |
| 111 | state_dict = get_model_state_dict( |
| 112 | model, |
| 113 | options=StateDictOptions( |
| 114 | full_state_dict=False, |
| 115 | cpu_offload=True, |
| 116 | ), |
| 117 | ) |
| 118 | DCP.load(state_dict=state_dict, checkpoint_id=path) |
| 119 | if rank == 0 and logger is not None: |
| 120 | logger.info(f' Loaded Model from {path}: ') |
| 121 | elif rank == 0: |
| 122 | state_dict = torch.load(path, map_location='cpu') |
| 123 | if lora: |
| 124 | model_dict = model.state_dict() |
| 125 | model_dict.update(state_dict) |
| 126 | if logger is not None: |
| 127 | logger.info("loaded parameters from LoRA.") |
| 128 | missing_keys, unexpected_keys = model.load_state_dict( |
| 129 | state_dict, strict=strict) |
| 130 | if logger is not None: |
| 131 | logger.info(f' Loaded Model from {path}: ' |
| 132 | f' Missing keys: {missing_keys} ' |
| 133 | f' Unexpected keys: {unexpected_keys} ') |
| 134 | |
| 135 | if dist.is_initialized(): |
| 136 | dist.barrier() |
| 137 | |
| 138 | @staticmethod |
| 139 | def load_optimizer(model, optimizer, path): |
no test coverage detected