Skip to content

fix(dflash): add sglang 0.5.14 post-load init to the DFlash target model#658

Open
jessiewei7 wants to merge 1 commit into
sgl-project:mainfrom
jessiewei7:pr/dflash-sglang0514
Open

fix(dflash): add sglang 0.5.14 post-load init to the DFlash target model#658
jessiewei7 wants to merge 1 commit into
sgl-project:mainfrom
jessiewei7:pr/dflash-sglang0514

Conversation

@jessiewei7

@jessiewei7 jessiewei7 commented Jul 7, 2026

Copy link
Copy Markdown

Summary

On current main (which pins sglang==0.5.14 since #605),
prepare_hidden_states.py --model-type dflash crashes immediately:

File ".../dflash_target_model.py", line 141, in _extend
    batch = ScheduleBatch.init_new(req_to_token_pool=self.model_runner.req_to_token_pool, ...)
File ".../sglang/srt/managers/schedule_batch.py", line 1866, in init_new
    device=req_to_token_pool.device,
AttributeError: 'NoneType' object has no attribute 'device'

Root cause

sglang 0.5.14 refactored ModelRunner: __init__ no longer allocates the
memory pool. Pool allocation moved to explicit post-load calls
(alloc_memory_pool()init_memory_pool() in
model_runner_kv_cache_mixin.py), normally invoked by the
Scheduler/TpModelWorker. SpecForge drives the ModelRunner directly through
its standalone SGLangRunner, so those calls never happen and
req_to_token_pool / token_to_kv_pool_allocator stay None.

#605 fixed exactly this for eagle3_target_model.py (added the three
post-load calls with a detailed comment) but did not touch
dflash_target_model.py, leaving the DFlash extraction path broken on the
very sglang version the repo now pins.

Fix

Mirror #605's eagle3 fix in SGLangDFlashTargetModel.from_pretrained():

model_runner.alloc_memory_pool()
model_runner.init_attention_backends()
model_runner.init_cuda_graphs()

(init_cuda_graphs() captures nothing here because DFlash sets
disable_cuda_graph=True; it still builds the EagerRunner that forward()
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 SGLangRunner itself would prevent a third recurrence — but that
changes 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 dflash on
GLM-5.2-FP8, multi-node tp8, sglang 0.5.14: the generation loop completed
(Processed 176 / Filtered 24 short) and hidden-state .ckpt.gz files were
written correctly. Without the patch, the run crashes at the first batch as
shown above.

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().
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

@jessiewei7

Copy link
Copy Markdown
Author

@FrankLeeeee PTAL — dflash training crashes on main with sglang 0.5.14 (same pattern as #605). Companion fixes: #663 #659 #665.

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