Repository navigation
[Bugfix] Support strided weights in sharded state saving - #61432
Open
mgehre-amd wants to merge 2 commits into
Open
mgehre-amd wants to merge 2 commits into
mgehre-amd wants to merge 2 commits into
Conversation
save_sharded_state serializes the in-memory state dict, i.e. weights after process_weights_after_loading. On ROCm, VLLM_ROCM_FP8_PADDING and VLLM_ROCM_MOE_PADDING (both on by default) offset weight row strides there via F.pad(...)[..., :-n], leaving non-contiguous views. _filter_subtensors flattened every tensor with view(-1), which raises on such views, and safetensors refuses to save non-contiguous tensors, so saving failed for any padded model. Compute a tensor's end address from its strides, and save a contiguous copy. The saved checkpoint holds the plain weights; loading copies them into freshly allocated dense parameters before process_weights_after_loading, which re-applies the padding, so the load side needs no change. Signed-off-by: Matthias Gehre <matthias.gehre@amd.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
6 tasks done
dllehr-amd
self-requested a review
October 20, 2026 14:06
This branch has not been deployed
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.
Purpose
save_sharded_stateserializes each worker's in-memorystate_dict(), i.e. theweights after
process_weights_after_loading. Some post-processing leaves aweight as a non-contiguous view with a widened row stride. On ROCm,
VLLM_ROCM_FP8_PADDING(on by default) does exactly that withF.pad(w, (0, n))[..., :-n].Saving such a model fails twice over:
ShardedStateLoader._filter_subtensorscomputes a tensor's end address withtensor.view(-1)[-1], which raises on a strided view:RuntimeError: view size is not compatible with input tensor's size and stride.safetensors.save_filerefuses non-contiguous tensors:ValueError: You are trying to save a non contiguous tensor.Test plan
1. FP8 model
Qwen/Qwen3-0.6B-FP8python examples/features/sharded_state/save_sharded_state_offline.py \ --model Qwen3-0.6B-FP8 --output /tmp/q3-fp8-shardedLoading that FP8 checkpoint back is a separate, pre-existing problem: it fails
with
Missing keys (...attn.q_scale, ...k_scale, ...v_scale, ...prob_scale)on main too, even with
VLLM_ROCM_FP8_PADDING=0, because those attentionscales are deleted in
process_weights_after_loadingand so never saved. Itis out of scope here.
**2. w4a16 model **
Qwen/Qwen3-0.6B:python examples/features/sharded_state/save_sharded_state_offline.py \ --model Qwen3-0.6B --output /tmp/q3-sharded python examples/features/sharded_state/load_sharded_state_offline.py \ --model /tmp/q3-sharded --load-format sharded_state \ --prompt "The capital of France is" --max-tokens 32 --temperature 0Not a duplicate
No duplicate PRs found.
AI assistance
AI assistance was used for this PR. I reviewed every changed line, and ran the
tests and the end-to-end round trip listed above.