[megatron] Support megatron CP/non-padding-free more tasks (#9516)

Jintao-Huang · web-flow · commit 5728eda28f43 · 2026-06-09T15:39:12.000+08:00
diff --git a/swift/megatron/trainers/base.py b/swift/megatron/trainers/base.py
@@ -32,16 +32,16 @@
                                   get_optimizer_param_scheduler, get_padding_to, init_persistent_async_worker,
                                   initialize_tp_communicators, load_mcore_checkpoint,
                                   logical_and_across_model_parallel_group, maybe_finalize_async_save,
-                                  prepare_mcore_model, reduce_max_stat_across_model_parallel_group,
-                                  save_mcore_checkpoint, should_disable_forward_pre_hook, warmup_jit_function,
-                                  wrap_model)
+                                  prepare_mcore_model, reconstruct_tensor_cp,
+                                  reduce_max_stat_across_model_parallel_group, save_mcore_checkpoint,
+                                  should_disable_forward_pre_hook, warmup_jit_function, wrap_model)
 from swift.template import Template
 from swift.trainers import dynamic_gradient_checkpointing
 from swift.trainers.utils import patch_modelscope_hub_timeout
 from swift.utils import (deep_getattr, gc_collect, get_current_device, get_last_valid_indices, get_logger, is_last_rank,
                          is_master, ms_logger_context)
 from .batch_sampler import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
-from .utils import TrainerState, build_streaming_dataloader
+from .utils import TrainerState, build_streaming_dataloader, prepare_batch
 
 try:
     from megatron.core.optimizer import param_group_identifier_keys
@@ -985,7 +985,6 @@ def _should_use_npu_generated_attention_mask(self, args) -> bool:
                 and getattr(args, 'attention_backend', None) != 'local' and getattr(args, 'use_flash_attn', False))
 
     def _prepare_batch(self, data, vp_stage=None, num_samples=None):
-        from .utils import prepare_batch
         return prepare_batch(self.args, data, vp_stage=vp_stage, num_samples=num_samples)
 
     def get_batch(self, data_iterator, vp_stage=None):
@@ -1014,11 +1013,19 @@ def _collect_config_info(self) -> Dict[str, str]:
         return {}
 
     def get_last_tokens(self, output_tensor, packed_seq_params=None, attention_mask=None, num_samples=None):
+        if self.args.context_parallel_size > 1:
+            output_tensor = reconstruct_tensor_cp(output_tensor, packed_seq_params, dim=1)
         if packed_seq_params is None:
-            last_token_idx = get_last_valid_indices((~attention_mask[:, 0, -1]).long())
+            # Compatible with attention_mask_2d
+            if attention_mask.dim() > 2:
+                attention_mask = (~attention_mask).sum(dim=(1, 2)) > 0
+            last_token_idx = get_last_valid_indices(attention_mask.long())
             last_tokens = output_tensor[torch.arange(output_tensor.shape[0]), last_token_idx]
         else:
             num_samples = num_samples or packed_seq_params.num_samples
-            last_token_idx = packed_seq_params.cu_seqlens_q[1:num_samples + 1] - 1
+            if self.args.context_parallel_size > 1:
+                last_token_idx = packed_seq_params.cu_seqlens_q[:num_samples] + packed_seq_params.seq_lens - 1
+            else:
+                last_token_idx = packed_seq_params.cu_seqlens_q[1:num_samples + 1] - 1
             last_tokens = output_tensor[0, last_token_idx]
         return last_tokens
diff --git a/swift/megatron/trainers/dpo_trainer.py b/swift/megatron/trainers/dpo_trainer.py
@@ -75,7 +75,6 @@ def loss_func(self, output_tensor: torch.Tensor, *, labels: torch.Tensor, packed
         return loss, metric
 
     def forward_step(self, data_iterator, model):
-        # Get the batch.
         unwrapped_model = model.module.module
         input_tensor = unwrapped_model.get_input_tensor()
         vp_stage = unwrapped_model.vp_stage
diff --git a/swift/megatron/trainers/embedding_trainer.py b/swift/megatron/trainers/embedding_trainer.py
@@ -14,17 +14,18 @@ class MegatronEmbeddingTrainer(BaseMegatronTrainer):
 
     def __init__(self, args, template):
         super().__init__(args, template)
-        if args.context_parallel_size > 1:
-            raise ValueError('Currently `task_type="embedding"` does not support context parallelism.')
-        if not args.padding_free:
-            raise ValueError('Currently, task_type embedding only supports padding_free.')
         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):
+    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)
+        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)
         loss = self._loss_func({'last_hidden_state': last_hidden_state}, labels)
@@ -33,11 +34,14 @@ def loss_func(self, output_tensor: torch.Tensor, *, labels: torch.Tensor, packed
         return loss, metric
 
     def forward_step(self, data_iterator, model):
-        # Get the batch.
         vp_stage = model.module.module.vp_stage
         data = self.get_batch(data_iterator, vp_stage)
         labels = data.pop('labels', None)
         output_tensor = model(**data)
-        packed_seq_params = data.get('packed_seq_params')
-        loss_func = partial(self.loss_func, labels=labels, packed_seq_params=packed_seq_params)
+        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
diff --git a/swift/megatron/trainers/kto_trainer.py b/swift/megatron/trainers/kto_trainer.py
@@ -110,7 +110,6 @@ def _get_input_tensor(input_tensor, is_KL: bool, is_ref: bool, length: int, dim:
         return res
 
     def forward_step(self, data_iterator, model):
-        # Get the batch.
         unwrapped_model = model.module.module
         input_tensor = unwrapped_model.get_input_tensor()
         vp_stage = unwrapped_model.vp_stage
diff --git a/swift/megatron/trainers/reranker_trainer.py b/swift/megatron/trainers/reranker_trainer.py
@@ -16,11 +16,6 @@ class MegatronRerankerTrainer(BaseMegatronTrainer):
 
     def __init__(self, args, template):
         super().__init__(args, template)
-        if args.context_parallel_size > 1:
-            raise ValueError('Currently `task_type="reranker/generative_reranker"` does not support '
-                             'context parallelism.')
-        if not args.padding_free:
-            raise ValueError('Currently, task_type reranker/generative_reranker only supports padding_free.')
         self._loss_func = loss_map[args.loss_type](args, self)
         self.eval_metrics = eval_metrics_map['reranker'](args, self)
 
@@ -36,9 +31,14 @@ def _get_listwise_reranker_preds(logits, labels):
         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):
+    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)
+        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)
@@ -60,11 +60,14 @@ def prepare_model(self):
             lm_model.tokenizer = self.template.tokenizer
 
     def forward_step(self, data_iterator, model):
-        # Get the batch.
         vp_stage = model.module.module.vp_stage
         data = self.get_batch(data_iterator, vp_stage)
         labels = data.pop('labels', None)
         output_tensor = model(**data)
-        packed_seq_params = data.get('packed_seq_params')
-        loss_func = partial(self.loss_func, labels=labels, packed_seq_params=packed_seq_params)
+        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
diff --git a/swift/megatron/trainers/reward_trainer.py b/swift/megatron/trainers/reward_trainer.py
@@ -11,10 +11,6 @@
 
 class MegatronRewardTrainer(MegatronRLHFTrainer):
 
-    def __init__(self, args, template):
-        super().__init__(args, template)
-        assert args.context_parallel_size == 1, 'Currently `rlhf_type="rm"` does not support context parallelism.'
-
     def loss_func(self, output_tensor, *, data):
         packed_seq_params = data.get('packed_seq_params')
         margin = data.pop('margin', None)
@@ -43,7 +39,6 @@ def loss_func(self, output_tensor, *, data):
         return loss, metric
 
     def forward_step(self, data_iterator, model):
-        # Get the batch.
         vp_stage = model.module.module.vp_stage
         data = self.get_batch(data_iterator, vp_stage)
         data.pop('loss_scale', None)
diff --git a/swift/megatron/trainers/trainer.py b/swift/megatron/trainers/trainer.py
@@ -19,8 +19,6 @@ class MegatronTrainer(BaseMegatronTrainer):
 
     def seq_cls_loss_func(self, output_tensor, *, labels: torch.Tensor, packed_seq_params=None, attention_mask=None):
         args = self.args
-        if args.context_parallel_size > 1:
-            raise ValueError('Currently `task_type="seq_cls"` does not support context parallelism.')
         logits = self.get_last_tokens(output_tensor, packed_seq_params, attention_mask)
         num_labels = args.num_labels
         acc = None
@@ -106,7 +104,6 @@ def _compute_channel_loss(self, losses, loss_mask, channels, packed_seq_params=N
         return new_metrics
 
     def forward_step(self, data_iterator, model):
-        # Get the batch.
         vp_stage = model.module.module.vp_stage
         data = self.get_batch(data_iterator, vp_stage)
         loss_scale = data.pop('loss_scale', None)
@@ -121,7 +118,8 @@ def forward_step(self, data_iterator, model):
                 self.seq_cls_loss_func,
                 labels=labels,
                 packed_seq_params=packed_seq_params,
-                attention_mask=data.get('attention_mask'))
+                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,
diff --git a/swift/megatron/trainers/utils.py b/swift/megatron/trainers/utils.py
@@ -11,6 +11,7 @@
 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()
@@ -312,7 +313,7 @@ class TrainerState:
     should_log: bool = False
 
     iteration: int = 0
-    consumed_train_samples = 0
+    consumed_train_samples: int = 0
     # compat transformers
     max_steps: Optional[int] = None
 
@@ -357,10 +358,10 @@ def prepare_batch(args, data, vp_stage=None, num_samples=None):
 
     Extracted from BaseMegatronTrainer._prepare_batch for reuse in ray workers.
     """
-    from swift.megatron.utils import get_batch_on_this_cp_rank, get_packed_seq_params
     batch = get_batch_on_this_pp_rank(args, data, vp_stage=vp_stage)
     if num_samples is None:
         num_samples = batch.pop('num_samples')
+    seq_lens = batch.pop('seq_lens', None)
     text_position_ids = batch.pop('text_position_ids', None)
     if text_position_ids is None:
         text_position_ids = batch.get('position_ids')
@@ -373,6 +374,8 @@ def prepare_batch(args, data, vp_stage=None, num_samples=None):
     if args.padding_free and text_position_ids is not None:
         batch['packed_seq_params'] = get_packed_seq_params(text_position_ids)
         batch['packed_seq_params'].num_samples = num_samples
+        if seq_lens is not None:
+            batch['packed_seq_params'].seq_lens = torch.tensor(seq_lens, device=text_position_ids.device)
     batch = get_batch_on_this_cp_rank(args, batch)
     return batch
 
diff --git a/swift/megatron/utils/__init__.py b/swift/megatron/utils/__init__.py
@@ -9,4 +9,5 @@
 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
+from .utils import (forward_step_helper, get_packed_seq_params, get_padding_to, prepare_mcore_model,
+                    reconstruct_tensor_cp)
diff --git a/swift/megatron/utils/convert_utils.py b/swift/megatron/utils/convert_utils.py
@@ -7,6 +7,7 @@
 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)
@@ -267,7 +268,6 @@ def test_convert_precision(args, hf_model, mg_model, template, test_convert_dtyp
             if mg_logits is not None:
                 mg_logits = gather_from_tensor_model_parallel_region(mg_logits)
         if args.context_parallel_size > 1:
-            from megatron.core.ssm.mamba_context_parallel import _undo_attention_load_balancing
             if mg_logits is not None:
                 mg_logits = gather_from_sequence_parallel_region(
                     mg_logits.transpose(0, 1), group=mpu.get_context_parallel_group())
diff --git a/swift/megatron/utils/megatron_lm_utils.py b/swift/megatron/utils/megatron_lm_utils.py
@@ -716,7 +716,7 @@ def get_batch_on_this_cp_rank(args, batch: Dict[str, Any]):
         for key, val in batch.items():
             if key not in keys:
                 continue
-            if args.task_type == 'seq_cls' and key == 'labels':
+            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)
diff --git a/swift/megatron/utils/utils.py b/swift/megatron/utils/utils.py
@@ -7,6 +7,7 @@
 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
@@ -237,3 +238,41 @@ def get_packed_seq_params(position_ids: torch.Tensor) -> PackedSeqParams:
         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
diff --git a/swift/template/base.py b/swift/template/base.py
@@ -1504,7 +1504,8 @@ def _handle_megatron_cp(self, batch: List[Dict[str, Any]]) -> None:
             input_ids = encoded['input_ids']
             padding_len = math.ceil(len(input_ids) / (cp_size * 2)) * (cp_size * 2) - len(input_ids)
             input_ids += [self.tokenizer.pad_token_id] * padding_len
-            encoded['labels'] += [-100] * padding_len
+            if encoded.get('labels') is not None:
+                encoded['labels'] += [-100] * padding_len
             if encoded.get('loss_scale') is not None:
                 encoded['loss_scale'] += [0] * padding_len
             if encoded.get('length') is not None:
@@ -1839,6 +1840,7 @@ def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[in
         assert self.tokenizer.pad_token_id is not None
         padding_side = self.padding_side if self.is_training else 'left'
         padding_right = padding_side == 'right'
+        real_seq_lens = [len(b['input_ids']) for b in batch] if self.use_megatron else None
         self._handle_megatron_cp(batch)
         if self.padding_free:
             batch[:] = [self.packing_row(batch)]
@@ -1900,7 +1902,6 @@ def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[in
                 res['position_ids'] = [torch.arange(seq_len, dtype=torch.int64) for seq_len in seq_lens]
 
         if self.use_megatron:
-            # For code simplicity, only the attention_backend 'flash' is supported here.
             if padding_to is not None:
                 padding_to = math.ceil(max(seq_lens) / padding_to) * padding_to
             if self.padding_free:
@@ -1918,7 +1919,7 @@ def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[in
                 res['attention_mask'] = torch.tril(torch.ones(
                     (len(seq_lens), seq_len, seq_len), dtype=torch.bool)).view(len(seq_lens), 1, seq_len, seq_len)
                 assert res['attention_mask'].dtype is torch.bool, f'attention_mask.dtype: {res["attention_mask"].dtype}'
-                for i, seq_len in enumerate(seq_lens):
+                for i, seq_len in enumerate(real_seq_lens):
                     res['attention_mask'][i, :, :, seq_len:] = 0
                 res['attention_mask'] = ~res['attention_mask']
 
@@ -1942,7 +1943,8 @@ def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[in
         res.update(self._data_collator_mm_data(batch))
         if not self.use_megatron and self.sequence_parallel_size > 1:
             res = self._sp_data_collator(res, padding_to, self.tokenizer, padding_side)
-
+        if self.use_megatron:
+            res['seq_lens'] = real_seq_lens  # CP locates the last token.
         return res
 
     def _pad_3d_position_ids(self,
