feat(subsystembenchmarks): add PyTorch Lightning checkpointing save benchmark - #994
feat(subsystembenchmarks): add PyTorch Lightning checkpointing save benchmark#994Yonghui-Lee wants to merge 3 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
|
E2E failed because of limited resources in the region. |
| - {axis: "strategy", strategy: "fsdp_sharded"} | ||
| - {axis: "strategy", strategy: "fsdp_full"} | ||
| - {axis: "strategy", strategy: "model_parallel_full", tensor_parallel_size: 4, data_parallel_size: 2, world_size: 8} | ||
| - {axis: "strategy", strategy: "model_parallel_sharded", tensor_parallel_size: 4, data_parallel_size: 2, world_size: 8} |
There was a problem hiding this comment.
Should tensor_parallel_size and data_parallel_size be added to cloudbuild/subsystembenchmarks/subsystembenchmarks_schema.json?
| super().__init__() | ||
| self.params = params | ||
|
|
||
| model_id = os.getenv("MODEL_ID", params.model_id) |
There was a problem hiding this comment.
Why not just params.model_id?
| rounds: 1 | ||
| baseline: | ||
| model_id: "meta-llama/Llama-3.1-8B" | ||
| strategy: "single" |
There was a problem hiding this comment.
Should the world_size be 1 for baseline?
| strategy: "single" | ||
|
|
||
| scenarios: | ||
| - scenario: "checkpoint_write" |
There was a problem hiding this comment.
Can we use name here for benchmark filtering?
| "measurement_window_end_unix_seconds": int(window_end), | ||
| } | ||
| ) | ||
| if result.extra_columns: |
There was a problem hiding this comment.
Would you mind extracting a func like publish_environment_metadata from gcsfs/tests/perf/subsystembenchmarks/dataloading/read_case.py and apply here too? including distributed_backend, compute_accelerator_type, machine_type, benchmark_source_commit_sha, requirements_override, requirements_resolved, config_sweep_axes_requested, etc.
| "checkpoint_physical_size_bytes": physical_size_bytes, | ||
| "checkpoint_strategy": params.strategy, | ||
| "measurement_window_start_unix_seconds": int(window_start), | ||
| "measurement_window_end_unix_seconds": int(window_end), |
There was a problem hiding this comment.
How about world_size, tensor_parallel_size, data_parallel_size?
|
|
||
| model_id: str | ||
| strategy: str # single, ddp, fsdp, model_parallel_* | ||
| world_size: int = 2 |
There was a problem hiding this comment.
Should world_size be default 1?
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).