diff --git a/docs/ai/encoder_cpu_wte_backward_bucket_width.md b/docs/ai/encoder_cpu_wte_backward_bucket_width.md new file mode 100644 index 0000000..08b2943 --- /dev/null +++ b/docs/ai/encoder_cpu_wte_backward_bucket_width.md @@ -0,0 +1,72 @@ +# Encoder CPU `dwte` backward skipped token rows + +**Symptom.** `make test-python` failed on all three +`tests/test_encoder_equivalence.py::test_backward_matches_torch` cases. The +position-embedding gradient (`dwpe`) matched PyTorch, but the token-embedding +gradient (`dwte`) stayed entirely zero for rows whose tokens appeared in the +batch. + +For the small fp32 case, PyTorch expected non-zero gradients for token rows +`[5, 12, 15, 16, 25, 26, 27, 28]`; the custom op returned no non-zero `dwte` +rows at all. + +## Cause + +The CPU `wte_backward_cpu` path derived each bucket's channel span with: + +```mojo +var c_per_warp = WARP_SIZE * width +``` + +That is the right shape for GPU terminology, but the CPU path is consuming +bucket metadata produced with the module-level bucket contract: + +```mojo +comptime WTE_C_PER_WARP = 32 * WTE_BWD_SIMD_WIDTH +``` + +On this MAX/Mojo CPU custom-op path, the mismatch made `c_len` evaluate to zero +for the Python equivalence cases, so the kernel never entered the accumulation +loop that writes `dwte`. + +## Fix + +Use `WTE_C_PER_WARP` directly in the CPU backward path when mapping a bucket's +`channel_group` to `[c_base, c_end)`. This keeps the consumer in lockstep with +`build_wte_buckets` and the Python test helper, both of which emit bucket groups +using the same constant. + +## Verification + +Commands run on Apple Silicon/macOS with Pixi-installed Mojo +`1.0.0b3.dev2026072606`: + +```bash +PATH=/Users/saisandeepkantareddy/.pixi/bin:$PATH \ +LLMM_DISABLE_MEF_CACHE=1 \ +/Users/saisandeepkantareddy/.pixi/bin/pixi -q run pytest -q \ + tests/test_encoder_equivalence.py +# 6 passed + +PATH=/Users/saisandeepkantareddy/.pixi/bin:$PATH \ +/Users/saisandeepkantareddy/.pixi/bin/pixi -q run pytest -q \ + tests/test_dataloader.py \ + tests/test_tokenizer_equivalence.py \ + tests/test_encoder_equivalence.py +# 12 passed + +PATH=/Users/saisandeepkantareddy/.pixi/bin:$PATH \ +make PIXI='/Users/saisandeepkantareddy/.pixi/bin/pixi -q' test-python +# 233 passed, 15 skipped + +PATH=/Users/saisandeepkantareddy/.pixi/bin:$PATH \ +make PIXI='/Users/saisandeepkantareddy/.pixi/bin/pixi -q' build +# passed +``` + +The Mojo build printed Crashpad warnings under the sandboxed macOS environment, +but exited successfully and produced `build/train_gpt2`. + +## AI use statement + +Written with AI assistance (Codex/Agent), directed by Sai Sandeep Kantareddy. diff --git a/llmm/encoder.mojo b/llmm/encoder.mojo index 7a6f81c..205d881 100644 --- a/llmm/encoder.mojo +++ b/llmm/encoder.mojo @@ -651,10 +651,9 @@ def wte_backward_cpu[ var token_idx = Int(info[2]) var channel_group = Int(info[3]) - var c_per_warp = WARP_SIZE * width - var c_base = channel_group * c_per_warp + var c_base = channel_group * WTE_C_PER_WARP - var c_end = min(c_base + c_per_warp, channels) + var c_end = min(c_base + WTE_C_PER_WARP, channels) var c_len = c_end - c_base if c_len > 0: