Skip to content

A trial of apply @fused to optimzier state offloading #2513

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.

Sign up for GitHub

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

Jump to bottom
Draft
wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

A trial of apply @fused to optimzier state offloading #2513

wants to merge 4 commits into from

Conversation

@zhenying-liu
Copy link
Contributor

Apply @jax.fused to optimizer state and parameter update. This can avoid memory copies while offloading the optimizer states to the host memory. The device will read and write the host memory directly with fine grained memory control.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk will be requested when the pull request is marked ready for review gobbleturk is a code owner

@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit will be requested when the pull request is marked ready for review khatwanimohit is a code owner

@shralex shralex Awaiting requested review from shralex shralex will be requested when the pull request is marked ready for review shralex is a code owner

@parambole parambole Awaiting requested review from parambole parambole will be requested when the pull request is marked ready for review parambole is a code owner

@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon will be requested when the pull request is marked ready for review bvandermoon is a code owner

@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 will be requested when the pull request is marked ready for review richjames0 is a code owner

@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla will be requested when the pull request is marked ready for review vipannalla is a code owner

@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan will be requested when the pull request is marked ready for review RissyRan is a code owner

@gagika gagika Awaiting requested review from gagika gagika will be requested when the pull request is marked ready for review gagika is a code owner

@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC will be requested when the pull request is marked ready for review SurbhiJainUSC is a code owner

@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo will be requested when the pull request is marked ready for review hengtaoguo is a code owner

@A9isha A9isha Awaiting requested review from A9isha A9isha will be requested when the pull request is marked ready for review A9isha is a code owner

@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei will be requested when the pull request is marked ready for review aireenmei is a code owner

@NuojCheng NuojCheng Awaiting requested review from NuojCheng NuojCheng will be requested when the pull request is marked ready for review NuojCheng is a code owner

@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 will be requested when the pull request is marked ready for review jiangjy1982 is a code owner

@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 will be requested when the pull request is marked ready for review suexu1025 is a code owner

At least 2 approving reviews are required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@zhenying-liu @khatwanimohit @NuojCheng