Just-Talk-Once-Communication-Efficient-Split-Federated-LLM-Fine-Tuning-on-Edge-Devices

This repo hosts a real-system research prototype for L-shaped split federated LLM fine-tuning.

• Real split chain already runs with 2 Jetson Nano CPU clients and 3 Huawei nova Android clients against one shared server-side model on the A800 (CUDA_VISIBLE_DEVICES=4)
• Both full and lora server-side training modes work with the real backend=mft client prefix
• Metrics are logged into CSV/JSONL tables on the server
• third_party/mobilefinetuner/ is the only MobileFineTuner tree modified in this repo
• The stable orchestration path is:
  • remote server + remote Nano clients via scripts/run_multi_nano_experiment.py
  • local Android launch via scripts/run_android_clients_only.py
  • one-shot wrapper orchestration via scripts/run_parallel_hybrid_experiment.py

Layout

• src/lshaped/: Python server, protocol, trainer, loopback client
• clients/cpp/: Flower-compatible C++ client for Nano/Android-side prefix upload
• configs/: runnable YAML configs
• scripts/: sync/bootstrap/run helpers
• docs/: stage plan and troubleshooting notes
• third_party/mobilefinetuner/: copied MobileFineTuner source tree used for incremental integration only

Server Environment

On <SERVER_ALIAS>:

cd <SERVER_REPO_ROOT>
<SERVER_PYTHON> -m pip install -U pip
<SERVER_PYTHON> -m pip install -r requirements-server.txt
<SERVER_PYTHON> -m pip install -e .

Local sync from this Windows workspace to the server:

powershell -ExecutionPolicy Bypass -File scripts\sync_to_server.ps1

Proven Runs

Loopback smoke test on the server:

cd <SERVER_REPO_ROOT>
bash scripts/run_loopback_smoke_remote.sh

Stable 5-client hybrid run, one shared server-side model, 2 Nano + 3 nova:

python scripts/run_parallel_hybrid_experiment.py `
  --base-config configs/a800_gemma_real_full_five_client_long.yaml `
  --nano-client-specs-json configs/dual_nano_clients.json `
  --android-client-specs-json configs/five_clients_android_only_mft.json `
  --run-label full20_wrapper `
  --shared-client-dataset-local-csv data/mmlu/official_mmlu_test_100.csv

LoRA variant:

python scripts/run_parallel_hybrid_experiment.py `
  --base-config configs/a800_gemma_real_lora_five_client_tuned.yaml `
  --nano-client-specs-json configs/dual_nano_clients.json `
  --android-client-specs-json configs/five_clients_android_only_mft.json `
  --run-label lora20_tuned_wrapper `
  --shared-client-dataset-local-csv data/mmlu/official_mmlu_test_100.csv

Nano Client Bootstrap

The helper below installs Nano-side gRPC/protobuf packages, syncs clients/cpp, and builds the client natively on the Nano.

Example from the server or any machine with Python + paramiko:

cd <SERVER_REPO_ROOT>
export NANO_PASSWORD=jetson
<SERVER_PYTHON> scripts/prepare_nano_cpp_client.py \
  --host <NANO_CLIENT_1_HOST> \
  --username jetson \
  --backend mock

Second Nano:

cd <SERVER_REPO_ROOT>
export NANO_PASSWORD=jetson
<SERVER_PYTHON> scripts/prepare_nano_cpp_client.py \
  --host <NANO_CLIENT_2_HOST> \
  --username jetson \
  --backend mock

Run the mock client on a Nano after build:

cd <CLIENT_REPO_ROOT>
./build/cpp_client_mock/lshaped_flower_client \
  --server_address <SERVER_HOST>:19080 \
  --client_id nano64_mock_cpu \
  --backend mock \
  --batch_size 2 \
  --max_seq_len 128 \
  --synthetic_samples 32 \
  --mock_hidden_size 128 \
  --metrics_path <CLIENT_REPO_ROOT>/outputs/client_metrics.csv

The same launch can be triggered from the server with:

cd <SERVER_REPO_ROOT>
export NANO_PASSWORD=jetson
<SERVER_PYTHON> scripts/run_nano_client.py \
  --host <NANO_CLIENT_1_HOST> \
  --username jetson \
  --client-id nano64_mock_cpu \
  --backend mock

The MFT-backed binary can be built with:

cd <SERVER_REPO_ROOT>
export NANO_PASSWORD=jetson
<SERVER_PYTHON> scripts/prepare_nano_cpp_client.py \
  --host <NANO_CLIENT_1_HOST> \
  --username jetson \
  --backend mft

Metrics

Server-side metrics table:

• outputs/<run_name>/metrics.csv
• outputs/<run_name>/metrics.jsonl

Logged fields include:

• loss
• contrastive_ppl_proxy
• accuracy
• transmitted_bytes
• client_encode_time_sec
• client_serialize_time_sec
• client_round_time_sec
• client_rss_mb
• gpu_mem_mb
• gpu_power_w

Client-side CSV on Nano includes per-batch upload bytes and local timing. client_power_w is currently -1.0 because Jetson power sampling is not wired yet.

Every wrapper run also writes:

• outputs/runs/<run_id>/summary_rounds.csv
• outputs/runs/<run_id>/summary_clients.csv
• outputs/runs/<run_id>/summary.json

contrastive_ppl_proxy is exp(loss) for the activation contrastive objective. It is not LM token perplexity.

Important Notes

• This repo does not modify the original <MOBILEFINETUNER_ROOT> tree. All MobileFineTuner changes are applied only to the copied tree under third_party/mobilefinetuner/.
• The current split point is split_layer=0, which means the client uploads scaled input embeddings as the cut-layer activation.
• Server responses do not contain gradients, cut activations, or parameter deltas. This keeps Phase 1 strictly L-shaped and prevents accidental U-shaped behavior.
• In the real-device runs, the actual client samples come from the exported flat MMLU CSV passed through --shared-client-dataset-local-csv. The dataset: block in the server YAML currently provides control-plane settings such as client ids, batch size, and max sequence length.