fix hiyouga#327

GeoffGao · Jul 21, 2023 · eb26e3a · eb26e3a
1 parent 3ed046a
commit eb26e3a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/glmtuner/extras/save_and_load.py b/src/glmtuner/extras/save_and_load.py
@@ -1,6 +1,6 @@
 import os
 import torch
-from typing import Dict
+from typing import Dict, Optional
 
 from transformers.trainer import WEIGHTS_NAME, WEIGHTS_INDEX_NAME
 from transformers.modeling_utils import load_sharded_checkpoint
@@ -12,12 +12,12 @@
 logger = get_logger(__name__)
 
 
-def get_state_dict(model: torch.nn.Module) -> Dict[str, torch.Tensor]: # get state dict containing trainable parameters
+def get_state_dict(model: torch.nn.Module, trainable_only: Optional[bool] = True) -> Dict[str, torch.Tensor]:
     state_dict = model.state_dict()
     filtered_state_dict = {}
 
     for k, v in model.named_parameters():
-        if v.requires_grad:
+        if (not trainable_only) or v.requires_grad:
             filtered_state_dict[k] = state_dict[k].cpu().clone().detach()
 
     return filtered_state_dict

diff --git a/src/glmtuner/tuner/core/trainer.py b/src/glmtuner/tuner/core/trainer.py
@@ -56,7 +56,7 @@ def _save(self, output_dir: Optional[str] = None, state_dict: Optional[Dict[str,
             backbone_model.config.use_cache = True
             backbone_model.save_pretrained(
                 output_dir,
-                state_dict=get_state_dict(backbone_model),
+                state_dict=get_state_dict(backbone_model, trainable_only=(self.finetuning_args.finetuning_type != "full")),
                 safe_serialization=self.args.save_safetensors
             )
             backbone_model.config.use_cache = False