diff --git a/python/freetoken/layers/linear.py b/python/freetoken/layers/linear.py index f707e0429..bf1b29b3b 100644 --- a/python/freetoken/layers/linear.py +++ b/python/freetoken/layers/linear.py @@ -59,10 +59,14 @@ def __init__( input_size: int, output_sizes: List[int], has_bias: bool, + local_output_sizes: List[int] | None = None, ): # check that all output sizes are divisible by tp_size tp_info = get_tp_info() - tp_output_sizes = [div_even(size, tp_info.size) for size in output_sizes] + if local_output_sizes is not None: + tp_output_sizes = local_output_sizes + else: + tp_output_sizes = [div_even(size, tp_info.size) for size in output_sizes] output_size = sum(output_sizes) tp_output_size = sum(tp_output_sizes) super().__init__(input_size, output_size, input_size, tp_output_size, has_bias) diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index d68d8ded5..271396cc8 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -219,6 +219,11 @@ def __init__( self.layer_id = layer_id self.offload_cache: OffloadMoeCache | None = None + def _maybe_all_reduce(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Offload MoE: each rank computes the full expert output (banks are not TP-sharded), + so no all-reduce is needed — the result is already the complete output.""" + return hidden_states + def forward( self, hidden_states: torch.Tensor, diff --git a/python/freetoken/models/quant_linear.py b/python/freetoken/models/quant_linear.py index 6ba67a9ee..9d9c3de80 100644 --- a/python/freetoken/models/quant_linear.py +++ b/python/freetoken/models/quant_linear.py @@ -12,8 +12,11 @@ def make_col_merged_quant(expert_quant: str, attn_quant: str, in_f: int, - output_sizes: list[int], has_bias: bool = False): + output_sizes: list[int], has_bias: bool = False, + local_output_sizes: list[int] | None = None): """Column-merged linear for a dense projection: block-fp8 / per-tensor-fp8 / nvfp4 / bf16.""" + if local_output_sizes is not None and (expert_quant != "none" or attn_quant != "none"): + raise NotImplementedError("local_output_sizes (GQA KV replication) not yet supported for quantized checkpoints") if expert_quant == "fp8_block": from freetoken.kernel.triton.fp8_block_linear import Fp8BlockColMerged @@ -28,7 +31,8 @@ def make_col_merged_quant(expert_quant: str, attn_quant: str, in_f: int, return Nvfp4DenseColMerged(in_f, output_sizes, has_bias) from freetoken.layers import LinearColParallelMerged - return LinearColParallelMerged(in_f, output_sizes, has_bias=has_bias) + return LinearColParallelMerged(in_f, output_sizes, has_bias=has_bias, + local_output_sizes=local_output_sizes) def make_replicated_quant(expert_quant: str, attn_quant: str, in_f: int, out_f: int, @@ -51,6 +55,21 @@ def make_replicated_quant(expert_quant: str, attn_quant: str, in_f: int, out_f: return LinearReplicated(in_f, out_f, has_bias=has_bias) +def make_row_parallel_quant(expert_quant: str, attn_quant: str, in_f: int, out_f: int, + has_bias: bool = False): + """Row-parallel linear for a dense projection: block-fp8 / per-tensor-fp8 / nvfp4 / bf16. + Shards the input dimension by tp_size and all-reduces on forward.""" + if expert_quant == "fp8_block": + raise NotImplementedError("row-parallel Fp8BlockLinear not yet implemented") + if attn_quant == "fp8_pertensor": + raise NotImplementedError("row-parallel Fp8PerTensorLinear not yet implemented") + if attn_quant == "nvfp4": + raise NotImplementedError("row-parallel Nvfp4DenseLinear not yet implemented") + from freetoken.layers import LinearOProj + + return LinearOProj(in_f, out_f, has_bias=has_bias) + + def make_replicated(config, in_f: int, out_f: int, has_bias: bool = False): """Config-driven replicated linear: ``Fp8BlockLinear`` under block-fp8, ``Fp8PerTensorLinear`` under per-tensor-fp8 attention, ``Nvfp4DenseLinear`` under nvfp4, else ``LinearReplicated``.""" @@ -60,19 +79,30 @@ def make_replicated(config, in_f: int, out_f: int, has_bias: bool = False): ) -def make_col_merged(config, in_f: int, output_sizes: list[int], has_bias: bool = False): +def make_col_merged(config, in_f: int, output_sizes: list[int], has_bias: bool = False, + local_output_sizes: list[int] | None = None): """Config-driven column-merged linear: ``Fp8BlockColMerged`` under block-fp8, ``Fp8PerTensorColMerged`` under per-tensor-fp8 attention, ``Nvfp4DenseColMerged`` under nvfp4, else ``LinearColParallelMerged``.""" return make_col_merged_quant( getattr(config, "expert_quant", "none"), getattr(config, "attn_quant", "none"), - in_f, output_sizes, has_bias, + in_f, output_sizes, has_bias, local_output_sizes=local_output_sizes, + ) + + +def make_row_parallel(config, in_f: int, out_f: int, has_bias: bool = False): + """Config-driven row-parallel linear: ``LinearOProj`` (bf16); quant variants TBD.""" + return make_row_parallel_quant( + getattr(config, "expert_quant", "none"), getattr(config, "attn_quant", "none"), + in_f, out_f, has_bias, ) __all__ = [ "make_col_merged_quant", "make_replicated_quant", + "make_row_parallel_quant", "make_replicated", "make_col_merged", + "make_row_parallel", ] diff --git a/python/freetoken/models/qwen3_5_moe/attention.py b/python/freetoken/models/qwen3_5_moe/attention.py index 2421264e9..3f953297d 100644 --- a/python/freetoken/models/qwen3_5_moe/attention.py +++ b/python/freetoken/models/qwen3_5_moe/attention.py @@ -4,11 +4,12 @@ import torch from freetoken.core import get_global_ctx +from freetoken.distributed import get_tp_info from freetoken.layers import BaseOP, GemmaRMSNorm from freetoken.layers.rotary import get_rope -from freetoken.utils import nvtx_annotate +from freetoken.utils import div_even, nvtx_annotate -from .quant_linear import make_col_merged, make_replicated +from .quant_linear import make_col_merged, make_row_parallel if TYPE_CHECKING: from freetoken.models.config import ModelConfig @@ -22,9 +23,6 @@ class Qwen3_5Attention(BaseOP): q, k = rope(q, k) # first rotary_dim dims attn = paged_attention(q, k, v) out = o_proj(attn * sigmoid(gate)) - - TP note: uses replicated linears (tp=1 correctness milestone); swap to - column/row-parallel for tensor parallelism later. """ def __init__(self, config: ModelConfig, layer_id: int): @@ -36,13 +34,21 @@ def __init__(self, config: ModelConfig, layer_id: int): self.qo_attn_dim = self.num_q * head_dim self.kv_attn_dim = self.num_kv * head_dim + tp = get_tp_info() # Fused q/k/v projection (one GEMM instead of three); q half is 2x for the # output gate. Split sizes: [num_q*head_dim*2, num_kv*head_dim, num_kv*head_dim]. - self._qkv_split = [self.num_q * head_dim * 2, self.kv_attn_dim, self.kv_attn_dim] + # _qkv_split is TP-local (for torch.split in forward); make_col_merged gets full sizes. # Block-fp8 (Fp8BlockColMerged) when the checkpoint is quantized, else bf16 # LinearColParallelMerged. q/k/v out dims are all /128, so the merged fp8 weight + # weight_scale_inv concatenate cleanly along the output dim. - self.qkv_proj = make_col_merged(config, config.hidden_size, self._qkv_split, has_bias=False) + full_split = [self.num_q * head_dim * 2, self.kv_attn_dim, self.kv_attn_dim] + self._local_num_q = div_even(self.num_q, tp.size) + self._local_num_kv = div_even(self.num_kv, tp.size, allow_replicate=True) + self._local_qo_attn_dim = self._local_num_q * head_dim + self._local_kv_attn_dim = self._local_num_kv * head_dim + self._qkv_split = [self._local_num_q * head_dim * 2, self._local_kv_attn_dim, self._local_kv_attn_dim] + local_qkv_sizes = [self._local_num_q * head_dim * 2, self._local_kv_attn_dim, self._local_kv_attn_dim] + self.qkv_proj = make_col_merged(config, config.hidden_size, full_split, has_bias=False, local_output_sizes=local_qkv_sizes) # Qwen3.5 uses Gemma-style (1+weight) RMSNorm; the weight loader bakes the +1 # into the stored weight (GemmaRMSNorm scales by the raw weight). self.q_norm = GemmaRMSNorm(head_dim, eps=config.rms_norm_eps) @@ -58,7 +64,7 @@ def __init__(self, config: ModelConfig, layer_id: int): else None ), ) - self.o_proj = make_replicated(config, self.qo_attn_dim, config.hidden_size, has_bias=False) + self.o_proj = make_row_parallel(config, self.qo_attn_dim, config.hidden_size, has_bias=False) def _project(self, x: torch.Tensor): """Returns (q, k, v, gate): q [N, num_q, head_dim] post qk-norm+rope, @@ -66,18 +72,18 @@ def _project(self, x: torch.Tensor): positions = get_global_ctx().batch.positions qkv = self.qkv_proj.forward(x) qg, k, v = torch.split(qkv, self._qkv_split, dim=-1) - qg = qg.view(-1, self.num_q, self.head_dim * 2) + qg = qg.view(-1, self._local_num_q, self.head_dim * 2) q = qg[..., : self.head_dim].contiguous() # [N, num_q, head_dim] - gate = qg[..., self.head_dim :].reshape(-1, self.qo_attn_dim) - k = k.view(-1, self.num_kv, self.head_dim).contiguous() + gate = qg[..., self.head_dim :].reshape(-1, self._local_qo_attn_dim) + k = k.reshape(-1, self._local_num_kv, self.head_dim).contiguous() v = v.contiguous() # split view has the qkv row stride; the KV store needs contiguous - q = self.q_norm.forward(q).reshape(-1, self.qo_attn_dim) - k = self.k_norm.forward(k).reshape(-1, self.kv_attn_dim) + q = self.q_norm.forward(q).reshape(-1, self._local_qo_attn_dim) + k = self.k_norm.forward(k).reshape(-1, self._local_kv_attn_dim) q, k = self.rotary.forward(positions, q, k) - return q.view(-1, self.num_q, self.head_dim), k, v, gate + return q.view(-1, self._local_num_q, self.head_dim), k, v, gate def _combine(self, attn_out: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - gated = attn_out.reshape(-1, self.qo_attn_dim) * torch.sigmoid(gate) + gated = attn_out.reshape(-1, self._local_qo_attn_dim) * torch.sigmoid(gate) return self.o_proj.forward(gated) @nvtx_annotate("MHA") diff --git a/python/freetoken/models/qwen3_5_moe/gdn.py b/python/freetoken/models/qwen3_5_moe/gdn.py index 2e7320051..4b02a1366 100644 --- a/python/freetoken/models/qwen3_5_moe/gdn.py +++ b/python/freetoken/models/qwen3_5_moe/gdn.py @@ -3,14 +3,16 @@ import torch import torch.nn.functional as F from freetoken.core import get_global_ctx +from freetoken.distributed import get_tp_info from freetoken.kernel.causal_conv1d import causal_conv1d_decode, causal_conv1d_varlen from freetoken.layers import BaseOP, LinearColParallelMerged +from freetoken.utils import div_even from freetoken.kernel.triton.fp8_block_linear import Fp8BlockColMerged from freetoken.kernel.triton.fp8_pertensor_linear import Fp8PerTensorColMerged from .gdn_kernels import gdn_decode_fla, gdn_prefill_chunk_fla -from .quant_linear import make_replicated_quant +from .quant_linear import make_row_parallel_quant class _DepthwiseConv1d(BaseOP): @@ -56,6 +58,7 @@ def __init__( attn_quant: str = "none", ): self.layer_id = layer_id + tp = get_tp_info() # The fla chunk/decode kernels read+write the recurrent state and the per-chunk h as # [V, K] while the LinearStatePool declares it [K, V]; these coincide (and the # hybrid-radix snapshot scatter h[h_row]->slot is a plain copy) only when the two head @@ -71,6 +74,12 @@ def __init__( self.value_dim = num_v_heads * head_v_dim self.conv_dim = 2 * self.key_dim + self.value_dim self.conv_kernel_size = conv_kernel_size + # TP-local head counts for forward (reshape/split); LinearColParallelMerged gets full sizes + self._local_num_k_heads = div_even(num_k_heads, tp.size) + self._local_num_v_heads = div_even(num_v_heads, tp.size, allow_replicate=True) + self._local_key_dim = self._local_num_k_heads * head_k_dim + self._local_value_dim = self._local_num_v_heads * head_v_dim + self._local_conv_dim = 2 * self._local_key_dim + self._local_value_dim # qkv|z carry a weight scale (block-fp8 weight_scale_inv, or per-tensor FP8 # weight_scale); b|a stay bf16. Both quant modes therefore split the four-way # fusion into an fp8 qkvz GEMM + a bf16 ba GEMM (matches sglang/vLLM). @@ -78,7 +87,8 @@ def __init__( self._pertensor_fp8 = attn_quant == "fp8_pertensor" self._fp8 = self._block_fp8 or self._pertensor_fp8 - self._in_proj_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads] + full_split = [self.conv_dim, self.value_dim, num_v_heads, num_v_heads] + self._in_proj_split = [div_even(s, tp.size) for s in full_split] # TP-local for torch.split if self._fp8: ColMerged = Fp8BlockColMerged if self._block_fp8 else Fp8PerTensorColMerged self.in_proj_qkvz = ColMerged( @@ -89,19 +99,19 @@ def __init__( ) else: # Fused input projection (one GEMM instead of four): qkv | z | b | a. - self.in_proj = LinearColParallelMerged(hidden_size, self._in_proj_split, has_bias=False) - self.conv1d = _DepthwiseConv1d(self.conv_dim, conv_kernel_size) + self.in_proj = LinearColParallelMerged(hidden_size, full_split, has_bias=False) + self.conv1d = _DepthwiseConv1d(self._local_conv_dim, conv_kernel_size) # Recurrence-gating params kept in fp32 (exp/softplus is precision-sensitive, # and the fla kernel reads them as fp32) -- matches HF/sglang, and avoids a # per-call .float() upcast in the decode wrapper. The weight loader exempts # *.A_log / *.dt_bias from the model-dtype downcast. - self.dt_bias = torch.empty(num_v_heads, dtype=torch.float32) - self.A_log = torch.empty(num_v_heads, dtype=torch.float32) + self.dt_bias = torch.empty(self._local_num_v_heads, dtype=torch.float32) + self.A_log = torch.empty(self._local_num_v_heads, dtype=torch.float32) self.norm = _GatedRMSNorm(head_v_dim, eps=rms_norm_eps) # out_proj follows the checkpoint quant: block-fp8 / per-tensor-fp8 / compressed-tensors # NVFP4 (W4A16) / bf16. in_proj_* stay bf16 in every mode (above), so a compressed-tensors # NVFP4 checkpoint (attn_quant=="nvfp4") only makes out_proj native FP4. - self.out_proj = make_replicated_quant( + self.out_proj = make_row_parallel_quant( expert_quant, attn_quant, self.value_dim, hidden_size, has_bias=False ) @@ -163,13 +173,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if self._fp8: qkvz = self.in_proj_qkvz.forward(hidden_states) - conv_in, z = torch.split(qkvz, [self.conv_dim, self.value_dim], dim=-1) + conv_in, z = torch.split(qkvz, [self._local_conv_dim, self._local_value_dim], dim=-1) ba = self.in_proj_ba.forward(hidden_states) - b, a = torch.split(ba, [self.num_v_heads, self.num_v_heads], dim=-1) + b, a = torch.split(ba, [self._local_num_v_heads, self._local_num_v_heads], dim=-1) else: proj = self.in_proj.forward(hidden_states) conv_in, z, b, a = torch.split(proj, self._in_proj_split, dim=-1) - z = z.reshape(total, self.num_v_heads, self.head_v_dim) + z = z.reshape(total, self._local_num_v_heads, self.head_v_dim) li = pool.local_index(self.layer_id) if batch.is_decode: @@ -178,10 +188,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: # no clone, no external l2norm). q/k stay at num_k_heads (kernel handles GQA). mixed = self._conv_decode(conv_in, fla.cache_indices, pool) # [B, conv_dim] B = mixed.shape[0] - qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1) - q = qf.reshape(1, B, self.num_k_heads, self.head_k_dim).to(dtype) - k = kf.reshape(1, B, self.num_k_heads, self.head_k_dim).to(dtype) - v = vf.reshape(1, B, self.num_v_heads, self.head_v_dim).to(dtype) + qf, kf, vf = torch.split(mixed, [self._local_key_dim, self._local_key_dim, self._local_value_dim], dim=-1) + q = qf.reshape(1, B, self._local_num_k_heads, self.head_k_dim).to(dtype) + k = kf.reshape(1, B, self._local_num_k_heads, self.head_k_dim).to(dtype) + v = vf.reshape(1, B, self._local_num_v_heads, self.head_v_dim).to(dtype) core_out = gdn_decode_fla( q, k, v, a, b, A_log=self.A_log, dt_bias=self.dt_bias, state_source=pool.recurrent_states[li], indices=fla.cache_indices, @@ -191,13 +201,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: mixed = self._conv_prefill( conv_in, pool, fla.cu_seqlens, fla.cache_indices, fla.has_initial_state) # fla chunk handles GQA in-kernel: q/k stay at num_k_heads, v at num_v_heads. - qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1) - q = qf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype) - k = kf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype) - v = vf.reshape(1, total, self.num_v_heads, self.head_v_dim).to(dtype) + qf, kf, vf = torch.split(mixed, [self._local_key_dim, self._local_key_dim, self._local_value_dim], dim=-1) + q = qf.reshape(1, total, self._local_num_k_heads, self.head_k_dim).to(dtype) + k = kf.reshape(1, total, self._local_num_k_heads, self.head_k_dim).to(dtype) + v = vf.reshape(1, total, self._local_num_v_heads, self.head_v_dim).to(dtype) g, beta = self._gate_params(a, b) - g = g.reshape(1, total, self.num_v_heads) - beta = beta.float().reshape(1, total, self.num_v_heads) + g = g.reshape(1, total, self._local_num_v_heads) + beta = beta.float().reshape(1, total, self._local_num_v_heads) # The chunk kernel reads + writes back initial_state[cache_indices] in place; # fresh sequences (cached_len==0) must start from a zeroed slot. if fla.fresh_state_indices is not None: diff --git a/python/freetoken/models/qwen3_5_moe/quant_linear.py b/python/freetoken/models/qwen3_5_moe/quant_linear.py index cd0443b67..c135d7c2a 100644 --- a/python/freetoken/models/qwen3_5_moe/quant_linear.py +++ b/python/freetoken/models/qwen3_5_moe/quant_linear.py @@ -9,11 +9,15 @@ make_col_merged_quant, make_replicated, make_replicated_quant, + make_row_parallel, + make_row_parallel_quant, ) __all__ = [ "make_col_merged_quant", "make_replicated_quant", + "make_row_parallel_quant", "make_replicated", "make_col_merged", + "make_row_parallel", ] diff --git a/python/freetoken/models/qwen3_5_moe/weight.py b/python/freetoken/models/qwen3_5_moe/weight.py index b34140890..2711e97d0 100644 --- a/python/freetoken/models/qwen3_5_moe/weight.py +++ b/python/freetoken/models/qwen3_5_moe/weight.py @@ -9,6 +9,7 @@ import torch from freetoken.distributed import get_tp_info from freetoken.kernel.triton.nvfp4_dequant import dequant_nvfp4 +from freetoken.utils import div_even from freetoken.models.loader import ( CT_SCALE_SUFFIXES, ShardReader, @@ -81,6 +82,77 @@ } +def _shard_tp(tensor: torch.Tensor, *, rank: int, world_size: int, dim: int) -> torch.Tensor: + """Simple chunk sharding along ``dim``.""" + if world_size == 1: + return tensor + return tensor.chunk(world_size, dim=dim)[rank].clone() + + +def _shard_tp_parts( + tensor: torch.Tensor, + part_sizes: tuple[int, ...], + *, + rank: int, + world_size: int, + local_part_sizes: tuple[int, ...] | None = None, +) -> torch.Tensor: + """Per-part column-parallel sharding (dim=0). Each part is sharded independently, + matching ``LinearColParallelMerged``'s per-output-size sharding. + + When ``local_part_sizes`` is provided, parts where ``local < full`` but + ``local * world != full`` are replicated by head index (``rank % num_heads``), + for GQA KV head replication when ``num_kv < tp_size``.""" + if world_size == 1: + return tensor + shards: list[torch.Tensor] = [] + offset = 0 + for i, full_size in enumerate(part_sizes): + chunk = tensor[offset:offset + full_size] + local_size = full_size // world_size if local_part_sizes is None else local_part_sizes[i] + if local_size >= full_size: + shards.append(chunk.clone()) + elif local_size * world_size == full_size: + shards.append(chunk.chunk(world_size, dim=0)[rank].clone()) + else: + num_heads = full_size // local_size + head_idx = rank * num_heads // world_size + shards.append(chunk[head_idx * local_size:(head_idx + 1) * local_size].clone()) + offset += full_size + return torch.cat(shards, dim=0) + + +def _maybe_shard(name: str, tensor: torch.Tensor) -> torch.Tensor: + """Apply TP sharding to a non-fused weight based on its state-dict key suffix. + Fused projections (qkv_proj, in_proj, gate_up_proj) are handled by ``_try_fuse`` + → ``_shard_tp_parts`` in ``iter_weights`` before reaching here. + conv1d.weight is also handled in ``iter_weights`` (needs config for sub-part sizes).""" + tp = get_tp_info() + if tp.size == 1: + return tensor + if name.endswith((".o_proj.weight", ".down_proj.weight", ".out_proj.weight")): + return _shard_tp(tensor, rank=tp.rank, world_size=tp.size, dim=1) + if name.endswith(("embed_tokens.weight", "lm_head.weight")): + return _shard_tp(tensor, rank=tp.rank, world_size=tp.size, dim=0) + if name.endswith(("A_log", "dt_bias")): + return _shard_tp(tensor, rank=tp.rank, world_size=tp.size, dim=0) + # Routed expert weights (3D stacked: [num_experts, ...]) — resident mode only. + # gate_up_proj is [num_experts, 2*intermediate, hidden] = [gate, up] fused on dim=1. + # Must shard gate and up independently (like _shard_tp_parts on dim=1). + # down_proj is [num_experts, hidden, intermediate] — single dim, simple chunk on dim=2. + if name.endswith("experts.gate_up_proj"): + intermediate = tensor.shape[1] // 2 + gate = tensor[:, :intermediate, :] + up = tensor[:, intermediate:, :] + return torch.cat([ + gate.chunk(tp.size, dim=1)[tp.rank].clone(), + up.chunk(tp.size, dim=1)[tp.rank].clone(), + ], dim=1) + if name.endswith("experts.down_proj"): + return _shard_tp(tensor, rank=tp.rank, world_size=tp.size, dim=2) + return tensor + + def _dequant_fp8_weight(weight: torch.Tensor, weight_scale: torch.Tensor) -> torch.Tensor: """Weight-only FP8 -> bf16 (per-tensor static scale). Activations stay bf16 (W8A16), which is at least as precise as the checkpoint's intended W8A8.""" @@ -154,8 +226,8 @@ def _is_gemma_norm(name: str) -> bool: def _try_fuse( name: str, tensor: torch.Tensor, buf: dict[str, dict[int, torch.Tensor]] -) -> tuple[str, torch.Tensor] | tuple[()] | None: - """buffer a fusion part; return merged ``(name, tensor)`` once all parts arrive, +) -> tuple[str, torch.Tensor, tuple[int, ...]] | tuple[()] | None: + """buffer a fusion part; return merged ``(name, tensor, part_sizes)`` once all parts arrive, ``()`` while incomplete, ``None`` if not a fusion part.""" for fused_suffix, parts in _FUSIONS.items(): for idx, part in enumerate(parts): @@ -165,7 +237,18 @@ def _try_fuse( slots[idx] = tensor if len(slots) == len(parts): del buf[key] - return key, torch.cat([slots[i] for i in range(len(parts))], dim=0) + tensors = [slots[i] for i in range(len(parts))] + part_sizes = tuple(t.shape[0] for t in tensors) + # in_proj_qkv is itself a fusion of [q, k, v] with sizes + # (key_dim, key_dim, value_dim). value_dim == z_size, and + # key_dim = (qkv_size - value_dim) // 2. Expand qkv into + # its sub-parts so _shard_tp_parts shards each independently. + if key.endswith(".in_proj.weight") and len(part_sizes) == 4: + qkv_size, z_size, b_size, a_size = part_sizes + value_dim = z_size + key_dim = (qkv_size - value_dim) // 2 + part_sizes = (key_dim, key_dim, value_dim, z_size, b_size, a_size) + return key, torch.cat(tensors, dim=0), part_sizes return () return None @@ -210,8 +293,6 @@ def iter_weights( ) return tp_info = get_tp_info() - if tp_info.size > 1: - raise NotImplementedError("qwen3_5_moe weight loading currently supports TP=1 only") # Pure-NVFP4 checkpoint (bf16 attn): the dense MLP projections (shared_expert) are still # stored as packed FP4 -- keep them native (W4A16) when dense_quant=="nvfp4" rather than @@ -259,7 +340,8 @@ def iter_weights( shared_buf=nvfp4_shared_buf, ) if emit is not _NOT_DENSE_NVFP4: - yield from emit + for ename, etensor in emit: + yield ename, _maybe_shard(ename, etensor) continue tensor = _load_maybe_quantized(f, raw_name, keyset) @@ -272,20 +354,43 @@ def iter_weights( if "gate" in slots and "up" in slots: merged = torch.cat([slots["gate"], slots["up"]], dim=0) del shared_buf[prefix] - yield f"{prefix}.mlp.shared_expert.gate_up_proj.weight", merged + merged_name = f"{prefix}.mlp.shared_expert.gate_up_proj.weight" + gate_size = slots["gate"].shape[0] + up_size = slots["up"].shape[0] + yield merged_name, _shard_tp_parts(merged, (gate_size, up_size), rank=tp_info.rank, world_size=tp_info.size) continue # fuse q/k/v -> qkv_proj and GDN in_proj_{qkv,z,b,a} -> in_proj fused = _try_fuse(name, tensor, fuse_buf) if fused is not None: if fused != (): # () means buffered, not yet complete - yield fused + fused_name, fused_tensor, part_sizes = fused + # For qkv_proj: compute local part sizes with KV head replication + # when num_kv_heads < tp_size (GQA replication). + local_part_sizes = None + if fused_name.endswith(".qkv_proj.weight"): + num_kv = config.num_kv_heads + head_dim = config.head_dim + local_kv = div_even(num_kv, tp_info.size, allow_replicate=True) * head_dim + local_q = div_even(part_sizes[0], tp_info.size) + local_part_sizes = (local_q, local_kv, local_kv) + yield fused_name, _shard_tp_parts( + fused_tensor, part_sizes, rank=tp_info.rank, world_size=tp_info.size, + local_part_sizes=local_part_sizes) continue if _is_gemma_norm(name): tensor = tensor + 1.0 # (1 + weight) baked into the stored weight - yield name, tensor + if name.endswith("conv1d.weight") and tp_info.size > 1: + g = config.linear_attention_group() + if g is not None: + key_dim = g.num_key_heads * g.key_head_dim + value_dim = g.num_value_heads * g.value_head_dim + tensor = _shard_tp_parts(tensor, (key_dim, key_dim, value_dim), + rank=tp_info.rank, world_size=tp_info.size) + + yield name, _maybe_shard(name, tensor) assert not shared_buf, f"Incomplete shared-expert merges: {list(shared_buf.keys())}" assert not nvfp4_shared_buf, f"Incomplete NVFP4 shared-expert merges: {list(nvfp4_shared_buf.keys())}" diff --git a/tests/models/test_qwen3_5_tp.py b/tests/models/test_qwen3_5_tp.py new file mode 100644 index 000000000..fdfbd48be --- /dev/null +++ b/tests/models/test_qwen3_5_tp.py @@ -0,0 +1,32 @@ +import torch +from freetoken.models.qwen3_5_moe.weight import _shard_tp, _shard_tp_parts + + +def test_shard_tp(): + t = torch.arange(32).reshape(8, 4) + s0, s1 = _shard_tp(t, rank=0, world_size=2, dim=0), _shard_tp(t, rank=1, world_size=2, dim=0) + assert s0.shape == (4, 4) and s1.shape == (4, 4) + torch.testing.assert_close(torch.cat([s0, s1]), t) + assert _shard_tp(t, rank=0, world_size=1, dim=0).equal(t) + assert _shard_tp(torch.arange(32).reshape(4, 8), rank=0, world_size=2, dim=1).shape == (4, 4) + + +def test_shard_tp_parts(): + t = torch.arange(48).reshape(12, 4) + s0 = _shard_tp_parts(t, (4, 4, 4), rank=0, world_size=2) + s1 = _shard_tp_parts(t, (4, 4, 4), rank=1, world_size=2) + assert s0.shape == (6, 4) + for i in range(3): + torch.testing.assert_close(torch.cat([s0[i*2:i*2+2], s1[i*2:i*2+2]]), t[i*4:i*4+4]) + + +def test_shard_tp_parts_replicate(): + t = torch.arange(16).reshape(8, 2) + local = (2, 4) + kw = dict(tensor=t, part_sizes=(4, 4), world_size=4, local_part_sizes=local) + s0, s1, s2, s3 = [_shard_tp_parts(rank=r, **kw) for r in range(4)] + assert s0.shape == (6, 2) + torch.testing.assert_close(s0[:2], s1[:2]) # ranks 0,1 share head 0 + torch.testing.assert_close(s2[:2], s3[:2]) # ranks 2,3 share head 1 + assert not torch.equal(s0[:2], s2[:2]) # different heads + torch.testing.assert_close(s0[2:], s2[2:]) # replicated part identical