RomanNum3ral
/
soup
mirror of https://github.com/razor-ai/soup.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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.
This commit is contained in:
Alpamys 2026-06-09 12:37:50 +05:00
parent ae6a18e7d7
commit f51331d637
11 changed files with 543 additions and 70 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
CHANGELOG.md
View File

@ -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

2
CONTRIBUTING.md
View File

@ -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
```

24
README.md
View File

@ -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).

2
pyproject.toml
View File

@ -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"

2
src/soup_cli/__init__.py
View File

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

15
src/soup_cli/config/schema.py
View File

@ -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")

54
src/soup_cli/trainer/sft.py
View File

@ -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.)

42
src/soup_cli/utils/multipack_trainer.py
View File

@ -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

14
tests/test_quant_menu.py
View File

@ -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"
# ---------------------------------------------------------------------------

17
tests/test_v0405_part_a.py
View File

@ -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"
# ---------------------------------------------------------------------------

410
tests/test_v07119.py Normal file
View File

@ -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__