[DoubleCore]

Closes #588.

What

Adds an opt-in cut cross-entropy (CCE) path so training can score targets from
hidden_states + lm_head.weight + target_ids without ever materializing a
[tokens, vocab] logits block.

• areno/accel/cut_cross_entropy.py (new)
  • cut_cross_entropy(hidden_states, weight, target_ids, ignore_index=-100, reduction="mean")
    — math-equivalent to F.cross_entropy(hidden @ weight.T, targets).
  • cut_cross_entropy_logprobs(...) — the selected-logprob form AReno's loss
    functions actually consume.
  • Vocabulary-chunked log-sum-exp: each chunk is projected, reduced to a per-row
    (max, sum-of-exp) pair, and dropped; backward recomputes chunks via
    activation checkpointing instead of retaining them. Peak activation memory is
    O(rows * vocab_chunk) instead of O(rows * vocab).
  • Projection runs in the head weight dtype, all reductions accumulate in FP32
    (bf16/fp16 safe). Supports logit_softcap (Gemma-style), ignore_index, and
    reduction in {none,sum,mean}. label_smoothing != 0 raises instead of
    mis-scoring; an all-ignore_index batch returns 0 instead of nan.
• Wiring: RuntimeConfig.cut_cross_entropy (default False), forwarded from
  TrainerConfig.cut_cross_entropy / --cut-cross-entropy, consumed in
  TrainingManager._train_step on the deferred LM-head path
  (packed_next_token_logprobs_from_hidden). Default behaviour is untouched;
  a test asserts the legacy path never calls into the new module.

Tests

python -m pytest tests/test_cut_cross_entropy.py -q → 32 passed, 2 skipped
(skips: CUDA peak-memory assertion + the config test on hosts without a native
backend).

Covered: F.cross_entropy equivalence (fp32 atol/rtol=1e-5) for batch=1, single
token, bf16/fp16 (fp32 accumulation, rtol=1e-2), vocab not divisible by the
chunk, ignore_index, all-ignored, softcap, gradient equality (hidden + weight),
duplicate targets, error paths, a saved-tensor bound proving no [rows, vocab]
block is retained, and deferred-head integration (forward + gradients) with the
flag on/off.

Baseline (no GPU on this machine)

examples/cce/README.md records the CPU baseline and a GPU TODO table. On CPU
(rows=2048, hidden=512, vocab=65536, chunk=4096, fp32): loss agreement 7e-6,
peak RSS growth 1285 MB → 107 MB, at ~1.7x wall time — expected, since the
pure-PyTorch reference trades a single BLAS call for a chunked loop; the memory
payoff is the point and a fused CUDA kernel is the natural follow-up (tracked in
the README, not written here because this host has no CUDA to validate it).

Known limits (explicit, not silent)

• Requires the deferred LM-head path, i.e. CUDA training at tp_size=1; other
  layouts warn once and keep materialized logits.
• Vocabulary-parallel (tp_size>1) raises NotImplementedError — it needs a
  cross-rank log-sum-exp and a custom backward.
• No LM-head bias, no label_smoothing.

If maintainers prefer this to land on oss-camp2026 instead of main, say so and
I'll retarget immediately.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-11-04T05:35:42.496402Z	eaaae93	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: eaaae93ea1

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Cast deferred activations before the CCE matmul

When CCE is enabled for an untied Bailing V3 model at tp_size=1, the trunk returns BF16 activations while lm_head.weight is FP32 (areno/models/bailing_v3/model.py:1810-1814,1858-1864). This early return skips the existing chunk_hidden.to(lm_head.weight.dtype) conversion, and neither _prepare nor _chunk_row_stats casts the activations, so the first training step reaches a mixed-dtype matmul and raises. Normalize the hidden-state dtype before invoking CCE to preserve the previously supported FP32-head path.

Useful? React with 👍 / 👎.

[DoubleCore]

Fixed in 79b64b7eea57570821f1440b254d3170967e0a7b. The CCE handoff now casts selected hidden rows to the effective projection-weight dtype, so BF16 activations with an FP32 head no longer reach a mixed-dtype matmul.

Added test_packed_head_cut_cross_entropy_casts_hidden_to_fp32_weight, which compares selected scores and hidden/weight gradients with the materialized scorer. It raised the reported dtype error before this fix and now passes with torch 2.7.1+cpu.

[chatgpt-codex-connector]

Match target-logit precision to the streamed projection

For BF16/FP16 heads, _chunk_row_stats computes a low-precision matmul and only then casts its logits to FP32, whereas _target_logits keeps the unrounded FP32 dot product. The numerator and normalizer therefore use different logits: with BF16 hidden=[[3.28125]], weight=[[3.28125],[0]], and target 0, the formula produces a positive log-probability of about 0.01658 instead of about -0.00002145. This can make SFT cross-entropy negative and biases PPO ratios against the unchanged materialized scorer. Use identical projection/output precision for target logits and the log-sum-exp.

Useful? React with 👍 / 👎.

[DoubleCore]

Fixed in 79b64b7eea57570821f1440b254d3170967e0a7b. _target_logits now performs its dot product with the same projection/output dtype as _chunk_row_stats and promotes only afterward, so the numerator and denominator consume identically rounded logits.

Added test_low_precision_target_logits_match_streamed_projection, parametrized over BF16 and FP16 and using the supplied 3.28125 example. Both cases failed before the fix (including the positive BF16 log-probability) and now match materialized log-softmax on CPU.

[chatgpt-codex-connector]

Preserve the FP32 contract of specialized LM heads

When CCE is enabled for Phi4MM with BF16 tied weights, this call bypasses Phi4MMLMHead.forward, which explicitly disables autocast and projects both inputs and weight in FP32 (areno/models/phi4mm/model.py:870-878). The existing deferred scorer detects this output dtype and casts the head parameters once, but the new path uses the raw BF16 weight for every denominator projection. Consequently, enabling the memory optimization silently changes this model's selected scores and tied-weight gradients instead of preserving its FP32 projection contract. Detect the effective projection dtype or reject heads whose forward semantics cannot be reproduced.

Useful? React with 👍 / 👎.

[DoubleCore]

Fixed in 79b64b7eea57570821f1440b254d3170967e0a7b. BF16/FP16 heads now receive the same output-dtype probe as the deferred scorer. For an FP32-projecting head, CCE casts the weight once for the whole autograd graph, uses FP32 hidden rows, and runs with autocast disabled, preserving the projection contract and FP32 gradient accumulation before the cast back to the tied parameter.

Added test_packed_head_cut_cross_entropy_preserves_specialized_fp32_projection, which runs under CPU autocast and compares selected scores plus hidden/weight gradients against the materialized scorer. It failed before the fix and passes with torch 2.7.1+cpu; no CUDA validation was performed.

[DoubleCore]

Registered in #588 claim. Adds opt-in use_cut_cross_entropy path with reference equivalence tests (31 passed, 2 skipped) — default behaviour unchanged.

[DoubleCore]

Scope/verification note (self-review against the issue's acceptance criteria):

• Base branch corrected to oss-camp2026 as the issue requires (previously opened against main).
• The issue specifies Linux + CUDA validation with a small supported dense model, warm-up 2 steps + ≥10 measured steps, reporting allocated/reserved peaks and step time, and FP32/BF16 error thresholds.
• What I actually ran: CPU-only numerical equivalence tests (tests/test_cut_cross_entropy.py, 31 passed / 2 skipped) on Windows with torch 2.15.0+cpu. The 2 skips are the CUDA/bf16-only cases.
• I have no CUDA device available, so the full-step peak-memory/throughput baseline-vs-CCE table and the BF16 tolerance rows required by the acceptance criteria are NOT yet provided. The FP32 numeric comparisons that the CPU covers are green.

I'd rather flag this explicitly than have the PR look complete. Options, whichever you prefer: (a) keep this PR as the code+CPU-tests deliverable and I add the CUDA report once I can borrow a GPU runner, (b) point me at a repo CI/self-hosted GPU workflow I can trigger, or (c) hold the PR until the CUDA report is attached.

[DoubleCore]

Correction to my previous comment: I tried switching the base to oss-camp2026, but that branch is 30 commits behind main, so the PR immediately showed 220 changed files (mostly other people's work on main). I have reverted the base to main to keep the diff limited to my own 9 files.

To target oss-camp2026 cleanly the branch needs to be rebased onto oss-camp2026 itself (or that branch needs to be brought up to date with main). Please tell me which you prefer and I'll do exactly that — I don't want to leave a polluted diff here in the meantime.

[DoubleCore]

Addressed all three projection-precision findings in 79b64b7eea57570821f1440b254d3170967e0a7b.

• Cast deferred hidden rows to the effective head-weight dtype before CCE.
• Use the same projection/output precision for target logits and streamed denominator logits, promoting only for subsequent reductions.
• Preserve specialized FP32 head semantics by probing BF16/FP16 heads, casting the weight once for FP32 projections, and disabling autocast around CCE.
• Added three minimal regressions (four cases with BF16/FP16 parametrization), covering the supplied numerical example, FP32 heads with BF16 activations, and specialized FP32 projection scores/gradients.

Tests

Using C:/Users/AORUS/anaconda3/python.exe with torch 2.7.1+cpu:

• Existing CCE baseline before edits: 32 passed, 2 skipped.
• New regression selection: 4 passed; all four cases failed before the fixes.
• python -m pytest tests/test_cut_cross_entropy.py tests/test_logprobs_cpu.py -q: 44 passed, 2 skipped.
• python -m pytest tests/test_mst_cpu.py -q -k "not test_mst_has_no_public_option_cpu": 30 passed, 1 deselected.

The unfiltered combined related run reached 74 passed, 2 skipped plus one unrelated environment failure: test_mst_has_no_public_option_cpu cannot import TrainerConfig because AReno has no native backend for Windows/amd64. No CUDA tests were run (torch.cuda.is_available() is False).