From f51331d637b486e8748d50556cf2ca61057a6d4e Mon Sep 17 00:00:00 2001 From: Alpamys Date: Tue, 9 Jun 2026 12:37:50 +0500 Subject: [PATCH] feat(quant,multipack): vision/audio Quant Menu + multipack FSDP sharding (v0.71.19) #81 Quant Menu for vision/audio modality - config/schema: drop the `modality != "text"` rejection in _validate_quant_menu_supported_tasks (mlx-backend gate retained) so the full Quant Menu (gptq/awq/hqq:Nbit/aqlm/eetq/mxfp4/fp8) applies to vision+audio. - trainer/sft: _setup_vision_transformers + _setup_audio_transformers call build_quantization_config_for_loader (strict superset of the inline 4bit/8bit BNB blocks they replaced); drop BitsAndBytesConfig import; retain the prepare_model_for_kbit_training gate on (4bit,8bit,mxfp4). #80 multipack DataLoader sharding under FSDP/DeepSpeed/DDP - utils/multipack_trainer: get_train_dataloader routes the multipack DataLoader through accelerator.prepare when num_processes > 1 so accelerate's BatchSamplerShard shards whole FFD bins across ranks (even_batches=True avoids the epoch-boundary collective hang; seed identical across ranks, no `+ rank`). Single-process path unchanged. Defence-in-depth guard against an unconfigured MagicMock num_processes. Tests: +37 in tests/test_v07119.py (13770 -> 13807). ruff clean. Supersedes the v0.40.5 vision/audio Quant-Menu and v0.40.4 multipack-FSDP known-limitations. Full multi-GPU validation remains an INFRA-BLOCKED QA item. --- CHANGELOG.md | 31 ++ CONTRIBUTING.md | 2 +- README.md | 24 +- pyproject.toml | 2 +- src/soup_cli/__init__.py | 2 +- src/soup_cli/config/schema.py | 15 +- src/soup_cli/trainer/sft.py | 54 ++-- src/soup_cli/utils/multipack_trainer.py | 42 ++- tests/test_quant_menu.py | 14 +- tests/test_v0405_part_a.py | 17 +- tests/test_v07119.py | 410 ++++++++++++++++++++++++ 11 files changed, 543 insertions(+), 70 deletions(-) create mode 100644 tests/test_v07119.py diff --git a/CHANGELOG.md b/CHANGELOG.md index cd4cc97..648ac89 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,37 @@ reproducing 70+ versions of notes. ## [Unreleased] +## [0.71.19] - 2026-06-09 + +### Added +- **Quant Menu for vision / audio modality** (closes #81). The Quant Menu + (`gptq` / `awq` / `hqq:Nbit` / `aqlm` / `eetq` / `mxfp4` / `fp8`) was rejected + by the config modality gate for `modality in {vision, audio}` — those paths + carried inline `BitsAndBytesConfig` blocks that handled only `4bit` / `8bit`. + v0.71.19 drops the gate (the mlx-backend gate is retained) and threads the + unified `build_quantization_config_for_loader` through + `_setup_vision_transformers` / `_setup_audio_transformers`, so multi-modal SFT + can train a LoRA on top of any pre-quantized base. The `4bit` / `8bit` config + shapes are byte-for-byte the same as the old inline blocks; `mxfp4` still + routes through `prepare_model_for_kbit_training`. Verified: the unified loader + returns the right config object for every format on both modalities, and + `_setup_vision_transformers` threads a `GPTQConfig` into + `AutoModelForVision2Seq.from_pretrained`. + +### Fixed +- **Multipack DataLoader sharding under FSDP / DeepSpeed ZeRO / DDP** (closes + #80). The multipack `get_train_dataloader` override built a raw `DataLoader` + and returned it directly, so under distribution every rank trained on the + **same** packed bins (no data sharding). It now routes the loader through + `accelerator.prepare(...)` when `num_processes > 1` — exactly what HF Trainer's + own `get_train_dataloader` does — so accelerate's `BatchSamplerShard` + round-robins whole bins across ranks (preserving the FFD packing) and + equalises per-rank batch counts. The single-process path is unchanged + (byte-for-byte the validated v0.40.4 raw-DataLoader behaviour). Verified live: + a single-GPU multipack SFT on SmolLM2-135M trains end-to-end (RTX 3050). Full + multi-GPU validation remains a QA item (no multi-GPU box); the distributed + routing is mocked-tested. + ## [0.71.18] - 2026-06-08 ### Added diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d72d9bd..05502ad 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -120,7 +120,7 @@ src/soup_cli/ templates/ - 17 built-in soup.yaml templates (YAML + manifest.json) with load_template loader (v0.39.0, +bco v0.40.0) ui/ - Web UI (FastAPI + HTML/JS SPA) -tests/ - Test suite (289 files, 13770 tests) +tests/ - Test suite (290 files, 13807 tests) examples/ - Real-world config examples and datasets ``` diff --git a/README.md b/README.md index c0108d4..2d00f5e 100644 --- a/README.md +++ b/README.md @@ -49,22 +49,16 @@ infrastructure instead of improving models. Soup fixes that. ## What's New -**v0.71.18 — Distill + agent depth.** Deeper distillation and agent-eval, validated live on -Windows + RTX 3050 (tiny-gpt2 + SmolLM2-135M): +**v0.71.19 — Quant-menu + multipack hardening.** Two fixes validated on Windows + RTX 3050: -- **MiniLLM true on-policy rollout** — `soup train --minillm-enabled --minillm-on-policy` replaces the - offline distribution blend with the real on-policy procedure (sample a fresh autoregressive rollout - from the per-token teacher/student mixture, then length-normalised reverse-KL). Tune the rollout with - `training.minillm_rollout_length`. -- **Cross-tokenizer ULD for fully-disjoint tokenizers** — `training.uld_strategy: wasserstein_aligned` - aligns the student and teacher token sequences over their decoded character spans, so you can distill - across genuinely different tokenizers (e.g. a GPT-2 BPE student from a Llama SentencePiece teacher). -- **`soup agent eval --sandbox`** — each heuristic-passing tool-call prediction is *executed* against a - generated mock of the endpoint in the RLVR sandbox and classified ok / tool_error / timeout / - arg_error (strong isolation on POSIX; subprocess + timeout + output cap + network guard everywhere). -- **`soup train --cloud modal`** — no local GPU? Render a self-contained Modal.com app from your - `soup.yaml` (config base64-embedded, no secrets) for serverless GPU training; plan-only by default, - `--cloud-submit` submits live. `pip install 'soup-cli[modal]'`. +- **Quant Menu for vision / audio** — the full quant menu (`gptq` / `awq` / `hqq:Nbit` / `aqlm` / + `eetq` / `mxfp4` / `fp8`) now works for `modality: vision` and `modality: audio`, not just text. + The vision/audio setup paths thread the same unified quantization loader as the text path, so you + can fine-tune a LoRA on top of a pre-quantized multi-modal base. +- **Multipack sharding under multi-GPU** — `training.multipack: true` now shards the FFD-packed bins + correctly across ranks under FSDP / DeepSpeed ZeRO / DDP (the DataLoader is routed through accelerate + exactly as HF Trainer does). The single-GPU path is unchanged. Full multi-GPU validation remains a + QA item; the routing is mocked-tested. Full history: [CHANGELOG.md](CHANGELOG.md) · [GitHub Releases](https://github.com/MakazhanAlpamys/Soup/releases). diff --git a/pyproject.toml b/pyproject.toml index ed7b5be..f9c9644 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "soup-cli" -version = "0.71.18" +version = "0.71.19" 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 7418ebb..e0b13a3 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.18" +__version__ = "0.71.19" diff --git a/src/soup_cli/config/schema.py b/src/soup_cli/config/schema.py index 22a4ac9..e6398dd 100644 --- a/src/soup_cli/config/schema.py +++ b/src/soup_cli/config/schema.py @@ -3823,9 +3823,12 @@ class SoupConfig(BaseModel): """v0.40.5 (#66) — Quant Menu (gptq/awq/hqq:Nbit/aqlm/eetq/mxfp4/fp8) is wired across every transformer-backend trainer (sft / dpo / grpo / kto / orpo / simpo / ipo / ppo / reward_model / pretrain / embedding / - bco). MLX backend still rejected (no equivalent kernels). Vision/audio - modality multi-trainer wiring deferred (mirrors v0.38.1 stub-then-live - pattern for non-text modalities). + bco). MLX backend still rejected (no equivalent kernels). + + v0.71.19 (#81) — vision / audio modality wiring landed: the SFT + ``_setup_vision_transformers`` / ``_setup_audio_transformers`` paths now + thread the unified ``build_quantization_config_for_loader``, so the + modality gate is dropped (only the mlx-backend gate remains). """ from soup_cli.utils.quant_menu import is_quant_menu_format @@ -3839,12 +3842,6 @@ class SoupConfig(BaseModel): "(no equivalent kernels). Use backend='transformers' or " "switch to quantization in {'4bit', '8bit', 'none'}." ) - if self.modality != "text": - raise ValueError( - f"quantization={quant!r} (Quant Menu) is wired for " - f"modality='text' only; got modality={self.modality!r}. " - "Vision/audio multi-modal wiring is tracked for a follow-up patch." - ) return self @model_validator(mode="after") diff --git a/src/soup_cli/trainer/sft.py b/src/soup_cli/trainer/sft.py index 49331f0..83d214f 100644 --- a/src/soup_cli/trainer/sft.py +++ b/src/soup_cli/trainer/sft.py @@ -848,7 +848,7 @@ class SFTTrainerWrapper: def _setup_vision_transformers(self, cfg, tcfg): """Load vision-language model via transformers (LLaMA-Vision, Qwen2-VL, etc.).""" from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training - from transformers import AutoModelForVision2Seq, AutoProcessor, BitsAndBytesConfig + from transformers import AutoModelForVision2Seq, AutoProcessor console.print(f"[dim]Loading vision processor: {cfg.base}[/]") self.processor = AutoProcessor.from_pretrained( @@ -856,19 +856,16 @@ class SFTTrainerWrapper: ) self.tokenizer = self.processor # SFTTrainer uses processing_class - # Quantization - bnb_config = None - if tcfg.quantization == "4bit": - from soup_cli.utils.gpu import get_compute_dtype + # Quantization (v0.71.19 #81) — unified Quant Menu loader. Replaces the + # inline BitsAndBytesConfig block so vision training gets the full menu + # (gptq / awq / hqq:Nbit / aqlm / eetq / mxfp4 / fp8 + bnb 4bit/8bit). + from soup_cli.utils.quant_menu import build_quantization_config_for_loader - bnb_config = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=get_compute_dtype(), - bnb_4bit_use_double_quant=True, - ) - elif tcfg.quantization == "8bit": - bnb_config = BitsAndBytesConfig(load_in_8bit=True) + quant_config_obj = build_quantization_config_for_loader( + tcfg=tcfg, + base=cfg.base, + console=console, + ) console.print(f"[dim]Loading vision model: {cfg.base}[/]") dev_map = "cpu" if self.device == "cpu" else "auto" @@ -876,8 +873,8 @@ class SFTTrainerWrapper: "trust_remote_code": self._trust_remote_code, "device_map": dev_map, } - if bnb_config: - model_kwargs["quantization_config"] = bnb_config + if quant_config_obj is not None: + model_kwargs["quantization_config"] = quant_config_obj self.model = AutoModelForVision2Seq.from_pretrained(cfg.base, **model_kwargs) @@ -943,7 +940,7 @@ class SFTTrainerWrapper: """Load audio-language model via transformers (Qwen2-Audio, Whisper, etc.).""" from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training from rich.panel import Panel as RichPanel - from transformers import AutoModel, AutoProcessor, BitsAndBytesConfig + from transformers import AutoModel, AutoProcessor console.print( RichPanel( @@ -962,19 +959,16 @@ class SFTTrainerWrapper: ) self.tokenizer = self.processor # SFTTrainer uses processing_class - # Quantization - bnb_config = None - if tcfg.quantization == "4bit": - from soup_cli.utils.gpu import get_compute_dtype + # Quantization (v0.71.19 #81) — unified Quant Menu loader. Replaces the + # inline BitsAndBytesConfig block so audio training gets the full menu + # (gptq / awq / hqq:Nbit / aqlm / eetq / mxfp4 / fp8 + bnb 4bit/8bit). + from soup_cli.utils.quant_menu import build_quantization_config_for_loader - bnb_config = BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=get_compute_dtype(), - bnb_4bit_use_double_quant=True, - ) - elif tcfg.quantization == "8bit": - bnb_config = BitsAndBytesConfig(load_in_8bit=True) + quant_config_obj = build_quantization_config_for_loader( + tcfg=tcfg, + base=cfg.base, + console=console, + ) console.print(f"[dim]Loading audio model: {cfg.base}[/]") dev_map = "cpu" if self.device == "cpu" else "auto" @@ -982,8 +976,8 @@ class SFTTrainerWrapper: "trust_remote_code": self._trust_remote_code, "device_map": dev_map, } - if bnb_config: - model_kwargs["quantization_config"] = bnb_config + if quant_config_obj is not None: + model_kwargs["quantization_config"] = quant_config_obj # Use AutoModel for audio models — AutoModelForCausalLM doesn't handle # audio-language architectures (Qwen2-Audio, Whisper, etc.) diff --git a/src/soup_cli/utils/multipack_trainer.py b/src/soup_cli/utils/multipack_trainer.py index 73b1d01..acaf0ee 100644 --- a/src/soup_cli/utils/multipack_trainer.py +++ b/src/soup_cli/utils/multipack_trainer.py @@ -167,6 +167,12 @@ def make_multipack_trainer_class(base_cls: type) -> type: delegate to ``super().get_train_dataloader()`` so the subclass is safe to instantiate even when multipack is disabled at runtime (matches the ``_get_train_sampler`` fallback policy). + + v0.71.19 #80 — when running distributed (``num_processes > 1`` via + FSDP / DeepSpeed ZeRO / DDP) the loader is routed through + ``self.accelerator.prepare`` so accelerate shards the packed bins + across ranks. The single-process path returns the raw DataLoader + unchanged. """ lengths = getattr(self, _LENGTHS_ATTR, None) max_seq = getattr(self, _MAX_SEQ_ATTR, None) @@ -210,7 +216,7 @@ def make_multipack_trainer_class(base_cls: type) -> type: ) data_collator = getattr(self, "data_collator", None) - return DataLoader( + loader = DataLoader( train_dataset, batch_sampler=sampler, collate_fn=data_collator, @@ -218,6 +224,40 @@ def make_multipack_trainer_class(base_cls: type) -> type: pin_memory=pin_memory, ) + # v0.71.19 #80 — under FSDP / DeepSpeed ZeRO / DDP the packed bins + # must be sharded across ranks, otherwise every rank trains on the + # SAME data. Route the DataLoader through ``accelerator.prepare`` — + # exactly what HF Trainer's own ``get_train_dataloader`` does — so + # accelerate wraps ``batch_sampler`` in a ``BatchSamplerShard`` that + # round-robins whole bins to each rank (preserving the FFD packing). + # With accelerate's default ``even_batches=True`` the per-rank batch + # count is equalised, which is the mechanism that avoids a + # collective-op hang at the epoch boundary. The ``seed`` is identical + # on every rank (set from the same config), so all ranks agree on the + # global bin order BEFORE sharding — do NOT add ``+ rank`` to it. + # Full multi-GPU validation is a QA issue (no multi-GPU box here); + # this path is mocked-tested. + # + # Gated on ``num_processes > 1`` so the single-process path stays on + # the raw DataLoader (byte-for-byte the validated v0.40.4 behaviour; + # HF Trainer's ``_prepare_inputs`` handles single-device placement). + # When no accelerator is present (older transformers / a direct + # unit-test construction) also fall back to the raw loader. The + # ``isinstance(int)`` + ``not isinstance(bool)`` guard is defence-in + # -depth: ``MagicMock().num_processes > 1`` is truthy and ``bool`` is + # an ``int`` subclass. + accelerator = getattr(self, "accelerator", None) + prepare = getattr(accelerator, "prepare", None) + num_processes = getattr(accelerator, "num_processes", 1) + if ( + callable(prepare) + and isinstance(num_processes, int) + and not isinstance(num_processes, bool) + and num_processes > 1 + ): + return prepare(loader) + return loader + MultipackTrainer.__name__ = f"Multipack{base_cls.__name__}" MultipackTrainer.__qualname__ = MultipackTrainer.__name__ return MultipackTrainer diff --git a/tests/test_quant_menu.py b/tests/test_quant_menu.py index d08748a..0fc101b 100644 --- a/tests/test_quant_menu.py +++ b/tests/test_quant_menu.py @@ -549,17 +549,21 @@ training: {quantization: hqq:4bit} """ ) - def test_quant_menu_with_vision_modality_rejected(self): - with pytest.raises(ValueError, match="modality|vision|v0.38.1"): - load_config_from_string( - """ + def test_quant_menu_with_vision_modality_now_accepted(self): + # v0.71.19 (#81) — vision / audio modality wiring landed: the SFT + # vision/audio setup paths thread build_quantization_config_for_loader, + # so the Quant Menu modality gate is dropped (only mlx still rejected). + cfg = load_config_from_string( + """ base: m task: sft modality: vision data: {train: d.jsonl, format: llava} training: {quantization: gptq} """ - ) + ) + assert cfg.modality == "vision" + assert cfg.training.quantization == "gptq" # --------------------------------------------------------------------------- diff --git a/tests/test_v0405_part_a.py b/tests/test_v0405_part_a.py index 58e8561..076c26f 100644 --- a/tests/test_v0405_part_a.py +++ b/tests/test_v0405_part_a.py @@ -78,19 +78,22 @@ class TestQuantMenuMultiTrainerMlxRejection: load_config_from_string(_build_yaml(task, "gptq", backend="mlx")) -class TestQuantMenuVisionStillSftOnly: - """Modality gate is independent — Quant Menu still text-only.""" +class TestQuantMenuVisionNowSupported: + """v0.71.19 (#81) — the modality gate was dropped: the SFT vision/audio + setup paths thread the unified Quant Menu loader, so vision/audio configs + with a quant-menu format now load (only the mlx-backend gate remains).""" - def test_vision_modality_rejected_even_after_widening(self): - with pytest.raises(ValueError, match="modality|vision"): - load_config_from_string( - """base: m + def test_vision_modality_now_accepted(self): + cfg = load_config_from_string( + """base: m task: sft modality: vision data: {train: d.jsonl, format: llava} training: {quantization: gptq} """ - ) + ) + assert cfg.modality == "vision" + assert cfg.training.quantization == "gptq" # --------------------------------------------------------------------------- diff --git a/tests/test_v07119.py b/tests/test_v07119.py new file mode 100644 index 0000000..28caad5 --- /dev/null +++ b/tests/test_v07119.py @@ -0,0 +1,410 @@ +"""Tests for v0.71.19 — Quant-menu + multipack hardening. + +Closes: + +* #81 — Quant Menu (gptq / awq / hqq:Nbit / aqlm / eetq / mxfp4 / fp8) was + rejected by the SoupConfig modality gate for ``modality in {vision, audio}``. + The vision / audio ``_setup_*`` paths in ``sft.py`` carried inline + ``BitsAndBytesConfig`` blocks (4bit / 8bit only). v0.71.19 drops the gate and + threads the unified ``build_quantization_config_for_loader`` through both + paths so multi-modal training can use the full quant menu. + +* #80 — the multipack ``get_train_dataloader`` override built a raw + ``DataLoader`` and returned it directly, so under FSDP / DeepSpeed ZeRO / DDP + every rank trained on the SAME packed bins (no data sharding). v0.71.19 routes + the DataLoader through ``accelerator.prepare`` when ``num_processes > 1`` so + accelerate's ``BatchSamplerShard`` shards whole bins across ranks. The + single-process path is unchanged (raw DataLoader, the validated v0.40.4 + behaviour). Full multi-GPU validation stays a QA issue; this is a mocked-env + test of the routing. +""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +from soup_cli.config.loader import load_config_from_string + +# Anchor source reads on the repo root derived from this file's location so the +# source-grep tests survive another test's ``monkeypatch.chdir`` (cwd leak) in +# the full suite — matches the v0.71.5 precedent. +_REPO = Path(__file__).resolve().parent.parent +_SFT_SRC = (_REPO / "src/soup_cli/trainer/sft.py").read_text(encoding="utf-8") +_MP_SRC = (_REPO / "src/soup_cli/utils/multipack_trainer.py").read_text(encoding="utf-8") +_SCHEMA_SRC = (_REPO / "src/soup_cli/config/schema.py").read_text(encoding="utf-8") + + +# --------------------------------------------------------------------------- +# #81 — Quant Menu for vision / audio modality (schema gate dropped) +# --------------------------------------------------------------------------- + + +class TestVisionAudioQuantMenuSchema: + """The ``modality != 'text'`` Quant Menu rejection is gone; vision / audio + configs with a quant-menu format now load.""" + + # Every quant-menu format the unified loader supports, on BOTH multi-modal + # modalities — covers the docstring-advertised menu incl. aqlm / mxfp4 / fp8. + @pytest.mark.parametrize("modality", ["vision", "audio"]) + @pytest.mark.parametrize( + "fmt", ["gptq", "awq", "hqq:4bit", "aqlm", "eetq", "mxfp4", "fp8"] + ) + def test_quant_menu_format_accepted(self, modality, fmt): + data_fmt = "llava" if modality == "vision" else "audio" + cfg = load_config_from_string( + f"base: m\ntask: sft\nmodality: {modality}\n" + f"data: {{train: d.jsonl, format: {data_fmt}}}\n" + f"training: {{quantization: {fmt}}}\n" + ) + assert cfg.modality == modality + assert cfg.training.quantization == fmt + + @pytest.mark.parametrize("modality", ["vision", "audio"]) + @pytest.mark.parametrize("fmt", ["4bit", "8bit"]) + def test_bnb_formats_still_accepted(self, modality, fmt): + # bnb 4bit/8bit were always universal — regression guard for both + # modalities (the inline-BNB blocks they replace handled only these). + data_fmt = "llava" if modality == "vision" else "audio" + cfg = load_config_from_string( + f"base: m\ntask: sft\nmodality: {modality}\n" + f"data: {{train: d.jsonl, format: {data_fmt}}}\n" + f"training: {{quantization: {fmt}}}\n" + ) + assert cfg.training.quantization == fmt + + def test_text_gptq_still_accepted(self): + # Text path unchanged — regression guard. + cfg = load_config_from_string( + """ +base: TheBloke/Llama-2-7B-GPTQ +task: sft +data: {train: d.jsonl} +training: {quantization: gptq} +""" + ) + assert cfg.modality == "text" + assert cfg.training.quantization == "gptq" + + def test_mlx_quant_menu_still_rejected(self): + # mlx backend gate is independent of modality and must still fire. + with pytest.raises(ValueError, match="mlx"): + load_config_from_string( + """ +base: m +task: sft +backend: mlx +data: {train: d.jsonl} +training: {quantization: hqq:4bit} +""" + ) + + +class TestVisionAudioQuantMenuWiring: + """Source-level proof that the vision / audio setup paths use the unified + quant-menu loader and dropped the inline ``BitsAndBytesConfig`` blocks.""" + + def test_vision_setup_uses_unified_loader(self): + # The vision setup method threads build_quantization_config_for_loader. + # Split on the ``def`` (not the call site in setup()) to isolate the body. + vision_block = _SFT_SRC.split("def _setup_vision_transformers")[1].split( + "def _prepare_vision_dataset" + )[0] + assert "build_quantization_config_for_loader" in vision_block + + def test_audio_setup_uses_unified_loader(self): + audio_block = _SFT_SRC.split("def _setup_audio_transformers")[1].split( + "def _prepare_audio_dataset" + )[0] + assert "build_quantization_config_for_loader" in audio_block + + def test_vision_setup_no_inline_bnb_config(self): + vision_block = _SFT_SRC.split("def _setup_vision_transformers")[1].split( + "def _prepare_vision_dataset" + )[0] + # The inline BitsAndBytesConfig construction is gone (the unified + # loader builds it lazily). Importing BitsAndBytesConfig in the method + # is also gone. + assert "BitsAndBytesConfig(" not in vision_block + assert "import BitsAndBytesConfig" not in vision_block + + def test_audio_setup_no_inline_bnb_config(self): + audio_block = _SFT_SRC.split("def _setup_audio_transformers")[1].split( + "def _prepare_audio_dataset" + )[0] + assert "BitsAndBytesConfig(" not in audio_block + assert "import BitsAndBytesConfig" not in audio_block + + def test_vision_setup_still_prepares_kbit(self): + # prepare_model_for_kbit_training is still gated on bnb formats so + # 4bit/8bit/mxfp4 still run through kbit-prep in vision/audio. + vision_block = _SFT_SRC.split("def _setup_vision_transformers")[1].split( + "def _prepare_vision_dataset" + )[0] + assert "prepare_model_for_kbit_training" in vision_block + + def test_schema_no_longer_rejects_non_text_modality(self): + gate = _SCHEMA_SRC.split("_validate_quant_menu_supported_tasks")[1].split( + "def _validate_preference_dispatcher" + )[0] + # The modality-specific rejection branch is removed. + assert "modality='text' only" not in gate + assert 'self.modality != "text"' not in gate + # The mlx gate is retained. + assert "mlx" in gate + + +# --------------------------------------------------------------------------- +# #80 — multipack get_train_dataloader hardening under FSDP / DeepSpeed / DDP +# --------------------------------------------------------------------------- + + +# Module-level sentinel: passed as ``accelerator`` to mean "the base class has +# no ``accelerator`` attribute at all" (vs ``accelerator=None``). +_NO_ACCEL = object() + + +def _torch_or_skip(): + try: + from torch.utils.data import DataLoader, Dataset + except ImportError: # pragma: no cover - torch always present in [dev] + pytest.skip("torch not installed") + return DataLoader, Dataset + + +def _make_base_class(*, accelerator): + """Return a fresh base class for ``make_multipack_trainer_class`` whose + instances expose ``train_dataset`` / ``data_collator`` / ``args`` and the + supplied ``accelerator`` (or no ``accelerator`` attr when ``_NO_ACCEL``).""" + _, dataset_cls = _torch_or_skip() + + class TinyDataset(dataset_cls): + def __len__(self): + return 12 + + def __getitem__(self, idx): + return {"x": idx} + + class _Base: + def __init__(self): + self.train_dataset = TinyDataset() + self.data_collator = None + self.args = MagicMock( + dataloader_num_workers=0, + dataloader_pin_memory=False, + dataloader_drop_last=False, + ) + if accelerator is not _NO_ACCEL: + self.accelerator = accelerator + + def get_train_dataloader(self): + return "super-dl" + + return _Base + + +def _attach(instance): + from soup_cli.utils.multipack_trainer import attach_multipack_state + + attach_multipack_state( + instance, + lengths=[10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 12, 24], + max_seq_len=128, + batch_size=2, + seed=7, + ) + + +class TestMultipackDistributedDataLoader: + def test_routes_through_accelerator_prepare_when_distributed(self): + dataloader_cls, _ = _torch_or_skip() + from soup_cli.utils.multipack_sampler import MultipackBatchSampler + from soup_cli.utils.multipack_trainer import make_multipack_trainer_class + + sentinel = object() + captured: dict = {} + + def _prepare(loader): + captured["loader"] = loader + return sentinel + + accel = MagicMock() + accel.num_processes = 2 # simulate FSDP / ZeRO / DDP + accel.prepare.side_effect = _prepare + + base_cls = _make_base_class(accelerator=accel) + sub = make_multipack_trainer_class(base_cls) + instance = sub() + _attach(instance) + + result = instance.get_train_dataloader() + + accel.prepare.assert_called_once() + assert result is sentinel + prepared = captured["loader"] + assert isinstance(prepared, dataloader_cls) + assert isinstance(prepared.batch_sampler, MultipackBatchSampler) + # The sampler is built with the ATTACHED seed (7), identical on every + # rank — sharding happens after, in accelerate's BatchSamplerShard. A + # regression that hardcoded seed=0 or added `+ rank` would break the + # cross-rank global-order invariant; assert the seed propagated. + assert prepared.batch_sampler._seed == 7 + # real_batches=False contract preserved: each pack is a flat list[int]. + first = next(iter(prepared.batch_sampler)) + assert isinstance(first, list) + assert all(isinstance(x, int) for x in first) + + def test_single_process_returns_raw_dataloader(self): + dataloader_cls, _ = _torch_or_skip() + from soup_cli.utils.multipack_sampler import MultipackBatchSampler + from soup_cli.utils.multipack_trainer import make_multipack_trainer_class + + accel = MagicMock() + accel.num_processes = 1 # single GPU / CPU + accel.prepare.side_effect = AssertionError("should not be called") + + base_cls = _make_base_class(accelerator=accel) + sub = make_multipack_trainer_class(base_cls) + instance = sub() + _attach(instance) + + dl = instance.get_train_dataloader() + accel.prepare.assert_not_called() + assert isinstance(dl, dataloader_cls) + assert isinstance(dl.batch_sampler, MultipackBatchSampler) + # Single-process path must keep the v0.40.4 flat-yield contract + # (real_batches=False → list[int] per pack), not silently flip to + # real_batches=True. + first = next(iter(dl.batch_sampler)) + assert isinstance(first, list) + assert all(isinstance(x, int) for x in first) + + def test_no_accelerator_attr_returns_raw_dataloader(self): + dataloader_cls, _ = _torch_or_skip() + from soup_cli.utils.multipack_sampler import MultipackBatchSampler + from soup_cli.utils.multipack_trainer import make_multipack_trainer_class + + # Build a base with NO accelerator attribute at all (older transformers + # / direct construction). + base_cls = _make_base_class(accelerator=_NO_ACCEL) + sub = make_multipack_trainer_class(base_cls) + instance = sub() + assert not hasattr(instance, "accelerator") + _attach(instance) + + dl = instance.get_train_dataloader() + assert isinstance(dl, dataloader_cls) + assert isinstance(dl.batch_sampler, MultipackBatchSampler) + + def test_accelerator_none_returns_raw_dataloader(self): + dataloader_cls, _ = _torch_or_skip() + from soup_cli.utils.multipack_sampler import MultipackBatchSampler + from soup_cli.utils.multipack_trainer import make_multipack_trainer_class + + base_cls = _make_base_class(accelerator=None) + sub = make_multipack_trainer_class(base_cls) + instance = sub() + _attach(instance) + + dl = instance.get_train_dataloader() + assert isinstance(dl, dataloader_cls) + assert isinstance(dl.batch_sampler, MultipackBatchSampler) + + def test_unconfigured_magicmock_accelerator_does_not_prepare(self): + # Defence-in-depth: an accelerator whose num_processes is a MagicMock + # (not a real int) must NOT route through prepare — `MagicMock() > 1` + # is truthy, so the isinstance(int) guard is required. + dataloader_cls, _ = _torch_or_skip() + from soup_cli.utils.multipack_trainer import make_multipack_trainer_class + + accel = MagicMock() # num_processes is an auto-MagicMock (not int) + accel.prepare.side_effect = AssertionError("should not be called") + + base_cls = _make_base_class(accelerator=accel) + sub = make_multipack_trainer_class(base_cls) + instance = sub() + _attach(instance) + + dl = instance.get_train_dataloader() + accel.prepare.assert_not_called() + assert isinstance(dl, dataloader_cls) + + def test_falls_back_to_super_when_state_missing(self): + from soup_cli.utils.multipack_trainer import make_multipack_trainer_class + + accel = MagicMock() + accel.num_processes = 4 + base_cls = _make_base_class(accelerator=accel) + sub = make_multipack_trainer_class(base_cls) + instance = sub() + # No attach_multipack_state — must delegate to super (no prepare). + assert instance.get_train_dataloader() == "super-dl" + accel.prepare.assert_not_called() + + def test_empty_lengths_falls_back_to_super(self): + # The `not lengths` arm of the state guard: attrs present but lengths is + # an empty list (unreachable via attach_multipack_state, which rejects + # it — so poke the attrs directly to exercise the defensive branch). + from soup_cli.utils.multipack_trainer import ( + _BATCH_SIZE_ATTR, + _LENGTHS_ATTR, + _MAX_SEQ_ATTR, + make_multipack_trainer_class, + ) + + accel = MagicMock() + accel.num_processes = 1 + base_cls = _make_base_class(accelerator=accel) + sub = make_multipack_trainer_class(base_cls) + instance = sub() + setattr(instance, _LENGTHS_ATTR, []) + setattr(instance, _MAX_SEQ_ATTR, 128) + setattr(instance, _BATCH_SIZE_ATTR, 2) + assert instance.get_train_dataloader() == "super-dl" + + def test_drop_last_forwarded_through_override(self): + # v0.40.4 H3 regression, behaviourally — `args.dataloader_drop_last` + # must reach the MultipackBatchSampler through the #80-refactored path. + _torch_or_skip() + from soup_cli.utils.multipack_trainer import make_multipack_trainer_class + + accel = MagicMock() + accel.num_processes = 1 # raw path so we can inspect the sampler + base_cls = _make_base_class(accelerator=accel) + sub = make_multipack_trainer_class(base_cls) + instance = sub() + instance.args.dataloader_drop_last = True + _attach(instance) + + dl = instance.get_train_dataloader() + assert dl.batch_sampler._drop_last is True + + +class TestMultipackSourceWiring: + def test_override_routes_through_accelerator_prepare(self): + # The override reads self.accelerator and routes the loader through + # accelerate's prepare under distribution. + assert 'getattr(self, "accelerator"' in _MP_SRC + assert "prepare(loader)" in _MP_SRC + # Gated on num_processes > 1 so the single-process path is unchanged. + assert "num_processes" in _MP_SRC + + def test_drop_last_still_forwarded(self): + # Regression guard from v0.40.4 H3 — must survive the #80 refactor. + assert 'getattr(args, "dataloader_drop_last"' in _MP_SRC + assert "drop_last=drop_last" in _MP_SRC + + +# --------------------------------------------------------------------------- +# Patch 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, 19), soup_cli.__version__