streamft 0.2.0__tar.gz → 0.2.1__tar.gz

{streamft-0.2.0/streamft.egg-info → streamft-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streamft
-Version: 0.2.0
+Version: 0.2.1
 Summary: Fine-tune models that do not fit: stream NF4 weights from disk, one transformer block at a time.
 Author-email: Gokhan Cava <byjameson@users.noreply.github.com>
 License: Apache-2.0
@@ -30,6 +30,8 @@ Dynamic: license-file
 
 # stream-finetune — fine-tune a 32B model on a 6 GB laptop GPU
 
+[![PyPI](https://img.shields.io/pypi/v/streamft)](https://pypi.org/project/streamft/) [![try it](https://img.shields.io/badge/try-will%20it%20fit%3F-6ee7b7)](https://huggingface.co/spaces/byjameson/will-it-fit)
+
 A small training engine that fine-tunes models that **do not fit in VRAM**. The NF4
 weights are mmap'd from disk, one transformer block is materialised on the GPU at a
 time, and the backward recomputes through it. The base stays frozen; a LoRA adapter is
@@ -61,7 +63,7 @@ and the correction of an earlier wrong claim ("8.6× / 100 tok/s") are in
 ## Quick start
 
 ```powershell
-pip install -e .            # or: pip install streamft (once it is on PyPI)
+pip install streamft
 
 # 0) will it fit? memory arithmetic + the measured reference, no weights downloaded
 streamft fit --model Qwen/Qwen2.5-14B-Instruct --gpu 6 --seq 512
@@ -142,30 +144,33 @@ Raw output: `results/gradcheck.json` (git-ignored; the numbers above are copied
 
 ## Compared with PEFT (matched runs)
 
-Same data, same LoRA (r=8, α=16, q/k/v/o), same AdamW lr 1e-4, batch 1, same prompt
+Same data, same LoRA (r=8, alpha=16, q/k/v/o), same AdamW lr 1e-4, batch 1, same prompt
 masking, same token-weighted masked eval on 8 held-out rows. PEFT uses
 `prepare_model_for_kbit_training` with gradient checkpointing (the standard QLoRA setup).
+Step times move by up to ~30% between runs on this laptop; quality is the stable number.
 
 **3B (fits entirely):**
 
 | | streamft (36 blocks resident) | PEFT / transformers |
 |---|---:|---:|
-| step time, median of 30 @ seq 512 | 0.938 s | **0.579 s** |
+| step time, median of 30 @ seq 512 | 0.94–1.24 s | **0.58–0.88 s** |
 | peak VRAM | **3.85 GB** | 4.90 GB |
 | eval before | 1.42087 | 1.42087 |
-| eval after | **1.30154** (−0.1193) | 1.32722 (−0.0936) |
+| eval after | **1.30720** (−0.1137) | 1.33002 (−0.0909) |
 
-**7B and 14B (do not fit comfortably):**
+**7B / 14B / 32B (do not fit comfortably), after the masked-gradient fix:**
 
-| model | step time, median @ seq 512 | peak VRAM | eval (start → after) |
+| model | step time, median @ seq 512 | peak VRAM | eval (start → after, 30 steps) |
 |---|---:|---:|---:|
-| 7B, streamft (28 blocks) | **1.70 s** | **5.51 GB** | 2.1928 → 1.1094 |
+| 7B, streamft (28 blocks) | **1.55 s** | **5.51 GB** | 2.1928 → **1.1055** |
 | 7B, PEFT | 8.96 s | 10.05 GB | 2.1914 → 1.1491 |
-| 14B, streamft (6 blocks) | **6.73 s** | **3.86 GB** | 2.3748 → 1.5312 (10 steps) |
-| 14B, PEFT | 18.56 s | 16.48 GB | 2.3748 (3 steps only) |
+| 14B, streamft (6 blocks) | **5.86 s** | **3.86 GB** | 2.3748 → **0.9742** |
+| 14B, PEFT | 18.56 s | 16.48 GB | 3 steps only |
+| 32B, streamft (2 blocks, seq 1024) | **12.48 s** | **3.90 GB** | 1.8755 → **1.0008** |
+| 32B, PEFT | fails before step 1 | — | — |
 
 At 7B the crossover has already happened: PEFT's 4.4 GB of weights plus activations no
-longer fit, so it spills (10.05 GB reserved) and pays 5.3× the step time. Note that
+longer fit, so it spills (10.05 GB reserved) and pays ~5.8× the step time. Note that
 streamft's 7B peak (5.51 GB) is *higher* than its 14B peak (3.86 GB) because the 7B run
 kept all 28 blocks resident while the 14B run kept 6 — the window is a choice, and the
 14B number is the one to copy for a 6 GB card.
@@ -174,17 +179,34 @@ kept all 28 blocks resident while the 14B run kept 6 — the window is a choice,
 into shared memory (19.45 GB reserved on a 6 GiB card), and `prepare_model_for_kbit_training`
 then dies allocating the fp32 norm copy: *CUDA out of memory. Tried to allocate 2.90 GiB;
 18.00 GiB is already allocated.* The streaming engine runs the same checkpoint at
-**12.5 s/step and 3.90 GB** (64/64 layers, held-out 1.8755 to 0.9959).
+**12.5 s/step and 3.90 GB** (64/64 layers).
 
 The 3B row is the honest answer to "is this faster?": **no** — on a model that fits, PEFT
-is 1.6× faster and uses 1.05 GB more. The 14B row is the design point: PEFT does not
-OOM on Windows, it **spills into shared system memory** (16.48 GB reserved on a 6.4 GB
-card) and pays 2.8× the step time; streamft stays at 3.86 GB. Over 30 steps and one seed
-the 3B quality difference is not meaningful.
+is faster and uses 1.05 GB more. The 14B/32B rows are the design point: PEFT does not
+OOM on Windows, it **spills into shared system memory** and pays 2.8–5.8× the step time.
+
+### Correction: masked gradient (25 Oct, after an external review)
+
+`chunked_lm_loss_and_backward_masked` divided every chunk by **that chunk's** valid-token
+count. The reported loss was still the correct mean, so the curves looked right, but the
+accumulated gradient over-weighted chunks with few unmasked tokens — i.e. the prompt-heavy
+start of an SFT row. Measured on a toy model with an uneven mask: **112–427% relative
+gradient error** at seq 512/1024 with 60% of the row masked (chunk 128), while balanced
+masks were unaffected.
+
+The fix divides every chunk by the **global** valid count (computed once, no-grad), so the
+gradient is exactly d(mean CE over valid tokens)/dh. Two tests pin it:
+`test_masked_chunked_gradient_matches_full_with_uneven_masks` and
+`test_masked_chunked_handles_fully_masked_chunks`.
+
+All SFT runs above were re-run with the fix. Quality held: 3B −0.114 (was −0.119),
+7B −1.087 (was −1.083), 32B −0.875 (was −0.880); 14B is now 30 steps (−1.40) instead of
+10. The numbers in the tables are the post-fix ones.
 
 Reproduce: `python scripts/compare_peft.py --model checkpoints/qwen3b_bnb4bit --steps 30`
-and `... --model checkpoints/qwen14b_nf4 --steps 3`.
+and `... --model checkpoints/qwen14b_nf4 --steps 30`.
 
+## What this does not do
 ## What this does not do
 
 - **2k context does not fit** on 6 GB. It is ~0.6 GB short: with `--activation-cache cpu`

{streamft-0.2.0 → streamft-0.2.1}/README.md

@@ -1,5 +1,7 @@
 # stream-finetune — fine-tune a 32B model on a 6 GB laptop GPU
 
+[![PyPI](https://img.shields.io/pypi/v/streamft)](https://pypi.org/project/streamft/) [![try it](https://img.shields.io/badge/try-will%20it%20fit%3F-6ee7b7)](https://huggingface.co/spaces/byjameson/will-it-fit)
+
 A small training engine that fine-tunes models that **do not fit in VRAM**. The NF4
 weights are mmap'd from disk, one transformer block is materialised on the GPU at a
 time, and the backward recomputes through it. The base stays frozen; a LoRA adapter is
@@ -31,7 +33,7 @@ and the correction of an earlier wrong claim ("8.6× / 100 tok/s") are in
 ## Quick start
 
 ```powershell
-pip install -e .            # or: pip install streamft (once it is on PyPI)
+pip install streamft
 
 # 0) will it fit? memory arithmetic + the measured reference, no weights downloaded
 streamft fit --model Qwen/Qwen2.5-14B-Instruct --gpu 6 --seq 512
@@ -112,30 +114,33 @@ Raw output: `results/gradcheck.json` (git-ignored; the numbers above are copied
 
 ## Compared with PEFT (matched runs)
 
-Same data, same LoRA (r=8, α=16, q/k/v/o), same AdamW lr 1e-4, batch 1, same prompt
+Same data, same LoRA (r=8, alpha=16, q/k/v/o), same AdamW lr 1e-4, batch 1, same prompt
 masking, same token-weighted masked eval on 8 held-out rows. PEFT uses
 `prepare_model_for_kbit_training` with gradient checkpointing (the standard QLoRA setup).
+Step times move by up to ~30% between runs on this laptop; quality is the stable number.
 
 **3B (fits entirely):**
 
 | | streamft (36 blocks resident) | PEFT / transformers |
 |---|---:|---:|
-| step time, median of 30 @ seq 512 | 0.938 s | **0.579 s** |
+| step time, median of 30 @ seq 512 | 0.94–1.24 s | **0.58–0.88 s** |
 | peak VRAM | **3.85 GB** | 4.90 GB |
 | eval before | 1.42087 | 1.42087 |
-| eval after | **1.30154** (−0.1193) | 1.32722 (−0.0936) |
+| eval after | **1.30720** (−0.1137) | 1.33002 (−0.0909) |
 
-**7B and 14B (do not fit comfortably):**
+**7B / 14B / 32B (do not fit comfortably), after the masked-gradient fix:**
 
-| model | step time, median @ seq 512 | peak VRAM | eval (start → after) |
+| model | step time, median @ seq 512 | peak VRAM | eval (start → after, 30 steps) |
 |---|---:|---:|---:|
-| 7B, streamft (28 blocks) | **1.70 s** | **5.51 GB** | 2.1928 → 1.1094 |
+| 7B, streamft (28 blocks) | **1.55 s** | **5.51 GB** | 2.1928 → **1.1055** |
 | 7B, PEFT | 8.96 s | 10.05 GB | 2.1914 → 1.1491 |
-| 14B, streamft (6 blocks) | **6.73 s** | **3.86 GB** | 2.3748 → 1.5312 (10 steps) |
-| 14B, PEFT | 18.56 s | 16.48 GB | 2.3748 (3 steps only) |
+| 14B, streamft (6 blocks) | **5.86 s** | **3.86 GB** | 2.3748 → **0.9742** |
+| 14B, PEFT | 18.56 s | 16.48 GB | 3 steps only |
+| 32B, streamft (2 blocks, seq 1024) | **12.48 s** | **3.90 GB** | 1.8755 → **1.0008** |
+| 32B, PEFT | fails before step 1 | — | — |
 
 At 7B the crossover has already happened: PEFT's 4.4 GB of weights plus activations no
-longer fit, so it spills (10.05 GB reserved) and pays 5.3× the step time. Note that
+longer fit, so it spills (10.05 GB reserved) and pays ~5.8× the step time. Note that
 streamft's 7B peak (5.51 GB) is *higher* than its 14B peak (3.86 GB) because the 7B run
 kept all 28 blocks resident while the 14B run kept 6 — the window is a choice, and the
 14B number is the one to copy for a 6 GB card.
@@ -144,17 +149,34 @@ kept all 28 blocks resident while the 14B run kept 6 — the window is a choice,
 into shared memory (19.45 GB reserved on a 6 GiB card), and `prepare_model_for_kbit_training`
 then dies allocating the fp32 norm copy: *CUDA out of memory. Tried to allocate 2.90 GiB;
 18.00 GiB is already allocated.* The streaming engine runs the same checkpoint at
-**12.5 s/step and 3.90 GB** (64/64 layers, held-out 1.8755 to 0.9959).
+**12.5 s/step and 3.90 GB** (64/64 layers).
 
 The 3B row is the honest answer to "is this faster?": **no** — on a model that fits, PEFT
-is 1.6× faster and uses 1.05 GB more. The 14B row is the design point: PEFT does not
-OOM on Windows, it **spills into shared system memory** (16.48 GB reserved on a 6.4 GB
-card) and pays 2.8× the step time; streamft stays at 3.86 GB. Over 30 steps and one seed
-the 3B quality difference is not meaningful.
+is faster and uses 1.05 GB more. The 14B/32B rows are the design point: PEFT does not
+OOM on Windows, it **spills into shared system memory** and pays 2.8–5.8× the step time.
+
+### Correction: masked gradient (25 Oct, after an external review)
+
+`chunked_lm_loss_and_backward_masked` divided every chunk by **that chunk's** valid-token
+count. The reported loss was still the correct mean, so the curves looked right, but the
+accumulated gradient over-weighted chunks with few unmasked tokens — i.e. the prompt-heavy
+start of an SFT row. Measured on a toy model with an uneven mask: **112–427% relative
+gradient error** at seq 512/1024 with 60% of the row masked (chunk 128), while balanced
+masks were unaffected.
+
+The fix divides every chunk by the **global** valid count (computed once, no-grad), so the
+gradient is exactly d(mean CE over valid tokens)/dh. Two tests pin it:
+`test_masked_chunked_gradient_matches_full_with_uneven_masks` and
+`test_masked_chunked_handles_fully_masked_chunks`.
+
+All SFT runs above were re-run with the fix. Quality held: 3B −0.114 (was −0.119),
+7B −1.087 (was −1.083), 32B −0.875 (was −0.880); 14B is now 30 steps (−1.40) instead of
+10. The numbers in the tables are the post-fix ones.
 
 Reproduce: `python scripts/compare_peft.py --model checkpoints/qwen3b_bnb4bit --steps 30`
-and `... --model checkpoints/qwen14b_nf4 --steps 3`.
+and `... --model checkpoints/qwen14b_nf4 --steps 30`.
 
+## What this does not do
 ## What this does not do
 
 - **2k context does not fit** on 6 GB. It is ~0.6 GB short: with `--activation-cache cpu`

{streamft-0.2.0 → streamft-0.2.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "streamft"
-version = "0.2.0"
+version = "0.2.1"
 description = "Fine-tune models that do not fit: stream NF4 weights from disk, one transformer block at a time."
 readme = "README.md"
 requires-python = ">=3.10"

{streamft-0.2.0 → streamft-0.2.1}/streamft/stream_quant_lora.py

@@ -22,27 +22,39 @@ from .stream_train import _layer_out, chunked_lm_loss_and_backward
 
 
 def chunked_lm_loss_and_backward_masked(head, h, ys, chunk=128, ignore_index=-100):
-    """Chunked CE with an ignore_index (for SFT prompt masking)."""
+    """Chunked CE with an ignore_index (for SFT prompt masking).
+
+    Every chunk contributes `loss_sum / valid_total` where valid_total is the number of
+    unmasked targets in the **whole sequence**, so the accumulated gradient is exactly
+    d(mean CE over valid tokens)/dh.
+
+    FIX (25.10.2026, reported by an external reviewer): the first version divided each
+    chunk by that chunk's own valid count. Losses were unaffected (the returned value
+    summed to the right mean) but the gradient over-weighted chunks with few valid
+    tokens, i.e. the prompt-heavy start of an SFT row. tests/test_streamft_smoke.py
+    pins the corrected behaviour with an uneven mask.
+    """
     B, T, _ = h.shape
+    tgt_all = ys.to(h.device)
+    valid_total = int((tgt_all != ignore_index).sum())
+    if valid_total == 0:
+        zero = torch.zeros((), device=h.device)
+        return zero, zero
     total = torch.zeros((), device=h.device, dtype=torch.float32)
-    valid_total = 0
     for a in range(0, T, chunk):
         b = min(a + chunk, T)
-        logits = head(h[:, a:b])
-        tgt = ys[:, a:b].reshape(-1).to(h.device)
-        valid = int((tgt != ignore_index).sum())
-        if valid == 0:
+        tgt = tgt_all[:, a:b].reshape(-1)
+        if int((tgt != ignore_index).sum()) == 0:
             continue
+        logits = head(h[:, a:b])
         loss_sum = torch.nn.functional.cross_entropy(
             logits.reshape(-1, logits.shape[-1]).float(), tgt,
             reduction="sum", ignore_index=ignore_index,
         )
-        (loss_sum / valid).backward()
+        (loss_sum / valid_total).backward()
         with torch.no_grad():
             total = total + loss_sum.detach()
-        valid_total += valid
-    denom = max(valid_total, 1)
-    return total / denom, total / denom
+    return total / valid_total, total / valid_total
 
 
 def base_model_of(model):

{streamft-0.2.0 → streamft-0.2.1/streamft.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streamft
-Version: 0.2.0
+Version: 0.2.1
 Summary: Fine-tune models that do not fit: stream NF4 weights from disk, one transformer block at a time.
 Author-email: Gokhan Cava <byjameson@users.noreply.github.com>
 License: Apache-2.0
@@ -30,6 +30,8 @@ Dynamic: license-file
 
 # stream-finetune — fine-tune a 32B model on a 6 GB laptop GPU
 
+[![PyPI](https://img.shields.io/pypi/v/streamft)](https://pypi.org/project/streamft/) [![try it](https://img.shields.io/badge/try-will%20it%20fit%3F-6ee7b7)](https://huggingface.co/spaces/byjameson/will-it-fit)
+
 A small training engine that fine-tunes models that **do not fit in VRAM**. The NF4
 weights are mmap'd from disk, one transformer block is materialised on the GPU at a
 time, and the backward recomputes through it. The base stays frozen; a LoRA adapter is
@@ -61,7 +63,7 @@ and the correction of an earlier wrong claim ("8.6× / 100 tok/s") are in
 ## Quick start
 
 ```powershell
-pip install -e .            # or: pip install streamft (once it is on PyPI)
+pip install streamft
 
 # 0) will it fit? memory arithmetic + the measured reference, no weights downloaded
 streamft fit --model Qwen/Qwen2.5-14B-Instruct --gpu 6 --seq 512
@@ -142,30 +144,33 @@ Raw output: `results/gradcheck.json` (git-ignored; the numbers above are copied
 
 ## Compared with PEFT (matched runs)
 
-Same data, same LoRA (r=8, α=16, q/k/v/o), same AdamW lr 1e-4, batch 1, same prompt
+Same data, same LoRA (r=8, alpha=16, q/k/v/o), same AdamW lr 1e-4, batch 1, same prompt
 masking, same token-weighted masked eval on 8 held-out rows. PEFT uses
 `prepare_model_for_kbit_training` with gradient checkpointing (the standard QLoRA setup).
+Step times move by up to ~30% between runs on this laptop; quality is the stable number.
 
 **3B (fits entirely):**
 
 | | streamft (36 blocks resident) | PEFT / transformers |
 |---|---:|---:|
-| step time, median of 30 @ seq 512 | 0.938 s | **0.579 s** |
+| step time, median of 30 @ seq 512 | 0.94–1.24 s | **0.58–0.88 s** |
 | peak VRAM | **3.85 GB** | 4.90 GB |
 | eval before | 1.42087 | 1.42087 |
-| eval after | **1.30154** (−0.1193) | 1.32722 (−0.0936) |
+| eval after | **1.30720** (−0.1137) | 1.33002 (−0.0909) |
 
-**7B and 14B (do not fit comfortably):**
+**7B / 14B / 32B (do not fit comfortably), after the masked-gradient fix:**
 
-| model | step time, median @ seq 512 | peak VRAM | eval (start → after) |
+| model | step time, median @ seq 512 | peak VRAM | eval (start → after, 30 steps) |
 |---|---:|---:|---:|
-| 7B, streamft (28 blocks) | **1.70 s** | **5.51 GB** | 2.1928 → 1.1094 |
+| 7B, streamft (28 blocks) | **1.55 s** | **5.51 GB** | 2.1928 → **1.1055** |
 | 7B, PEFT | 8.96 s | 10.05 GB | 2.1914 → 1.1491 |
-| 14B, streamft (6 blocks) | **6.73 s** | **3.86 GB** | 2.3748 → 1.5312 (10 steps) |
-| 14B, PEFT | 18.56 s | 16.48 GB | 2.3748 (3 steps only) |
+| 14B, streamft (6 blocks) | **5.86 s** | **3.86 GB** | 2.3748 → **0.9742** |
+| 14B, PEFT | 18.56 s | 16.48 GB | 3 steps only |
+| 32B, streamft (2 blocks, seq 1024) | **12.48 s** | **3.90 GB** | 1.8755 → **1.0008** |
+| 32B, PEFT | fails before step 1 | — | — |
 
 At 7B the crossover has already happened: PEFT's 4.4 GB of weights plus activations no
-longer fit, so it spills (10.05 GB reserved) and pays 5.3× the step time. Note that
+longer fit, so it spills (10.05 GB reserved) and pays ~5.8× the step time. Note that
 streamft's 7B peak (5.51 GB) is *higher* than its 14B peak (3.86 GB) because the 7B run
 kept all 28 blocks resident while the 14B run kept 6 — the window is a choice, and the
 14B number is the one to copy for a 6 GB card.
@@ -174,17 +179,34 @@ kept all 28 blocks resident while the 14B run kept 6 — the window is a choice,
 into shared memory (19.45 GB reserved on a 6 GiB card), and `prepare_model_for_kbit_training`
 then dies allocating the fp32 norm copy: *CUDA out of memory. Tried to allocate 2.90 GiB;
 18.00 GiB is already allocated.* The streaming engine runs the same checkpoint at
-**12.5 s/step and 3.90 GB** (64/64 layers, held-out 1.8755 to 0.9959).
+**12.5 s/step and 3.90 GB** (64/64 layers).
 
 The 3B row is the honest answer to "is this faster?": **no** — on a model that fits, PEFT
-is 1.6× faster and uses 1.05 GB more. The 14B row is the design point: PEFT does not
-OOM on Windows, it **spills into shared system memory** (16.48 GB reserved on a 6.4 GB
-card) and pays 2.8× the step time; streamft stays at 3.86 GB. Over 30 steps and one seed
-the 3B quality difference is not meaningful.
+is faster and uses 1.05 GB more. The 14B/32B rows are the design point: PEFT does not
+OOM on Windows, it **spills into shared system memory** and pays 2.8–5.8× the step time.
+
+### Correction: masked gradient (25 Oct, after an external review)
+
+`chunked_lm_loss_and_backward_masked` divided every chunk by **that chunk's** valid-token
+count. The reported loss was still the correct mean, so the curves looked right, but the
+accumulated gradient over-weighted chunks with few unmasked tokens — i.e. the prompt-heavy
+start of an SFT row. Measured on a toy model with an uneven mask: **112–427% relative
+gradient error** at seq 512/1024 with 60% of the row masked (chunk 128), while balanced
+masks were unaffected.
+
+The fix divides every chunk by the **global** valid count (computed once, no-grad), so the
+gradient is exactly d(mean CE over valid tokens)/dh. Two tests pin it:
+`test_masked_chunked_gradient_matches_full_with_uneven_masks` and
+`test_masked_chunked_handles_fully_masked_chunks`.
+
+All SFT runs above were re-run with the fix. Quality held: 3B −0.114 (was −0.119),
+7B −1.087 (was −1.083), 32B −0.875 (was −0.880); 14B is now 30 steps (−1.40) instead of
+10. The numbers in the tables are the post-fix ones.
 
 Reproduce: `python scripts/compare_peft.py --model checkpoints/qwen3b_bnb4bit --steps 30`
-and `... --model checkpoints/qwen14b_nf4 --steps 3`.
+and `... --model checkpoints/qwen14b_nf4 --steps 30`.
 
+## What this does not do
 ## What this does not do
 
 - **2k context does not fit** on 6 GB. It is ~0.6 GB short: with `--activation-cache cpu`

{streamft-0.2.0 → streamft-0.2.1}/tests/test_streamft_smoke.py

@@ -112,3 +112,53 @@ def test_chunked_ce_gradient_matches_full_mean():
 
 def test_cli_is_importable():
     assert callable(cli.main)
+
+
+# ------------------------------------------- masked (SFT) chunked loss
+
+def test_masked_chunked_gradient_matches_full_with_uneven_masks():
+    """The 25.10 fix: chunks with different valid counts must not be re-weighted.
+
+    With chunk=3 and a prompt mask over the first 4 positions, chunk 0 has zero valid
+    targets, chunk 1 has two, later chunks have three. The old implementation divided
+    each chunk by its own count, which inflated the sparse chunks' gradient.
+    """
+    from streamft.stream_quant_lora import chunked_lm_loss_and_backward_masked
+
+    torch.manual_seed(0)
+    d, T, V = 12, 10, 7
+    head = nn.Linear(d, V, bias=False)
+    ys = torch.randint(0, V, (1, T))
+    ys[0, :4] = -100     # prompt
+    ys[0, 7] = -100      # a masked middle position
+
+    h1 = torch.randn(1, T, d, requires_grad=True)
+    full = F.cross_entropy(head(h1).reshape(-1, V), ys.reshape(-1),
+                           ignore_index=-100, reduction="mean")
+    full.backward()
+
+    h2 = torch.randn(1, T, d, requires_grad=True)
+    with torch.no_grad():
+        h2.copy_(h1.detach())
+    _, mean = chunked_lm_loss_and_backward_masked(head, h2, ys, chunk=3)
+
+    assert math.isclose(float(mean.detach()), float(full.detach()), rel_tol=1e-6)
+    assert torch.allclose(h1.grad, h2.grad, atol=1e-6)
+
+
+def test_masked_chunked_handles_fully_masked_chunks():
+    """A chunk with no valid targets must contribute nothing at all."""
+    from streamft.stream_quant_lora import chunked_lm_loss_and_backward_masked
+
+    torch.manual_seed(0)
+    d, T, V = 8, 6, 5
+    head = nn.Linear(d, V, bias=False)
+    ys = torch.randint(0, V, (1, T))
+    ys[0, :3] = -100
+    h = torch.randn(1, T, d, requires_grad=True)
+    _, mean = chunked_lm_loss_and_backward_masked(head, h, ys, chunk=2)
+    logits = head(h).reshape(-1, V)
+    tgt = ys.reshape(-1)
+    mask = tgt != -100
+    ref = F.cross_entropy(logits[mask], tgt[mask], reduction="mean")
+    assert math.isclose(float(mean.detach()), float(ref.detach()), rel_tol=1e-6)