English | 中文
FlagTrain is part of FlagOS, an open-source system software stack for diverse AI chips. FlagOS connects models, systems, and chips through software, reducing the cost of porting and maintaining models across hardware platforms.
FlagTrain is a library of training operators and compute components implemented in Triton. It targets mainstream training frameworks and key components, with reusable training computation implementations that can be optimized to accelerate large-model training across multiple hardware backends.
For training framework developers, FlagTrain aims to provide Triton implementations aligned with the relevant interfaces and computational semantics, reducing repeated development and hardware adaptation work. For operator developers, FlagTrain provides a collaborative platform for organizing implementations, tests, and performance optimization around real training workloads.
The project is in its initial development stage. The repository contains documentation and a component directory skeleton; no operator implementations are available yet, and component and hardware compatibility have not been validated.
- Training computation in Triton: Develop reusable operators for forward computation, backpropagation, and related compute modules in training.
- Fusion and performance optimization: Optimize operator fusion, memory access, and parallel computation for training workloads to reduce intermediate data transfers and execution overhead.
- Multiple hardware backends: Build on compatible Triton compilation and runtime environments to adapt functionality and tune performance across hardware backends.
- Upstream collaboration: Align with the interfaces, computational semantics, and usage requirements of target frameworks and components, and keep pace with upstream changes.
- Verifiable implementations: Provide correctness tests and performance benchmarks for operators and compute components, with clear applicability and validation scope.
| Framework or Component | Areas of Focus |
|---|---|
| Megatron | Key compute modules and related operators for large-model training |
| DeepSpeed | Operators and fused computation for training optimization |
| Transformer Engine | Transformer compute modules and operators for mixed-precision and low-precision training |
These projects form the initial scope. Support will be documented by operator, data type, component version, and hardware backend in the compatibility notes.
FlagTrain is developed in collaboration with operator libraries, compilers, and training frameworks in the FlagOS ecosystem, with distinct responsibilities:
- FlagGems: Provides general-purpose operator implementations. FlagTrain prioritizes collaboration with FlagGems on reusable foundational operators to reduce duplicate implementations.
- FlagTrain: Focuses on operators, fused computation, and compute components for training, organizing their implementations, tests, and performance optimization.
- Training frameworks and upstream components: Provide use cases and references for interfaces and computational semantics. FlagTrain develops and validates Triton implementations for their key computation paths.
Specific invocation and integration methods will be described in the relevant module documentation and usage examples.
FlagTrain/
├── src/flag_train/
│ ├── megatron/ # Megatron-related computation implementations
│ ├── deepspeed/ # DeepSpeed-related computation implementations
│ └── transformer_engine/ # Transformer Engine-related computation implementations
├── tests/ # Correctness and interface compatibility tests by component
├── benchmark/ # Performance benchmarks by component
├── examples/ # Usage and integration examples by component
├── docs/ # Installation, development, and compatibility notes
├── tools/ # Development and maintenance tools
├── CONTRIBUTING.md
├── SECURITY.md
├── LICENSE
└── pyproject.toml
The three component directories are intended to house their respective Triton computation implementations and interface wrappers. Within each component, files will be organized by operator or compute module; forward computation, backward computation, and invocation wrappers can live together. Shared functionality will be extracted when a clear need for reuse emerges.
Each operator or compute component implementation should include a corresponding reference implementation, test cases, and a method for evaluating performance.
- Cover representative input shapes, data types, and edge cases.
- Validate results against the reference implementation and specify numerical error tolerances.
- Validate gradient computation for implementations involving backpropagation.
- Record the component versions, software environment, and hardware backend used in testing.
- Compare implementations under the same hardware, input, and precision conditions.
- Describe the reference implementation, warm-up procedure, and timing method.
- Report execution time, throughput, or device memory usage as appropriate for the computation.
- Distinguish individual operator performance from end-to-end training gains.
See Getting Started to install the project skeleton locally and the development guide for development conventions. Operator usage examples and test commands will be provided with the first implementations.
Contributions to training operators, component adaptation, performance optimization, tests, and documentation are welcome.
- See the contribution guide for the contribution process.
- Report issues and suggest features through GitHub Issues.
- Report security issues privately by following the security reporting instructions.
This project is licensed under the Apache License 2.0. Any included third-party code must retain its original license and copyright notices.