__init__.py 561 B

1234567891011121314151617181920
  1. # Copyright (c) Microsoft Corporation.
  2. # SPDX-License-Identifier: Apache-2.0
  3. # DeepSpeed Team
  4. from .reshape_meg_2d import reshape_meg_2d_parallel
  5. from .deepspeed_checkpoint import DeepSpeedCheckpoint
  6. from .utils import (get_layer_ckpt_name_for_rank, get_model_ckpt_name_for_rank, get_zero_ckpt_name_for_rank)
  7. from .reshape_utils import (merge_state)
  8. from .reshape_3d_utils import (model_3d_desc, get_model_3d_descriptor)
  9. from .zero_checkpoint import ZeROCheckpoint
  10. from .universal_checkpoint import enable_universal_checkpoint
  11. from .constants import *