diff --git a/docs/cli.md b/docs/cli.md index ff4af382..48dfbf66 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, 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 | | `--max-seq-len-override` | from checkpoint | Max sequence length | @@ -80,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 | @@ -88,6 +91,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 | @@ -172,4 +179,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..657a33d0 100644 --- a/docs/models.md +++ b/docs/models.md @@ -31,6 +31,32 @@ 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. By default rank `i` uses `cuda:i`; an explicit +`--gpu` list maps its `i`th entry to rank `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. + +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/config.py b/python/freetoken/engine/config.py index 543012f3..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 @@ -72,7 +76,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..10c89e4c 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 — @@ -579,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, ) @@ -1046,6 +1096,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 @@ -1211,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/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..17c2fbc9 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,36 +303,72 @@ 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 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. @@ -302,27 +376,22 @@ 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) + 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 return placed 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/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/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/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/python/freetoken/server/args.py b/python/freetoken/server/args.py index a71b6819..870c0b7d 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, @@ -495,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( @@ -633,7 +654,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/dsv4/test_dsv4_tensor_parallel.py b/tests/dsv4/test_dsv4_tensor_parallel.py new file mode 100644 index 00000000..26b1371b --- /dev/null +++ b/tests/dsv4/test_dsv4_tensor_parallel.py @@ -0,0 +1,247 @@ +"""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 + +from types import SimpleNamespace + +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, _place_dsfp4 + +# 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) + + +@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)) + + +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) 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 diff --git a/tests/server/test_parser_auto_selection.py b/tests/server/test_parser_auto_selection.py index 78b1663e..7596da05 100644 --- a/tests/server/test_parser_auto_selection.py +++ b/tests/server/test_parser_auto_selection.py @@ -94,3 +94,30 @@ 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", + "--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