mirror of https://github.com/razor-ai/soup.git
756 lines
28 KiB
Python
756 lines
28 KiB
Python
"""Tests for v0.15.0 Performance + Long-context features.
|
|
|
|
Covers: Liger Kernel, FlashAttention, FSDP2, Ring FlashAttention, long-context (RoPE).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from soup_cli.config.schema import TEMPLATES, SoupConfig
|
|
|
|
# ─── Liger Kernel Tests ─────────────────────────────────────────────────
|
|
|
|
|
|
class TestLigerConfig:
|
|
"""Test Liger Kernel configuration in SoupConfig."""
|
|
|
|
def test_use_liger_default_false(self):
|
|
cfg = SoupConfig(base="test/model", data={"train": "./data.jsonl"})
|
|
assert cfg.training.use_liger is False
|
|
|
|
def test_use_liger_enabled(self):
|
|
cfg = SoupConfig(
|
|
base="test/model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"use_liger": True},
|
|
)
|
|
assert cfg.training.use_liger is True
|
|
|
|
|
|
class TestLigerValidation:
|
|
"""Test Liger Kernel validation logic."""
|
|
|
|
def test_validate_liger_disabled_returns_empty(self):
|
|
from soup_cli.utils.liger import validate_liger_config
|
|
|
|
errors = validate_liger_config(False, "transformers", "cuda")
|
|
assert errors == []
|
|
|
|
def test_validate_liger_not_installed(self):
|
|
from soup_cli.utils.liger import validate_liger_config
|
|
|
|
with patch("soup_cli.utils.liger.check_liger_available", return_value=False):
|
|
errors = validate_liger_config(True, "transformers", "cuda")
|
|
assert any("not installed" in err for err in errors)
|
|
|
|
def test_validate_liger_unsloth_incompatible(self):
|
|
from soup_cli.utils.liger import validate_liger_config
|
|
|
|
errors = validate_liger_config(True, "unsloth", "cuda")
|
|
assert any("unsloth" in err.lower() for err in errors)
|
|
|
|
def test_validate_liger_cpu_incompatible(self):
|
|
from soup_cli.utils.liger import validate_liger_config
|
|
|
|
errors = validate_liger_config(True, "transformers", "cpu")
|
|
assert any("CUDA" in err for err in errors)
|
|
|
|
def test_validate_liger_valid_config(self):
|
|
from soup_cli.utils.liger import validate_liger_config
|
|
|
|
with patch("soup_cli.utils.liger.check_liger_available", return_value=True):
|
|
errors = validate_liger_config(True, "transformers", "cuda")
|
|
assert errors == []
|
|
|
|
|
|
class TestLigerDetection:
|
|
"""Test Liger Kernel availability detection."""
|
|
|
|
def test_check_liger_available_not_installed(self):
|
|
from soup_cli.utils.liger import check_liger_available
|
|
|
|
with patch.dict("sys.modules", {"liger_kernel": None}):
|
|
# When import fails, should return False
|
|
result = check_liger_available()
|
|
# Result depends on actual environment; just verify it's bool
|
|
assert isinstance(result, bool)
|
|
|
|
def test_get_liger_version_not_installed(self):
|
|
from soup_cli.utils.liger import get_liger_version
|
|
|
|
with patch("soup_cli.utils.liger.check_liger_available", return_value=False):
|
|
# get_liger_version does its own import attempt
|
|
result = get_liger_version()
|
|
assert result is None or isinstance(result, str)
|
|
|
|
def test_apply_liger_kernel_not_available(self):
|
|
from soup_cli.utils.liger import apply_liger_kernel
|
|
|
|
with patch("soup_cli.utils.liger.check_liger_available", return_value=False):
|
|
result = apply_liger_kernel("meta-llama/Llama-3.1-8B")
|
|
assert result is False
|
|
|
|
def test_apply_liger_manual_unknown_model(self):
|
|
from soup_cli.utils.liger import _apply_liger_manual
|
|
|
|
result = _apply_liger_manual("completely-unknown-model")
|
|
assert result is False
|
|
|
|
|
|
# ─── FlashAttention Tests ────────────────────────────────────────────────
|
|
|
|
|
|
class TestFlashAttnConfig:
|
|
"""Test FlashAttention configuration in SoupConfig."""
|
|
|
|
def test_use_flash_attn_default_false(self):
|
|
cfg = SoupConfig(base="test/model", data={"train": "./data.jsonl"})
|
|
assert cfg.training.use_flash_attn is False
|
|
|
|
def test_use_flash_attn_enabled(self):
|
|
cfg = SoupConfig(
|
|
base="test/model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"use_flash_attn": True},
|
|
)
|
|
assert cfg.training.use_flash_attn is True
|
|
|
|
|
|
class TestFlashAttnDetection:
|
|
"""Test FlashAttention detection and validation."""
|
|
|
|
def test_check_flash_attn_available_no_cuda(self):
|
|
with patch("soup_cli.utils.flash_attn.check_flash_attn_available") as mock_detect:
|
|
mock_detect.return_value = None
|
|
result = mock_detect()
|
|
assert result is None
|
|
|
|
def test_get_attn_implementation_disabled(self):
|
|
from soup_cli.utils.flash_attn import get_attn_implementation
|
|
|
|
result = get_attn_implementation(use_flash_attn=False, device="cuda")
|
|
assert result is None
|
|
|
|
def test_get_attn_implementation_cpu(self):
|
|
from soup_cli.utils.flash_attn import get_attn_implementation
|
|
|
|
result = get_attn_implementation(use_flash_attn=True, device="cpu")
|
|
assert result is None
|
|
|
|
def test_validate_flash_attn_disabled(self):
|
|
from soup_cli.utils.flash_attn import validate_flash_attn_config
|
|
|
|
errors = validate_flash_attn_config(False, "transformers", "cuda")
|
|
assert errors == []
|
|
|
|
def test_validate_flash_attn_cpu_error(self):
|
|
from soup_cli.utils.flash_attn import validate_flash_attn_config
|
|
|
|
errors = validate_flash_attn_config(True, "transformers", "cpu")
|
|
assert any("CUDA" in err for err in errors)
|
|
|
|
def test_validate_flash_attn_unsloth_no_error(self):
|
|
"""Unsloth handles FlashAttention internally — not an error."""
|
|
from soup_cli.utils.flash_attn import validate_flash_attn_config
|
|
|
|
# Even with unsloth, the only error should be about availability (not backend)
|
|
errors = validate_flash_attn_config(True, "unsloth", "cuda")
|
|
assert not any("unsloth" in err.lower() for err in errors)
|
|
|
|
def test_get_flash_attn_version_not_installed(self):
|
|
from soup_cli.utils.flash_attn import get_flash_attn_version
|
|
|
|
result = get_flash_attn_version()
|
|
assert result is None or isinstance(result, str)
|
|
|
|
def test_flash_attn_versions_constant(self):
|
|
from soup_cli.utils.flash_attn import FLASH_ATTN_VERSIONS
|
|
|
|
assert "flash_attention_2" in FLASH_ATTN_VERSIONS
|
|
assert "flash_attention_3" in FLASH_ATTN_VERSIONS
|
|
|
|
|
|
# ─── FSDP2 Tests ────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestFSDPConfig:
|
|
"""Test FSDP2 configuration and presets."""
|
|
|
|
def test_fsdp_full_shard_preset(self):
|
|
from soup_cli.utils.fsdp import get_fsdp_config
|
|
|
|
config = get_fsdp_config("full_shard")
|
|
assert "full_shard" in config["fsdp"]
|
|
assert "auto_wrap" in config["fsdp"]
|
|
|
|
def test_fsdp_shard_grad_preset(self):
|
|
from soup_cli.utils.fsdp import get_fsdp_config
|
|
|
|
config = get_fsdp_config("shard_grad")
|
|
assert "shard_grad_op" in config["fsdp"]
|
|
|
|
def test_fsdp_full_offload_preset(self):
|
|
from soup_cli.utils.fsdp import get_fsdp_config
|
|
|
|
config = get_fsdp_config("full_offload")
|
|
assert "offload" in config["fsdp"]
|
|
|
|
def test_fsdp_unknown_preset_raises(self):
|
|
from soup_cli.utils.fsdp import get_fsdp_config
|
|
|
|
with pytest.raises(ValueError, match="Unknown FSDP config"):
|
|
get_fsdp_config("nonexistent")
|
|
|
|
def test_fsdp_training_args_keys(self):
|
|
from soup_cli.utils.fsdp import get_fsdp_training_args
|
|
|
|
kwargs = get_fsdp_training_args("full_shard")
|
|
assert "fsdp" in kwargs
|
|
assert "fsdp_config" in kwargs
|
|
|
|
def test_fsdp_config_deep_copy(self):
|
|
"""get_fsdp_config should return a deep copy (no shared state)."""
|
|
from soup_cli.utils.fsdp import get_fsdp_config
|
|
|
|
config1 = get_fsdp_config("full_shard")
|
|
config2 = get_fsdp_config("full_shard")
|
|
config1["fsdp"] = "modified"
|
|
assert config2["fsdp"] != "modified"
|
|
|
|
def test_fsdp_configs_dict(self):
|
|
from soup_cli.utils.fsdp import FSDP_CONFIGS
|
|
|
|
assert "full_shard" in FSDP_CONFIGS
|
|
assert "shard_grad" in FSDP_CONFIGS
|
|
assert "full_offload" in FSDP_CONFIGS
|
|
|
|
|
|
class TestFSDPValidation:
|
|
"""Test FSDP2 validation logic."""
|
|
|
|
def test_validate_fsdp_disabled(self):
|
|
from soup_cli.utils.fsdp import validate_fsdp_config
|
|
|
|
errors = validate_fsdp_config(None, None, "transformers", "cuda")
|
|
assert errors == []
|
|
|
|
def test_validate_fsdp_with_deepspeed_conflict(self):
|
|
from soup_cli.utils.fsdp import validate_fsdp_config
|
|
|
|
errors = validate_fsdp_config("full_shard", "/tmp/ds.json", "transformers", "cuda")
|
|
assert any("DeepSpeed" in err for err in errors)
|
|
|
|
def test_validate_fsdp_cpu_error(self):
|
|
from soup_cli.utils.fsdp import validate_fsdp_config
|
|
|
|
errors = validate_fsdp_config("full_shard", None, "transformers", "cpu")
|
|
assert any("CUDA" in err for err in errors)
|
|
|
|
def test_validate_fsdp_unsloth_error(self):
|
|
from soup_cli.utils.fsdp import validate_fsdp_config
|
|
|
|
errors = validate_fsdp_config("full_shard", None, "unsloth", "cuda")
|
|
assert any("unsloth" in err.lower() for err in errors)
|
|
|
|
def test_validate_fsdp_unknown_preset(self):
|
|
from soup_cli.utils.fsdp import validate_fsdp_config
|
|
|
|
errors = validate_fsdp_config("invalid_preset", None, "transformers", "cuda")
|
|
assert any("Unknown FSDP preset" in err for err in errors)
|
|
|
|
|
|
class TestFSDPAvailability:
|
|
"""Test FSDP2 availability detection."""
|
|
|
|
def test_is_fsdp_available_returns_bool(self):
|
|
from soup_cli.utils.fsdp import is_fsdp_available
|
|
|
|
result = is_fsdp_available()
|
|
assert isinstance(result, bool)
|
|
|
|
|
|
# ─── Ring FlashAttention Tests ───────────────────────────────────────────
|
|
|
|
|
|
class TestRingAttentionConfig:
|
|
"""Test Ring FlashAttention configuration."""
|
|
|
|
def test_use_ring_attention_default_false(self):
|
|
cfg = SoupConfig(base="test/model", data={"train": "./data.jsonl"})
|
|
assert cfg.training.use_ring_attention is False
|
|
|
|
def test_use_ring_attention_enabled(self):
|
|
cfg = SoupConfig(
|
|
base="test/model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"use_ring_attention": True},
|
|
)
|
|
assert cfg.training.use_ring_attention is True
|
|
|
|
|
|
class TestRingAttentionUtils:
|
|
"""Test Ring FlashAttention utility functions."""
|
|
|
|
def test_get_sequence_parallel_size_single_gpu(self):
|
|
from soup_cli.utils.ring_attention import get_sequence_parallel_size
|
|
|
|
result = get_sequence_parallel_size(gpu_count=1, max_length=131072)
|
|
assert result == 1
|
|
|
|
def test_get_sequence_parallel_size_short_sequence(self):
|
|
from soup_cli.utils.ring_attention import get_sequence_parallel_size
|
|
|
|
result = get_sequence_parallel_size(gpu_count=4, max_length=2048)
|
|
assert result == 1
|
|
|
|
def test_get_sequence_parallel_size_long_sequence(self):
|
|
from soup_cli.utils.ring_attention import get_sequence_parallel_size
|
|
|
|
result = get_sequence_parallel_size(gpu_count=8, max_length=131072)
|
|
assert result >= 2
|
|
# Should be power of 2
|
|
assert result & (result - 1) == 0
|
|
|
|
def test_get_sequence_parallel_size_power_of_two(self):
|
|
from soup_cli.utils.ring_attention import get_sequence_parallel_size
|
|
|
|
result = get_sequence_parallel_size(gpu_count=6, max_length=131072)
|
|
# Should be power of 2, max <= gpu_count
|
|
assert result in (1, 2, 4)
|
|
|
|
|
|
class TestRingAttentionValidation:
|
|
"""Test Ring FlashAttention validation."""
|
|
|
|
def test_validate_disabled(self):
|
|
from soup_cli.utils.ring_attention import validate_ring_attention_config
|
|
|
|
errors = validate_ring_attention_config(False, "cuda", 131072)
|
|
assert errors == []
|
|
|
|
def test_validate_cpu_error(self):
|
|
from soup_cli.utils.ring_attention import validate_ring_attention_config
|
|
|
|
errors = validate_ring_attention_config(True, "cpu", 131072)
|
|
assert any("CUDA" in err for err in errors)
|
|
|
|
def test_validate_short_sequence_warning(self):
|
|
from soup_cli.utils.ring_attention import validate_ring_attention_config
|
|
|
|
errors = validate_ring_attention_config(True, "cuda", 2048)
|
|
assert any("8192" in err for err in errors)
|
|
|
|
def test_check_ring_attention_available_returns_bool(self):
|
|
from soup_cli.utils.ring_attention import check_ring_attention_available
|
|
|
|
result = check_ring_attention_available()
|
|
assert isinstance(result, bool)
|
|
|
|
def test_get_ring_attention_version_not_installed(self):
|
|
from soup_cli.utils.ring_attention import get_ring_attention_version
|
|
|
|
result = get_ring_attention_version()
|
|
assert result is None or isinstance(result, str)
|
|
|
|
|
|
# ─── Long-Context Tests ─────────────────────────────────────────────────
|
|
|
|
|
|
class TestLongContextConfig:
|
|
"""Test long-context configuration in SoupConfig."""
|
|
|
|
def test_rope_scaling_type_default_none(self):
|
|
cfg = SoupConfig(base="test/model", data={"train": "./data.jsonl"})
|
|
assert cfg.training.rope_scaling_type is None
|
|
|
|
def test_rope_scaling_type_dynamic(self):
|
|
cfg = SoupConfig(
|
|
base="test/model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"rope_scaling_type": "dynamic"},
|
|
)
|
|
assert cfg.training.rope_scaling_type == "dynamic"
|
|
|
|
def test_rope_scaling_type_invalid_rejected(self):
|
|
from pydantic import ValidationError
|
|
|
|
with pytest.raises(ValidationError):
|
|
SoupConfig(
|
|
base="test/model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"rope_scaling_type": "invalid_type"},
|
|
)
|
|
|
|
def test_gradient_checkpointing_default_false(self):
|
|
cfg = SoupConfig(base="test/model", data={"train": "./data.jsonl"})
|
|
assert cfg.training.gradient_checkpointing is False
|
|
|
|
def test_gradient_checkpointing_enabled(self):
|
|
cfg = SoupConfig(
|
|
base="test/model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"gradient_checkpointing": True},
|
|
)
|
|
assert cfg.training.gradient_checkpointing is True
|
|
|
|
|
|
class TestLongContextUtils:
|
|
"""Test long-context utility functions."""
|
|
|
|
def test_rope_scaling_types_constant(self):
|
|
from soup_cli.utils.long_context import ROPE_SCALING_TYPES
|
|
|
|
assert "linear" in ROPE_SCALING_TYPES
|
|
assert "dynamic" in ROPE_SCALING_TYPES
|
|
assert "yarn" in ROPE_SCALING_TYPES
|
|
assert "longrope" in ROPE_SCALING_TYPES
|
|
|
|
def test_model_default_context_llama3(self):
|
|
from soup_cli.utils.long_context import get_model_default_context
|
|
|
|
ctx = get_model_default_context("meta-llama/Llama-3.1-8B")
|
|
assert ctx == 8192
|
|
|
|
def test_model_default_context_mistral(self):
|
|
from soup_cli.utils.long_context import get_model_default_context
|
|
|
|
ctx = get_model_default_context("mistralai/Mistral-7B-v0.1")
|
|
assert ctx == 32768
|
|
|
|
def test_model_default_context_unknown(self):
|
|
from soup_cli.utils.long_context import get_model_default_context
|
|
|
|
ctx = get_model_default_context("unknown/model-xyz")
|
|
assert ctx == 4096 # Conservative default
|
|
|
|
def test_get_rope_scaling_config_linear(self):
|
|
from soup_cli.utils.long_context import get_rope_scaling_config
|
|
|
|
config = get_rope_scaling_config("linear", 131072, 8192)
|
|
assert config["type"] == "linear"
|
|
assert config["factor"] == pytest.approx(16.0)
|
|
|
|
def test_get_rope_scaling_config_dynamic(self):
|
|
from soup_cli.utils.long_context import get_rope_scaling_config
|
|
|
|
config = get_rope_scaling_config("dynamic", 65536, 8192)
|
|
assert config["type"] == "dynamic"
|
|
assert config["factor"] == pytest.approx(8.0)
|
|
|
|
def test_get_rope_scaling_config_yarn(self):
|
|
from soup_cli.utils.long_context import get_rope_scaling_config
|
|
|
|
config = get_rope_scaling_config("yarn", 32768, 8192)
|
|
assert config["type"] == "yarn"
|
|
assert config["factor"] == pytest.approx(4.0)
|
|
assert config["original_max_position_embeddings"] == 8192
|
|
|
|
def test_get_rope_scaling_config_longrope(self):
|
|
from soup_cli.utils.long_context import get_rope_scaling_config
|
|
|
|
config = get_rope_scaling_config("longrope", 32768, 8192)
|
|
assert config["type"] == "longrope"
|
|
|
|
def test_get_rope_scaling_config_no_scaling_needed(self):
|
|
from soup_cli.utils.long_context import get_rope_scaling_config
|
|
|
|
config = get_rope_scaling_config("linear", 4096, 8192)
|
|
assert config == {}
|
|
|
|
def test_get_rope_scaling_config_factor_as_target(self):
|
|
"""When target_length < original_length and > 1.0, treat as factor."""
|
|
from soup_cli.utils.long_context import get_rope_scaling_config
|
|
|
|
config = get_rope_scaling_config("linear", 4.0, 4096)
|
|
assert config["type"] == "linear"
|
|
assert config["factor"] == pytest.approx(4.0)
|
|
|
|
def test_get_rope_scaling_config_dynamic_factor(self):
|
|
"""Dynamic scaling with factor-style argument."""
|
|
from soup_cli.utils.long_context import get_rope_scaling_config
|
|
|
|
config = get_rope_scaling_config("dynamic", 2.0, 8192)
|
|
assert config["type"] == "dynamic"
|
|
assert config["factor"] == pytest.approx(2.0)
|
|
|
|
def test_get_rope_scaling_config_invalid_type(self):
|
|
from soup_cli.utils.long_context import get_rope_scaling_config
|
|
|
|
with pytest.raises(ValueError, match="Unknown RoPE scaling type"):
|
|
get_rope_scaling_config("invalid", 131072, 8192)
|
|
|
|
def test_apply_long_context_config_modifies_model(self):
|
|
from soup_cli.utils.long_context import apply_long_context_config
|
|
|
|
model_config = MagicMock()
|
|
model_config.max_position_embeddings = 8192
|
|
model_config.rope_scaling = None
|
|
|
|
rope_config = apply_long_context_config(
|
|
model_config, target_length=131072, rope_scaling_type="dynamic",
|
|
model_name="meta-llama/Llama-3.1-8B",
|
|
)
|
|
assert rope_config is not None
|
|
assert rope_config["type"] == "dynamic"
|
|
assert model_config.max_position_embeddings == 131072
|
|
|
|
def test_apply_long_context_config_no_scaling_needed(self):
|
|
from soup_cli.utils.long_context import apply_long_context_config
|
|
|
|
model_config = MagicMock()
|
|
model_config.max_position_embeddings = 131072
|
|
|
|
result = apply_long_context_config(
|
|
model_config, target_length=8192, rope_scaling_type="dynamic",
|
|
model_name="test/model",
|
|
)
|
|
assert result is None
|
|
|
|
|
|
class TestLongContextValidation:
|
|
"""Test long-context validation logic."""
|
|
|
|
def test_validate_no_rope_scaling(self):
|
|
from soup_cli.utils.long_context import validate_long_context_config
|
|
|
|
errors = validate_long_context_config(2048, None, False)
|
|
assert errors == []
|
|
|
|
def test_validate_invalid_rope_type(self):
|
|
from soup_cli.utils.long_context import validate_long_context_config
|
|
|
|
errors = validate_long_context_config(131072, "invalid_type", True)
|
|
assert any("Unknown RoPE" in err for err in errors)
|
|
|
|
def test_validate_long_context_no_gradient_checkpointing(self):
|
|
from soup_cli.utils.long_context import validate_long_context_config
|
|
|
|
errors = validate_long_context_config(131072, "dynamic", False)
|
|
assert any("gradient checkpointing" in err.lower() for err in errors)
|
|
|
|
def test_validate_long_context_with_gradient_checkpointing(self):
|
|
from soup_cli.utils.long_context import validate_long_context_config
|
|
|
|
errors = validate_long_context_config(131072, "dynamic", True)
|
|
assert not any("gradient checkpointing" in err.lower() for err in errors)
|
|
|
|
|
|
# ─── Template Tests ──────────────────────────────────────────────────────
|
|
|
|
|
|
class TestMaxLengthBounds:
|
|
"""Test max_length has proper bounds (security fix H3+M3)."""
|
|
|
|
def test_max_length_too_small_rejected(self):
|
|
from pydantic import ValidationError
|
|
|
|
with pytest.raises(ValidationError):
|
|
SoupConfig(
|
|
base="test/model",
|
|
data={"train": "./data.jsonl", "max_length": 0},
|
|
)
|
|
|
|
def test_max_length_too_large_rejected(self):
|
|
from pydantic import ValidationError
|
|
|
|
with pytest.raises(ValidationError):
|
|
SoupConfig(
|
|
base="test/model",
|
|
data={"train": "./data.jsonl", "max_length": 2000000},
|
|
)
|
|
|
|
def test_max_length_valid_128k(self):
|
|
cfg = SoupConfig(
|
|
base="test/model",
|
|
data={"train": "./data.jsonl", "max_length": 131072},
|
|
)
|
|
assert cfg.data.max_length == 131072
|
|
|
|
|
|
class TestLongContextTemplate:
|
|
"""Test that the longcontext template exists and is valid."""
|
|
|
|
def test_longcontext_template_exists(self):
|
|
assert "longcontext" in TEMPLATES
|
|
|
|
def test_longcontext_template_has_max_length(self):
|
|
assert "131072" in TEMPLATES["longcontext"]
|
|
|
|
def test_longcontext_template_has_rope_scaling(self):
|
|
assert "rope_scaling_type" in TEMPLATES["longcontext"]
|
|
|
|
def test_longcontext_template_has_gradient_checkpointing(self):
|
|
assert "gradient_checkpointing" in TEMPLATES["longcontext"]
|
|
|
|
def test_longcontext_template_has_flash_attn(self):
|
|
assert "use_flash_attn" in TEMPLATES["longcontext"]
|
|
|
|
def test_longcontext_template_count(self):
|
|
"""Should now have 17 templates (16 + bco added in v0.40.0)."""
|
|
assert len(TEMPLATES) == 17
|
|
|
|
|
|
# ─── Trainer fsdp_config Parameter Tests ─────────────────────────────────
|
|
|
|
|
|
class TestTrainerFSDPParam:
|
|
"""Test that all trainers accept fsdp_config parameter."""
|
|
|
|
def test_sft_trainer_accepts_fsdp_config(self):
|
|
from soup_cli.trainer.sft import SFTTrainerWrapper
|
|
|
|
cfg = SoupConfig(base="test/model", data={"train": "./data.jsonl"})
|
|
wrapper = SFTTrainerWrapper(cfg, fsdp_config={"fsdp": "full_shard auto_wrap"})
|
|
assert wrapper.fsdp_config == {"fsdp": "full_shard auto_wrap"}
|
|
|
|
def test_dpo_trainer_accepts_fsdp_config(self):
|
|
from soup_cli.trainer.dpo import DPOTrainerWrapper
|
|
|
|
cfg = SoupConfig(base="test/model", task="dpo", data={"train": "./data.jsonl"})
|
|
wrapper = DPOTrainerWrapper(cfg, fsdp_config={"fsdp": "full_shard"})
|
|
assert wrapper.fsdp_config == {"fsdp": "full_shard"}
|
|
|
|
def test_grpo_trainer_accepts_fsdp_config(self):
|
|
from soup_cli.trainer.grpo import GRPOTrainerWrapper
|
|
|
|
cfg = SoupConfig(base="test/model", task="grpo", data={"train": "./data.jsonl"})
|
|
wrapper = GRPOTrainerWrapper(cfg, fsdp_config={"fsdp": "full_shard"})
|
|
assert wrapper.fsdp_config == {"fsdp": "full_shard"}
|
|
|
|
def test_kto_trainer_accepts_fsdp_config(self):
|
|
from soup_cli.trainer.kto import KTOTrainerWrapper
|
|
|
|
cfg = SoupConfig(base="test/model", task="kto", data={"train": "./data.jsonl"})
|
|
wrapper = KTOTrainerWrapper(cfg, fsdp_config={"fsdp": "full_shard"})
|
|
assert wrapper.fsdp_config == {"fsdp": "full_shard"}
|
|
|
|
def test_orpo_trainer_accepts_fsdp_config(self):
|
|
from soup_cli.trainer.orpo import ORPOTrainerWrapper
|
|
|
|
cfg = SoupConfig(base="test/model", task="orpo", data={"train": "./data.jsonl"})
|
|
wrapper = ORPOTrainerWrapper(cfg, fsdp_config={"fsdp": "full_shard"})
|
|
assert wrapper.fsdp_config == {"fsdp": "full_shard"}
|
|
|
|
def test_simpo_trainer_accepts_fsdp_config(self):
|
|
from soup_cli.trainer.simpo import SimPOTrainerWrapper
|
|
|
|
cfg = SoupConfig(base="test/model", task="simpo", data={"train": "./data.jsonl"})
|
|
wrapper = SimPOTrainerWrapper(cfg, fsdp_config={"fsdp": "full_shard"})
|
|
assert wrapper.fsdp_config == {"fsdp": "full_shard"}
|
|
|
|
def test_ipo_trainer_accepts_fsdp_config(self):
|
|
from soup_cli.trainer.ipo import IPOTrainerWrapper
|
|
|
|
cfg = SoupConfig(base="test/model", task="ipo", data={"train": "./data.jsonl"})
|
|
wrapper = IPOTrainerWrapper(cfg, fsdp_config={"fsdp": "full_shard"})
|
|
assert wrapper.fsdp_config == {"fsdp": "full_shard"}
|
|
|
|
def test_pretrain_trainer_accepts_fsdp_config(self):
|
|
from soup_cli.trainer.pretrain import PretrainTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="test/model", task="pretrain",
|
|
data={"train": "./data.jsonl", "format": "plaintext"},
|
|
)
|
|
wrapper = PretrainTrainerWrapper(cfg, fsdp_config={"fsdp": "full_shard"})
|
|
assert wrapper.fsdp_config == {"fsdp": "full_shard"}
|
|
|
|
def test_reward_model_trainer_accepts_fsdp_config(self):
|
|
from soup_cli.trainer.reward_model import RewardModelTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="test/model", task="reward_model",
|
|
data={"train": "./data.jsonl"},
|
|
)
|
|
wrapper = RewardModelTrainerWrapper(cfg, fsdp_config={"fsdp": "full_shard"})
|
|
assert wrapper.fsdp_config == {"fsdp": "full_shard"}
|
|
|
|
def test_ppo_trainer_accepts_fsdp_config(self):
|
|
from soup_cli.trainer.ppo import PPOTrainerWrapper
|
|
|
|
cfg = SoupConfig(base="test/model", task="ppo", data={"train": "./data.jsonl"})
|
|
wrapper = PPOTrainerWrapper(cfg, fsdp_config={"fsdp": "full_shard"})
|
|
assert wrapper.fsdp_config == {"fsdp": "full_shard"}
|
|
|
|
def test_sft_trainer_fsdp_config_default_none(self):
|
|
from soup_cli.trainer.sft import SFTTrainerWrapper
|
|
|
|
cfg = SoupConfig(base="test/model", data={"train": "./data.jsonl"})
|
|
wrapper = SFTTrainerWrapper(cfg)
|
|
assert wrapper.fsdp_config is None
|
|
|
|
|
|
# ─── Combined Config Tests ───────────────────────────────────────────────
|
|
|
|
|
|
class TestCombinedPerformanceConfig:
|
|
"""Test multiple performance features enabled together."""
|
|
|
|
def test_all_features_enabled(self):
|
|
cfg = SoupConfig(
|
|
base="test/model",
|
|
data={"train": "./data.jsonl", "max_length": 131072},
|
|
training={
|
|
"use_liger": True,
|
|
"use_flash_attn": True,
|
|
"use_ring_attention": True,
|
|
"rope_scaling_type": "dynamic",
|
|
"gradient_checkpointing": True,
|
|
},
|
|
)
|
|
assert cfg.training.use_liger is True
|
|
assert cfg.training.use_flash_attn is True
|
|
assert cfg.training.use_ring_attention is True
|
|
assert cfg.training.rope_scaling_type == "dynamic"
|
|
assert cfg.training.gradient_checkpointing is True
|
|
assert cfg.data.max_length == 131072
|
|
|
|
def test_config_serialization_roundtrip(self):
|
|
cfg = SoupConfig(
|
|
base="test/model",
|
|
data={"train": "./data.jsonl"},
|
|
training={
|
|
"use_liger": True,
|
|
"use_flash_attn": True,
|
|
"gradient_checkpointing": True,
|
|
"rope_scaling_type": "yarn",
|
|
},
|
|
)
|
|
dumped = cfg.model_dump()
|
|
assert dumped["training"]["use_liger"] is True
|
|
assert dumped["training"]["use_flash_attn"] is True
|
|
assert dumped["training"]["gradient_checkpointing"] is True
|
|
assert dumped["training"]["rope_scaling_type"] == "yarn"
|
|
|
|
def test_rope_scaling_all_types_accepted(self):
|
|
for scaling_type in ("linear", "dynamic", "yarn", "longrope"):
|
|
cfg = SoupConfig(
|
|
base="test/model",
|
|
data={"train": "./data.jsonl"},
|
|
training={"rope_scaling_type": scaling_type},
|
|
)
|
|
assert cfg.training.rope_scaling_type == scaling_type
|
|
|
|
|
|
# ─── Train Command --fsdp Flag Tests ────────────────────────────────────
|
|
|
|
|
|
class TestTrainCommandFSDPFlag:
|
|
"""Test --fsdp flag in train command."""
|
|
|
|
def test_fsdp_configs_accessible_from_fsdp_module(self):
|
|
from soup_cli.utils.fsdp import FSDP_CONFIGS
|
|
|
|
assert len(FSDP_CONFIGS) == 3
|
|
|
|
def test_fsdp_full_shard_has_correct_keys(self):
|
|
from soup_cli.utils.fsdp import FSDP_FULL_SHARD
|
|
|
|
assert "fsdp" in FSDP_FULL_SHARD
|
|
assert "fsdp_config" in FSDP_FULL_SHARD
|
|
assert "use_orig_params" in FSDP_FULL_SHARD["fsdp_config"]
|
|
assert FSDP_FULL_SHARD["fsdp_config"]["use_orig_params"] is True
|