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

examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_padding_free.yaml

@@ -20,7 +20,7 @@ train_dataset: data/v1_sft_demo.yaml
 output_dir: outputs/test_fsdp2
 micro_batch_size: 4
 batching_strategy: dynamic_padding_free
-flash_attn: flash_attention2
+flash_attn: flash_attention_2
 cutoff_len: 2048
 learning_rate: 1.0e-4
 max_steps: 10

examples/v1/train_batching_strategy/train_full_fsdp2_padding_free.yaml

@@ -20,7 +20,7 @@ train_dataset: data/v1_sft_demo.yaml
 output_dir: outputs/test_fsdp2
 micro_batch_size: 4
 batching_strategy: padding_free
-flash_attn: flash_attention2
+flash_attn: flash_attention_2
 cutoff_len: 2048
 learning_rate: 1.0e-4
 max_steps: 10

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 = [
+            global_position_ids = [torch.empty_like(position_ids) for _ in range(sp_world_size)]
-                torch.empty_like(position_ids) for _ in range(get_ulysses_sequence_parallel_world_size(self.spg))
-            ]
             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 = [
+            global_attention_mask = [torch.empty_like(attention_mask) for _ in range(sp_world_size)]
-                torch.empty_like(attention_mask) for _ in range(get_ulysses_sequence_parallel_world_size(self.spg))
-            ]
             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)
+    packed["attention_mask"] = None
-    return {key: torch.tensor(value).unsqueeze(0) for key, value in packed.items()}
+    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")

tests_v1/core/utils/test_batching.py

@@ -32,6 +32,7 @@ def _make_model_input(length: int, start: int = 0):
         "attention_mask": [1] * length,
         "labels": input_ids.copy(),
         "loss_weights": [1.0] * length,
+        "position_ids": list(range(1, length + 1)),
     }
 
 
@@ -62,7 +63,7 @@ def test_padding_free():
     assert len(batch) == 1
     assert batch[0]["input_ids"].shape == (1, 5)
     assert batch[0]["input_ids"].tolist() == [[0, 1, 10, 11, 12]]
-    assert batch[0]["attention_mask"].tolist() == [[1, 1, 1, 1, 1]]
+    assert batch[0]["attention_mask"] is None
     assert batch[0]["position_ids"].tolist() == [[0, 1, 0, 1, 2]]
     assert batch[0]["labels"].tolist() == [[0, 1, IGNORE_INDEX, 11, 12]]
     assert batch[0]["loss_weights"].tolist() == [[1.0, 1.0, 0.0, 1.0, 1.0]]
@@ -115,6 +116,8 @@ def test_dynamic_batching():
     assert len(batch) == 1
     assert batch[0]["input_ids"].shape == (3, 6)
     assert batch[0]["input_ids"].tolist()[0] == [0, 1, 2, 0, 0, 0]
+    assert batch[0]["position_ids"].shape == (3, 6)
+    assert batch[0]["position_ids"].tolist()[0] == [1, 2, 3, 0, 0, 0]
     assert len(buffer) == 3
 
 
@@ -201,6 +204,7 @@ def test_normal_batching():
     batch = next(iter(batch_generator))
     assert len(batch) == 2
     assert batch[0]["input_ids"].shape == (4, 10)
+    assert batch[0]["position_ids"].shape == (4, 10)
 
 
 def test_dynamic_padding_free():
@@ -280,8 +284,8 @@ def test_dynamic_padding_free():
         ]  # Sample 3
     ]
 
-    # Verify attention_mask
+    # Verify attention_mask: padding-free relies on reset-style position_ids instead of a dense mask.
-    assert packed_batch["attention_mask"].tolist() == [[1] * 15]
+    assert packed_batch["attention_mask"] is None
 
     # Verify position_ids
     assert packed_batch["position_ids"].tolist() == [