()
| 338 | ): |
| 339 | # run saving in separate functions to save memory |
| 340 | def _save_model(): |
| 341 | save_dtype = { |
| 342 | "fp16": torch.float16, |
| 343 | "bf16": torch.bfloat16, |
| 344 | "tf32": torch.float, |
| 345 | }[args.precision] |
| 346 | if getattr(args, "only_save_trainable", False): |
| 347 | model_trainable_params = model.get_trainable_params() |
| 348 | model_trainable_params = ['.'.join([_ for _ in key.split('.') if not _.startswith('_')]) |
| 349 | for key in model_trainable_params.keys()] |
| 350 | consolidated_model_state_dict = { |
| 351 | "model": {key: val.to(save_dtype) for key, val in model.state_dict().items() if key in model_trainable_params}, |
| 352 | } |
| 353 | else: |
| 354 | consolidated_model_state_dict = { |
| 355 | "model": {key: val.to(save_dtype) for key, val in model.state_dict().items()}, |
| 356 | } |
| 357 | |
| 358 | model_save_path = os.path.join( |
| 359 | save_dir, |
| 360 | f"consolidated.{mp_rank:02d}-of-{mp_world_size:02d}.model.pth", |
| 361 | ) |
| 362 | if fs_init.get_data_parallel_rank() == 0: |
| 363 | torch.save(consolidated_model_state_dict, model_save_path) |
| 364 | |
| 365 | # Tokenizer |
| 366 | if dist.get_rank() == 0: |
| 367 | model.tokenizer.save(save_dir) |
| 368 | |
| 369 | # Model Args |
| 370 | if dist.get_rank() == 0: |
| 371 | model_args_save_path = os.path.join( |
| 372 | save_dir, |
| 373 | f"config.json", |
| 374 | ) |
| 375 | with open(model_args_save_path, 'w') as f: |
| 376 | json.dump(dataclasses.asdict(model.llma.args), f, indent=2) |
| 377 | |
| 378 | # Meta Information |
| 379 | if dist.get_rank() == 0: |
| 380 | model_meta_save_path = os.path.join( |
| 381 | save_dir, |
| 382 | f"meta.json", |
| 383 | ) |
| 384 | with open(model_meta_save_path, 'w') as f: |
| 385 | json.dump({"llama_type": model.llama_type}, f, indent=2) |
| 386 | |
| 387 | _save_model() |
| 388 | print("model saved") |
no test coverage detected