# epfml/powersgd architecture

> Architecture diagram of the GitHub repository epfml/powersgd, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-27). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/epfml/powersgd
- Repository: https://github.com/epfml/powersgd

## Overview

PowerSGD is both a reusable distributed-gradient compression library and a paper-experiment codebase. The main user workflow is to compute model gradients, aggregate them across workers with low-rank PowerSGD compression, then apply the optimizer step. The research branch runs configured training and evaluation experiments across image-classification and language-modeling tasks, with multiple reducers and experiment schedules. The graph reflects the README and sampled files; some research-code wiring is documented by the tree and README but not confirmed in the excerpts.

Main entry points: [`paper-code/train.py`](https://github.com/epfml/powersgd/blob/master/paper-code/train.py) (Experiment trainer), [`powersgd/__init__.py`](https://github.com/epfml/powersgd/blob/master/powersgd/__init__.py) (Optimizer step), [`powersgd/powersgd.py`](https://github.com/epfml/powersgd/blob/master/powersgd/powersgd.py) (PowerSGD aggregator), [`powersgd/utils.py`](https://github.com/epfml/powersgd/blob/master/powersgd/utils.py) (Distributed utilities)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_api["PowerSGD API"]
  node_optimizer["Optimizer step<br/>[__init__.py]"]
end

subgraph group_algorithm["Compression Runtime"]
  node_compressor["PowerSGD aggregator<br/>[powersgd.py]"]
  node_orthogonal["Orthogonalization"]
  node_collectives["Distributed utilities<br/>[utils.py]"]
end

subgraph group_research["Research Training"]
  node_paper_train["Experiment trainer<br/>[train.py]"]
  node_reducers["Gradient reducers"]
  node_mean["Metric accumulator"]
  node_cluster["Cluster management<br/>[cluster.py]"]
end

subgraph group_experiments["Tasks And Experiments"]
  node_cifar["CIFAR task<br/>[cifar.py]"]
  node_architectures["CIFAR models"]
  node_language["Language modeling"]
  node_schedules["Experiment schedules"]
end

node_user(("Training caller"))
node_workers(("Distributed workers"))
node_cloud["Google Cloud"]

node_user -->|"calls"| node_optimizer
node_optimizer -->|"aggregates gradients"| node_compressor
node_compressor -->|"orthogonalizes factors"| node_orthogonal
node_compressor -->|"all-reduces factors"| node_collectives
node_collectives -->|"communicates"| node_workers
node_optimizer -->|"steps optimizer"| node_user
node_paper_train -->|"trains or evaluates"| node_cifar
node_paper_train -->|"trains or evaluates"| node_language
node_cifar -->|"builds model"| node_architectures
node_paper_train -->|"uses configured reducer"| node_reducers
node_paper_train -->|"aggregates metrics"| node_mean
node_schedules -->|"configures experiments"| node_paper_train
node_cluster -->|"manages instances"| node_cloud

click node_optimizer "https://github.com/epfml/powersgd/blob/master/powersgd/__init__.py"
click node_compressor "https://github.com/epfml/powersgd/blob/master/powersgd/powersgd.py"
click node_orthogonal "https://github.com/epfml/powersgd/blob/master/powersgd/orthogonalization.py"
click node_collectives "https://github.com/epfml/powersgd/blob/master/powersgd/utils.py"
click node_paper_train "https://github.com/epfml/powersgd/blob/master/paper-code/train.py"
click node_reducers "https://github.com/epfml/powersgd/blob/master/paper-code/gradient_reducers.py"
click node_mean "https://github.com/epfml/powersgd/blob/master/paper-code/mean_accumulator.py"
click node_cifar "https://github.com/epfml/powersgd/blob/master/paper-code/tasks/cifar.py"
click node_architectures "https://github.com/epfml/powersgd/tree/master/paper-code/tasks/cifar_architectures"
click node_language "https://github.com/epfml/powersgd/tree/master/paper-code/tasks/language_modeling"
click node_schedules "https://github.com/epfml/powersgd/tree/master/paper-code/schedule"
click node_cluster "https://github.com/epfml/powersgd/blob/master/paper-code/environment/cluster.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_optimizer toneBlue
class node_compressor,node_orthogonal,node_collectives toneAmber
class node_paper_train,node_reducers,node_mean,node_cluster toneMint
class node_cifar,node_architectures,node_language,node_schedules,node_workers toneRose
class node_user,node_cloud toneIndigo
```

## Components

### PowerSGD API

- [Optimizer step](https://github.com/epfml/powersgd/blob/master/powersgd/__init__.py): `powersgd/__init__.py`

### Compression Runtime

- [PowerSGD aggregator](https://github.com/epfml/powersgd/blob/master/powersgd/powersgd.py): `powersgd/powersgd.py`
- [Orthogonalization](https://github.com/epfml/powersgd/blob/master/powersgd/orthogonalization.py): `powersgd/orthogonalization.py`
- [Distributed utilities](https://github.com/epfml/powersgd/blob/master/powersgd/utils.py): `powersgd/utils.py`

### Research Training

- [Experiment trainer](https://github.com/epfml/powersgd/blob/master/paper-code/train.py): `paper-code/train.py`
- [Gradient reducers](https://github.com/epfml/powersgd/blob/master/paper-code/gradient_reducers.py): `paper-code/gradient_reducers.py`
- [Metric accumulator](https://github.com/epfml/powersgd/blob/master/paper-code/mean_accumulator.py): `paper-code/mean_accumulator.py`
- [Cluster management](https://github.com/epfml/powersgd/blob/master/paper-code/environment/cluster.py): `paper-code/environment/cluster.py`

### Tasks And Experiments

- [CIFAR task](https://github.com/epfml/powersgd/blob/master/paper-code/tasks/cifar.py): `paper-code/tasks/cifar.py`
- [CIFAR models](https://github.com/epfml/powersgd/tree/master/paper-code/tasks/cifar_architectures): `paper-code/tasks/cifar_architectures`
- [Language modeling](https://github.com/epfml/powersgd/tree/master/paper-code/tasks/language_modeling): `paper-code/tasks/language_modeling`
- [Experiment schedules](https://github.com/epfml/powersgd/tree/master/paper-code/schedule): `paper-code/schedule`

### Other components

- **Training caller**
- **Distributed workers**
- **Google Cloud**

## Connections

- Training caller → Optimizer step: calls
- Optimizer step → PowerSGD aggregator: aggregates gradients
- PowerSGD aggregator → Orthogonalization: orthogonalizes factors
- PowerSGD aggregator → Distributed utilities: all-reduces factors
- Distributed utilities → Distributed workers: communicates
- Optimizer step → Training caller: steps optimizer
- Experiment trainer → CIFAR task: trains or evaluates
- Experiment trainer → Language modeling: trains or evaluates
- CIFAR task → CIFAR models: builds model
- Experiment trainer → Gradient reducers: uses configured reducer
- Experiment trainer → Metric accumulator: aggregates metrics
- Experiment schedules → Experiment trainer: configures experiments
- Cluster management → Google Cloud: manages instances