feat(subsystembenchmarks): add PyTorch Lightning checkpointing save benchmark - #994
feat(subsystembenchmarks): add PyTorch Lightning checkpointing save benchmark#994Yonghui-Lee wants to merge 2 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new checkpointing subsystem benchmark group (checkpointing/pytorch_lightning) to measure checkpoint write performance using PyTorch Lightning and various training strategies (DDP, FSDP, Model Parallel) on CPU-simulated environments. It includes the necessary driver, configuration, and test files, along with updates to Cloud Build scripts and schemas to support checkpointing metrics. The review feedback highlights two key improvement opportunities: preventing a potential ZeroDivisionError in checkpoint_case.py when the durations list is empty, and initializing the dummy model's linear layer directly with dtype=torch.bfloat16 in driver.py to avoid memory overhead and potential Out-Of-Memory (OOM) errors during initialization.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #994 +/- ##
=======================================
Coverage 89.68% 89.69%
=======================================
Files 16 16
Lines 3579 3581 +2
=======================================
+ Hits 3210 3212 +2
Misses 369 369 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
…layer in benchmarks
| _MACHINE_TYPE: "c4-standard-192" # CPU now; point at a GPU type later (auto-detect handles backend) | ||
| _REQUIREMENTS: "" | ||
| _SWEEP_AXES: "" | ||
| _STRATEGY: "" |
There was a problem hiding this comment.
Could we avoid adding a checkpoint-specific _STRATEGY substitution?
If individual-case execution is operationally necessary, a generic case filter shared by all subsystem benchmarks would be a cleaner API.
| {"name": "checkpoint_physical_size_bytes", "type": "INTEGER", "description": "Total physical size in bytes of the saved checkpoint files."}, | ||
| {"name": "checkpoint_strategy", "type": "STRING", "description": "Checkpointing strategy used (single, ddp, fsdp, etc.)."}, | ||
| {"name": "checkpoint_write_throughput_mean_bytes_per_second", "type": "FLOAT", "description": "Mean write throughput in bytes per second for the checkpoint save operation."}, | ||
| {"name": "model_size_mb", "type": "INTEGER", "description": "Simulated model size parameter in Megabytes."} |
There was a problem hiding this comment.
Is this useful to have model_size_mb. where model_size_mb describes the requested size of the synthetic Linear rather than the real model or persisted checkpoint. How about model id instead?
| return torch.randn(self.in_features) | ||
|
|
||
|
|
||
| class DummyModel(L.LightningModule): |
There was a problem hiding this comment.
A single very large Linear exercises Lightning’s save API, but it does not preserve the Llama workload’s tensor count and shapes, layerwise FSDP wrapping, architecture-specific tensor-parallel plan, or optimizer-state layout.
Shall we use the real model, so it can be a bridge between microbenchmarks and macrobenchmarks.
| return durations | ||
|
|
||
|
|
||
| class PLCheckpointDriver: |
There was a problem hiding this comment.
Is params.rounds ignored?
| ) | ||
| logging.info("Physical checkpoint size: %d bytes", physical_size_bytes) | ||
| except Exception as e: | ||
| logging.warning("Could not measure physical checkpoint size: %s", e) |
There was a problem hiding this comment.
Why not fail loudly here?
| model = model.to(torch.bfloat16) | ||
|
|
||
| # Trigger save checkpoint | ||
| trainer.fit(model, train_dataloaders=dataloader) |
There was a problem hiding this comment.
Why perform a untimed automatic checkpoint before the measured save
2 participants
Summary
This PR implements the PyTorch Lightning checkpointing save subsystem benchmark to measure checkpoint write performance under various strategies in CPU-simulated environments.
Detailed Changes
PLCheckpointDriverwhich configures a dummy PyTorch Model and Dataset.torch.multiprocessing.spawn, synchronizes ranks usingdist.barrier()to accurately time the checkpoint write, and aggregates durations.single: Single device checkpointing.ddp: Distributed Data Parallel.fsdp_sharded: Fully Sharded Data Parallel with sharded checkpoint state dicts.fsdp_full: FSDP with consolidated full checkpoint state dicts.model_parallel_full/model_parallel_sharded: Combined Tensor Parallel (TP=4) and Data Parallel (DP=2) strategies (up to 8 ranks).