Skip to content

fix(lora): let a resumed disaggregated run reach the rollout engine - #2708

Closed
Arist12 wants to merge 1 commit into
radixark:mainfrom
Arist12:fix/lora-disaggregated-resume-adapter
Closed

Arist12 wants to merge 1 commit into
radixark:mainfrom
Arist12:fix/lora-disaggregated-resume-adapter

Conversation

@Arist12

@Arist12 Arist12 commented Aug 22, 2026

Copy link
Copy Markdown
Contributor

Part of #2705.

Problem

A disaggregated LoRA run resumed with --lora-adapter-path can fail before its first rollout: native-only checkpoints cannot be preloaded by SGLang, partial best-effort PEFT output can look preloadable, and a preloaded adapter collides with the mandatory initial publish.

Change

Normal train drivers start resumed rollout engines from the base model and let the mandatory initial weight publish install the trainer's restored adapter. Publish single LoRA adapters with SGLang's existing upsert=True contract so fresh startup and later refreshes use one path without unload churn. --debug-rollout-only, which skips the publish, remains a blocker before this PR is ready.

Validation

  • 28 focused fast tests covering complete, partial and native-only resume directories
  • Full Qwen3-30B-A3B, 4 trainer + 4 rollout MI350X: restored four EP shards (384 tensors each), optimizer, scheduler and global-dataset cursor from iter_29; continued at rollout 30, completed two optimizer steps, saved iter_30, and republished the trained adapter (grad_norm=0.0152, train/rollout abs diff 0.0210-0.0237)
  • Full Qwen3-4B, 4+4 MI350X: a 30-rollout run resumed at iter_4 and iter_14, restored 288 tensors plus optimizer/scheduler/data cursor, and completed through iter_29
  • Dense held-out DAPO accuracy 0.711 → 0.789, truncation 0.211 → 0.078; train/rollout abs diff mean 0.01265, max 0.01419
  • Checkpoints at 4/14/29 each contain one finite TP0/PP0 native shard and distinct adapter hashes

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

A resume passed --lora-adapter-path straight to SGLang as lora_paths, but that
directory holds miles native shards plus a best-effort PEFT export. When the
export was skipped or written in a layout the serving loader does not expect,
the engine failed to start.

Start every engine from the base model instead. train.py and train_async.py both
publish the trainer adapter unconditionally before the first rollout, so the
adapter arrives over the normal weight-sync path and native-only and
PEFT-producing resumes take the same route.

That publish now upserts rather than unload-then-load, so an adapter is never
briefly absent from a live engine. Upsert is only wired up on the from_distributed
route, so the colocated from_tensors path keeps its explicit unload.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@Arist12
Arist12 force-pushed the fix/lora-disaggregated-resume-adapter branch from f7fe762 to 7665877 Compare September 1, 2026 20:30
@Arist12

Arist12 commented Sep 3, 2026

Copy link
Copy Markdown
Contributor Author

Closing against the refactor. The update_weight/ half of this patch no longer exists — #2754/#2756 deleted the unload-then-reload contract it modified and replaced it with register-then-stream. The one issue that survives is unrelated to weight sync: sglang_engine.py still passes --lora-adapter-path straight through as a PEFT lora_paths entry at engine startup, which cannot work for a directory holding miles' native per-rank shards. That is SGLang-LoRA integration code, so it will come back as its own PR once that refactor lands.

@Arist12 Arist12 closed this Sep 3, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant