diff --git a/CMakeLists.txt b/CMakeLists.txt index 035492ed..88911e84 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -6,7 +6,7 @@ option(USE_OMP "Use OpenMP as backend for Eigen" ON) option(USE_NCCL "Build project for distributed running on CUDA using NCCL" ON) option(BUILD_TEST "Build InfiniTrain tests" OFF) -project(infini_train VERSION 0.6.0 LANGUAGES CXX) +project(infini_train VERSION 0.7.0 LANGUAGES CXX) set(CMAKE_CXX_STANDARD 20) set(CMAKE_CXX_STANDARD_REQUIRED ON) diff --git a/README.md b/README.md index 8f856700..25114aed 100644 --- a/README.md +++ b/README.md @@ -1,105 +1,110 @@ -# InfiniTrain - -[![CI](https://github.com/InfiniTensor/InfiniTrain/actions/workflows/format-check.yaml/badge.svg)]( -https://github.com/InfiniTensor/InfiniTrain/actions -) -[![Issues](https://img.shields.io/github/issues/InfiniTensor/InfiniTrain)]( -https://github.com/InfiniTensor/InfiniTrain/issues -) -[![PR](https://img.shields.io/github/issues-pr/InfiniTensor/InfiniTrain)]( -https://github.com/InfiniTensor/InfiniTrain/pulls -) -[![License](https://img.shields.io/github/license/InfiniTensor/InfiniTrain)]( -https://github.com/InfiniTensor/InfiniTrain/blob/master/LICENSE -) - -A from-scratch C++ training framework for large-scale models with multi-dimensional distributed parallelism. - -## ๐Ÿš€ Quick Start - -### System Requirements - -#### Hardware Requirements - -- **Recommended**: NVIDIA Ampere-class GPUs (A100/A800) or newer - -#### Software Requirements - -- **CUDA / NCCL**: Latest stable versions -- **gcc / g++**: Version **13+** -- **CMake**: Version **3.13+** - -### Installation - -```bash -mkdir build -cd build -cmake .. -DUSE_CUDA=ON -DUSE_NCCL=ON -make -j -``` - -Build Options: - -- `USE_CUDA=ON` - - Enable CUDA backend support. - -- `USE_NCCL=ON` - - Enable NCCL-based distributed communication. - -> Both options are optional and can be disabled for CPU-only builds. - +# InfiniTrain + +[![CI](https://github.com/InfiniTensor/InfiniTrain/actions/workflows/format-check.yaml/badge.svg)]( +https://github.com/InfiniTensor/InfiniTrain/actions +) +[![Issues](https://img.shields.io/github/issues/InfiniTensor/InfiniTrain)]( +https://github.com/InfiniTensor/InfiniTrain/issues +) +[![PR](https://img.shields.io/github/issues-pr/InfiniTensor/InfiniTrain)]( +https://github.com/InfiniTensor/InfiniTrain/pulls +) +[![License](https://img.shields.io/github/license/InfiniTensor/InfiniTrain)]( +https://github.com/InfiniTensor/InfiniTrain/blob/master/LICENSE +) + +A from-scratch C++ training framework for large-scale models with multi-dimensional distributed parallelism. + +## ๐Ÿš€ Quick Start + +### System Requirements + +#### Hardware Requirements + +- **Recommended**: NVIDIA Ampere-class GPUs (A100/A800) or newer + +#### Software Requirements + +- **CUDA / NCCL**: Latest stable versions +- **gcc / g++**: Version **13+** +- **CMake**: Version **3.13+** + +### Installation + +```bash +mkdir build +cd build +cmake .. -DUSE_CUDA=ON -DUSE_NCCL=ON +make -j +``` + +Build Options: + +- `USE_CUDA=ON` + + Enable CUDA backend support. + +- `USE_NCCL=ON` + + Enable NCCL-based distributed communication. + +> Both options are optional and can be disabled for CPU-only builds. + ## โœจ InfiniTrain Overview - -### โœ” Support Matrix - -| Category | Feature | Description | Status | -| ------------------------- | ------------------------------- | ---------------------------------------------------- | -------------- | -| Model Support | GPT-2 | Decoder-only Transformer language model | โœ” Supported | -| | LLaMA 3 | Modern LLaMA-family Transformer architecture | โœ” Supported | -| | Qwen3-8B | Qwen3 8B language model with QK norm and GQA | โœ” Supported | -| | DeepSeek-V3 | Large-scale MoE-based language model | ๐Ÿ—“ Planned | -| Precision | Multiple Data Type | FP32, BF16 | โœ” Supported | -| | Mixed Precision | Autocast-based BF16 compute with FP32 accumulation | โœ” Supported | -| Distributed Training | Data Parallel (DP) | Parameter-server-style data parallelism | โœ” Supported | -| | Distributed Data Parallel (DDP) | Collective-based data parallelism | โœ” Supported | -| | Tensor Parallelism (TP) | Intra-layer tensor sharding | โœ” Supported | -| | Sequence Parallelism (SP) | Sequence dimension sharding | โœ” Supported | -| | Pipeline Parallelism (PP) | GPipe, 1F1B scheduling, Virtual Pipeline (vPP) | โœ” Supported | -| | Hybrid Parallelism | Arbitrary combination of DDP + TP + SP + PP | โœ” Supported | -| Core Components | Multi-backend | CPU and CUDA execution backends | โœ” Supported | -| | Multi-node Distributed Training | Distributed execution across multiple nodes | โœ” Supported | -| | Transformer Abstraction | Generic Transformer structure abstraction | โœ” Supported | -| | Backend Registries | Device / CCL / dtype abstraction and registration | โœ” Supported | -| | Kernel Dispatcher | Kernel registration and dynamic dispatch mechanism | โœ” Supported | -| | Autograd | Automatic differentiation engine | โœ” Supported | -| | Autocast | Automatic mixed precision runtime | โœ” Supported | -| | Checkpointing | Training checkpoint save and restore | ๐Ÿ—“ Planned | -| Fine-tuning | LoRA | Memory-efficient fine-tuning with merge / unmerge | โœ” Supported | -| Memory Optimizations | ZeRO Stage-1 | Sharded optimizer states for DDP | โœ” Supported | -| | ZeRO Stage-2 | Sharded gradients across DDP ranks | โœ” Supported | -| | Activation Recomputation | Recompute activations to reduce memory usage | ๐Ÿ—“ Planned | -| Performance Optimizations | Computeโ€“Comm Overlap | Explicit scheduling to hide communication latency | โœ” Supported | -| | DDP Gradient Bucketing | Deferred and bucketed gradient synchronization | โœ” Supported | -| Execution Mode | Training Mode | Full forwardโ€“backward training with autograd | โœ” Supported | -| | `no_grad` Inference | Forward-only execution without gradient tracking | โœ” Supported | -| Debugging & Tooling | Built-in Profiler | Kernel-level performance profiling | โœ” Supported | -| | Precision Alignment Checker | Function / Module precision checks and E2E loss diff | โœ” Supported | -| | CTest + GTest Infrastructure | Automated unit tests with CTest integration | โœ” Supported | -| | Automated Benchmarking | One-click execution, log analysis and Feishu export | โœ” Supported | - -## ๐Ÿ‹๏ธ Training - -Each model in the `example/` directory is compiled into an independent executable. -For example, the `llama3` example produces a binary named `llama3`. - -To view available runtime options: - + +### โœ” Support Matrix + +| Category | Feature | Description | Status | +| ------------------------- | ------------------------------- | ---------------------------------------------------- | ----------- | +| Model Support | GPT-2 | Decoder-only Transformer language model | โœ” Supported | +| | LLaMA 3 | Modern LLaMA-family Transformer architecture | โœ” Supported | +| | Qwen3-8B | Qwen3 8B language model with QK norm and GQA | โœ” Supported | +| | TinyMixtral (MoE) | Single-GPU and distributed EP training validation | โœ” Supported | +| | DeepSeek-V3 | Large-scale MoE-based language model | ๐Ÿ—“ Planned | +| Precision | Multiple Data Type | FP32, BF16 | โœ” Supported | +| | Mixed Precision | Autocast-based BF16 compute with FP32 accumulation | โœ” Supported | +| Distributed Training | Data Parallel (DP) | Parameter-server-style data parallelism | โœ” Supported | +| | Distributed Data Parallel (DDP) | Collective-based data parallelism | โœ” Supported | +| | Tensor Parallelism (TP) | Intra-layer tensor sharding | โœ” Supported | +| | Sequence Parallelism (SP) | Sequence dimension sharding | โœ” Supported | +| | Pipeline Parallelism (PP) | GPipe, 1F1B scheduling, Virtual Pipeline (vPP) | โœ” Supported | +| | Expert Parallelism (EP) | Expert sharding and AllGather-based token dispatch | โœ” Supported | +| | Hybrid Parallelism | Arbitrary combination of DDP + TP + SP + PP | โœ” Supported | +| | Multi-process Launch | `infini_run` and thread-based launch modes | โœ” Supported | +| Core Components | Multi-backend | CPU and CUDA execution backends | โœ” Supported | +| | Multi-node Distributed Training | Distributed execution across multiple nodes | โœ” Supported | +| | Transformer Abstraction | Generic Transformer structure abstraction | โœ” Supported | +| | Backend Registries | Device / CCL / dtype abstraction and registration | โœ” Supported | +| | PrivateUse1 Backend Extension | External runtime / CCL / kernel registration | โœ” Supported | +| | Kernel Dispatcher | Kernel registration and dynamic dispatch mechanism | โœ” Supported | +| | Autograd | Automatic differentiation engine | โœ” Supported | +| | Autocast | Automatic mixed precision runtime | โœ” Supported | +| | Checkpointing | Custom format; model/optimizer/trainer state resume | โœ” Supported | +| | Learning Rate Schedulers | PyTorch-style API; Megatron-style warmup / decay | โœ” Supported | +| Fine-tuning | LoRA | Memory-efficient fine-tuning with merge / unmerge | โœ” Supported | +| Memory Optimizations | ZeRO Stage-1 | Sharded optimizer states for DDP | โœ” Supported | +| | ZeRO Stage-2 | Sharded gradients across DDP ranks | โœ” Supported | +| | Activation Recomputation | Reentrant/non-reentrant; none/uniform/block | โœ” Supported | +| Performance Optimizations | Computeโ€“Comm Overlap | Explicit scheduling to hide communication latency | โœ” Supported | +| | DDP Gradient Bucketing | Deferred and bucketed gradient synchronization | โœ” Supported | +| Execution Mode | Training Mode | Full forwardโ€“backward training with autograd | โœ” Supported | +| | `no_grad` Inference | Forward-only execution without gradient tracking | โœ” Supported | +| Debugging & Tooling | Built-in Profiler | Kernel-level performance profiling | โœ” Supported | +| | Precision Alignment Checker | Function / Module precision checks and E2E loss diff | โœ” Supported | +| | CTest + GTest Infrastructure | Automated unit tests with CTest integration | โœ” Supported | +| | Automated Benchmarking | One-click execution, log analysis and Feishu export | โœ” Supported | + +## ๐Ÿ‹๏ธ Training + +Each model in the `example/` directory is compiled into an independent executable. +For example, the `llama3` example produces a binary named `llama3`. + +To view available runtime options: + ```bash ./build/llama3 --help -``` - +``` + ### Getting Started #### Prepare Datasets and Weights @@ -247,9 +252,9 @@ training world size is: ```text world_size = nnodes ร— nproc_per_node ร— nthread_per_process ``` - -### Parallelism Strategies - + +### Parallelism Strategies + #### Distributed Data Parallelism (DDP) For a direct launch with TP and PP disabled, the following starts eight @@ -265,70 +270,114 @@ size after accounting for tensor and pipeline parallelism: ```text data_parallel_size = world_size / (tensor_parallel ร— pipeline_parallel) ``` - -#### Tensor Parallelism (TP) - -```bash ---tensor_parallel 4 # 4-way tensor parallelism ---sequence_parallel # Enable sequence parallelism (requires TP > 1) -``` - -#### Pipeline Parallelism (PP) - -```bash ---pipeline_parallel 8 # 8 pipeline stages ---virtual_pipeline_parallel 4 # Virtual pipeline for better load balancing -``` - -#### Combining Parallelism Strategies - -Multiple parallelism strategies (DDP, TP, SP, PP) can be freely combined to scale training across devices and nodes. - -## ๐Ÿ—บ Roadmap - -- **2025/03/10** โ€” InfiniTrain **v0.1.0** - - Initial framework prototype with MNIST CPU training. - -- **2025/04/30** โ€” InfiniTrain **v0.3.0** - - Added Autograd support and GPT-2 training on CPU/CUDA. - -- **2025/07/09** โ€” InfiniTrain **v0.4.0** - - Introduced kernel registration, LLaMA training on CPU/CUDA, BF16 precision, and Data Parallelism. - -- **2025/12/31** โ€” InfiniTrain **v0.5.0** - - Added Autocast, multi-dimensional distributed parallelism - (DDP, TP, SP, PP with GPipe / 1F1B / vPP), - multi-node training, `no_grad` mode, - and communicationโ€“computation overlap with bucketed gradient synchronization. - -- **2026/06/08** โ€” InfiniTrain **v0.6.0** - - Added loss alignment tooling for Function / Module level precision checks - and end-to-end loss comparison, with a unified hook mechanism. - - Added memory optimizations for DDP training and Autograd execution. - ZeRO Stage-1 shards optimizer states across DDP ranks, while ZeRO Stage-2 - further shards gradients. Autograd Tensor release timing was also optimized - to reduce peak memory usage. - - Introduced LoRA fine-tuning with `merge` / `unmerge` support for efficient - training and inference-time weight merging. - - Refactored core backend abstractions around device, communication, and - low-precision dtype registration. The framework layer now uses - `DeviceGuard`, `CclGroupGuard`, and backend-registered FP16 / BF16 native - types to avoid hardware-specialized framework code. - - Introduced a generic Transformer structure abstraction backed by - `TransformerConfig`, providing a common foundation for GPT-2 and LLaMA 3 - style model construction. - - Improved BF16 training performance through autocast and elementwise kernel - optimizations. - - Integrated a CTest + GTest based testing infrastructure to strengthen the + +#### Tensor Parallelism (TP) + +```bash +--tensor_parallel 4 # 4-way tensor parallelism +--sequence_parallel # Enable sequence parallelism (requires TP > 1) +``` + +#### Pipeline Parallelism (PP) + +```bash +--pipeline_parallel 8 # 8 pipeline stages +--virtual_pipeline_parallel 4 # Virtual pipeline for better load balancing +``` + +#### Expert Parallelism (EP) + +MoE training supports expert sharding across ranks with an +`AllGather`-based token dispatcher. The distributed TinyMixtral example +validates cross-rank token dispatch, local expert computation, and gradient +synchronization. Dense and MoE communication topology views provide the +corresponding process groups. + +#### Combining Parallelism Strategies + +Multiple parallelism strategies (DDP, TP, SP, PP) can be freely combined to scale training across devices and nodes. +MoE models additionally support EP; supported combinations depend on the model and configuration. + +## ๐Ÿ—บ Roadmap + +- **2025/03/10** โ€” InfiniTrain **v0.1.0** + + Initial framework prototype with MNIST CPU training. + +- **2025/04/30** โ€” InfiniTrain **v0.3.0** + + Added Autograd support and GPT-2 training on CPU/CUDA. + +- **2025/07/09** โ€” InfiniTrain **v0.4.0** + + Introduced kernel registration, LLaMA training on CPU/CUDA, BF16 precision, and Data Parallelism. + +- **2025/12/31** โ€” InfiniTrain **v0.5.0** + + Added Autocast, multi-dimensional distributed parallelism + (DDP, TP, SP, PP with GPipe / 1F1B / vPP), + multi-node training, `no_grad` mode, + and communicationโ€“computation overlap with bucketed gradient synchronization. + +- **2026/06/08** โ€” InfiniTrain **v0.6.0** + + Added loss alignment tooling for Function / Module level precision checks + and end-to-end loss comparison, with a unified hook mechanism. + + Added memory optimizations for DDP training and Autograd execution. + ZeRO Stage-1 shards optimizer states across DDP ranks, while ZeRO Stage-2 + further shards gradients. Autograd Tensor release timing was also optimized + to reduce peak memory usage. + + Introduced LoRA fine-tuning with `merge` / `unmerge` support for efficient + training and inference-time weight merging. + + Refactored core backend abstractions around device, communication, and + low-precision dtype registration. The framework layer now uses + `DeviceGuard`, `CclGroupGuard`, and backend-registered FP16 / BF16 native + types to avoid hardware-specialized framework code. + + Introduced a generic Transformer structure abstraction backed by + `TransformerConfig`, providing a common foundation for GPT-2 and LLaMA 3 + style model construction. + + Improved BF16 training performance through autocast and elementwise kernel + optimizations. + + Integrated a CTest + GTest based testing infrastructure to strengthen the framework's automated test workflow. + +- InfiniTrain **v0.7.0** + + Added training checkpointing with a custom checkpoint format, supporting + model weights, optimizer state, and trainer state save/load, with validated + resume-from-checkpoint workflows. Additional formats and recovery scenarios + are planned. ([#146](https://github.com/InfiniTensor/InfiniTrain/pull/146)) + + Introduced PyTorch-style learning rate schedulers and Megatron-inspired + training helpers for common warmup and decay schedules. + ([#151](https://github.com/InfiniTensor/InfiniTrain/pull/151)) + + Added MoE training support with `TopKRouter`, `Experts`, + `MoETokenDispatcher`, and `MoELayer`, and extended it with Expert + Parallelism (EP), including Dense / MoE topology views, expert sharding, + AllGather-based token dispatch, and gradient synchronization. Validated + end-to-end TinyMixtral training on both single-GPU and distributed + configurations. + ([#153](https://github.com/InfiniTensor/InfiniTrain/pull/153), + [#246](https://github.com/InfiniTensor/InfiniTrain/pull/246), + [#248](https://github.com/InfiniTensor/InfiniTrain/pull/248), + [#191](https://github.com/InfiniTensor/InfiniTrain/pull/191)) + + Added torchrun-style multi-process distributed launching via `infini_run`, + while retaining the existing multi-thread launch mode. + ([#184](https://github.com/InfiniTensor/InfiniTrain/pull/184)) + + Introduced a `PrivateUse1` backend extension point for external runtime, + CCL, and kernel registration without vendor-specific framework changes. + ([#206](https://github.com/InfiniTensor/InfiniTrain/pull/206)) + + Added activation checkpointing aligned with `torch.utils.checkpoint`, + supporting reentrant and non-reentrant (default) modes and `none`, `uniform`, + and `block` Transformer recomputation strategies. + ([#154](https://github.com/InfiniTensor/InfiniTrain/pull/154)) \ No newline at end of file