Skip to content

[REQUEST] Activation Checkpoint Prefetch  #1575

Description

@ckddls1321

Is your feature request related to a problem? Please describe.
Activation prefetch features to enlarge batch size on middle-size(100B~1T) of models

  • From DeepSpeedExamples repo, GPU throughput with Activation checkpoint in CPU matters.
  • On A100 server pods, Activation checkpoint in CPU perform worse because of synchronization (HtoD memcpy or all-gather when partitioned activation).
  • For Tensor Parallelism, it would be better to save activation checkpoint in partitioned tensor.

Describe the solution you'd like
Activation prefetch during backward re-computation based on configuration
Furthermore, to increase batch size in tensor parallelism, asynchronous all-gather prefetched partitioned activation is needed.

Describe alternatives you've considered
For large scale GPU pods more than 128 GPU would not be problem.
It would be appreciated if we have candidates of GPU cluster size based on model configuration (10B, 50B, 100B, 1T).

Activity

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

Metadata

Metadata

Assignees

Labels

enhancementNew feature or request

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions