🚀 Distributed LoRA Trainer

Reproducible Infrastructure for Fine-Tuning Large Language Models
DDP · FSDP · DeepSpeed · QLoRA · W&B · Airflow · Kubernetes

This project provides a production-grade, reproducible training infrastructure for fine-tuning LLMs with parameter-efficient adapters (LoRA/QLoRA). It’s designed for:

• Scalable distributed training across multi-GPU/multi-node clusters.
• Experiment reproducibility with config-driven workflows.
• High throughput with sharded streaming datasets + optimized I/O.
• Research enablement: rapid prototyping of new training methods.

• ⚡ Distributed Backends: PyTorch DDP, FSDP, DeepSpeed (Zero-2/3).
• 💾 Parameter-Efficient Fine-Tuning: LoRA + 4-bit QLoRA (bitsandbytes).
• 📊 Experiment Tracking: Automated logging, checkpointing, and sweeps with Weights & Biases.
• 🗂 Streaming Datasets: Shard-aware Hugging Face datasets with per-rank partitioning.
• 🧩 Config-Driven: YAML + .env toggles for reproducibility.
• ☸️ Orchestration: Airflow DAGs + Kubernetes CronJobs for scheduled training.
• 🛠 Resilience: Standardized checkpoints, monitoring, and failure recovery.

ddp-lora-trainer/
│
├── configs/                # YAML configs
├── src/                    # Core training code
├── scripts/                # Launchers & sweep runners
├── sweeps/                 # W&B sweep configs
├── orchestration/
│   ├── airflow/            # Airflow DAGs
│   └── k8s/                # Kubernetes manifests
├── deepspeed/              # ZeRO configs
├── Dockerfile
├── README.md
└── .env.example

git clone https://github.com/<your-username>/ddp-lora-trainer.git
cd ddp-lora-trainer
pip install -r requirements.txt

bash scripts/launch_local.sh 4 configs/base.yaml

python -m src.train --config configs/base.yaml   data.use_hf_streaming=true data.hf_dataset=wikitext data.hf_split=train

python -m src.train --config configs/base.yaml   distributed.backend=deepspeed distributed.deepspeed_config=./deepspeed/zero2.json

python -m src.train --config configs/base.yaml   distributed.backend=fsdp fsdp.sharding_strategy=full fsdp.mixed_precision=bf16

python -m src.train --config configs/base.yaml   quantization.load_in_4bit=true quantization.compute_dtype=bf16

Run training with W&B logging:

WANDB_PROJECT=ddp-lora-trainer WANDB_API_KEY=your_key bash scripts/launch_local.sh 4 configs/base.yaml

Run a sweep:

bash scripts/wandb_sweep.sh

• Located in orchestration/airflow/dags/train_ddp_lora.py.
• Configure via Airflow Variables (ddp_lora_repo_path, hf_dataset, etc.).

• orchestration/k8s/cronjob.yaml schedules recurring training jobs.
• Requires GPU node pool + wandb-secret + PVC ddp-lora-pvc.

python scripts/bench_efficiency.py --config configs/base.yaml --gpus 8

Generates scaling efficiency plots (tokens/sec vs GPUs).

Apache-2.0

⚡ Maintainer: Tharun Jagarlamudi