Closed MingjiHan99 closed 2 months ago
From my understanding, Checkpoint Broadcasting might be beneficial only when the storage throughput is limited under certain circumstances. Maybe it's better to add some more guidance on when to use this feature.
Let's keep pushing forward 💪!
Summary
We improved
vescale.checkpoint
with the following new features for fast checkpointing (where front three features are built-in techniques without necessitating manual activation):Saving Plan Caching: During training, the program may save model and optimizer checkpoints every n steps. Once a saving plan is created, it remains unchanged as long as the model does. We implemented plan caching to avoid regenerating the plan when checkpointing a model or optimizer multiple times, reducing unnecessary compute and communication costs. As of 05/30/2024, PyTorch DCP does not support plan caching.
Saving Plan Load-Balancing: In data parallel training, models are replicated across GPUs with different data parallel ranks but the same pipeline and tensor parallel ranks. Existing PyTorch DCP (as of 05/30/2024) deduplicates replicated tensors using a simple algorithm, causing GPUs with data parallel rank 0 to save the entire model, leading to load imbalance. We implemented a load-balancing algorithm to address this issue when deduplicating model tensors.
D2H Tensor Copying via Pinned Memory: When copying tensors from GPU to host memory,
vescale.checkpoint
uses pinned host memory, reducing memory allocation costs each time a checkpoint is saved. As of 05/30/2024, PyTorch DCP does not support pinned memory.Checkpoint Broadcasting: In data parallel training, models are replicated across GPUs with different data parallel ranks but the same pipeline and tensor parallel ranks. If
broadcast_checkpoint
is enabled,vescale.checkpoint.load
lets GPUs with data parallel rank 0 to load the model and broadcast it to other GPUs with higher data parallel ranks. If GPUs are connected with NCCL and I/O bandwidth is fully utilized, broadcasting model tensors speeds up checkpoint loading compared to all GPUs loading models from persistent storage. E.g.:Asynchronous Checkpointing: When
vescale.checkpoint.save
is called, it first generates a saving plan and then synchronously copies tensors from GPU to host memory. Ifasync_checkpoint
is enabled, the training program can continue after the D2H copying, whilevescale.checkpoint.save
continues to serialize tensors and dump the checkpoint to persistent storage asynchronously without blocking training. As of 05/30/2024, PyTorch DCP does not support asynchronous checkpointing. E.g.:Acknowledgement
We sincerely appreciate all contributors including but not limited to @shanesyy-1992 @raywan-110 @lazychao @AHEADer @MingjiHan99