Skip to content

fix(qwen4): freeze forward-only QSA indexer - #203

Open
leooop-al wants to merge 1 commit into
modelscope:mainfrom
leooop-al:fix/freeze_fwd_only_qsa_indexer_params
Open

leooop-al wants to merge 1 commit into
modelscope:mainfrom
leooop-al:fix/freeze_fwd_only_qsa_indexer_params

Conversation

@leooop-al

Copy link
Copy Markdown

Summary

Freeze QSAIndexer parameters after construction so they are excluded from
optimizer groups and Megatron DDP gradient buckets.

Problem

All QSA selection paths run under torch.no_grad() and produce discrete boolean
masks or integer token indices. There is currently no auxiliary QSA indexer loss,
so index_qk_proj, q_layernorm, and k_layernorm cannot receive gradients from
the language-model loss.

Despite that, these parameters default to requires_grad=True. Megatron DDP then
adds them to gradient buckets and installs backward hooks that never fire. With
overlapped gradient reduction enabled, the next zero_grad_buffer() fails because
the bucket ready count is incomplete.

A two-node reproduction reported exactly these missing parameters:

first-batch hook mismatch: ready=31 expected=34 missing=3
indexer.index_qk_proj.weight
indexer.k_layernorm.weight
indexer.q_layernorm.weight

Fix

Call self.requires_grad_(False) after all QSAIndexer parameters are created.
This preserves checkpoint loading and forward selection while keeping the
forward-only parameters out of optimizer state and DDP gradient buckets.

This is preferable to synthesizing zero gradients or disabling
overlap_grad_reduce, because neither workaround changes the fact that the
current QSA indexer has no gradient path.

Tests

  • Assert every QSA indexer parameter has requires_grad=False and grad is None.
  • Assert all pretrained indexer parameters remain present in state_dict().
  • Run the existing forward-only and QSA selection tests.

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