[gf239]

Summary

fused_linear_cross_entropy_forward runs its three GEMMs (logits, grad_input, grad_weight) on every row.

Rows whose target is ignore_index add nothing to the loss or the gradients.

In SFT they are often most of the batch: masked prompts, padding, image tokens.

This PR runs the projections only on the other rows.

chunk_size still comes from all BT rows, so skipped rows remove whole chunks and the logits memory budget stays the same.

Rows are skipped only when that removes a chunk; otherwise the current code path runs.

nonzero() replaces the existing .item(), so there is still one host sync.

Under torch.compile the current code path runs.

Training step time, forward + backward, RTX 4090:

Setup	Alpaca (65% labelled)	LLaVA + Qwen2.5-VL (16.5%)	Function calling (5.5%)
Qwen3-0.6B, full	-11.4%	-27.9%	-31.4%
Qwen3-1.7B, full	-6.6%	-16.0%	-18.1%
Qwen3-4B, full	-2.5%	-5.9%	-6.8%
Qwen3-8B, LoRA	-1.1%	-2.5%	-2.9%
Gemma-3-12B, QLoRA	-2.2%	-5.3%	-6.0%
Qwen3-32B, QLoRA	-0.5%	-1.3%	-1.5%

No change when every token is labelled.

The fused loss alone gets about 1 / (labelled share) faster: 1.3x at 75%, 2x at 50%, 7.8-9.0x at 10%.

Testing Done

• Hardware Type: NVIDIA GeForce RTX 4090
• run make test to ensure correctness (5140 passed; the 22 failures fail the same way on main in this environment)
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence (168 passed; the 2 failures fail the same way on main in this environment)

Details

1. Ignored rows keep what they get today: a zero grad_input row, 0 per-token loss, z-loss and accuracy, and -1 as predicted token. Per-token outputs are scattered back to all BT rows.

2. Each chunk's GEMMs read the full V x H weight, so the saving comes from removing chunks. When no chunk is removed (for example, only the shifted last label of each sequence is ignored), the contiguous slices run as before and results are bit-identical to main.

3. With grad_input rows no longer contiguous, the dX GEMM writes into one reused chunk buffer (mm(out=)) and index_copy_ moves the rows.

4. Numerics: in bf16, the loss matches main to 1e-6 (random weights) and to 2e-5 relative (public 4-bit checkpoints); gradients are bit-identical when chunk boundaries line up and otherwise differ by bf16 rounding (up to 1.3e-3 relative). The difference comes from the last, shorter chunk.

5. Memory: transient peak +2 to +4 MB (row index, one chunk x H buffer).

6. torch.compile: torch._dynamo.explain gives the same graphs and breaks on main and this PR, and compiled results equal eager, with capture_scalar_outputs on and off.

7. Benchmark: LIGER_FLCE_BENCH_IGNORE_FRACTION=0.5 python benchmark/scripts/benchmark_fused_linear_cross_entropy.py sets that share of targets to ignore_index. On llama_3_8b shapes it gives 0.52-0.53x of main's full time for liger and liger-fp32-accum, memory within 2 MB.

Tests

New in test/transformers/test_fused_linear_cross_entropy.py (54 cases):

• test_correctness_ignored_rows: prompt-first, scattered, all and no ignored rows; mean, sum, none (per-token upstream gradient); bias on and off; against torch.
• test_ignored_rows_do_not_enter_projection: NaN in ignored rows must not reach loss or gradients. Fails on main: the NaN reaches grad_weight through 0 x NaN in the dW GEMM.
• test_per_token_outputs_with_ignored_rows: loss, z-loss, accuracy and predicted tokens with reduction="none".

In test/ops/test_fused_linear_cross_entropy.py:

• New test_ignored_rows_skip_whole_chunks_only pins the rule (0, 1, 8, 24, 40 ignored rows of 40 give 3, 3, 2, 1, 0 logits GEMMs).
• test_chunk_storage_released_before_next_chunk_and_final_cast and test_reduction_none_reenters_with_reusable_buffers ignore the whole first chunk, so they now expect one chunk fewer.

Results:

Run	Result
FLCE test files, this PR	344 passed
New transformers tests on main	51 passed, 3 failed (NaN test)
New and updated ops tests on main	2 passed, 35 failed (chunk counts)
make test, this PR	5140 passed, 22 failed, 1142 skipped, 21 xfailed
make test-convergence, this PR	168 passed, 2 failed, 10 skipped, 4 xfailed, 2 xpassed
The 24 failing tests, on main	24 failed

The 24 tests that fail on both: test_mlp.py (19), test_fused_neighborhood_attention.py (2), test_fused_linear_jsd.py::test_correctness_with_ignore_index (1), fp32 mini_qwen2_vl multimodal and fp32 mini_llava with logits (2).

test/cutedsl/test_fused_linear_cross_entropy.py: 11 passed, 226 skipped (native CuTe DSL FLCE needs SM100), same on main.

Environment: torch 2.14.1+cu132, triton 3.8.0, transformers 5.20.2, ruff 0.15.22.

FLCE alone (BT = 8192, bf16, forward + backward, ms)

lm_head trained:

Model (H, V)	75% labelled	50%	25%	10%
Llama-3.2-1B (2048, 128256)	98.4 -> 74.7	98.0 -> 50.2	97.2 -> 25.4	96.9 -> 12.4
Qwen3-1.7B (2048, 151936)	191.0 -> 144.4	190.2 -> 96.5	189.7 -> 48.9	189.3 -> 21.9
Gemma-3-4B (2560, 262208)	417.1 -> 314.5	415.8 -> 210.2	414.6 -> 106.0	413.9 -> 47.0
Qwen3-8B (4096, 151936)	226.2 -> 171.1	225.2 -> 114.0	224.5 -> 57.9	224.2 -> 28.4
Llama-3.1-8B (4096, 128256)	203.5 -> 153.6	203.1 -> 102.9	202.5 -> 52.3	202.1 -> 25.2

lm_head frozen (LoRA):

Model (H, V)	75% labelled	50%	25%	10%
Llama-3.2-1B (2048, 128256)	59.4 -> 45.6	58.4 -> 30.1	57.8 -> 15.3	57.5 -> 7.2
Qwen3-1.7B (2048, 151936)	100.3 -> 76.2	99.3 -> 50.7	98.8 -> 25.6	98.3 -> 11.3
Gemma-3-4B (2560, 262208)	223.6 -> 168.9	222.5 -> 112.8	221.3 -> 56.7	220.5 -> 24.6
Qwen3-8B (4096, 151936)	137.9 -> 105.4	135.6 -> 69.8	133.8 -> 34.3	133.1 -> 15.9
Llama-3.1-8B (4096, 128256)	122.3 -> 93.0	120.8 -> 61.9	119.1 -> 30.5	117.4 -> 14.0

With every row labelled, all shapes are within 1.2% of main (same code path apart from nonzero()).

All training-step results (ms)

Setup (tokens per step)	92% labelled	77%	65%	16.5%	9.1%	5.5%
Qwen3-0.6B, full (8192)	552 -> 538	551 -> 510	551 -> 488	549 -> 396	549 -> 382	548 -> 376
Qwen3-1.7B, full (8192)	974 -> 960	974 -> 933	973 -> 909	972 -> 816	972 -> 802	972 -> 795
Qwen3-4B, full (8192)	2096 -> 2086	2096 -> 2062	2095 -> 2044	2094 -> 1970	2094 -> 1956	2094 -> 1952
Qwen3-8B, LoRA (8192)	4273 -> 4264	4271 -> 4242	4271 -> 4226	4268 -> 4160	4269 -> 4149	4269 -> 4144
Qwen3-14B, QLoRA (4096)	3613 -> 3607	3613 -> 3595	3613 -> 3583	3611 -> 3539	3611 -> 3533	3611 -> 3529
Gemma-3-12B, QLoRA (4096)	4600 -> 4578	4601 -> 4534	4600 -> 4499	4598 -> 4356	4598 -> 4334	4597 -> 4323
Gemma-3-27B, QLoRA (4096)	7705 -> 7693	7705 -> 7656	7705 -> 7627	7704 -> 7519	7704 -> 7500	7703 -> 7492
Qwen3-32B, QLoRA (2048)	4104 -> 4100	4108 -> 4096	4107 -> 4088	4108 -> 4056	4108 -> 4052	4108 -> 4048

With every token labelled, all setups are within 0.03% of main.

Setup: Liger-patched Hugging Face models, SDPA, gradient checkpointing, bf16, optimizer step not included.
Full and LoRA runs use random weights from the public configs; QLoRA runs load the unsloth/*-bnb-4bit checkpoints with lm_head in bf16.
LoRA: r = 16 on all linear layers.
Labels: the first (1 - share) of each sequence is -100.
Qwen3-32B runs at 2048 tokens because 4096 does not fit in 24 GB.
Medians of 5 (3 for the QLoRA runs) interleaved main/PR steps in one process, GPU clocks locked at 2520/10501 MHz.

Labelled share of common SFT datasets

Dataset	Labelled tokens
tatsu-lab/alpaca	65%
HuggingFaceH4/ultrachat_200k (multi-turn)	77%
trl-lib/tldr (prompt-completion)	9.1%
NousResearch/hermes-function-calling-v1	5.5%
bespokelabs/Bespoke-Stratos-17k (reasoning)	92%
HuggingFaceH4/llava-instruct-mix-vsft with Qwen2.5-VL image tokens	16.5%

Qwen3 chat template, assistant tokens labelled, everything else -100 (TRL completion_only_loss / assistant_only_loss, LLaMA-Factory train_on_prompt=False, Axolotl train_on_inputs: false).
Token-weighted over 2000 samples (300 for LLaVA), packed.
Padding to the longest sample lowers these shares further.

Generated with Claude Code and Human In The Loop 🙈