chore: import upstream snapshot with attribution

swift/megatron/__init__.py

@@ -0,0 +1,48 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+try:
+    from transformers.utils import is_torch_npu_available
+
+    if is_torch_npu_available():
+        # Enable Megatron on Ascend NPU
+        import mindspeed.megatron_adaptor  # F401
+    from .init import init_megatron_env
+    init_megatron_env()
+except Exception:
+    # allows lint pass.
+    raise
+
+from typing import TYPE_CHECKING
+
+from swift.utils.import_utils import _LazyModule
+
+if TYPE_CHECKING:
+    from .arguments import (MegatronArguments, MegatronExportArguments, MegatronPretrainArguments,
+                            MegatronRLHFArguments, MegatronSftArguments)
+    from .convert import convert_hf2mcore, convert_mcore2hf
+    from .model import get_mcore_model
+    from .pipelines import megatron_export_main, megatron_pretrain_main, megatron_rlhf_main, megatron_sft_main
+    from .trainers import MegatronDPOTrainer, MegatronTrainer
+    from .utils import initialize_megatron, prepare_mcore_model
+else:
+    _import_structure = {
+        'pipelines': ['megatron_sft_main', 'megatron_pretrain_main', 'megatron_rlhf_main', 'megatron_export_main'],
+        'convert': ['convert_hf2mcore', 'convert_mcore2hf'],
+        'utils': ['prepare_mcore_model', 'initialize_megatron'],
+        'arguments': [
+            'MegatronSftArguments', 'MegatronPretrainArguments', 'MegatronRLHFArguments', 'MegatronExportArguments',
+            'MegatronArguments'
+        ],
+        'model': ['get_mcore_model'],
+        'trainers': ['MegatronTrainer', 'MegatronDPOTrainer'],
+    }
+
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

swift/megatron/arguments/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .export_args import MegatronExportArguments
+from .megatron_args import MegatronArguments
+from .pretrain_args import MegatronPretrainArguments
+from .rlhf_args import MegatronRLHFArguments
+from .sft_args import MegatronSftArguments

swift/megatron/arguments/export_args.py

@@ -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

swift/megatron/arguments/megatron_base_args.py

@@ -0,0 +1,47 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+from dataclasses import dataclass
+
+from swift.arguments import BaseArguments
+from swift.utils import get_logger
+from .megatron_args import MegatronArguments, RLHFMegatronArgumentsMixin
+
+logger = get_logger()
+
+
+@dataclass
+class MegatronBaseArguments(MegatronArguments, BaseArguments):
+
+    # true for ray pipeline to skip distributed init
+    skip_megatron_init: bool = False
+
+    def _init_distributed(self):
+        if self.skip_megatron_init:
+            return
+        super()._init_distributed()
+
+    def _init_output_dir(self):
+        if self.skip_megatron_init:
+            if self.output_dir is None:
+                self.output_dir = f'megatron_output/{self.model_suffix}'
+            os.makedirs(self.output_dir, exist_ok=True)
+            return
+        super()._init_output_dir()
+
+    def _init_megatron_args(self):
+        MegatronArguments.__post_init__(self)
+
+    def __post_init__(self):
+        self.sequence_parallel_size = self.context_parallel_size
+        if self.packing:
+            self.padding_free = True
+        BaseArguments.__post_init__(self)
+        self.seq_length = self.packing_length or self.max_length
+        if self.skip_megatron_init:
+            RLHFMegatronArgumentsMixin.__post_init__(self)
+        else:
+            self._init_megatron_args()
+        if self.streaming:
+            if self.dataloader_num_workers > 1:
+                self.dataloader_num_workers = 1
+                logger.info('Using streaming dataset, setting args.dataloader_num_workers to 1.')

swift/megatron/arguments/pretrain_args.py

@@ -0,0 +1,9 @@
+from dataclasses import dataclass
+
+from .sft_args import MegatronSftArguments
+
+
+@dataclass
+class MegatronPretrainArguments(MegatronSftArguments):
+    use_chat_template: bool = False
+    loss_scale: str = 'all'

swift/megatron/arguments/rlhf_args.py

@@ -0,0 +1,29 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from dataclasses import dataclass
+from typing import Literal, Optional
+
+from .sft_args import MegatronSftArguments
+
+
+@dataclass
+class MegatronRLHFArguments(MegatronSftArguments):
+    rlhf_type: Literal['dpo', 'kto', 'grpo', 'gkd', 'rm'] = 'dpo'
+    loss_scale: str = 'last_round'
+    truncation_strategy: Optional[Literal['delete', 'left', 'right', 'split', None]] = None
+
+    calculate_per_token_loss: bool = False
+
+    def __post_init__(self):
+        if self.rlhf_type == 'rm':
+            self.task_type = 'seq_cls'
+            self.num_labels = 1
+        self._init_truncation_strategy()
+        super().__post_init__()
+
+    def _init_truncation_strategy(self):
+        if self.truncation_strategy is not None:
+            return
+        if self.rlhf_type == 'grpo':
+            self.truncation_strategy = 'left'
+        else:
+            self.truncation_strategy = 'delete'

swift/megatron/arguments/sft_args.py

@@ -0,0 +1,53 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import json
+import os
+from dataclasses import dataclass
+
+from swift.utils import add_version_to_work_dir, get_logger, init_process_group, is_last_rank, to_abspath
+from .megatron_base_args import MegatronBaseArguments
+
+logger = get_logger()
+
+
+@dataclass
+class MegatronSftArguments(MegatronBaseArguments):
+    add_version: bool = True
+    create_checkpoint_symlink: bool = False
+    load_args: bool = False
+
+    def _init_output_dir(self):
+        init_process_group(backend=self.ddp_backend, timeout=self.ddp_timeout)
+        if self.output_dir is None:
+            self.output_dir = f'megatron_output/{self.model_suffix}'
+        self.output_dir = to_abspath(self.output_dir)
+        if self.add_version:
+            self.output_dir = add_version_to_work_dir(self.output_dir)
+            logger.info(f'args.output_dir: {self.output_dir}')
+        os.makedirs(self.output_dir, exist_ok=True)
+
+    def _init_ckpt_dir(self, adapters=None):
+        super()._init_ckpt_dir(adapters)
+        if self.ckpt_dir and self.model is None:
+            args_path = os.path.join(self.ckpt_dir, 'args.json')
+            if not os.path.exists(args_path):
+                return
+            with open(args_path, 'r', encoding='utf-8') as f:
+                old_args = json.load(f)
+            self.model = old_args.get('model')
+
+    def _init_megatron_args(self):
+        self._init_output_dir()
+        super()._init_megatron_args()
+
+    def __post_init__(self):
+        self.mcore_model = to_abspath(self.mcore_model, check_path_exist=True)
+        super().__post_init__()
+        if len(self.dataset) == 0 and len(self.cached_dataset) == 0:
+            raise ValueError(f'self.dataset: {self.dataset}, self.cached_dataset: {self.cached_dataset}. '
+                             'Please input the training dataset.')
+        if self.tensorboard_dir is None and self.output_dir is not None:
+            self.tensorboard_dir = f'{self.output_dir}/runs'
+        self.tensorboard_dir = to_abspath(self.tensorboard_dir)
+        if self.mcore_model is None and self.model is None and not self.perform_initialization:
+            raise ValueError('You did not pass `--mcore_model/--model` to read weights, so you need to set '
+                             '`--perform_initialization true` to allow the model to initialize weights properly.')

swift/megatron/callbacks/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .base import MegatronCallback
+from .mapping import megatron_callbacks_map

swift/megatron/callbacks/base.py

@@ -0,0 +1,46 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from swift.megatron.trainers import BaseMegatronTrainer
+
+
+class MegatronCallback:
+
+    def __init__(self, trainer: 'BaseMegatronTrainer'):
+        self.trainer = trainer
+        self.args = trainer.args
+        self.state = trainer.state
+
+    def on_train_begin(self):
+        pass
+
+    def on_train_end(self):
+        pass
+
+    def on_step_begin(self):
+        pass
+
+    def on_step_end(self):
+        pass
+
+    def on_log(self, logs):
+        pass
+
+    def on_eval_begin(self):
+        pass
+
+    def on_eval_end(self):
+        pass
+
+    def on_eval_step(self):
+        pass
+
+    def on_save(self, output_dir):
+        """Called after save_checkpoint() returns.
+
+        Note: When async_save is enabled, the checkpoint may not be fully
+        written to disk yet. Use only for non-I/O-dependent logic, or
+        ensure async_save is disabled if you need to read the checkpoint.
+        """
+        pass

swift/megatron/callbacks/default_flow.py

@@ -0,0 +1,43 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import gc
+
+from .base import MegatronCallback
+
+
+class DefaultFlowCallback(MegatronCallback):
+
+    def on_train_begin(self):
+        args = self.args
+        if args.manual_gc:
+            gc.disable()
+            gc.collect()
+
+    def on_step_end(self):
+        args = self.args
+        state = self.state
+
+        state.consumed_train_samples += args.global_batch_size
+
+        if state.iteration == 1 or state.iteration % args.logging_steps == 0:
+            state.should_log = True
+        if args.eval_steps and state.iteration % args.eval_steps == 0 and args.eval_iters > 0:
+            state.should_eval = True
+        if args.save_steps and state.iteration % args.save_steps == 0:
+            state.should_save = True
+
+        if state.iteration >= args.train_iters:
+            if args.eval_iters > 0:
+                state.should_eval = True
+            state.should_save = True
+        if args.manual_gc and args.manual_gc_steps != 0 and state.iteration % args.manual_gc_steps == 0:
+            gc.collect()
+
+    def on_eval_begin(self):
+        args = self.args
+        if args.manual_gc and args.manual_gc_eval:
+            gc.collect()
+
+    def on_eval_end(self):
+        args = self.args
+        if args.manual_gc and args.manual_gc_eval:
+            gc.collect(generation=0)

swift/megatron/callbacks/mapping.py

@@ -0,0 +1,14 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .default_flow import DefaultFlowCallback
+from .print import PrintCallback
+from .swanlab import SwanlabCallback
+from .tensorboard import TensorboardCallback
+from .wandb import WandbCallback
+
+megatron_callbacks_map = {
+    'print': PrintCallback,
+    'default_flow': DefaultFlowCallback,
+    'swanlab': SwanlabCallback,
+    'wandb': WandbCallback,
+    'tensorboard': TensorboardCallback,
+}

swift/megatron/callbacks/print.py

@@ -0,0 +1,69 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+import time
+import torch
+from tqdm import tqdm
+
+from swift.megatron.utils import reduce_max_stat_across_model_parallel_group
+from swift.utils import JsonlWriter, format_time, get_logger, is_last_rank
+from .base import MegatronCallback
+
+logger = get_logger()
+
+
+class PrintCallback(MegatronCallback):
+
+    def __init__(self, trainer):
+        super().__init__(trainer)
+        self.training_bar = None
+        self.eval_bar = None
+        self.jsonl_writer = None
+        self.is_write_rank = is_last_rank()
+
+    def on_train_begin(self):
+        self.training_bar = tqdm(
+            total=self.args.train_iters, dynamic_ncols=True, disable=not self.is_write_rank, desc='Train: ')
+        self.start_step = self.state.iteration
+        self.training_bar.update(self.state.iteration)
+        self.current_step = self.state.iteration
+        self.start_time = time.time()
+        logging_path = os.path.join(self.args.output_dir, 'logging.jsonl')
+        logger.info(f'logging_path: {logging_path}')
+        self.jsonl_writer = JsonlWriter(logging_path, enable_async=True, write_on_rank='last')
+
+    def on_train_end(self):
+        self.training_bar.close()
+        self.training_bar = None
+
+    def on_step_end(self):
+        n_step = self.state.iteration - self.current_step
+        self.current_step = self.state.iteration
+        self.training_bar.update(n_step)
+
+    def on_eval_begin(self):
+        self.eval_bar = tqdm(
+            total=self.args.eval_iters, dynamic_ncols=True, disable=not self.is_write_rank, desc='Evaluate: ')
+
+    def on_eval_end(self):
+        self.eval_bar.close()
+        self.eval_bar = None
+
+    def on_eval_step(self):
+        self.eval_bar.update()
+
+    def on_log(self, logs):
+        state = self.state
+        args = self.args
+        logs['iteration'] = f'{state.iteration}/{args.train_iters}'
+        elapsed = time.time() - self.start_time
+        logs['elapsed_time'] = format_time(elapsed)
+        n_steps = state.iteration - self.start_step
+        train_speed = elapsed / n_steps if n_steps > 0 else 0.0
+        logs['remaining_time'] = format_time((args.train_iters - state.iteration) * train_speed)
+        memory = reduce_max_stat_across_model_parallel_group(torch.cuda.max_memory_reserved() / 1024**3)
+        logs['memory(GiB)'] = round(memory, 2)
+        logs['train_speed(s/it)'] = round(train_speed, 6)
+        logs = {k: round(v, 8) if isinstance(v, float) else v for k, v in logs.items()}
+        self.jsonl_writer.append(logs)
+        if self.is_write_rank:
+            self.training_bar.write(str(logs))

swift/megatron/callbacks/swanlab.py

@@ -0,0 +1,35 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+
+from swift.utils import check_json_format, is_last_rank
+from .base import MegatronCallback
+from .utils import rewrite_logs
+
+
+class SwanlabCallback(MegatronCallback):
+
+    def __init__(self, trainer):
+        super().__init__(trainer)
+        args = self.args
+        self.config = check_json_format(vars(args))
+        if args.swanlab_exp_name is None:
+            args.swanlab_exp_name = args.output_dir
+        self.save_dir = os.path.join(args.output_dir, 'swanlab')
+        self.writer = None
+        self.setup()
+
+    def setup(self):
+        import swanlab
+        args = self.args
+        if is_last_rank():
+            swanlab.init(
+                logdir=self.save_dir,
+                experiment_name=args.swanlab_exp_name,
+                project=args.swanlab_project,
+                config=self.config)
+            self.writer = swanlab
+
+    def on_log(self, logs):
+        logs = rewrite_logs(logs)
+        if is_last_rank():
+            self.writer.log(logs, step=self.state.iteration)

swift/megatron/callbacks/tensorboard.py

@@ -0,0 +1,32 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from swift.utils import check_json_format, is_last_rank
+from .base import MegatronCallback
+from .utils import rewrite_logs
+
+
+class TensorboardCallback(MegatronCallback):
+
+    def __init__(self, trainer):
+        super().__init__(trainer)
+        args = self.args
+        self.config = check_json_format(vars(args))
+        self.save_dir = args.tensorboard_dir
+        if self.save_dir is None:
+            self.save_dir = f'{args.output_dir}/runs'
+        from torch.utils.tensorboard import SummaryWriter
+        self.writer = None
+        if is_last_rank():
+            self.writer = SummaryWriter(log_dir=self.save_dir, max_queue=args.tensorboard_queue_size)
+            for k, v in self.config.items():
+                self.writer.add_text(k, str(v), global_step=self.state.iteration)
+
+    def on_log(self, logs):
+        logs = rewrite_logs(logs)
+        if self.writer:
+            for k, v in logs.items():
+                self.writer.add_scalar(k, v, self.state.iteration)
+
+    def on_train_end(self):
+        if self.writer:
+            self.writer.close()
+            self.writer = None

swift/megatron/callbacks/utils.py

@@ -0,0 +1,19 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+
+def rewrite_logs(logs):
+    new_logs = {}
+    for k, v in logs.items():
+        if isinstance(v, str):
+            continue
+        k = k.replace('/', '_')
+        if k.startswith('eval_'):
+            k = k[len('eval_'):]
+            k = f'eval/{k}'
+        elif k.startswith('test_'):
+            k = k[len('test_'):]
+            k = f'test/{k}'
+        else:
+            k = f'train/{k}'
+        new_logs[k] = v
+    return new_logs

swift/megatron/callbacks/wandb.py

@@ -0,0 +1,31 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+
+from swift.utils import check_json_format, is_last_rank
+from .base import MegatronCallback
+from .utils import rewrite_logs
+
+
+class WandbCallback(MegatronCallback):
+
+    def __init__(self, trainer):
+        super().__init__(trainer)
+        args = self.args
+        self.config = check_json_format(vars(args))
+        if args.wandb_exp_name is None:
+            args.wandb_exp_name = args.output_dir
+        self.save_dir = args.output_dir
+        self.writer = None
+        self.setup()
+
+    def setup(self):
+        import wandb
+        args = self.args
+        if is_last_rank():
+            wandb.init(dir=self.save_dir, name=args.wandb_exp_name, project=args.wandb_project, config=self.config)
+            self.writer = wandb
+
+    def on_log(self, logs):
+        logs = rewrite_logs(logs)
+        if is_last_rank():
+            self.writer.log(logs, step=self.state.iteration)

swift/megatron/convert.py

@@ -0,0 +1,120 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+import math
+import os
+import shutil
+import torch
+from transformers.utils import strtobool
+
+from swift.arguments import ExportArguments
+from swift.pipelines import prepare_model_template
+from swift.utils import get_logger, get_n_params_grads, is_master
+from .arguments import MegatronArguments
+from .model import get_mcore_model
+from .utils import (load_mcore_checkpoint, patch_torch_dist_shard, prepare_mcore_model, save_mcore_checkpoint,
+                    test_convert_precision)
+
+logger = get_logger()
+
+convert_kwargs = {
+    'use_cpu_initialization': True,
+    'no_save_optim': True,
+    'no_save_rng': True,
+    'no_load_optim': True,
+    'no_load_rng': True,
+    'finetune': True,
+    'attention_backend': 'unfused',
+    'padding_free': False,
+    'recompute_granularity': 'none',  # deepseek-v4
+}
+
+
+def convert_hf2mcore(args: ExportArguments) -> None:
+    args.experts_impl = 'eager'  # Compatible with transformers 5.4.0
+    hf_model, template = prepare_model_template(args, patch_offload=not args.test_convert_precision)
+    processor = template.processor
+    if args.thread_count is None:
+        checkpoint_size = sum(get_n_params_grads(hf_model)[0]) * torch.finfo(args.torch_dtype).bits // 8e9
+        args.thread_count = max(math.ceil(checkpoint_size / 10), 2)  # 10GB
+    patch_torch_dist_shard(args.thread_count)
+
+    hf_config = processor.model_info.config
+    current_convert_kwargs = convert_kwargs.copy()
+    if args.model_info.is_moe_model:
+        current_convert_kwargs['moe_grouped_gemm'] = True
+    megatron_args = MegatronArguments(
+        model=args.model,
+        model_type=args.model_type,
+        **current_convert_kwargs,
+        output_dir=args.output_dir,
+        torch_dtype=args.torch_dtype)
+
+    mg_model = get_mcore_model(megatron_args, hf_config)[0]
+    logger.info('Megatron model created successfully.')
+    bridge = mg_model.config.bridge
+    bridge.load_weights([mg_model], args.model_info.model_dir)
+    logger.info('Successfully transferred HF model weights to MG model.')
+    _test_convert_precision = strtobool(os.getenv('SWIFT_TEST_CONVERT_PRECISION', '0'))
+    if not _test_convert_precision:
+        args.save_args()
+        logger.info('Saving the model...')
+        save_mcore_checkpoint(megatron_args, [mg_model])
+    # Place it at the end to avoid test_convert_precision affecting precision.
+    if args.test_convert_precision:
+        test_convert_precision(megatron_args, hf_model, mg_model, template, test_convert_dtype=args.test_convert_dtype)
+
+
+def convert_mcore2hf(args: ExportArguments) -> None:
+    args.experts_impl = 'eager'
+    _, template = prepare_model_template(args, load_model=False)
+    processor = template.processor
+
+    hf_config = processor.model_info.config
+    current_convert_kwargs = convert_kwargs.copy()
+    if args.model_info.is_moe_model:
+        current_convert_kwargs['moe_grouped_gemm'] = True
+    extra_config = MegatronArguments.load_args_config(args.mcore_adapter or args.mcore_model)
+    extra_config['mcore_adapter'] = args.mcore_adapter
+    if args.mcore_model is not None:
+        extra_config['mcore_model'] = args.mcore_model
+    current_convert_kwargs.update(extra_config)
+    megatron_args = MegatronArguments(
+        model=args.model,
+        model_type=args.model_type,
+        **current_convert_kwargs,
+        output_dir=args.output_dir if args.to_mcore else None,
+        torch_dtype=args.torch_dtype)
+
+    mg_model = get_mcore_model(megatron_args, hf_config)[0]
+    if megatron_args.mcore_model is None:
+        raise ValueError('Please specify `--mcore_model`.')
+    load_mcore_checkpoint(megatron_args, [mg_model], load_arg='mcore_model')
+    if megatron_args.mcore_adapter is not None:
+        peft_model = prepare_mcore_model(megatron_args, mg_model)
+        load_mcore_checkpoint(megatron_args, [mg_model], load_arg='mcore_adapter')
+        logger.info('Merge LoRA...')
+        mg_model = peft_model.merge_and_unload()
+    logger.info('Megatron model created successfully.')
+    if args.to_hf:
+        bridge = mg_model.config.bridge
+        logger.info('Converting weights and saving the model...')
+        bridge.save_weights([mg_model], args.output_dir, args=megatron_args, processor=processor)
+        if is_master():
+            if args.ckpt_dir:
+                args_path = os.path.join(args.ckpt_dir, 'args.json')
+                shutil.copy(args_path, os.path.join(args.output_dir, 'args.json'))
+            else:
+                args.save_args(args.output_dir)
+        if args.test_convert_precision:
+            hf_model, template = prepare_model_template(args, model=args.output_dir)
+            test_convert_precision(
+                megatron_args, hf_model, mg_model, template, test_convert_dtype=args.test_convert_dtype)
+    elif args.to_mcore:
+        if args.thread_count is None:
+            checkpoint_size = sum(get_n_params_grads(mg_model)[0]) * torch.finfo(args.torch_dtype).bits // 8e9
+            args.thread_count = max(math.ceil(checkpoint_size / 10), 2)  # 10GB
+        patch_torch_dist_shard(args.thread_count)
+
+        args.save_args()
+        logger.info('Saving the model...')
+        save_mcore_checkpoint(megatron_args, [mg_model])

swift/megatron/init.py

@@ -0,0 +1,224 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import concurrent.futures
+import importlib.metadata
+import logging
+import os
+import torch
+import torch.distributed as dist
+from contextlib import contextmanager
+from copy import copy, deepcopy
+from packaging import version
+from tqdm import tqdm
+from transformers.modeling_utils import custom_object_save
+from transformers.utils import is_torch_npu_available
+from transformers.utils.versions import require_version
+
+from swift.model import get_model_processor, save_checkpoint
+from swift.utils import (HfConfigFactory, disable_safe_ddp_context_use_barrier, get_logger, get_modules_to_not_convert,
+                         get_multimodal_target_regex, is_master, split_list)
+
+logger = get_logger()
+
+
+def _patch__batched_p2p_ops():
+    from megatron.core.pipeline_parallel import p2p_communication
+
+    _batched_p2p_ops_origin = p2p_communication._batched_p2p_ops
+
+    def _batched_p2p_ops(**kwargs):
+        kwargs['group'] = None
+        return _batched_p2p_ops_origin(**kwargs)
+
+    p2p_communication._batched_p2p_ops = _batched_p2p_ops
+
+
+def _patch_torch_FileSystemReader():
+    from torch.distributed.checkpoint.filesystem import FileSystemReader
+    from torch.futures import Future
+    _origin_read_data = FileSystemReader.read_data
+    _origin__slice_file = FileSystemReader._slice_file
+    READER_MAX_WORKERS = int(os.environ.get('MCORE_READER_MAX_WORKERS', '16'))
+
+    @contextmanager
+    def _patch__slice_file(prog_bar):
+
+        def _slice_file(self, *args, **kwargs):
+            prog_bar.update()
+            return _origin__slice_file(self, *args, **kwargs)
+
+        FileSystemReader._slice_file = _slice_file
+        try:
+            yield
+        finally:
+            FileSystemReader._slice_file = _origin__slice_file
+
+    def read_data(self, plan, planner):
+
+        def _worker(plan_shard):
+            _origin_read_data(self, plan_shard, planner)
+
+        prog_bar = tqdm(total=len(plan.items), dynamic_ncols=True, desc='Loading: ')
+        plan_shards = split_list(plan.items, READER_MAX_WORKERS, contiguous=False)
+        with _patch__slice_file(prog_bar):
+            with concurrent.futures.ThreadPoolExecutor(max_workers=READER_MAX_WORKERS) as pool:
+                futures = []
+                for i in range(READER_MAX_WORKERS):
+                    plan_shard = copy(plan)
+                    plan_shard.items = plan_shards[i]
+                    futures.append(pool.submit(_worker, plan_shard))
+                concurrent.futures.wait(futures)
+        prog_bar.close()
+        fut: Future = Future()
+        fut.set_result(None)
+        return fut
+
+    FileSystemReader.read_data = read_data
+
+
+def _patch_validate_non_overlapping_shards_metadata():
+    # too slow
+    from torch.distributed._shard.sharded_tensor import api
+    from torch.distributed._shard.sharding_spec import api as api2
+    from torch.distributed.checkpoint import default_planner
+
+    def validate_non_overlapping_shards_metadata(*args, **kwargs):
+        pass
+
+    api.validate_non_overlapping_shards_metadata = validate_non_overlapping_shards_metadata
+    api2.validate_non_overlapping_shards_metadata = validate_non_overlapping_shards_metadata
+
+    def _validate_global_plan(*args, **kwargs):
+        return True
+
+    default_planner._validate_global_plan = _validate_global_plan
+
+
+def _patch_unified_memory():
+    if is_torch_npu_available():
+        return
+
+    from torch.utils import cpp_extension
+    load_inline = cpp_extension.load_inline
+
+    def _new_load_inline(*args, **kwargs):
+        name = kwargs.get('name')
+        if name == 'managed_alloc_runtime':
+            raise RuntimeError
+        return load_inline(*args, **kwargs)
+
+    # not create unified memory mempool
+    cpp_extension.load_inline = _new_load_inline
+    try:
+        from megatron.core.inference import unified_memory
+    except Exception:
+        pass
+    finally:
+        cpp_extension.load_inline = load_inline
+
+
+def _patch_mcore_bridge():
+    require_version('mcore-bridge>=1.4.0', 'please install mcore-bridge via `pip install mcore-bridge -U`')
+    import mcore_bridge
+    from mcore_bridge import GPTBridge
+    logger.info(f'mcore_bridge.__version__: {mcore_bridge.__version__}')
+    origin_save_weights = GPTBridge.save_weights
+
+    def save_weights(
+        self,
+        mg_models,
+        output_dir: str,
+        peft_format: bool = False,
+        max_shard_size: str = '5GB',
+        args=None,
+        processor=None,
+    ) -> None:
+        origin_save_weights(self, mg_models, output_dir, peft_format=peft_format, max_shard_size=max_shard_size)
+        if processor is None or args is None:
+            return
+        hf_config = self.config.hf_config
+        hf_config = deepcopy(hf_config)
+        if is_master() and not hasattr(self, 'hf_model'):
+            if hasattr(self, 'get_hf_meta_model'):
+                self.hf_model = self.get_hf_meta_model()
+                self.hf_model.model_meta = processor.model_meta
+                self.hf_model.model_info = processor.model_info
+            else:
+                with torch.device('meta'), disable_safe_ddp_context_use_barrier():
+                    self.hf_model = get_model_processor(
+                        args.model_dir, model_type=args.model_type, return_dummy_model=True)[0]
+
+        if is_master():
+            if peft_format:
+                peft_config = copy(mg_models[0].peft_config[self._adapter_name])
+                if self.config.task_type == 'seq_cls':
+                    peft_config.task_type = 'SEQ_CLS'
+                if self.is_multimodal and 'all-linear' in args.target_modules:
+                    peft_config.target_modules = get_multimodal_target_regex(
+                        self.hf_model,
+                        freeze_llm=args.freeze_llm,
+                        freeze_vit=args.freeze_vit,
+                        freeze_aligner=args.freeze_aligner,
+                        include_embedding='all-embedding' in args.target_modules,
+                        exclude_router='all-router' not in args.target_modules)
+                else:
+                    assert not isinstance(peft_config.target_modules, str), (
+                        'target_regex is not currently supported for LoRA conversion. Please set `--merge_lora true`.')
+                    peft_config.target_modules = self._peft_target_modules
+                peft_config.modules_to_save = self._peft_modules_to_save
+                peft_config.save_pretrained(output_dir)
+            else:
+                config = self.config
+                llm_config = HfConfigFactory.get_text_config(hf_config)
+                if config.mtp_num_layers:
+                    for key in ['num_nextn_predict_layers', 'mtp_num_hidden_layers']:
+                        if hasattr(llm_config, key):
+                            setattr(llm_config, key, config.mtp_num_layers)
+                            break
+                    else:
+                        llm_config.num_nextn_predict_layers = config.mtp_num_layers
+                HfConfigFactory.del_config_attr(hf_config, 'quantization_config')
+                expert_dtype = None
+                if config.fp8 is not None and config.fp8_recipe == 'blockwise' and config.fp8_param:
+                    from transformers.utils.quantization_config import FineGrainedFP8Config
+                    modules_to_not_convert = get_modules_to_not_convert(self.hf_model)
+                    if hasattr(self, '_fp8_skip_modules'):
+                        modules_to_not_convert = (modules_to_not_convert or []) + list(self._fp8_skip_modules)
+                    hf_config.quantization_config = FineGrainedFP8Config(modules_to_not_convert=modules_to_not_convert)
+                    expert_dtype = 'fp8'
+                if args.model_type == 'deepseek_v4':
+                    HfConfigFactory.set_config_attr(hf_config, 'expert_dtype', expert_dtype)
+                hf_config.save_pretrained(output_dir)
+                if getattr(self.hf_model, '_auto_class') is not None:
+                    try:
+                        custom_object_save(self.hf_model, output_dir, config=hf_config)
+                    except FileNotFoundError as e:
+                        logger.error(f'custom_object_save Error: {e}')
+                save_checkpoint(
+                    None,
+                    processor,
+                    output_dir,
+                    model_dirs=[args.model_dir],
+                    additional_saved_files=self.hf_model.model_meta.additional_saved_files)
+            logger.info(f'Successfully saved `safetensors` model weights in `{output_dir}`.')
+        dist.barrier()  # Ensure all weights are saved completely
+
+    GPTBridge.save_weights = save_weights
+
+
+def init_megatron_env():
+    os.environ.pop('VLLM_USE_MODELSCOPE', None)
+    logging_level = logging.root.level
+    _patch_unified_memory()
+    _patch__batched_p2p_ops()
+    logging.root.setLevel(logging_level)  # revert logger level
+    try:
+        _patch_torch_FileSystemReader()
+    except Exception:
+        logger.warning('Failed to patch FileSystemReader.')
+    try:
+        _patch_validate_non_overlapping_shards_metadata()
+    except Exception:
+        logger.warning('Patch validate_non_overlapping_shards_metadata failed.')
+        pass
+    import megatron.core
+    logger.info(f'megatron.core.__version__: {megatron.core.__version__}')

swift/megatron/model/__init__.py

@@ -0,0 +1 @@
+from .utils import get_mcore_model

swift/megatron/model/utils.py

@@ -0,0 +1,295 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from dataclasses import fields
+from mcore_bridge import ModelConfig
+from mcore_bridge import get_mcore_model as _get_mcore_model
+from mcore_bridge import hf_to_mcore_config
+from transformers.utils import is_torch_npu_available
+from typing import Any, Generator, Optional, Tuple
+
+from swift.utils import get_logger
+
+logger = get_logger()
+
+
+def _check_attention_backend(args, config):
+    """Validate attention backend compatibility with configuration."""
+    attention_backend = config.attention_backend.name
+    if attention_backend == 'flash' and config.softmax_type == 'learnable':
+        raise ValueError(f'Attention backend "{attention_backend}" does not support learnable softmax_type.')
+
+
+def _check_padding_free(args, config):
+    """Validate and adjust padding_free setting based on configuration constraints."""
+    if not args.padding_free:
+        return
+
+    attention_backend = config.attention_backend.name
+    message = None
+
+    if config.experimental_attention_variant == 'dsa':
+        message = 'DSA is not supported in padding-free mode'
+    elif attention_backend == 'unfused':
+        message = f'Attention backend "{attention_backend}" is not supported in padding-free mode'
+
+    if message:
+        logger.warning(f'{message}. Setting args.padding_free to False.')
+        args.padding_free = False
+
+
+def get_mcore_model_config(args, hf_config):
+    kwargs = hf_to_mcore_config(hf_config)
+    kwargs['mcore_model_type'] = args.megatron_model_meta.model_type
+    kwargs['hf_config'] = hf_config
+    for f in fields(ModelConfig):
+        key, value = f.name, getattr(args, f.name, None)
+        if value is None or isinstance(value, (list, tuple)) and len(value) == 0:
+            continue
+        kwargs[key] = value
+
+    if args.task_type == 'seq_cls':
+        args.problem_type = args.problem_type or getattr(hf_config, 'problem_type', None)
+        logger.info(f'args.problem_type: {args.problem_type}')
+
+    kwargs['params_dtype'] = args.torch_dtype
+    kwargs['num_layers_in_first_pipeline_stage'] = args.decoder_first_pipeline_num_layers
+    kwargs['num_layers_in_last_pipeline_stage'] = args.decoder_last_pipeline_num_layers
+    kwargs['fp4_param'] = args.fp4_param_gather
+    kwargs['fp8_param'] = args.fp8_param_gather
+    swiglu = kwargs.get('swiglu', True)
+    add_bias_linear = kwargs.get('add_bias_linear', False)
+    num_moe_experts = kwargs.get('num_moe_experts', None)
+    position_embedding_type = kwargs.get('position_embedding_type', 'rope')
+    if position_embedding_type != 'rope':
+        kwargs['apply_rope_fusion'] = False
+    if not swiglu and not add_bias_linear:
+        kwargs['bias_activation_fusion'] = False
+    if add_bias_linear and num_moe_experts and args.moe_grouped_gemm:
+        kwargs['bias_dropout_fusion'] = False
+    if num_moe_experts is None:
+        kwargs['expert_model_parallel_size'] = 1
+        kwargs['expert_tensor_parallel_size'] = 1
+
+    if args.router_replay_mode != 'disabled':
+        kwargs['moe_enable_routing_replay'] = True
+    if args.megatron_extra_kwargs:
+        kwargs.update(args.megatron_extra_kwargs)
+    config = ModelConfig(**kwargs)
+    if is_torch_npu_available() and getattr(args, 'attention_backend', 'flash') != 'local':
+        setattr(config, 'use_flash_attn', True)
+    _check_attention_backend(args, config)
+    _check_padding_free(args, config)
+    return config
+
+
+class MegatronBridgeBackend:
+    """Adapter for NVIDIA ``megatron.bridge.AutoBridge``.
+
+    Limitations:
+    - LoRA / PEFT loading is not yet supported.
+    - MLLM is not yet supported
+    - FP8 export is not yet supported.
+    """
+
+    def __init__(self, auto_bridge: Any, hf_config: Optional[Any] = None):
+        self._bridge = auto_bridge
+        self._hf_config = hf_config
+
+    @classmethod
+    def from_hf_config(cls, hf_config) -> 'MegatronBridgeBackend':
+        from megatron.bridge.models.conversion.auto_bridge import AutoBridge
+        return cls(AutoBridge.from_hf_config(hf_config), hf_config)
+
+    def load_weights(self, models, hf_model_dir, peft_format=False, adapter_name='default', converter=None):
+        if peft_format:
+            raise NotImplementedError('LoRA loading via megatron-bridge backend is not yet supported. '
+                                      'Please use bridge_backend="mcore-bridge" for LoRA training.')
+        if converter is not None:
+            logger.warning('converter is not supported by megatron-bridge backend, ignoring.')
+        self._bridge.load_hf_weights(models, hf_path=hf_model_dir)
+
+    def export_weights(self,
+                       models,
+                       target_device=None,
+                       only_master_rank=False,
+                       peft_format=False,
+                       adapter_name='default',
+                       converter=None,
+                       tqdm_desc='Exporting: ',
+                       disable_tqdm=True,
+                       _is_saving=False) -> Generator[Tuple[str, 'torch.Tensor'], None, None]:
+        if peft_format:
+            raise NotImplementedError('LoRA export via megatron-bridge backend is not yet supported. '
+                                      'Please use bridge_backend="mcore-bridge" for LoRA training.')
+        cpu = (target_device == 'cpu')
+        for weight_tuple in self._bridge.export_hf_weights(models, cpu=cpu):
+            key = weight_tuple.param_name
+            tensor = weight_tuple.weight
+            if converter is not None and tensor is not None:
+                kv = converter(key, tensor)
+                if kv is None:
+                    continue
+                key, tensor = kv
+            yield key, tensor
+
+    def save_weights(self,
+                     models,
+                     output_dir,
+                     peft_format=False,
+                     max_shard_size='5GB',
+                     args=None,
+                     processor=None) -> None:
+        if peft_format:
+            raise NotImplementedError('LoRA saving via megatron-bridge backend is not yet supported. '
+                                      'Please use bridge_backend="mcore-bridge" for LoRA training.')
+
+        # 1. Save weights via megatron-bridge (safetensors format)
+        self._bridge.save_hf_weights(models, path=output_dir)
+
+        # 2. Save HF config and tokenizer on rank 0.
+        # We use the original HF config (not the one reconstructed by megatron-bridge)
+        # because the bridge's config-only path may drop fields like num_attention_heads.
+        import torch.distributed as dist
+        is_master = (not dist.is_initialized()) or dist.get_rank() == 0
+        if is_master and args is not None and self._hf_config is not None:
+            from copy import deepcopy
+
+            from swift.model import save_checkpoint
+            from swift.utils import HfConfigFactory
+            hf_config = deepcopy(self._hf_config)
+            llm_config = HfConfigFactory.get_text_config(hf_config)
+
+            # MTP: write back num_nextn_predict_layers
+            mtp_num_layers = getattr(args, 'mtp_num_layers', None)
+            if mtp_num_layers:
+                for key in ['num_nextn_predict_layers', 'mtp_num_hidden_layers']:
+                    if hasattr(llm_config, key):
+                        setattr(llm_config, key, mtp_num_layers)
+                        break
+                else:
+                    llm_config.num_nextn_predict_layers = mtp_num_layers
+
+            HfConfigFactory.del_config_attr(hf_config, 'quantization_config')
+
+            # FP8: write back quantization_config
+            expert_dtype = None
+            fp8_format = getattr(args, 'fp8_format', None)
+            fp8_recipe = getattr(args, 'fp8_recipe', 'delayed')
+            fp8_param = getattr(args, 'fp8_param_gather', False)
+            if fp8_format is not None and fp8_recipe == 'blockwise' and fp8_param:
+                from transformers.utils.quantization_config import FineGrainedFP8Config
+                hf_config.quantization_config = FineGrainedFP8Config()
+                expert_dtype = 'fp8'
+            if getattr(args, 'model_type', None) == 'deepseek_v4':
+                HfConfigFactory.set_config_attr(hf_config, 'expert_dtype', expert_dtype)
+
+            hf_config.save_pretrained(output_dir)
+            if processor is not None:
+                additional_saved_files = getattr(getattr(processor, 'model_meta', None), 'additional_saved_files', None)
+                save_checkpoint(
+                    None,
+                    processor,
+                    output_dir,
+                    model_dirs=[args.model_dir],
+                    additional_saved_files=additional_saved_files)
+        if dist.is_initialized():
+            dist.barrier()
+
+
+def get_mcore_model(args, hf_config):
+    bridge_backend = args.bridge_backend
+    if bridge_backend == 'megatron-bridge':
+        return _get_megatron_bridge_model(args, hf_config)
+    config = get_mcore_model_config(args, hf_config)
+    models = _get_mcore_model(config)
+
+    return models
+
+
+def _get_megatron_bridge_model(args, hf_config):
+    import dataclasses
+
+    backend = MegatronBridgeBackend.from_hf_config(hf_config)
+    auto_bridge = backend._bridge
+
+    # Validate model support via AutoBridge.supports()
+    from megatron.bridge.models.conversion.auto_bridge import AutoBridge
+    if not AutoBridge.supports(hf_config):
+        raise ValueError(f'Model {getattr(hf_config, "model_type", "unknown")} is not supported by '
+                         f'megatron-bridge. Please use bridge_backend="mcore-bridge" or check '
+                         f'AutoBridge.list_supported_models() for supported architectures.')
+
+    # --- Step 1: Get provider (GPTModelProvider, which extends TransformerConfig) ---
+    provider = auto_bridge.to_megatron_provider(load_weights=False)
+
+    # --- Step 2: Build overrides from args ---
+    # Auto-match: iterate over provider's dataclass fields and pick up matching args fields.
+    # This mirrors mcore-bridge's get_mcore_model_config which does:
+    #   for f in fields(ModelConfig): kwargs[f.name] = getattr(args, f.name, None)
+    overrides = {}
+    provider_fields = {f.name for f in dataclasses.fields(provider)}
+    for field_name in provider_fields:
+        value = getattr(args, field_name, None)
+        if value is None or (isinstance(value, (list, tuple)) and len(value) == 0):
+            continue
+        overrides[field_name] = value
+
+    # Explicit field name mappings (args name → provider field name)
+    explicit_mappings = {
+        'decoder_first_pipeline_num_layers': 'num_layers_in_first_pipeline_stage',
+        'decoder_last_pipeline_num_layers': 'num_layers_in_last_pipeline_stage',
+    }
+    for args_key, provider_key in explicit_mappings.items():
+        value = getattr(args, args_key, None)
+        if value is not None and provider_key in provider_fields:
+            overrides[provider_key] = value
+
+    # dtype
+    dtype = getattr(args, 'torch_dtype', None)
+
+    # MoE: if no experts, force EP/ETP to 1
+    # num_moe_experts comes from HF config (parsed by AutoBridge into provider),
+    # not from args — so check provider too.
+    num_moe_experts = overrides.get('num_moe_experts') or getattr(provider, 'num_moe_experts', None)
+    if num_moe_experts is None:
+        overrides['expert_model_parallel_size'] = 1
+        overrides['expert_tensor_parallel_size'] = 1
+
+    # Router replay
+    if getattr(args, 'router_replay_mode', 'disabled') != 'disabled':
+        if 'moe_enable_routing_replay' in provider_fields:
+            overrides['moe_enable_routing_replay'] = True
+
+    # megatron_extra_kwargs (user-specified raw overrides)
+    if getattr(args, 'megatron_extra_kwargs', None):
+        overrides.update(args.megatron_extra_kwargs)
+
+    # padding_free requires variable_seq_lengths=True so that RotaryEmbedding
+    # generates freqs matching the actual packed sequence length (cu_seqlens[-1])
+    # instead of the fixed seq_length. Without this, mcore-bridge's patcher
+    # use_batched_rope check fails and falls back to the original
+    # _apply_rotary_pos_emb_thd which calls torch.split on a padded tensor.
+    if getattr(args, 'padding_free', False) and 'variable_seq_lengths' in provider_fields:
+        overrides['variable_seq_lengths'] = True
+
+    # --- Step 3: Apply overrides and finalize ---
+    provider.apply_overrides_and_finalize(dtype=dtype, overrides=overrides)
+
+    import torch.nn.functional as F
+    provider.swiglu = (provider.gated_linear_unit and provider.activation_func is F.silu)
+
+    # --- Step 4: Create raw models (no DDP/Float16 wrapping) ---
+    # swift's wrap_model handles DDP/Float16 wrapping with the correct DDP config from args.
+    models = provider.provide_distributed_model(
+        wrap_with_ddp=False,
+        mixed_precision_wrapper=None,
+        use_cpu_initialization=getattr(args, 'use_cpu_initialization', False),
+    )
+    if not isinstance(models, list):
+        models = [models]
+
+    # --- Step 5: Attach backend to model.config.bridge ---
+    for model in models:
+        model.config.bridge = backend
+
+    logger.info('Created Megatron model via megatron-bridge backend')
+    return models

swift/megatron/pipelines/__init__.py

@@ -0,0 +1,22 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import TYPE_CHECKING
+
+from swift.utils.import_utils import _LazyModule
+
+if TYPE_CHECKING:
+    from .export import megatron_export_main
+    from .train import megatron_pretrain_main, megatron_rlhf_main, megatron_sft_main
+else:
+    _import_structure = {
+        'train': ['megatron_pretrain_main', 'megatron_rlhf_main', 'megatron_sft_main'],
+        'export': ['megatron_export_main'],
+    }
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

swift/megatron/pipelines/export/__init__.py

@@ -0,0 +1 @@
+from .export import megatron_export_main

swift/megatron/pipelines/export/export.py

@@ -0,0 +1,132 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+import shutil
+import torch.distributed as dist
+from transformers.utils import strtobool
+from typing import List, Optional, Union
+
+from swift.megatron.arguments import MegatronExportArguments
+from swift.megatron.convert import test_convert_precision
+from swift.megatron.model import get_mcore_model
+from swift.megatron.utils import load_mcore_checkpoint, prepare_mcore_model, save_mcore_checkpoint
+from swift.pipelines import SwiftPipeline, prepare_model_template
+from swift.utils import disable_safe_ddp_context_use_barrier, get_logger, is_master
+
+logger = get_logger()
+
+
+class MegatronExport(SwiftPipeline):
+    args_class = MegatronExportArguments
+    args: args_class
+
+    def run(self):
+        os.environ['DISABLE_MP_DDP'] = 'true'
+        args = self.args
+        if args.to_hf:
+            self.convert_mcore2hf()
+        elif args.to_mcore:
+            self.convert_hf2mcore()
+
+    def convert_mcore2hf(self) -> None:
+        args = self.args
+        args.experts_impl = 'eager'
+        download_model = args.model is not None
+        _, template = prepare_model_template(args, load_model=False, download_model=download_model)
+        self.processor = template.processor
+        hf_config = self.processor.model_info.config
+        mg_model = get_mcore_model(args, hf_config)[0]
+        logger.info('Megatron model created successfully.')
+        bridge = mg_model.config.bridge
+        if args.mcore_model is not None:
+            load_mcore_checkpoint(args, [mg_model], load_arg='mcore_model')
+        elif args.model is not None:
+            bridge.load_weights([mg_model], args.model_info.model_dir)
+        else:
+            raise ValueError('Please specify `--mcore_model` or `--model`.')
+        if args.adapters or args.mcore_adapter is not None:
+            peft_model = prepare_mcore_model(args, mg_model)
+            if args.mcore_adapter is not None:
+                load_mcore_checkpoint(args, [mg_model], load_arg='mcore_adapter')
+            elif args.adapters:
+                assert len(args.adapters) == 1, 'Currently only support one adapter'
+                bridge.load_weights([mg_model], args.adapters[0], peft_format=True)
+            if args.merge_lora:
+                logger.info('Merge LoRA...')
+                mg_model = peft_model.merge_and_unload()
+        logger.info('Converting weights and saving the model...')
+        save_peft_format = args.tuner_type == 'lora' and not args.merge_lora
+        bridge.save_weights(
+            [mg_model],
+            args.output_dir,
+            peft_format=save_peft_format,
+            args=args,
+            processor=self.processor,
+        )
+        if is_master():
+            if args.ckpt_dir:
+                args_path = os.path.join(args.ckpt_dir, 'args.json')
+                shutil.copy(args_path, os.path.join(args.output_dir, 'args.json'))
+            else:
+                args.save_args(args.output_dir)
+        if args.test_convert_precision:
+            with disable_safe_ddp_context_use_barrier():
+                if save_peft_format:
+                    kwargs = {'adapters': [args.output_dir]}
+                else:
+                    kwargs = {'model': args.output_dir, 'torch_dtype': None, 'adapters': []}
+                device_map = args.device_map or 'auto'
+                hf_model, template = prepare_model_template(
+                    args, device_map=device_map, **kwargs) if is_master() else (None, template)
+            test_convert_precision(args, hf_model, mg_model, template, test_convert_dtype=args.test_convert_dtype)
+            dist.barrier()
+
+    def convert_hf2mcore(self) -> None:
+        args = self.args
+        args.experts_impl = 'eager'
+        download_model = args.model is not None
+        _, template = prepare_model_template(args, load_model=False, download_model=download_model)
+        self.processor = template.processor
+        hf_config = self.processor.model_info.config
+        mg_model = get_mcore_model(args, hf_config)[0]
+        logger.info('Megatron model created successfully.')
+        bridge = mg_model.config.bridge
+        if args.model is not None:
+            bridge.load_weights([mg_model], args.model_info.model_dir)
+        elif args.mcore_model is not None:
+            load_mcore_checkpoint(args, [mg_model], load_arg='mcore_model')
+        else:
+            raise ValueError('Please specify `--mcore_model` or `--model`.')
+        dist.barrier()
+        if args.adapters or args.mcore_adapter is not None:
+            peft_model = prepare_mcore_model(args, mg_model)
+            if args.adapters:
+                assert len(args.adapters) == 1, 'Currently only support one adapter'
+                bridge.load_weights([mg_model], args.adapters[0], peft_format=True)
+            elif args.mcore_adapter is not None:
+                load_mcore_checkpoint(args, [mg_model], load_arg='mcore_adapter')
+            if args.merge_lora:
+                logger.info('Merge LoRA...')
+                mg_model = peft_model.merge_and_unload()
+        logger.info('Successfully transferred HF model weights to MG model.')
+        _test_convert_precision = strtobool(os.getenv('SWIFT_TEST_CONVERT_PRECISION', '0'))
+        if not _test_convert_precision:
+            args.save_args(args.output_dir)
+            logger.info('Saving the model...')
+            save_peft_format = args.tuner_type == 'lora' and not args.merge_lora
+            save_mcore_checkpoint(args, [mg_model], peft_format=save_peft_format)
+        # hf_model does not support loading args.mcore_adapter, so test_convert_precision cannot be performed
+        support_convert_precision = args.mcore_adapter is None
+        if args.test_convert_precision:
+            if support_convert_precision:
+                with disable_safe_ddp_context_use_barrier():
+                    device_map = args.device_map or 'auto'
+                    hf_model, template = prepare_model_template(
+                        args, device_map=device_map) if is_master() else (None, template)
+                test_convert_precision(args, hf_model, mg_model, template, test_convert_dtype=args.test_convert_dtype)
+                dist.barrier()
+            else:
+                logger.warning('Skip test_convert_precision because `--mcore_adapter` is specified.')
+
+
+def megatron_export_main(args: Optional[Union[List[str], MegatronExportArguments]] = None):
+    return MegatronExport(args).main()

swift/megatron/pipelines/train/__init__.py

@@ -0,0 +1,24 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import TYPE_CHECKING
+
+from swift.utils.import_utils import _LazyModule
+
+if TYPE_CHECKING:
+    from .pretrain import megatron_pretrain_main
+    from .rlhf import megatron_rlhf_main
+    from .sft import megatron_sft_main
+else:
+    _import_structure = {
+        'pretrain': ['megatron_pretrain_main'],
+        'rlhf': ['megatron_rlhf_main'],
+        'sft': ['megatron_sft_main'],
+    }
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

swift/megatron/pipelines/train/pretrain.py

@@ -0,0 +1,17 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import List, Optional, Union
+
+from swift.megatron.arguments import MegatronPretrainArguments
+from swift.utils import get_logger
+from .sft import MegatronSft
+
+logger = get_logger()
+
+
+class MegatronPretrain(MegatronSft):
+    args_class = MegatronPretrainArguments
+    args: args_class
+
+
+def megatron_pretrain_main(args: Optional[Union[List[str], MegatronPretrainArguments]] = None):
+    return MegatronPretrain(args).main()

swift/megatron/pipelines/train/rlhf.py

@@ -0,0 +1,73 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import importlib
+from typing import List, Optional, Union
+
+from swift.megatron.arguments import MegatronRLHFArguments
+from swift.pipelines.train import prepare_kto_dataset
+from swift.utils import get_current_device, get_logger, is_last_rank
+from .sft import MegatronSft
+
+logger = get_logger()
+
+
+class MegatronRLHF(MegatronSft):
+    args_class = MegatronRLHFArguments
+    args: args_class
+
+    def prepare_trainer(self):
+        args = self.args
+        trainer_mapping = {
+            'dpo': 'MegatronDPOTrainer',
+            'gkd': 'MegatronGKDTrainer',
+            'grpo': 'MegatronGRPOTrainer',
+            'kto': 'MegatronKTOTrainer',
+            'rm': 'MegatronRewardTrainer'
+        }
+        module = importlib.import_module('swift.megatron.trainers')
+        trainer_cls_name = trainer_mapping.get(args.rlhf_type)
+        if trainer_cls_name is None:
+            raise ValueError(f'The current Megatron-SWIFT does not support rlhf_type: {args.rlhf_type}.')
+        trainer_cls = getattr(module, trainer_cls_name)
+        kwargs = {}
+        if args.rlhf_type in ('grpo', 'gkd'):
+            kwargs['vllm_client'] = self._prepare_vllm_client()
+        return trainer_cls(args, self.template, **kwargs)
+
+    def _prepare_template(self) -> None:
+        super()._prepare_template()
+        model_mapping = {'grpo': 'train', 'gkd': 'train', 'kto': 'kto'}
+        self.template.set_mode(model_mapping.get(self.args.rlhf_type, 'rlhf'))
+
+    def _get_dataset(self):
+        args = self.args
+        train_dataset, val_dataset = super()._get_dataset()
+        if args.rlhf_type == 'kto':
+            train_dataset, val_dataset = prepare_kto_dataset(args, train_dataset, val_dataset)
+        return train_dataset, val_dataset
+
+    def _prepare_vllm_client(self):
+        # Only prepare vLLM client for server mode
+        if self.args.rlhf_type not in ('grpo', 'gkd') or self.args.vllm_mode != 'server':
+            return None
+        # GKD may not use vLLM (off-policy mode)
+        if not getattr(self.args, 'use_vllm', False):
+            return None
+
+        from swift.rlhf_trainers.vllm_client import VLLMClient
+        vllm_client = None
+        if is_last_rank():
+            logger.info('Start connecting to vLLM server')
+            vllm_client = VLLMClient(
+                base_urls=self.args.vllm_server_base_url,
+                hosts=self.args.vllm_server_host,
+                server_ports=self.args.vllm_server_port,
+                group_ports=self.args.vllm_server_group_port,
+                connection_timeout=self.args.vllm_server_timeout)
+            vllm_client.close_communicator()
+            vllm_client.init_communicator(device=get_current_device())
+            logger.info('Connected to vLLM server')
+        return vllm_client
+
+
+def megatron_rlhf_main(args: Optional[Union[List[str], MegatronRLHFArguments]] = None):
+    return MegatronRLHF(args).main()

swift/megatron/pipelines/train/sft.py

@@ -0,0 +1,97 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+import torch
+import torch.distributed as dist
+from dataclasses import asdict
+from transformers.utils import is_torch_npu_available
+from typing import List, Optional, Union
+
+from swift.megatron.arguments import MegatronSftArguments
+from swift.megatron.trainers import MegatronEmbeddingTrainer, MegatronRerankerTrainer, MegatronTrainer
+from swift.pipelines import SwiftSft
+from swift.utils import append_to_jsonl, get_logger, is_last_rank, plot_images
+
+if is_torch_npu_available():
+    # Enable Megatron on Ascend NPU
+    from mindspeed.megatron_adaptor import repatch
+
+    from swift.model.npu_patcher import patch_mindspeed_te_cp_implementation
+else:
+    repatch = None
+    patch_mindspeed_te_cp_implementation = None
+
+logger = get_logger()
+
+
+class MegatronSft(SwiftSft):
+    args_class = MegatronSftArguments
+    args: args_class
+
+    def prepare_trainer(self):
+        args = self.args
+        if args.task_type == 'embedding':
+            return MegatronEmbeddingTrainer(self.args, self.template)
+        elif args.task_type in {'reranker', 'generative_reranker'}:
+            return MegatronRerankerTrainer(self.args, self.template)
+        else:
+            return MegatronTrainer(self.args, self.template)
+
+    def _set_seed(self):
+        pass
+
+    def __init__(self, args: Optional[Union[List[str], MegatronSftArguments]] = None) -> None:
+        self.train_msg = {}
+        super(SwiftSft, self).__init__(args)
+        args = self.args
+        if repatch is not None:
+            megatron_args = asdict(self.args)
+            if args.attention_backend != 'local':
+                # MindSpeed requires passing `use_flash_attn` to Megatron
+                # to enable flash attention on Ascend NPU.
+                args.use_flash_attn = True
+                megatron_args['use_flash_attn'] = True
+            patch_mindspeed_te_cp_implementation(megatron_args)
+            repatch(megatron_args)
+        template_cls = args.template_meta.template_cls
+        if args.model_meta.is_multimodal and template_cls and template_cls.use_model:
+            kwargs = {'return_dummy_model': True}
+        else:
+            kwargs = {'load_model': False}
+        with torch.device('meta'):
+            self.model, self.processor = args.get_model_processor(**kwargs, download_model=args.mcore_model is None)
+        self._prepare_template()
+        args.save_args(args.output_dir)
+        self.template.use_megatron = True
+
+    def run(self):
+        args = self.args
+        train_dataset, val_dataset = self._prepare_dataset()
+        args.init_iters(train_dataset, val_dataset)
+        trainer = self.prepare_trainer()
+        try:
+            trainer.train(train_dataset, val_dataset)
+        finally:
+            state = trainer.state
+            self._handle_trainer_state(trainer, is_last_rank())
+            self.train_msg.update({
+                'last_model_checkpoint': state.last_model_checkpoint,
+                'best_model_checkpoint': state.best_model_checkpoint,
+                'best_metric': state.best_metric,
+            })
+            # Visualization
+            if is_last_rank():
+                images_dir = os.path.join(args.output_dir, 'images')
+                logger.info(f'images_dir: {images_dir}')
+                plot_images(images_dir, args.tensorboard_dir)
+
+                jsonl_path = os.path.join(args.output_dir, 'logging.jsonl')
+                append_to_jsonl(jsonl_path, self.train_msg, strict=False, write_on_rank='last')
+        # Exceptions may cause the process to hang, preventing the exception from being propagated.
+        # Therefore, destroy_process_group() should not be placed inside the finally block.
+        if dist.is_initialized():
+            dist.destroy_process_group()
+        return self.train_msg
+
+
+def megatron_sft_main(args: Optional[Union[List[str], MegatronSftArguments]] = None):
+    return MegatronSft(args).main()

swift/megatron/trainers/__init__.py

@@ -0,0 +1,38 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import TYPE_CHECKING
+
+from swift.utils.import_utils import _LazyModule
+
+if TYPE_CHECKING:
+    from .base import BaseMegatronTrainer
+    from .dpo_trainer import MegatronDPOTrainer
+    from .embedding_trainer import MegatronEmbeddingTrainer
+    from .gkd_trainer import MegatronGKDTrainer
+    from .grpo_trainer import MegatronGRPOTrainer
+    from .kto_trainer import MegatronKTOTrainer
+    from .reranker_trainer import MegatronRerankerTrainer
+    from .reward_trainer import MegatronRewardTrainer
+    from .rollout_mixin import MegatronRolloutMixin
+    from .trainer import MegatronTrainer
+else:
+    _import_structure = {
+        'dpo_trainer': ['MegatronDPOTrainer'],
+        'gkd_trainer': ['MegatronGKDTrainer'],
+        'grpo_trainer': ['MegatronGRPOTrainer'],
+        'kto_trainer': ['MegatronKTOTrainer'],
+        'reward_trainer': ['MegatronRewardTrainer'],
+        'rollout_mixin': ['MegatronRolloutMixin'],
+        'embedding_trainer': ['MegatronEmbeddingTrainer'],
+        'reranker_trainer': ['MegatronRerankerTrainer'],
+        'trainer': ['MegatronTrainer'],
+        'base': ['BaseMegatronTrainer'],
+    }
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()['__file__'],
+        _import_structure,
+        module_spec=__spec__,
+        extra_objects={},
+    )

swift/megatron/trainers/batch_sampler.py

@@ -0,0 +1,155 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+
+from swift.utils import get_logger
+
+logger = get_logger()
+
+
+# Code borrowed from megatron-lm
+class MegatronPretrainingSampler:
+
+    def __init__(self,
+                 total_samples,
+                 consumed_samples,
+                 micro_batch_size,
+                 data_parallel_rank,
+                 data_parallel_size,
+                 drop_last=True):
+        # Keep a copy of input params for later use.
+        self.total_samples = total_samples
+        self.consumed_samples = consumed_samples
+        self.micro_batch_size = micro_batch_size
+        self.data_parallel_rank = data_parallel_rank
+        self.micro_batch_times_data_parallel_size = \
+            self.micro_batch_size * data_parallel_size
+        self.drop_last = drop_last
+
+        # Sanity checks.
+        assert self.total_samples > 0, \
+            'no sample to consume: {}'.format(self.total_samples)
+        assert self.consumed_samples < self.total_samples, \
+            'no samples left to consume: {}, {}'.format(self.consumed_samples,
+                                                        self.total_samples)
+        assert self.micro_batch_size > 0
+        assert data_parallel_size > 0
+        assert self.data_parallel_rank < data_parallel_size, \
+            'data_parallel_rank should be smaller than data size: {}, ' \
+            '{}'.format(self.data_parallel_rank, data_parallel_size)
+
+    def __len__(self):
+        return self.total_samples
+
+    def get_start_end_idx(self):
+        start_idx = self.data_parallel_rank * self.micro_batch_size
+        end_idx = start_idx + self.micro_batch_size
+        return start_idx, end_idx
+
+    def __iter__(self):
+        batch = []
+        # Last batch will be dropped if drop_last is not set False
+        for idx in range(self.consumed_samples, self.total_samples):
+            batch.append(idx)
+            if len(batch) == self.micro_batch_times_data_parallel_size:
+                start_idx, end_idx = self.get_start_end_idx()
+                yield batch[start_idx:end_idx]
+                batch = []
+
+        # Check the last partial batch and see drop_last is set
+        if len(batch) > 0 and not self.drop_last:
+            start_idx, end_idx = self.get_start_end_idx()
+            yield batch[start_idx:end_idx]
+
+
+# Code borrowed from megatron-lm
+class MegatronPretrainingRandomSampler:
+
+    def __init__(
+        self,
+        dataset,
+        total_samples,
+        consumed_samples,
+        micro_batch_size,
+        data_parallel_rank,
+        data_parallel_size,
+        data_sharding,
+        shuffle: bool = True,
+        group_by_length: bool = False,
+    ):
+        # Keep a copy of input params for later use.
+        self.dataset = dataset
+        self.total_samples = total_samples
+        self.consumed_samples = consumed_samples
+        self.micro_batch_size = micro_batch_size
+        self.data_parallel_rank = data_parallel_rank
+        self.data_parallel_size = data_parallel_size
+        if group_by_length:
+            if data_sharding:
+                data_sharding = False
+                logger.warning('`group_by_length=True` is incompatible with `data_sharding=True`. '
+                               'Setting `data_sharding=False` to enable length grouping.')
+            if not shuffle:
+                raise ValueError('shuffle must be True when group_by_length is True')
+        self.data_sharding = data_sharding
+        self.shuffle = shuffle
+        self.group_by_length = group_by_length
+        self.lengths = self.dataset['lengths'] if group_by_length else None
+        if self.lengths is not None:
+            self.lengths = [max(length) if isinstance(length, list) else length for length in self.lengths]
+        self.micro_batch_times_data_parallel_size = self.micro_batch_size * data_parallel_size
+        self.last_batch_size = self.total_samples % self.micro_batch_times_data_parallel_size
+
+        # Sanity checks.
+        assert self.total_samples > 0, 'no sample to consume: {}'.format(self.total_samples)
+        assert self.micro_batch_size > 0
+        assert data_parallel_size > 0
+        assert self.data_parallel_rank < data_parallel_size, (
+            'data_parallel_rank should be smaller than data size: {}, '
+            '{}'.format(self.data_parallel_rank, data_parallel_size))
+
+    def __len__(self):
+        return self.total_samples
+
+    def __iter__(self):
+        active_total_samples = self.total_samples - self.last_batch_size
+        self.epoch = self.consumed_samples // active_total_samples
+        current_epoch_samples = self.consumed_samples % active_total_samples
+        assert current_epoch_samples % self.micro_batch_times_data_parallel_size == 0
+
+        if self.shuffle:
+            # data sharding and random sampling
+            if self.data_sharding:
+                bucket_size = (self.total_samples // self.micro_batch_times_data_parallel_size) * self.micro_batch_size
+                bucket_offset = current_epoch_samples // self.data_parallel_size
+                start_idx = self.data_parallel_rank * bucket_size
+
+                g = torch.Generator()
+                g.manual_seed(self.epoch)
+                random_idx = torch.randperm(bucket_size, generator=g).tolist()
+                idx_range = [start_idx + x for x in random_idx[bucket_offset:]]
+            else:
+                full_bucket_size = (self.total_samples // self.micro_batch_size) * self.micro_batch_size
+                full_bucket_offset = current_epoch_samples
+                g = torch.Generator()
+                g.manual_seed(self.epoch)
+                if self.group_by_length:
+                    from transformers.trainer_pt_utils import get_length_grouped_indices
+                    idx_range_total = get_length_grouped_indices(
+                        self.lengths, self.micro_batch_times_data_parallel_size, generator=g)
+                else:
+                    idx_range_total = torch.randperm(full_bucket_size, generator=g).tolist()
+                idx_range_active = idx_range_total[full_bucket_offset:]
+                idx_range = idx_range_active[self.data_parallel_rank::self.data_parallel_size]
+        else:
+            full_bucket_size = (self.total_samples // self.micro_batch_size) * self.micro_batch_size
+            full_bucket_offset = current_epoch_samples
+            idx_range = range(full_bucket_offset + self.data_parallel_rank, full_bucket_size, self.data_parallel_size)
+
+        batch = []
+        # Last batch if not complete will be dropped.
+        for idx in idx_range:
+            batch.append(idx)
+            if len(batch) == self.micro_batch_size:
+                self.consumed_samples += self.micro_batch_times_data_parallel_size
+                yield batch
+                batch = []

swift/megatron/trainers/dpo_trainer.py

@@ -0,0 +1,94 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+from collections import namedtuple
+from functools import partial
+from megatron.core import mpu
+from torch.distributed.nn import all_reduce
+
+from swift.rlhf_trainers import DPOTrainer
+from swift.utils import get_current_device, get_logger
+from .rlhf_mixin import MegatronRLHFTrainer
+
+logger = get_logger()
+
+
+class DummyDPOTrainer(DPOTrainer):
+    # For reusing the dpo_loss function implemented in Swift's DPOTrainer.
+    def __init__(self, args):
+        self.accelerator = namedtuple('Accelerator', ['device'])(device=get_current_device())
+        self.f_alpha_divergence_coef = 1.
+        self.f_divergence_params = {'alpha_divergence_coef': self.f_alpha_divergence_coef}
+        self.reference_free = args.reference_free
+        self.label_smoothing = args.label_smoothing
+        self.f_divergence_type = args.f_divergence_type
+        self.loss_type = args.loss_type
+        self.beta = args.beta
+
+
+class MegatronDPOTrainer(MegatronRLHFTrainer):
+
+    def __init__(self, args, template):
+        super().__init__(args, template)
+        self.dummy_dpo_trainer = DummyDPOTrainer(args)
+
+    def loss_func(self, output_tensor: torch.Tensor, *, labels: torch.Tensor, packed_seq_params):
+        ref_output_tensor = output_tensor[:output_tensor.shape[0] // 2].detach()
+        output_tensor = output_tensor[output_tensor.shape[0] // 2:]
+        args = self.args
+        num_samples = labels.shape[0] if packed_seq_params is None else packed_seq_params.seq_lens.shape[0]
+
+        logps = self.get_logps(output_tensor, labels, packed_seq_params)
+        ref_logps = self.get_logps(ref_output_tensor, labels, packed_seq_params)
+        loss, chosen_rewards, rejected_rewards = self.dummy_dpo_trainer.dpo_loss(
+            logps[:num_samples // 2],
+            logps[num_samples // 2:],
+            ref_logps[:num_samples // 2],
+            ref_logps[num_samples // 2:],
+        )
+        if args.rpo_alpha:
+            loss_mask = labels != -100
+            if args.padding_free:
+                num_tokens = packed_seq_params.cu_seqlens_q[num_samples // 2] // args.context_parallel_size
+                loss_mask[:, num_tokens:] = 0
+            else:
+                loss_mask[num_samples // 2:] = 0
+            nll_loss = torch.concat([torch.sum(output_tensor * loss_mask)[None], loss_mask.sum()[None]])
+            if args.context_parallel_size > 1:
+                nll_loss = all_reduce(nll_loss, group=mpu.get_context_parallel_group())
+            nll_loss = nll_loss[0] / nll_loss[1]
+            loss = loss + args.rpo_alpha * nll_loss
+        loss = loss.mean()
+        metric = {
+            'loss': loss.detach().clone(),
+            'logps/chosen': logps[:num_samples // 2].mean(),
+            'logps/rejected': logps[num_samples // 2:].mean(),
+            'rewards/chosen': chosen_rewards.mean(),
+            'rewards/rejected': rejected_rewards.mean(),
+            'rewards/accuracies': (chosen_rewards > rejected_rewards).float().mean(),
+            'rewards/margins': (chosen_rewards - rejected_rewards).mean(),
+        }
+        if args.rpo_alpha:
+            metric['nll_loss'] = nll_loss.detach()
+        metric = self._all_reduce_metric(metric)
+        # fix megatron-lm bug
+        loss = loss / mpu.get_context_parallel_world_size()
+        return loss, metric
+
+    def forward_step(self, data_iterator, model):
+        unwrapped_model = model.module.module
+        input_tensor = unwrapped_model.get_input_tensor()
+        vp_stage = unwrapped_model.vp_stage
+        data = self.get_batch(data_iterator, vp_stage)
+        data.pop('loss_scale', None)
+        # ref_model
+        with torch.no_grad(), self.null_ref_context() as ref_models:
+            ref_model = ref_models[vp_stage or 0]
+            if input_tensor is not None:
+                ref_model.set_input_tensor(input_tensor[:input_tensor.shape[0] // 2].detach())
+            ref_output_tensor = ref_model(**data)
+
+        if input_tensor is not None:
+            unwrapped_model.set_input_tensor(input_tensor[input_tensor.shape[0] // 2:])
+        output_tensor = model(**data)
+        return torch.concat([ref_output_tensor, output_tensor], dim=0), partial(
+            self.loss_func, labels=data.get('labels'), packed_seq_params=data.get('packed_seq_params'))

swift/megatron/trainers/embedding_trainer.py

@@ -0,0 +1,62 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch.nn
+import torch.nn.functional as F
+from functools import partial
+
+from swift.loss import loss_map
+from swift.metrics import eval_metrics_map
+from swift.utils import get_logger
+from .base import BaseMegatronTrainer
+
+logger = get_logger()
+
+
+class MegatronEmbeddingTrainer(BaseMegatronTrainer):
+
+    def __init__(self, args, template):
+        super().__init__(args, template)
+        self._loss_func = loss_map[args.loss_type](args, self)
+        eval_metric = 'infonce' if args.loss_type == 'infonce' else 'paired'
+        self.eval_metrics = eval_metrics_map[eval_metric](args, self)
+
+    def loss_func(self,
+                  output_tensor: torch.Tensor,
+                  *,
+                  labels: torch.Tensor,
+                  packed_seq_params=None,
+                  attention_mask=None):
+        training = self.unwrapped_models[0].training
+        last_hidden_state = self.get_last_tokens(output_tensor, packed_seq_params, attention_mask)
+        if not training:
+            self.eval_metrics.update(last_hidden_state.detach(), labels)
+        mrl_dims = self.args.mrl_dims
+        if mrl_dims:
+            # Matryoshka Representation Learning: compute loss on each truncated dimension
+            # and aggregate with the corresponding weights.
+            loss = None
+            for dim, weight in mrl_dims.items():
+                if dim > last_hidden_state.shape[-1]:
+                    logger.warning_once(f'MRL: skipping dimension {dim} because it exceeds the model hidden size '
+                                        f'({last_hidden_state.shape[-1]}).')
+                    continue
+                sliced = F.normalize(last_hidden_state[..., :dim], p=2, dim=-1)
+                cur_loss = weight * self._loss_func({'last_hidden_state': sliced}, labels)
+                loss = cur_loss if loss is None else loss + cur_loss
+        else:
+            loss = self._loss_func({'last_hidden_state': last_hidden_state}, labels)
+        metric = {'loss': loss.detach().clone()}
+        metric = self._all_reduce_metric(metric)
+        return loss, metric
+
+    def forward_step(self, data_iterator, model):
+        vp_stage = model.module.module.vp_stage
+        data = self.get_batch(data_iterator, vp_stage)
+        labels = data.pop('labels', None)
+        output_tensor = model(**data)
+        loss_func = partial(
+            self.loss_func,
+            labels=labels,
+            packed_seq_params=data.get('packed_seq_params'),
+            attention_mask=data.get('attention_mask')
+            if data.get('attention_mask') is not None else data.get('attention_mask_2d'))
+        return output_tensor, loss_func

swift/megatron/trainers/gkd_trainer.py

@@ -0,0 +1,440 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import os
+import random
+import torch
+import torch.nn.functional as F
+from contextlib import contextmanager
+from functools import partial
+from mcore_bridge import set_random_seed
+from megatron.core import mpu
+from megatron.core.rerun_state_machine import RerunDataIterator
+from transformers.utils import ContextManagers
+from typing import Dict, List, Optional
+
+from swift.megatron.arguments import MegatronArguments
+from swift.rl_core.data import GKDSample
+from swift.rl_core.resample import resample_encode_failed_inputs
+from swift.rlhf_trainers.gkd_helpers import (assemble_teacher_output, build_opsd_samples, build_teacher_requests,
+                                             encode_gkd_samples, fetch_teacher_parsed_by_routing)
+from swift.rlhf_trainers.gkd_loss import DataSource, TeacherOutput, gkd_loss
+from swift.template import Template
+from swift.utils import get_logger, to_device
+from ..utils import forward_step_helper, get_padding_to
+from .gkd_utils import cp_reduce, tp_gather_topk, vocab_parallel_topk
+from .rlhf_mixin import MegatronRLHFTrainer
+from .rollout_mixin import MegatronRolloutMixin
+from .utils import gather_object
+from .vocab_parallel_utils import vocab_parallel_kl_div, vocab_parallel_log_softmax
+
+logger = get_logger()
+
+
+class MegatronGKDTrainer(MegatronRolloutMixin, MegatronRLHFTrainer):
+
+    sample_cls = GKDSample
+
+    def __init__(self, args: MegatronArguments, template, **kwargs):
+        self.vllm_client = kwargs.pop('vllm_client', None)
+
+        # GKD-specific parameters
+        self.beta = args.beta  # JSD interpolation coefficient
+        self.temperature = args.temperature
+        self.lmbda = args.lmbda  # On-policy probability
+        self.args = args
+        self._setup_teacher()
+        if self._teacher_use_disable_adapter:
+            logger.info('Self-distillation mode: using disable_adapter() for fixed teacher (no extra model)')
+        self.sft_alpha = getattr(args, 'sft_alpha', 0.0)  # Weight for SFT loss
+
+        # GKD top-k logits configuration
+        self.gkd_logits_topk = getattr(args, 'gkd_logits_topk', None)
+
+        self.use_vllm = getattr(args, 'use_vllm', False)
+        self.steps_per_generation = args.steps_per_generation
+        self.generation_batch_size = args.generation_batch_size
+        super().__init__(args, template)
+
+        if self.use_teacher_api:
+            logger.info(f'Using teacher model API for logprobs, top_logprobs={self.gkd_logits_topk}')
+
+        # Get device for data processing
+        self.device = torch.cuda.current_device()
+
+        # Initialize vLLM rollout engine if on-policy generation is enabled
+        self._init_rollout_engine()
+
+        # Truncation strategy for handling sequences that exceed max_length
+        self.truncation_strategy = args.truncation_strategy
+        self.max_completion_length = args.max_completion_length
+
+        self.resample_data_iterator = None
+        self._buffered_inputs = None
+
+        self._prepare_logging()
+
+    def train(self, train_dataset, val_dataset):
+        if self.truncation_strategy == 'delete':
+            self.resample_data_iterator = self._init_resample_data_iterator(train_dataset)
+        super().train(train_dataset, val_dataset)
+
+    def prepare_model(self):
+        super().prepare_model()
+        if self.use_teacher_api:
+            logger.info('Skipping local teacher model loading - using external API for teacher logprobs')
+        elif self._is_self_distillation:
+            logger.info('Self-distillation mode: using student model as teacher (no separate teacher loaded)')
+        self._load_teacher_model()
+
+    @contextmanager
+    def _template_context(self, template: Template, max_length: Optional[int] = None):
+        """Context manager to temporarily modify max_length constraint from template."""
+        original_max_length = template.max_length
+        template.max_length = max_length
+        try:
+            yield
+        finally:
+            template.max_length = original_max_length
+
+    def _build_teacher_requests(self, samples: List[GKDSample]):
+        if not self.use_teacher_api:
+            return []
+        return build_teacher_requests(samples, self.template)
+
+    def _encode_samples(self, samples: List[GKDSample]) -> Dict[str, torch.Tensor]:
+        template = self.template
+        args = self.args
+
+        with self._template_context(template):
+            student_encoded_list, teacher_encoded_list, has_opsd = encode_gkd_samples(samples, template)
+
+        padding_to = get_padding_to(args)
+        encoded_batch = to_device(template.data_collator(student_encoded_list, padding_to=padding_to), self.device)
+        if has_opsd:
+            teacher_model_inputs = to_device(
+                template.data_collator(teacher_encoded_list, padding_to=padding_to), self.device)
+        else:
+            teacher_model_inputs = encoded_batch.copy()
+        encoded_batch['teacher_model_inputs'] = teacher_model_inputs
+        return encoded_batch
+
+    def _get_random_num(self) -> float:
+        """Generate a deterministic random number consistent across all processes.
+
+        Uses an isolated Random instance with seed based on args.seed + step counter
+
+        Returns:
+            float: A random number in the range [0.0, 1.0).
+        """
+        seed = int(getattr(self.args, 'seed', 0))
+        seed += int(self._step)
+        rng = random.Random(seed)
+        return rng.random()
+
+    def _determine_data_source(self) -> DataSource:
+        """Determine data source for current step based on GKD algorithm.
+
+        GKD training mode selection logic:
+        1. With probability lmbda: On-Policy (student generates)
+        2. Otherwise: Off-Policy (use dataset responses)
+
+        Returns:
+            DataSource enum indicating which source to use.
+        """
+        random_num = self._get_random_num()
+
+        if random_num < self.lmbda:
+            # Mode 1: On-Policy learning, student model generates responses
+            if self.use_vllm:
+                return DataSource.STUDENT
+            else:
+                # If vLLM not enabled, fall back to dataset
+                logger.warning_once('On-policy mode triggered but use_vllm=False. '
+                                    'Falling back to dataset responses. Enable vLLM for on-policy generation.')
+                return DataSource.DATASET
+        else:
+            # Mode 2: Off-Policy learning, use dataset responses
+            return DataSource.DATASET
+
+    def _init_resample_data_iterator(self, train_dataset):
+        """Initialize an independent data iterator for resampling.
+
+        Uses a different seed (args.seed + 1) to avoid overlapping with training samples.
+
+        Args:
+            train_dataset: The training dataset to create the resample iterator from.
+
+        Returns:
+            The resample data iterator (first element of the iterator tuple).
+        """
+        args = self.args
+        resample_seed = getattr(args, 'seed', 42) + 1
+        try:
+            set_random_seed(
+                resample_seed,
+                args.data_parallel_random_init,
+                args.te_rng_tracker,
+            )
+            resample_data_iterator = self._prepare_data_iterator(train_dataset, use_origin_cyclic=True)[0]
+        finally:
+            set_random_seed(
+                args.seed,
+                args.data_parallel_random_init,
+                args.te_rng_tracker,
+            )
+        return resample_data_iterator
+
+    def resample_encode_failed_inputs(self, inputs: List[Dict], max_resample_rounds: int = 10) -> List[Dict]:
+        """Attempt to encode each input. If encoding fails, resample until we have enough valid samples.
+
+        Args:
+            inputs: List of input data samples
+            max_resample_rounds: Maximum number of resample rounds
+
+        Returns:
+            List of successfully encoded input samples with the same length as inputs
+        """
+        return resample_encode_failed_inputs(
+            self.template,
+            self.resample_data_iterator,
+            inputs,
+            max_resample_rounds=max_resample_rounds,
+            strip_response=False,
+        )
+
+    def _assemble_teacher_outputs(self, encoded_batches: List[Dict]) -> None:
+        for encoded_batch in encoded_batches:
+            parsed = encoded_batch.pop('_teacher_parsed')
+            teacher_model_inputs = encoded_batch['teacher_model_inputs']
+            teacher_out = assemble_teacher_output(
+                parsed,
+                teacher_model_inputs=teacher_model_inputs,
+                topk=self.gkd_logits_topk,
+                template_padding_free=self.template.padding_free,
+                device=self.device,
+            )
+            if teacher_out.labels is not None:
+                teacher_out.labels = torch.roll(teacher_out.labels, shifts=-1, dims=-1)
+            encoded_batch['teacher_output'] = teacher_out
+
+    def _compute_teacher_logits(self, encoded_batches: List[Dict], vp_stage: Optional[int] = None) -> None:
+        if self.use_teacher_api:
+            self._assemble_teacher_outputs(encoded_batches)
+            return
+        if self._is_self_distillation:
+            # Self-distillation teacher == current student weights. Computing it here (at batch
+            # preparation, once per steps_per_generation cycle) would reuse stale weights across the
+            # cycle's train steps. Defer to _replace_data_iterator so each train step recomputes the
+            # teacher with up-to-date student weights (weights are constant within a train step).
+            return
+        self._compute_teacher_logits_local(encoded_batches, vp_stage)
+
+    def _compute_teacher_logits_local(self, encoded_batches: List[Dict], vp_stage: Optional[int] = None) -> None:
+        """Compute teacher_output for each micro-batch via a local forward.
+
+        Handles both a separate fixed teacher and self-distillation (teacher == current student
+        weights). For self-distillation the caller is responsible for invoking this per train step
+        so the weights are current.
+        """
+        topk = self.gkd_logits_topk
+        if self._is_self_distillation:
+            teacher_model = self.unwrapped_models[vp_stage or 0]
+            adapter_contexts = []
+            if self._teacher_use_disable_adapter:
+                adapter_contexts = [m.disable_adapter() for m in self.peft_models]
+            outer_context = ContextManagers(adapter_contexts)
+        else:
+            teacher_model = self.teacher_models[vp_stage or 0]
+            outer_context = self.load_teacher_model_context()
+
+        with torch.no_grad(), outer_context:
+            for encoded_batch in encoded_batches:
+                teacher_model_inputs = encoded_batch['teacher_model_inputs']
+                teacher_batch = {
+                    k: v.clone() if isinstance(v, torch.Tensor) else v
+                    for k, v in teacher_model_inputs.items()
+                }
+                teacher_data = self._prepare_batch(teacher_batch, vp_stage)
+                teacher_data.pop('loss_scale', None)
+                teacher_labels = teacher_data.pop('labels', None)
+                teacher_logits = forward_step_helper(teacher_model, teacher_data)
+                if teacher_logits is not None:
+                    teacher_logits = teacher_logits.detach()
+
+                if topk is not None and teacher_logits is not None:
+                    topk_logits, topk_indices = vocab_parallel_topk(teacher_logits, k=topk)
+                    teacher_out = TeacherOutput(topk_logprobs=topk_logits, topk_indices=topk_indices)
+                else:
+                    teacher_out = TeacherOutput(full_logits=teacher_logits)
+
+                teacher_out.labels = teacher_labels
+                encoded_batch['teacher_output'] = teacher_out
+
+    def _generate_and_score_completions(self, inputs: List[Dict]) -> List[Dict]:
+        """Unified rollout → teacher → encode pipeline (mirrors Megatron GRPO).
+
+        Stages: determine data source → to_samples → (student) generate → teacher
+        requests/logprobs → encode micro-batches → teacher logits. Returns the flat
+        list of encoded micro-batches (length == total microbatches).
+        """
+        data_source = self._determine_data_source()
+
+        # Convert to samples (resample operates on dict, to_samples after)
+        samples = self.to_samples(inputs)
+
+        if data_source == DataSource.STUDENT:
+            local_batch = self._get_local_rollout_batch(samples)
+            local_batch = self._generate_completions(local_batch)
+            samples = self._gather_rollout_results(local_batch)
+            self._log_completions_from_samples(samples)
+
+        # Teacher API: build requests from samples, fetch logprobs
+        local_parsed = None
+        if self.use_teacher_api:
+            build_opsd_samples(samples)
+            teacher_requests = self._build_teacher_requests(samples)
+            if teacher_requests:
+                local_parsed = fetch_teacher_parsed_by_routing(
+                    samples,
+                    teacher_requests,
+                    self.teacher_configs,
+                    self.teacher_clients,
+                    gather_fn=self._gather_teacher_requests,
+                    infer_fn=lambda handle, client: self._infer_teacher_requests(
+                        handle, topk=self.gkd_logits_topk, teacher_client=client),
+                    scatter_fn=self._scatter_teacher_parsed,
+                    is_main_process=self.is_main_process,
+                    tag_key=self.args.teacher_tag_key)
+
+        # Encode micro-batches
+        total_microbatches = self.args.num_microbatches * self.steps_per_generation
+        micro_batch_size = len(samples) // total_microbatches
+        assert micro_batch_size == self.args.micro_batch_size
+        all_encoded_batches = []
+        for i in range(total_microbatches):
+            start_idx = i * micro_batch_size
+            end_idx = start_idx + micro_batch_size
+            sample_slice = samples[start_idx:end_idx]
+            encoded_batch = self._encode_samples(sample_slice)
+            encoded_batch['data_source'] = data_source
+            if local_parsed is not None:
+                encoded_batch['_teacher_parsed'] = local_parsed[start_idx:end_idx]
+            all_encoded_batches.append(encoded_batch)
+        self._compute_teacher_logits(all_encoded_batches)
+        return all_encoded_batches
+
+    def _replace_data_iterator(self, data_iterator):
+        num_microbatches = self.args.num_microbatches
+        steps_per_generation = self.steps_per_generation
+
+        if self._step % steps_per_generation == 0:
+            total_microbatches = num_microbatches * steps_per_generation
+            global_batch = []
+            for _ in range(total_microbatches):
+                raw_batch = next(data_iterator)
+                if self.truncation_strategy == 'delete' and self.resample_data_iterator is not None:
+                    raw_batch = self.resample_encode_failed_inputs(raw_batch)
+                global_batch.extend(raw_batch)
+
+            all_encoded_batches = self._generate_and_score_completions(global_batch)
+            self._buffered_inputs = [
+                all_encoded_batches[i * num_microbatches:(i + 1) * num_microbatches]
+                for i in range(steps_per_generation)
+            ]
+
+        step_idx = self._step % steps_per_generation
+        encoded_batches = self._buffered_inputs[step_idx]
+
+        # Self-distillation teacher == current student weights. Recompute per train step (weights are
+        # constant within a step) instead of once per generation cycle, so it tracks student updates
+        # across steps_per_generation. Runs outside the pipeline schedule, so PP > 1 is supported.
+        if self._is_self_distillation:
+            self._compute_teacher_logits_local(encoded_batches)
+
+        self._step += 1
+
+        return RerunDataIterator(iter(encoded_batches))
+
+    def loss_func(self,
+                  output_tensor: torch.Tensor,
+                  *,
+                  labels: torch.Tensor,
+                  teacher_output: TeacherOutput,
+                  data_source: DataSource = DataSource.DATASET):
+        """Compute GKD loss (JSD + optional SFT loss)."""
+        student_logits = output_tensor
+
+        jsd_total, jsd_num_valid = gkd_loss(
+            student_logits,
+            teacher_output,
+            labels,
+            self.beta,
+            self.temperature,
+            gather_fn=tp_gather_topk,
+            log_softmax_fn=vocab_parallel_log_softmax,
+            kl_div_fn=vocab_parallel_kl_div)
+        jsd_loss_val = cp_reduce(jsd_total, jsd_num_valid, cp_size=self.args.context_parallel_size)
+
+        loss = jsd_loss_val
+
+        # Add SFT loss if enabled (skip for student-generated responses)
+        sft_loss = None
+        if self.sft_alpha > 0 and data_source != DataSource.STUDENT:
+            args = self.args
+            logits_sbv = student_logits.transpose(0, 1).contiguous()
+            model = self.unwrapped_models[0]
+            if hasattr(model, 'language_model'):
+                model = model.language_model
+            per_token_loss = model.compute_language_model_loss(labels, logits_sbv)
+            loss_mask = labels != -100
+            sft_loss_sum = (per_token_loss * loss_mask).sum()
+            sft_loss_count = loss_mask.sum().float()
+
+            # All-reduce across CP group for correct averaging
+            if args.context_parallel_size > 1:
+                sft_stats = torch.stack([sft_loss_sum, sft_loss_count])
+                torch.distributed.all_reduce(
+                    sft_stats, op=torch.distributed.ReduceOp.SUM, group=mpu.get_context_parallel_group())
+                sft_loss_sum, sft_loss_count = sft_stats[0], sft_stats[1]
+
+            sft_loss = sft_loss_sum / sft_loss_count
+
+            loss = loss + self.sft_alpha * sft_loss
+
+        metric = {'loss': loss.detach().clone()}
+        if sft_loss is not None:
+            metric['jsd_loss'] = jsd_loss_val.detach().clone()
+            metric['sft_loss'] = sft_loss.detach().clone()
+        metric = self._all_reduce_metric(metric)
+
+        loss = loss / mpu.get_context_parallel_world_size()
+
+        # Flush completion logs at generation cycle boundaries.
+        if (self._step - 1) % self.steps_per_generation == 0:
+            self._flush_log_completions()
+
+        return loss, metric
+
+    def forward_step(self, data_iterator, model):
+        unwrapped_model = model.module.module
+        input_tensor = unwrapped_model.get_input_tensor()
+        vp_stage = unwrapped_model.vp_stage
+
+        data = next(data_iterator)
+        data_source = data.pop('data_source', DataSource.DATASET)
+        teacher_output = data.pop('teacher_output')
+        data.pop('teacher_model_inputs', None)  # consumed by _compute_teacher_logits; not needed for student forward
+        data = self._prepare_batch(data, vp_stage)
+
+        data.pop('loss_scale', None)
+        labels = data.pop('labels', None)
+
+        if input_tensor is not None:
+            unwrapped_model.set_input_tensor(input_tensor)
+        student_output = model(**data)
+
+        return student_output, partial(
+            self.loss_func,
+            labels=labels,
+            teacher_output=teacher_output,
+            data_source=data_source,
+        )

swift/megatron/trainers/gkd_utils.py

@@ -0,0 +1,67 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Megatron-specific GKD utilities: TP-aware gather/topk and CP reduce."""
+import torch
+from megatron.core import mpu
+
+
+def vocab_parallel_topk(logits: torch.Tensor, k: int) -> tuple:
+    """Global top-k from vocab-parallel sharded logits. TP=1 → plain torch.topk."""
+    tp_size = mpu.get_tensor_model_parallel_world_size()
+    if tp_size == 1:
+        return torch.topk(logits, k=k, dim=-1)
+
+    tp_rank = mpu.get_tensor_model_parallel_rank()
+    tp_group = mpu.get_tensor_model_parallel_group()
+    partition_vocab_size = logits.shape[-1]
+
+    local_topk_vals, local_topk_ids = torch.topk(logits, k=k, dim=-1)
+    local_topk_ids = local_topk_ids + tp_rank * partition_vocab_size
+
+    gathered_vals = [torch.empty_like(local_topk_vals) for _ in range(tp_size)]
+    gathered_ids = [torch.empty_like(local_topk_ids) for _ in range(tp_size)]
+    torch.distributed.all_gather(gathered_vals, local_topk_vals, group=tp_group)
+    torch.distributed.all_gather(gathered_ids, local_topk_ids, group=tp_group)
+
+    all_vals = torch.cat(gathered_vals, dim=-1)
+    all_ids = torch.cat(gathered_ids, dim=-1)
+    global_topk_vals, sel = torch.topk(all_vals, k=k, dim=-1)
+    global_topk_ids = torch.gather(all_ids, dim=-1, index=sel)
+    return global_topk_vals, global_topk_ids
+
+
+def tp_gather_topk(logits: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
+    """Gather logits at global top-k indices with TP-aware partitioning."""
+    tp_size = mpu.get_tensor_model_parallel_world_size()
+    if tp_size == 1:
+        return torch.gather(logits, dim=-1, index=indices)
+
+    tp_rank = mpu.get_tensor_model_parallel_rank()
+    partition_vocab_size = logits.shape[-1]
+    vocab_start = tp_rank * partition_vocab_size
+
+    in_range = (indices >= vocab_start) & (indices < vocab_start + partition_vocab_size)
+    local_indices = (indices - vocab_start).clamp(0, partition_vocab_size - 1)
+    gathered = torch.gather(logits, dim=-1, index=local_indices)
+    gathered = gathered.masked_fill(~in_range, float('-inf'))
+
+    gathered_for_reduce = gathered.detach()
+    torch.distributed.all_reduce(
+        gathered_for_reduce, op=torch.distributed.ReduceOp.MAX, group=mpu.get_tensor_model_parallel_group())
+    return torch.where(in_range, gathered, gathered_for_reduce)
+
+
+def cp_reduce(total_loss, num_valid, *, cp_size):
+    """Normalize total_loss by num_valid with CP all-reduce when cp_size > 1."""
+    num_valid_f = num_valid.float() if isinstance(num_valid, torch.Tensor) else torch.tensor(
+        float(num_valid), device=total_loss.device)
+    if cp_size > 1:
+        torch.distributed.all_reduce(
+            num_valid_f, op=torch.distributed.ReduceOp.SUM, group=mpu.get_context_parallel_group())
+        torch.distributed.all_reduce(
+            total_loss, op=torch.distributed.ReduceOp.SUM, group=mpu.get_context_parallel_group())
+    # Avoid the host-device sync from ``if num_valid_f == 0`` (num_valid_f is a GPU scalar,
+    # so a Python bool() forces a .item() stream sync every step). Keep the zero-valid
+    # guard on-device with torch.where: num_valid==0 -> loss 0, else total_loss/num_valid.
+    safe = torch.where(num_valid_f == 0, torch.ones_like(num_valid_f), num_valid_f)
+    loss = total_loss / safe
+    return torch.where(num_valid_f == 0, torch.zeros_like(loss), loss)

swift/megatron/trainers/kto_trainer.py

@@ -0,0 +1,165 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+from collections import namedtuple
+from functools import partial
+from megatron.core import mpu
+from trl import KTOTrainer
+from typing import Any, Dict
+
+from swift.utils import get_current_device, get_logger
+from .rlhf_mixin import MegatronRLHFTrainer
+
+logger = get_logger()
+
+
+class DummyKTOTrainer(KTOTrainer):
+    # For reusing the kto_loss function in TRL.
+
+    def gather_for_metrics(self, input_data, *args, **kwargs):
+        output_tensors = torch.empty(
+            mpu.get_data_parallel_world_size() * input_data.numel(),
+            dtype=input_data.dtype,
+            device=input_data.device,
+        )
+        torch.distributed.all_gather_into_tensor(output_tensors, input_data, group=mpu.get_data_parallel_group())
+        return output_tensors
+
+    def __init__(self, args):
+        self.accelerator = namedtuple('Accelerator', ['device', 'gather_for_metrics'])(
+            device=get_current_device(), gather_for_metrics=self.gather_for_metrics)
+        self.loss_type = args.loss_type
+        self.beta = args.beta
+        self.desirable_weight = args.desirable_weight
+        self.undesirable_weight = args.undesirable_weight
+        self.calculate_KL = args.calculate_KL
+
+
+class MegatronKTOTrainer(MegatronRLHFTrainer):
+
+    def __init__(self, args, template):
+        super().__init__(args, template)
+        self.dummy_kto_trainer = DummyKTOTrainer(args)
+
+    def _kto_get_logps(self, output_tensor, data, is_KL: bool, is_ref: bool, length: int):
+        output = self._get_input_tensor(output_tensor, is_KL, is_ref, length, dim=1)
+        return self.get_logps(output, data['labels'], data.get('packed_seq_params'))
+
+    def _get_kto_length(self, data: Dict[str, Any]) -> int:
+        if 'packed_seq_params' in data:
+            return data['packed_seq_params'].cu_seqlens_q[-1] // self.args.context_parallel_size
+        else:
+            return data['position_ids'].shape[-1]
+
+    def loss_func(self, output_tensor, *, data, kl_data, label):
+        length = self._get_kto_length(data)
+        policy_logps = self._kto_get_logps(output_tensor, data, False, False, length)
+        ref_logps = self._kto_get_logps(output_tensor, data, False, True, length)
+        if self.args.calculate_KL:
+            policy_KL_logps = self._kto_get_logps(output_tensor, kl_data, True, False, length)
+            ref_KL_logps = self._kto_get_logps(output_tensor, kl_data, True, True, length)
+        else:
+            policy_KL_logps, ref_KL_logps = None, None
+        label = output_tensor.new_tensor(label, dtype=torch.bool)
+        policy_chosen_logps = policy_logps[label]
+        policy_rejected_logps = policy_logps[~label]
+        ref_chosen_logps = ref_logps[label]
+        ref_rejected_logps = ref_logps[~label]
+
+        loss, chosen_rewards, rejected_rewards, kl = self.dummy_kto_trainer.kto_loss(
+            policy_chosen_logps,
+            policy_rejected_logps,
+            policy_KL_logps,
+            ref_chosen_logps,
+            ref_rejected_logps,
+            ref_KL_logps,
+        )
+
+        loss = loss.mean()
+        mean_metric = {
+            'loss': loss.detach().clone(),
+            'kl': kl.squeeze().detach(),
+        }
+        metric = self._all_reduce_metric(mean_metric)
+        chosen_count = chosen_rewards.shape[0]
+        rejected_count = rejected_rewards.shape[0]
+        sum_metric = {
+            'logps/chosen': loss.new_tensor([policy_chosen_logps.detach().nansum(), chosen_count]),
+            'logps/rejected': loss.new_tensor([policy_rejected_logps.detach().nansum(), rejected_count]),
+            'rewards/chosen': loss.new_tensor([chosen_rewards.nansum(), chosen_count]),
+            'rewards/rejected': loss.new_tensor([rejected_rewards.nansum(), rejected_count]),
+        }
+        metric.update(self._all_reduce_metric(sum_metric, torch.distributed.ReduceOp.SUM))
+        # fix megatron-lm bug
+        loss = loss / mpu.get_context_parallel_world_size()
+        return loss, metric
+
+    @staticmethod
+    def _get_input_tensor(input_tensor, is_KL: bool, is_ref: bool, length: int, dim: int):
+        # policy, ref, policy_KL, ref_KL
+        total_length = input_tensor.shape[dim]
+        KL_length = (total_length - 2 * length) // 2
+        slice_list = [0, length, 2 * length, total_length - KL_length, total_length]
+        idx = is_KL * 2 + is_ref
+        slice_ = (slice(None), ) * dim + (slice(slice_list[idx], slice_list[idx + 1]), )
+        res = input_tensor[slice_]
+        if is_KL or is_ref:
+            res = res.detach()
+        return res
+
+    def forward_step(self, data_iterator, model):
+        unwrapped_model = model.module.module
+        input_tensor = unwrapped_model.get_input_tensor()
+        vp_stage = unwrapped_model.vp_stage
+        # not support loss_scale
+        data, kl_data = self.get_batch(data_iterator, vp_stage)
+        label = data.pop('label')
+        data.pop('loss_scale', None)
+        kl_data.pop('loss_scale', None)
+
+        length = self._get_kto_length(data)
+        if self.args.sequence_parallel:
+            length //= mpu.get_tensor_model_parallel_world_size()
+        with torch.no_grad(), self.null_ref_context() as ref_models:
+            ref_model = ref_models[vp_stage or 0]
+            if self.args.calculate_KL:
+                if input_tensor is not None:
+                    ref_model.set_input_tensor(self._get_input_tensor(input_tensor, True, True, length, 0))
+                ref_KL_output_tensor = ref_model(**kl_data)
+
+            if input_tensor is not None:
+                ref_model.set_input_tensor(self._get_input_tensor(input_tensor, False, True, length, 0))
+            ref_output_tensor = ref_model(**data)
+
+        if self.args.calculate_KL:
+            with torch.no_grad():
+                if input_tensor is not None:
+                    unwrapped_model.set_input_tensor(self._get_input_tensor(input_tensor, True, False, length, 0))
+                KL_output_tensor = model(**kl_data)
+
+        if input_tensor is not None:
+            unwrapped_model.set_input_tensor(self._get_input_tensor(input_tensor, False, False, length, 0))
+        output_tensor = model(**data)
+        is_pp_last_stage = mpu.is_pipeline_last_stage(ignore_virtual=False, vp_stage=vp_stage)
+        dim = 1 if is_pp_last_stage else 0
+        if self.args.calculate_KL:
+            res = torch.concat([output_tensor, ref_output_tensor, KL_output_tensor, ref_KL_output_tensor], dim=dim)
+        else:
+            res = torch.concat([output_tensor, ref_output_tensor], dim=dim)
+        return res, partial(self.loss_func, data=data, kl_data=kl_data, label=label)
+
+    def _prepare_batch(self, data, vp_stage=None):
+        res = []
+        for key in ['completion_', 'KL_completion_']:
+            _data = {k[len(key):]: v for k, v in data.items() if k.startswith(key)}
+            if not self.args.calculate_KL and key == 'KL_completion_':
+                _data = {}
+            else:
+                _data = super()._prepare_batch(_data, vp_stage)
+            res.append(_data)
+        res[0]['label'] = data['label']
+        return res
+
+    def _log_callback(self, logs, n_steps):
+        super()._log_callback(logs, n_steps)
+        if 'rewards/chosen' in logs and 'rewards/rejected' in logs:
+            logs['rewards/margins'] = logs['rewards/chosen'] - logs['rewards/rejected']

swift/megatron/trainers/reranker_trainer.py

@@ -0,0 +1,73 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch.nn
+from collections import namedtuple
+from functools import partial
+
+from swift.loss import loss_map
+from swift.metrics import eval_metrics_map
+from swift.utils import get_logger
+from .base import BaseMegatronTrainer
+
+logger = get_logger()
+ModelOutputs = namedtuple('ModelOutputs', ['logits'])
+
+
+class MegatronRerankerTrainer(BaseMegatronTrainer):
+
+    def __init__(self, args, template):
+        super().__init__(args, template)
+        self._loss_func = loss_map[args.loss_type](args, self)
+        self.eval_metrics = eval_metrics_map['reranker'](args, self)
+
+    @staticmethod
+    def _get_listwise_reranker_preds(logits, labels):
+        positive_indices = torch.nonzero(labels == 1, as_tuple=False).squeeze(-1).tolist()
+        positive_indices.append(labels.shape[0])
+        preds = []
+        for i in range(len(positive_indices) - 1):
+            start, end = positive_indices[i], positive_indices[i + 1]
+            preds.append(logits[start:end].argmax())
+        preds = torch.stack(preds)
+        labels = torch.tensor([0] * (len(positive_indices) - 1), device=preds.device)
+        return preds, labels
+
+    def loss_func(self,
+                  output_tensor: torch.Tensor,
+                  *,
+                  labels: torch.Tensor,
+                  packed_seq_params=None,
+                  attention_mask=None):
+        training = self.unwrapped_models[0].training
+        logits = self.get_last_tokens(output_tensor, packed_seq_params, attention_mask)
+        loss = self._loss_func(ModelOutputs(logits=logits), labels)
+        args = self.args
+        logits_detach = logits.detach().squeeze(-1)
+        if not training:
+            self.eval_metrics.update(logits_detach, labels)
+        if args.loss_type == 'listwise_reranker':
+            preds, labels = self._get_listwise_reranker_preds(logits_detach, labels)
+        else:
+            preds = (logits_detach.detach() > 0).long()
+        acc = (preds == labels).float().mean()
+        metric = {'loss': loss.detach().clone(), 'acc': acc}
+        metric = self._all_reduce_metric(metric)
+        return loss, metric
+
+    def prepare_model(self):
+        super().prepare_model()
+        for model in self.unwrapped_models:
+            lm_model = model.language_model if hasattr(model, 'language_model') else model
+            lm_model.tokenizer = self.template.tokenizer
+
+    def forward_step(self, data_iterator, model):
+        vp_stage = model.module.module.vp_stage
+        data = self.get_batch(data_iterator, vp_stage)
+        labels = data.pop('labels', None)
+        output_tensor = model(**data)
+        loss_func = partial(
+            self.loss_func,
+            labels=labels,
+            packed_seq_params=data.get('packed_seq_params'),
+            attention_mask=data.get('attention_mask')
+            if data.get('attention_mask') is not None else data.get('attention_mask_2d'))
+        return output_tensor, loss_func

swift/megatron/trainers/reward_trainer.py

@@ -0,0 +1,46 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+from functools import partial
+from torch import nn
+
+from swift.utils import get_logger
+from .rlhf_mixin import MegatronRLHFTrainer
+
+logger = get_logger()
+
+
+class MegatronRewardTrainer(MegatronRLHFTrainer):
+
+    def loss_func(self, output_tensor, *, data):
+        packed_seq_params = data.get('packed_seq_params')
+        margin = data.pop('margin', None)
+        num_samples = output_tensor.shape[0] if packed_seq_params is None else packed_seq_params.seq_lens.shape[0]
+        rewards = self.get_last_tokens(output_tensor, packed_seq_params, data.get('attention_mask'))
+        rewards_chosen, rewards_rejected = torch.split(rewards, num_samples // 2, dim=0)
+        if margin is not None:
+            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - margin).mean()
+        else:
+            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()
+        if self.args.center_rewards_coefficient is not None:
+            center_rewards_loss = self.args.center_rewards_coefficient * torch.mean(
+                (rewards_chosen + rewards_rejected)**2)
+            loss += center_rewards_loss
+        rewards_chosen, rewards_rejected = rewards_chosen.detach(), rewards_rejected.detach()
+        metric = {
+            'loss': loss.detach().clone(),
+            'rewards/chosen': rewards_chosen.mean(),
+            'rewards/rejected': rewards_rejected.mean(),
+            'rewards/accuracies': (rewards_chosen > rewards_rejected).float().mean(),
+            'rewards/margins': (rewards_chosen - rewards_rejected).mean(),
+        }
+        if self.args.center_rewards_coefficient is not None:
+            metric['center_rewards_loss'] = center_rewards_loss.detach()
+        metric = self._all_reduce_metric(metric)
+        return loss, metric
+
+    def forward_step(self, data_iterator, model):
+        vp_stage = model.module.module.vp_stage
+        data = self.get_batch(data_iterator, vp_stage)
+        data.pop('loss_scale', None)
+        output_tensor = model(**data)
+        return output_tensor, partial(self.loss_func, data=data)

swift/megatron/trainers/rlhf_mixin.py

@@ -0,0 +1,104 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from contextlib import contextmanager
+from megatron.core import mpu
+from torch.distributed.nn import all_reduce
+from transformers.utils import ContextManagers
+
+from swift.megatron.model import get_mcore_model
+from swift.megatron.utils import load_mcore_checkpoint
+from swift.rlhf_trainers.utils import identity_data_collator
+from swift.utils import get_current_device, get_logger, safe_snapshot_download
+from .base import BaseMegatronTrainer
+from .utils import compute_per_token_logps_fn, reconstruct_tensor_cp
+
+logger = get_logger()
+
+
+class MegatronRLHFTrainer(BaseMegatronTrainer):
+
+    def _load_checkpoint(self):
+        args = self.args
+        if args.mcore_ref_model is not None:
+            load_mcore_checkpoint(args, self.ref_models, load_arg='mcore_ref_model')
+        if args.mcore_ref_adapter is not None:
+            load_mcore_checkpoint(args, self.wrapped_models, load_arg='mcore_ref_adapter')
+        super()._load_checkpoint()
+
+    def prepare_model(self):
+        super().prepare_model()
+        args = self.args
+        self.ref_models = []
+        if args.tuner_type == 'full' and args.rlhf_type not in ['rm', 'gkd']:
+            self.ref_models = get_mcore_model(args, self.template.config)
+        for ref_model in self.ref_models:
+            if not args.use_cpu_initialization:
+                ref_model.to(get_current_device())
+            ref_model.requires_grad_(False)
+            ref_model.eval()
+        if self.ref_models and args.mcore_ref_model is None:
+            ref_model_id_or_path = args.ref_model or args.model
+            ref_model_dir = safe_snapshot_download(ref_model_id_or_path, use_hf=args.use_hf, hub_token=args.hub_token)
+            self.bridge.load_weights(self.ref_models, ref_model_dir)
+        if args.tuner_type == 'lora' and args.ref_adapters and args.mcore_ref_adapter is None:
+            assert len(args.ref_adapters) == 1, 'Currently only support one adapter.'
+            self.bridge.load_weights(
+                self.ref_models, args.ref_adapters[0], peft_format=True, adapter_name='ref_adapter')
+
+    def _get_data_collator(self):
+        if self.args.rlhf_type in ('grpo', 'gkd'):
+            return identity_data_collator
+        return super()._get_data_collator()
+
+    @contextmanager
+    def null_ref_context(self):
+        args = self.args
+        contexts = []
+        has_ref_adapter = bool(args.mcore_ref_adapter or args.ref_adapters)
+        if args.tuner_type == 'full':
+            ref_models = self.ref_models
+        else:
+            if not has_ref_adapter:
+                for m in self.peft_models:
+                    contexts.append(m.disable_adapter())
+            ref_models = self.unwrapped_models
+        with ContextManagers(contexts):
+            if has_ref_adapter:
+                for m in self.peft_models:
+                    m.set_adapter('ref_adapter')
+            yield ref_models
+            if has_ref_adapter:
+                for m in self.peft_models:
+                    m.set_adapter('default')
+
+    def get_logps(self, output_tensor, labels, packed_seq_params, per_token=False):
+        args = self.args
+        per_token_logps = -output_tensor
+        loss_mask = labels != -100
+        per_token_logps = per_token_logps * loss_mask
+        num_samples = packed_seq_params.seq_lens.shape[0] if packed_seq_params is not None else labels.shape[0]
+        if per_token:
+            if args.context_parallel_size > 1:
+                per_token_logps = reconstruct_tensor_cp(args.context_parallel_size, per_token_logps, packed_seq_params,
+                                                        num_samples)
+            return per_token_logps
+
+        if args.padding_free:
+            cu_seqlens = packed_seq_params.cu_seqlens_q[:num_samples + 1] // args.context_parallel_size
+            all_logps = per_token_logps.new_zeros((num_samples, ))
+            for i in range(num_samples):
+                start, end = cu_seqlens[i], cu_seqlens[i + 1]
+                all_logps[i] = per_token_logps[:, start:end].sum()
+        else:
+            all_logps = per_token_logps.sum(-1)
+        if args.context_parallel_size > 1:
+            all_logps = all_reduce(all_logps, group=mpu.get_context_parallel_group())
+        return all_logps
+
+    def compute_per_token_logps(self, model, data_iterator, no_grad=True, temperature=1.0):
+        return compute_per_token_logps_fn(
+            model,
+            self.args,
+            data_iterator,
+            temperature=temperature,
+            no_grad=no_grad,
+            enable_routing_replay=self.enable_routing_replay)

swift/megatron/trainers/rollout_mixin.py

@@ -0,0 +1,979 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""
+Megatron Rollout Mixin - Provides vLLM integration for on-policy generation.
+
+This mixin extracts common vLLM rollout functionality from MegatronGRPOTrainer
+to be reused by GKD and other trainers that need online generation.
+"""
+import base64
+import inspect
+import json
+import os
+import re
+import time
+import torch
+import uuid
+from accelerate.utils import broadcast_object_list
+from collections import OrderedDict, defaultdict, deque
+from contextlib import contextmanager, nullcontext
+from copy import copy
+from dacite import from_dict
+from dataclasses import asdict
+from megatron.core import mpu
+from transformers import AutoConfig
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from swift.infer_engine.protocol import RequestConfig, RolloutInferRequest, RolloutOutput
+from swift.megatron.model import get_mcore_model
+from swift.rl_core.data import OnPolicySample
+from swift.rlhf_trainers.base_rollout_mixin import BaseRolloutTrainerMixin
+from swift.rlhf_trainers.gkd_helpers import resolve_dynamic_opd_self_distillation
+from swift.rlhf_trainers.utils import (VLLM_LORA_INT_ID, VLLM_LORA_NAME, VLLM_LORA_PATH, FlattenedTensorBucket,
+                                       TensorLoRARequest, add_base_layer_suffix_by_param_names, aggressive_empty_cache,
+                                       check_vllm_version_ge, expand_vllm_param_name_aliases, finish_vllm_weight_reload,
+                                       parse_prompt_logprobs, patch_vllm_load_adapter,
+                                       patch_vllm_moe_model_weight_loader, profiling_context, profiling_decorator,
+                                       set_expandable_segments, vllm_supports_lora_load_inplace)
+from swift.rlhf_trainers.vllm_client import VLLMInferClient
+from swift.rollout import MultiTurnScheduler, invoke_async_hook, multi_turns, run_multi_turn
+from swift.utils import (JsonlWriter, get_current_device, get_logger, is_last_rank, is_vllm_available, remove_response,
+                         synchronize, to_device)
+from .utils import (gather_object, load_megatron_model_to_gpu, load_megatron_optimizer, offload_megatron_model_to_cpu,
+                    offload_megatron_optimizer)
+
+DataType = List[Dict[str, Union[torch.Tensor, Any]]]
+logger = get_logger()
+
+
+def create_rollout_group(trainer) -> torch.distributed.ProcessGroup:
+    """
+    Get or create the rollout process group (TP×PP×CP).
+
+    This is a shared function used by both MegatronRolloutMixin and MegatronGRPOTrainer.
+
+    The rollout group is used for:
+    1. Data slicing: distributing rollout data across ranks with same data samples
+    2. Gather operations: collecting results from ranks with same data samples
+
+    Note: Groups are created per data parallel index, containing TP×PP×CP ranks each.
+    This follows Megatron's data_iterator logic where same data_parallel_rank processes
+    identical data samples.
+
+    Key insight: ranks with the SAME data parallel index process the SAME data samples
+    and must coordinate for rollout data distribution.
+    Megatron rank order: TP → CP → EP → DP → PP
+
+    Args:
+        trainer: Trainer instance with _rollout_group and _rollout_groups_created attributes
+
+    Returns:
+        The rollout process group for this rank
+    """
+    if trainer._rollout_group is not None:
+        return trainer._rollout_group
+
+    cp_size = mpu.get_context_parallel_world_size()
+    if cp_size == 1:
+        # No CP, use the standard MODEL_PARALLEL_GROUP
+        trainer._rollout_group = mpu.get_model_parallel_group()
+        return trainer._rollout_group
+
+    # Use RankGenerator to create rollout groups following Megatron-LM logic
+    global_rank = torch.distributed.get_rank()
+
+    # Get parallel dimensions
+    tp_size = mpu.get_tensor_model_parallel_world_size()
+    pp_size = mpu.get_pipeline_model_parallel_world_size()
+    dp_size = mpu.get_data_parallel_world_size()
+
+    # Create RankGenerator following Megatron-LM pattern
+    # Order: tp-cp-ep-dp-pp (default in Megatron-LM)
+    decoder_rank_generator = mpu.RankGenerator(
+        tp=tp_size,
+        ep=1,
+        dp=dp_size,
+        pp=pp_size,
+        cp=cp_size,
+        order='tp-cp-ep-dp-pp',
+        rank_offset=0,
+    )
+
+    # Create rollout groups based on data consistency from data_iterator
+    # Same data_parallel_rank processes same data - group ranks with same DP index
+    if not trainer._rollout_groups_created:
+        # Use 'tp-cp-ep-pp' to get groups with same DP index (DP is excluded from variation)
+        dp_groups = decoder_rank_generator.get_ranks('tp-cp-ep-pp')
+        for dp_group_ranks in dp_groups:
+            # Sort for consistency
+            dp_group_ranks = sorted(dp_group_ranks)
+            group = torch.distributed.new_group(ranks=dp_group_ranks, group_desc='ROLLOUT_GROUP')
+
+            if global_rank in dp_group_ranks:
+                trainer._rollout_group = group
+        trainer._rollout_groups_created = True
+
+    return trainer._rollout_group
+
+
+class MegatronRolloutMixin(BaseRolloutTrainerMixin):
+
+    # Per-sample container class; subclasses override (GRPOSample / GKDSample).
+    sample_cls = OnPolicySample
+
+    def _init_rollout_params(self):
+        """Initialize rollout generation parameters."""
+        args = self.args
+        # distributed params
+        self.world_size = torch.distributed.get_world_size()
+        self.process_index = torch.distributed.get_rank()
+        self.is_main_process = is_last_rank()
+        self.device = get_current_device()
+
+        # sampling params
+        self.temperature = getattr(args, 'temperature', 1.0)
+        self.max_completion_length = args.max_completion_length
+        structured_outputs_regex = getattr(args, 'structured_outputs_regex', None)
+
+        self.request_config = RequestConfig(
+            n=1,
+            max_tokens=args.max_completion_length,
+            temperature=args.temperature,
+            top_p=getattr(args, 'top_p', 1.0),
+            top_k=getattr(args, 'top_k', -1),
+            repetition_penalty=getattr(args, 'repetition_penalty', 1.0),
+            stop=getattr(args, 'stop_words', None),
+            return_details=True,
+            logprobs=True,
+            structured_outputs_regex=structured_outputs_regex)
+
+        self._last_loaded_step = -1
+        self._step = 0
+        self._rollout_group = None  # Lazily initialized rollout group (TP×PP×CP)
+        self._rollout_groups_created = False  # Flag for group creation (all ranks must create together)
+        self._bridge = None
+
+    def _get_rollout_group(self):
+        """Get or create the rollout process group (TP×PP×CP)."""
+        return create_rollout_group(self)
+
+    def _setup_teacher(self) -> None:
+        """Resolve teacher mode from args and init the API client when applicable.
+
+        Sets ``teacher_model_server`` / ``use_teacher_api`` / ``_is_self_distillation`` /
+        ``_teacher_use_disable_adapter`` / ``offload_teacher_model`` / ``_has_teacher``.
+        Must be called before ``_init_rollout_engine`` (the API client lives on last rank).
+        """
+        args = self.args
+        self.teacher_model_server = getattr(args, 'teacher_model_server', None)
+        self.use_teacher_api = self.teacher_model_server is not None
+        self.offload_teacher_model = args.offload_teacher_model
+        self._teacher_use_disable_adapter = getattr(args, '_teacher_use_disable_adapter', False)
+        self._is_self_distillation = (args.teacher_model is None and self.teacher_model_server is None)
+        self._has_teacher_explicit = (
+            args.teacher_model is not None or self.teacher_model_server is not None
+            or self._teacher_use_disable_adapter)
+        self._is_dynamic_self_distillation = resolve_dynamic_opd_self_distillation(
+            has_teacher_explicit=self._has_teacher_explicit,
+            is_self_distillation=self._is_self_distillation,
+        )
+        self._has_teacher = self._has_teacher_explicit or self._is_dynamic_self_distillation
+        self.teacher_models = None
+        self.teacher_configs: list = []
+        self.teacher_clients: list = []
+        if self.use_teacher_api:
+            from swift.rlhf_trainers.gkd_helpers import parse_teacher_model_server
+            self.teacher_configs = parse_teacher_model_server(self.teacher_model_server)
+            if is_last_rank():
+                self.teacher_clients = [VLLMInferClient(base_urls=[cfg.url]) for cfg in self.teacher_configs]
+
+    def _load_teacher_model(self) -> None:
+        """Load the separate local teacher mcore model (called from ``prepare_model``).
+
+        No-op for the API path, dynamic self-distillation, and same-model LoRA
+        (disable_adapter) — those reuse the student weights or an external server.
+        """
+        if self.use_teacher_api or self._is_self_distillation or self._teacher_use_disable_adapter:
+            return
+        args = self.args
+        vp_size = getattr(args, 'virtual_pipeline_model_parallel_size', None)
+        assert vp_size is None or vp_size == 1, 'Teacher distillation does not support VPP.'
+        self.teacher_hf_config = AutoConfig.from_pretrained(args.teacher_model_dir, trust_remote_code=True)
+        self.teacher_models = get_mcore_model(args, self.teacher_hf_config)
+        self.teacher_config = self.teacher_models[0].config
+        if not args.use_cpu_initialization:
+            for teacher_model in self.teacher_models:
+                teacher_model.cuda(torch.cuda.current_device())
+        for teacher_model in self.teacher_models:
+            teacher_model.requires_grad_(False)
+            teacher_model.eval()
+        self.teacher_config.bridge.load_weights(self.teacher_models, args.teacher_model_dir)
+        if self.offload_teacher_model:
+            self._offload_teacher_models()
+            logger.info('Teacher models offloaded to CPU to save GPU memory')
+
+    def _offload_teacher_models(self) -> None:
+        if self.teacher_models and not self.use_teacher_api:
+            offload_megatron_model_to_cpu(self.teacher_models)
+
+    def _load_teacher_models_to_gpu(self) -> None:
+        if self.teacher_models and not self.use_teacher_api:
+            load_megatron_model_to_gpu(self.teacher_models, load_grad=False)
+
+    @contextmanager
+    def load_teacher_model_context(self):
+        """Load the teacher to GPU for a forward and offload after (when offloading is on)."""
+        if not self.offload_teacher_model or self.teacher_models is None:
+            yield
+            return
+        self._load_teacher_models_to_gpu()
+        try:
+            yield
+        finally:
+            self._offload_teacher_models()
+
+    def _gather_teacher_requests(self, requests: List[RolloutInferRequest]) -> Dict[str, Any]:
+        """Phase 1 (all ranks, collective): gather this teacher's rollout-group-rank-0 requests.
+
+        Only rollout-group rank 0 contributes (others hold TP/PP/CP replicas of the same data);
+        contributions are tagged with ``dp_rank`` so segments order by DP rank regardless of world
+        layout or empty (zero-routed) subsets. Returns a handle with the per-DP-rank segments (for
+        the main-process infer) plus this DP rank's offset/length (for the later slice). The offset
+        is the prefix sum of preceding DP ranks' lengths, not ``dp_rank * n``: under multi-teacher
+        routing each DP rank's subset may differ in length, so equal-length slicing misaligns.
+        """
+        rollout_group = self._get_rollout_group()
+        rollout_rank = torch.distributed.get_rank(group=rollout_group)
+        dp_rank = mpu.get_data_parallel_rank()
+        contribution = (dp_rank, list(requests)) if rollout_rank == 0 else None
+
+        world_size = torch.distributed.get_world_size()
+        all_contributions = [None] * world_size
+        torch.distributed.all_gather_object(all_contributions, contribution)
+
+        segments_by_dp = {dp: reqs for c in all_contributions if c is not None for dp, reqs in [c]}
+        dp_ranks_sorted = sorted(segments_by_dp)
+        offset = sum(len(segments_by_dp[dp]) for dp in dp_ranks_sorted if dp < dp_rank)
+        flat_global = [req for dp in dp_ranks_sorted for req in segments_by_dp[dp]]
+        return {'flat_global': flat_global, 'offset': offset, 'n_local': len(requests)}
+
+    def _infer_teacher_requests(self, handle: Dict[str, Any], topk: int, teacher_client: Optional[Any] = None):
+        """Phase 2 (main process only, no collective): run the teacher HTTP infer.
+
+        Safe to call concurrently across teachers (distinct clients, no collective inside).
+        """
+        if not handle['flat_global']:  # no sample routed to this teacher: skip the empty HTTP call
+            return []
+        client = teacher_client if teacher_client is not None else self.teacher_clients[0]
+        request_config = RequestConfig(prompt_logprobs=topk, max_tokens=1, temperature=0.0)
+        responses = client.infer(handle['flat_global'], request_config=request_config, use_tqdm=False)
+        return [parse_prompt_logprobs(r, topk=topk) for r in responses]
+
+    def _scatter_teacher_parsed(self, handle: Dict[str, Any], parsed_global):
+        """Phase 3 (all ranks, collective): broadcast the parsed result and slice this rank's part."""
+        world_size = torch.distributed.get_world_size()
+        obj_list = [parsed_global]
+        torch.distributed.broadcast_object_list(obj_list, src=world_size - 1)
+        parsed_global = obj_list[0]
+        offset, n = handle['offset'], handle['n_local']
+        return parsed_global[offset:offset + n]
+
+    def _fetch_teacher_parsed_logprobs(self,
+                                       requests: List[RolloutInferRequest],
+                                       topk: int,
+                                       teacher_client: Optional[Any] = None):
+        """Combined gather→infer→broadcast for a single teacher (serial); returns this rank's slice."""
+        handle = self._gather_teacher_requests(requests)
+        parsed_global = self._infer_teacher_requests(handle, topk, teacher_client) \
+            if self.is_main_process else None
+        return self._scatter_teacher_parsed(handle, parsed_global)
+
+    def _get_local_rollout_batch(self, samples: List[OnPolicySample]) -> List[OnPolicySample]:
+        """Split batch within rollout group for distributed vLLM generation.
+
+        The batch is evenly split across the rollout group (TP×PP×CP ranks with
+        the same DP index). This is the base implementation that simply splits
+        the batch without repetition.
+
+        Subclasses (e.g., GRPO) may override this to implement custom logic like
+        repeating each prompt num_generations times.
+
+        Note: In Megatron, batch size should always be divisible by rollout group size.
+        This is ensured by global_batch_size = micro_batch_size * num_microbatches * dp_size,
+        where rollout_group_size = tp_size * pp_size * cp_size, and world_size = dp_size * rollout_group_size.
+
+        Args:
+            samples: Full batch of data samples
+
+        Returns:
+            Local slice of samples for this rank to process
+        """
+        rollout_group = self._get_rollout_group()
+        rollout_rank = torch.distributed.get_rank(group=rollout_group)
+        rollout_group_size = torch.distributed.get_world_size(group=rollout_group)
+
+        total_batch_size = len(samples)
+        assert total_batch_size % rollout_group_size == 0, \
+            f'Batch size ({total_batch_size}) must be divisible by rollout group size ({rollout_group_size})'
+
+        per_device_batch_size = total_batch_size // rollout_group_size
+        start_idx = rollout_rank * per_device_batch_size
+        end_idx = start_idx + per_device_batch_size
+
+        return samples[start_idx:end_idx]
+
+    def _gather_rollout_results(self, samples: List[OnPolicySample]) -> List[OnPolicySample]:
+        """Gather rollout results from all ranks in the rollout group.
+
+        Args:
+            samples: Local rollout results from this rank
+
+        Returns:
+            Gathered results from all ranks in the rollout group
+        """
+        rollout_group = self._get_rollout_group()
+        return gather_object(samples, group=rollout_group)
+
+    def _init_rollout_engine(self):
+        """Initialize vLLM engine for rollout generation."""
+        args = self.args
+        self._init_rollout_params()
+
+        self.vllm_mode = args.vllm_mode
+        self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization
+        self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size
+        self.use_vllm = args.use_vllm
+        self.vllm_use_async_engine = False
+        self.enable_offload = False
+        self.vllm_version_ge_0_10_2 = check_vllm_version_ge('0.10.2')
+        self.rollout_enable_lora = False
+        self.enable_server_multi_turn = False
+        self.base_sync_done = False
+        self._cached_vllm_param_names = None
+
+        self._prepare_scheduler()
+
+        if not args.use_vllm:
+            return
+
+        if args.rlhf_type == 'gkd' and args.lmbda == 0:
+            return
+
+        if not is_vllm_available():
+            raise ImportError('vLLM is not available and `use_vllm` is set to True. '
+                              'Please install vLLM with `pip install vllm -U` to use it.')
+
+        if self.vllm_mode == 'server':
+            # Server mode uses external vLLM server
+            if self.is_main_process:
+                self.vllm_client.get_engine_type()
+                self.vllm_client.reset_mm_cache()
+                enable_lora = [self.vllm_client.enable_lora]
+                enable_multi_turn = [self.vllm_client.enable_multi_turn]
+            else:
+                enable_lora = [False]
+                enable_multi_turn = [False]
+            self.rollout_enable_lora = broadcast_object_list(enable_lora, from_process=self.world_size - 1)[0]
+            self.enable_server_multi_turn = broadcast_object_list(
+                enable_multi_turn, from_process=self.world_size - 1)[0]
+        elif self.vllm_mode == 'colocate':
+            if self.world_size % self.vllm_tensor_parallel_size != 0:
+                raise ValueError(f'vllm_tensor_parallel_size ({self.vllm_tensor_parallel_size}) must divide world size '
+                                 f'({self.world_size}) evenly.')
+
+            self.enable_offload = args.offload_model or args.offload_optimizer
+            context = self.offload_context if self.enable_offload else nullcontext
+
+            with context():
+                set_expandable_segments(False)
+                self.engine = self._prepare_vllm_engine()
+                self.engine.engine.reset_mm_cache()
+                if args.sleep_level > 0:
+                    self.engine.engine.sleep(args.sleep_level)
+                set_expandable_segments(True)
+        else:
+            raise ValueError(f'Invalid vllm_mode: {self.vllm_mode}')
+
+    def _prepare_scheduler(self):
+        """Prepare multi-turn scheduler (shared by GRPO and GKD)."""
+        args = self.args
+
+        self.multi_turn_scheduler = None
+        if not hasattr(args, 'multi_turn_scheduler'):
+            return
+
+        if args.multi_turn_scheduler:
+            tokenizer = getattr(self, 'processing_class', None) or getattr(self.template, 'tokenizer', None)
+            if isinstance(args.multi_turn_scheduler, str):
+                assert args.multi_turn_scheduler in multi_turns
+                scheduler_kwargs = {'max_turns': args.max_turns, 'tokenizer': tokenizer}
+                gym_env = getattr(args, 'gym_env', None)
+                if gym_env is not None:
+                    scheduler_kwargs['gym_env'] = gym_env
+                multi_turn_scheduler = multi_turns[args.multi_turn_scheduler](**scheduler_kwargs)
+                self.multi_turn_scheduler: MultiTurnScheduler = multi_turn_scheduler
+            else:
+                assert isinstance(args.multi_turn_scheduler, MultiTurnScheduler)
+                self.multi_turn_scheduler: MultiTurnScheduler = args.multi_turn_scheduler
+
+    def _prepare_vllm_engine(self):
+        """Create and configure vLLM engine for colocate mode."""
+        from vllm.distributed import parallel_state as vllm_ps
+
+        from swift.infer_engine import GRPOVllmEngine
+
+        args = self.args
+        per_device_batch_size = getattr(args, 'per_device_generation_batch_size', args.micro_batch_size)
+        max_num_seqs = args.vllm_max_num_seqs or per_device_batch_size * self.vllm_tensor_parallel_size
+
+        vllm_template = copy(self.template)
+        vllm_template.padding_free = False
+        vllm_template.sequence_parallel_size = 1
+
+        logprobs_mode = 'processed_logprobs' if self.vllm_version_ge_0_10_2 else None
+
+        vllm_engine_kwargs = args.vllm_engine_kwargs or {}
+        load_format = vllm_engine_kwargs.pop('load_format', 'auto')
+
+        if self.args.router_replay_mode == 'R3':
+            assert check_vllm_version_ge('0.14.0'), \
+                'The enable_return_routed_experts attribute is not supported. Please upgrade vllm to 0.14.0 or higher'
+            vllm_engine_kwargs['enable_return_routed_experts'] = True
+            # https://github.com/vllm-project/vllm/pull/39917
+            import vllm
+            from packaging import version
+            vllm_version = vllm.__version__
+            if vllm_version is not None and version.parse('0.21.0rc1') <= version.parse(vllm_version) <= version.parse(
+                    '0.21.0'):
+                vllm_engine_kwargs.setdefault('async_scheduling', False)
+
+        enable_lora = False
+        max_loras = 1
+        max_lora_rank = args.lora_rank
+        if args.tuner_type == 'lora' and args.vllm_enable_lora:
+            enable_lora = True
+            self.rollout_enable_lora = True
+            patch_vllm_load_adapter()
+            logger.info(f'Enabled vLLM LoRA adapter sync with max_lora_rank={args.lora_rank}')
+
+        engine = GRPOVllmEngine(
+            args.model_info.model_dir,
+            torch_dtype=args.torch_dtype,
+            model_type=args.model_type,
+            use_async_engine=False,
+            tensor_parallel_size=self.vllm_tensor_parallel_size,
+            gpu_memory_utilization=self.vllm_gpu_memory_utilization,
+            enable_prefix_caching=args.vllm_enable_prefix_caching,
+            max_num_seqs=max_num_seqs,
+            enforce_eager=args.vllm_enforce_eager,
+            limit_mm_per_prompt=args.vllm_limit_mm_per_prompt,
+            enable_sleep_mode=args.sleep_level > 0,
+            max_model_len=args.vllm_max_model_len,
+            seed=self.process_index // self.vllm_tensor_parallel_size,
+            disable_cascade_attn=args.vllm_disable_cascade_attn,
+            load_format=load_format,
+            mm_processor_cache_gb=args.vllm_mm_processor_cache_gb,
+            template=vllm_template,
+            distributed_executor_backend='external_launcher',
+            enable_lora=enable_lora,
+            max_loras=max_loras,
+            max_lora_rank=max_lora_rank,
+            engine_kwargs=vllm_engine_kwargs,
+            logprobs_mode=logprobs_mode)
+
+        if self.vllm_tensor_parallel_size > 1:
+            self.vllm_tp_group = vllm_ps.get_tp_group().device_group
+
+        return engine
+
+    @profiling_decorator
+    def _move_model_to_vllm(self):
+        """Synchronize model weights to vLLM engine.
+
+        - Full sync: when tuner_type != 'lora' (e.g. full, lora_llm), or first sync
+          (base_sync_done=False), or sleep_level==2, or rollout_enable_lora is disabled.
+        - Adapter-only sync: when tuner_type == 'lora' with rollout_enable_lora=True and
+          base weights have already been synced.
+        """
+        args = self.args
+        tuner_type = args.tuner_type
+
+        if (tuner_type != 'lora' or (not self.base_sync_done or args.sleep_level == 2) or not self.rollout_enable_lora):
+            self._move_full_model_to_vllm()
+        else:
+            self._move_adapter_to_vllm()
+
+        self._reset_vllm_cache()
+
+    def _reset_vllm_cache(self):
+        # Reset prefix cache and encoder cache
+        vllm_ge_16 = check_vllm_version_ge('0.16')
+        if self.vllm_mode == 'server' and self.is_main_process:
+            self.vllm_client.reset_prefix_cache()
+            if vllm_ge_16:
+                self.vllm_client.reset_encoder_cache()
+        elif self.vllm_mode == 'colocate':
+            self.engine.engine.reset_prefix_cache()
+            if vllm_ge_16:
+                self.engine.engine.reset_encoder_cache()
+
+    def _move_full_model_to_vllm(self):
+        """Transfer full model weights to vLLM engine.
+
+        For LoRA training (tuner_type == 'lora'):
+        - When rollout_enable_lora=False: merge LoRA into base, export merged weights, then unmerge.
+        - When rollout_enable_lora=True: export base weights only (no merge needed),
+          then follow up with adapter-only sync via _move_adapter_to_vllm.
+        For lora_llm: always merge LLM LoRA into exported dense weights (vLLM has no separate
+        adapter pass for this tuner_type; see _move_model_to_vllm).
+        """
+        is_lora_training = self.args.tuner_type in ('lora', 'lora_llm')
+        is_pure_lora = self.args.tuner_type == 'lora'
+        should_merge = not self.rollout_enable_lora
+        if self.args.tuner_type == 'lora_llm' and self.rollout_enable_lora:
+            logger.warning('lora_llm is not supported with vllm_enable_lora=True. plz set vllm_enable_lora to False')
+
+        try:
+            if should_merge:
+                self.merge_lora_adapters()
+
+            self._export_and_load_weights()
+
+        finally:
+            if should_merge:
+                self.unmerge_lora_adapters()
+
+        if is_lora_training:
+            self.base_sync_done = True
+            if self.rollout_enable_lora and is_pure_lora:
+                self._move_adapter_to_vllm()
+
+    def _move_adapter_to_vllm(self):
+        """Transfer only LoRA adapter weights to vLLM engine.
+
+        Uses bridge.export_weights(peft_format=True) to export LoRA delta weights.
+        Yielded names follow PEFT convention: 'base_model.model.<hf_path>.lora_A.weight'.
+        """
+        target_device = 'cpu' if self.args.offload_bridge else None
+
+        with profiling_context(self, 'export_adapter_weights'):
+            adapter_iterator = self.bridge.export_weights(
+                self.unwrapped_models, target_device=target_device, peft_format=True)
+            lora_params = OrderedDict()
+            for name, tensor in adapter_iterator:
+                lora_params[name] = tensor.detach()
+
+        peft_config = self.unwrapped_models[0].peft_config.get('default', None)
+
+        if self.vllm_mode == 'colocate':
+            req_kw = dict(
+                lora_name=VLLM_LORA_NAME,
+                lora_int_id=VLLM_LORA_INT_ID,
+                lora_path=VLLM_LORA_PATH,
+                peft_config=asdict(peft_config),
+                lora_tensors=lora_params,
+            )
+            if vllm_supports_lora_load_inplace():
+                req_kw['load_inplace'] = True
+            lora_request = TensorLoRARequest(**req_kw)
+            self.engine.engine.add_lora(lora_request)
+        elif self.vllm_mode == 'server' and self.is_main_process:
+            bucket = FlattenedTensorBucket(named_tensors=list(lora_params.items()))
+            metadatas = bucket.get_metadata()
+            flattened_tensor = bucket.get_flattened_tensor()
+            self.vllm_client.update_adapter_flattened_param(peft_config, metadatas, flattened_tensor)
+            del bucket, metadatas, flattened_tensor
+
+        del lora_params
+
+    def _export_and_load_weights(self):
+        """Export weights from Megatron and load to vLLM."""
+        target_device = 'cpu' if self.args.offload_bridge else None
+
+        with profiling_context(self, 'export_weights'):
+            weight_iterator = self.bridge.export_weights(self.unwrapped_models, target_device=target_device)
+
+        if self.rollout_enable_lora:
+            vllm_param_names = self._get_vllm_param_names_for_mapping()
+            if vllm_param_names:
+                weight_iterator = add_base_layer_suffix_by_param_names(weight_iterator, vllm_param_names)
+
+        if self.vllm_mode == 'colocate':
+            llm_model = self.engine.inner_model
+            patch_vllm_moe_model_weight_loader(llm_model)
+            llm_model.load_weights(weight_iterator)
+            _model_config = self.engine.engine.model_config
+            finish_vllm_weight_reload(llm_model, model_config=_model_config, target_device=self.device)
+        elif self.vllm_mode == 'server':
+            self._load_weights_to_server_in_buckets(weight_iterator)
+            if self.is_main_process:
+                self.vllm_client.process_weights_after_loading()
+
+    def _get_vllm_param_names_for_mapping(self):
+        """Get vLLM runtime parameter names for base_layer mapping.
+
+        Returns an alias-expanded set so bridge/HF names can match vLLM packed names.
+        """
+        if self.vllm_mode == 'colocate':
+            llm_model = self.engine.inner_model
+            raw_names = set(dict[Any, Any](llm_model.named_parameters()).keys())
+            return expand_vllm_param_name_aliases(raw_names)
+
+        if self.vllm_mode != 'server' or not self.is_main_process:
+            return None
+
+        if self._cached_vllm_param_names is None:
+            keys = self.vllm_client.get_model_state_keys()
+            self._cached_vllm_param_names = expand_vllm_param_name_aliases(set(keys))
+        return self._cached_vllm_param_names
+
+    def _load_weights_to_server_in_buckets(self, weight_iterator):
+        """Load weights to vLLM server in buckets."""
+        bucket_size_mb = int(os.environ.get('SWIFT_UPDATE_WEIGHTS_BUCKET_SIZE', 512))
+        bucket_size_bytes = bucket_size_mb * 1024 * 1024
+
+        current_bucket = []
+        current_size = 0
+
+        for name, param in weight_iterator:
+            param_size = param.numel() * param.element_size()
+            current_bucket.append((name, param))
+            current_size += param_size
+
+            if current_size > bucket_size_bytes and current_bucket:
+                self._sync_bucket_to_server(current_bucket)
+                current_bucket = []
+                current_size = 0
+
+        if current_bucket:
+            self._sync_bucket_to_server(current_bucket)
+
+    def _sync_bucket_to_server(self, bucket_params: List[Tuple[str, torch.Tensor]]):
+        """Synchronize a bucket of parameters to vLLM server."""
+        if not bucket_params or not self.is_main_process:
+            return
+
+        # Ensure all async GPU ops (e.g. TP all-gather on NCCL stream from bridge.export_weights)
+        # are complete before .copy_() reads param data on the default stream.
+        synchronize()
+
+        bucket = FlattenedTensorBucket(named_tensors=bucket_params)
+        metadatas = bucket.get_metadata()
+        flattened_tensor = bucket.get_flattened_tensor()
+
+        self.vllm_client.update_flattened_params(metadatas, flattened_tensor)
+
+        del bucket, metadatas, flattened_tensor
+
+    @profiling_decorator
+    def _generate_completions(self, samples: List[OnPolicySample]) -> List[OnPolicySample]:
+        """Generate completions for a batch using vLLM engine.
+
+        Args:
+            samples: List of OnPolicySample carrying messages + rollout fields
+
+        Returns:
+            Batch with rollout completion results merged in
+        """
+        samples = self._preprocess_inputs(samples)
+
+        # Wake up engine if sleeping (colocate mode)
+        if self.vllm_mode == 'colocate' and self.engine.inner_model_executor.is_sleeping:
+            wake_up_params = inspect.signature(self.engine.engine.wake_up).parameters
+            kwargs = {'tags': ['weights']} if 'tags' in wake_up_params else {}
+            aggressive_empty_cache()
+            self.engine.engine.wake_up(**kwargs)
+
+        # Load model weights if needed
+        if self._step != self._last_loaded_step or self.args.sleep_level == 2:
+            self._move_model_to_vllm()
+            self._last_loaded_step = self._step
+
+        context = self.offload_context if self.enable_offload else nullcontext
+        with context():
+            if (self.vllm_mode == 'colocate' and self.engine.inner_model_executor.is_sleeping
+                    and 'tags' in inspect.signature(self.engine.engine.wake_up).parameters):
+                aggressive_empty_cache()
+                set_expandable_segments(False)
+                self.engine.engine.wake_up(tags=['kv_cache'])
+
+            multi_turn_scheduler = getattr(self, 'multi_turn_scheduler', None)
+            colocate_multi_turn = (
+                multi_turn_scheduler is not None and not getattr(self, 'enable_server_multi_turn', False))
+
+            if colocate_multi_turn:
+                requests = self.samples2requests(samples)
+                invoke_async_hook(multi_turn_scheduler.on_trajectory_start(requests))
+                request_config = self._get_request_config()
+                outputs: List[RolloutOutput] = self._rollout_requests(requests, request_config)
+                outputs = run_multi_turn(
+                    requests=requests,
+                    first_turn_outputs=outputs,
+                    scheduler=multi_turn_scheduler,
+                    rollout_fn=lambda reqs, cfg: self._rollout_requests(reqs, cfg),
+                    request_config=request_config,
+                    max_turns=self.args.max_turns,
+                    gather_fn=lambda x: gather_object(x, group=self._get_rollout_group()),
+                )
+            else:
+                # Single-turn rollout (or server multi-turn handled by the engine).
+                outputs: List[RolloutOutput] = self._rollout(samples)
+
+            # Sleep to release memory
+            if self.vllm_mode == 'colocate' and self.args.sleep_level > 0:
+                self.engine.engine.reset_prefix_cache()
+                self.engine.engine.sleep(level=self.args.sleep_level)
+                aggressive_empty_cache()
+                set_expandable_segments(True)
+
+            samples = self._postprocess_rollout_outputs(samples, outputs)
+
+        return samples
+
+    def _rollout_requests(self, requests: List[RolloutInferRequest],
+                          request_config: RequestConfig) -> List[RolloutOutput]:
+        """Continuation rollout taking already-prepared ``RolloutInferRequest`` objects.
+
+        Used by the multi-turn driver (:func:`swift.rollout.run_multi_turn`) on
+        every turn after the first. ``_set_inputs_system`` is skipped because the
+        system prompt is already encoded into ``requests[i].messages``.
+        """
+        if self.vllm_mode == 'server':
+            return self._server_rollout(requests, request_config)
+        elif self.vllm_mode == 'colocate':
+            return self._colocate_rollout(requests, request_config)
+        raise ValueError(f'Invalid vllm_mode: {self.vllm_mode}')
+
+    def _rollout(self, samples: List[OnPolicySample]) -> List[RolloutOutput]:
+        """Execute rollout using vLLM engine (system already injected by _preprocess_inputs)."""
+        request_config = self._get_request_config()
+
+        if self.vllm_mode == 'server':
+            return self._server_rollout(samples, request_config)
+        elif self.vllm_mode == 'colocate':
+            return self._colocate_rollout(samples, request_config)
+
+    def _get_request_config(self) -> RequestConfig:
+        """Get request config with proper seed for distributed TP groups."""
+        request_config = copy(self.request_config)
+
+        if self.vllm_mode == 'colocate' and self.vllm_tensor_parallel_size > 1:
+            batch_size = getattr(self.args, 'per_device_generation_batch_size', self.args.micro_batch_size)
+            batch_size *= self.vllm_tensor_parallel_size
+            request_config.seed = batch_size * (self.process_index // self.vllm_tensor_parallel_size)
+
+        return request_config
+
+    def _server_rollout(self, samples: Union[List[OnPolicySample], List[RolloutInferRequest]],
+                        request_config: RequestConfig) -> List[RolloutOutput]:
+        """Perform rollout using vLLM server mode."""
+        infer_requests = self.samples2requests(samples)
+
+        all_requests = gather_object(infer_requests)
+        all_requests_lengths = gather_object([len(infer_requests)])
+
+        if not any(requests for requests in all_requests):
+            return []
+
+        if self.is_main_process:
+            all_outputs: List[RolloutOutput] = self.vllm_client.infer(
+                infer_requests=all_requests, request_config=request_config)
+            if len(all_outputs) != len(all_requests):
+                # Per-turn-split multi-turn (`dynamic_num_samples`) is HF-only
+                raise NotImplementedError(
+                    'Per-turn-split multi-turn (dynamic_num_samples) is not supported on Megatron — '
+                    f'server returned {len(all_outputs)} outputs for {len(all_requests)} requests. '
+                    'Return one RolloutOutput per request from MultiTurnScheduler.run (combine turns '
+                    'inside response_token_ids: List[List[int]]), or use the HF trainer.')
+        else:
+            all_outputs = [None] * len(all_requests)
+
+        all_outputs = broadcast_object_list(all_outputs, from_process=self.world_size - 1)
+        start_idx = sum(all_requests_lengths[:self.process_index])
+        end_idx = start_idx + all_requests_lengths[self.process_index]
+        outputs = all_outputs[start_idx:end_idx]
+
+        return outputs
+
+    def _colocate_rollout(self, samples: Union[List[OnPolicySample], List[RolloutInferRequest]],
+                          request_config: RequestConfig) -> List[RolloutOutput]:
+        """Perform co-located rollout with vLLM engine."""
+        # Normalize samples (first turn) / RolloutInferRequest (continuation turns)
+        # into engine-ready requests.
+        samples = self.samples2requests(samples)
+        start_idx = 0
+        end_idx = len(samples)
+
+        # Handle vLLM tensor parallelism
+        if self.vllm_tensor_parallel_size > 1:
+            local_rank_in_group = torch.distributed.get_rank(group=self.vllm_tp_group)
+            local_input_length = len(samples)
+            all_input_lengths = [None] * self.vllm_tensor_parallel_size
+            torch.distributed.all_gather_object(all_input_lengths, local_input_length, group=self.vllm_tp_group)
+
+            start_idx = sum(all_input_lengths[:local_rank_in_group])
+            end_idx = start_idx + all_input_lengths[local_rank_in_group]
+
+            gathered = [None for _ in range(self.vllm_tensor_parallel_size)]
+            torch.distributed.all_gather_object(gathered, samples, group=self.vllm_tp_group)
+            samples = [p for sublist in gathered for p in sublist]
+
+        outputs: List[RolloutOutput] = self.engine.infer(
+            infer_requests=samples, request_config=request_config, use_tqdm=False)
+
+        if self.vllm_tensor_parallel_size > 1:
+            # R3 router replay: vLLM's routing capturer host cache only exists on
+            # TP rank 0, so non-primary TP ranks have routed_experts=None in outputs.
+            # Broadcast routed_experts from TP primary to all TP ranks.
+            if getattr(self.args, 'router_replay_mode', None) == 'R3':
+                routed_experts_list = [output.response.choices[0].routed_experts for output in outputs]
+                tp_primary_global_rank = torch.distributed.get_global_rank(self.vllm_tp_group, 0)
+                torch.distributed.broadcast_object_list(
+                    routed_experts_list, src=tp_primary_global_rank, group=self.vllm_tp_group)
+                if local_rank_in_group != 0:
+                    for output, experts in zip(outputs, routed_experts_list):
+                        output.response.choices[0].routed_experts = experts
+            outputs = outputs[start_idx:end_idx]
+
+        return outputs
+
+    def _preprocess_inputs(self, samples: List[OnPolicySample]) -> List[OnPolicySample]:
+        """Preprocess samples before rollout inference.
+
+        Unified pre-processing (mirrors HF RolloutTrainerMixin._preprocess_inputs):
+        1. Insert default system message if absent
+        2. Assign unique request_id (for rollout tracking)
+        3. Strip any prior assistant response (prompt-only for generation)
+        """
+        samples = self._set_inputs_system(samples)
+        for s in samples:
+            if not s.request_id:
+                s.request_id = f'chatcmpl-{str(uuid.uuid4().hex)}'
+            remove_response(s.messages)
+        return samples
+
+    def samples2requests(self, samples: Union[List[OnPolicySample],
+                                              List[RolloutInferRequest]]) -> List[RolloutInferRequest]:
+        """Convert samples into RolloutInferRequest objects.
+
+        Already-built ``RolloutInferRequest`` (multi-turn continuation) pass
+        through unchanged. The per-sample mapping lives in
+        ``OnPolicySample.to_infer_request``.
+        """
+        if not samples:
+            return []
+        requests_list = []
+        for data in samples:
+            if isinstance(data, RolloutInferRequest):
+                requests_list.append(data)
+            else:
+                include_extra = bool(getattr(self.args, 'vllm_server_pass_dataset', False)) or bool(
+                    getattr(self, 'multi_turn_scheduler', None))
+                requests_list.append(data.to_infer_request(include_extra=include_extra))
+        return requests_list
+
+    def _log_completions_from_samples(self, samples: List[OnPolicySample]) -> None:
+        """Log prompts/completions from sample messages (post-filtering).
+
+        Unlike ``_log_rollout`` (which needs raw RolloutOutput objects), this
+        method reads completions directly from ``sample.messages[-1]['content']``,
+        so it works after ``_dynamic_sampling`` has filtered / resampled samples.
+        """
+        if not self.log_completions:
+            return
+        messages = gather_object([s.messages for s in samples])
+        completions = []
+        for s in samples:
+            content = s.messages[-1]['content'] if s.messages else ''
+            if isinstance(content, str):
+                completions.append(content)
+            elif isinstance(content, list):
+                completions.append(self.template.safe_decode(content))
+            elif isinstance(content, dict) and 'input_ids' in content:
+                completions.append(self.template.safe_decode(content['input_ids']))
+            else:
+                completions.append(str(content))
+        completions = gather_object(completions)
+        self._logs['prompt'].extend(self._apply_chat_template_to_messages_list(messages))
+        self._logs['completion'].extend(completions)
+
+    def _prepare_logging(self):
+        """Initialize logging infrastructure (shared by GRPO and GKD)."""
+        args = self.args
+        self.log_completions = getattr(args, 'log_completions', False)
+        self.wandb_log_unique_prompts = getattr(args, 'wandb_log_unique_prompts', False)
+        self.jsonl_writer = JsonlWriter(os.path.join(args.output_dir, 'completions.jsonl'), write_on_rank='last')
+        self._last_logged_step = -1
+        self._logs = {
+            'prompt': deque(),
+            'completion': deque(),
+        }
+
+    def _flush_log_completions(self):
+        """Flush accumulated completion logs to jsonl/wandb/swanlab."""
+        if not (self.log_completions and self.is_main_process and len(self._logs['prompt']) > 0):
+            return
+        if self._step == self._last_logged_step:
+            return
+        self._last_logged_step = self._step
+
+        table = self._build_log_table()
+        self.jsonl_writer.append(table)
+        for val in self._logs.values():
+            if isinstance(val, deque):
+                val.clear()
+            elif isinstance(val, defaultdict):
+                for d in val.values():
+                    d.clear()
+
+        args = self.args
+        if 'wandb' in args.report_to:
+            import pandas as pd
+            import wandb
+            df = pd.DataFrame(table)
+            if self.wandb_log_unique_prompts:
+                df = df.drop_duplicates(subset=['prompt'])
+            wandb.log({'completions': wandb.Table(dataframe=df)}, commit=False)
+        if 'swanlab' in args.report_to:
+            import swanlab
+            headers = list(table.keys())
+            rows = [[table[h][i] for h in headers] for i in range(len(table.get('gen_step', table['prompt'])))]
+            swanlab.log({'completions': swanlab.echarts.Table().add(headers, rows)})
+
+    def _build_log_table(self) -> Dict[str, list]:
+        """Build the completion log table. Subclasses extend with extra columns (rewards/advantages)."""
+        return {
+            'gen_step': [self._step - 1] * len(self._logs['prompt']),
+            'prompt': list(self._logs['prompt']),
+            'completion': list(self._logs['completion']),
+        }
+
+    def _apply_chat_template_to_messages_list(self, messages_list):
+        """Convert messages list to prompt text using template."""
+        from swift.template import TemplateInputs
+        prompts_text = []
+        for messages in messages_list:
+            remove_response(messages)
+            template_inputs = TemplateInputs.from_dict({'messages': messages})
+            res = self.template.encode(template_inputs)
+            prompts_text.append(self.template.safe_decode(res['input_ids']))
+        return prompts_text
+
+    @contextmanager
+    def offload_context(self):
+        """Context manager for model/optimizer offloading during vLLM generation."""
+        if self.args.offload_model:
+            offload_megatron_model_to_cpu(self.wrapped_models)
+            if hasattr(self, 'ref_models') and self.ref_models:
+                offload_megatron_model_to_cpu(self.ref_models)
+        if getattr(self, 'optimizer', None) and self.args.offload_optimizer:
+            offload_megatron_optimizer(self.optimizer)
+
+        try:
+            yield
+        finally:
+            if self.args.offload_model:
+                load_megatron_model_to_gpu(self.wrapped_models)
+                if hasattr(self, 'ref_models') and self.ref_models:
+                    load_megatron_model_to_gpu(self.ref_models)
+            if getattr(self, 'optimizer', None) and self.args.offload_optimizer:
+                load_megatron_optimizer(self.optimizer)

swift/megatron/trainers/trainer.py

@@ -0,0 +1,130 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+import torch.distributed as dist
+import torch.nn
+from collections import defaultdict
+from functools import partial
+from megatron.core import mpu
+from torch.distributed.nn import all_reduce
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from typing import List, Optional
+
+from swift.utils import get_logger
+from .base import BaseMegatronTrainer
+
+logger = get_logger()
+
+
+class MegatronTrainer(BaseMegatronTrainer):
+
+    def seq_cls_loss_func(self, output_tensor, *, labels: torch.Tensor, packed_seq_params=None, attention_mask=None):
+        args = self.args
+        logits = self.get_last_tokens(output_tensor, packed_seq_params, attention_mask)
+        num_labels = args.num_labels
+        acc = None
+        if args.problem_type == 'regression':
+            loss_fct = MSELoss()
+            if num_labels == 1:
+                loss = loss_fct(logits.squeeze(), labels.squeeze())
+            else:
+                loss = loss_fct(logits, labels)
+        elif args.problem_type == 'single_label_classification':
+            loss_fct = CrossEntropyLoss()
+            logits = logits.view(-1, num_labels)
+            labels = labels.view(-1)
+            loss = loss_fct(logits, labels)
+            acc = (logits.detach().argmax(dim=-1) == labels).float().mean()
+        elif args.problem_type == 'multi_label_classification':
+            loss_fct = BCEWithLogitsLoss()
+            loss = loss_fct(logits, labels)
+            preds = logits.sigmoid() > 0.5
+            acc = (labels == preds).all(dim=-1).float().mean()
+        metric = {'loss': loss.detach().clone()}
+        if acc is not None:
+            metric['acc'] = acc
+        metric = self._all_reduce_metric(metric)
+        return loss, metric
+
+    def loss_func(self,
+                  output_tensor: torch.Tensor,
+                  *,
+                  labels: torch.Tensor,
+                  loss_scale: Optional[torch.Tensor] = None,
+                  channels: Optional[List[str]] = None,
+                  packed_seq_params=None):
+        args = self.args
+
+        losses = output_tensor.float()
+        loss_mask = labels != -100
+        if args.enable_dft_loss:
+            losses = losses * torch.exp(-losses.detach())
+        if loss_scale is not None:
+            losses = losses * loss_scale
+        loss = torch.cat([torch.sum(losses * loss_mask).view(1), loss_mask.sum().view(1)])
+
+        # Reduce loss for logging.
+        reporting_loss = loss.detach().clone()
+        torch.distributed.all_reduce(reporting_loss, group=mpu.get_data_parallel_group(with_context_parallel=True))
+
+        lm_loss = loss[0]
+        lm_loss = lm_loss.clone()
+        local_num_tokens = loss[1].detach().clone().to(torch.int)
+        metrics = {'loss': reporting_loss}
+        if args.enable_channel_loss:
+            metrics.update(self._compute_channel_loss(losses, loss_mask, channels, packed_seq_params))
+        return (lm_loss, local_num_tokens, metrics)
+
+    def _compute_channel_loss(self, losses, loss_mask, channels, packed_seq_params=None):
+        args = self.args
+        metrics = defaultdict(lambda: torch.tensor([0.0, 0.0], dtype=torch.float32, device=torch.cuda.current_device()))
+        if args.padding_free:
+            num_samples = packed_seq_params.seq_lens.shape[0]
+            cu_seqlens = packed_seq_params.cu_seqlens_q[:num_samples + 1] // args.context_parallel_size
+            for i in range(cu_seqlens.shape[0] - 1):
+                channel = None if channels is None else channels[i]
+                slice_ = slice(cu_seqlens[i], cu_seqlens[i + 1])
+                c_loss = losses[0, slice_][loss_mask[0, slice_]]
+                metrics[f'loss_{channel}'][0] += c_loss.detach().sum()
+                metrics[f'loss_{channel}'][1] += c_loss.shape[0]
+        else:
+            for i in range(losses.shape[0]):
+                channel = None if channels is None else channels[i]
+                c_loss = losses[i][loss_mask[i]]
+                metrics[f'loss_{channel}'][0] += c_loss.detach().sum()
+                metrics[f'loss_{channel}'][1] += c_loss.shape[0]
+
+        # Synchronize keys to avoid getting stuck.
+        dp_cp_group = mpu.get_data_parallel_group(with_context_parallel=True)
+        all_keys = [None] * torch.distributed.get_world_size(group=dp_cp_group)
+        dist.all_gather_object(all_keys, list(metrics.keys()), group=dp_cp_group)
+        new_metrics = {}
+        for key in sorted(set().union(*all_keys)):
+            new_metrics[key] = metrics[key]
+        new_metrics = self._all_reduce_metric(new_metrics, torch.distributed.ReduceOp.SUM, group=dp_cp_group)
+        return new_metrics
+
+    def forward_step(self, data_iterator, model):
+        vp_stage = model.module.module.vp_stage
+        data = self.get_batch(data_iterator, vp_stage)
+        loss_scale = data.pop('loss_scale', None)
+        channels = data.pop('channel', None)
+        labels = data.get('labels')
+        if self.args.task_type == 'seq_cls':
+            data.pop('labels', None)
+        output_tensor = model(**data)
+        packed_seq_params = data.get('packed_seq_params')
+        if self.args.task_type == 'seq_cls':
+            loss_func = partial(
+                self.seq_cls_loss_func,
+                labels=labels,
+                packed_seq_params=packed_seq_params,
+                attention_mask=data.get('attention_mask')
+                if data.get('attention_mask') is not None else data.get('attention_mask_2d'))
+        else:
+            loss_func = partial(
+                self.loss_func,
+                labels=labels,
+                loss_scale=loss_scale,
+                channels=channels,
+                packed_seq_params=packed_seq_params)
+        return output_tensor, loss_func

swift/megatron/trainers/utils.py

@@ -0,0 +1,525 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import gc
+import torch
+from accelerate.utils import gather as hf_gather
+from accelerate.utils import gather_object as hf_gather_object
+from contextlib import nullcontext
+from dataclasses import dataclass
+from megatron.core import mpu
+from megatron.core.distributed import DistributedDataParallel as DDP
+from megatron.core.optimizer import ChainedOptimizer
+from typing import Any, Optional
+
+from swift.dataloader import DataLoaderDispatcher
+from swift.megatron.utils import get_batch_on_this_cp_rank, get_packed_seq_params
+from swift.utils import empty_cache, get_current_device, get_logger, to_device
+
+logger = get_logger()
+
+
+def get_batch_on_this_pp_rank(args, data, vp_stage=None):
+    if args.task_type == 'causal_lm':
+        data['labels'] = torch.roll(data['labels'], -1, dims=-1)
+        if 'loss_scale' in data:
+            data['loss_scale'] = torch.roll(data['loss_scale'], -1, dims=-1)
+    batch = to_device(data, get_current_device(), non_blocking=True)
+    if args.pipeline_model_parallel_size == 1:
+        return batch
+    is_pp_last_stage = mpu.is_pipeline_last_stage(ignore_virtual=False, vp_stage=vp_stage)
+    if not is_pp_last_stage:
+        batch['labels'] = None
+        if 'loss_scale' in batch:
+            batch['loss_scale'] = None
+
+    return batch
+
+
+def gather(tensor, group: Optional[torch.distributed.ProcessGroup] = None):
+    if group is None:
+        return hf_gather(tensor)
+    size = torch.distributed.get_world_size(group=group)
+    output = [torch.empty_like(tensor) for _ in range(size)]
+    torch.distributed.all_gather(output, tensor, group=group, async_op=False)
+
+    return torch.cat(output, dim=0)
+
+
+def gather_object(object: Any, group: Optional[torch.distributed.ProcessGroup] = None):
+    if group is None:
+        return hf_gather_object(object)
+    size = torch.distributed.get_world_size(group=group)
+    output_objects = [None for _ in range(size)]
+    torch.distributed.all_gather_object(output_objects, object, group=group)
+    return [x for y in output_objects for x in y]
+
+
+# code borrowed from verl
+@torch.no_grad()
+def load_megatron_model_to_gpu(models, load_grad=True, load_frozen_params=True):
+    for model_chunk in models:
+        if isinstance(model_chunk, DDP):
+            model_chunk_all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers]
+            for buffers in model_chunk_all_buffers:
+                for buffer in buffers:
+                    # sometimes, we don't want to load grad for pure inference
+                    if load_grad and hasattr(buffer, 'grad_data_size'):
+                        current_storage_size = buffer.grad_data.storage().size()
+                        if current_storage_size == 0 or current_storage_size == buffer.grad_data_size:
+                            buffer.grad_data.storage().resize_(buffer.grad_data_size)
+                            buffer.grad_data.zero_()
+                        else:
+                            # Non-standard layers (e.g. GatedDeltaNet) may have grad
+                            # buffers with mismatched storage size; skip resize and
+                            # zero in-place with current storage.
+                            buffer.grad_data.zero_()
+
+                    if buffer.param_data.storage().size() == 0:
+                        buffer.param_data.storage().resize_(buffer.param_data_size)
+                        # copy data from cpu to cuda
+                        buffer.param_data.copy_(buffer.param_data.cpu_data, non_blocking=True)
+
+            if load_frozen_params:
+                device_id = get_current_device()
+                for param in model_chunk.module.parameters():
+                    if not param.requires_grad and param.device.type == 'cpu':
+                        param.data = param.data.to(device_id, non_blocking=True)
+        else:
+            # we need this for ref module
+            device_id = get_current_device()
+            for _, param in model_chunk.named_parameters():
+                param.data = param.data.to(device_id, non_blocking=True)
+                if param.grad is not None:
+                    param.grad = param.grad.to(device_id, non_blocking=True)
+    gc.collect()
+    empty_cache()
+
+
+@torch.no_grad()
+def offload_megatron_model_to_cpu(models):
+    """
+    In megatron, the model and optimizer storage are:
+    - bf16 parameter data chunked in model parallel group
+    - fp32 grad chunked in model parallel group
+    - fp32 main_parameter chunked in model and dp group
+    - fp32 optimizer state chunked in model and dp group
+
+    When using LoRA, frozen base model parameters are NOT managed by DDP buffers.
+    They must be offloaded separately via direct param iteration.
+    """
+    for model_chunk in models:
+        if isinstance(model_chunk, DDP):
+            model_chunk_all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers]
+            for buffers in model_chunk_all_buffers:
+                for buffer in buffers:
+                    # offload parameters
+                    if buffer.param_data.storage().size() > 0:
+                        existing = getattr(buffer.param_data, 'cpu_data', None)
+                        if existing is None:
+                            buffer.param_data.cpu_data = torch.empty(
+                                buffer.param_data.size(),
+                                dtype=buffer.param_data.dtype,
+                                device='cpu',
+                                pin_memory=True,
+                            )
+                            buffer.param_data_size = buffer.param_data.storage().size()
+                        else:
+                            assert existing.shape == buffer.param_data.shape, (
+                                f'cpu_data shape {tuple(existing.shape)} != '
+                                f'param_data shape {tuple(buffer.param_data.shape)}; '
+                                'reallocating would reintroduce the 2x peak.')
+                            assert existing.dtype == buffer.param_data.dtype, (
+                                f'cpu_data dtype {existing.dtype} != '
+                                f'param_data dtype {buffer.param_data.dtype}; '
+                                'reallocating would reintroduce the 2x peak.')
+                        # Synchronous D2H copy into the preexisting pinned
+                        # buffer; must complete before resize_(0) frees the
+                        # GPU storage.
+                        buffer.param_data.cpu_data.copy_(buffer.param_data.data, non_blocking=False)
+                        buffer.param_data.storage().resize_(0)
+
+                    assert buffer.param_data_size == buffer.param_data.cpu_data.storage().size()
+
+                    if buffer.grad_data.storage().size() > 0:
+                        # if the grad_data size is already zero, we assume that it is already offloaded
+                        buffer.grad_data_size = buffer.grad_data.storage().size()
+                        buffer.grad_data.storage().resize_(0)
+
+            for param in model_chunk.module.parameters():
+                if not param.requires_grad and param.device.type != 'cpu':
+                    param.data = param.data.to('cpu', non_blocking=True)
+        else:
+            # we need this for ref module
+            for _, param in model_chunk.named_parameters():
+                param.data = param.data.to('cpu', non_blocking=True)
+                if param.grad is not None:
+                    param.grad = param.grad.to('cpu', non_blocking=True)
+    gc.collect()
+    empty_cache()
+
+
+@torch.no_grad()
+def load_megatron_copy_params(optimizers):
+    """
+    Load optimizer parameters back to GPU. Handles ChainedOptimizer.
+
+    Args:
+        optimizers: Optimizer or ChainedOptimizer instance.
+    """
+
+    def _iter_opts(opt):
+        if isinstance(opt, ChainedOptimizer):
+            return opt.chained_optimizers
+        return [opt]
+
+    def load_tensor_to_gpu(tensor):
+        if tensor is None:
+            return
+        device_id = get_current_device()
+        tensor.data = tensor.data.to(device_id, non_blocking=True)
+
+    def load_group_to_gpu(group):
+        if group is None:
+            return
+
+        if isinstance(group, list):
+            for param_group in group:
+                if isinstance(param_group, list):
+                    for param in param_group:
+                        load_tensor_to_gpu(param)
+                else:
+                    load_tensor_to_gpu(param_group)
+        else:
+            load_tensor_to_gpu(group)
+
+    # Load all parameter groups to GPU for each underlying optimizer
+
+    for _opt in _iter_opts(optimizers):
+        if hasattr(_opt, 'shard_fp32_from_float16_groups'):
+            load_group_to_gpu(_opt.shard_fp32_from_float16_groups)
+
+
+@torch.no_grad()
+def offload_megatron_copy_params(optimizers):
+    """
+    Offload optimizer parameters to CPU. Supports both Megatron optimizers
+    and `ChainedOptimizer`, which wraps a list of underlying optimizers.
+
+    Args:
+        optimizers: The optimizer or ChainedOptimizer instance.
+    """
+
+    def _iter_opts(opt):
+        if isinstance(opt, ChainedOptimizer):
+            return opt.chained_optimizers
+        return [opt]
+
+    def offload_tensor_to_cpu(tensor):
+        if tensor is None:
+            return
+        tensor.data = tensor.data.to('cpu', non_blocking=True)
+
+    def offload_group_to_cpu(group):
+        if group is None:
+            return
+
+        if isinstance(group, list):
+            for param_group in group:
+                if isinstance(param_group, list):
+                    for param in param_group:
+                        offload_tensor_to_cpu(param)
+                else:
+                    offload_tensor_to_cpu(param_group)
+        else:
+            offload_tensor_to_cpu(group)
+
+    # Offload all parameter groups to CPU for each underlying optimizer
+
+    for _opt in _iter_opts(optimizers):
+        if hasattr(_opt, 'shard_fp32_from_float16_groups'):
+            offload_group_to_cpu(_opt.shard_fp32_from_float16_groups)
+
+
+@torch.no_grad()
+def load_megatron_optimizer(optimizers):
+
+    def _iter_opts(opt):
+        if isinstance(opt, ChainedOptimizer):
+            return opt.chained_optimizers
+        return [opt]
+
+    for _opt in _iter_opts(optimizers):
+        load_megatron_copy_params(_opt)
+        if _opt.optimizer is not None:
+            # if we are using HybridDeviceOptimizer, we need to only move gpu optimizer state to gpu
+            if hasattr(_opt.optimizer, '_move_new_state_to_right_device'):
+                _opt.optimizer._move_new_state_to_right_device()
+            else:
+                opt_state_dict_values = _opt.optimizer.state.values()
+                for v in opt_state_dict_values:
+                    if 'exp_avg' in v:
+                        v['exp_avg'] = v['exp_avg'].to(get_current_device(), non_blocking=True)
+                    if 'exp_avg_sq' in v:
+                        v['exp_avg_sq'] = v['exp_avg_sq'].to(get_current_device(), non_blocking=True)
+        gc.collect()
+        empty_cache()
+
+
+@torch.no_grad()
+def offload_megatron_optimizer(optimizers):
+
+    def _iter_opts(opt):
+        if isinstance(opt, ChainedOptimizer):
+            return opt.chained_optimizers
+        return [opt]
+
+    for _opt in _iter_opts(optimizers):
+        offload_megatron_copy_params(_opt)
+        # worker may hold zero parameter when enabling custom pipeline layout
+        if _opt.optimizer is not None:
+            # HybridDeviceOptimizer: offload all sub-optimizer states to CPU
+            hdo = _opt.optimizer
+            if all(hasattr(hdo, attr) for attr in ('sub_optimizers', 'inner_param_to_orig_param', 'state')):
+                for optimizer in hdo.sub_optimizers:
+                    for param, state in optimizer.state.items():
+                        for k, v in state.items():
+                            if not isinstance(v, torch.Tensor):
+                                continue
+                            orig_param = hdo.inner_param_to_orig_param.get(param, param)
+                            hdo.state[orig_param][k] = state[k] = v.to('cpu')
+            else:
+                opt_state_dict_values = _opt.optimizer.state.values()
+                for v in opt_state_dict_values:
+                    if 'exp_avg' in v:
+                        v['exp_avg'] = v['exp_avg'].to('cpu', non_blocking=True)
+                    if 'exp_avg_sq' in v:
+                        v['exp_avg_sq'] = v['exp_avg_sq'].to('cpu', non_blocking=True)
+        gc.collect()
+        empty_cache()
+
+
+def log_gpu_memory(prefix: str = '', info_once: bool = False):
+    log_msg = (f'{prefix} GPU memory: {torch.cuda.memory_allocated() / 1024**3:.2f}GB allocated, '
+               f'{torch.cuda.memory_reserved() / 1024**3:.2f}GB reserved')
+    if info_once:
+        logger.info_once(log_msg, hash_id=prefix)
+    else:
+        logger.info(log_msg)
+
+
+@dataclass
+class TrainerState:
+    should_save: bool = False
+    should_eval: bool = False
+    should_log: bool = False
+
+    iteration: int = 0
+    consumed_train_samples: int = 0
+    # compat transformers
+    max_steps: Optional[int] = None
+
+    best_metric: Optional[float] = None
+    best_global_step: Optional[int] = None
+    last_model_checkpoint: Optional[str] = None
+    best_model_checkpoint: Optional[str] = None
+
+    @property
+    def global_step(self) -> int:
+        return self.iteration
+
+
+class MegatronDataLoaderDispatcher(DataLoaderDispatcher):
+
+    @property
+    def group(self):
+        return mpu.get_data_parallel_group()
+
+
+def build_streaming_dataloader(args, dataset, collate_fn):
+    base_dataloader = torch.utils.data.DataLoader(
+        dataset,
+        num_workers=args.dataloader_num_workers,
+        pin_memory=args.dataloader_pin_memory,
+        collate_fn=collate_fn,
+        batch_size=args.micro_batch_size,
+        prefetch_factor=args.dataloader_prefetch_factor if args.dataloader_num_workers > 0 else None,
+        persistent_workers=args.dataloader_persistent_workers if args.dataloader_num_workers > 0 else False,
+    )
+    return MegatronDataLoaderDispatcher(base_dataloader)
+
+
+_NPU_ATTENTION_MASK_2D_MODEL_TYPES = {'qwen3_5', 'qwen3_5_moe'}
+
+
+def _should_use_npu_generated_attention_mask(args) -> bool:
+    from transformers.utils import is_torch_npu_available
+    if not is_torch_npu_available():
+        return False
+    if args.task_type != 'causal_lm' or args.padding_free:
+        return False
+    if getattr(args, 'attention_backend', None) == 'local':
+        return False
+    return bool(getattr(args, 'use_flash_attn', False))
+
+
+def _prepare_npu_generated_attention_mask(batch, *, keep_attention_mask_2d: bool) -> None:
+    if keep_attention_mask_2d:
+        attention_mask = batch.get('attention_mask')
+        if 'attention_mask_2d' not in batch and attention_mask is not None:
+            batch['attention_mask_2d'] = (attention_mask == 0).sum(dim=(1, 2)) > 0
+    else:
+        batch.pop('attention_mask_2d', None)
+    batch['attention_mask'] = None
+
+
+def prepare_batch(args, data, vp_stage=None):
+    """Prepare a micro-batch for Megatron forward: PP slicing, packed_seq_params, CP slicing.
+
+    Extracted from BaseMegatronTrainer._prepare_batch for reuse in ray workers.
+    """
+    batch = get_batch_on_this_pp_rank(args, data, vp_stage=vp_stage)
+    seq_lens = batch.pop('seq_lens', None)
+    # Consider compatibility and security.
+    num_samples = batch.pop('num_samples', None)
+    if seq_lens is not None:
+        if num_samples is not None:
+            assert num_samples == len(seq_lens), (
+                f"'num_samples' ({num_samples}) is inconsistent with len(seq_lens) ({len(seq_lens)}).")
+        num_samples = len(seq_lens)
+    text_position_ids = batch.pop('text_position_ids', None)
+    if text_position_ids is None:
+        text_position_ids = batch.get('position_ids')
+    if _should_use_npu_generated_attention_mask(args):
+        _prepare_npu_generated_attention_mask(
+            batch, keep_attention_mask_2d=getattr(args, 'model_type', None) in _NPU_ATTENTION_MASK_2D_MODEL_TYPES)
+    else:
+        batch.pop('attention_mask_2d', None)
+    if args.padding_free and text_position_ids is not None:
+        batch['packed_seq_params'] = get_packed_seq_params(args, text_position_ids)
+        if seq_lens is not None:
+            batch['packed_seq_params'].seq_lens = torch.tensor(seq_lens, device=text_position_ids.device)
+        if num_samples is not None:
+            batch['packed_seq_params'].num_samples = num_samples
+    batch.setdefault('attention_mask', None)
+    batch = get_batch_on_this_cp_rank(args, batch)
+    return batch
+
+
+def compute_per_token_logps_fn(model, args, data_iterator, temperature=1.0, no_grad=True, enable_routing_replay=False):
+    """Forward pass → logits → temperature-scaled per-token logps.
+
+    Returns:
+        (per_token_logps, routing_topk_idx) — either may be None on non-last PP stages.
+    """
+    from swift.megatron.utils import (RouterReplayHelper, forward_step_helper, get_local_topk_idx_for_current_rank,
+                                      get_router_replay_data, set_router_replay_data)
+    from .vocab_parallel_utils import compute_logps_and_entropy_from_logits
+
+    data = prepare_batch(args, next(data_iterator))
+    data.pop('loss_scale', None)
+    labels = data.get('labels')
+
+    routing_topk_idx = None
+    global_topk_idx = data.pop('routed_experts', None)
+    if enable_routing_replay and RouterReplayHelper.is_replay_forward_action(model.config):
+        assert global_topk_idx is not None, 'When router_replay_mode = R3, routed_experts must be in data'
+        routing_topk_idx = get_local_topk_idx_for_current_rank(global_topk_idx, model.config,
+                                                               data.get('packed_seq_params'))
+        set_router_replay_data(routing_topk_idx, model.config)
+
+    data_for_forward = {k: v for k, v in data.items() if k != 'labels'}
+    context = torch.no_grad() if no_grad else nullcontext()
+    is_training = model.training
+    if is_training:
+        model.eval()
+    try:
+        with context:
+            output_tensor = forward_step_helper(model, data_for_forward)
+    finally:
+        if is_training:
+            model.train()
+
+    if enable_routing_replay and RouterReplayHelper.is_r2_record_action(model.config):
+        routing_topk_idx = get_router_replay_data(model.config)
+
+    if labels is None or output_tensor is None:
+        return None, routing_topk_idx
+
+    if temperature != 1.0:
+        output_tensor.div_(temperature)
+    per_token_logps, _ = compute_logps_and_entropy_from_logits(output_tensor, labels)
+
+    packed_seq_params = data.get('packed_seq_params')
+    if packed_seq_params is not None:
+        num_samples = packed_seq_params.seq_lens.shape[0]
+    else:
+        input_ids = data.get('input_ids')
+        num_samples = input_ids.shape[0] if input_ids is not None else labels.shape[0]
+
+    if args.context_parallel_size > 1:
+        per_token_logps = reconstruct_tensor_cp(args.context_parallel_size, per_token_logps, packed_seq_params,
+                                                num_samples)
+    return per_token_logps, routing_topk_idx
+
+
+def reconstruct_tensor_cp(cp_size, tensor, packed_seq_params, num_samples):
+    """In CP mode, all_gather and reconstruct full tensor sequences."""
+    cp_rank = mpu.get_context_parallel_rank()
+
+    # All-gather across CP ranks
+    output_list = [torch.empty_like(tensor) for _ in range(cp_size)]
+    torch.distributed.all_gather(output_list, tensor.contiguous(), group=mpu.get_context_parallel_group())
+    output_list[cp_rank] = tensor
+
+    if packed_seq_params is not None:
+        cu_seqlens_full = packed_seq_params.cu_seqlens_q
+        cu_seqlens_cp = cu_seqlens_full // cp_size
+
+        # Calculate total packed length
+        total_packed_len = cu_seqlens_full[num_samples].item()
+        output_full = tensor.new_zeros(1, total_packed_len)
+
+        # Reconstruct each sequence
+        for i in range(num_samples):
+            start_full = cu_seqlens_full[i].item()
+            end_full = cu_seqlens_full[i + 1].item()
+            seq_len = end_full - start_full
+
+            # Length of each chunk after CP split
+            chunk_len = seq_len // cp_size
+            half_chunk = chunk_len // 2
+
+            # Concatenate from each CP rank's output (load-balanced split)
+            for j in range(cp_size):
+                o = output_list[j][0]
+                start_cp = cu_seqlens_cp[i].item()
+                o0 = o[start_cp:start_cp + half_chunk]
+                o1 = o[start_cp + half_chunk:start_cp + chunk_len]
+
+                # Place back to full sequence
+                output_full[0, start_full + j * half_chunk:start_full + (j + 1) * half_chunk] = o0
+                output_full[0, end_full - (j + 1) * half_chunk:end_full - j * half_chunk] = o1
+    else:
+        # non-padding_free mode: [batch_size, seq_len/cp_size] -> [batch_size, seq_len]
+        # Each CP rank has chunks split with load-balanced pattern (2*cp_size chunks)
+        batch_size = tensor.shape[0]
+        seq_len_per_cp = tensor.shape[1]
+        full_seq_len = seq_len_per_cp * cp_size
+        output_full = tensor.new_zeros(batch_size, full_seq_len)
+
+        # Each CP rank j holds chunks j and (2*cp_size - j - 1) from the original 2*cp_size split
+        # Reconstruct the full sequence by placing chunks back in correct positions
+        chunk_len = full_seq_len // (2 * cp_size)
+
+        for j in range(cp_size):
+            o = output_list[j]
+            # This rank holds 2 chunks: chunk j and chunk (2*cp_size - j - 1)
+            half_len = seq_len_per_cp // 2
+            o0 = o[:, :half_len]
+            o1 = o[:, half_len:]
+
+            # Place chunk j at position j * chunk_len
+            output_full[:, j * chunk_len:(j + 1) * chunk_len] = o0
+            reverse_idx = 2 * cp_size - j - 1
+            output_full[:, reverse_idx * chunk_len:(reverse_idx + 1) * chunk_len] = o1
+
+    return output_full

swift/megatron/trainers/vocab_parallel_utils.py

@@ -0,0 +1,236 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Vocabulary-parallel utilities for Tensor Parallelism.
+
+This module provides utilities for computing log_softmax, entropy, KL divergence,
+and other operations across vocab-parallel sharded tensors in Tensor Parallelism (TP).
+
+When using TP, the vocabulary dimension is sharded across TP ranks. These utilities
+correctly handle the distributed computation by:
+1. Finding global max via all_reduce (for numerical stability)
+2. Computing sum of exp via all_reduce (for normalization)
+3. All-reducing partial sums for final results
+"""
+
+import torch
+from megatron.core import mpu
+from typing import Optional, Tuple
+
+
+def vocab_parallel_log_softmax(logits: torch.Tensor) -> torch.Tensor:
+    """Compute log_softmax across vocab-parallel sharded logits.
+
+    When using Tensor Parallelism, vocab is sharded across TP ranks.
+    This function correctly computes log_softmax by:
+    1. Finding global max via all_reduce
+    2. Computing sum of exp via all_reduce
+    3. Computing log_softmax using the global statistics
+
+    Args:
+        logits: Logits tensor [..., partition_vocab_size]
+
+    Returns:
+        log_softmax tensor [..., partition_vocab_size]
+    """
+    tp_size = mpu.get_tensor_model_parallel_world_size()
+
+    if tp_size == 1:
+        return torch.nn.functional.log_softmax(logits, dim=-1)
+
+    tp_group = mpu.get_tensor_model_parallel_group()
+
+    # Step 1: Find global max for numerical stability
+    logits_max = logits.max(dim=-1, keepdim=True)[0]
+    torch.distributed.all_reduce(logits_max, op=torch.distributed.ReduceOp.MAX, group=tp_group)
+
+    # Step 2: Compute exp(logits - max) and sum across all TP ranks
+    exp_logits = torch.exp(logits - logits_max)
+    sum_exp = exp_logits.sum(dim=-1, keepdim=True)
+    torch.distributed.all_reduce(sum_exp, op=torch.distributed.ReduceOp.SUM, group=tp_group)
+
+    # Step 3: Compute log_softmax
+    log_softmax = logits - logits_max - torch.log(sum_exp)
+
+    return log_softmax
+
+
+def vocab_parallel_entropy(log_probs: torch.Tensor, chunk_size: int = 512) -> torch.Tensor:
+    """Compute entropy from pre-computed vocab-parallel sharded log probabilities.
+
+    When using Tensor Parallelism, vocab is sharded across TP ranks.
+    This function correctly computes entropy by:
+    1. Computing partial entropy = -sum(exp(log_p) * log_p) on each rank's partition
+    2. All-reducing the partial entropies to get the global sum.
+
+    Entropy is computed in chunks to reduce memory usage.
+
+    Args:
+        log_probs: Pre-computed log probabilities tensor [..., partition_vocab_size]
+        chunk_size: Number of tokens to process per chunk (default: 512)
+
+    Returns:
+        Entropy tensor [...] (scalar per position)
+    """
+    tp_group = mpu.get_tensor_model_parallel_group()
+    tp_size = mpu.get_tensor_model_parallel_world_size()
+
+    # Flatten all but the last dimension for chunked processing
+    original_shape = log_probs.shape[:-1]
+    vocab_size = log_probs.shape[-1]
+    log_probs_flat = log_probs.view(-1, vocab_size)  # [total_tokens, partition_vocab_size]
+    total_tokens = log_probs_flat.shape[0]
+
+    entropies_list = []
+    for start_idx in range(0, total_tokens, chunk_size):
+        end_idx = min(start_idx + chunk_size, total_tokens)
+        log_probs_chunk = log_probs_flat[start_idx:end_idx]  # [chunk_size, partition_vocab_size]
+
+        # Compute partial entropy on this rank's vocab partition
+        # entropy = -sum(p * log_p) = -sum(exp(log_p) * log_p)
+        probs = torch.exp(log_probs_chunk)
+        partial_entropy = -(probs * log_probs_chunk).sum(dim=-1)  # [chunk_size]
+
+        # All-reduce to get global entropy if using TP
+        if tp_size > 1:
+            torch.distributed.all_reduce(partial_entropy, op=torch.distributed.ReduceOp.SUM, group=tp_group)
+
+        entropies_list.append(partial_entropy)
+
+    # Concatenate all chunks and reshape back
+    entropies = torch.cat(entropies_list, dim=0)
+    entropies = entropies.view(original_shape)
+
+    return entropies
+
+
+def vocab_parallel_kl_div(input_log_probs: torch.Tensor, target_log_probs: torch.Tensor) -> torch.Tensor:
+    """Compute KL divergence for vocab-parallel sharded log probabilities.
+
+    KL(target || input) = sum(target_prob * (target_log_prob - input_log_prob))
+                        = sum(exp(target_log_prob) * (target_log_prob - input_log_prob))
+
+    Since both log_probs are sharded across TP, we compute the partial sum
+    on each rank and then all_reduce to get the global sum.
+
+    Args:
+        input_log_probs: Input log probabilities [..., partition_vocab_size]
+        target_log_probs: Target log probabilities [..., partition_vocab_size]
+
+    Returns:
+        KL divergence per position [...], already reduced across TP
+    """
+    tp_group = mpu.get_tensor_model_parallel_group()
+
+    # Compute partial KL on this rank's vocab partition
+    target_probs = torch.exp(target_log_probs)
+    partial_kl = (target_probs * (target_log_probs - input_log_probs)).sum(dim=-1)
+
+    if mpu.get_tensor_model_parallel_world_size() > 1:
+        tp_group = mpu.get_tensor_model_parallel_group()
+        torch.distributed.all_reduce(partial_kl, op=torch.distributed.ReduceOp.SUM, group=tp_group)
+
+    return partial_kl
+
+
+def vocab_parallel_gather_logps(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+    log_probs: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Gather log probabilities for target labels from vocab-parallel logits.
+
+    When using TP, each rank only has a partition of the vocabulary. This function:
+    1. Computes log_softmax with vocab-parallel support (if log_probs not provided)
+    2. Gathers log probs for target tokens
+    3. All-reduces to combine results from all TP ranks
+
+    Args:
+        logits: Logits tensor [batch, seq, partition_vocab_size]
+        labels: Token labels [batch, seq], -100 for masked positions
+        log_probs: Pre-computed log_softmax (optional, to avoid recomputation)
+
+    Returns:
+        per_token_logps: [batch, seq] log probabilities for target tokens
+    """
+    tp_size = mpu.get_tensor_model_parallel_world_size()
+    tp_rank = mpu.get_tensor_model_parallel_rank()
+
+    # Compute log_probs if not provided
+    if log_probs is None:
+        log_probs = vocab_parallel_log_softmax(logits)
+
+    # Get the local vocab range for this TP rank
+    partition_vocab_size = log_probs.shape[-1]
+    if tp_size > 1:
+        vocab_start = tp_rank * partition_vocab_size
+        vocab_end = vocab_start + partition_vocab_size
+        # Check which labels fall within this TP rank's vocab partition
+        local_labels = labels - vocab_start
+        # Mask for labels within local vocab range
+        in_range_mask = (labels >= vocab_start) & (labels < vocab_end)
+        # Clamp local_labels to valid range for gather
+        local_labels = local_labels.clamp(min=0, max=partition_vocab_size - 1)
+    else:
+        local_labels = labels
+        in_range_mask = torch.ones_like(labels, dtype=torch.bool)
+        local_labels = local_labels.clamp(min=0)
+
+    # Gather log probs for target tokens
+    gathered_logps = torch.gather(log_probs, dim=-1, index=local_labels.unsqueeze(-1)).squeeze(-1)
+
+    # For TP: only the rank that owns the target token has the correct log prob
+    # Other ranks have incorrect values, so we need to zero them out
+    if tp_size > 1:
+        gathered_logps = gathered_logps * in_range_mask.float()
+        # All-reduce to sum contributions from all ranks
+        # (only one rank has non-zero value for each token)
+        torch.distributed.all_reduce(
+            gathered_logps, op=torch.distributed.ReduceOp.SUM, group=mpu.get_tensor_model_parallel_group())
+
+    # Apply loss mask (labels == -100 are masked)
+    loss_mask = labels != -100
+    per_token_logps = gathered_logps * loss_mask.float()
+
+    return per_token_logps
+
+
+def compute_logps_and_entropy_from_logits(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+    compute_entropy: bool = False,
+    entropy_chunk_size: int = 512,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """Compute per-token log probabilities and optionally entropy from logits.
+
+    This is a unified method that efficiently computes both logps and entropy
+    in a single pass when both are needed, sharing the log_softmax computation.
+
+    Note: In Megatron, labels are already shifted (via torch.roll in get_batch_on_this_tp_rank),
+    so logits and labels are already aligned. No additional shift is needed here.
+
+    Temperature scaling should be applied by the caller before invoking this function,
+    so that this function remains a pure computation without side effects on the input.
+
+    Args:
+        logits: Logits tensor [batch, seq, partition_vocab_size] or [1, total_tokens, partition_vocab_size].
+                Should be pre-scaled by temperature if needed.
+        labels: Token labels [batch, seq] or [1, total_tokens], -100 for masked positions
+        compute_entropy: Whether to compute entropy (default: False)
+        entropy_chunk_size: Chunk size for entropy computation (default: 512)
+
+    Returns:
+        Tuple of:
+            - per_token_logps: [batch, seq] or [1, total_tokens] log probabilities for target tokens
+            - per_token_entropy: Same shape as per_token_logps, or None if compute_entropy=False
+    """
+    # Compute log_softmax (shared for both logps and entropy)
+    log_probs = vocab_parallel_log_softmax(logits)
+
+    # Gather logps for target tokens
+    per_token_logps = vocab_parallel_gather_logps(logits, labels, log_probs=log_probs)
+
+    # Compute entropy if requested (reuse log_probs to avoid redundant computation)
+    per_token_entropy = None
+    if compute_entropy:
+        per_token_entropy = vocab_parallel_entropy(log_probs, chunk_size=entropy_chunk_size)
+
+    return per_token_logps, per_token_entropy

swift/megatron/utils/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+from .convert_utils import test_convert_precision
+from .megatron_lm_utils import (disable_forward_pre_hook, enable_forward_pre_hook, get_batch_on_this_cp_rank,
+                                get_optimizer_param_scheduler, init_persistent_async_worker, initialize_megatron,
+                                initialize_tp_communicators, load_mcore_checkpoint, maybe_finalize_async_save,
+                                save_mcore_checkpoint, should_disable_forward_pre_hook, warmup_jit_function, wrap_model)
+from .parallel_utils import logical_and_across_model_parallel_group, reduce_max_stat_across_model_parallel_group
+from .patcher import patch_merge_fn, patch_torch_dist_shard
+from .router_replay_utils import (RouterReplayHelper, apply_router_replay_patch, get_local_topk_idx_for_current_rank,
+                                  get_router_replay_data, set_router_replay_data)
+from .utils import (forward_step_helper, get_packed_seq_params, get_padding_to, prepare_mcore_model,
+                    reconstruct_tensor_cp)

swift/megatron/utils/convert_utils.py

@@ -0,0 +1,303 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+
+import math
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from contextlib import contextmanager, nullcontext
+from megatron.core import mpu
+from megatron.core.extensions.transformer_engine import TEDotProductAttention
+from megatron.core.ssm.mamba_context_parallel import _undo_attention_load_balancing
+from megatron.core.tensor_parallel import VocabParallelEmbedding
+from megatron.core.tensor_parallel.mappings import (gather_from_sequence_parallel_region,
+                                                    gather_from_tensor_model_parallel_region)
+from typing import Any, Dict
+
+from swift.utils import HfConfigFactory, get_logger, to_device, to_float_dtype
+from .megatron_lm_utils import get_batch_on_this_cp_rank
+from .utils import forward_step_helper, get_packed_seq_params, get_padding_to
+
+logger = get_logger()
+
+
+def _test_params_sum(model):
+    total_sum = 0
+    zero_count = 0
+    n_parameter = 0
+    for n, p in model.named_parameters():
+        n_parameter += 1
+        sum_ = p.to(device='cuda', dtype=torch.float32).abs().sum().cpu().item()
+        if sum_ == 0 and '.lora_B.' not in n:
+            zero_count += 1
+            logger.warning(f'n: {n}, sum: {sum_}')
+        elif math.isnan(sum_) or math.isinf(sum_) or sum_ > 1e10:
+            logger.warning(f'n: {n}, sum: {sum_}')
+        else:
+            total_sum += sum_
+    cond = mpu.get_data_parallel_rank() == 0
+    logger.info_if(f'n_parameter: {n_parameter}', cond=cond)
+    logger.info_if(f'total_sum: {total_sum}', cond=cond)
+    logger.info_if(f'zero_count: {zero_count}', cond=cond)
+
+
+def _find_modules(model, recurse: bool = True, prefix='', ignore_modules=None):
+    ignore_modules = ignore_modules or []
+    for k in ignore_modules:
+        if prefix.startswith(k):
+            return []
+    else:
+        named_children = list(model.named_children())
+
+    modules = []
+    for n, module in named_children:
+        if module.__class__ is nn.ModuleList:
+            modules += _find_modules(module, False, prefix=f'{prefix}{n}.', ignore_modules=ignore_modules)
+        elif recurse:
+            modules += _find_modules(module, prefix=f'{prefix}{n}.', ignore_modules=ignore_modules)
+        else:
+            modules.append(module)
+    if not named_children:
+        modules.append(model)
+    return modules
+
+
+@contextmanager
+def _model_cpu_forward_context(modules,
+                               torch_dtype=None,
+                               compute_device=None,
+                               share_embedding: bool = False,
+                               target_device='cpu'):
+    for module in modules:
+        try:
+            origin_torch_dtype = next(module.parameters()).dtype
+        except StopIteration:
+            pass
+        else:
+            break
+    embeddings = None
+    if share_embedding:
+        embeddings = [module for module in modules if isinstance(module, (nn.Embedding, VocabParallelEmbedding))]
+
+    def _to_cuda_hook(module, args):
+        if compute_device is not None or torch_dtype is not None:
+            module.to(device=compute_device, dtype=torch_dtype)
+            args = to_float_dtype(args, dtype=torch_dtype)
+        return args
+
+    def _to_cpu_hook(module, args, output):
+        if share_embedding and module in embeddings or 'rotaryemb' in module.__class__.__name__.lower():
+            return
+        module.to(device=target_device, dtype=origin_torch_dtype)
+
+    hooks = []
+    for module in modules:
+        hooks.append(module.register_forward_pre_hook(_to_cuda_hook))
+        hooks.append(module.register_forward_hook(_to_cpu_hook))
+    try:
+        yield
+    finally:
+        for hook in hooks:
+            hook.remove()
+
+
+def get_examples(mm_type: str) -> Dict[str, Any]:
+    if mm_type == 'image':
+        data = {
+            'messages': [{
+                'role': 'user',
+                'content': '<image>describe the image.'
+            }, {
+                'role':
+                'assistant',
+                'content':
+                'The image depicts a close-up of a kitten with striking features. '
+                'The kitten has a white and gray coat with distinct black stripes, '
+                'particularly noticeable on its face and ears. Its eyes are large '
+                'and expressive, with a captivating blue hue that stands out against '
+                "the darker fur around them. The kitten's nose is small and pink, "
+                'and it has long, delicate whiskers extending from either side of its mouth. '
+                "The background is blurred, drawing attention to the kitten's face and "
+                'making it the focal point of the image. The overall impression is '
+                'one of cuteness and charm.'
+            }],
+            'images': ['http://modelscope-open.oss-cn-hangzhou.aliyuncs.com/images/cat.png']
+        }
+    elif mm_type == 'audio':
+        data = {
+            'messages': [{
+                'role': 'user',
+                'content': '<audio>Caption the audio.'
+            }, {
+                'role': 'assistant',
+                'content': "The audio contains a male voice speaking the phrase '今天天气真好呀' in Mandarin."
+            }],
+            'audios': ['http://modelscope-open.oss-cn-hangzhou.aliyuncs.com/images/weather.wav']
+        }
+    else:  # text
+        data = {
+            'messages': [
+                {
+                    'role': 'user',
+                    'content': 'Introduction to ms-swift.'
+                },
+                {
+                    'role':
+                    'assistant',
+                    'content':
+                    'ms-swift is an official framework provided by the ModelScope community for fine-tuning '
+                    'and deploying large language models and multi-modal large models.'
+                },
+            ]
+        }
+    return data
+
+
+def broadcast_mg_logits(mg_logits=None, src_rank=None):
+    if not dist.is_initialized():
+        return
+    rank = dist.get_rank()
+    if src_rank is None:
+        src_rank = dist.get_world_size() - 1
+    if rank == src_rank:
+        meta = [tuple(mg_logits.shape), str(mg_logits.dtype).split('.', 1)[1]]
+    else:
+        meta = [None, None]
+
+    dist.broadcast_object_list(meta, src=src_rank)
+    shape, dtype = meta
+    dtype = getattr(torch, dtype)
+
+    if rank != src_rank:
+        mg_logits = torch.empty(shape, dtype=dtype, device='cuda')
+
+    dist.broadcast(mg_logits, src=src_rank)
+
+    return mg_logits
+
+
+@contextmanager
+def _patch_attention_fp32(compute_dtype):
+    forward = TEDotProductAttention.forward
+
+    def new_forward(self, query_layer, key_layer, value_layer, *args, **kwargs):
+        torch_dtype = query_layer.dtype
+        query_layer = query_layer.to(compute_dtype)
+        key_layer = key_layer.to(compute_dtype)
+        value_layer = value_layer.to(compute_dtype)
+        res = forward(self, query_layer, key_layer, value_layer, *args, **kwargs)
+        res = res.to(dtype=torch_dtype)
+        return res
+
+    TEDotProductAttention.forward = new_forward
+    try:
+        yield
+    finally:
+        TEDotProductAttention.forward = forward
+
+
+def test_convert_precision(args, hf_model, mg_model, template, test_convert_dtype=None):
+    if test_convert_dtype is None:
+        test_convert_dtype = getattr(args, 'test_convert_dtype', torch.float32)
+    template.set_mode('train')
+    _test_params_sum(mg_model)
+
+    config = mg_model.config
+    is_multimodal = config.is_multimodal
+    if is_multimodal:
+        test_mm_type = getattr(config, 'test_mm_type', 'image')
+    else:
+        test_mm_type = 'text'
+    mg_language_model = mg_model.language_model if is_multimodal else mg_model
+    if mg_language_model.config.fp8 is not None:
+        raise ValueError('fp8 models currently do not support testing convert_precision. '
+                         'Please set `--test_convert_precision false`.')
+    share_embedding = mg_language_model.share_embeddings_and_output_weights
+    if hf_model is not None:
+        hf_model.eval()
+        if dist.get_world_size() == 1:
+            _test_params_sum(hf_model)
+        inputs = template.encode(get_examples(test_mm_type), return_length=True)
+        hf_inputs = to_device(template.data_collator([inputs]), 'cuda')
+        template.register_post_encode_hook([hf_model])
+        HfConfigFactory.set_config_attr(hf_model.config, 'use_cache', False)
+        model_arch = hf_model.model_meta.model_arch
+        ignore_modules = (model_arch.vision_tower + model_arch.aligner) if is_multimodal else []
+        hf_modules = _find_modules(hf_model, ignore_modules=ignore_modules)
+        with torch.inference_mode(), _model_cpu_forward_context(
+                hf_modules, test_convert_dtype, share_embedding=share_embedding):
+            hf_inputs.pop('text_position_ids', None)
+            hf_logits = hf_model(**hf_inputs).logits
+            hf_logits = hf_logits.to('cuda')
+        hf_model.to('cpu')
+
+    template.use_megatron = True
+    inputs = [
+        template.encode(get_examples(test_mm_type), return_length=True) for _ in range(2 if args.padding_free else 1)
+    ]
+    mg_inputs = to_device(template.data_collator(inputs, padding_to=get_padding_to(args)), 'cuda')
+    mg_model.eval()
+    # thd
+    text_position_ids = mg_inputs.pop('text_position_ids', None)
+    if text_position_ids is None:
+        text_position_ids = mg_inputs.get('position_ids')
+    if args.padding_free:
+        mg_inputs['packed_seq_params'] = get_packed_seq_params(args, text_position_ids)
+    mg_language_model.config.fp8 = None  # compat fp8
+    mg_modules = _find_modules(mg_language_model, ignore_modules=['visual'])
+    for key in ['labels', 'seq_lens', 'attention_mask_2d']:
+        mg_inputs.pop(key, None)
+    mg_inputs = get_batch_on_this_cp_rank(args, mg_inputs)
+    _param = next(mg_language_model.parameters())
+    mg_dtype = _param.dtype
+    mg_device = _param.device
+    if args.model_type == 'minimax_m2':
+        # router to bfloat16 (expert_bias). No need to do this when actually training.
+        for n, m in mg_language_model.named_modules():
+            if n.endswith('router'):
+                m.to(mg_dtype)
+    if getattr(config, 'enable_hyper_connections', False):
+        for param in mg_language_model.decoder.parameters(recurse=False):
+            param.data = param.data.cuda()
+    attention_context = (
+        _patch_attention_fp32(mg_dtype) if args.attention_backend.name in {'flash', 'fused'} else nullcontext())
+    with torch.inference_mode(), _model_cpu_forward_context(
+            mg_modules, test_convert_dtype, 'cuda', share_embedding=share_embedding,
+            target_device=mg_device), attention_context:
+        mg_logits = forward_step_helper(mg_model, mg_inputs, dtype=test_convert_dtype)
+        if args.tensor_model_parallel_size > 1 and args.task_type != 'seq_cls':
+            if mg_logits is not None:
+                mg_logits = gather_from_tensor_model_parallel_region(mg_logits)
+        if args.context_parallel_size > 1:
+            if mg_logits is not None:
+                mg_logits = gather_from_sequence_parallel_region(
+                    mg_logits.transpose(0, 1), group=mpu.get_context_parallel_group())
+                # Contiguous CP already gathers ranks in original token order, so the
+                # zigzag un-balancing must be skipped (it is only correct for zigzag split).
+                if getattr(args, 'cp_partition_mode', 'zigzag') != 'contiguous':
+                    mg_logits = _undo_attention_load_balancing(mg_logits, args.context_parallel_size)
+                mg_logits = mg_logits.transpose(0, 1)
+    mg_logits = broadcast_mg_logits(mg_logits)
+    if hf_model is None:
+        return
+    if args.task_type == 'seq_cls':
+        mg_logits = mg_logits[:, -1]
+        mean_diff = (mg_logits - hf_logits).abs().mean().item()
+        max_diff = (mg_logits - hf_logits).abs().max().item()
+        print(f'mean_diff: {mean_diff}, max_diff: {max_diff}')
+    else:
+        mg_logits = mg_logits[:, :hf_logits.shape[1]]
+        token_mean_diff = (mg_logits - hf_logits).abs().mean(dim=-1)
+        mean_diff = token_mean_diff.mean().item()
+        max_diff = (mg_logits - hf_logits).abs().max().item()
+        loss_mask = (torch.roll(hf_inputs['labels'], -1) != -100)
+        mean_diff_with_loss = token_mean_diff[loss_mask].mean().item()
+        max_diff_with_loss = (mg_logits - hf_logits)[loss_mask].abs().max().item()
+        print(f'token_mean_diff: {token_mean_diff}')
+        print(f'mean_diff: {mean_diff}, max_diff: {max_diff}')
+        print(f'mean_diff (with loss): {mean_diff_with_loss}, max_diff (with loss): {max_diff_with_loss} '
+              '(Please check that mean_diff (with loss) is less than 0.1).')
+        hf_tokens = hf_logits.argmax(-1)
+        mg_tokens = mg_logits.argmax(-1)
+        print(f'hf_tokens: {hf_tokens[0].tolist()}\nmg_tokens: {mg_tokens[0].tolist()}')
+        print(f'token_diff: {(hf_tokens != mg_tokens).sum().item()}')
+        print(f'token_diff (with loss): {(hf_tokens[loss_mask] != mg_tokens[loss_mask]).sum().item()}')

swift/megatron/utils/megatron_lm_utils.py

@@ -0,0 +1,743 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+# Parts of the functions in this file are code borrowed from NVIDIA/Megatron-LM
+import copy
+import dataclasses
+import megatron.core
+import numpy as np
+import os
+import random
+import torch
+from argparse import Namespace
+from contextlib import contextmanager
+from datetime import timedelta
+from mcore_bridge import set_random_seed, split_cp_inputs, unwrap_model
+from megatron.core import dist_checkpointing, mpu, parallel_state, tensor_parallel
+from megatron.core.dist_checkpointing.mapping import ShardedObject
+from megatron.core.dist_checkpointing.serialization import (get_default_load_sharded_strategy,
+                                                            get_default_save_sharded_strategy)
+from megatron.core.dist_checkpointing.strategies.async_utils import AsyncCallsQueue, AsyncRequest
+from megatron.core.dist_checkpointing.strategies.fully_parallel import (FullyParallelLoadStrategyWrapper,
+                                                                        FullyParallelSaveStrategyWrapper)
+from megatron.core.dist_checkpointing.strategies.torch import TorchDistLoadShardedStrategy, TorchDistSaveShardedStrategy
+from megatron.core.distributed import DistributedDataParallel as DDP
+from megatron.core.distributed import DistributedDataParallelConfig
+from megatron.core.fusions.fused_bias_dropout import bias_dropout_add_fused_train
+from megatron.core.fusions.fused_bias_gelu import bias_gelu
+from megatron.core.fusions.fused_bias_swiglu import bias_swiglu
+from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
+from megatron.core.transformer.module import Float16Module
+from megatron.core.utils import get_torch_version, is_te_min_version, is_torch_min_version
+from packaging import version
+from transformers.utils import is_torch_npu_available
+from typing import Any, Dict, Optional
+
+from swift.utils import check_json_format, get_logger, init_process_group, is_master, set_device
+from .patcher import patch_merge_fn
+
+logger = get_logger()
+
+mcore_017 = version.parse(megatron.core.__version__) >= version.parse('0.17.0rc0')
+
+
+@contextmanager
+def _patch_megatron_timeout(distributed_timeout_minutes):
+    origin_create_group = parallel_state.create_group
+
+    def create_group(ranks=None, timeout=None, *_args, **kwargs):
+        if timeout is None:
+            timeout = timedelta(minutes=distributed_timeout_minutes)
+        return origin_create_group(ranks, timeout, *_args, **kwargs)
+
+    parallel_state.create_group = create_group
+    try:
+        yield
+    finally:
+        parallel_state.create_group = origin_create_group
+
+
+def _initialize_mpu(args):
+    """Initialize torch.distributed and core model parallel."""
+    if not torch.distributed.is_initialized():
+        set_device()
+        init_process_group(args.ddp_backend, args.ddp_timeout)
+    args.rank = torch.distributed.get_rank()
+    args.world_size = torch.distributed.get_world_size()
+
+    if mpu.model_parallel_is_initialized():
+        logger.info('model parallel is already initialized')
+    else:
+        distributed_timeout_minutes = args.ddp_timeout // 60
+        with _patch_megatron_timeout(distributed_timeout_minutes):
+            mpu.initialize_model_parallel(
+                args.tensor_model_parallel_size,
+                args.pipeline_model_parallel_size,
+                args.virtual_pipeline_model_parallel_size,
+                context_parallel_size=args.context_parallel_size,
+                expert_model_parallel_size=args.expert_model_parallel_size,
+                expert_tensor_parallel_size=args.expert_tensor_parallel_size,
+                distributed_timeout_minutes=distributed_timeout_minutes,
+            )
+        if is_master():
+            logger.info(f'TP: {args.tensor_model_parallel_size}, PP: {args.pipeline_model_parallel_size}, '
+                        f'VPP: {args.virtual_pipeline_model_parallel_size}, CP: {args.context_parallel_size}, '
+                        f'EP: {args.expert_model_parallel_size}, ETP: {args.expert_tensor_parallel_size}')
+
+
+def initialize_megatron(args):
+    # Pytorch distributed.
+    _initialize_mpu(args)
+
+    # Random seeds for reproducibility.
+    logger.info(f'Setting random seeds to {args.seed}.')
+    set_random_seed(args.seed, args.data_parallel_random_init, args.te_rng_tracker)
+
+    # Setup MoE aux loss scale value.
+    if args.model_info.is_moe_model:
+        from megatron.core.transformer.moe.router import MoEAuxLossAutoScaler
+        MoEAuxLossAutoScaler.set_loss_scale(torch.ones(1, device=torch.cuda.current_device()))
+
+
+def _get_rng_state():
+    """Collect rng state across data parallel ranks."""
+    rng_state = {
+        'random_rng_state': random.getstate(),
+        'np_rng_state': np.random.get_state(),
+        'torch_rng_state': torch.get_rng_state(),
+        'cuda_rng_state': torch.cuda.get_rng_state(),
+        'rng_tracker_states': tensor_parallel.get_cuda_rng_tracker().get_states()
+    }
+
+    # data_parallel_random_init False
+    rng_state_list = [rng_state]
+
+    pp_rank = mpu.get_pipeline_model_parallel_rank()
+    pp_size = mpu.get_pipeline_model_parallel_world_size()
+    tp_rank = mpu.get_tensor_model_parallel_rank()
+    tp_size = mpu.get_tensor_model_parallel_world_size()
+    rng_state_list = ShardedObject(
+        'rng_state',
+        rng_state_list, (pp_size, tp_size), (pp_rank, tp_rank),
+        replica_id=mpu.get_data_parallel_rank(with_context_parallel=True))
+    return rng_state_list
+
+
+def _generate_state_dict(args,
+                         models,
+                         optimizer=None,
+                         opt_param_scheduler=None,
+                         rng_state=None,
+                         iteration=None,
+                         model_sd_kwargs=None,
+                         optim_sd_kwargs=None):
+    model_sd_kwargs = model_sd_kwargs or {}
+    state_dict = {
+        'args': Namespace(**check_json_format(vars(args))),
+        'checkpoint_version': 3.0,
+    }
+    if iteration is not None:
+        state_dict['iteration'] = iteration
+    for i, m in enumerate(models):
+        key = 'model'
+        if len(models) > 1:
+            key = f'model{i}'
+        model_sd = models[i].sharded_state_dict(**model_sd_kwargs)
+        state_dict[key] = model_sd
+
+    if not args.no_save_optim:
+        if optimizer is not None:
+            state_dict['optimizer'] = _optimizer_sharded_state_dict(optimizer, state_dict, optim_sd_kwargs or {})
+        if opt_param_scheduler is not None:
+            state_dict['opt_param_scheduler'] = opt_param_scheduler.state_dict()
+
+    if not args.no_save_rng and rng_state is not None:
+        state_dict['rng_state'] = rng_state
+    return state_dict
+
+
+def _optimizer_sharded_state_dict(optimizer, state_dict, optim_sd_kwargs):
+    if is_torch_npu_available():
+        from swift.model.npu_patch.megatron_checkpoint import optimizer_sharded_state_dict
+        return optimizer_sharded_state_dict(optimizer, state_dict, **optim_sd_kwargs)
+    return optimizer.sharded_state_dict(state_dict, **optim_sd_kwargs)
+
+
+def _load_optimizer_state_dict(optimizer, state_dict):
+    if is_torch_npu_available():
+        from swift.model.npu_patch.megatron_checkpoint import load_optimizer_state_dict
+        load_optimizer_state_dict(optimizer, state_dict)
+        return
+    optimizer.load_state_dict(state_dict)
+
+
+def _filter_adapter_state_dict(state_dict, peft_format: bool, adapter_name: str = 'default'):
+    """
+    When peft_format is True, keep only the PEFT format state_dict;
+    when False, remove the PEFT format state_dict.
+
+    This function ensures it is called when tuner_type != 'full'.
+    """
+    if 'model' in state_dict:
+        n_models = 1
+    else:
+        n_models = 0
+        while f'model{n_models}' in state_dict:
+            n_models += 1
+
+    for i in range(n_models):
+        if i == 0 and n_models == 1:
+            model_key = 'model'
+        else:
+            model_key = f'model{i}'
+        new_state_dict = {}
+        state_dict_model = state_dict[model_key]
+        for k, v in state_dict_model.items():
+            if peft_format:
+                if '.lora_A.' in k or '.lora_B.' in k or '.modules_to_save.' in k:
+                    new_state_dict[k] = v
+            else:
+                if '.lora_A.' in k or '.lora_B.' in k or 'original_module.' in k:
+                    continue
+                k = k.replace('base_layer.', '')
+                k = k.replace(f'modules_to_save.{adapter_name}.', '')
+                v.key = v.key.replace('base_layer.', '')
+                v.key = v.key.replace(f'modules_to_save.{adapter_name}.', '')
+                new_state_dict[k] = v
+        state_dict[model_key] = new_state_dict
+
+
+def _preprocess_common_before_consistancy_check(common_state_dict):
+    # Convert args key of type namespace to dictionary
+    preprocessed_common_state_dict = copy.deepcopy(common_state_dict)
+    preprocessed_common_state_dict['args'] = vars(preprocessed_common_state_dict['args'])
+    # Remove rank and local rank from state dict if it exists, since they are expected to be different
+    preprocessed_common_state_dict['args'].pop('local_rank', None)
+    preprocessed_common_state_dict['args'].pop('rank', None)
+    return preprocessed_common_state_dict
+
+
+def get_sharded_sd_metadata(args):
+    sharded_sd_metadata = {'singleton_local_shards': False, 'chained_optim_avoid_prefix': True}
+    force_pre_mcore_014 = not is_torch_min_version('2.6a0')
+    if force_pre_mcore_014 and not args.dist_ckpt_save_pre_mcore_014:
+        args.dist_ckpt_save_pre_mcore_014 = True
+        logger.warning(f'PyTorch version {get_torch_version()} below 2.6 detected.'
+                       f' Forcing dist_ckpt_save_pre_mcore_014 behavior.')
+
+    if args.dist_ckpt_save_pre_mcore_014:
+        sharded_sd_metadata['distrib_optim_sharding_type'] = 'fully_sharded_model_space'
+    else:
+        if args.dist_ckpt_optim_fully_reshardable:
+            sharded_sd_metadata['distrib_optim_sharding_type'] = 'fully_reshardable'
+            sharded_sd_metadata[
+                'distrib_optim_fully_reshardable_mem_efficient'] = args.distrib_optim_fully_reshardable_mem_efficient
+        else:
+            sharded_sd_metadata['distrib_optim_sharding_type'] = 'dp_reshardable'
+    return sharded_sd_metadata
+
+
+def save_mcore_checkpoint(
+    args,
+    models,
+    optimizer=None,
+    opt_param_scheduler=None,
+    iteration=1,
+    output_dir: Optional[str] = None,
+    peft_format: bool = False,
+):
+    if output_dir is None:
+        output_dir = args.output_dir
+    models = unwrap_model(models)
+    rng_state = _get_rng_state() if models else None
+    checkpoint_dir = os.path.join(output_dir, f'iter_{iteration:07d}')
+    sharded_sd_metadata = get_sharded_sd_metadata(args)
+    os.makedirs(checkpoint_dir, exist_ok=True)
+
+    state_dict = _generate_state_dict(
+        args,
+        models,
+        optimizer,
+        opt_param_scheduler,
+        rng_state,
+        iteration=iteration,
+        model_sd_kwargs={'metadata': sharded_sd_metadata},
+        optim_sd_kwargs={'metadata': sharded_sd_metadata},
+    )
+    _filter_adapter_state_dict(state_dict, peft_format)
+    if mcore_017:
+        save_strategy = TorchDistSaveShardedStrategy()
+    else:
+        save_strategy = get_default_save_sharded_strategy()
+    save_strategy = FullyParallelSaveStrategyWrapper(
+        save_strategy,
+        mpu.get_data_parallel_group(with_context_parallel=True),
+    )
+    kwargs = {'content_metadata': sharded_sd_metadata}
+    async_save = args.async_save
+    if not models:  # save GPU memory
+        assert 'optimizer' not in state_dict
+        async_save = False
+        common_path = os.path.join(checkpoint_dir, 'common.pt')
+        if is_master():
+            state_dict.update(kwargs)
+            torch.save(state_dict, common_path)
+        async_save_request = None
+    else:
+        async_save_request = dist_checkpointing.save(
+            state_dict,
+            checkpoint_dir,
+            save_strategy,
+            async_sharded_save=async_save,
+            validate_access_integrity=True,
+            preprocess_common_before_consistancy_check=_preprocess_common_before_consistancy_check,
+            **kwargs)
+
+    if not async_save:
+        assert async_save_request is None
+        # Wait so everyone is done (necessary)
+        if torch.distributed.is_initialized():
+            torch.distributed.barrier()
+
+    if is_master():
+
+        tracker_path = os.path.join(output_dir, 'latest_checkpointed_iteration.txt')
+        try:
+            from megatron.core.msc_utils import open_file
+        except ImportError:
+            open_file = open
+        with open_file(tracker_path, 'w') as f:
+            f.write(str(iteration))
+
+        def iter_finalize_fn():
+            logger.info(f'Successfully saved Megatron model weights in `{output_dir}`.')
+
+        if async_save:
+            assert async_save_request is not None
+            async_save_request.add_finalize_fn(iter_finalize_fn)
+        else:
+            iter_finalize_fn()
+
+    if async_save:
+        schedule_async_save(async_save_request)
+
+
+# Singleton manager of async calls
+# The default is `TemporalAsyncCaller`
+_async_calls_queue = AsyncCallsQueue()
+
+
+def init_persistent_async_worker():
+    global _async_calls_queue
+    # Recreate the async_calls_queue for persistent worker
+    # This duplicate step is for backward compatiblity
+    _async_calls_queue = AsyncCallsQueue(persistent=True)
+
+
+def schedule_async_save(async_request: AsyncRequest):
+    """Schedule the async save request.
+
+    Args:
+        async_request (AsyncRequest): the async save request.
+    """
+    _async_calls_queue.schedule_async_request(async_request)
+
+
+def maybe_finalize_async_save(args, blocking: bool = False, terminate=False):
+    """Finalizes active async save calls.
+
+    Args:
+        blocking (bool, optional): if True, will wait until all active requests
+            are done. Otherwise, finalizes only the async request that already
+            finished. Defaults to False.
+        terminate (bool, optional): if True, the asynchronous queue will
+                be closed as the last action of this function.
+    """
+    if not args.async_save:
+        return
+
+    _async_calls_queue.maybe_finalize_async_calls(blocking, no_dist=False)
+
+    if terminate:
+        _async_calls_queue.close()
+
+
+def is_empty_async_queue() -> bool:
+    """Check if async calls queue is empty. This result is consistent across ranks."""
+    return _async_calls_queue.get_num_unfinalized_calls() == 0
+
+
+def _load_iteration(tracker_path: str):
+    if not os.path.exists(tracker_path):
+        return 0
+    with open(tracker_path, 'r') as f:
+        iteration = int(f.read())
+    # Get the max iteration retrieved across the ranks.
+    if torch.distributed.is_initialized():
+        iters_cuda = torch.tensor([iteration], dtype=torch.long, device='cuda')
+        torch.distributed.all_reduce(iters_cuda, op=torch.distributed.ReduceOp.MAX)
+        iteration = iters_cuda[0].item()
+    return iteration
+
+
+def load_mcore_checkpoint(args,
+                          ddp_models: list,
+                          optimizer=None,
+                          opt_param_scheduler=None,
+                          load_arg: str = 'mcore_model',
+                          adapter_name: str = 'default'):
+    if load_arg in {'mcore_adapter', 'mcore_ref_adapter'}:
+        peft_format = True
+    else:
+        # 'mcore_model', 'mcore_ref_model'
+        peft_format = False
+    load_dir = getattr(args, load_arg)
+
+    no_load_optim = args.no_load_optim
+    no_load_rng = args.no_load_rng
+    finetune = args.finetune
+    if not peft_format and args.tuner_type != 'full':
+        # When training with LoRA and loading the base model
+        no_load_optim = True
+        no_load_rng = True
+        finetune = True
+    models = unwrap_model(ddp_models)
+    tracker_path = os.path.join(load_dir, 'latest_checkpointed_iteration.txt')
+    iteration = _load_iteration(tracker_path)
+    checkpoint_dir = os.path.join(load_dir, f'iter_{iteration:07d}')
+    state_dict = dist_checkpointing.load_common_state_dict(checkpoint_dir)
+
+    ckpt_tp_pp = (
+        state_dict['args'].tensor_model_parallel_size,
+        state_dict['args'].pipeline_model_parallel_size,
+    )
+    run_tp_pp = (
+        args.tensor_model_parallel_size,
+        args.pipeline_model_parallel_size,
+    )
+    mismatch_msg = f'(TP, PP) mismatch after resume ({run_tp_pp} vs {ckpt_tp_pp} from checkpoint)'
+    # Determine if RNG state will be loaded
+    if (ckpt_tp_pp == run_tp_pp and not finetune and not no_load_rng
+            and not getattr(state_dict['args'], 'no_save_rng', False)):
+        gen_sd_rng_state = _get_rng_state()  # we can load the rng state
+    else:
+        gen_sd_rng_state = None
+        if ckpt_tp_pp != run_tp_pp:
+            logger.info(f'{mismatch_msg}: RNG state will be ignored')
+    sharded_sd_metadata = state_dict.get('content_metadata')
+    if (not finetune and not no_load_optim and not getattr(state_dict['args'], 'no_save_optim', False)):
+        gen_sd_optim = optimizer
+        gen_sd_opt_param_scheduler = opt_param_scheduler
+
+        if (args.use_distributed_optimizer and ckpt_tp_pp != run_tp_pp
+                and (sharded_sd_metadata or {}).get('distrib_optim_sharding_type') not in {
+                    'fully_reshardable',
+                    'fully_sharded_model_space',
+                    'fsdp_dtensor',
+                }):
+            raise RuntimeError(f'{mismatch_msg}: not supported for DistributedOptimizer')
+    else:
+        gen_sd_optim, gen_sd_opt_param_scheduler = None, None
+    optim_sd_kwargs = dict(metadata=sharded_sd_metadata, is_loading=True)
+    model_sd_kwargs = dict(metadata=sharded_sd_metadata)
+    # TODO: check no_save_optim
+    sharded_state_dict = _generate_state_dict(
+        args,
+        models,
+        gen_sd_optim,
+        gen_sd_opt_param_scheduler,
+        gen_sd_rng_state,
+        iteration=iteration,
+        model_sd_kwargs=model_sd_kwargs,
+        optim_sd_kwargs=optim_sd_kwargs)
+    _filter_adapter_state_dict(sharded_state_dict, peft_format, adapter_name=adapter_name)
+    model_keys = [k for k in sharded_state_dict.keys() if k.startswith('model')]  # compat vpp
+    for k in model_keys:
+        patch_merge_fn(sharded_state_dict[k])
+    if mcore_017:
+        load_strategy = TorchDistLoadShardedStrategy()
+    else:
+        load_strategy = get_default_load_sharded_strategy(checkpoint_dir)
+
+    load_strategy = FullyParallelLoadStrategyWrapper(load_strategy,
+                                                     mpu.get_data_parallel_group(with_context_parallel=True))
+    state_dict = dist_checkpointing.load(sharded_state_dict, checkpoint_dir, load_strategy)
+
+    if finetune:
+        iteration = 0
+    if 'args' in state_dict and not finetune:
+        args.consumed_train_samples = getattr(state_dict['args'], 'consumed_train_samples', 0)
+
+    if len(ddp_models) == 1:
+        ddp_models[0].load_state_dict(state_dict['model'], strict=False)
+    else:
+        for i, m in enumerate(ddp_models):
+            if f'model{i}' not in state_dict:
+                continue
+            m.load_state_dict(state_dict[f'model{i}'])
+
+    if not finetune and not no_load_optim:
+        if optimizer is not None:
+            _load_optimizer_state_dict(optimizer, state_dict['optimizer'])
+        if opt_param_scheduler is not None:
+            opt_param_scheduler.load_state_dict(state_dict['opt_param_scheduler'])
+    elif (args.fp16 or args.bf16) and optimizer is not None:
+        optimizer.reload_model_params()
+
+    if not finetune and not no_load_rng:
+        if 'rng_state' in state_dict:
+            rng_state = state_dict['rng_state']
+            if args.data_parallel_random_init:
+                rng_state = rng_state[mpu.get_data_parallel_rank()]
+            else:
+                rng_state = rng_state[0]
+            random.setstate(rng_state['random_rng_state'])
+            np.random.set_state(rng_state['np_rng_state'])
+            torch.set_rng_state(rng_state['torch_rng_state'])
+            torch.cuda.set_rng_state(rng_state['cuda_rng_state'])
+            tensor_parallel.get_cuda_rng_tracker().set_states(rng_state['rng_tracker_states'])
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+
+    logger.info(f'Successfully loaded Megatron model weights from: {load_dir}')
+    return iteration
+
+
+def wrap_model(args, models, wrap_with_ddp: bool = True):
+    # Set tensor model parallel attributes if not set.
+    # Only parameters that are already tensor model parallel have these
+    # attributes set for them. We should make sure the default attributes
+    # are set for all params so the optimizer can use them.
+    for m in models:
+        for param in m.parameters():
+            tensor_parallel.set_defaults_if_not_set_tensor_model_parallel_attributes(param)
+        m.cuda(torch.cuda.current_device())
+    # Fp16
+    config = models[0].config
+    if args.fp16 or args.bf16:
+        models = [Float16Module(config, model_module) for model_module in models]
+
+    # DDP
+    if not wrap_with_ddp:
+        return models
+
+    kwargs = {}
+    for f in dataclasses.fields(DistributedDataParallelConfig):
+        if hasattr(args, f.name):
+            kwargs[f.name] = getattr(args, f.name)
+
+    # compat: SWIFT keeps the user-facing Megatron-LM arg name, while MCore
+    # DistributedDataParallelConfig expects grad_reduce_in_fp32.
+    if hasattr(args, 'accumulate_allreduce_grads_in_fp32'):
+        kwargs['grad_reduce_in_fp32'] = args.accumulate_allreduce_grads_in_fp32
+    kwargs['check_for_nan_in_grad'] = True
+    ddp_config = DistributedDataParallelConfig(**kwargs)
+
+    # In the Megatron FSDP and DDP use path, we need to initialize the bucket size.
+    # If bucket_size is not provided as an input, use sane default.
+    # If using very large dp_sizes, make buckets larger to ensure that chunks used in NCCL
+    # ring-reduce implementations are large enough to remain bandwidth-bound rather than
+    # latency-bound.
+    if ddp_config.bucket_size is None:
+        ddp_config.bucket_size = max(40000000, 1000000 * mpu.get_data_parallel_world_size(with_context_parallel=True))
+    # Set bucket_size to infinity if overlap_grad_reduce is False.
+    if not ddp_config.overlap_grad_reduce:
+        ddp_config.bucket_size = None
+
+    # Setup stream for DDP initialization with proper synchronization.
+    ddp_stream = torch.cuda.Stream()
+    ddp_stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(ddp_stream):
+        models = [
+            DDP(
+                config=config,
+                ddp_config=ddp_config,
+                module=model_chunk,
+                # Turn off bucketing for model_chunk 2 onwards, since communication for these
+                # model chunks is overlapped with compute anyway.
+                disable_bucketing=(model_chunk_idx > 0) or args.overlap_param_gather_with_optimizer_step,
+            ) for (model_chunk_idx, model_chunk) in enumerate(models)
+        ]
+    # Ensure DDP initialization completes before proceeding on the default stream.
+    torch.cuda.current_stream().wait_stream(ddp_stream)
+
+    # Broadcast params from data parallel src rank to other data parallel ranks.
+    if args.data_parallel_random_init:
+        for m in models:
+            m.broadcast_params()
+
+    return models
+
+
+def get_optimizer_param_scheduler(args, optimizer):
+    # Iteration-based training.
+    if args.lr_decay_iters is None:
+        args.lr_decay_iters = args.train_iters
+    lr_decay_steps = args.lr_decay_iters * args.global_batch_size
+    wd_incr_steps = args.train_iters * args.global_batch_size
+    wsd_decay_steps = None
+    if args.lr_wsd_decay_iters is not None:
+        wsd_decay_steps = args.lr_wsd_decay_iters * args.global_batch_size
+    if args.lr_warmup_fraction is not None:
+        lr_warmup_steps = args.lr_warmup_fraction * lr_decay_steps
+    else:
+        lr_warmup_steps = args.lr_warmup_iters * args.global_batch_size
+
+    opt_param_scheduler = OptimizerParamScheduler(
+        optimizer,
+        init_lr=args.lr_warmup_init,
+        max_lr=args.lr,
+        min_lr=args.min_lr,
+        lr_warmup_steps=lr_warmup_steps,
+        lr_decay_steps=lr_decay_steps,
+        lr_decay_style=args.lr_decay_style,
+        start_wd=args.start_weight_decay,
+        end_wd=args.end_weight_decay,
+        wd_incr_steps=wd_incr_steps,
+        wd_incr_style=args.weight_decay_incr_style,
+        wsd_decay_steps=wsd_decay_steps,
+        lr_wsd_decay_style=args.lr_wsd_decay_style,
+    )
+
+    return opt_param_scheduler
+
+
+def should_disable_forward_pre_hook(args):
+    """Block forward pre-hook for certain configurations."""
+    return args.use_distributed_optimizer and args.overlap_param_gather
+
+
+def enable_forward_pre_hook(model_chunks):
+    for model_chunk in model_chunks:
+        assert isinstance(model_chunk, DDP)
+        model_chunk.enable_forward_pre_hook()
+
+
+def disable_forward_pre_hook(model_chunks, param_sync=True):
+    for model_chunk in model_chunks:
+        assert isinstance(model_chunk, DDP)
+        model_chunk.disable_forward_pre_hook(param_sync=param_sync)
+
+
+def initialize_tp_communicators(args, config):
+    """initializing the communicators with user buffers for high-performance tensor-model-parallel
+    communication overlap"""
+    from transformer_engine.pytorch import module as te_module
+    input_shape = [
+        (args.seq_length * args.micro_batch_size) // args.context_parallel_size,
+        config.hidden_size,
+    ]
+
+    if is_te_min_version('2.7.0'):
+        UserBufferQuantizationMode = te_module.base.UserBufferQuantizationMode
+        quantization_modes = [UserBufferQuantizationMode.FP8 if args.fp8 else UserBufferQuantizationMode.NONE]
+        if args.fp8 is not None and args.first_last_layers_bf16 and (args.num_layers_at_start_in_bf16 > 0
+                                                                     or args.num_layers_at_end_in_bf16 > 0):
+            quantization_modes.append(UserBufferQuantizationMode.NONE)
+        # The process group with the target bootstrap backend is created in Transformer Engine.
+        te_module.base.initialize_ub(
+            shape=input_shape,
+            tp_size=args.tensor_model_parallel_size,
+            quantization_modes=quantization_modes,
+        )
+    elif is_te_min_version('1.9.0'):
+        # The process group with the target bootstrap backend is created in Transformer Engine.
+        te_module.base.initialize_ub(
+            shape=input_shape,
+            tp_size=args.tensor_model_parallel_size,
+            use_fp8=(args.fp8 is not None),
+        )
+
+
+def warmup_jit_function(config, args):
+    if args.bf16:
+        dtype = torch.bfloat16
+    elif args.fp16:
+        dtype = torch.float16
+    else:
+        dtype = torch.float32
+
+    bias = torch.rand(config.ffn_hidden_size // config.tensor_model_parallel_size, dtype=dtype, device='cuda')
+    input_tensor = torch.rand(
+        (
+            args.seq_length // config.context_parallel_size,
+            args.micro_batch_size,
+            config.ffn_hidden_size // config.tensor_model_parallel_size,
+        ),
+        dtype=dtype,
+        device='cuda',
+    )
+    # Warmup JIT fusions with the input_tensor grad_enable state of both forward
+    # prop and recomputation
+    for bias_grad, input_grad in zip([True, True], [False, True]):
+        bias.requires_grad, input_tensor.requires_grad = bias_grad, input_grad
+        for _ in range(5):
+            if config.swiglu:
+                output = bias_swiglu(input_tensor, bias)
+            else:
+                output = bias_gelu(bias, input_tensor)
+    del bias, input_tensor, output
+
+    # Warmup fused bias+dropout+add
+    if config.sequence_parallel:
+        seq_length = args.seq_length // mpu.get_tensor_model_parallel_world_size()
+    else:
+        seq_length = args.seq_length
+    input_tensor = torch.rand(
+        (seq_length // config.context_parallel_size, args.micro_batch_size, config.hidden_size),
+        dtype=dtype,
+        device='cuda',
+    )
+    residual = torch.rand(
+        (seq_length // config.context_parallel_size, args.micro_batch_size, config.hidden_size),
+        dtype=dtype,
+        device='cuda',
+    )
+    bias = torch.rand((config.hidden_size), dtype=dtype, device='cuda').expand_as(residual)
+    dropout_rate = 0.1
+    # Warmup JIT fusions with the input_tensor grad_enable state of both forward
+    # prop and recomputation
+    for input_grad, bias_grad, residual_grad in zip([False, True], [True, True], [True, True]):
+        input_tensor.requires_grad = input_grad
+        bias.requires_grad = bias_grad
+        residual.requires_grad = residual_grad
+        for _ in range(5):
+            output = bias_dropout_add_fused_train([input_tensor, bias], residual, dropout_rate)
+    del bias, input_tensor, residual, output
+    torch.cuda.empty_cache()
+
+
+def get_batch_on_this_cp_rank(args, batch: Dict[str, Any]):
+    """Slice batch input along sequence dimension into multiple chunks,
+    which are parallelized across GPUs in a context parallel group.
+    """
+
+    # With causal masking, each token only attends to its prior tokens. Simply split
+    # sequence into CP chunks can result in severe load imbalance. That's to say, chunks
+    # at the end of sequence have bigger workload than others. To address this issue,
+    # we split sequence into 2*CP ranks. Assuming CP=2, we then get 4 chunks, chunk_0
+    # and chunk_3 are assigned to GPU0, chunk_1 and chunk_2 are assigned to GPU1, so
+    # that we can get balanced workload among GPUs in a context parallel group.
+    cp_size = mpu.get_context_parallel_world_size()
+    if cp_size > 1:
+        keys = ['labels', 'position_ids', 'loss_scale']
+        if not args.is_multimodal:
+            # Multimodal models will handle CP in input_embeds.
+            keys.append('input_ids')
+
+        packed_seq_params = batch.get('packed_seq_params')
+        cp_partition_mode = getattr(args, 'cp_partition_mode', 'zigzag')
+        kwargs = {}
+        if cp_partition_mode == 'contiguous':
+            kwargs['cp_partition_mode'] = 'contiguous'
+        for key, val in batch.items():
+            if key not in keys:
+                continue
+            if args.task_type in ('seq_cls', 'embedding', 'generative_reranker') and key == 'labels':
+                continue
+            if val is not None:
+                batch[key] = split_cp_inputs(val, getattr(packed_seq_params, 'cu_seqlens_q', None), -1, **kwargs)
+        attention_mask = batch.get('attention_mask')
+        if is_torch_npu_available() and attention_mask is not None and attention_mask.ndim >= 4:
+            batch['attention_mask'] = split_cp_inputs(attention_mask, getattr(packed_seq_params, 'cu_seqlens_q', None),
+                                                      -2, **kwargs)
+
+    return batch

swift/megatron/utils/parallel_utils.py

@@ -0,0 +1,26 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+from megatron.core import mpu
+
+
+def reduce_max_stat_across_model_parallel_group(stat: float) -> float:
+    """
+    Ranks without an optimizer will have no grad_norm or num_zeros_in_grad stats.
+    We need to ensure the logging and writer rank has those values.
+    This function reduces a stat tensor across the model parallel group.
+
+    We use an all_reduce max since the values have already been summed across optimizer ranks where possible
+    """
+    stat = torch.tensor([stat], dtype=torch.float32, device=torch.cuda.current_device())
+    torch.distributed.all_reduce(stat, op=torch.distributed.ReduceOp.MAX, group=mpu.get_model_parallel_group())
+    return stat.item()
+
+
+def logical_and_across_model_parallel_group(input: bool) -> bool:
+    """
+    This function gathers a bool value across the model parallel group
+    """
+    input = int(bool(input))
+    input = torch.tensor([input], dtype=torch.int, device=torch.cuda.current_device())
+    torch.distributed.all_reduce(input, op=torch.distributed.ReduceOp.MIN, group=mpu.get_model_parallel_group())
+    return bool(input.item())

swift/megatron/utils/patcher.py

@@ -0,0 +1,41 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+from megatron.core.dist_checkpointing.mapping import ShardedTensorFactory
+from megatron.core.dist_checkpointing.strategies.torch import TorchDistSaveShardedStrategy
+
+from swift.utils import get_logger
+
+logger = get_logger()
+
+
+def patch_torch_dist_shard(thread_count):
+    __init__ = TorchDistSaveShardedStrategy.__init__
+
+    def __new_init__(*args, **kwargs):
+        kwargs['thread_count'] = thread_count
+        return __init__(*args, **kwargs)
+
+    TorchDistSaveShardedStrategy.__init__ = __new_init__
+
+
+def patch_merge_fn(state_dict_model):
+    # https://github.com/NVIDIA/Megatron-LM/issues/1380
+
+    def sh_ten_merge_fn(sub_state_dict):
+        with torch.no_grad():
+            shared_storage = sub_state_dict[0].untyped_storage()
+            if all(shared_storage.data_ptr() == tensor.untyped_storage().data_ptr() for tensor in sub_state_dict):
+                element_size = sub_state_dict[0].element_size()
+                total_numel = sum(tensor.numel() for tensor in sub_state_dict)
+                if shared_storage.nbytes() == total_numel * element_size:
+                    dim_0 = sum(tensor.shape[0] for tensor in sub_state_dict)
+                    shape = (dim_0, ) + sub_state_dict[0].shape[1:]
+                    combined_tensor = torch.empty(
+                        shape, dtype=sub_state_dict[0].dtype,
+                        device=sub_state_dict[0].device).set_(shared_storage, 0, shape)
+                    return combined_tensor
+            return torch.cat(sub_state_dict)
+
+    for v in state_dict_model.values():
+        if isinstance(v, ShardedTensorFactory) and 'apply_swiglu_sharded_factory' in v.merge_fn.__qualname__:
+            v.merge_fn = sh_ten_merge_fn

swift/megatron/utils/router_replay_utils.py

@@ -0,0 +1,215 @@
+"""
+Router Replay Utilities
+Utilities for handling router replay functionality in Megatron models.
+"""
+
+import torch
+from mcore_bridge import split_cp_inputs
+from megatron.core import mpu
+from megatron.core.tensor_parallel import scatter_to_sequence_parallel_region
+from megatron.core.transformer.transformer_block import get_num_layers_to_build
+from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
+
+from swift.utils import get_logger
+from swift.utils.torch_utils import get_current_device
+
+logger = get_logger()
+
+try:
+    from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
+    from megatron.core.transformer.moe.token_dispatcher import MoEAlltoAllTokenDispatcher
+    ROUTER_REPLAY_AVAILABLE = True
+except ImportError:
+    logger.warning('RouterReplay not available in current megatron-core version')
+    RouterReplay = None
+    RouterReplayAction = None
+    MoEAlltoAllTokenDispatcher = None
+    ROUTER_REPLAY_AVAILABLE = False
+
+device_name = get_current_device()
+
+
+def is_moe_layer(tf_config, layer_idx):
+    moe_layer_freq = getattr(tf_config, 'moe_layer_freq', None)
+    if isinstance(moe_layer_freq, int):
+        return layer_idx % moe_layer_freq == 0
+    elif isinstance(moe_layer_freq, list):
+        return moe_layer_freq[layer_idx] == 1
+    else:
+        raise ValueError(f'Unsupported moe_layer_freq type: {type(moe_layer_freq)}')
+
+
+def get_moe_num_layers_to_build(tf_config, vp_stage=None, pp_rank=None):
+    """Count the number of MoE layers assigned to the current rank.
+    When ``moe_layer_freq`` is 1 or unset, every transformer layer is an MoE
+    layer, so the count equals the total layer count. Otherwise only layers
+    whose global index satisfies the frequency predicate are counted.
+    Args:
+        config: Megatron TransformerConfig providing layer layout information.
+        vp_stage: Virtual-pipeline stage index (None defaults to current).
+        pp_rank: Pipeline-parallel rank (None defaults to current).
+    Returns:
+        Number of MoE layers on the specified rank/stage.
+    """
+    total_layers = get_num_layers_to_build(tf_config, vp_stage=vp_stage, pp_rank=pp_rank)
+
+    layer_offset = get_transformer_layer_offset(tf_config, vp_stage=vp_stage)
+    local_global_indices = range(layer_offset, layer_offset + total_layers)
+
+    num_moe_layers = sum(1 for idx in local_global_indices if is_moe_layer(tf_config, idx))
+
+    return num_moe_layers
+
+
+def get_local_layer_range(tf_config, vp_rank=None, only_moe_layer=True):
+    vp_size = tf_config.virtual_pipeline_model_parallel_size
+    if vp_size is not None:
+        vp_rank = 0 if vp_rank is None else vp_rank
+        offset = 0
+        for pre_vp_stage in range(vp_size):
+            if pre_vp_stage == vp_rank:
+                break
+            num_layers_to_build = get_moe_num_layers_to_build(
+                tf_config, pre_vp_stage) if only_moe_layer else get_num_layers_to_build(tf_config, pre_vp_stage)
+            offset += num_layers_to_build
+    else:
+        offset = 0
+    count = get_moe_num_layers_to_build(tf_config, vp_rank) if only_moe_layer else get_num_layers_to_build(
+        tf_config, vp_rank)
+    return offset, count
+
+
+def get_local_topk_idx_for_current_rank(global_topk_idx, tf_config, packed_seq_params=None):
+    if global_topk_idx is None:
+        return None
+    # 1. pp slice
+    # For the hybrid model, global_topk_idx contains data from all layers
+    # because vLLM reports routed_experts across all transformer layers(including dense).
+    # However megatron only has routers for MoE layers.
+    # So local_topk_idx should filter only data from the MoE layer.
+    layer_offset = get_transformer_layer_offset(tf_config, vp_stage=0)
+    offset, count = get_local_layer_range(
+        tf_config, tf_config.virtual_pipeline_model_parallel_size, only_moe_layer=False)
+    num_layers = offset + count
+    moe_layer_idx = torch.tensor([
+        layer_idx for layer_idx in range(layer_offset, layer_offset + num_layers) if is_moe_layer(tf_config, layer_idx)
+    ],
+                                 dtype=torch.long,
+                                 device=global_topk_idx.device)
+    local_topk_idx = torch.index_select(global_topk_idx, dim=2, index=moe_layer_idx)
+    # 2. cp slice
+    cp_size = mpu.get_context_parallel_world_size()
+    if cp_size > 1:
+        local_topk_idx = split_cp_inputs(local_topk_idx, getattr(packed_seq_params, 'cu_seqlens_q', None), 1)
+    # 3. sp slice
+    local_topk_idx = scatter_to_sequence_parallel_region(local_topk_idx.transpose(0, 1)).transpose(0, 1)
+    return local_topk_idx
+
+
+def get_router_replay_data(tf_config, vp_rank=None):
+    router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank)
+    layers_topk_idx = []
+    for router in router_instances_list:
+        layers_topk_idx.append(router.recorded_topk_idx.to(torch.uint8))
+    # layer_num, seq_len, topk -> 1, seq_len, layer_num, topk
+    layers_topk_idx = torch.stack(layers_topk_idx).transpose(0, 1).unsqueeze(0).to(device_name)
+    return layers_topk_idx
+
+
+def set_router_replay_data(layers_topk_idx, tf_config, vp_rank=None):
+    # bs, seq_len, layer_num, topk -> layer_num, total_seq_len, topk
+    layers_topk_idx_reshape = layers_topk_idx.flatten(0, 1).transpose(0, 1).to(device_name)
+    offset, count = get_local_layer_range(tf_config, vp_rank)
+    router_instances_list = RouterReplay.global_router_replay_instances[offset:offset + count]
+    for i, router in enumerate(router_instances_list):
+        router.set_target_indices(layers_topk_idx_reshape[i + offset].to(torch.int64))
+
+
+class RouterReplayHelper:
+    """Helper class to query router replay state and locate local RouterReplay instances."""
+
+    @staticmethod
+    def get_micro_batch_router_list(tf_config, vp_rank=None):
+        """
+        Return the list of RouterReplay instances corresponding to the current micro-batch and local
+        (pp_rank, vp_stage) layer range.
+
+        When virtual pipeline (VPP) is enabled, the local range for the PP rank is expanded to include
+        all VP stages by multiplying the per-VP count by vp_size. The returned slice is taken from the
+        global RouterReplay.global_router_replay_instances list.
+
+        Args:
+            tf_config: Configuration object used to compute layer assignments.
+            vp_rank (Optional[int]): Explicit virtual pipeline stage to query. If None, the current VP
+                rank from Megatron parallel state is used when available.
+        Returns:
+            list: A contiguous sublist of RouterReplay.router_instances for the local layer range.
+        """
+        offset, count = get_local_layer_range(tf_config, vp_rank)
+        router_instances_list = RouterReplay.global_router_replay_instances[offset:offset + count]
+        return router_instances_list
+
+    @staticmethod
+    def is_r2_record_action(tf_config, vp_rank=None) -> bool:
+        """Return True if the current router_replay_action is RECORD (R2) for the local router instances.
+
+        This inspects the first local RouterReplay instance's router_replay_action and compares it to
+        RouterReplayAction.RECORD.
+        """
+        router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank)
+        return (router_instances_list and router_instances_list[0].router_replay_action == RouterReplayAction.RECORD)
+
+    @staticmethod
+    def is_replay_forward_action(tf_config, vp_rank=None) -> bool:
+        """Return True if the current router_replay_action is REPLAY_FORWARD for the local router instances.
+
+        This inspects the first local RouterReplay instance's router_replay_action and compares it to
+        RouterReplayAction.REPLAY_FORWARD.
+        """
+        router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank)
+        return (router_instances_list
+                and router_instances_list[0].router_replay_action == RouterReplayAction.REPLAY_FORWARD)
+
+    @staticmethod
+    def is_replay_backward_action(tf_config, vp_rank=None) -> bool:
+        """Return True if the current router_replay_action is REPLAY_BACKWARD for the local router instances.
+
+        This inspects the first local RouterReplay instance's router_replay_action and compares it to
+        RouterReplayAction.REPLAY_BACKWARD.
+        """
+        router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank)
+        return (router_instances_list
+                and router_instances_list[0].router_replay_action == RouterReplayAction.REPLAY_BACKWARD)
+
+
+def apply_router_replay_patch():
+    """
+    Applies the monkey patch for MoE Router Replay functionality.
+    """
+    logger.info('Applying Router Replay Patch...')
+
+    assert ROUTER_REPLAY_AVAILABLE, \
+        'The routing replay is not supported. Please upgrade megatron-core to 0.16.0 or higher'
+
+    # Patch MoEAlltoAllTokenDispatcher.preprocess to handle router replay
+    # When router replay is enabled, duplicate indices in top_indices can cause
+    # routing_map.sum() < num_tokens * topk, leading to split size mismatch in alltoall.
+    if MoEAlltoAllTokenDispatcher is not None and not hasattr(MoEAlltoAllTokenDispatcher, '_preprocess_patched'):
+        original_preprocess = MoEAlltoAllTokenDispatcher.preprocess
+
+        def patched_preprocess(self, routing_map):
+            """Patched preprocess that handles router replay correctly for alltoall dispatcher."""
+            # Call original preprocess
+            result = original_preprocess(self, routing_map)
+            # Fix num_out_tokens when router replay is enabled
+            if (getattr(self.config, 'moe_enable_routing_replay', False) and not self.drop_and_pad
+                    and self.config.moe_expert_capacity_factor is None
+                    and not (getattr(self.config, 'moe_router_padding_for_quantization', None)
+                             or getattr(self.config, 'moe_router_padding_for_fp8', None))):
+                # With router replay, duplicate indices can reduce the actual routed
+                # token count, so derive it from the routing map instead.
+                self.num_out_tokens = int(routing_map.sum().item())
+            return result
+
+        MoEAlltoAllTokenDispatcher.preprocess = patched_preprocess
+        MoEAlltoAllTokenDispatcher._preprocess_patched = True

swift/megatron/utils/utils.py

@@ -0,0 +1,286 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import re
+import torch
+import torch.distributed as dist
+from megatron.core import mpu
+from megatron.core.extensions.transformer_engine import TEGroupedLinear, TELayerNormColumnParallelLinear, TELinear
+from megatron.core.inference.communication_utils import recv_from_prev_pipeline_rank_, send_to_next_pipeline_rank
+from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.ssm.mamba_context_parallel import _undo_attention_load_balancing
+from megatron.core.transformer.moe.router import TopKRouter
+from torch import nn
+from transformers.utils import is_torch_npu_available
+
+from swift.tuners import LoraConfig, Swift
+from swift.utils import (activate_parameters, deep_getattr, find_layers, freeze_parameters, get_logger,
+                         get_model_parameter_info)
+from swift.utils import get_packed_seq_params as _get_packed_seq_params
+
+logger = get_logger()
+
+
+def find_all_linears(model, extra_layers=None):
+
+    def _cond(name, module):
+        if (extra_layers and isinstance(module, tuple(extra_layers))) or name != 'output_layer' and isinstance(
+                module, (TELinear, TELayerNormColumnParallelLinear, TEGroupedLinear, nn.Linear)):
+            return True
+        return False
+
+    return find_layers(model, _cond)
+
+
+def find_router(model):
+    return find_layers(model, lambda name, module: isinstance(module, TopKRouter))
+
+
+def find_embedding(model):
+    return find_layers(model, lambda name, module: isinstance(module, LanguageModelEmbedding))
+
+
+def get_multimodal_target_regex(
+    args,
+    model,
+    *,
+    freeze_llm: bool = False,
+    freeze_vit: bool = True,
+    freeze_aligner: bool = True,
+    include_embedding: bool = False,
+    include_router: bool = False,
+) -> str:
+    megatron_model_meta = args.megatron_model_meta
+    modules = []
+    visual_cls = megatron_model_meta.visual_cls
+    vision_tower = [f'visual.{vit}' for vit in visual_cls._vision_tower]
+    aligner = [f'visual.{aligner}' for aligner in visual_cls._aligner]
+    if not freeze_llm:
+        modules.append('language_model')
+    if not freeze_vit:
+        modules += vision_tower
+    if not freeze_aligner:
+        modules += aligner
+    assert len(modules) > 0, f'modules: {modules}'
+    extra_layers = []
+    if include_embedding:
+        extra_layers.append(LanguageModelEmbedding)
+    if include_router:
+        extra_layers.append(TopKRouter)
+
+    res = []
+    for module in modules:
+        rejected_modules = []
+        if not freeze_vit:
+            for _aligner in aligner:
+                if _aligner.startswith(f'{module}.'):
+                    rejected_modules.append(_aligner)
+
+        sub_module = deep_getattr(model, module)
+        if sub_module is None:
+            continue
+        target_modules = find_all_linears(sub_module, extra_layers)
+        if not target_modules:
+            continue
+        target_modules = [tm for tm in target_modules if tm]
+        target_pattern = rf'.*\.({"|".join(target_modules)})' if target_modules else ''
+        rejected_pattern = rf'(?!({"|".join(rejected_modules)}))' if rejected_modules else ''
+        res.append(rf'{rejected_pattern}{re.escape(module)}(?=\.){target_pattern}')
+
+    return rf'^({"|".join(res)})$'
+
+
+def get_target_modules(args, model):
+    if isinstance(args.target_modules, str):
+        return args.target_modules
+    target_modules = args.target_modules.copy()
+    if 'all-linear' in target_modules:
+        if args.is_multimodal and not args.language_model_only:
+            if args.tuner_type == 'lora_llm':
+                kwargs = {
+                    'freeze_llm': False,
+                    'freeze_vit': True,
+                    'freeze_aligner': True,
+                }
+            else:  # lora
+                kwargs = {
+                    'freeze_llm': args.freeze_llm,
+                    'freeze_vit': args.freeze_vit,
+                    'freeze_aligner': args.freeze_aligner,
+                }
+            return get_multimodal_target_regex(
+                args,
+                model,
+                include_embedding='all-embedding' in target_modules,
+                include_router='all-router' in target_modules,
+                **kwargs,
+            )
+        else:
+            target_modules.remove('all-linear')
+            target_modules += find_all_linears(model)
+    if 'all-embedding' in target_modules:
+        target_modules.remove('all-embedding')
+        target_modules += find_embedding(model)
+    if 'all-router' in target_modules:
+        target_modules.remove('all-router')
+        target_modules += find_router(model)
+    return target_modules
+
+
+def get_modules_to_save(args, model):
+    if args.task_type == 'seq_cls':
+        args.modules_to_save.append('output_layer')
+    modules_to_save = args.modules_to_save.copy()
+    if 'all-embedding' in args.modules_to_save:
+        modules_to_save.remove('all-embedding')
+        modules_to_save += find_embedding(model)
+    return modules_to_save
+
+
+def prepare_adapter(args, model):
+    target_modules = get_target_modules(args, model)
+    modules_to_save = get_modules_to_save(args, model)
+    lora_kwargs = {
+        'r': args.lora_rank,
+        'target_modules': target_modules,
+        'lora_alpha': args.lora_alpha,
+        'lora_dropout': args.lora_dropout,
+        'bias': args.lora_bias,
+        'modules_to_save': modules_to_save,
+        'use_rslora': args.use_rslora,
+    }
+    lora_config = LoraConfig(task_type='CAUSAL_LM', lora_dtype=args.lora_dtype, **lora_kwargs)
+    logger.info(f'lora_config: {lora_config}')
+    model = Swift.prepare_model(model, lora_config)
+    if args.mcore_ref_adapter or args.ref_adapters:
+        model.add_adapter('ref_adapter', lora_config)
+        model.base_model._cast_adapter_dtype(adapter_name='ref_adapter', autocast_adapter_dtype=True)
+        for n, p in model.named_parameters():
+            if '.ref_adapter.' in n:
+                p.requires_grad = False
+    return model
+
+
+def _prepare_full_vit(args, model):
+    megatron_model_meta = args.megatron_model_meta
+    visual_cls = megatron_model_meta.visual_cls
+    vision_tower = [f'visual.{vit}' for vit in visual_cls._vision_tower]
+    aligner = [f'visual.{aligner}' for aligner in visual_cls._aligner]
+    for module_prefix in vision_tower + aligner:
+        module = deep_getattr(model, module_prefix)
+        if module is not None:
+            module.requires_grad_(True)
+
+
+def prepare_mcore_model(args, model):
+    if args.tuner_type == 'full':
+        freeze_parameters(model, args.freeze_parameters_ratio, args.freeze_parameters, args.freeze_parameters_regex)
+        if args.trainable_parameters or args.trainable_parameters_regex:
+            activate_parameters(model, args.trainable_parameters, args.trainable_parameters_regex)
+    elif args.tuner_type in {'lora', 'lora_llm'}:
+        model = prepare_adapter(args, model)
+        if args.tuner_type == 'lora_llm':
+            _prepare_full_vit(args, model)
+    logger.info(f'model: {model}')
+    logger.info_if(
+        f'[rank{dist.get_rank()}] model_parameter_info: {get_model_parameter_info(model)}',
+        cond=mpu.get_data_parallel_rank() == 0)
+    return model
+
+
+def forward_step_helper(model, inputs, dtype=None):
+    config = model.config
+    dtype = dtype or config.params_dtype
+    if not mpu.is_pipeline_first_stage():
+        recv_shape_buffer = torch.empty((3, ), device=torch.cuda.current_device(), dtype=torch.int64)
+        recv_from_prev_pipeline_rank_(recv_shape_buffer)
+        recv_buffer = torch.empty(recv_shape_buffer.tolist(), device=torch.cuda.current_device(), dtype=dtype)
+        recv_from_prev_pipeline_rank_(recv_buffer)
+        model.set_input_tensor(recv_buffer)
+    output_tensor = model(**inputs)
+    if not mpu.is_pipeline_last_stage():
+        recv_shape_buffer = torch.tensor(output_tensor.shape, device=torch.cuda.current_device(), dtype=torch.int64)
+        send_to_next_pipeline_rank(recv_shape_buffer)
+        send_to_next_pipeline_rank(output_tensor)
+        output_tensor = None
+
+    return output_tensor
+
+
+def get_padding_to(args):
+    padding_to = None
+    if args.tensor_model_parallel_size > 1 and args.sequence_parallel:
+        padding_to = args.tensor_model_parallel_size
+    if args.context_parallel_size > 1:
+        padding_to = (padding_to or 1) * args.context_parallel_size
+    origin_padding_to = padding_to
+    fp8_format = getattr(args, 'fp8_format', None) or getattr(args, 'fp8', None)
+    fp4_format = getattr(args, 'fp4_format', None) or getattr(args, 'fp4', None)
+    if args.fp8_recipe == 'blockwise':
+        padding_to = (padding_to or 1) * 128
+    elif args.fp8_recipe == 'mxfp8':
+        # MXFP8 uses a block size of 32. Under sequence parallel, the sequence is
+        # split across TP ranks, so each per-rank shard (seq_len / TP) must itself
+        # be divisible by 32. Pad the total length to TP * 32 to guarantee this.
+        padding_to = (padding_to or 1) * 32
+    elif fp8_format is not None or fp4_format is not None:
+        padding_to = (padding_to or 1) * 16
+    if args.attention_backend == 'fused':
+        padding_to = max(padding_to or 1, ((origin_padding_to) or 1) * 64)
+    return padding_to
+
+
+def get_packed_seq_params(args, position_ids: torch.Tensor) -> PackedSeqParams:
+    params = _get_packed_seq_params(position_ids)
+    packed = PackedSeqParams(
+        cu_seqlens_q=params['cu_seq_lens_q'],
+        cu_seqlens_kv=params['cu_seq_lens_k'],
+        max_seqlen_q=params['max_length_q'],
+        max_seqlen_kv=params['max_length_k'],
+        qkv_format='thd',
+    )
+    if hasattr(packed, 'cp_partition_mode'):
+        packed.cp_partition_mode = args.cp_partition_mode
+
+    if is_torch_npu_available():
+        packed.cu_seqlens_q_padded = params['cu_seq_lens_q']
+        packed.cu_seqlens_kv_padded = params['cu_seq_lens_k']
+
+    return packed
+
+
+def reconstruct_tensor_cp(tensor, packed_seq_params, dim=1) -> torch.Tensor:
+    """In CP mode, all-gather and undo the load-balanced (zigzag) chunking
+    produced by ``split_cp_inputs``, restoring the full sequence in original
+    token order along ``dim``.
+
+    Args:
+        tensor: CP-sharded local tensor whose sequence dim is at ``dim``.
+        packed_seq_params: ``PackedSeqParams`` for THD inputs, or ``None`` for
+            regular ``[B, S, ...]`` inputs.
+        dim: Sequence dimension index of ``tensor`` (default: 1).
+
+    Returns:
+        torch.Tensor: Full-sequence tensor with the same shape as ``tensor``
+        except the size at ``dim`` is multiplied by ``cp_size``.
+    """
+
+    cp_size = mpu.get_context_parallel_world_size()
+    if cp_size <= 1:
+        return tensor
+
+    cp_rank = mpu.get_context_parallel_rank()
+    cp_group = mpu.get_context_parallel_group()
+
+    # All-gather across CP ranks (preserve local autograd graph for `tensor`).
+    output_list = [torch.empty_like(tensor) for _ in range(cp_size)]
+    torch.distributed.all_gather(output_list, tensor.contiguous(), group=cp_group)
+    output_list[cp_rank] = tensor
+    gathered = torch.cat(output_list, dim=dim)
+
+    # `_undo_attention_load_balancing` assumes sequence dim is 0; transpose if needed.
+    if dim != 0:
+        gathered = gathered.transpose(0, dim).contiguous()
+    out = _undo_attention_load_balancing(gathered, cp_size, packed_seq_params)
+    if dim != 0:
+        out = out.transpose(0, dim).contiguous()
+    return out