Skip to content

feat: support dense Qwen3 weight conversion - #114

Merged
chaokunyang merged 1 commit into
inclusionAI:mainfrom
Le8r0nJames:zjw/qwen3-dense-register
Aug 7, 2026
Merged

feat: support dense Qwen3 weight conversion#114
chaokunyang merged 1 commit into
inclusionAI:mainfrom
Le8r0nJames:zjw/qwen3-dense-register

Conversation

@Le8r0nJames

Copy link
Copy Markdown
Contributor

What does this PR do?

Qwen3ForCausalLM is not registered in awex/models, so it falls back to the
default converters. Those follow the bailing-style attention naming, and
SGlangToHFWeightConverter._convert_layer_norm_param raises
NotImplementedError on self_attn.q_norm.weight. The failure happens while
the reader is building its parameter metadata, so the reader never publishes
infer_conf and the writer waits for it indefinitely: the run stalls rather
than reporting the unsupported name.

Dense and MoE Qwen3 share the layout the existing qwen3_moe converters
already target:

  • canonical self_attn.{q,k,v,o}_proj with per-head self_attn.{q,k}_norm
  • GQA head grouping in the fused Megatron linear_qkv
  • gate/up projections that sglang serves fused through
    MergedColumnParallelLinear while the train side reports them split

The dense model only lacks the expert parameters, which those converters emit
nothing for, so registering the architecture against them is enough and the
base converters are untouched. This mirrors ling_linear.py, which reuses
ling.py's converters for its own architectures.

Dense Qwen3 also ties its output weights, which surfaced a second gap. Three
places add the lm_head.weight alias for tied embeddings, and they disagree:

engine gate
reader/weights_reader.py none
meta/train_meta_resolver.py none
meta/infer_meta_resolver.py engine_name == "vllm"

The reader therefore adds the alias on sglang while the inference metadata does
not declare it, and the transfer plan reports lm_head.weight as missing. This
PR drops the engine gate so the third site matches the other two.

Verification

  • awex/tests passes in full, on top of four new cases for the dense
    registration and four for the tied-embedding alias.
  • Negative check against current main: asking the registry for
    Qwen3ForCausalLM logs Model Qwen3ForCausalLM not found, using default strategy, returns the base SGlangToHFWeightConverter, and
    self_attn.q_norm.weight raises NotImplementedError. The new tests fail
    without this change and the tied-embedding case fails for sglang while
    passing for vllm.
  • A 20-step colocated GRPO run on Qwen3-0.6B (sglang inference, Megatron
    training, AWEX weight transfer) completes cleanly: no Unsupported layer norm parameter name, no wait on infer_conf, no missing lm_head.weight, and the
    "using default strategy" fallback no longer appears.

Not covered by unit tests: the Megatron-side dense path. It is exercised by the
end-to-end run above but has no test of its own yet.

Related issues

None.

Does this PR introduce any user-facing change?

  • Does this PR introduce any public API change?
  • Does this PR introduce any binary protocol compatibility change?

No signatures change and no new configuration is required; Qwen3ForCausalLM
is picked up by the existing registry discovery.

Worth noting for review: for a tied-embedding model served by sglang, the
published inference metadata now contains an additional lm_head.weight entry.
That is not a new protocol field. The reader already materialises this alias
regardless of engine, so the change makes the declared metadata agree with what
is actually transferred; previously the two sides disagreed and the transfer
plan failed.

Qwen3ForCausalLM was not registered, so it fell back to the default
converters and raised NotImplementedError on self_attn.q_norm.weight while
the reader was building its parameter metadata. The reader then never
published infer_conf and the writer waited for it indefinitely.

Dense and MoE Qwen3 share the layout the existing converters target:
canonical self_attn.{q,k,v,o}_proj with per-head self_attn.{q,k}_norm, GQA
head grouping in the fused Megatron linear_qkv, and gate/up projections
that sglang serves fused. The dense model only lacks the expert
parameters, which those converters emit nothing for, so registering the
architecture against them is enough.

Dense Qwen3 also ties its output weights, which exposed a second gap: the
lm_head alias for tied embeddings was only added for vLLM. sglang ties
them the same way while the training side still publishes lm_head, and the
reader adds the same alias on its side, so without it the transfer plan
reports lm_head.weight as missing.
@chaokunyang
chaokunyang merged commit fe6635e into inclusionAI:main Aug 7, 2026
3 checks passed
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.

2 participants