DeepSpeed
799120e7 - Universal checkpoint for zero stage 1 (#2284)

Commit

2 years ago

Universal checkpoint for zero stage 1 (#2284) * Refactor universal checkpointing and tensor fragments * Formatting * Support zero stage1; Expand TP dim * Remove debug prints * Detect sharded optimizer state * Format fixes * Encode reshaping guide * More symbolic constants Co-authored-by: Michael Wyatt <michaelwyatt@microsoft.com>

References

#2284 - Universal checkpoint for zero stage 1

Author

tjruwase

Parents

906b4a02

Files11

deepspeed
- checkpoint
  - __init__.py
  - constants.py
  - deepspeed_checkpoint.py
  - reshape_3d_utils.py
  - reshape_utils.py
  - universal_checkpoint.py
  - zero_checkpoint.py
- runtime
  - bf16_optimizer.py
  - engine.py
  - zero
    - stage_1_and_2.py
- utils
  - tensor_fragment.py

DeepSpeed 799120e7 - Universal checkpoint for zero stage 1 (#2284)

DeepSpeed
799120e7 - Universal checkpoint for zero stage 1 (#2284)