From 5b9644acc6ae4cc9d16562d7c5d42404b8c78c0e Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Thu, 22 Oct 2026 23:10:52 +0000 Subject: [PATCH 1/2] feat: add --embed-weights host for a pinned host embedding table --- docs/cli.md | 1 + python/freetoken/engine/config.py | 2 + python/freetoken/engine/engine.py | 81 +++++++++++++++++-- .../kernel/csrc/ggml/freetoken/ggml_ops.cu | 9 ++- .../kernel/csrc/include/freetoken/pinned.h | 40 +++++++++ .../kernel/csrc/jit/fast_index_copy.cuh | 32 +------- python/freetoken/kernel/csrc/jit/index.cu | 11 ++- python/freetoken/kernel/ggml.py | 2 +- python/freetoken/kernel/index.py | 2 +- python/freetoken/server/args.py | 11 +++ 10 files changed, 146 insertions(+), 45 deletions(-) create mode 100644 python/freetoken/kernel/csrc/include/freetoken/pinned.h diff --git a/docs/cli.md b/docs/cli.md index 138683b011..feda4f7a34 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -71,6 +71,7 @@ ft serve --model ... --gpu GPU-9e8d7c6b # the same card by UUID (a unique prefi |---|---|---| | `--memory-ratio` | 0.9 | Fraction of free VRAM the engine may use (weights + MoE cache + KV) | | `--num-pages` / `--num-tokens` | auto | KV capacity override in pages / tokens (mutually exclusive; auto sizes from VRAM left after weights and MoE cache) | +| `--embed-weights` | gpu | Where the input token embedding table lives. `host` keeps it in pinned host RAM and each lookup reads its rows over PCIe, so its VRAM goes to the expert cache and KV. Refused when the table is tied to the LM head | | `--page-size` | 1 | KV page size; DSV4 forces 128, the TRTLLM backend needs 16/32/64, SWA models require 1 | | `--cache-type` | radix | `radix` (prefix reuse; SWA/GDN-aware variants picked automatically) or `naive` | | `--attention-backend`, `--attn` | auto | `trtllm`/`fi`/`fa`/`triton`/`dsv4_sparse`/`dsa`; `prefill,decode` pair allowed; auto picks per model + GPU | diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 23afd35735..77aaa3eefe 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -34,6 +34,8 @@ class EngineConfig: quant_backend: str | None = None # PLE table backend: "disk" (default) reads rows from the checkpoint files per fill, "pinned" preloads the table into page-locked host RAM. ple_backend: str = "disk" + # --embed-weights: "gpu" (default) keeps the input embedding tables in VRAM; "host" keeps the untied ones in pinned host RAM, read zero-copy by the lookup. + embed_weights: str = "gpu" # DeepSeek-V4.1 Decoder SWA Bounded Replay: "bounded" (default) runs the decoder layers on each # prompt's last window; "exact" runs them on every token. swa_decoder_replay: str = "bounded" diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 7982095ac4..1de67b471a 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -15,7 +15,7 @@ from freetoken.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info from freetoken.gpu_select import gpu_identity from freetoken.checkpoint.ftw import ftw_tensor_names, is_ftw_checkpoint -from freetoken.layers import OffloadMoELayer, iter_moe_layers, set_rope_device +from freetoken.layers import BaseOP, OffloadMoELayer, ParallelLMHead, VocabParallelEmbedding, iter_moe_layers, set_rope_device from freetoken.layers.quantization import LayerKind, QuantBackend, QuantKind, finalize_quant, set_quant_backend from freetoken.moe.offload_cache import iter_offload_moe_layers from freetoken.mm.config import ENCODER_SECTIONS @@ -285,10 +285,67 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att ) +def _untied_embeddings(model: BaseOP) -> list[VocabParallelEmbedding]: + """Embedding tables under ``model`` that no LM head reads as its weight.""" + found: list[VocabParallelEmbedding] = [] + tied: set[int] = set() + seen: set[int] = set() + + def walk(op: BaseOP) -> None: + if id(op) in seen: + return + seen.add(id(op)) + if isinstance(op, ParallelLMHead): + if op.tied_embedding is not None: + tied.add(id(op.tied_embedding)) + return + if isinstance(op, VocabParallelEmbedding): + found.append(op) + for value in vars(op).values(): + if isinstance(value, BaseOP): + walk(value) + elif isinstance(value, (list, tuple)): + for item in value: + if isinstance(item, BaseOP): + walk(item) + + walk(model) + return [op for op in found if id(op) not in tied] + + +def _host_embedding_keys(model: BaseOP, embed_weights: str) -> frozenset[str]: + """State-dict keys of the tables --embed-weights host keeps in pinned host RAM.""" + if embed_weights != "host": + return frozenset() + tables = {id(op.weight) for op in _untied_embeddings(model)} + keys = frozenset(key for key, t in model.state_dict().items() if id(t) in tables) + if not keys: + raise ValueError( + "--embed-weights host: this model's embedding table is tied to its LM head, which reads the whole table " + "every step; serve it with --embed-weights gpu" + ) + return keys + + +def _pinned_copy(key: str, t: torch.Tensor) -> torch.Tensor: + from freetoken.kernel.pinned import alloc_pinned_tensor + + try: + host = alloc_pinned_tensor(*t.shape, dtype=t.dtype) + except RuntimeError as exc: + raise PinFailed( + f"cudaHostAlloc refused {t.nbytes / 2**20:.0f} MiB for {key} (--embed-weights host); " + "free host RAM or serve with --embed-weights gpu" + ) from exc + host.copy_(t) + return host + + def _make_dummy_weight_state_dict( model_state: Dict[str, torch.Tensor], *, device: torch.device, + host_keys: frozenset[str] = frozenset(), ) -> Dict[str, torch.Tensor]: state_dict: Dict[str, torch.Tensor] = {} fp8_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2) @@ -311,6 +368,8 @@ def _make_dummy_weight_state_dict( state_dict[key] = torch.full(param.shape, 127, dtype=param.dtype, device=device) else: state_dict[key] = torch.zeros(param.shape, dtype=param.dtype, device=device) + if key in host_keys: + state_dict[key] = _pinned_copy(key, state_dict[key]) return state_dict @@ -344,6 +403,7 @@ def _materialize_loaded_weight_state_dict( weights: Iterable[Tuple[str, torch.Tensor]], *, device: torch.device, + host_keys: frozenset[str] = frozenset(), ) -> Dict[str, torch.Tensor]: state_dict: Dict[str, torch.Tensor] = {} for key, weight in weights: @@ -353,6 +413,8 @@ def _materialize_loaded_weight_state_dict( if key.endswith(".input_scale"): continue state_dict[key] = weight.to(device=device) + elif key in host_keys: + state_dict[key] = _pinned_copy(key, weight.to(dtype=expected.dtype)) else: state_dict[key] = weight.to(device=device, dtype=expected.dtype) return state_dict @@ -418,13 +480,16 @@ def __init__(self, config: EngineConfig): self._post_weights_free = post_weights_free self.moe_offload_cache = None self.cpu_moe_executor = None - # Host-side auxiliary stores (qwen4_exp's pinned PLE table): after the weights so a + # Host-side auxiliary stores (--embed-weights host tables, qwen4_exp's pinned PLE table): after the weights so a # load failure is not masked, before the MoE offload cache so the bank residency # planning sees the pin quota the table already spent. - self._host_tables_bytes = 0 + host_embeds = [op for op in _untied_embeddings(self.model) if not op.weight.is_cuda] + self._host_tables_bytes = sum(op.weight.nbytes for op in host_embeds) + if host_embeds: + logger.info_rank0(f"Embedding tables in pinned host RAM (--embed-weights host): {mem_GB(self._host_tables_bytes)}") if hasattr(self.model, "load_host_tables"): with _weight_load_context(): - self._host_tables_bytes = int(self.model.load_host_tables(config) or 0) + self._host_tables_bytes += int(self.model.load_host_tables(config) or 0) if is_offload_moe_strategy(config.moe_strategy): self._init_offload_moe_cache(config) if hasattr(self.model, "prepare_for_runtime"): @@ -574,16 +639,17 @@ def _load_weights(self, config: EngineConfig) -> None: resident = [layer for layer in iter_moe_layers(self.model) if not isinstance(layer, OffloadMoELayer)] if resident and not config.use_dummy_weight and is_ftw_checkpoint(config.model_path) and not ftw_tensor_names(config.model_path, "experts_bank"): raise ValueError(f"{config.model_path} keeps its routed experts as dense weights (an older `ft checkpoint --moe-backend fused`); reconvert it with `ft checkpoint`") + host_keys = _host_embedding_keys(self.model, config.embed_weights) with _weight_load_context(): - self.model.load_state_dict(self._load_weight_state_dict(config)) + self.model.load_state_dict(self._load_weight_state_dict(config, host_keys)) finalize_quant(self.model) if resident: self._load_resident_experts(config, resident) - def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tensor]: + def _load_weight_state_dict(self, config: EngineConfig, host_keys: frozenset[str]) -> Dict[str, torch.Tensor]: model_state = self.model.state_dict() if config.use_dummy_weight: - return _make_dummy_weight_state_dict(model_state, device=self.device) + return _make_dummy_weight_state_dict(model_state, device=self.device, host_keys=host_keys) # _materialize casts each loaded tensor to its model-param dtype (model_state), so # models declaring per-tensor dtypes (e.g. DSV4's mixed fp8/fp32/bf16) are preserved; # MoE models exclude routed experts (they load as expert banks, not dense weights). @@ -596,6 +662,7 @@ def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tenso include_vision=bool(config.active_encoders), ), device=self.device, + host_keys=host_keys, ) def _load_resident_experts(self, config: EngineConfig, layers) -> None: diff --git a/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu index 170cbcac71..444b154c61 100644 --- a/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu +++ b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu @@ -10,6 +10,8 @@ #include "mmq.cuh" #include "mmvq.cuh" #include "quantize.cuh" +// after the ggml headers, so their HIP macros win over hip_compat.h +#include #include #include @@ -436,7 +438,8 @@ void op_dequantize(TensorView w, int64_t type_id, int64_t k, TensorView out) { void op_get_rows(TensorView w, int64_t type_id, int64_t k, TensorView ids, TensorView out) { RuntimeCheck(is_dtype(ids.dtype(), kDLInt, 32) && ids.ndim() == 1 && ids.IsContiguous(), "ids must be a contiguous int32 vector"); RuntimeCheck(w.ndim() == 2, "get_rows takes w[rows, row_bytes]"); - const DeviceGuard guard(w.device()); + // w may be a pinned host table read zero-copy, so the launch device comes from out + const DeviceGuard guard(out.device()); const ggml_type type = (ggml_type) type_id; check_weight(w, type, k); const int64_t n = ids.size(0); @@ -446,8 +449,8 @@ void op_get_rows(TensorView w, int64_t type_id, int64_t k, TensorView ids, Tenso return; } const size_t nb01 = w.stride(0), esz = otype == GGML_TYPE_F32 ? 4 : 2; - get_rows_cuda(w.data_ptr(), type, (const int32_t *) ids.data_ptr(), out.data_ptr(), otype, k, nb01, nb01 * w.size(0), nb01 * w.size(0), - n, 1, 1, 4, 4 * n, 4 * n, k * esz, k * esz * n, k * esz * n, current_stream(w.device())); + get_rows_cuda(host::device_alias(w.data_ptr(), w.device()), type, (const int32_t *) ids.data_ptr(), out.data_ptr(), otype, k, nb01, nb01 * w.size(0), nb01 * w.size(0), + n, 1, 1, 4, 4 * n, 4 * n, k * esz, k * esz * n, k * esz * n, current_stream(out.device())); } bool op_uses_mmvq(int64_t type_id, int64_t m) { return ggml_cuda_should_use_mmvq((ggml_type) type_id, device_cc(), m); } diff --git a/python/freetoken/kernel/csrc/include/freetoken/pinned.h b/python/freetoken/kernel/csrc/include/freetoken/pinned.h new file mode 100644 index 0000000000..7c744de4a2 --- /dev/null +++ b/python/freetoken/kernel/csrc/include/freetoken/pinned.h @@ -0,0 +1,40 @@ +#pragma once + +#include +#include + +#include + +namespace host { + +// Pinned host memory is GPU-dereferenceable at its host VA only where UVA identity +// holds (Linux). On Windows/WDDM, cudaHostRegister'd memory maps to a different device +// address, so host-resident tensors are translated where their pointer enters kernel params. +// Cached once per process: FreeToken pins one CUDA device per process (set at engine launch). +inline bool host_ptr_identity() { + static const bool identity = [] { + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess) { + return false; // fail closed: translate (and surface errors), don't assume identity + } + int uva = 0, reg = 0; + cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device); + cudaDeviceGetAttribute(®, cudaDevAttrCanUseHostPointerForRegisteredMem, device); + return uva == 1 && reg == 1; + }(); + return identity; +} + +inline void* device_alias(void* ptr, DLDevice dev) { + if (dev.device_type == kDLCUDA || dev.device_type == kDLROCM || host_ptr_identity()) { + return ptr; + } + void* mapped = nullptr; + const auto err = cudaHostGetDevicePointer(&mapped, ptr, 0); + RuntimeCheck(err == cudaSuccess, + "host tensor must be pinned+mapped (cudaHostGetDevicePointer: ", + cudaGetErrorString(err), ")"); + return mapped; +} + +} // namespace host diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index 279b15ae75..053f0b5512 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -1,3 +1,4 @@ +#include #include #include #include @@ -162,37 +163,6 @@ __always_inline __device__ void store_vec(void* __restrict__ dst, const Tp& vec) -// Pinned host memory is GPU-dereferenceable at its host VA only where UVA identity -// holds (Linux). On Windows/WDDM, cudaHostRegister'd memory maps to a different device -// address, so host-resident tensors are translated here -- the one point their pointer -// enters kernel params. Cached once per process: FreeToken pins one CUDA device per -// process (set at engine launch). -inline bool host_ptr_identity() { - static const bool identity = [] { - int device = 0; - if (cudaGetDevice(&device) != cudaSuccess) { - return false; // fail closed: translate (and surface errors), don't assume identity - } - int uva = 0, reg = 0; - cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device); - cudaDeviceGetAttribute(®, cudaDevAttrCanUseHostPointerForRegisteredMem, device); - return uva == 1 && reg == 1; - }(); - return identity; -} - -inline void* device_alias(void* ptr, DLDevice dev) { - if (dev.device_type == kDLCUDA || dev.device_type == kDLROCM || host_ptr_identity()) { - return ptr; - } - void* mapped = nullptr; - const auto err = cudaHostGetDevicePointer(&mapped, ptr, 0); - host::RuntimeCheck(err == cudaSuccess, - "fast_index_copy: host tensor must be pinned+mapped (cudaHostGetDevicePointer: ", - cudaGetErrorString(err), ")"); - return mapped; -} - struct IndexKernelParams { void* __restrict__ dst; const void* __restrict__ indices_dst; diff --git a/python/freetoken/kernel/csrc/jit/index.cu b/python/freetoken/kernel/csrc/jit/index.cu index 4832e0a460..f197ce0874 100644 --- a/python/freetoken/kernel/csrc/jit/index.cu +++ b/python/freetoken/kernel/csrc/jit/index.cu @@ -1,3 +1,4 @@ +#include #include #include #include @@ -114,9 +115,10 @@ struct IndexKernel { auto weights_dtype_ = SymbolicDType{}; auto indices_dtype_ = SymbolicDType{}; + // the table may live in pinned host memory, read zero-copy TensorMatcher({-1, D}) // .with_dtype(weights_dtype_) - .with_device(device_) + .with_device() .verify(weights); TensorMatcher({L, D}) // .with_dtype(weights_dtype_) @@ -128,6 +130,10 @@ struct IndexKernel { .verify(indices); const auto device = device_.unwrap(); + const auto weights_device = weights.device(); + RuntimeCheck(weights_device.device_type != device.device_type || + weights_device.device_id == device.device_id, + "IndexKernel: weights are on another GPU than the output."); const auto use_int32 = indices_dtype_.unwrap().bits == 32; const auto num_indices = L.unwrap(); const auto entry_size = dtype_bytes(weights_dtype_.unwrap()) * D.unwrap(); @@ -144,7 +150,8 @@ struct IndexKernel { const auto num_blocks = div_ceil(num_warps, kWarpPerBlock); const auto params = IndexKernelParams{ .output = static_cast(output.data_ptr()), - .weight = static_cast(weights.data_ptr()), + .weight = static_cast( + device_alias(weights.data_ptr(), weights_device)), .indice = indices.data_ptr(), .num_warps = num_warps, }; diff --git a/python/freetoken/kernel/ggml.py b/python/freetoken/kernel/ggml.py index 7a2765c5c0..589bc45042 100644 --- a/python/freetoken/kernel/ggml.py +++ b/python/freetoken/kernel/ggml.py @@ -149,7 +149,7 @@ def dequantize(w: torch.Tensor, ggml_type: int, k: int, dtype: torch.dtype = tor def get_rows(w: torch.Tensor, ggml_type: int, k: int, ids: torch.Tensor, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: - out = torch.empty(ids.shape[0], k, dtype=dtype, device=w.device) + out = torch.empty(ids.shape[0], k, dtype=dtype, device=ids.device) _module().get_rows(w, ggml_type, k, ids, out) return out diff --git a/python/freetoken/kernel/index.py b/python/freetoken/kernel/index.py index 95e61ff47d..cbc29e2853 100644 --- a/python/freetoken/kernel/index.py +++ b/python/freetoken/kernel/index.py @@ -46,7 +46,7 @@ def indexing( vocab_range: Tuple[int, int] | None = None, # (start, length) ) -> torch.Tensor: if output is None: - output = weights.new_empty(indices.shape[0], weights.shape[1]) + output = weights.new_empty(indices.shape[0], weights.shape[1], device=indices.device) element_size = weights.shape[1] * weights.element_size() module = _jit_index_module(element_size, num_splits=num_splits_for(element_size)) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index efe3b8c0bd..b3eea9f8a0 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -689,6 +689,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--embed-weights", + default=ServerArgs.embed_weights, + choices=["gpu", "host"], + help=( + "Where the input token embedding table lives. 'gpu' (default) keeps it in VRAM; 'host' keeps it in " + "pinned host RAM and each lookup reads its rows over PCIe, so its VRAM goes to the expert cache and KV. " + "Not for a table tied to the LM head." + ), + ) + parser.add_argument( "--swa-decoder-replay", default=ServerArgs.swa_decoder_replay, From 1d9a79c8498b227233419369350f7f21539665cd Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Fri, 23 Oct 2026 05:24:25 +0000 Subject: [PATCH 2/2] refactor: check the tie from the model config and share the op walk --- python/freetoken/engine/engine.py | 49 +++++-------------- python/freetoken/layers/__init__.py | 3 +- python/freetoken/layers/base.py | 22 ++++++++- .../freetoken/layers/quantization/method.py | 19 +------ 4 files changed, 38 insertions(+), 55 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 1de67b471a..2de7f47a3e 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -15,7 +15,7 @@ from freetoken.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info from freetoken.gpu_select import gpu_identity from freetoken.checkpoint.ftw import ftw_tensor_names, is_ftw_checkpoint -from freetoken.layers import BaseOP, OffloadMoELayer, ParallelLMHead, VocabParallelEmbedding, iter_moe_layers, set_rope_device +from freetoken.layers import OffloadMoELayer, VocabParallelEmbedding, iter_moe_layers, iter_ops, set_rope_device from freetoken.layers.quantization import LayerKind, QuantBackend, QuantKind, finalize_quant, set_quant_backend from freetoken.moe.offload_cache import iter_offload_moe_layers from freetoken.mm.config import ENCODER_SECTIONS @@ -285,45 +285,22 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att ) -def _untied_embeddings(model: BaseOP) -> list[VocabParallelEmbedding]: - """Embedding tables under ``model`` that no LM head reads as its weight.""" - found: list[VocabParallelEmbedding] = [] - tied: set[int] = set() - seen: set[int] = set() - - def walk(op: BaseOP) -> None: - if id(op) in seen: - return - seen.add(id(op)) - if isinstance(op, ParallelLMHead): - if op.tied_embedding is not None: - tied.add(id(op.tied_embedding)) - return - if isinstance(op, VocabParallelEmbedding): - found.append(op) - for value in vars(op).values(): - if isinstance(value, BaseOP): - walk(value) - elif isinstance(value, (list, tuple)): - for item in value: - if isinstance(item, BaseOP): - walk(item) - - walk(model) - return [op for op in found if id(op) not in tied] - - -def _host_embedding_keys(model: BaseOP, embed_weights: str) -> frozenset[str]: +def _host_embedding_keys(model, config: EngineConfig) -> frozenset[str]: """State-dict keys of the tables --embed-weights host keeps in pinned host RAM.""" - if embed_weights != "host": + if config.embed_weights != "host": return frozenset() - tables = {id(op.weight) for op in _untied_embeddings(model)} - keys = frozenset(key for key, t in model.state_dict().items() if id(t) in tables) - if not keys: + if config.model_config.tie_word_embeddings: raise ValueError( "--embed-weights host: this model's embedding table is tied to its LM head, which reads the whole table " "every step; serve it with --embed-weights gpu" ) + tables = {id(op.weight) for op in iter_ops(model) if type(op) is VocabParallelEmbedding} + keys = frozenset(key for key, t in model.state_dict().items() if id(t) in tables) + if not keys: + raise ValueError( + f"--embed-weights host: {type(model).__name__} has no plain embedding table to move (a DF11 table stays " + "on the GPU); serve it with --embed-weights gpu" + ) return keys @@ -483,7 +460,7 @@ def __init__(self, config: EngineConfig): # Host-side auxiliary stores (--embed-weights host tables, qwen4_exp's pinned PLE table): after the weights so a # load failure is not masked, before the MoE offload cache so the bank residency # planning sees the pin quota the table already spent. - host_embeds = [op for op in _untied_embeddings(self.model) if not op.weight.is_cuda] + host_embeds = [op for op in iter_ops(self.model) if type(op) is VocabParallelEmbedding and not op.weight.is_cuda] self._host_tables_bytes = sum(op.weight.nbytes for op in host_embeds) if host_embeds: logger.info_rank0(f"Embedding tables in pinned host RAM (--embed-weights host): {mem_GB(self._host_tables_bytes)}") @@ -639,7 +616,7 @@ def _load_weights(self, config: EngineConfig) -> None: resident = [layer for layer in iter_moe_layers(self.model) if not isinstance(layer, OffloadMoELayer)] if resident and not config.use_dummy_weight and is_ftw_checkpoint(config.model_path) and not ftw_tensor_names(config.model_path, "experts_bank"): raise ValueError(f"{config.model_path} keeps its routed experts as dense weights (an older `ft checkpoint --moe-backend fused`); reconvert it with `ft checkpoint`") - host_keys = _host_embedding_keys(self.model, config.embed_weights) + host_keys = _host_embedding_keys(self.model, config) with _weight_load_context(): self.model.load_state_dict(self._load_weight_state_dict(config, host_keys)) finalize_quant(self.model) diff --git a/python/freetoken/layers/__init__.py b/python/freetoken/layers/__init__.py index a07484620e..590b58abbe 100644 --- a/python/freetoken/layers/__init__.py +++ b/python/freetoken/layers/__init__.py @@ -7,7 +7,7 @@ swiglu_clamp_and_mul, swigluoai_and_mul, ) -from .base import BaseOP, OPList, StateLessOP +from .base import BaseOP, OPList, StateLessOP, iter_ops from .embedding import ParallelLMHead, VocabParallelEmbedding from .linear import ( LinearColParallelMerged, @@ -39,6 +39,7 @@ "BaseOP", "StateLessOP", "OPList", + "iter_ops", "VocabParallelEmbedding", "ParallelLMHead", "LinearColParallelMerged", diff --git a/python/freetoken/layers/base.py b/python/freetoken/layers/base.py index b1e939d98b..a05f6636d4 100644 --- a/python/freetoken/layers/base.py +++ b/python/freetoken/layers/base.py @@ -1,7 +1,7 @@ from __future__ import annotations from abc import abstractmethod -from typing import Any, Dict, Generic, List, TypeAlias, TypeVar +from typing import Any, Dict, Generic, Iterator, List, TypeAlias, TypeVar import torch @@ -97,3 +97,23 @@ def load_state_dict( if not _internal and state_dict: raise RuntimeError(f"Unexpected keys in state_dict: {list(state_dict.keys())}") + + +def iter_ops(root: BaseOP) -> Iterator[BaseOP]: + """Each op under ``root`` once, ``root`` first; an op's attributes are read after it is yielded.""" + seen: set[int] = set() + + def walk(op: BaseOP) -> Iterator[BaseOP]: + if id(op) in seen: + return + seen.add(id(op)) + yield op + for value in vars(op).values(): + if isinstance(value, BaseOP): + yield from walk(value) + elif isinstance(value, (list, tuple)): + for item in value: + if isinstance(item, BaseOP): + yield from walk(item) + + return walk(root) diff --git a/python/freetoken/layers/quantization/method.py b/python/freetoken/layers/quantization/method.py index 2a079539f1..9e90637577 100644 --- a/python/freetoken/layers/quantization/method.py +++ b/python/freetoken/layers/quantization/method.py @@ -63,27 +63,12 @@ def scheme(self): def finalize_quant(root: Any) -> int: """Call ``quant_method.finalize(layer)`` on every layer under ``root``; returns the count.""" - from freetoken.layers.base import BaseOP + from freetoken.layers.base import iter_ops - seen: set[int] = set() count = 0 - - def walk(op: Any) -> None: - nonlocal count - if id(op) in seen: - return - seen.add(id(op)) + for op in iter_ops(root): method = getattr(op, "quant_method", None) if method is not None: method.finalize(op) count += 1 - for value in vars(op).values(): - if isinstance(value, BaseOP): - walk(value) - elif isinstance(value, (list, tuple)): - for item in value: - if isinstance(item, BaseOP): - walk(item) - - walk(root) return count