report scuba events for detailed sparse static memory info #5029
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Summary: X-link: facebookresearch/FBGEMM#2042 Add granular sparse static memory breakdown metrics for TBE to enable validation of planner estimates against runtime memory usage. This implementation separates static sparse memory (weights, optimizer states, cache) from ephemeral memory (activations, IO buffers, gradients) and provides per-component HBM/UVM categorization. The existing `tbe.total_hbm_usage` aggregates all memory without distinguishing between persistent storage and ephemeral buffers, making it difficult to identify and validate static sparse parameter estimates. ## Changes ### 1. New Scuba Metrics (`tbe_stats_reporters.py`) Added 10 granular memory metrics to `SyncBatchODSStatsReporter`: **HBM metrics:** - `tbe.hbm.sparse_params` - Embedding weights in HBM - `tbe.hbm.optimizer_states` - Momentum states in HBM - `tbe.hbm.cache` - Cache storage in HBM - `tbe.hbm.total_static_sparse` - Total static memory in HBM - `tbe.hbm.ephemeral` - Ephemeral memory in HBM (activations, temp buffers, etc.) **UVM metrics:** (same structure for UVM) ### 2. Memory Categorization Logic (`split_table_batched_embeddings_ops_training.py`) - Added helper methods: - `_get_tensor_memory()` - Get tensor memory size - `_categorize_memory_by_location()` - Categorize tensors into HBM/UVM - Refactored `_report_tbe_mem_usage()` with clean list-based tensor grouping ### 3. Memory Components **Static Sparse:** - Weights: `weights_dev`, `weights_host`, `weights_uvm` - Optimizer: `momentum1_dev/host/uvm`, `momentum2_dev/host/uvm` - Cache: `lxu_cache_weights`, `lxu_cache_state`, `lxu_state`, cache aux data **Ephemeral (calculated):** - `ephemeral = total_mem_usage - static_sparse` - Includes IO buffers, activations, gradients detailed analysis revealed QPS drop when enabling additional logging. head to head comparison of time it takes to do the logging reveals 4x increase in duration (see (https://fburl.com/scuba/tbe_stats_runtime/y85ur4k9) - avg QPS across 4 runs w/o logging: 246k vs. avg QPS across 4 runs w/ added logging: 243k (~1.2% QPS drop) ran the following models w/o added logging: - aps-icvrbase-tbe-dump-test-old-1-d234a33214 - aps-icvrbase-tbe-dump-test-old-2-f5d7f5d97a - aps-icvrbase-tbe-dump-test-old-3-92aa2d14c3 - aps-icvrbase-tbe-dump-test-old-timed-9ac1869846 ran the following models w/ added logging: - aps-icvrbase-tbe-dump-test-new-1-fcb93df6a6 - aps-icvrbase-tbe-dump-test-new-2-3f15ec3a29 - aps-icvrbase-tbe-dump-test-new-3-211d3c3f01 - aps-icvrbase-tbe-dump-test-new-timed-6e6a932849 Differential Revision: D84624978
✅ Deploy Preview for pytorch-fbgemm-docs ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
|
@ashuaibi7 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D84624978. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/2042
Add granular sparse static memory breakdown metrics for TBE to enable validation of planner estimates against runtime memory usage. This implementation separates static sparse memory (weights, optimizer states, cache) from ephemeral memory (activations, IO buffers, gradients) and provides per-component HBM/UVM categorization. The existing
tbe.total_hbm_usageaggregates all memory without distinguishing between persistent storage and ephemeral buffers, making it difficult to identify and validate static sparse parameter estimates.Changes
1. New Scuba Metrics (
tbe_stats_reporters.py)Added 10 granular memory metrics to
SyncBatchODSStatsReporter:HBM metrics:
tbe.hbm.sparse_params- Embedding weights in HBMtbe.hbm.optimizer_states- Momentum states in HBMtbe.hbm.cache- Cache storage in HBMtbe.hbm.total_static_sparse- Total static memory in HBMtbe.hbm.ephemeral- Ephemeral memory in HBM (activations, temp buffers, etc.)UVM metrics: (same structure for UVM)
2. Memory Categorization Logic (
split_table_batched_embeddings_ops_training.py)_get_tensor_memory()- Get tensor memory size_categorize_memory_by_location()- Categorize tensors into HBM/UVM_report_tbe_mem_usage()with clean list-based tensor grouping3. Memory Components
Static Sparse:
weights_dev,weights_host,weights_uvmmomentum1_dev/host/uvm,momentum2_dev/host/uvmlxu_cache_weights,lxu_cache_state,lxu_state, cache aux dataEphemeral (calculated):
ephemeral = total_mem_usage - static_sparsedetailed analysis revealed QPS drop when enabling additional logging. head to head comparison of time it takes to do the logging reveals 4x increase in duration (see (https://fburl.com/scuba/tbe_stats_runtime/y85ur4k9)
ran the following models w/o added logging:
ran the following models w/ added logging:
Differential Revision: D84624978