FlexModel: A Framework for Interpretability of Distributed Large
  Language Models
Matthew Choi Muhammad Adil Asif John Willes David Emerson
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.