fix(megatron): compact actor logits before projection

[liuhao-labs]

Summary

Long masked prompts still produce full [T, V] vocabulary logits before Slime selects trainable tokens, causing the projection and fp32 conversion peaks in #2253. Add opt-in --compact-actor-logits to project only the N loss-masked rows for eligible built-in actor policy/SFT paths.

A Slime-owned CompactLogitsGPTModel is selected only for eligible actor post-process chunks. Log-probability/entropy computation accepts [N, V], builds compact top-p masks, and scatters scalar results back to the existing response/CP layout. Empty masks retain zero-gradient connections.

Critics, custom providers/hooks/losses, MTP, TIS/mismatch metrics, debug capture, rollout log-probs, PPO reward-KL, score centering, MuP scaling, and custom Megatron output processors preserve the full-logit path. The optimization is disabled by default and can be enabled with --compact-actor-logits for eligible built-in paths.

Normal forward and combined 1F1B both resolve the model's _postprocess method dynamically. The context covers the pipeline schedule call. Model state-dict keys and checkpoint ownership are unchanged.

Fixes #2253.

Validation

• 134 related local CPU tests passed against main 58b08786, covering compact model dispatch, provider ownership, VLM/plain critic isolation, compact/full loss and gradient parity with and without top-p replay, top-p/reference handling, empty masks, vocab-parallel log-probability/entropy, and import-time collective rebinding.
• That suite includes 12 checkpoint/optimizer tests and the upstream isolated train_one_step fixture. The new test_model_provider_critic.py fixture is compatible with the actor-only compact import and is registered in the permanent CPU CI matrix.
• Direct registered compact-model runner: 35 passed, including score-centering, MuP, and current Megatron output-processor fallback.
• Repository-pinned pre-commit checks, workflow regeneration, compile checks, and git diff --check passed.
• Real TP=2/NCCL coverage in tests/test_ppo_logprob_entropy_gpu.py compares selected logits plus hidden/output-weight gradients with sequence parallel off/on. All 13 CUDA cases are skipped on this CPU-only local Python environment and still require the self-hosted GPU job.
• Combined-schedule unit coverage checks dispatch through a schedule node; source at the current pinned Megatron commit 4b4acac9 was audited for the actual PostProcessNode call. Full CUDA scheduler integration is not claimed as locally tested.
• Straw-backed integration runs in the registered upstream CI jobs; local ARM has no compatible prebuilt straw-queue wheel.
• Preserved Add sunabako agent E2E and align top-p score centering #2457's complete-support score-centering residual and tests, removal of the obsolete 8-node example, and the Sunabako E2E/container setup. 82 available single-process score-centering tests passed locally; disk-backed Straw cases remain CI validation.

[liuhao-labs]

CPU/pre-commit CI is green (58 successful checks). Could a maintainer add run-ci-changed for the modified tests/test_ppo_logprob_entropy_gpu.py? It contains the new real TP=2/NCCL forward-and-gradient parity coverage for sequence parallel off/on, which this local single-device environment cannot execute.

[liuhao-labs]

Rebased this PR onto the latest main (3778dbf) and regenerated the workflow from .github/workflows/pr-test.yml.j2. The new CI run is queued on head 1a2a046.

Local validation is still green: the focused CPU suite passes (80 passed), the new test passes via the direct CI-style runner (31 passed), and Ruff/Black/isort/diff checks pass. This host cannot run the TP=2/NCCL coverage in tests/test_ppo_logprob_entropy_gpu.py; could a maintainer please add run-ci-changed again so the self-hosted changed-test job executes it?

[liuhao-labs]

Rebased onto the current v0.4 mainline and resolved the provider-test conflict, retaining the upstream MLA compatibility coverage. The compact path now falls back for score centering, MuP scaling, and custom Megatron output processors, with focused regression tests. The upstream isolated checkpoint/optimizer fixture was also updated for the new schedule context.

Local validation: 117 CPU tests plus 12 checkpoint/optimizer tests passed, and the direct compact-model runner has 35 passes. The full refreshed PR Test run is green with 90 successful jobs and no failures; pre-commit also passed.

Could a maintainer please add run-ci-changed to run the real TP=2/NCCL projection and gradient parity coverage in tests/test_ppo_logprob_entropy_gpu.py? CPU CI is complete, and the PR is mergeable and ready for review.

[liuhao-labs]

Rechecked and synchronized with main 22a89f9, including #2450's VLM critic-head ownership fix and #2453's Megatron collective rebinding. The actor-only compact projection, fallback rules, schedule context, and SP gradient contract remain unchanged.

The new upstream critic tests needed import stubs for the existing MLA module and this PR's compact module. Their fixture now also fails if a critic provider attempts actor-compaction eligibility, and the file is registered in the permanent CPU matrix. Local related coverage is 132 passed; the critic file also passes via its direct CI-style runner.

The refreshed PR Test run completed successfully: 91 successful jobs, 7 conditional skips, 0 failures, including the new critic job. Pre-commit passed. Real TP=2/NCCL parity remains pending the previously requested GPU CI authorization; these CPU results do not replace it.

[liuhao-labs]

Resolved the #2457 conflict against main 58b08786 and pushed ddfff568. The obsolete 8-node example remains deleted. The complete-support score-centering alpha = w(1) / residual = coefficients - alpha[row_ids] * p.exp() path and its upstream tests are preserved exactly; score-centering configurations continue to use the full-logit fallback.

Added compact/full top-p replay output and gradient parity to the existing response-span tests. Local validation: 134 related CPU tests, 82 single-process score-centering tests (including the nine new centered-scalar regressions), and 18 non-disk distributed score-centering gradient tests passed. This does not replace the pending real TP=2/NCCL test.

The new PR Test run completed successfully with 101 successful jobs and 7 conditional skips; pre-commit also passed. The workflow template/YAML preserve the Sunabako 8-GPU/35-minute E2E, container lifecycle/proxy/cache setup and evidence artifacts, plus this PR's critic and compact CPU registrations. The PR is mergeable again. The previously recorded GPU authorization dependency is unchanged.