From 9d4f9c61e452a8ea0715750cde572d192aa066ac Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Wed, 1 Jul 2026 20:10:46 +0900 Subject: [PATCH] feat(qwen3-next): add pipeline-parallel stage support Add Qwen3NextStageModel and Qwen3NextStageExecutor so qwen3_next checkpoints shard across pipeline-parallel stages, mirroring the existing qwen3_5 stage plumbing. Each stage loads only its layer range, builds and advances the correct per-layer cache type (gated-delta ArraysCache for linear layers, KVCache for attention layers), and derives the SSM and attention masks from the stage-local layers using global layer indices. Registers the Qwen3Next StageFamily variant with name mapping, PP-capable set, and ModelType dispatch. --- .../stage_executor/family_registry_tests.rs | 13 + .../pipeline/stage_executor/mod.rs | 14 + .../pipeline/stage_executor/qwen3_next.rs | 102 ++++++ src/models/qwen3_next.rs | 310 ++++++++++++++++++ 4 files changed, 439 insertions(+) create mode 100644 src/distributed/pipeline/stage_executor/qwen3_next.rs diff --git a/src/distributed/pipeline/stage_executor/family_registry_tests.rs b/src/distributed/pipeline/stage_executor/family_registry_tests.rs index 770863995..1d787a435 100644 --- a/src/distributed/pipeline/stage_executor/family_registry_tests.rs +++ b/src/distributed/pipeline/stage_executor/family_registry_tests.rs @@ -64,6 +64,19 @@ fn new_issue_345_families_are_registered() { } } +#[test] +fn qwen3_next_family_is_registered() { + // Floor test for the Qwen3-Next pipeline-parallel stage family. Removing + // it (or its stable name) without bumping the pipeline capability protocol + // version would break running multi-host qwen3-next deployments. + let families: HashSet = supported_families().iter().copied().collect(); + assert!( + families.contains(&StageFamily::Qwen3Next), + "StageFamily::Qwen3Next must appear in supported_families()", + ); + assert_eq!(StageFamily::Qwen3Next.name(), "qwen3_next"); +} + #[test] fn all_stage_family_variants_have_stable_names() { // If a new variant is added, this match must be updated to give it a diff --git a/src/distributed/pipeline/stage_executor/mod.rs b/src/distributed/pipeline/stage_executor/mod.rs index 687473b08..248e91188 100644 --- a/src/distributed/pipeline/stage_executor/mod.rs +++ b/src/distributed/pipeline/stage_executor/mod.rs @@ -41,6 +41,7 @@ mod mixtral; mod nemotron_h; mod qwen3; mod qwen35; +mod qwen3_next; #[cfg(test)] #[path = "family_registry_tests.rs"] @@ -69,6 +70,7 @@ use mistral::MistralStageExecutor; use mixtral::MixtralStageExecutor; use nemotron_h::NemotronHStageExecutor; use qwen3::Qwen3StageExecutor; +use qwen3_next::Qwen3NextStageExecutor; use qwen35::Qwen35StageExecutor; /// Input payload for a single stage-local forward. @@ -238,6 +240,7 @@ pub enum StageFamily { Glm4MoeLite, GlmMoeDsa, Qwen3, + Qwen3Next, Qwen35, Qwen35Vlm, Qwen35Moe, @@ -267,6 +270,7 @@ impl StageFamily { Self::Glm4MoeLite => "glm4_moe_lite", Self::GlmMoeDsa => "glm_moe_dsa", Self::Qwen3 => "qwen3", + Self::Qwen3Next => "qwen3_next", Self::Qwen35 => "qwen3_5", Self::Qwen35Vlm => "qwen3_5_vlm", Self::Qwen35Moe => "qwen3_5_moe", @@ -318,6 +322,7 @@ pub fn supported_families() -> &'static [StageFamily] { StageFamily::Qwen35Moe, StageFamily::Qwen35MoeVlm, StageFamily::Qwen35Vlm, + StageFamily::Qwen3Next, ]; FAMILIES } @@ -348,6 +353,7 @@ fn resolve_stage_family(model_dir: &Path) -> Result { ModelType::Glm4MoeLite => StageFamily::Glm4MoeLite, ModelType::GlmMoeDsa => StageFamily::GlmMoeDsa, ModelType::Qwen3 => StageFamily::Qwen3, + ModelType::Qwen3Next => StageFamily::Qwen3Next, ModelType::Qwen35 => StageFamily::Qwen35, ModelType::Qwen35VLM => StageFamily::Qwen35Vlm, ModelType::Qwen35Moe => StageFamily::Qwen35Moe, @@ -483,6 +489,14 @@ fn load_family_backend( stage_index, )?)) } + StageFamily::Qwen3Next => { + ensure_no_adapter(adapter_path, family)?; + Ok(Box::new(Qwen3NextStageExecutor::load( + model_dir, + filter, + stage_index, + )?)) + } StageFamily::Qwen35 | StageFamily::Qwen35Vlm | StageFamily::Qwen35Moe diff --git a/src/distributed/pipeline/stage_executor/qwen3_next.rs b/src/distributed/pipeline/stage_executor/qwen3_next.rs new file mode 100644 index 000000000..54f1b3980 --- /dev/null +++ b/src/distributed/pipeline/stage_executor/qwen3_next.rs @@ -0,0 +1,102 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Pipeline-parallel stage executor for Qwen3-Next. +//! +//! Qwen3-Next is a hybrid model that interleaves linear (GatedDeltaNet, +//! `GatedDeltaCache`) and full-attention (`KVCache`) layers. This executor is +//! the sibling of the Qwen 3.5 stage executor (`qwen35.rs`): it wraps a +//! [`Qwen3NextStageModel`] holding only the stage's layer range and keeps the +//! stage-local heterogeneous cache vector in a +//! [`PointerOwnedCacheStore`], keyed off the identity of the external +//! `KVCache` handle slice the runtime hands in. +//! +//! The external handles carry only the sequence offset back to admission and +//! scheduler code; the real conv-state / SSM-state / KV tensors live in the +//! internal [`Qwen3NextCache`] vector, one entry per LOCAL layer, whose +//! `Linear`/`Attention` variant is chosen from the GLOBAL layer index so the +//! cache type and mask match the single-process path exactly. +//! +//! Like the Qwen 3.5 executor, the speculative-decoding hooks are not +//! propagated across the stage boundary: mlxcel's pipeline-parallel runner +//! does not support speculative decoding for any family today. + +use std::path::Path; + +use anyhow::Result; +use mlxcel_core::layers::KVCache; +use mlxcel_core::{MlxArray, copy}; + +use crate::models::qwen3_next::{Qwen3NextCache, Qwen3NextStageModel}; + +use super::common::PointerOwnedCacheStore; +use super::{FamilyStageExecutor, LayerFilter, StageExecutionInput, StageExecutionOutput}; + +pub struct Qwen3NextStageExecutor { + model: Qwen3NextStageModel, + cache_store: PointerOwnedCacheStore, +} + +impl Qwen3NextStageExecutor { + pub fn load(model_dir: &Path, filter: &LayerFilter, stage_index: usize) -> Result { + Ok(Self { + model: Qwen3NextStageModel::load(model_dir, filter, stage_index) + .map_err(anyhow::Error::msg)?, + cache_store: PointerOwnedCacheStore::default(), + }) + } +} + +impl FamilyStageExecutor for Qwen3NextStageExecutor { + fn make_caches(&self) -> Vec { + (0..self.model.num_layers()) + .map(|_| KVCache::new()) + .collect() + } + + fn release_caches(&self, caches: &[KVCache]) { + self.cache_store.release_caches(caches); + } + + fn execute( + &self, + input: StageExecutionInput<'_>, + caches: &mut [KVCache], + _mask: Option<&MlxArray>, + ) -> Result { + let mut internal_caches = self.cache_store.caches_for_sequence( + caches, + || self.model.make_caches(), + Qwen3NextCache::offset, + "qwen3-next sequence cache entry must exist", + ); + + let output = match input { + StageExecutionInput::TokenIds(input_ids) => self + .model + .execute_from_token_ids(input_ids, &mut internal_caches), + StageExecutionInput::HiddenStates(hidden_states) => self + .model + .execute_from_hidden_states(copy(hidden_states), &mut internal_caches), + } + .map_err(anyhow::Error::msg)?; + + PointerOwnedCacheStore::sync_external_offsets( + caches, + &internal_caches, + Qwen3NextCache::offset, + ); + Ok(output) + } +} diff --git a/src/models/qwen3_next.rs b/src/models/qwen3_next.rs index 66ba77173..ca9c2f869 100644 --- a/src/models/qwen3_next.rs +++ b/src/models/qwen3_next.rs @@ -32,6 +32,9 @@ mod helpers; #[path = "qwen3_next_helpers_tests.rs"] mod helper_tests; +use crate::distributed::pipeline::LayerFilter; +use crate::distributed::pipeline::StageExecutionOutput; +use crate::distributed::pipeline::partial_loading::filter_weight_map; use crate::models::gated_delta::{ GatedDeltaCache, RMSNormGated, gated_delta_update, scaled_fast_rms_norm_no_weight, }; @@ -1601,6 +1604,313 @@ impl LanguageModel for Qwen3NextModel { } } +// Pipeline-parallel stage model. +// +// Stage-local counterpart of [`Qwen3NextModel`], modeled on the Qwen 3.5 +// `Qwen35StageModel`. It loads only the layers in `filter.layer_range`, +// accepts token IDs on the entry stage (the stage that hosts the embedding +// table) and hidden states on intermediate / final stages, and emits hidden +// states on non-final stages or final logits on the last stage. +// +// The load-bearing detail is the hybrid cache. qwen3-next interleaves linear +// (GatedDeltaNet, `GatedDeltaCache`) and full-attention (`KVCache`) layers. +// Each stage builds and advances the cache variant that matches its LOCAL +// layer, derived from the GLOBAL layer index via `config.is_linear_layer` +// (carried on `DecoderLayer::is_linear`), so a stage that starts partway +// through the network still assigns the correct cache type per layer. The +// causal attention mask is derived from the first full-attention layer that +// is actually present in the stage, and linear layers receive no mask (an +// all-valid SSM step), exactly as the single-process `Qwen35Model` +// `forward_internal` path does. +pub(crate) struct Qwen3NextStageModel { + filter: LayerFilter, + embed_tokens: Option, + layers: Vec, + norm: Option, + lm_head: Option, +} + +/// Select the causal-mask offset for a stage from the first full-attention +/// layer present in the stage. Linear (GatedDeltaNet) layers do not track a +/// KV offset, so the mask offset comes from the earliest attention cache in +/// the stage; a stage with no attention layer (all linear) uses offset 0. +/// +/// `is_linear[i]` marks stage-local layer `i` as a linear (GatedDeltaNet) +/// layer. `caches` is the stage-local cache vector in the same order. +fn stage_attention_offset(is_linear: &[bool], caches: &[Qwen3NextCache]) -> i32 { + is_linear + .iter() + .zip(caches.iter()) + .find_map(|(&linear, cache)| (!linear).then_some(cache.offset())) + .unwrap_or(0) +} + +impl Qwen3NextStageModel { + pub(crate) fn load( + model_dir: &Path, + filter: &LayerFilter, + stage_index: usize, + ) -> Result { + let config_path = model_dir.join("config.json"); + let config_str = std::fs::read_to_string(&config_path) + .map_err(|e| format!("Failed to read config.json: {}", e))?; + let config: Qwen3NextConfig = serde_json::from_str(&config_str) + .map_err(|e| format!("Failed to parse config.json: {}", e))?; + + let mut weights = crate::models::load_text_weights(model_dir, None)?; + weights = Qwen3NextModel::sanitize_weights(weights, &config); + + // When embeddings are tied, the final stage produces logits from the + // embedding table, so it must keep `model.embed_tokens.weight` through + // the layer-range filter even though `filter.has_embedding` is false. + let mut effective_filter = filter.clone(); + if config.tie_word_embeddings && filter.has_lm_head { + effective_filter.has_embedding = true; + } + filter_weight_map(&mut weights, &effective_filter); + Self::from_filtered_weights(&weights, &config, filter, stage_index) + } + + fn from_filtered_weights( + weights: &WeightMap, + config: &Qwen3NextConfig, + filter: &LayerFilter, + stage_index: usize, + ) -> Result { + let group_size = config.group_size(); + let bits = config.bits(); + + let load_embeddings = + filter.has_embedding || (config.tie_word_embeddings && filter.has_lm_head); + let embed_tokens = if load_embeddings { + Some(UnifiedEmbedding::from_weights( + weights, + "model.embed_tokens", + group_size, + bits, + )?) + } else { + None + }; + + let mut layers = Vec::with_capacity(filter.num_layers()); + for layer_idx in filter.layer_range.clone() { + layers.push(DecoderLayer::from_weights(weights, config, layer_idx)?); + } + + if layers.is_empty() { + return Err(format!( + "stage {} did not load any layers from range {}..{}", + stage_index, filter.layer_range.start, filter.layer_range.end + )); + } + + let norm = if filter.has_lm_head { + Some(RMSNorm::new( + weights + .get("model.norm.weight") + .map(|w| mlxcel_core::copy(w)) + .ok_or_else(|| "Missing model.norm.weight".to_string())?, + config.rms_norm_eps, + )) + } else { + None + }; + + let lm_head = if filter.has_lm_head && !config.tie_word_embeddings { + Some(UnifiedLinear::from_weights( + weights, "lm_head", group_size, bits, + )?) + } else { + None + }; + + Ok(Self { + filter: filter.clone(), + embed_tokens, + layers, + norm, + lm_head, + }) + } + + pub(crate) fn num_layers(&self) -> usize { + self.layers.len() + } + + /// Build the stage-local hybrid cache vector: one entry per local layer, + /// `Linear`/`Attention` chosen by that layer's `is_linear` flag (which is + /// itself derived from the GLOBAL layer index at load time). + pub(crate) fn make_caches(&self) -> Vec { + self.layers + .iter() + .map(|layer| { + if layer.is_linear { + Qwen3NextCache::Linear(GatedDeltaCache::new()) + } else { + Qwen3NextCache::Attention(KVCache::new()) + } + }) + .collect() + } + + pub(crate) fn execute_from_token_ids( + &self, + input_ids: &MlxArray, + caches: &mut [Qwen3NextCache], + ) -> Result { + let hidden = self + .embed_tokens + .as_ref() + .ok_or_else(|| { + "stage does not host embeddings; hidden-state input required".to_string() + })? + .forward(input_ids); + self.execute_hidden(hidden, caches) + } + + pub(crate) fn execute_from_hidden_states( + &self, + hidden_states: UniquePtr, + caches: &mut [Qwen3NextCache], + ) -> Result { + if self.filter.has_embedding { + return Err("entry stage expects token IDs, not hidden states".to_string()); + } + self.execute_hidden(hidden_states, caches) + } + + fn execute_hidden( + &self, + mut hidden: UniquePtr, + caches: &mut [Qwen3NextCache], + ) -> Result { + if caches.len() != self.layers.len() { + return Err(format!( + "stage cache count mismatch: expected {}, got {}", + self.layers.len(), + caches.len() + )); + } + + let shape = mlxcel_core::array_shape(hidden.as_ref().unwrap()); + let seq_len = shape[1]; + let fa_mask = if seq_len > 1 { + let is_linear: Vec = self.layers.iter().map(|layer| layer.is_linear).collect(); + let offset = stage_attention_offset(&is_linear, caches); + Some(create_causal_mask(seq_len, offset)) + } else { + None + }; + + for (layer, cache) in self.layers.iter().zip(caches.iter_mut()) { + let mask = if layer.is_linear { + None + } else { + fa_mask.as_deref() + }; + hidden = layer.forward(hidden.as_ref().unwrap(), mask, cache); + } + + let hidden = if let Some(norm) = &self.norm { + norm.forward(hidden.as_ref().unwrap()) + } else { + hidden + }; + + if self.filter.has_lm_head { + let logits = if let Some(lm_head) = &self.lm_head { + lm_head.forward(&hidden) + } else { + self.embed_tokens + .as_ref() + .ok_or_else(|| { + "final tied-word-embedding stage missing embeddings".to_string() + })? + .as_linear(&hidden) + }; + Ok(StageExecutionOutput::Logits(logits)) + } else { + Ok(StageExecutionOutput::HiddenStates(hidden)) + } + } +} + +#[cfg(test)] +mod stage_tests { + use super::{Qwen3NextCache, Qwen3NextConfig, stage_attention_offset}; + use crate::models::gated_delta::GatedDeltaCache; + use mlxcel_core::layers::KVCache; + + /// Minimal config that exercises `is_linear_layer` / stage layer typing. + /// Field values are placeholders; only the layer-typing knobs matter. + fn tiny_config(full_attention_interval: usize, num_hidden_layers: usize) -> Qwen3NextConfig { + serde_json::from_value(serde_json::json!({ + "model_type": "qwen3_next", + "hidden_size": 16, + "num_hidden_layers": num_hidden_layers, + "intermediate_size": 32, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 4, + "linear_num_value_heads": 4, + "linear_num_key_heads": 2, + "linear_key_head_dim": 4, + "linear_value_head_dim": 4, + "linear_conv_kernel_dim": 4, + "num_experts": 0, + "num_experts_per_tok": 0, + "decoder_sparse_step": 1, + "moe_intermediate_size": 0, + "shared_expert_intermediate_size": 0, + "full_attention_interval": full_attention_interval, + "vocab_size": 100 + })) + .expect("tiny qwen3-next config") + } + + #[test] + fn stage_layer_typing_uses_global_layer_indices() { + let cfg = tiny_config(4, 8); + // full_attention_interval == 4: layer i is full attention iff + // (i + 1) % 4 == 0, i.e. layer 3 and 7. Everything else is linear. + // A stage covering GLOBAL layers 2..6 must type its LOCAL layers as + // [linear(2), attention(3), linear(4), linear(5)] — not renumbered + // from zero, which would wrongly type local layer 3 as attention. + let kinds: Vec = (2..6).map(|i| cfg.is_linear_layer(i)).collect(); + assert_eq!(kinds, vec![true, false, true, true]); + // First stage 0..4 -> [linear, linear, linear, attention]. + let head: Vec = (0..4).map(|i| cfg.is_linear_layer(i)).collect(); + assert_eq!(head, vec![true, true, true, false]); + } + + #[test] + fn stage_attention_offset_uses_first_stage_local_attention_layer() { + // Local layers: [linear, linear, linear, attention]. + let is_linear = [true, true, true, false]; + let mut kv = KVCache::new(); + kv.offset = 9; + let caches = vec![ + Qwen3NextCache::Linear(GatedDeltaCache::new()), + Qwen3NextCache::Linear(GatedDeltaCache::new()), + Qwen3NextCache::Linear(GatedDeltaCache::new()), + Qwen3NextCache::Attention(kv), + ]; + assert_eq!(stage_attention_offset(&is_linear, &caches), 9); + } + + #[test] + fn stage_attention_offset_defaults_to_zero_for_all_linear_stage() { + let is_linear = [true, true, true]; + let caches = vec![ + Qwen3NextCache::Linear(GatedDeltaCache::new()), + Qwen3NextCache::Linear(GatedDeltaCache::new()), + Qwen3NextCache::Linear(GatedDeltaCache::new()), + ]; + assert_eq!(stage_attention_offset(&is_linear, &caches), 0); + } +} + #[cfg(test)] mod cache_tests { use super::Qwen3NextCache;