SplitFM is an open-source framework for Split Parameter-Efficient Fine-Tuning (i.e., SplitLoRA) and Inference (i.e., SplitInfer) for foundation models. SplitLoRA combines the benefits of data privacy protection from Federated Learning (FL) and model partition-based computational offloading from Split Learning (SL), while SplitInfer leverages cloud computing resources to enable large foundation model inference on resource-constrained edge devices without compromising privacy by avoiding data transmission to high-performance servers. We hope SplitFM provides a solid foundation for research focused on advancing the deployment of foundation models in resource-limited, data-sensitive edge network environments.
Supported foundation models:
This repository is based on LoRA.
SplitLoRA contains the source code of the Python package loralib and a example of how to integrate it with PyTorch models, GPT2-s. We only support PyTorch for now. In the future, we will integrate more open source LLMs and more tasks into the SplitLoRA framework
-
The source code of the Python package loralib
-
LoRA fine-tuning implementation of large language models
-
LoRA fine-tuning implementation of large language model under
SplitLoRA framework
We have verified in the environment below:
-
OS: Ubuntu 22.04
-
Python: 3.10.0
| torch 2.9.1 |
transformers 4.36.2 |
spacy | tqdm | tensorboard | progress |
|---|
Note: You still need the original pre-trained checkpoint from Hugging Face to use the LoRA checkpoints.
- Clone the repo and set up the environment.
conda create -n SplitFM python=3.10 -y
conda activate SplitFM- Navigate to the
examplesdirectory and install the required packages.
cd SplitLoRA/examples
pip install -r requirements.txt- Download the necessary pre-trained models, datasets, and evaluation scripts.
# Download pre-trained GPT-2 checkpoints
bash download_pretrained_checkpoints.sh
# Prepare datasets
bash create_datasets.sh
# Download evaluation scripts
cd ./eval
bash download_evalscript.sh
cd ..Our implementation is based on the fine-tuning code for GPT-2 in Hugging Face. There are several directories in this repo:
- src/ contains the source code used for data processing, training, and decoding.
- eval/ contains the code for task-specific evaluation scripts.
- data/ contains the raw data we used in our experiments.
- vocab/ contains the GPT-2 vocabulary files.
| Argument | Description | Default |
|---|---|---|
--train_batch_size |
Training batch size. | 4 |
--grad_acc |
Number of gradient accumulation steps. | 1 |
--seq_len |
Sequence length. | 512 |
--model_card |
Path to the model configuration file. | gpt2.md |
--init_checkpoint |
Path to the initial checkpoint file for model initialization. | ./pretrained_checkpoints/gpt2-medium-pytorch_model.bin.bin. |
--platform |
Execution platform. | local |
--lr |
Learning rate. | 0.0002 |
--max_epoch |
Maximum number of training epochs. | 5 |
--lora_dim |
The dimension of LoRA (Local-Regional Attention). | 4 |
--lora_alpha |
Alpha hyperparameter for LoRA. | 32 |
--lora_dropout |
Dropout rate for LoRA. | 0.1 |
--work_dir |
Working directory where the models and log files are saved. | ./trained_models/GPT2_M/e2e |
For a full list of arguments, please refer to the source code arguments parser.
- Train GPT-2 Medium with SplitLoRA
In the examples directory, run:
python -m torch.distributed.launch --nproc_per_node=1 --use_env src/gpt2_ft_sfl.py \
--train_data0 ./data/e2e/train0.jsonl \
--train_data1 ./data/e2e/train1.jsonl \
--train_data2 ./data/e2e/train2.jsonl \
--valid_data ./data/e2e/valid.jsonl \
--train_batch_size 4 \
--grad_acc 1 \
--valid_batch_size 4 \
--seq_len 512 \
--model_card gpt2.md \
--init_checkpoint ./pretrained_checkpoints/gpt2-medium-pytorch_model.bin \
--platform local \
--clip 0.0 \
--lr 0.0002 \
--weight_decay 0.01 \
--correct_bias \
--adam_beta2 0.999 \
--scheduler linear \
--warmup_step 500 \
--max_epoch 5 \
--save_interval 400000 \
--lora_dim 4 \
--lora_alpha 32 \
--lora_dropout 0.1 \
--label_smooth 0.1 \
--work_dir ./trained_models/GPT2_M/e2e \
--random_seed 40- Generate outputs from the trained model using beam search:
python -m torch.distributed.launch --nproc_per_node=1 src/gpt2_beam.py \
--data ./data/e2e/test.jsonl \
--batch_size 1 \
--seq_len 512 \
--eval_len 64 \
--model_card gpt2.md \
--init_checkpoint ./trained_models/GPT2_M/e2e/{model.name.pt} \
--platform local \
--lora_dim 4 \
--lora_alpha 32 \
--beam 10 \
--length_penalty 0.8 \
--no_repeat_ngram_size 4 \
--repetition_penalty 1.0 \
--eos_token_id 628 \
--work_dir ./trained_models/GPT2_M/e2e \
--output_file predict.26289.b10p08r4.jsonl- Decode outputs from step (2)
python src/gpt2_decode.py \
--vocab ./vocab \
--sample_file ./trained_models/GPT2_M/e2e/predict.26289.b10p08r4.jsonl \
--input_file ./data/e2e/test_formatted.jsonl \
--output_ref_file e2e_ref.txt \
--output_pred_file e2e_pred.txt- Run evaluation on E2E test set
python eval/e2e/measure_scores.py e2e_ref.txt e2e_pred.txt -pSplitInfer is a framework designed to facilitate the split inference of large foundation models, enabling efficient deployment on resource-constrained edge devices while preserving data privacy. By leveraging cloud computing resources, SplitInfer allows for the inference of large models without transmitting sensitive data to high-performance servers. This framework supports popular foundation models such as GPT-2, Llama3, Qwen2-VL, and DeepSeek-R1, providing a flexible and efficient solution for edge computing environments.
OpenAI-compatible HTTP serving is available for inference-only use cases. See SplitInfer/API_SERVER_USAGE_EN.md and SplitInfer/API_SERVER_USAGE_ZH.md for the server guides. Model weights can be configured with CLI parameters, environment variables, or a JSON config file.
We have verified SplitInfer in the following environment:
- OS: Ubuntu 18.04
- Python: 3.8.20
- torch: 2.4.1+cuda118
- transformers: 4.46.3
Note: You still need the original pre-trained checkpoint from Hugging Face or ModelScope to use the model checkpoints.
-
Download the corresponding model files from Hugging Face or ModelScope. For example, using ModelScope download Qwen2-VL-7B-Instruct model:
pip install modelscope modelscope download --model Qwen/Qwen2-VL-7B-Instruct
-
Replace the GPU and model file paths in the corresponding model folder:
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2" # Set the visible CUDA devices (GPUs) for PyTorch model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B" # Your model path
-
Modify the
input_sentenceto ask the model specific questions. You can also adjust the loop count to change the number of tokens generated by the model:input_sentence = "Who is Crayon Shinchan?\n" input_sentence = "Who is the most wealthy person in the world?\n"
-
Run the split model inference demo:
python infer_splitmodel.py
After waiting for 2-3 minutes, you will see the model's response in the command line window.
The repository contains the following directories:
- infer_modelsplit.py: A simple demo for split model inference.
- modelsplit.py: The definition file for the split model, modified based on the model definition files in the transformers library.
- utils.py: Functions for loading model parameters and printing model parameter counts are stored here.
SplitInfer currently supports the following models:
- GPT-2: A widely used language model for text generation tasks.
- Llama3: A high-performance language model optimized for various NLP tasks.
- Qwen2-VL: A vision-language model capable of handling both text and image inputs.
- DeepSeek-R1: A state-of-the-art model designed for complex reasoning and generation tasks.
To run inference with a split model, follow these steps:
- Set up the environment: Ensure all dependencies are installed and the model files are downloaded.
- Configure the model: Update the GPU and model paths in the configuration file.
- Run the inference script: Execute the
infer_splitmodel.pyscript to start the inference process.
We plan to expand SplitInfer to support more foundation models and tasks, further enhancing its flexibility and usability in edge computing environments. Stay tuned for updates!
-SplitLoRA: A Split Parameter-Efficient Fine-Tuning Framework for Large Language Models [Link]
-Efficient Parallel Split Learning Over Resource-Constrained Wireless Edge Networks [Link]
-AdaptSFL: Adaptive Split Federated Learning in Resource-constrained Edge Networks [Link]
-Split Learning in 6G Edge Networks [Link]
-FedSN: A Federated Learning Framework over Heterogeneous LEO Satellite Networks [Link]
-Pushing Large Language Models to the 6G Edge: Vision, Challenges, and Opportunities [Link]
-Automated Federated Pipeline for Parameter-Efficient Fine-Tuning of Large Language Models [Link]
-Hierarchical Split Federated Learning: Convergence Analysis and System Optimization [Link]
-LEO-Split: A Semi-Supervised Split Learning Framework over LEO Satellite Networks [Link]
If you've found SplitFM useful for your project, please cite our paper.
[2025/01/09] 🔥 We are excited to announce the release of the SplitFM v1.1.0 version for Llama3 . This code allows you to efficiently train the Llama3 model by leveraging split learning techniques. You can access and review the code at the following link: [Code]