FlexFlow Train is a deep learning framework that automatically parallelizes and distributes neural network training across GPU/CPU clusters. The system ingests models from PyTorch, Keras, ONNX, or its own C++ API, converts them into a computation graph, and then applies a multi-stage optimization pipeline: graph-level substitutions (pattern matching and rewriting), series-parallel decomposition, cost estimation, and machine mapping via MCMC or dynamic programming. The optimized parallel graph is executed on the Legion/Realm runtime, coordinating distributed tasks across devices. The codebase is organized as a C++ library with Python bindings, comprising ~3500 files including kernels, a compiler, graph utilities, and model definitions. While the core architecture demonstrates sophisticated engineering with extensive tests and reusable graph abstractions, the repository carries significant complexity from heavy template metaprogramming, two execution backends, and partially implemented CPU/HIP kernels, making it more suitable for integration by teams with C++ and distributed systems expertise.

The LLM identified 2 files worth highlighting for reuse: