chore: import upstream snapshot with attribution

2026-07-13 13:34:58 +08:00
commit a203934033
1368 changed files with 175001 additions and 0 deletions
@@ -0,0 +1,70 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+from dataclasses import dataclass
+from typing import Literal, Optional
+
+from swift.utils import HfConfigFactory, get_logger, to_abspath
+from .megatron_args import MegatronArguments
+from .megatron_base_args import MegatronBaseArguments
+
+logger = get_logger()
+
+
+@dataclass
+class MegatronExportArguments(MegatronBaseArguments):
+    to_mcore: bool = False
+    to_hf: bool = False
+    test_convert_precision: bool = False
+    padding_free: bool = False
+    attention_backend: str = 'unfused'
+    test_convert_dtype: str = 'float32'
+    recompute_granularity: Literal['selective', 'full', 'none'] = 'none'
+
+    exist_ok: bool = False
+    merge_lora: Optional[bool] = None
+
+    def _init_output_dir(self):
+        if self.output_dir is None:
+            ckpt_dir = self.ckpt_dir or f'./{self.model_suffix}'
+            ckpt_dir, ckpt_name = os.path.split(ckpt_dir)
+            if self.to_mcore:
+                suffix = 'mcore'
+            elif self.to_hf:
+                suffix = 'hf'
+            self.output_dir = os.path.join(ckpt_dir, f'{ckpt_name}-{suffix}')
+
+        self.output_dir = to_abspath(self.output_dir)
+        if not self.exist_ok and os.path.exists(self.output_dir):
+            raise FileExistsError(f'args.output_dir: `{self.output_dir}` already exists.')
+        logger.info(f'args.output_dir: `{self.output_dir}`')
+
+    def _init_megatron_args(self):
+        self._init_output_dir()
+        self.test_convert_dtype = HfConfigFactory.to_torch_dtype(self.test_convert_dtype)
+        extra_config = MegatronArguments.load_args_config(self.ckpt_dir)
+        extra_config['mcore_adapter'] = self.mcore_adapter
+        if self.mcore_model:
+            extra_config['mcore_model'] = self.mcore_model
+        for k, v in extra_config.items():
+            setattr(self, k, v)
+        if self.to_hf or self.to_mcore:
+            self._init_convert()
+            if self.model_info.is_moe_model is not None and self.tensor_model_parallel_size > 1:
+                self.sequence_parallel = True
+                logger.info('Settting args.sequence_parallel: True')
+            if self.merge_lora is None:
+                self.merge_lora = self.to_hf
+        super()._init_megatron_args()
+
+    def _init_convert(self):
+        convert_kwargs = {
+            'no_save_optim': True,
+            'no_save_rng': True,
+            'no_load_optim': True,
+            'no_load_rng': True,
+            'finetune': True,
+        }
+        for k, v in convert_kwargs.items():
+            setattr(self, k, v)
+        if self.model_info.is_moe_model:
+            self.moe_grouped_gemm = True