[Bugfix][DeepseekV4] Guard expert loader against UnboundLocalError#42835
[Bugfix][DeepseekV4] Guard expert loader against UnboundLocalError#42835varjoranta wants to merge 2 commits into
Conversation
`DeepseekV4Model.load_weights` adds `name_mapped` to `loaded_params`
unconditionally after the `for mapping in expert_mapping` loop, but
`name_mapped` is only bound inside the loop, after the
`if weight_name not in name: continue` guard:
for mapping in expert_mapping:
param_name, weight_name, expert_id, shard_id = mapping
if weight_name not in name:
continue
name_mapped = name.replace(weight_name, param_name)
...
loaded_params.add(name_mapped) # UnboundLocalError
If a `.experts.` weight matches no entry in `expert_mapping` (the loop
body never reaches the assignment for any mapping), `name_mapped` is
never bound and `loaded_params.add(name_mapped)` raises
UnboundLocalError, aborting checkpoint load.
This is reachable with non-canonical DSV4 checkpoints whose expert
tensor names differ from the model's `expert_mapping` (e.g.
custom-quantized variants). It surfaced loading a custom 3-bit
quantized DSV4-Flash checkpoint during the safetensors load.
Fix: initialize `name_mapped = None` before the loop and skip the
weight when no mapping matched (consistent with the other
unmatched-weight branches in this method, which `continue`) instead
of raising.
The affected path is `load_weights`, which requires a fully
instantiated model + distributed init to exercise; vLLM's existing
DSV4 tests (`tests/models/test_deepseek_v4_mega_moe.py`) are
CUDA-gated and cover only pure helpers, so there is no unit-test
seam for this path. The change is a self-contained
unbound-variable guard; happy to add a CUDA integration test if
preferred.
Closes vllm-project#42769
Signed-off-by: Hannu Varjoranta <hannu@varjosoft.com>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
mergify
Bot
added
deepseek
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request addresses a potential UnboundLocalError in the load_weights method of the DeepseekV4 model by initializing name_mapped and skipping parameters that do not match any expert mapping. The reviewer suggests refining this logic by introducing a success flag to track whether the weight was actually loaded for the current rank, ensuring that loaded_params only includes successfully initialized parameters.
| and loaded_weight.dtype == torch.float8_e8m0fnu | ||
| ): | ||
| loaded_weight = loaded_weight.view(torch.uint8) | ||
| name_mapped = None |
There was a problem hiding this comment.
Initializing name_mapped to None prevents the UnboundLocalError, but we should also track whether any expert mapping was successfully loaded. If weight_loader returns False for all mappings (e.g., because the expert is not assigned to the current rank), the parameter should not be added to loaded_params. Initializing a success flag here allows for a more robust check after the loop.
name_mapped = None
success = False| if name_mapped is None: | ||
| # No expert mapping matched (e.g. a non-canonical | ||
| # checkpoint whose expert tensor names differ from | ||
| # this model's expert_mapping); nothing was loaded | ||
| # for this weight, so skip it instead of raising | ||
| # UnboundLocalError. | ||
| continue |
There was a problem hiding this comment.
Instead of only checking if name_mapped is None, we should check if success is True. If success is False, it means either no mapping matched the weight name or the weight_loader did not actually load any data for this rank. In either case, we should skip adding the parameter to loaded_params to avoid incorrectly marking it as initialized when it might be empty on this rank.
| if name_mapped is None: | |
| # No expert mapping matched (e.g. a non-canonical | |
| # checkpoint whose expert tensor names differ from | |
| # this model's expert_mapping); nothing was loaded | |
| # for this weight, so skip it instead of raising | |
| # UnboundLocalError. | |
| continue | |
| if not success: | |
| # No expert mapping matched or the weight loader | |
| # did not load this weight (e.g. not for this rank); | |
| # nothing was loaded for this weight, so skip it | |
| # instead of raising UnboundLocalError or incorrectly | |
| # marking the parameter as loaded. | |
| continue |
A non-None name_mapped no longer implies the weight was loaded: when every expert mapping is attempted but the loader returns False for this rank, name_mapped is set yet nothing was loaded. Track an explicit success flag and skip (continue) unless a mapping actually loaded, so loaded_params only records weights that were really applied. Signed-off-by: Hannu Varjoranta <hannu@varjosoft.com>
1 participant
Suggested fix for #42769.
DeepseekV4Model.load_weightsaddsname_mappedtoloaded_paramsunconditionally after thefor mapping in expert_mappingloop, butname_mappedis only bound inside the loop, after theif weight_name not in name: continueguard. A.experts.weight that matches no entry inexpert_mappingleavesname_mappedunbound, soloaded_params.add(name_mapped)raisesUnboundLocalErrorand aborts checkpoint load.This initializes
name_mapped = Nonebefore the loop and skips the weight when no mapping matched (consistent with the other unmatched-weight branches in this method, whichcontinue), instead of raising. Full repro and trace are in #42769.The affected path is
load_weights, which requires a fully instantiated model plus distributed init to exercise; vLLM's existing DSV4 tests are CUDA-gated and cover only pure helpers, so there is no unit-test seam. This is a self-contained unbound-variable guard; glad to add a CUDA integration test if a maintainer prefers.Closes #42769
Update
Addressed the automated review in
cd5cd496de.The initial fix avoided
UnboundLocalErrorby initializingname_mapped, but that was not enough: if the expert mapping loop runs and everyweight_loader(..., return_success=True)call returnsFalse, the code could still add a parameter name toloaded_paramseven though no data was loaded for this rank.The follow-up commit tracks an explicit
successflag and only records the parameter when the loader reports that it actually loaded the weight. Successful canonical loads are unchanged; this only hardens the unsuccessful mapping / non-local expert path soloaded_paramsremains accurate.