[YLlllllllllll]

What does this PR do?

Make chunk_loss / npu cross-entropy work with activation offload.

Previously, ChunkLoss used torch.func.grad_and_value, which rejects any
active saved_tensors_hooks — including the hooks installed by
enable_activation. Enabling both paths failed closed.

This PR replaces that transform with an equivalent first-order
torch.autograd.grad VJP, keeps the short-lived chunk CE graph on device,
and lets only the long-lived precomputed grads go through the outer offload
hooks. An opt-in cross_entropy_loss_release_cache flag can reclaim the
transient CE allocator cache before model backward on memory-constrained
profiles.

Checklist Before Starting

• Search for relative PRs/issues and link here: none found on GitHub main; extracted from internal NPU/act-offload bring-up (kept separate from Ascend 910B4 Mojo overlay work)
• PR title follows [{modules}] {type}: {description} format

Test

• New unit tests in tests/ops/test_chunk_loss_saved_tensor_hooks.py:
  • chunk loss under synthetic saved-tensor hooks preserves loss/grads
  • production chunk_loss_function under custom_save_on_cpu
  • opt-in release_cache runs only after outer pack completes
  • factory binds cross_entropy_loss_release_cache
• python3 -m py_compile on touched Python files
• Validated in internal NPU training profiles that combine
  chunk_loss/npu with train.accelerator.offload_config.enable_activation

API and Usage Example

model:
  ops_implementation:
    cross_entropy_loss_implementation: chunk_loss   # or npu
    cross_entropy_loss_release_cache: false          # set true on tight HBM profiles
train:
  accelerator:
    offload_config:
      enable_activation: true

Design & Code Changes

1. veomni/ops/kernels/cross_entropy/chunk_loss.py
   • Drop torch.func; compute per-chunk VJP with torch.autograd.grad
   • Shadow outer activation-offload hooks with identity hooks for the
     short-lived CE graph so vocab-sized logits are not forced through CPU
   • After ChunkLoss.apply, optionally sync + empty_cache when
     release_cache=true
2. veomni/ops/kernels/cross_entropy/__init__.py
   • Thread cross_entropy_loss_release_cache from ops config into the
     chunk-loss dispatch
3. veomni/arguments/arguments_types.py + docs
   • Document the new opt-in flag
4. Tests
   • Regression coverage for hooks / offload / release_cache

This change is hardware-agnostic and does not include Ascend 910B4 Mojo
compat overlays, validators, or GDR pad logic.

Checklist Before Submitting

• Read the Contribute Guide
• Applied pre-commit checks (targeted compile + diff --check)
• Added/updated documentation
• Added tests (tests/ops/test_chunk_loss_saved_tensor_hooks.py)

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[YLlllllllllll]

@Luosuu All CI checks passed on eba3837 (3 green) — ready to merge when you have a moment.

[YLlllllllllll]

@Coach257 CI is all green and your approval is already on eba3837 — ready to merge when you have a moment. Thanks!

[cursor]

🤖 Thanks for the contribution! Cursor has started reviewing this PR and will post the full review shortly.

  

_{Sent by Cursor Automation: VeOmni-Review-Agent(Outer)}

[cursor]

Review summary — approve (with non-blocking follow-ups)

Baseline: main@1f586482f200, ahead 1, behind 9. Reviewed read-only (no checkout of the PR branch; PR file versions were fetched over tracked paths, tested, then reverted).

1. Does it solve a real problem? Yes — verified

The incompatibility is real and reproducible on this repo's pinned torch (2.11.0):

RuntimeError: torch.func.{grad, vjp, jacrev, hessian} don't yet support saved tensor hooks.

Any active saved_tensors_hooks — which is exactly what custom_save_on_cpu (activation offload) installs — makes the old torch.func.grad_and_value path in ChunkLoss.forward raise. Replacing it with torch.enable_grad() + torch.autograd.grad on detached leaves computes the same first-order VJP and is hook-compatible.

Evidence: with only chunk_loss.py reverted to main, test_chunk_loss_supports_saved_tensor_hooks_and_preserves_gradients and test_full_chunk_loss_runs_under_activation_offload_hook both fail; with the PR applied the whole new file passes (8 passed). A wider run of tests/ops/test_chunk_loss_saved_tensor_hooks.py test_chunk_logprobs.py test_chunk_topk_distill.py test_seqcls_loss.py gives 22 passed / 12 skipped, the single failure being the Liger/Triton GPU test that also fails on main in a CPU-only environment.

2. Comments / docs / tests — complete

docs/usage/arguments.md and docs/design/kernel_selection.md both document the new cross_entropy_loss_release_cache field including its throughput tradeoff. The tests are genuinely behavioral rather than tautological: gradient parity against an unchunked reference, an exact pack/unpack call count ({"pack": 2, "unpack": 2}) proving only the two long-lived saved tensors reach the outer hook, and an explicit ordering assertion that synchronize precedes empty_cache. The inline comments explaining why the cleanup lives after apply and why the del is needed before rebinding are the useful kind.

3. Cleanliness / conflicts

ruff check veomni tests and ruff format --check are clean. Three follow-ups, none blocking:

• Rebase. 9 commits behind, and both veomni/ops/kernels/cross_entropy/ and veomni/arguments/arguments_types.py have moved on main since the merge base (9ff48586, 542f792a, df65f4fd).
• Overlaps #943. That PR edits ce_loss_func's loss_kwargs_chunks construction in this same file, a few lines away from this PR's chunk_loss_function hunk. Whichever lands second will need a manual rebase; worth cross-linking.
• Scope. cross_entropy_loss_release_cache is an independent memory/throughput knob riding along with a compatibility fix. It is opt-in and defaults to false, so the risk is low, but splitting it would make both changes easier to bisect and would let the compatibility fix land without introducing a new public config field.

One behavioral note worth adding to the design doc: the nested identity saved_tensors_hooks is the right call, but its effect is that the largest transient — the vocab-sized FP32 logits — is now explicitly excluded from activation offload. That is precisely the memory an offload user most wants moved, and release_cache is the mitigation. Stating that tradeoff explicitly would help users choose.

  

_{Sent by Cursor Automation: VeOmni-Review-Agent(Outer)}

[cursor]

release_cache now has two resolution points with two different timings. Here the factory snapshots the ops config at OpSlot bind time and freezes it into the partial; that explicitly-bound False then shadows the lazy fallback at lines 397-399, so a later set_ops_config is ignored on this path. The install_loss_mapping("chunk_loss") path (line 505) assigns the bare _chunk_loss_dispatch, so it resolves lazily per call and does observe a later config change.

Since the lazy fallback already covers both paths correctly, consider dropping this factory change and keeping return _chunk_loss_dispatch. That removes the divergence, and also avoids returning a functools.partial here — which makes _current_cross_entropy_name() fall back to repr() instead of a readable kernel name in the startup ops log.

[cursor]

release_cache is only threaded through the plain chunk_loss_function path. The return_log_probs and teacher_topk_* distillation branches in _chunk_loss_dispatch also compute chunked logits but never reach this cleanup, so a user who sets cross_entropy_loss_release_cache: true on a distillation or log-probs run gets a silent no-op. Either forward the flag to those helpers, or note the limitation in the field help text — which currently says only that it affects chunk_loss/npu.

[YLlllllllllll]

Closed accidentally while recreating the head branch to rebase past the fork's ByteDance-author push hook. Rebased tip is now in #1004 (same change, on latest main).