(model)
| 280 | |
| 281 | |
| 282 | def load_deepspeed_state(model): |
| 283 | model = utils.unwrap_model(model) |
| 284 | args = get_args() |
| 285 | load_dir = args.load |
| 286 | if os.path.isdir(load_dir): |
| 287 | model_state_paths = glob(os.path.join(load_dir, "*model_states.pt")) |
| 288 | assert len(model_state_paths) == 1, ( |
| 289 | "only support loading deepspeed checkpoint of model parallel size 1" |
| 290 | ", but got {}".format(model_state_paths) |
| 291 | ) |
| 292 | model_state_path = model_state_paths[0] |
| 293 | else: |
| 294 | model_state_path = load_dir |
| 295 | state_dict = torch.load(model_state_path, map_location="cpu") |
| 296 | state_dict = state_dict["module"] |
| 297 | |
| 298 | model[0].load_state_dict(state_dict, strict=True) |
| 299 | |
| 300 | |
| 301 | def load_checkpoint( |
nothing calls this directly
no test coverage detected