cp: feat(qwen3_5): port dense Qwen3.5 to a custom-model(2557) into r0.5.0#2663
Merged
Conversation
…tion (#2557) * feat(qwen3_5): port dense Qwen3.5 to a native custom-model implementation Build the dense Qwen3.5 model as a custom implementation (its own backbone on the shared Block + CPAwareGatedDeltaNet + fp32 SSMGate built at __init__), mirroring Qwen3.5-MoE, instead of mutating the HF model at runtime via patch_hf_model class-swaps + an __getattr__ monkeypatch. Removes patch_hf_model/apply_model_runtime_patches, Qwen3_5TextModelPP, and decoder_layer.py. The native fp32 SSMGate keeps A_log/dt_bias in fp32 storage (true master weights) under a bf16 bulk dtype via its own fp32 FSDP group. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * fix(qwen3_5): align dense VLM lm_head dtype and refresh stale tests The native dense backbone is swapped in after HF's super().__init__, but HF post_init -> init_weights -> initialize_weights already ran against the inherited text model, casting lm_head to the bulk (bf16) dtype. Realign lm_head to the backbone dtype after the swap so final hidden states and lm_head agree, and guard the backbone init_weights(buffer_device=...) call against the still-HF text model during post_init (try/except, mirroring the MoE sibling). Refresh tests for the native-fp32-SSMGate design: dense/VLM adapters now route both A_log and dt_bias into the _fp32_params holder; add coverage for cast_model_to_dtype(skip_modules=...) and the MoE adapter fp32 route/strip helpers. Pin the tiny MTP test config to fp32 to match HF parity. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * test(qwen3_5): cover native dense backbone and fp32 SSMGate machinery Add CPU unit tests for the custom-model building blocks that replaced the old runtime patch_hf_model: SSMGate / install_ssm_gate / the _SSMGateParam descriptor (fp32 param relocation + attribute resolution), _resolve_ssm_dtype, Fp32SafeQwen3_5TextRotaryEmbedding (fp32 inv_freq survives a bf16 cast), and Qwen3_5DenseTextBackbone forward-shape smoke tests (full/linear/mixed layer types, inputs_embeds, KV-cache guard). Addresses the review's coverage gap for the new native implementation. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * style(qwen3_5): apply ruff format to cp_linear_attn and moe parallelizer CI ruff format --check (ruff 0.9.10) flagged trailing blank lines and a wrapped call that fit on one line in the dense custom-model changes. No behavior change. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * fix(qwen3_5): detect layer_types hybrids so CP works with SDPA (no TE) Qwen3.5 dense+MoE are hybrid (linear_attention gated-delta + full_attention) but were not detected as hybrid because their config exposes the structure via per-layer ``layer_types`` rather than ``layers_block_type`` / ``is_hybrid_model``. Without hybrid detection, ``ModelSupports.supports_cp`` took the "requires TE attention" branch, wrongly forcing ``backend.attn=te`` for context-parallel (cp_size>1). ``_is_hybrid`` now also returns True when a config's ``layer_types`` contains "linear_attention" and drills into a nested ``text_config`` (VLM configs), so ``supports_cp`` allows ``attn in (te, sdpa)`` for these models and dense CP works with SDPA again. Verified on 8xH100 (dense 4B, fp32, medpix): cp_size=1 and cp_size=2 both train 20 steps with CP1==CP2 parity (step-0 loss diff 2e-4, max 0.034). Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * fix(qwen3_5): guard fp32-holder sharding for non-module blocks; fix stale tests L0_Unit_Tests_CPU surfaced two regressions from the dense custom-model port: - _shard_fp32_param_holders unconditionally called block.named_modules(), which broke moe/test_parallelizer.py's plain (non-nn.Module) DummyBlock stubs. Guard with hasattr(block, "named_modules") and return an empty set — a block that cannot enumerate submodules cannot hold _fp32_params. - test_auto_model TestModelRuntimePatches asserted Qwen3.5 runtime-patch dispatch, but the port emptied _MODEL_RUNTIME_PATCHES (models build their CP/fp32-gate modules at construction). Rewrite the dispatch/dedup tests to inject a temporary registry entry so they cover the mechanism without a built-in entry. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * fix(qwen3_5): TP/PP-robust attention reshape, AC-safe Block._mlp, native TP plan registration - Qwen3NextAttention.forward: infer local head count via -1 in q/k/v reshape so colwise-sharded q/k/v_proj (TP) yield the correct per-rank head count instead of asserting the full (unsharded) head count. - Block._mlp: unwrap CheckpointWrapper (_checkpoint_wrapped_module) before the MLP/MoE isinstance dispatch so activation checkpointing does not break the dense (MLP) path. - optimized_tp_plans: register the native Qwen3_5ForConditionalGeneration and Qwen3_5ForCausalLM qualnames against _parallelize_qwen3_5_vlm so the custom dense model routes to the Qwen3.5 TP plan. Verified: dense 4B TP2+PP4 (AC on, fp32 storage, pipeline bf16) trains; MoE small fp32 loss unchanged (12.8736->12.8792->12.8772). Signed-off-by: HuiyingLi <willwin.lee@gmail.com> * test(qwen3_5): pin rope_fusion=False and run linear-attn forward on CUDA The dense backbone + MTP unit tests built BackendConfig without setting rope_fusion, which defaults to (HAVE_TE and torch.cuda.is_available()). That is False on CPU dev boxes but True on the GPU CI runner, where Qwen3.5's gated-query RoPE hits TE's fused rope kernel and fails with 'expected 4D tensor'. Every Qwen3.5 recipe sets rope_fusion=false, so pin it in the test backends to exercise the supported (non-fused) path deterministically. Also make test_forward_shape_mixed_layers build the backbone + inputs on CUDA (skipping when unavailable): the linear_attention block runs FLA's gated-delta-rule / causal_conv1d CUDA kernels, which require CUDA tensors and were previously masked by the earlier rope failure. Signed-off-by: HuiyingLi <willwin.lee@gmail.com> --------- Signed-off-by: HuiyingLi <willwin.lee@gmail.com> (cherry picked from commit 28541ee) Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
Contributor
Author
|
/ok to test f312292 |
Contributor
Author
|
/ok to test f312292 |
HuiyingLi
approved these changes
Jun 20, 2026
HuiyingLi
changed the title
feat(qwen3_5): port dense Qwen3.5 to a native custom-model implementation (2557) into r0.5.0feat(qwen3_5): port dense Qwen3.5 to a custom-model(2557) into r0.5.0
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
2 participants
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.
Cherry-pick of #2557 (28541ee) onto
r0.5.0.Why this is on the release branch: #2557 was not originally labeled
r0.5.0. It is being backported as a prerequisite for #2484 (preserve fp32A_login Qwen3.5-MoE / Qwen3-Next GatedDeltaNet). #2484 builds on the generalized_fp32_paramsholder mechanism and the native Qwen3.5 model that #2557 introduces, so it cannot be cherry-picked alone. The #2484 cherry-pick (#2664) is stacked on this branch and should merge after this one.Release-branch resolutions
optimized_tp_plans.py: kept both r0.5.0's Falcon-H1 TP entries (from fix(distributed): register Falcon-H1 TP plan to fix 34B PEFT OOM #2589's cp cp:fix(distributed): register Falcon-H1 TP plan to fix 34B PEFT OOM (2589)intor0.5.0#2605) and feat(qwen3_5): port dense Qwen3.5 to a native custom-model implementation #2557's new native-Qwen3.5 qualname registrations — non-overlapping additions.tests/unit_tests/models/qwen3_5/test_qwen3_5_mtp.py: kept r0.5.0'sdispatcher="torch"(post-refactor(moe): remove enable_deepep, switch failing ep recipes to hybridep #2630, already on r0.5.0) and added feat(qwen3_5): port dense Qwen3.5 to a native custom-model implementation #2557'srope_fusion=False; dropped the obsoleteenable_deepep=False(superseded bydispatcher="torch", mirroring cp:refactor(moe): remove enable_deepep, switch failing ep recipes to hybridep (2630)intor0.5.0#2644's migration).test_qwen3_5_pp.py(theQwen3_5TextModelPPit covered is removed by the port; supersedes r0.5.0's fix(qwen3_5): make dense VLM pipeline-parallel safe #2524 cp, which feat(qwen3_5): port dense Qwen3.5 to a native custom-model implementation #2557 was built on top of onmain).examples/vlm_finetune/qwen3_5/qwen3_6_27b_medpix_cp{1,2}_2k_2node_100steps.yaml.Validation (CPU/local):
py_compileclean (20 files); nativeqwen3_5.model+optimized_tp_plansimport cleanly; 267 CPU unit tests pass (full qwen3_5 suite,moe/test_parallelizer,_transformersauto_model + capabilities,cast_model_to_dtype). GPU parity not re-run locally — see #2557's own 8×H100 validation.