pfl - Private Federated Learning Simulation Framework

repository·develop·Indexed 18 days ago

https://github.com/apple/pfl-research

A Python framework developed by Apple for simulating Private Federated Learning (PFL). It supports efficient, scalable simulations with privacy-preserving features like Differential Privacy across various models (Neural Networks, GBDTs) and frameworks including PyTorch, TensorFlow, and MLX. The library includes official benchmarks for image classification (CIFAR10), language models (StackOverflow), and the FLAIR dataset.

Tokens
38.6K
Snippets
111
Records
163
Agent score
62%

What's inside pfl

  1. Overview of the pfl framework

    develop

    pfl is a Python framework developed by Apple designed for researchers to run efficient simulations of privacy-preserving federated learning (FL) and disseminate research results.

    Key Capabilities:

    • Rapid Prototyping: Quickly test PFL use cases with existing models and data.
    • Scalable Simulations: Supports multiple levels of distributed training across processes, GPUs, and machines.
    • Framework Flexibility: Provides APIs to express new ideas in models, algorithms, federated datasets, and privacy mechanisms.
    • Multi-Framework Support: Compatible with both PyTorch and TensorFlow.
    • Diverse Model Support: Supports neural networks and other models like GBDTs (Gradient Boosted Decision Trees).
    • Privacy Integration: Tight integration with local and central differential privacy mechanisms.
    • Unified Benchmarks: Provides vetted datasets compatible with both TensorFlow and PyTorch.

    Note: This framework is intended for research simulations and is not designed for third-party production FL deployments, though simulation results can inform actual deployments.

  2. Overview of `pfl` features

    develop

    The pfl framework is designed for simulating Private Federated Learning (PFL) research. Key capabilities include:

    • Fast Simulations: Supports multiple levels of distributed training across processes, GPUs, and machines.
    • Framework Support: Native support for both PyTorch and TensorFlow, as well as MLX (for Apple silicon).
    • Model Flexibility: Supports neural networks and other models like GBDTs (Gradient Boosted Decision Trees).
    • Privacy Integration: Tight integration with local and central Differential Privacy (DP) mechanisms.
    • Unified Benchmarks: Vetted datasets available for both PyTorch and TensorFlow.

    Note: pfl is intended for research simulations and is not intended for third-party production FL deployments.

  3. Overview of pfl-research simulation framework

    develop

    The pfl-research repository provides a simulation framework designed to accelerate research in Private Federated Learning (PFL). It contains benchmarking code used to compare different Federated Learning (FL) frameworks, specifically as presented in the NeurIPS 2024 benchmark and dataset track paper: "pfl-research: simulation framework for accelerating research in Private Federated Learning" (arXiv:2404.06430).

    This repository specifically implements the experimental setups for Table 1 and Table 2 of the aforementioned publication. Detailed setup descriptions can be found in Appendix D of the paper and within the individual directory structures of the repository.

  4. Understand the MDM repository structure

    develop

    The MDM research code is organized into three main components:

    • mdm/: Contains the core algorithmic implementation of the MDM model within the pfl-research framework.
    • mdm_paper/: Contains research-specific execution code:
      • training/: Python scripts for running MDM parameter inference on CIFAR-10 and FEMNIST.
      • notebooks/: Jupyter notebooks for result visualization and generating paper plots.
    • mdm_utils/: Provides utility functions for training, including argument parsers and dataset helpers.
  5. Access code for 'Improved Modelling of Federated Datasets using Mixtures-of-Dirichlet-Multinomials'

    develop
    This repository contains the implementation for the research paper "Improved Modelling of Federated Datasets using Mixtures-of-Dirichlet-Multinomials". It includes the code required to execute all experiments described in the paper and the post-processing scripts used to generate the results and plots presented in the publication.
  6. Available PFL official benchmarks

    develop

    The pfl package provides several official benchmarks for simulating different scenarios:

    • image_classification: Training small CNNs on the CIFAR10 dataset.
    • lm: Training transformer models on the StackOverflow dataset.
    • flair: Training ResNet18 on the FLAIR dataset.
  7. Optimize Central Evaluation performance

    develop

    Central evaluation using CentralEvaluationCallback can be a bottleneck. Use these three strategies to minimize compute time:

    1. Reduce Frequency: Use evaluation_frequency to run evaluation less often than every central iteration.
    2. Increase Batch Size: Set a larger batch size in ModelHyperParams specifically for evaluation to speed up processing.
    3. Shard Evaluation: In distributed simulations, pfl can shard the evaluation workload across available GPUs.
  8. Optimize multi-process training on a single GPU

    develop

    If your models are small or user datasets are small, you can achieve speedups by running multiple processes that share a single GPU. This helps offset the overhead caused by federated learning.

    Best Practices:

    • The optimal number of processes per GPU is typically between 1 and 5.
    • For PyTorch, use torchrun with --nproc_per_node set higher than the number of available GPUs. For example, if you have 8 GPUs and want 4 processes per GPU, set --nproc_per_node=32.
    • GPU Memory Warning: You must allow multiple processes to allocate memory on the GPU simultaneously. You can switch the GPU's compute mode to "Default" using nvidia-smi -c 0 (use -i <index> to target specific GPUs).
    # Example: 2 processes per GPU using 2 GPUs (4 total processes)
    export PFL_WORKER_ADDRESSES=localhost:8000,localhost:8001,localhost:8002,localhost:8003
    PFL_WORKER_RANK=0 python train.py &
    PFL_WORKER_RANK=1 python train.py &
    PFL_WORKER_RANK=2 python train.py &
    PFL_WORKER_RANK=3 python train.py &
  9. Implement custom privacy mechanisms using PrivacyMechanism

    develop
    The pfl.privacy.privacy_mechanism module provides the abstract base classes for defining differential privacy (DP) mechanisms. To implement a custom mechanism, you should inherit from the base classes provided in this module. Note that NoPrivacy and NormClippingOnly are specialized implementations that do not add noise and are excluded from the general mechanism interface.
  10. Understand the Public API vs Internal modules

    develop

    The pfl project follows semantic versioning. Breaking changes are only permitted within the pfl.internal namespace.

    • Public API: Everything in the pfl module except for pfl.internal.
    • Internal API: Everything inside pfl.internal. You can make breaking changes here without violating semantic versioning.

    To maintain framework encapsulation, code for specific frameworks (e.g., PyTorch) is split into:

    • A dataset module (pfl.data.pytorch)
    • A model module (pfl.model.pytorch)
    • An ops module (pfl.internal.ops.pytorch_ops)
    • A bridge module (pfl.internal.bridge.pytorch)

    To use the currently selected framework's operations dynamically, use the selector module.

    from pfl.internal.ops.selector import get_default_framework_module as ops
    
    # Use the selected framework's ops
    ops().get_shape(tensor)
  11. Choose the right Federated Dataset type

    develop

    The choice of dataset class depends on your data size and loading requirements:

    • pfl.data.dataset.Dataset: Fastest option. Use this if all data fits into RAM.
    • pfl.data.tensorflow.TFFederatedDataset: Use this for TensorFlow models if you need lazy loading from disk or heavy preprocessing to enable parallelization.
    • pfl.data.pytorch.PyTorchFederatedDataset: Use this for PyTorch models if you need lazy loading from disk or heavy preprocessing to enable parallelization.
  12. Use Backend classes for simulation environments

    develop

    The pfl.aggregate module provides two primary backend abstractions for managing simulation environments:

    1. pfl.aggregate.base.Backend: The base class defining the interface for all backends.
    2. pfl.aggregate.simulate.SimulatedBackend: A concrete implementation used for running simulations within a controlled environment.

    Developers should use SimulatedBackend when they need to orchestrate federated learning tasks in a simulated setting.