In this tutorial, we demonstrate how to federate fine-tuning of a large language model using LoRA without ever centralizing private text data. We simulate multiple organizations as virtual clients and show how each client adapts a shared base model locally while exchanging only lightweight LoRA adapter parameters. By combining Flower’s federated learning simulation engine with parameter-efficient fine-tuning, we demonstrate a practical, scalable approach for organizations that want to customize LLMs on sensitive data while preserving privacy and reducing communication and compute costs. Check out the FULL CODES here.
],
2: [
“Fleet ops: preventive maintenance reduces downtime; prioritize vehicles with repeated fault codes.”,
“Dispatch note: optimize routes by time windows and driver hours to reduce empty miles.”,
“Safety policy: enforce rest breaks and log inspections before long-haul trips.”,
“Inventory update: track spare parts usage; reorder thresholds should reflect lead time and seasonality.”,
“Customer SLA: late deliveries require proactive notifications and documented root cause.”
],
}
for cid in list(CLIENT_TEXTS.keys()):
base = CLIENT_TEXTS[cid]
CLIENT_TEXTS[cid] = base + [f”Q: Summarize this for leadership. A: {t}” for t in base]
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
bnb_config: Optional[BitsAndBytesConfig] = None
if DEVICE == “cuda”:
compute_dtype = torch.bfloat16 if torch.cuda.get_device_capability(0)[0] >= 8 else torch.float16
bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type=”nf4″, bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=compute_dtype)
if “gpt2” in MODEL_ID.lower():
TARGET_MODULES = [“c_attn”, “c_proj”]
else:
TARGET_MODULES = [“q_proj”, “k_proj”, “v_proj”, “o_proj”]
LORA_R = 16
LORA_ALPHA = 32
LORA_DROPOUT = 0.05
lora_config = LoraConfig(r=LORA_R, lora_alpha=LORA_ALPHA, lora_dropout=LORA_DROPOUT, bias=”none”, task_type=”CAUSAL_LM”, target_modules=TARGET_MODULES)
def model_primary_device(model) -> torch.device:
return next(model.parameters()).device
def build_model_with_lora():
if DEVICE == “cuda”:
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map=”auto”, quantization_config=bnb_config, torch_dtype=”auto”)
model = prepare_model_for_kbit_training(model)
else:
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
model.to(“cpu”)
model = get_peft_model(model, lora_config)
model.train()
return model
def make_dataset(texts: List[str]) -> Dataset:
ds = Dataset.from_dict({“text”: texts})
def tok(batch):
return tokenizer(batch[“text”], truncation=True, max_length=MAX_LEN, padding=”max_length”)
ds = ds.map(tok, batched=True, remove_columns=[“text”])
return ds
collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
def lora_state_keys(model) -> List[str]:
sd = model.state_dict()
keys = sorted([k for k in sd.keys() if “lora_” in k])
if not keys:
raise RuntimeError(“No LoRA keys found. Your model might not have the target_modules specified. ” f”Current TARGET_MODULES={TARGET_MODULES}, MODEL_ID={MODEL_ID}”)
return keys
def get_lora_ndarrays(model) -> List[np.ndarray]:
sd = model.state_dict()
keys = lora_state_keys(model)
return [sd[k].detach().float().cpu().numpy() for k in keys]
def set_lora_ndarrays(model, arrays: List[np.ndarray]) -> None:
keys = lora_state_keys(model)
if len(keys) != len(arrays):
raise ValueError(f”Mismatch: got {len(arrays)} arrays but expected {len(keys)}.”)
sd = model.state_dict()
for k, arr in zip(keys, arrays):
t = torch.from_numpy(arr).to(sd[k].device).to(sd[k].dtype)
sd[k].copy_(t)
def cosine_warmup_lr(step: int, total_steps: int, base_lr: float, warmup_steps: int) -> float:
if step float:
model.eval()
dl = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=False, collate_fn=collator)
losses = []
dev = model_primary_device(model)
for i, batch in enumerate(dl):
if i >= max_batches:
break
batch = {k: v.to(dev) for k, v in batch.items()}
out = model(**batch, labels=batch[“input_ids”])
losses.append(float(out.loss.detach().cpu()))
model.train()
return float(np.mean(losses)) if losses else float(“nan”)
def train_one_client_round(model, ds: Dataset, epochs: int, lr: float, grad_accum: int, warmup_steps: int) -> Tuple[float, int]:
dl = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collator)
total_steps = max(1, (len(dl) * epochs) // max(1, grad_accum))
step = 0
optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=WEIGHT_DECAY)
optimizer.zero_grad(set_to_none=True)
running = []
examples = 0
dev = model_primary_device(model)
for _ in range(epochs):
for bi, batch in enumerate(dl):
batch = {k: v.to(dev) for k, v in batch.items()}
out = model(**batch, labels=batch[“input_ids”])
loss = out.loss / grad_accum
loss.backward()
running.append(float(loss.detach().cpu()) * grad_accum)
examples += batch[“input_ids”].shape[0]
if (bi + 1) % grad_accum == 0:
lr_t = cosine_warmup_lr(step, total_steps, lr, warmup_steps)
for pg in optimizer.param_groups:
pg[“lr”] = lr_t
optimizer.step()
optimizer.zero_grad(set_to_none=True)
step += 1
if step % LOG_EVERY == 0:
print(f” step={step}/{total_steps} loss={np.mean(running[-LOG_EVERY:]):.4f} lr={lr_t:.2e}”)
return float(np.mean(running)) if running else float(“nan”), examples
We set up the full execution environment and define all global configurations required for the experiment. We prepare the private client text silos, tokenizer, LoRA configuration, and model-loading logic so they automatically adapt to CPU or GPU availability. We also establish all helper utilities that enable parameter-efficient fine-tuning and safe device handling across federated clients. Check out the FULL CODES here.
def __init__(self, cid: int):
self.cid = cid
self._model = None
self._ds_train = None
self._ds_eval = None
def _ensure(self):
if self._model is None:
print(f”[Client {self.cid}] Loading model + LoRA (MODEL_ID={MODEL_ID})…”)
self._model = build_model_with_lora()
texts = CLIENT_TEXTS[self.cid].copy()
random.shuffle(texts)
split = max(1, int(0.8 * len(texts)))
self._ds_train = make_dataset(texts[:split])
self._ds_eval = make_dataset(texts[split:])
def get_parameters(self, config):
self._ensure()
return get_lora_ndarrays(self._model)
def fit(self, parameters, config):
self._ensure()
set_lora_ndarrays(self._model, parameters)
loss_before = eval_loss(self._model, self._ds_eval, max_batches=10)
print(f”[Client {self.cid}] eval_loss_before={loss_before:.4f}”)
train_loss, n_examples = train_one_client_round(self._model, self._ds_train, epochs=int(config.get(“local_epochs”, LOCAL_EPOCHS)), lr=float(config.get(“lr”, LR)), grad_accum=int(config.get(“grad_accum”, GRAD_ACCUM)), warmup_steps=int(config.get(“warmup_steps”, WARMUP_STEPS)))
loss_after = eval_loss(self._model, self._ds_eval, max_batches=10)
print(f”[Client {self.cid}] train_loss={train_loss:.4f} eval_loss_after={loss_after:.4f}”)
new_params = get_lora_ndarrays(self._model)
metrics = {“eval_loss_before”: loss_before, “eval_loss_after”: loss_after, “train_loss”: train_loss}
return new_params, n_examples, metrics
def evaluate(self, parameters, config):
self._ensure()
set_lora_ndarrays(self._model, parameters)
loss = eval_loss(self._model, self._ds_eval, max_batches=20)
return float(loss), len(self._ds_eval), {“eval_loss”: float(loss)}
def client_fn(context: Context):
cid = None
try:
cid = int(context.node_config.get(“partition-id”))
except Exception:
try:
cid = int(context.node_id)
except Exception:
cid = 0
return FedLoRAClient(cid).to_client()
We define the federated client logic that simulates independent organizations participating in training. We initialize a LoRA-augmented language model per client and ensure that local datasets remain isolated. We implement client-side training, evaluation, and parameter exchange while exposing only LoRA adapter weights to the server.
return {“local_epochs”: LOCAL_EPOCHS, “lr”: LR, “grad_accum”: GRAD_ACCUM, “warmup_steps”: WARMUP_STEPS}
strategy = fl.server.strategy.FedAvg(fraction_fit=1.0, fraction_evaluate=1.0, min_fit_clients=NUM_CLIENTS, min_evaluate_clients=NUM_CLIENTS, min_available_clients=NUM_CLIENTS, on_fit_config_fn=fit_config)
print(“\nStarting Flower simulation…\n”)
client_resources = {“num_cpus”: 2, “num_gpus”: 0.0}
if DEVICE == “cuda”:
client_resources = {“num_cpus”: 2, “num_gpus”: 0.25}
history = fl.simulation.start_simulation(client_fn=client_fn, num_clients=NUM_CLIENTS, config=fl.server.ServerConfig(num_rounds=ROUNDS), strategy=strategy, client_resources=client_resources, ray_init_args={“include_dashboard”: False, “ignore_reinit_error”: True})
print(“\nSimulation done.”)
We configure the federated learning strategy and orchestrate the global training process. We specify how many clients participate, how parameters are aggregated, and how training rounds are scheduled. We then launch the Flower simulation to coordinate communication and aggregation across all virtual clients. Check out the FULL CODES here.
demo_model.eval()
prompt = “Summarize this internal note for leadership in 2 bullets:\nDispatch note: optimize routes by time windows and driver hours to reduce empty miles.\n\nAnswer:”
inputs = tokenizer(prompt, return_tensors=”pt”)
dev = model_primary_device(demo_model)
inputs = {k: v.to(dev) for k, v in inputs.items()}
with torch.no_grad():
out = demo_model.generate(**inputs, max_new_tokens=80, do_sample=True, temperature=0.8, top_p=0.95, repetition_penalty=1.05, eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id)
print(“\n=== Generation output ===\n”)
print(tokenizer.decode(out[0], skip_special_tokens=True))
We load a final LoRA-augmented model instance to demonstrate inference after federated training. We prepare a realistic prompt and run text generation using the same architecture employed during training. We verify that the pipeline executes end-to-end by producing coherent, task-aligned outputs.
print(history.__dict__.keys())
We inspect the training artifacts and simulation outputs produced by the federated run. We examine the returned history object to confirm that rounds, metrics, and aggregation completed successfully. We use this step to validate the overall integrity and reproducibility of the federated fine-tuning workflow.
In conclusion, we showed that federated fine-tuning of LLMs is not only a research concept but something we can run end-to-end in a Colab environment today. We successfully coordinated client-side LoRA training, server-side aggregation, and evaluation without sharing raw text or full model weights. This workflow highlights how federated learning, when paired with modern PEFT techniques, enables privacy-preserving adaptation of generative models and provides a strong foundation for extending the system toward personalization, robustness, and real-world enterprise deployment.
Check out the FULL CODES here. Also, feel free to follow us on Twitter and don’t forget to join our 100k+ ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.
