fix(dflash): add sglang 0.5.14 post-load init to the DFlash target model#658
Open
jessiewei7 wants to merge 1 commit into
Open
fix(dflash): add sglang 0.5.14 post-load init to the DFlash target model#658jessiewei7 wants to merge 1 commit into
jessiewei7 wants to merge 1 commit into
Conversation
PR sgl-project#605 upgraded sglang to 0.5.14, where ModelRunner no longer allocates the memory pool during init; it added the required post-load calls (alloc_memory_pool / init_attention_backends / init_cuda_graphs) to the eagle3 target model but missed the DFlash path. As a result, prepare_hidden_states.py --model-type dflash crashes in ScheduleBatch.init_new() with AttributeError: NoneType has no attribute device. Mirror the same three calls in SGLangDFlashTargetModel.from_pretrained().
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
Author
|
@FrankLeeeee PTAL — dflash training crashes on main with sglang 0.5.14 (same pattern as #605). Companion fixes: #663 #659 #665. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
On current
main(which pinssglang==0.5.14since #605),prepare_hidden_states.py --model-type dflashcrashes immediately:Root cause
sglang 0.5.14 refactored
ModelRunner:__init__no longer allocates thememory pool. Pool allocation moved to explicit post-load calls
(
alloc_memory_pool()→init_memory_pool()inmodel_runner_kv_cache_mixin.py), normally invoked by theScheduler/TpModelWorker. SpecForge drives the
ModelRunnerdirectly throughits standalone
SGLangRunner, so those calls never happen andreq_to_token_pool/token_to_kv_pool_allocatorstayNone.#605 fixed exactly this for
eagle3_target_model.py(added the threepost-load calls with a detailed comment) but did not touch
dflash_target_model.py, leaving the DFlash extraction path broken on thevery sglang version the repo now pins.
Fix
Mirror #605's eagle3 fix in
SGLangDFlashTargetModel.from_pretrained():(
init_cuda_graphs()captures nothing here because DFlash setsdisable_cuda_graph=True; it still builds the EagerRunner thatforward()needs.) 4-line diff, identical pattern to the already-merged eagle3 change.
Since this is the second target model needing these exact calls, hoisting
them into
SGLangRunneritself would prevent a third recurrence — but thatchanges the pattern #605 established, so this PR keeps the minimal call-site
fix and leaves that refactor as a possible follow-up.
Validation
Patched and ran
prepare_hidden_states.py --model-type dflashonGLM-5.2-FP8, multi-node tp8, sglang 0.5.14: the generation loop completed
(Processed 176 / Filtered 24 short) and hidden-state
.ckpt.gzfiles werewritten correctly. Without the patch, the run crashes at the first batch as
shown above.