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.

翻译：摘要：随着大语言模型规模的增长（现已包含数十亿参数），其训练与部署所需的硬件门槛也随之提升。尽管现有工具支持模型并行化与分布式训练，但针对可解释性与负责任人工智能等关键技术所需的深层模型交互，仍要求研究者具备分布式计算的专业知识。这往往阻碍了具备机器学习专长但分布式计算背景有限的研究者参与相关工作。针对这一挑战，我们提出FlexModel——一款提供简化接口的软件包，支持与跨多GPU及多节点配置的分布式模型进行交互。该库兼容现有模型分布式库，可封装PyTorch模型，并通过暴露用户可注册的HookFunctions，实现与分布式模型内部结构的直接交互，从而弥合分布式与单设备模型范式之间的鸿沟。核心而言，FlexModel通过降低模型交互门槛提升可访问性，推动大规模神经网络领域的包容性研究。本程序包已开源至https://github.com/VectorInstitute/flex_model。