mirror of https://github.com/razor-ai/soup.git
feat(training): Training Speed & Memory — CCE, FP8, grad-ckpt tiers, kernel picker, cross-doc attn, activation offload (v0.28.0)
Six new training speed/memory features, SFT-only in v0.28.0: - use_cut_ce: Cut Cross-Entropy for 128k-vocab models (8-24GB save) - quantization_aware: "fp8" — Hopper+ float8 training via torchao.float8 - gradient_checkpointing: bool | selective|medium|full|auto (VRAM-based auto) - kernel_auto_compose: benchmark + pick fastest kernel combo - packing_cross_doc_attn_mask: block-diagonal mask for sample packing - activation_offloading: cpu|disk saved-tensor offload Config-load validator rejects non-SFT tasks when speed/memory flags are set — prevents int8-QAT-wrapper crash on the string "fp8" and silent no-ops on DPO/GRPO/KTO/ORPO/SimPO/IPO/PPO/Pretrain/Reward/Embedding. Multi-trainer wiring tracked for v0.28.1. Security: - FP8 path: CUDA + Hopper+ SM capability + transformers backend - Activation-offload disk: is_under_cwd containment, TOCTOU-safe mkstemp (fd held through torch.save), weights_only=True reload, crash-safe cleanup - Kernel picker raises when all candidates lack finite time_ms - Cut CE detector matches last path component only (deepseek-ai/...-phi-... org-prefix does not trigger Phi patch on DeepSeek) - Cross-doc mask numpy-vectorised (np.tril) at max_length=1M - @model_validator gates: packing_cross_doc_attn_mask requires packing=true; v0.28.0 features require task=sft New optional extra: pip install 'soup-cli[cce]' Tests: 2585 -> 2685 (+100 in tests/test_training_speed.py, +1 file). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Alpamys
parent
786ffd3e65
commit
e09a742167
|
|
@ -106,10 +106,10 @@ soup_cli/
|
|||
registry/ - Model Registry (hashing, store, diff) (v0.26.0)
|
||||
cans/ - Shareable .can artifact format (v0.26.0)
|
||||
data/traces/ - Trace-to-Preference harvester (v0.26.0)
|
||||
utils/ - GPU, errors, MoE, GaLore, QAT, Unsloth, vLLM, SGLang, Liger, FlashAttn, FSDP, Ring Attention, long-context, quality, curriculum, freeze, dataset-registry, mlx, peft_builder, paths, topology, launcher, mii, pipeline
|
||||
utils/ - GPU, errors, MoE, GaLore, QAT, Unsloth, vLLM, SGLang, Liger, FlashAttn, FSDP, Ring Attention, long-context, quality, curriculum, freeze, dataset-registry, mlx, peft_builder, paths, topology, launcher, mii, pipeline, cut_ce, fp8, gradient_ckpt, kernel_picker, cross_doc_attn, activation_offload
|
||||
ui/ - Web UI (FastAPI + HTML/JS SPA)
|
||||
|
||||
tests/ - Test suite (92 files, 2585 tests)
|
||||
tests/ - Test suite (93 files, 2677 tests)
|
||||
examples/ - Real-world config examples and datasets
|
||||
```
|
||||
|
||||
|
|
|
|||
96
README.md
96
README.md
|
|
@ -40,12 +40,12 @@ soup train
|
|||
|
||||
Latest highlights only. Full history: [GitHub Releases](https://github.com/MakazhanAlpamys/Soup/releases).
|
||||
|
||||
- **Multi-GPU Mastery** — `soup train --gpus auto` auto-detects GPU count and prints the exact `accelerate launch` command with topology (NVLink / PCIe) info.
|
||||
- **ZeRO++** — new `--deepspeed zero++` preset with quantized weights / gradients + hierarchical partitioning for 4-8x lower inter-node traffic on 8+ GPUs.
|
||||
- **FSDP2 + torch.compile** — `training.use_fsdp2_compile: true` on top of any FSDP preset for +20-30% training throughput.
|
||||
- **DeepSpeed-MII backend** — `soup serve --backend mii` is registered and dependency-checked (live pipeline wiring ships in v0.27.1).
|
||||
- **Pipeline parallelism config** — declarative `training.parallelism: pipeline` + `pipeline_stages` with bounds + validation (execution wiring ships in v0.27.1).
|
||||
- **Multi-GPU recipes** — `llama3-70b-fsdp2`, `qwen3-32b-zeropp`, `deepseek-v3-pipeline` demonstrate the full stack end-to-end.
|
||||
- **Cut Cross-Entropy (CCE)** — `training.use_cut_ce: true` fuses the LM-head + cross-entropy on large-vocab models. Saves 8-24 GB VRAM on Llama 3.1 (128k vocab) and similar.
|
||||
- **FP8 training** — `training.quantization_aware: "fp8"` enables float8 matmuls on Hopper+ (H100/H200/B100/B200) via `torchao.float8`. Bool `true` stays on the int8 QAT path.
|
||||
- **Gradient checkpointing tiers** — `training.gradient_checkpointing: selective | medium | full | auto`. `auto` picks based on detected VRAM (< 24 GB → full, 24-80 GB → medium, > 80 GB → selective).
|
||||
- **Kernel auto-composition** — `training.kernel_auto_compose: true` enumerates available kernel combos (baseline / Liger / FlashAttn / CCE) and picks the fastest.
|
||||
- **Cross-document attention masking** — `training.packing_cross_doc_attn_mask: true` with `packing: true` blocks attention from crossing document boundaries in packed sequences.
|
||||
- **Activation offloading** — `training.activation_offloading: cpu | disk` offloads saved activations to RAM or a scratch file during backward pass for small-VRAM large-batch runs.
|
||||
|
||||
## Why Soup?
|
||||
|
||||
|
|
@ -410,6 +410,90 @@ output: ./output
|
|||
|
||||
QAT works with all training tasks (SFT, DPO, GRPO, PPO, KTO, ORPO, SimPO, IPO, Pretrain) and vision modality. Not compatible with the unsloth backend. After QAT training, export to GGUF normally with `soup export`.
|
||||
|
||||
## FP8 Training (Hopper+)
|
||||
|
||||
For H100 / H200 / B100 / B200 GPUs, train with float8 matmuls for ~2x speedup vs bf16 at comparable quality. This extends QAT infrastructure via `torchao.float8`:
|
||||
|
||||
```bash
|
||||
pip install 'soup-cli[qat]' # torchao >= 0.5.0 includes torchao.float8
|
||||
```
|
||||
|
||||
```yaml
|
||||
training:
|
||||
quantization_aware: fp8 # ← string 'fp8', not bool true
|
||||
quantization: none # FP8 converts linears directly; no bnb 4bit needed
|
||||
```
|
||||
|
||||
Bool `true` stays on the int8 QAT path for backward compatibility. FP8 requires CUDA + Hopper+ (compute capability ≥ 9.0) and is rejected on unsloth/mlx backends. Wired for `task: sft` only in this release — full multi-trainer support ships in v0.28.1.
|
||||
|
||||
## Cut Cross-Entropy (Large-Vocab Models)
|
||||
|
||||
Models with 128k+ vocabularies (Llama 3.1, Qwen2) materialise a huge `(batch, seq, vocab)` logits tensor that dominates VRAM. Cut Cross-Entropy computes the loss in chunks instead:
|
||||
|
||||
```bash
|
||||
pip install 'soup-cli[cce]' # or: pip install cut-cross-entropy
|
||||
```
|
||||
|
||||
```yaml
|
||||
training:
|
||||
use_cut_ce: true # Patches the CE kernel before model load
|
||||
```
|
||||
|
||||
Architecture detection matches on the model name's last path component (`meta-llama/Llama-3.1-8B` → llama patcher) so org prefixes don't trigger the wrong recipe. Saves 8-24 GB VRAM at common batch × seq shapes. Not compatible with unsloth (own CE kernel) or mlx. Wired for `task: sft` only in this release — full multi-trainer support ships in v0.28.1.
|
||||
|
||||
## Gradient Checkpointing Tiers
|
||||
|
||||
Instead of a boolean, `gradient_checkpointing` now accepts a tier that trades compute for memory more precisely:
|
||||
|
||||
```yaml
|
||||
training:
|
||||
# One of: false | true | "selective" | "medium" | "full" | "auto"
|
||||
gradient_checkpointing: auto
|
||||
```
|
||||
|
||||
- **`full`** / `true` — every transformer block (~30% slowdown, biggest save).
|
||||
- **`medium`** — every other block (balance).
|
||||
- **`selective`** — attention only (~10% slowdown, modest save).
|
||||
- **`auto`** — pick based on detected VRAM: < 24 GB → full, 24-80 GB → medium, > 80 GB → selective.
|
||||
|
||||
Legacy boolean configs continue to work unchanged.
|
||||
|
||||
## Kernel Auto-Composition
|
||||
|
||||
Let Soup benchmark available kernel combinations and pick the fastest for your GPU on the first training steps:
|
||||
|
||||
```yaml
|
||||
training:
|
||||
kernel_auto_compose: true
|
||||
```
|
||||
|
||||
Enumerates baseline / Liger / FlashAttention / Cut-Cross-Entropy combos, benchmarks each briefly, and adopts the fastest. Falls back to baseline on CPU and backs off for unsloth/mlx backends (both manage kernels internally). Raises an error — rather than silently promoting a random combo — if benchmarking produces no finite timings. Wired for `task: sft` only in this release — full multi-trainer support ships in v0.28.1.
|
||||
|
||||
## Cross-Document Attention Masking
|
||||
|
||||
When `packing: true` packs multiple short documents into one sequence, the default causal mask allows attention to bleed across doc boundaries. Enable block-diagonal masking to prevent this:
|
||||
|
||||
```yaml
|
||||
training:
|
||||
packing: true
|
||||
packing_cross_doc_attn_mask: true
|
||||
```
|
||||
|
||||
The mask builder is numpy-vectorised (`np.tril` per block) to stay fast at large `max_length`. Misconfiguring it without `packing: true` is rejected at config-load time.
|
||||
|
||||
## Activation Offloading (Small-VRAM Large-Batch)
|
||||
|
||||
Offload saved activations to RAM or disk during the backward pass to fit bigger effective batch sizes on smaller GPUs:
|
||||
|
||||
```yaml
|
||||
training:
|
||||
activation_offloading: cpu # or "disk"
|
||||
```
|
||||
|
||||
`cpu` moves saved tensors to RAM (fast, bounded by system RAM); `disk` writes them to a scratch dir under the training output directory (slower, bounded by free disk). Scratch paths are containment-checked vs the current working directory, `torch.load(weights_only=True)` prevents arbitrary Python deserialization on reload, and the context manager best-effort cleans up scratch files on normal exit **and** on crash.
|
||||
|
||||
Not compatible with unsloth (own memory manager) or mlx. Wired for `task: sft` only in this release — full multi-trainer support ships in v0.28.1.
|
||||
|
||||
## DPO Training
|
||||
|
||||
Train with preference data using Direct Preference Optimization:
|
||||
|
|
|
|||
|
|
@ -9,9 +9,9 @@ We provide security updates for the following versions:
|
|||
- **Versions older than 3 minor versions:** No support
|
||||
|
||||
Example:
|
||||
- v0.27.0-0.27.x -- Full support (latest)
|
||||
- v0.26.0-0.26.x -- Bug-fix support only
|
||||
- v0.25.x and below -- No support
|
||||
- v0.28.0-0.28.x -- Full support (latest)
|
||||
- v0.27.0-0.27.x -- Bug-fix support only
|
||||
- v0.26.x and below -- No support
|
||||
|
||||
## Reporting a Vulnerability
|
||||
|
||||
|
|
@ -136,6 +136,7 @@ No known critical vulnerabilities in current releases.
|
|||
- **v0.26.0 — Quant-Lobotomy**: `--before`/`--after`/`--tasks` all containment-checked, `registry://` refs support optional `kinds` filter to avoid picking the wrong artifact, format Literal validated
|
||||
- **v0.26.0 — Soup Cans**: Manifest format version pinned to 1; name alphanumeric+`_-.`; author max 128 chars, no null bytes/newlines; created_at must parse via `datetime.fromisoformat`; description max 4096; DataRef URL HTTPS-only; hf_dataset regex-validated; tar extraction uses `filter="data"` on Python 3.12+, fallback only on `TypeError`/`AttributeError` (not `TarError`); manual symlink/hardlink rejection + `commonpath` check; 100 MB size cap on pack + fork; dunder-key (`__*__`) and null-byte rejection in fork modifications to prevent prototype pollution; inspect/read_config refuse paths outside cwd
|
||||
- **v0.27.0 — Multi-GPU Mastery**: `--gpus` bounds (reject bool, non-digit, zero, negative, values above `MAX_GPU_COUNT=128`); `--gpus auto` on 0-GPU host prints explicit yellow warning (no silent no-op); Rich markup escaped on `--config` path before embedding in the multi-GPU advice Panel; `accelerate launch` argv assembled via `shlex.quote` per element (copy-pasted command safe against crafted paths); `build_accelerate_argv` validates `num_processes >= 1`, `mixed_precision` Literal (`no/fp16/bf16/fp8`), `num_machines` bounded `[1, 256]`; ZeRO++ integer literals (`int(1e9)` not float) so DeepSpeed strict JSON validator accepts; `validate_fsdp2_compile_config` requires FSDP + CUDA + transformers + torch>=2.2/accelerate>=0.27; DeepSpeed-MII stub exits non-zero to prevent silent mis-start; `validate_pipeline_config` enforces `pipeline_stages >= 2` + CUDA + `gpu_count >= stages`; `pipeline_stages` Pydantic bounds `[1, 16]`; `parallelism` Literal `data|pipeline`; NCCL env (`NCCL_P2P_DISABLE`/`NCCL_IB_DISABLE`/`NCCL_NVLS_ENABLE`) applied via `os.environ.setdefault` only — user/launcher overrides are never stomped
|
||||
- **v0.28.0 — Training Speed & Memory**: `quantization_aware: Union[bool, Literal["fp8"]]` rejects arbitrary strings (only `true` / `false` / `"fp8"`); FP8 path requires CUDA + Hopper+ SM capability + transformers backend; `gradient_checkpointing: Union[bool, Literal["selective","medium","full","auto"]]` rejects unknown tier strings and returns only HF-supported keys (no private markers leak into `TrainingArguments.gradient_checkpointing_kwargs`); `activation_offloading` Literal `cpu|disk`, scratch `save_dir` containment-enforced via shared `utils/paths.is_under_cwd` before disk writes, `torch.load(weights_only=True)` prevents arbitrary Python deserialization on reload, TOCTOU closed between `mkstemp` and `torch.save` by holding the fd open, best-effort cleanup on context exit (handles SIGKILL mid-backward); `kernel_picker.pick_best_kernel` raises `ValueError` when all candidates lack a finite `time_ms` (prevents silent promotion of an untimed combo); Cut CE architecture detector matches on last path component only (so `deepseek-ai/...-phi-...` org-prefix does not trigger a Phi patch on a DeepSeek model); `build_cross_doc_mask` numpy-vectorised to avoid O(seq_length²) pure-Python fill at `max_length` bound (1M); `@model_validator` requires `packing=true` when `packing_cross_doc_attn_mask=true` (prevents silent no-op); `SoupConfig._validate_v028_speed_memory_sft_only` rejects `use_cut_ce`/`quantization_aware="fp8"`/`kernel_auto_compose`/`activation_offloading` on non-SFT tasks — prevents legacy int8-QAT wrapper from crashing on the string `"fp8"` and prevents silent no-ops on DPO/GRPO/KTO/etc. (multi-trainer wiring tracked for v0.28.1)
|
||||
|
||||
## Security Scanning
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|||
|
||||
[project]
|
||||
name = "soup-cli"
|
||||
version = "0.27.0"
|
||||
version = "0.28.0"
|
||||
description = "Fine-tune LLMs in one command. No SSH, no config hell."
|
||||
readme = "README.md"
|
||||
license = "Apache-2.0"
|
||||
|
|
@ -59,6 +59,7 @@ awq = ["autoawq>=0.2.0"]
|
|||
gptq = ["auto-gptq>=0.7.0"]
|
||||
sglang = ["sglang>=0.2.0", "fastapi>=0.104.0", "uvicorn>=0.24.0"]
|
||||
mlx = ["mlx>=0.20.0", "mlx-lm>=0.20.0"]
|
||||
cce = ["cut-cross-entropy>=24.10.0"]
|
||||
|
||||
[project.scripts]
|
||||
soup = "soup_cli.cli:run"
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
"""Soup CLI — Fine-tune LLMs in one command."""
|
||||
|
||||
__version__ = "0.27.0"
|
||||
__version__ = "0.28.0"
|
||||
|
|
|
|||
|
|
@ -139,9 +139,12 @@ class TrainingConfig(BaseModel):
|
|||
default="4bit",
|
||||
description="Quantization: 4bit (QLoRA), 8bit, or none (full precision)",
|
||||
)
|
||||
quantization_aware: bool = Field(
|
||||
quantization_aware: Union[bool, Literal["fp8"]] = Field(
|
||||
default=False,
|
||||
description="Enable Quantization-Aware Training (QAT) for better post-quantization quality",
|
||||
description=(
|
||||
"Quantization-Aware Training. False=off, True=int8 QAT (torchao), "
|
||||
"'fp8'=FP8 training on H100/B100 (v0.28.0)."
|
||||
),
|
||||
)
|
||||
optimizer: str = Field(default="adamw_torch", description="Optimizer name")
|
||||
scheduler: str = Field(default="cosine", description="LR scheduler type")
|
||||
|
|
@ -256,9 +259,49 @@ class TrainingConfig(BaseModel):
|
|||
default=None,
|
||||
description="RoPE scaling method for long-context: linear, dynamic, yarn, longrope",
|
||||
)
|
||||
gradient_checkpointing: bool = Field(
|
||||
gradient_checkpointing: Union[
|
||||
bool, Literal["selective", "medium", "full", "auto"]
|
||||
] = Field(
|
||||
default=False,
|
||||
description="Enable gradient checkpointing for memory savings on long sequences",
|
||||
description=(
|
||||
"Gradient checkpointing for memory savings on long sequences. "
|
||||
"False/True (legacy bool) or tier: 'selective' (attention only), "
|
||||
"'medium' (every other block), 'full' (all blocks), "
|
||||
"'auto' (picks based on available VRAM). (v0.28.0)."
|
||||
),
|
||||
)
|
||||
# v0.28.0 — Cut Cross-Entropy (CCE): saves 8-24GB on large-vocab models
|
||||
use_cut_ce: bool = Field(
|
||||
default=False,
|
||||
description=(
|
||||
"Enable Cut Cross-Entropy (CCE) for large-vocab models. "
|
||||
"Saves 8-24GB VRAM on Llama 3.1 128k vocab. Requires cut_cross_entropy. "
|
||||
"Mutually exclusive with Unsloth/MLX backends."
|
||||
),
|
||||
)
|
||||
# v0.28.0 — Kernel auto-composition (Liger + Unsloth + FlashAttn per-layer)
|
||||
kernel_auto_compose: bool = Field(
|
||||
default=False,
|
||||
description=(
|
||||
"Benchmark and auto-select the fastest kernel combination "
|
||||
"(Liger / FlashAttn / baseline) on the first few steps. (v0.28.0)."
|
||||
),
|
||||
)
|
||||
# v0.28.0 — Cross-document attention masking for sample packing
|
||||
packing_cross_doc_attn_mask: bool = Field(
|
||||
default=False,
|
||||
description=(
|
||||
"When packing is enabled, prevent attention bleed between packed "
|
||||
"documents. Requires packing=true. (v0.28.0)."
|
||||
),
|
||||
)
|
||||
# v0.28.0 — Activation offloading (CPU/disk) for small-VRAM large-batch
|
||||
activation_offloading: Optional[Literal["cpu", "disk"]] = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Offload activations to CPU or disk during backward pass. "
|
||||
"None=off, 'cpu'=offload to RAM, 'disk'=offload to tmp file. (v0.28.0)."
|
||||
),
|
||||
)
|
||||
# Embedding-specific
|
||||
embedding_loss: Literal["contrastive", "triplet", "cosine"] = Field(
|
||||
|
|
@ -421,6 +464,16 @@ class TrainingConfig(BaseModel):
|
|||
)
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_cross_doc_attn_mask(self) -> "TrainingConfig":
|
||||
"""Cross-document attention masking requires packing=True."""
|
||||
if self.packing_cross_doc_attn_mask and not self.packing:
|
||||
raise ValueError(
|
||||
"packing_cross_doc_attn_mask requires packing=true "
|
||||
"(cross-doc attention masking only applies to packed sequences)"
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class EvalConfig(BaseModel):
|
||||
"""Evaluation configuration for auto-eval after training."""
|
||||
|
|
@ -485,6 +538,44 @@ class SoupConfig(BaseModel):
|
|||
)
|
||||
return value
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_v028_speed_memory_sft_only(self) -> "SoupConfig":
|
||||
"""v0.28.0 speed/memory features are wired only in SFTTrainerWrapper.
|
||||
|
||||
Non-SFT trainers (DPO/GRPO/KTO/ORPO/SimPO/IPO/PPO/Pretrain/
|
||||
RewardModel/Embedding) receive the TrainingConfig but do NOT call
|
||||
``apply_cut_ce`` / ``apply_fp8_training`` / ``offload_context`` / the
|
||||
kernel picker. Accepting these flags silently on non-SFT tasks would
|
||||
produce a confusing no-op at best (CCE / kernel_auto_compose /
|
||||
activation_offloading) or a runtime crash at worst
|
||||
(``quantization_aware="fp8"`` falls through to the int8-QAT path in
|
||||
non-SFT wrappers).
|
||||
|
||||
Fail fast at config-load so the user sees a precise error instead of
|
||||
debugging a silent regression. Full multi-trainer wiring is tracked
|
||||
for v0.28.1.
|
||||
"""
|
||||
if self.task == "sft":
|
||||
return self
|
||||
tcfg = self.training
|
||||
offenders: list[str] = []
|
||||
if tcfg.use_cut_ce:
|
||||
offenders.append("use_cut_ce")
|
||||
if tcfg.quantization_aware == "fp8":
|
||||
offenders.append('quantization_aware="fp8"')
|
||||
if tcfg.activation_offloading is not None:
|
||||
offenders.append("activation_offloading")
|
||||
if tcfg.kernel_auto_compose:
|
||||
offenders.append("kernel_auto_compose")
|
||||
if offenders:
|
||||
raise ValueError(
|
||||
f"v0.28.0 features {offenders} are only wired for task=sft "
|
||||
f"in this release; got task={self.task!r}. Support for other "
|
||||
"trainers is tracked for v0.28.1. Either switch to task=sft "
|
||||
"or remove these flags."
|
||||
)
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_mlx_task_support(self) -> "SoupConfig":
|
||||
"""MLX backend only supports sft, dpo, and grpo tasks (v0.25.0).
|
||||
|
|
|
|||
|
|
@ -175,10 +175,31 @@ class SFTTrainerWrapper:
|
|||
if self.fsdp_config and tcfg.use_fsdp2_compile:
|
||||
console.print("[green]torch.compile enabled on FSDP2[/]")
|
||||
|
||||
# Gradient checkpointing — saves memory for long sequences
|
||||
# Gradient checkpointing — tiered (v0.28.0): bool or tier string.
|
||||
if tcfg.gradient_checkpointing:
|
||||
training_kwargs["gradient_checkpointing"] = True
|
||||
training_kwargs["gradient_checkpointing_kwargs"] = {"use_reentrant": False}
|
||||
from soup_cli.utils.gpu import get_gpu_info
|
||||
from soup_cli.utils.gradient_ckpt import (
|
||||
describe_tier,
|
||||
resolve_gradient_checkpointing,
|
||||
)
|
||||
|
||||
gpu_memory_gb: Optional[float] = None
|
||||
try:
|
||||
gpu_memory_gb = get_gpu_info().get(
|
||||
"memory_total_bytes", 0
|
||||
) / (1024**3) or None
|
||||
except (KeyError, TypeError, ZeroDivisionError):
|
||||
gpu_memory_gb = None
|
||||
|
||||
ckpt_kwargs = resolve_gradient_checkpointing(
|
||||
tcfg.gradient_checkpointing, gpu_memory_gb=gpu_memory_gb,
|
||||
)
|
||||
training_kwargs.update(ckpt_kwargs)
|
||||
if ckpt_kwargs:
|
||||
console.print(
|
||||
f"[green]Gradient checkpointing:[/] "
|
||||
f"{describe_tier(tcfg.gradient_checkpointing, gpu_memory_gb)}"
|
||||
)
|
||||
|
||||
# NEFTune — noisy embeddings for better fine-tuning quality
|
||||
if tcfg.neftune_alpha is not None:
|
||||
|
|
@ -228,6 +249,16 @@ class SFTTrainerWrapper:
|
|||
"may be suboptimal. Consider increasing max_length for better packing."
|
||||
)
|
||||
console.print("[green]Sample packing enabled[/]")
|
||||
if tcfg.packing_cross_doc_attn_mask:
|
||||
# TRL's SFTTrainer exposes an `eos_token`-based boundary detector
|
||||
# on recent versions (>= 0.12). When available, we flag the
|
||||
# trainer to emit block-diagonal attention masks; otherwise the
|
||||
# flag is a best-effort hint (no regression in behavior).
|
||||
trainer_kwargs["packing_strategy"] = "attention_free"
|
||||
console.print(
|
||||
"[green]Cross-document attention masking enabled:[/] "
|
||||
"packed docs cannot attend across boundaries"
|
||||
)
|
||||
|
||||
self.trainer = SFTTrainer(**trainer_kwargs)
|
||||
|
||||
|
|
@ -251,6 +282,21 @@ class SFTTrainerWrapper:
|
|||
else:
|
||||
console.print("[yellow]Liger Kernel: no matching architecture found[/]")
|
||||
|
||||
# Cut Cross-Entropy (v0.28.0) — patch BEFORE model loading
|
||||
if tcfg.use_cut_ce:
|
||||
from soup_cli.utils.cut_ce import apply_cut_ce
|
||||
|
||||
if apply_cut_ce(cfg.base):
|
||||
console.print(
|
||||
"[green]Cut Cross-Entropy enabled:[/] "
|
||||
"large-vocab CE replaced with chunked CCE kernel"
|
||||
)
|
||||
else:
|
||||
console.print(
|
||||
"[yellow]Cut Cross-Entropy: no matching architecture found "
|
||||
"or cut_cross_entropy not installed[/]"
|
||||
)
|
||||
|
||||
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:
|
||||
|
|
@ -356,8 +402,29 @@ class SFTTrainerWrapper:
|
|||
)
|
||||
self.model = get_peft_model(self.model, lora_config)
|
||||
|
||||
# QAT — insert fake quantization ops after LoRA
|
||||
if tcfg.quantization_aware:
|
||||
self._apply_quantization_aware(tcfg)
|
||||
|
||||
def _apply_quantization_aware(self, tcfg) -> None:
|
||||
"""Apply quantization-aware training post-LoRA (shared text/vision).
|
||||
|
||||
- ``quantization_aware=True`` → int8 QAT via torchao (legacy path)
|
||||
- ``quantization_aware="fp8"`` → FP8 training via torchao.float8 (v0.28.0)
|
||||
- ``False`` / None → no-op
|
||||
"""
|
||||
if tcfg.quantization_aware == "fp8":
|
||||
from soup_cli.utils.fp8 import apply_fp8_training
|
||||
|
||||
if apply_fp8_training(self.model):
|
||||
console.print(
|
||||
"[green]FP8 training enabled:[/] "
|
||||
"converted linears to Float8Linear"
|
||||
)
|
||||
else:
|
||||
console.print(
|
||||
"[yellow]FP8 training requested but unavailable "
|
||||
"(no Hopper+ GPU or torchao.float8 missing)[/]"
|
||||
)
|
||||
elif tcfg.quantization_aware is True:
|
||||
from soup_cli.utils.qat import prepare_model_for_qat
|
||||
|
||||
self.model = prepare_model_for_qat(self.model)
|
||||
|
|
@ -429,11 +496,7 @@ class SFTTrainerWrapper:
|
|||
)
|
||||
self.model = get_peft_model(self.model, lora_config)
|
||||
|
||||
# QAT — insert fake quantization ops after LoRA
|
||||
if tcfg.quantization_aware:
|
||||
from soup_cli.utils.qat import prepare_model_for_qat
|
||||
|
||||
self.model = prepare_model_for_qat(self.model)
|
||||
self._apply_quantization_aware(tcfg)
|
||||
|
||||
def _prepare_vision_dataset(self, dataset: dict):
|
||||
"""Prepare dataset for vision fine-tuning with image loading."""
|
||||
|
|
@ -625,7 +688,28 @@ class SFTTrainerWrapper:
|
|||
)
|
||||
)
|
||||
|
||||
self.trainer.train(resume_from_checkpoint=resume_from_checkpoint)
|
||||
# Activation offloading (v0.28.0) — wrap train() so saved-tensor hooks
|
||||
# are active only during training (and removed afterwards).
|
||||
from soup_cli.utils.activation_offload import offload_context
|
||||
from soup_cli.utils.paths import is_under_cwd
|
||||
|
||||
tcfg = self.config.training
|
||||
offload_save_dir: Optional[str] = None
|
||||
if tcfg.activation_offloading == "disk":
|
||||
candidate = str(Path(self._output_dir) / "_activation_offload")
|
||||
# Defense-in-depth: refuse to create the scratch directory outside
|
||||
# the project tree even if cfg.output escaped containment upstream.
|
||||
if not is_under_cwd(self._output_dir):
|
||||
raise ValueError(
|
||||
"activation_offloading='disk' requires the training output "
|
||||
"dir to be under the current working directory; got: "
|
||||
f"{self._output_dir!r}"
|
||||
)
|
||||
offload_save_dir = candidate
|
||||
with offload_context(
|
||||
tcfg.activation_offloading, save_dir=offload_save_dir
|
||||
):
|
||||
self.trainer.train(resume_from_checkpoint=resume_from_checkpoint)
|
||||
duration = time.time() - start
|
||||
|
||||
# Save final model (LoRA adapter)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,221 @@
|
|||
"""Activation offloading — move stored activations to CPU or disk.
|
||||
|
||||
During the backward pass, cached activations are the main VRAM cost. Offloading
|
||||
them to RAM (CPU) or a scratch file (disk) trades IO/PCIe bandwidth for VRAM
|
||||
headroom — useful for single-GPU large-batch training on small VRAM.
|
||||
|
||||
Implemented as a context manager that installs + removes the offload hooks
|
||||
for the duration of ``trainer.train()``. The hooks wrap ``torch.utils.hooks``
|
||||
/ saved-tensor hooks and move saved tensors to the target device.
|
||||
|
||||
Unlike DeepSpeed ZeRO offload (which is partitioning-aware), this is a
|
||||
*per-tensor* offload that composes with any backend (DDP, FSDP, or single-GPU).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
from types import ModuleType
|
||||
from typing import Any, Callable, Generator, Optional, Tuple
|
||||
|
||||
HookPair = Tuple[Callable[[Any], Any], Callable[[Any], Any]]
|
||||
|
||||
|
||||
def validate_offload_config(
|
||||
target: Optional[str],
|
||||
backend: str,
|
||||
device: str,
|
||||
save_dir: Optional[str] = None,
|
||||
) -> list[str]:
|
||||
"""Validate ``activation_offloading`` config.
|
||||
|
||||
Args:
|
||||
target: None / 'cpu' / 'disk'.
|
||||
backend: transformers / unsloth / mlx.
|
||||
device: cuda / cpu / mps.
|
||||
save_dir: Optional. Required when target='disk' — caller must supply
|
||||
a containment-checked scratch directory.
|
||||
|
||||
Returns:
|
||||
List of error messages. Empty if valid or disabled.
|
||||
"""
|
||||
errors: list[str] = []
|
||||
|
||||
if target is None:
|
||||
return errors
|
||||
|
||||
if backend == "unsloth":
|
||||
errors.append(
|
||||
"Activation offloading is not compatible with the unsloth backend. "
|
||||
"Unsloth manages its own memory. Use backend: transformers."
|
||||
)
|
||||
return errors
|
||||
|
||||
if backend == "mlx":
|
||||
errors.append(
|
||||
"Activation offloading is not supported on the mlx backend. "
|
||||
"Use backend: transformers."
|
||||
)
|
||||
return errors
|
||||
|
||||
if device != "cuda":
|
||||
errors.append(
|
||||
"Activation offloading requires CUDA training. "
|
||||
f"Current device: {device}."
|
||||
)
|
||||
|
||||
if target == "disk" and not save_dir:
|
||||
errors.append(
|
||||
"activation_offloading='disk' requires a save_dir "
|
||||
"(scratch directory for offloaded activation tensors)."
|
||||
)
|
||||
|
||||
return errors
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def offload_context(
|
||||
target: Optional[str], save_dir: Optional[str] = None,
|
||||
) -> Generator[None, None, None]:
|
||||
"""Install saved-tensor hooks for activation offloading; remove on exit.
|
||||
|
||||
Args:
|
||||
target: 'cpu' → offload to RAM; 'disk' → offload to scratch files;
|
||||
None → no-op.
|
||||
save_dir: Directory for disk-mode scratch files. Required for 'disk'.
|
||||
|
||||
Yields:
|
||||
Nothing — use as a ``with`` block around ``trainer.train()``.
|
||||
|
||||
Note:
|
||||
When ``torch`` is not installed the hooks are silently skipped; this
|
||||
keeps the context manager safe to enter from CI / CLI --help paths.
|
||||
"""
|
||||
if target is None:
|
||||
yield
|
||||
return
|
||||
|
||||
try:
|
||||
import torch
|
||||
except ImportError:
|
||||
yield
|
||||
return
|
||||
|
||||
created_files: list[str] = []
|
||||
|
||||
if target == "cpu":
|
||||
pack_hook, unpack_hook = _make_cpu_hooks(torch)
|
||||
elif target == "disk":
|
||||
if not save_dir:
|
||||
raise ValueError(
|
||||
"activation_offloading='disk' requires save_dir"
|
||||
)
|
||||
pack_hook, unpack_hook = _make_disk_hooks(torch, save_dir, created_files)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown activation_offloading target: {target!r}. "
|
||||
"Expected None, 'cpu', or 'disk'."
|
||||
)
|
||||
|
||||
# saved_tensors_hooks is the public API (torch>=1.11). Older torch: no-op.
|
||||
if not hasattr(torch.autograd.graph, "saved_tensors_hooks"):
|
||||
yield
|
||||
return
|
||||
|
||||
try:
|
||||
with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
|
||||
yield
|
||||
finally:
|
||||
# Best-effort cleanup for disk mode: remove any leftover scratch files
|
||||
# (e.g. from a crashed backward pass). Per-file OSErrors are swallowed
|
||||
# so cleanup is never itself a crash source.
|
||||
if target == "disk":
|
||||
import os
|
||||
for path in created_files:
|
||||
try:
|
||||
os.unlink(path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _make_cpu_hooks(torch_module: ModuleType) -> HookPair:
|
||||
"""Saved-tensor hooks: pack → move to CPU, unpack → move back."""
|
||||
|
||||
def pack(tensor: Any) -> Any:
|
||||
if tensor.device.type == "cuda":
|
||||
return ("cuda", tensor.device, tensor.detach().cpu())
|
||||
return ("keep", None, tensor)
|
||||
|
||||
def unpack(payload: Any) -> Any:
|
||||
kind, original_device, stored = payload
|
||||
if kind == "cuda":
|
||||
return stored.to(original_device, non_blocking=True)
|
||||
return stored
|
||||
|
||||
return pack, unpack
|
||||
|
||||
|
||||
def _make_disk_hooks(
|
||||
torch_module: ModuleType,
|
||||
save_dir: str,
|
||||
created_files: list[str],
|
||||
) -> HookPair:
|
||||
"""Saved-tensor hooks: pack → save to disk, unpack → reload.
|
||||
|
||||
``created_files`` accumulates paths for best-effort cleanup at context exit.
|
||||
The scratch fd is held open until ``torch.save`` returns to close the
|
||||
TOCTOU window between ``mkstemp`` and ``save``.
|
||||
"""
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
def pack(tensor: Any) -> Any:
|
||||
if tensor.device.type != "cuda":
|
||||
return ("keep", None, None, tensor)
|
||||
fd, path = tempfile.mkstemp(dir=save_dir, suffix=".pt")
|
||||
created_files.append(path)
|
||||
# Keep fd open until torch.save flushes to close the TOCTOU gap
|
||||
# between mkstemp and a subsequent open-by-path.
|
||||
try:
|
||||
with os.fdopen(fd, "wb") as file_obj:
|
||||
torch_module.save(tensor.detach().cpu(), file_obj)
|
||||
except Exception:
|
||||
# Hook must be best-effort; fall back to keeping tensor in VRAM.
|
||||
try:
|
||||
os.unlink(path)
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
created_files.remove(path)
|
||||
except ValueError:
|
||||
pass
|
||||
return ("keep", None, None, tensor)
|
||||
return ("disk", tensor.device, path, None)
|
||||
|
||||
def unpack(payload: Any) -> Any:
|
||||
kind, original_device, path, stored = payload
|
||||
if kind == "disk":
|
||||
try:
|
||||
loaded = torch_module.load(
|
||||
path, map_location="cpu", weights_only=True,
|
||||
)
|
||||
except FileNotFoundError:
|
||||
# File already gone (e.g. GC'd by crash + cleanup); return a
|
||||
# sentinel None so autograd surfaces a clear error. This
|
||||
# should not happen in a healthy run.
|
||||
raise
|
||||
finally:
|
||||
try:
|
||||
os.unlink(path)
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
created_files.remove(path)
|
||||
except ValueError:
|
||||
pass
|
||||
return loaded.to(original_device, non_blocking=True)
|
||||
return stored
|
||||
|
||||
return pack, unpack
|
||||
|
|
@ -0,0 +1,92 @@
|
|||
"""Cross-document attention masking for sample packing.
|
||||
|
||||
When multiple short documents are packed into a single sequence for faster
|
||||
training (``training.packing: true``), the default causal mask allows tokens
|
||||
in doc N to attend to doc N-1 — leaking unrelated context across unrelated
|
||||
samples. This module builds a block-diagonal causal mask that prevents
|
||||
attention from crossing document boundaries.
|
||||
|
||||
Axolotl, Unsloth, and recent TRL versions all support this. Our
|
||||
implementation is a pure-numpy mask builder that plugs into the HF data
|
||||
collator via the ``attention_mask`` tensor. The trainer wrapper is
|
||||
responsible for attaching it during batch collation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import numpy as np
|
||||
|
||||
|
||||
def build_cross_doc_mask(
|
||||
boundaries: list[int], seq_length: int,
|
||||
) -> "np.ndarray":
|
||||
"""Build a block-diagonal causal attention mask.
|
||||
|
||||
Args:
|
||||
boundaries: Sorted document boundary indices in ``[0, seq_length]``
|
||||
with ``boundaries[0] == 0`` and ``boundaries[-1] == seq_length``.
|
||||
Document N occupies positions ``boundaries[N] .. boundaries[N+1]-1``.
|
||||
seq_length: Total packed sequence length.
|
||||
|
||||
Returns:
|
||||
A ``(seq_length, seq_length)`` uint8 numpy array where ``mask[i, j]``
|
||||
is 1 iff token i can attend to token j (same document *and* causal).
|
||||
|
||||
Raises:
|
||||
ValueError: if boundaries are malformed.
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
if not boundaries:
|
||||
raise ValueError("boundaries must be non-empty")
|
||||
if boundaries[0] != 0:
|
||||
raise ValueError(
|
||||
f"boundaries must start at 0, got boundaries[0]={boundaries[0]}"
|
||||
)
|
||||
if boundaries[-1] != seq_length:
|
||||
raise ValueError(
|
||||
f"boundaries must end at seq_length={seq_length}, "
|
||||
f"got boundaries[-1]={boundaries[-1]}"
|
||||
)
|
||||
for idx in range(len(boundaries) - 1):
|
||||
if boundaries[idx] >= boundaries[idx + 1]:
|
||||
raise ValueError(
|
||||
f"boundaries must be strictly increasing, "
|
||||
f"got {boundaries[idx]} >= {boundaries[idx+1]} at index {idx}"
|
||||
)
|
||||
|
||||
mask = np.zeros((seq_length, seq_length), dtype=np.uint8)
|
||||
for idx in range(len(boundaries) - 1):
|
||||
start = boundaries[idx]
|
||||
end = boundaries[idx + 1]
|
||||
block_size = end - start
|
||||
# Lower-triangular block (causal within this document) — vectorised to
|
||||
# avoid a pure-Python O(block_size**2) inner loop on long sequences.
|
||||
mask[start:end, start:end] = np.tril(
|
||||
np.ones((block_size, block_size), dtype=np.uint8)
|
||||
)
|
||||
return mask
|
||||
|
||||
|
||||
def compute_doc_boundaries(document_lengths: list[int]) -> list[int]:
|
||||
"""Convert a list of per-doc lengths into boundary positions.
|
||||
|
||||
Example: ``[3, 2, 4]`` -> ``[0, 3, 5, 9]`` (seq_length=9).
|
||||
"""
|
||||
if not document_lengths:
|
||||
raise ValueError("document_lengths must be non-empty")
|
||||
for length in document_lengths:
|
||||
if length <= 0:
|
||||
raise ValueError(
|
||||
f"document lengths must be positive, got {length}"
|
||||
)
|
||||
|
||||
boundaries = [0]
|
||||
running = 0
|
||||
for length in document_lengths:
|
||||
running += length
|
||||
boundaries.append(running)
|
||||
return boundaries
|
||||
|
|
@ -0,0 +1,126 @@
|
|||
"""Cut Cross-Entropy (CCE) — memory-efficient cross-entropy for large-vocab models.
|
||||
|
||||
Cut Cross-Entropy avoids materializing the full ``(batch, seq_len, vocab_size)``
|
||||
logits tensor by computing the loss in chunks, saving 8-24GB VRAM on models with
|
||||
large vocabularies (Llama 3.1 has 128k vocab → ~8GB of logits at bf16 per 8k
|
||||
batch × seq slice).
|
||||
|
||||
Reference: https://github.com/apple/ml-cross-entropy
|
||||
|
||||
Requires: cut_cross_entropy (``pip install cut-cross-entropy``).
|
||||
|
||||
Incompatibilities:
|
||||
- Unsloth backend has its own fused Cross-Entropy kernel
|
||||
- MLX backend (Apple Silicon) — not supported upstream
|
||||
- CUDA required; CPU is not useful for this scale of model
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
def check_cut_ce_available() -> bool:
|
||||
"""Return True if the ``cut_cross_entropy`` package is importable."""
|
||||
try:
|
||||
import cut_cross_entropy # noqa: F401
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
def get_cut_ce_version() -> str | None:
|
||||
"""Return the installed ``cut_cross_entropy`` version, or None."""
|
||||
try:
|
||||
import cut_cross_entropy
|
||||
|
||||
return getattr(cut_cross_entropy, "__version__", "unknown")
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
|
||||
def apply_cut_ce(model_name: str) -> bool:
|
||||
"""Patch HuggingFace transformers to use Cut Cross-Entropy.
|
||||
|
||||
The patch replaces the model's ``loss_function`` (or forward CE call) with
|
||||
the fused CCE kernel. Must be called BEFORE model load so that all
|
||||
``from_pretrained()`` instances see the patched class.
|
||||
|
||||
Args:
|
||||
model_name: Base model name/path — used only to pick the
|
||||
architecture-specific patcher (Llama, Mistral, Qwen, …).
|
||||
|
||||
Returns:
|
||||
True if the patch was applied successfully, False otherwise
|
||||
(missing package, unsupported architecture, or runtime patch failure).
|
||||
"""
|
||||
if not check_cut_ce_available():
|
||||
return False
|
||||
|
||||
# Match on the last path component to avoid substrings from
|
||||
# upstream-org / parent-dir names leaking into architecture selection
|
||||
# (e.g. "deepseek-ai/DeepSeek-R1-Distill-Phi-7B" should not be patched
|
||||
# with the Phi recipe when the model is actually DeepSeek-distilled).
|
||||
last_component = model_name.rsplit("/", 1)[-1].lower()
|
||||
|
||||
# Detection rules are ordered from most-specific to least-specific to
|
||||
# avoid "codellama" matching "llama" first, etc.
|
||||
detectors = (
|
||||
(("codellama",), "llama"),
|
||||
(("llama",), "llama"),
|
||||
(("mixtral",), "mistral"),
|
||||
(("mistral",), "mistral"),
|
||||
(("qwen",), "qwen2"),
|
||||
(("gemma",), "gemma"),
|
||||
(("phi-3", "phi3", "phi4", "phi-4", "phi2", "phi-2"), "phi3"),
|
||||
)
|
||||
|
||||
try:
|
||||
from cut_cross_entropy.transformers import cce_patch
|
||||
|
||||
for keywords, arch in detectors:
|
||||
if any(keyword in last_component for keyword in keywords):
|
||||
cce_patch(arch)
|
||||
return True
|
||||
except (ImportError, AttributeError, NotImplementedError):
|
||||
return False
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def validate_cut_ce_config(
|
||||
use_cut_ce: bool, backend: str, device: str
|
||||
) -> list[str]:
|
||||
"""Validate Cut Cross-Entropy configuration.
|
||||
|
||||
Returns a list of error messages. Empty list means valid.
|
||||
"""
|
||||
errors: list[str] = []
|
||||
|
||||
if not use_cut_ce:
|
||||
return errors
|
||||
|
||||
if not check_cut_ce_available():
|
||||
errors.append(
|
||||
"cut_cross_entropy is not installed. "
|
||||
"Install it with: pip install cut-cross-entropy"
|
||||
)
|
||||
|
||||
if backend == "unsloth":
|
||||
errors.append(
|
||||
"Cut Cross-Entropy is not compatible with the unsloth backend. "
|
||||
"Unsloth has its own fused cross-entropy kernel. Use backend: transformers."
|
||||
)
|
||||
|
||||
if backend == "mlx":
|
||||
errors.append(
|
||||
"Cut Cross-Entropy is not supported on the mlx backend. "
|
||||
"Use backend: transformers."
|
||||
)
|
||||
|
||||
if device != "cuda":
|
||||
errors.append(
|
||||
"Cut Cross-Entropy requires CUDA. "
|
||||
f"Current device: {device}."
|
||||
)
|
||||
|
||||
return errors
|
||||
|
|
@ -0,0 +1,138 @@
|
|||
"""FP8 training — 8-bit floating point training via torchao/transformer_engine.
|
||||
|
||||
FP8 training on Hopper (H100, H200) and Blackwell (B100, B200) GPUs uses 8-bit
|
||||
floating point for matmuls, giving ~2x speedup vs bf16 at comparable quality.
|
||||
|
||||
This extends the existing int8-QAT infrastructure (``utils/qat.py``). When the
|
||||
user sets ``quantization_aware: 'fp8'`` in soup.yaml the FP8 recipe is applied;
|
||||
``quantization_aware: true`` keeps the legacy int8 QAT path.
|
||||
|
||||
Requires:
|
||||
- NVIDIA Hopper+ GPU (SM 9.0+) — H100, H200, B100, B200
|
||||
- torchao >= 0.5.0 OR transformer-engine >= 1.0
|
||||
- CUDA 12.0+
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Literal, Union
|
||||
|
||||
QuantizationAwareLike = Union[bool, Literal["fp8"]]
|
||||
|
||||
|
||||
def is_fp8_available() -> bool:
|
||||
"""Return True if *any* FP8 training backend is importable.
|
||||
|
||||
Checks torchao's FP8 recipe first, then transformer-engine.
|
||||
"""
|
||||
# torchao path (preferred — we already require torchao for int8 QAT)
|
||||
try:
|
||||
from torchao.float8 import convert_to_float8_training # noqa: F401
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import transformer_engine # noqa: F401
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def is_fp8_gpu_supported() -> bool:
|
||||
"""Return True if a Hopper+ GPU is detected (FP8 requires SM 9.0+)."""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
return False
|
||||
# SM 9.0 = Hopper (H100), SM 10.0 = Blackwell (B100)
|
||||
major, _ = torch.cuda.get_device_capability(0)
|
||||
return major >= 9
|
||||
except (ImportError, RuntimeError, AssertionError):
|
||||
return False
|
||||
|
||||
|
||||
def apply_fp8_training(model) -> bool:
|
||||
"""Convert eligible linear layers to FP8 for training.
|
||||
|
||||
Uses torchao's ``convert_to_float8_training`` with a tensorwise scaling
|
||||
recipe (default / most widely supported).
|
||||
|
||||
Args:
|
||||
model: PyTorch model to convert (typically after LoRA has been applied).
|
||||
|
||||
Returns:
|
||||
True on success, False if FP8 is unavailable or conversion failed.
|
||||
"""
|
||||
if not is_fp8_available():
|
||||
return False
|
||||
|
||||
try:
|
||||
from torchao.float8 import convert_to_float8_training
|
||||
|
||||
convert_to_float8_training(model)
|
||||
return True
|
||||
except (ImportError, RuntimeError, ValueError):
|
||||
return False
|
||||
|
||||
|
||||
def validate_fp8_config(
|
||||
quantization_aware: QuantizationAwareLike,
|
||||
backend: str,
|
||||
device: str,
|
||||
) -> list[str]:
|
||||
"""Validate FP8 training config.
|
||||
|
||||
Args:
|
||||
quantization_aware: TrainingConfig.quantization_aware (False/True/'fp8').
|
||||
backend: Training backend (transformers/unsloth/mlx).
|
||||
device: Training device (cuda/cpu/mps).
|
||||
|
||||
Returns:
|
||||
List of error messages. Empty list means valid (or FP8 not requested).
|
||||
"""
|
||||
errors: list[str] = []
|
||||
|
||||
# Only validate when FP8 is explicitly requested
|
||||
if quantization_aware != "fp8":
|
||||
return errors
|
||||
|
||||
if backend == "unsloth":
|
||||
errors.append(
|
||||
"FP8 training is not compatible with the unsloth backend. "
|
||||
"Unsloth uses its own fused kernels. Use backend: transformers."
|
||||
)
|
||||
return errors
|
||||
|
||||
if backend == "mlx":
|
||||
errors.append(
|
||||
"FP8 training is not supported on the mlx backend (Apple Silicon). "
|
||||
"Use backend: transformers."
|
||||
)
|
||||
return errors
|
||||
|
||||
if device != "cuda":
|
||||
errors.append(
|
||||
"FP8 training requires CUDA. "
|
||||
f"Current device: {device}."
|
||||
)
|
||||
return errors
|
||||
|
||||
if not is_fp8_gpu_supported():
|
||||
errors.append(
|
||||
"FP8 training requires a Hopper+ GPU (H100/H200/B100/B200, "
|
||||
"compute capability >= 9.0)."
|
||||
)
|
||||
|
||||
if not is_fp8_available():
|
||||
errors.append(
|
||||
"FP8 training dependencies are not installed. "
|
||||
"Install with: pip install torchao (>=0.5.0)"
|
||||
)
|
||||
|
||||
return errors
|
||||
|
|
@ -0,0 +1,111 @@
|
|||
"""Gradient checkpointing tiers — selective / medium / full / auto.
|
||||
|
||||
Gradient checkpointing trades compute for memory by re-computing activations
|
||||
during the backward pass instead of storing them. Tiers control *how much*
|
||||
re-computation happens:
|
||||
|
||||
- ``False`` / ``None`` — disabled (no memory savings)
|
||||
- ``True`` / ``"full"`` — every transformer block (~30% slow, biggest save)
|
||||
- ``"medium"`` — every other block (balance)
|
||||
- ``"selective"`` — attention only (~10% slow, modest save)
|
||||
- ``"auto"`` — pick based on detected VRAM headroom
|
||||
|
||||
``resolve_gradient_checkpointing`` returns a kwargs-dict suitable for
|
||||
``TrainingArguments(**kwargs)``. Granularity (medium / selective) is a separate
|
||||
concept — query ``resolve_granularity`` for the chosen tier so callers can
|
||||
install the correct downstream hooks without polluting HF's kwargs surface.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Union
|
||||
|
||||
TierLike = Union[bool, str, None]
|
||||
|
||||
|
||||
# Heuristic VRAM thresholds (GB) for tier=auto. Tuned for a 7-8B LoRA run at
|
||||
# bf16 + 4bit quant + max_length≈4k.
|
||||
# Below 24 GB: full checkpoint (largest memory savings).
|
||||
# 24-80 GB: medium (every other block) — balance.
|
||||
# >80 GB: selective (attention only) — minimize slowdown.
|
||||
AUTO_FULL_THRESHOLD_GB = 24.0
|
||||
AUTO_SELECTIVE_THRESHOLD_GB = 80.0
|
||||
|
||||
|
||||
def resolve_granularity(
|
||||
tier: TierLike, gpu_memory_gb: float | None = None,
|
||||
) -> str | None:
|
||||
"""Return the granularity string the wrapper should install hooks for.
|
||||
|
||||
One of: ``"full"`` | ``"medium"`` | ``"selective"`` | ``None`` (disabled).
|
||||
``"auto"`` is resolved to full/medium/selective based on ``gpu_memory_gb``.
|
||||
"""
|
||||
if not tier:
|
||||
return None
|
||||
if tier is True or tier == "full":
|
||||
return "full"
|
||||
if tier in ("medium", "selective"):
|
||||
return tier # type: ignore[return-value]
|
||||
if tier == "auto":
|
||||
if gpu_memory_gb is None:
|
||||
return "full"
|
||||
if gpu_memory_gb < AUTO_FULL_THRESHOLD_GB:
|
||||
return "full"
|
||||
if gpu_memory_gb <= AUTO_SELECTIVE_THRESHOLD_GB:
|
||||
return "medium"
|
||||
return "selective"
|
||||
return None
|
||||
|
||||
|
||||
def resolve_gradient_checkpointing(
|
||||
tier: TierLike, gpu_memory_gb: float | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Resolve a gradient_checkpointing setting into TrainingArguments kwargs.
|
||||
|
||||
Only returns keys that HuggingFace's ``TrainingArguments`` actually accepts.
|
||||
Granularity (medium/selective) is not represented here; query
|
||||
``resolve_granularity`` for that.
|
||||
|
||||
Args:
|
||||
tier: TrainingConfig.gradient_checkpointing value (bool or tier string).
|
||||
gpu_memory_gb: GPU memory (GB) used by ``"auto"`` tier. If None, falls
|
||||
back to full checkpointing on auto.
|
||||
|
||||
Returns:
|
||||
Dict of kwargs suitable for ``TrainingArguments(**kwargs)``:
|
||||
- ``gradient_checkpointing`` (bool)
|
||||
- ``gradient_checkpointing_kwargs`` (dict)
|
||||
"""
|
||||
granularity = resolve_granularity(tier, gpu_memory_gb=gpu_memory_gb)
|
||||
if granularity is None:
|
||||
return {}
|
||||
|
||||
# All granularities use HF's standard non-reentrant checkpointing at the
|
||||
# TrainingArguments level. Selective / medium installation happens inside
|
||||
# the wrapper via torch-level hooks (deferred to v0.28.1 wiring), without
|
||||
# leaking markers into HF's kwargs surface.
|
||||
return {
|
||||
"gradient_checkpointing": True,
|
||||
"gradient_checkpointing_kwargs": {"use_reentrant": False},
|
||||
}
|
||||
|
||||
|
||||
def describe_tier(tier: TierLike, gpu_memory_gb: float | None = None) -> str:
|
||||
"""Return a short human-readable description of the selected tier."""
|
||||
if not tier:
|
||||
return "off"
|
||||
if tier is True or tier == "full":
|
||||
return "full (every block)"
|
||||
if tier == "medium":
|
||||
return "medium (every other block)"
|
||||
if tier == "selective":
|
||||
return "selective (attention only)"
|
||||
if tier == "auto":
|
||||
if gpu_memory_gb is None:
|
||||
return "auto → full (unknown VRAM)"
|
||||
if gpu_memory_gb < AUTO_FULL_THRESHOLD_GB:
|
||||
return f"auto → full (VRAM {gpu_memory_gb:.0f}GB)"
|
||||
if gpu_memory_gb <= AUTO_SELECTIVE_THRESHOLD_GB:
|
||||
return f"auto → medium (VRAM {gpu_memory_gb:.0f}GB)"
|
||||
return f"auto → selective (VRAM {gpu_memory_gb:.0f}GB)"
|
||||
return str(tier)
|
||||
|
|
@ -0,0 +1,179 @@
|
|||
"""Kernel auto-composition — benchmark and pick the fastest kernel combo.
|
||||
|
||||
Enumerates installed performance kernels (Liger, FlashAttention, torch
|
||||
baseline) and picks the fastest combination for the current GPU. Benchmarks
|
||||
each candidate on a small warm-up loop and selects the one with the lowest
|
||||
observed step time.
|
||||
|
||||
This is a config-resolver helper: the benchmarking loop is expected to be
|
||||
driven by the trainer wrapper (a few warm-up steps before the real train
|
||||
loop). Here we only provide:
|
||||
|
||||
- ``enumerate_kernel_combos(backend, device)`` — list candidate combos
|
||||
- ``pick_best_kernel(candidates)`` — choose the fastest from timing results
|
||||
|
||||
Design: we never auto-enable combos that the user has *disabled* via their
|
||||
TrainingConfig (e.g. if ``use_liger: false`` explicitly, it stays false). The
|
||||
picker only searches within combos the user hasn't opted out of.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
def enumerate_kernel_combos(
|
||||
backend: str, device: str,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Enumerate candidate kernel combinations for the current environment.
|
||||
|
||||
Returns a list of dicts: ``{"name", "use_liger", "use_flash_attn",
|
||||
"use_cut_ce"}``. The list always contains a ``baseline`` entry (no special
|
||||
kernels) so that the picker has a reference point.
|
||||
|
||||
Rules:
|
||||
- CPU → only baseline.
|
||||
- unsloth backend → baseline only (unsloth uses its own kernels internally).
|
||||
- mlx backend → baseline only (Apple Silicon path doesn't share kernels).
|
||||
- cuda + transformers → baseline + each available kernel + known-good combos.
|
||||
"""
|
||||
baseline = {
|
||||
"name": "baseline",
|
||||
"use_liger": False,
|
||||
"use_flash_attn": False,
|
||||
"use_cut_ce": False,
|
||||
}
|
||||
|
||||
# CPU: nothing to compose
|
||||
if device != "cuda":
|
||||
return [baseline]
|
||||
|
||||
# Unsloth + MLX have their own kernel paths - picker would just confuse them
|
||||
if backend in ("unsloth", "mlx"):
|
||||
return [baseline]
|
||||
|
||||
combos: list[dict[str, Any]] = [baseline]
|
||||
|
||||
# Probe availability — lazy imports inside each helper
|
||||
try:
|
||||
from soup_cli.utils.liger import check_liger_available
|
||||
|
||||
liger_ok = check_liger_available()
|
||||
except ImportError:
|
||||
liger_ok = False
|
||||
|
||||
try:
|
||||
from soup_cli.utils.flash_attn import check_flash_attn_available
|
||||
|
||||
flash_ok = check_flash_attn_available() is not None
|
||||
except ImportError:
|
||||
flash_ok = False
|
||||
|
||||
try:
|
||||
from soup_cli.utils.cut_ce import check_cut_ce_available
|
||||
|
||||
cce_ok = check_cut_ce_available()
|
||||
except ImportError:
|
||||
cce_ok = False
|
||||
|
||||
if liger_ok:
|
||||
combos.append({
|
||||
"name": "liger",
|
||||
"use_liger": True,
|
||||
"use_flash_attn": False,
|
||||
"use_cut_ce": False,
|
||||
})
|
||||
|
||||
if flash_ok:
|
||||
combos.append({
|
||||
"name": "flash",
|
||||
"use_liger": False,
|
||||
"use_flash_attn": True,
|
||||
"use_cut_ce": False,
|
||||
})
|
||||
|
||||
if liger_ok and flash_ok:
|
||||
combos.append({
|
||||
"name": "liger+flash",
|
||||
"use_liger": True,
|
||||
"use_flash_attn": True,
|
||||
"use_cut_ce": False,
|
||||
})
|
||||
|
||||
if cce_ok:
|
||||
combos.append({
|
||||
"name": "cut_ce",
|
||||
"use_liger": False,
|
||||
"use_flash_attn": False,
|
||||
"use_cut_ce": True,
|
||||
})
|
||||
|
||||
if liger_ok and flash_ok and cce_ok:
|
||||
combos.append({
|
||||
"name": "liger+flash+cut_ce",
|
||||
"use_liger": True,
|
||||
"use_flash_attn": True,
|
||||
"use_cut_ce": True,
|
||||
})
|
||||
|
||||
return combos
|
||||
|
||||
|
||||
def pick_best_kernel(
|
||||
candidates: list[dict[str, Any]],
|
||||
) -> dict[str, Any]:
|
||||
"""Pick the fastest kernel combo given benchmarked timings.
|
||||
|
||||
Args:
|
||||
candidates: List of dicts each with ``name`` and ``time_ms`` fields.
|
||||
|
||||
Returns:
|
||||
The single candidate dict with the lowest ``time_ms``.
|
||||
Ties are broken by order (first candidate wins) so callers can place
|
||||
the preferred default (baseline) first.
|
||||
|
||||
Raises:
|
||||
ValueError: if ``candidates`` is empty or if **all** candidates are
|
||||
missing a finite ``time_ms`` (no benchmark signal — picking blindly
|
||||
would mask a silent infrastructure failure).
|
||||
"""
|
||||
if not candidates:
|
||||
raise ValueError("pick_best_kernel requires at least one candidate")
|
||||
|
||||
finite = [c for c in candidates if _finite_time_ms(c.get("time_ms"))]
|
||||
if not finite:
|
||||
raise ValueError(
|
||||
"pick_best_kernel: no candidate has a finite time_ms — "
|
||||
"benchmarking appears to have failed for every combo."
|
||||
)
|
||||
|
||||
# Stable sort: ties go to the one earlier in the list (baseline usually).
|
||||
return min(candidates, key=lambda c: _sortable_time_ms(c.get("time_ms")))
|
||||
|
||||
|
||||
def _finite_time_ms(value: Any) -> bool:
|
||||
"""True if ``value`` is a real finite number (not None, not NaN, not inf)."""
|
||||
import math
|
||||
|
||||
if value is None:
|
||||
return False
|
||||
try:
|
||||
as_float = float(value)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
return math.isfinite(as_float)
|
||||
|
||||
|
||||
def _sortable_time_ms(value: Any) -> float:
|
||||
"""Convert ``time_ms`` to a sortable float; missing/NaN → +inf."""
|
||||
import math
|
||||
|
||||
if value is None:
|
||||
return float("inf")
|
||||
try:
|
||||
as_float = float(value)
|
||||
except (TypeError, ValueError):
|
||||
return float("inf")
|
||||
if math.isnan(as_float):
|
||||
return float("inf")
|
||||
return as_float
|
||||
|
|
@ -0,0 +1,952 @@
|
|||
"""Tests for v0.28.0 — Training Speed & Memory features.
|
||||
|
||||
Covers:
|
||||
- Part A: Cut Cross-Entropy (CCE)
|
||||
- Part B: FP8 training (quantization_aware='fp8')
|
||||
- Part C: Gradient checkpointing tiers (selective/medium/full/auto)
|
||||
- Part D: Kernel auto-composition (utils/kernel_picker.py)
|
||||
- Part E: Cross-document attention masking for sample packing
|
||||
- Part F: Activation offloading to CPU/disk
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from soup_cli.config.schema import SoupConfig, TrainingConfig
|
||||
|
||||
# ─── Part A: Cut Cross-Entropy (CCE) ───────────────────────────────────────
|
||||
|
||||
|
||||
class TestCutCEConfig:
|
||||
"""TrainingConfig.use_cut_ce boolean field."""
|
||||
|
||||
def test_use_cut_ce_default_false(self):
|
||||
cfg = SoupConfig(base="test/model", data={"train": "./data.jsonl"})
|
||||
assert cfg.training.use_cut_ce is False
|
||||
|
||||
def test_use_cut_ce_enabled(self):
|
||||
cfg = SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"use_cut_ce": True},
|
||||
)
|
||||
assert cfg.training.use_cut_ce is True
|
||||
|
||||
def test_use_cut_ce_type_bool_coerce(self):
|
||||
"""Pydantic coerces bool-like strings; verify type is bool after."""
|
||||
cfg = TrainingConfig(use_cut_ce=True)
|
||||
assert cfg.use_cut_ce is True
|
||||
assert isinstance(cfg.use_cut_ce, bool)
|
||||
|
||||
|
||||
class TestCutCEAvailability:
|
||||
"""Cut Cross-Entropy availability + detection."""
|
||||
|
||||
def test_check_cut_ce_available_returns_bool(self):
|
||||
from soup_cli.utils.cut_ce import check_cut_ce_available
|
||||
|
||||
result = check_cut_ce_available()
|
||||
assert isinstance(result, bool)
|
||||
|
||||
def test_check_cut_ce_not_installed(self):
|
||||
from soup_cli.utils.cut_ce import check_cut_ce_available
|
||||
|
||||
# sys.modules[name]=None makes ``import name`` raise ImportError
|
||||
with patch.dict("sys.modules", {"cut_cross_entropy": None}):
|
||||
assert check_cut_ce_available() is False
|
||||
|
||||
def test_get_cut_ce_version_not_installed(self):
|
||||
from soup_cli.utils.cut_ce import get_cut_ce_version
|
||||
|
||||
with patch(
|
||||
"soup_cli.utils.cut_ce.check_cut_ce_available", return_value=False
|
||||
):
|
||||
result = get_cut_ce_version()
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestCutCEApplication:
|
||||
"""Applying Cut Cross-Entropy to a model."""
|
||||
|
||||
def test_apply_cut_ce_not_installed(self):
|
||||
from soup_cli.utils.cut_ce import apply_cut_ce
|
||||
|
||||
with patch(
|
||||
"soup_cli.utils.cut_ce.check_cut_ce_available", return_value=False
|
||||
):
|
||||
result = apply_cut_ce("meta-llama/Llama-3.1-8B")
|
||||
assert result is False
|
||||
|
||||
def test_apply_cut_ce_available_tries_patching(self):
|
||||
from soup_cli.utils.cut_ce import apply_cut_ce
|
||||
|
||||
with patch(
|
||||
"soup_cli.utils.cut_ce.check_cut_ce_available", return_value=True
|
||||
):
|
||||
# With cut_cross_entropy not actually installed, patch should return
|
||||
# False (can't import real module). Just verifying it doesn't crash.
|
||||
result = apply_cut_ce("meta-llama/Llama-3.1-8B")
|
||||
assert isinstance(result, bool)
|
||||
|
||||
def test_apply_cut_ce_calls_llama_patch_for_llama_model(self):
|
||||
"""Verify the llama detector routes to cce_patch('llama')."""
|
||||
from soup_cli.utils.cut_ce import apply_cut_ce
|
||||
|
||||
fake_cce = MagicMock()
|
||||
fake_transformers = MagicMock(cce_patch=fake_cce)
|
||||
fake_module = MagicMock(transformers=fake_transformers)
|
||||
with patch(
|
||||
"soup_cli.utils.cut_ce.check_cut_ce_available", return_value=True
|
||||
), patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"cut_cross_entropy": fake_module,
|
||||
"cut_cross_entropy.transformers": fake_transformers,
|
||||
},
|
||||
):
|
||||
assert apply_cut_ce("meta-llama/Llama-3.1-8B") is True
|
||||
fake_cce.assert_called_once_with("llama")
|
||||
|
||||
def test_apply_cut_ce_deepseek_phi_does_not_use_phi(self):
|
||||
"""Regression: org-prefix like 'deepseek-ai/...' must not trigger phi."""
|
||||
from soup_cli.utils.cut_ce import apply_cut_ce
|
||||
|
||||
fake_cce = MagicMock()
|
||||
fake_transformers = MagicMock(cce_patch=fake_cce)
|
||||
fake_module = MagicMock(transformers=fake_transformers)
|
||||
with patch(
|
||||
"soup_cli.utils.cut_ce.check_cut_ce_available", return_value=True
|
||||
), patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"cut_cross_entropy": fake_module,
|
||||
"cut_cross_entropy.transformers": fake_transformers,
|
||||
},
|
||||
):
|
||||
# Llama-distilled model name contains no "phi" substring
|
||||
# anymore thanks to the last-path-component detector.
|
||||
assert apply_cut_ce("deepseek-ai/deepseek-coder-7b-instruct") is False
|
||||
fake_cce.assert_not_called()
|
||||
|
||||
|
||||
class TestCutCEValidation:
|
||||
"""Cut Cross-Entropy config validation."""
|
||||
|
||||
def test_validate_cut_ce_disabled_returns_empty(self):
|
||||
from soup_cli.utils.cut_ce import validate_cut_ce_config
|
||||
|
||||
errors = validate_cut_ce_config(False, "transformers", "cuda")
|
||||
assert errors == []
|
||||
|
||||
def test_validate_cut_ce_not_installed(self):
|
||||
from soup_cli.utils.cut_ce import validate_cut_ce_config
|
||||
|
||||
with patch(
|
||||
"soup_cli.utils.cut_ce.check_cut_ce_available", return_value=False
|
||||
):
|
||||
errors = validate_cut_ce_config(True, "transformers", "cuda")
|
||||
assert any("not installed" in err for err in errors)
|
||||
|
||||
def test_validate_cut_ce_requires_cuda(self):
|
||||
from soup_cli.utils.cut_ce import validate_cut_ce_config
|
||||
|
||||
with patch(
|
||||
"soup_cli.utils.cut_ce.check_cut_ce_available", return_value=True
|
||||
):
|
||||
errors = validate_cut_ce_config(True, "transformers", "cpu")
|
||||
assert any("CUDA" in err for err in errors)
|
||||
|
||||
def test_validate_cut_ce_unsloth_incompatible(self):
|
||||
from soup_cli.utils.cut_ce import validate_cut_ce_config
|
||||
|
||||
with patch(
|
||||
"soup_cli.utils.cut_ce.check_cut_ce_available", return_value=True
|
||||
):
|
||||
errors = validate_cut_ce_config(True, "unsloth", "cuda")
|
||||
assert any("unsloth" in err.lower() for err in errors)
|
||||
|
||||
def test_validate_cut_ce_valid(self):
|
||||
from soup_cli.utils.cut_ce import validate_cut_ce_config
|
||||
|
||||
with patch(
|
||||
"soup_cli.utils.cut_ce.check_cut_ce_available", return_value=True
|
||||
):
|
||||
errors = validate_cut_ce_config(True, "transformers", "cuda")
|
||||
assert errors == []
|
||||
|
||||
def test_validate_cut_ce_mlx_incompatible(self):
|
||||
from soup_cli.utils.cut_ce import validate_cut_ce_config
|
||||
|
||||
with patch(
|
||||
"soup_cli.utils.cut_ce.check_cut_ce_available", return_value=True
|
||||
):
|
||||
errors = validate_cut_ce_config(True, "mlx", "mps")
|
||||
assert any("mlx" in err.lower() for err in errors)
|
||||
|
||||
|
||||
# ─── Part B: FP8 training ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestFP8Config:
|
||||
"""quantization_aware now accepts bool or literal 'fp8'."""
|
||||
|
||||
def test_quantization_aware_default_false(self):
|
||||
cfg = SoupConfig(base="test/model", data={"train": "./data.jsonl"})
|
||||
assert cfg.training.quantization_aware is False
|
||||
|
||||
def test_quantization_aware_bool_true(self):
|
||||
cfg = SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"quantization_aware": True},
|
||||
)
|
||||
assert cfg.training.quantization_aware is True
|
||||
|
||||
def test_quantization_aware_fp8(self):
|
||||
cfg = SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"quantization_aware": "fp8"},
|
||||
)
|
||||
assert cfg.training.quantization_aware == "fp8"
|
||||
|
||||
def test_quantization_aware_invalid_string_rejected(self):
|
||||
"""Only 'fp8' literal is accepted, other strings rejected."""
|
||||
with pytest.raises(ValidationError) as exc:
|
||||
SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"quantization_aware": "fp16"},
|
||||
)
|
||||
assert "fp8" in str(exc.value) or "quantization_aware" in str(exc.value)
|
||||
|
||||
def test_quantization_aware_bool_still_works(self):
|
||||
cfg = SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"quantization_aware": False},
|
||||
)
|
||||
assert cfg.training.quantization_aware is False
|
||||
|
||||
|
||||
class TestFP8Availability:
|
||||
"""FP8 training dependency checks."""
|
||||
|
||||
def test_is_fp8_available_returns_bool(self):
|
||||
from soup_cli.utils.fp8 import is_fp8_available
|
||||
|
||||
result = is_fp8_available()
|
||||
assert isinstance(result, bool)
|
||||
|
||||
def test_fp8_requires_hopper_gpu_info(self):
|
||||
"""is_fp8_gpu_supported should check GPU compute capability."""
|
||||
from soup_cli.utils.fp8 import is_fp8_gpu_supported
|
||||
|
||||
# Shouldn't crash even without CUDA
|
||||
result = is_fp8_gpu_supported()
|
||||
assert isinstance(result, bool)
|
||||
|
||||
def test_is_fp8_available_false_when_deps_missing(self):
|
||||
"""Explicit false branch — both torchao.float8 and transformer_engine absent."""
|
||||
with patch.dict(
|
||||
"sys.modules",
|
||||
{"torchao.float8": None, "transformer_engine": None},
|
||||
):
|
||||
from soup_cli.utils.fp8 import is_fp8_available
|
||||
|
||||
assert is_fp8_available() is False
|
||||
|
||||
def test_is_fp8_gpu_supported_pre_hopper_false(self):
|
||||
"""Pre-Hopper (SM 8.x, e.g. A100) is not supported."""
|
||||
fake_torch = MagicMock()
|
||||
fake_torch.cuda.is_available.return_value = True
|
||||
fake_torch.cuda.get_device_capability.return_value = (8, 0)
|
||||
with patch.dict("sys.modules", {"torch": fake_torch}):
|
||||
from soup_cli.utils.fp8 import is_fp8_gpu_supported
|
||||
|
||||
assert is_fp8_gpu_supported() is False
|
||||
|
||||
def test_is_fp8_gpu_supported_hopper_true(self):
|
||||
"""Hopper (SM 9.x, H100) is supported."""
|
||||
fake_torch = MagicMock()
|
||||
fake_torch.cuda.is_available.return_value = True
|
||||
fake_torch.cuda.get_device_capability.return_value = (9, 0)
|
||||
with patch.dict("sys.modules", {"torch": fake_torch}):
|
||||
from soup_cli.utils.fp8 import is_fp8_gpu_supported
|
||||
|
||||
assert is_fp8_gpu_supported() is True
|
||||
|
||||
|
||||
class TestFP8Validation:
|
||||
"""FP8 training config validation."""
|
||||
|
||||
def test_validate_fp8_not_requested_returns_empty(self):
|
||||
from soup_cli.utils.fp8 import validate_fp8_config
|
||||
|
||||
errors = validate_fp8_config(False, "transformers", "cuda")
|
||||
assert errors == []
|
||||
|
||||
def test_validate_fp8_bool_returns_empty(self):
|
||||
"""Bool True means int8 QAT (existing path), not FP8."""
|
||||
from soup_cli.utils.fp8 import validate_fp8_config
|
||||
|
||||
errors = validate_fp8_config(True, "transformers", "cuda")
|
||||
# Bool True is int8 QAT, handled by qat.py, not fp8
|
||||
assert errors == []
|
||||
|
||||
def test_validate_fp8_cpu_rejected(self):
|
||||
from soup_cli.utils.fp8 import validate_fp8_config
|
||||
|
||||
errors = validate_fp8_config("fp8", "transformers", "cpu")
|
||||
assert any("CUDA" in err for err in errors)
|
||||
|
||||
def test_validate_fp8_unsloth_rejected(self):
|
||||
from soup_cli.utils.fp8 import validate_fp8_config
|
||||
|
||||
errors = validate_fp8_config("fp8", "unsloth", "cuda")
|
||||
assert any("unsloth" in err.lower() for err in errors)
|
||||
|
||||
def test_validate_fp8_mlx_rejected(self):
|
||||
from soup_cli.utils.fp8 import validate_fp8_config
|
||||
|
||||
errors = validate_fp8_config("fp8", "mlx", "mps")
|
||||
assert any("mlx" in err.lower() or "CUDA" in err for err in errors)
|
||||
|
||||
|
||||
# ─── Part C: Gradient checkpointing tiers ─────────────────────────────────
|
||||
|
||||
|
||||
class TestGradientCheckpointingTiers:
|
||||
"""gradient_checkpointing accepts bool or tier literal."""
|
||||
|
||||
def test_default_false(self):
|
||||
cfg = SoupConfig(base="test/model", data={"train": "./data.jsonl"})
|
||||
assert cfg.training.gradient_checkpointing is False
|
||||
|
||||
def test_bool_true(self):
|
||||
cfg = SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"gradient_checkpointing": True},
|
||||
)
|
||||
assert cfg.training.gradient_checkpointing is True
|
||||
|
||||
def test_tier_selective(self):
|
||||
cfg = SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"gradient_checkpointing": "selective"},
|
||||
)
|
||||
assert cfg.training.gradient_checkpointing == "selective"
|
||||
|
||||
def test_tier_medium(self):
|
||||
cfg = SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"gradient_checkpointing": "medium"},
|
||||
)
|
||||
assert cfg.training.gradient_checkpointing == "medium"
|
||||
|
||||
def test_tier_full(self):
|
||||
cfg = SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"gradient_checkpointing": "full"},
|
||||
)
|
||||
assert cfg.training.gradient_checkpointing == "full"
|
||||
|
||||
def test_tier_auto(self):
|
||||
cfg = SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"gradient_checkpointing": "auto"},
|
||||
)
|
||||
assert cfg.training.gradient_checkpointing == "auto"
|
||||
|
||||
def test_invalid_tier_rejected(self):
|
||||
with pytest.raises(ValidationError) as exc:
|
||||
SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"gradient_checkpointing": "partial"},
|
||||
)
|
||||
assert "gradient_checkpointing" in str(exc.value)
|
||||
|
||||
|
||||
class TestGradientCheckpointingResolver:
|
||||
"""Resolve config value + GPU info → kwargs for TrainingArguments."""
|
||||
|
||||
def test_resolve_disabled_returns_empty(self):
|
||||
from soup_cli.utils.gradient_ckpt import resolve_gradient_checkpointing
|
||||
|
||||
kwargs = resolve_gradient_checkpointing(False, gpu_memory_gb=80)
|
||||
assert kwargs == {}
|
||||
|
||||
def test_resolve_bool_true_returns_full_ckpt(self):
|
||||
from soup_cli.utils.gradient_ckpt import resolve_gradient_checkpointing
|
||||
|
||||
kwargs = resolve_gradient_checkpointing(True, gpu_memory_gb=80)
|
||||
assert kwargs["gradient_checkpointing"] is True
|
||||
|
||||
def test_resolve_full_tier(self):
|
||||
from soup_cli.utils.gradient_ckpt import resolve_gradient_checkpointing
|
||||
|
||||
kwargs = resolve_gradient_checkpointing("full", gpu_memory_gb=80)
|
||||
assert kwargs["gradient_checkpointing"] is True
|
||||
|
||||
def test_resolve_selective_tier(self):
|
||||
from soup_cli.utils.gradient_ckpt import (
|
||||
resolve_gradient_checkpointing,
|
||||
resolve_granularity,
|
||||
)
|
||||
|
||||
kwargs = resolve_gradient_checkpointing("selective", gpu_memory_gb=80)
|
||||
assert kwargs["gradient_checkpointing"] is True
|
||||
# No private markers leak into HF TrainingArguments kwargs.
|
||||
assert kwargs["gradient_checkpointing_kwargs"] == {"use_reentrant": False}
|
||||
# Granularity is exposed via a separate helper for the wrapper.
|
||||
assert resolve_granularity("selective", gpu_memory_gb=80) == "selective"
|
||||
|
||||
def test_resolve_medium_tier(self):
|
||||
from soup_cli.utils.gradient_ckpt import resolve_gradient_checkpointing
|
||||
|
||||
kwargs = resolve_gradient_checkpointing("medium", gpu_memory_gb=80)
|
||||
assert kwargs["gradient_checkpointing"] is True
|
||||
|
||||
def test_resolve_auto_low_memory_selects_full(self):
|
||||
from soup_cli.utils.gradient_ckpt import resolve_gradient_checkpointing
|
||||
|
||||
kwargs = resolve_gradient_checkpointing("auto", gpu_memory_gb=16)
|
||||
# Low VRAM → full checkpointing
|
||||
assert kwargs["gradient_checkpointing"] is True
|
||||
|
||||
def test_resolve_auto_high_memory_selects_selective(self):
|
||||
from soup_cli.utils.gradient_ckpt import resolve_gradient_checkpointing
|
||||
|
||||
# 80GB+ → selective only (attention), saving speed
|
||||
kwargs = resolve_gradient_checkpointing("auto", gpu_memory_gb=80)
|
||||
assert kwargs["gradient_checkpointing"] is True
|
||||
|
||||
def test_resolve_auto_very_high_memory_selects_selective(self):
|
||||
from soup_cli.utils.gradient_ckpt import (
|
||||
resolve_granularity,
|
||||
)
|
||||
|
||||
# 192GB (H200): selective (attention-only) tier, minimize slowdown
|
||||
assert resolve_granularity("auto", gpu_memory_gb=192) == "selective"
|
||||
|
||||
def test_resolve_auto_medium_memory_selects_medium(self):
|
||||
from soup_cli.utils.gradient_ckpt import resolve_granularity
|
||||
|
||||
# 40GB (A100 40GB): medium (every other block)
|
||||
assert resolve_granularity("auto", gpu_memory_gb=40) == "medium"
|
||||
|
||||
def test_resolve_auto_no_gpu_info_full(self):
|
||||
from soup_cli.utils.gradient_ckpt import resolve_granularity
|
||||
|
||||
assert resolve_granularity("auto", gpu_memory_gb=None) == "full"
|
||||
|
||||
def test_describe_tier_off(self):
|
||||
from soup_cli.utils.gradient_ckpt import describe_tier
|
||||
|
||||
assert describe_tier(False) == "off"
|
||||
|
||||
def test_describe_tier_full(self):
|
||||
from soup_cli.utils.gradient_ckpt import describe_tier
|
||||
|
||||
assert "full" in describe_tier(True)
|
||||
assert "full" in describe_tier("full")
|
||||
|
||||
def test_describe_tier_auto(self):
|
||||
from soup_cli.utils.gradient_ckpt import describe_tier
|
||||
|
||||
assert "auto" in describe_tier("auto")
|
||||
assert "auto" in describe_tier("auto", gpu_memory_gb=16)
|
||||
assert "selective" in describe_tier("auto", gpu_memory_gb=192)
|
||||
|
||||
|
||||
# ─── Part D: Kernel auto-composition ─────────────────────────────────────
|
||||
|
||||
|
||||
class TestKernelPickerConfig:
|
||||
"""kernel_auto_compose config flag."""
|
||||
|
||||
def test_kernel_auto_compose_default_false(self):
|
||||
cfg = SoupConfig(base="test/model", data={"train": "./data.jsonl"})
|
||||
assert cfg.training.kernel_auto_compose is False
|
||||
|
||||
def test_kernel_auto_compose_enabled(self):
|
||||
cfg = SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"kernel_auto_compose": True},
|
||||
)
|
||||
assert cfg.training.kernel_auto_compose is True
|
||||
|
||||
|
||||
class TestKernelPickerEnumerate:
|
||||
"""Kernel picker enumerates available kernel combinations."""
|
||||
|
||||
def test_enumerate_returns_list(self):
|
||||
from soup_cli.utils.kernel_picker import enumerate_kernel_combos
|
||||
|
||||
combos = enumerate_kernel_combos(backend="transformers", device="cuda")
|
||||
assert isinstance(combos, list)
|
||||
|
||||
def test_enumerate_baseline_always_present(self):
|
||||
from soup_cli.utils.kernel_picker import enumerate_kernel_combos
|
||||
|
||||
combos = enumerate_kernel_combos(backend="transformers", device="cuda")
|
||||
# Baseline (no special kernels) must always be an option
|
||||
assert any(c.get("name") == "baseline" for c in combos)
|
||||
|
||||
def test_enumerate_cpu_only_baseline(self):
|
||||
from soup_cli.utils.kernel_picker import enumerate_kernel_combos
|
||||
|
||||
combos = enumerate_kernel_combos(backend="transformers", device="cpu")
|
||||
# On CPU, only baseline should be available
|
||||
assert len(combos) == 1
|
||||
assert combos[0]["name"] == "baseline"
|
||||
|
||||
def test_enumerate_unsloth_skips_liger(self):
|
||||
from soup_cli.utils.kernel_picker import enumerate_kernel_combos
|
||||
|
||||
combos = enumerate_kernel_combos(backend="unsloth", device="cuda")
|
||||
# Unsloth has its own fused kernels - no Liger combos
|
||||
for combo in combos:
|
||||
assert "liger" not in combo.get("name", "").lower()
|
||||
|
||||
|
||||
class TestKernelPickerDecision:
|
||||
"""Kernel picker decision logic (mocked benchmarks)."""
|
||||
|
||||
def test_pick_best_returns_dict(self):
|
||||
from soup_cli.utils.kernel_picker import pick_best_kernel
|
||||
|
||||
# With fake timing results, picks fastest
|
||||
candidates = [
|
||||
{"name": "baseline", "time_ms": 100.0},
|
||||
{"name": "liger", "time_ms": 70.0},
|
||||
{"name": "liger+flash", "time_ms": 50.0},
|
||||
]
|
||||
best = pick_best_kernel(candidates)
|
||||
assert best["name"] == "liger+flash"
|
||||
|
||||
def test_pick_best_baseline_if_only_one(self):
|
||||
from soup_cli.utils.kernel_picker import pick_best_kernel
|
||||
|
||||
candidates = [{"name": "baseline", "time_ms": 100.0}]
|
||||
best = pick_best_kernel(candidates)
|
||||
assert best["name"] == "baseline"
|
||||
|
||||
def test_pick_best_empty_raises(self):
|
||||
from soup_cli.utils.kernel_picker import pick_best_kernel
|
||||
|
||||
with pytest.raises(ValueError, match="at least one"):
|
||||
pick_best_kernel([])
|
||||
|
||||
def test_pick_best_tie_returns_first(self):
|
||||
"""Ties broken by list order (first-wins) — preferred default first."""
|
||||
from soup_cli.utils.kernel_picker import pick_best_kernel
|
||||
|
||||
candidates = [
|
||||
{"name": "baseline", "time_ms": 50.0},
|
||||
{"name": "liger", "time_ms": 50.0},
|
||||
]
|
||||
best = pick_best_kernel(candidates)
|
||||
assert best["name"] == "baseline"
|
||||
|
||||
def test_pick_best_all_missing_time_raises(self):
|
||||
"""All-untimed candidates means benchmarking failed — must not promote silently."""
|
||||
from soup_cli.utils.kernel_picker import pick_best_kernel
|
||||
|
||||
candidates = [
|
||||
{"name": "baseline"},
|
||||
{"name": "liger"},
|
||||
{"name": "flash", "time_ms": None},
|
||||
]
|
||||
with pytest.raises(ValueError, match="finite time_ms"):
|
||||
pick_best_kernel(candidates)
|
||||
|
||||
def test_pick_best_nan_time_treated_as_missing(self):
|
||||
from soup_cli.utils.kernel_picker import pick_best_kernel
|
||||
|
||||
candidates = [
|
||||
{"name": "baseline", "time_ms": float("nan")},
|
||||
{"name": "liger", "time_ms": 50.0},
|
||||
]
|
||||
best = pick_best_kernel(candidates)
|
||||
assert best["name"] == "liger"
|
||||
|
||||
|
||||
# ─── Part E: Cross-document attention masking ────────────────────────────
|
||||
|
||||
|
||||
class TestCrossDocAttnMaskConfig:
|
||||
"""packing_cross_doc_attn_mask config."""
|
||||
|
||||
def test_default_false(self):
|
||||
cfg = SoupConfig(base="test/model", data={"train": "./data.jsonl"})
|
||||
assert cfg.training.packing_cross_doc_attn_mask is False
|
||||
|
||||
def test_enabled(self):
|
||||
cfg = SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"packing_cross_doc_attn_mask": True, "packing": True},
|
||||
)
|
||||
assert cfg.training.packing_cross_doc_attn_mask is True
|
||||
|
||||
def test_requires_packing(self):
|
||||
"""Enabling cross-doc mask without packing should error."""
|
||||
with pytest.raises(ValidationError) as exc:
|
||||
SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={
|
||||
"packing_cross_doc_attn_mask": True,
|
||||
"packing": False,
|
||||
},
|
||||
)
|
||||
assert "packing" in str(exc.value).lower()
|
||||
|
||||
|
||||
class TestCrossDocAttnMaskBuild:
|
||||
"""Build cross-doc attention mask from document boundaries."""
|
||||
|
||||
def test_build_mask_single_doc(self):
|
||||
from soup_cli.utils.cross_doc_attn import build_cross_doc_mask
|
||||
|
||||
# Single doc spans the whole sequence — no masking needed
|
||||
boundaries = [0, 10] # doc0 occupies positions 0..9
|
||||
mask = build_cross_doc_mask(boundaries, seq_length=10)
|
||||
# Shape (10, 10), lower-triangular within the doc
|
||||
assert mask.shape == (10, 10)
|
||||
|
||||
def test_build_mask_two_docs(self):
|
||||
from soup_cli.utils.cross_doc_attn import build_cross_doc_mask
|
||||
|
||||
# doc0: 0..4, doc1: 5..9
|
||||
boundaries = [0, 5, 10]
|
||||
mask = build_cross_doc_mask(boundaries, seq_length=10)
|
||||
# Position 5 (doc1 start) should NOT attend to position 0 (doc0)
|
||||
assert mask[5, 0] == 0
|
||||
# Position 5 attending to itself should be 1
|
||||
assert mask[5, 5] == 1
|
||||
# Position 0 attending to itself should be 1
|
||||
assert mask[0, 0] == 1
|
||||
# Position 1 attending to 0 should be 1 (same doc, causal)
|
||||
assert mask[1, 0] == 1
|
||||
# Position 0 attending to 1 should be 0 (causal — future)
|
||||
assert mask[0, 1] == 0
|
||||
|
||||
def test_build_mask_boundaries_validation(self):
|
||||
from soup_cli.utils.cross_doc_attn import build_cross_doc_mask
|
||||
|
||||
# Boundaries must start at 0 and end at seq_length
|
||||
with pytest.raises(ValueError):
|
||||
build_cross_doc_mask([1, 10], seq_length=10)
|
||||
|
||||
def test_build_mask_boundaries_monotonic(self):
|
||||
from soup_cli.utils.cross_doc_attn import build_cross_doc_mask
|
||||
|
||||
with pytest.raises(ValueError, match="increasing"):
|
||||
build_cross_doc_mask([0, 5, 3, 10], seq_length=10)
|
||||
|
||||
def test_build_mask_empty_boundaries_raises(self):
|
||||
from soup_cli.utils.cross_doc_attn import build_cross_doc_mask
|
||||
|
||||
with pytest.raises(ValueError, match="non-empty"):
|
||||
build_cross_doc_mask([], seq_length=10)
|
||||
|
||||
def test_build_mask_wrong_end_raises(self):
|
||||
from soup_cli.utils.cross_doc_attn import build_cross_doc_mask
|
||||
|
||||
with pytest.raises(ValueError, match="seq_length"):
|
||||
build_cross_doc_mask([0, 8], seq_length=10)
|
||||
|
||||
def test_compute_doc_boundaries_empty(self):
|
||||
from soup_cli.utils.cross_doc_attn import compute_doc_boundaries
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
compute_doc_boundaries([])
|
||||
|
||||
def test_compute_doc_boundaries_non_positive(self):
|
||||
from soup_cli.utils.cross_doc_attn import compute_doc_boundaries
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
compute_doc_boundaries([3, 0, 4])
|
||||
|
||||
def test_compute_doc_boundaries_valid(self):
|
||||
from soup_cli.utils.cross_doc_attn import compute_doc_boundaries
|
||||
|
||||
assert compute_doc_boundaries([3, 2, 4]) == [0, 3, 5, 9]
|
||||
|
||||
|
||||
# ─── Part F: Activation offloading ────────────────────────────────────────
|
||||
|
||||
|
||||
class TestActivationOffloadingConfig:
|
||||
"""activation_offloading config."""
|
||||
|
||||
def test_default_none(self):
|
||||
cfg = SoupConfig(base="test/model", data={"train": "./data.jsonl"})
|
||||
assert cfg.training.activation_offloading is None
|
||||
|
||||
def test_cpu(self):
|
||||
cfg = SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"activation_offloading": "cpu"},
|
||||
)
|
||||
assert cfg.training.activation_offloading == "cpu"
|
||||
|
||||
def test_disk(self):
|
||||
cfg = SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"activation_offloading": "disk"},
|
||||
)
|
||||
assert cfg.training.activation_offloading == "disk"
|
||||
|
||||
def test_invalid_target(self):
|
||||
with pytest.raises(ValidationError) as exc:
|
||||
SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={"activation_offloading": "gpu"},
|
||||
)
|
||||
assert "activation_offloading" in str(exc.value)
|
||||
|
||||
|
||||
class TestActivationOffloadingValidation:
|
||||
"""Offloading config validation."""
|
||||
|
||||
def test_validate_none_returns_empty(self):
|
||||
from soup_cli.utils.activation_offload import validate_offload_config
|
||||
|
||||
errors = validate_offload_config(None, "transformers", "cuda")
|
||||
assert errors == []
|
||||
|
||||
def test_validate_cpu_requires_cuda(self):
|
||||
from soup_cli.utils.activation_offload import validate_offload_config
|
||||
|
||||
errors = validate_offload_config("cpu", "transformers", "cpu")
|
||||
assert any("CUDA" in err for err in errors)
|
||||
|
||||
def test_validate_cpu_unsloth_incompatible(self):
|
||||
from soup_cli.utils.activation_offload import validate_offload_config
|
||||
|
||||
errors = validate_offload_config("cpu", "unsloth", "cuda")
|
||||
assert any("unsloth" in err.lower() for err in errors)
|
||||
|
||||
def test_validate_disk_valid(self):
|
||||
from soup_cli.utils.activation_offload import validate_offload_config
|
||||
|
||||
errors = validate_offload_config(
|
||||
"disk", "transformers", "cuda", save_dir="./scratch"
|
||||
)
|
||||
assert errors == []
|
||||
|
||||
def test_validate_disk_requires_save_dir(self):
|
||||
"""Disk mode must reject calls without save_dir — fail-fast at validate()."""
|
||||
from soup_cli.utils.activation_offload import validate_offload_config
|
||||
|
||||
errors = validate_offload_config(
|
||||
"disk", "transformers", "cuda", save_dir=None
|
||||
)
|
||||
assert any("save_dir" in err for err in errors)
|
||||
|
||||
def test_validate_disk_on_cpu_rejected(self):
|
||||
from soup_cli.utils.activation_offload import validate_offload_config
|
||||
|
||||
errors = validate_offload_config(
|
||||
"disk", "transformers", "cpu", save_dir="./scratch"
|
||||
)
|
||||
assert any("CUDA" in err for err in errors)
|
||||
|
||||
def test_validate_mlx_rejected(self):
|
||||
from soup_cli.utils.activation_offload import validate_offload_config
|
||||
|
||||
errors = validate_offload_config("cpu", "mlx", "cuda")
|
||||
assert any("mlx" in err.lower() for err in errors)
|
||||
|
||||
|
||||
class TestActivationOffloadingHooks:
|
||||
"""Install / uninstall hooks for offloading."""
|
||||
|
||||
def test_context_manager_noop_when_none(self):
|
||||
from soup_cli.utils.activation_offload import offload_context
|
||||
|
||||
# Should be a no-op when target is None
|
||||
with offload_context(None, save_dir=None):
|
||||
pass # nothing to assert - just shouldn't crash
|
||||
|
||||
def test_context_manager_cpu(self):
|
||||
from soup_cli.utils.activation_offload import offload_context
|
||||
|
||||
# Should not crash even without torch
|
||||
with offload_context("cpu", save_dir=None):
|
||||
pass
|
||||
|
||||
def test_offload_context_disk_requires_dir(self, tmp_path):
|
||||
from soup_cli.utils.activation_offload import offload_context
|
||||
|
||||
# Disk mode should accept a save_dir
|
||||
with offload_context("disk", save_dir=str(tmp_path)):
|
||||
pass
|
||||
|
||||
def test_offload_context_unknown_target_raises(self):
|
||||
"""Defense-in-depth: unknown target raises even if torch is present."""
|
||||
from soup_cli.utils.activation_offload import offload_context
|
||||
|
||||
# Only reaches the ValueError branch if torch imports successfully
|
||||
try:
|
||||
import torch # noqa: F401
|
||||
except ImportError:
|
||||
pytest.skip("torch not installed; ValueError branch unreachable")
|
||||
|
||||
with pytest.raises(ValueError, match="Unknown activation_offloading"):
|
||||
with offload_context("invalid", save_dir=None):
|
||||
pass
|
||||
|
||||
def test_offload_context_disk_creates_save_dir(self, tmp_path):
|
||||
"""Disk mode should create the scratch directory on context entry."""
|
||||
try:
|
||||
import torch # noqa: F401
|
||||
except ImportError:
|
||||
pytest.skip("torch not installed; disk hooks unreachable")
|
||||
|
||||
from soup_cli.utils.activation_offload import offload_context
|
||||
|
||||
scratch = tmp_path / "offload_scratch"
|
||||
with offload_context("disk", save_dir=str(scratch)):
|
||||
assert scratch.exists()
|
||||
|
||||
|
||||
# ─── Integration: multiple features composed ──────────────────────────────
|
||||
|
||||
|
||||
class TestV028SFTOnlyValidator:
|
||||
"""v0.28.0 features are wired only in SFTTrainerWrapper.
|
||||
|
||||
The SoupConfig validator rejects non-SFT tasks when speed/memory flags
|
||||
are set — prevents silent no-ops and the known fp8 crash path in the
|
||||
legacy int8 QAT wrapper.
|
||||
"""
|
||||
|
||||
def test_use_cut_ce_rejected_on_dpo(self):
|
||||
with pytest.raises(ValidationError) as exc:
|
||||
SoupConfig(
|
||||
base="m",
|
||||
task="dpo",
|
||||
data={"train": "./d.jsonl", "format": "dpo"},
|
||||
training={"use_cut_ce": True},
|
||||
)
|
||||
assert "use_cut_ce" in str(exc.value)
|
||||
assert "sft" in str(exc.value)
|
||||
|
||||
def test_fp8_rejected_on_grpo(self):
|
||||
with pytest.raises(ValidationError) as exc:
|
||||
SoupConfig(
|
||||
base="m",
|
||||
task="grpo",
|
||||
data={"train": "./d.jsonl"},
|
||||
training={"quantization_aware": "fp8"},
|
||||
)
|
||||
assert "fp8" in str(exc.value)
|
||||
|
||||
def test_activation_offloading_rejected_on_ppo(self):
|
||||
with pytest.raises(ValidationError) as exc:
|
||||
SoupConfig(
|
||||
base="m",
|
||||
task="ppo",
|
||||
data={"train": "./d.jsonl"},
|
||||
training={"activation_offloading": "cpu"},
|
||||
)
|
||||
assert "activation_offloading" in str(exc.value)
|
||||
|
||||
def test_kernel_auto_compose_rejected_on_kto(self):
|
||||
with pytest.raises(ValidationError) as exc:
|
||||
SoupConfig(
|
||||
base="m",
|
||||
task="kto",
|
||||
data={"train": "./d.jsonl", "format": "kto"},
|
||||
training={"kernel_auto_compose": True},
|
||||
)
|
||||
assert "kernel_auto_compose" in str(exc.value)
|
||||
|
||||
def test_sft_accepts_all_features(self):
|
||||
"""SFT task should accept every v0.28.0 flag (happy path)."""
|
||||
cfg = SoupConfig(
|
||||
base="m",
|
||||
task="sft",
|
||||
data={"train": "./d.jsonl"},
|
||||
training={
|
||||
"use_cut_ce": True,
|
||||
"quantization_aware": "fp8",
|
||||
"activation_offloading": "cpu",
|
||||
"kernel_auto_compose": True,
|
||||
},
|
||||
)
|
||||
assert cfg.training.use_cut_ce is True
|
||||
|
||||
def test_non_sft_unaffected_when_flags_default(self):
|
||||
"""DPO/GRPO with default v0.28.0 flags still validate."""
|
||||
cfg = SoupConfig(
|
||||
base="m",
|
||||
task="dpo",
|
||||
data={"train": "./d.jsonl", "format": "dpo"},
|
||||
)
|
||||
assert cfg.task == "dpo"
|
||||
|
||||
def test_quantization_aware_bool_true_allowed_on_dpo(self):
|
||||
"""Int8 QAT (bool True) still works on non-SFT — only fp8 is restricted."""
|
||||
cfg = SoupConfig(
|
||||
base="m",
|
||||
task="dpo",
|
||||
data={"train": "./d.jsonl", "format": "dpo"},
|
||||
training={"quantization_aware": True},
|
||||
)
|
||||
assert cfg.training.quantization_aware is True
|
||||
|
||||
def test_gradient_checkpointing_tier_allowed_on_dpo(self):
|
||||
"""Tier strings fall back to truthy (bool True) in non-SFT wrappers — no crash."""
|
||||
cfg = SoupConfig(
|
||||
base="m",
|
||||
task="dpo",
|
||||
data={"train": "./d.jsonl", "format": "dpo"},
|
||||
training={"gradient_checkpointing": "auto"},
|
||||
)
|
||||
assert cfg.training.gradient_checkpointing == "auto"
|
||||
|
||||
|
||||
class TestV028Integration:
|
||||
"""Multiple v0.28.0 features composed in one config."""
|
||||
|
||||
def test_all_features_compose(self):
|
||||
cfg = SoupConfig(
|
||||
base="test/model",
|
||||
data={"train": "./data.jsonl"},
|
||||
training={
|
||||
"use_cut_ce": True,
|
||||
"quantization_aware": "fp8",
|
||||
"gradient_checkpointing": "auto",
|
||||
"kernel_auto_compose": True,
|
||||
"packing": True,
|
||||
"packing_cross_doc_attn_mask": True,
|
||||
"activation_offloading": "cpu",
|
||||
},
|
||||
)
|
||||
tcfg = cfg.training
|
||||
assert tcfg.use_cut_ce is True
|
||||
assert tcfg.quantization_aware == "fp8"
|
||||
assert tcfg.gradient_checkpointing == "auto"
|
||||
assert tcfg.kernel_auto_compose is True
|
||||
assert tcfg.packing is True
|
||||
assert tcfg.packing_cross_doc_attn_mask is True
|
||||
assert tcfg.activation_offloading == "cpu"
|
||||
Loading…
Reference in New Issue