feat(precision,rollout): live fp8/nvfp4 + vLLM sleep + openenv rollout + apple-adapter + delinearize-llama4 (v0.71.21)

Closes #141, #124, #125, #228, #97.

- #141: apply_fp8_attention (torchao float8 on attention projections, Hopper
  gate) + apply_nvfp4 (NVFP4Config, Blackwell gate); partial-conversion honesty;
  wired into the v0.28 speed/memory pipeline with yellow-advisory degrade.
- #124: vllm_sleep_mode live - create_vllm_engine(sleep_mode=True) +
  vllm_sleep_cycle ctx (wake in finally) + TRL GRPOConfig hook probe.
- #125: openenv rollout fully live via training.rollout_func module:fn
  resolver; rows replace the prompt dataset; art/ruler/nemo_gym honest dep
  gates + _EXTERNAL_ROLLOUT_RUNNERS seam. Real GRPO train on SmolLM2-135M.
- #228: convert_apple_adapter live - PEFT LoRA <-> mlx-lm (both matrices
  transpose, bf16 upcast, adapters.safetensors + num_layers, npz legacy read,
  np.ascontiguousarray fix for safetensors non-contiguous mangling);
  *-to-apple upstream-gated exit 3.
- #97: delinearize-llama4 live - [E*din,dout] -> [E,din,dout] per shard,
  config.json expert-count probe + --num-experts, sidecar copy, atomic writes.

Review waves: 3 HIGH + ~8 MEDIUM + ~12 LOW fixed.
Tests: 13874 -> 14084 (+210 in tests/test_v07121.py).
Full suite: 13967 passed, 117 skipped. ruff clean.
This commit is contained in:
Alpamys 2026-06-10 16:29:52 +05:00
parent a4dfbb308c
commit ed5fc3a8b3
25 changed files with 4008 additions and 135 deletions

View File

@ -12,6 +12,68 @@ reproducing 70+ versions of notes.
## [Unreleased]
## [0.71.21] - 2026-06-10
### Added
- **Precision & rollout lift (BETA, hw-gated)** — lifts five deferred
`NotImplementedError` stubs to live code.
- **FP8 attention + NVFP4 (closes #141).** `training.fp8_attention: true`
now converts the model's attention projections (q/k/v/o + fused qkv
variants) to FP8 training modules via torchao's
`convert_to_float8_training` with an attention-only `module_filter_fn`
(Hopper SM ≥ 9.0 gate); `training.nvfp4: true` quantises via torchao's
`NVFP4Config` (Blackwell SM ≥ 10.0 gate). Both are wired into the v0.28
speed/memory pipeline and degrade to a visible yellow advisory when the
gate fires — a conversion failing partway raises an honest "model may be
PARTIALLY converted" error rather than silently training on a
half-converted model.
- **vLLM sleep mode (closes #124).** `training.vllm_sleep_mode: true` is
live: `create_vllm_engine(sleep_mode=True)` sets
`AsyncEngineArgs.enable_sleep_mode` (vLLM ≥ 0.7 gate with a friendly
upgrade message), the new `vllm_sleep_cycle(engine, level=1|2)` context
manager wraps the optimisation step (wake in `finally`), and the GRPO
trainer threads the flag into TRL's `GRPOConfig` when the installed TRL
exposes the hook (advisory otherwise).
- **Multi-turn agent rollout launchers (closes #125).** `soup train` with
`task: grpo` + `training.rollout_backend: openenv` +
`training.rollout_func: my_module:fn` now runs a LIVE rollout: the
resolver imports the operator's callable (same trusted-code policy as
`data.prompt_strategy`), feeds it the dataset prompts as seeds, and the
returned `{prompt, answer?}` rows replace the prompt dataset. Rows are
normalised (extra keys stripped, message-list prompts deep-copied,
non-string answers rejected loudly). `art` / `ruler` / `nemo_gym` raise a
friendly ImportError when the backend package is missing and an honest
BETA gate when present (injectable `_EXTERNAL_ROLLOUT_RUNNERS` seam).
Validated by a real GRPO + openenv rollout train on SmolLM2-135M.
- **Apple-adapter conversion (closes #228).** `soup apple-adapter` is live
for `hf-to-mlx` / `mlx-to-hf`: PEFT LoRA safetensors ↔ mlx-lm adapters
with both matrices transposed (`lora_A [r,in]` ↔ `lora_a [in,r]`),
bf16 sources upcast via the torch loader, `adapters.safetensors` +
`num_layers` emitted for mlx-lm's `load_adapters`, rank/alpha/dropout
carried through, legacy `adapters.npz` still read, optional v0.60
Merkle-root signing. The `*-to-apple` directions stay upstream-gated
(no published FoundationModels adapter spec). Validated by a real bf16
PEFT adapter round-tripping with numeric equality.
- **Llama-4 expert delinearization (closes #97).** `soup
delinearize-llama4` now runs a live torch runtime: fused 2-D expert
tensors `[E*dim_in, dim_out]` reshape to 3-D `[E, dim_in, dim_out]`
(expert count from `config.json` or `--num-experts`), other tensors pass
through, JSON sidecars are copied, writes are atomic. `--plan-only`
keeps the old render-and-exit flow.
### Fixed
- `safetensors.numpy.save` silently mangles non-contiguous (transposed)
arrays — the apple-adapter writer now makes every array C-contiguous
first (caught by the new round-trip assertions).
### Known limitations
- fp8_attention / nvfp4 / vllm_sleep_mode are BETA hardware-gated — the
converters and gates ship validated via capability probes and fake-module
dispatch tests, but end-to-end runs need a Hopper/Blackwell GPU + torchao
(or vLLM ≥ 0.7), none of which exist on the maintainer's RTX 3050 /
Windows box. The `art` / `ruler` / `nemo_gym` rollout adapters are
honestly BETA-gated until validated against the upstream packages.
## [0.71.20] - 2026-06-09
### Added

View File

@ -49,20 +49,23 @@ infrastructure instead of improving models. Soup fixes that.
## What's New
**v0.71.20 — Modality II trainers (BETA, hardware-gated).** Three v0.52.0 schema stubs become real
trainers:
**v0.71.21 — Precision & rollout lift (BETA, hardware-gated).** Five long-deferred stubs go live:
- **TTS fine-tuning**`task: tts` now trains a real text-to-speech model. The five families
(Orpheus / Sesame-CSM / Llasa / Spark / Oute) are decoder LMs, so a TTS fine-tune is next-token
cross-entropy over `[text][audio-codec-token]` chat sequences. Pre-encode your audio to codec tokens
offline and train with `data.format: chat` — Soup adds per-family emotion control + codec
special-token registration. (The encode-at-train-time path is dependency-gated on each family's codec.)
- **BitNet 1.58-bit**`soup export --format bitnet | tq1_0` runs a real llama.cpp TQ1_0 ternary
export, and `quantization: bitnet_1.58` routes to a live BitNet trainer (gated on `onebitllms`).
- **MoE expert quant**`training.moe_expert_quant: nf4 | int8_rowwise` quantizes just the
fused-MoE expert layers with bitsandbytes (attention + router stay full precision), and
`training.train_router_only: true` freezes the experts to train only the gating router. Validated
live on an RTX 3050.
- **Multi-turn agent rollouts for GRPO**`training.rollout_backend: openenv` +
`training.rollout_func: my_module:my_fn` runs a live multi-turn rollout at the start of GRPO
training; the rows your function returns replace the prompt dataset. ART / RULER / NeMo-Gym
adapters ship behind honest dependency gates.
- **FP8 attention + NVFP4 training**`training.fp8_attention: true` converts the attention
projections to torchao float8 (Hopper-gated) and `training.nvfp4: true` applies torchao NVFP4
quantization (Blackwell-gated). Unsupported hardware degrades to a clear advisory, never a crash.
- **vLLM sleep mode for RL**`training.vllm_sleep_mode: true` puts the vLLM engine on standby
between GRPO rollouts (vLLM ≥ 0.7), freeing VRAM for the training step.
- **Apple-adapter conversion is live**`soup apple-adapter <dir> --direction hf-to-mlx | mlx-to-hf`
converts PEFT LoRA safetensors ↔ mlx-lm adapters with a numerically-equal round trip
(rank / scale / dropout / num_layers carried).
- **Llama-4 expert delinearization is live**`soup delinearize-llama4 <src> --target <out>`
reshapes fused `[E*din, dout]` expert weights to `[E, din, dout]` shard-by-shard and copies the
JSON sidecars so the target stays loadable.
Full history: [CHANGELOG.md](CHANGELOG.md) &middot; [GitHub Releases](https://github.com/MakazhanAlpamys/Soup/releases).

View File

@ -641,10 +641,10 @@ Closed catalog (`MappingProxyType`) of ready-to-edit YAML / JSON. Output path cw
## Llama 4 Delinearizer
```bash
soup delinearize-llama4 ./llama4-checkpoint --target ./out-delinearized --yes
soup delinearize-llama4 ./llama4-checkpoint --target ./out-delinearized [--num-experts N] [--plan-only]
```
Plans Llama 4 expert-weight reshape for export. v0.44.0 ships the planner; live runtime in v0.44.1. `is_llama4_model` uses a word-boundary regex matching the `is_gemma4_model` pattern — `ungemma-llama-4ish` is rejected.
LIVE (v0.71.21): reshapes fused Llama-4 expert weights `[E*din, dout]``[E, din, dout]` shard-by-shard (atomic writes, per-shard 16 GiB cap, cwd containment) and copies the JSON sidecars so the target stays loadable. The expert count defaults from `config.json` (`text_config.num_local_experts`); pass `--num-experts` when the config doesn't carry it (exit 2 otherwise). `--plan-only` keeps the original preview flow and writes nothing. `is_llama4_model` uses a word-boundary regex matching the `is_gemma4_model` pattern — `ungemma-llama-4ish` is rejected.
## Ctrl+C Graceful Save

View File

@ -24,6 +24,7 @@ soup push --model ./output --repo user/name Upload to HuggingFace
soup push --model ./output --repo user/name --collection user/coll-abc123 Add to HF Collection
soup merge --adapter ./output Merge LoRA with base model
soup merge-sharded-fsdp-weights ./shards -o merged.safetensors Consolidate FSDP shards into one safetensors (v0.71.14; --plan-only previews)
soup delinearize-llama4 ./src --target ./out [--num-experts N] [--plan-only] Live Llama-4 fused-expert reshape [E*din,dout] -> [E,din,dout] + sidecar copy (v0.71.21)
soup export --model ./output --format gguf Export to GGUF (Ollama)
soup export --model ./output --deploy ollama Export GGUF + auto-deploy to Ollama
soup export --model ./output --format onnx Export to ONNX
@ -225,7 +226,7 @@ soup lock show / soup lock check Show + drift-check (exit 3 on drif
soup compile <program.py> --eval <suite> [--optimizer mipro|gepa|textgrad|copro|bootstrap_fewshot] [--plan-only] DSPy / GEPA / TextGrad prompt-program compiler — live (v0.71.13; pip install 'soup-cli[compile]')
soup distill-prompt --traces <jsonl> --teacher <m> --student <m> --strategy sft|preference|kl [--provider ollama|anthropic|vllm] [--base-url <url>] [--temperature F] [--max-rows N] Distill prompt-heavy traces via a live teacher (v0.71.13)
soup compile-tools <spec.json|yaml> --eval <jsonl> [--optimizer textgrad|gepa] [--plan-only] TextGrad / GEPA tool-schema optimiser — live (v0.71.13; pip install 'soup-cli[compile]')
soup apple-adapter <source-dir> --direction hf-to-mlx|mlx-to-hf|hf-to-apple|mlx-to-apple --output <dir> [--sign] HF / MLX / Apple FoundationModels adapter conversion (v0.68.0)
soup apple-adapter <source-dir> --direction hf-to-mlx|mlx-to-hf|hf-to-apple|mlx-to-apple --output <dir> [--sign] [--plan-only] PEFT LoRA <-> mlx-lm adapter conversion — live (v0.71.21; *-to-apple upstream-gated exit 3)
soup local-rl init --db <path> Create personal-LLM flywheel SQLite schema (v0.68.0)
soup local-rl status --db <path> Print interactions / thumbs-up / thumbs-down counters
soup local-rl record --db <path> --prompt <q> --response <r> --thumb up|down Append thumbs record

View File

@ -468,12 +468,13 @@ soup serve --model ./output --kv-cache-type q8_0 # 8-bit quantized KV cache
- vLLM / SGLang serve wiring is still tracked under [#140](https://github.com/MakazhanAlpamys/Soup/issues/140) (`infra-blocked`).
## FP8 Attention + NVFP4 + Native `unsloth_bnb_4bit` (v0.53.0)
## FP8 Attention + NVFP4 + Native `unsloth_bnb_4bit`
Three new TrainingConfig bools extend the v0.28.0 FP8 menu:
Three TrainingConfig bools extend the v0.28.0 FP8 menu. `fp8_attention` and `nvfp4` are LIVE
torchao converters as of v0.71.21 (hardware-gated):
- `fp8_attention: true` — requires `quantization_aware: fp8` AND a non-MLX backend. Targets axolotl parity for FP8 attention on Hopper+ GPUs.
- `nvfp4: true` — Blackwell-only FP4 training. Gated to non-MLX + `modality: text`; the SM ≥ 12.0 runtime check fires at trainer construction.
- `fp8_attention: true` — requires `quantization_aware: fp8` AND a non-MLX backend. Converts the attention projections (q/k/v/o and fused variants) to torchao float8 training on Hopper+ GPUs. Missing torchao or a pre-Hopper GPU degrades to a clear advisory; a conversion-phase failure raises an honest "model may be PARTIALLY converted" error instead of training on a half-converted model.
- `nvfp4: true` — Blackwell-only FP4 training via torchao `NVFP4Config` + `quantize_`. Gated to non-MLX + `modality: text`; the SM ≥ 10 runtime check fires at trainer construction.
- `unsloth_bnb_4bit: true` — promotes "Unsloth Dynamic 4-bit" from an implicit `backend=unsloth + quantization=4bit` combo to a named flag. Mutual rejection of inconsistent combos at config load.
Cross-validator ordering picks the most actionable error: `quantization_aware='fp8'` prerequisite fires before the MLX rejection on `fp8_attention`, so a YAML missing both surfaces the deeper issue first.

View File

@ -484,9 +484,10 @@ training:
grpo_fp16: true # FP16 RL (unsloth parity)
# Long-context + memory-efficient RL
long_context_grpo: true # wires Tiled MLP when available
vllm_sleep_mode: true # between-rollouts vLLM standby
# Multi-turn agent rollout
rollout_backend: art # one of: art / ruler / nemo_gym / openenv
vllm_sleep_mode: true # between-rollouts vLLM standby — LIVE (vLLM >= 0.7)
# Multi-turn agent rollout — openenv is LIVE: your function's rows replace the prompt dataset
rollout_backend: openenv # one of: art / ruler / nemo_gym / openenv
rollout_func: my_module:my_rollout # module:function resolver (openenv; trusted operator code)
# Stability / efficiency knobs
ref_model_ema_alpha: 0.99 # EMA sync policy → reference
replay_buffer_size: 2048

View File

@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "soup-cli"
version = "0.71.20"
version = "0.71.21"
description = "Fine-tune and post-train LLMs in one command. No SSH, no config hell."
readme = "README.md"
license = "Apache-2.0"

View File

@ -1,3 +1,3 @@
"""Soup CLI — Fine-tune and post-train LLMs in one command."""
__version__ = "0.71.20"
__version__ = "0.71.21"

View File

@ -61,13 +61,37 @@ def apple_adapter_cmd(
if plan_only:
return
# v0.71.21 #228 — live conversion. Exit codes: 2 = validation /
# missing-input, 3 = upstream-gated (Apple spec not yet public),
# 1 = missing dependency, 0 = success.
try:
convert_apple_adapter(plan)
except NotImplementedError as exc:
report = convert_apple_adapter(plan)
except RuntimeError as exc:
console.print(
Panel(
f"[yellow]{escape(str(exc))}[/]",
title="Live apple-adapter deferred",
title="apple-adapter — upstream-gated",
)
)
raise typer.Exit(3) from exc
except (ValueError, FileNotFoundError) as exc:
console.print(f"[red]{escape(str(exc))}[/]")
raise typer.Exit(2) from exc
except ImportError as exc:
console.print(f"[red]{escape(str(exc))}[/]")
raise typer.Exit(1) from exc
skipped_note = (
f"\nSkipped (non-LoRA): [bold]{len(report.skipped_keys)}[/]"
if report.skipped_keys
else ""
)
console.print(
Panel(
f"Converted LoRA keys: [bold]{report.converted_keys}[/]\n"
f"Output: [bold]{escape(report.output_dir)}[/]\n"
f"Signed: [bold]{report.signed}[/]{skipped_note}",
title="soup apple-adapter — converted",
border_style="green",
)
)

View File

@ -1,13 +1,24 @@
"""v0.44.0 Part D — `soup delinearize-llama4` command (schema stub)."""
"""`soup delinearize-llama4` — Llama 4 expert-weight reshape.
v0.44.0 shipped the planner; v0.71.21 (#97) lifts the live torch runtime
(load each shard reshape fused 2-D expert weights to 3-D atomic write
to ``--target``). ``--plan-only`` keeps the old render-and-exit flow.
"""
from __future__ import annotations
from typing import Optional
import typer
from rich.console import Console
from rich.markup import escape
from rich.panel import Panel
from soup_cli.utils.delinearize_llama4 import is_llama4_model, plan_delinearize
from soup_cli.utils.delinearize_llama4 import (
is_llama4_model,
plan_delinearize,
run_delinearize,
)
console = Console()
@ -28,15 +39,26 @@ def delinearize_llama4(
"--model-id",
help="Optional model id; warn if it doesn't look like a Llama 4 model.",
),
yes: bool = typer.Option(
num_experts: Optional[int] = typer.Option(
None,
"--num-experts",
help=(
"Expert count for the reshape. Defaults to "
"config.json's num_local_experts."
),
),
plan_only: bool = typer.Option(
False,
"--yes",
help="Acknowledge that live runtime lands in v0.44.1 (plan-only now).",
"--plan-only",
help="Render the plan and exit without reshaping (v0.44.0 flow).",
),
) -> None:
"""Plan Llama 4 expert-weight delinearization for export.
"""Delinearize Llama 4 fused expert weights to 3-D form.
v0.44.0 ships the planner; live torch-side reshape lands in v0.44.1.
Live since v0.71.21 (#97): reshapes ``...experts.gate_up_proj`` /
``...experts.down_proj`` from ``[E*dim_in, dim_out]`` to
``[E, dim_in, dim_out]`` and copies JSON sidecars so the target stays
a loadable checkpoint.
"""
if model_id is not None and not is_llama4_model(model_id):
console.print(
@ -48,17 +70,36 @@ def delinearize_llama4(
except (ValueError, FileNotFoundError) as exc:
console.print(f"[red]{escape(str(exc))}[/]")
raise typer.Exit(code=2) from exc
body = (
f"Weight files: {len(plan.weight_files)}\n"
f"Source dir: {escape(plan.source_dir)}\n"
f"Target dir: {escape(plan.target_dir)}\n\n"
"Live delinearization runtime lands in v0.44.1 — plan-only for now."
)
console.print(
Panel(body, title="Llama 4 Delinearization Plan", border_style="cyan")
)
if not yes:
console.print(
"[yellow]Pass --yes to acknowledge the deferred runtime.[/]"
Panel(
f"Weight files: {len(plan.weight_files)}\n"
f"Source dir: {escape(plan.source_dir)}\n"
f"Target dir: {escape(plan.target_dir)}",
title="Llama 4 Delinearization Plan",
border_style="cyan",
)
raise typer.Exit(code=0)
)
if plan_only:
return
try:
result = run_delinearize(plan, num_experts=num_experts)
except (TypeError, ValueError, FileNotFoundError) as exc:
console.print(f"[red]{escape(str(exc))}[/]")
raise typer.Exit(code=2) from exc
except ImportError as exc:
console.print(f"[red]{escape(str(exc))}[/]")
raise typer.Exit(code=1) from exc
console.print(
Panel(
f"Files written: {len(result.files_written)}\n"
f"Reshaped keys: {result.reshaped_keys}\n"
f"Already 3-D: {result.already_3d_keys}\n"
f"Passthrough keys: {result.passthrough_keys}\n"
f"Sidecars copied: {result.sidecars_copied}\n"
f"Target dir: {escape(result.target_dir)}",
title="Llama 4 Delinearization — done",
border_style="green",
)
)

View File

@ -1096,7 +1096,7 @@ class TrainingConfig(BaseModel):
"in v0.50.1."
),
)
# v0.50.0 Part C — Multi-turn agent rollout backend
# v0.50.0 Part C — Multi-turn agent rollout backend (live v0.71.21 #125)
rollout_backend: Optional[Literal[
"art", "ruler", "nemo_gym", "openenv"
]] = Field(
@ -1104,10 +1104,31 @@ class TrainingConfig(BaseModel):
description=(
"Multi-turn agent rollout backend (unsloth / axolotl parity): "
"art (OpenPipe ART) / ruler / nemo_gym / openenv. "
"Requires task='grpo'. Schema-only in v0.50.0; live launcher "
"wired in v0.50.1."
"Requires task='grpo'. openenv runs live (v0.71.21 #125) via "
"training.rollout_func; art/ruler/nemo_gym are lazy-import "
"gated."
),
)
# v0.71.21 #125 — user-supplied OpenEnv rollout callable.
rollout_func: Optional[str] = Field(
default=None,
description=(
"OpenEnv rollout function as 'module.path:function_name' "
"(v0.71.21 #125). Requires rollout_backend='openenv'. The "
"callable receives the seed prompts list and returns rollout "
"rows ({'prompt': str|messages, 'answer'?: str}) that replace "
"the GRPO prompt dataset. Trusted-input policy: names "
"operator-controlled code (mirrors data.prompt_strategy)."
),
)
@field_validator("rollout_func", mode="before")
@classmethod
def _validate_rollout_func_field(cls, value):
"""v0.71.21 #125 — module:fn shape validation at config load."""
from soup_cli.utils.agent_rollout import validate_rollout_func
return validate_rollout_func(value)
# v0.50.0 Part D — GRPO stability / efficiency knobs (axolotl + unsloth).
# All schema-only in v0.50.0; live trainer callbacks wired in v0.50.1.
ref_model_ema_alpha: Optional[float] = Field(
@ -3609,7 +3630,17 @@ class SoupConfig(BaseModel):
@model_validator(mode="after")
def _validate_rollout_backend(self) -> "SoupConfig":
"""v0.50.0 Part C — ``rollout_backend`` requires task='grpo' and a
non-mlx backend. Live launcher wired in v0.50.1."""
non-mlx backend. Live launcher wired in v0.71.21 (#125):
openenv requires ``rollout_func`` and ``rollout_func`` is
openenv-only (silent-no-op footgun rejection)."""
if (
self.training.rollout_func is not None
and self.training.rollout_backend != "openenv"
):
raise ValueError(
"rollout_func requires rollout_backend='openenv'; got "
f"rollout_backend={self.training.rollout_backend!r}"
)
if self.training.rollout_backend is None:
return self
if self.task != "grpo":
@ -3620,6 +3651,14 @@ class SoupConfig(BaseModel):
raise ValueError(
"rollout_backend is not supported on backend=mlx in v0.50.0"
)
if (
self.training.rollout_backend == "openenv"
and self.training.rollout_func is None
):
raise ValueError(
"rollout_backend='openenv' requires training.rollout_func "
"('module.path:function_name')"
)
return self
@model_validator(mode="after")

View File

@ -292,6 +292,30 @@ class GRPOTrainerWrapper:
# --- Dataset ---
# GRPO expects prompts — extract from messages or use prompt field
train_data = _prepare_grpo_dataset(dataset["train"])
# v0.71.21 #125 — multi-turn agent rollout backend. The backend
# receives the dataset prompts as seeds; its rows REPLACE the
# prompt dataset (the env is the data source).
if tcfg.rollout_backend is not None:
from soup_cli.utils.agent_rollout import launch_rollout
rollout_result = launch_rollout(
tcfg.rollout_backend,
prompts=[row["prompt"] for row in train_data],
rollout_func=tcfg.rollout_func,
model=self.model,
tokenizer=self.tokenizer,
reward_fn=reward_fn,
)
train_data = _prepare_grpo_dataset(
[dict(row) for row in rollout_result.rows]
)
console.print(
f"[green]Rollout backend '{tcfg.rollout_backend}':[/] "
f"{len(train_data)} prompts collected "
"(replacing dataset prompts)"
)
train_ds = Dataset.from_list(train_data)
eval_ds = None
if "val" in dataset and dataset["val"]:
@ -346,6 +370,22 @@ class GRPOTrainerWrapper:
"max_completion_length": cfg.data.max_length,
}
# v0.71.21 #124 — vLLM sleep mode: set TRL's GRPOConfig hook when the
# installed TRL exposes it; otherwise print a friendly advisory
# (Soup's own vLLM engine factory honors sleep_mode=True).
if tcfg.vllm_sleep_mode:
import inspect as _inspect
from soup_cli.utils.grpo_long_context import (
maybe_enable_trl_sleep_mode,
)
maybe_enable_trl_sleep_mode(
grpo_kwargs,
_inspect.signature(GRPOConfig).parameters,
console,
)
# CPU support: set use_cpu and prevent empty generations
if self.device == "cpu":
import inspect as _inspect

View File

@ -1,6 +1,6 @@
"""v0.53.0 Part D — Train-time advanced precision schema helpers.
"""v0.53.0 Part D — Train-time advanced precision (schema + live wiring).
Three new TrainingConfig surfaces ship this release (schema-only):
Three TrainingConfig surfaces ship here:
* ``fp8_attention: bool`` extend the v0.28.0 FP8 menu to apply FP8 to
attention (axolotl-parity flag). Requires ``quantization_aware='fp8'``.
@ -11,11 +11,93 @@ Three new TrainingConfig surfaces ship this release (schema-only):
+ ``quantization='4bit'``). When True, requires ``backend='unsloth'`` and
``quantization='4bit'``.
Live wiring lands in v0.53.1 (mirrors v0.50.0 / v0.52.0 stub-then-live).
v0.71.21 #141 lifts the two ``apply_*`` stubs to live, BETA hw-gated code:
* :func:`apply_fp8_attention` converts the attention-projection linears to
torchao ``Float8Linear`` training modules (Hopper+ gate, SM >= 9.0).
* :func:`apply_nvfp4` routes the model through torchao's ``NVFP4Config``
quantisation (Blackwell gate SM 10.0 datacenter B100/B200/GB200 or
SM 12.0 consumer RTX 50-series).
Both raise friendly ``RuntimeError`` on missing hardware / torchao instead
of silently no-opping; the trainer-side wiring in
``utils/v028_features.apply_v028_speed_memory`` degrades those to yellow
advisories so a training kick-off never crashes on instrumentation.
"""
from __future__ import annotations
# Attention-projection module names (last FQN component). Covers the
# separate-QKV Llama/Mistral/Qwen/Phi shape, GPT-2's fused ``c_attn``,
# Phi-3 / GPT-NeoX fused variants, and encoder-style ``out_proj``.
_ATTENTION_PROJ_NAMES: frozenset[str] = frozenset({
"q_proj",
"k_proj",
"v_proj",
"o_proj",
"out_proj",
"qkv_proj",
"c_attn",
"query_key_value",
"Wqkv",
})
# NVIDIA Blackwell compute capability: SM 10.0 (B100 / B200 / GB200
# datacenter) and SM 12.0 (RTX 50-series consumer). Anything >= 10 is
# Blackwell-family; Hopper is SM 9.x.
_BLACKWELL_MIN_CC_MAJOR = 10
def is_attention_projection(fqn: object) -> bool:
"""Return True when ``fqn``'s last component names an attention projection.
Defensive surface returns False (never raises) on non-string / empty /
null-byte input, matching the project model-detection policy
(``is_gemma4_model`` / ``is_known_vlm_base``).
"""
if not isinstance(fqn, str) or not fqn or "\x00" in fqn:
return False
if len(fqn) > 4096: # defensive cap, mirrors sibling detectors
return False
return fqn.rsplit(".", 1)[-1] in _ATTENTION_PROJ_NAMES
def is_blackwell_gpu() -> bool:
"""Return True when a Blackwell-family GPU (SM >= 10.0) is detected."""
try:
import torch
if not torch.cuda.is_available():
return False
major, _minor = torch.cuda.get_device_capability(0)
return major >= _BLACKWELL_MIN_CC_MAJOR
except (ImportError, RuntimeError, AssertionError):
return False
def _is_float8_linear(module: object) -> bool:
"""Class-name probe for torchao's ``Float8Linear`` (no torchao import)."""
return type(module).__name__ == "Float8Linear"
def _torchao_available() -> bool:
"""True when torchao is importable.
Checks ``sys.modules`` first so the ``sys.modules['torchao'] = None``
test-stub idiom (v0.27.0 MII policy) and injected fake modules both
resolve correctly ``find_spec`` raises ValueError on an in-module
fake whose ``__spec__`` is None.
"""
import importlib.util
import sys
if "torchao" in sys.modules:
return sys.modules["torchao"] is not None
try:
return importlib.util.find_spec("torchao") is not None
except (ModuleNotFoundError, ValueError):
return False
def validate_fp8_attention_compat(
*,
@ -127,17 +209,147 @@ def validate_unsloth_bnb_4bit_compat(
)
def apply_fp8_attention() -> None:
"""Live FP8-attention wiring — deferred to v0.53.1."""
raise NotImplementedError(
"fp8_attention live wiring deferred to v0.53.1. Schema accepts the "
"flag but no torchao FP8 attention swap is registered yet."
)
def apply_fp8_attention(model: object, *, recipe: str = "tensorwise") -> int:
"""Convert attention-projection linears to FP8 training modules.
Live since v0.71.21 (#141). Walks ``model.named_modules()``, collects
every ``nn.Linear`` whose FQN ends in an attention-projection name
(q/k/v/o + fused qkv variants), and converts the not-yet-converted ones
via torchao's ``convert_to_float8_training`` with an attention-only
``module_filter_fn``. Projections that are already ``Float8Linear``
(e.g. because the base ``quantization_aware='fp8'`` pass converted the
whole model) are counted but not re-wrapped.
Returns:
The number of attention projections that are FP8 after the call.
Raises:
TypeError: ``model`` is None or ``recipe`` is not a string.
ValueError: ``recipe`` is empty, or the model has no attention
projections at all (silent-no-op footgun).
RuntimeError: torchao is missing or the GPU is not Hopper+
(BETA hw gate friendly message, never a silent no-op).
"""
if model is None:
raise TypeError("model must not be None")
if isinstance(recipe, bool) or not isinstance(recipe, str):
raise TypeError(f"recipe must be a string, got {type(recipe).__name__}")
if not recipe:
raise ValueError("recipe must be a non-empty string")
if "\x00" in recipe:
raise ValueError("recipe must not contain null bytes")
# Probe the torchao float8 path SPECIFICALLY — fp8.is_fp8_available()
# also accepts transformer_engine, which cannot serve this converter
# (review fix: an NGC container with TE but no torchao must hit the
# friendly gate, not an uncaught ImportError below).
if not _torchao_available():
raise RuntimeError(
"fp8_attention requires torchao's float8 recipe "
"(pip install 'torchao>=0.5.0')."
)
from soup_cli.utils.fp8 import is_fp8_gpu_supported
if not is_fp8_gpu_supported():
raise RuntimeError(
"fp8_attention requires a Hopper+ GPU (H100/H200/B100/B200, "
"compute capability >= 9.0)."
)
import torch.nn as nn
already_converted: list[str] = []
pending: list[str] = []
for fqn, module in model.named_modules():
if not is_attention_projection(fqn):
continue
if _is_float8_linear(module):
already_converted.append(fqn)
elif isinstance(module, nn.Linear):
pending.append(fqn)
if not already_converted and not pending:
raise ValueError(
"fp8_attention found no attention projections on this model "
"(expected q_proj/k_proj/v_proj/o_proj or a fused qkv variant) "
"— refusing the silent no-op."
)
if pending:
try:
from torchao.float8 import convert_to_float8_training
from torchao.float8.config import Float8LinearConfig
except ImportError as exc:
raise RuntimeError(
"fp8_attention requires torchao's float8 recipe "
"(pip install 'torchao>=0.5.0')."
) from exc
pending_set = frozenset(pending)
config = Float8LinearConfig.from_recipe_name(recipe)
try:
convert_to_float8_training(
model,
config=config,
module_filter_fn=lambda _mod, name: name in pending_set,
)
except Exception as exc: # noqa: BLE001 — in-place mutation honesty
raise RuntimeError(
"fp8_attention conversion failed partway — the model may "
"be PARTIALLY converted; restart training without the "
f"flag ({type(exc).__name__}: {exc})"
) from exc
return len(already_converted) + len(pending)
def apply_nvfp4() -> None:
"""Live NVFP4 wiring — deferred to v0.53.1."""
raise NotImplementedError(
"NVFP4 live wiring deferred to v0.53.1. Schema accepts the flag "
"but no Blackwell-FP4 quant prep is registered yet."
def apply_nvfp4(model: object) -> int:
"""Quantise ``model`` with torchao's NVFP4 scheme (Blackwell-only).
Live since v0.71.21 (#141). Routes through the same
``torchao.quantization.NVFP4Config`` surface as the v0.53.1
``soup export --format torchao`` path, gated on a Blackwell GPU
(SM 10.0 datacenter / SM 12.0 consumer).
Returns:
The number of ``nn.Linear`` modules torchao targeted (advisory
count, taken before the in-place ``quantize_`` call).
Raises:
TypeError: ``model`` is None.
RuntimeError: non-Blackwell GPU, torchao missing, or the installed
torchao does not expose ``NVFP4Config``.
"""
if model is None:
raise TypeError("model must not be None")
if not is_blackwell_gpu():
raise RuntimeError(
"NVFP4 requires a Blackwell GPU (B100/B200/GB200 at SM 10.0 or "
"RTX 50-series at SM 12.0); no Blackwell device detected."
)
try:
from torchao import quantization as ao_q
except ImportError as exc:
raise RuntimeError(
"NVFP4 requires torchao (pip install 'torchao>=0.7.0')."
) from exc
if not hasattr(ao_q, "NVFP4Config") or not hasattr(ao_q, "quantize_"):
raise RuntimeError(
"torchao does not expose NVFP4Config / quantize_; upgrade "
"torchao (pip install -U torchao)."
)
import torch.nn as nn
linear_count = sum(
1 for _fqn, module in model.named_modules()
if isinstance(module, nn.Linear)
)
try:
ao_q.quantize_(model, ao_q.NVFP4Config())
except Exception as exc: # noqa: BLE001 — in-place mutation honesty
raise RuntimeError(
"NVFP4 quantisation failed partway — the model may be "
"PARTIALLY quantised; restart training without the flag "
f"({type(exc).__name__}: {exc})"
) from exc
return linear_count

View File

@ -1,4 +1,4 @@
"""Multi-turn agent rollout backends — v0.50.0 Part C.
"""Multi-turn agent rollout backends — v0.50.0 Part C / v0.71.21 #125 live.
Closed allowlist for multi-turn agent RL rollout backends shipped for
unsloth + axolotl parity:
@ -8,25 +8,64 @@ unsloth + axolotl parity:
- ``nemo_gym`` NVIDIA NeMo Gym (unsloth, axolotl)
- ``openenv`` Generic ``rollout_func`` protocol (unsloth, axolotl)
Schema-only in v0.50.0; live launcher wiring deferred to v0.50.1 (mirrors
v0.27.0 MII / v0.37.0 multipack / v0.41.0 LLaMA Pro / v0.45.0 plugins /
v0.46.0 deploy autopilot / v0.48.0 curriculum / v0.49.0 LongLoRA
stub-then-live pattern).
v0.71.21 (#125) lifts :func:`launch_rollout` to live:
- ``openenv`` runs fully live on CPU ``training.rollout_func`` names a
user-supplied ``module.path:function_name`` callable (mirrors the
v0.42.0 ``data.prompt_strategy`` design) that receives the seed prompts
and returns rollout rows (``{"prompt": ..., "answer"?: ...}``). The
rows replace the GRPO prompt dataset.
- ``art`` / ``ruler`` / ``nemo_gym`` are lazy-import gated: a friendly
ImportError names the pip package when it is missing; when present, an
honest "adapter not yet validated" RuntimeError fires instead of
shipping unrun integration code (mirrors the v0.71.20 TTS-codec /
onebitllms BETA-gate policy). ``_EXTERNAL_ROLLOUT_RUNNERS`` is the
injectable runner seam (mirrors ``cloud/modal._MODAL_SUBMIT_OVERRIDE``).
Security:
- Closed allowlist; arbitrary string at schema level rejected.
- ``_BACKEND_METADATA`` wrapped in MappingProxyType (matches v0.36.0
_REGISTRY policy).
- ``validate_rollout_backend`` rejects empty / null-byte / non-string /
oversize inputs.
oversize inputs; ``rollout_func`` is regex-validated (``module:fn``
shape, no null bytes / oversize) before any import is attempted.
- Rollout rows are capped at ``_MAX_ROLLOUT_ROWS`` and normalised to
``{prompt, answer?}``-only dicts so a misbehaving rollout callable
cannot smuggle arbitrary payloads into the training dataset.
"""
from __future__ import annotations
import importlib
import importlib.util
import re
import types
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Optional
_MAX_BACKEND_NAME_LEN = 32
_MAX_ROLLOUT_FUNC_LEN = 260
_MAX_ROLLOUT_ROWS = 100_000
_MAX_ROLLOUT_STEPS = 100_000
# ``module.path:function_name`` — mirrors v0.42.0 data.prompt_strategy.
_ROLLOUT_FUNC_RE = re.compile(
r"^[A-Za-z_][A-Za-z0-9_.]{0,127}:[A-Za-z_][A-Za-z0-9_]{0,127}$"
)
# Import names per external backend (pip name differs — see
# ``required_package`` on the spec).
_EXTERNAL_IMPORT_NAMES = types.MappingProxyType({
"art": "art",
"ruler": "ruler",
"nemo_gym": "nemo_gym",
})
# Injectable runner seam for the external backends — tests and advanced
# operators may install ``name -> callable(**kwargs) -> rows`` here.
# NOT a public API (mirrors cloud/modal._MODAL_SUBMIT_OVERRIDE).
_EXTERNAL_ROLLOUT_RUNNERS: dict[str, Callable[..., Any]] = {}
SUPPORTED_ROLLOUT_BACKENDS: frozenset[str] = frozenset({
"art",
@ -69,7 +108,7 @@ _BACKEND_METADATA = types.MappingProxyType({
name="openenv",
description="Generic OpenEnv rollout_func protocol",
required_package=None,
live_wired=False,
live_wired=True, # v0.71.21 #125 — fully live on CPU.
),
})
@ -116,15 +155,240 @@ def list_rollout_backends() -> tuple[str, ...]:
return tuple(sorted(SUPPORTED_ROLLOUT_BACKENDS))
def launch_rollout(name: str) -> None:
"""Live launcher for the rollout backend — deferred to v0.50.1.
def validate_rollout_func(value: object) -> Optional[str]:
"""Validate a ``module.path:function_name`` rollout-func spec.
Planned v0.50.1 signature:
``launch_rollout(name, *, prompts, model, reward_fn, max_steps)``.
``None`` passes through (field unset). Mirrors the v0.42.0
``validate_prompt_strategy`` policy with rollout-func-named errors.
"""
validate_rollout_backend(name)
raise NotImplementedError(
f"rollout_backend={name!r} live launcher deferred to v0.50.1. "
"The schema accepts the value but the actual rollout/optimization "
"loop is not yet wired."
if value is None:
return None
if not isinstance(value, str):
raise ValueError(
f"rollout_func must be a string, got {type(value).__name__}"
)
if not value:
raise ValueError("rollout_func must not be empty")
if "\x00" in value:
raise ValueError("rollout_func must not contain null bytes")
if len(value) > _MAX_ROLLOUT_FUNC_LEN:
raise ValueError(
f"rollout_func must be <= {_MAX_ROLLOUT_FUNC_LEN} chars"
)
if not _ROLLOUT_FUNC_RE.match(value):
raise ValueError(
"rollout_func must match 'module.path:function_name' "
"(letters / digits / underscore / dot only)."
)
return value
def resolve_rollout_func(spec: str) -> Callable[..., Any]:
"""Resolve a validated ``module:fn`` spec into the rollout callable.
Mirrors the v0.53.7 ``resolve_prompt_strategy`` runtime resolver.
Trusted-input policy applies: the spec names operator-controlled code
(same trust level as ``data.prompt_strategy`` / custom reward files).
"""
validated = validate_rollout_func(spec)
if validated is None:
raise ValueError("rollout_func must not be None")
module_path, _, fn_name = validated.partition(":")
try:
module = importlib.import_module(module_path)
except ImportError as exc:
raise ValueError(
f"rollout_func module {module_path!r} could not be imported: {exc}"
) from exc
if not hasattr(module, fn_name):
raise ValueError(
f"rollout_func {validated!r}: module {module_path!r} has no "
f"attribute {fn_name!r}"
)
fn = getattr(module, fn_name)
if not callable(fn):
raise ValueError(
f"rollout_func {validated!r}: resolved attribute is not callable "
f"({type(fn).__name__})"
)
return fn
def _spec_exists(name: str) -> bool:
"""``importlib.util.find_spec`` wrapper that never raises.
Mirrors the v0.71.14 ``kv_cache._spec_exists`` policy
``ModuleNotFoundError`` / ``ValueError`` from a half-installed parent
package count as "not available".
"""
try:
return importlib.util.find_spec(name) is not None
except (ModuleNotFoundError, ValueError):
return False
def _normalise_rollout_rows(raw: object, backend: str) -> tuple[dict, ...]:
"""Validate + normalise rollout output to ``{prompt, answer?}`` dicts.
Message-list prompts are deep-copied (the rollout callable must not be
able to mutate training rows through a retained alias immutability
policy); a non-string ``answer`` is rejected LOUDLY, matching the
prompt-side rejection (a silently dropped answer starves the accuracy
reward of its reference).
"""
if raw is None or isinstance(raw, (str, bytes)) or not hasattr(raw, "__iter__"):
raise ValueError(
f"rollout backend {backend!r} must return an iterable of row "
f"mappings, got {type(raw).__name__}"
)
rows: list[dict] = []
for index, row in enumerate(raw):
if len(rows) >= _MAX_ROLLOUT_ROWS:
raise ValueError(
f"rollout backend {backend!r} produced more than "
f"{_MAX_ROLLOUT_ROWS} rows"
)
if not isinstance(row, Mapping):
raise ValueError(
f"rollout row {index} must be a mapping, "
f"got {type(row).__name__}"
)
prompt = row.get("prompt")
if isinstance(prompt, str):
if not prompt:
raise ValueError(f"rollout row {index} has an empty prompt")
elif isinstance(prompt, list):
if not prompt:
raise ValueError(f"rollout row {index} has an empty prompt")
# Break the alias with the rollout callable's return value.
prompt = [
dict(message) if isinstance(message, Mapping) else message
for message in prompt
]
else:
raise ValueError(
f"rollout row {index} must carry a 'prompt' (str or message "
f"list), got {type(prompt).__name__}"
)
normalised: dict = {"prompt": prompt}
answer = row.get("answer")
if answer is not None:
if not isinstance(answer, str):
raise ValueError(
f"rollout row {index} 'answer' must be a string, "
f"got {type(answer).__name__}"
)
if answer:
normalised["answer"] = answer
rows.append(normalised)
if not rows:
raise ValueError(
f"rollout backend {backend!r} produced no rows — refusing to "
"train on an empty prompt set"
)
return tuple(rows)
@dataclass(frozen=True)
class RolloutResult:
"""Normalised rollout output: ``rows`` feed the GRPO prompt dataset."""
backend: str
rows: tuple[dict, ...]
def __post_init__(self) -> None:
object.__setattr__(
self, "backend", validate_rollout_backend(self.backend)
)
if not isinstance(self.rows, tuple):
raise TypeError("rows must be a tuple")
# Re-validate through the shared normaliser so a directly
# constructed RolloutResult cannot smuggle malformed rows.
object.__setattr__(
self, "rows", _normalise_rollout_rows(self.rows, self.backend)
)
def launch_rollout(
name: str,
*,
prompts: Optional[Sequence[Any]] = None,
rollout_func: Optional[str] = None,
model: Any = None,
tokenizer: Any = None,
reward_fn: Any = None,
max_steps: int = 64,
) -> RolloutResult:
"""Run the named rollout backend and return its normalised result.
Live since v0.71.21 (#125):
- ``openenv`` resolves ``rollout_func`` (``module:fn``) and calls it
with the seed ``prompts`` list. The callable returns an iterable of
``{"prompt": str | messages, "answer"?: str}`` rows.
- ``art`` / ``ruler`` / ``nemo_gym`` raise a friendly ImportError when
the backend package is missing, and an honest "not yet validated"
RuntimeError when it is present (lazy-import + advisory the live
integration adapters are dep-gated; inject a runner via the
``_EXTERNAL_ROLLOUT_RUNNERS`` seam to use them today).
``model`` / ``tokenizer`` / ``reward_fn`` are forwarded to injected
external runners (forward-compat per the v0.50.0 planned signature);
the openenv contract intentionally stays minimal (prompts only).
"""
normalised_name = validate_rollout_backend(name)
if isinstance(max_steps, bool) or not isinstance(max_steps, int):
raise TypeError(
f"max_steps must be an int, got {type(max_steps).__name__}"
)
if not 1 <= max_steps <= _MAX_ROLLOUT_STEPS:
raise ValueError(
f"max_steps must be in [1, {_MAX_ROLLOUT_STEPS}], got {max_steps}"
)
if prompts is not None and (
isinstance(prompts, (str, bytes)) or not isinstance(prompts, Sequence)
):
raise TypeError(
f"prompts must be a sequence, got {type(prompts).__name__}"
)
seed_prompts = list(prompts or [])
if normalised_name == "openenv":
if rollout_func is None:
raise ValueError(
"rollout_backend='openenv' requires training.rollout_func "
"('module.path:function_name')"
)
fn = resolve_rollout_func(rollout_func)
raw = fn(seed_prompts)
rows = _normalise_rollout_rows(raw, normalised_name)
return RolloutResult(backend=normalised_name, rows=rows)
# External backends: injected runner seam first, then lazy-import gate.
runner = _EXTERNAL_ROLLOUT_RUNNERS.get(normalised_name)
if runner is not None:
raw = runner(
prompts=seed_prompts,
model=model,
tokenizer=tokenizer,
reward_fn=reward_fn,
max_steps=max_steps,
)
rows = _normalise_rollout_rows(raw, normalised_name)
return RolloutResult(backend=normalised_name, rows=rows)
import_name = _EXTERNAL_IMPORT_NAMES[normalised_name]
spec = _BACKEND_METADATA[normalised_name]
if not _spec_exists(import_name):
raise ImportError(
f"rollout_backend={normalised_name!r} requires the "
f"{spec.required_package!r} package: "
f"pip install {spec.required_package}"
)
raise RuntimeError(
f"rollout_backend={normalised_name!r}: the {spec.required_package!r} "
"integration adapter is not yet validated on any environment "
"(BETA gate). Use rollout_backend='openenv' with "
"training.rollout_func, or inject a runner via "
"soup_cli.utils.agent_rollout._EXTERNAL_ROLLOUT_RUNNERS."
)

View File

@ -1,28 +1,77 @@
"""``soup apple-adapter`` — HF / PEFT ↔ MLX ↔ Apple FoundationModels (v0.68.0 Part D).
Schema-only release: live conversion (HF safetensors -> MLX npz / Apple
FoundationModels adapter blob) lands in v0.68.1. Reuses v0.60 Part B
signing for the optional ``--sign`` path.
v0.71.21 (#228) lifts ``convert_apple_adapter`` to live:
- ``hf-to-mlx`` reads a PEFT LoRA adapter (``adapter_model.safetensors``)
and writes an mlx-lm-shaped ``adapters.safetensors`` +
``adapter_config.json`` (with ``num_layers`` derived from the converted
keys mlx-lm's ``load_adapters`` reads both unconditionally). PEFT
stores ``...lora_A.weight`` as ``[r, in]`` / ``...lora_B.weight`` as
``[out, r]``; mlx-lm's ``LoRALinear`` computes ``(x @ lora_a) @ lora_b``
with ``lora_a [in, r]`` / ``lora_b [r, out]``, so both matrices
transpose on the way through. Pure numpy no mlx import needed, so the
conversion runs on any OS; loading the artifact into mlx-lm itself
requires Apple hardware (documented BETA gate). bf16 adapters are
upcast to float32 via the torch loader (numpy has no bf16).
- ``mlx-to-hf`` reverses the conversion (``adapters.safetensors`` or
legacy ``adapters.npz`` ``adapter_model.safetensors`` + a PEFT-style
``adapter_config.json``).
- ``hf-to-apple`` / ``mlx-to-apple`` raise a friendly upstream-gate
RuntimeError: Apple has not published a stable FoundationModels adapter
spec refusing to export wrong-shaped weights (per #228 fix path).
- ``sign=True`` reuses v0.60 Part B Merkle-root signing to emit a
``.soup-signature.json`` next to the converted adapter.
Public surface:
- ``SUPPORTED_ADAPTER_DIRECTIONS`` closed frozenset
- ``validate_direction(name)`` bool-first / null-byte / case-insensitive
- ``validate_source_adapter(path)`` cwd containment + directory check + symlink reject
- ``AppleAdapterPlan`` frozen dataclass
- ``build_apple_adapter_plan(...)`` factory
- ``convert_apple_adapter(plan)`` NotImplementedError stub w/ v0.68.1 marker
- ``AppleAdapterPlan`` frozen dataclass + ``build_apple_adapter_plan(...)``
- ``hf_key_to_mlx`` / ``mlx_key_to_hf`` LoRA key mapping
- ``convert_hf_to_mlx_arrays`` / ``convert_mlx_to_hf_arrays`` pure kernels
- ``convert_apple_adapter(plan)`` LIVE; returns ``ConversionReport``
"""
from __future__ import annotations
import json
import logging
import os
import re
import stat
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Optional
from soup_cli.utils.paths import is_under_cwd
from soup_cli.utils.paths import (
atomic_write_bytes,
atomic_write_text,
enforce_under_cwd_and_no_symlink,
is_under_cwd,
)
SUPPORTED_ADAPTER_DIRECTIONS: frozenset = frozenset(
if TYPE_CHECKING: # pragma: no cover
from collections.abc import Mapping
logger = logging.getLogger(__name__)
# PEFT LoRA weight keys: ``base_model.model.<path>.lora_A.weight`` (the
# leading prefix is optional — some exporters strip it).
_HF_LORA_KEY_RE = re.compile(
r"^(?:base_model\.model\.)?(?P<path>.+)\.(?P<matrix>lora_[AB])\.weight$"
)
# mlx-lm LoRA keys: ``<path>.lora_a`` / ``<path>.lora_b``.
_MLX_LORA_KEY_RE = re.compile(r"^(?P<path>.+)\.(?P<matrix>lora_[ab])$")
# Decoder-layer index inside an mlx key (``model.layers.N.…`` /
# GPT-2-style ``transformer.h.N.…``) — used to derive the ``num_layers``
# field mlx-lm's load_adapters reads.
_LAYER_INDEX_RE = re.compile(r"(?:^|\.)(?:layers|h)\.(\d+)\.")
# Adapters are small; a multi-GiB "adapter" is a red flag, not a use case.
# The cap also bounds DECOMPRESSED npz arrays (zip bomb defence).
_MAX_ADAPTER_FILE_BYTES = 4 * 1024**3 # 4 GiB
SUPPORTED_ADAPTER_DIRECTIONS: frozenset[str] = frozenset(
{"hf-to-mlx", "mlx-to-hf", "hf-to-apple", "mlx-to-apple"}
)
@ -30,6 +79,7 @@ _MAX_DIRECTION_LEN = 32
def validate_direction(name: object) -> str:
"""Canonicalise a conversion direction against the closed allowlist."""
if isinstance(name, bool):
raise TypeError("direction must not be bool")
if not isinstance(name, str):
@ -67,16 +117,16 @@ def validate_source_adapter(path: object) -> str:
)
if os.path.lexists(path):
try:
st = os.lstat(path)
link_stat = os.lstat(path)
except OSError as exc:
raise ValueError(
f"source_dir unreadable: {type(exc).__name__}"
) from exc
if stat.S_ISLNK(st.st_mode):
if stat.S_ISLNK(link_stat.st_mode):
raise ValueError(
"source_dir must not be a symlink (TOCTOU defence)"
)
if not stat.S_ISDIR(st.st_mode):
if not stat.S_ISDIR(link_stat.st_mode):
raise ValueError("source_dir must be a directory")
return os.path.realpath(path)
@ -95,6 +145,8 @@ def _validate_output_dir(path: object) -> str:
@dataclass(frozen=True)
class AppleAdapterPlan:
"""Resolved conversion plan (validated in ``__post_init__``)."""
source_dir: str
output_dir: str
direction: str
@ -117,6 +169,7 @@ def build_apple_adapter_plan(
direction: str,
sign: bool = False,
) -> AppleAdapterPlan:
"""Build a validated :class:`AppleAdapterPlan` from raw CLI inputs."""
return AppleAdapterPlan(
source_dir=source_dir,
output_dir=output_dir,
@ -125,12 +178,386 @@ def build_apple_adapter_plan(
)
def convert_apple_adapter(plan: AppleAdapterPlan) -> None:
"""Live conversion. Deferred to v0.68.1."""
@dataclass(frozen=True)
class ConversionReport:
"""Outcome of a live adapter conversion (v0.71.21 #228)."""
direction: str
output_dir: str
converted_keys: int
skipped_keys: tuple[str, ...]
signed: bool
def __post_init__(self) -> None:
validate_direction(self.direction)
if isinstance(self.converted_keys, bool) or not isinstance(
self.converted_keys, int
):
raise TypeError("converted_keys must be an int")
if self.converted_keys < 0:
raise ValueError("converted_keys must be >= 0")
if not isinstance(self.skipped_keys, tuple):
raise TypeError("skipped_keys must be a tuple")
if not isinstance(self.signed, bool):
raise TypeError("signed must be bool")
def hf_key_to_mlx(key: object) -> Optional[str]:
"""Map a PEFT LoRA key to its mlx-lm name (None for non-LoRA keys)."""
if not isinstance(key, str):
return None
match = _HF_LORA_KEY_RE.match(key)
if match is None:
return None
matrix = "lora_a" if match.group("matrix") == "lora_A" else "lora_b"
return f"{match.group('path')}.{matrix}"
def mlx_key_to_hf(key: object) -> Optional[str]:
"""Map an mlx-lm LoRA key back to its PEFT name (None for non-LoRA)."""
if not isinstance(key, str):
return None
match = _MLX_LORA_KEY_RE.match(key)
if match is None:
return None
matrix = "lora_A" if match.group("matrix") == "lora_a" else "lora_B"
return f"base_model.model.{match.group('path')}.{matrix}.weight"
def convert_hf_to_mlx_arrays(
arrays: "Mapping[str, Any]",
) -> tuple[dict[str, Any], tuple[str, ...]]:
"""Rename + transpose PEFT LoRA arrays into mlx-lm shape.
Returns ``(converted, skipped_keys)``. Non-LoRA keys (embeddings,
``modules_to_save`` etc.) are skipped mlx-lm adapters carry only the
LoRA matrices. Raises ``ValueError`` when no LoRA keys exist at all.
"""
converted: dict[str, Any] = {}
skipped: list[str] = []
for key, value in arrays.items():
mlx_key = hf_key_to_mlx(key)
if mlx_key is None:
skipped.append(key)
continue
converted[mlx_key] = value.T # [r, in] -> [in, r] / [out, r] -> [r, out]
if not converted:
raise ValueError(
"no LoRA keys found in the source adapter (expected "
"'...lora_A.weight' / '...lora_B.weight' PEFT keys)"
)
return converted, tuple(skipped)
def convert_mlx_to_hf_arrays(
arrays: "Mapping[str, Any]",
) -> tuple[dict[str, Any], tuple[str, ...]]:
"""Rename + transpose mlx-lm LoRA arrays back into PEFT shape."""
converted: dict[str, Any] = {}
skipped: list[str] = []
for key, value in arrays.items():
hf_key = mlx_key_to_hf(key)
if hf_key is None:
skipped.append(key)
continue
converted[hf_key] = value.T
if not converted:
raise ValueError(
"no LoRA keys found in the source adapter (expected "
"'...lora_a' / '...lora_b' mlx-lm keys)"
)
return converted, tuple(skipped)
def _read_adapter_file(path: str, field: str) -> str:
"""Per-file symlink rejection + size cap before any open (TOCTOU)."""
try:
file_stat = os.lstat(path)
except OSError as exc:
raise FileNotFoundError(
f"{field} not found: {os.path.basename(path)}"
) from exc
if stat.S_ISLNK(file_stat.st_mode):
raise ValueError(f"{field} must not be a symlink (TOCTOU defence)")
if not stat.S_ISREG(file_stat.st_mode):
raise ValueError(f"{field} must be a regular file")
if file_stat.st_size > _MAX_ADAPTER_FILE_BYTES:
raise ValueError(
f"{field} exceeds the {_MAX_ADAPTER_FILE_BYTES // 1024**3} GiB "
"adapter cap"
)
return path
def _load_safetensors_arrays(path: str, field: str) -> dict[str, Any]:
"""Load a safetensors file as numpy arrays (bf16 upcast via torch).
``safetensors.numpy.load_file`` raises ``SafetensorError`` (a direct
``Exception`` subclass) on bf16 tensors the overwhelmingly common
PEFT adapter dtype so the fallback decision catches broadly, then
the torch loader either succeeds (bf16 upcast) or proves the file is
genuinely corrupt (friendly ``ValueError``).
"""
_read_adapter_file(path, field)
try:
from safetensors.numpy import load_file as np_load_file
except ImportError as exc: # pragma: no cover — safetensors in [train]
raise ImportError(
"apple-adapter conversion requires safetensors "
"(pip install safetensors)"
) from exc
try:
return dict(np_load_file(path))
except Exception as np_exc: # noqa: BLE001 — SafetensorError is a bare Exception
# bf16 tensors are not representable in numpy — fall back to the
# torch loader and upcast to float32 (documented precision note).
try:
import torch # noqa: F401
from safetensors.torch import load_file as torch_load_file
except ImportError as exc:
raise ImportError(
"this adapter holds non-numpy dtypes (likely bf16); "
"converting it requires torch "
"(pip install 'soup-cli[train]')"
) from exc
try:
tensors = torch_load_file(path)
except Exception as exc: # noqa: BLE001 — corrupt file, not a dtype issue
raise ValueError(
f"{field} is not a valid safetensors file: "
f"{type(np_exc).__name__}"
) from exc
logger.warning(
"adapter %s holds non-numpy dtypes (likely bf16); upcasting "
"to float32 for conversion",
os.path.basename(path),
)
return {
key: tensor.float().numpy()
for key, tensor in tensors.items()
}
def _load_npz_arrays(path: str, field: str) -> dict[str, Any]:
"""Load a legacy mlx-lm ``adapters.npz`` as a plain dict of arrays.
The 4 GiB cap is re-applied to the DECOMPRESSED arrays the on-disk
cap in ``_read_adapter_file`` bounds only the compressed container
(zip bomb defence).
"""
_read_adapter_file(path, field)
import numpy as np
try:
with np.load(path, allow_pickle=False) as bundle:
arrays = {key: bundle[key] for key in bundle.files}
except Exception as exc: # noqa: BLE001 — BadZipFile etc. are bare Exception
raise ValueError(
f"{field} is not a valid npz file: {type(exc).__name__}"
) from exc
total_bytes = sum(int(getattr(arr, "nbytes", 0)) for arr in arrays.values())
if total_bytes > _MAX_ADAPTER_FILE_BYTES:
raise ValueError(
f"{field} decompresses past the "
f"{_MAX_ADAPTER_FILE_BYTES // 1024**3} GiB adapter cap"
)
return arrays
def _read_source_config(source_dir: str) -> dict[str, Any]:
"""Best-effort read of the source adapter_config.json (never raises)."""
config_path = os.path.join(source_dir, "adapter_config.json")
try:
_read_adapter_file(config_path, "adapter_config.json")
with open(config_path, encoding="utf-8") as handle:
loaded = json.load(handle)
return loaded if isinstance(loaded, dict) else {}
except (OSError, ValueError, json.JSONDecodeError):
return {}
def _save_safetensors(arrays: dict[str, Any], output_path: str) -> None:
"""Serialise arrays to safetensors bytes and write atomically.
The converted matrices are ``.T`` views (reversed strides)
safetensors serialises the raw base buffer, silently mangling
non-contiguous input, so every array is made C-contiguous first
(bug caught by the v0.71.21 review-wave round-trip assertions).
"""
import numpy as np
from safetensors.numpy import save as st_save
contiguous = {
key: np.ascontiguousarray(value) for key, value in arrays.items()
}
atomic_write_bytes(st_save(contiguous), output_path, field="output")
def _infer_rank(hf_arrays: dict[str, Any]) -> Optional[int]:
"""Infer the LoRA rank from any ``lora_A.weight`` matrix ([r, in])."""
for key, value in hf_arrays.items():
if key.endswith(".lora_A.weight") and getattr(value, "ndim", 0) == 2:
return int(value.shape[0])
return None
def _infer_num_layers(mlx_keys: "Mapping[str, Any]") -> Optional[int]:
"""Derive ``num_layers`` (max decoder index + 1) from converted keys.
mlx-lm's ``load_adapters`` reads ``config.num_layers`` unconditionally
before wiring LoRA layers, so the emitted adapter_config.json must
carry it whenever it is derivable.
"""
max_index = -1
for key in mlx_keys:
match = _LAYER_INDEX_RE.search(key)
if match is not None:
max_index = max(max_index, int(match.group(1)))
return max_index + 1 if max_index >= 0 else None
def _source_dropout(source_config: dict[str, Any]) -> float:
"""Carry the source PEFT ``lora_dropout`` through (default 0.0)."""
dropout = source_config.get("lora_dropout")
if (
isinstance(dropout, (int, float))
and not isinstance(dropout, bool)
and 0.0 <= float(dropout) < 1.0
):
return float(dropout)
return 0.0
def _convert_hf_to_mlx(plan: AppleAdapterPlan) -> tuple[int, tuple[str, ...]]:
source_file = os.path.join(plan.source_dir, "adapter_model.safetensors")
if not os.path.lexists(source_file):
if os.path.lexists(os.path.join(plan.source_dir, "adapter_model.bin")):
raise ValueError(
"adapter_model.bin (pickle) is not supported — re-save the "
"adapter as safetensors first (v0.57.0 policy)"
)
raise FileNotFoundError(
"adapter_model.safetensors not found in the source adapter"
)
arrays = _load_safetensors_arrays(source_file, "adapter_model.safetensors")
converted, skipped = convert_hf_to_mlx_arrays(arrays)
_save_safetensors(
converted, os.path.join(plan.output_dir, "adapters.safetensors")
)
source_config = _read_source_config(plan.source_dir)
rank = source_config.get("r")
if not isinstance(rank, int) or isinstance(rank, bool):
rank = _infer_rank(arrays)
alpha = source_config.get("lora_alpha")
mlx_config: dict[str, Any] = {
"fine_tune_type": "lora",
"soup_converted_from": "peft",
}
num_layers = _infer_num_layers(converted)
if num_layers is not None:
mlx_config["num_layers"] = num_layers
lora_parameters: dict[str, Any] = {
"dropout": _source_dropout(source_config)
}
if isinstance(rank, int) and rank > 0:
lora_parameters["rank"] = rank
if isinstance(alpha, (int, float)) and not isinstance(alpha, bool):
lora_parameters["scale"] = float(alpha) / float(rank)
mlx_config["lora_parameters"] = lora_parameters
atomic_write_text(
json.dumps(mlx_config, indent=2),
os.path.join(plan.output_dir, "adapter_config.json"),
field="output",
)
return len(converted), skipped
def _convert_mlx_to_hf(plan: AppleAdapterPlan) -> tuple[int, tuple[str, ...]]:
st_file = os.path.join(plan.source_dir, "adapters.safetensors")
npz_file = os.path.join(plan.source_dir, "adapters.npz")
if os.path.lexists(st_file):
arrays = _load_safetensors_arrays(st_file, "adapters.safetensors")
elif os.path.lexists(npz_file):
arrays = _load_npz_arrays(npz_file, "adapters.npz")
else:
raise FileNotFoundError(
"no adapters.npz or adapters.safetensors found in the source "
"adapter (expected an mlx-lm adapter directory)"
)
converted, skipped = convert_mlx_to_hf_arrays(arrays)
_save_safetensors(
converted, os.path.join(plan.output_dir, "adapter_model.safetensors")
)
source_config = _read_source_config(plan.source_dir)
rank = _infer_rank(converted)
target_modules = sorted({
key[: -len(".lora_A.weight")].rsplit(".", 1)[-1]
for key in converted
if key.endswith(".lora_A.weight")
})
hf_config: dict[str, Any] = {
"peft_type": "LORA",
"task_type": "CAUSAL_LM",
"soup_converted_from": "mlx",
"target_modules": target_modules,
}
if isinstance(rank, int) and rank > 0:
hf_config["r"] = rank
lora_params = source_config.get("lora_parameters")
scale = (
lora_params.get("scale") if isinstance(lora_params, dict) else None
)
if isinstance(scale, (int, float)) and not isinstance(scale, bool):
hf_config["lora_alpha"] = float(scale) * rank
atomic_write_text(
json.dumps(hf_config, indent=2),
os.path.join(plan.output_dir, "adapter_config.json"),
field="output",
)
return len(converted), skipped
def convert_apple_adapter(plan: AppleAdapterPlan) -> ConversionReport:
"""Run the live adapter conversion described by ``plan``.
Live since v0.71.21 (#228) for the ``hf-to-mlx`` / ``mlx-to-hf``
directions; the two ``*-to-apple`` directions raise a friendly
upstream-gate RuntimeError until Apple publishes a stable
FoundationModels adapter spec.
"""
if not isinstance(plan, AppleAdapterPlan):
raise TypeError("plan must be AppleAdapterPlan")
raise NotImplementedError(
"apple-adapter live conversion is deferred to v0.68.1"
if plan.direction in ("hf-to-apple", "mlx-to-apple"):
raise RuntimeError(
"the Apple FoundationModels adapter format has no stable public "
"spec yet — refusing to export wrong-shaped weights. Track "
"https://developer.apple.com/documentation/foundationmodels "
"for the published format."
)
enforce_under_cwd_and_no_symlink(plan.output_dir, "output_dir")
os.makedirs(plan.output_dir, exist_ok=True)
if plan.direction == "hf-to-mlx":
converted_count, skipped = _convert_hf_to_mlx(plan)
else: # mlx-to-hf — directions are a closed allowlist
converted_count, skipped = _convert_mlx_to_hf(plan)
signed = False
if plan.sign:
from soup_cli.utils.adapter_sign import sign_adapter
sign_adapter(plan.output_dir)
signed = True
return ConversionReport(
direction=plan.direction,
output_dir=plan.output_dir,
converted_keys=converted_count,
skipped_keys=skipped,
signed=signed,
)
@ -139,6 +566,11 @@ __all__ = [
"validate_direction",
"validate_source_adapter",
"AppleAdapterPlan",
"ConversionReport",
"build_apple_adapter_plan",
"hf_key_to_mlx",
"mlx_key_to_hf",
"convert_hf_to_mlx_arrays",
"convert_mlx_to_hf_arrays",
"convert_apple_adapter",
]

View File

@ -1,35 +1,83 @@
"""v0.44.0 Part D — `soup delinearize-llama4` weight reshape stub.
"""v0.44.0 Part D — `soup delinearize-llama4` weight reshape.
Llama 4 ships with linearised expert weights that some downstream backends
expect in 3-D form. This module declares the planned reshape; live runtime
deferred to v0.44.1 (mirrors the project's stub-then-live pattern).
Llama 4 ships its fused-MoE expert weights as parameters on
``Llama4TextExperts`` (``...feed_forward.experts.gate_up_proj`` /
``...experts.down_proj``). Checkpoints exported through tooling that
flattens parameters carry those tensors in *linearised* 2-D form
``[num_experts * dim_in, dim_out]``; HF transformers (and most downstream
backends) expect the 3-D ``[num_experts, dim_in, dim_out]`` layout.
v0.71.21 (#97) lifts the v0.44.0 planner to a live runtime:
:func:`run_delinearize` loads each ``.safetensors`` shard (torch loader
handles bf16), reshapes the fused expert tensors row-major to 3-D, passes
every other tensor through unchanged, and writes the shards to the target
directory (atomic write via the shared ``paths.atomic_write_bytes``
helper, mirroring the v0.71.14 FSDP-consolidate policy). JSON sidecars
(``config.json``, tokenizer files, the safetensors index) are copied so
the target stays a loadable checkpoint. Per-expert numbered keys
(``...experts.0.gate_proj.weight`` Mixtral-style, already unfused) are
intentionally NOT matched.
"""
from __future__ import annotations
import json
import os
import re
import stat
from dataclasses import dataclass
from typing import List, Tuple
from typing import Any, List, Optional, Tuple
from soup_cli.utils.paths import is_under_cwd
from soup_cli.utils.paths import (
atomic_write_bytes,
atomic_write_text,
enforce_under_cwd_and_no_symlink,
is_under_cwd,
)
# Restrict to canonical Llama 4 model id shape; reject crafted names.
_LLAMA4_RE = re.compile(r"(?i)(?:^|[^a-z0-9])llama-?4(?:[^a-z0-9]|$)")
# Fused-expert parameter keys on Llama4TextExperts. Numbered per-expert
# keys (``.experts.0.gate_proj``) are already unfused — never matched.
_EXPERT_FUSED_KEY_RE = re.compile(
r"\.experts\.(?:gate_up_proj|down_proj)(?:\.weight)?$"
)
# Mirrors the v0.71.14 fsdp_consolidate per-shard cap.
_MAX_WEIGHT_FILE_BYTES = 16 * 1024**3 # 16 GiB
_MAX_SIDECAR_BYTES = 256 * 1024**2 # 256 MiB (tokenizer.json scale)
_MAX_NUM_EXPERTS = 4096
@dataclass(frozen=True)
class DelinearizePlan:
"""Planned weights to reshape, source path, target path.
`weight_files` is a `tuple` for genuine immutability (matches the
project frozen-collection policy).
project frozen-collection policy). ``__post_init__`` re-runs the
containment checks so a directly constructed plan cannot bypass
``plan_delinearize`` and route writes outside cwd.
"""
source_dir: str
target_dir: str
weight_files: Tuple[str, ...]
def __post_init__(self) -> None:
for field_name in ("source_dir", "target_dir"):
value = getattr(self, field_name)
if not isinstance(value, str) or not value:
raise ValueError(f"{field_name} must be a non-empty str")
if "\x00" in value:
raise ValueError(f"{field_name} contains NUL byte")
if not is_under_cwd(value):
raise ValueError(
f"{field_name} is outside cwd: {os.path.basename(value)}"
)
if not isinstance(self.weight_files, tuple):
raise TypeError("weight_files must be a tuple")
def is_llama4_model(name: str) -> bool:
"""Return True iff `name` looks like a Llama 4 family model."""
@ -77,3 +125,246 @@ def plan_delinearize(source_dir: str, target_dir: str) -> DelinearizePlan:
target_dir=os.path.realpath(target_dir),
weight_files=tuple(files),
)
# --- v0.71.21 #97 — live runtime -------------------------------------------
@dataclass(frozen=True)
class DelinearizeResult:
"""Outcome of a live delinearize run (v0.71.21 #97)."""
source_dir: str
target_dir: str
files_written: Tuple[str, ...]
reshaped_keys: int
passthrough_keys: int
already_3d_keys: int
sidecars_copied: int = 0
def __post_init__(self) -> None:
if not isinstance(self.files_written, tuple):
raise TypeError("files_written must be a tuple")
for field_name in (
"reshaped_keys",
"passthrough_keys",
"already_3d_keys",
"sidecars_copied",
):
value = getattr(self, field_name)
if isinstance(value, bool) or not isinstance(value, int):
raise TypeError(f"{field_name} must be an int")
if value < 0:
raise ValueError(f"{field_name} must be >= 0")
def is_expert_weight_key(key: object) -> bool:
"""True for fused Llama-4 expert keys (``...experts.gate_up_proj``).
Defensive surface returns False (never raises) on non-string /
empty / null-byte input. Numbered per-expert keys are NOT matched
(those are already unfused, Mixtral-style).
"""
if not isinstance(key, str) or not key or "\x00" in key:
return False
return bool(_EXPERT_FUSED_KEY_RE.search(key))
def _validate_num_experts(value: object) -> int:
"""Bounds-check ``num_experts`` (bool-rejected, [1, 4096])."""
if isinstance(value, bool) or not isinstance(value, int):
raise TypeError(
f"num_experts must be an int, got {type(value).__name__}"
)
if not 1 <= value <= _MAX_NUM_EXPERTS:
raise ValueError(
f"num_experts must be in [1, {_MAX_NUM_EXPERTS}], got {value}"
)
return value
def read_num_experts(source_dir: str) -> Optional[int]:
"""Probe ``config.json`` for the expert count (None when absent).
Checks ``text_config.num_local_experts`` (HF Llama4TextConfig),
then top-level ``num_local_experts`` / ``num_experts``.
"""
config_path = os.path.join(source_dir, "config.json")
try:
file_stat = os.lstat(config_path)
if stat.S_ISLNK(file_stat.st_mode) or not stat.S_ISREG(file_stat.st_mode):
return None
if file_stat.st_size > _MAX_SIDECAR_BYTES:
return None
with open(config_path, encoding="utf-8") as handle:
config = json.load(handle)
except (OSError, ValueError, json.JSONDecodeError):
return None
if not isinstance(config, dict):
return None
candidates: List[Any] = []
text_config = config.get("text_config")
if isinstance(text_config, dict):
candidates.append(text_config.get("num_local_experts"))
candidates.append(text_config.get("num_experts"))
candidates.append(config.get("num_local_experts"))
candidates.append(config.get("num_experts"))
for candidate in candidates:
if (
isinstance(candidate, int)
and not isinstance(candidate, bool)
and 1 <= candidate <= _MAX_NUM_EXPERTS
):
return candidate
return None
def delinearize_tensor(tensor: Any, *, num_experts: int) -> Tuple[Any, str]:
"""Reshape a linearised 2-D expert tensor to 3-D.
``[num_experts * dim_in, dim_out]`` reshapes row-major to
``[num_experts, dim_in, dim_out]``; an already-3-D tensor passes
through unchanged (``status='already_3d'``).
Returns ``(tensor, status)`` with status ``'reshaped'`` or
``'already_3d'``.
"""
_validate_num_experts(num_experts)
ndim = getattr(tensor, "ndim", None)
if ndim == 3:
return tensor, "already_3d"
if ndim != 2:
raise ValueError(
f"expert tensor must be 2-D (linearised) or 3-D, got {ndim}-D"
)
rows = int(tensor.shape[0])
if rows % num_experts != 0:
raise ValueError(
f"expert tensor dim 0 ({rows}) is not divisible by "
f"num_experts={num_experts}"
)
return tensor.reshape(num_experts, rows // num_experts, tensor.shape[1]), "reshaped"
# JSON sidecars copied so the target stays a loadable checkpoint.
_SIDECAR_SUFFIXES = (".json",)
def _copy_json_sidecars(source_dir: str, target_dir: str) -> int:
"""Best-effort copy of top-level ``.json`` sidecars (config/tokenizer)."""
copied = 0
for entry in sorted(os.listdir(source_dir)):
if not entry.endswith(_SIDECAR_SUFFIXES):
continue
source_path = os.path.join(source_dir, entry)
try:
file_stat = os.lstat(source_path)
if stat.S_ISLNK(file_stat.st_mode) or not stat.S_ISREG(file_stat.st_mode):
continue
if file_stat.st_size > _MAX_SIDECAR_BYTES:
continue
with open(source_path, encoding="utf-8") as handle:
text = handle.read()
atomic_write_text(
text, os.path.join(target_dir, entry), field="target_dir",
)
copied += 1
except (OSError, ValueError, UnicodeDecodeError):
continue
return copied
def run_delinearize(
plan: DelinearizePlan,
*,
num_experts: Optional[int] = None,
) -> DelinearizeResult:
"""Run the live Llama-4 expert delinearization described by ``plan``.
Live since v0.71.21 (#97). ``num_experts`` defaults to the value read
from the source ``config.json``; an explicit argument wins. Raises a
friendly ``ValueError`` naming ``--num-experts`` when neither is
available, and names the offending key when a fused expert tensor's
leading dim is not divisible by the expert count.
"""
if not isinstance(plan, DelinearizePlan):
raise TypeError("plan must be a DelinearizePlan")
if num_experts is None:
num_experts = read_num_experts(plan.source_dir)
if num_experts is None:
raise ValueError(
"could not determine the expert count from config.json — "
"pass --num-experts explicitly"
)
_validate_num_experts(num_experts)
try:
import torch # noqa: F401
from safetensors.torch import load_file
from safetensors.torch import save as st_save
except ImportError as exc:
raise ImportError(
"delinearize-llama4 requires torch + safetensors "
"(pip install 'soup-cli[train]')"
) from exc
# Containment BEFORE makedirs — a directly constructed plan must not
# be able to create directories outside cwd (defence-in-depth on top
# of DelinearizePlan.__post_init__).
if not is_under_cwd(plan.target_dir):
raise ValueError(
f"target_dir is outside cwd: {os.path.basename(plan.target_dir)}"
)
os.makedirs(plan.target_dir, exist_ok=True)
reshaped = 0
passthrough = 0
already_3d = 0
written: List[str] = []
for name in plan.weight_files:
source_path = os.path.join(plan.source_dir, name)
enforce_under_cwd_and_no_symlink(source_path, "source file")
size = os.path.getsize(source_path)
if size > _MAX_WEIGHT_FILE_BYTES:
raise ValueError(
f"{name} exceeds the "
f"{_MAX_WEIGHT_FILE_BYTES // 1024**3} GiB per-file cap"
)
try:
tensors = load_file(source_path)
except Exception as exc: # noqa: BLE001 — SafetensorError is bare Exception
raise ValueError(
f"{name} is not a valid safetensors file: "
f"{type(exc).__name__}"
) from exc
out_tensors = {}
for key, tensor in tensors.items():
if is_expert_weight_key(key):
try:
out_tensor, status = delinearize_tensor(
tensor, num_experts=num_experts
)
except ValueError as exc:
raise ValueError(f"{key}: {exc}") from exc
out_tensors[key] = out_tensor
if status == "reshaped":
reshaped += 1
else:
already_3d += 1
else:
out_tensors[key] = tensor
passthrough += 1
target_path = os.path.join(plan.target_dir, name)
atomic_write_bytes(st_save(out_tensors), target_path, field="target_dir")
written.append(name)
sidecars_copied = _copy_json_sidecars(plan.source_dir, plan.target_dir)
return DelinearizeResult(
source_dir=plan.source_dir,
target_dir=plan.target_dir,
files_written=tuple(written),
reshaped_keys=reshaped,
passthrough_keys=passthrough,
already_3d_keys=already_3d,
sidecars_copied=sidecars_copied,
)

View File

@ -19,9 +19,20 @@ Security:
from __future__ import annotations
import contextlib
import logging
from collections.abc import Collection, Iterator
from typing import Any
logger = logging.getLogger(__name__)
# Backends where vLLM is genuinely available (mirrors v0.30.0 vllm.py policy).
_VLLM_SUPPORTED_BACKENDS: frozenset[str] = frozenset({"transformers", "unsloth"})
# vLLM grew ``AsyncEngineArgs.enable_sleep_mode`` + ``engine.sleep()`` /
# ``wake_up()`` in the 0.7 line.
_MIN_VLLM_SLEEP_VERSION: tuple[int, ...] = (0, 7)
def validate_long_context_grpo_compat(
*,
@ -82,15 +93,139 @@ def validate_vllm_sleep_mode_compat(*, backend: str) -> None:
)
def apply_vllm_sleep_mode(engine_args: object) -> None: # noqa: ARG001
"""Live wiring for vLLM sleep mode — deferred to v0.50.1.
def _parse_version_tuple(value: str) -> tuple[int, ...]:
"""Parse leading-int dot-chunks: ``"0.7.0.dev0"`` -> ``(0, 7, 0)``.
The schema flag is accepted in v0.50.0 but the actual
``AsyncEngineArgs.enable_sleep_mode=True`` plumbing is wired in
v0.50.1 (mirrors v0.27.0 MII stub-then-live pattern).
Mirrors the v0.40.1 ``_version_ge`` leading-int policy a chunk that
does not start with digits terminates the parse, so pre-release
suffixes never break the comparison.
"""
raise NotImplementedError(
"vllm_sleep_mode live wiring deferred to v0.50.1; "
"schema flag accepts the value but rollout-time sleep/wake "
"is not yet wired into the vLLM engine factory."
)
if not isinstance(value, str):
return ()
parts: list[int] = []
for chunk in value.split("."):
digits = ""
for char in chunk:
if char.isdigit():
digits += char
else:
break
if not digits:
break
parts.append(int(digits))
return tuple(parts)
def _installed_vllm_version() -> tuple[int, ...] | None:
"""Return the installed vLLM version tuple, or None when missing."""
try:
import vllm
except ImportError:
return None
return _parse_version_tuple(getattr(vllm, "__version__", "") or "")
def vllm_supports_sleep_mode() -> bool:
"""True when vLLM is installed and >= 0.7 (sleep/wake API present)."""
version = _installed_vllm_version()
return version is not None and version >= _MIN_VLLM_SLEEP_VERSION
def apply_vllm_sleep_mode(engine_args: object) -> object:
"""Enable vLLM sleep mode on an ``AsyncEngineArgs``-like object.
Live since v0.71.21 (#124): sets ``enable_sleep_mode=True`` so the
engine can be put on standby between rollouts (``engine.sleep()`` /
``engine.wake_up()`` see :func:`vllm_sleep_cycle`).
Raises:
TypeError: ``engine_args`` is None.
RuntimeError: vLLM is missing or older than 0.7 (friendly,
actionable message BETA gate; vLLM is not installable on
every platform, e.g. Windows).
"""
if engine_args is None:
raise TypeError("engine_args must not be None")
version = _installed_vllm_version()
if version is None:
raise RuntimeError(
"vllm_sleep_mode requires vLLM >= 0.7 but vLLM is not "
"installed (pip install 'soup-cli[serve-fast]')."
)
if version < _MIN_VLLM_SLEEP_VERSION:
found = ".".join(str(part) for part in version) or "unknown"
raise RuntimeError(
f"vllm_sleep_mode requires vLLM >= 0.7; found {found}. "
"Upgrade with: pip install -U vllm"
)
engine_args.enable_sleep_mode = True # type: ignore[attr-defined]
return engine_args
@contextlib.contextmanager
def vllm_sleep_cycle(engine: object, *, level: int = 1) -> Iterator[object]:
"""Put a vLLM engine to sleep for the duration of the ``with`` body.
Live since v0.71.21 (#124). Intended to wrap the optimisation step
between rollouts: ``with vllm_sleep_cycle(engine): optimizer_step()``.
The engine is woken in a ``finally`` block so an exception inside the
body never leaves it asleep. An engine without the sleep/wake API logs
a WARNING and the body runs unchanged memory instrumentation must
never crash a training loop (matches the v0.71.11 callback policy).
``level`` is vLLM's sleep level: 1 = offload weights to CPU, 2 = also
discard the KV cache. Bool-rejected and bounded to [1, 2].
"""
if isinstance(level, bool) or not isinstance(level, int):
raise TypeError(f"level must be an int, got {type(level).__name__}")
if not 1 <= level <= 2:
raise ValueError(f"level must be 1 or 2, got {level}")
sleep_fn = getattr(engine, "sleep", None)
wake_fn = getattr(engine, "wake_up", None)
if not callable(sleep_fn) or not callable(wake_fn):
logger.warning(
"vllm_sleep_cycle: engine %s does not expose sleep()/wake_up() "
"— running the body without engine standby.",
type(engine).__name__,
)
yield engine
return
sleep_fn(level=level)
try:
yield engine
finally:
wake_fn()
def maybe_enable_trl_sleep_mode(
grpo_kwargs: "dict[str, Any]",
grpo_param_names: "Collection[str]",
console: object = None,
) -> bool:
"""Thread ``vllm_sleep_mode`` into TRL's GRPOConfig when it has the hook.
Live since v0.71.21 (#124). TRL manages its own vLLM engine inside
``GRPOTrainer`` (server / colocate modes), so Soup cannot reach the
engine to call ``sleep()``/``wake_up()`` directly. When the installed
TRL exposes a ``vllm_enable_sleep_mode`` GRPOConfig parameter it is
set here; otherwise a yellow advisory explains that Soup's own vLLM
engine factory (``create_vllm_engine(sleep_mode=True)``) honors the
flag for serve / custom rollout loops.
Returns True when the TRL kwarg was set.
"""
if "vllm_enable_sleep_mode" in grpo_param_names:
grpo_kwargs["vllm_enable_sleep_mode"] = True
if console is not None:
console.print(
"[green]vLLM sleep mode enabled via TRL GRPOConfig[/]"
)
return True
if console is not None:
console.print(
"[yellow]vllm_sleep_mode: the installed TRL GRPOConfig does not "
"expose a sleep-mode option; Soup's vLLM engine factory honors "
"sleep_mode=True for serve / custom rollout loops. Upgrade trl "
"when it ships the hook.[/]"
)
return False

View File

@ -84,6 +84,31 @@ def apply_v028_speed_memory(
"compatible linears", style="yellow",
)
# --- FP8 attention (v0.71.21 #141) ---------------------------------------
# Key added only when the flag is set — keeps the legacy 3-key dict
# contract on the no-features path (test_part_c exact-equality).
if getattr(tcfg, "fp8_attention", False):
recipe = getattr(tcfg, "fp8_recipe", "tensorwise")
try:
from soup_cli.utils.advanced_precision import apply_fp8_attention
converted = apply_fp8_attention(model, recipe=recipe)
applied["fp8_attention"] = True
_say(f"FP8 attention enabled ({converted} projections)")
except (RuntimeError, ValueError, TypeError) as exc:
applied["fp8_attention"] = False
_say(f"FP8 attention: {exc}", style="yellow")
# --- NVFP4 (v0.71.21 #141 — Blackwell-only) ------------------------------
if getattr(tcfg, "nvfp4", False):
try:
from soup_cli.utils.advanced_precision import apply_nvfp4
targeted = apply_nvfp4(model)
applied["nvfp4"] = True
_say(f"NVFP4 quantisation applied ({targeted} linears)")
except (RuntimeError, ValueError, TypeError) as exc:
applied["nvfp4"] = False
_say(f"NVFP4: {exc}", style="yellow")
# --- Kernel auto-compose -------------------------------------------------
if getattr(tcfg, "kernel_auto_compose", False):
picked_name = _bench_and_pick_kernel(

View File

@ -41,6 +41,7 @@ def create_vllm_engine(
num_speculative_tokens: int = 5,
enable_prefix_caching: bool = False,
quantization: Optional[str] = None,
sleep_mode: bool = False,
):
"""Create a vLLM AsyncLLMEngine for serving.
@ -54,6 +55,8 @@ def create_vllm_engine(
dtype: Data type for model weights.
enable_prefix_caching: Enable vLLM's automatic prefix cache — big
win for RAG / agent workloads with shared system prompts.
sleep_mode: Enable vLLM sleep/standby support (v0.71.21 #124 —
requires vLLM >= 0.7; raises a friendly RuntimeError otherwise).
Returns:
(engine, engine_model_name) tuple.
@ -111,6 +114,13 @@ def create_vllm_engine(
engine_args.speculative_model = speculative_model
engine_args.num_speculative_tokens = num_speculative_tokens
# v0.71.21 #124 — vLLM sleep/standby support (engine.sleep()/wake_up()).
# apply_vllm_sleep_mode raises a friendly RuntimeError on vLLM < 0.7.
if sleep_mode:
from soup_cli.utils.grpo_long_context import apply_vllm_sleep_mode
apply_vllm_sleep_mode(engine_args)
engine = AsyncLLMEngine.from_engine_args(engine_args)
return engine, engine_model_name

View File

@ -277,6 +277,7 @@ def test_plan_delinearize_outside_cwd_rejected(tmp_path, monkeypatch):
def test_cli_delinearize_llama4(tmp_path, monkeypatch):
"""v0.71.21 #97 — the runtime is live; --plan-only keeps the old flow."""
monkeypatch.chdir(tmp_path)
src = tmp_path / "model"
src.mkdir()
@ -290,7 +291,7 @@ def test_cli_delinearize_llama4(tmp_path, monkeypatch):
str(src),
"--target",
str(target),
"--yes",
"--plan-only",
],
)
assert result.exit_code == 0, (result.output, repr(result.exception))

View File

@ -143,13 +143,15 @@ def test_vllm_sleep_mode_null_byte_backend():
# ---------------------------------------------------------------------------
# Deferred live wiring
# Live wiring (lifted in v0.71.21 #124 — vLLM-present gate)
# ---------------------------------------------------------------------------
def test_apply_vllm_sleep_mode_deferred():
"""tdd-guide MEDIUM fix: assert v0.50.1 marker in error msg."""
with pytest.raises(NotImplementedError, match=r"v0\.50\.1"):
def test_apply_vllm_sleep_mode_live_gated():
"""v0.71.21 #124 lifted the stub — friendly gate when vLLM is absent
or too old; sets enable_sleep_mode=True on modern vLLM (covered in
test_v07121.py with a fake vllm module)."""
with pytest.raises(RuntimeError, match=r"vLLM >= 0\.7"):
apply_vllm_sleep_mode(object())

View File

@ -80,12 +80,17 @@ def test_required_rollout_package_unknown_raises():
required_rollout_package("trlx")
@pytest.mark.parametrize("name", ["art", "ruler", "nemo_gym", "openenv"])
def test_rollout_backend_spec_not_live_wired(name):
"""tdd-guide LOW fix: exercise live_wired field."""
@pytest.mark.parametrize("name", ["art", "ruler", "nemo_gym"])
def test_external_rollout_backends_stay_gated(name):
"""v0.71.21 #125 — external backends stay lazy-import gated."""
assert get_rollout_backend_spec(name).live_wired is False
def test_openenv_rollout_backend_live_wired():
"""v0.71.21 #125 — openenv runs fully live via rollout_func."""
assert get_rollout_backend_spec("openenv").live_wired is True
def test_get_rollout_backend_spec_frozen():
spec = get_rollout_backend_spec("art")
assert isinstance(spec, RolloutBackendSpec)
@ -108,12 +113,21 @@ def test_metadata_immutable():
agent_rollout._BACKEND_METADATA["evil"] = None # type: ignore[index]
@pytest.mark.parametrize("name", ["art", "ruler", "nemo_gym", "openenv"])
def test_launch_rollout_deferred(name):
with pytest.raises(NotImplementedError, match="v0.50.1"):
@pytest.mark.parametrize("name", ["art", "ruler", "nemo_gym"])
def test_launch_rollout_external_gated(name):
"""v0.71.21 #125 lifted the stub — external backends now raise a
friendly ImportError (package missing) or an honest BETA RuntimeError
(package present, adapter not yet validated)."""
with pytest.raises((ImportError, RuntimeError), match="rollout"):
launch_rollout(name)
def test_launch_rollout_openenv_requires_func():
"""v0.71.21 #125 — openenv is live and requires rollout_func."""
with pytest.raises(ValueError, match="rollout_func"):
launch_rollout("openenv")
def test_launch_rollout_unknown_validation_first():
with pytest.raises(ValueError, match="not supported"):
launch_rollout("trlx")

View File

@ -412,11 +412,12 @@ class TestFP8Attention:
"training: {fp8_attention: true}\n"
)
def test_apply_deferred(self):
def test_apply_live_gated(self):
"""v0.71.21 #141 lifted the stub — now a friendly hw/dep gate."""
from soup_cli.utils.advanced_precision import apply_fp8_attention
with pytest.raises(NotImplementedError, match="v0.53.1"):
apply_fp8_attention()
with pytest.raises((RuntimeError, ValueError), match="(?i)torchao|hopper"):
apply_fp8_attention(object())
class TestNVFP4:
@ -495,11 +496,12 @@ class TestNVFP4:
"training: {nvfp4: true}\n"
)
def test_apply_deferred(self):
def test_apply_live_gated(self):
"""v0.71.21 #141 lifted the stub — now a friendly Blackwell gate."""
from soup_cli.utils.advanced_precision import apply_nvfp4
with pytest.raises(NotImplementedError, match="v0.53.1"):
apply_nvfp4()
with pytest.raises(RuntimeError, match="Blackwell"):
apply_nvfp4(object())
class TestUnslothBNB4Bit:

View File

@ -178,10 +178,13 @@ class TestAppleAdapterPlan:
)
class TestConvertDeferred:
def test_raises_v068_1(
class TestConvertLive:
def test_missing_weights_friendly(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
"""v0.71.21 #228 lifted the stub — live conversion now runs; a
source dir without weights surfaces a friendly FileNotFoundError
(full happy-path round-trip coverage lives in test_v07121.py)."""
from soup_cli.utils.apple_adapter import (
build_apple_adapter_plan,
convert_apple_adapter,
@ -197,7 +200,7 @@ class TestConvertDeferred:
direction="hf-to-mlx",
sign=False,
)
with pytest.raises(NotImplementedError, match="v0.68.1"):
with pytest.raises(FileNotFoundError, match="adapter_model"):
convert_apple_adapter(plan)
def test_non_plan_rejected(self) -> None:
@ -263,9 +266,11 @@ class TestCli:
)
assert result.exit_code == 2
def test_live_exits_3(
def test_live_missing_weights_exits_2(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
"""v0.71.21 #228 — the runner is live; a weight-less source dir is
a validation failure (exit 2), not the old deferred exit 3."""
from soup_cli.cli import app
monkeypatch.chdir(tmp_path)
@ -284,7 +289,7 @@ class TestCli:
"out",
],
)
assert result.exit_code == 3, (result.output, repr(result.exception))
assert result.exit_code == 2, (result.output, repr(result.exception))
class TestSourceWiring:

2268
tests/test_v07121.py Normal file

File diff suppressed because it is too large Load Diff