Daiz
28817d3ee3
## Motivation Qwen3.5 MoE models (e.g., `Qwen3.5-397B-A17B-6bit`) are now supported by `mlx-lm` via `qwen3_5_moe` model type, but exo lacks tensor parallel sharding support for this architecture. This prevents running large Qwen3.5 models across multiple nodes. Qwen3.5 uses a GatedDeltaNet hybrid attention mechanism similar to Qwen3-Next, but with a different projection layout — separate `in_proj_qkv`, `in_proj_z`, `in_proj_b`, `in_proj_a` instead of Qwen3-Next's combined `in_proj_qkvz` and `in_proj_ba`. This requires architecture-aware sharding logic. ## Changes (evan summary) - enable qwen3_5 dense + moe tensor parallelism from config - defensively skip evalling _cache.keys if it doesn't exist - ignore kwargs in qwen35 pipeline masking and ensure pipeline segments match global model parameters for mask creation - add sharding for qwen3_5 moe linear attention - added another 6 million model cards ## Why It Works Qwen3.5's GatedDeltaNet has an `in_proj_qkv` linear layer with three concatenated sections: `[q(key_dim), k(key_dim), v(value_dim)]`. A naive contiguous split (`segments=1`) would slice across section boundaries, corrupting q/k/v values and producing garbled output. By passing `segments=[key_dim, key_dim + key_dim]` to `shard_linear()`, each section is split independently before distributing across devices. This ensures every rank receives correctly aligned q, k, and v components. The remaining separate projections (`in_proj_z`, `in_proj_b`, `in_proj_a`) and the MoE layers follow the same `all_to_sharded` / `sharded_to_all` pattern already used for Qwen3-Next. Some pipeline splits didn't include an ssm layer or a linear layer resulting in a subset of the model acting like it shouldn't create the appropriate masks for the next layer - we patch the model to manually create such masks. ## Test Plan tensor sharded 2,3,4 models & pipeline sharded 2,3,4 with simple eval. --------- Co-authored-by: hw <hw@hwStudio1.local> Co-authored-by: Ryuichi Leo Takashige <leo@exolabs.net> Co-authored-by: Evan <evanev7@gmail.com>
154 lines
4.0 KiB
Python
154 lines
4.0 KiB
Python
from dataclasses import dataclass
|
|
from typing import Any, Optional
|
|
|
|
import mlx.core as mx
|
|
import mlx.nn as nn
|
|
|
|
from .cache import ArraysCache, KVCache
|
|
from .qwen3_next import (
|
|
Qwen3NextAttention as Attention,
|
|
Qwen3NextMLP as MLP,
|
|
Qwen3NextRMSNormGated as RMSNormGated,
|
|
Qwen3NextSparseMoeBlock,
|
|
)
|
|
|
|
SparseMoeBlock = Qwen3NextSparseMoeBlock
|
|
from .switch_layers import SwitchGLU
|
|
|
|
@dataclass
|
|
class TextModelArgs:
|
|
model_type: str
|
|
hidden_size: int
|
|
intermediate_size: int
|
|
num_hidden_layers: int
|
|
num_attention_heads: int
|
|
rms_norm_eps: float
|
|
vocab_size: int
|
|
num_key_value_heads: int
|
|
max_position_embeddings: int
|
|
linear_num_value_heads: int
|
|
linear_num_key_heads: int
|
|
linear_key_head_dim: int
|
|
linear_value_head_dim: int
|
|
linear_conv_kernel_dim: int
|
|
tie_word_embeddings: bool
|
|
attention_bias: bool
|
|
head_dim: Optional[int]
|
|
full_attention_interval: int
|
|
num_experts: int
|
|
num_experts_per_tok: int
|
|
decoder_sparse_step: int
|
|
shared_expert_intermediate_size: int
|
|
moe_intermediate_size: int
|
|
norm_topk_prob: bool
|
|
rope_parameters: Optional[dict[str, Any]]
|
|
partial_rotary_factor: float
|
|
rope_theta: float
|
|
rope_scaling: Optional[dict[str, Any]]
|
|
|
|
@classmethod
|
|
def from_dict(cls, params: dict[str, Any]) -> TextModelArgs: ...
|
|
def __post_init__(self) -> None: ...
|
|
|
|
class GatedDeltaNet(nn.Module):
|
|
hidden_size: int
|
|
num_v_heads: int
|
|
num_k_heads: int
|
|
head_k_dim: int
|
|
head_v_dim: int
|
|
key_dim: int
|
|
value_dim: int
|
|
conv_kernel_size: int
|
|
conv_dim: int
|
|
conv1d: nn.Conv1d
|
|
in_proj_qkv: nn.Linear
|
|
in_proj_z: nn.Linear
|
|
in_proj_b: nn.Linear
|
|
in_proj_a: nn.Linear
|
|
dt_bias: mx.array
|
|
A_log: mx.array
|
|
norm: RMSNormGated
|
|
out_proj: nn.Linear
|
|
|
|
def __init__(self, config: TextModelArgs) -> None: ...
|
|
def __call__(
|
|
self,
|
|
inputs: mx.array,
|
|
mask: Optional[mx.array] = None,
|
|
cache: Optional[Any] = None,
|
|
) -> mx.array: ...
|
|
|
|
class DecoderLayer(nn.Module):
|
|
is_linear: bool
|
|
linear_attn: GatedDeltaNet
|
|
self_attn: Attention
|
|
input_layernorm: nn.RMSNorm
|
|
post_attention_layernorm: nn.RMSNorm
|
|
mlp: MLP | SparseMoeBlock
|
|
|
|
def __init__(self, args: TextModelArgs, layer_idx: int) -> None: ...
|
|
def __call__(
|
|
self,
|
|
x: mx.array,
|
|
mask: Optional[mx.array] = None,
|
|
cache: Optional[Any] = None,
|
|
) -> mx.array: ...
|
|
|
|
class Qwen3_5TextModel(nn.Module):
|
|
embed_tokens: nn.Embedding
|
|
layers: list[DecoderLayer]
|
|
norm: nn.RMSNorm
|
|
ssm_idx: int
|
|
fa_idx: int
|
|
|
|
def __init__(self, args: TextModelArgs) -> None: ...
|
|
def __call__(
|
|
self,
|
|
inputs: mx.array,
|
|
cache: Optional[Any] = None,
|
|
input_embeddings: Optional[mx.array] = None,
|
|
) -> mx.array: ...
|
|
|
|
class TextModel(nn.Module):
|
|
args: TextModelArgs
|
|
model_type: str
|
|
model: Qwen3_5TextModel
|
|
lm_head: nn.Linear
|
|
|
|
def __init__(self, args: TextModelArgs) -> None: ...
|
|
def __call__(
|
|
self,
|
|
inputs: mx.array,
|
|
cache: Optional[Any] = None,
|
|
input_embeddings: Optional[mx.array] = None,
|
|
) -> mx.array: ...
|
|
@property
|
|
def layers(self) -> list[DecoderLayer]: ...
|
|
def make_cache(self) -> list[ArraysCache | KVCache]: ...
|
|
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
|
|
|
|
@dataclass
|
|
class ModelArgs:
|
|
model_type: str
|
|
text_config: dict[str, Any]
|
|
|
|
@classmethod
|
|
def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
|
|
|
|
class Model(nn.Module):
|
|
args: ModelArgs
|
|
model_type: str
|
|
language_model: TextModel
|
|
|
|
def __init__(self, args: ModelArgs) -> None: ...
|
|
def __call__(
|
|
self,
|
|
inputs: mx.array,
|
|
cache: Optional[Any] = None,
|
|
input_embeddings: Optional[mx.array] = None,
|
|
) -> mx.array: ...
|
|
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
|
|
@property
|
|
def layers(self) -> list[DecoderLayer]: ...
|
|
def make_cache(self) -> list[ArraysCache | KVCache]: ...
|