Save a model checkpoint.
(iteration, model, optimizer, lr_scheduler)
| 108 | |
| 109 | |
| 110 | def save_checkpoint(iteration, model, optimizer, lr_scheduler): |
| 111 | """Save a model checkpoint.""" |
| 112 | args = get_args() |
| 113 | |
| 114 | # Only rank zero of the data parallel writes to the disk. |
| 115 | if not args.deepspeed: |
| 116 | model = utils.unwrap_model(model) |
| 117 | |
| 118 | print_rank_0( |
| 119 | "saving checkpoint at iteration {:7d} to {}".format(iteration, args.save) |
| 120 | ) |
| 121 | |
| 122 | if ( |
| 123 | not torch.distributed.is_initialized() |
| 124 | or mpu.get_data_parallel_rank() == 0 |
| 125 | or args.deepspeed |
| 126 | ): |
| 127 | |
| 128 | # Arguments, iteration, and model. |
| 129 | state_dict = {} |
| 130 | state_dict["args"] = args |
| 131 | state_dict["checkpoint_version"] = 3.0 |
| 132 | state_dict["iteration"] = iteration |
| 133 | state_dict["tokens"] = args.consumed_train_tokens |
| 134 | |
| 135 | # DeepSpeed saves the model/optimizer/scheduler |
| 136 | if not args.deepspeed: |
| 137 | if len(model) == 1: |
| 138 | state_dict["model"] = model[0].state_dict_for_save_checkpoint() |
| 139 | else: |
| 140 | for i in range(len(model)): |
| 141 | mpu.set_virtual_pipeline_model_parallel_rank(i) |
| 142 | state_dict["model%d" % i] = model[ |
| 143 | i |
| 144 | ].state_dict_for_save_checkpoint() |
| 145 | |
| 146 | # Optimizer stuff. |
| 147 | if not args.no_save_optim: |
| 148 | if optimizer is not None: |
| 149 | state_dict["optimizer"] = optimizer.state_dict() |
| 150 | if lr_scheduler is not None: |
| 151 | state_dict["lr_scheduler"] = lr_scheduler.state_dict() |
| 152 | |
| 153 | # RNG states. |
| 154 | if not args.no_save_rng: |
| 155 | state_dict["random_rng_state"] = random.getstate() |
| 156 | state_dict["np_rng_state"] = np.random.get_state() |
| 157 | state_dict["torch_rng_state"] = torch.get_rng_state() |
| 158 | state_dict["cuda_rng_state"] = torch.cuda.get_rng_state() |
| 159 | state_dict["rng_tracker_states"] = mpu.get_cuda_rng_tracker().get_states() |
| 160 | |
| 161 | # Save. |
| 162 | checkpoint_name = get_checkpoint_name(args.save, iteration) |
| 163 | if not args.deepspeed: |
| 164 | ensure_directory_exists(checkpoint_name) |
| 165 | torch.save(state_dict, checkpoint_name) |
| 166 | |
| 167 | if args.deepspeed: |
no test coverage detected