fix(glm5next): skip PP-missing layers in FP8 dequant load helpers - #56
Open
lalalabreeze wants to merge 1 commit into
Open
fix(glm5next): skip PP-missing layers in FP8 dequant load helpers#56lalalabreeze wants to merge 1 commit into
lalalabreeze wants to merge 1 commit into
Conversation
With pipeline parallelism every rank iterates the full checkpoint and layers held by other stages are built as PPMissingLayer (no parameters). The normal load paths guard with is_pp_missing_parameter, but the two FP8 dequant helpers (_try_load_fp8_indexer_wk, _try_load_fp8_attn_proj) did not, so loading a block-FP8 checkpoint with --pipeline-parallel-size > 1 crashed with a KeyError when looking up params of layers owned by other ranks. Mirror deepseek_v2: collect pp_missing_layer_names and skip those weights before buffering. Co-authored-by: ZCode <noreply@zcode.dev>
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.
Description
Fixes a crash when loading a GLM-5.3-Flash block-FP8 checkpoint with
--pipeline-parallel-size > 1(e.g. TP=4 + PP=2).With PP, every rank iterates the full checkpoint and layers held by other stages are built as
PPMissingLayer(no parameters). The normal load paths inGlm5NextModel.load_weightsguard withis_pp_missing_parameter, but the two FP8 dequant helpers added for glm5next lacked this protection:_try_load_fp8_indexer_wk: crashed withKeyErroratparams_dict[fused_name](wk_weights_projparam does not exist on this rank)._try_load_fp8_attn_proj: theif target_s in params_dict: return Falseguard passes for missing layers (the whole layer is absent), then crashed withKeyErroratparams_dict[target_w].The fix mirrors the existing pattern in
deepseek_v2.py(_try_load_fp8_indexer_wkthere already takespp_missing_layer_names): collectget_pp_missing_layer_names(self)once perload_weightsand skip matching weights before buffering.mtp.pycallers are updated for the new signature (the MTP model builds all layers, so the list is empty there and behavior is unchanged).Not a duplicate
vllm-project/vllmdoes not containvllm/models/glm5next/(fork-specific code), so there is no upstream PR to backport from.Testing
tests/models/test_glm5next_fp8_pp_load.py— verifies both helpers skip PP-missing layers (noKeyError, nothing buffered) and still dequantize + load owned layers correctly.python -m py_compile vllm/models/glm5next/nvidia/model.py vllm/models/glm5next/nvidia/mtp.py tests/models/test_glm5next_fp8_pp_load.py— passruff check/ruff formaton the three touched files — passpytest tests/models/test_glm5next_fp8_pp_load.py— to be run before merge (environment install was interrupted; submitter will confirm results here)--tensor-parallel-size 4 --pipeline-parallel-size 2on a block-FP8 GLM-5.3-Flash checkpoint previously raisedKeyErrorduring weight load; with this change the missing-layer weights are skipped like every other path.AI assistance
This change was developed with AI assistance (ZCode). The submitting human has reviewed every changed line and is responsible for running the tests above before merge.