[Amix29]

Refs #840

What

This remains a draft: the implementation and review fixes are present, while the issue-required CUDA acceptance grid remains deliberately unclaimed.

• carries the exact pre-flight plan through streaming setup for preference tasks;
• defers DPO/ORPO/SimPO/KTO measurement until the real TRL trainer exists;
• measures trainer.compute_loss through the same temporary Accelerate native-AMP/FP32-output forward wrapper that the real training loop prepares, then restores the original dispatch;
• never truncates an overlength trainer batch and refuses partial batches whose effective row count is below the plan;
• treats a non-finite probe loss as a failed measurement;
• snapshots and restores TRL metric stores, cumulative token count, Python/NumPy/CPU/CUDA RNG state, model mode and forward dispatch;
• fails closed at the start of all four train() methods if a deferred probe was not consumed;
• shares the existing measured-fit policy with SFT and closes the stream runtime on every setup/probe failure.

Real-device regression added after G1

Two gpu rows now build real tiny streamed Llama trainers for DPO and KTO, use their actual TRL batches, and drive measure_loss_step_peak_bytes through the production AMP wrapper. A warm-up creates the persistent per-thread cuBLAS workspaces; the test then measures seq 32 and seq 64 from the same resident baseline, requires peak growth with shape, and bounds only the step's incremental allocation by the fitted prediction. It also pins correct shape metadata, positive duration and gradient cleanup.

The CPU contract now verifies that reset_peak_memory_stats runs before the measured loss step, so CI catches a stale high-water mark even without CUDA.

Deliberate non-change

The schema still refuses training.stream_vram_probe on preference tasks. #840 requires the CUDA grid before enabling losses one by one; this M4 Max cannot supply that evidence.

Validation on Apple M4 Max

• six-file review set: 136 passed, 21 expected skips;
• full Ruff and git diff --check clean.

Hardware follow-up required before ready-for-review

• rerun the two corrected DPO/KTO instrument rows on the named CUDA card;
• complete Layer streaming: extend training.stream_vram_probe to dpo/orpo/simpo/kto by measuring the loss's own step, not a causal-LM stand-in #840's probe-vs-real-step grid after Fix preference max length enforcement across TRL wrappers #1196: DPO, ORPO, SimPO and KTO, at least two shapes each, one at max_length >= 4096, KTO batch >= 2, recording probe peak, real-step peak, prediction and cost;
• widen the schema/docs/changelog gate only for losses whose grid is conservative.

No CUDA result is claimed from this Mac.

[MakazhanAlpamys]

Advisory review of the draft at 632c3fd8. Keep it in draft. Route 1 from #840 is the right route, and leaving the schema gate, docs and fragment until the grid exists is right too. But the instrument under-reads the real step, and that is the direction #840's criterion 2 forbids.

What holds:

• The probe builds the real loss, not an imitation. Forward hooks show the same forwards and shapes as the real step:
  • DPO: a policy pass with grad plus the ref adapter without grad, at [2, seq];
  • KTO: KL, policy, ref and ref-KL;
  • ORPO/SimPO: one concatenated forward.
• Dropping the rejected half of the batch is caught.
• Nothing new is reachable by a user: the schema still refuses the probe for task != sft, so no preference plan is ever built.

Measured on an RTX 5070 8 GB (sm_120), torch 2.14.0+cu130, transformers 5.17.0, trl 0.29.1, peft 0.20.0, accelerate 1.15.0. The model is a tiny Llama (vocab 32000, hidden 256, 2 layers), streamed in bf16 with LoRA r4 on q/v, fed rows that fill max_length. I compared the probe's peak with the real first training_step peak, both max_memory_allocated, with no other GPU work on the card. One toy model and two shapes, so these are a new baseline, not comparable with the old 3050 record.

1. Blocking, the core defect: the probe does not run the forward the Trainer will use. It calls trainer.compute_loss(...) from setup(), before _inner_training_loop runs accelerator.prepare(self.model). Under bf16 native AMP, prepare replaces model.forward with autocast plus convert_outputs_to_fp32 (accelerate/accelerator.py:1824-1835). So the probe sees bf16 logits with autocast off, while the real step has fp32 logits with autocast on. Probe peak / real-step peak:

loss	seq 512	seq 2048
DPO	0.664	0.666
ORPO	0.405	0.372
SimPO	0.404	0.372
KTO (batch 2)	0.523	0.508

• Control: the existing SFT causal-LM probe on the same box and model reads 0.997 / 1.112.
• Attribution: your probe with model.forward wrapped exactly as accelerate wraps it reads 0.94-0.99. The wrapped DPO probe reads 646.3 MB, the same as the real step's 646.3 MB, so the gap is entirely the missing wrapper.
• Fix: run the probe through the forward the Trainer will use, then restore it, and pin it with a CUDA test asserting that the probe and the real step see the same logits dtype and autocast state. The logits term dominates here, and it is larger still on 128k-152k vocabularies.

2. Blocking: the stretch truncates, and the truncation hides real overruns. _resize_probe_tensor cuts batches longer than max_length. But on trl 0.29.1, Soup's SimPO and KTO wrappers never call enforce_preference_sequence_limit (DPO and ORPO call it 3 times each), and TRL's own truncation lets rows run past max_length:

• CPO truncates the response to max_length - longer_response_length, which overruns whenever the prompt is longer than the longer response;
• KTO slices with a negative bound once the prompt alone reaches max_length.

Real steps against the probe:

loss	max_length	real step	probe	probe / real
SimPO	512	[2, 1025]	[2, 512]	0.218
SimPO	2048	[2, 4097]	[2, 2048]	0.190
KTO	512	[2, 1024]	[2, 512]	0.282
KTO	2048	[2, 4096]	[2, 2048]	0.259

The truncated SimPO batch had no completion tokens left, so its log-probs were NaN, and the instrument accepted a NaN loss as a valid measurement. Never truncate: either measure the batch as TRL produced it, or refuse and name the overrun. Treat a non-finite loss as failed. The overrun itself predates this PR; I am filing it separately.

3. Blocking: the probe writes into the trainer state the run then reports. All four compute_loss implementations store metrics.

• A one-step DPO run logs num_tokens 2048 with the probe and 1024 without it, and that counter is cumulative, so the error is permanent.
• ORPO/SimPO's first logged rewards are the mean of the probe batch and the real step.

Snapshot and restore the metric stores and the token counter around the probe. Also observed, not claimed as harm: the CPU RNG advances across the probe, which runs after the Trainer's seeding point.

4. Blocking: the deferred gate fails open. With the probe on, a predicted miss up to 4x is demoted from refusal to advice. The plan is set in one place and consumed in another, and nothing checks that it was consumed. So a wrapper without the hook, or with the hook in a branch that does not run, trains with no gate at all: disabling the DPO hook, or moving it before materialize_meta_adapter_copy, leaves every test green. Refuse in train() while _pending_stream_vram_probe is still set, and replace the source-grep test with a behavioural one through the real setup(); the tests/test_v07204.py harness can do that.

5. Blocking before this leaves draft: the new instrument has no direct tests. The one hook test monkeypatches it out. These mutations all survive:

• no backward;
• peak read from max_memory_reserved;
• dropping _zero_probe_grads;
• a zero-filled attention-mask extension;
• not closing the stream runtime when building the batch raises;
• routing preference back to the causal-LM step.

Mirror the real-device class in test_issue349_measured_fit.py:435-535, and add a dispatch test.

Non-blocking:

• The plan's rows is only a label: the measured row count is whatever the first dataloader batch holds, which is fewer than batch_size on a small dataset. Assert the batch's leading dimension.
• Every probe run leaves ~33.5 MB resident: the autograd thread's cuBLAS workspace, freed by torch._C._cuda_clearCublasWorkspaces(). It predates this PR, since the SFT probe does the same, but the grid should account for it.
• Nits: the SFT messages now read "a sft streaming step at 1 model rows". The deleted comments on the three-outcome policy (None / failed / oom) were its only documentation.

Before this leaves draft: fix 1-5, then run #840's grid on a named card, publish it under benchmarks/, and widen the gate one loss at a time, together with the docs and the fragment.

[Amix29]

Implemented review blockers 1–5 at bb5b5595 after rebasing onto current main.

The probe now uses Accelerate’s training-time AMP/output wrapper and restores it; refuses overlength and partial batches; rejects non-finite loss; restores metrics, cumulative tokens, RNG and model state; and all four trainers fail closed if setup leaves a pending plan. Direct instrument and real-setup regressions were added.

Local focused result: 111 passed, 35 skipped, Ruff clean, diff-check clean. The PR remains draft because the named-CUDA acceptance grid required by #840 cannot be executed or claimed on this M4 Max.

[MakazhanAlpamys]

@Amix29 thanks, this is a solid round. Blockers 1-5 from my review of 632c3fd8 are closed, and I checked each one by running it rather than by reading the diff. I ran everything on CPU: Windows, torch 2.14.0 with CUDA hidden, transformers 5.17.0, trl 0.29.1, peft 0.20.0, accelerate 1.15.0. What follows is what I executed at bb5b5595. Keep it in draft: the CUDA grid is still the gate, as you say.

What holds, by execution

• 1, the forward. I checked _trainer_forward_for_probe (stream_setup.py:388) against a real Accelerator(mixed_precision="bf16", cpu=True). Your wrapped forward and a twin wrapped by accelerate's own prepare_model behave the same:

  • both run with autocast on, compute bf16 logits and return fp32, with identical values;
  • with "no" precision, both stay plain fp32.

  It also puts back the class forward, a forward already set on the instance, and the forward after an exception. The memory effect itself (the 0.37-0.67 under-read) can only be confirmed on the grid.

• 2, overlength, partial and non-finite batches. I built real trainers with max_length: 64 and a 100-token prompt:

  • DPO (2, 35) and ORPO (1, 35)/(1, 34) are measured;
  • SimPO chosen (1, 104) and KTO completion (2, 102) are refused with "received a sequence of length 104, above the configured data.max_length=64; refusing to truncate …";
  • the stream runtime is closed once in each refused case.

• 3, trainer state. I ran _run_pending_stream_vram_probe on real DPO/ORPO/SimPO/KTO wrappers built through the real setup(), then a full train(). The run logged exactly the same history as a control without the probe, at every logged step: DPO 8, ORPO 8, SimPO 8, KTO 4 (DPO's first num_tokens is 7.0 in both). With _preserve_preference_probe_state switched off, DPO logs num_tokens 134 against 7, and the first-step metrics of the other three are contaminated. So the restore is what does the work, and the attribute names match real trl 0.29.1.

• 4, fail closed. I injected a plan through the real setup(), at the point where the pre-flight produces one (_stream_budget_lines patched, since CPU never builds a plan):

  • all four trainers consume it through the loss instrument only, leave nothing pending, and train;
  • with the DPO hook deleted, train() refuses with your message, so a lost hook can no longer train without a gate.

• 5, instrument tests. Five of the six mutations I listed last time are now killed: reserved peak, no grad zeroing, zero-filled mask, runtime not closed, and routing the pending probe to the causal-LM instrument.

Tests I ran

• Your six files (test_issue840…, test_issue349…, test_issue649…, test_issue654…, test_issue1003…, test_v07204.py): 126 passed, 20 skipped, the same 146 tests as your 111 / 35 on the M4 Max. Also 126 / 20 on the tree merged with current main.
• Ruff and git diff --check are clean.
• The new test_issue840 file on main: 13 failed, 1 skipped.
• Mutations, one at a time against those six files: 12 killed, 6 survived. The survivors are below.

Small things to fold in before the grid run. All of them can be tested on CPU.

1. _preserve_preference_probe_state adds attributes (stream_setup.py:426-445). copy.deepcopy(missing) returns a new object: deepcopy(object()) is sentinel is False. So if value is missing never fires, and any name the trainer did not have comes back as a bare object(). On the real trainers:

   • DPOTrainer gains _stored_metrics;
   • CPOTrainer (SimPO), ORPOTrainer and KTOTrainer gain _metrics and _total_train_tokens.

   It is harmless on trl 0.29.1 (the logs above are identical), but it is the same "the probe writes trainer state" class as blocker 3. Fix: snapshot only the names in vars(trainer), and delete on exit the ones that were absent. Then assert set(vars(trainer)) is unchanged in test_pending_probe_executes_trainer_compute_loss.

2. No backward still survives (layer_stream_runtime.py:1928). test_loss_instrument_uses_allocated_peak_runs_backward_and_clears_grads asserts that the grads are None afterwards, which is just as true when backward() never ran. Count a parameter hook, or assert that the grads were populated before _zero_probe_grads.

3. The hook and the setup dispatch are not pinned.

   • Deleting self._run_pending_stream_vram_probe() (dpo.py:254) leaves every test green. It fails closed, but every preference probe run would then refuse and no test would say why.
   • Changing if cfg.task == "sft": (stream_setup.py:1145) to always take the causal-LM branch also stays green. The preference plan is then "measured" by the SFT stand-in at setup, and training proceeds.

   test_v07204.py:402 injects the plan after setup(), so it cannot see either change. Add one test per trainer through the real setup(), with _stream_budget_lines returning a plan, asserting that the loss instrument is called once, the causal-LM instrument never, nothing is left pending, and train() runs. That kills both mutations. The harness is about 60 lines on top of _tiny_llama_dir: patch _stream_budget_lines to return a _ProbePlan, and record both instruments.

4. The train/eval mode restore is unpinned (stream_setup.py:1493). Start the model in eval() in test_pending_probe_executes_trainer_compute_loss and assert it is still in eval afterwards.

5. The AMP wrapper is pinned only by a gpu test (test_issue840…py:313). Removing the wrapper, or its restore, passes every CPU test. The same test also works on CPU: patch get_mixed_precision_context_manager to torch.autocast("cpu", dtype=torch.bfloat16), use a CPU Linear, and check torch.is_autocast_enabled("cpu"). I ran exactly that: it passes on your head and fails on both mutations. It needs no Accelerator and leaves no global state behind.

6. Two nits from last time are still open.

   • The SFT refusal still reads "one sft step at 1 model rows x seq 512" (stream_setup.py:1532-1533), and that message is user-visible today.
   • The comments explaining the None / failed / oom policy are still gone from _finish_stream_vram_probe.

What has to happen before a full review

1. Items 1-5 above.
2. Layer streaming: extend training.stream_vram_probe to dpo/orpo/simpo/kto by measuring the loss's own step, not a causal-LM stand-in #840's grid on a named CUDA card: DPO, ORPO, SimPO and KTO, at least two shapes each with one at seq >= 4096, KTO at batch >= 2, and for each row the probe peak against the real first-step peak, the prediction and the probe cost. You can't run it on the M4 Max. The maintainer runs it on the RTX 5070 when the card is free.
3. Order matters: Fix preference max length enforcement across TRL wrappers #1196 (restoring max_length enforcement for SimPO/KTO, closes SimPO and KTO rows run past max_length on trl 0.29: the two wrappers never call enforce_preference_sequence_limit #1185) should be merged before that run. Until then, SimPO and KTO rows with long prompts are refused by your overlength check, by design, instead of being measured.
4. Then widen the schema gate one loss at a time, together with the field description, the docs, a fragment and the gate record under benchmarks/, as Layer streaming: extend training.stream_vram_probe to dpo/orpo/simpo/kto by measuring the loss's own step, not a causal-LM stand-in #840 asks.

All 13 required checks pass on bb5b5595. The branch merges cleanly with main, #1143, #1196, #1080 and #1178.

[Amix29]

Folded in all six CPU-review follow-ups at 0a3e6474 after rebasing onto current main.

• _preserve_preference_probe_state now snapshots only attributes actually present in vars(trainer) and removes probe-created ones; the regression pins the complete attribute set.
• The loss-instrument test now proves backward populated gradients before cleanup.
• A real-setup() integration covers DPO, ORPO, SimPO, and KTO: it injects a plan through _stream_budget_lines, requires the loss instrument exactly once, rejects the causal-LM instrument, verifies the plan was consumed, and reaches each wrapper's real train() entry point. Its CPU device state is isolated from Accelerate's process-wide MPS singleton.
• Model eval/train mode restoration is pinned.
• The AMP-forward contract now runs on CPU bf16 autocast, so wrapper use and restoration are covered without a CUDA skip.
• Restored the None/failed/OOM policy comments and fixed the user-facing singular 1 model row wording.

Local M4 Max result on the six review files: 116 passed, 34 skipped. Ruff (including preview E30) and git diff --check are clean.

Keeping the PR in draft: I have not run or claimed #840's named-CUDA acceptance grid, and the requested #1196-before-grid ordering still applies.

[MakazhanAlpamys]

@Amix29 thank you. This round closes all six CPU follow-ups from my review of bb5b5595. I checked each one by breaking the fix and watching a test go red, not by reading the diff.

• The state snapshot now keys off vars(trainer).
• The backward is proven by grads that were really populated before cleanup.
• The per-trainer test goes through the real setup(). It fails for exactly the wrapper whose hook I deleted, and for all four when setup is forced onto the causal-LM probe.
• The AMP-forward contract now runs on CPU instead of hiding behind a gpu marker.

Everything below ran at 0a3e6474 on CPU (Windows, torch 2.14.0 with CUDA hidden, transformers 5.17.0, trl 0.29.1, peft 0.20.0, accelerate 1.15.0).

Tests I ran

• Your six files (test_issue840…, test_issue349…, test_issue649…, test_issue654…, test_issue1003…, test_v07204.py): 131 passed, 19 skipped. That is the same 150 tests as your 116 / 34 on the M4 Max.
• The same files on your head merged with current main (20ab4668, which now includes fix(layer-stream): prefetch embed_tokens at the backward tail, not the next step's prime (#975) #1178): 131 passed, 19 skipped, and the merge is clean.
• fix(trainer): wire LoRA+ into the remaining thirteen trainers (#745) #1080 changed the same four wrappers after your base. On the merged tree:
  • its LoRA+ tests (test_issue724_loraplus_wiring_coverage.py): 108 passed;
  • test_issue745_*, test_trl_preference_config_contract.py, test_kto.py, test_orpo.py and test_simpo.py: 153 passed.
• The new tests on main: all 19 fail, plus the changed assertion in test_issue349.
• Ruff and git diff --check are clean. All 19 new tests pass in all nine CI cells, and all 13 required checks are green.

Mutations: 28 against those six files, 24 killed

• The six items:
  • the old deepcopy-sentinel snapshot;
  • no backward;
  • each wrapper's hook deleted (4 mutations);
  • setup always taking the causal-LM probe;
  • no train/eval restore;
  • no AMP wrapper;
  • no forward restore;
  • "1 model rows".
• The earlier blockers, re-checked on the rebased code:
  • truncation, partial batch and non-finite loss;
  • no metric/RNG snapshot;
  • the fail-closed check as a no-op, and DPO skipping it;
  • peak read from reserved;
  • no grad cleanup and a zero-filled mask;
  • runtime not closed and plan never cleared;
  • KL not stretched and a chosen/rejected pair counted as one row.

Three small gaps, none blocking. Four mutations survive, and three of them are restore paths that nothing pins yet:

1. stream_setup.py:444-445: deleting the elif name in vars(trainer): delattr(...) branch passes everything. No test's compute_loss creates a name the trainer did not have before the probe, and that is the only case this branch exists for.
2. :415-416: putting back a forward that was already set on the model instance is unpinned. Replacing model.forward = instance_forward with delattr(model, "forward") passes.
3. :412-418: the restore when the measured step raises is unpinned. Replacing try/finally with plain sequencing passes.

These three tests kill all three mutations. On your head they pass, 3/3:

from types import SimpleNamespace

import pytest

torch = pytest.importorskip("torch")


def test_an_attribute_the_probe_creates_is_removed_again(monkeypatch):
    from soup_cli.trainer.stream_setup import StreamingSetupMixin, _ProbePlan
    from soup_cli.utils.layer_stream_runtime import StepPeak

    class Runtime:
        def close(self):
            pass

    class Trainer:
        def __init__(self):
            self._metrics = {"train": {"loss": [1.0]}}

        def get_train_dataloader(self):
            return [{"input_ids": torch.tensor([[1, 2, 3], [1, 4, 5]]),
                     "attention_mask": torch.ones((2, 3), dtype=torch.long)}]

        def compute_loss(self, model, batch):
            self._stored_metrics = {"train": {"rewards": [9.0]}}  # created by the loss
            return model.weight.sum()

    class Wrapper(StreamingSetupMixin):
        pass

    wrapper = Wrapper()
    wrapper.model = torch.nn.Linear(4, 4, bias=False)
    wrapper.trainer = Trainer()
    wrapper.device = "cuda"
    wrapper._stream_runtime = Runtime()
    wrapper._pending_stream_vram_probe = _ProbePlan(
        task="dpo", batch_size=1, rows=2, seq_len=8, vocab_size=64,
        predicted_bytes=100, available_bytes=1_000,
    )

    def fake_measure(model, *, step, rows, seq_len, device):
        step()
        return StepPeak(peak_bytes=200, reserved_bytes=220, seconds=0.01,
                        rows=rows, seq_len=seq_len)

    monkeypatch.setattr(
        "soup_cli.utils.layer_stream_runtime.measure_loss_step_peak_bytes", fake_measure
    )
    before = set(vars(wrapper.trainer))
    wrapper._run_pending_stream_vram_probe()
    assert "_stored_metrics" not in vars(wrapper.trainer)
    assert set(vars(wrapper.trainer)) == before


def _autocast_cpu(monkeypatch):
    monkeypatch.setattr(
        "accelerate.utils.modeling.get_mixed_precision_context_manager",
        lambda *_args, **_kwargs: torch.autocast("cpu", dtype=torch.bfloat16),
    )
    return SimpleNamespace(accelerator=SimpleNamespace(native_amp=True, autocast_handler=None))


def test_an_instance_forward_is_put_back(monkeypatch):
    from soup_cli.trainer.stream_setup import _trainer_forward_for_probe

    trainer = _autocast_cpu(monkeypatch)
    model = torch.nn.Linear(8, 8)

    def instance_forward(x):
        return torch.nn.Linear.forward(model, x)

    model.forward = instance_forward
    with _trainer_forward_for_probe(trainer, model):
        assert model(torch.randn(2, 8)).dtype == torch.float32
        assert vars(model)["forward"] is not instance_forward
    assert vars(model)["forward"] is instance_forward


def test_the_forward_is_put_back_when_the_step_raises(monkeypatch):
    from soup_cli.trainer.stream_setup import _trainer_forward_for_probe

    trainer = _autocast_cpu(monkeypatch)
    model = torch.nn.Linear(8, 8)
    with pytest.raises(RuntimeError, match="step failed"):
        with _trainer_forward_for_probe(trainer, model):
            assert "forward" in vars(model)
            raise RuntimeError("step failed")
    assert "forward" not in vars(model)

The fourth survivor is cosmetic: labels[:, original:] = resized_ids[:, original:] (:361) can be dropped without a failure, because in every fixture the last label equals the last input id. One fixture row ending in -100 would pin it, if you want it pinned.

For later, not for now: when the gate widens, the refusal at schema.py:5890 still justifies itself with "the probe measures a causal-LM step". After this PR, a preference task would be measured through its own loss and would stay refused only because it has not been measured. That message changes together with the field description, as #840 already asks.

What remains before this leaves draft. None of it is yours to do until the grid exists.

1. Fix preference max length enforcement across TRL wrappers #1196 has to land first. It restores max_length for SimPO and KTO, and it still has open review items. On the tree merged with Fix preference max length enforcement across TRL wrappers #1196, your hook runs after its cap in simpo.py and kto.py, so the probe will measure the capped batch. That is the right order. DPO and ORPO already cap in their wrappers, so those rows do not depend on Fix preference max length enforcement across TRL wrappers #1196.
2. The CUDA part, on the RTX 5070:
   • the 19 gpu-marked tests in these six files (pytest -m gpu --no-cov -v on the six files: 5 in test_issue349, 14 in test_v07204);
   • Layer streaming: extend training.stream_vram_probe to dpo/orpo/simpo/kto by measuring the loss's own step, not a causal-LM stand-in #840's grid: DPO, ORPO, SimPO and KTO, at least two shapes each with one at max_length >= 4096, KTO at batch >= 2 with its KL forward in the measured step, and rows that genuinely fill max_length (checked on the collated batch).
   • Each row gets the probe peak against the real first-step peak (both max_memory_allocated), the prediction and the probe cost, published under benchmarks/ with the card and the trl version.
   • The criterion is Layer streaming: extend training.stream_vram_probe to dpo/orpo/simpo/kto by measuring the loss's own step, not a causal-LM stand-in #840's: the probe never reads below the real step at any row.
   • A grid harness has to force the probe on, because the schema refuses it for these tasks. I checked that setting cfg.training.stream_vram_probe = True after load_config_from_string sticks, since the model does not re-validate on assignment.
3. Then widen the gate in this PR, one passing loss at a time, with the field description, the docs (docs/performance-and-quantization.md:606 and :681), a test that the message and the description agree, and a fragment.

If you fold the three tests in before step 2, the card only has to run the final head once.

[Amix29]

Folded in all three suggested restore-path tests plus the cosmetic label-tail pin at 3a62c0ec, after rebasing onto current main.

• a trainer attribute created by compute_loss is removed after the probe;
• an existing instance-level forward is restored by identity;
• forward is restored when the measured step raises;
• stretched labels ending in -100 are explicitly overwritten from the stretched token ids.

Six-file local result: 120 passed, 34 skipped on the M4 Max. Ruff, preview E30 and git diff --check are clean. The PR stays draft; #1196 and the named-CUDA grid remain external prerequisites, and no CUDA result is claimed.

[MakazhanAlpamys]

@Amix29 thank you. 3a62c0ec closes the three restore-path gaps and the label-tail pin from my review of 0a3e6474. I checked each one by breaking the code and watching the new test go red:

• deleting the elif name in vars(trainer): delattr(...) branch (stream_setup.py:444-445) now fails test_an_attribute_the_probe_creates_is_removed_again;
• putting back an instance forward with delattr instead of restoring it (:415-416) now fails test_an_instance_forward_is_put_back;
• replacing the try/finally around the step with plain sequencing (:412) now fails test_the_forward_is_put_back_when_the_step_raises;
• dropping labels[:, original:] = resized_ids[:, original:] (:361) now fails test_stretched_labels_follow_ids_when_the_original_row_ends_ignored.

Your label test does one more thing than I asked for. It also checks that the original columns keep their -100. So overwriting the whole row with the ids (original = 0, :358) fails too, and the prompt mask is now pinned as well as the tail.

I also checked the rebase. git range-diff shows your three earlier commits unchanged, so the only change since my last review is this test-only commit.

Everything below ran at 3a62c0ec on CPU (Windows, torch 2.14.0 with CUDA hidden, transformers 5.17.0, trl 0.29.1, peft 0.20.0, accelerate 1.15.0).

Tests I ran

• Your six files (test_issue840…, test_issue349…, test_issue649…, test_issue654…, test_issue1003…, test_v07204.py): 135 passed, 19 skipped. That is the same 154 tests as your 120 / 34 on the M4 Max.
• The same files on your head merged with current main (df106781): 135 passed, 19 skipped. The merge is clean.
• Three open PRs touch the same files, so I ran each one merged together with yours:
  • with feat(stream): make NF4 fused forward checkpoint-visible (#842) #1143 (layer_stream_runtime.py), both PRs' review files: 327 passed, 42 skipped;
  • with Fix preference max length enforcement across TRL wrappers #1196 (simpo.py and kto.py), plus test_trl_preference_config_contract.py, test_kto.py and test_simpo.py: 242 passed, 19 skipped. On that tree your hook still runs after Fix preference max length enforcement across TRL wrappers #1196's max_length cap in both wrappers;
  • with fix(train): evaluate the validation split, or stop withholding it (#1223) #1313 (all four wrappers), plus its three test_issue1223_* files: 255 passed, 19 skipped.
• The new tests on main: all 23 fail (18 in test_issue840, including your four new ones, and 5 in test_v07204), plus the changed assertion in test_issue349.
• Ruff and git diff --check are clean.

Mutations: 29 against your six files, 29 killed. That covers the six items from my review of bb5b5595 and the earlier blockers, re-run on the rebased code:

• the snapshot, the backward, each wrapper's hook, the setup dispatch, the train/eval restore, the AMP wrapper and its restores, and the "1 model row" wording;
• truncation, partial batches, non-finite loss, the metric/RNG snapshot, the fail-closed check, the allocated-vs-reserved peak, grad cleanup, the mask fill, runtime closing, clearing the plan, the KL stretch and the pair row count;
• the two label mutations above.

Nothing survives. On 0a3e6474 four did, and each one is now caught by the test you added for it.

CI. The required checks on 3a62c0ec have been queued since 11:13Z and have not started yet, so there is no CI result for this head.

What remains before this leaves draft. Nothing from you until the grid exists, and the order has not changed:

1. Fix preference max length enforcement across TRL wrappers #1196 lands first. It still has open review items on its side, and the SimPO and KTO rows depend on it. The DPO and ORPO rows do not.
2. The CUDA part, on the RTX 5070:
   • the 19 gpu tests in your six files (pytest -m gpu --no-cov -v);
   • Layer streaming: extend training.stream_vram_probe to dpo/orpo/simpo/kto by measuring the loss's own step, not a causal-LM stand-in #840's grid: DPO, ORPO, SimPO and KTO, at least two shapes each with one at max_length >= 4096, KTO at batch >= 2, and rows that genuinely fill max_length. Each row records the probe peak against the real first-step peak, the prediction and the probe cost.
3. Then the gate widens in this PR, one passing loss at a time. That change carries the refusal in schema.py (currently :5992-6008, including its "the probe measures a causal-LM step" wording), the field description, the docs (docs/performance-and-quantization.md:370, :609 and :684 say "sft only" today), a test that the message and the description agree, and a fragment.

[MakazhanAlpamys]

@Amix29, G1 is done on the card. RTX 5070, on 3a62c0ec merged with main: pytest -m gpu on the six files gives 19 passed and none skipped (the 5 #349 real-hardware probe tests and the 14 test_v07204 DPO/ORPO/SimPO/KTO tests).

G2, the #840 probe-vs-real-step grid and its gate record, is still to come. It needs the harness, which is not in the tree yet, and it waits for #1196 as agreed, so nothing changes on your side for now. When the harness is ready, the grid runs on the same card.

[MakazhanAlpamys]

@Amix29, a follow-up on G1. I re-ran it today on the RTX 5070, on 3a62c0ec merged with current main (452d071f). The merge is clean.

• pytest -m gpu on your six files gives 19 passed, none skipped or failed: 5 in test_issue349, 14 in test_v07204. That matches yesterday's run on the older main.

I also ran the same 19 with your source changes reverted to main, meaning all six files under src/, with your tests kept. They still pass, 19 of 19. On that reverted tree the CPU half of your six files fails 24: 18 in test_issue840, 5 in test_v07204, and the changed assertion in test_issue349. So the revert was real.

That is expected, and it is not a gap in your tests. The 19 GPU tests predate this PR:

• the five in test_issue349 call the causal-LM probe, measure_step_peak_bytes, which you did not change;
• the fourteen in test_v07204 go through the real setup() with no probe planned, so your hook returns at once, and they call TRL's train() directly, so the wrapper's fail-closed check is never reached.

So G1 shows that your change breaks nothing in a real streamed DPO, ORPO, SimPO or KTO step on CUDA. It does not show measure_loss_step_peak_bytes working on a device: every test of it so far fakes torch.cuda, and its first real run will be G2.

G2 is unchanged. The #840 grid waits for #1196, which is still in review, and for the harness. Nothing is needed from you until then.

One request that would make the GPU run count for this PR: a @pytest.mark.gpu test that drives measure_loss_step_peak_bytes for real on a tiny preference model (one DPO and one KTO step is enough), asserting the measured peak is positive and within the fitted bound. Today every test of the new instrument fakes torch.cuda, so G1 cannot fail on this PR's code; with that test it would, and G2 would no longer be the first real run of the feature.

[Amix29]

Added the requested real-device preference-probe test at 4df267a8.

The two new gpu rows build real tiny streamed Llama trainers for DPO and KTO, take their actual TRL dataloader batches, stretch them to the configured seq 64, run measure_loss_step_peak_bytes through the same Accelerate AMP forward wrapper as production, and assert: successful positive peak, peak within the fitted prediction, correct shape metadata, positive duration, and no surviving gradients. The runtime/model references are released between rows.

Local M4 Max validation: the six review files give 135 passed, 21 expected skips; the final focused pair gives 87 passed, 16 expected skips. Full Ruff and git diff --check are clean. The two new CUDA rows were collected but skipped here and still need the maintainer RTX run; no CUDA result is claimed.

[MakazhanAlpamys]

@Amix29 thank you for turning the GPU request around the same day. The test drives the real thing in the right order: a wrapper built through the real setup(), the real TRL dataloader batch, then _stretch_preference_probe_batch, _preference_probe_rows, the Accelerate AMP forward and measure_loss_step_peak_bytes, exactly as _run_pending_stream_vram_probe does it. It also releases the runtime between rows. You were right not to claim a CUDA result. I ran it on the RTX 5070, and it fails there, but not because of your instrument. The bound I asked for is what fails.

What I ran. RTX 5070 Laptop (sm_120), torch 2.14.0+cu130, trl 0.29.1, peft 0.20.0, accelerate 1.15.0, on 4df267a8 merged with current main (dccc8680). The merge is clean.

• CPU, your six review files: 135 passed, 21 skipped, the same as your M4 Max count. Ruff and git diff --check are clean.
• pytest -m gpu on the six files: 19 passed, 2 failed. The 19 are the tests that were already there. The two failures are your new rows, at tests/test_v07204.py:709:
  • dpo: measured 67734016 bytes against fitted bound 13934816
  • kto: measured 67855872 bytes against fitted bound 13934816

Why. On this card, cuBLAS holds a 32 MiB workspace per handle, the allocator counts it, and a training step uses two handles: the forward thread's and the autograd thread's. The fitted formula has no term for that. For a model this small, its prediction is almost entirely STREAM_FIXED_SLACK_BYTES = 13_500_000 (layer_stream.py:185), plus 434,816 bytes of modelled terms. I checked the attribution three ways:

• Before the DPO probe, 33,731,072 bytes are allocated: 32 MiB plus 176,640 for the model. Afterwards, with the grads cleared, the count is exactly 33,554,432 bytes (32 MiB) higher. That is the autograd thread's workspace, created by the first backward.
• With the backward removed, the peak falls to 34,060,288, which leaves one workspace.
• With CUBLAS_WORKSPACE_CONFIG=:4096:1 (4 MiB per handle), the peak is 9,013,760 (DPO) and 8,971,776 (KTO), 0.65 of the prediction, and your two rows pass unchanged. The difference from the default run is exactly 56.0 MiB, which is 2 x (32 - 4) MiB.

So the step itself costs about 0.6 MB, and the instrument reads it correctly. The assertion cannot hold at this shape on this card, and "within the fitted bound" was my wording, so that part is on me.

A replacement I ran on the card. It keeps your structure and your per-call assertions, adds one warm-up probe so the workspaces are already allocated, and then takes the resident allocation as the baseline. It measures at seq 32 and at seq 64, and it bounds the growth above the resident set by the fitted prediction:

@pytest.mark.gpu(reason="the preference-loss peak instrument needs CUDA")
class TestPreferenceLossProbeOnRealHardware:
    """Drive #840's new instrument through real TRL DPO and KTO losses on CUDA.

    The first backward in a process creates the autograd thread's cuBLAS
    workspace (32 MiB per handle on this card) and it stays allocated, so a
    first measurement also counts memory the step does not own. One warm-up
    probe makes every later measurement start from the same resident set.
    """

    @pytest.mark.parametrize("task", ("dpo", "kto"))
    def test_real_loss_peak_moves_with_the_shape_and_fits_the_fitted_bound(
        self, tmp_path, monkeypatch, task
    ):
        import gc

        import torch

        from soup_cli.trainer.stream_setup import (
            _preference_probe_rows,
            _stretch_preference_probe_batch,
            _trainer_forward_for_probe,
        )
        from soup_cli.utils import layer_stream as layer_stream_module
        from soup_cli.utils.layer_stream_runtime import measure_loss_step_peak_bytes

        predicted = []
        real_estimate = layer_stream_module.estimate_stream_peak_vram

        def capture_estimate(**kwargs):
            value = real_estimate(**kwargs)
            predicted.append(value)
            return value

        monkeypatch.setattr(layer_stream_module, "estimate_stream_peak_vram", capture_estimate)
        wrapper = trainer = model = raw = None
        try:
            wrapper, _, _ = _build_streamed_wrapper(tmp_path, monkeypatch, task=task, n_layers=2)
            assert len(predicted) == 1
            trainer = wrapper.trainer
            model = wrapper.model
            raw = next(iter(trainer.get_train_dataloader()))

            def probe(seq_len):
                batch = _stretch_preference_probe_batch(raw, seq_len=seq_len)
                rows = _preference_probe_rows(batch)
                batch = _batch_on(model, batch)
                model.train()
                with _trainer_forward_for_probe(trainer, model):
                    peak = measure_loss_step_peak_bytes(
                        model,
                        step=lambda: trainer.compute_loss(model, batch),
                        rows=rows,
                        seq_len=seq_len,
                        device="cuda",
                    )
                assert peak is not None and not peak.oom and not peak.failed
                assert peak.rows == rows and peak.seq_len == seq_len and peak.seconds > 0
                assert all(parameter.grad is None for parameter in model.parameters())
                return peak

            probe(64)  # warm-up: creates the per-thread cuBLAS workspaces
            torch.cuda.synchronize()
            resident = torch.cuda.memory_allocated()
            short = probe(32)
            full = probe(64)
            assert resident < short.peak_bytes < full.peak_bytes, (
                f"{task}: resident {resident}, seq 32 {short.peak_bytes}, "
                f"seq 64 {full.peak_bytes}"
            )
            assert full.peak_bytes - resident <= predicted[0], (
                f"{task}: the step added {full.peak_bytes - resident} bytes above "
                f"the resident {resident}, against a fitted bound of {predicted[0]}"
            )
        finally:
            if wrapper is not None:
                wrapper._close_stream_runtime()
            raw = model = trainer = wrapper = None
            gc.collect()
            torch.cuda.empty_cache()

It replaces your class in tests/test_v07204.py as it stands, because it uses only that module's _build_streamed_wrapper and _batch_on. Ruff passes with the repo config. On the card:

• At your head: 2 passed, in two separate runs. DPO: resident 67,283,968, seq 32 67,556,864, seq 64 67,735,552, so the step adds 451,584 bytes against the 13,934,816 bound. KTO adds 413,696.
• It fails when your source is taken from main (ImportError), and it fails under three single mutations of measure_loss_step_peak_bytes:
  • grads not cleared: the grads survive;
  • peak read from max_memory_reserved: seq 32 and seq 64 both read 69,206,016;
  • no reset_peak_memory_stats: both read 234,898,432, the high-water mark left by setup.
• Two mutations pass it on the card, and the CPU tests catch both. No backward is caught by test_loss_instrument_uses_allocated_peak_runs_backward_and_clears_grads, and no AMP forward wrapper by the AMP contract test and the two restore tests. I re-ran both mutations on CPU today, and each turned those tests red.

For comparison, your current class fails under every one of those mutations, but only because it already fails at the head. So on this card it cannot tell a working instrument from a broken one yet.

One small CPU gap, optional. Deleting the reset_peak_memory_stats call (layer_stream_runtime.py:1959) passes all six CPU files (135 passed). The replacement above catches it on the card. If you also want CI to catch it, this test passes at your head and fails without the call (['step'] == ['reset', 'step']):

def test_loss_instrument_resets_the_peak_counter_before_the_step(monkeypatch):
    import torch

    from soup_cli.utils.layer_stream_runtime import measure_loss_step_peak_bytes

    events = []
    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
    monkeypatch.setattr(torch.cuda, "synchronize", lambda *args: None)
    monkeypatch.setattr(
        torch.cuda, "reset_peak_memory_stats", lambda *args: events.append("reset")
    )
    monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None)
    monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda *args: 123)
    monkeypatch.setattr(torch.cuda, "max_memory_reserved", lambda *args: 456)
    model = torch.nn.Linear(4, 1)
    x = torch.randn(3, 4)

    def step():
        events.append("step")
        return model(x).square().mean()

    peak = measure_loss_step_peak_bytes(model, step=step, rows=3, seq_len=1, device="cuda")
    assert peak is not None and not peak.failed
    assert events == ["reset", "step"]

CI. The required checks on 4df267a8 have been queued since 14:53Z and have not started. The new rows would skip there anyway, because CI has no GPU.

What remains before this leaves draft.

1. The GPU class above, or your own version of it. Once it is pushed, the -m gpu set needs one more run on the card.
2. G2, Layer streaming: extend training.stream_vram_probe to dpo/orpo/simpo/kto by measuring the loss's own step, not a causal-LM stand-in #840's probe-vs-real-step grid, still waits for Fix preference max length enforcement across TRL wrappers #1196. Fix preference max length enforcement across TRL wrappers #1196 is still in review with changes requested, and the harness does not exist yet. Nothing else is needed from you for that part.
3. Then the gate widens one loss at a time, together with the schema message, the field description, the docs and a fragment.

[Amix29]

Replaced the two failing CUDA assertions with the resident-baseline protocol you validated, at head dfc4531b.

The DPO/KTO rows now warm up once to create the persistent cuBLAS workspaces, measure the resident allocation, require peak growth from seq 32 to seq 64, and compare only full.peak_bytes - resident with the fitted bound. The existing per-probe shape, duration and gradient-cleanup assertions remain.

I also added the optional CPU contract that proves reset_peak_memory_stats occurs before the loss step, closing that surviving mutation in CI.

Local M4 Max validation across the six review files: 136 passed, 21 expected skips. Full Ruff and git diff --check are clean. The corrected DPO/KTO GPU rows remain unclaimed here and are ready for the RTX rerun.

[MakazhanAlpamys]

@Amix29, thank you for turning the GPU test around the same day, and for using the resident-baseline version as written. dfc4531b has no open question for you, and the PR stays in draft for the reasons below, none of which is yours to do.

What I checked at dfc4531b, on CPU (Windows, CUDA hidden), merged with current main (4b5f4a3c, the merge is clean):

• git diff --stat 3a62c0ec dfc4531b touches only tests/test_issue840_preference_vram_probe.py (+28) and tests/test_v07204.py (+86). src/ is unchanged, so the earlier mutation results still describe the source.
• Your six review files: 136 passed, 21 skipped, the same count as your M4 Max run. The 21 skips are the 19 existing gpu tests plus your two rewritten DPO/KTO rows.
• Ruff and git diff --check are clean.
• The PR's tests against main's source: 25 fail.
• New CPU test test_loss_instrument_resets_the_peak_counter_before_the_step: I deleted torch.cuda.reset_peak_memory_stats(device) inside measure_loss_step_peak_bytes (layer_stream_runtime.py:2229) and it fails (1 failed, 18 passed in that file), so that surviving mutation is closed in CI.
• The GPU class TestPreferenceLossProbeOnRealHardware matches the one in the earlier comment; the only differences are line wrapping. I can't run -m gpu here, so its result on the RTX 5070 is still to be re-run on the card.
• All required checks are green on this head (lint, mlx-smoke, transformers-floor, pytorch-smoke and the nine test cells); GitHub reports the PR as mergeable and clean.

What has to happen before it leaves draft. Nothing from you until these exist:

1. Fix preference max length enforcement across TRL wrappers #1196 lands first (it restores max_length for SimPO and KTO). It is still open with changes requested.
2. The CUDA run on the RTX 5070: the 21 gpu rows on the six files, including your two new ones, then Layer streaming: extend training.stream_vram_probe to dpo/orpo/simpo/kto by measuring the loss's own step, not a causal-LM stand-in #840's grid (DPO, ORPO, SimPO, KTO; at least two shapes each with one at max_length >= 4096; KTO at batch >= 2), published under benchmarks/ with the card and the trl version. The harness for the grid is not in the tree yet.
3. Then the gate widens in this PR, one loss at a time, together with the schema.py refusal wording, the field description, the docs lines that say "sft only", a test that the message and the description agree, and a fragment.

Thank you for how carefully this has been iterated: every round has closed what was asked with a test that goes red when the code is broken.

[Amix29]

Rebased this PR onto current origin/main (1d52d3a3). The only textual conflict was in src/soup_cli/trainer/simpo.py; the resolution preserves the newer #1208 empty-completion guards and then runs the pending streaming VRAM probe.

New head: 4d5ce353

Local validation:

• 158 passed, 25 skipped across the six preference/VRAM regression files plus test_issue1208_simpo_zero_token_completion.py
• standard Ruff scope clean
• Ruff preview E30 clean on all files changed by this PR
• git diff --check origin/main...HEAD clean

No CUDA result is claimed; this remains the CPU/MPS/MLX-side implementation and regression coverage.

[MakazhanAlpamys]

@Amix29, thank you for the rebase. You resolved the simpo.py conflict with #1500 in the order I would have picked: the #1208 refusals come first, because they are a cheap data check, and the probe hook comes after them.

What I checked at 4d5ce353 (CPU, Windows, CUDA hidden, trl 0.29.1):

• git range-diff against dfc4531b shows your last three commits unchanged. The first three differ only in context lines: the SimPO trains to NaN, behind a logged loss of 0.0, when chosen is longer than max_length #1208 calls in simpo.py, plus one import line each in stream_setup.py and tests/test_v07204.py. None of your own added or removed lines changed, so the earlier mutation results still describe the code.
• Your six review files plus tests/test_issue1208_simpo_zero_token_completion.py: 162 passed, 21 skipped. These are the same 183 tests as your 158 / 25. The 21 skips are the gpu rows.
• A trial merge of this head with main and with Fix preference max length enforcement across TRL wrappers #1196's head (4a3edc94) was clean. On that tree, simpo.py runs the max_length cap first, then the SimPO trains to NaN, behind a logged loss of 0.0, when chosen is longer than max_length #1208 refusal, then your hook, and kto.py runs the cap and then the hook. Those seven files, plus test_trl_preference_config_contract.py, test_kto.py and test_simpo.py, give 288 passed, 21 skipped. ruff is clean.
• The quick CI on this head is green (lint and test (ubuntu-latest, 3.12)), and GitHub no longer reports a conflict. A merge of the current main into this head is still clean.

What still has to happen before this leaves draft. None of it is owed by you today.

1. Fix preference max length enforcement across TRL wrappers #1196 lands. Three small items are left on it, and it merges cleanly with your branch.
2. The CUDA run on the RTX 5070:
   • the 21 gpu rows of your six files, including the two resident-baseline DPO/KTO rows, which have not yet run on a card in their current form;
   • Layer streaming: extend training.stream_vram_probe to dpo/orpo/simpo/kto by measuring the loss's own step, not a causal-LM stand-in #840's grid: DPO, ORPO, SimPO and KTO; at least two shapes each, one of them at max_length >= 4096; KTO at batch >= 2; rows that really fill max_length. Each cell records the probe peak, the real first-step peak, the prediction and the probe's cost, and the grid is published under benchmarks/.
3. Then the gate widens in this PR, one passing loss at a time. That step brings the schema.py wording, the field description, the "sft only" docs lines, a test that the message and the description agree, and a fragment.

You don't need to rebase again before #1196 lands. Thank you for keeping every round this careful.