Skip to content

fix(glm5next): skip PP-missing layers in FP8 dequant load helpers - #56

Open
lalalabreeze wants to merge 1 commit into
wtdcode:masterfrom
lalalabreeze:fix/glm5next-fp8-pp-missing-layer
Open

fix(glm5next): skip PP-missing layers in FP8 dequant load helpers#56
lalalabreeze wants to merge 1 commit into
wtdcode:masterfrom
lalalabreeze:fix/glm5next-fp8-pp-missing-layer

Conversation

@lalalabreeze

Copy link
Copy Markdown

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 in Glm5NextModel.load_weights guard with is_pp_missing_parameter, but the two FP8 dequant helpers added for glm5next lacked this protection:

  • _try_load_fp8_indexer_wk: crashed with KeyError at params_dict[fused_name] (wk_weights_proj param does not exist on this rank).
  • _try_load_fp8_attn_proj: the if target_s in params_dict: return False guard passes for missing layers (the whole layer is absent), then crashed with KeyError at params_dict[target_w].

The fix mirrors the existing pattern in deepseek_v2.py (_try_load_fp8_indexer_wk there already takes pp_missing_layer_names): collect get_pp_missing_layer_names(self) once per load_weights and skip matching weights before buffering. mtp.py callers 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

Testing

  • New unit tests: tests/models/test_glm5next_fp8_pp_load.py — verifies both helpers skip PP-missing layers (no KeyError, nothing buffered) and still dequantize + load owned layers correctly.
  • Commands run:
    • python -m py_compile vllm/models/glm5next/nvidia/model.py vllm/models/glm5next/nvidia/mtp.py tests/models/test_glm5next_fp8_pp_load.py — pass
    • ruff check / ruff format on the three touched files — pass
    • pytest tests/models/test_glm5next_fp8_pp_load.py — to be run before merge (environment install was interrupted; submitter will confirm results here)
  • Reproduction of the original failure: loading with --tensor-parallel-size 4 --pipeline-parallel-size 2 on a block-FP8 GLM-5.3-Flash checkpoint previously raised KeyError during 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.

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>
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