Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/utilities/deepspeed.py (Python)

Code before:
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Modified script from https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/utils/zero_to_fp32.py

This script extracts fp32 consolidated weights from a zero 2 and 3 DeepSpeed checkpoints. It gets
copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
the future. Once extracted, the weights don't require DeepSpeed and can be used in any
application. Additionally the script has been modified to ensure we keep the lightning state inside the state dict
for being able to run Model.load_from_checkpoint('...').

Example usage within the Lightning checkpoint directory where 'latest' is found:

>>> from pytorch_lightning.utilities.deepspeed import convert_zero_checkpoint_to_fp32_state_dict

# Lightning deepspeed has saved a directory instead of a file

>>> save_path = "lightning_logs/version_0/checkpoints/epoch=0-step=0.ckpt/"
>>> output_path = "lightning_model.pt"
>>> convert_zero_checkpoint_to_fp32_state_dict(save_path, output_path)
Saving fp32 state dict to lightning_model.pt
"""

import os

import torch

from pytorch_lightning.utilities import _DEEPSPEED_AVAILABLE

if _DEEPSPEED_AVAILABLE:
    from deepspeed.utils.zero_to_fp32 import (
        get_fp32_state_dict_from_zero_checkpoint,
        get_model_state_file,
        get_optim_files,
    )

CPU_DEVICE = torch.device("cpu")


def ds_checkpoint_dir(checkpoint_dir: str, tag: str = None):
    if tag is None:
        latest_path = os.path.join(checkpoint_dir, "latest")
        if os.path.isfile(latest_path):
            with open(latest_path) as fd:
                tag = fd.read().strip()

Review comment (suggestion):
```suggestion
from pytorch_lightning.utilities.deepspeed import convert_zero_checkpoint_to_fp32_state_dict

Lightning deepspeed has saved a directory instead of a file

save_path = "lightning_logs/version_0/checkpoints/epoch=0-step=0.ckpt/"
output_path = "lightning_model.pt"
convert_zero_checkpoint_to_fp32_state_dict(save_path, output_path)
```

Code after:
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Modified script from https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/utils/zero_to_fp32.py

This script extracts fp32 consolidated weights from a zero 2 and 3 DeepSpeed checkpoints. It gets
copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
the future. Once extracted, the weights don't require DeepSpeed and can be used in any
application. Additionally the script has been modified to ensure we keep the lightning state inside the state dict
for being able to run Model.load_from_checkpoint('...').

Example usage within the Lightning checkpoint directory where 'latest' is found:

>>> from pytorch_lightning.utilities.deepspeed import convert_zero_checkpoint_to_fp32_state_dict # doctest: +SKIP

# Lightning deepspeed has saved a directory instead of a file

>>> save_path = "lightning_logs/version_0/checkpoints/epoch=0-step=0.ckpt/" # doctest: +SKIP
>>> output_path = "lightning_model.pt" # doctest: +SKIP
>>> convert_zero_checkpoint_to_fp32_state_dict(save_path, output_path) # doctest: +SKIP
Saving fp32 state dict to lightning_model.pt
"""

import os

import torch

from pytorch_lightning.utilities import _DEEPSPEED_AVAILABLE

if _DEEPSPEED_AVAILABLE:
    from deepspeed.utils.zero_to_fp32 import (
        get_fp32_state_dict_from_zero_checkpoint,
        get_model_state_file,
        get_optim_files,
    )

CPU_DEVICE = torch.device("cpu")


def ds_checkpoint_dir(checkpoint_dir: str, tag: str = None):
    if tag is None:
        latest_path = os.path.join(checkpoint_dir, "latest")
        if os.path.isfile(latest_path):
            with open(latest_path) as fd:
                tag = fd.read().strip()