From b7615dbdc9f95c169361dacad1185852c4d28521 Mon Sep 17 00:00:00 2001 From: jiaqiw09 <60021713+jiaqiw09@users.noreply.github.com> Date: Thu, 25 Jun 2026 20:05:56 +0800 Subject: [PATCH] [v1] Fix device mesh, fix lora for reward model and fix sp (#10555) --- ...train_full_fsdp2_dynamic_padding_free.yaml | 2 +- .../train_full_fsdp2_padding_free.yaml | 2 +- src/llamafactory/v1/accelerator/interface.py | 9 +++++ .../model_plugins/parallelization/ulysses.py | 39 +++++++++++++------ .../v1/plugins/model_plugins/peft.py | 12 +++++- .../v1/plugins/trainer_plugins/batching.py | 4 +- tests_v1/core/utils/test_batching.py | 10 +++-- 7 files changed, 58 insertions(+), 20 deletions(-) diff --git a/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_padding_free.yaml b/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_padding_free.yaml index aa7ab54b0..9898a50d8 100644 --- a/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_padding_free.yaml +++ b/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 diff --git a/examples/v1/train_batching_strategy/train_full_fsdp2_padding_free.yaml b/examples/v1/train_batching_strategy/train_full_fsdp2_padding_free.yaml index 2f96a065e..ec5babca0 100644 --- a/examples/v1/train_batching_strategy/train_full_fsdp2_padding_free.yaml +++ b/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 diff --git a/src/llamafactory/v1/accelerator/interface.py b/src/llamafactory/v1/accelerator/interface.py index f32b4dc6b..a43885a55 100644 --- a/src/llamafactory/v1/accelerator/interface.py +++ b/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( diff --git a/src/llamafactory/v1/plugins/model_plugins/parallelization/ulysses.py b/src/llamafactory/v1/plugins/model_plugins/parallelization/ulysses.py index 4bf568e30..2f3dd18b7 100644 --- a/src/llamafactory/v1/plugins/model_plugins/parallelization/ulysses.py +++ b/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, diff --git a/src/llamafactory/v1/plugins/model_plugins/peft.py b/src/llamafactory/v1/plugins/model_plugins/peft.py index 2ef2035e1..9552de0b4 100644 --- a/src/llamafactory/v1/plugins/model_plugins/peft.py +++ b/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), diff --git a/src/llamafactory/v1/plugins/trainer_plugins/batching.py b/src/llamafactory/v1/plugins/trainer_plugins/batching.py index 23746a6d4..715f4cc5c 100644 --- a/src/llamafactory/v1/plugins/trainer_plugins/batching.py +++ b/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") diff --git a/tests_v1/core/utils/test_batching.py b/tests_v1/core/utils/test_batching.py index ef01e50ee..153931f9c 100644 --- a/tests_v1/core/utils/test_batching.py +++ b/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 - assert packed_batch["attention_mask"].tolist() == [[1] * 15] + # Verify attention_mask: padding-free relies on reset-style position_ids instead of a dense mask. + assert packed_batch["attention_mask"] is None # Verify position_ids assert packed_batch["position_ids"].tolist() == [