Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Patch RNG for better memory utilization in dropout layers #5710

Merged
merged 6 commits into from
Oct 24, 2023

Conversation

yeounoh
Copy link
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
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
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 added the perf label Oct 23, 2023
@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
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
18 checks passed
ghpvnist pushed a commit to ghpvnist/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
Labels
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants