[jiapingW]

Move MTP to canonical SDPA and a checkpointed, chunked LM-head objective. Normalize the old chunk-size option into objective_options and reject conflicts. Update native offline/NPU recipes, remove the old MTP model/training wrappers and duplicate exporters, and retain historical configuration/weight import through codecs and mappings.

Coverage: CPU/CUDA FP32/BF16 loss and gradient parity, bounded chunk recomputation, frozen-backbone/empty-mask cases, GPU save/resume and pinned-SGLang compatibility. Old MTP recipes must use draft/initialization; cross-implementation optimizer resume is unsupported. Other algorithms keep legacy paths for unported backends/SP capabilities. The merge CLI delegates to the canonical publisher. NPU hardware execution remains unverified. Existing local model validation used Torch 2.13.0+cu130 and Transformers 5.12.1; the repository's 5.19.0 Transformers pin was not exercised locally.

Base: #1163.

[maocheng23]

Requesting changes. There are three P1s:

• MTP capture asks SGLang for a layer it can't tap;
• offline eval accuracy is always 0;
• a measured regression against current main.

One P2 is below; the rest are inline.

I reviewed e349f01b on CPU with the pinned torch 2.14.1 / transformers 5.19.0. On the full CPU suite, every test that fails at this head also fails at #1163's head (c7005264). The causes are sglang/torchvision missing locally, and the transformers 5.19 Qwen3RotaryEmbedding(device=) break. So this PR adds no CPU failure, and it fixes test_aggregate_packages_stay_doc_only. The GPU PR Test workflow only runs for PRs into main. So none of the CUDA tests in the description ran in CI. A workflow_dispatch run on this branch would cover them. I re-ran the probes behind each comment on this head.

What already works:

• Historical checkpoints now merge back. Thanks for the rotary fix. I checked it at this head:
  • an old config decodes to MultiAxisRoPESpec(1e7, 0.25, [11, 11, 10], interleaved) against the real Qwen3.5-4B config, and passes the export check;
  • merge_mtp_to_base.py succeeds on default and mrope targets;
  • logits match post-#1014 main bit for bit;
  • the regenerated fixture agrees with my own regeneration from main within 2e-6.
• Old mtp_offline_v1 caches load unchanged. On the same record, the old and new paths give a bit-identical loss, with gradients within 3e-7 (eager and SDPA).
• Loss normalization is unchanged. The old strategy already returned (loss * denom, denom).
• Resuming across implementations fails loudly, and old strategy: mtp recipes fail with a migration hint.
• No reference to a deleted symbol is left in the repo.

P1: MTP capture now requests the final target layer. Inline on algorithms/mtp/providers.py. Both MTP recipes are affected.

P1: offline eval accuracy is always 0. Inline on objectives/mtp.py:87. The best checkpoint never moves, so anyone exporting *-best silently gets the first eval step.

P1: #1015's supervised-row path is lost. Inline on objectives/mtp.py.

P2: fsdp_sharding: NO_SHARD can't train MTP any more.

• The trainer passes the wrapped model to CanonicalTrainStrategy, which reads model.draft_model (training/strategies/canonical.py:19,27).
• FSDP forwards that attribute; DDP doesn't. So every rank raises AttributeError: 'DistributedDataParallel' object has no attribute 'draft_model' (CPU/gloo, 1 and 2 ranks).
• The default SHARD_GRAD_OP is fine, and legacy MTP trained under DDP.
• The code is from #1160. Looking through the wrapper (model if hasattr(model, "draft_model") else model.module) fixes it. With that change the whole workflow runs under DDP, and the merged export re-imports to the trained weights.

Questions for maintainers. Each changes a user-visible contract or numerics:

1. merge_mtp_to_base.py now symlinks unchanged shards into the HF cache by default, and copies only .json/.txt/.model/.tiktoken. The old script wrote a self-contained copy. Should the wrapper default to --copy-target? I'd say yes; specforge export can keep linking.
2. Cross-entropy now runs in fp32 (ce_l1_loss). The old path ran F.cross_entropy on bf16 logits.
   • On Qwen3.5's vocabulary, the old bf16 CE is off by up to 1.4% per row against an fp64 reference: 1.6e-4 on the mean loss, 1.4e-3 on dlogits. fp32 is 2–3 orders closer.
   • The cost: each 4096-row chunk holds an extra ~4 GB fp32 copy.
   • I'd keep fp32, say so in the description, and add a bf16 parity check against main.
3. export/checkpoint_io.py is deleted. If #1161 keeps --to on main's exporters, they need this file and tests/test_runtime/test_export.py.
4. draft: OLD.json now decodes any legacy architecture, not only MTP, and only MTP is tested. Is that in scope here, or should it move to its own PR?

What I'd need to approve:

1. MTP capture stops requesting the final layer, with a test on the server capture contract.
2. Offline eval reads the canonical accuracy (or the name is unified), and MTP without an initialization source fails instead of training from random weights.
3. Supervised-row lm_head + CE ported to MTPObjective, or the regression stated with numbers.
4. A rebase on current main (07e13e1d). There are still 4 conflicts: core/mtp.py, mtp/qwen3_5.py, strategies/base.py, test_mtp.py.
5. Answers to the four questions, recorded in an MTP change list in training.md.

I've written each remaining fix as a patch, and they apply cleanly on this head. On CPU, the full suite fails the same tests as the head, and the new tests pass. I can share the patches if useful.

[maocheng23]

[P1] Port #1015: run lm_head + CE only on supervised rows.

Main now selects the supervised rows before this loop (#1015, in 224b36bd). This objective runs the head on all B × (L-1) rows and zeroes the unsupervised ones, which is main's cost before #1015.

#1015 measured that difference on H100 with this geometry:

• UltraChat: 1.51× slower, with peak memory 9.9 → 6.6 GiB;
• ShareGPT: 1.41× slower;
• up to 2.5× slower at 8% supervision.

Suggest:

• Have the objective declare a host-side input: supervised row indices from loss_mask, computed before .to(device) in CanonicalTrainStrategy._forward. Then the forward doesn't sync. Fall back to nonzero otherwise.
• index_select hidden, labels and weights before chunking.
• For an empty set, still run the head on an empty chunk. That way every rank's backward reaches the same parameters, since DDP runs with find_unused_parameters=False. Note that checkpointed_chunk_reduce raises on a zero-length input today.
• Port perf(mtp): run lm_head and cross-entropy only on supervised positions #1015's five tests:
  • sparse parity;
  • zero grads when nothing is supervised;
  • host rows match in-model rows;
  • rows are passed from the strategy;
  • no CUDA sync.

I prototyped this; the patch applies cleanly on this head:

• The objective declares a host_inputs hook that CanonicalTrainStrategy calls before the copy.
• One shared function computes the rows on the host or the device, from MTPSpec's offsets.

On CPU, eager and SDPA, chunk sizes 0/1/3/4096:

• Loss, accuracy and gradients match the every-row path.
• A batch with nothing supervised gives zero (not None) grads.
• Host rows equal the in-objective lookup bit for bit.
• lm_head only sees the supervised rows.

The no-sync check needs CUDA, which I didn't have.

[maocheng23]

[P2] Make fully masked rows backend-independent.

objectives/mtp.py:39-43 masks query-side padding, so every padded row reaches SDPA with no visible key.

• Eager zeroes those rows explicitly. SDPA leaves them to the kernel.
• What's been checked: CPU returns 0; CUDA only at small shapes; NPU not at all.
• The bundled NPU recipe (batch 2, 32k) hits this every step.
• The old path never did. Its mask was causal only (_make_causal_mask), so every row could see at least itself.

gqa.py:148 zeroes the forward, but not the backward. To check, I wrapped SDPA so it returns NaN on rows with no visible key, which some backends are allowed to do. Then the loss stays finite, but 7 parameters get NaN gradients: o_proj, q_proj, input_layernorm, fusion.fc, both fusion norms, and embed_tokens. With the change below, none do, and the SDPA/eager parity tests still pass.

Giving those rows one visible key and zeroing them afterwards matches eager on any backend:

empty = ~attention_mask.any(dim=-1, keepdim=True)
output = F.scaled_dot_product_attention(
    query, _repeat_kv(key, groups), _repeat_kv(value, groups),
    attn_mask=attention_mask | empty,
    dropout_p=dropout_p if training else 0.0, is_causal=False, scale=scale,
)
return output.masked_fill(empty, 0)

[maocheng23]

[P2] Check the native target conflict when the run is resolved, not only at export.

Here, and for any edited preset, both the draft config and the target config are already known. Today the only check is at export (native.py:53-60).

So draft: OLD.json plus the preset's from_target mapping trains to the end, and only then fails to export. That is the path training.md:783 and the preset README point users to.

Suggest: in resolve_run, when the architecture has a native converter, run convert_config and the same target comparison. It only reads configs.

I tried this by moving the comparison in export_native into one function and calling it from resolve_run with the target config the CLI already passes. Tested with the real Qwen3.5-4B config:

• The preset resolves.
• A preset with a different θ fails at resolve, with the same message export would give.
• The raw config.json text config and the to_dict() one the CLI passes give the same results.

[maocheng23]

Thanks for fixing the two sentences here. One thing is still missing: a short change list for MTP users. Each of these is a breaking change for someone coming from main:

• Old recipes now fail.
• The old JSON moved to tests/fixtures/legacy_mtp/original_config.json.
• The merge script links shards by default.
• model.embedding_key and lm_head_key no longer apply. The NPU recipe still sets them, with a comment.
• An MTP with its own head (share_lm_head: false) trains, but can't be exported.
• Optimizer state doesn't resume across implementations.
• Resume now rejects a changed objective_chunk_size.

[maocheng23]

nit: this case still passes if backend="sdpa" is dropped somewhere and eager runs, because it only compares losses and gradients. Wrapping torch.nn.functional.scaled_dot_product_attention with patch.object(..., wraps=...) and asserting it was called would pin it.

[maocheng23]

[P1] MTP capture now asks SGLang for the final target layer, which its hook can't tap.

The parent returned _CAPTURE_LAYER_IDS = [1], with the reason in a comment: MTP keeps no aux tensor, but the capture plan needs a non-empty layer list. This now delegates to the canonical capture_layers, which returns [num_hidden_layers - 1] for a final-hidden draft.

For qwen3.5-4b-mtp-npu.yaml, the server flags change from --return-aux-hidden-state-layer-ids 1 to 31.

SGLang's DFLASH hook taps layer k+1, so 31 becomes layer 32 on a 32-layer model. Reading the code, that's an IndexError at server start. Offline prepare requires the server's aux ids to match the contract (offline_capture/client.py:62-73), so both MTP recipes are affected. I haven't run a live server.

capture_layers comes from #1160, so the fix could go there and propagate up. Suggest restoring [1] for final-hidden drafts in capture_layers, keeping the reason, and asserting resolve_server_capture_contract(...).aux_layer_ids == (1,) for both MTP recipes.

I checked the consumer: with aux_feature=None the request asks only for last_hidden, so the extra configured layer is never transferred.

[maocheng23]

[P2] Run directories now resolve three different ways.

Main's merge and model.draft_checkpoint_path accepted a run directory. Here:

• Resume (CheckpointManager.resolve_resume_dir) takes the single *-latest, or falls back to the highest step of one run id.
• Export (read_export) requires exactly one *-latest, so a run root without one fails with "expected one latest pointer". CheckpointManager._point skips the link on filesystems without symlinks, so our own runs can look like this.
• initialization.checkpoint, which training.md:783 points old recipes to, rejects any run root, even one with -latest. For a canonical run directory it says "checkpoint has no config.json; supply legacy_config explicitly".

read_export comes from #1161 and the initialization path from #1160, but this PR is what sends MTP users down both. Suggest one resolver in the checkpoint layer, used by all three.

I tried moving resolve_resume_dir into specforge/checkpoint/resolve.py and calling it from read_export and the initialization path:

• the merge script then accepts all 5 input shapes it took on the parent (step dir, state file, root with and without -latest, file://);
• initialization.checkpoint accepts a run root;
• the existing checkpoint-manager and resume tests pass.

[maocheng23]

[P1] Offline eval reads accuracy under names this result doesn't use, so eval/avg_acc is always 0.

The evaluator takes accuracy from ratio_metrics["acc"] or from metrics["accuracy"] (eval/evaluator.py:94-105). CanonicalTrainStrategy returns metrics={} and this ratio as "accuracy". So both eval accuracy keys fall back to 0.0 (:160-165):

• eval/avg_acc;
• eval/simulated_acc_len, which CheckpointManager uses to pick the best checkpoint (training/checkpoint.py:42).

In a CPU end-to-end run:

• training accuracy was 0.034 / 0.019, eval accuracy 0.0;
• specforge-run-best stayed at step 1;
• the parent's legacy MTP reported 0.0177.

All canonical objectives use this name (objectives/base.py:25), so the fix belongs in the evaluator, or in one shared name. With the evaluator also reading ratio_metrics["accuracy"], the same run reports 0.0185 and the best pointer moves.

[maocheng23]

[P2] This message says MTP requires initialization, but nothing enforces it once draft is set.

enabled() turns on the canonical path for draft alone, and build_bundle then calls initialize_draft(from_target=None). A migrated recipe that forgets from_target therefore trains the MTP layer from random weights, without a warning.

I checked this with a tiny Qwen3.5-shaped target: fusion.fc.weight differs from the target's mtp.fc.weight. The parent's _init_from_native_mtp refused to start this way.

Suggest raising in build_bundle for MTP when there is neither an initialization source nor a resume, and accepting an explicit initialization: {} for training from scratch. That is the training-only spot; offline capture planning also goes through resolve_run and doesn't need initialization.

[maocheng23]

[P2] No shipped tool prepares these features.

training.md:753-756 calls this a complete recipe. But scripts/prepare_hidden_states.py --strategy mtp builds a config without draft/initialization, so it hits the "MTP now requires draft and initialization" error. Before this PR it failed on the default flex_attention instead.

Suggest adding --draft to the prepare script, using a backend every canonical objective accepts (eager) for capture planning, and naming the command here.

With that change, the native preset resolves to a plan with dflash capture, aux layer (1,), and input_ids/loss_mask/target_last_hidden_states. I haven't run the capture itself, since that needs SGLang.