How thudm/swissarmytransformer works
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: transformer.py Transformer backbone · deepspeed_training.py Training runtime · autoregressive_sampling.py Autoregressive sampling · registry.py Model registry
Components 24
User API
- Arguments and setup
sat/arguments.py - SAT public API
sat/__init__.py
Model Architecture
- Model registry
sat/model/registry.py - Model resources
sat/resources/download.py - Model variants
sat/model/official - Base model
sat/model/base_model.py - Transformer backbone
sat/model/transformer.py - Model mixins
sat/model/mixins.py - Fine-tuning adapters
sat/model/finetune
Data and Tokenization
- Data configuration
sat/data_utils/configure_data.py - Dataset loaders
sat/data_utils/datasets.py - Tokenizers
sat/tokenization
Training and Generation
- Training runtime
sat/training/deepspeed_training.py - Checkpoint I/O
sat/training/model_io.py - Autoregressive sampling
sat/generation/autoregressive_sampling.py - Sampling strategies
sat/generation/sampling_strategies
Parallel and Efficient Compute
- Model parallelism
sat/mpu - Quantized kernels
sat/quantization/kernels.py - Attention operations
sat/ops/memory_efficient_attention.py
Other components
- Application author
- Pretrained weights
- Training data
- Checkpoint storage
- Model output
Connections 28
- 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