fix(domino): resume the base-anchor curriculum from the checkpoint step

[khazic]

Problem

DominoTrainingModel drives the base-anchor curriculum off a counter it increments itself:

self._forward_count = 0          # __init__
...
self._forward_count += 1         # forward()
lambda_base = self._current_lambda_base()

That counter is a plain Python attribute on the training wrapper. It is not a buffer, it is not in any state_dict, and nothing saves or restores it. Every resume therefore rebuilds the wrapper at step 0 and restarts the curriculum at lambda_base_start.

The trainer already knows the real step: optimizer_steps_total is restored from the checkpoint's trainer state, and _resume_optimizer_steps is plumbed into drafter_train_config for exactly this purpose. Today only the LR scheduler consumes it.

Why it matters

The Domino loss is

loss = (1 - lambda_base) * final_loss + lambda_base * base_loss

and only final_loss depends on the correction head (prefix_gru, embed_proj). The head's gradient therefore scales exactly with 1 - lambda_base. Restarting the curriculum on an already-trained drafter cuts that learning signal by the full decay factor and takes a whole domino_lambda_base_decay_steps window to recover, while the backbone is re-trained under a curriculum it had already finished. Nothing fails loudly: domino_lambda_base is reported in the diagnostics and simply reads ~1.0 again.

With the default domino_lambda_base_decay_steps=2000, the first step after a resume gives the correction head 1/2000 of its previous gradient.

Fix

Seed the curriculum from _resume_optimizer_steps in DominoTrainerBackend.setup_optimizer, next to where the LR scheduler already resumes from the same value. The counter is renamed to _curriculum_step and gets an explicit set_curriculum_step() entry point so the intent is visible at the call site.

Validation

Ran a before/after repro on both main and this branch. It walks a fresh wrapper to optimizer step 100 with domino_lambda_base_decay_steps=100, then simulates a process restart that resumes the drafter at the same step, and reports the curriculum weight plus the correction head's gradient norm. Both sides use identical seeds and identical weights, so the only variable is the curriculum weight.

Before/after repro output

=================== before: main (333b754) ===================
domino_lambda_base_decay_steps = 100
resume from optimizer step     = 100

[phase 1] fresh run, walk the curriculum to step 100
          curriculum_step=101 lambda_base=0.0000 head_grad_norm=0.878218

[phase 2] process restart, resume the drafter at step 100
          curriculum_step=1 lambda_base=0.9900 head_grad_norm=0.008782

[verdict]
          lambda_base across the restart: 0.0000 -> 0.9900
          correction-head gradient retained: 1.00%
          RESULT: curriculum RESTARTED (bug present)

=================== after: this branch ===================
domino_lambda_base_decay_steps = 100
resume from optimizer step     = 100

[phase 1] fresh run, walk the curriculum to step 100
          curriculum_step=101 lambda_base=0.0000 head_grad_norm=0.878218

[phase 2] process restart, resume the drafter at step 100
          curriculum_step=101 lambda_base=0.0000 head_grad_norm=0.878218

[verdict]
          lambda_base across the restart: 0.0000 -> 0.0000
          correction-head gradient retained: 100.00%
          RESULT: curriculum CONTINUED (fixed)

Tests

Four new tests in tests/integration/test_domino_backend_contract.py:

• test_domino_curriculum_step_is_seedable covers the curriculum weight across a seeded restart.
• test_domino_correction_head_gradient_scales_with_curriculum pins the exact 1 - lambda_base gradient scaling that makes the regression matter.
• test_domino_setup_optimizer_seeds_curriculum_from_resume_steps covers the resume path.
• test_domino_setup_optimizer_without_resume_starts_curriculum_at_zero covers the fresh-start path.

Full CPU suite (tests/integration tests/compat tests/config tests/examples) run on both sides:

Test suite before/after

### BASELINE (origin/main, 333b754) ###
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[DSPARK-veomni-npu-veomni_lm_head_full]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[DFLASH-veomni-npu-veomni_lm_head_sparse]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[EAGLE3-veomni-cuda-veomni_lm_head_full]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[EAGLE1-fsdp-npu-engine_full_param]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[DOMINO-fsdp2-cuda-engine_full_param]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_transfer_waits_after_actor_update
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_async_publish_sets_pending_ref_and_waits_before_next_publish
FAILED tests/integration/test_dspark_trainer_backend.py::test_dspark_checkpoint_preserves_source_config_and_vllm_weight_names
FAILED tests/integration/test_verl_npu_vllm_compat.py::test_factory_fused_moe_survives_verl_npu_patch_import
9 failed, 245 passed, 2 warnings in 33.60s

### THIS BRANCH ###
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[DSPARK-veomni-npu-veomni_lm_head_full]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[DFLASH-veomni-npu-veomni_lm_head_sparse]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[EAGLE3-veomni-cuda-veomni_lm_head_full]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[EAGLE1-fsdp-npu-engine_full_param]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_sync_defers_for_all_lm_head_drafters[DOMINO-fsdp2-cuda-engine_full_param]
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_target_head_transfer_waits_after_actor_update
FAILED tests/integration/test_drafter_runtime_control_contract.py::test_async_publish_sets_pending_ref_and_waits_before_next_publish
FAILED tests/integration/test_dspark_trainer_backend.py::test_dspark_checkpoint_preserves_source_config_and_vllm_weight_names
FAILED tests/integration/test_verl_npu_vllm_compat.py::test_factory_fused_moe_survives_verl_npu_patch_import
9 failed, 249 passed, 2 warnings in 68.17s

The same 9 tests fail on main and on this branch. They need optional dependencies (VeOmni, the NPU vLLM stack) that are absent in this environment, so they are pre-existing and unrelated. The +4 on this branch are the new tests above.