diff --git a/modelopt/torch/speculative/plugins/modeling_final_norm.py b/modelopt/torch/speculative/plugins/modeling_final_norm.py index 718d591b662..973a04b3763 100644 --- a/modelopt/torch/speculative/plugins/modeling_final_norm.py +++ b/modelopt/torch/speculative/plugins/modeling_final_norm.py @@ -90,6 +90,10 @@ def extra_repr(self): "deepseek_v3": "rmsnorm", "kimi_k2": "rmsnorm", # Kimi-K2 / K2-Thinking (DeepSeek-V3 arch) report model_type "kimi_k2" "kimi_k25": "rmsnorm", # Kimi-K2.5 / K2.6 / K2.7 all report model_type "kimi_k25" + # Kimi-K3 is a VLM: its OUTER config reports model_type "kimi_k3", but the text + # backbone this table is keyed on reports "kimi_linear". Keying on "kimi_k3" would + # silently miss and drop the final norm. + "kimi_linear": "rmsnorm", # M3's final norm is always gemma-style; map it here too so a config that lost its # use_gemma_norm flag still gets the correct flavor instead of silently dropping the +1. "minimax_m3_vl_text": "gemma_rmsnorm", diff --git a/modelopt/torch/utils/loss_mask.py b/modelopt/torch/utils/loss_mask.py index 93ade59577e..4ff5e5dc1d4 100644 --- a/modelopt/torch/utils/loss_mask.py +++ b/modelopt/torch/utils/loss_mask.py @@ -79,20 +79,34 @@ def get_loss_mask_recovery(tokenizer) -> LossMaskRecovery | None: _KIMI_ROLE_MARKERS = ("<|im_user|>", "<|im_assistant|>", "<|im_system|>") +# Kimi-K3's XTML structural tags (see the Kimi-K3 section below). Declared here +# because ``_kimi_detect`` uses them to hand K3 tokenizers off to the K3 recovery. +_K3_MARKERS = ("<|open|>", "<|close|>", "<|sep|>", "<|end_of_msg|>") -def _kimi_detect(tokenizer) -> bool: - """Whether ``tokenizer`` defines Kimi's chat role markers as real tokens.""" + +def _has_all_tokens(tokenizer, tokens) -> bool: + """Whether ``tokenizer`` maps every one of ``tokens`` to a real (non-unk) id.""" unk = getattr(tokenizer, "unk_token_id", None) try: - ids = [ - tokenizer.convert_tokens_to_ids(t) - for t in (*_KIMI_ROLE_MARKERS, "<|im_middle|>", "<|im_end|>") - ] + ids = [tokenizer.convert_tokens_to_ids(t) for t in tokens] except Exception: return False return all(i is not None and i != unk for i in ids) +def _kimi_detect(tokenizer) -> bool: + """Whether ``tokenizer`` defines Kimi's chat role markers as real tokens. + + K3 keeps the K2 ``<|im_*|>`` markers for back-compat but *also* defines the XTML + structural tags that classic Kimi lacks. When those are present the tokenizer is + K3, so defer to the ``kimi_k3`` recovery whose ``compute`` understands the XTML + turn layout; matching here would silently produce an empty mask. + """ + if not _has_all_tokens(tokenizer, (*_KIMI_ROLE_MARKERS, "<|im_middle|>", "<|im_end|>")): + return False + return not _has_all_tokens(tokenizer, _K3_MARKERS) + + def _kimi_compute(tokenizer, input_ids) -> torch.Tensor: """Recover the assistant-content mask from already-tokenized Kimi chat ids. @@ -136,3 +150,77 @@ def _kimi_compute(tokenizer, input_ids) -> torch.Tensor: register_loss_mask_recovery( LossMaskRecovery(name="kimi", detect=_kimi_detect, compute=_kimi_compute) ) + + +# --------------------------------------------------------------------------- +# Kimi-K3 +# +# K3 replaces the K2/K2.5 <|im_*|> turn markers with an XTML tag format: +# <|open|> {tag} {attrs..} <|sep|> {content} <|close|> {tag} <|sep|> [<|end_of_msg|>] +# Tag names and attribute values (including the role) are ORDINARY text tokens; +# only open/close/sep/end_of_msg are special tokens. An assistant turn reads +# <|open|> message role assistant <|sep|> {content} <|close|> message <|sep|> +# and its content may nest further tags (``think``, ``response``), so the content +# span is found by tracking open/close depth rather than by scanning for the next +# marker. +# --------------------------------------------------------------------------- + + +def _k3_detect(tokenizer) -> bool: + """Whether ``tokenizer`` defines K3's XTML structural markers as real tokens.""" + return _has_all_tokens(tokenizer, _K3_MARKERS) + + +def _k3_compute(tokenizer, input_ids) -> torch.Tensor: + """Recover the assistant-content mask from already-tokenized K3 chat ids. + + Marks the content of each ``message role=assistant`` turn -- from the token after + the message-open ``<|sep|>`` up to (excluding) the matching ``<|close|>`` -- which + includes the nested ``think``/``response`` sub-tags the model generates. Role + headers, other roles, and the turn-closing tokens stay unmasked, matching the + ``{% generation %}`` span a fast tokenizer would report. + """ + ids = input_ids.tolist() if hasattr(input_ids, "tolist") else list(input_ids) + open_id = tokenizer.convert_tokens_to_ids("<|open|>") + close_id = tokenizer.convert_tokens_to_ids("<|close|>") + sep_id = tokenizer.convert_tokens_to_ids("<|sep|>") + + n = len(ids) + mask = [0] * n + i = 0 + while i < n: + if ids[i] != open_id: + i += 1 + continue + # The header tokens sit between this <|open|> and its <|sep|>. + j = i + 1 + while j < n and ids[j] != sep_id: + j += 1 + if j >= n: + break + header = tokenizer.decode(ids[i + 1 : j]).lower() + if "message" not in header or "assistant" not in header: + i = j + 1 + continue + # Content = [after this <|sep|>, matching <|close|>), skipping nested tags. + start = j + 1 + depth = 1 + k = start + while k < n: + if ids[k] == open_id: + depth += 1 + elif ids[k] == close_id: + depth -= 1 + if depth == 0: + break + k += 1 + for t in range(start, k): + mask[t] = 1 + i = k + 1 + + return torch.tensor(mask, dtype=torch.long) + + +register_loss_mask_recovery( + LossMaskRecovery(name="kimi_k3", detect=_k3_detect, compute=_k3_compute) +) diff --git a/tests/unit/torch/speculative/plugins/test_modeling_final_norm.py b/tests/unit/torch/speculative/plugins/test_modeling_final_norm.py index 882e1c85416..4a18648599a 100644 --- a/tests/unit/torch/speculative/plugins/test_modeling_final_norm.py +++ b/tests/unit/torch/speculative/plugins/test_modeling_final_norm.py @@ -37,6 +37,11 @@ ("deepseek_v3", "rmsnorm"), ("kimi_k2", "rmsnorm"), ("kimi_k25", "rmsnorm"), + # Kimi-K3's text backbone reports "kimi_linear"; the outer VLM config's + # "kimi_k3" is deliberately NOT a key, since this table is keyed on the + # resolved text config. + ("kimi_linear", "rmsnorm"), + ("kimi_k3", None), # gpt_oss is intentionally DISABLED: its RMSNorm forward/weight dtype differs from the # Llama variant we reuse, so it must not resolve to a norm type until a matching class # is added. diff --git a/tests/unit/torch/utils/test_loss_mask.py b/tests/unit/torch/utils/test_loss_mask.py index a98edec7896..ec619e1303e 100644 --- a/tests/unit/torch/utils/test_loss_mask.py +++ b/tests/unit/torch/utils/test_loss_mask.py @@ -142,6 +142,106 @@ def test_kimi_mask_accepts_list_input(): assert mask.tolist() == [0, 0, 0, 1, 1, 0] +# --------------------------------------------------------------------------- +# Kimi-K3 (XTML chat format) +# --------------------------------------------------------------------------- + +_K3_MARKER_IDS = {"<|open|>": 10, "<|close|>": 11, "<|sep|>": 12, "<|end_of_msg|>": 13} +# Header/content tokens are ORDINARY text tokens in K3, so the recovery decodes them. +_K3_WORDS = {100: "message", 101: "role", 102: "assistant", 103: "user", 104: "think"} + + +class FakeK3Tokenizer: + """K3 tokenizer: XTML structural markers, plain-text tag names.""" + + is_fast = False + unk_token_id = 999 + + def convert_tokens_to_ids(self, token): + return _K3_MARKER_IDS.get(token, self.unk_token_id) + + def decode(self, ids): + return " ".join(_K3_WORDS.get(i, "") for i in ids) + + +class FakeK3WithLegacyMarkersTokenizer(FakeK3Tokenizer): + """Real K3 keeps the K2 ``<|im_*|>`` markers too; K3 must still win.""" + + def convert_tokens_to_ids(self, token): + if token in _K3_MARKER_IDS: + return _K3_MARKER_IDS[token] + return _MARKER_IDS.get(token, self.unk_token_id) + + +def _k3_turn(role_id, content_ids): + # <|open|> message role {role} <|sep|> {content} <|close|> message <|sep|> + return [ + _K3_MARKER_IDS["<|open|>"], + 100, # "message" + 101, # "role" + role_id, + _K3_MARKER_IDS["<|sep|>"], + *content_ids, + _K3_MARKER_IDS["<|close|>"], + 100, + _K3_MARKER_IDS["<|sep|>"], + ] + + +def test_k3_recovery_is_registered(): + recovery = get_loss_mask_recovery(FakeK3Tokenizer()) + assert recovery is not None + assert recovery.name == "kimi_k3" + + +def test_k3_tokenizer_does_not_match_the_k2_recovery(): + """K3 defines both marker sets; the K2 recovery must defer or the mask is empty.""" + recovery = get_loss_mask_recovery(FakeK3WithLegacyMarkersTokenizer()) + assert recovery is not None + assert recovery.name == "kimi_k3" + + +def test_k3_mask_marks_only_assistant_content(): + tok = FakeK3Tokenizer() + ids = _k3_turn(103, [200]) + _k3_turn(102, [300, 301]) # user turn, assistant turn + mask = get_loss_mask_recovery(tok).compute(tok, torch.tensor(ids)) + + marked = {i for i, v in enumerate(mask.tolist()) if v == 1} + assert marked == {i for i, v in enumerate(ids) if v in (300, 301)} + + +def test_k3_mask_includes_nested_tags(): + """Assistant content nests ``think``/``response`` sub-tags; they belong to the span.""" + tok = FakeK3Tokenizer() + nested = [ + 200, + _K3_MARKER_IDS["<|open|>"], + 104, # "think" + _K3_MARKER_IDS["<|sep|>"], + 201, + _K3_MARKER_IDS["<|close|>"], + 104, + _K3_MARKER_IDS["<|sep|>"], + 202, + ] + ids = _k3_turn(102, nested) + mask = get_loss_mask_recovery(tok).compute(tok, torch.tensor(ids)).tolist() + + # Content spans everything between the header <|sep|> and the matching <|close|>. + start = 5 + end = len(ids) - 3 + assert mask[:start] == [0] * start + assert mask[start:end] == [1] * (end - start) + assert mask[end:] == [0] * (len(ids) - end) + + +def test_k3_mask_accepts_list_input(): + tok = FakeK3Tokenizer() + ids = _k3_turn(102, [300, 301]) + mask = get_loss_mask_recovery(tok).compute(tok, ids) # plain list, not a tensor + assert mask.tolist() == [0, 0, 0, 0, 0, 1, 1, 0, 0, 0] + + def test_register_and_lookup_custom_recovery(restore_registry): sentinel = object() diff --git a/tools/launcher/examples/moonshotai/Kimi-K3/hf_streaming_dspark_multi_node.yaml b/tools/launcher/examples/moonshotai/Kimi-K3/hf_streaming_dspark_multi_node.yaml new file mode 100644 index 00000000000..06268efd847 --- /dev/null +++ b/tools/launcher/examples/moonshotai/Kimi-K3/hf_streaming_dspark_multi_node.yaml @@ -0,0 +1,137 @@ +# DSpark streaming speculative-decoding training for Kimi-K3 (multi-node). +# DSpark = the DFlash backbone + a lightweight Markov head + a confidence head, +# generating a causal block semi-autoregressively; see dspark.yaml for the head +# and loss config. Runs the shared streaming pipeline +# (common/eagle3/train_eagle_streaming.sh) with the K3-specific base, draft dims, +# mask token and sliding window; trained from scratch. A starting point for +# reproduction — tune node counts, batch, steps and serve limits for your cluster. +# +# Kimi-K3 specifics this yaml encodes (each was a silent failure mode): +# * Loss mask. K3 replaces the K2/K2.5 <|im_*|> turn markers with an XTML tag +# format (<|open|> message role assistant <|sep|> ... <|close|>) and ships a +# slow tiktoken tokenizer, so answer_only_loss depends on the `kimi_k3` +# loss-mask recovery. Without it the legacy `kimi` recovery matches (K3 keeps +# the old markers for back-compat) and returns an ALL-ZERO mask. +# * Final norm. K3 is a VLM whose outer config reports model_type "kimi_k3" +# while its text backbone reports "kimi_linear"; the final-norm table is keyed +# on the text config, so the `kimi_linear` entry is what stops FakeBaseModel +# from silently building no norm and reconstructing target logits from an +# un-normed hidden. +# * EAGLE_CAPTURE_IDS ends at 93 == num_hidden_layers, the TRUE final hidden. +# K3 has a block-residual backbone (attn_res_block_size=12) whose output is +# produced after the layer loop by attn_res, so capturing that id requires +# vllm#50815. Without it the "final" capture is prefix_sum + hidden_states, +# which is missing the block-residual combination and skews the KD teacher. +# * The DSpark draft does NOT inherit the base dims, so they are set explicitly +# below. Note these are draft-side choices, not mirrors of the backbone: K3's +# text config carries no rope_theta at all (so the draft's value is a free +# hyperparameter), and its FFN is 33792 dense / 3072 MoE, neither of which the +# draft copies. +# * Serving K3 needs TP8, and this pipeline runs one serve replica per node, so +# the serve nodes must have >= 8 GPUs. +# +# Run ON the cluster login node (paramiko can't reach it through the login proxy): +# export SLURM_HOST=localhost SLURM_ACCOUNT= \ +# SLURM_PARTITION= \ +# SLURM_HF_LOCAL= \ +# SLURM_JOB_DIR= \ +# NEMORUN_HOME=$PWD +# uv run launch.py --yaml examples/moonshotai/Kimi-K3/hf_streaming_dspark_multi_node.yaml \ +# identity=$HOME/.ssh/id_ecdsa detach=True --yes +# +# The export lands in /scratchspace/export. + +job_name: Kimi-K3_DSpark_streaming_multi_node +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: /hf-local/moonshotai/Kimi-K3 + + # Build /scratchspace/data/train.jsonl. Point data.data_path at the full + # Spec-Decoding-Dataset-v2 corpus to reproduce. + task_0: + script: common/eagle3/make_dataset.sh + args: + - -f modules/Model-Optimizer/examples/dataset/example_data_config.yaml + - --full-conversations + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 + + task_1: + script: common/eagle3/train_eagle_streaming.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dspark.yaml + - model.model_name_or_path=<> + - model.use_fake_base_for_offline=true + - model.trust_remote_code=true + - data.mode=streaming + - data.data_path=/scratchspace/data/train.jsonl + - training.output_dir=/scratchspace/dspark + - training.training_seq_len=4096 + - training.disable_tqdm=true + - training.ar_validate_steps=500000 + - training.num_train_epochs=1 + - training.per_device_train_batch_size=4 + - training.gradient_accumulation_steps=1 + - training.save_steps=1000 + - training.logging_steps=20 + - training.learning_rate=1.0e-4 + - training.warmup_steps=2000 + # K3's slow tokenizer can't emit assistant masks; the `kimi_k3` loss-mask + # recovery reconstructs them from token ids (see header). + - training.answer_only_loss=true + # The vLLM serve container has no tensorboard -> trainer init crash. + - training.report_to=none + # Draft dims are NOT inherited from the base — set them explicitly or the + # draft is silently built with the Qwen3Config defaults. + - dflash.dflash_architecture_config.num_hidden_layers=6 + - dflash.dflash_architecture_config.num_key_value_heads=8 + - dflash.dflash_architecture_config.intermediate_size=12288 + # K3's text config exposes no rope_theta, so this is a draft-side choice + # rather than a value inherited from the backbone. + - dflash.dflash_architecture_config.rope_theta=10000 + # Semi-AR generation block (dspark.yaml ships 16; the Kimi backbone uses 8). + - dflash.dflash_block_size=8 + # K3 has no dedicated mask token; 163606 is a reserved slot. + - dflash.dflash_mask_token_id=163606 + # Sliding-window attention over the context; the block itself stays + # bidirectional. Must be >= dflash_block_size. + - dflash.dflash_swa_window_size=1024 + environment: + - HF_MODEL_CKPT: <> + # 6 evenly-spaced aux capture ids + the true final hidden (93 == the base's + # num_hidden_layers). The final id requires vllm#50815 on a block-residual + # backbone; see header. + - EAGLE_CAPTURE_IDS: "[2,19,36,52,69,86,93]" + - SERVE_NODES: "8" + - SERVE_TP: "8" + - STREAMING_NUM_WORKERS: "4" + # K3's custom-modeling base needs trust_remote_code at export and serve. + - EXPORT_EXTRA_ARGS: "--trust_remote_code" + - SERVE_EXTRA_ARGS: "--trust-remote-code" + - SERVE_MAX_MODEL_LEN: "4160" + - SERVE_MAX_NUM_SEQS: "16" + - SERVE_GPU_MEM_UTIL: "0.9" + # K3 is large; allow a long load before the trainer gives up on the serve. + - SERVE_READY_TIMEOUT: "6000" + - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1200" + - VLLM_ENGINE_ITERATION_TIMEOUT_S: "1200" + # RDMA transport is UCX (InfiniBand) by default. On AWS EFA, uncomment: + # - NIXL_BACKENDS: "LIBFABRIC" + # - FI_PROVIDER: "efa" + # - NCCL_IB_DISABLE: "1" + slurm_config: + _factory_: "slurm_factory" + nodes: 9 + ntasks_per_node: 1 + gpus_per_node: 8 + # vLLM build with native Kimi-K3 support (vllm/models/kimi_k3) and the + # final-layer aux-capture fix (vllm#50815). + container: