FedML Open Source

FedNLP contains 12 datasets designed for training models on four primary NLP tasks:

• text classification
• span extraction
• sequence tagging
• sequence-to-sequence

FedLLM is an MLOps-supported training pipeline designed to help users build Large Language Models (LLMs) on proprietary or private data. It provides minimalist examples for efficient LLM training, fine-tuning, and evaluation using both conventional (centralized) and federated learning approaches.

Key features include:

• Centralized Training: Follows transformers.Trainer logic, integrates with peft for LoRA, supports DeepSpeed, and uses HuggingFace datasets.
• Federated Training: Implements cross-silo federated training/fine-tuning using the FedML platform.
• Default Models: Uses Pythia 7B by default, with support for Llama 2.

FedML provides examples for various Federated Learning (FL) scenarios and machine learning frameworks.

Supported Scenarios

• Cross-silo Federated Learning: Designed for cross-organization or cross-account training.
• Cross-device Federated Learning: Designed for scenarios like smartphones (e.g., using the Beehive platform).
• FL Simulation: Includes single-process simulation, MPI-based simulation, and NCCL-based simulation (the fastest).

Supported Frameworks and Backends

• ML Frameworks: PyTorch, TensorFlow, JAX (with Haiku), and MXNet.
• Communication Backends: MPI, gRPC, PyTorch RPC, and MQTT + S3.
• Secure Aggregation: Supports multi-party computation via light_sec_agg_example.

FedML DevOps utilizes Jenkins for pipeline automation. The infrastructure is organized into three primary components:

1. dockerfile: Contains all Dockerfile definitions required to build the FedML training client, server, and model serving components.
2. k8s: Contains deployment files and Helm Charts used to deploy the training client, server, and model serving components to Kubernetes clusters.
3. scripts: Contains auxiliary scripts used during the Docker image build process.

This guide demonstrates how to build and run LLaMA models using TensorRT-LLM across various configurations, including single GPU, single node multi-GPU, and multi-node multi-GPU setups.

The implementation relies on three primary scripts located in the examples/llama directory:

• build.py: Used to build the TensorRT engines required to run the LLaMA model.
• run.py: Used to perform inference on input text.
• summarize.py: Used to summarize articles from the cnn_dailymail dataset using the model.

The TensorRT-LLM LLaMA implementation supports single GPU, single node multi-GPU, and multi-node multi-GPU configurations. It supports various data types and optimizations including FP16, FP8, INT8 & INT4 Weight-Only quantization, FP8 KV CACHE, Tensor Parallelism, and STRONGLY TYPED modes.

The implementation consists of three primary components:

• build.py: Used to build the TensorRT engine(s) required for the LLaMA model.
• run.py: Used to execute inference on input text.
• summarize.py: Used to perform summarization tasks on the cnn_dailymail dataset.

FedML provides an AllReduce-based Federated Learning simulation framework that supports Multi-Process simulation. This mode allows users to maximize GPU utilization by assigning multiple clients to different GPUs.

The framework consists of three primary roles:

1. Client: Represents the individual user in federated learning who optimizes local models on local datasets.
2. Server: Responsible for aggregating models from all LocalAggregator instances and broadcasting the global model. It uses NCCL acceleration for communication and schedules/assigns specific clients to different LocalAggregator instances (GPUs) in each round.
3. LocalAggregator (GPU): Responsible for simulating assigned clients sequentially, performing a local reduce of those clients' models, and communicating the local reduced model to the server.

FedML is a unified machine learning library designed for running training and deployment at any scale. It provides APIs for several key machine learning paradigms:

• Distributed Training: Accelerate model training using the Lightweight Cheetah framework.
• Federated Learning (FL) Simulation: Simulate FL environments using single-process, MPI-based, or NCCL-based (fastest) simulators.
• Cross-silo Federated Learning: Support for training across different organizations or accounts, including a Python-based edge SDK.
• Cross-device Federated Learning: Support for smartphones and IoT devices via edge SDKs for Android, iOS, and embedded Linux.
• Model Serving: Optimized for edge AI inference.
• MLOps: A complete machine learning operation pipeline for AI running anywhere.

The FedML framework provides built-in APIs and references for a wide variety of datasets across different domains to facilitate experimentation. Supported domains include:

• Computer Vision (CV): MNIST, cifar10, cifar100, fed_cifar100, fed_emnist, cinic10, ImageNet, and Landmarks.
• Natural Language Processing (NLP): shakespeare, fed_shakespeare, and stackoverflow.
• Transportation: AutonomousDriving.
• Finance: lending_club_loan and NUS_WIDE.
• Other: UCI, Synthetic, and edge_case_examples (specifically tailored for the paper "Attack of the Tails: Yes, You Really Can Backdoor Federated Learning").