[v1] Fix device mesh, fix lora for reward model and fix sp (#10555)

src/llamafactory/v1/accelerator/interface.py

@@ -68,6 +68,11 @@ class DistributedStrategy:
         if not helper.is_distributed():
             self.mp_shard_size = 1
         elif self.mp_shard_size is None:
+            if helper.get_world_size() % self.mp_replicate_size != 0:
+                raise ValueError(
+                    f"world_size ({helper.get_world_size()}) must be divisible by "
+                    f"mp_replicate_size ({self.mp_replicate_size})."
+                )
             self.mp_shard_size = helper.get_world_size() // self.mp_replicate_size
         elif self.mp_replicate_size * self.mp_shard_size != helper.get_world_size():
             raise ValueError(
@@ -78,6 +83,10 @@ class DistributedStrategy:
         if not helper.is_distributed():
             self.dp_size = 1
         elif self.dp_size is None:
+            if helper.get_world_size() % self.cp_size != 0:
+                raise ValueError(
+                    f"world_size ({helper.get_world_size()}) must be divisible by cp_size ({self.cp_size})."
+                )
             self.dp_size = helper.get_world_size() // self.cp_size
         elif self.dp_size * self.cp_size != helper.get_world_size():
             raise ValueError(

src/llamafactory/v1/plugins/model_plugins/parallelization/ulysses.py

@@ -81,7 +81,7 @@ class UlyssesAttention(torch.nn.Module):
         query: Tensor,
         key: Tensor,
         value: Tensor,
-        attention_mask: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
         query_length: int,
         dropout_p=0.0,
         softmax_scale=None,
@@ -122,25 +122,42 @@ class UlyssesAttention(torch.nn.Module):
         if softmax_scale is None:
             softmax_scale = q.shape[-1] ** -0.5
 
+        sp_world_size = get_ulysses_sequence_parallel_world_size(self.spg)
+        local_position_ids = position_ids
+
         if position_ids is not None:
-            global_position_ids = [
-                torch.empty_like(position_ids) for _ in range(get_ulysses_sequence_parallel_world_size(self.spg))
-            ]
+            global_position_ids = [torch.empty_like(position_ids) for _ in range(sp_world_size)]
             dist.all_gather(global_position_ids, position_ids, group=self.spg)
             position_ids = torch.cat(global_position_ids, dim=-1).contiguous()
-            attention_mask = None
-        else:
+
+        # HF may turn an all-ones local attention_mask into None before this
+        # function. Under CP, different ranks can then disagree: some local
+        # shards still contain padding and keep a mask, while others see None.
+        # Synchronize that boolean first so every rank takes the same collective
+        # path below.
+        has_attention_mask = torch.tensor([attention_mask is not None], dtype=torch.int64, device=query.device)
+        global_has_attention_mask = [torch.empty_like(has_attention_mask) for _ in range(sp_world_size)]
+        dist.all_gather(global_has_attention_mask, has_attention_mask, group=self.spg)
+
+        # Padded path: at least one shard has real padding, so rebuild the full
+        # sequence mask for all ranks. Ranks whose local mask was optimized away
+        # contribute an all-ones shard.
+        if torch.any(torch.stack(global_has_attention_mask)):
             if attention_mask is None:
-                attention_mask = torch.ones(q.shape[0], q.shape[1], dtype=torch.int64, device=q.device)
+                if local_position_ids is not None:
+                    attention_mask = torch.ones_like(local_position_ids, dtype=torch.int64)
+                else:
+                    attention_mask = torch.ones(query.shape[0], query.shape[1], dtype=torch.int64, device=query.device)
             else:
                 attention_mask = attention_mask.to(torch.int64)
 
-            global_attention_mask = [
-                torch.empty_like(attention_mask) for _ in range(get_ulysses_sequence_parallel_world_size(self.spg))
-            ]
+            global_attention_mask = [torch.empty_like(attention_mask) for _ in range(sp_world_size)]
             dist.all_gather(global_attention_mask, attention_mask, group=self.spg)
-            attention_mask = torch.cat(global_attention_mask, dim=1)
+            attention_mask = torch.cat(global_attention_mask, dim=1).contiguous()
 
+        # Packed/dense path: no rank has a mask, so leave attention_mask as None.
+        # HF can then use position_ids for padding-free packed varlen attention,
+        # or dense flash attention when position_ids are monotonic.
         context_layer = self.attn_fn(
             q,
             k,

src/llamafactory/v1/plugins/model_plugins/peft.py

@@ -79,7 +79,7 @@ class PeftPlugin(BasePlugin):
 
 def _find_all_linear_modules(model: HFModel) -> list[str]:
     r"""Find all available modules to apply LoRA."""
-    forbidden_modules = {"lm_head", "output_layer", "output"}
+    forbidden_modules = {"lm_head", "output_layer", "output", "score", "classifier"}
     module_names = set()
     for name, module in model.named_modules():
         if any(forbidden_module in name for forbidden_module in forbidden_modules):
@@ -167,8 +167,16 @@ def get_lora_model(model: HFModel, config: LoraConfigDict, is_train: bool = Fals
 
     logger.info_rank0(f"LoRA target modules: {target_modules}")
 
+    cls_name = model.__class__.__name__
+    if cls_name.endswith("ForTokenClassification"):
+        task_type = TaskType.TOKEN_CLS
+    elif cls_name.endswith("ForSequenceClassification"):
+        task_type = TaskType.SEQ_CLS
+    else:
+        task_type = TaskType.CAUSAL_LM
+
     peft_config = LoraConfig(
-        task_type=TaskType.CAUSAL_LM,
+        task_type=task_type,
         inference_mode=not is_train,
         r=config.get("r", 8),
         lora_alpha=config.get("lora_alpha", 16),

src/llamafactory/v1/plugins/trainer_plugins/batching.py

@@ -149,8 +149,8 @@ def _pack_padding_free_samples(samples: list[ModelInput], cutoff_len: int) -> Ba
         return None
 
     packed["position_ids"] = position_ids
-    packed["attention_mask"] = [1] * len(position_ids)
-    return {key: torch.tensor(value).unsqueeze(0) for key, value in packed.items()}
+    packed["attention_mask"] = None
+    return {key: None if value is None else torch.tensor(value).unsqueeze(0) for key, value in packed.items()}
 
 
 @BatchingPlugin("padding_free").register("get_data_provider_batch_size")