SCALE-tuner

SCALE-tuner is an analytical collective-communication tuner built on top of ASTRA-sim. It reproduces and analyzes the security-aware collective algorithm selection proposed in our HPCA 2026 paper:

SCALE: Tackling Communication Bottlenecks in Confidential Distributed Machine Learning

This repository focuses on the tuning aspect of SCALE: how encryption overhead fundamentally changes the optimal choice of collective communication algorithms (e.g., Ring, Tree, Hybrid) under confidential computing.

---

Installation

SCALE-tuner directly depends on ASTRA-sim.

Please follow the official ASTRA-sim installation guide as-is:

https://github.com/astra-sim/astra-sim/blob/master/README.md

Make sure ASTRA-sim builds and runs correctly before proceeding.

---

Usage

Benchmark Command

Run bandwidth benchmarks for collective operations.

# Unsecure (baseline) mode - no encryption overhead
python run.py benchmark unsecure allreduce
python run.py benchmark unsecure reducescatter
python run.py benchmark unsecure all  # Run all collectives

# SCALE (secure) mode - with encryption modeling
python run.py benchmark scale allreduce --otp=1GB

# Save results to database
python run.py benchmark scale allreduce --otp=1GB --save

Options:

• --otp=SIZE: OTP buffer size (e.g., 512MB, 4GB). Only valid in scale mode.
• --topo={ring,tree}: Collective topology. Only valid in scale mode.
• --sizes=LIST: Custom comm sizes (e.g., 4MB,8MB,16MB)
• --quick: Use smaller size range for quick testing
• --output={table,csv,aggregated}: Output format
• --save: Persist results to database
• --no-progress: Disable progress output

Sweep Command

Sweep over message sizes and OTP sizes to explore the design space.

# Sweep all collectives
python run.py sweep all --msg-range=1KB:1GB --otp-range=64MB:4GB --save

Tune Command

Get algorithm recommendations based on cached results.

# Generate an algorithm recommendation according to msg size and OTP size
python run.py tune allreduce --msg-size=512MB --otp=4GB

Generate Best Practice Table

Reproduce the optimal algorithm selection table from the paper:

# Generate recommendations for all collectives across message sizes, assuming 4GB OTP
./scripts/generate_scale_best_practice.sh 4GB

Export Command

Export cached results to CSV files for analysis.

# Export with specific OTP size for tuning
python run.py export --otp=4GB

This creates CSV files in the results/ directory:

• unsecure_benchmark.csv - Baseline benchmark results
• scale_benchmark.csv - SCALE benchmark results (ring first, then tree)
• tuning.csv - Recommended topology per collective at each message size

Configuration

Security Modes

Mode	Description	OTP	Topology
unsecure	Baseline without encryption	N/A	N/A
scale	SCALE secure communication	Configurable	ring/tree

Supported Collectives

• allreduce
• reducescatter
• allgather
• alltoall

OTP Sizes

Common sizes: 1KB, 1MB, 16MB, 64MB, 256MB, 512MB, 1GB, 4GB

The OTP (One-Time Pad) buffer size affects the encryption overhead modeling, please refer to the paper for the OTP life cycle.

---

Citation

If you use SCALE-tuner in your research, please cite:

@inproceedings{scale2026,
  title={SCALE: Tackling Communication Bottlenecks in Confidential Distributed Machine Learning},
  author={Park, Joongun and Wang, Yongqin and Xu, Huan and Wu, Hanjiang and Li, Mengyuan and Krishna, Tushar},
  booktitle={HPCA},
  year={2026}
}

---

License

MIT License