Lightning-AI / Lightning-AI/lightning-thunder
Remove "block" and "layer" bucketing modes from FSDP implementation
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
The experimental FSDP bucketing modes, specifically "block" and "layer", have not shown performance improvements in recent benchmarks. In fact, the `--bucketing_mode none` option currently yields the best performance for models like Llama 3 8B on DGX-H100.
We've known this for a while, but haven't taken any action.
*Benchmark Results (Llama-3-8B, DGX-H100, FSDP zero2):*
* `--bucketing_mode none`:
```
Average iter time: 801.54 ms
Tokens/s: 81753.72
Tokens/s/GPU: 10219.22
TFLOP/s: 4734.59
```
* `--bucketing_mode block`:
```
Average iter time: 801.98 ms
Tokens/s: 81711.28
Tokens/s/GPU: 10213.91
TFLOP/s: 4732.13
```
* `--bucketing_mode layer`:
```
Average iter time: 809.66 ms
Tokens/s: 80900.08
Tokens/s/GPU: 10112.51
TFLOP/s: 4685.15
```
As seen from the results, none provides a slightly better iteration time and throughput compared to block, and significantly better performance than layer.
**Reasons for removal**:
1. Performance: The "block" and "layer" bucketing strategies do not offer a performance advantage over no bucketing.
2. Maintenance Cost: These options are becoming expensive to maintain, especially considering the significant changes planned for Thunder's autograd implementation outlined in #2102. We should focus our efforts on the most performant and maintainable options.
**Proposed action**
Remove the `block` and `layer` options for `--bucketing_mode` in Thunder's FSDP implementation. This will simplify the codebase and reduce the maintenance burden.
**Benchmark Command Used:**
```shell
for bm in none block layer ; do torchrun --nproc_per_node=8 thunder/benchmarks/benchmark_litgpt.py --model Llama-3-8B --compile thunder --distributed_mode fsdp --bucketing_mode $bm ; done
```
**Affected code:**
* Relevant FSDP implementation files where the bucketing logic resides (`torch_autograd.py` and `fsdp.py` files)
* Argument parsing for `--bucketing_mode` in the `benchmark_litgpt.py` file.
* Any related documentation or tests.
**User impact:**
Users will no longer be able to select "block" or "layer" bucketing. The default and only effective option for FSDP parameter communication bucketing will be "none" (or its equivalent if renamed), so this option will be removed entirely. This change is expected to simplify the user experience by removing underperforming experimental features.
cc @crcrpar @ali-alshaar7 @lantiga @t-vi
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with the bucketing logic in torch_autograd.py and fsdp.py, then trace --bucketing_mode parsing in thunder/benchmarks/benchmark_litgpt.py. Run the listed benchmark command and inspect related tests and documentation. Done means the block and layer modes are removed consistently, leaving none as the supported option and updating affected coverage and guidance.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100