mirror of https://github.com/razor-ai/soup.git
v0.15.0: performance + long-context fine-tuning
Features: - Liger Kernel integration (fused RMSNorm, SwiGLU, CrossEntropy, RoPE) - FlashAttention v2/v3 auto-detection via attn_implementation - FSDP2 support alongside DeepSpeed (--fsdp flag: full_shard, shard_grad, full_offload) - Ring FlashAttention for sequence parallelism across GPUs - 128k+ context fine-tuning with RoPE scaling (linear, dynamic, yarn, longrope) - Gradient checkpointing for memory-efficient long-sequence training - New 'longcontext' template: soup init --template longcontext Security: - rope_scaling_type validated via Literal constraint (no free-form strings) - max_length bounded ge=64, le=1048576 (prevents OOM/corruption) - FSDP config key allowlist prevents injection of unexpected TrainingArguments - Liger exception handling narrowed (no silent CUDA error swallowing) New files: utils/liger.py, utils/flash_attn.py, utils/fsdp.py, utils/ring_attention.py, utils/long_context.py, tests/test_performance.py Modified: all 10 trainers (fsdp_config support), schema.py, train.py, pyproject.toml 91 new tests (1182 total), 58.5% coverage, ruff clean
This commit is contained in:
Alpamys
parent
e82685c555
commit
fd2513aa8b
|
|
@ -1,12 +1,12 @@
|
|||
# Soup CLI — Project CLAUDE.md
|
||||
|
||||
Soup is a CLI-first LLM fine-tuning tool (v0.14.3). Python 3.9+, MIT license.
|
||||
Soup is a CLI-first LLM fine-tuning tool (v0.15.0). Python 3.9+, MIT license.
|
||||
|
||||
## Build & Development
|
||||
|
||||
```bash
|
||||
pip install -e ".[dev]" # Install editable + test deps
|
||||
pytest tests/ -v --tb=short # Run all tests (1091 tests)
|
||||
pytest tests/ -v --tb=short # Run all tests (1182 tests)
|
||||
ruff check soup_cli/ tests/ # Lint (must pass before commit)
|
||||
ruff check --fix soup_cli/ tests/ # Auto-fix lint issues
|
||||
```
|
||||
|
|
@ -71,6 +71,11 @@ soup_cli/
|
|||
vllm.py # AsyncLLMEngine backend (2-4x inference throughput)
|
||||
galore.py # GaLore optimizer config + validation
|
||||
moe.py # MoE model detection, ScatterMoE LoRA target modules
|
||||
liger.py # Liger Kernel detection + fused ops (RMSNorm, SwiGLU, etc.)
|
||||
flash_attn.py # FlashAttention v2/v3 auto-detection
|
||||
fsdp.py # FSDP2 config templates (full_shard, shard_grad, offload)
|
||||
ring_attention.py # Ring FlashAttention for sequence parallelism
|
||||
long_context.py # RoPE scaling for 128k+ context fine-tuning
|
||||
constants.py # APP_NAME, paths, default chat template
|
||||
tests/ # 47 test files, 1014 tests
|
||||
examples/
|
||||
|
|
@ -82,7 +87,7 @@ examples/
|
|||
|
||||
```
|
||||
soup init # Create config (interactive or --template)
|
||||
soup train # Main training (--config, --resume, --wandb, --tensorboard, --deepspeed, --yes)
|
||||
soup train # Main training (--config, --resume, --wandb, --tensorboard, --deepspeed, --fsdp, --yes)
|
||||
soup infer # Batch inference (--model, --input, --output)
|
||||
soup chat # Terminal chat with model
|
||||
soup serve # OpenAI-compatible inference server (--backend transformers|vllm)
|
||||
|
|
@ -115,10 +120,10 @@ soup version # Show version (--full for details)
|
|||
|
||||
- **SoupConfig**: base (required), task (sft/dpo/kto/orpo/simpo/ipo/grpo/ppo/reward_model/pretrain), modality (text/vision), backend (transformers/unsloth), data, training, output
|
||||
- **DataConfig**: train, format (alpaca/sharegpt/chatml/dpo/kto/llava/sharegpt4v/plaintext/auto), val_split, max_length, image_dir
|
||||
- **TrainingConfig**: epochs, lr, batch_size (int or "auto"), quantization (4bit/8bit/none), quantization_aware, optimizer, scheduler, dpo_beta, kto_beta, orpo_beta, simpo_gamma, cpo_alpha, ipo_tau, grpo_beta, num_generations, reward_fn, ppo_epochs, ppo_clip_ratio, ppo_kl_penalty, reward_model, loraplus_lr_ratio, use_galore, galore_rank, galore_update_proj_gap, galore_scale, moe_lora, moe_aux_loss_coeff
|
||||
- **TrainingConfig**: epochs, lr, batch_size (int or "auto"), quantization (4bit/8bit/none), quantization_aware, optimizer, scheduler, dpo_beta, kto_beta, orpo_beta, simpo_gamma, cpo_alpha, ipo_tau, grpo_beta, num_generations, reward_fn, ppo_epochs, ppo_clip_ratio, ppo_kl_penalty, reward_model, loraplus_lr_ratio, use_galore, galore_rank, galore_update_proj_gap, galore_scale, moe_lora, moe_aux_loss_coeff, use_liger, use_flash_attn, use_ring_attention, rope_scaling_type, gradient_checkpointing
|
||||
- **LoraConfig**: r, alpha, dropout, target_modules, use_dora
|
||||
|
||||
12 built-in templates: chat, code, medical, reasoning, vision, kto, orpo, simpo, ipo, rlhf, pretrain, moe.
|
||||
13 built-in templates: chat, code, medical, reasoning, vision, kto, orpo, simpo, ipo, rlhf, pretrain, moe, longcontext.
|
||||
|
||||
## Training Tasks
|
||||
|
||||
|
|
@ -166,6 +171,10 @@ soup version # Show version (--full for details)
|
|||
- **Data validate**: auto-detects format when `--format` not specified (v0.14.3)
|
||||
- **Data stats**: UTF-8 stdout redirect on Windows for plotext histograms (v0.14.3)
|
||||
- **Web UI**: `--show-token` flag + auth token documented in `--help` (v0.14.3)
|
||||
- **RoPE scaling**: `rope_scaling_type` validated via Literal constraint (v0.15.0)
|
||||
- **max_length bounds**: ge=64, le=1048576 prevents OOM/corruption from extreme values (v0.15.0)
|
||||
- **FSDP config**: key allowlist prevents injection of unexpected TrainingArguments (v0.15.0)
|
||||
- **Liger Kernel**: exception handling narrowed to prevent silent CUDA error swallowing (v0.15.0)
|
||||
|
||||
## Code Conventions
|
||||
|
||||
|
|
@ -192,6 +201,8 @@ soup version # Show version (--full for details)
|
|||
- `[fast]`: unsloth (2-5x training speedup)
|
||||
- `[vision]`: Pillow
|
||||
- `[qat]`: torchao
|
||||
- `[liger]`: liger-kernel (fused ops)
|
||||
- `[ring-attn]`: ring-flash-attn (sequence parallelism)
|
||||
- `[ui]`: FastAPI + uvicorn + static SPA
|
||||
- `[dev]`: pytest, ruff, pytest-cov, httpx
|
||||
|
||||
|
|
@ -229,7 +240,7 @@ soup version # Show version (--full for details)
|
|||
15. **Tag**: `git tag v0.X.Y && git push origin v0.X.Y`
|
||||
16. **Release**: `gh release create v0.X.Y` with changelog (What's New, Install/Upgrade)
|
||||
|
||||
## Tests (48 test files, 1091 tests)
|
||||
## Tests (49 test files, 1182 tests)
|
||||
|
||||
| File | Covers |
|
||||
|------|--------|
|
||||
|
|
@ -280,3 +291,4 @@ soup version # Show version (--full for details)
|
|||
| test_moe.py | MoE detection, ScatterMoE LoRA targets, MoE info extraction |
|
||||
| test_bugfixes.py | v0.10.1-v0.14.3 regression fixes |
|
||||
| test_cli_subprocess.py | Subprocess CLI tests: entry point, encoding, paths, platform regressions |
|
||||
| test_performance.py | Liger Kernel, FlashAttention, FSDP2, Ring Attention, long-context, RoPE scaling |
|
||||
|
|
|
|||
53
README.md
53
README.md
|
|
@ -76,6 +76,7 @@ soup init --template ipo # IPO regularized preference
|
|||
soup init --template rlhf # full RLHF pipeline (SFT→RM→PPO)
|
||||
soup init --template pretrain # continued pre-training on raw text
|
||||
soup init --template moe # MoE fine-tuning (ScatterMoE LoRA)
|
||||
soup init --template longcontext # 128k+ context fine-tuning
|
||||
```
|
||||
|
||||
### 3. Train
|
||||
|
|
@ -779,22 +780,55 @@ soup diff --model-a ./base --model-b ./finetuned --prompts test_prompts.jsonl
|
|||
soup diff --model-a ./a --model-b ./b --prompts prompts.txt --output results.jsonl
|
||||
```
|
||||
|
||||
## Multi-GPU / DeepSpeed
|
||||
## Multi-GPU / DeepSpeed / FSDP
|
||||
|
||||
Train on multiple GPUs with DeepSpeed:
|
||||
Train on multiple GPUs with DeepSpeed or PyTorch FSDP2:
|
||||
|
||||
```bash
|
||||
# ZeRO Stage 2 (recommended for most cases)
|
||||
# DeepSpeed ZeRO Stage 2 (recommended for most cases)
|
||||
soup train --config soup.yaml --deepspeed zero2
|
||||
|
||||
# ZeRO Stage 3 (for very large models)
|
||||
# DeepSpeed ZeRO Stage 3 (for very large models)
|
||||
soup train --config soup.yaml --deepspeed zero3
|
||||
|
||||
# ZeRO Stage 2 with CPU offload (memory-constrained)
|
||||
# DeepSpeed ZeRO Stage 2 with CPU offload (memory-constrained)
|
||||
soup train --config soup.yaml --deepspeed zero2_offload
|
||||
|
||||
# Custom DeepSpeed config
|
||||
soup train --config soup.yaml --deepspeed ./my_ds_config.json
|
||||
# FSDP2 Full Shard (native PyTorch, like ZeRO-3)
|
||||
soup train --config soup.yaml --fsdp fsdp_full_shard
|
||||
|
||||
# FSDP2 Shard Grad Op (like ZeRO-2)
|
||||
soup train --config soup.yaml --fsdp fsdp_shard_grad
|
||||
|
||||
# FSDP2 Full Shard with CPU offload
|
||||
soup train --config soup.yaml --fsdp fsdp_full_offload
|
||||
```
|
||||
|
||||
## Performance + Long-Context
|
||||
|
||||
Optimize training throughput and extend context windows:
|
||||
|
||||
```yaml
|
||||
# soup.yaml — performance options
|
||||
training:
|
||||
use_liger: true # Liger Kernel fused ops (20-60% memory savings)
|
||||
use_flash_attn: true # FlashAttention v2/v3 auto-detection
|
||||
gradient_checkpointing: true # Required for long sequences
|
||||
|
||||
# Long-context (128k+ tokens)
|
||||
rope_scaling_type: dynamic # RoPE scaling: linear, dynamic, yarn, longrope
|
||||
# use_ring_attention: true # Sequence parallelism across GPUs
|
||||
|
||||
data:
|
||||
max_length: 131072 # Up to 1M tokens supported
|
||||
```
|
||||
|
||||
Install optional performance packages:
|
||||
|
||||
```bash
|
||||
pip install 'soup-cli[liger]' # Liger Kernel fused operations
|
||||
pip install flash-attn --no-build-isolation # FlashAttention
|
||||
pip install 'soup-cli[ring-attn]' # Ring FlashAttention (sequence parallelism)
|
||||
```
|
||||
|
||||
## Quickstart Demo
|
||||
|
|
@ -975,9 +1009,10 @@ soup eval --model ./output --benchmarks mmlu --run-id run_20260223_143052_a1b2
|
|||
## All Commands
|
||||
|
||||
```
|
||||
soup init [--template chat|code|...|pretrain|moe] Create config
|
||||
soup init [--template chat|code|...|moe|longcontext] Create config
|
||||
soup train --config soup.yaml Start training
|
||||
soup train --config soup.yaml --tensorboard Train with TensorBoard logging
|
||||
soup train --config soup.yaml --fsdp fsdp_full_shard Train with FSDP2
|
||||
soup infer --model ./output --input p.jsonl Batch inference
|
||||
soup chat --model ./output Interactive chat
|
||||
soup push --model ./output --repo user/name Upload to HuggingFace
|
||||
|
|
@ -1070,6 +1105,8 @@ Soup works with **any** of the **340,000+** text-generation models on [HuggingFa
|
|||
| `data` | `pip install 'soup-cli[data]'` | Deduplication (MinHash via datasketch) |
|
||||
| `eval` | `pip install 'soup-cli[eval]'` | Benchmark evaluation (lm-evaluation-harness) |
|
||||
| `deepspeed` | `pip install 'soup-cli[deepspeed]'` | Multi-GPU training (DeepSpeed ZeRO) |
|
||||
| `liger` | `pip install 'soup-cli[liger]'` | Liger Kernel fused ops (20-60% memory savings) |
|
||||
| `ring-attn` | `pip install 'soup-cli[ring-attn]'` | Ring FlashAttention (sequence parallelism) |
|
||||
| `dev` | `pip install 'soup-cli[dev]'` | Tests + linting (pytest, ruff) |
|
||||
|
||||
## Development
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|||
|
||||
[project]
|
||||
name = "soup-cli"
|
||||
version = "0.14.3"
|
||||
version = "0.15.0"
|
||||
description = "Fine-tune LLMs in one command. No SSH, no config hell."
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
|
|
@ -50,6 +50,8 @@ deepspeed = ["deepspeed>=0.12.0"]
|
|||
fast = ["unsloth>=2024.8"]
|
||||
vision = ["Pillow>=9.0.0"]
|
||||
qat = ["torchao>=0.4.0"]
|
||||
liger = ["liger-kernel>=0.3.0"]
|
||||
ring-attn = ["ring-flash-attn>=0.1.0"]
|
||||
|
||||
[project.scripts]
|
||||
soup = "soup_cli.cli:run"
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
"""Soup CLI — Fine-tune LLMs in one command."""
|
||||
|
||||
__version__ = "0.14.3"
|
||||
__version__ = "0.15.0"
|
||||
|
|
|
|||
|
|
@ -54,6 +54,11 @@ def train(
|
|||
"--deepspeed",
|
||||
help="Enable DeepSpeed: zero2, zero3, zero2_offload, or path to config JSON",
|
||||
),
|
||||
fsdp: str = typer.Option(
|
||||
None,
|
||||
"--fsdp",
|
||||
help="Enable FSDP2: fsdp_full_shard, fsdp_shard_grad, or fsdp_full_offload",
|
||||
),
|
||||
yes: bool = typer.Option(
|
||||
False,
|
||||
"--yes",
|
||||
|
|
@ -128,6 +133,20 @@ def train(
|
|||
if ds_config_path:
|
||||
console.print(f"[green]DeepSpeed enabled:[/] {deepspeed}")
|
||||
|
||||
# --- FSDP2 setup ---
|
||||
fsdp_kwargs = None
|
||||
if fsdp:
|
||||
from soup_cli.utils.fsdp import FSDP_CONFIGS, get_fsdp_training_args
|
||||
|
||||
if fsdp not in FSDP_CONFIGS:
|
||||
console.print(
|
||||
f"[red]Invalid FSDP preset: {fsdp}[/]\n"
|
||||
f"Options: {', '.join(FSDP_CONFIGS.keys())}"
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
fsdp_kwargs = get_fsdp_training_args(fsdp)
|
||||
console.print(f"[green]FSDP2 enabled:[/] {fsdp}")
|
||||
|
||||
# Detect hardware
|
||||
device, device_name = detect_device()
|
||||
gpu_info = get_gpu_info()
|
||||
|
|
@ -185,6 +204,69 @@ def train(
|
|||
if qat_errors:
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Validate FSDP configuration
|
||||
if fsdp:
|
||||
from soup_cli.utils.fsdp import validate_fsdp_config
|
||||
|
||||
fsdp_errors = validate_fsdp_config(
|
||||
fsdp_preset=fsdp,
|
||||
deepspeed_config=ds_config_path,
|
||||
backend=cfg.backend,
|
||||
device=device,
|
||||
)
|
||||
for err in fsdp_errors:
|
||||
console.print(f"[red]FSDP error:[/] {err}")
|
||||
if fsdp_errors:
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Validate Liger Kernel configuration
|
||||
if cfg.training.use_liger:
|
||||
from soup_cli.utils.liger import validate_liger_config
|
||||
|
||||
liger_errors = validate_liger_config(
|
||||
cfg.training.use_liger, cfg.backend, device,
|
||||
)
|
||||
for err in liger_errors:
|
||||
console.print(f"[red]Liger error:[/] {err}")
|
||||
if liger_errors:
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Validate FlashAttention configuration
|
||||
if cfg.training.use_flash_attn:
|
||||
from soup_cli.utils.flash_attn import validate_flash_attn_config
|
||||
|
||||
fa_errors = validate_flash_attn_config(
|
||||
cfg.training.use_flash_attn, cfg.backend, device,
|
||||
)
|
||||
for err in fa_errors:
|
||||
console.print(f"[red]FlashAttention error:[/] {err}")
|
||||
if fa_errors:
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Validate Ring FlashAttention configuration
|
||||
if cfg.training.use_ring_attention:
|
||||
from soup_cli.utils.ring_attention import validate_ring_attention_config
|
||||
|
||||
ring_errors = validate_ring_attention_config(
|
||||
cfg.training.use_ring_attention, device, cfg.data.max_length,
|
||||
)
|
||||
for err in ring_errors:
|
||||
console.print(f"[red]Ring Attention error:[/] {err}")
|
||||
if ring_errors:
|
||||
raise typer.Exit(1)
|
||||
|
||||
# Validate long-context configuration
|
||||
if cfg.training.rope_scaling_type:
|
||||
from soup_cli.utils.long_context import validate_long_context_config
|
||||
|
||||
ctx_errors = validate_long_context_config(
|
||||
cfg.data.max_length,
|
||||
cfg.training.rope_scaling_type,
|
||||
cfg.training.gradient_checkpointing,
|
||||
)
|
||||
for err in ctx_errors:
|
||||
console.print(f"[yellow]Long-context warning:[/] {err}")
|
||||
|
||||
# Suggest unsloth if available but not being used
|
||||
if cfg.backend == "transformers":
|
||||
from soup_cli.utils.unsloth import is_unsloth_available
|
||||
|
|
@ -236,64 +318,50 @@ def train(
|
|||
else:
|
||||
report_to = "none"
|
||||
console.print("[dim]Setting up model + trainer...[/]")
|
||||
trainer_kwargs = {
|
||||
"device": device,
|
||||
"report_to": report_to,
|
||||
"deepspeed_config": ds_config_path,
|
||||
"fsdp_config": fsdp_kwargs,
|
||||
}
|
||||
if cfg.task == "dpo":
|
||||
from soup_cli.trainer.dpo import DPOTrainerWrapper
|
||||
|
||||
trainer_wrapper = DPOTrainerWrapper(
|
||||
cfg, device=device, report_to=report_to, deepspeed_config=ds_config_path,
|
||||
)
|
||||
trainer_wrapper = DPOTrainerWrapper(cfg, **trainer_kwargs)
|
||||
elif cfg.task == "grpo":
|
||||
from soup_cli.trainer.grpo import GRPOTrainerWrapper
|
||||
|
||||
trainer_wrapper = GRPOTrainerWrapper(
|
||||
cfg, device=device, report_to=report_to, deepspeed_config=ds_config_path,
|
||||
)
|
||||
trainer_wrapper = GRPOTrainerWrapper(cfg, **trainer_kwargs)
|
||||
elif cfg.task == "ppo":
|
||||
from soup_cli.trainer.ppo import PPOTrainerWrapper
|
||||
|
||||
trainer_wrapper = PPOTrainerWrapper(
|
||||
cfg, device=device, report_to=report_to, deepspeed_config=ds_config_path,
|
||||
)
|
||||
trainer_wrapper = PPOTrainerWrapper(cfg, **trainer_kwargs)
|
||||
elif cfg.task == "kto":
|
||||
from soup_cli.trainer.kto import KTOTrainerWrapper
|
||||
|
||||
trainer_wrapper = KTOTrainerWrapper(
|
||||
cfg, device=device, report_to=report_to, deepspeed_config=ds_config_path,
|
||||
)
|
||||
trainer_wrapper = KTOTrainerWrapper(cfg, **trainer_kwargs)
|
||||
elif cfg.task == "orpo":
|
||||
from soup_cli.trainer.orpo import ORPOTrainerWrapper
|
||||
|
||||
trainer_wrapper = ORPOTrainerWrapper(
|
||||
cfg, device=device, report_to=report_to, deepspeed_config=ds_config_path,
|
||||
)
|
||||
trainer_wrapper = ORPOTrainerWrapper(cfg, **trainer_kwargs)
|
||||
elif cfg.task == "simpo":
|
||||
from soup_cli.trainer.simpo import SimPOTrainerWrapper
|
||||
|
||||
trainer_wrapper = SimPOTrainerWrapper(
|
||||
cfg, device=device, report_to=report_to, deepspeed_config=ds_config_path,
|
||||
)
|
||||
trainer_wrapper = SimPOTrainerWrapper(cfg, **trainer_kwargs)
|
||||
elif cfg.task == "ipo":
|
||||
from soup_cli.trainer.ipo import IPOTrainerWrapper
|
||||
|
||||
trainer_wrapper = IPOTrainerWrapper(
|
||||
cfg, device=device, report_to=report_to, deepspeed_config=ds_config_path,
|
||||
)
|
||||
trainer_wrapper = IPOTrainerWrapper(cfg, **trainer_kwargs)
|
||||
elif cfg.task == "reward_model":
|
||||
from soup_cli.trainer.reward_model import RewardModelTrainerWrapper
|
||||
|
||||
trainer_wrapper = RewardModelTrainerWrapper(
|
||||
cfg, device=device, report_to=report_to, deepspeed_config=ds_config_path,
|
||||
)
|
||||
trainer_wrapper = RewardModelTrainerWrapper(cfg, **trainer_kwargs)
|
||||
elif cfg.task == "pretrain":
|
||||
from soup_cli.trainer.pretrain import PretrainTrainerWrapper
|
||||
|
||||
trainer_wrapper = PretrainTrainerWrapper(
|
||||
cfg, device=device, report_to=report_to, deepspeed_config=ds_config_path,
|
||||
)
|
||||
trainer_wrapper = PretrainTrainerWrapper(cfg, **trainer_kwargs)
|
||||
else:
|
||||
trainer_wrapper = SFTTrainerWrapper(
|
||||
cfg, device=device, report_to=report_to, deepspeed_config=ds_config_path,
|
||||
)
|
||||
trainer_wrapper = SFTTrainerWrapper(cfg, **trainer_kwargs)
|
||||
trainer_wrapper.setup(dataset)
|
||||
|
||||
# Train with live display and experiment tracking
|
||||
|
|
|
|||
|
|
@ -30,7 +30,10 @@ class DataConfig(BaseModel):
|
|||
description="Data format",
|
||||
)
|
||||
val_split: float = Field(default=0.1, ge=0.0, le=0.5, description="Validation split ratio")
|
||||
max_length: int = Field(default=2048, description="Max sequence length in tokens")
|
||||
max_length: int = Field(
|
||||
default=2048, ge=64, le=1048576,
|
||||
description="Max sequence length in tokens",
|
||||
)
|
||||
image_dir: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Base directory for resolving relative image paths in vision datasets",
|
||||
|
|
@ -139,6 +142,30 @@ class TrainingConfig(BaseModel):
|
|||
ge=0,
|
||||
description="Auxiliary load-balancing loss coefficient for MoE models",
|
||||
)
|
||||
# Performance — Liger Kernel (fused operations)
|
||||
use_liger: bool = Field(
|
||||
default=False,
|
||||
description="Enable Liger Kernel fused operations (20-60% memory savings, 20-40% speedup)",
|
||||
)
|
||||
# Performance — FlashAttention
|
||||
use_flash_attn: bool = Field(
|
||||
default=False,
|
||||
description="Enable FlashAttention (auto-detects v2/v3/v4 for faster attention)",
|
||||
)
|
||||
# Performance — Ring FlashAttention (sequence parallelism)
|
||||
use_ring_attention: bool = Field(
|
||||
default=False,
|
||||
description="Enable Ring FlashAttention for sequence parallelism across GPUs",
|
||||
)
|
||||
# Long-context — RoPE scaling
|
||||
rope_scaling_type: Optional[Literal["linear", "dynamic", "yarn", "longrope"]] = Field(
|
||||
default=None,
|
||||
description="RoPE scaling method for long-context: linear, dynamic, yarn, longrope",
|
||||
)
|
||||
gradient_checkpointing: bool = Field(
|
||||
default=False,
|
||||
description="Enable gradient checkpointing for memory savings on long sequences",
|
||||
)
|
||||
|
||||
|
||||
class SoupConfig(BaseModel):
|
||||
|
|
@ -494,6 +521,40 @@ training:
|
|||
moe_aux_loss_coeff: 0.01
|
||||
|
||||
output: ./output
|
||||
""",
|
||||
"longcontext": """# Soup template: Long-Context Fine-tuning (128k+)
|
||||
# Extend model context window for long-document understanding
|
||||
#
|
||||
# Uses RoPE scaling + gradient checkpointing + FlashAttention for 128k tokens.
|
||||
# Optionally enable Liger Kernel for additional memory savings.
|
||||
|
||||
base: meta-llama/Llama-3.1-8B-Instruct
|
||||
task: sft
|
||||
# backend: unsloth # 2-5x faster, pip install 'soup-cli[fast]'
|
||||
|
||||
data:
|
||||
train: ./data/long_context_train.jsonl
|
||||
format: alpaca
|
||||
val_split: 0.05
|
||||
max_length: 131072
|
||||
|
||||
training:
|
||||
epochs: 1
|
||||
lr: 5e-6
|
||||
batch_size: 1
|
||||
gradient_accumulation_steps: 16
|
||||
lora:
|
||||
r: 64
|
||||
alpha: 16
|
||||
target_modules: auto
|
||||
quantization: 4bit
|
||||
gradient_checkpointing: true
|
||||
rope_scaling_type: dynamic
|
||||
use_flash_attn: true
|
||||
# use_liger: true # pip install 'soup-cli[liger]' for fused ops
|
||||
# use_ring_attention: true # Multi-GPU sequence parallelism
|
||||
|
||||
output: ./output_longctx
|
||||
""",
|
||||
"rlhf": """# Soup template: Full RLHF Pipeline (SFT + Reward Model + PPO)
|
||||
# Three-stage training: 1) SFT warmup, 2) Reward model, 3) PPO alignment
|
||||
|
|
|
|||
|
|
@ -27,11 +27,13 @@ class DPOTrainerWrapper:
|
|||
device: str = "cuda",
|
||||
report_to: str = "none",
|
||||
deepspeed_config: Optional[str] = None,
|
||||
fsdp_config: Optional[dict] = None,
|
||||
):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.report_to = report_to
|
||||
self.deepspeed_config = deepspeed_config
|
||||
self.fsdp_config = fsdp_config
|
||||
self.model = None
|
||||
self.ref_model = None
|
||||
self.tokenizer = None
|
||||
|
|
@ -122,6 +124,7 @@ class DPOTrainerWrapper:
|
|||
report_to=self.report_to,
|
||||
remove_unused_columns=False,
|
||||
deepspeed=self.deepspeed_config,
|
||||
**(self.fsdp_config or {}),
|
||||
beta=tcfg.dpo_beta,
|
||||
max_length=cfg.data.max_length,
|
||||
max_prompt_length=cfg.data.max_length // 2,
|
||||
|
|
|
|||
|
|
@ -29,11 +29,13 @@ class GRPOTrainerWrapper:
|
|||
device: str = "cuda",
|
||||
report_to: str = "none",
|
||||
deepspeed_config: Optional[str] = None,
|
||||
fsdp_config: Optional[dict] = None,
|
||||
):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.report_to = report_to
|
||||
self.deepspeed_config = deepspeed_config
|
||||
self.fsdp_config = fsdp_config
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.trainer = None
|
||||
|
|
@ -154,6 +156,7 @@ class GRPOTrainerWrapper:
|
|||
"report_to": self.report_to,
|
||||
"remove_unused_columns": False,
|
||||
"deepspeed": self.deepspeed_config,
|
||||
**(self.fsdp_config or {}),
|
||||
"beta": tcfg.grpo_beta,
|
||||
"num_generations": tcfg.num_generations,
|
||||
"max_completion_length": cfg.data.max_length,
|
||||
|
|
|
|||
|
|
@ -32,11 +32,13 @@ class IPOTrainerWrapper:
|
|||
device: str = "cuda",
|
||||
report_to: str = "none",
|
||||
deepspeed_config: Optional[str] = None,
|
||||
fsdp_config: Optional[dict] = None,
|
||||
):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.report_to = report_to
|
||||
self.deepspeed_config = deepspeed_config
|
||||
self.fsdp_config = fsdp_config
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.trainer = None
|
||||
|
|
@ -124,6 +126,7 @@ class IPOTrainerWrapper:
|
|||
report_to=self.report_to,
|
||||
remove_unused_columns=False,
|
||||
deepspeed=self.deepspeed_config,
|
||||
**(self.fsdp_config or {}),
|
||||
loss_type="ipo",
|
||||
beta=tcfg.ipo_tau,
|
||||
max_length=cfg.data.max_length,
|
||||
|
|
|
|||
|
|
@ -28,11 +28,13 @@ class KTOTrainerWrapper:
|
|||
device: str = "cuda",
|
||||
report_to: str = "none",
|
||||
deepspeed_config: Optional[str] = None,
|
||||
fsdp_config: Optional[dict] = None,
|
||||
):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.report_to = report_to
|
||||
self.deepspeed_config = deepspeed_config
|
||||
self.fsdp_config = fsdp_config
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.trainer = None
|
||||
|
|
@ -121,6 +123,7 @@ class KTOTrainerWrapper:
|
|||
report_to=self.report_to,
|
||||
remove_unused_columns=False,
|
||||
deepspeed=self.deepspeed_config,
|
||||
**(self.fsdp_config or {}),
|
||||
beta=tcfg.kto_beta,
|
||||
max_length=cfg.data.max_length,
|
||||
max_prompt_length=cfg.data.max_length // 2,
|
||||
|
|
|
|||
|
|
@ -31,11 +31,13 @@ class ORPOTrainerWrapper:
|
|||
device: str = "cuda",
|
||||
report_to: str = "none",
|
||||
deepspeed_config: Optional[str] = None,
|
||||
fsdp_config: Optional[dict] = None,
|
||||
):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.report_to = report_to
|
||||
self.deepspeed_config = deepspeed_config
|
||||
self.fsdp_config = fsdp_config
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.trainer = None
|
||||
|
|
@ -122,6 +124,7 @@ class ORPOTrainerWrapper:
|
|||
report_to=self.report_to,
|
||||
remove_unused_columns=False,
|
||||
deepspeed=self.deepspeed_config,
|
||||
**(self.fsdp_config or {}),
|
||||
beta=tcfg.orpo_beta,
|
||||
max_length=cfg.data.max_length,
|
||||
max_prompt_length=cfg.data.max_length // 2,
|
||||
|
|
|
|||
|
|
@ -39,11 +39,13 @@ class PPOTrainerWrapper:
|
|||
device: str = "cuda",
|
||||
report_to: str = "none",
|
||||
deepspeed_config: Optional[str] = None,
|
||||
fsdp_config: Optional[dict] = None,
|
||||
):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.report_to = report_to
|
||||
self.deepspeed_config = deepspeed_config
|
||||
self.fsdp_config = fsdp_config
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.trainer = None
|
||||
|
|
@ -134,6 +136,10 @@ class PPOTrainerWrapper:
|
|||
"learning_rate": tcfg.lr,
|
||||
}
|
||||
|
||||
# FSDP2 — alternative to DeepSpeed
|
||||
if self.fsdp_config:
|
||||
ppo_kwargs.update(self.fsdp_config)
|
||||
|
||||
ppo_params = inspect.signature(ppo_config_cls).parameters
|
||||
|
||||
# trl renamed ppo_epochs -> num_ppo_epochs in newer versions
|
||||
|
|
|
|||
|
|
@ -27,11 +27,13 @@ class PretrainTrainerWrapper:
|
|||
device: str = "cuda",
|
||||
report_to: str = "none",
|
||||
deepspeed_config: Optional[str] = None,
|
||||
fsdp_config: Optional[dict] = None,
|
||||
):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.report_to = report_to
|
||||
self.deepspeed_config = deepspeed_config
|
||||
self.fsdp_config = fsdp_config
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.trainer = None
|
||||
|
|
@ -123,6 +125,10 @@ class PretrainTrainerWrapper:
|
|||
"deepspeed": self.deepspeed_config,
|
||||
}
|
||||
|
||||
# FSDP2 — alternative to DeepSpeed
|
||||
if self.fsdp_config:
|
||||
training_kwargs.update(self.fsdp_config)
|
||||
|
||||
# LoRA+ — different learning rates for A and B matrices
|
||||
if tcfg.loraplus_lr_ratio is not None:
|
||||
training_kwargs["loraplus_lr_ratio"] = tcfg.loraplus_lr_ratio
|
||||
|
|
|
|||
|
|
@ -35,11 +35,13 @@ class RewardModelTrainerWrapper:
|
|||
device: str = "cuda",
|
||||
report_to: str = "none",
|
||||
deepspeed_config: Optional[str] = None,
|
||||
fsdp_config: Optional[dict] = None,
|
||||
):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.report_to = report_to
|
||||
self.deepspeed_config = deepspeed_config
|
||||
self.fsdp_config = fsdp_config
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.trainer = None
|
||||
|
|
@ -127,6 +129,7 @@ class RewardModelTrainerWrapper:
|
|||
report_to=self.report_to,
|
||||
remove_unused_columns=False,
|
||||
deepspeed=self.deepspeed_config,
|
||||
**(self.fsdp_config or {}),
|
||||
max_length=cfg.data.max_length,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -21,11 +21,13 @@ class SFTTrainerWrapper:
|
|||
device: str = "cuda",
|
||||
report_to: str = "none",
|
||||
deepspeed_config: Optional[str] = None,
|
||||
fsdp_config: Optional[dict] = None,
|
||||
):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.report_to = report_to
|
||||
self.deepspeed_config = deepspeed_config
|
||||
self.fsdp_config = fsdp_config
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.trainer = None
|
||||
|
|
@ -138,6 +140,19 @@ class SFTTrainerWrapper:
|
|||
"deepspeed": self.deepspeed_config,
|
||||
}
|
||||
|
||||
# FSDP2 — alternative to DeepSpeed
|
||||
if self.fsdp_config:
|
||||
allowed_fsdp_keys = {"fsdp", "fsdp_config"}
|
||||
unexpected = set(self.fsdp_config.keys()) - allowed_fsdp_keys
|
||||
if unexpected:
|
||||
raise ValueError(f"Unexpected FSDP config keys: {unexpected}")
|
||||
training_kwargs.update(self.fsdp_config)
|
||||
|
||||
# Gradient checkpointing — saves memory for long sequences
|
||||
if tcfg.gradient_checkpointing:
|
||||
training_kwargs["gradient_checkpointing"] = True
|
||||
training_kwargs["gradient_checkpointing_kwargs"] = {"use_reentrant": False}
|
||||
|
||||
# LoRA+ — different learning rates for A and B matrices
|
||||
if tcfg.loraplus_lr_ratio is not None:
|
||||
training_kwargs["loraplus_lr_ratio"] = tcfg.loraplus_lr_ratio
|
||||
|
|
@ -182,6 +197,17 @@ class SFTTrainerWrapper:
|
|||
|
||||
from soup_cli.utils.moe import detect_moe_model, get_moe_target_modules
|
||||
|
||||
# Liger Kernel — apply fused ops BEFORE model loading
|
||||
if tcfg.use_liger:
|
||||
from soup_cli.utils.liger import apply_liger_kernel
|
||||
|
||||
if apply_liger_kernel(cfg.base):
|
||||
console.print(
|
||||
"[green]Liger Kernel enabled:[/] fused RMSNorm, SwiGLU, CrossEntropy, RoPE"
|
||||
)
|
||||
else:
|
||||
console.print("[yellow]Liger Kernel: no matching architecture found[/]")
|
||||
|
||||
console.print(f"[dim]Loading tokenizer: {cfg.base}[/]")
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(cfg.base, trust_remote_code=True)
|
||||
if self.tokenizer.pad_token is None:
|
||||
|
|
@ -208,8 +234,33 @@ class SFTTrainerWrapper:
|
|||
if bnb_config:
|
||||
model_kwargs["quantization_config"] = bnb_config
|
||||
|
||||
# FlashAttention — set attn_implementation for faster attention
|
||||
if tcfg.use_flash_attn:
|
||||
from soup_cli.utils.flash_attn import get_attn_implementation
|
||||
|
||||
attn_impl = get_attn_implementation(tcfg.use_flash_attn, self.device)
|
||||
if attn_impl:
|
||||
model_kwargs["attn_implementation"] = attn_impl
|
||||
console.print(f"[green]FlashAttention enabled:[/] {attn_impl}")
|
||||
|
||||
self.model = AutoModelForCausalLM.from_pretrained(cfg.base, **model_kwargs)
|
||||
|
||||
# Long-context — apply RoPE scaling after model load
|
||||
if tcfg.rope_scaling_type:
|
||||
from soup_cli.utils.long_context import apply_long_context_config
|
||||
|
||||
rope_config = apply_long_context_config(
|
||||
self.model.config,
|
||||
target_length=cfg.data.max_length,
|
||||
rope_scaling_type=tcfg.rope_scaling_type,
|
||||
model_name=cfg.base,
|
||||
)
|
||||
if rope_config:
|
||||
console.print(
|
||||
f"[green]Long-context enabled:[/] RoPE {tcfg.rope_scaling_type} "
|
||||
f"scaling to {cfg.data.max_length} tokens"
|
||||
)
|
||||
|
||||
# MoE aux loss for load balancing
|
||||
is_moe = detect_moe_model(self.model)
|
||||
if is_moe and tcfg.moe_aux_loss_coeff > 0:
|
||||
|
|
|
|||
|
|
@ -32,11 +32,13 @@ class SimPOTrainerWrapper:
|
|||
device: str = "cuda",
|
||||
report_to: str = "none",
|
||||
deepspeed_config: Optional[str] = None,
|
||||
fsdp_config: Optional[dict] = None,
|
||||
):
|
||||
self.config = config
|
||||
self.device = device
|
||||
self.report_to = report_to
|
||||
self.deepspeed_config = deepspeed_config
|
||||
self.fsdp_config = fsdp_config
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.trainer = None
|
||||
|
|
@ -123,6 +125,7 @@ class SimPOTrainerWrapper:
|
|||
report_to=self.report_to,
|
||||
remove_unused_columns=False,
|
||||
deepspeed=self.deepspeed_config,
|
||||
**(self.fsdp_config or {}),
|
||||
loss_type="simpo",
|
||||
cpo_alpha=tcfg.cpo_alpha,
|
||||
simpo_gamma=tcfg.simpo_gamma,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,129 @@
|
|||
"""FlashAttention auto-detection and configuration.
|
||||
|
||||
Detects FlashAttention availability (v2/v3/v4) and configures models
|
||||
to use the best available attention implementation automatically.
|
||||
|
||||
FlashAttention provides 2-4x speedup and significant memory savings
|
||||
for long sequences by avoiding materializing the full attention matrix.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
# Ordered by preference (newest first)
|
||||
FLASH_ATTN_VERSIONS = ("flash_attention_3", "flash_attention_2")
|
||||
|
||||
|
||||
def detect_flash_attention() -> str | None:
|
||||
"""Detect the best available FlashAttention implementation.
|
||||
|
||||
Returns:
|
||||
The attention implementation string for model_kwargs, or None if unavailable.
|
||||
One of: "flash_attention_3", "flash_attention_2", None.
|
||||
"""
|
||||
# FlashAttention requires CUDA
|
||||
try:
|
||||
import torch
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
return None
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
# Check FlashAttention 3 (Hopper architecture, H100+)
|
||||
try:
|
||||
import flash_attn # noqa: F401
|
||||
|
||||
version = getattr(flash_attn, "__version__", "0.0.0")
|
||||
major = int(version.split(".")[0])
|
||||
if major >= 3:
|
||||
return "flash_attention_3"
|
||||
except (ImportError, ValueError, IndexError):
|
||||
pass
|
||||
|
||||
# Check FlashAttention 2
|
||||
try:
|
||||
from transformers.utils import is_flash_attn_2_available
|
||||
|
||||
if is_flash_attn_2_available():
|
||||
return "flash_attention_2"
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Direct import check for flash_attn 2.x
|
||||
try:
|
||||
import flash_attn # noqa: F401
|
||||
|
||||
version = getattr(flash_attn, "__version__", "0.0.0")
|
||||
major = int(version.split(".")[0])
|
||||
if major >= 2:
|
||||
return "flash_attention_2"
|
||||
except (ImportError, ValueError, IndexError):
|
||||
pass
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def get_flash_attn_version() -> str | None:
|
||||
"""Return the installed flash-attn package version, or None."""
|
||||
try:
|
||||
import flash_attn
|
||||
|
||||
return getattr(flash_attn, "__version__", "unknown")
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
|
||||
def get_attn_implementation(use_flash_attn: bool, device: str) -> str | None:
|
||||
"""Get the best attention implementation to use.
|
||||
|
||||
Args:
|
||||
use_flash_attn: Whether FlashAttention is requested in config.
|
||||
device: Training device (cuda/cpu/mps).
|
||||
|
||||
Returns:
|
||||
Attention implementation string for from_pretrained(), or None for default.
|
||||
"""
|
||||
if not use_flash_attn:
|
||||
return None
|
||||
|
||||
if device != "cuda":
|
||||
return None
|
||||
|
||||
return detect_flash_attention()
|
||||
|
||||
|
||||
def validate_flash_attn_config(
|
||||
use_flash_attn: bool, backend: str, device: str,
|
||||
) -> list[str]:
|
||||
"""Validate FlashAttention configuration and return error messages.
|
||||
|
||||
Args:
|
||||
use_flash_attn: Whether FlashAttention is requested.
|
||||
backend: Training backend (transformers/unsloth).
|
||||
device: Training device (cuda/cpu/mps).
|
||||
|
||||
Returns:
|
||||
List of error messages. Empty list means valid.
|
||||
"""
|
||||
errors: list[str] = []
|
||||
|
||||
if not use_flash_attn:
|
||||
return errors
|
||||
|
||||
# Unsloth handles FlashAttention internally — no validation needed
|
||||
if backend == "unsloth":
|
||||
return errors
|
||||
|
||||
if device != "cuda":
|
||||
errors.append(
|
||||
"FlashAttention requires CUDA. "
|
||||
f"Current device: {device}."
|
||||
)
|
||||
|
||||
if device == "cuda" and detect_flash_attention() is None:
|
||||
errors.append(
|
||||
"FlashAttention is not available. "
|
||||
"Install it with: pip install flash-attn --no-build-isolation"
|
||||
)
|
||||
|
||||
return errors
|
||||
|
|
@ -0,0 +1,174 @@
|
|||
"""FSDP2 (Fully Sharded Data Parallel) configuration templates.
|
||||
|
||||
FSDP2 is PyTorch's native distributed training solution, an alternative to DeepSpeed.
|
||||
It shards model parameters, gradients, and optimizer states across GPUs with
|
||||
tighter integration into PyTorch's autograd engine.
|
||||
|
||||
FSDP2 advantages over DeepSpeed:
|
||||
- Native PyTorch (no external dependency)
|
||||
- Better composability with torch.compile
|
||||
- Simpler configuration for most use cases
|
||||
- Built-in mixed precision via torch.amp
|
||||
|
||||
Requires: torch >= 2.2.0, accelerate >= 0.27.0
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
|
||||
# FSDP2 Full Shard: shards params + gradients + optimizer states (like ZeRO-3)
|
||||
FSDP_FULL_SHARD = {
|
||||
"fsdp": "full_shard auto_wrap",
|
||||
"fsdp_config": {
|
||||
"backward_prefetch": "backward_pre",
|
||||
"forward_prefetch": True,
|
||||
"use_orig_params": True,
|
||||
"limit_all_gathers": True,
|
||||
"sync_module_states": True,
|
||||
},
|
||||
}
|
||||
|
||||
# FSDP2 Shard Grad Op: shards gradients + optimizer states only (like ZeRO-2)
|
||||
FSDP_SHARD_GRAD_OP = {
|
||||
"fsdp": "shard_grad_op auto_wrap",
|
||||
"fsdp_config": {
|
||||
"backward_prefetch": "backward_pre",
|
||||
"forward_prefetch": True,
|
||||
"use_orig_params": True,
|
||||
"limit_all_gathers": True,
|
||||
"sync_module_states": True,
|
||||
},
|
||||
}
|
||||
|
||||
# FSDP2 Full Shard with CPU offload (memory-constrained setups)
|
||||
FSDP_FULL_SHARD_OFFLOAD = {
|
||||
"fsdp": "full_shard auto_wrap offload",
|
||||
"fsdp_config": {
|
||||
"backward_prefetch": "backward_pre",
|
||||
"forward_prefetch": True,
|
||||
"use_orig_params": True,
|
||||
"limit_all_gathers": True,
|
||||
"sync_module_states": True,
|
||||
},
|
||||
}
|
||||
|
||||
FSDP_CONFIGS = {
|
||||
"fsdp_full_shard": FSDP_FULL_SHARD,
|
||||
"fsdp_shard_grad": FSDP_SHARD_GRAD_OP,
|
||||
"fsdp_full_offload": FSDP_FULL_SHARD_OFFLOAD,
|
||||
}
|
||||
|
||||
|
||||
def get_fsdp_config(preset: str) -> dict:
|
||||
"""Get FSDP config dict by preset name.
|
||||
|
||||
Args:
|
||||
preset: One of 'fsdp_full_shard', 'fsdp_shard_grad', 'fsdp_full_offload'.
|
||||
|
||||
Returns:
|
||||
Deep copy of the FSDP config dict.
|
||||
|
||||
Raises:
|
||||
ValueError: If preset is not recognized.
|
||||
"""
|
||||
if preset not in FSDP_CONFIGS:
|
||||
raise ValueError(
|
||||
f"Unknown FSDP config: {preset}. "
|
||||
f"Options: {', '.join(FSDP_CONFIGS.keys())}"
|
||||
)
|
||||
return copy.deepcopy(FSDP_CONFIGS[preset])
|
||||
|
||||
|
||||
def get_fsdp_training_args(preset: str) -> dict:
|
||||
"""Get FSDP kwargs to pass to TrainingArguments.
|
||||
|
||||
Args:
|
||||
preset: FSDP preset name.
|
||||
|
||||
Returns:
|
||||
Dict of kwargs to unpack into TrainingArguments.
|
||||
"""
|
||||
config = get_fsdp_config(preset)
|
||||
return {
|
||||
"fsdp": config["fsdp"],
|
||||
"fsdp_config": config["fsdp_config"],
|
||||
}
|
||||
|
||||
|
||||
def is_fsdp_available() -> bool:
|
||||
"""Check if FSDP2 requirements are met (torch >= 2.2, accelerate >= 0.27)."""
|
||||
try:
|
||||
import torch
|
||||
|
||||
parts = torch.__version__.split(".")[:2]
|
||||
torch_version = tuple(
|
||||
int(p.split("+")[0].split("a")[0].split("b")[0].split("rc")[0])
|
||||
for p in parts
|
||||
)
|
||||
if torch_version < (2, 2):
|
||||
return False
|
||||
except (ImportError, ValueError):
|
||||
return False
|
||||
|
||||
try:
|
||||
import accelerate # noqa: F401
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
def validate_fsdp_config(
|
||||
fsdp_preset: str | None,
|
||||
deepspeed_config: str | None,
|
||||
backend: str,
|
||||
device: str,
|
||||
) -> list[str]:
|
||||
"""Validate FSDP configuration and return error messages.
|
||||
|
||||
Args:
|
||||
fsdp_preset: FSDP preset name, or None if not using FSDP.
|
||||
deepspeed_config: DeepSpeed config path, or None.
|
||||
backend: Training backend (transformers/unsloth).
|
||||
device: Training device (cuda/cpu/mps).
|
||||
|
||||
Returns:
|
||||
List of error messages. Empty list means valid.
|
||||
"""
|
||||
errors: list[str] = []
|
||||
|
||||
if not fsdp_preset:
|
||||
return errors
|
||||
|
||||
if deepspeed_config:
|
||||
errors.append(
|
||||
"Cannot use FSDP and DeepSpeed together. Choose one: "
|
||||
"--fsdp or --deepspeed."
|
||||
)
|
||||
|
||||
if device != "cuda":
|
||||
errors.append(
|
||||
"FSDP requires CUDA GPUs. "
|
||||
f"Current device: {device}."
|
||||
)
|
||||
|
||||
if backend == "unsloth":
|
||||
errors.append(
|
||||
"FSDP is not compatible with the unsloth backend. "
|
||||
"Use backend: transformers."
|
||||
)
|
||||
|
||||
if not is_fsdp_available():
|
||||
errors.append(
|
||||
"FSDP2 requires torch >= 2.2.0 and accelerate >= 0.27.0. "
|
||||
"Upgrade with: pip install -U torch accelerate"
|
||||
)
|
||||
|
||||
if fsdp_preset not in FSDP_CONFIGS:
|
||||
errors.append(
|
||||
f"Unknown FSDP preset: {fsdp_preset}. "
|
||||
f"Options: {', '.join(FSDP_CONFIGS.keys())}"
|
||||
)
|
||||
|
||||
return errors
|
||||
|
|
@ -0,0 +1,133 @@
|
|||
"""Liger Kernel — fused operations for faster, memory-efficient training.
|
||||
|
||||
Liger Kernel provides fused CUDA kernels (RMSNorm, SwiGLU, CrossEntropy, RoPE, etc.)
|
||||
that replace standard HuggingFace operations with optimized fused versions.
|
||||
This can yield 20-60% memory savings and 20-40% throughput improvement.
|
||||
|
||||
Requires: liger-kernel >= 0.3.0
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
def is_liger_available() -> bool:
|
||||
"""Check if liger-kernel is installed."""
|
||||
try:
|
||||
import liger_kernel # noqa: F401
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
def get_liger_version() -> str | None:
|
||||
"""Return liger-kernel version string, or None if not installed."""
|
||||
try:
|
||||
import liger_kernel
|
||||
|
||||
return getattr(liger_kernel, "__version__", "unknown")
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
|
||||
def apply_liger_kernel(model_name: str) -> bool:
|
||||
"""Apply Liger Kernel fused operations for the given model architecture.
|
||||
|
||||
Patches the model class in-place so that all subsequent model instantiations
|
||||
use fused kernels (RMSNorm, SwiGLU, CrossEntropy, RoPE, FusedLinearCrossEntropy).
|
||||
|
||||
This must be called BEFORE loading the model.
|
||||
|
||||
Args:
|
||||
model_name: HuggingFace model name/path (used to detect architecture).
|
||||
|
||||
Returns:
|
||||
True if Liger Kernel was applied, False otherwise.
|
||||
"""
|
||||
if not is_liger_available():
|
||||
return False
|
||||
|
||||
model_lower = model_name.lower()
|
||||
|
||||
try:
|
||||
from liger_kernel.transformers import (
|
||||
AutoLigerKernelForCausalLM,
|
||||
)
|
||||
|
||||
# AutoLigerKernelForCausalLM handles architecture detection automatically
|
||||
AutoLigerKernelForCausalLM._apply_liger_kernel(model_name)
|
||||
return True
|
||||
except (ImportError, AttributeError, NotImplementedError):
|
||||
# Fallback: try manual patching for known architectures
|
||||
return _apply_liger_manual(model_lower)
|
||||
|
||||
|
||||
def _apply_liger_manual(model_lower: str) -> bool:
|
||||
"""Manually apply Liger Kernel patches for known model architectures."""
|
||||
try:
|
||||
if "llama" in model_lower or "codellama" in model_lower:
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_llama
|
||||
|
||||
apply_liger_kernel_to_llama()
|
||||
return True
|
||||
elif "mistral" in model_lower or "mixtral" in model_lower:
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_mistral
|
||||
|
||||
apply_liger_kernel_to_mistral()
|
||||
return True
|
||||
elif "gemma" in model_lower:
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_gemma2
|
||||
|
||||
apply_liger_kernel_to_gemma2()
|
||||
return True
|
||||
elif "qwen" in model_lower:
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_qwen2
|
||||
|
||||
apply_liger_kernel_to_qwen2()
|
||||
return True
|
||||
elif "phi" in model_lower:
|
||||
from liger_kernel.transformers import apply_liger_kernel_to_phi3
|
||||
|
||||
apply_liger_kernel_to_phi3()
|
||||
return True
|
||||
except (ImportError, AttributeError):
|
||||
pass
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def validate_liger_config(use_liger: bool, backend: str, device: str) -> list[str]:
|
||||
"""Validate Liger Kernel configuration and return error messages.
|
||||
|
||||
Args:
|
||||
use_liger: Whether Liger Kernel is requested.
|
||||
backend: Training backend (transformers/unsloth).
|
||||
device: Training device (cuda/cpu/mps).
|
||||
|
||||
Returns:
|
||||
List of error messages. Empty list means valid.
|
||||
"""
|
||||
errors: list[str] = []
|
||||
|
||||
if not use_liger:
|
||||
return errors
|
||||
|
||||
if not is_liger_available():
|
||||
errors.append(
|
||||
"liger-kernel is not installed. "
|
||||
"Install it with: pip install 'soup-cli[liger]'"
|
||||
)
|
||||
|
||||
if backend == "unsloth":
|
||||
errors.append(
|
||||
"Liger Kernel is not compatible with the unsloth backend. "
|
||||
"Unsloth has its own fused kernels. Use backend: transformers."
|
||||
)
|
||||
|
||||
if device != "cuda":
|
||||
errors.append(
|
||||
"Liger Kernel requires CUDA. "
|
||||
f"Current device: {device}. Use a GPU for Liger Kernel."
|
||||
)
|
||||
|
||||
return errors
|
||||
|
|
@ -0,0 +1,182 @@
|
|||
"""Long-context fine-tuning utilities — 128k+ token support.
|
||||
|
||||
Configures RoPE (Rotary Position Embedding) scaling to extend model context
|
||||
windows beyond their pre-training length. Supports multiple scaling strategies:
|
||||
|
||||
- linear: Simple linear interpolation (PI) — good baseline
|
||||
- dynamic: NTK-aware Dynamic scaling — better for large extensions
|
||||
- yarn: YaRN (Yet another RoPE extensioN) — best quality for 4-8x extension
|
||||
- longrope: LongRoPE — progressive extension with search-based factors
|
||||
|
||||
Also handles gradient checkpointing configuration for memory efficiency
|
||||
when training on very long sequences.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
# Supported RoPE scaling methods
|
||||
ROPE_SCALING_TYPES = ("linear", "dynamic", "yarn", "longrope")
|
||||
|
||||
# Default context lengths for known model families
|
||||
MODEL_DEFAULT_CONTEXT: dict[str, int] = {
|
||||
"llama-3": 8192,
|
||||
"llama-2": 4096,
|
||||
"mistral": 32768,
|
||||
"mixtral": 32768,
|
||||
"qwen2": 32768,
|
||||
"qwen3": 32768,
|
||||
"phi-3": 4096,
|
||||
"phi-4": 16384,
|
||||
"gemma": 8192,
|
||||
"gemma-2": 8192,
|
||||
"deepseek": 4096,
|
||||
"codellama": 16384,
|
||||
}
|
||||
|
||||
|
||||
def get_model_default_context(model_name: str) -> int:
|
||||
"""Estimate the default context length for a model based on its name.
|
||||
|
||||
Args:
|
||||
model_name: HuggingFace model name/path.
|
||||
|
||||
Returns:
|
||||
Estimated default context length in tokens.
|
||||
"""
|
||||
model_lower = model_name.lower()
|
||||
for family, ctx_len in MODEL_DEFAULT_CONTEXT.items():
|
||||
if family in model_lower:
|
||||
return ctx_len
|
||||
# Conservative default for unknown models
|
||||
return 4096
|
||||
|
||||
|
||||
def get_rope_scaling_config(
|
||||
scaling_type: str,
|
||||
target_length: int,
|
||||
original_length: int,
|
||||
) -> dict:
|
||||
"""Build RoPE scaling configuration for extending context.
|
||||
|
||||
Args:
|
||||
scaling_type: One of 'linear', 'dynamic', 'yarn', 'longrope'.
|
||||
target_length: Desired context length (e.g., 131072 for 128k).
|
||||
original_length: Model's pre-trained context length.
|
||||
|
||||
Returns:
|
||||
Dict to pass as `rope_scaling` in model config.
|
||||
|
||||
Raises:
|
||||
ValueError: If scaling_type is not supported.
|
||||
"""
|
||||
if scaling_type not in ROPE_SCALING_TYPES:
|
||||
raise ValueError(
|
||||
f"Unknown RoPE scaling type: {scaling_type}. "
|
||||
f"Options: {', '.join(ROPE_SCALING_TYPES)}"
|
||||
)
|
||||
|
||||
factor = target_length / original_length
|
||||
if factor <= 1.0:
|
||||
# No scaling needed — target is within original context
|
||||
return {}
|
||||
|
||||
if scaling_type == "linear":
|
||||
return {
|
||||
"type": "linear",
|
||||
"factor": float(factor),
|
||||
}
|
||||
elif scaling_type == "dynamic":
|
||||
return {
|
||||
"type": "dynamic",
|
||||
"factor": float(factor),
|
||||
}
|
||||
elif scaling_type == "yarn":
|
||||
return {
|
||||
"type": "yarn",
|
||||
"factor": float(factor),
|
||||
"original_max_position_embeddings": original_length,
|
||||
}
|
||||
else: # longrope — guaranteed by Literal constraint in schema
|
||||
return {
|
||||
"type": "longrope",
|
||||
"factor": float(factor),
|
||||
"original_max_position_embeddings": original_length,
|
||||
}
|
||||
|
||||
|
||||
def apply_long_context_config(
|
||||
model_config,
|
||||
target_length: int,
|
||||
rope_scaling_type: str = "dynamic",
|
||||
model_name: str = "",
|
||||
) -> dict | None:
|
||||
"""Apply long-context configuration to a model config object.
|
||||
|
||||
Modifies the model config to extend the context window using RoPE scaling
|
||||
and returns the scaling config that was applied.
|
||||
|
||||
Args:
|
||||
model_config: The model's config object (from model.config).
|
||||
target_length: Desired max sequence length.
|
||||
rope_scaling_type: RoPE scaling strategy.
|
||||
model_name: Model name (for estimating original context length).
|
||||
|
||||
Returns:
|
||||
The rope_scaling dict that was applied, or None if no scaling needed.
|
||||
"""
|
||||
# Determine original context length
|
||||
original_length = getattr(
|
||||
model_config,
|
||||
"max_position_embeddings",
|
||||
get_model_default_context(model_name),
|
||||
)
|
||||
|
||||
if target_length <= original_length:
|
||||
return None
|
||||
|
||||
rope_config = get_rope_scaling_config(
|
||||
scaling_type=rope_scaling_type,
|
||||
target_length=target_length,
|
||||
original_length=original_length,
|
||||
)
|
||||
|
||||
if not rope_config:
|
||||
return None
|
||||
|
||||
# Apply to model config
|
||||
model_config.rope_scaling = rope_config
|
||||
model_config.max_position_embeddings = target_length
|
||||
|
||||
return rope_config
|
||||
|
||||
|
||||
def validate_long_context_config(
|
||||
max_length: int,
|
||||
rope_scaling_type: str | None,
|
||||
use_gradient_checkpointing: bool,
|
||||
) -> list[str]:
|
||||
"""Validate long-context configuration.
|
||||
|
||||
Args:
|
||||
max_length: Target sequence length.
|
||||
rope_scaling_type: RoPE scaling type, or None if not specified.
|
||||
use_gradient_checkpointing: Whether gradient checkpointing is enabled.
|
||||
|
||||
Returns:
|
||||
List of warning/error messages.
|
||||
"""
|
||||
errors: list[str] = []
|
||||
|
||||
if rope_scaling_type and rope_scaling_type not in ROPE_SCALING_TYPES:
|
||||
errors.append(
|
||||
f"Unknown RoPE scaling type: {rope_scaling_type}. "
|
||||
f"Options: {', '.join(ROPE_SCALING_TYPES)}"
|
||||
)
|
||||
|
||||
if max_length >= 65536 and not use_gradient_checkpointing:
|
||||
errors.append(
|
||||
f"Training with max_length={max_length} without gradient checkpointing "
|
||||
"will likely cause OOM. Set gradient_checkpointing: true in config."
|
||||
)
|
||||
|
||||
return errors
|
||||
|
|
@ -0,0 +1,121 @@
|
|||
"""Ring FlashAttention — sequence parallelism for ultra-long contexts.
|
||||
|
||||
Ring FlashAttention distributes a single long sequence across multiple GPUs,
|
||||
splitting the sequence into chunks and using ring communication to compute
|
||||
attention without materializing the full attention matrix on any single GPU.
|
||||
|
||||
This enables training on sequences much longer than a single GPU's memory
|
||||
allows (e.g., 128k-1M+ tokens).
|
||||
|
||||
Requires: ring-flash-attn >= 0.1.0 OR transformers >= 4.43.0 (built-in SP support)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
def is_ring_attention_available() -> bool:
|
||||
"""Check if Ring FlashAttention is available.
|
||||
|
||||
Checks for either:
|
||||
1. The ring-flash-attn package
|
||||
2. Transformers' built-in sequence parallelism support
|
||||
"""
|
||||
try:
|
||||
import ring_flash_attn # noqa: F401
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
def get_ring_attention_version() -> str | None:
|
||||
"""Return ring-flash-attn version, or None if not installed."""
|
||||
try:
|
||||
import ring_flash_attn
|
||||
|
||||
return getattr(ring_flash_attn, "__version__", "unknown")
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
|
||||
def get_sequence_parallel_size(gpu_count: int, max_length: int) -> int:
|
||||
"""Calculate optimal sequence parallel size.
|
||||
|
||||
Distributes sequences across GPUs when sequence length exceeds per-GPU capacity.
|
||||
|
||||
Args:
|
||||
gpu_count: Number of available GPUs.
|
||||
max_length: Target sequence length in tokens.
|
||||
|
||||
Returns:
|
||||
Number of GPUs to use for sequence parallelism (1 = no SP).
|
||||
"""
|
||||
# Sequence parallelism is beneficial for long sequences
|
||||
# Rule of thumb: use SP when sequence > 32k tokens per GPU
|
||||
tokens_per_gpu_threshold = 32768
|
||||
|
||||
if gpu_count <= 1:
|
||||
return 1
|
||||
|
||||
if max_length <= tokens_per_gpu_threshold:
|
||||
return 1
|
||||
|
||||
# Use power-of-2 SP sizes for efficient communication
|
||||
sp_size = 1
|
||||
while sp_size * 2 <= gpu_count and max_length // (sp_size * 2) > 4096:
|
||||
sp_size *= 2
|
||||
|
||||
return sp_size
|
||||
|
||||
|
||||
def validate_ring_attention_config(
|
||||
use_ring_attention: bool,
|
||||
device: str,
|
||||
max_length: int,
|
||||
) -> list[str]:
|
||||
"""Validate Ring FlashAttention configuration.
|
||||
|
||||
Args:
|
||||
use_ring_attention: Whether ring attention is requested.
|
||||
device: Training device (cuda/cpu/mps).
|
||||
max_length: Max sequence length in tokens.
|
||||
|
||||
Returns:
|
||||
List of error messages. Empty list means valid.
|
||||
"""
|
||||
errors: list[str] = []
|
||||
|
||||
if not use_ring_attention:
|
||||
return errors
|
||||
|
||||
if device != "cuda":
|
||||
errors.append(
|
||||
"Ring FlashAttention requires CUDA GPUs. "
|
||||
f"Current device: {device}."
|
||||
)
|
||||
|
||||
if not is_ring_attention_available():
|
||||
errors.append(
|
||||
"Ring FlashAttention is not available. "
|
||||
"Install it with: pip install ring-flash-attn"
|
||||
)
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
gpu_count = torch.cuda.device_count() if torch.cuda.is_available() else 0
|
||||
if gpu_count < 2:
|
||||
errors.append(
|
||||
"Ring FlashAttention requires at least 2 GPUs for sequence parallelism. "
|
||||
f"Found: {gpu_count} GPU(s)."
|
||||
)
|
||||
except ImportError:
|
||||
errors.append("PyTorch is required for Ring FlashAttention.")
|
||||
|
||||
if max_length < 8192:
|
||||
errors.append(
|
||||
f"Ring FlashAttention is designed for long sequences (>= 8192 tokens). "
|
||||
f"Current max_length: {max_length}. Consider increasing max_length."
|
||||
)
|
||||
|
||||
return errors
|
||||
|
|
@ -0,0 +1,739 @@
|
|||
"""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.is_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.is_liger_available", return_value=True):
|
||||
errors = validate_liger_config(True, "transformers", "cuda")
|
||||
assert errors == []
|
||||
|
||||
|
||||
class TestLigerDetection:
|
||||
"""Test Liger Kernel availability detection."""
|
||||
|
||||
def test_is_liger_available_not_installed(self):
|
||||
from soup_cli.utils.liger import is_liger_available
|
||||
|
||||
with patch.dict("sys.modules", {"liger_kernel": None}):
|
||||
# When import fails, should return False
|
||||
result = is_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.is_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.is_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_detect_flash_attention_no_cuda(self):
|
||||
with patch("soup_cli.utils.flash_attn.detect_flash_attention") 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("fsdp_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("fsdp_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("fsdp_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("fsdp_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("fsdp_full_shard")
|
||||
config2 = get_fsdp_config("fsdp_full_shard")
|
||||
config1["fsdp"] = "modified"
|
||||
assert config2["fsdp"] != "modified"
|
||||
|
||||
def test_fsdp_configs_dict(self):
|
||||
from soup_cli.utils.fsdp import FSDP_CONFIGS
|
||||
|
||||
assert "fsdp_full_shard" in FSDP_CONFIGS
|
||||
assert "fsdp_shard_grad" in FSDP_CONFIGS
|
||||
assert "fsdp_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("fsdp_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("fsdp_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("fsdp_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_is_ring_attention_available_returns_bool(self):
|
||||
from soup_cli.utils.ring_attention import is_ring_attention_available
|
||||
|
||||
result = is_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_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 13 templates (12 + longcontext)."""
|
||||
assert len(TEMPLATES) == 13
|
||||
|
||||
|
||||
# ─── 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
|
||||
Loading…
Reference in New Issue