Fix semi-sync race conditions and optimize memory usage #2598

sarckk · 2024-11-27T19:10:16Z

Summary:
This PR introduces a set of fixes and optimizations for TorchRec's semi-synchronous training pipeline (where embeddings are fetched in advance and in parallel with backward):

Fixes memory safety issues causing CUDA Illegal Memory Access errors and NaNs during training, especially at high memory utilisations. The fix involves using calling record_stream for tensors allocated in one CUDA stream and used in another (an example is embedding tensors which are looked up in an embedding stream but later accessed in the main stream).
Optimizes memory allocations to reduce memory fragmentation observed at high memory utilizations causing significant performance degradation due to expensive defragmentation calls (thanks to @che-sh for initial observation and analysis). Optimizations include:
- Freeing context objects and cached module outputs as early as possible to save memory
- Moving small tensor allocations earlier to minimize fragmentation
- Using a single stream per embedding module (instead of even and odd streams) as memory allocations by PyTorch's CUDACachingAllocator are associated with streams, meaning freed memory blocks in one stream cannot be used by other streams. By using more streams, we effectively decrease memory available to each stream, making defrags more likely.

This is joint work with @che-sh (optimizations to reduce memory fragmentation) and @dstaay-fb (record_stream fixes for embeddings)

Differential Revision: D64220706

facebook-github-bot · 2024-11-27T19:10:35Z

This pull request was exported from Phabricator. Differential Revision: D64220706

facebook-github-bot · 2024-11-27T19:13:29Z

This pull request was exported from Phabricator. Differential Revision: D64220706

Summary: This PR introduces a set of fixes and optimizations for TorchRec's semi-synchronous training pipeline (where embeddings are fetched in advance and in parallel with backward): - Fixes memory safety issues causing CUDA Illegal Memory Access errors and NaNs during training, especially at high memory utilisations. The fix involves using calling [`record_stream`](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) for tensors allocated in one CUDA stream and used in another (an example is embedding tensors which are looked up in an embedding stream but later accessed in the main stream). - Optimizes memory allocations to reduce memory fragmentation observed at high memory utilizations causing significant performance degradation due to expensive defragmentation calls (thanks to che-sh for initial observation and analysis). Optimizations include: - Freeing context objects and cached module outputs as early as possible to save memory - Moving small tensor allocations earlier to minimize fragmentation - Using a single stream per embedding module (instead of even and odd streams) as memory allocations by PyTorch's CUDACachingAllocator are associated with streams, meaning freed memory blocks in one stream cannot be used by other streams. By using more streams, we effectively decrease memory available to each stream, making defrags more likely. This is joint work with che-sh (optimizations to reduce memory fragmentation) and dstaay-fb (record_stream fixes for embeddings) Reviewed By: che-sh Differential Revision: D64220706

facebook-github-bot · 2024-12-05T20:21:47Z

This pull request was exported from Phabricator. Differential Revision: D64220706

Summary: This PR introduces a set of fixes and optimizations for TorchRec's semi-synchronous training pipeline (where embeddings are fetched in advance and in parallel with backward): - Fixes memory safety issues causing CUDA Illegal Memory Access errors and NaNs during training, especially at high memory utilisations. The fix involves using calling [`record_stream`](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) for tensors allocated in one CUDA stream and used in another (an example is embedding tensors which are looked up in an embedding stream but later accessed in the main stream). - Optimizes memory allocations to reduce memory fragmentation observed at high memory utilizations causing significant performance degradation due to expensive defragmentation calls (thanks to che-sh for initial observation and analysis). Optimizations include: - Freeing context objects and cached module outputs as early as possible to save memory - Moving small tensor allocations earlier to minimize fragmentation - Using a single stream per embedding module (instead of even and odd streams) as memory allocations by PyTorch's CUDACachingAllocator are associated with streams, meaning freed memory blocks in one stream cannot be used by other streams. By using more streams, we effectively decrease memory available to each stream, making defrags more likely. This is joint work with che-sh (optimizations to reduce memory fragmentation) and dstaay-fb (record_stream fixes for embeddings) Reviewed By: che-sh Differential Revision: D64220706

facebook-github-bot · 2024-12-10T02:20:10Z

This pull request was exported from Phabricator. Differential Revision: D64220706

Summary: This PR introduces a set of fixes and optimizations for TorchRec's semi-synchronous training pipeline (where embeddings are fetched in advance and in parallel with backward): - Fixes memory safety issues causing CUDA Illegal Memory Access errors and NaNs during training, especially at high memory utilisations. The fix involves using calling [`record_stream`](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) for tensors allocated in one CUDA stream and used in another (an example is embedding tensors which are looked up in an embedding stream but later accessed in the main stream). - Optimizes memory allocations to reduce memory fragmentation observed at high memory utilizations causing significant performance degradation due to expensive defragmentation calls (thanks to che-sh for initial observation and analysis). Optimizations include: - Freeing context objects and cached module outputs as early as possible to save memory - Moving small tensor allocations earlier to minimize fragmentation - Using a single stream per embedding module (instead of even and odd streams) as memory allocations by PyTorch's CUDACachingAllocator are associated with streams, meaning freed memory blocks in one stream cannot be used by other streams. By using more streams, we effectively decrease memory available to each stream, making defrags more likely. This is joint work with che-sh (optimizations to reduce memory fragmentation) and dstaay-fb (record_stream fixes for embeddings) Reviewed By: che-sh Differential Revision: D64220706

facebook-github-bot · 2024-12-10T21:25:53Z

This pull request was exported from Phabricator. Differential Revision: D64220706

Summary: This PR introduces a set of fixes and optimizations for TorchRec's semi-synchronous training pipeline (where embeddings are fetched in advance and in parallel with backward): - Fixes memory safety issues causing CUDA Illegal Memory Access errors and NaNs during training, especially at high memory utilisations. The fix involves using calling [`record_stream`](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) for tensors allocated in one CUDA stream and used in another (an example is embedding tensors which are looked up in an embedding stream but later accessed in the main stream). - Optimizes memory allocations to reduce memory fragmentation observed at high memory utilizations causing significant performance degradation due to expensive defragmentation calls (thanks to che-sh for initial observation and analysis). Optimizations include: - Freeing context objects and cached module outputs as early as possible to save memory - Moving small tensor allocations earlier to minimize fragmentation - Using a single stream per embedding module (instead of even and odd streams) as memory allocations by PyTorch's CUDACachingAllocator are associated with streams, meaning freed memory blocks in one stream cannot be used by other streams. By using more streams, we effectively decrease memory available to each stream, making defrags more likely. This is joint work with che-sh (optimizations to reduce memory fragmentation) and dstaay-fb (record_stream fixes for embeddings) Reviewed By: che-sh Differential Revision: D64220706

facebook-github-bot · 2024-12-10T21:26:22Z

This pull request was exported from Phabricator. Differential Revision: D64220706

Summary: This PR introduces a set of fixes and optimizations for TorchRec's semi-synchronous training pipeline (where embeddings are fetched in advance and in parallel with backward): - Fixes memory safety issues causing CUDA Illegal Memory Access errors and NaNs during training, especially at high memory utilisations. The fix involves using calling [`record_stream`](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) for tensors allocated in one CUDA stream and used in another (an example is embedding tensors which are looked up in an embedding stream but later accessed in the main stream). - Optimizes memory allocations to reduce memory fragmentation observed at high memory utilizations causing significant performance degradation due to expensive defragmentation calls (thanks to che-sh for initial observation and analysis). Optimizations include: - Freeing context objects and cached module outputs as early as possible to save memory - Moving small tensor allocations earlier to minimize fragmentation - Using a single stream per embedding module (instead of even and odd streams) as memory allocations by PyTorch's CUDACachingAllocator are associated with streams, meaning freed memory blocks in one stream cannot be used by other streams. By using more streams, we effectively decrease memory available to each stream, making defrags more likely. This is joint work with che-sh (optimizations to reduce memory fragmentation) and dstaay-fb (record_stream fixes for embeddings) Reviewed By: che-sh Differential Revision: D64220706

facebook-github-bot · 2024-12-10T23:39:09Z

This pull request was exported from Phabricator. Differential Revision: D64220706

Summary: This PR introduces a set of fixes and optimizations for TorchRec's semi-synchronous training pipeline (where embeddings are fetched in advance and in parallel with backward): - Fixes memory safety issues causing CUDA Illegal Memory Access errors and NaNs during training, especially at high memory utilisations. The fix involves using calling [`record_stream`](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) for tensors allocated in one CUDA stream and used in another (an example is embedding tensors which are looked up in an embedding stream but later accessed in the main stream). - Optimizes memory allocations to reduce memory fragmentation observed at high memory utilizations causing significant performance degradation due to expensive defragmentation calls (thanks to che-sh for initial observation and analysis). Optimizations include: - Freeing context objects and cached module outputs as early as possible to save memory - Moving small tensor allocations earlier to minimize fragmentation - Using a single stream per embedding module (instead of even and odd streams) as memory allocations by PyTorch's CUDACachingAllocator are associated with streams, meaning freed memory blocks in one stream cannot be used by other streams. By using more streams, we effectively decrease memory available to each stream, making defrags more likely. This is joint work with che-sh (optimizations to reduce memory fragmentation) and dstaay-fb (record_stream fixes for embeddings) Reviewed By: che-sh Differential Revision: D64220706

facebook-github-bot · 2024-12-10T23:39:26Z

This pull request was exported from Phabricator. Differential Revision: D64220706

Summary: This PR introduces a set of fixes and optimizations for TorchRec's semi-synchronous training pipeline (where embeddings are fetched in advance and in parallel with backward): - Fixes memory safety issues causing CUDA Illegal Memory Access errors and NaNs during training, especially at high memory utilisations. The fix involves using calling [`record_stream`](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) for tensors allocated in one CUDA stream and used in another (an example is embedding tensors which are looked up in an embedding stream but later accessed in the main stream). - Optimizes memory allocations to reduce memory fragmentation observed at high memory utilizations causing significant performance degradation due to expensive defragmentation calls (thanks to che-sh for initial observation and analysis). Optimizations include: - Freeing context objects and cached module outputs as early as possible to save memory - Moving small tensor allocations earlier to minimize fragmentation - Using a single stream per embedding module (instead of even and odd streams) as memory allocations by PyTorch's CUDACachingAllocator are associated with streams, meaning freed memory blocks in one stream cannot be used by other streams. By using more streams, we effectively decrease memory available to each stream, making defrags more likely. This is joint work with che-sh (optimizations to reduce memory fragmentation) and dstaay-fb (record_stream fixes for embeddings) Reviewed By: che-sh Differential Revision: D64220706

facebook-github-bot · 2024-12-11T00:30:09Z

This pull request was exported from Phabricator. Differential Revision: D64220706

Summary: This PR introduces a set of fixes and optimizations for TorchRec's semi-synchronous training pipeline (where embeddings are fetched in advance and in parallel with backward): - Fixes memory safety issues causing CUDA Illegal Memory Access errors and NaNs during training, especially at high memory utilisations. The fix involves using calling [`record_stream`](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) for tensors allocated in one CUDA stream and used in another (an example is embedding tensors which are looked up in an embedding stream but later accessed in the main stream). - Optimizes memory allocations to reduce memory fragmentation observed at high memory utilizations causing significant performance degradation due to expensive defragmentation calls (thanks to che-sh for initial observation and analysis). Optimizations include: - Freeing context objects and cached module outputs as early as possible to save memory - Moving small tensor allocations earlier to minimize fragmentation - Using a single stream per embedding module (instead of even and odd streams) as memory allocations by PyTorch's CUDACachingAllocator are associated with streams, meaning freed memory blocks in one stream cannot be used by other streams. By using more streams, we effectively decrease memory available to each stream, making defrags more likely. This is joint work with che-sh (optimizations to reduce memory fragmentation) and dstaay-fb (record_stream fixes for embeddings) Reviewed By: che-sh Differential Revision: D64220706

facebook-github-bot · 2024-12-11T00:34:35Z

This pull request was exported from Phabricator. Differential Revision: D64220706

Summary: This PR introduces a set of fixes and optimizations for TorchRec's semi-synchronous training pipeline (where embeddings are fetched in advance and in parallel with backward): - Fixes memory safety issues causing CUDA Illegal Memory Access errors and NaNs during training, especially at high memory utilisations. The fix involves using calling [`record_stream`](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) for tensors allocated in one CUDA stream and used in another (an example is embedding tensors which are looked up in an embedding stream but later accessed in the main stream). - Optimizes memory allocations to reduce memory fragmentation observed at high memory utilizations causing significant performance degradation due to expensive defragmentation calls (thanks to che-sh for initial observation and analysis). Optimizations include: - Freeing context objects and cached module outputs as early as possible to save memory - Moving small tensor allocations earlier to minimize fragmentation - Using a single stream per embedding module (instead of even and odd streams) as memory allocations by PyTorch's CUDACachingAllocator are associated with streams, meaning freed memory blocks in one stream cannot be used by other streams. By using more streams, we effectively decrease memory available to each stream, making defrags more likely. This is joint work with che-sh (optimizations to reduce memory fragmentation) and dstaay-fb (record_stream fixes for embeddings) Reviewed By: che-sh Differential Revision: D64220706

facebook-github-bot · 2024-12-11T00:38:52Z

This pull request was exported from Phabricator. Differential Revision: D64220706

Summary: This PR introduces a set of fixes and optimizations for TorchRec's semi-synchronous training pipeline (where embeddings are fetched in advance and in parallel with backward): - Fixes memory safety issues causing CUDA Illegal Memory Access errors and NaNs during training, especially at high memory utilisations. The fix involves using calling [`record_stream`](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) for tensors allocated in one CUDA stream and used in another (an example is embedding tensors which are looked up in an embedding stream but later accessed in the main stream). - Optimizes memory allocations to reduce memory fragmentation observed at high memory utilizations causing significant performance degradation due to expensive defragmentation calls (thanks to che-sh for initial observation and analysis). Optimizations include: - Freeing context objects and cached module outputs as early as possible to save memory - Moving small tensor allocations earlier to minimize fragmentation - Using a single stream per embedding module (instead of even and odd streams) as memory allocations by PyTorch's CUDACachingAllocator are associated with streams, meaning freed memory blocks in one stream cannot be used by other streams. By using more streams, we effectively decrease memory available to each stream, making defrags more likely. This is joint work with che-sh (optimizations to reduce memory fragmentation) and dstaay-fb (record_stream fixes for embeddings) Reviewed By: che-sh Differential Revision: D64220706

facebook-github-bot · 2024-12-11T22:12:57Z

This pull request was exported from Phabricator. Differential Revision: D64220706

Summary: This PR introduces a set of fixes and optimizations for TorchRec's semi-synchronous training pipeline (where embeddings are fetched in advance and in parallel with backward): - Fixes memory safety issues causing CUDA Illegal Memory Access errors and NaNs during training, especially at high memory utilisations. The fix involves using calling [`record_stream`](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) for tensors allocated in one CUDA stream and used in another (an example is embedding tensors which are looked up in an embedding stream but later accessed in the main stream). - Optimizes memory allocations to reduce memory fragmentation observed at high memory utilizations causing significant performance degradation due to expensive defragmentation calls (thanks to che-sh for initial observation and analysis). Optimizations include: - Freeing context objects and cached module outputs as early as possible to save memory - Moving small tensor allocations earlier to minimize fragmentation - Using a single stream per embedding module (instead of even and odd streams) as memory allocations by PyTorch's CUDACachingAllocator are associated with streams, meaning freed memory blocks in one stream cannot be used by other streams. By using more streams, we effectively decrease memory available to each stream, making defrags more likely. This is joint work with che-sh (optimizations to reduce memory fragmentation) and dstaay-fb (record_stream fixes for embeddings) Reviewed By: che-sh Differential Revision: D64220706

facebook-github-bot · 2024-12-12T06:04:45Z

This pull request was exported from Phabricator. Differential Revision: D64220706

Summary: This PR introduces a set of fixes and optimizations for TorchRec's semi-synchronous training pipeline (where embeddings are fetched in advance and in parallel with backward): - Fixes memory safety issues causing CUDA Illegal Memory Access errors and NaNs during training, especially at high memory utilisations. The fix involves using calling [`record_stream`](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) for tensors allocated in one CUDA stream and used in another (an example is embedding tensors which are looked up in an embedding stream but later accessed in the main stream). - Optimizes memory allocations to reduce memory fragmentation observed at high memory utilizations causing significant performance degradation due to expensive defragmentation calls (thanks to che-sh for initial observation and analysis). Optimizations include: - Freeing context objects and cached module outputs as early as possible to save memory - Moving small tensor allocations earlier to minimize fragmentation - Using a single stream per embedding module (instead of even and odd streams) as memory allocations by PyTorch's CUDACachingAllocator are associated with streams, meaning freed memory blocks in one stream cannot be used by other streams. By using more streams, we effectively decrease memory available to each stream, making defrags more likely. This is joint work with che-sh (optimizations to reduce memory fragmentation) and dstaay-fb (record_stream fixes for embeddings) Reviewed By: che-sh Differential Revision: D64220706

facebook-github-bot · 2024-12-12T06:18:11Z

This pull request was exported from Phabricator. Differential Revision: D64220706

Summary: pytorch#2598 causes failures when running on cpu: ``` RuntimeError: unknown parameter type ``` This is because `record_stream` doesn't accept CPU streams. This diff forward fixes the issue. Differential Revision: D67171994

Summary: Pull Request resolved: #2634 #2598 causes failures when running on cpu: ``` RuntimeError: unknown parameter type ``` This is because `record_stream` doesn't accept CPU streams. This diff forward fixes the issue. Differential Revision: D67171994 fbshipit-source-id: a9ffddfcd978d343a58fb039a4fcd89c336ceee6

Summary: pytorch#2598 (D64220706) causes failures when using other accelerators that do not support CUDA. Making the stream contexts hardware agnostic. Differential Revision: D67363141

Summary: Pull Request resolved: #2644 #2598 (D64220706) causes failures when using other accelerators that do not support CUDA. Making the stream contexts hardware agnostic. Reviewed By: hpnhxxwn, iamzainhuda Differential Revision: D67363141 fbshipit-source-id: fc2c6fec1dcbbe15f0385e299b666207d2d9a8f5

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 27, 2024

facebook-github-bot added the fb-exported label Nov 27, 2024

sarckk force-pushed the export-D64220706 branch from 82a5651 to ab693d3 Compare November 27, 2024 19:12

sarckk force-pushed the export-D64220706 branch from ab693d3 to fb4c04e Compare December 5, 2024 20:21

sarckk force-pushed the export-D64220706 branch from fb4c04e to 4608336 Compare December 10, 2024 02:19

sarckk force-pushed the export-D64220706 branch from 4608336 to 5233e64 Compare December 10, 2024 21:25

sarckk force-pushed the export-D64220706 branch from 5233e64 to 75c3c62 Compare December 10, 2024 21:26

sarckk force-pushed the export-D64220706 branch from 75c3c62 to c5484d9 Compare December 10, 2024 23:38

sarckk force-pushed the export-D64220706 branch from c5484d9 to 8040d2f Compare December 10, 2024 23:39

sarckk force-pushed the export-D64220706 branch from 8040d2f to a08ad09 Compare December 11, 2024 00:29

sarckk force-pushed the export-D64220706 branch from a08ad09 to 3979ce5 Compare December 11, 2024 00:34

sarckk force-pushed the export-D64220706 branch from 3979ce5 to 98db96f Compare December 11, 2024 00:38

sarckk force-pushed the export-D64220706 branch from 98db96f to 13a7d2d Compare December 11, 2024 22:12

sarckk force-pushed the export-D64220706 branch from 13a7d2d to 4174768 Compare December 12, 2024 06:04

sarckk force-pushed the export-D64220706 branch from 4174768 to 3996038 Compare December 12, 2024 06:18

facebook-github-bot closed this in 575e081 Dec 12, 2024

sarckk mentioned this pull request Dec 13, 2024

Fix invalid record_stream arg for Request awaitable #2634

Closed

sarckk mentioned this pull request Dec 17, 2024

Make streams device-agnostic #2644

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix semi-sync race conditions and optimize memory usage #2598

Fix semi-sync race conditions and optimize memory usage #2598

sarckk commented Nov 27, 2024 •

edited

Loading

facebook-github-bot commented Nov 27, 2024

facebook-github-bot commented Nov 27, 2024

facebook-github-bot commented Dec 5, 2024

facebook-github-bot commented Dec 10, 2024

facebook-github-bot commented Dec 10, 2024

facebook-github-bot commented Dec 10, 2024

facebook-github-bot commented Dec 10, 2024

facebook-github-bot commented Dec 10, 2024

facebook-github-bot commented Dec 11, 2024

facebook-github-bot commented Dec 11, 2024

facebook-github-bot commented Dec 11, 2024

facebook-github-bot commented Dec 11, 2024

facebook-github-bot commented Dec 12, 2024

facebook-github-bot commented Dec 12, 2024

Fix semi-sync race conditions and optimize memory usage #2598

Fix semi-sync race conditions and optimize memory usage #2598

Conversation

sarckk commented Nov 27, 2024 • edited Loading

facebook-github-bot commented Nov 27, 2024

facebook-github-bot commented Nov 27, 2024

facebook-github-bot commented Dec 5, 2024

facebook-github-bot commented Dec 10, 2024

facebook-github-bot commented Dec 10, 2024

facebook-github-bot commented Dec 10, 2024

facebook-github-bot commented Dec 10, 2024

facebook-github-bot commented Dec 10, 2024

facebook-github-bot commented Dec 11, 2024

facebook-github-bot commented Dec 11, 2024

facebook-github-bot commented Dec 11, 2024

facebook-github-bot commented Dec 11, 2024

facebook-github-bot commented Dec 12, 2024

facebook-github-bot commented Dec 12, 2024

sarckk commented Nov 27, 2024 •

edited

Loading