FlexModel: A Framework for Interpretability of Distributed Large Language Models
On the map: Mechanistic interpretability → Knowledge and representations inside LLMs
Abstract (arXiv)
With the growth of large language models, now incorporating billions of parameters, the hardware prerequisites for their training and deployment have seen a corresponding increase. Although existing tools facilitate model parallelization and distributed training, deeper model interactions, crucial for interpretability and responsible AI techniques, still demand thorough knowledge of distributed computing. This often hinders contributions from researchers with machine learning expertise but limited distributed computing background. Addressing this challenge, we present FlexModel, a software package providing a streamlined interface for engaging with models distributed across multi-GPU and multi-node configurations. The library is compatible with existing model distribution libraries and encapsulates PyTorch models. It exposes user-registerable HookFunctions to facilitate straightforward interaction with distributed model internals, bridging the gap between distributed and single-device model paradigms. Primarily, FlexModel enhances accessibility by democratizing model interactions and promotes more inclusive research in the domain of large-scale neural networks. The package is found at https://github.com/VectorInstitute/flex_model.
Key facts. Cited by 1 paper in the map (data through 14 Oct 2026). Placed in Knowledge and representations inside LLMs, one of 13 directions in Mechanistic interpretability. Builds on 21 of its 29 references that are in the map. First posted on arXiv 5 Dec 2023.
Builds on
21 of 29 references.
Closest ancestors · 1
Same direction, newest first
Other references · 20
- 2020Language Models are Few-Shot Learnerscited by 32,419
- 2023LLaMA: Open and Efficient Foundation Language Modelscited by 14,132
- 2023Llama 2: Open Foundation and Fine-Tuned Chat Modelscited by 13,028
- 2020Scaling Laws for Neural Language Modelscited by 6,586
- 2022Large Language Models are Zero-Shot Reasonerscited by 4,374
- 2016Pointer Sentinel Mixture Modelscited by 3,096
- 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,868
- 2016Understanding intermediate layers using linear classifier probescited by 1,353
- 2019ZeRO: Memory Optimizations Toward Training Trillion Parameter Modelscited by 1,301
- 2022Toy Models of Superpositioncited by 971
- 2022In-context Learning and Induction Headscited by 964
- 2023Progress measures for grokking via mechanistic interpretabilitycited by 839
8 more
- 2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 715
- 2021BBQ: A Hand-Built Bias Benchmark for Question Answeringcited by 701
- 2023PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallelcited by 660
- 2023The Internal State of an LLM Knows When It's Lyingcited by 659
- 2021Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LMcited by 632
- 2023Larger language models do in-context learning differentlycited by 394
- 2022Maieutic Prompting: Logically Consistent Reasoning with Recursive Explanationscited by 159
- 2023Soft-prompt Tuning for Large Language Models to Evaluate Biascited by 7
Built on by
Citing papers per year: 2026: 1.
1 paper.
Similar, not cited
Close in wording, no citation either way
- 2025FlexOlmo: Open Language Models for Flexible Data Use
- 2026Distributed Interpretability and Control for Large Language Models
- 2025An Explorative Study on Distributed Computing Techniques in Training and Inference of Large Language Models ↗
- 2025Characterizing Communication Patterns in Distributed Large Language Model Inference
- 2026Mechanistic Interpretability for Large Language Model Alignment: Progress, Challenges, and Future Directions
- 2025Speed Always Wins: A Survey on Efficient Architectures for Large Language Models
- 2026Heterogeneous Parallelism for Multimodal Large Language Model Training
- 2026Towards Intrinsic Interpretability of Large Language Models:A Survey of Design Principles and Architectures