[Feature Request] Add Anchored Supervised Fine-Tuning (ASFT) as an SFT algorithm

# [Feature Request] Add Anchored Supervised Fine-Tuning (ASFT) as an SFT algorithm

## Is your feature request related to a problem?

Standard cross-entropy (CE) SFT reliably causes **catastrophic forgetting**:
because CE treats every token equally and pushes the model to fit the SFT set as
hard as possible, capabilities the base model already had (instruction
following, code, general knowledge) degrade during fine-tuning. NeMo-RL's only
SFT loss today is plain NLL/CE (`NLLLossFn`), so users have no built-in way to
fine-tune on a task while preserving base-model behavior.

## Describe the solution you'd like

Add **Anchored Supervised Fine-Tuning (ASFT)** ("Anchored Supervised
Fine-Tuning", [arXiv:2509.23753](https://arxiv.org/abs/2509.23753), ICLR 2026)
as a first-class SFT algorithm alongside SFT/DPO/GRPO.

ASFT augments ordinary SFT with two lightweight ideas:

1. **DFT reweighting** ([arXiv:2508.05629](https://arxiv.org/abs/2508.05629)):
weight each response token's NLL by the model's own **detached** probability
of that token. This counteracts the implicit `1/p` weighting of the SFT
gradient, so confidently-correct tokens are not over-optimized.
2. **KL anchoring**: a per-token KL penalty that keeps the policy close to the
**frozen base (pretrained) model**.

```
L_ASFT = E_t[ sg(p_θ(y_t))·(−log p_θ(y_t)) + λ·KL(π_θ || π_ref)_t ]
```

over response tokens only, where `sg` is stop-gradient and `λ = kl_weight`. It
maps cleanly onto existing infrastructure: the loss is a `LossType.TOKEN_LEVEL`
/ `LossInputType.LOGPROB` function (like `NLLLossFn`), the KL uses the existing
Schulman `calculate_kl` estimator on the correct-token logprobs, and the frozen
reference is obtained exactly like DPO (`init_reference_model=True` +
`get_reference_policy_logprobs`) — **no new worker machinery required**.

## Why we want it

- **Directly addresses forgetting** — the number 1 pain point of SFT — without needing
preference data (unlike DPO) or a reward model / rollouts (unlike GRPO). It is
a drop-in replacement for CE SFT.
- **Cheap**: one extra reference forward per step (the reference is a CPU↔GPU
weight-swap, so no 2× GPU memory), and it reuses the DPO reference path and the
fused-linear-logprobs path already in the codebase.
- **Complements the SFT family**: gives users a knob (`kl_weight`) to trade task
fit against base-capability retention.

## We have a preliminary solution that validated the benefits
Branch: https://github.com/pengdurice/RL/tree/peng-asft-v1

> **Internal adoption (motivation).** Beyond the open reproduction above, we have
> been using ASFT in our own in-house SFT pipeline on proprietary data and models,
> where it consistently delivers the same reduced-forgetting benefit. Those
> results aren't shareable, so the fully-open reproduction here stands as the
> public, verifiable evidence — the internal usage is what motivated us to
> upstream it.

A working prototype is implemented and has been validated end-to-end on
**Qwen3-8B (Megatron backend)**:

- `ASFTLossFn` (+ config/data-dict) in `nemo_rl/algorithms/loss/loss_functions.py`
- `nemo_rl/algorithms/asft.py` (`setup` / `asft_train`, derived from `sft.py` +
a DPO-style reference-logprob precompute)
- `examples/run_asft.py`, `examples/configs/asft.yaml`, and a Megatron recipe

**Validation (fully open-source data, apples-to-apples).** ASFT (`kl_weight=0.01`,
LR 5e-6) vs a plain-CE baseline, trained with **byte-identical policy + data —
only the loss differs**, on a balanced 4-domain open 10k set (2.5k each:
OpenMathInstruct-2 / OpenCodeInstruct / WebInstructSub / ultrachat_200k), 1
epoch. Evaluated with lm-evaluation-harness (vLLM) across 11 open benchmarks
spanning six domains (generative tasks scored with the chat template; MC tasks
scored raw).

**Generative task performance (what a deployed model is actually used for) —
ASFT wins across the board:**

| benchmark (domain) | CE | ASFT | **ASFT − CE** |
| --- | --- | --- | --- |
| IFEval — prompt-level strict (instruction following) | 0.493 | **0.547** | **+0.054** |
| IFEval — inst-level strict | 0.627 | **0.662** | **+0.035** |
| MATH-500 (hard math) | 0.386 | **0.434** | **+0.048** |
| HumanEval — pass@1 (code) | 0.774 | **0.829** | **+0.055** |
| GSM8K (easy math, in-domain) | 0.842 | 0.824 | −0.018 |

**Medical + general knowledge (MC probes) — tied (ASFT preserves knowledge):**

| benchmark | CE | ASFT | **ASFT − CE** |
| --- | --- | --- | --- |
| MedQA (USMLE 4-opt) | 0.639 | **0.648** | +0.009 |
| MedMCQA | 0.595 | 0.597 | +0.002 |
| MMLU | 0.730 | 0.730 | ±0.000 |

**Commonsense / truthfulness (MC probes) — slight CE edge:**

| benchmark | CE | ASFT | **ASFT − CE** |
| --- | --- | --- | --- |
| HellaSwag (acc_norm) | 0.761 | 0.721 | −0.040 |
| WinoGrande | 0.716 | 0.668 | −0.047 |
| TruthfulQA (mc2) | 0.527 | 0.506 | −0.021 |

**Takeaway.** ASFT's advantage is concentrated in **generative task
performance** — instruction following (+5.4% / +3.5%), hard math (+4.8%), and
code (+5.5%), consistent same-direction gains of ~+0.05 across three independent
domains — **while fully preserving medical and general knowledge** (MedQA /
MedMCQA / MMLU all tied). On the latent-knowledge MC probes it is not a universal
win: CE edges ASFT on commonsense (HellaSwag / WinoGrande) and TruthfulQA by
~0.02–0.05, because ASFT's anchor keeps it close to the base model (e.g.
WinoGrande base 0.678 → ASFT 0.668, essentially unchanged) while CE drifted
upward. Net: **ASFT makes the model behave better without sacrificing what it
knows** — reproducing the paper's core finding (largest edge on instruction
following) and extending it across math, code, and medical retention on fully
open data.

## Proposed scope for the PR

- `ASFTLossFn` + `asft.py` + `run_asft.py` + base `asft.yaml`.
- One example recipe (`examples/configs/recipes/llm/asft-qwen3-8b-1n8g-megatron.yaml`)
+ paired nightly test (`tests/test_suites/llm/…`) registered in `nightly.txt`.
- A unit agreement test for the loss (mirroring the SFT/DPO loss tests) and a
short `docs/guides/asft.md`.

## Additional context

- The KL anchor uses the **target-token Schulman estimator** (matching the
per-token-logprob infra and DPO's reference pattern), not the paper's
full-vocabulary KL; the regularization is qualitatively the same, though the
numeric `kl_weight` sweet spot may differ from a full-vocab implementation.
- Stability guidance from the study: `kl_weight ≈ 0.005–0.03` and `LR ≤ 5e-6`
(ASFT collapses at 5e-5; `kl_weight ≥ 0.05` over-regularizes).
- References: ASFT [arXiv:2509.23753](https://arxiv.org/abs/2509.23753); DFT
reweighting [arXiv:2508.05629](https://arxiv.org/abs/2508.05629).