requirements.txt

@@ -19,7 +19,6 @@ sentencepiece
 tiktoken
 modelscope>=1.14.0
 hf-transfer
-safetensors<=0.5.3
 # python
 fire
 omegaconf

src/llamafactory/hparams/data_args.py

@@ -35,7 +35,7 @@ class DataArguments:
         default=None,
         metadata={"help": "The name of dataset(s) to use for evaluation. Use commas to separate multiple datasets."},
     )
-    dataset_dir: str = field(
+    dataset_dir: Union[str, dict] = field(
         default="data",
         metadata={"help": "Path to the folder containing the datasets."},
     )

src/llamafactory/train/callbacks.py

@@ -73,7 +73,7 @@ def fix_valuehead_checkpoint(
     if safe_serialization:
         path_to_checkpoint = os.path.join(output_dir, SAFE_WEIGHTS_NAME)
         with safe_open(path_to_checkpoint, framework="pt", device="cpu") as f:
-            state_dict: dict[str, torch.Tensor] = {key: f.get_tensor(key).clone() for key in f.keys()}
+            state_dict: dict[str, torch.Tensor] = {key: f.get_tensor(key) for key in f.keys()}
     else:
         path_to_checkpoint = os.path.join(output_dir, WEIGHTS_NAME)
         state_dict: dict[str, torch.Tensor] = torch.load(path_to_checkpoint, map_location="cpu", weights_only=True)