Fix standalone ZeRO conversion script

Signed-off-by: gaoxiaomo <165135449+gaoxiaomo@users.noreply.github.com>

Author: gaoxiaomo

deepspeed/utils/zero_to_torch.py

@@ -1,9 +1,14 @@
+#!/usr/bin/env python
+
 # SPDX-License-Identifier: Apache-2.0
 # DeepSpeed Team
 
 import argparse
 
-import deepspeed.utils.zero_to_fp32 as zero_to_fp32
+if __package__:
+    from . import zero_to_fp32
+else:
+    import zero_to_fp32
 
 
 def main(args=None):

tests/unit/checkpoint/test_convert_checkpoint.py

@@ -3,12 +3,17 @@
 
 # DeepSpeed Team
 
+import shutil
+import subprocess
+import sys
+
 import torch
 import torch.nn as nn
 import pytest
 
 import deepspeed
 import deepspeed.utils.zero_to_fp32 as zero_to_fp32
+import deepspeed.utils.zero_to_torch as zero_to_torch
 from deepspeed.utils.zero_to_fp32 import (convert_zero_checkpoint_to_fp32_state_dict,
                                           convert_zero_checkpoint_to_state_dict, to_torch_tensor)
 from deepspeed.utils.zero_to_torch import main as zero_to_torch_main
@@ -67,6 +72,31 @@ def record_conversion(checkpoint_dir, output_dir, **kwargs):
     assert call["max_shard_size"] == "1GB"
 
 
+def test_zero_to_torch_standalone_uses_local_converter(tmp_path):
+    script_path = tmp_path / "zero_to_torch.py"
+    shutil.copyfile(zero_to_torch.__file__, script_path)
+    (tmp_path / "zero_to_fp32.py").write_text(
+        "from pathlib import Path\n"
+        "OUTPUT_DTYPE_NAMES = {'fp16': None}\n"
+        "debug = False\n"
+        "def convert_zero_checkpoint_to_state_dict(checkpoint_dir, output_dir, **kwargs):\n"
+        "    Path(output_dir).write_text(kwargs['dtype'], encoding='utf-8')\n",
+        encoding="utf-8")
+    marker_path = tmp_path / "converter.txt"
+
+    subprocess.run([
+        sys.executable,
+        str(script_path),
+        "checkpoint",
+        str(marker_path),
+        "--dtype",
+        "fp16",
+    ], check=True)
+
+    assert script_path.read_text(encoding="utf-8").startswith("#!/usr/bin/env python\n")
+    assert marker_path.read_text(encoding="utf-8") == "fp16"
+
+
 class ModelWithSharedWeights(nn.Module):
 
     def __init__(self):