Skip to content

Conversation

@MayDomine
Copy link
Collaborator

add async Offload Activation for CheckpointBlock

@MayDomine MayDomine requested review from Achazwl, a710128 and zkh2016 and removed request for a710128 and zkh2016 August 29, 2023 03:10
from .global_var import config
from .checkpointing import CheckpointBlockContext

from .distributed import all_gather, broadcast, all_reduce, send_activations, recv_activations
Copy link
Collaborator

Choose a reason for hiding this comment

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

这些没用到,可以去掉

ctypes.CDLL(os.path.join(path, file_so))


def find_pre_module_helper(m):
Copy link
Collaborator

Choose a reason for hiding this comment

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

find_pre_offload_module_helper

module = CheckpointBlock(module)

module._mode = "ZERO"
module._mode = "ZERO" if module._mode == "BLOCK" else module._mode
Copy link
Collaborator

Choose a reason for hiding this comment

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

这里可以直接去掉,现在不需要区分ZERO和BLOCK了

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.

3 participants