| 214 | return loss.detach() / self.args.gradient_accumulation_steps |
| 215 | |
| 216 | def _save(self, output_dir: Optional[str] = None, state_dict=None): |
| 217 | # If we are executing this function, we are the process zero, so we don't check for that. |
| 218 | output_dir = output_dir if output_dir is not None else self.args.output_dir |
| 219 | os.makedirs(output_dir, exist_ok=True) |
| 220 | logger.info(f"Saving model checkpoint to {output_dir}") |
| 221 | |
| 222 | supported_classes = (PreTrainedModel,) if not is_peft_available() else (PreTrainedModel, PeftModel) |
| 223 | # Save a trained model and configuration using `save_pretrained()`. |
| 224 | # They can then be reloaded using `from_pretrained()` |
| 225 | if not isinstance(self.model, supported_classes): |
| 226 | if state_dict is None: |
| 227 | state_dict = self.model.state_dict() |
| 228 | |
| 229 | if isinstance(unwrap_model(self.model), supported_classes): |
| 230 | unwrap_model(self.model).save_pretrained( |
| 231 | output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors |
| 232 | ) |
| 233 | else: |
| 234 | logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.") |
| 235 | if self.args.save_safetensors: |
| 236 | safetensors.torch.save_file( |
| 237 | state_dict, os.path.join(output_dir, SAFE_WEIGHTS_NAME), metadata={"format": "pt"} |
| 238 | ) |
| 239 | else: |
| 240 | torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME)) |
| 241 | else: |
| 242 | |
| 243 | self.model.save_pretrained( |
| 244 | output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors |
| 245 | ) |
| 246 | |
| 247 | if self.tokenizer is not None: |
| 248 | self.tokenizer.save_pretrained(output_dir) |
| 249 | |
| 250 | # Good practice: save your training arguments together with the trained model |
| 251 | torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME)) |