From 67cf58944c5d54559d7196926341228d71731014 Mon Sep 17 00:00:00 2001 From: Carlos Alvarado Date: Sat, 22 Aug 2026 21:10:38 -0400 Subject: [PATCH 1/6] feat(dsv4): add tensor-parallel DeepSeek-V4 runtime Shard MLA heads, vocabulary, shared experts, and routed FP4 expert banks across TP ranks. Give every shard independent storage, size collectives for skewed expert loading, partition CPU workers, and avoid redundant checkpoint reads. This is the foundation that makes DeepSeek-V4-Flash fit and start reliably on 4x RTX A4000. --- docs/cli.md | 3 +- docs/models.md | 21 ++ python/freetoken/engine/config.py | 7 +- python/freetoken/engine/engine.py | 51 ++++- python/freetoken/models/deepseek_v4/args.py | 41 ++++ .../freetoken/models/deepseek_v4/attention.py | 23 +- python/freetoken/models/deepseek_v4/model.py | 56 ++++- python/freetoken/models/deepseek_v4/moe.py | 34 ++- .../freetoken/models/deepseek_v4/parallel.py | 114 ++++++++++ python/freetoken/models/deepseek_v4/weight.py | 207 ++++++++++++------ python/freetoken/moe/cpu_executor.py | 57 ++++- python/freetoken/server/args.py | 12 + tests/dsv4/test_dsv4_tensor_parallel.py | 150 +++++++++++++ 13 files changed, 680 insertions(+), 96 deletions(-) create mode 100644 python/freetoken/models/deepseek_v4/parallel.py create mode 100644 tests/dsv4/test_dsv4_tensor_parallel.py diff --git a/docs/cli.md b/docs/cli.md index ff4af382..34302326 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -41,6 +41,8 @@ parsers all resolve automatically from the checkpoint and the GPU. | `--host` | 127.0.0.1 | Bind address | | `--port` | 1919 | Bind port | | `--gpu` | GPU 0 | GPU to run on: a UUID from `nvidia-smi -L` or an `nvidia-smi` index; see [below](#choosing-a-gpu) | +| `--tensor-parallel-size`, `--tp-size` | 1 | GPUs to shard the model over; one scheduler process per rank, rank `i` on `cuda:i` | +| `--distributed-timeout` | 1800 | Seconds a TP collective waits before failing; sized for the per-rank skew in weight-load time | | `--max-running-requests` | 4 | Max concurrently running requests | | `--max-output-tokens` | 32768 | Default output budget for requests that omit one | | `--max-seq-len-override` | from checkpoint | Max sequence length | @@ -172,4 +174,3 @@ profile that `ft serve --moe-backend auto` and `--moe-hybrid-max-fetch -1` then - What to measure: `--dtype`, `--model`, `--formats`, `--isa`. - `--threshold` (default 2.0) sets the call: recommend hybrid when CPU bandwidth beats PCIe by that factor. - diff --git a/docs/models.md b/docs/models.md index e4850a12..4381a695 100644 --- a/docs/models.md +++ b/docs/models.md @@ -31,6 +31,27 @@ for them; other checkpoints of the same architectures work too. `offload`, upgraded to `hybrid` when a cached `ft bench bw` profile recommends it. +## Tensor parallelism + +`ft serve --tensor-parallel-size N` shards the model over N GPUs on one host. One +scheduler process runs per rank, and rank `i` uses `cuda:i`. + +DeepSeek-V4 shards: + +- MLA query and output heads (`wq_b` column-parallel, `wo_a` by output group, + `wo_b` row-parallel with one all-reduce per block) +- the MoE intermediate dimension, for both the shared expert and the routed FP4 + experts, so the host expert banks divide by N instead of being replicated +- the vocabulary, for the embedding table and the output head + +It keeps replicated: the MLA latent KV path (`wkv` and the paged KV pools — every +head reads the same latent KV), the compressors and the Lightning Indexer (every +rank must select the same blocks), and the router. + +Constraints for DeepSeek-V4: N must divide `o_groups` (8), so N is 1, 2, 4 or 8. +The KV pool is replicated, so its cost per GPU does not fall with N; the weights +and the expert banks do. + ## Notes - `ft checkpoint` conversion is optional — it pre-converts a checkpoint into diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 543012f3..0bb45827 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -72,7 +72,12 @@ class EngineConfig: # ratio default above. A runtime cache rebuild sets this (num_swa_pages) to pin the window # regardless of the full anchor; the ratio is the startup default and the fallback. swa_num_pages_override: int | None = None - distributed_timeout: float = 60.0 + # Collective timeout for the TP process group. The FIRST collective is the post-load + # free-memory all-reduce, and ranks reach it minutes apart on a large MoE checkpoint + # (each reads its own expert shards, at its own speed), so a short timeout kills the + # server during a healthy startup rather than protecting it. Sized for that load skew; + # lower it with --distributed-timeout if a hung rank should fail faster. + distributed_timeout: float = 1800.0 use_dummy_weight: bool = False use_pynccl: bool = True max_seq_len_override: int | None = None diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index cd6505d2..220c2b99 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -246,8 +246,18 @@ def _make_dummy_weight_state_dict( ) -> Dict[str, torch.Tensor]: state_dict: Dict[str, torch.Tensor] = {} fp8_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2) + e8m0 = getattr(torch, "float8_e8m0fnu", None) for key, param in model_state.items(): - if param.dtype in fp8_dtypes: + if e8m0 is not None and param.dtype == e8m0: + # e8m0 is a bare exponent code (value = 2^(code-127)), so 127 is scale 1.0. + # It reports is_floating_point, but there is no normal_ kernel for it -- and a + # RANDOM exponent would scale a block by up to 2^127 anyway, so 1.0 is both the + # only fillable value and the only sane one. Without this a --dummy-weight run + # of any block-scaled model dies with "normal_kernel_cuda not implemented". + t = torch.empty(param.shape, dtype=param.dtype, device=device) + t.view(torch.uint8).fill_(127) + state_dict[key] = t + elif param.dtype in fp8_dtypes: # torch.randn is not implemented for fp8; fill via a uint8 view with small # codes (avoid NaN/inf fp8 encodings). Lets dummy-weight startup work for # block-fp8 models (the dense fp8 linears are fp8 regardless of moe_backend). @@ -290,11 +300,37 @@ class ForwardOutput(NamedTuple): copy_done_event: torch.cuda.Event +def _share_cpu_threads_across_ranks(tp_size: int) -> None: + """Split the machine's CPU threads across the TP ranks, before any tensor work. + + torch sizes its intra-op pool from the WHOLE machine, and every rank does the same, + so N ranks each start a machine-sized pool and oversubscribe every core N times. + Measured on a 40-core / 80-thread host at TP=4: 117 threads per rank, 468 in total, + ~6x the hardware, all of it contending through the weight load -- which is thread- + bound (host-side copies into the expert banks), not disk-bound. The NVMe sat at + 187 MB/s while the CPU ran at 92% user. + + The engine already narrows threads once the pinned CPU MoE pool exists; this covers + everything BEFORE that, the weight load included. Honour an explicit OMP_NUM_THREADS. + """ + if tp_size <= 1 or os.environ.get("OMP_NUM_THREADS"): + return + total = os.cpu_count() or tp_size + per_rank = max(1, total // tp_size) + if per_rank < torch.get_num_threads(): + torch.set_num_threads(per_rank) + logger.info_rank0( + f"torch intra-op threads: {total} -> {per_rank} per rank " + f"({tp_size} ranks share this host)" + ) + + class Engine: def __init__(self, config: EngineConfig): assert not torch.cuda.is_initialized() set_tp_info(rank=config.tp_info.rank, size=config.tp_info.size) _ensure_expandable_segments() # before the first CUDA allocation below + _share_cpu_threads_across_ranks(config.tp_info.size) from freetoken.gpu_select import bind_assigned_gpu @@ -325,6 +361,19 @@ def __init__(self, config: EngineConfig): self.model.load_state_dict(self._load_weight_state_dict(config)) post_weights_free = self._sync_get_memory()[0] self._weights_bytes = self._baseline_free - post_weights_free + # What the parameters declare vs what the load actually cost. Every byte of the gap + # is resident-but-unaccounted (staging buffers, per-layer scratch) and comes straight + # out of the MoE-cache + KV budget, so surface it instead of leaving it to arithmetic + # on the cache-sizing line. + declared = sum(t.numel() * t.element_size() for t in self.model.state_dict().values()) + logger.info_rank0( + f"Weights: {mem_GB(declared)} declared by parameters, " + f"{mem_GB(self._weights_bytes)} measured on device " + f"(torch allocator holds {mem_GB(torch.cuda.memory_allocated(self.device))} live, " + f"{mem_GB(torch.cuda.memory_reserved(self.device))} reserved; " + f"peak {mem_GB(torch.cuda.max_memory_reserved(self.device))}). " + "A gap between 'reserved' and 'measured' is held OUTSIDE the allocator." + ) # Pool-budget baseline for the desktop cache sliders: free VRAM after the weights are # resident but before ANY runtime cache pool (MoE expert cache below, KV pages, GDN # state) is allocated. This is the stable "if all free VRAM went to one pool" budget — diff --git a/python/freetoken/models/deepseek_v4/args.py b/python/freetoken/models/deepseek_v4/args.py index b75d2800..6e5239f3 100644 --- a/python/freetoken/models/deepseek_v4/args.py +++ b/python/freetoken/models/deepseek_v4/args.py @@ -69,15 +69,56 @@ class DeepseekV4Args: hc_sinkhorn_iters: int = 20 hc_eps: float = 1e-6 + # ----- dSpark (semi-autoregressive speculative decoding) ----- + # The checkpoint's ``mtp.{0..n_mtp_layers-1}`` stack. Unlike a one-token-per-step MTP + # head, dSpark proposes a whole block: ``dspark_block_size`` tokens per draft pass, + # scored by a Markov head and gated by a confidence head. Each mtp layer is a FULL + # DSV4 block, so enabling it costs n_mtp_layers x n_routed_experts more expert banks. + dspark_block_size: int = 0 + dspark_markov_rank: int = 0 + dspark_noise_token_id: int = -1 + # Which target layers feed the drafter its hidden state. + dspark_target_layer_ids: Tuple[int, ...] = () + # Runtime opt-in, set by the server from --speculative-dspark. The fields above only + # describe what the CHECKPOINT ships; this says whether to pay for it. Off by default + # because the drafter's own routed experts enlarge the host banks and the GPU cache. + dspark_enabled: bool = False + def __post_init__(self) -> None: # JSON lists -> tuple so the dataclass stays hashable / immutable-ish. if isinstance(self.compress_ratios, list): self.compress_ratios = tuple(self.compress_ratios) + if isinstance(self.dspark_target_layer_ids, list): + self.dspark_target_layer_ids = tuple(self.dspark_target_layer_ids) @property def nope_head_dim(self) -> int: return self.head_dim - self.rope_head_dim + @property + def has_dspark(self) -> bool: + """Does this checkpoint ship a usable dSpark drafter?""" + return ( + self.dspark_block_size > 1 + and self.n_mtp_layers > 0 + and self.dspark_markov_rank > 0 + ) + + @property + def n_draft_layers(self) -> int: + """Drafter layers actually built. Zero unless dSpark is both shipped and enabled.""" + return self.n_mtp_layers if (self.has_dspark and self.dspark_enabled) else 0 + + @property + def n_moe_layers(self) -> int: + """MoE layers the expert banks and the offload cache must cover. + + The drafter's layers are appended AFTER the target's, so a draft layer keeps the + cache-facing id ``n_layers + k``. Every layer-indexed structure (host banks, GPU + slot cache, KV pools) then addresses target and draft layers the same way. + """ + return self.n_layers + self.n_draft_layers + def _config_path(model_path: str) -> str: """Locate the authors' ModelArgs JSON inside the checkpoint directory.""" diff --git a/python/freetoken/models/deepseek_v4/attention.py b/python/freetoken/models/deepseek_v4/attention.py index 357585ee..b266b4e2 100644 --- a/python/freetoken/models/deepseek_v4/attention.py +++ b/python/freetoken/models/deepseek_v4/attention.py @@ -8,6 +8,7 @@ from torch import nn from freetoken.core import get_global_ctx +from freetoken.distributed import DistributedCommunicator from freetoken.kernel.triton.dsv4.fp8_linear import act_quant_fp8_inplace from freetoken.kernel.triton.dsv4.norm import rms_norm @@ -15,6 +16,7 @@ from .compress import Compressor, Indexer from .layers import Linear, RMSNorm, get_compress_topk_idxs, get_window_topk_idxs from .ops import apply_rotary_emb, apply_rotary_emb_decode, get_freqs_cis +from .parallel import div_tp, tp_size class Attention(nn.Module): @@ -30,17 +32,23 @@ def __init__(self, layer_id: int, args: DeepseekV4Args): super().__init__() self.layer_id = layer_id self.dim = args.dim - self.n_heads = args.n_heads self.q_lora_rank = args.q_lora_rank self.o_lora_rank = args.o_lora_rank self.head_dim = args.head_dim self.rope_head_dim = args.rope_head_dim - self.n_groups = args.o_groups self.window_size = args.window_size self.compress_ratio = args.compress_ratios[layer_id] self.eps = args.norm_eps + # Head parallelism: a rank owns whole o_groups, and therefore whole heads. The + # per-group head width (wo_a_k) is a property of one group, so it does NOT shard. + self.tp_size = tp_size() + self.n_heads = div_tp(args.n_heads, "n_heads") + self.n_groups = div_tp(args.o_groups, "o_groups") + self._comm = DistributedCommunicator() if self.tp_size > 1 else None + self.attn_sink = nn.Parameter(torch.empty(self.n_heads, dtype=torch.float32), requires_grad=False) + # wq_a / wkv / kv_norm stay replicated: the MLA latent KV is shared by every head. self.wq_a = Linear(self.dim, self.q_lora_rank, kind="fp8") self.q_norm = RMSNorm(self.q_lora_rank, self.eps) self.wq_b = Linear(self.q_lora_rank, self.n_heads * self.head_dim, kind="fp8") @@ -48,8 +56,9 @@ def __init__(self, layer_id: int, args: DeepseekV4Args): self.kv_norm = RMSNorm(self.head_dim, self.eps) # wo_a: dequantized to bf16 (reference runs a bf16 grouped-output einsum). wo_a_rows = self.n_groups * args.o_lora_rank - wo_a_k = self.n_heads * self.head_dim // self.n_groups + wo_a_k = args.n_heads * self.head_dim // args.o_groups self.wo_a = nn.Parameter(torch.empty(wo_a_rows, wo_a_k, dtype=torch.bfloat16), requires_grad=False) + # Row-parallel: each rank consumes its own groups and contributes a partial sum. self.wo_b = Linear(self.n_groups * args.o_lora_rank, self.dim, kind="fp8") self.softmax_scale = self.head_dim ** -0.5 @@ -99,10 +108,16 @@ def reset(self) -> None: def _wo(self, o: torch.Tensor, bsz: int, seqlen: int) -> torch.Tensor: + # Under TP, ``o`` holds only this rank's heads, so the grouped einsum and wo_b + # produce a PARTIAL sum over the full output dim. One all-reduce completes it -- + # the single attention-side collective per block. o = o.reshape(bsz, seqlen, self.n_groups, -1) wo_a = self.wo_a.view(self.n_groups, self.o_lora_rank, -1) o = torch.einsum("bsgd,grd->bsgr", o, wo_a).flatten(2) - return self.wo_b(o) + o = self.wo_b(o) + if self._comm is not None: + o = self._comm.all_reduce(o) + return o def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: int): """One request's prefill work inside a (possibly ragged) batch: persist its window KV, diff --git a/python/freetoken/models/deepseek_v4/model.py b/python/freetoken/models/deepseek_v4/model.py index c00d20cc..df75e172 100644 --- a/python/freetoken/models/deepseek_v4/model.py +++ b/python/freetoken/models/deepseek_v4/model.py @@ -26,6 +26,7 @@ from torch import nn from freetoken.core import get_global_ctx +from freetoken.distributed import DistributedCommunicator from freetoken.kernel.triton.dsv4.hc import hc_post_combine, hc_pre_combine from freetoken.kernel.triton.dsv4.sinkhorn import hc_split_sinkhorn from freetoken.models.blocks import BaseLLMModel @@ -45,6 +46,7 @@ get_window_topk_idxs, ) from .moe import Expert, Gate # noqa: F401 +from .parallel import div_tp, tp_info, validate_tp class Block(nn.Module): @@ -130,15 +132,30 @@ def decode_step(self, x, pos, rows, cmp_stage_cap, input_ids, wctx=None): class Transformer(nn.Module): def __init__(self, args: DeepseekV4Args): super().__init__() + # Check every tensor-parallel split before building a single layer, so a bad + # --tensor-parallel-size fails with one clear message, not a reshape deep in a + # forward. embed / head / norm / hyper-connections stay replicated. + validate_tp(args) self.args = args self.norm_eps = args.norm_eps self.hc_eps = args.hc_eps self.hc_mult = hc_mult = args.hc_mult - self.embed = nn.Embedding(args.vocab_size, args.dim) + # Vocabulary parallelism. The embedding table and the output head are the two + # largest replicated tensors (3.0 GiB together), so each rank keeps one + # contiguous vocabulary block: the embedding all-reduces its masked lookup, the + # head all-gathers its logit slice. + rank, tp = tp_info() + self.tp_size = tp + self.vocab_size = args.vocab_size + vocab_local = div_tp(args.vocab_size, "vocab_size") + self.vocab_start = rank * vocab_local + self.vocab_local = vocab_local + self._comm = DistributedCommunicator() if tp > 1 else None + self.embed = nn.Embedding(vocab_local, args.dim) self.embed.weight.requires_grad_(False) self.layers = nn.ModuleList([Block(i, args) for i in range(args.n_layers)]) self.norm = RMSNorm(args.dim, self.norm_eps) - self.head = nn.Parameter(torch.empty(args.vocab_size, args.dim, dtype=torch.bfloat16), requires_grad=False) + self.head = nn.Parameter(torch.empty(vocab_local, args.dim, dtype=torch.bfloat16), requires_grad=False) hc_dim = hc_mult * args.dim self.hc_head_fn = nn.Parameter(torch.empty(hc_mult, hc_dim, dtype=torch.float32), requires_grad=False) self.hc_head_base = nn.Parameter(torch.empty(hc_mult, dtype=torch.float32), requires_grad=False) @@ -148,6 +165,33 @@ def bind(self, pool, device: torch.device) -> None: for layer in self.layers: layer.attn.bind(pool, device) + def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: + """Vocabulary-parallel lookup. Rows outside this rank's block contribute zero, so + the all-reduce sums exactly one real row per token. Index arithmetic only -- no + host sync, so a captured decode graph replays it.""" + if self._comm is None: + return self.embed(input_ids) + local = input_ids - self.vocab_start + outside = (local < 0) | (local >= self.vocab_local) + y = self.embed(local.masked_fill(outside, 0)) + y = y.masked_fill(outside.unsqueeze(-1), 0) + return self._comm.all_reduce(y) + + def logits(self, h: torch.Tensor) -> torch.Tensor: + """Head projection over this rank's vocabulary block, then an all-gather back to + the full [rows, vocab]. all_gather concatenates on dim 0, so the gathered tensor + is rank-major: one row block per rank, which is the vocabulary order.""" + local = F.linear(h, self.head) # [rows, vocab_local] + if self._comm is None: + return local + rows = local.shape[0] + gathered = self._comm.all_gather(local) # [tp*rows, vocab_local] + if rows == 1: + return gathered.view(1, -1)[:, : self.vocab_size] + gathered = gathered.view(self.tp_size, rows, self.vocab_local) + gathered = gathered.permute(1, 0, 2).reshape(rows, self.tp_size * self.vocab_local) + return gathered[:, : self.vocab_size] + def hc_head(self, x): shape, dtype = x.size(), x.dtype dim = self.args.dim @@ -172,13 +216,13 @@ def prefill_batched( # metadata; ``flat_positions`` [T] is the scheduler-staged batch.positions (per-token # ABSOLUTE position); ``last_indices`` [B] is the flattened index of each request's # final token -> its next-token logits row. - h = self.embed(input_ids) + h = self.embed_tokens(input_ids) h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1) for layer in self.layers: h = layer.prefill_batched(h, input_ids, segments, flat_positions) h = self.hc_head(h) h = self.norm(h) - return F.linear(h[0, last_indices], self.head) # [B, vocab] + return self.logits(h[0, last_indices]) # [B, vocab] def decode( self, input_ids: torch.Tensor, pos: torch.Tensor, cmp_stage_cap: int @@ -194,7 +238,7 @@ def decode( # in-flight replay (it mutates only the live map). B = input_ids.size(0) rows = torch.arange(B, device=input_ids.device) - h = self.embed(input_ids) + h = self.embed_tokens(input_ids) h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1) # Hoist the layer-invariant per-step decode tensors (shared window-ring global slots): # resolved ONCE off the attention metadata (recomputed per call, so a capture records @@ -206,7 +250,7 @@ def decode( h = layer.decode_step(h, pos, rows, cmp_stage_cap, input_ids, wctx) h = self.hc_head(h) h = self.norm(h) - return F.linear(h[:, -1], self.head) + return self.logits(h[:, -1]) class DeepseekV4ForCausalLM(BaseLLMModel): diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py index d7e555a9..022aafb1 100644 --- a/python/freetoken/models/deepseek_v4/moe.py +++ b/python/freetoken/models/deepseek_v4/moe.py @@ -7,12 +7,14 @@ import torch.nn.functional as F from torch import nn +from freetoken.distributed import DistributedCommunicator from freetoken.kernel.triton.dsv4.bf16_linear import bf16_linear_fp32 from freetoken.kernel.triton.dsv4.swiglu import fused_swiglu from freetoken.layers import OffloadMoELayer from .args import DeepseekV4Args from .layers import Linear +from .parallel import div_tp, tp_size class Gate(nn.Module): @@ -56,13 +58,19 @@ def forward(self, x: torch.Tensor, input_ids: torch.Tensor): class Expert(nn.Module): - """Dense SwiGLU expert (the shared expert; routed experts are offloaded FP4).""" + """Dense SwiGLU expert (the shared expert; routed experts are offloaded FP4). + + Under TP the intermediate dimension splits: ``w1``/``w3`` are column-parallel and + ``w2`` is row-parallel, so the output is a partial sum. ``MoE.forward`` owns the + single all-reduce that completes it together with the routed half. + """ def __init__(self, dim: int, inter_dim: int, swiglu_limit: float): super().__init__() - self.w1 = Linear(dim, inter_dim, kind="fp8") - self.w2 = Linear(inter_dim, dim, kind="fp8") - self.w3 = Linear(dim, inter_dim, kind="fp8") + inter_local = div_tp(inter_dim, "moe_inter_dim", multiple_of=128) + self.w1 = Linear(dim, inter_local, kind="fp8") + self.w2 = Linear(inter_local, dim, kind="fp8") + self.w3 = Linear(dim, inter_local, kind="fp8") self.swiglu_limit = swiglu_limit def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -88,6 +96,16 @@ def __init__(self, layer_id: int, args: DeepseekV4Args): ) self.swiglu_limit = args.swiglu_limit + def _maybe_all_reduce(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Suppress the base class's per-call collective. + + The routed and the shared expert are both partial sums over the same output + dim, so DSV4 adds them first and all-reduces ONCE in ``MoE.forward``. Letting + the base reduce here would cost a second collective per layer and would also + double-count the shared expert. + """ + return hidden_states + def _prefill_routed( self, hidden_states: torch.Tensor, @@ -133,6 +151,7 @@ def __init__(self, layer_id: int, args: DeepseekV4Args): self.gate = Gate(layer_id, args) self.shared_experts = Expert(args.dim, args.moe_inter_dim, args.swiglu_limit) self.experts = DSV4OffloadMoELayer(layer_id, args) + self._comm = DistributedCommunicator() if tp_size() > 1 else None def forward(self, x: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: shape = x.size() @@ -147,4 +166,9 @@ def forward(self, x: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor: routed = self.experts.routed_forward( x, weights.float().contiguous(), indices.to(torch.int32).contiguous() ) - return (routed + shared).view(shape) + out = routed + shared + # Both halves are partial sums over the split intermediate dim; one collective + # completes the layer (see DSV4OffloadMoELayer._maybe_all_reduce). + if self._comm is not None: + out = self._comm.all_reduce(out) + return out.view(shape) diff --git a/python/freetoken/models/deepseek_v4/parallel.py b/python/freetoken/models/deepseek_v4/parallel.py new file mode 100644 index 00000000..4459e0bd --- /dev/null +++ b/python/freetoken/models/deepseek_v4/parallel.py @@ -0,0 +1,114 @@ +"""Tensor-parallel helpers for DeepSeek-V4-Flash. + +What DSV4 shards, and why: + +* **MLA query/output heads.** ``wq_b`` is column-parallel over ``n_heads``, + ``wo_a`` is sharded over the ``o_groups`` output groups, and ``wo_b`` is + row-parallel with one all-reduce at the end of the block. ``o_groups`` is the + binding constraint: a group owns ``n_heads // o_groups`` heads, so a rank must + own whole groups. +* **MoE intermediate dimension.** The routed FP4 expert banks and the shared + SwiGLU expert both split ``moe_inter_dim``. This is the memory win that makes + the model fit: the host expert banks divide by the TP size, so 4 ranks hold + the same 143 GB in total that 1 rank holds today. + +What stays REPLICATED, and why: + +* **The latent KV path** (``wkv``, ``kv_norm``, the paged pools). MLA keeps one + latent KV per token that every head reads, so there is nothing to split; the + KV pool cost per rank is unchanged. +* **The compressors and the Lightning Indexer.** They select which blocks the + attention reads. Every rank must select the SAME blocks. Replicating them + guarantees that without a collective on the selection path. +* **The router (``Gate``) and the hyper-connection mixers.** Small, and the + routing decision must agree across ranks. +""" + +from __future__ import annotations + +import torch + +from freetoken.distributed import try_get_tp_info + + +def tp_info() -> tuple[int, int]: + """``(rank, size)`` for this process; ``(0, 1)`` before the engine sets it.""" + info = try_get_tp_info() + return (0, 1) if info is None else (info.rank, info.size) + + +def tp_size() -> int: + return tp_info()[1] + + +def tp_rank() -> int: + return tp_info()[0] + + +def div_tp(total: int, what: str, *, multiple_of: int = 1) -> int: + """Per-rank size of ``total``. Fails loudly when the split is not exact. + + ``multiple_of`` guards the block-quantized weights: a 128x128 FP8 scale block + and a 32-element FP4 scale block cannot be cut in half. + """ + size = tp_size() + if total % size != 0: + raise ValueError( + f"DeepSeek-V4 tensor parallelism: {what}={total} is not divisible by " + f"--tensor-parallel-size {size}" + ) + local = total // size + if local % multiple_of != 0: + raise ValueError( + f"DeepSeek-V4 tensor parallelism: {what}={total} split over {size} ranks " + f"gives {local}, which is not a multiple of {multiple_of} (quantization " + f"block size)" + ) + return local + + +def shard(t: torch.Tensor, dim: int) -> torch.Tensor: + """This rank's slice of ``t`` along ``dim``, in its OWN storage. + + The copy is the point. ``narrow`` returns a view, and on dim 0 that view is already + contiguous, so ``.contiguous()`` hands it straight back -- keeping the whole parent + tensor alive behind a 1/N-sized view. Every rank then pays for the full weight it + just sharded: measured at 5.1 GiB per GPU on DeepSeek-V4-Flash at TP=4, silently + charged to the model and taken out of the KV and expert-cache budget. + """ + rank, size = tp_info() + if size == 1: + return t + total = t.shape[dim] + if total % size != 0: + raise ValueError( + f"cannot shard a tensor of shape {tuple(t.shape)} on dim {dim} over {size} ranks" + ) + step = total // size + return t.narrow(dim, rank * step, step).clone(memory_format=torch.contiguous_format) + + +def validate_tp(args) -> None: + """Check every DSV4 split up front, so a bad ``--tensor-parallel-size`` fails + at config time with one clear message instead of at the first bad reshape.""" + size = tp_size() + if size == 1: + return + div_tp(args.o_groups, "o_groups") + div_tp(args.n_heads, "n_heads") + if args.n_heads % args.o_groups != 0: + raise ValueError( + f"DeepSeek-V4 expects n_heads ({args.n_heads}) to be a multiple of " + f"o_groups ({args.o_groups})" + ) + # wq_b output rows and the FP8 128x128 scale grid. + div_tp(args.n_heads * args.head_dim, "n_heads*head_dim", multiple_of=128) + # wo_b input columns. + div_tp(args.o_groups * args.o_lora_rank, "o_groups*o_lora_rank", multiple_of=128) + # Shared expert (FP8 128-blocks) and routed FP4 experts (32-element scale blocks). + div_tp(args.moe_inter_dim, "moe_inter_dim", multiple_of=128) + # Vocabulary-parallel embedding + head. + div_tp(args.vocab_size, "vocab_size") + + +__all__ = ["div_tp", "shard", "tp_info", "tp_rank", "tp_size", "validate_tp"] diff --git a/python/freetoken/models/deepseek_v4/weight.py b/python/freetoken/models/deepseek_v4/weight.py index f5a3fd17..eb8c9c62 100644 --- a/python/freetoken/models/deepseek_v4/weight.py +++ b/python/freetoken/models/deepseek_v4/weight.py @@ -22,6 +22,7 @@ from freetoken.models.loader import drop_page_cache from .args import DeepseekV4Args, load_args +from .parallel import div_tp, shard, tp_info class _ShardReader: @@ -99,31 +100,44 @@ def iter_weights( def get(name: str) -> torch.Tensor: return reader.get(name) - def linear(prefix: str): - yield f"{prefix}.weight", get(f"{prefix}.weight") + def linear(prefix: str, split: int | None = None): + """Yield one linear's tensors, sharded for this rank. + + ``split=0`` is column-parallel (split the output rows), ``split=1`` is + row-parallel (split the input columns), ``None`` replicates. The 128x128 FP8 + ``scale`` grid splits on the SAME axis as its weight, so the two stay aligned. + """ + w = get(f"{prefix}.weight") + yield f"{prefix}.weight", w if split is None else shard(w, split) if reader.has(f"{prefix}.scale"): - yield f"{prefix}.scale", get(f"{prefix}.scale") + s = get(f"{prefix}.scale") + yield f"{prefix}.scale", s if split is None else shard(s, split) try: - yield "embed.weight", get("embed.weight") + # Vocabulary-parallel: each rank keeps one contiguous block of rows. + yield "embed.weight", shard(get("embed.weight"), 0) yield "norm.weight", get("norm.weight") - yield "head", get("head.weight") + yield "head", shard(get("head.weight"), 0) for nm in ("hc_head_fn", "hc_head_base", "hc_head_scale"): yield nm, get(nm) for L in range(args.n_layers): a = f"layers.{L}.attn" + # wq_a / wkv / the norms stay replicated: MLA keeps ONE latent KV per token + # that every head reads, so there is nothing to split on that path. yield from linear(f"{a}.wq_a") yield f"{a}.q_norm.weight", get(f"{a}.q_norm.weight") - yield from linear(f"{a}.wq_b") + yield from linear(f"{a}.wq_b", split=0) # column-parallel over heads yield from linear(f"{a}.wkv") yield f"{a}.kv_norm.weight", get(f"{a}.kv_norm.weight") # wo_a: FP8 in the checkpoint, dequantized to bf16 (reference bf16 einsum). - yield f"{a}.wo_a", _dequant_fp8_block( + # Rows are o_groups blocks of o_lora_rank, so a dim-0 split hands each rank + # whole groups -- matching the heads its wq_b shard produced. + yield f"{a}.wo_a", shard(_dequant_fp8_block( get(f"{a}.wo_a.weight"), get(f"{a}.wo_a.scale") - ) - yield from linear(f"{a}.wo_b") - yield f"{a}.attn_sink", get(f"{a}.attn_sink") + ), 0) + yield from linear(f"{a}.wo_b", split=1) # row-parallel; all-reduced in _wo + yield f"{a}.attn_sink", shard(get(f"{a}.attn_sink"), 0) ratio = args.compress_ratios[L] if ratio: @@ -151,8 +165,10 @@ def linear(prefix: str): yield f"{g}.tid2eid", get(f"{g}.tid2eid") else: yield f"{g}.bias", get(f"{g}.bias") - for proj in ("w1", "w2", "w3"): - yield from linear(f"layers.{L}.ffn.shared_experts.{proj}") + # Shared expert: the intermediate dim splits (w1/w3 column, w2 row); the + # partial sum is all-reduced together with the routed half in MoE.forward. + for proj, split in (("w1", 0), ("w2", 1), ("w3", 0)): + yield from linear(f"layers.{L}.ffn.shared_experts.{proj}", split=split) for nm in ( "hc_attn_fn", "hc_ffn_fn", "hc_attn_base", @@ -171,6 +187,35 @@ def linear(prefix: str): r"(?Pw1|w2|w3)\.(?Pweight|scale)$" ) +# The dSpark drafter's own routed experts. Its layers are appended after the target's, +# so ``mtp.k`` becomes bank layer ``n_layers + k`` and every layer-indexed structure +# (host banks, GPU slot cache) addresses draft and target layers identically. +_MTP_EXPERT_RE = re.compile( + r"^mtp\.(?P\d+)\.ffn\.experts\.(?P\d+)\." + r"(?Pw1|w2|w3)\.(?Pweight|scale)$" +) + + +def _expert_key(name: str, args: DeepseekV4Args): + """``(bank_layer, expert, proj, kind)`` for a routed-expert tensor, else None. + + Returns None for anything this run does not serve: a target layer past ``n_layers`` + and, when dSpark is off, every ``mtp.*`` expert. + """ + m = _EXPERT_RE.match(name) + if m is not None: + layer = int(m.group("layer")) + if layer >= args.n_layers: # a trailing MTP layer in the target namespace + return None + return layer, int(m.group("expert")), m.group("proj"), m.group("kind") + m = _MTP_EXPERT_RE.match(name) + if m is None: + return None + k = int(m.group("mtp")) + if k >= args.n_draft_layers: # dSpark off, or a layer beyond n_mtp_layers + return None + return args.n_layers + k, int(m.group("expert")), m.group("proj"), m.group("kind") + def load_dsfp4_expert_sources( model_path: str, args: DeepseekV4Args, *, layer_sink=None @@ -191,39 +236,39 @@ def load_dsfp4_expert_sources( folder = model_path weight_map = _weight_map(folder) - L, E = args.n_layers, args.n_routed_experts - H, I = args.dim, args.moe_inter_dim + # L covers the target layers plus, when dSpark is enabled, the drafter's own. + L, E = args.n_moe_layers, args.n_routed_experts - for shard in sorted(set(weight_map.values())): - drop_page_cache(os.path.join(folder, shard)) + for shard_file in sorted(set(weight_map.values())): + drop_page_cache(os.path.join(folder, shard_file)) - shards: dict[str, list[tuple[str, re.Match]]] = collections.defaultdict(list) - for name, shard in weight_map.items(): - m = _EXPERT_RE.match(name) - if m is None: - continue - if int(m.group("layer")) >= L: # skip the MTP layer (index L) - continue - shards[shard].append((name, m)) + shards: dict[str, list[tuple[str, tuple]]] = collections.defaultdict(list) + for name, shard_file in weight_map.items(): + key = _expert_key(name, args) + if key is not None: + shards[shard_file].append((name, key)) - e8m0 = torch.float8_e8m0fnu - specs = { # alloc UNPINNED, fill, then pin-after-fill (skips slow cudaHostAlloc zero-fill) - "gate_up_packed": ((E, 2 * I, H // 2), torch.uint8), - "gate_up_scale": ((E, 2 * I, H // 32), e8m0), - "down_packed": ((E, H, I // 2), torch.uint8), - "down_scale": ((E, H, I // 32), e8m0), - } + specs, I, i_lo = _expert_bank_specs(args) hb = alloc_layer_banks(specs, L) banks = {name: [b.tensor for b in hb[name]] for name in specs} + # Under TP a rank keeps only rows [i_lo, i_lo+I) of w1/w3, and those rows are + # CONTIGUOUS in the file -- so read just them instead of reading the whole tensor and + # throwing (N-1)/N of it away. w2 needs a COLUMN slice, whose rows are strided, so it + # is still read whole. w1+w3 are two thirds of the expert bytes, so at TP=4 this cuts + # a rank's read (and its copy work) roughly in half. + sliced_read = I != args.moe_inter_dim + def _load(sink) -> int: tracker = LayerCompletionTracker(E * 6, hb, sink) # {w1,w2,w3} x {weight,scale} x experts placed = 0 - for shard in tqdm(sorted(shards), desc="Loading DSV4 FP4 experts"): - path = os.path.join(folder, shard) + for shard_file in tqdm(sorted(shards), desc="Loading DSV4 FP4 experts"): + path = os.path.join(folder, shard_file) with safetensors.safe_open(path, framework="pt", device="cpu") as f: - for name, m in shards[shard]: - layer = _place_dsfp4(banks, name, f.get_tensor(name), I) + for name, key in shards[shard_file]: + rows_ready = sliced_read and key[2] in ("w1", "w3") + t = f.get_slice(name)[i_lo:i_lo + I] if rows_ready else f.get_tensor(name) + layer = _place_dsfp4(banks, key, t, I, i_lo, rows_ready) tracker.note(layer) placed += 1 drop_page_cache(path) @@ -244,15 +289,8 @@ def dummy_dsfp4_expert_sources(args: DeepseekV4Args) -> dict[str, list[torch.Ten """Fabricate the 4 ds_fp4 banks for --dummy-weight (no checkpoint on disk).""" from freetoken.moe.host_banks import alloc_layer_banks, pin_banks - L, E = args.n_layers, args.n_routed_experts - H, I = args.dim, args.moe_inter_dim - e8m0 = torch.float8_e8m0fnu - specs = { - "gate_up_packed": ((E, 2 * I, H // 2), torch.uint8), - "gate_up_scale": ((E, 2 * I, H // 32), e8m0), - "down_packed": ((E, H, I // 2), torch.uint8), - "down_scale": ((E, H, I // 32), e8m0), - } + L = args.n_moe_layers + specs, _I, _i_lo = _expert_bank_specs(args) hb = alloc_layer_banks(specs, L) banks = {name: [b.tensor for b in hb[name]] for name in specs} for t in banks["gate_up_packed"]: # packed e2m1; scales stay 0 (valid e8m0) @@ -265,30 +303,66 @@ def dummy_dsfp4_expert_sources(args: DeepseekV4Args) -> dict[str, list[torch.Ten def is_expert_tensor(name: str) -> bool: """Predicate for the common parallel reader: is this a routed-expert tensor?""" - return _EXPERT_RE.match(name) is not None + return _EXPERT_RE.match(name) is not None or _MTP_EXPERT_RE.match(name) is not None -def _place_dsfp4(banks: dict, name: str, t: torch.Tensor, I: int) -> int: +def _expert_bank_specs(args: DeepseekV4Args) -> tuple[dict, int, int]: + """The 4 routed-expert bank specs for THIS rank, plus its slice of ``moe_inter_dim``. + + Tensor parallelism splits the intermediate dim, so each rank allocates and reads + only its own ``I // tp`` rows. This is what divides the 143 GB of host expert banks + across the ranks instead of replicating them. + """ + E, H = args.n_routed_experts, args.dim + # FP4 packs 2 values per byte and carries one e8m0 scale per 32 values, so a rank's + # slice must stay a multiple of 32 on the intermediate axis. + i_local = div_tp(args.moe_inter_dim, "moe_inter_dim", multiple_of=32) + i_lo = tp_info()[0] * i_local + e8m0 = torch.float8_e8m0fnu + specs = { # alloc UNPINNED, fill, then pin-after-fill (skips slow cudaHostAlloc zero-fill) + "gate_up_packed": ((E, 2 * i_local, H // 2), torch.uint8), + "gate_up_scale": ((E, 2 * i_local, H // 32), e8m0), + "down_packed": ((E, H, i_local // 2), torch.uint8), + "down_scale": ((E, H, i_local // 32), e8m0), + } + return specs, i_local, i_lo + + +def _place_dsfp4( + banks: dict, key: tuple, t: torch.Tensor, I: int, i_lo: int = 0, + rows_ready: bool = False, +) -> int: """Copy one expert tensor into its layer/expert slot (shared by serial + parallel - readers): w1->gate_up[:,:I], w3->gate_up[:,I:], w2->down. Returns the layer index.""" - m = _EXPERT_RE.match(name) - layer, expert = int(m.group("layer")), int(m.group("expert")) - proj, kind = m.group("proj"), m.group("kind") + readers): w1->gate_up[:,:I], w3->gate_up[:,I:], w2->down. Returns the layer index. + + ``key`` is ``_expert_key``'s ``(bank_layer, expert, proj, kind)``, so a drafter + expert lands in bank layer ``n_layers + k`` with no special case here. + + ``I`` is this rank's intermediate width and ``i_lo`` its offset into the checkpoint's + full width, so each rank keeps only its own rows of w1/w3 and its own columns of w2. + ``rows_ready`` says the caller already read JUST those rows (the serial reader slices + w1/w3 at the file), so they must not be sliced a second time. + """ + layer, expert, proj, kind = key + + def rows(x: torch.Tensor) -> torch.Tensor: + return x if rows_ready else x[i_lo:i_lo + I] + if kind == "weight": t = t.view(torch.uint8) if proj == "w1": - banks["gate_up_packed"][layer][expert, :I] = t + banks["gate_up_packed"][layer][expert, :I] = rows(t) elif proj == "w3": - banks["gate_up_packed"][layer][expert, I:] = t - else: # w2 -> down - banks["down_packed"][layer][expert] = t - else: # scale (e8m0) + banks["gate_up_packed"][layer][expert, I:] = rows(t) + else: # w2 -> down; the intermediate axis is packed 2-per-byte + banks["down_packed"][layer][expert] = t[:, i_lo // 2:(i_lo + I) // 2] + else: # scale (e8m0), one per 32 values on the intermediate axis if proj == "w1": - banks["gate_up_scale"][layer][expert, :I] = t + banks["gate_up_scale"][layer][expert, :I] = rows(t) elif proj == "w3": - banks["gate_up_scale"][layer][expert, I:] = t + banks["gate_up_scale"][layer][expert, I:] = rows(t) else: - banks["down_scale"][layer][expert] = t + banks["down_scale"][layer][expert] = t[:, i_lo // 32:(i_lo + I) // 32] return layer @@ -302,27 +376,20 @@ def load_dsfp4_expert_sources_parallel( from freetoken.models.weight import iter_expert_tensors_parallel from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, alloc_layer_banks - L, E = args.n_layers, args.n_routed_experts - H, I = args.dim, args.moe_inter_dim - e8m0 = torch.float8_e8m0fnu - specs = { - "gate_up_packed": ((E, 2 * I, H // 2), torch.uint8), - "gate_up_scale": ((E, 2 * I, H // 32), e8m0), - "down_packed": ((E, H, I // 2), torch.uint8), - "down_scale": ((E, H, I // 32), e8m0), - } + # L covers the target layers plus, when dSpark is enabled, the drafter's own. + L, E = args.n_moe_layers, args.n_routed_experts + specs, I, i_lo = _expert_bank_specs(args) hb = alloc_layer_banks(specs, L) # lazy host banks (unpinned) banks = {name: [b.tensor for b in hb[name]] for name in specs} def _is_expert(name: str) -> bool: - m = _EXPERT_RE.match(name) - return m is not None and int(m.group("layer")) < L # skip the MTP layer (index L) + return _expert_key(name, args) is not None def _load(sink) -> int: tracker = LayerCompletionTracker(E * 6, hb, sink) placed = 0 for name, t in iter_expert_tensors_parallel(model_path, _is_expert, workers=workers, chunk=chunk): - layer = _place_dsfp4(banks, name, t, I) + layer = _place_dsfp4(banks, _expert_key(name, args), t, I, i_lo) tracker.note(layer) placed += 1 return placed diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index b96205aa..a1936637 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -88,8 +88,51 @@ def compiled_extension_supports(activation: str) -> bool: return _ACT_IDS[activation] <= getattr(_cpu_moe, "max_generic_act_id", lambda: 2)() +def _tp_core_share(cores: list[int]) -> list[int]: + """This rank's slice of ``cores`` when several TP ranks share one host. + + Every rank sees the same affinity mask, so without this each rank pins its whole + worker pool to the SAME physical cores: N ranks then oversubscribe those cores N + times and the spin-barrier collapses. Ranks take contiguous blocks rather than + striding, which keeps one rank's threads on one socket. + + A leftover from an uneven split stays unused: giving it to some ranks and not + others would desynchronize the per-layer CPU expert compute that every rank + finishes before the same collective. + """ + from freetoken.distributed import try_get_tp_info + + info = try_get_tp_info() + if info is None or info.size <= 1: + return cores + per = len(cores) // info.size + if per < 1: # fewer cores than ranks: let the OS schedule them + return cores + return cores[info.rank * per: (info.rank + 1) * per] + + +def _smt_siblings_of(cores: list[int]) -> list[int]: + """Every logical CPU that shares a physical core with one of ``cores``.""" + out: list[int] = [] + for cpu in cores: + try: + with open(f"/sys/devices/system/cpu/cpu{cpu}/topology/thread_siblings_list") as f: + spec = f.read().strip() + except OSError: + out.append(cpu) + continue + for part in spec.split(","): + if "-" in part: + lo, hi = part.split("-") + out.extend(range(int(lo), int(hi) + 1)) + elif part: + out.append(int(part)) + return sorted(set(out)) + + def physical_core_cpus() -> list[int]: - """One logical CPU per physical core, restricted to this process's affinity. + """One logical CPU per physical core, restricted to this process's affinity and, + under tensor parallelism, to this rank's share of them. MoE decode is memory-bandwidth-bound, so SMT siblings only contend for the same core's load ports without adding bandwidth. Picking one logical CPU per @@ -112,7 +155,7 @@ def physical_core_cpus() -> list[int]: if key not in seen: seen.add(key) reps.append(cpu) - return reps or allowed or [0] + return _tp_core_share(reps or allowed or [0]) def resolve_threads_and_affinity(requested: int) -> tuple[int, list[int]]: @@ -127,12 +170,10 @@ def resolve_threads_and_affinity(requested: int) -> tuple[int, list[int]]: reps = physical_core_cpus() if requested and requested > 0: n = int(requested) - try: - allowed = sorted(os.sched_getaffinity(0)) - except AttributeError: - allowed = list(range(os.cpu_count() or 1)) - # physical-core reps first, then the rest of the logical CPUs. - order = reps + [c for c in allowed if c not in set(reps)] + # physical-core reps first, then the SMT siblings OF THOSE CORES. Taking the + # siblings of this rank's own cores (rather than the rest of the affinity mask) + # keeps an explicit thread count from spilling onto another TP rank's cores. + order = reps + [c for c in _smt_siblings_of(reps) if c not in set(reps)] if not order: order = [0] core_ids = [order[i % len(order)] for i in range(n)] diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index a71b6819..d8181468 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -240,6 +240,18 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--distributed-timeout", + type=float, + default=ServerArgs.distributed_timeout, + help=( + "Seconds a tensor-parallel collective waits before failing. The first one is " + "the post-load memory all-reduce, which ranks reach minutes apart on a large " + "MoE checkpoint, so the default is sized for that load skew; lower it to make " + "a hung rank fail faster." + ), + ) + parser.add_argument( "--max-running-requests", type=int, diff --git a/tests/dsv4/test_dsv4_tensor_parallel.py b/tests/dsv4/test_dsv4_tensor_parallel.py new file mode 100644 index 00000000..dae12372 --- /dev/null +++ b/tests/dsv4/test_dsv4_tensor_parallel.py @@ -0,0 +1,150 @@ +"""DSV4 tensor parallelism: the shard contract, on the meta device. + +Every sharded parameter must be exactly ``1/tp`` of its TP=1 shape on exactly ONE axis, +and the routed FP4 expert banks must divide the same way -- that division is what lets +N ranks hold the same host bank bytes that one rank holds today, instead of N copies. +Replicated tensors (the MLA latent KV path, the compressors, the Lightning Indexer, the +router) must keep their full shape, because every rank reads the same latent KV and must +select the same blocks. + +CPU-only: shapes are read off a meta-device build, so no weights and no GPU. +""" + +from __future__ import annotations + +import pytest +import torch + +import freetoken.distributed.info as info_mod +from freetoken.distributed import DistributedInfo +from freetoken.models.deepseek_v4.args import DeepseekV4Args +from freetoken.models.deepseek_v4.weight import _expert_bank_specs + +# o_groups=8 bounds the split: a rank must own whole output groups. +TP_SIZES = (2, 4, 8) + +REPLICATED = (".wkv", ".kv_norm", ".q_norm", ".wq_a", ".compressor.", ".indexer.", "hc_") + + +@pytest.fixture +def args() -> DeepseekV4Args: + # Shipped DSV4-Flash widths (every divisibility rule keys off these), on a 4-layer + # stack so the meta build stays cheap. The ratio pattern mirrors tests/dsv4's: + # one uncompressed layer, one indexed (ratio 4) layer, one ratio-128 layer. + return DeepseekV4Args( + max_batch_size=1, + max_seq_len=4096, + n_layers=4, + n_hash_layers=1, + compress_ratios=(0, 4, 128, 4), + ) + + +def _set_tp(size: int, rank: int = 0) -> None: + # set_tp_info() is write-once by design; these tests walk several sizes in one process. + info_mod._TP_INFO = DistributedInfo(rank, size) + + +def _shapes(args: DeepseekV4Args, tp: int, rank: int = 0) -> dict[str, tuple[int, ...]]: + _set_tp(tp, rank) + from freetoken.models.deepseek_v4.model import Transformer + + with torch.device("meta"): + model = Transformer(args) + return {n: tuple(p.shape) for n, p in model.named_parameters()} + + +@pytest.fixture(autouse=True) +def _restore_tp(): + yield + _set_tp(1) + + +@pytest.mark.parametrize("tp", TP_SIZES) +def test_every_shard_is_a_clean_split_on_one_axis(args, tp): + base = _shapes(args, 1) + got = _shapes(args, tp) + assert got.keys() == base.keys(), "TP must not add or drop parameters" + + sharded = [n for n in got if got[n] != base[n]] + assert sharded, f"tp={tp} sharded nothing" + for name in sharded: + a, b = base[name], got[name] + axes = [i for i in range(len(a)) if a[i] != b[i]] + assert len(axes) == 1, f"{name}: {a} -> {b} splits on {len(axes)} axes, want 1" + assert a[axes[0]] == b[axes[0]] * tp, f"{name}: {a} -> {b} is not a 1/{tp} split" + + +@pytest.mark.parametrize("tp", TP_SIZES) +def test_replicated_tensors_keep_their_full_shape(args, tp): + base = _shapes(args, 1) + got = _shapes(args, tp) + for name, shape in got.items(): + if any(marker in name for marker in REPLICATED): + assert shape == base[name], f"{name} must stay replicated under TP" + + +@pytest.mark.parametrize("tp", TP_SIZES) +def test_ranks_cover_the_vocabulary_exactly_once(args, tp): + _set_tp(tp) + from freetoken.models.deepseek_v4.model import Transformer + + covered = 0 + for rank in range(tp): + _set_tp(tp, rank) + with torch.device("meta"): + model = Transformer(args) + assert model.vocab_start == covered + covered += model.vocab_local + assert covered == args.vocab_size + + +@pytest.mark.parametrize("tp", TP_SIZES) +def test_expert_banks_divide_and_tile_the_intermediate_dim(args, tp): + _set_tp(1) + full, full_i, _ = _expert_bank_specs(args) + + covered = 0 + for rank in range(tp): + _set_tp(tp, rank) + specs, i_local, i_lo = _expert_bank_specs(args) + assert i_lo == covered, "rank slices must tile the intermediate dim with no gap" + covered += i_local + assert i_local * tp == full_i + # A rank's bank bytes are exactly 1/tp of the whole -- the memory win. + for name, (shape, dtype) in specs.items(): + whole = torch.Size(full[name][0]).numel() + part = torch.Size(shape).numel() + assert part * tp == whole, f"{name}: {shape} is not 1/{tp} of {full[name][0]}" + assert dtype == full[name][1] + assert covered == full_i + + +@pytest.mark.parametrize("dim", [0, 1]) +def test_a_shard_does_not_keep_its_parent_alive(dim): + """A shard must own its storage. + + ``narrow`` returns a view, and a dim-0 view is already contiguous, so a + ``.contiguous()`` there hands the view straight back and pins the whole parent. + Every rank then pays for the full tensor it just sharded -- 5.1 GiB per GPU on + DSV4-Flash at TP=4, charged to the model and taken out of the cache budget. + """ + from freetoken.models.deepseek_v4.parallel import shard + + _set_tp(4, rank=1) + parent = torch.zeros(64, 64) + piece = shard(parent, dim) + assert piece.shape[dim] == 16 + assert piece.is_contiguous() + assert piece.untyped_storage().nbytes() == piece.numel() * piece.element_size(), ( + "shard is a view into the parent's storage" + ) + assert piece.data_ptr() != parent.data_ptr() + + +def test_a_split_that_does_not_divide_o_groups_fails_loudly(args): + from freetoken.models.deepseek_v4.parallel import validate_tp + + _set_tp(16) # o_groups == 8, so a rank cannot own a whole group + with pytest.raises(ValueError, match="o_groups"): + validate_tp(args) From 6151d7f0022fb3585b9bec2852e734b6d15cfc79 Mon Sep 17 00:00:00 2001 From: Nhan Nguyen Date: Tue, 25 Aug 2026 02:03:26 +0000 Subject: [PATCH 2/6] fix(server): align TP GPU mapping after rebase --- docs/cli.md | 2 +- docs/models.md | 3 ++- python/freetoken/server/args.py | 4 +++- tests/server/test_parser_auto_selection.py | 21 +++++++++++++++++++++ 4 files changed, 27 insertions(+), 3 deletions(-) diff --git a/docs/cli.md b/docs/cli.md index 34302326..339c4d22 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -41,7 +41,7 @@ parsers all resolve automatically from the checkpoint and the GPU. | `--host` | 127.0.0.1 | Bind address | | `--port` | 1919 | Bind port | | `--gpu` | GPU 0 | GPU to run on: a UUID from `nvidia-smi -L` or an `nvidia-smi` index; see [below](#choosing-a-gpu) | -| `--tensor-parallel-size`, `--tp-size` | 1 | GPUs to shard the model over; one scheduler process per rank, rank `i` on `cuda:i` | +| `--tensor-parallel-size`, `--tp-size` | 1 | GPUs to shard the model over; one scheduler process per rank, mapped in `--gpu` order (default: rank `i` on `cuda:i`) | | `--distributed-timeout` | 1800 | Seconds a TP collective waits before failing; sized for the per-rank skew in weight-load time | | `--max-running-requests` | 4 | Max concurrently running requests | | `--max-output-tokens` | 32768 | Default output budget for requests that omit one | diff --git a/docs/models.md b/docs/models.md index 4381a695..6f6e68a8 100644 --- a/docs/models.md +++ b/docs/models.md @@ -34,7 +34,8 @@ for them; other checkpoints of the same architectures work too. ## Tensor parallelism `ft serve --tensor-parallel-size N` shards the model over N GPUs on one host. One -scheduler process runs per rank, and rank `i` uses `cuda:i`. +scheduler process runs per rank. By default rank `i` uses `cuda:i`; an explicit +`--gpu` list maps its `i`th entry to rank `i`. DeepSeek-V4 shards: diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index d8181468..eb8f706e 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -645,7 +645,9 @@ def _infer_reasoning_parser(model_path: str) -> str | None: # reject a too-long list here with a clear reason, not as a dead rank later if len(kwargs["gpu"]) not in (0, kwargs["tensor_parallel_size"]): if kwargs["tensor_parallel_size"] == 1 and len(kwargs["gpu"]) > 1: - parser.error("tensor parallelism is not supported yet: --gpu takes one entry") + parser.error( + "--gpu has multiple entries; set --tensor-parallel-size to the same count" + ) parser.error( f"--gpu has {len(kwargs['gpu'])} entries but --tensor-parallel-size is " f"{kwargs['tensor_parallel_size']}; give one entry per TP rank" diff --git a/tests/server/test_parser_auto_selection.py b/tests/server/test_parser_auto_selection.py index 78b1663e..58f9563d 100644 --- a/tests/server/test_parser_auto_selection.py +++ b/tests/server/test_parser_auto_selection.py @@ -94,3 +94,24 @@ def test_an_explicit_choice_beats_inference(): pinned, _ = parse_args(["--model", ANON_PATH, "--reasoning-parser", "qwen3"]) assert off.reasoning_parser is None assert pinned.reasoning_parser == "qwen3" + + +def test_tp_gpu_list_and_distributed_timeout_reach_server_args(): + config = _Config({"architectures": ["DeepseekV4ForCausalLM"], "torch_dtype": "bfloat16"}) + with patch("freetoken.utils.cached_load_hf_config", lambda _path: config): + args, _ = parse_args( + [ + "--model", + ANON_PATH, + "--tp-size", + "4", + "--gpu", + "3,1,0,2", + "--distributed-timeout", + "900", + ] + ) + + assert args.tp_info.size == 4 + assert args.gpu == ("3", "1", "0", "2") + assert args.distributed_timeout == 900 From ec5481f79a574ecdf9d15265e7773b6f796a5185 Mon Sep 17 00:00:00 2001 From: Nhan Nguyen Date: Tue, 25 Aug 2026 02:10:58 +0000 Subject: [PATCH 3/6] test(dsv4): verify TP expert bank contents --- tests/dsv4/test_dsv4_tensor_parallel.py | 83 ++++++++++++++++++++++++- 1 file changed, 82 insertions(+), 1 deletion(-) diff --git a/tests/dsv4/test_dsv4_tensor_parallel.py b/tests/dsv4/test_dsv4_tensor_parallel.py index dae12372..72f65efe 100644 --- a/tests/dsv4/test_dsv4_tensor_parallel.py +++ b/tests/dsv4/test_dsv4_tensor_parallel.py @@ -18,7 +18,7 @@ import freetoken.distributed.info as info_mod from freetoken.distributed import DistributedInfo from freetoken.models.deepseek_v4.args import DeepseekV4Args -from freetoken.models.deepseek_v4.weight import _expert_bank_specs +from freetoken.models.deepseek_v4.weight import _expert_bank_specs, _place_dsfp4 # o_groups=8 bounds the split: a rank must own whole output groups. TP_SIZES = (2, 4, 8) @@ -148,3 +148,84 @@ def test_a_split_that_does_not_divide_o_groups_fails_loudly(args): _set_tp(16) # o_groups == 8, so a rank cannot own a whole group with pytest.raises(ValueError, match="o_groups"): validate_tp(args) + + +@pytest.mark.parametrize("rank", range(4)) +def test_expert_loader_places_exact_rank_slice(rank): + """The bank contents, not only their shapes, must tile the packed I axis.""" + + full_i, local_i, hidden = 128, 32, 64 + i_lo = rank * local_i + e8m0 = torch.float8_e8m0fnu + + def payload(shape, offset, dtype=torch.int8): + raw = (torch.arange(torch.Size(shape).numel(), dtype=torch.int64) + offset) % 251 + return raw.to(torch.uint8).reshape(shape).view(dtype) + + w1 = payload((full_i, hidden // 2), 1) + w3 = payload((full_i, hidden // 2), 17) + w2 = payload((hidden, full_i // 2), 33) + s1 = payload((full_i, hidden // 32), 49, e8m0) + s3 = payload((full_i, hidden // 32), 65, e8m0) + s2 = payload((hidden, full_i // 32), 81, e8m0) + banks = { + "gate_up_packed": [torch.zeros(1, 2 * local_i, hidden // 2, dtype=torch.uint8)], + "gate_up_scale": [torch.zeros(1, 2 * local_i, hidden // 32, dtype=e8m0)], + "down_packed": [torch.zeros(1, hidden, local_i // 2, dtype=torch.uint8)], + "down_scale": [torch.zeros(1, hidden, local_i // 32, dtype=e8m0)], + } + + for proj, kind, tensor in ( + ("w1", "weight", w1), + ("w3", "weight", w3), + ("w2", "weight", w2), + ("w1", "scale", s1), + ("w3", "scale", s3), + ("w2", "scale", s2), + ): + _place_dsfp4(banks, (0, 0, proj, kind), tensor, local_i, i_lo) + + assert torch.equal( + banks["gate_up_packed"][0][0, :local_i], + w1.view(torch.uint8)[i_lo:i_lo + local_i], + ) + assert torch.equal( + banks["gate_up_packed"][0][0, local_i:], + w3.view(torch.uint8)[i_lo:i_lo + local_i], + ) + assert torch.equal( + banks["down_packed"][0][0], + w2.view(torch.uint8)[:, i_lo // 2:(i_lo + local_i) // 2], + ) + assert torch.equal( + banks["gate_up_scale"][0][0, :local_i].view(torch.uint8), + s1[i_lo:i_lo + local_i].view(torch.uint8), + ) + assert torch.equal( + banks["gate_up_scale"][0][0, local_i:].view(torch.uint8), + s3[i_lo:i_lo + local_i].view(torch.uint8), + ) + assert torch.equal( + banks["down_scale"][0][0].view(torch.uint8), + s2[:, i_lo // 32:(i_lo + local_i) // 32].view(torch.uint8), + ) + + +def test_serial_expert_loader_does_not_slice_pre_sliced_rows_twice(): + full_i, local_i, hidden, i_lo = 128, 32, 64, 64 + full = torch.arange(full_i * (hidden // 2), dtype=torch.int64) + full = (full % 251).to(torch.uint8).reshape(full_i, hidden // 2).view(torch.int8) + rank_rows = full[i_lo:i_lo + local_i] + bank = torch.zeros(1, 2 * local_i, hidden // 2, dtype=torch.uint8) + banks = {"gate_up_packed": [bank]} + + _place_dsfp4( + banks, + (0, 0, "w1", "weight"), + rank_rows, + local_i, + i_lo, + rows_ready=True, + ) + + assert torch.equal(bank[0, :local_i], rank_rows.view(torch.uint8)) From 13abef7c2cead08d28eadfe04660a647af38fdad Mon Sep 17 00:00:00 2001 From: Nhan Nguyen Date: Tue, 25 Aug 2026 02:14:46 +0000 Subject: [PATCH 4/6] fix(dsv4): reject TP1 FTW layout under TP --- docs/models.md | 4 ++++ python/freetoken/engine/engine.py | 10 ++++++++++ tests/dsv4/test_dsv4_tensor_parallel.py | 16 ++++++++++++++++ 3 files changed, 30 insertions(+) diff --git a/docs/models.md b/docs/models.md index 6f6e68a8..657a33d0 100644 --- a/docs/models.md +++ b/docs/models.md @@ -53,6 +53,10 @@ Constraints for DeepSeek-V4: N must divide `o_groups` (8), so N is 1, 2, 4 or 8. The KV pool is replicated, so its cost per GPU does not fall with N; the weights and the expert banks do. +DeepSeek-V4 TP currently loads the original safetensors checkpoint. An FTW +conversion records the TP=1 dense and expert-bank layout and is rejected for +`N > 1` instead of failing later with rank-local shape mismatches. + ## Notes - `ft checkpoint` conversion is optional — it pre-converts a checkpoint into diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 220c2b99..ca68b0f5 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -1095,6 +1095,16 @@ def _adjust_dsv4_config(config: EngineConfig, override) -> None: page_size to the window page P, forces single-chunk prefill, and clamps cuda_graph_bs/max_bs to the DSV4 decode batch size. """ + if config.tp_info.size > 1: + from freetoken.checkpoint.ftw import is_ftw_checkpoint + + if is_ftw_checkpoint(config.model_path): + raise ValueError( + "DeepSeek-V4 tensor parallelism currently requires the original " + "safetensors checkpoint. FTW stores the TP=1 dense weights and expert " + "banks, so it cannot be replayed into rank-local DSV4 shapes; use the " + "raw checkpoint or run with --tp-size 1." + ) model_config = config.model_config model_config.dsv4_args.max_seq_len = config.max_seq_len model_config.dsv4_args.max_batch_size = config.max_running_req + 1 # +1 dummy diff --git a/tests/dsv4/test_dsv4_tensor_parallel.py b/tests/dsv4/test_dsv4_tensor_parallel.py index 72f65efe..26b1371b 100644 --- a/tests/dsv4/test_dsv4_tensor_parallel.py +++ b/tests/dsv4/test_dsv4_tensor_parallel.py @@ -12,6 +12,8 @@ from __future__ import annotations +from types import SimpleNamespace + import pytest import torch @@ -229,3 +231,17 @@ def test_serial_expert_loader_does_not_slice_pre_sliced_rows_twice(): ) assert torch.equal(bank[0, :local_i], rank_rows.view(torch.uint8)) + + +def test_tp_rejects_ftw_tp1_layout_before_model_setup(tmp_path): + from freetoken.checkpoint.ftw import INDEX_NAME + from freetoken.engine.engine import _adjust_dsv4_config + + (tmp_path / INDEX_NAME).write_text("{}", encoding="utf-8") + config = SimpleNamespace( + model_path=str(tmp_path), + tp_info=SimpleNamespace(size=4), + ) + + with pytest.raises(ValueError, match="FTW stores the TP=1"): + _adjust_dsv4_config(config, lambda _name, _value: None) From a9cdf5879a291c4f6d14cdfeb5d6a673ee257552 Mon Sep 17 00:00:00 2001 From: Nhan Nguyen Date: Tue, 25 Aug 2026 03:36:14 +0000 Subject: [PATCH 5/6] feat(moe): allow pageable CPU-layer banks --- docs/cli.md | 4 ++++ python/freetoken/moe/host_banks.py | 6 +++++- tests/moe/test_offload.py | 19 +++++++++++++++++++ 3 files changed, 28 insertions(+), 1 deletion(-) diff --git a/docs/cli.md b/docs/cli.md index 339c4d22..1bc8ecf3 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -90,6 +90,10 @@ See [models.md](models.md#moe-backends) for what each backend does. | `--moe-prefill-hit-d2d` | off | Prefill: copy cache-hit experts device-side, stream only misses (CUDA >= 13) | | `--disable-moe-prefill-overlap` | overlap on | Disable the two-buffer prefill copy overlap | +On hosts where OS-locked and CUDA-registered memory share one quota, combine +`--moe-cpu-layers` with `FREETOKEN_SKIP_BANK_LOCK=1`. CPU-only layer banks then +remain pageable instead of consuming the quota needed by GPU-fetch layers. + ### API behaviour | Flag | Default | Meaning | diff --git a/python/freetoken/moe/host_banks.py b/python/freetoken/moe/host_banks.py index d7af348a..5f57cbb0 100644 --- a/python/freetoken/moe/host_banks.py +++ b/python/freetoken/moe/host_banks.py @@ -152,9 +152,13 @@ def release(self) -> None: def lock(self) -> None: """mlock the (now-filled) buffer: resident without CUDA pin quota, but no device address -- only the CPU executor can serve a locked layer. - Lock after fill, or the lazy mmap faults+zero-fills every page. A failed lock (RLIMIT_MEMLOCK) warns once and leaves the bank PAGEABLE, which every consumer treats the same.""" + Lock after fill, or the lazy mmap faults+zero-fills every page. A failed lock (RLIMIT_MEMLOCK) warns once and leaves the bank PAGEABLE, which every consumer treats the same. ``FREETOKEN_SKIP_BANK_LOCK=1`` deliberately leaves CPU-only layers pageable on hosts where ``mlock`` and ``cudaHostRegister`` consume the same quota.""" if self._locked or self._pinned: # cudaHostRegister already page-locks return + if os.environ.get("FREETOKEN_SKIP_BANK_LOCK", "").strip().lower() in ( + "1", "true", "yes", "on", + ): + return global _os_lock_failed if _os_lock_failed: return # the quota is exhausted for good; skip the syscall spam diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 422ca867..ea13f559 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -868,3 +868,22 @@ def boom(addr, nbytes): with hb.PinPipeline() as pins: pins(1, {"gate_up": hb.HostBank((4,), torch.uint8)}) assert plan2.actual == {1: hb.HostResidency.PAGEABLE.value} + + +def test_skip_bank_lock_leaves_cpu_layer_pageable(monkeypatch): + """A shared OS/CUDA lock quota can reserve locking for GPU-fetch layers.""" + import freetoken.moe.host_banks as hb + + def unexpected_lock(addr, nbytes): + raise AssertionError("FREETOKEN_SKIP_BANK_LOCK must bypass mlock") + + monkeypatch.setattr(hb, "_os_lock", unexpected_lock) + monkeypatch.setenv("FREETOKEN_SKIP_BANK_LOCK", "1") + labels = [hb.HostResidency.LOCKED.value] + banks = {"gate_up": [hb.HostBank((4,), torch.uint8)]} + + with hb.requested_residency(labels) as plan: + hb.pin_banks(banks) + + assert plan.actual == {0: hb.HostResidency.PAGEABLE.value} + assert banks["gate_up"][0].residency == hb.HostResidency.PAGEABLE From c066b38e1c9981c8c8adb0efb328049fe87aa8c0 Mon Sep 17 00:00:00 2001 From: Nhan Nguyen Date: Tue, 25 Aug 2026 05:06:55 +0000 Subject: [PATCH 6/6] feat(moe): bound parallel expert prefetch --- docs/cli.md | 1 + python/freetoken/engine/config.py | 10 +++-- python/freetoken/engine/engine.py | 2 + python/freetoken/models/deepseek_v4/weight.py | 6 ++- python/freetoken/models/weight.py | 3 ++ python/freetoken/moe/expert_banks.py | 45 +++++++++++-------- python/freetoken/server/args.py | 9 ++++ tests/server/test_parser_auto_selection.py | 6 +++ 8 files changed, 59 insertions(+), 23 deletions(-) diff --git a/docs/cli.md b/docs/cli.md index 1bc8ecf3..48dfbf66 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -82,6 +82,7 @@ See [models.md](models.md#moe-backends) for what each backend does. | Flag | Default | Meaning | |---|---|---| | `--moe-backend` | auto | `fused`/`offload`/`cpu`/`hybrid`; auto → offload, or hybrid with a `ft bench bw` profile | +| `--expert-load` / `--expert-prefetch` | auto / 2 | Expert-bank reader; parallel prefetch controls whole shards queued ahead of placement (use 1 to reduce TP host-RAM peak) | | `--moe-cache-size` / `--moe-cache-rate` / `--moe-cache-auto` | auto | GPU expert-cache size as slots / fraction of all experts / sized from free VRAM (mutually exclusive; auto is enabled by default for offload-family backends) | | `--kv-reserve-tokens` | 8192 | KV token floor reserved before `--moe-cache-auto` fills experts | | `--moe-cpu-threads` | physical cores | CPU worker threads for the cpu/hybrid executor | diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 0bb45827..3a2d0b61 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -24,10 +24,14 @@ class EngineConfig: # NVFP4 routed-expert GEMM backend (--nvfp4-backend): auto|marlin|flashinfer|triton. nvfp4_backend: str = "triton" # Expert-bank host load (--expert-load): auto|serial|parallel. "auto" reads scattered - # experts in parallel but falls back to serial when free RAM can't cover the banks + the - # parallel reader's extra (non-reclaimable) whole-shard buffer; "serial" forces the - # low-memory reclaimable read; "parallel" forces the fast read. + # experts in parallel but falls back to serial when free RAM can't cover the banks + + # bounded (non-reclaimable) whole-shard read-ahead; "serial" forces the low-memory + # reclaimable read; "parallel" forces the fast read. expert_load: str = "auto" + # Whole checkpoint shards queued ahead of placement by parallel expert + # readers. One preserves I/O/placement overlap with lower peak host RAM; + # larger values can smooth uneven shard-placement time on roomy hosts. + expert_prefetch: int = 2 moe_cache_size: int = 0 moe_cache_rate: float | None = None moe_cache_auto: bool = False diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index ca68b0f5..10c89e4c 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -628,6 +628,7 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: dtype=self.dtype, dummy=config.use_dummy_weight, parallel=expert_parallel, + prefetch=config.expert_prefetch, decode_target=("cpu" if decode_target in ("cpu", "hybrid") else "gpu"), layer_residency=requested_residency, ) @@ -1270,6 +1271,7 @@ def _auto_cpu_layers(config: EngineConfig, num_moe_layers: int) -> frozenset[int "moe_prefill_overlap": True, "moe_prefill_hit_d2d": False, "expert_load": "auto", + "expert_prefetch": 2, } diff --git a/python/freetoken/models/deepseek_v4/weight.py b/python/freetoken/models/deepseek_v4/weight.py index eb8c9c62..17c2fbc9 100644 --- a/python/freetoken/models/deepseek_v4/weight.py +++ b/python/freetoken/models/deepseek_v4/weight.py @@ -368,7 +368,7 @@ def rows(x: torch.Tensor) -> torch.Tensor: def load_dsfp4_expert_sources_parallel( model_path: str, args: DeepseekV4Args, *, workers: int = 8, chunk: int = 8 << 20, - layer_sink=None, + prefetch: int = 2, layer_sink=None, ) -> dict[str, list[torch.Tensor]]: """parallel path: same banks as load_dsfp4_expert_sources, filled from the common chunked multi-threaded O_DIRECT reader instead of serial per-shard safe_open. @@ -388,7 +388,9 @@ def _is_expert(name: str) -> bool: def _load(sink) -> int: tracker = LayerCompletionTracker(E * 6, hb, sink) placed = 0 - for name, t in iter_expert_tensors_parallel(model_path, _is_expert, workers=workers, chunk=chunk): + for name, t in iter_expert_tensors_parallel( + model_path, _is_expert, workers=workers, chunk=chunk, prefetch=prefetch + ): layer = _place_dsfp4(banks, _expert_key(name, args), t, I, i_lo) tracker.note(layer) placed += 1 diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index 6a34f3b9..4a71a505 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -90,6 +90,9 @@ def iter_expert_tensors_parallel( Peak host memory is ~(prefetch+1) shards + the banks the caller fills. Order is shard-then-header order (NOT global), so the consumer must place by ``name``. """ + if prefetch < 1: + raise ValueError(f"expert prefetch must be >= 1, got {prefetch}") + from freetoken.utils.hf import download_hf_weight model_path = download_hf_weight(model_path) # resolve hub id -> local (parity w/ serial) diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index 8b6116ba..868310cf 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -16,6 +16,7 @@ from __future__ import annotations import glob +import inspect import os import threading from dataclasses import dataclass, field @@ -252,7 +253,7 @@ def _q4_0_banks(model_path, model_config, device, dtype, dummy, parallel=False, ) -def _dsfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None) -> ExpertBanks: +def _dsfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None, prefetch=2) -> ExpertBanks: args = model_config.dsv4_args assert args is not None, "ds_fp4 expert banks require dsv4_args on the model config" # DeepSeek-FP4: packed e2m1 + e8m0 per-32 block scales, no global scale -> 4 banks, @@ -267,7 +268,8 @@ def _dsfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, from freetoken.models.deepseek_v4.weight import load_dsfp4_expert_sources_parallel banks = load_dsfp4_expert_sources_parallel( - model_path, args, workers=workers, chunk=chunk, layer_sink=sink + model_path, args, workers=workers, chunk=chunk, prefetch=prefetch, + layer_sink=sink, ) else: from freetoken.models.deepseek_v4.weight import load_dsfp4_expert_sources @@ -304,7 +306,7 @@ def _model_setup_override(model_config): } -def _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel, workers, chunk, decode_target="gpu", layer_sink=None) -> ExpertBanks: +def _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel, workers, chunk, decode_target="gpu", layer_sink=None, prefetch=2) -> ExpertBanks: """Dispatch to the model's setup-override or the per-quant provider. ``parallel=True`` is the parallel read; a provider that hasn't implemented it raises NotImplementedError (the caller falls back to serial). ``decode_target`` lets the cpu backend force CPU-readable @@ -313,8 +315,6 @@ def _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel materialize-and-write path (``ExpertBanks.streamed`` reports which happened).""" setup = _model_setup_override(model_config) if setup is not None: - import inspect - params = inspect.signature(setup).parameters supports_parallel = "parallel" in params if parallel and not supports_parallel: @@ -326,6 +326,8 @@ def _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel kw = dict(device=device, dtype=dtype, dummy=dummy) if supports_parallel: kw.update(parallel=parallel, workers=workers, chunk=chunk) + if "prefetch" in params: + kw["prefetch"] = prefetch if "decode_target" in params: kw["decode_target"] = decode_target if "layer_sink" in params and layer_sink is not None: @@ -338,18 +340,23 @@ def _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel f"no expert-bank provider for expert_quant={expert_quant!r} " f"(known: {sorted(_PROVIDERS)})" ) - return _PROVIDERS[expert_quant]( + provider = _PROVIDERS[expert_quant] + extra = {"prefetch": prefetch} if "prefetch" in inspect.signature(provider).parameters else {} + return provider( model_path, model_config, device, dtype, dummy, parallel=parallel, workers=workers, chunk=chunk, decode_target=decode_target, - layer_sink=layer_sink, + layer_sink=layer_sink, **extra, ) -def _host_ram_fits_parallel(model_path: str) -> bool: - """Best-effort: can free host RAM hold the expert banks plus the parallel reader's one - extra (non-reclaimable) whole-shard buffer? Unknown (non-local path / no /proc) -> True, - i.e. keep the fast path. Banks ~= checkpoint size (experts dominate); transient ~= the - largest shard. Uses MemAvailable (counts reclaimable cache) -- the OOM-relevant figure.""" +def _host_ram_fits_parallel(model_path: str, prefetch: int = 2) -> bool: + """Best-effort: can free host RAM hold banks plus the parallel read-ahead? + + The consumer retains its current whole-shard buffer while the bounded queue + holds up to ``prefetch`` more, so transient residency is at most + ``(prefetch + 1) * largest_shard``. Unknown inputs conservatively keep the + fast path. Uses MemAvailable, which includes reclaimable page cache. + """ avail = None try: with open("/proc/meminfo") as f: @@ -370,7 +377,7 @@ def _host_ram_fits_parallel(model_path: str) -> bool: sizes = [os.path.getsize(p) for p in glob.glob(os.path.join(model_path, "*.safetensors"))] if not sizes: return True - return avail > sum(sizes) + max(sizes) + return avail > sum(sizes) + (prefetch + 1) * max(sizes) def ftw_bank_bytes(model_path: str) -> int | None: @@ -415,6 +422,7 @@ def load_expert_banks( parallel: bool | None = None, workers: int = 8, chunk: int = _PARALLEL_CHUNK, + prefetch: int = 2, decode_target: str = "gpu", layer_sink=None, layer_residency: list[str] | None = None, @@ -431,6 +439,7 @@ def load_expert_banks( ``parallel`` overrides the slow-path auto-pick: ``None`` = auto (production), ``True`` / ``False`` = force parallel / serial (used by the loader benchmark and the converter). + ``prefetch`` bounds whole shards queued ahead of placement by supporting parallel readers. ``layer_sink`` (the converter only): forwarded to whichever provider is picked; a provider only engages it (and reports ``ExpertBanks.streamed=True``) for its own @@ -465,9 +474,9 @@ def load_expert_banks( # Low-RAM fallback: the parallel reader holds whole-shard ANONYMOUS buffers # (non-reclaimable) on top of the ~bank-sized resident set, so on a memory-tight box # it OOMs where the serial path (reclaimable file mmap) survives. Drop to serial when - # free RAM can't cover the banks + one shard's transient. (--expert-load serial/parallel - # bypass this by forcing ``parallel`` explicitly.) - if parallel and not _host_ram_fits_parallel(model_path): + # free RAM can't cover the banks + the bounded shard read-ahead. (--expert-load + # serial/parallel bypass this by forcing ``parallel`` explicitly.) + if parallel and not _host_ram_fits_parallel(model_path, prefetch): logger.warning_rank0( "expert banks: low free RAM -> serial build (avoids parallel-reader OOM; " "override with --expert-load parallel)" @@ -483,13 +492,13 @@ def load_expert_banks( with requested_residency(layer_residency) as residency_plan: try: banks = _build_expert_banks(model_path, model_config, device, dtype, dummy, parallel, workers, chunk, - decode_target, layer_sink) + decode_target, layer_sink, prefetch) except NotImplementedError as exc: if not parallel: raise logger.warning_rank0(f"parallel reader unavailable ({exc}); falling back to serial build") banks = _build_expert_banks(model_path, model_config, device, dtype, dummy, False, workers, chunk, - decode_target, layer_sink) + decode_target, layer_sink, prefetch) return _echo_residency(banks, layer_residency, residency_plan) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index eb8f706e..870c0b7d 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -507,6 +507,15 @@ def _infer_reasoning_parser(model_path: str) -> str | None: "low-memory reclaimable read (slower); 'parallel' forces the fast read." ), ) + parser.add_argument( + "--expert-prefetch", + type=_positive_int, + default=ServerArgs.expert_prefetch, + help=( + "Whole checkpoint shards queued ahead of placement by parallel expert " + "loading (default: 2). Use 1 on host-memory-constrained TP deployments." + ), + ) moe_cache_group = parser.add_mutually_exclusive_group() moe_cache_group.add_argument( diff --git a/tests/server/test_parser_auto_selection.py b/tests/server/test_parser_auto_selection.py index 58f9563d..7596da05 100644 --- a/tests/server/test_parser_auto_selection.py +++ b/tests/server/test_parser_auto_selection.py @@ -109,9 +109,15 @@ def test_tp_gpu_list_and_distributed_timeout_reach_server_args(): "3,1,0,2", "--distributed-timeout", "900", + "--expert-load", + "parallel", + "--expert-prefetch", + "1", ] ) assert args.tp_info.size == 4 assert args.gpu == ("3", "1", "0", "2") assert args.distributed_timeout == 900 + assert args.expert_load == "parallel" + assert args.expert_prefetch == 1