chore: import upstream snapshot with attribution

swift/cli/_megatron/export.py

@@ -0,0 +1,7 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+
+if __name__ == '__main__':
+    os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1')
+    from swift.megatron import megatron_export_main
+    megatron_export_main()

swift/cli/_megatron/main.py

@@ -0,0 +1,19 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import Dict
+
+from ..main import cli_main as swift_cli_main
+
+ROUTE_MAPPING: Dict[str, str] = {
+    'pt': 'swift.cli._megatron.pt',
+    'sft': 'swift.cli._megatron.sft',
+    'rlhf': 'swift.cli._megatron.rlhf',
+    'export': 'swift.cli._megatron.export',
+}
+
+
+def cli_main():
+    return swift_cli_main(ROUTE_MAPPING, is_megatron=True)
+
+
+if __name__ == '__main__':
+    cli_main()

swift/cli/_megatron/pt.py

@@ -0,0 +1,7 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+
+if __name__ == '__main__':
+    os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1')
+    from swift.megatron import megatron_pretrain_main
+    megatron_pretrain_main()

swift/cli/_megatron/rlhf.py

@@ -0,0 +1,24 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+import sys
+
+
+def _use_ray() -> bool:
+    if '--use_ray' not in sys.argv:
+        return False
+    idx = sys.argv.index('--use_ray')
+    sys.argv.pop(idx)
+    if idx < len(sys.argv) and sys.argv[idx].lower() in ('true', 'false'):
+        val = sys.argv.pop(idx).lower() == 'true'
+        return val
+    return True
+
+
+if __name__ == '__main__':
+    if _use_ray():
+        from swift.ray.megatron.pipeline import main as ray_main
+        ray_main()
+    else:
+        os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1')
+        from swift.megatron import megatron_rlhf_main
+        megatron_rlhf_main()

swift/cli/_megatron/sft.py

@@ -0,0 +1,7 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+
+if __name__ == '__main__':
+    os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1')
+    from swift.megatron import megatron_sft_main
+    megatron_sft_main()

swift/cli/app.py

@@ -0,0 +1,5 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.pipelines import app_main
+
+if __name__ == '__main__':
+    app_main()

swift/cli/deploy.py

@@ -0,0 +1,5 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.pipelines import deploy_main
+
+if __name__ == '__main__':
+    deploy_main()

swift/cli/eval.py

@@ -0,0 +1,5 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.pipelines import eval_main
+
+if __name__ == '__main__':
+    eval_main()

swift/cli/export.py

@@ -0,0 +1,5 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.pipelines import export_main
+
+if __name__ == '__main__':
+    export_main()

swift/cli/infer.py

@@ -0,0 +1,5 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.pipelines import infer_main
+
+if __name__ == '__main__':
+    infer_main()

swift/cli/main.py

@@ -0,0 +1,106 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import importlib.util
+import json
+import os
+import subprocess
+import sys
+import yaml
+from typing import Dict, List, Optional
+
+from swift.utils import get_logger
+
+logger = get_logger()
+
+ROUTE_MAPPING: Dict[str, str] = {
+    'pt': 'swift.cli.pt',
+    'sft': 'swift.cli.sft',
+    'infer': 'swift.cli.infer',
+    'merge-lora': 'swift.cli.merge_lora',
+    'web-ui': 'swift.cli.web_ui',
+    'deploy': 'swift.cli.deploy',
+    'rollout': 'swift.cli.rollout',
+    'rlhf': 'swift.cli.rlhf',
+    'sample': 'swift.cli.sample',
+    'export': 'swift.cli.export',
+    'eval': 'swift.cli.eval',
+    'app': 'swift.cli.app',
+}
+
+
+def use_torchrun() -> bool:
+    nproc_per_node = os.getenv('NPROC_PER_NODE')
+    nnodes = os.getenv('NNODES')
+    if nproc_per_node is None and nnodes is None:
+        return False
+    return True
+
+
+def parse_yaml_args(argv):
+    if not argv:
+        return
+    config = None
+    if argv[0].endswith('.json'):
+        with open(argv[0], 'r', encoding='utf-8') as f:
+            config = json.load(f)
+    elif argv[0].endswith('.yaml') or argv[0].endswith('.yml'):
+        with open(argv[0], 'r', encoding='utf-8') as f:
+            config = yaml.safe_load(f)
+    if config is None:
+        return
+    # Used for saving configurations
+    os.environ['SWIFT_CONFIG_FILE'] = argv[0]
+
+    env = config.pop('ENV', None)
+    if env:
+        for k, v in env.items():
+            if k not in os.environ:
+                os.environ[k] = str(v)
+            elif str(v) != os.environ[k]:
+                logger.warning(f'{k} is already set in environment, using `{os.environ[k]}` instead of `{v}`')
+    config_argv = []
+    for k, v in config.items():
+        config_argv.append(f'--{k}')
+        if isinstance(v, list):
+            config_argv += v
+        else:
+            if isinstance(v, dict):
+                v = json.dumps(v, ensure_ascii=False)
+            else:
+                v = str(v)
+            config_argv.append(v)
+    argv[0:1] = config_argv
+
+
+def get_torchrun_args() -> Optional[List[str]]:
+    if not use_torchrun():
+        return
+    torchrun_args = []
+    for env_key in ['NPROC_PER_NODE', 'MASTER_PORT', 'NNODES', 'NODE_RANK', 'MASTER_ADDR']:
+        env_val = os.getenv(env_key)
+        if env_val is None:
+            continue
+        torchrun_args += [f'--{env_key.lower()}', env_val]
+    return torchrun_args
+
+
+def cli_main(route_mapping: Optional[Dict[str, str]] = None, is_megatron: bool = False) -> None:
+    route_mapping = route_mapping or ROUTE_MAPPING
+    argv = sys.argv[1:]
+    method_name = argv[0].replace('_', '-')
+    argv = argv[1:]
+    file_path = importlib.util.find_spec(route_mapping[method_name]).origin
+    parse_yaml_args(argv)
+    torchrun_args = get_torchrun_args()
+    python_cmd = sys.executable
+    if torchrun_args is None or (not is_megatron and method_name not in {'pt', 'sft', 'rlhf', 'infer'}):
+        args = [python_cmd, file_path, *argv]
+    else:
+        args = [python_cmd, '-m', 'torch.distributed.run', *torchrun_args, file_path, *argv]
+    print(f"run sh: `{' '.join(args)}`", flush=True)
+    result = subprocess.run(args)
+    if result.returncode != 0:
+        sys.exit(result.returncode)
+
+
+if __name__ == '__main__':
+    cli_main()

swift/cli/merge_lora.py

@@ -0,0 +1,15 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.arguments import ExportArguments
+from swift.pipelines import SwiftPipeline, merge_lora
+
+
+class SwiftMergeLoRA(SwiftPipeline):
+    args_class = ExportArguments
+    args: args_class
+
+    def run(self):
+        merge_lora(self.args)
+
+
+if __name__ == '__main__':
+    SwiftMergeLoRA().main()

swift/cli/pt.py

@@ -0,0 +1,7 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+if __name__ == '__main__':
+    from swift.cli.utils import try_use_single_device_mode
+    try_use_single_device_mode()
+    from swift.pipelines import pretrain_main
+    pretrain_main()

swift/cli/rlhf.py

@@ -0,0 +1,7 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+if __name__ == '__main__':
+    from swift.cli.utils import try_use_single_device_mode
+    try_use_single_device_mode()
+    from swift.pipelines import rlhf_main
+    rlhf_main()

swift/cli/rollout.py

@@ -0,0 +1,5 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.pipelines import rollout_main
+
+if __name__ == '__main__':
+    rollout_main()

swift/cli/sample.py

@@ -0,0 +1,7 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+if __name__ == '__main__':
+    from swift.ray_utils import try_init_ray
+    try_init_ray()
+    from swift.pipelines import sampling_main
+    sampling_main()

swift/cli/sft.py

@@ -0,0 +1,20 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+
+def try_init_unsloth():
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--tuner_backend', type=str, default='peft')
+    args, _ = parser.parse_known_args()
+    if args.tuner_backend == 'unsloth':
+        import unsloth
+
+
+if __name__ == '__main__':
+    from swift.cli.utils import try_use_single_device_mode
+    try_use_single_device_mode()
+    try_init_unsloth()
+    from swift.ray_utils import try_init_ray
+    try_init_ray()
+    from swift.pipelines import sft_main
+    sft_main()

swift/cli/utils.py

@@ -0,0 +1,14 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+
+
+def try_use_single_device_mode():
+    if os.environ.get('SWIFT_SINGLE_DEVICE_MODE', '0') == '1':
+        visible_devices = os.environ.get('CUDA_VISIBLE_DEVICES')
+        local_rank = os.environ.get('LOCAL_RANK')
+        if local_rank is None or not visible_devices:
+            return
+        visible_devices = visible_devices.split(',')
+        visible_device = visible_devices[int(local_rank)]
+        os.environ['CUDA_VISIBLE_DEVICES'] = str(visible_device)
+        os.environ['LOCAL_RANK'] = '0'

swift/cli/web_ui.py

@@ -0,0 +1,5 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.ui import webui_main
+
+if __name__ == '__main__':
+    webui_main()