Welcome to FlexModel’s documentation!

FlexModel is a wrapper for Pytorch models which exposes powerful primitives for model surgery and introspection.

Getting Started

• Introduction
• Examples
• Demos

Check-out our examples for single and multi-gpu, which use both megatron-style layers and PyTorch’s FSDP wrapper. Additionally, we have a demo which includes code for induction head identification in Llama-2-70b.

Key Python API:

• Core: Wrapper and HookFunction
• Distributed: Backend, mappings and strategies
• Traverse: Pytree traversal utility