[maocheng23]

Motivation

DSpine can reuse the existing optional Liger kernels and avoid redundant objective work during accumulation prepasses. These execution optimizations are separated from the architecture, objective and runtime integration.

Changes

• Honor model.use_liger_kernel through the explicit RMSNorm/SwiGLU kernel factories, preserving checkpoint keys and avoiding global Transformers patching.
• Flatten proposal rows for the frozen vocabulary projection and restore the original output shape.
• Compute only backbone CE during the graph-free accumulation prepass, skipping teacher projection, L1, top-k, refinement and alignment work. Full objectives still run during replay.
• Add Liger wiring/compatibility tests, projection loss/gradient parity tests, and CE-only prepass/collective-count regressions. Document the performance controls.

Liger remains opt-in. The flattened projection and cheaper prepass are the DSpine implementation paths already present in original #965. Floating-point rounding may differ across GEMM layouts; no universal speedup is claimed.

Stack

DSpine 4/4, extracted from original #965 (b32657405d7df33f049cc638da3d739c240e82b4), within workstream 8 of #942.

Order: #965 -> #1005 -> #1006 -> this PR -> #966 -> #967. Base: #1006 / codex/pr965-03-training; review only the base-relative diff.

This PR's complete file tree is identical to the original #965 tree. No tests or behavior were removed by the split. #966 and #967 are restacked on top with their respective file trees also unchanged. Original tips remain on backup branches; original authorship is preserved. The adapted integration originates in #942 and credits #917; this does not close #917.

Validation

• Both performance modules passed: 8 test methods, zero failures/errors, 2 optional-dependency/CUDA skips and 3 CUDA subtest skips.
• CPU projection loss/gradient comparisons and CE-prepass equivalence tests ran.
• The original identical tree passed all 9 targeted modules in fresh processes: 78 test methods, zero failures/errors, 6 reported skips (including CUDA subtests).
• Earlier split layers were independently validated, including CPU Gloo/DDP and CPU FSDP BF16 tests. Changed-file repository pre-commit checks passed for each layer.
• Environment: Python 3.12.13, PyTorch 2.13.0, Transformers 5.12.1, macOS arm64. CUDA and Liger are unavailable locally; actual Liger-kernel/GPU numerical tests remain unvalidated by this run.
• No training, downstream accuracy or end-to-end performance benchmark was launched. Historical measurements remain in Optimize distributed draft training and decouple attention range from direction #942 with their original scope.

OMP_NUM_THREADS=1 PYTHONPATH=. python -m unittest tests.test_modeling.test_dspine_liger -v
OMP_NUM_THREADS=1 PYTHONPATH=. python -m unittest tests.test_modeling.test_dspine_optimizations -v