Skip to content

Patch RNG for better memory utilization in dropout layers#5710

Merged
yeounoh merged 6 commits intomasterfrom
patch_rng_xla
Oct 24, 2023
Merged

Patch RNG for better memory utilization in dropout layers#5710
yeounoh merged 6 commits intomasterfrom
patch_rng_xla

Conversation

@yeounoh
Copy link
Copy Markdown
Contributor

@yeounoh yeounoh commented Oct 17, 2023

The patch has landed in OpenXLA openxla/xla#6015, we can enable this now, to reduce memory pressure in nn.Dropout layer.

DO_NOT_MERGE until the next pin update.

@yeounoh yeounoh requested a review from JackCaoG October 17, 2023 20:24
@yeounoh yeounoh self-assigned this Oct 17, 2023
@yeounoh yeounoh changed the title Patch RNG Patch RNG for better memory utilization in dropout layers Oct 17, 2023
@yeounoh yeounoh marked this pull request as draft October 17, 2023 20:25
@yeounoh yeounoh added the DO_NOT_MERGE Not for merging. label Oct 17, 2023
@JackCaoG
Copy link
Copy Markdown
Collaborator

do we want to add a test to check that int8 is used for rng? I think we can check the HLO generated,

@yeounoh
Copy link
Copy Markdown
Contributor Author

yeounoh commented Oct 20, 2023

do we want to add a test to check that int8 is used for rng? I think we can check the HLO generated,

Hey Jack, yea good idea, will add before merging.

@yeounoh yeounoh removed the DO_NOT_MERGE Not for merging. label Oct 23, 2023
@yeounoh yeounoh marked this pull request as ready for review October 23, 2023 22:14
@yeounoh yeounoh force-pushed the patch_rng_xla branch 3 times, most recently from df85d09 to e26a3ee Compare October 23, 2023 22:33
@yeounoh
Copy link
Copy Markdown
Contributor Author

yeounoh commented Oct 23, 2023

do we want to add a test to check that int8 is used for rng? I think we can check the HLO generated,

Hey Jack, yea good idea, will add before merging.

Verified locally with the new unit test -- cc @JackCaoG

@yeounoh yeounoh merged commit 8f121c3 into master Oct 24, 2023
ghpvnist pushed a commit to ghpvnist/pytorch-xla that referenced this pull request Oct 31, 2023
* Enable u8 rng-bit-generator with downcast

* Use BF16 values if downcast for uniform dist is set.
mbzomowski pushed a commit to mbzomowski-test-org/xla that referenced this pull request Nov 16, 2023
* Enable u8 rng-bit-generator with downcast

* Use BF16 values if downcast for uniform dist is set.
chunnienc pushed a commit to chunnienc/xla that referenced this pull request Dec 14, 2023
* Enable u8 rng-bit-generator with downcast

* Use BF16 values if downcast for uniform dist is set.
golechwierowicz pushed a commit that referenced this pull request Jan 12, 2024
* Enable u8 rng-bit-generator with downcast

* Use BF16 values if downcast for uniform dist is set.
bhavya01 pushed a commit that referenced this pull request Apr 22, 2024
* Enable u8 rng-bit-generator with downcast

* Use BF16 values if downcast for uniform dist is set.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants