[mustafayildirim]

Purpose

Fixes #55087. GLM-5.3 KDA merges the q/k/v short-conv weights into one cached tensor (_merged_conv_weight), built lazily on first forward. Online-RL workflows warm up with dummy weights (cache built), then refit q_conv1d/k_conv1d/v_conv1d — but the cached merge is never invalidated, so inference silently keeps using the pre-refit convolution weights.

Approach

Weight loading dispatches on the param's stamped weight_loader (getattr(param, "weight_loader", default_weight_loader) in AutoWeightsLoader). So the cache can be invalidated at exactly the point the weights change: in Glm5NextLinearAttention.__init__, wrap each conv weight's stamped loader so any write clears _merged_conv_weight. The next forward rebuilds it from the current params. No new hooks, no changes to the reload machinery.

Coverage / limitations

• ✅ Initial load (loaders run before first forward; cache built after anyway)
• ✅ Checkpoint-format refit: reload_weights → model.load_weights → stamped loaders → cache cleared
• ✅ Layerwise reload: re-dispatches through the same stamped loaders
• ❌ Kernel-format refit (param.copy_ bypass) writes params without loaders — the cache is NOT invalidated there. That path also skips every other derived-weight rebuild in the tree (same class of limitation); noted for follow-up if it matters in practice.

A _version-based approach was considered and rejected: the default loader writes via param.data.copy_, which bumps neither param._version nor param.data._version (verified on torch 2.10/ROCm), so version tracking cannot see the main refit path.

Verification

Ad-hoc functional test on MI350X against a real Glm5NextLinearAttention (gloo TP1, minimal Glm5NextConfig):

Step	Result
Build layer, populate merge cache	OK
Refit q_conv1d through its real stamped weight loader	cache invalidated
Next merge uses refit values	OK
No refit → cache persists (no spurious invalidation)	OK

Pre-commit: all hooks pass. No GPU-unit test added (KDA layer requires a full VllmConfig + mamba cache plumbing; happy to add one if maintainers want it).

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mustafayildirim]

Ping for review — fixes the #55087 stale-weights bug for GLM-5.3 KDA in online-RL refit loops. Verified on MI350X against a real Glm5NextLinearAttention: cache invalidates on any loader-mediated weight write, rebuilds from the new values, and no spurious invalidation on untouched layers. The gap (raw param.copy_ refits bypassing loaders) is documented in the PR.

@gdhyi @dllehr-amd

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @mustafayildirim.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork