From ed5fc3a8b3d7724debdf8b289134f397faa5d5fd Mon Sep 17 00:00:00 2001 From: Alpamys Date: Wed, 10 Jun 2026 16:29:52 +0500 Subject: [PATCH] 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. --- CHANGELOG.md | 62 + README.md | 29 +- docs/backends-and-ops.md | 4 +- docs/commands.md | 3 +- docs/performance-and-quantization.md | 9 +- docs/training.md | 7 +- pyproject.toml | 2 +- src/soup_cli/__init__.py | 2 +- src/soup_cli/commands/apple_adapter.py | 30 +- src/soup_cli/commands/delinearize_llama4.py | 79 +- src/soup_cli/config/schema.py | 47 +- src/soup_cli/trainer/grpo.py | 40 + src/soup_cli/utils/advanced_precision.py | 240 +- src/soup_cli/utils/agent_rollout.py | 296 ++- src/soup_cli/utils/apple_adapter.py | 462 +++- src/soup_cli/utils/delinearize_llama4.py | 305 ++- src/soup_cli/utils/grpo_long_context.py | 155 +- src/soup_cli/utils/v028_features.py | 25 + src/soup_cli/utils/vllm.py | 10 + tests/test_v0440_part_d.py | 3 +- tests/test_v0500_part_b.py | 10 +- tests/test_v0500_part_c.py | 26 +- tests/test_v0530.py | 14 +- tests/test_v0680_part_d.py | 15 +- tests/test_v07121.py | 2268 +++++++++++++++++++ 25 files changed, 4008 insertions(+), 135 deletions(-) create mode 100644 tests/test_v07121.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c5b45d7..6270219 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/README.md b/README.md index 0e74f73..5c53bbd 100644 --- a/README.md +++ b/README.md @@ -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 --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 --target ` + 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) · [GitHub Releases](https://github.com/MakazhanAlpamys/Soup/releases). diff --git a/docs/backends-and-ops.md b/docs/backends-and-ops.md index 7087a86..437ea28 100644 --- a/docs/backends-and-ops.md +++ b/docs/backends-and-ops.md @@ -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 diff --git a/docs/commands.md b/docs/commands.md index 0a99b00..a5f2694 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -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 --eval [--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 --teacher --student --strategy sft|preference|kl [--provider ollama|anthropic|vllm] [--base-url ] [--temperature F] [--max-rows N] Distill prompt-heavy traces via a live teacher (v0.71.13) soup compile-tools --eval [--optimizer textgrad|gepa] [--plan-only] TextGrad / GEPA tool-schema optimiser — live (v0.71.13; pip install 'soup-cli[compile]') -soup apple-adapter --direction hf-to-mlx|mlx-to-hf|hf-to-apple|mlx-to-apple --output [--sign] HF / MLX / Apple FoundationModels adapter conversion (v0.68.0) +soup apple-adapter --direction hf-to-mlx|mlx-to-hf|hf-to-apple|mlx-to-apple --output [--sign] [--plan-only] PEFT LoRA <-> mlx-lm adapter conversion — live (v0.71.21; *-to-apple upstream-gated exit 3) soup local-rl init --db Create personal-LLM flywheel SQLite schema (v0.68.0) soup local-rl status --db Print interactions / thumbs-up / thumbs-down counters soup local-rl record --db --prompt --response --thumb up|down Append thumbs record diff --git a/docs/performance-and-quantization.md b/docs/performance-and-quantization.md index 6186d83..6f8c957 100644 --- a/docs/performance-and-quantization.md +++ b/docs/performance-and-quantization.md @@ -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. diff --git a/docs/training.md b/docs/training.md index e1fe2ee..66cca2a 100644 --- a/docs/training.md +++ b/docs/training.md @@ -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 diff --git a/pyproject.toml b/pyproject.toml index 9688a44..3dc78b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/src/soup_cli/__init__.py b/src/soup_cli/__init__.py index e745f81..b9f617f 100644 --- a/src/soup_cli/__init__.py +++ b/src/soup_cli/__init__.py @@ -1,3 +1,3 @@ """Soup CLI — Fine-tune and post-train LLMs in one command.""" -__version__ = "0.71.20" +__version__ = "0.71.21" diff --git a/src/soup_cli/commands/apple_adapter.py b/src/soup_cli/commands/apple_adapter.py index c31ae00..5862ec6 100644 --- a/src/soup_cli/commands/apple_adapter.py +++ b/src/soup_cli/commands/apple_adapter.py @@ -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", + ) + ) diff --git a/src/soup_cli/commands/delinearize_llama4.py b/src/soup_cli/commands/delinearize_llama4.py index 859718c..cb1a7d1 100644 --- a/src/soup_cli/commands/delinearize_llama4.py +++ b/src/soup_cli/commands/delinearize_llama4.py @@ -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", + ) + ) diff --git a/src/soup_cli/config/schema.py b/src/soup_cli/config/schema.py index e6398dd..fafaf54 100644 --- a/src/soup_cli/config/schema.py +++ b/src/soup_cli/config/schema.py @@ -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") diff --git a/src/soup_cli/trainer/grpo.py b/src/soup_cli/trainer/grpo.py index d31f02a..6ec109d 100644 --- a/src/soup_cli/trainer/grpo.py +++ b/src/soup_cli/trainer/grpo.py @@ -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 diff --git a/src/soup_cli/utils/advanced_precision.py b/src/soup_cli/utils/advanced_precision.py index abc1f38..cc86e20 100644 --- a/src/soup_cli/utils/advanced_precision.py +++ b/src/soup_cli/utils/advanced_precision.py @@ -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 diff --git a/src/soup_cli/utils/agent_rollout.py b/src/soup_cli/utils/agent_rollout.py index 7639f65..a9c7bfc 100644 --- a/src/soup_cli/utils/agent_rollout.py +++ b/src/soup_cli/utils/agent_rollout.py @@ -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." ) diff --git a/src/soup_cli/utils/apple_adapter.py b/src/soup_cli/utils/apple_adapter.py index 40dd811..1c8c265 100644 --- a/src/soup_cli/utils/apple_adapter.py +++ b/src/soup_cli/utils/apple_adapter.py @@ -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..lora_A.weight`` (the +# leading prefix is optional — some exporters strip it). +_HF_LORA_KEY_RE = re.compile( + r"^(?:base_model\.model\.)?(?P.+)\.(?Plora_[AB])\.weight$" +) +# mlx-lm LoRA keys: ``.lora_a`` / ``.lora_b``. +_MLX_LORA_KEY_RE = re.compile(r"^(?P.+)\.(?Plora_[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", ] diff --git a/src/soup_cli/utils/delinearize_llama4.py b/src/soup_cli/utils/delinearize_llama4.py index 68cfe25..a2b304e 100644 --- a/src/soup_cli/utils/delinearize_llama4.py +++ b/src/soup_cli/utils/delinearize_llama4.py @@ -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, + ) diff --git a/src/soup_cli/utils/grpo_long_context.py b/src/soup_cli/utils/grpo_long_context.py index 00f7a8a..c24cd00 100644 --- a/src/soup_cli/utils/grpo_long_context.py +++ b/src/soup_cli/utils/grpo_long_context.py @@ -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 diff --git a/src/soup_cli/utils/v028_features.py b/src/soup_cli/utils/v028_features.py index e272220..1ee2748 100644 --- a/src/soup_cli/utils/v028_features.py +++ b/src/soup_cli/utils/v028_features.py @@ -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( diff --git a/src/soup_cli/utils/vllm.py b/src/soup_cli/utils/vllm.py index 19f7f73..87b8db9 100644 --- a/src/soup_cli/utils/vllm.py +++ b/src/soup_cli/utils/vllm.py @@ -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 diff --git a/tests/test_v0440_part_d.py b/tests/test_v0440_part_d.py index d1fc6b4..bbe3217 100644 --- a/tests/test_v0440_part_d.py +++ b/tests/test_v0440_part_d.py @@ -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)) diff --git a/tests/test_v0500_part_b.py b/tests/test_v0500_part_b.py index d67c6af..761d8ad 100644 --- a/tests/test_v0500_part_b.py +++ b/tests/test_v0500_part_b.py @@ -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()) diff --git a/tests/test_v0500_part_c.py b/tests/test_v0500_part_c.py index f95363d..1592dca 100644 --- a/tests/test_v0500_part_c.py +++ b/tests/test_v0500_part_c.py @@ -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") diff --git a/tests/test_v0530.py b/tests/test_v0530.py index 2645ac2..f0507f7 100644 --- a/tests/test_v0530.py +++ b/tests/test_v0530.py @@ -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: diff --git a/tests/test_v0680_part_d.py b/tests/test_v0680_part_d.py index b8613f8..0e70652 100644 --- a/tests/test_v0680_part_d.py +++ b/tests/test_v0680_part_d.py @@ -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: diff --git a/tests/test_v07121.py b/tests/test_v07121.py new file mode 100644 index 0000000..79ba584 --- /dev/null +++ b/tests/test_v07121.py @@ -0,0 +1,2268 @@ +"""v0.71.21 "Precision & rollout lift" tests. + +Closes #141 (live fp8_attention + NVFP4), #124 (vLLM sleep mode), +#125 (multi-turn agent rollout launchers), #228 (apple-adapter converter), +#97 (delinearize-llama4 runtime). + +Hardware notes: fp8_attention / nvfp4 / vllm-sleep are BETA hw-gated — the +gates + dispatch are exercised here via monkeypatched capability probes and +fake torchao / vllm modules (sys.modules swap via ``monkeypatch.setitem``, +the established test_v0531_142 pattern). The apple-adapter conversion and +the delinearize runtime are fully CPU-validatable and run on real arrays. +""" + +from __future__ import annotations + +import itertools +import json +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +from typer.testing import CliRunner + +_SRC = Path(__file__).resolve().parent.parent / "src" / "soup_cli" + +runner = CliRunner() + + +def _torch_or_skip(): + return pytest.importorskip("torch") + + +# --------------------------------------------------------------------------- +# #141 — attention-projection detection +# --------------------------------------------------------------------------- + + +class TestIsAttentionProjection: + def test_canonical_projections(self): + from soup_cli.utils.advanced_precision import is_attention_projection + + for name in ("q_proj", "k_proj", "v_proj", "o_proj"): + assert is_attention_projection(f"model.layers.0.self_attn.{name}") + + def test_fused_variants(self): + from soup_cli.utils.advanced_precision import is_attention_projection + + assert is_attention_projection("transformer.h.0.attn.c_attn") + assert is_attention_projection("model.layers.3.self_attn.qkv_proj") + assert is_attention_projection( + "transformer.layers.1.attention.query_key_value" + ) + + def test_mlp_projections_rejected(self): + from soup_cli.utils.advanced_precision import is_attention_projection + + assert not is_attention_projection("model.layers.0.mlp.gate_proj") + assert not is_attention_projection("model.layers.0.mlp.down_proj") + assert not is_attention_projection("lm_head") + + def test_defensive_inputs(self): + from soup_cli.utils.advanced_precision import is_attention_projection + + assert not is_attention_projection("") + assert not is_attention_projection(None) # type: ignore[arg-type] + assert not is_attention_projection(123) # type: ignore[arg-type] + assert not is_attention_projection("a\x00q_proj") + + def test_substring_not_enough(self): + """Last-component match only — 'my_q_proj_extra' is not q_proj.""" + from soup_cli.utils.advanced_precision import is_attention_projection + + assert not is_attention_projection("model.layers.0.my_q_proj_extra") + + +class TestIsBlackwellGpu: + def test_false_without_cuda(self, monkeypatch): + torch = _torch_or_skip() + from soup_cli.utils.advanced_precision import is_blackwell_gpu + + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + assert is_blackwell_gpu() is False + + @pytest.mark.parametrize( + ("capability", "expected"), + [((8, 6), False), ((9, 0), False), ((10, 0), True), ((12, 0), True)], + ) + def test_capability_matrix(self, monkeypatch, capability, expected): + torch = _torch_or_skip() + from soup_cli.utils.advanced_precision import is_blackwell_gpu + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr( + torch.cuda, "get_device_capability", lambda _i=0: capability + ) + assert is_blackwell_gpu() is expected + + +# --------------------------------------------------------------------------- +# #141 — apply_fp8_attention (live) +# --------------------------------------------------------------------------- + + +def _tiny_attn_model(with_attention: bool = True): + """A tiny real-torch model with q/k/v/o projections + an MLP linear.""" + torch = _torch_or_skip() + nn = torch.nn + + class _Attn(nn.Module): + def __init__(self): + super().__init__() + self.q_proj = nn.Linear(8, 8) + self.k_proj = nn.Linear(8, 8) + self.v_proj = nn.Linear(8, 8) + self.o_proj = nn.Linear(8, 8) + + class _Mlp(nn.Module): + def __init__(self): + super().__init__() + self.gate_proj = nn.Linear(8, 8) + + class _Block(nn.Module): + def __init__(self): + super().__init__() + if with_attention: + self.self_attn = _Attn() + self.mlp = _Mlp() + + class _Model(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([_Block()]) + + return _Model() + + +def _install_fake_torchao_float8(monkeypatch, record): + """Inject a fake ``torchao.float8`` that performs a real module swap.""" + torch = _torch_or_skip() + nn = torch.nn + + class _F8(nn.Linear): + pass + + _F8.__name__ = "Float8Linear" + + class _Cfg: + @classmethod + def from_recipe_name(cls, name): + record["recipe"] = name + return cls() + + def _convert(model, config=None, module_filter_fn=None): + record["config"] = config + swapped = [] + modules = dict(model.named_modules()) + for fqn, mod in list(modules.items()): + if not isinstance(mod, nn.Linear): + continue + if type(mod).__name__ == "Float8Linear": + continue + if module_filter_fn is not None and not module_filter_fn(mod, fqn): + continue + parent_fqn, _, child = fqn.rpartition(".") + parent = modules[parent_fqn] if parent_fqn else model + new = _F8(mod.in_features, mod.out_features, bias=mod.bias is not None) + setattr(parent, child, new) + swapped.append(fqn) + record["swapped"] = swapped + return model + + fake_cfg_mod = types.ModuleType("torchao.float8.config") + fake_cfg_mod.Float8LinearConfig = _Cfg + fake_f8 = types.ModuleType("torchao.float8") + fake_f8.convert_to_float8_training = _convert + fake_f8.config = fake_cfg_mod + fake_root = types.ModuleType("torchao") + fake_root.float8 = fake_f8 + monkeypatch.setitem(sys.modules, "torchao", fake_root) + monkeypatch.setitem(sys.modules, "torchao.float8", fake_f8) + monkeypatch.setitem(sys.modules, "torchao.float8.config", fake_cfg_mod) + return _F8 + + +def _enable_fp8_gates(monkeypatch): + # The torchao probe checks sys.modules first (v0.27.0 None-stub + # idiom), so a bare fake root is enough for the gate; conversion + # tests install the full fake via _install_fake_torchao_float8. + if "torchao" not in sys.modules: + monkeypatch.setitem(sys.modules, "torchao", types.ModuleType("torchao")) + monkeypatch.setattr("soup_cli.utils.fp8.is_fp8_gpu_supported", lambda: True) + + +class TestApplyFp8Attention: + def test_none_model_rejected(self): + from soup_cli.utils.advanced_precision import apply_fp8_attention + + with pytest.raises(TypeError, match="model"): + apply_fp8_attention(None) + + def test_no_torchao_friendly_gate(self, monkeypatch): + from soup_cli.utils.advanced_precision import apply_fp8_attention + + # v0.27.0 None-stub idiom — the probe must treat this as absent. + monkeypatch.setitem(sys.modules, "torchao", None) + with pytest.raises(RuntimeError, match="torchao"): + apply_fp8_attention(_tiny_attn_model()) + + def test_non_hopper_friendly_gate(self, monkeypatch): + from soup_cli.utils.advanced_precision import apply_fp8_attention + + monkeypatch.setitem( + sys.modules, "torchao", types.ModuleType("torchao") + ) + monkeypatch.setattr( + "soup_cli.utils.fp8.is_fp8_gpu_supported", lambda: False + ) + with pytest.raises(RuntimeError, match="Hopper"): + apply_fp8_attention(_tiny_attn_model()) + + @pytest.mark.parametrize("recipe", [True, "", None, 123, "row\x00wise"]) + def test_bad_recipe_rejected(self, monkeypatch, recipe): + from soup_cli.utils.advanced_precision import apply_fp8_attention + + _enable_fp8_gates(monkeypatch) + with pytest.raises((TypeError, ValueError), match="recipe"): + apply_fp8_attention(_tiny_attn_model(), recipe=recipe) + + def test_converts_only_attention(self, monkeypatch): + from soup_cli.utils.advanced_precision import apply_fp8_attention + + record: dict = {} + _install_fake_torchao_float8(monkeypatch, record) + _enable_fp8_gates(monkeypatch) + model = _tiny_attn_model() + count = apply_fp8_attention(model, recipe="rowwise") + assert count == 4 + assert record["recipe"] == "rowwise" + block = model.layers[0] + assert type(block.self_attn.q_proj).__name__ == "Float8Linear" + assert type(block.self_attn.o_proj).__name__ == "Float8Linear" + # MLP linears untouched. + assert type(block.mlp.gate_proj).__name__ == "Linear" + + def test_already_converted_counted_not_reswapped(self, monkeypatch): + from soup_cli.utils.advanced_precision import apply_fp8_attention + + record: dict = {} + f8_cls = _install_fake_torchao_float8(monkeypatch, record) + _enable_fp8_gates(monkeypatch) + model = _tiny_attn_model() + # Pre-convert q_proj manually. + model.layers[0].self_attn.q_proj = f8_cls(8, 8) + count = apply_fp8_attention(model) + assert count == 4 # all four attention projections are now Float8 + assert "self_attn.q_proj" not in " ".join(record["swapped"]) + assert len(record["swapped"]) == 3 + + def test_no_attention_projections_value_error(self, monkeypatch): + from soup_cli.utils.advanced_precision import apply_fp8_attention + + record: dict = {} + _install_fake_torchao_float8(monkeypatch, record) + _enable_fp8_gates(monkeypatch) + with pytest.raises(ValueError, match="no attention projections"): + apply_fp8_attention(_tiny_attn_model(with_attention=False)) + + +# --------------------------------------------------------------------------- +# #141 — apply_nvfp4 (live, Blackwell-gated) +# --------------------------------------------------------------------------- + + +def _install_fake_torchao_quant(monkeypatch, record, *, with_nvfp4=True): + fake_q = types.ModuleType("torchao.quantization") + + class _NVFP4Config: + pass + + def _quantize(model, config): + record["model"] = model + record["config"] = config + + fake_q.quantize_ = _quantize + if with_nvfp4: + fake_q.NVFP4Config = _NVFP4Config + fake_root = types.ModuleType("torchao") + fake_root.quantization = fake_q + monkeypatch.setitem(sys.modules, "torchao", fake_root) + monkeypatch.setitem(sys.modules, "torchao.quantization", fake_q) + + +class TestApplyNvfp4: + def test_none_model_rejected(self): + from soup_cli.utils.advanced_precision import apply_nvfp4 + + with pytest.raises(TypeError, match="model"): + apply_nvfp4(None) + + def test_non_blackwell_friendly_gate(self, monkeypatch): + from soup_cli.utils.advanced_precision import apply_nvfp4 + + monkeypatch.setattr( + "soup_cli.utils.advanced_precision.is_blackwell_gpu", lambda: False + ) + with pytest.raises(RuntimeError, match="Blackwell"): + apply_nvfp4(_tiny_attn_model()) + + def test_missing_torchao_friendly(self, monkeypatch): + from soup_cli.utils.advanced_precision import apply_nvfp4 + + monkeypatch.setattr( + "soup_cli.utils.advanced_precision.is_blackwell_gpu", lambda: True + ) + monkeypatch.setitem(sys.modules, "torchao", None) + with pytest.raises(RuntimeError, match="torchao"): + apply_nvfp4(_tiny_attn_model()) + + def test_old_torchao_missing_nvfp4config(self, monkeypatch): + from soup_cli.utils.advanced_precision import apply_nvfp4 + + record: dict = {} + _install_fake_torchao_quant(monkeypatch, record, with_nvfp4=False) + monkeypatch.setattr( + "soup_cli.utils.advanced_precision.is_blackwell_gpu", lambda: True + ) + with pytest.raises(RuntimeError, match="NVFP4Config"): + apply_nvfp4(_tiny_attn_model()) + + def test_happy_path_quantizes(self, monkeypatch): + from soup_cli.utils.advanced_precision import apply_nvfp4 + + record: dict = {} + _install_fake_torchao_quant(monkeypatch, record) + monkeypatch.setattr( + "soup_cli.utils.advanced_precision.is_blackwell_gpu", lambda: True + ) + model = _tiny_attn_model() + count = apply_nvfp4(model) + assert count == 5 # 4 attention + 1 mlp linear + assert record["model"] is model + assert type(record["config"]).__name__ == "_NVFP4Config" + + +# --------------------------------------------------------------------------- +# #141 — apply_v028_speed_memory wiring +# --------------------------------------------------------------------------- + + +class TestV028PrecisionWiring: + def _tcfg(self, **kwargs): + base = { + "use_cut_ce": False, + "quantization_aware": False, + "kernel_auto_compose": False, + } + base.update(kwargs) + return SimpleNamespace(**base) + + def test_no_features_dict_unchanged(self): + """Back-compat regression: 3-key exact dict on the no-features path.""" + from soup_cli.utils.v028_features import apply_v028_speed_memory + + result = apply_v028_speed_memory( + model=object(), tcfg=self._tcfg(), base_model="x/y", + ) + assert result == { + "cut_ce": False, "fp8": False, "kernel_auto_compose": False, + } + + def test_fp8_attention_gate_failure_degrades(self, monkeypatch): + from soup_cli.utils.v028_features import apply_v028_speed_memory + + def _gate(model, **kwargs): + raise RuntimeError("no Hopper") + + # Hermetic: force the gate (review fix — the bare call passed only + # because the host lacks torchao/Hopper). + monkeypatch.setattr( + "soup_cli.utils.advanced_precision.apply_fp8_attention", _gate + ) + result = apply_v028_speed_memory( + model=object(), + tcfg=self._tcfg(quantization_aware="fp8", fp8_attention=True), + base_model="x/y", + ) + assert result["fp8_attention"] is False + + def test_fp8_attention_applied(self, monkeypatch): + from soup_cli.utils.v028_features import apply_v028_speed_memory + + monkeypatch.setattr( + "soup_cli.utils.advanced_precision.apply_fp8_attention", + lambda model, recipe="tensorwise": 4, + ) + result = apply_v028_speed_memory( + model=object(), + tcfg=self._tcfg(quantization_aware="fp8", fp8_attention=True), + base_model="x/y", + ) + assert result["fp8_attention"] is True + + def test_nvfp4_gate_failure_degrades(self, monkeypatch): + from soup_cli.utils.v028_features import apply_v028_speed_memory + + def _gate(model, **kwargs): + raise RuntimeError("no Blackwell") + + monkeypatch.setattr( + "soup_cli.utils.advanced_precision.apply_nvfp4", _gate + ) + result = apply_v028_speed_memory( + model=object(), tcfg=self._tcfg(nvfp4=True), base_model="x/y", + ) + assert result["nvfp4"] is False + + def test_nvfp4_applied(self, monkeypatch): + from soup_cli.utils.v028_features import apply_v028_speed_memory + + monkeypatch.setattr( + "soup_cli.utils.advanced_precision.apply_nvfp4", lambda model: 2, + ) + result = apply_v028_speed_memory( + model=object(), tcfg=self._tcfg(nvfp4=True), base_model="x/y", + ) + assert result["nvfp4"] is True + + +# --------------------------------------------------------------------------- +# #124 — vLLM sleep mode (live) +# --------------------------------------------------------------------------- + + +def _install_fake_vllm(monkeypatch, version="0.8.1"): + fake = types.ModuleType("vllm") + fake.__version__ = version + monkeypatch.setitem(sys.modules, "vllm", fake) + return fake + + +class TestParseVersionTuple: + @pytest.mark.parametrize( + ("raw", "expected"), + [ + ("0.8.1", (0, 8, 1)), + ("0.7.0.dev0", (0, 7, 0)), + ("0.6", (0, 6)), + ("garbage", ()), + ("", ()), + ], + ) + def test_matrix(self, raw, expected): + from soup_cli.utils.grpo_long_context import _parse_version_tuple + + assert _parse_version_tuple(raw) == expected + + +class TestVllmSupportsSleepMode: + def test_false_when_vllm_missing(self, monkeypatch): + from soup_cli.utils.grpo_long_context import vllm_supports_sleep_mode + + monkeypatch.setitem(sys.modules, "vllm", None) + assert vllm_supports_sleep_mode() is False + + def test_true_on_modern_vllm(self, monkeypatch): + from soup_cli.utils.grpo_long_context import vllm_supports_sleep_mode + + _install_fake_vllm(monkeypatch, "0.8.1") + assert vllm_supports_sleep_mode() is True + + def test_false_on_old_vllm(self, monkeypatch): + from soup_cli.utils.grpo_long_context import vllm_supports_sleep_mode + + _install_fake_vllm(monkeypatch, "0.6.2") + assert vllm_supports_sleep_mode() is False + + +class TestApplyVllmSleepMode: + def test_none_rejected(self): + from soup_cli.utils.grpo_long_context import apply_vllm_sleep_mode + + with pytest.raises(TypeError, match="engine_args"): + apply_vllm_sleep_mode(None) + + def test_missing_vllm_friendly(self, monkeypatch): + from soup_cli.utils.grpo_long_context import apply_vllm_sleep_mode + + monkeypatch.setitem(sys.modules, "vllm", None) + with pytest.raises(RuntimeError, match=r"vLLM >= 0\.7"): + apply_vllm_sleep_mode(SimpleNamespace()) + + def test_old_vllm_friendly(self, monkeypatch): + from soup_cli.utils.grpo_long_context import apply_vllm_sleep_mode + + _install_fake_vllm(monkeypatch, "0.6.2") + with pytest.raises(RuntimeError, match=r"0\.6\.2"): + apply_vllm_sleep_mode(SimpleNamespace()) + + def test_happy_sets_enable_sleep_mode(self, monkeypatch): + from soup_cli.utils.grpo_long_context import apply_vllm_sleep_mode + + _install_fake_vllm(monkeypatch, "0.8.1") + args = SimpleNamespace() + out = apply_vllm_sleep_mode(args) + assert out is args + assert args.enable_sleep_mode is True + + +class TestVllmSleepCycle: + def test_sleep_then_wake(self): + from soup_cli.utils.grpo_long_context import vllm_sleep_cycle + + calls: list = [] + + class _Engine: + def sleep(self, level=1): + calls.append(("sleep", level)) + + def wake_up(self): + calls.append(("wake_up",)) + + with vllm_sleep_cycle(_Engine()): + calls.append(("body",)) + assert calls == [("sleep", 1), ("body",), ("wake_up",)] + + def test_wakes_even_on_body_exception(self): + from soup_cli.utils.grpo_long_context import vllm_sleep_cycle + + calls: list = [] + + class _Engine: + def sleep(self, level=1): + calls.append("sleep") + + def wake_up(self): + calls.append("wake_up") + + with pytest.raises(RuntimeError, match="boom"): + with vllm_sleep_cycle(_Engine()): + raise RuntimeError("boom") + assert calls == ["sleep", "wake_up"] + + def test_engine_without_sleep_warns_not_crashes(self, caplog): + import logging + + from soup_cli.utils.grpo_long_context import vllm_sleep_cycle + + with caplog.at_level(logging.WARNING, logger="soup_cli.utils.grpo_long_context"): + with vllm_sleep_cycle(object()): + pass + assert any("sleep" in rec.message for rec in caplog.records) + + +class TestMaybeEnableTrlSleepMode: + def test_sets_kwarg_when_trl_exposes_it(self): + from soup_cli.utils.grpo_long_context import maybe_enable_trl_sleep_mode + + kwargs: dict = {} + ok = maybe_enable_trl_sleep_mode( + kwargs, ("output_dir", "vllm_enable_sleep_mode"), None, + ) + assert ok is True + assert kwargs["vllm_enable_sleep_mode"] is True + + def test_advisory_when_trl_lacks_hook(self): + import io + + from rich.console import Console + + from soup_cli.utils.grpo_long_context import maybe_enable_trl_sleep_mode + + buf = io.StringIO() + kwargs: dict = {} + ok = maybe_enable_trl_sleep_mode( + kwargs, ("output_dir",), Console(file=buf), + ) + assert ok is False + assert kwargs == {} + assert "sleep" in buf.getvalue().lower() + + +class TestSleepModeSourceWiring: + def test_create_vllm_engine_accepts_sleep_mode(self): + source = (_SRC / "utils" / "vllm.py").read_text(encoding="utf-8") + assert "sleep_mode" in source + assert "apply_vllm_sleep_mode" in source + + def test_grpo_trainer_wires_sleep_mode(self): + source = (_SRC / "trainer" / "grpo.py").read_text(encoding="utf-8") + assert "maybe_enable_trl_sleep_mode" in source + + +# --------------------------------------------------------------------------- +# #125 — rollout launchers +# --------------------------------------------------------------------------- + +_MODULE_COUNTER = itertools.count() + + +def _write_rollout_module(tmp_path, monkeypatch, body: str) -> str: + """Write an importable rollout module; return the 'module:fn' spec.""" + name = f"soup_test_rollout_{next(_MODULE_COUNTER)}" + (tmp_path / f"{name}.py").write_text(body, encoding="utf-8") + monkeypatch.syspath_prepend(str(tmp_path)) + return f"{name}:rollout" + + +class TestValidateRolloutFunc: + def test_none_passthrough(self): + from soup_cli.utils.agent_rollout import validate_rollout_func + + assert validate_rollout_func(None) is None + + def test_happy(self): + from soup_cli.utils.agent_rollout import validate_rollout_func + + assert validate_rollout_func("my_mod.sub:my_fn") == "my_mod.sub:my_fn" + + @pytest.mark.parametrize( + "bad", + ["", "no-colon", "mod:", ":fn", "mod:fn:extra", "a b:c", "mod\x00:fn", 7], + ) + def test_rejection_matrix(self, bad): + from soup_cli.utils.agent_rollout import validate_rollout_func + + with pytest.raises(ValueError, match="rollout_func"): + validate_rollout_func(bad) + + def test_oversize_rejected(self): + from soup_cli.utils.agent_rollout import validate_rollout_func + + with pytest.raises(ValueError, match="rollout_func"): + validate_rollout_func("a" * 300 + ":fn") + + +class TestResolveRolloutFunc: + def test_resolves_callable(self, tmp_path, monkeypatch): + from soup_cli.utils.agent_rollout import resolve_rollout_func + + spec = _write_rollout_module( + tmp_path, monkeypatch, "def rollout(prompts):\n return []\n" + ) + fn = resolve_rollout_func(spec) + assert callable(fn) + + def test_missing_module_friendly(self): + from soup_cli.utils.agent_rollout import resolve_rollout_func + + with pytest.raises(ValueError, match="could not be imported"): + resolve_rollout_func("definitely_not_a_module_xyz:fn") + + def test_missing_attr_friendly(self, tmp_path, monkeypatch): + from soup_cli.utils.agent_rollout import resolve_rollout_func + + spec = _write_rollout_module( + tmp_path, monkeypatch, "def rollout(prompts):\n return []\n" + ) + module_name = spec.split(":")[0] + with pytest.raises(ValueError, match="no attribute"): + resolve_rollout_func(f"{module_name}:nope") + + def test_non_callable_rejected(self, tmp_path, monkeypatch): + from soup_cli.utils.agent_rollout import resolve_rollout_func + + spec = _write_rollout_module(tmp_path, monkeypatch, "rollout = 42\n") + with pytest.raises(ValueError, match="not callable"): + resolve_rollout_func(spec) + + +class TestRolloutResult: + def test_frozen(self): + import dataclasses + + from soup_cli.utils.agent_rollout import RolloutResult + + result = RolloutResult(backend="openenv", rows=({"prompt": "x"},)) + with pytest.raises(dataclasses.FrozenInstanceError): + result.backend = "art" # type: ignore[misc] + + def test_rows_must_be_tuple(self): + from soup_cli.utils.agent_rollout import RolloutResult + + with pytest.raises(TypeError, match="tuple"): + RolloutResult(backend="openenv", rows=[{"prompt": "x"}]) # type: ignore[arg-type] + + def test_unknown_backend_rejected(self): + from soup_cli.utils.agent_rollout import RolloutResult + + with pytest.raises(ValueError, match="not supported"): + RolloutResult(backend="trlx", rows=({"prompt": "x"},)) + + def test_row_without_prompt_rejected(self): + from soup_cli.utils.agent_rollout import RolloutResult + + with pytest.raises(ValueError, match="prompt"): + RolloutResult(backend="openenv", rows=({"answer": "y"},)) + + def test_row_cap(self): + from soup_cli.utils.agent_rollout import ( + _MAX_ROLLOUT_ROWS, + RolloutResult, + ) + + rows = tuple({"prompt": "x"} for _ in range(_MAX_ROLLOUT_ROWS + 1)) + with pytest.raises(ValueError, match="rows"): + RolloutResult(backend="openenv", rows=rows) + + +class TestLaunchRolloutOpenenv: + def test_happy_path(self, tmp_path, monkeypatch): + from soup_cli.utils.agent_rollout import launch_rollout + + spec = _write_rollout_module( + tmp_path, + monkeypatch, + "def rollout(prompts):\n" + " return [\n" + " {'prompt': p, 'answer': 'a'} for p in prompts\n" + " ] + [{'prompt': 'extra'}]\n", + ) + result = launch_rollout( + "openenv", prompts=["hello", "world"], rollout_func=spec, + ) + assert result.backend == "openenv" + assert len(result.rows) == 3 + assert result.rows[0]["prompt"] == "hello" + assert result.rows[0]["answer"] == "a" + assert "answer" not in result.rows[2] + + def test_message_list_prompts_preserved(self, tmp_path, monkeypatch): + from soup_cli.utils.agent_rollout import launch_rollout + + spec = _write_rollout_module( + tmp_path, + monkeypatch, + "def rollout(prompts):\n" + " return [{'prompt': [{'role': 'user', 'content': 'hi'}]}]\n", + ) + result = launch_rollout("openenv", prompts=[], rollout_func=spec) + assert result.rows[0]["prompt"] == [{"role": "user", "content": "hi"}] + + def test_openenv_requires_rollout_func(self): + from soup_cli.utils.agent_rollout import launch_rollout + + with pytest.raises(ValueError, match="rollout_func"): + launch_rollout("openenv", prompts=["x"]) + + def test_empty_rows_rejected(self, tmp_path, monkeypatch): + from soup_cli.utils.agent_rollout import launch_rollout + + spec = _write_rollout_module( + tmp_path, monkeypatch, "def rollout(prompts):\n return []\n" + ) + with pytest.raises(ValueError, match="no rows"): + launch_rollout("openenv", prompts=["x"], rollout_func=spec) + + def test_non_iterable_rows_rejected(self, tmp_path, monkeypatch): + from soup_cli.utils.agent_rollout import launch_rollout + + spec = _write_rollout_module( + tmp_path, monkeypatch, "def rollout(prompts):\n return 42\n" + ) + with pytest.raises(ValueError, match="iterable"): + launch_rollout("openenv", prompts=["x"], rollout_func=spec) + + def test_non_mapping_row_rejected(self, tmp_path, monkeypatch): + from soup_cli.utils.agent_rollout import launch_rollout + + spec = _write_rollout_module( + tmp_path, monkeypatch, "def rollout(prompts):\n return ['x']\n" + ) + with pytest.raises(ValueError, match="mapping"): + launch_rollout("openenv", prompts=["x"], rollout_func=spec) + + @pytest.mark.parametrize("bad_steps", [True, 0, -3, "many"]) + def test_max_steps_validation(self, tmp_path, monkeypatch, bad_steps): + from soup_cli.utils.agent_rollout import launch_rollout + + spec = _write_rollout_module( + tmp_path, + monkeypatch, + "def rollout(prompts):\n return [{'prompt': 'x'}]\n", + ) + with pytest.raises((TypeError, ValueError), match="max_steps"): + launch_rollout( + "openenv", prompts=["x"], rollout_func=spec, + max_steps=bad_steps, + ) + + def test_prompts_string_rejected(self, tmp_path, monkeypatch): + from soup_cli.utils.agent_rollout import launch_rollout + + spec = _write_rollout_module( + tmp_path, + monkeypatch, + "def rollout(prompts):\n return [{'prompt': 'x'}]\n", + ) + with pytest.raises(TypeError, match="prompts"): + launch_rollout("openenv", prompts="hi", rollout_func=spec) + + +class TestLaunchRolloutExternal: + @pytest.mark.parametrize( + ("name", "pkg"), + [("art", "openpipe-art"), ("ruler", "ruler-eval"), ("nemo_gym", "nemo-gym")], + ) + def test_missing_package_friendly(self, monkeypatch, name, pkg): + from soup_cli.utils import agent_rollout + + monkeypatch.setattr(agent_rollout, "_spec_exists", lambda _n: False) + with pytest.raises(ImportError, match=pkg): + agent_rollout.launch_rollout(name, prompts=["x"]) + + def test_present_package_honest_gate(self, monkeypatch): + from soup_cli.utils import agent_rollout + + monkeypatch.setattr(agent_rollout, "_spec_exists", lambda _n: True) + with pytest.raises(RuntimeError, match="openenv"): + agent_rollout.launch_rollout("art", prompts=["x"]) + + def test_runner_override_seam(self, monkeypatch): + from soup_cli.utils import agent_rollout + + def _fake_runner(**kwargs): + return [{"prompt": "from-art", "answer": "ok"}] + + monkeypatch.setitem( + agent_rollout._EXTERNAL_ROLLOUT_RUNNERS, "art", _fake_runner + ) + result = agent_rollout.launch_rollout("art", prompts=["x"]) + assert result.backend == "art" + assert result.rows[0]["prompt"] == "from-art" + + def test_unknown_backend_still_validated_first(self): + from soup_cli.utils.agent_rollout import launch_rollout + + with pytest.raises(ValueError, match="not supported"): + launch_rollout("trlx", prompts=["x"]) + + +class TestRolloutLiveWiredFlags: + def test_openenv_live_wired(self): + from soup_cli.utils.agent_rollout import get_rollout_backend_spec + + assert get_rollout_backend_spec("openenv").live_wired is True + + @pytest.mark.parametrize("name", ["art", "ruler", "nemo_gym"]) + def test_external_backends_stay_gated(self, name): + from soup_cli.utils.agent_rollout import get_rollout_backend_spec + + assert get_rollout_backend_spec(name).live_wired is False + + +class TestRolloutSchema: + def test_openenv_with_func_happy(self): + from soup_cli.config.loader import load_config_from_string + + cfg = load_config_from_string( + "base: a/b\n" + "task: grpo\n" + "data: {train: x.jsonl}\n" + "training: {rollout_backend: openenv, rollout_func: 'my_mod:my_fn'}\n" + ) + assert cfg.training.rollout_func == "my_mod:my_fn" + + def test_func_without_backend_rejected(self): + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="openenv"): + load_config_from_string( + "base: a/b\n" + "task: grpo\n" + "data: {train: x.jsonl}\n" + "training: {rollout_func: 'my_mod:my_fn'}\n" + ) + + def test_openenv_without_func_rejected(self): + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="rollout_func"): + load_config_from_string( + "base: a/b\n" + "task: grpo\n" + "data: {train: x.jsonl}\n" + "training: {rollout_backend: openenv}\n" + ) + + def test_func_with_art_rejected(self): + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="openenv"): + load_config_from_string( + "base: a/b\n" + "task: grpo\n" + "data: {train: x.jsonl}\n" + "training: {rollout_backend: art, rollout_func: 'my_mod:my_fn'}\n" + ) + + def test_bad_func_shape_rejected(self): + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="rollout_func"): + load_config_from_string( + "base: a/b\n" + "task: grpo\n" + "data: {train: x.jsonl}\n" + "training: {rollout_backend: openenv, rollout_func: 'no-colon'}\n" + ) + + def test_default_none(self): + from soup_cli.config.schema import TrainingConfig + + assert TrainingConfig().rollout_func is None + + +class TestRolloutGrpoWiring: + def test_grpo_trainer_launches_rollout(self): + source = (_SRC / "trainer" / "grpo.py").read_text(encoding="utf-8") + assert "launch_rollout" in source + assert "rollout_backend" in source + + +# --------------------------------------------------------------------------- +# #228 — apple-adapter live converter +# --------------------------------------------------------------------------- + + +class TestLoraKeyMapping: + def test_hf_to_mlx_strips_prefix_and_lowers(self): + from soup_cli.utils.apple_adapter import hf_key_to_mlx + + key = "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight" + assert hf_key_to_mlx(key) == "model.layers.0.self_attn.q_proj.lora_a" + + def test_hf_to_mlx_lora_b(self): + from soup_cli.utils.apple_adapter import hf_key_to_mlx + + key = "base_model.model.model.layers.2.mlp.gate_proj.lora_B.weight" + assert hf_key_to_mlx(key) == "model.layers.2.mlp.gate_proj.lora_b" + + def test_non_lora_key_returns_none(self): + from soup_cli.utils.apple_adapter import hf_key_to_mlx + + assert hf_key_to_mlx("base_model.model.model.embed_tokens.weight") is None + + def test_mlx_to_hf_round_trip(self): + from soup_cli.utils.apple_adapter import hf_key_to_mlx, mlx_key_to_hf + + original = "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight" + assert mlx_key_to_hf(hf_key_to_mlx(original)) == original + + def test_mlx_non_lora_returns_none(self): + from soup_cli.utils.apple_adapter import mlx_key_to_hf + + assert mlx_key_to_hf("model.layers.0.self_attn.q_proj.weight") is None + + +class TestConvertArrays: + def _hf_arrays(self): + rng = np.random.default_rng(0) + return { + "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": + rng.standard_normal((4, 16)).astype(np.float32), + "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight": + rng.standard_normal((16, 4)).astype(np.float32), + } + + def test_hf_to_mlx_transposes(self): + from soup_cli.utils.apple_adapter import convert_hf_to_mlx_arrays + + converted, skipped = convert_hf_to_mlx_arrays(self._hf_arrays()) + assert skipped == () + lora_a = converted["model.layers.0.self_attn.q_proj.lora_a"] + lora_b = converted["model.layers.0.self_attn.q_proj.lora_b"] + assert lora_a.shape == (16, 4) # [in, r] + assert lora_b.shape == (4, 16) # [r, out] + + def test_non_lora_keys_skipped(self): + from soup_cli.utils.apple_adapter import convert_hf_to_mlx_arrays + + arrays = self._hf_arrays() + arrays["base_model.model.model.embed_tokens.weight"] = np.zeros( + (4, 4), dtype=np.float32 + ) + converted, skipped = convert_hf_to_mlx_arrays(arrays) + assert len(converted) == 2 + assert skipped == ("base_model.model.model.embed_tokens.weight",) + + def test_zero_lora_keys_rejected(self): + from soup_cli.utils.apple_adapter import convert_hf_to_mlx_arrays + + with pytest.raises(ValueError, match="no LoRA"): + convert_hf_to_mlx_arrays( + {"embed_tokens.weight": np.zeros((2, 2), dtype=np.float32)} + ) + + def test_mlx_to_hf_round_trip_values(self): + from soup_cli.utils.apple_adapter import ( + convert_hf_to_mlx_arrays, + convert_mlx_to_hf_arrays, + ) + + original = self._hf_arrays() + mlx_arrays, _ = convert_hf_to_mlx_arrays(original) + back, skipped = convert_mlx_to_hf_arrays(mlx_arrays) + assert skipped == () + for key, value in original.items(): + np.testing.assert_allclose(back[key], value) + + +def _write_peft_adapter(adapter_dir: Path) -> dict: + """Write a synthetic PEFT LoRA adapter; return its arrays.""" + from safetensors.numpy import save_file + + rng = np.random.default_rng(7) + arrays = { + "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": + rng.standard_normal((4, 16)).astype(np.float32), + "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight": + rng.standard_normal((16, 4)).astype(np.float32), + "base_model.model.model.layers.0.self_attn.v_proj.lora_A.weight": + rng.standard_normal((4, 16)).astype(np.float32), + "base_model.model.model.layers.0.self_attn.v_proj.lora_B.weight": + rng.standard_normal((16, 4)).astype(np.float32), + } + adapter_dir.mkdir(parents=True, exist_ok=True) + save_file(arrays, str(adapter_dir / "adapter_model.safetensors")) + (adapter_dir / "adapter_config.json").write_text( + json.dumps({ + "r": 4, + "lora_alpha": 8, + "peft_type": "LORA", + "base_model_name_or_path": "tiny/base", + }), + encoding="utf-8", + ) + return arrays + + +class TestConvertAppleAdapterLive: + def test_hf_to_mlx_writes_safetensors(self, tmp_path, monkeypatch): + """mlx-lm's load_adapters reads adapters.safetensors + num_layers + (review fix — the legacy npz artifact is unloadable by current + mlx-lm).""" + from safetensors.numpy import load_file + + from soup_cli.utils.apple_adapter import ( + build_apple_adapter_plan, + convert_apple_adapter, + ) + + monkeypatch.chdir(tmp_path) + arrays = _write_peft_adapter(tmp_path / "adapter") + plan = build_apple_adapter_plan( + source_dir="adapter", output_dir="out", direction="hf-to-mlx", + ) + report = convert_apple_adapter(plan) + assert report.converted_keys == 4 + st_path = tmp_path / "out" / "adapters.safetensors" + assert st_path.is_file() + loaded = load_file(str(st_path)) + key = "model.layers.0.self_attn.q_proj.lora_a" + assert key in loaded + np.testing.assert_allclose( + loaded[key], + arrays[ + "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight" + ].T, + ) + config = json.loads( + (tmp_path / "out" / "adapter_config.json").read_text(encoding="utf-8") + ) + assert config["fine_tune_type"] == "lora" + assert config["num_layers"] == 1 # max layer index + 1 + assert config["lora_parameters"]["rank"] == 4 + assert config["lora_parameters"]["scale"] == 2.0 # alpha 8 / r 4 + + def test_full_round_trip(self, tmp_path, monkeypatch): + from safetensors.numpy import load_file + + from soup_cli.utils.apple_adapter import ( + build_apple_adapter_plan, + convert_apple_adapter, + ) + + monkeypatch.chdir(tmp_path) + arrays = _write_peft_adapter(tmp_path / "adapter") + convert_apple_adapter( + build_apple_adapter_plan( + source_dir="adapter", output_dir="mlx", direction="hf-to-mlx", + ) + ) + report = convert_apple_adapter( + build_apple_adapter_plan( + source_dir="mlx", output_dir="hf2", direction="mlx-to-hf", + ) + ) + assert report.converted_keys == 4 + back = load_file(str(tmp_path / "hf2" / "adapter_model.safetensors")) + for key, value in arrays.items(): + np.testing.assert_allclose(back[key], value) + config = json.loads( + (tmp_path / "hf2" / "adapter_config.json").read_text(encoding="utf-8") + ) + assert config["peft_type"] == "LORA" + assert config["r"] == 4 + # scale (2.0) * rank (4) reconstructs the source lora_alpha — and + # the round trip exercises the adapters.safetensors input branch + # (hf-to-mlx now emits safetensors). + assert config["lora_alpha"] == 8 + + def test_sign_writes_signature(self, tmp_path, monkeypatch): + from soup_cli.utils.apple_adapter import ( + build_apple_adapter_plan, + convert_apple_adapter, + ) + + monkeypatch.chdir(tmp_path) + _write_peft_adapter(tmp_path / "adapter") + report = convert_apple_adapter( + build_apple_adapter_plan( + source_dir="adapter", output_dir="out", + direction="hf-to-mlx", sign=True, + ) + ) + assert report.signed is True + assert (tmp_path / "out" / ".soup-signature.json").is_file() + + @pytest.mark.parametrize("direction", ["hf-to-apple", "mlx-to-apple"]) + def test_apple_directions_upstream_gated(self, tmp_path, monkeypatch, direction): + from soup_cli.utils.apple_adapter import ( + build_apple_adapter_plan, + convert_apple_adapter, + ) + + monkeypatch.chdir(tmp_path) + _write_peft_adapter(tmp_path / "adapter") + plan = build_apple_adapter_plan( + source_dir="adapter", output_dir="out", direction=direction, + ) + with pytest.raises(RuntimeError, match="FoundationModels"): + convert_apple_adapter(plan) + + def test_missing_safetensors_friendly(self, tmp_path, monkeypatch): + from soup_cli.utils.apple_adapter import ( + build_apple_adapter_plan, + convert_apple_adapter, + ) + + monkeypatch.chdir(tmp_path) + (tmp_path / "adapter").mkdir() + plan = build_apple_adapter_plan( + source_dir="adapter", output_dir="out", direction="hf-to-mlx", + ) + with pytest.raises(FileNotFoundError, match="adapter_model"): + convert_apple_adapter(plan) + + def test_bin_adapter_rejected(self, tmp_path, monkeypatch): + from soup_cli.utils.apple_adapter import ( + build_apple_adapter_plan, + convert_apple_adapter, + ) + + monkeypatch.chdir(tmp_path) + adapter = tmp_path / "adapter" + adapter.mkdir() + (adapter / "adapter_model.bin").write_bytes(b"\x80\x02") + plan = build_apple_adapter_plan( + source_dir="adapter", output_dir="out", direction="hf-to-mlx", + ) + with pytest.raises(ValueError, match="safetensors"): + convert_apple_adapter(plan) + + def test_output_outside_cwd_rejected(self, tmp_path, monkeypatch): + from soup_cli.utils.apple_adapter import ( + build_apple_adapter_plan, + convert_apple_adapter, + ) + + workdir = tmp_path / "work" + workdir.mkdir() + monkeypatch.chdir(workdir) + _write_peft_adapter(workdir / "adapter") + plan = build_apple_adapter_plan( + source_dir="adapter", + output_dir=str(tmp_path / "outside"), + direction="hf-to-mlx", + ) + with pytest.raises(ValueError, match="cwd"): + convert_apple_adapter(plan) + + def test_mlx_source_missing_files_friendly(self, tmp_path, monkeypatch): + from soup_cli.utils.apple_adapter import ( + build_apple_adapter_plan, + convert_apple_adapter, + ) + + monkeypatch.chdir(tmp_path) + (tmp_path / "mlx").mkdir() + plan = build_apple_adapter_plan( + source_dir="mlx", output_dir="out", direction="mlx-to-hf", + ) + with pytest.raises(FileNotFoundError, match="adapters"): + convert_apple_adapter(plan) + + def test_non_plan_rejected(self): + from soup_cli.utils.apple_adapter import convert_apple_adapter + + with pytest.raises(TypeError, match="AppleAdapterPlan"): + convert_apple_adapter({}) # type: ignore[arg-type] + + def test_report_frozen(self): + import dataclasses + + from soup_cli.utils.apple_adapter import ConversionReport + + report = ConversionReport( + direction="hf-to-mlx", output_dir="out", + converted_keys=2, skipped_keys=(), signed=False, + ) + with pytest.raises(dataclasses.FrozenInstanceError): + report.direction = "mlx-to-hf" # type: ignore[misc] + + +class TestAppleAdapterCli: + def test_live_conversion_exit_0(self, tmp_path, monkeypatch): + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + _write_peft_adapter(tmp_path / "adapter") + result = runner.invoke( + app, + [ + "apple-adapter", "adapter", + "--direction", "hf-to-mlx", + "--output", "out", + ], + ) + assert result.exit_code == 0, (result.output, repr(result.exception)) + assert (tmp_path / "out" / "adapters.safetensors").is_file() + + def test_apple_direction_exit_3(self, tmp_path, monkeypatch): + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + _write_peft_adapter(tmp_path / "adapter") + result = runner.invoke( + app, + [ + "apple-adapter", "adapter", + "--direction", "hf-to-apple", + "--output", "out", + ], + ) + assert result.exit_code == 3, (result.output, repr(result.exception)) + + def test_missing_weights_exit_2(self, tmp_path, monkeypatch): + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + (tmp_path / "adapter").mkdir() + result = runner.invoke( + app, + [ + "apple-adapter", "adapter", + "--direction", "hf-to-mlx", + "--output", "out", + ], + ) + assert result.exit_code == 2, (result.output, repr(result.exception)) + + +# --------------------------------------------------------------------------- +# #97 — delinearize-llama4 runtime +# --------------------------------------------------------------------------- + + +class TestIsExpertWeightKey: + @pytest.mark.parametrize( + "key", + [ + "language_model.model.layers.0.feed_forward.experts.gate_up_proj", + "language_model.model.layers.3.feed_forward.experts.down_proj", + "model.layers.0.feed_forward.experts.gate_up_proj.weight", + ], + ) + def test_fused_expert_keys(self, key): + from soup_cli.utils.delinearize_llama4 import is_expert_weight_key + + assert is_expert_weight_key(key) + + @pytest.mark.parametrize( + "key", + [ + # Per-expert numbered keys are already unfused (Mixtral-style). + "model.layers.0.feed_forward.experts.0.gate_proj.weight", + "model.layers.0.feed_forward.router.weight", + "model.layers.0.self_attn.q_proj.weight", + "", + None, + 123, + ], + ) + def test_non_fused_keys(self, key): + from soup_cli.utils.delinearize_llama4 import is_expert_weight_key + + assert not is_expert_weight_key(key) + + +class TestDelinearizeTensor: + def test_reshapes_2d_to_3d(self): + torch = _torch_or_skip() + from soup_cli.utils.delinearize_llama4 import delinearize_tensor + + tensor = torch.arange(24, dtype=torch.float32).reshape(6, 4) + out, status = delinearize_tensor(tensor, num_experts=3) + assert status == "reshaped" + assert tuple(out.shape) == (3, 2, 4) + assert torch.equal(out[0], tensor[0:2]) + + def test_3d_passthrough(self): + torch = _torch_or_skip() + from soup_cli.utils.delinearize_llama4 import delinearize_tensor + + tensor = torch.zeros(3, 2, 4) + out, status = delinearize_tensor(tensor, num_experts=3) + assert status == "already_3d" + assert out is tensor + + def test_not_divisible_rejected(self): + torch = _torch_or_skip() + from soup_cli.utils.delinearize_llama4 import delinearize_tensor + + tensor = torch.zeros(7, 4) + with pytest.raises(ValueError, match="divisible"): + delinearize_tensor(tensor, num_experts=3) + + def test_1d_rejected(self): + torch = _torch_or_skip() + from soup_cli.utils.delinearize_llama4 import delinearize_tensor + + with pytest.raises(ValueError, match="2-D"): + delinearize_tensor(torch.zeros(8), num_experts=2) + + +class TestReadNumExperts: + def test_text_config_nested(self, tmp_path): + from soup_cli.utils.delinearize_llama4 import read_num_experts + + (tmp_path / "config.json").write_text( + json.dumps({"text_config": {"num_local_experts": 16}}), + encoding="utf-8", + ) + assert read_num_experts(str(tmp_path)) == 16 + + def test_top_level(self, tmp_path): + from soup_cli.utils.delinearize_llama4 import read_num_experts + + (tmp_path / "config.json").write_text( + json.dumps({"num_local_experts": 8}), encoding="utf-8" + ) + assert read_num_experts(str(tmp_path)) == 8 + + def test_missing_returns_none(self, tmp_path): + from soup_cli.utils.delinearize_llama4 import read_num_experts + + (tmp_path / "config.json").write_text("{}", encoding="utf-8") + assert read_num_experts(str(tmp_path)) is None + + def test_no_config_returns_none(self, tmp_path): + from soup_cli.utils.delinearize_llama4 import read_num_experts + + assert read_num_experts(str(tmp_path)) is None + + def test_bool_value_rejected(self, tmp_path): + from soup_cli.utils.delinearize_llama4 import read_num_experts + + (tmp_path / "config.json").write_text( + json.dumps({"num_local_experts": True}), encoding="utf-8" + ) + assert read_num_experts(str(tmp_path)) is None + + +def _write_llama4_stub(source: Path, *, num_experts: int = 4) -> dict: + """Write a stub Llama-4-shaped checkpoint; return the tensor dict.""" + torch = _torch_or_skip() + from safetensors.torch import save_file + + tensors = { + "language_model.model.layers.0.feed_forward.experts.gate_up_proj": + torch.arange(48, dtype=torch.float32).reshape(8, 6), + "language_model.model.layers.0.feed_forward.experts.down_proj": + torch.arange(48, dtype=torch.float32).reshape(12, 4), + "language_model.model.layers.0.self_attn.q_proj.weight": + torch.ones(4, 4), + "language_model.model.layers.1.feed_forward.experts.gate_up_proj": + torch.zeros(num_experts, 2, 6), + } + source.mkdir(parents=True, exist_ok=True) + save_file(tensors, str(source / "model.safetensors")) + (source / "config.json").write_text( + json.dumps({"text_config": {"num_local_experts": num_experts}}), + encoding="utf-8", + ) + return tensors + + +class TestRunDelinearize: + def test_happy_path(self, tmp_path, monkeypatch): + torch = _torch_or_skip() + from safetensors.torch import load_file + + from soup_cli.utils.delinearize_llama4 import ( + plan_delinearize, + run_delinearize, + ) + + monkeypatch.chdir(tmp_path) + tensors = _write_llama4_stub(tmp_path / "src") + (tmp_path / "out").mkdir() + plan = plan_delinearize("src", "out") + result = run_delinearize(plan) + assert result.reshaped_keys == 2 + assert result.passthrough_keys == 1 + assert result.already_3d_keys == 1 + out = load_file(str(tmp_path / "out" / "model.safetensors")) + gate_up = out[ + "language_model.model.layers.0.feed_forward.experts.gate_up_proj" + ] + assert tuple(gate_up.shape) == (4, 2, 6) + original = tensors[ + "language_model.model.layers.0.feed_forward.experts.gate_up_proj" + ] + assert torch.equal(gate_up, original.reshape(4, 2, 6)) + # Sidecar config copied for a loadable target checkpoint. + assert (tmp_path / "out" / "config.json").is_file() + + def test_explicit_num_experts_overrides(self, tmp_path, monkeypatch): + from safetensors.torch import load_file + + from soup_cli.utils.delinearize_llama4 import ( + plan_delinearize, + run_delinearize, + ) + + monkeypatch.chdir(tmp_path) + _write_llama4_stub(tmp_path / "src", num_experts=4) + (tmp_path / "out").mkdir() + plan = plan_delinearize("src", "out") + result = run_delinearize(plan, num_experts=2) + assert result.reshaped_keys == 2 + out = load_file(str(tmp_path / "out" / "model.safetensors")) + gate_up = out[ + "language_model.model.layers.0.feed_forward.experts.gate_up_proj" + ] + assert tuple(gate_up.shape) == (2, 4, 6) + + def test_missing_num_experts_friendly(self, tmp_path, monkeypatch): + torch = _torch_or_skip() + from safetensors.torch import save_file + + from soup_cli.utils.delinearize_llama4 import ( + plan_delinearize, + run_delinearize, + ) + + monkeypatch.chdir(tmp_path) + src = tmp_path / "src" + src.mkdir() + save_file({"x": torch.zeros(2, 2)}, str(src / "model.safetensors")) + (tmp_path / "out").mkdir() + plan = plan_delinearize("src", "out") + with pytest.raises(ValueError, match="--num-experts"): + run_delinearize(plan) + + def test_non_divisible_names_key(self, tmp_path, monkeypatch): + torch = _torch_or_skip() + from safetensors.torch import save_file + + from soup_cli.utils.delinearize_llama4 import ( + plan_delinearize, + run_delinearize, + ) + + monkeypatch.chdir(tmp_path) + src = tmp_path / "src" + src.mkdir() + save_file( + { + "model.layers.0.feed_forward.experts.gate_up_proj": + torch.zeros(7, 4), + }, + str(src / "model.safetensors"), + ) + (tmp_path / "out").mkdir() + plan = plan_delinearize("src", "out") + with pytest.raises(ValueError, match="gate_up_proj"): + run_delinearize(plan, num_experts=4) + + @pytest.mark.parametrize("bad", [True, 0, -2, "four", 1_000_000]) + def test_num_experts_bounds(self, tmp_path, monkeypatch, bad): + from soup_cli.utils.delinearize_llama4 import ( + plan_delinearize, + run_delinearize, + ) + + monkeypatch.chdir(tmp_path) + _write_llama4_stub(tmp_path / "src") + (tmp_path / "out").mkdir() + plan = plan_delinearize("src", "out") + with pytest.raises((TypeError, ValueError), match="num_experts"): + run_delinearize(plan, num_experts=bad) + + def test_non_plan_rejected(self): + from soup_cli.utils.delinearize_llama4 import run_delinearize + + with pytest.raises(TypeError, match="DelinearizePlan"): + run_delinearize({}) # type: ignore[arg-type] + + def test_result_frozen(self): + import dataclasses + + from soup_cli.utils.delinearize_llama4 import DelinearizeResult + + result = DelinearizeResult( + source_dir="a", target_dir="b", files_written=("x",), + reshaped_keys=1, passthrough_keys=0, already_3d_keys=0, + ) + with pytest.raises(dataclasses.FrozenInstanceError): + result.reshaped_keys = 9 # type: ignore[misc] + + +class TestDelinearizeCli: + def test_live_run_exit_0(self, tmp_path, monkeypatch): + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + _write_llama4_stub(tmp_path / "src") + (tmp_path / "out").mkdir() + result = runner.invoke( + app, ["delinearize-llama4", "src", "--target", "out"], + ) + assert result.exit_code == 0, (result.output, repr(result.exception)) + assert (tmp_path / "out" / "model.safetensors").is_file() + + def test_plan_only_writes_nothing(self, tmp_path, monkeypatch): + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + _write_llama4_stub(tmp_path / "src") + (tmp_path / "out").mkdir() + result = runner.invoke( + app, + ["delinearize-llama4", "src", "--target", "out", "--plan-only"], + ) + assert result.exit_code == 0, (result.output, repr(result.exception)) + assert not (tmp_path / "out" / "model.safetensors").exists() + + def test_missing_num_experts_exit_2(self, tmp_path, monkeypatch): + torch = _torch_or_skip() + from safetensors.torch import save_file + + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + src = tmp_path / "src" + src.mkdir() + save_file({"x": torch.zeros(2, 2)}, str(src / "model.safetensors")) + (tmp_path / "out").mkdir() + result = runner.invoke( + app, ["delinearize-llama4", "src", "--target", "out"], + ) + assert result.exit_code == 2, (result.output, repr(result.exception)) + assert "--num-experts" in result.output + + +# --------------------------------------------------------------------------- +# Cross-cutting invariants +# --------------------------------------------------------------------------- + + +class TestPatchInvariants: + def test_version_bumped(self): + import soup_cli + + parts = tuple(int(p) for p in soup_cli.__version__.split(".")[:3]) + assert parts >= (0, 71, 21) + + @pytest.mark.parametrize( + "rel", + [ + "utils/advanced_precision.py", + "utils/grpo_long_context.py", + "utils/agent_rollout.py", + "utils/apple_adapter.py", + "utils/delinearize_llama4.py", + ], + ) + def test_no_heavy_top_level_imports(self, rel): + # Full-source scan — TYPE_CHECKING-block imports are indented so + # the column-0 patterns can never match them; splitting the source + # there would blind the guard for everything below the block + # (review fix). + source = (_SRC / rel).read_text(encoding="utf-8") + for heavy in ( + "torch", + "transformers", + "vllm", + "torchao", + "numpy", + "safetensors", + "mlx", + "peft", + ): + assert f"\nimport {heavy}\n" not in source + assert f"\nimport {heavy} " not in source + assert f"\nfrom {heavy}" not in source + + +# --------------------------------------------------------------------------- +# Review-fix follow-ups (v0.71.21 review wave) +# --------------------------------------------------------------------------- + + +class TestReviewFollowupsPrecision: + """#141 review fixes — torchao probe, recipe NUL, partial-conversion.""" + + def test_out_proj_and_wqkv_in_allowlist(self): + from soup_cli.utils.advanced_precision import is_attention_projection + + assert is_attention_projection("encoder.layers.0.self_attn.out_proj") + assert is_attention_projection("transformer.blocks.0.attn.Wqkv") + + def test_is_blackwell_gpu_runtime_error_false(self, monkeypatch): + torch = _torch_or_skip() + from soup_cli.utils.advanced_precision import is_blackwell_gpu + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + + def _boom(_i=0): + raise RuntimeError("driver mismatch") + + monkeypatch.setattr(torch.cuda, "get_device_capability", _boom) + assert is_blackwell_gpu() is False + + def test_already_converted_fast_path_no_float8_import(self, monkeypatch): + """All projections pre-converted -> count returned WITHOUT importing + torchao.float8 (the bare fake root has no float8 attr).""" + torch = _torch_or_skip() + nn = torch.nn + from soup_cli.utils.advanced_precision import apply_fp8_attention + + class _F8(nn.Linear): + pass + + _F8.__name__ = "Float8Linear" + _enable_fp8_gates(monkeypatch) + model = _tiny_attn_model() + attn = model.layers[0].self_attn + for name in ("q_proj", "k_proj", "v_proj", "o_proj"): + setattr(attn, name, _F8(8, 8)) + assert apply_fp8_attention(model) == 4 + + def test_partial_conversion_failure_is_honest(self, monkeypatch): + """torchao crashing mid-conversion surfaces 'PARTIALLY converted'.""" + from soup_cli.utils.advanced_precision import apply_fp8_attention + + record: dict = {} + _install_fake_torchao_float8(monkeypatch, record) + _enable_fp8_gates(monkeypatch) + + def _boom(model, config=None, module_filter_fn=None): + raise ValueError("unsupported dim") + + sys.modules["torchao.float8"].convert_to_float8_training = _boom + with pytest.raises(RuntimeError, match="PARTIALLY"): + apply_fp8_attention(_tiny_attn_model()) + + def test_nvfp4_missing_quantize_friendly(self, monkeypatch): + from soup_cli.utils.advanced_precision import apply_nvfp4 + + fake_q = types.ModuleType("torchao.quantization") + + class _NVFP4Config: + pass + + fake_q.NVFP4Config = _NVFP4Config # no quantize_ + fake_root = types.ModuleType("torchao") + fake_root.quantization = fake_q + monkeypatch.setitem(sys.modules, "torchao", fake_root) + monkeypatch.setitem(sys.modules, "torchao.quantization", fake_q) + monkeypatch.setattr( + "soup_cli.utils.advanced_precision.is_blackwell_gpu", lambda: True + ) + with pytest.raises(RuntimeError, match="quantize_"): + apply_nvfp4(_tiny_attn_model()) + + def test_v028_degrade_is_hermetic(self, monkeypatch): + """Force the degrade path by patching the converters directly — + environment-independent (review fix: the original test relied on + the host lacking torchao/Hopper).""" + from soup_cli.config.loader import load_config_from_string + from soup_cli.utils import advanced_precision, v028_features + + def _gate(model, **kwargs): + raise RuntimeError("gate fired") + + # v028_features lazy-imports the converters at call time, so the + # patch targets the defining module. + monkeypatch.setattr(advanced_precision, "apply_fp8_attention", _gate) + monkeypatch.setattr(advanced_precision, "apply_nvfp4", _gate) + cfg = load_config_from_string( + "base: test-llama\n" + "task: sft\n" + "training:\n" + " quantization_aware: fp8\n" + " fp8_attention: true\n" + " nvfp4: true\n" + "data:\n" + " train: data.jsonl\n" + "output: ./out\n" + ) + applied = v028_features.apply_v028_speed_memory( + model=object(), tcfg=cfg.training, base_model=cfg.base, + console=None, device="cpu", backend="transformers", + ) + assert applied.get("fp8_attention") is False + assert applied.get("nvfp4") is False + + +class TestReviewFollowupsSleepMode: + """#124 review fixes — level validation + version edge cases.""" + + def test_parse_version_tuple_none(self): + from soup_cli.utils.grpo_long_context import _parse_version_tuple + + assert _parse_version_tuple(None) == () # type: ignore[arg-type] + + def test_vllm_without_version_attr_false(self, monkeypatch): + from soup_cli.utils.grpo_long_context import vllm_supports_sleep_mode + + fake = types.ModuleType("vllm") # no __version__ + monkeypatch.setitem(sys.modules, "vllm", fake) + assert vllm_supports_sleep_mode() is False + + def test_exact_floor_0_7_0_supported(self, monkeypatch): + from soup_cli.utils.grpo_long_context import vllm_supports_sleep_mode + + _install_fake_vllm(monkeypatch, "0.7.0") + assert vllm_supports_sleep_mode() is True + + def test_sleep_cycle_level_2_forwarded(self): + from soup_cli.utils.grpo_long_context import vllm_sleep_cycle + + calls: list = [] + engine = SimpleNamespace( + sleep=lambda level: calls.append(("sleep", level)), + wake_up=lambda: calls.append(("wake", None)), + ) + with vllm_sleep_cycle(engine, level=2): + pass + assert calls == [("sleep", 2), ("wake", None)] + + @pytest.mark.parametrize("level", [True, 0, 3, "1"]) + def test_sleep_cycle_bad_level_rejected(self, level): + from soup_cli.utils.grpo_long_context import vllm_sleep_cycle + + engine = SimpleNamespace(sleep=lambda level: None, wake_up=lambda: None) + with pytest.raises((TypeError, ValueError), match="level"): + with vllm_sleep_cycle(engine, level=level): + pass + + def test_sleep_without_wake_up_warns_and_runs(self, caplog): + import logging + + from soup_cli.utils.grpo_long_context import vllm_sleep_cycle + + engine = SimpleNamespace(sleep=lambda level: None) # no wake_up + ran = [] + with caplog.at_level( + logging.WARNING, logger="soup_cli.utils.grpo_long_context" + ): + with vllm_sleep_cycle(engine): + ran.append(True) + assert ran == [True] + assert any("sleep" in rec.message for rec in caplog.records) + + +class TestReviewFollowupsRollout: + """#125 review fixes — smuggle-strip, loud answers, boundaries.""" + + def test_extra_keys_stripped(self, tmp_path, monkeypatch): + from soup_cli.utils.agent_rollout import launch_rollout + + spec = _write_rollout_module( + tmp_path, + monkeypatch, + "def rollout(prompts):\n" + " return [{'prompt': 'x', 'evil': 'payload', 'answer': 'a'}]\n", + ) + result = launch_rollout("openenv", prompts=["p"], rollout_func=spec) + assert set(result.rows[0]) == {"prompt", "answer"} + + def test_non_str_answer_rejected_loudly(self, tmp_path, monkeypatch): + from soup_cli.utils.agent_rollout import launch_rollout + + spec = _write_rollout_module( + tmp_path, + monkeypatch, + "def rollout(prompts):\n" + " return [{'prompt': 'x', 'answer': 42}]\n", + ) + with pytest.raises(ValueError, match="answer"): + launch_rollout("openenv", prompts=["p"], rollout_func=spec) + + def test_empty_prompt_rejected(self, tmp_path, monkeypatch): + from soup_cli.utils.agent_rollout import launch_rollout + + spec = _write_rollout_module( + tmp_path, + monkeypatch, + "def rollout(prompts):\n return [{'prompt': ''}]\n", + ) + with pytest.raises(ValueError, match="empty prompt"): + launch_rollout("openenv", prompts=["p"], rollout_func=spec) + + def test_empty_rows_result_rejected(self): + from soup_cli.utils.agent_rollout import RolloutResult + + with pytest.raises(ValueError, match="no rows"): + RolloutResult(backend="openenv", rows=()) + + def test_result_stores_canonical_backend(self): + from soup_cli.utils.agent_rollout import RolloutResult + + result = RolloutResult(backend="ART", rows=({"prompt": "x"},)) + assert result.backend == "art" + + def test_message_list_prompts_alias_broken(self): + """Mutating the rollout callable's retained message dict must not + leak into the normalised training rows (immutability policy).""" + from soup_cli.utils.agent_rollout import _normalise_rollout_rows + + message = {"role": "user", "content": "hi"} + rows = _normalise_rollout_rows([{"prompt": [message]}], "openenv") + message["content"] = "MUTATED" + assert rows[0]["prompt"][0]["content"] == "hi" + + def test_validate_rollout_func_bool_rejected(self): + from soup_cli.utils.agent_rollout import validate_rollout_func + + # ValueError (not TypeError) so the Pydantic field_validator wraps + # it into a ValidationError (mode="before" validators re-raise + # TypeError raw — v2 convention). + with pytest.raises(ValueError, match="rollout_func"): + validate_rollout_func(True) # type: ignore[arg-type] + + @pytest.mark.parametrize( + ("max_steps", "ok"), [(100_000, True), (100_001, False)] + ) + def test_max_steps_exact_boundary( + self, tmp_path, monkeypatch, max_steps, ok + ): + from soup_cli.utils.agent_rollout import launch_rollout + + spec = _write_rollout_module( + tmp_path, + monkeypatch, + "def rollout(prompts):\n return [{'prompt': 'x'}]\n", + ) + if ok: + result = launch_rollout( + "openenv", prompts=["p"], rollout_func=spec, + max_steps=max_steps, + ) + assert len(result.rows) == 1 + else: + with pytest.raises(ValueError, match="max_steps"): + launch_rollout( + "openenv", prompts=["p"], rollout_func=spec, + max_steps=max_steps, + ) + + def test_external_runner_receives_all_kwargs(self, monkeypatch): + from soup_cli.utils import agent_rollout + + seen: dict = {} + + def _runner(**kwargs): + seen.update(kwargs) + return [{"prompt": "from-runner"}] + + monkeypatch.setitem( + agent_rollout._EXTERNAL_ROLLOUT_RUNNERS, "art", _runner + ) + model, tokenizer, reward = object(), object(), object() + agent_rollout.launch_rollout( + "art", prompts=["p"], model=model, tokenizer=tokenizer, + reward_fn=reward, max_steps=7, + ) + assert seen["prompts"] == ["p"] + assert seen["model"] is model + assert seen["tokenizer"] is tokenizer + assert seen["reward_fn"] is reward + assert seen["max_steps"] == 7 + + +class TestReviewFollowupsAppleAdapter: + """#228 review fixes — symlink TOCTOU, corrupt files, config carry.""" + + def test_infer_num_layers_gpt2_style(self): + """GPT-2-style ``transformer.h.N`` paths also derive num_layers + (caught by the real bf16 PEFT adapter smoke).""" + from soup_cli.utils.apple_adapter import _infer_num_layers + + assert _infer_num_layers({ + "transformer.h.4.attn.c_attn.lora_a": None, + "transformer.h.0.attn.c_attn.lora_b": None, + }) == 5 + assert _infer_num_layers({"model.layers.2.q_proj.lora_a": None}) == 3 + assert _infer_num_layers({"no_layer_key.lora_a": None}) is None + + def test_prefixless_hf_key_mapped(self): + from soup_cli.utils.apple_adapter import hf_key_to_mlx, mlx_key_to_hf + + assert ( + hf_key_to_mlx("model.layers.0.self_attn.q_proj.lora_A.weight") + == "model.layers.0.self_attn.q_proj.lora_a" + ) + # The reverse direction re-adds the canonical prefix (documented + # asymmetry — prefix-less inputs do not byte-round-trip). + assert mlx_key_to_hf("model.layers.0.self_attn.q_proj.lora_a") == ( + "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight" + ) + + def test_lora_dropout_carried_through(self, tmp_path, monkeypatch): + from soup_cli.utils.apple_adapter import ( + build_apple_adapter_plan, + convert_apple_adapter, + ) + + monkeypatch.chdir(tmp_path) + _write_peft_adapter(tmp_path / "adapter") + config_path = tmp_path / "adapter" / "adapter_config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config["lora_dropout"] = 0.05 + config_path.write_text(json.dumps(config), encoding="utf-8") + convert_apple_adapter( + build_apple_adapter_plan( + source_dir="adapter", output_dir="out", direction="hf-to-mlx", + ) + ) + out_config = json.loads( + (tmp_path / "out" / "adapter_config.json").read_text( + encoding="utf-8" + ) + ) + assert out_config["lora_parameters"]["dropout"] == 0.05 + + def test_legacy_npz_input_still_loads(self, tmp_path, monkeypatch): + from safetensors.numpy import load_file + + from soup_cli.utils.apple_adapter import ( + build_apple_adapter_plan, + convert_apple_adapter, + ) + + monkeypatch.chdir(tmp_path) + mlx_dir = tmp_path / "mlx" + mlx_dir.mkdir() + rng = np.random.default_rng(3) + np.savez( + str(mlx_dir / "adapters.npz"), + **{ + "model.layers.0.self_attn.q_proj.lora_a": + rng.standard_normal((16, 4)).astype(np.float32), + "model.layers.0.self_attn.q_proj.lora_b": + rng.standard_normal((4, 16)).astype(np.float32), + }, + ) + report = convert_apple_adapter( + build_apple_adapter_plan( + source_dir="mlx", output_dir="hf", direction="mlx-to-hf", + ) + ) + assert report.converted_keys == 2 + back = load_file(str(tmp_path / "hf" / "adapter_model.safetensors")) + assert ( + "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight" + in back + ) + + def test_corrupt_safetensors_value_error(self, tmp_path, monkeypatch): + pytest.importorskip("torch") + from soup_cli.utils.apple_adapter import ( + build_apple_adapter_plan, + convert_apple_adapter, + ) + + monkeypatch.chdir(tmp_path) + adapter = tmp_path / "adapter" + adapter.mkdir() + (adapter / "adapter_model.safetensors").write_bytes(b"not-safetensors") + plan = build_apple_adapter_plan( + source_dir="adapter", output_dir="out", direction="hf-to-mlx", + ) + with pytest.raises(ValueError, match="not a valid safetensors"): + convert_apple_adapter(plan) + + def test_corrupt_npz_value_error(self, tmp_path, monkeypatch): + from soup_cli.utils.apple_adapter import ( + build_apple_adapter_plan, + convert_apple_adapter, + ) + + monkeypatch.chdir(tmp_path) + mlx_dir = tmp_path / "mlx" + mlx_dir.mkdir() + (mlx_dir / "adapters.npz").write_bytes(b"not-a-zip") + plan = build_apple_adapter_plan( + source_dir="mlx", output_dir="out", direction="mlx-to-hf", + ) + with pytest.raises(ValueError, match="not a valid npz"): + convert_apple_adapter(plan) + + def test_npz_decompression_cap(self, tmp_path, monkeypatch): + """The 4 GiB cap re-applies to DECOMPRESSED arrays (zip bomb).""" + from soup_cli.utils import apple_adapter + + monkeypatch.chdir(tmp_path) + mlx_dir = tmp_path / "mlx" + mlx_dir.mkdir() + zeros = np.zeros((128, 128), dtype=np.float32) # 64 KiB, compresses tiny + with open(mlx_dir / "adapters.npz", "wb") as handle: + np.savez_compressed( + handle, **{"model.layers.0.self_attn.q_proj.lora_a": zeros} + ) + on_disk = (mlx_dir / "adapters.npz").stat().st_size + monkeypatch.setattr( + apple_adapter, "_MAX_ADAPTER_FILE_BYTES", on_disk + 1024 + ) + plan = apple_adapter.build_apple_adapter_plan( + source_dir="mlx", output_dir="out", direction="mlx-to-hf", + ) + with pytest.raises(ValueError, match="decompresses past"): + apple_adapter.convert_apple_adapter(plan) + + def test_adapter_size_cap_branch(self, tmp_path, monkeypatch): + from soup_cli.utils import apple_adapter + + monkeypatch.chdir(tmp_path) + _write_peft_adapter(tmp_path / "adapter") + monkeypatch.setattr(apple_adapter, "_MAX_ADAPTER_FILE_BYTES", 16) + plan = apple_adapter.build_apple_adapter_plan( + source_dir="adapter", output_dir="out", direction="hf-to-mlx", + ) + with pytest.raises(ValueError, match="cap"): + apple_adapter.convert_apple_adapter(plan) + + @pytest.mark.skipif(sys.platform == "win32", reason="POSIX symlink") + def test_symlinked_weights_rejected(self, tmp_path, monkeypatch): + import os + + from soup_cli.utils.apple_adapter import ( + build_apple_adapter_plan, + convert_apple_adapter, + ) + + monkeypatch.chdir(tmp_path) + real = tmp_path / "real" + _write_peft_adapter(real) + adapter = tmp_path / "adapter" + adapter.mkdir() + os.symlink( + str(real / "adapter_model.safetensors"), + str(adapter / "adapter_model.safetensors"), + ) + plan = build_apple_adapter_plan( + source_dir="adapter", output_dir="out", direction="hf-to-mlx", + ) + with pytest.raises(ValueError, match="symlink"): + convert_apple_adapter(plan) + + @pytest.mark.parametrize( + ("field", "value", "exc"), + [ + ("converted_keys", -1, ValueError), + ("converted_keys", True, TypeError), + ("skipped_keys", ["a"], TypeError), + ("signed", "yes", TypeError), + ], + ) + def test_conversion_report_post_init(self, field, value, exc): + from soup_cli.utils.apple_adapter import ConversionReport + + kwargs = { + "direction": "hf-to-mlx", + "output_dir": "out", + "converted_keys": 2, + "skipped_keys": (), + "signed": False, + } + kwargs[field] = value + with pytest.raises(exc): + ConversionReport(**kwargs) + + +class TestReviewFollowupsDelinearize: + """#97 review fixes — plan containment, corrupt shards, boundaries.""" + + @pytest.mark.parametrize( + ("value", "ok"), [(4096, True), (4097, False), (1, True), (0, False)] + ) + def test_num_experts_exact_boundary(self, value, ok): + from soup_cli.utils.delinearize_llama4 import _validate_num_experts + + if ok: + assert _validate_num_experts(value) == value + else: + with pytest.raises(ValueError, match="num_experts"): + _validate_num_experts(value) + + def test_plan_direct_construction_outside_cwd_rejected( + self, tmp_path, monkeypatch + ): + from soup_cli.utils.delinearize_llama4 import DelinearizePlan + + workdir = tmp_path / "work" + workdir.mkdir() + monkeypatch.chdir(workdir) + with pytest.raises(ValueError, match="outside cwd"): + DelinearizePlan( + source_dir=str(workdir / "src"), + target_dir=str(tmp_path / "outside"), + weight_files=("model.safetensors",), + ) + + def test_corrupt_shard_value_error(self, tmp_path, monkeypatch): + pytest.importorskip("torch") + from soup_cli.utils.delinearize_llama4 import ( + plan_delinearize, + run_delinearize, + ) + + monkeypatch.chdir(tmp_path) + src = tmp_path / "src" + src.mkdir() + (src / "model.safetensors").write_bytes(b"garbage") + (tmp_path / "out").mkdir() + plan = plan_delinearize("src", "out") + with pytest.raises(ValueError, match="not a valid safetensors"): + run_delinearize(plan, num_experts=4) + + def test_weight_file_size_cap_branch(self, tmp_path, monkeypatch): + pytest.importorskip("torch") + from soup_cli.utils import delinearize_llama4 + + monkeypatch.chdir(tmp_path) + _write_llama4_stub(tmp_path / "src") + (tmp_path / "out").mkdir() + plan = delinearize_llama4.plan_delinearize("src", "out") + monkeypatch.setattr(delinearize_llama4, "_MAX_WEIGHT_FILE_BYTES", 16) + with pytest.raises(ValueError, match="cap"): + delinearize_llama4.run_delinearize(plan) + + def test_multi_shard_run(self, tmp_path, monkeypatch): + torch = _torch_or_skip() + from safetensors.torch import save_file + + from soup_cli.utils.delinearize_llama4 import ( + plan_delinearize, + run_delinearize, + ) + + monkeypatch.chdir(tmp_path) + _write_llama4_stub(tmp_path / "src") + save_file( + { + "language_model.model.layers.2.feed_forward.experts.down_proj": + torch.zeros(8, 4), + }, + str(tmp_path / "src" / "model-00002.safetensors"), + ) + (tmp_path / "out").mkdir() + result = run_delinearize(plan_delinearize("src", "out")) + assert len(result.files_written) == 2 + assert result.sidecars_copied >= 1 # config.json + + @pytest.mark.parametrize( + ("config", "expected"), + [ + ({"num_experts": 16}, 16), # top-level fallback key + ({"text_config": {"num_local_experts": 5000}}, None), # OOB + ], + ) + def test_read_num_experts_fallbacks( + self, tmp_path, monkeypatch, config, expected + ): + from soup_cli.utils.delinearize_llama4 import read_num_experts + + monkeypatch.chdir(tmp_path) + src = tmp_path / "src" + src.mkdir() + (src / "config.json").write_text(json.dumps(config), encoding="utf-8") + assert read_num_experts(str(src)) == expected + + @pytest.mark.skipif(sys.platform == "win32", reason="POSIX symlink") + def test_symlinked_shard_rejected(self, tmp_path, monkeypatch): + import os + + pytest.importorskip("torch") + from soup_cli.utils.delinearize_llama4 import ( + plan_delinearize, + run_delinearize, + ) + + monkeypatch.chdir(tmp_path) + real = tmp_path / "real" + _write_llama4_stub(real) + src = tmp_path / "src" + src.mkdir() + os.symlink( + str(real / "model.safetensors"), str(src / "model.safetensors") + ) + (tmp_path / "out").mkdir() + plan = plan_delinearize("src", "out") + with pytest.raises(ValueError, match="symlink"): + run_delinearize(plan, num_experts=4) + + @pytest.mark.skipif(sys.platform == "win32", reason="POSIX symlink") + def test_symlinked_config_json_skipped(self, tmp_path, monkeypatch): + import os + + from soup_cli.utils.delinearize_llama4 import read_num_experts + + monkeypatch.chdir(tmp_path) + real = tmp_path / "real" + real.mkdir() + (real / "config.json").write_text( + json.dumps({"num_local_experts": 4}), encoding="utf-8" + ) + src = tmp_path / "src" + src.mkdir() + os.symlink(str(real / "config.json"), str(src / "config.json")) + assert read_num_experts(str(src)) is None + + def test_cli_num_experts_flag_passthrough(self, tmp_path, monkeypatch): + pytest.importorskip("torch") + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + src = tmp_path / "src" + _write_llama4_stub(src) + (src / "config.json").unlink() # force the flag to matter + (tmp_path / "out").mkdir() + result = runner.invoke( + app, + [ + "delinearize-llama4", "src", + "--target", "out", + "--num-experts", "4", + ], + ) + assert result.exit_code == 0, (result.output, repr(result.exception)) + assert (tmp_path / "out" / "model.safetensors").is_file()