Lightning-AI / Lightning-AI/lightning-thunder

Remove "block" and "layer" bucketing modes from FSDP implementation

Open
#2,144 2 comments 1 reaction 0 assignees View on GitHub
distributed enhancement
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.