[VanshSharmaPES]

Draft. On @MakazhanAlpamys's RTX 5070 run this removes about 19-23 ms of embedding wait per step (about 22 ms down to about 0.15 ms on the disk tier), on a fixture with a 501 MiB shared slot. That is still only about 0.1% of a cold 70B step, so this is latency hygiene rather than a throughput fix. This implements option (b) from #1174.

What I changed

LargeLayerBufferPool holds embed_tokens and an untied lm_head in one shared slot. _prime() (the forward-pre-hook) reloaded embed_tokens right as the next step's forward started, a full head-sized H2D copy with nothing to overlap it against (#975).

I made StreamPrefetcher also fire that same load at the tail of the backward pass, when the checkpoint recompute for decoder layer 0 confirms the step's backward walk is done. _prime() keeps its unconditional load as a fallback (step 0, and any forward that follows a backward that never reached layer 0); from the second step on it is a same-owner no-op.

Why the refill is safe

Two separate guards, and only the second protects the GPU:

• CPU side: lm_head's backward Node sits between the loss and every decoder layer, so the autograd engine has already dispatched it, and the saved-tensor version check happens at dispatch. The embedding's own backward reads indices, never the weight values, so it never checks this slot's version.
• GPU side: when the CPU reaches layer 0, the head's backward GEMM can still be queued on the compute stream, reading the slot. What orders the refill after it is stream.wait_stream(torch.cuda.current_stream()) in LargeLayerBufferPool.load_async. Without it the refill overwrites bytes the queued backward has not read, and the gradients come out silently wrong. That wait is now the only thing protecting this refill point, and a GPU-only test pins it.

Measured (by @MakazhanAlpamys, RTX 5070 Laptop, not by me)

On an untied Llama with hidden 2048, 8 layers, vocab 128256 (a 501 MiB slot), decoder NF4, seq 512, batch 1:

configuration	embed wait() median, before to after	step median, before to after
disk tier, sync every step	22.22 to 0.15 ms	306.6 to 291.6 ms
disk tier, no per-step sync	22.67 to 0.19 ms	316.9 to 296.3 ms
RAM tier (pinned), sync every step	18.57 to 0.13 ms	257.7 to 242.4 ms

The embedding wait goes to essentially zero in every configuration. The 2-6.5% step-time change is real but close to that fixture's noise floor (identical code differs by about 6 ms between labels). The head's forward-tail wait (10.6-14.6 ms) and the decoder waits are unchanged. Option (c) was not measured.

Their correctness runs: untied SFT is bit-identical to main on CPU fp32 and CUDA bf16; untied DPO and KTO with #1061 are bit-identical to #1061 alone and to resident.

What to know

• Preference losses get part of the benefit. The prefetch fires once per backward, so DPO's no-grad reference forward still pays a real embedding load at _prime, and so do KTO's extra forwards: 1 of 2 embedding waits removed per DPO step, 1 of 4 per KTO step.
• One extra load per train(). Steady state is 2 large-layer loads per step in both trees, but the last backward's prefetch is never consumed by a later forward, so main makes 4 loads over 2 steps and this makes 5.
• Tied checkpoints have no shared slot, so the callback does nothing for them.

What I verified

• ruff check clean.
• On CPU (pytest -k "stream or 1049 or 975 or docs or changelog"): 602 passed, 35 skipped, 0 failed, after merging main.
• tests/test_issue975_backward_tail_embed_prefetch.py: the backward-tail fires, the slot already holds embed_tokens after step 1, 2 steps complete with no autograd error, and the real-load count is 5.
• The changelog fragment and the docs sentence are in.

What I have not run

The GPU-only test for the wait_stream guard (test_refill_waits_for_a_read_still_queued_on_the_compute_stream). My torch is a CPU build, so it skips on my machine. It follows the version @MakazhanAlpamys verified on the 5070: 5 of 5 pass with the guard, 5 of 5 fail with it removed.

[MakazhanAlpamys]

Advisory review of the draft at 5223a0de.

Thank you for turning (b) into code so quickly. I ran the measurement you could not, and the result is good news for (b). Everything below was run on the RTX 5070 Laptop GPU (8 GB, driver 616.92, Windows 11, Python 3.12.10, torch 2.14.0+cu130, bitsandbytes 0.50.2, transformers 5.17.0, peft 0.20.0, trl 0.29.1), on your head merged with current main (it merges cleanly, and cleanly with #1061's head too).

What holds

Correctness. I ran real trainer.train() steps (4 per arm, identical LoRA init loaded into every arm, dropout 0) and compared the loss tensor of every step, every LoRA gradient at on_pre_optimizer_step, and the final adapter weights with torch.equal:

• CPU fp32: untied SFT on this PR is bit-identical to main and to a resident run (48 of 48 gradient tensors). The same holds for untied SFT at gradient_accumulation_steps: 2, for untied ORPO and SimPO (also equal to resident), and for tied SFT and tied DPO. Tied checkpoints have no large pool, so the new callback is None there.
• Untied DPO and KTO crash the same way on main and on this PR: #1049, in step 1's backward, before your code runs. You were right about that. With #1061 merged in, untied DPO and KTO on #1061 + this PR are bit-identical to #1061 alone and to resident.
• CUDA bf16: untied SFT, two runs of main and two of this PR, all four bit-identical. Untied DPO on #1061 + this PR is bit-identical to #1061 alone, and tied SFT is unchanged.

Scheduling. I traced every large-layer load. The prefetch fires exactly once per backward. From step 2 on, _prime's load is a same-owner no-op. With gradient_accumulation_steps: 2 it fires at the tail of every micro-batch's backward, which is right, because every micro-batch's forward starts with the lookup. A no-grad forward never fires it, and an eval forward after training finds the embedding already in the slot.

Tests vs mutations (CPU, as CI runs them): reverting the wiring at layer_stream_runtime.py:2520 fails 2 of your 4 tests. So do deleting the call at :1055 and dropping the direction check at :1050. The last one is worth knowing about: the prefetch then fires in the forward, and the flag silently blocks the backward one. Removing the flag reset in prime() fails 1. Firing the refill before the head's backward (a grad hook on the head output) fails 3, with the modified by an inplace operation error, so test_two_steps... does its job. Calling the prefetch twice, or dropping the once-per-pass flag, survives. Both are equivalent mutants, because load_async's owner check makes a repeat a no-op. Against main your file gives 2 failed / 2 passed.

The measurement #1174 asked for (arms a and b). The cold synthetic-70B fixture from #975 does not fit what I could run here, so I built a smaller one with the committed synth_checkpoint.py: untied Llama, hidden 2048, 8 layers, vocab 128256. That makes the shared slot 501.0 MiB bf16. Decoder NF4, seq 512, batch 1, LoRA r8 on q/k/v/o, PagedAdamW8bit. I used the committed stream_probe.Instruments with layer0_wait's labelling, so these are the same wait() brackets as gate-971 §8.

Both arms ran in one process. Arm A is your tree with backward_tail_prefetch set to None, which is exactly the old prefetcher, and a run on unmodified main confirms that. Arm B is your change. The schedule: 3 warm-up steps discarded, then five interleaved A/B blocks, each with 1 transition step discarded and 2 measured. That gives 10 measured steps per arm.

configuration	embed wait() median, A → B	step median, A → B
disk tier (AsyncDiskSource, read_ahead 2), sync every step	22.22 → 0.15 ms	306.6 → 291.6 ms
same, repeated	22.22 → 0.16 ms	302.8 → 296.9 ms
disk tier, no per-step sync	22.67 → 0.19 ms	316.9 → 296.3 ms
RAM tier (pinned), sync every step	18.57 → 0.13 ms	257.7 → 242.4 ms
RAM tier, no per-step sync	18.72 → 0.12 ms	252.1 → 242.3 ms
control: unmodified main, both labels	22.31 / 22.61 ms	310.6 / 316.7 ms

The embedding's wait() goes to essentially zero in every configuration. On the disk tier the step's total wait drops from 38.08 to 15.33 ms. Step time falls by 5.9–20.6 ms (2–6.5%) in all five runs. Identical code differs by about 6 ms between labels, though, so the step-time gain is real but close to the noise floor on this fixture. The head's forward-tail wait (10.6–14.6 ms) and the decoder waits (~1.0 ms) are unchanged, as expected. I did not measure (c).

Blocking

1. Three of the four new tests fail on any machine with a CUDA GPU. _build_untied_wrapper pins device="cpu" (test file line 45) and then calls trainer.train(). TrainingArguments still picks CUDA, so the batches land on cuda:0 and the lookup fails with Expected all tensors to be on the same device, but got index is on cuda:0, different from other tensors on cpu. On the 5070 the file gives 3 failed / 1 passed. CI is green only because its runners have no GPU. None of these tests compares floats, so the line-55 reason for pinning the device does not apply. Fix: device="cuda" if cuda_available() else "cpu", which is what _build_streamed_wrapper and #1061's train-step test do. With that one change, all 4 pass on the 5070, and reverting the wiring still fails 2 of them there.

2. The safety comment names the wrong guard, and the guard that actually protects the new refill point is untested. The comment at :1037-1048 rests on the version counter: "dispatch, not GPU completion, is what the saved-tensor version counter checks against". That is true of the CPU-side check. It says nothing about the GPU reading the slot. When the CPU reaches layer 0, the head's backward GEMM can still be queued on the compute stream. What orders the refill after it is stream.wait_stream(torch.cuda.current_stream()) in LargeLayerBufferPool.load_async (:951). When layer 0 is reloaded during the backward, the same call at :860 is also queued earlier on the same copy stream.

   I checked this by execution. The setup: an untied 2-layer fixture with 2 buffers, so layer 0 is not reloaded in the backward; a torch.cuda._sleep(2e8) enqueued from a grad hook on the head's output, so the compute stream lags the CPU inside the backward (the sleep was still running when your prefetch fired, in steps 2–4); 4 real trainer steps; bf16. The results:

   • this PR with the lag is bit-identical to main
   • main with the lag and :951's wait removed is still bit-identical, because on main nothing refills the slot inside the backward
   • this PR with the lag and :951's wait removed diverges from step 2 on: layer-0 v_proj lora_A max |Δ| 0.068, q_proj lora_B 1.3e-3, and the losses follow

   So after this PR, :951 is the only thing standing between the new refill point and silently wrong gradients on a GPU. Please:

   • rewrite the comment to say that.
   • add a CUDA-only test (skipif without CUDA) built on that setup. It should compare against the same run with backward_tail_prefetch=None, and it must fail when :951's wait is removed. It will only run on machines with a GPU, but it is the one test that pins the hazard this change introduces.

3. The load-count test asserts the right number for the wrong reason. Lines 137–139 say "5, not 6: one fewer than two full steps … would cost". Line 126 says "without the backward-tail prefetch turning into an extra, unplanned load". On main, the same 2 steps make 4 real large-layer loads, not 6: the test fails there with assert 4 == 5. Both trees make 2 loads per step in steady state. This PR adds exactly one per train(), the last backward's prefetch, which nothing uses unless a forward follows. I measured 8 → 9 at 4 steps, and 16 → 17 at GA 2 and for DPO. Please rewrite the comment to say main = 4, PR = 5, and why.

4. Before this leaves draft:

   • add a changelog fragment, changelog.d/<current baseline>/1178.changed.md, crediting (#975 by @VanshSharmaPES in #1178).
   • add one sentence to docs/performance-and-quantization.md:354, which still says the embedding "is loaded for the model input". It is now reloaded at the tail of the previous step's backward.
   • state the measured magnitude at the top of the description, as #975's acceptance criteria ask: about 19–23 ms of embedding wait removed per step on this fixture, still about 0.1% of a cold 70B step.

Non-blocking

• Preference losses get part of the benefit. The prefetch fires once per backward. DPO's no-grad reference forward still pays a real embedding load at _prime, and so do KTO's three extra forwards. That is 1 of 2 embedding waits removed per DPO step, and 1 of 4 per KTO step. Worth one line in the description.
• What is left for (c). With (b) in, the remaining large-layer wait is the head's forward-tail load, 10.6–14.6 ms per step here. That is now what (c)'s second slot would buy, for 501 MiB here and about 1 GiB on an 8B-class head. The embedding's bytes in the slot are dead once the lookup has run, so the head load could also move earlier in the forward, at zero VRAM. I would measure that before (c). This belongs on #1174, not in this PR.
• The last step's copy is the only load that nothing ever waits on. It is harmless while the slot outlives it. If you want every load to be consumed again, synchronize the copy stream in StreamRuntime.close().
• backward_tail_prefetched is redundant. Layer 0 is recomputed once per backward. Keeping it as documentation is fine.
• _prime's comment. "a frozen embedding with no LoRA reaching it" would be clearer as "a backward that never reaches layer 0". That is the actual condition under which the fallback is the only load.

Nits

• Line 101 calls large_pool.load_async(embed_key, runtime.source, None). That re-tests the pool's existing owner check, not _prime. The load-count test already covers _prime.
• Once #1061 lands, _build_streamed_wrapper(..., tie=False) can replace the local helper. The two PRs merge cleanly, and together they are bit-exact for DPO and KTO.
• For whoever re-runs layer0_wait.py for #975: after this change, the embedding's copy bracket is recorded in the previous step's copies, and the last step carries one extra. Compare stalls with gate-971 §8, not the copy lists.

[VanshSharmaPES]

Thanks for the thorough review and for running the measurement I couldn't. I've pushed two commits so far and am working through the rest.

fa95add: blocking item 1. _build_untied_wrapper now picks "cuda" if cuda_available() else "cpu" instead of hardcoding "cpu", the same way _build_streamed_wrapper does, so the batches and the model land on the same device. I kept the MPS skip, because without CUDA the helper still falls back to the CPU and a Mac runner hits the same mismatch that broke macOS CI. I reworded its comment to say that instead of the float-pinning reason, which you're right doesn't apply here. All 4 tests pass on CPU for me, but I can't run the CUDA branch on a 4 GB card, so that half rests on your 5070 run.

2be7473: the comment half of blocking item 2. I rewrote the comment in advance() so it names both guards and says which one protects the GPU. The version counter covers the CPU-side check. stream.wait_stream(torch.cuda.current_stream()) in LargeLayerBufferPool.load_async is what orders the refill after the head's backward GEMM, and it says that removing it gives silently wrong gradients from step 2 on. Comment only, no behaviour change.

Still open: the CUDA-only test that fails when that wait is removed (I'll write it from your setup, but I can't run it here, so I'll say so in the PR), the load-count comment (main = 4, this PR = 5), the _prime comment wording and the line-101 assertion, the changelog fragment, the docs sentence, and the measured magnitude at the top of the description.

[MakazhanAlpamys]

Advisory review of the draft at 2be74739.

Thank you for turning both pushes around quickly. I re-ran everything on your head merged with main at dd9a09fd, on the same RTX 5070 as last time, and with CUDA hidden, which is what CI sees. While I was reviewing, #1031 landed in layer_stream_runtime.py. Your branch merges cleanly with it, so I repeated the test file, the mutations and the streaming suite on 1c053a8c, and everything below still holds.

Closed

Item 1: the tests on a CUDA machine. With the 5070 visible your file now gives 4 passed (41 s), and 4 passed with CUDA hidden. The mutation results are the same on both:

• reverting the wiring at :2530 fails 2 tests;
• deleting the call at :1065 fails 2;
• dropping the direction check fails 2;
• removing the flag reset in prime() fails 1;
• firing the refill before the head's backward fails 3.

Against main the file gives 2 failed / 2 passed. The 31 streaming test files give 1253 passed, 13 skipped and 0 failed on the 5070, and 1193 passed, 73 skipped and 0 failed with CUDA hidden.

Item 2, the comment half. The new comment names the right guard, and I re-checked what it says by execution rather than by reading. The setup was the 2-layer untied fixture with 2 buffers, bf16, and a torch.cuda._sleep(1e9) enqueued from a grad hook on the head's output. The sleep makes the compute stream lag the host inside the backward; at steps 2–4 it was still running when your prefetch fired. Over 4 real trainer steps:

• this PR is bit-identical to main in two runs: 32 of 32 LoRA gradients, every loss, and the final adapter;
• with the wait_stream at :951 removed, both runs diverge at step 2: all 8 gradients, worst layer-0 v_proj lora_B |Δ| 0.40, and the losses follow from step 3;
• main with the same wait removed stays bit-identical, because nothing refills the slot inside its backward.

So the comment is right that :951 is the only thing protecting the GPU.

Blocking (before this leaves draft)

1. The CUDA-only test for :951 is still missing, and the recipe I gave you last time would be a flaky test. With the wait removed, all 4 tests pass on the 5070 as well as on CPU, so nothing detects that removal on any machine.

   The flakiness comes from host load. I first ran the lag recipe with a smaller sleep while the host was busy with another test run. In 2 of 3 lag runs the sleep had already finished at every prefetch, so the race never opened, and "wait removed" came out bit-identical to your PR. A test built that way would have passed with the guard deleted.

   A pool-level test does not depend on how fast the host is. I checked this one on the 5070: 5 of 5 pass with the guard and 5 of 5 fail with it removed, about a second each, and it held while the host was loaded:

   class _PinnedSource:
       def __init__(self, tensors):
           self.tensors = tensors
   
       def get(self, idx, key):
           return self.tensors[key]
   
   
   @pytest.mark.gpu
   def test_refill_waits_for_a_read_still_queued_on_the_compute_stream():
       import torch
   
       from soup_cli.utils.layer_stream_runtime import LargeLayerBufferPool
   
       shape = (1024, 1024)
       pool = LargeLayerBufferPool(
           {"lm_head": (shape, "float32"), "embed_tokens": (shape, "float32")},
           {"lm_head": 0, "embed_tokens": 1},
           device="cuda",
       )
       head = torch.full(shape, 1.0).pin_memory()
       source = _PinnedSource({"lm_head": head, "embed_tokens": torch.full(shape, 2.0).pin_memory()})
       copy_stream = torch.cuda.Stream()
       pool.load_async("lm_head", source, copy_stream)
       slot = pool.wait("lm_head")
       torch.cuda.synchronize()
       torch.cuda._sleep(int(1e9))  # the compute stream lags, like a queued head backward
       seen = slot.clone()  # a read of the slot queued behind it
       pool.load_async("embed_tokens", source, copy_stream)  # the #975 refill
       torch.cuda.synchronize()
       assert torch.equal(seen, head.cuda())

   The host issues the refill microseconds after the read is queued, while the sleep takes about a second. So the race is always open, and the test only passes if load_async orders its copy after the read.

   If you also want a trainer-level test, it has to assert that the race actually opened, for example by recording event.query() on the sleep when the prefetch fires. Otherwise a slow host turns it into a test that always passes.

   About your note that you can't run the CUDA side on a 4 GB card: this is not a VRAM limit. The fixture is 2 layers × hidden 64, and the test above needs about 12 MB. Your 3050 can run both, as long as your torch is a CUDA build. Check with python -c "import torch; print(torch.cuda.is_available())"; on Windows, pip installs the CPU wheel unless you pass the PyTorch index.

2. The load-count comment (test lines 132 and 137–145) is unchanged. The assertion is right but the reasoning is not: on main the same 2 steps make 4 real loads, and the file fails there with assert 4 == 5. This PR adds exactly one load per train(), the last backward's prefetch. Please say main = 4, PR = 5, and why.

3. Changelog, docs and description.

   • Add the changelog fragment. Since #1126 the baseline is 0.75.1, so the file is changelog.d/0.75.1/1178.changed.md, crediting (#975 by @VanshSharmaPES in #1178).
   • Add the sentence at docs/performance-and-quantization.md:354, which still says the embedding "is loaded for the model input".
   • Update the description. It still says "performance NOT measured"; the measured numbers are in my last review (embedding wait 19–23 ms → ~0.15 ms on that fixture). It also still gives only the version-counter argument, which your new comment now says does not protect the GPU.

Non-blocking

• The items from my last review are unchanged: preference losses get part of the benefit; the head's forward-tail wait is what (c) would now buy; a copy-stream sync in close(); the redundant flag; and the _prime comment wording at :2587–2591. You listed the last one as still to do.

Nits

• The test module's docstring (lines 6–11) still gives only the CPU-side argument, "cannot invalidate a saved-tensor version check that already ran". Please point it at the GPU guard as well.
• The new comment says gradients come out wrong "from the second step on". That matches what I measured: step 1 never raced in any of the 7 lag runs. But the reason is that the first step's host-side overhead outlasts the lag, not anything in the code, and step 1's refill is issued the same way. I would drop the step qualifier.
• Line 107 still re-tests the pool's existing owner check rather than _prime.

[VanshSharmaPES]

Thanks for the second pass, and for catching that the lag recipe I'd planned would be flaky on a busy host. I've pushed bc13d61 and a merge of main (77dc04e).

Blocking 1, the CUDA-only test. I added your pool-level test as written: the compute stream is held back with torch.cuda._sleep, a read of the slot is queued behind it, then the refill is issued, and the read must still see the head's bytes. It's @pytest.mark.gpu. I haven't run it. The torch in my venv is the CPU wheel (2.14.0+cpu, is_available() is False), so it skips here and I'm relying on your 5-of-5 and 5-of-5 results. I'll try it on my 3050 once I've installed a CUDA build. I wrote the docstring to say why the pool-level version is used and that a trainer-level one would need to assert the race really opened.

Blocking 2, the load-count comment. Rewritten: main makes 4 real loads over the two steps (the cold embed and the head at step 0, then the embed and the head at step 1), this PR makes 5, and the fifth is the last backward's prefetch, which no later forward consumes. Steady state is 2 loads per step in both. I also removed the assertion that only re-tested the pool's owner check.

Blocking 3.

• Changelog fragment at changelog.d/0.75.1/1178.changed.md, crediting (#975 by @VanshSharmaPES in #1178). My branch predated 0.75.1, so test_repository_fragments_match_the_newest_release rejected it. That's why I merged main rather than only adding the file.
• One sentence added at docs/performance-and-quantization.md, saying the embedding is reloaded at the tail of the step's backward.
• I have not yet updated the PR description. I'll do that next, with your measured numbers at the top and the DPO/KTO partial-benefit line.

Nits, taken:

• The test module's docstring now names the GPU guard as well as the CPU one.
• Dropped "from the second step on" from the wait_stream comment, since the step qualifier reflects host overhead rather than the code.
• Reworded _prime's comment to "a forward that follows a backward that never reached layer 0".

Checks on the merged tree (CPU): ruff is clean, the changelog validator passes, and the streaming, docs and changelog suites give 602 passed, 35 skipped, 0 failed. Merging main brought in #1061's fix, which is why tests/test_issue1049_untied_streamed_preference.py is now there.

Not taken: the copy-stream sync in StreamRuntime.close() and removing backward_tail_prefetched. Both are fine to leave, and I'd rather not widen this PR.

[MakazhanAlpamys]

@VanshSharmaPES, thank you. Every item from the last review is addressed. Three things stand out:

• The GPU test is the pool-level version as written, and its docstring records why a trainer-level one would have to prove the race opened. That note stops someone from swapping it for the flaky variant later.
• The load-count comment is now right. I confirmed it by running your file against main, where it fails with assert 4 == 5.
• Dropping the owner-check block lost no coverage. I checked that by mutation (below).

The description now leads with the measured magnitude, the fragment and the docs sentence are in, and merging main was the right way to pick up the 0.75.1 fragment baseline.

What I ran. All of it was on CPU with CUDA_VISIBLE_DEVICES=-1, importing soup_cli from the checkout:

• Your head, 77dc04e3. The test file gives 4 passed, 1 skipped. The skipped one is the GPU test, with the conftest's "needs a CUDA device" reason.
• Your head merged with current main (210298c8), which merges cleanly. I ran every test file that touches the prefetcher, the large pool or install_streaming (32 files), plus the changelog-fragment tests: 1049 passed, 90 skipped, 0 failed. ruff check is clean.
• Two mutations on the merged tree:
  • backward_tail_prefetch=None in install_streaming, which reverts the wiring: 2 failed.
  • _prime() forced to make a real embedding copy every step, the case the removed block described: 1 failed, test_prime_no_op_saves_one_reload_from_the_second_step_on. The load count covers it.
• Your file on unmodified main. 2 failed (AttributeError: ... 'backward_tail_prefetched', and assert 4 == 5), 2 passed, 1 skipped.
• One more check, for The single shared large-layer slot forces a full embedding-matrix H2D every step on an untied checkpoint (~33 ms for 500 MiB, measured at 16 GB/s) #975. I recorded every source.get() call over 3 steps of a 4-layer untied fixture, on main and on your branch. The sequences are identical, except that yours has one more embedding get at the end of train(). Your change moves when the embedding is fetched, not the order any source sees. That is why The single shared large-layer slot forces a full embedding-matrix H2D every step on an untied checkpoint (~33 ms for 500 MiB, measured at 16 GB/s) #975's race-battery criterion does not apply to this PR.

#1141 and #1143, the other two open drafts in layer_stream_runtime.py, both merge cleanly with your branch. #1141 changes none of the test_v07204 helpers your file imports.

Before this leaves draft: one run of the GPU test, both ways. That is the only open item. The test matches the recipe line for line; only the embedding tensor got its own name. It is not deselected by default (only smoke is), so a plain pytest runs it on any CUDA machine. But nobody has run the committed file on a GPU yet. My run was CPU-only, and the five-and-five you cite was measured on the recipe before it went into your file. On a CUDA build of torch:

pytest tests/test_issue975_backward_tail_embed_prefetch.py -m gpu --no-cov -p no:cacheprovider

1. Run it five times. It should give 1 passed each time.
2. Comment out stream.wait_stream(torch.cuda.current_stream()) in LargeLayerBufferPool.load_async and run it five more times. It should give 1 failed each time.
3. Restore the line and run the whole file once on the GPU. It should give 5 passed.

Please post the counts, the card and the torch build. The test needs about 12 MB, so the 3050 is plenty. On Windows, pip installs the CPU wheel unless you point it at the PyTorch CUDA index for a CUDA version your driver supports; nvidia-smi shows the highest one. If getting a CUDA build onto that card turns out to be more trouble than it is worth, say so here. The 5070 is busy today, and a run will be scheduled on it when it is free.

After that, mark it ready for review. The two non-blocking items you left, the copy-stream sync in close() and the redundant flag, are fine as they are.

CI on 77dc04e3 is still queued: all 13 required checks are pending.

[VanshSharmaPES]

Here are the GPU counts you asked for, on 77dc04e.

Card and build: RTX 3050 Laptop GPU, torch 2.14.0+cu130 (a separate venv from my CPU one), driver 592.82.

1. Guard in place, pytest tests/test_issue975_backward_tail_embed_prefetch.py -m gpu --no-cov -p no:cacheprovider, five runs: 1 passed each time (5 of 5).
2. stream.wait_stream(torch.cuda.current_stream()) removed from LargeLayerBufferPool.load_async, five runs: 1 failed each time (5 of 5, assert False). I removed it by patching load_async at runtime with the line taken out, not by editing the file, so the tree stayed clean. It is the same removal.
3. Guard restored, whole file on the GPU: 5 passed (the four trainer tests plus the GPU one).

Earlier I got a first mutation attempt wrong: the script raised before writing, so its three passes were against unmodified code and proved nothing. I redid it and only counted runs where the guard was really gone.

I've also opened the prototype for moving the head's load earlier as a separate draft, #1258, branched from this head as you asked. Nothing in it changes what this PR does.

[MakazhanAlpamys]

@VanshSharmaPES, thank you for getting a CUDA build onto the 3050 and running this yourself. Thank you also for saying that your first mutation attempt ran against unmodified code, instead of counting its three passes. That note is why the other counts can be trusted.

Your run closes the one item left open on 2026-09-24. The committed GPU test:

• passed 5 of 5 times with the guard in place;
• failed 5 of 5 times with stream.wait_stream(torch.cuda.current_stream()) removed;
• and the whole file gave 5 passed on the GPU.

You also named the card, the torch build and the driver. Patching load_async at runtime is an acceptable way to remove the line. The five failures are what show the patch took effect, because a patch that changed nothing would have passed.

Current state, checked today:

• All 13 required checks pass on 77dc04e3.
• On CPU the file gives 4 passed and 1 skipped (the GPU test).
• The branch still merges cleanly with main at afd9a917, 24 commits past your last merge.

The GPU evidence on record is your run. No GPU was available for this check, so I did not re-run it.

The next step is yours: mark it ready for review. #1258 is branched from this head, so it can only land after this one.

[MakazhanAlpamys]

@VanshSharmaPES, thank you for marking this ready, and for the 3050 run that made it possible.

What is good here, specifically

• The change is exactly option (b) and nothing more. It is one callback and one flag in StreamPrefetcher, the wiring in install_streaming, and _prime's load kept as the fallback. Nothing in async_disk_source.py changes. I checked that the fallback is still needed: deleting _prime's load fails 3 of your 4 tests with large-layer scheduler bug: slot holds None.

• test_two_steps_complete_without_an_autograd_version_error does its job. I added a grad hook on the head's output that fires the refill before the head's backward, and 3 of the 4 tests failed, all with modified by an inplace operation.

• The load-count claim holds exactly. I recorded the order of real loads on a 4-layer untied fixture over 2 steps. It is main's order with the embedding load moved to the end of each backward, plus the one trailing load you describe:

  main: embed L0 L1 L2 L3 head L1 L0 | embed L0 L2 L3 head L1 L0
  PR:   embed L0 L1 L2 L3 head L1 L0 embed | L0 L2 L3 head L1 L0 embed
  

One blocking item: the tests pin that the refill fires in the backward, not where

layer_stream_runtime.py:1061 is the and idx == 0 in advance(). With that line removed, the refill fires at the first recompute that walks downward instead of at layer 0. On the file's 2-layer fixture those are the same call, so all 4 tests still pass (4 passed, 1 skipped). So do the 46 test files that exercise layer streaming (1495 passed, 93 skipped, the same as without the change).

On any model with three or more layers, the two placements differ. The same 4-layer recording with that line removed gives:

embed L0 L1 L2 L3 head L1 embed L0 | L0 L2 L3 head L1 embed L0

The head-sized copy is now issued ahead of layer 0's backward load. On CUDA, that load then queues behind it on the copy stream both pools share, so the wait moves into the backward instead of going away. It also changes the order of source.get() calls, and #975 dropped its race-battery criterion on the grounds that this PR leaves that order unchanged. No test would notice either change.

A 4-layer version of your trainer test catches it. Paste it into TestBackwardTailEmbedPrefetch, after test_prime_no_op_saves_one_reload_from_the_second_step_on:

    def test_the_refill_fires_once_per_backward_at_layer_0(self, tmp_path, monkeypatch):
        """Where in the backward the refill fires, not only that it fires.

        On two layers the first recompute that walks downward IS layer 0, so a
        refill moved to the start of the backward passes every other test here.
        On a real model that would queue the whole head-sized copy in front of
        the backward's decoder copies on the shared copy stream, and change the
        order of `source.get()` calls, which #975 relies on being unchanged.
        Four layers tell the two apart.
        """
        wrapper, _, _ = _build_untied_wrapper(tmp_path, monkeypatch, n_layers=4)
        prefetcher = wrapper._stream_runtime.prefetcher
        refill = prefetcher.backward_tail_prefetch
        fired_at = []

        def _recording_refill():
            fired_at.append((prefetcher.direction, prefetcher.prev))
            refill()

        prefetcher.backward_tail_prefetch = _recording_refill
        wrapper.trainer.args.max_steps = 2
        wrapper.trainer.train()
        wrapper._close_stream_runtime()
        # (direction, layer) at each call: once per step, walking down, at layer 0.
        assert fired_at == [(-1, 0), (-1, 0)], fired_at

Here is what I got with it in your file, on CPU, with your branch merged with current main:

• without mutation: 5 passed, 1 skipped;
• and idx == 0 removed: 1 failed, [(-1, 2), (-1, 2)] == [(-1, 0), (-1, 0)];
• direction check dropped: 3 failed (2 before);
• flag reset removed from prime(): 2 failed (1 before);
• once-per-pass flag check removed: still passes. That mutant is equivalent, as you and the maintainer found.

The method is clean under ruff check and ruff format. It also passes with #1258's four commits applied on top (5 passed, 1 skipped).

One nit

The module docstring (lines 23-25) says _build_streamed_wrapper in test_v07204.py "never passes tie through to _tiny_llama_dir". Since #1061, which your merge of main brought in, it takes tie= and passes it through (test_v07204.py:349-371). Please correct that sentence, or switch the helper to _build_streamed_wrapper(..., task="sft", tie=False).

What I ran

Everything ran on CPU with CUDA_VISIBLE_DEVICES=-1, with soup_cli imported from the checkout. I ran nothing on a GPU.

• Your head 77dc04e3: your file gives 4 passed, 1 skipped (the GPU test).
• Your head merged with main at 964ddc67, 43 commits ahead of your last merge. It merges cleanly, and still does with today's a7cf9c40.
  • your file plus the changelog-fragment tests: 25 passed, 1 skipped;
  • the 46 test files that exercise layer streaming: 1495 passed, 93 skipped, 0 failed;
  • ruff check: clean.
• Your file on unmodified main: 2 failed (AttributeError: ... 'backward_tail_prefetched' and assert 4 == 5), 2 passed, 1 skipped.
• Mutations of layer_stream_runtime.py, run against your file:

mutation	result
wiring reverted (:2671)	2 failed
refill call deleted, flag kept (:1065)	2 failed
direction check dropped, or inverted (:1060)	2 failed each
flag reset removed from prime() (:1012)	1 failed
_prefetch_embed loads the head's key (:2663)	2 failed
refill also fired before the head's backward	3 failed
_prime forced to copy every step (:2740)	1 failed
_prime's fallback load deleted	3 failed
once-per-pass flag check removed (:1062)	passes (equivalent)
wait_stream removed from LargeLayerBufferPool.load_async (:951)	passes on CPU; only the GPU test pins it
and idx == 0 removed (:1061)	passes: the blocking item above

On the GPU side, your 3050 run is the evidence for the wait_stream guard. LargeLayerBufferPool is byte-identical between 77dc04e3 and the merge with current main, so I am not asking for another run of it. Once the new test is in, one run of the whole file on your 3050 would confirm that the new trainer-level test passes on CUDA like the other four. It should give 6 passed:

pytest tests/test_issue975_backward_tail_embed_prefetch.py --no-cov -p no:cacheprovider

For #1258, once this lands

If this PR is squash-merged, a plain git rebase origin/main on #1258 stops at this PR's first commit, 3dcd7e70, with conflicts in layer_stream_runtime.py and the #975 test file. Merging main into #1258 conflicts in the docs paragraph and the runtime file. Replaying only #1258's own four commits is clean:

git rebase --onto origin/main 77dc04e3

I tried that on a local squash of this PR, and the runtime file and both test files came out byte-identical to 427befb0. The test above touches only the #975 file, which #1258 does not modify, so adding it does not change that.

CI: all 13 required checks pass on 77dc04e3. They will need to run again on the head that adds the test.

[VanshSharmaPES]

Thanks for the exact repro on two layers, and for the test to tell the two placements apart. Pushed a5cc33c and a merge of main (301a80c).

The blocking item. I added your test verbatim, on a 4-layer fixture: it records (direction, layer) at every backward-tail call and asserts it fires walking down, at layer 0, once per step. With and idx == 0 removed from the trigger it now fails 1, [(-1, 2), (-1, 2)] != [(-1, 0), (-1, 0)]. I also re-checked two of the other mutations against the updated file: the once-per-pass flag check is still an equivalent mutant (passes), and removing the flag reset in prime() still fails 2 — both match your counts.

The nit. The module docstring no longer claims _build_streamed_wrapper never passes tie through; it's said since #1061 that it predates that and keeps its own helper.

What I ran, on 301a80c:

• ruff check and ruff format --check both clean.
• CPU: the file gives 26 passed, 1 skipped (with the changelog-fragment tests included, matching your count).
• pytest -k stream: 531 passed, 35 skipped, 0 failed.
• Merged main at a7cf9c40, matching what you tested against; it merged cleanly.
• GPU (RTX 3050, torch 2.14.0+cu130): 6 passed, both before and after the merge — the run you asked for.

For #1258: once this is squash-merged, I'll rebase with git rebase --onto origin/main <this PR's squash commit> as you laid out, rather than a plain rebase or a merge of main.

[MakazhanAlpamys]

Merged as 20ab4668. Thank you, @VanshSharmaPES: you traced the stall to the shared slot rather than to the disk-read ordering the issue first blamed, then committed a GPU test for the one guard that keeps the new refill safe and ran it both ways on your own 3050. On CPU, with your branch merged with main, your file gives 5 passed and 1 skipped (the GPU test). 10 of 12 mutations of layer_stream_runtime.py fail it; the two survivors are an equivalent flag check and the GPU-only wait_stream. The 47 streaming test files give 1529 passed, 93 skipped, 0 failed.