Skip to content

pipeline aware cpu offload #1886

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

Open
wants to merge 3 commits into
base: main
Choose a base branch
from

Conversation

liuzhenhai93
Copy link

@liuzhenhai93 liuzhenhai93 commented Jun 17, 2025

Description

pipeline aware cpu offload

Type of change

  • New feature (non-breaking change which adds functionality)

Changes

Please list the changes introduced in this PR:

  • pipeline aware cpu offload

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

self._b_event = PipelineOffloadManager.get_instance()._b_event
self.do_offload = offload

def is_first_last_layer(self):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The naming is ambiguous. It returns true when the current layer is the last layer and the current vpp chunk is the first one or the last one, is my understanding correct?

return self.cur_backward_chunk().tensor_pop(saved_state)


OFFLOAD_TAG = "offloading_mlp_input"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we only support offloading mlp input?

timmoon10 and others added 2 commits June 17, 2025 20:37
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: liuzhenhai93 <liuzhenhai93@outlook.com>
Signed-off-by: liuzhenhai93 <liuzhenhai93@outlook.com>
@pggPL
Copy link
Collaborator

pggPL commented Jun 17, 2025

Hi, thank you for the PR. We are working on some bigger changes in CPU Offload so I think we will need to sync. I reached out to @lhb8125.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants