# thudm/swissarmytransformer architecture

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

- Interactive diagram: https://gitdiagram.com/thudm/swissarmytransformer
- Repository: https://github.com/thudm/swissarmytransformer

## Overview

SwissArmyTransformer (SAT) is a Python library for building, adapting, training, and generating with Transformer models. Its main user workflow configures distributed execution, loads a pretrained model and tokenizer, then either runs inference/generation or supplies data and a forward step to train or fine-tune. The shared Transformer backbone, model-specific implementations, mixins, model-parallel layers, and optional quantization support customization and scaling. README examples document the principal API flow; several internal connections are inferred only where the documented workflow or sampled training code supports them.

Main entry points: [`sat/model/transformer.py`](https://github.com/thudm/swissarmytransformer/blob/main/sat/model/transformer.py) (Transformer backbone), [`sat/training/deepspeed_training.py`](https://github.com/thudm/swissarmytransformer/blob/main/sat/training/deepspeed_training.py) (Training runtime), [`sat/generation/autoregressive_sampling.py`](https://github.com/thudm/swissarmytransformer/blob/main/sat/generation/autoregressive_sampling.py) (Autoregressive sampling), [`sat/model/registry.py`](https://github.com/thudm/swissarmytransformer/blob/main/sat/model/registry.py) (Model registry)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_api["User API"]
  node_args["Arguments and setup<br/>[arguments.py]"]
  node_public_api["SAT public API<br/>[__init__.py]"]
end

subgraph group_models["Model Architecture"]
  node_registry["Model registry<br/>[registry.py]"]
  node_resources["Model resources<br/>[download.py]"]
  node_official["Model variants"]
  node_base_model["Base model<br/>[base_model.py]"]
  node_transformer["Transformer backbone<br/>[transformer.py]"]
  node_mixins["Model mixins<br/>[mixins.py]"]
  node_finetune["Fine-tuning adapters"]
end

subgraph group_data["Data and Tokenization"]
  node_data_config["Data configuration<br/>[configure_data.py]"]
  node_datasets["Dataset loaders<br/>[datasets.py]"]
  node_tokenizers["Tokenizers"]
end

subgraph group_runtime["Training and Generation"]
  node_training["Training runtime"]
  node_model_io["Checkpoint I/O<br/>[model_io.py]"]
  node_generation["Autoregressive sampling"]
  node_strategies["Sampling strategies"]
end

subgraph group_scale["Parallel and Efficient Compute"]
  node_mpu["Model parallelism"]
  node_quantization["Quantized kernels<br/>[kernels.py]"]
  node_attention_ops["Attention operations"]
end

node_user(("Application author"))
node_pretrained_store[("Pretrained weights")]
node_training_data[("Training data")]
node_checkpoint_store[("Checkpoint storage")]
node_result["Model output"]

node_user -->|"configures"| node_args
node_user -->|"uses"| node_public_api
node_args -->|"supplies settings"| node_public_api
node_public_api -->|"loads model"| node_registry
node_registry -->|"selects variant"| node_official
node_registry -->|"resolves weights"| node_resources
node_resources -->|"downloads"| node_pretrained_store
node_official -->|"specializes"| node_base_model
node_base_model -->|"builds backbone"| node_transformer
node_base_model -->|"applies hooks"| node_mixins
node_user -->|"adds adapters"| node_finetune
node_finetune -->|"extends model"| node_mixins
node_user -->|"requests tokenizer"| node_tokenizers
node_user -->|"configures input"| node_data_config
node_data_config -->|"creates dataset"| node_datasets
node_datasets -->|"reads"| node_training_data
node_user -->|"starts training"| node_training
node_training -->|"consumes batches"| node_datasets
node_training -->|"loads and saves"| node_model_io
node_model_io -->|"reads and writes"| node_checkpoint_store
node_training -->|"runs forward step"| node_transformer
node_user -->|"requests generation"| node_generation
node_generation -->|"queries model"| node_transformer
node_generation -->|"dispatches decoding"| node_strategies
node_generation -->|"returns tokens"| node_result
node_transformer -->|"uses parallel layers"| node_mpu
node_transformer -.->|"uses attention kernels"| node_attention_ops
node_transformer -.->|"uses quantized layers"| node_quantization

click node_args "https://github.com/thudm/swissarmytransformer/blob/main/sat/arguments.py"
click node_public_api "https://github.com/thudm/swissarmytransformer/blob/main/sat/__init__.py"
click node_registry "https://github.com/thudm/swissarmytransformer/blob/main/sat/model/registry.py"
click node_resources "https://github.com/thudm/swissarmytransformer/blob/main/sat/resources/download.py"
click node_official "https://github.com/thudm/swissarmytransformer/tree/main/sat/model/official"
click node_base_model "https://github.com/thudm/swissarmytransformer/blob/main/sat/model/base_model.py"
click node_transformer "https://github.com/thudm/swissarmytransformer/blob/main/sat/model/transformer.py"
click node_mixins "https://github.com/thudm/swissarmytransformer/blob/main/sat/model/mixins.py"
click node_finetune "https://github.com/thudm/swissarmytransformer/tree/main/sat/model/finetune"
click node_data_config "https://github.com/thudm/swissarmytransformer/blob/main/sat/data_utils/configure_data.py"
click node_datasets "https://github.com/thudm/swissarmytransformer/blob/main/sat/data_utils/datasets.py"
click node_tokenizers "https://github.com/thudm/swissarmytransformer/tree/main/sat/tokenization"
click node_training "https://github.com/thudm/swissarmytransformer/blob/main/sat/training/deepspeed_training.py"
click node_model_io "https://github.com/thudm/swissarmytransformer/blob/main/sat/training/model_io.py"
click node_generation "https://github.com/thudm/swissarmytransformer/blob/main/sat/generation/autoregressive_sampling.py"
click node_strategies "https://github.com/thudm/swissarmytransformer/tree/main/sat/generation/sampling_strategies"
click node_mpu "https://github.com/thudm/swissarmytransformer/tree/main/sat/mpu"
click node_quantization "https://github.com/thudm/swissarmytransformer/blob/main/sat/quantization/kernels.py"
click node_attention_ops "https://github.com/thudm/swissarmytransformer/blob/main/sat/ops/memory_efficient_attention.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_args,node_public_api toneBlue
class node_registry,node_resources,node_official,node_base_model,node_transformer,node_mixins,node_finetune,node_pretrained_store,node_training_data,node_checkpoint_store toneAmber
class node_data_config,node_datasets,node_tokenizers toneMint
class node_training,node_model_io,node_generation,node_strategies toneRose
class node_mpu,node_quantization,node_attention_ops,node_user,node_result toneIndigo
```

## Components

### User API

- [Arguments and setup](https://github.com/thudm/swissarmytransformer/blob/main/sat/arguments.py): `sat/arguments.py`
- [SAT public API](https://github.com/thudm/swissarmytransformer/blob/main/sat/__init__.py): `sat/__init__.py`

### Model Architecture

- [Model registry](https://github.com/thudm/swissarmytransformer/blob/main/sat/model/registry.py): `sat/model/registry.py`
- [Model resources](https://github.com/thudm/swissarmytransformer/blob/main/sat/resources/download.py): `sat/resources/download.py`
- [Model variants](https://github.com/thudm/swissarmytransformer/tree/main/sat/model/official): `sat/model/official`
- [Base model](https://github.com/thudm/swissarmytransformer/blob/main/sat/model/base_model.py): `sat/model/base_model.py`
- [Transformer backbone](https://github.com/thudm/swissarmytransformer/blob/main/sat/model/transformer.py): `sat/model/transformer.py`
- [Model mixins](https://github.com/thudm/swissarmytransformer/blob/main/sat/model/mixins.py): `sat/model/mixins.py`
- [Fine-tuning adapters](https://github.com/thudm/swissarmytransformer/tree/main/sat/model/finetune): `sat/model/finetune`

### Data and Tokenization

- [Data configuration](https://github.com/thudm/swissarmytransformer/blob/main/sat/data_utils/configure_data.py): `sat/data_utils/configure_data.py`
- [Dataset loaders](https://github.com/thudm/swissarmytransformer/blob/main/sat/data_utils/datasets.py): `sat/data_utils/datasets.py`
- [Tokenizers](https://github.com/thudm/swissarmytransformer/tree/main/sat/tokenization): `sat/tokenization`

### Training and Generation

- [Training runtime](https://github.com/thudm/swissarmytransformer/blob/main/sat/training/deepspeed_training.py): `sat/training/deepspeed_training.py`
- [Checkpoint I/O](https://github.com/thudm/swissarmytransformer/blob/main/sat/training/model_io.py): `sat/training/model_io.py`
- [Autoregressive sampling](https://github.com/thudm/swissarmytransformer/blob/main/sat/generation/autoregressive_sampling.py): `sat/generation/autoregressive_sampling.py`
- [Sampling strategies](https://github.com/thudm/swissarmytransformer/tree/main/sat/generation/sampling_strategies): `sat/generation/sampling_strategies`

### Parallel and Efficient Compute

- [Model parallelism](https://github.com/thudm/swissarmytransformer/tree/main/sat/mpu): `sat/mpu`
- [Quantized kernels](https://github.com/thudm/swissarmytransformer/blob/main/sat/quantization/kernels.py): `sat/quantization/kernels.py`
- [Attention operations](https://github.com/thudm/swissarmytransformer/blob/main/sat/ops/memory_efficient_attention.py): `sat/ops/memory_efficient_attention.py`

### Other components

- **Application author**
- **Pretrained weights**
- **Training data**
- **Checkpoint storage**
- **Model output**

## Connections

- Application author → Arguments and setup: configures
- Application author → SAT public API: uses
- Arguments and setup → SAT public API: supplies settings
- SAT public API → Model registry: loads model
- Model registry → Model variants: selects variant
- Model registry → Model resources: resolves weights
- Model resources → Pretrained weights: downloads
- Model variants → Base model: specializes
- Base model → Transformer backbone: builds backbone
- Base model → Model mixins: applies hooks
- Application author → Fine-tuning adapters: adds adapters
- Fine-tuning adapters → Model mixins: extends model
- Application author → Tokenizers: requests tokenizer
- Application author → Data configuration: configures input
- Data configuration → Dataset loaders: creates dataset
- Dataset loaders → Training data: reads
- Application author → Training runtime: starts training
- Training runtime → Dataset loaders: consumes batches
- Training runtime → Checkpoint I/O: loads and saves
- Checkpoint I/O → Checkpoint storage: reads and writes
- Training runtime → Transformer backbone: runs forward step
- Application author → Autoregressive sampling: requests generation
- Autoregressive sampling → Transformer backbone: queries model
- Autoregressive sampling → Sampling strategies: dispatches decoding
- Autoregressive sampling → Model output: returns tokens
- Transformer backbone → Model parallelism: uses parallel layers
- Transformer backbone → Attention operations: uses attention kernels
- Transformer backbone → Quantized kernels: uses quantized layers