77 lines
2.0 KiB
Markdown
77 lines
2.0 KiB
Markdown
# distributed/
|
|
|
|
C implementation of memory-efficient distributed training.
|
|
|
|
## Python counterpart
|
|
|
|
**DeepSpeed** - https://github.com/microsoft/DeepSpeed
|
|
|
|
## Why DeepSpeed?
|
|
|
|
DeepSpeed enables training models that don't fit in GPU memory through:
|
|
|
|
- ZeRO optimizer stages (partition optimizer states, gradients, parameters)
|
|
- Gradient checkpointing
|
|
- Mixed precision training
|
|
- Pipeline parallelism
|
|
|
|
Alternatives (Megatron-LM, FairScale, FSDP) solve similar problems differently. DeepSpeed is most widely used for LoRA training.
|
|
|
|
## Scope
|
|
|
|
This module implements:
|
|
|
|
1. **Memory-efficient optimizer** - Partition Adam states across GPUs
|
|
2. **Gradient checkpointing** - Recompute activations to save memory
|
|
3. **Mixed precision** - FP16/BF16 training with FP32 master weights
|
|
4. **Communication primitives** - AllReduce, AllGather for gradients
|
|
|
|
## Getting Started
|
|
|
|
### Dependencies
|
|
|
|
- tensor/ module (foundation)
|
|
- MPI or NCCL for multi-GPU
|
|
- Optional: CUDA-aware MPI
|
|
|
|
### Key Files to Study
|
|
|
|
From DeepSpeed:
|
|
- `deepspeed/runtime/zero/` - ZeRO implementation
|
|
- `deepspeed/runtime/activation_checkpointing/` - Gradient checkpointing
|
|
|
|
### Core API
|
|
|
|
```c
|
|
// Initialize distributed
|
|
fd_init(int* argc, char*** argv);
|
|
int fd_rank(void);
|
|
int fd_world_size(void);
|
|
|
|
// ZeRO optimizer wrapper
|
|
fd_zero_optimizer* fd_zero_wrap(tm_adamw* opt, int stage);
|
|
void fd_zero_step(fd_zero_optimizer* opt);
|
|
|
|
// Gradient checkpointing
|
|
tm_tensor* fd_checkpoint(tm_tensor* (*fn)(tm_tensor*), tm_tensor* input);
|
|
|
|
// Communication
|
|
void fd_all_reduce(tm_tensor* t);
|
|
void fd_broadcast(tm_tensor* t, int root);
|
|
```
|
|
|
|
## Implementation Order
|
|
|
|
1. Single-GPU gradient checkpointing
|
|
2. Mixed precision training
|
|
3. Multi-GPU communication (MPI)
|
|
4. ZeRO Stage 1 (optimizer state partitioning)
|
|
5. ZeRO Stage 2 (gradient partitioning)
|
|
6. ZeRO Stage 3 (parameter partitioning)
|
|
|
|
## References
|
|
|
|
- [ZeRO Paper](https://arxiv.org/abs/1910.02054)
|
|
- [DeepSpeed Documentation](https://www.deepspeed.ai/)
|
|
- [PyTorch FSDP](https://pytorch.org/docs/stable/fsdp.html)
|