Spaces:
Running
Running
File size: 2,765 Bytes
ec19476 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
# Pipeline Parallelism Scheduler and Visualizer
This tool simulates and visualizes pipeline parallelism scheduling strategies, focusing on the 1F1B (One-Forward-One-Backward) scheduling algorithm commonly used in distributed deep learning.
## Usage
### Example Output
```bash
python pipeline.py --num-stages 4 --num-batches 8
```
### Command Line Interface
| Option | Short | Description |
|--------|-------|-------------|
| `--config` | `-c` | Path to config file (JSON or YAML) |
| `--num-stages` | `-s` | Number of pipeline stages (devices) |
| `--num-batches` | `-b` | Number of micro-batches |
| `--forward-times` | `-f` | Time for forward pass at each stage (space-separated list) |
| `--backward-times` | `-bw` | Time for backward pass at each stage (space-separated list) |
| `--output` | `-o` | Output file path for visualization |
| `--no-visualization` | | Skip visualization generation |
| `--p2p-time`| | P2P communication time of PP |
### Using Configuration Files
You can use either JSON or YAML configuration files:
Example JSON configuration (sample_config.json):
```json
{
"num_stages": 6,
"num_batches": 12,
"forward_times": [0.8, 1.0, 1.2, 1.0, 0.9, 1.1],
"backward_times": [1.6, 2.0, 2.4, 2.0, 1.8, 2.2],
"output_file": "pipeline_1f1b_custom.png"
}
```
Example YAML configuration (sample_config.yaml):
```yaml
# Pipeline Parallelism Configuration
num_stages: 5
num_batches: 8
forward_times:
- 0.9
- 1.1
- 1.0
- 0.8
- 1.2
backward_times:
- 1.8
- 2.2
- 2.0
- 1.6
- 2.4
output_file: "pipeline_1f1b_yaml.png"
```
## About Pipeline Parallelism
Pipeline parallelism is a distributed deep learning training strategy that splits model layers across multiple devices. Each device processes a different stage of the neural network, creating a pipeline where multiple micro-batches can be processed simultaneously.
The 1F1B (One-Forward-One-Backward) scheduling algorithm is an efficient strategy for pipeline parallelism that balances throughput with memory usage. It follows these phases:
1. **Warmup Phase**: Forward passes for the first several micro-batches
2. **Steady State**: Each device alternates between forward and backward passes
3. **Cooldown Phase**: Backward passes to complete the computation for remaining micro-batches
The "bubble rate" metric measures the inefficiency in the pipeline, representing the percentage of time devices spend idle waiting for dependencies.
## References
- PipeDream: Generalized Pipeline Parallelism for DNN Training (SOSP'19)
- GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism (NeurIPS'19)
- Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism |