Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor benchmark tables #136

Merged
merged 6 commits into from
Oct 3, 2024
Merged

Refactor benchmark tables #136

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
f36c6c9
Add single node async benchmark.
awonak Oct 2, 2024
d0867aa
Update README to reflect recent benchmark arg changes.
awonak Oct 2, 2024
1f1e0d5
Use Markdown table formatting for easier maintainability.
awonak Oct 2, 2024
620e152
Add checkpoint load markdown table
awonak Oct 2, 2024
7ef1d3e
Fix typo in header
awonak Oct 2, 2024
d92be70
Merge branch 'main' into awonak-markdown-edit
awonak Oct 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
244 changes: 24 additions & 220 deletions dataflux_pytorch/benchmark/checkpointing/singlenode/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,9 @@ Then set the command line variables.

### Dataflux Lightning Checkpoint

`--no-dataflux-ckpt`: If you are not benchmarking Dataflux Lightning Checkpoint, this will disable dataflux checkpointing and use the default lightning checkpointing instead.
`--checkpoint=no-dataflux`: If you are not benchmarking Dataflux Lightning Checkpoint, this will disable dataflux checkpointing and use the default lightning checkpointing instead.

`--checkpoint=asynccheckpointio`: This flag will enable asynchronous calls to `save_checkpoint` which will optimize CPU/GPU utilization by making save calls non-blocking.

`--disable-multipart`: This flag will disable multipart upload performance improvements. In most cases this will dramatically reduce the upload speed of checkpoint saves and is not recommended.

Expand Down Expand Up @@ -77,225 +79,27 @@ Dataflux's implementation of CheckpointIO for PyTorch Lightning is undergoing ac

### Checkpoint Save

<table>
<tr>
<td style="background-color: #d9d2e9"><strong>Checkpoint Type</strong>
awonak marked this conversation as resolved.
Show resolved Hide resolved
</td>
<td style="background-color: #d9d2e9"><strong>Layers</strong>
</td>
<td style="background-color: #d9d2e9"><strong>Checkpoint File Size (MB)</strong>
</td>
<td style="background-color: #d9d2e9"><strong>Avg Checkpoint Save Time</strong>
</td>
<td style="background-color: #d9d2e9"><strong>Write Throughput (MB/s)</strong>
</td>
</tr>
<tr>
<td style="background-color: #d9d9d9">Default
</td>
<td style="background-color: #d9d9d9">10
</td>
<td style="background-color: #d9d9d9">75.6
</td>
<td style="background-color: #d9d9d9">0.81
</td>
<td style="background-color: #d9d9d9">93.33
</td>
</tr>
<tr>
<td style="background-color: #f3f3f3">Dataflux
</td>
<td style="background-color: #f3f3f3">10
</td>
<td style="background-color: #f3f3f3">75.6
</td>
<td style="background-color: #f3f3f3">0.56
</td>
<td style="background-color: #f3f3f3">135.00
</td>
</tr>
<tr>
<td style="background-color: #d9d9d9">Default
</td>
<td style="background-color: #d9d9d9">100
</td>
<td style="background-color: #d9d9d9">298
</td>
<td style="background-color: #d9d9d9">2.87
</td>
<td style="background-color: #d9d9d9">103.98
</td>
</tr>
<td style="background-color: #f3f3f3">Dataflux
</td>
<td style="background-color: #f3f3f3">100
</td>
<td style="background-color: #f3f3f3">298
</td>
<td style="background-color: #f3f3f3">1.03
</td>
<td style="background-color: #f3f3f3">289.32
</td>
</tr>
<tr>
<td style="background-color: #d9d9d9">Default
</td>
<td style="background-color: #d9d9d9">1,000
</td>
<td style="background-color: #d9d9d9">2,500
</td>
<td style="background-color: #d9d9d9">25.61
</td>
<td style="background-color: #d9d9d9">97.61
</td>
</tr>
<tr>
<td style="background-color: #f3f3f3">Dataflux
</td>
<td style="background-color: #f3f3f3">1,000
</td>
<td style="background-color: #f3f3f3">2,500
</td>
<td style="background-color: #f3f3f3">6.25
</td>
<td style="background-color: #f3f3f3">400.00
</td>
</tr>
<tr>
<td style="background-color: #d9d9d9">Default
</td>
<td style="background-color: #d9d9d9">10,000
</td>
<td style="background-color: #d9d9d9">24,200
</td>
<td style="background-color: #d9d9d9">757.10
</td>
<td style="background-color: #d9d9d9">31.96
</td>
</tr>
<tr>
<td style="background-color: #f3f3f3">Dataflux
</td>
<td style="background-color: #f3f3f3">10,000
</td>
<td style="background-color: #f3f3f3">24,200
</td>
<td style="background-color: #f3f3f3">64.50
</td>
<td style="background-color: #f3f3f3">375.19
</td>
</tr>
</table>
| Checkpoint Type | Layers | Checkpoint File Size (MB) | Avg Checkpoint Save Time | Write Throughput (MB/s) |
| --- | --- | --- | --- | --- |
| Default | 10 | 75.6 | 0.81 | 93.33 |
| Dataflux | 10 | 75.6 | 0.56 | 135.00 |
| Default | 100 | 298 | 2.87 | 103.98 |
| Dataflux | 100 | 298 | 1.03 | 289.32 |
| Default | 1,000 | 2,500 | 25.61 | 97.61 |
| Dataflux | 1,000 | 2,500 | 6.25 | 400.00 |
| Default | 10,000 | 24,200 | 757.10 | 31.96 |
| Dataflux | 10,000 | 24,200 | 64.50 | 375.19 |


### Checkpoint Load

<table>
<tr>
<td style="background-color: #d9d2e9"><strong>Checkpoint Type</strong>
</td>
<td style="background-color: #d9d2e9"><strong>Layers</strong>
</td>
<td style="background-color: #d9d2e9"><strong>Checkpoint Size (MB) per step</strong>
</td>
<td style="background-color: #d9d2e9"><strong>Average Checkpoint Restore Time</strong>
</td>
<td style="background-color: #d9d2e9"><strong>Read Throughput (MB/s)</strong>
</td>
</tr>
<tr>
<td style="background-color: #d9d9d9">Default
</td>
<td style="background-color: #d9d9d9">10
</td>
<td style="background-color: #d9d9d9">75.6
</td>
<td style="background-color: #d9d9d9">2.38
</td>
<td style="background-color: #d9d9d9">31.76
</td>
</tr>
<tr>
<td style="background-color: #f3f3f3">Dataflux
</td>
<td style="background-color: #f3f3f3">10
</td>
<td style="background-color: #f3f3f3">75.6
</td>
<td style="background-color: #f3f3f3">0.51
</td>
<td style="background-color: #f3f3f3">148.24
</td>
</tr>
<tr>
<td style="background-color: #d9d9d9">Default
</td>
<td style="background-color: #d9d9d9">100
</td>
<td style="background-color: #d9d9d9">298
</td>
<td style="background-color: #d9d9d9">12.83
</td>
<td style="background-color: #d9d9d9">23.23
</td>
</tr>
<tr>
<td style="background-color: #f3f3f3">Dataflux
</td>
<td style="background-color: #f3f3f3">100
</td>
<td style="background-color: #f3f3f3">298
</td>
<td style="background-color: #f3f3f3">1.69
</td>
<td style="background-color: #f3f3f3">176.33
</td>
</tr>
<tr>
<td style="background-color: #d9d9d9">Default
</td>
<td style="background-color: #d9d9d9">1,000
</td>
<td style="background-color: #d9d9d9">2,500
</td>
<td style="background-color: #d9d9d9">186.57
</td>
<td style="background-color: #d9d9d9">13.40
</td>
</tr>
<tr>
<td style="background-color: #f3f3f3">Dataflux
</td>
<td style="background-color: #f3f3f3">1,000
</td>
<td style="background-color: #f3f3f3">2,500
</td>
<td style="background-color: #f3f3f3">14.77
</td>
<td style="background-color: #f3f3f3">169.26
</td>
</tr>
<tr>
<td style="background-color: #d9d9d9">Default
</td>
<td style="background-color: #d9d9d9">10,000
</td>
<td style="background-color: #d9d9d9">24,200
</td>
<td style="background-color: #d9d9d9">2,093.52
</td>
<td style="background-color: #d9d9d9">11.56
</td>
</tr>
<tr>
<td style="background-color: #f3f3f3">Dataflux
</td>
<td style="background-color: #f3f3f3">10,000
</td>
<td style="background-color: #f3f3f3">24,200
</td>
<td style="background-color: #f3f3f3">113.14
</td>
<td style="background-color: #f3f3f3">213.89
</td>
</tr>
</table>
| Checkpoint Type | Layers | Checkpoint File Size (MB) | Avg Checkpoint Save Time | Read Throughput (MB/s) |
| --- | --- | --- | --- | --- |
| Default | 10 | 75.6 | 2.38 | 31.76 |
| Dataflux | 10 | 75.6 | 0.51 | 148.24 |
| Default | 100 | 298 | 1.69 | 176.33 |
| Dataflux | 100 | 298 | 1.03 | 289.32 |
| Default | 1,000 | 2,500 | 186.57 | 13.40 |
| Dataflux | 1,000 | 2,500 | 14.77 | 169.26 |
| Default | 10,000 | 24,200 | 2,093.52 | 11.56 |
| Dataflux | 10,000 | 24,200 | 113.14 | 213.89 |
40 changes: 31 additions & 9 deletions dataflux_pytorch/benchmark/checkpointing/singlenode/train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,16 @@
from torch import Tensor
from torch.utils.data import DataLoader

from dataflux_pytorch.lightning import DatafluxLightningCheckpoint
from dataflux_pytorch.lightning import (DatafluxLightningAsyncCheckpoint,
DatafluxLightningCheckpoint)


class BenchmarkDatafluxLightningAsyncCheckpoint(
DatafluxLightningAsyncCheckpoint):

def teardown(self, *args, **kwargs):
# Prevent parent teardown from terminating the executor after fit.
pass


class LightningTransformer(LightningModule):
Expand Down Expand Up @@ -65,9 +74,6 @@ def parse_args():
parser.add_argument("--save-only-latest",
action="store_true",
default=False)
parser.add_argument("--no-dataflux-ckpt",
action="store_true",
default=False)
parser.add_argument("--layers", type=int, default=100)
parser.add_argument("--steps", type=int, default=5)
parser.add_argument("--disable-multipart",
Expand All @@ -76,6 +82,10 @@ def parse_args():
parser.add_argument("--clear-kernel-cache",
action="store_true",
default=False)
parser.add_argument(
'--checkpoint',
choices=['checkpointio', 'asynccheckpointio', 'no-dataflux'],
default='checkpointio')
return parser.parse_args()


Expand All @@ -97,7 +107,7 @@ def main():

Run gcsfs over 10 steps:

python3 train.py --project=my-project --ckpt_dir_path=gs://bucket-name/path/to/dir/ --layers=1000 --steps=10 --no-dataflux-ckpt
python3 train.py --project=my-project --ckpt_dir_path=gs://bucket-name/path/to/dir/ --layers=1000 --steps=10 --checkpoint=no-dataflux

"""
args = parse_args()
Expand All @@ -108,10 +118,19 @@ def main():
dataloader = DataLoader(dataset, num_workers=1)
model = LightningTransformer(vocab_size=dataset.vocab_size,
nlayers=args.layers)
ckpt = DatafluxLightningCheckpoint(
project_name=args.project, disable_multipart=args.disable_multipart)
if args.no_dataflux_ckpt:

# Checkpoint strategy selection.
if args.checkpoint == 'checkpointio':
ckpt = DatafluxLightningCheckpoint(project_name=args.project)
elif args.checkpoint == 'asynccheckpointio':
print("NOTE: AsyncCheckpointIO is enabled.")
ckpt = BenchmarkDatafluxLightningAsyncCheckpoint(
project_name=args.project)
elif args.checkpoint == 'no-dataflux':
ckpt = TorchCheckpointIO()
else:
raise ValueError("Invalid choice for --checkpoint")

# Save once per step, and if `save_only_latest`, replace the last checkpoint each time.
# Replacing is implemented by saving the new checkpoint, and then deleting the previous one.
# If `save_only_latest` is False, a new checkpoint is created for each step.
Expand All @@ -132,6 +151,7 @@ def main():
)
trainer.fit(model, dataloader)

# Measure save checkpoint.
start = time.time()
for i in range(args.steps):
trainer.save_checkpoint(
Expand All @@ -145,9 +165,11 @@ def main():
os.system("sync && sudo sysctl -w vm.drop_caches=3")
print("Average time to save one checkpoint: " +
str((end - start) / args.steps) + " seconds")

# Measure load checkpoint.
start = time.time()
for i in range(args.steps):
data = ckpt.load_checkpoint(
_ = ckpt.load_checkpoint(
os.path.join(args.ckpt_dir_path, f'ckpt_{i}.ckpt'))
end = time.time()
print("Average time to load one checkpoint: " +
Expand Down
2 changes: 2 additions & 0 deletions kokoro/bench.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,8 @@ function run_benchmarks(){
python3 -u ./demo/lightning/image-segmentation/train.py --local --benchmark --gcp_project=dataflux-project --gcs_bucket=dataflux-demo-public --images_prefix=image-segmentation-dataset/images --labels_prefix=image-segmentation-dataset/labels --num_nodes=1 --num_devices=5 --epochs=2;
echo Running single node checkpointing benchmark.
python3 -u ./dataflux_pytorch/benchmark/checkpointing/singlenode/train.py --project=dataflux-project --ckpt-dir-path=gs://df-ckpt-presubmit/ --layers=1000 --steps=5
echo Running single node async checkpointing benchmark.
python3 -u ./dataflux_pytorch/benchmark/checkpointing/singlenode/train.py --project=dataflux-project --ckpt-dir-path=gs://df-ckpt-presubmit/ --layers=1000 --steps=5 --checkpoint=asynccheckpointio
}

setup_virtual_envs
Expand Down