mirror of
https://github.com/hiyouga/LLaMA-Factory.git
synced 2026-08-17 13:35:44 +08:00
[v1] add FSDPTurbo EP/EFSDP plugin for MoE training (#10676)
This commit is contained in:
@@ -20,6 +20,7 @@ from llamafactory.v1.accelerator.interface import DistributedInterface
|
||||
from llamafactory.v1.config.model_args import ModelArguments
|
||||
from llamafactory.v1.config.training_args import TrainingArguments
|
||||
from llamafactory.v1.core.model_engine import ModelEngine
|
||||
from llamafactory.v1.plugins.model_plugins.parallelization import ulysses
|
||||
from llamafactory.v1.plugins.model_plugins.parallelization.sequence_parallel import (
|
||||
SequenceParallelModelPlugin,
|
||||
sequence_parallel_loss,
|
||||
@@ -28,6 +29,39 @@ from llamafactory.v1.utils.env import find_available_port
|
||||
from llamafactory.v1.utils.pytest import dist_env
|
||||
|
||||
|
||||
def test_qwen3_5_broadcast_position_ids_keep_packed_boundaries(monkeypatch: pytest.MonkeyPatch):
|
||||
local_position_ids = torch.tensor([[0, 1, 0]])
|
||||
remote_position_ids = torch.tensor([[1, 2, 3]])
|
||||
mrope_position_ids = local_position_ids.unsqueeze(0).expand(3, -1, -1)
|
||||
captured = {}
|
||||
|
||||
monkeypatch.setattr(ulysses.SeqAllToAll4D, "apply", lambda _, tensor, *__: tensor)
|
||||
monkeypatch.setattr(ulysses, "get_ulysses_sequence_parallel_world_size", lambda _: 2)
|
||||
|
||||
def fake_all_gather(outputs, tensor, **_):
|
||||
outputs[0].copy_(tensor)
|
||||
outputs[1].copy_(remote_position_ids if tensor.shape == local_position_ids.shape else tensor)
|
||||
|
||||
def fake_attention(query, _key, _value, _attention_mask, **kwargs):
|
||||
captured["position_ids"] = kwargs["position_ids"]
|
||||
return query
|
||||
|
||||
monkeypatch.setattr(ulysses.dist, "all_gather", fake_all_gather)
|
||||
attention = ulysses.UlyssesAttention(sequence_process_group=object(), attn_fn=fake_attention)
|
||||
hidden_states = torch.zeros(1, 3, 2, 4)
|
||||
|
||||
attention(hidden_states, hidden_states, hidden_states, None, 6, position_ids=mrope_position_ids)
|
||||
|
||||
assert captured["position_ids"].tolist() == [[0, 1, 0, 1, 2, 3]]
|
||||
assert captured["position_ids"].is_contiguous()
|
||||
|
||||
|
||||
def test_true_mrope_position_ids_are_not_used_as_packed_boundaries():
|
||||
mrope_position_ids = torch.tensor([[[0, 1, 2]], [[0, 1, 1]], [[0, 1, 0]]])
|
||||
|
||||
assert ulysses._get_text_position_ids(mrope_position_ids) is None
|
||||
|
||||
|
||||
def _test_sequence_parallel_loss(
|
||||
local_rank: int, world_size: int, master_port: int, cp_size: int, dp_size: int, batch_size: int
|
||||
):
|
||||
|
||||
Reference in New Issue
Block a user