diff --git a/examples/speculative_decoding/eagle_utils.py b/examples/speculative_decoding/eagle_utils.py index 721b981eaae..5593af45b09 100644 --- a/examples/speculative_decoding/eagle_utils.py +++ b/examples/speculative_decoding/eagle_utils.py @@ -78,6 +78,13 @@ def make_speculative_data_module( if mode == "streaming": # ``train_len`` right-truncates during tokenization and is also the collator's # pad target; caller must ensure ``train_len <= vllm.max_model_len``. + # The streaming dataset tokenizes via ``tokenizer.apply_chat_template`` (no + # chat_template arg), so a custom template (e.g. one carrying {% generation %} + # tags for answer_only_loss) must be installed on the tokenizer here — unlike + # the online path, which threads ``chat_template`` straight into the collator. + if chat_template is not None: + tokenizer.chat_template = chat_template + print_rank_0("Installed custom chat template on tokenizer for streaming.") print_rank_0(f"Streaming hidden states from {data_args.streaming_server_url}") from modelopt.torch.speculative.plugins.hf_streaming_dataset import ( EagleVllmStreamingConfig, diff --git a/tools/launcher/examples/Qwen/Qwen3-30B-A3B/hf_streaming_dflash_multi_node.yaml b/tools/launcher/examples/Qwen/Qwen3-30B-A3B/hf_streaming_dflash_multi_node.yaml new file mode 100644 index 00000000000..daefe6223d9 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-30B-A3B/hf_streaming_dflash_multi_node.yaml @@ -0,0 +1,124 @@ +# DFlash streaming speculative decoding pipeline for Qwen3-30B-A3B — MULTI-NODE. +# +# Same streaming transport / dispatch as hf_streaming_eagle3_multi_node.yaml: task_1 +# splits N nodes into K serve replicas + (N-K) DDP trainers via SERVE_NODES; hidden +# states move serve -> trainer over NIXL RDMA. DFlash just consumes a different set of +# captured layers and trains a block-diffusion draft instead of an autoregressive one. +# See common/eagle3/train_eagle_streaming.sh for dispatch, rendezvous, and sharding. +# +# Qwen3-30B-A3B notes (vs the dense Qwen3-8B example): +# - MoE: 128 experts, 8 active/token, shipped in BF16 (NOT quantized, unlike +# gpt-oss MXFP4). The serve still passes --no-enable-flashinfer-autotune +# unconditionally; it is a no-op cost here since there is no FP4 re-tuning. +# - 48 hidden layers, 5 draft layers -> build_target_layer_ids(48,5)=[1,12,23,34,45] +# (the draft's fc input) plus the final layer for self-logit distillation. vLLM's +# capture ids are those +1 -> [2,13,24,35,46], plus final layer 48. +# - dflash_mask_token_id: Qwen3 has no tokenizer mask token (would fall back to +# None), so set it explicitly to a reserved/unused vocab slot (151669), same as +# the Qwen3-8B DFlash example (shared Qwen3 tokenizer). +# - We use a fake base (lm_head + embed_tokens only): streaming trains only the +# DFlash draft on hidden states from the live serve, so the trainer never needs +# the 30B base transformer layers. Matches the gpt-oss streaming examples. +# +# 3-step pipeline: +# task_0: Build input conversations (jsonl) +# task_1: Streaming train — 2 serve nodes (2 GPU, TP=2) + 2 trainer nodes (2 GPU) +# task_2: vLLM smoke test with DFlash speculative decoding +# +# Usage: +# uv run launch.py --yaml examples/Qwen/Qwen3-30B-A3B/hf_streaming_dflash_multi_node.yaml --yes + +job_name: Qwen3-30B-A3B_DFlash_streaming_multi_node +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: /hf-local/Qwen/Qwen3-30B-A3B + + # Step 1: Build input conversations + 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: 1 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 + + # Step 2: Streaming DFlash training — 2 serve replicas (TP=2) + 2 trainer nodes (2 GPU each). + # DFlash extracts 5 target layers (build_target_layer_ids(48,5)=[1,12,23,34,45], the + # draft's fc input) plus the final layer for self-logit distillation. vLLM's capture + # ids are those +1 -> [2,13,24,35,46], plus final layer 48. + task_1: + script: common/eagle3/train_eagle_streaming.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dflash.yaml + - model.model_name_or_path=<> + # Streaming trains only the DFlash draft on hidden states from the live serve; the + # trainer never runs the base transformer layers, so load a fake base (lm_head + + # embed_tokens only). Consistent with the gpt-oss streaming examples. + - model.use_fake_base_for_offline=true + - data.mode=streaming + - data.data_path=/scratchspace/data/train.jsonl + - training.output_dir=/scratchspace/dflash + - training.training_seq_len=4096 + - training.disable_tqdm=true + # Streaming corpus is prompt-only (the serve generates the response and we + # capture its hidden states), so there is no assistant span to mask -> train + # over the full sequence, same as the EAGLE3 streaming example. + - training.answer_only_loss=false + - training.num_train_epochs=1 + - training.max_steps=2000 + # dflash.yaml sets report_to=tensorboard, which hard-fails if tensorboard + # isn't in the serve container; the streaming trainer doesn't need it. + - training.report_to=none + - dflash.dflash_block_size=16 + - dflash.dflash_num_anchors=512 + - dflash.dflash_loss_decay_factor=7 + # Qwen3 has no tokenizer mask token; use a reserved vocab slot (151669). + - dflash.dflash_mask_token_id=151669 + - dflash.dflash_architecture_config.num_hidden_layers=5 + environment: + - HF_MODEL_CKPT: <> + # No spaces: nemo_run emits `export FOO=value` unquoted. + - EAGLE_CAPTURE_IDS: "[2,13,24,35,46,48]" + - SERVE_TP: "2" + # K serve replica nodes (Slurm nodes 0..K-1); the rest are trainers. + - SERVE_NODES: "2" + # Per-rank in-flight fetches; keep low so the cold MoE serve isn't flooded past its execute-model timeout. + - STREAMING_NUM_WORKERS: "4" + # Qwen3-30B-A3B default ctx is 40960 -> full-len KV cache is wasteful for the + # 30B serve; we train at seq_len 4096, so cap it. + - SERVE_MAX_MODEL_LEN: "8192" + # DFlash uses a custom modeling file; export must trust remote code. + - EXPORT_EXTRA_ARGS: "--trust_remote_code" + slurm_config: + _factory_: "slurm_factory" + nodes: 4 + ntasks_per_node: 1 + gpus_per_node: 2 + container: vllm/vllm-openai:latest + + # Step 3: vLLM smoke test (DFlash, uses exported checkpoint from training) + task_2: + script: common/specdec/vllm_smoke_test.sh + environment: + - HF_MODEL_CKPT: <> + - DRAFT_MODEL: /scratchspace/export + - SPEC_METHOD: "dflash" + - NUM_SPEC_TOKENS: "7" + - MIN_ACCEPTANCE_LENGTH: "1.2" + # Qwen3-30B-A3B is ~60 GB in BF16 -> does not fit one 80 GB GPU; serve the + # smoke-test target with TP=2 (gpus_per_node: 2 below). + - TP_SIZE: "2" + slurm_config: + _factory_: "slurm_factory" + container: "vllm/vllm-openai:nightly" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 diff --git a/tools/launcher/examples/Qwen/Qwen3-30B-A3B/hf_streaming_eagle3_multi_node.yaml b/tools/launcher/examples/Qwen/Qwen3-30B-A3B/hf_streaming_eagle3_multi_node.yaml new file mode 100644 index 00000000000..8a7a7f2ac20 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-30B-A3B/hf_streaming_eagle3_multi_node.yaml @@ -0,0 +1,110 @@ +# EAGLE3 streaming speculative decoding pipeline for Qwen3-30B-A3B — MULTI-NODE. +# +# task_1 splits N nodes into K serve replicas + (N-K) DDP trainers via SERVE_NODES; +# see common/eagle3/train_eagle_streaming.sh for dispatch, rendezvous, and sharding. +# +# Qwen3-30B-A3B notes (vs the dense Qwen3-8B example): +# - MoE: 128 experts, 8 active/token, shipped in BF16 (NOT quantized, unlike +# gpt-oss MXFP4). The serve still passes --no-enable-flashinfer-autotune +# unconditionally; it is a no-op cost here since there is no FP4 re-tuning. +# - 48 hidden layers -> default_eagle_aux_layer_ids(48)=[1,23,44]; vLLM capture ids +# are each +1 plus the final layer 48 -> EAGLE_CAPTURE_IDS="[2,24,45,48]". +# - sliding_window is null (use_sliding_window=false), so the draft has no +# sliding-window leak. We still use a fake base (lm_head + embed_tokens only): +# streaming trains only the EAGLE3 draft on hidden states from the live serve, +# so the trainer never needs the 30B base transformer layers. This matches the +# gpt-oss streaming examples and keeps the draft config off Qwen3MoeConfig. +# - qwen3_moe is native in recent transformers/vLLM (no --trust-remote-code needed). +# +# 3-step pipeline: +# task_0: Build input conversations (jsonl) +# task_1: Streaming train — 2 serve nodes (2 GPU, TP=2) + 2 trainer nodes (2 GPU) +# task_2: Benchmark — evaluate speculative decoding speedup via VLLM +# +# Usage: +# uv run launch.py --yaml examples/Qwen/Qwen3-30B-A3B/hf_streaming_eagle3_multi_node.yaml --yes + +job_name: Qwen3-30B-A3B_EAGLE3_streaming_multi_node +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: /hf-local/Qwen/Qwen3-30B-A3B + + # Step 1: Build input conversations + 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: 1 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 + + # Step 2: Streaming EAGLE3 training — 2 serve replicas (TP=2) + 2 trainer nodes (2 GPU each). + # Capture ids: default_eagle_aux_layer_ids(48)=[1,23,44] +1, plus final layer 48. + task_1: + script: common/eagle3/train_eagle_streaming.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/eagle3.yaml + - model.model_name_or_path=<> + # Streaming trains only the EAGLE3 draft on hidden states from the live serve; the + # trainer never runs the base transformer layers, so load a fake base (lm_head + + # embed_tokens only). Consistent with the gpt-oss streaming examples. + - model.use_fake_base_for_offline=true + - data.mode=streaming + - data.data_path=/scratchspace/data/train.jsonl + - training.output_dir=/scratchspace/eagle3 + - training.training_seq_len=4096 + - training.disable_tqdm=true + - training.ar_validate_steps=500000 + - training.num_train_epochs=1 + - training.max_steps=2000 + - eagle.eagle_use_torch_compile=false + environment: + - HF_MODEL_CKPT: <> + # No spaces: nemo_run emits `export FOO=value` unquoted. + - EAGLE_CAPTURE_IDS: "[2,24,45,48]" + - SERVE_TP: "2" + # K serve replica nodes (Slurm nodes 0..K-1); the rest are trainers. + - SERVE_NODES: "2" + # Per-rank in-flight fetches; keep low so the cold MoE serve isn't flooded past its execute-model timeout (kills EngineCore). + - STREAMING_NUM_WORKERS: "4" + # Qwen3-30B-A3B default ctx is 40960 -> full-len KV cache is wasteful for the + # 30B serve; we train at seq_len 4096, so cap it. + - SERVE_MAX_MODEL_LEN: "8192" + slurm_config: + _factory_: "slurm_factory" + nodes: 4 + ntasks_per_node: 1 + gpus_per_node: 2 + container: vllm/vllm-openai:latest + + # Step 3: Benchmark speculative decoding (VLLM backend) + task_2: + script: common/specdec_bench/quick_check.sh + args: + - --draft_model_dir /scratchspace/export + - --draft_length 3 + - --output_length 4096 + - --engine VLLM + # Qwen3-30B-A3B is ~60 GB in BF16 -> does not fit one 80 GB GPU; serve the + # benchmark target with TP=2 (gpus_per_node: 2 below). + - --tp_size 2 + - --ep_size 1 + - --speculative_algorithm EAGLE3 + - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl + - --concurrency 32 + environment: + - HF_MODEL_CKPT: <> + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 2 + container: vllm/vllm-openai:latest diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/eagle3_quick_check.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/eagle3_quick_check.yaml index dec6f2989f5..7078255eb9c 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/eagle3_quick_check.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/eagle3_quick_check.yaml @@ -111,7 +111,7 @@ pipeline: - --ep_size 1 - --speculative_algorithm EAGLE3 - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl - - --concurrency 1 + - --concurrency 32 environment: - HF_MODEL_CKPT: <> slurm_config: diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_offline_eagle3.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_offline_eagle3.yaml index 24068c4bb3e..f0a99514a10 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_offline_eagle3.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_offline_eagle3.yaml @@ -93,7 +93,7 @@ pipeline: - --ep_size 1 - --speculative_algorithm EAGLE3 - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl - - --concurrency 1 + - --concurrency 32 environment: - HF_MODEL_CKPT: <> slurm_config: diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_offline_eagle3_ptq.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_offline_eagle3_ptq.yaml index bdd54397d19..58427e67ac2 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_offline_eagle3_ptq.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_offline_eagle3_ptq.yaml @@ -117,7 +117,7 @@ pipeline: - --ep_size 1 - --speculative_algorithm EAGLE3 - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl - - --concurrency 1 + - --concurrency 32 environment: - HF_MODEL_CKPT: <> slurm_config: diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_eagle3.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_eagle3.yaml index 969a865f35f..89a578ad57c 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_eagle3.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_eagle3.yaml @@ -62,7 +62,7 @@ pipeline: - --ep_size 1 - --speculative_algorithm EAGLE3 - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl - - --concurrency 1 + - --concurrency 32 environment: - HF_MODEL_CKPT: <> slurm_config: diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_eagle3.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_eagle3.yaml index 7a0d9e3834e..59656709d93 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_eagle3.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_eagle3.yaml @@ -70,7 +70,7 @@ pipeline: - --ep_size 1 - --speculative_algorithm EAGLE3 - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl - - --concurrency 1 + - --concurrency 32 environment: - HF_MODEL_CKPT: <> slurm_config: diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_eagle3_multi_node.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_eagle3_multi_node.yaml index 413c5bb7424..ccf8a9c8f8c 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_eagle3_multi_node.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_streaming_eagle3_multi_node.yaml @@ -77,7 +77,7 @@ pipeline: - --ep_size 1 - --speculative_algorithm EAGLE3 - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl - - --concurrency 1 + - --concurrency 32 environment: - HF_MODEL_CKPT: <> slurm_config: diff --git a/tools/launcher/examples/moonshotai/Kimi-K2.5/hf_streaming_dflash_multi_node.yaml b/tools/launcher/examples/moonshotai/Kimi-K2.5/hf_streaming_dflash_multi_node.yaml index 47d1bc0bef3..3801012112e 100644 --- a/tools/launcher/examples/moonshotai/Kimi-K2.5/hf_streaming_dflash_multi_node.yaml +++ b/tools/launcher/examples/moonshotai/Kimi-K2.5/hf_streaming_dflash_multi_node.yaml @@ -61,7 +61,7 @@ pipeline: - training.disable_tqdm=true - training.ar_validate_steps=500000 - training.num_train_epochs=1 - - training.max_steps=500 + - training.max_steps=2000 # Kimi's slow tokenizer can't emit assistant masks the standard way; the mask # is recovered from token ids (modelopt.torch.utils.loss_mask). - training.answer_only_loss=true diff --git a/tools/launcher/examples/moonshotai/Kimi-K2.5/hf_streaming_eagle3_multi_node.yaml b/tools/launcher/examples/moonshotai/Kimi-K2.5/hf_streaming_eagle3_multi_node.yaml index 06cbd1a8af9..966843b751b 100644 --- a/tools/launcher/examples/moonshotai/Kimi-K2.5/hf_streaming_eagle3_multi_node.yaml +++ b/tools/launcher/examples/moonshotai/Kimi-K2.5/hf_streaming_eagle3_multi_node.yaml @@ -52,7 +52,7 @@ pipeline: - training.disable_tqdm=true - training.ar_validate_steps=500000 - training.num_train_epochs=1 - - training.max_steps=500 + - training.max_steps=2000 - eagle.eagle_use_torch_compile=false environment: - HF_MODEL_CKPT: <> diff --git a/tools/launcher/examples/openai/gpt-oss-20b/chat_template_train.jinja b/tools/launcher/examples/openai/gpt-oss-20b/chat_template_train.jinja new file mode 100644 index 00000000000..18a95aa92eb --- /dev/null +++ b/tools/launcher/examples/openai/gpt-oss-20b/chat_template_train.jinja @@ -0,0 +1,331 @@ +{#- + In addition to the normal inputs of `messages` and `tools`, this template also accepts the + following kwargs: + - "builtin_tools": A list, can contain "browser" and/or "python". + - "model_identity": A string that optionally describes the model identity. + - "reasoning_effort": A string that describes the reasoning effort, defaults to "medium". + #} + +{#- Tool Definition Rendering ============================================== #} +{%- macro render_typescript_type(param_spec, required_params, is_nullable=false) -%} + {%- if param_spec.type == "array" -%} + {%- if param_spec['items'] -%} + {%- if param_spec['items']['type'] == "string" -%} + {{- "string[]" }} + {%- elif param_spec['items']['type'] == "number" -%} + {{- "number[]" }} + {%- elif param_spec['items']['type'] == "integer" -%} + {{- "number[]" }} + {%- elif param_spec['items']['type'] == "boolean" -%} + {{- "boolean[]" }} + {%- else -%} + {%- set inner_type = render_typescript_type(param_spec['items'], required_params) -%} + {%- if inner_type == "object | object" or inner_type|length > 50 -%} + {{- "any[]" }} + {%- else -%} + {{- inner_type + "[]" }} + {%- endif -%} + {%- endif -%} + {%- if param_spec.nullable -%} + {{- " | null" }} + {%- endif -%} + {%- else -%} + {{- "any[]" }} + {%- if param_spec.nullable -%} + {{- " | null" }} + {%- endif -%} + {%- endif -%} + {%- elif param_spec.type is defined and param_spec.type is iterable and param_spec.type is not string and param_spec.type is not mapping and param_spec.type[0] is defined -%} + {#- Handle array of types like ["object", "object"] from Union[dict, list] #} + {%- if param_spec.type | length > 1 -%} + {{- param_spec.type | join(" | ") }} + {%- else -%} + {{- param_spec.type[0] }} + {%- endif -%} + {%- elif param_spec.oneOf -%} + {#- Handle oneOf schemas - check for complex unions and fallback to any #} + {%- set has_object_variants = false -%} + {%- for variant in param_spec.oneOf -%} + {%- if variant.type == "object" -%} + {%- set has_object_variants = true -%} + {%- endif -%} + {%- endfor -%} + {%- if has_object_variants and param_spec.oneOf|length > 1 -%} + {{- "any" }} + {%- else -%} + {%- for variant in param_spec.oneOf -%} + {{- render_typescript_type(variant, required_params) -}} + {%- if variant.description %} + {{- "// " + variant.description }} + {%- endif -%} + {%- if variant.default is defined %} + {{ "// default: " + variant.default|tojson }} + {%- endif -%} + {%- if not loop.last %} + {{- " | " }} + {% endif -%} + {%- endfor -%} + {%- endif -%} + {%- elif param_spec.type == "string" -%} + {%- if param_spec.enum -%} + {{- '"' + param_spec.enum|join('" | "') + '"' -}} + {%- else -%} + {{- "string" }} + {%- if param_spec.nullable %} + {{- " | null" }} + {%- endif -%} + {%- endif -%} + {%- elif param_spec.type == "number" -%} + {{- "number" }} + {%- elif param_spec.type == "integer" -%} + {{- "number" }} + {%- elif param_spec.type == "boolean" -%} + {{- "boolean" }} + + {%- elif param_spec.type == "object" -%} + {%- if param_spec.properties -%} + {{- "{\n" }} + {%- for prop_name, prop_spec in param_spec.properties.items() -%} + {{- prop_name -}} + {%- if prop_name not in (param_spec.required or []) -%} + {{- "?" }} + {%- endif -%} + {{- ": " }} + {{ render_typescript_type(prop_spec, param_spec.required or []) }} + {%- if not loop.last -%} + {{-", " }} + {%- endif -%} + {%- endfor -%} + {{- "}" }} + {%- else -%} + {{- "object" }} + {%- endif -%} + {%- else -%} + {{- "any" }} + {%- endif -%} +{%- endmacro -%} + +{%- macro render_tool_namespace(namespace_name, tools) -%} + {{- "## " + namespace_name + "\n\n" }} + {{- "namespace " + namespace_name + " {\n\n" }} + {%- for tool in tools %} + {%- set tool = tool.function %} + {{- "// " + tool.description + "\n" }} + {{- "type "+ tool.name + " = " }} + {%- if tool.parameters and tool.parameters.properties %} + {{- "(_: {\n" }} + {%- for param_name, param_spec in tool.parameters.properties.items() %} + {%- if param_spec.description %} + {{- "// " + param_spec.description + "\n" }} + {%- endif %} + {{- param_name }} + {%- if param_name not in (tool.parameters.required or []) -%} + {{- "?" }} + {%- endif -%} + {{- ": " }} + {{- render_typescript_type(param_spec, tool.parameters.required or []) }} + {%- if param_spec.default is defined -%} + {%- if param_spec.enum %} + {{- ", // default: " + param_spec.default }} + {%- elif param_spec.oneOf %} + {{- "// default: " + param_spec.default }} + {%- else %} + {{- ", // default: " + param_spec.default|tojson }} + {%- endif -%} + {%- endif -%} + {%- if not loop.last %} + {{- ",\n" }} + {%- else %} + {{- ",\n" }} + {%- endif -%} + {%- endfor %} + {{- "}) => any;\n\n" }} + {%- else -%} + {{- "() => any;\n\n" }} + {%- endif -%} + {%- endfor %} + {{- "} // namespace " + namespace_name }} +{%- endmacro -%} + +{%- macro render_builtin_tools(browser_tool, python_tool) -%} + {%- if browser_tool %} + {{- "## browser\n\n" }} + {{- "// Tool for browsing.\n" }} + {{- "// The `cursor` appears in brackets before each browsing display: `[{cursor}]`.\n" }} + {{- "// Cite information from the tool using the following format:\n" }} + {{- "// `【{cursor}†L{line_start}(-L{line_end})?】`, for example: `【6†L9-L11】` or `【8†L3】`.\n" }} + {{- "// Do not quote more than 10 words directly from the tool output.\n" }} + {{- "// sources=web (default: web)\n" }} + {{- "namespace browser {\n\n" }} + {{- "// Searches for information related to `query` and displays `topn` results.\n" }} + {{- "type search = (_: {\n" }} + {{- "query: string,\n" }} + {{- "topn?: number, // default: 10\n" }} + {{- "source?: string,\n" }} + {{- "}) => any;\n\n" }} + {{- "// Opens the link `id` from the page indicated by `cursor` starting at line number `loc`, showing `num_lines` lines.\n" }} + {{- "// Valid link ids are displayed with the formatting: `【{id}†.*】`.\n" }} + {{- "// If `cursor` is not provided, the most recent page is implied.\n" }} + {{- "// If `id` is a string, it is treated as a fully qualified URL associated with `source`.\n" }} + {{- "// If `loc` is not provided, the viewport will be positioned at the beginning of the document or centered on the most relevant passage, if available.\n" }} + {{- "// Use this function without `id` to scroll to a new location of an opened page.\n" }} + {{- "type open = (_: {\n" }} + {{- "id?: number | string, // default: -1\n" }} + {{- "cursor?: number, // default: -1\n" }} + {{- "loc?: number, // default: -1\n" }} + {{- "num_lines?: number, // default: -1\n" }} + {{- "view_source?: boolean, // default: false\n" }} + {{- "source?: string,\n" }} + {{- "}) => any;\n\n" }} + {{- "// Finds exact matches of `pattern` in the current page, or the page given by `cursor`.\n" }} + {{- "type find = (_: {\n" }} + {{- "pattern: string,\n" }} + {{- "cursor?: number, // default: -1\n" }} + {{- "}) => any;\n\n" }} + {{- "} // namespace browser\n\n" }} + {%- endif -%} + + {%- if python_tool %} + {{- "## python\n\n" }} + {{- "Use this tool to execute Python code in your chain of thought. The code will not be shown to the user. This tool should be used for internal reasoning, but not for code that is intended to be visible to the user (e.g. when creating plots, tables, or files).\n\n" }} + {{- "When you send a message containing Python code to python, it will be executed in a stateful Jupyter notebook environment. python will respond with the output of the execution or time out after 120.0 seconds. The drive at '/mnt/data' can be used to save and persist user files. Internet access for this session is UNKNOWN. Depends on the cluster.\n\n" }} + {%- endif -%} +{%- endmacro -%} + +{#- System Message Construction ============================================ #} +{%- macro build_system_message() -%} + {%- if model_identity is not defined %} + {%- set model_identity = "You are ChatGPT, a large language model trained by OpenAI." %} + {%- endif %} + {{- model_identity + "\n" }} + {{- "Knowledge cutoff: 2024-06\n" }} + {{- "Current date: " + strftime_now("%Y-%m-%d") + "\n\n" }} + {%- if reasoning_effort is not defined %} + {%- set reasoning_effort = "medium" %} + {%- endif %} + {{- "Reasoning: " + reasoning_effort + "\n\n" }} + {%- if builtin_tools %} + {{- "# Tools\n\n" }} + {%- set available_builtin_tools = namespace(browser=false, python=false) %} + {%- for tool in builtin_tools %} + {%- if tool == "browser" %} + {%- set available_builtin_tools.browser = true %} + {%- elif tool == "python" %} + {%- set available_builtin_tools.python = true %} + {%- endif %} + {%- endfor %} + {{- render_builtin_tools(available_builtin_tools.browser, available_builtin_tools.python) }} + {%- endif -%} + {{- "# Valid channels: analysis, commentary, final. Channel must be included for every message." }} + {%- if tools -%} + {{- "\nCalls to these tools must go to the commentary channel: 'functions'." }} + {%- endif -%} +{%- endmacro -%} + +{#- Main Template Logic ================================================= #} +{#- Set defaults #} + +{#- Render system message #} +{{- "<|start|>system<|message|>" }} +{{- build_system_message() }} +{{- "<|end|>" }} + +{#- Extract developer message #} +{%- if messages[0].role == "developer" or messages[0].role == "system" %} + {%- set developer_message = messages[0].content %} + {%- set loop_messages = messages[1:] %} +{%- else %} + {%- set developer_message = "" %} + {%- set loop_messages = messages %} +{%- endif %} + +{#- Render developer message #} +{%- if developer_message or tools %} + {{- "<|start|>developer<|message|>" }} + {%- if developer_message %} + {{- "# Instructions\n\n" }} + {{- developer_message }} + {{- "\n\n" }} + {%- endif %} + {%- if tools -%} + {{- "# Tools\n\n" }} + {{- render_tool_namespace("functions", tools) }} + {%- endif -%} + {{- "<|end|>" }} +{%- endif %} + +{#- Render messages #} +{%- set last_tool_call = namespace(name=none) %} +{%- for message in loop_messages -%} + {#- At this point only assistant/user/tool messages should remain #} + {%- if message.role == 'assistant' -%} + {#- Checks to ensure the messages are being passed in the format we expect #} + {%- if "content" in message %} + {%- if "<|channel|>analysis<|message|>" in message.content or "<|channel|>final<|message|>" in message.content %} + {{- raise_exception("You have passed a message containing <|channel|> tags in the content field. Instead of doing this, you should pass analysis messages (the string between '<|message|>' and '<|end|>') in the 'thinking' field, and final messages (the string between '<|message|>' and '<|end|>') in the 'content' field.") }} + {%- endif %} + {%- endif %} + {%- if "thinking" in message %} + {%- if "<|channel|>analysis<|message|>" in message.thinking or "<|channel|>final<|message|>" in message.thinking %} + {{- raise_exception("You have passed a message containing <|channel|> tags in the thinking field. Instead of doing this, you should pass analysis messages (the string between '<|message|>' and '<|end|>') in the 'thinking' field, and final messages (the string between '<|message|>' and '<|end|>') in the 'content' field.") }} + {%- endif %} + {%- endif %} + {%- if "tool_calls" in message %} + {#- We need very careful handling here - we want to drop the tool call analysis message if the model #} + {#- has output a later <|final|> message, but otherwise we want to retain it. This is the only case #} + {#- when we render CoT/analysis messages in inference. #} + {%- set future_final_message = namespace(found=false) %} + {%- for future_message in loop_messages[loop.index:] %} + {%- if future_message.role == 'assistant' and "tool_calls" not in future_message %} + {%- set future_final_message.found = true %} + {%- endif %} + {%- endfor %} + {#- We assume max 1 tool call per message, and so we infer the tool call name #} + {#- in "tool" messages from the most recent assistant tool call name #} + {%- set tool_call = message.tool_calls[0] %} + {%- if tool_call.function %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {%- if message.content and message.thinking %} + {{- raise_exception("Cannot pass both content and thinking in an assistant message with tool calls! Put the analysis message in one or the other, but not both.") }} + {%- elif message.content and not future_final_message.found %} + {{- "<|start|>assistant<|channel|>analysis<|message|>" + message.content + "<|end|>" }} + {%- elif message.thinking and not future_final_message.found %} + {{- "<|start|>assistant<|channel|>analysis<|message|>" + message.thinking + "<|end|>" }} + {%- endif %} + {{- "<|start|>assistant to=" }} + {{- "functions." + tool_call.name + "<|channel|>commentary " }} + {{- (tool_call.content_type if tool_call.content_type is defined else "json") + "<|message|>" }} + {{- tool_call.arguments|tojson }} + {{- "<|call|>" }} + {%- set last_tool_call.name = tool_call.name %} + {%- elif loop.last and not add_generation_prompt %} + {#- Only render the CoT if the final turn is an assistant turn and add_generation_prompt is false #} + {#- This is a situation that should only occur in training, never in inference. #} + {%- if "thinking" in message %} + {{- "<|start|>assistant<|channel|>analysis<|message|>" -}}{%- generation -%}{{- message.thinking -}}{%- endgeneration -%}{{- "<|end|>" }} + {%- endif %} + {#- <|return|> indicates the end of generation, but <|end|> does not #} + {#- <|return|> should never be an input to the model, but we include it as the final token #} + {#- when training, so the model learns to emit it. #} + {{- "<|start|>assistant<|channel|>final<|message|>" -}}{%- generation -%}{{- message.content -}}{%- endgeneration -%}{{- "<|return|>" }} + {%- else %} + {#- CoT is dropped during all previous turns, so we never render it for inference #} + {{- "<|start|>assistant<|channel|>final<|message|>" -}}{%- generation -%}{{- message.content -}}{%- endgeneration -%}{{- "<|end|>" }} + {%- set last_tool_call.name = none %} + {%- endif %} + {%- elif message.role == 'tool' -%} + {%- if last_tool_call.name is none %} + {{- raise_exception("Message has tool role, but there was no previous assistant message with a tool call!") }} + {%- endif %} + {{- "<|start|>functions." + last_tool_call.name }} + {{- " to=assistant<|channel|>commentary<|message|>" + message.content|tojson + "<|end|>" }} + {%- elif message.role == 'user' -%} + {{- "<|start|>user<|message|>" + message.content + "<|end|>" }} + {%- endif -%} +{%- endfor -%} + +{#- Generation prompt #} +{%- if add_generation_prompt -%} +<|start|>assistant +{%- endif -%} \ No newline at end of file diff --git a/tools/launcher/examples/openai/gpt-oss-20b/hf_streaming_dflash_multi_node.yaml b/tools/launcher/examples/openai/gpt-oss-20b/hf_streaming_dflash_multi_node.yaml new file mode 100644 index 00000000000..f17604410fa --- /dev/null +++ b/tools/launcher/examples/openai/gpt-oss-20b/hf_streaming_dflash_multi_node.yaml @@ -0,0 +1,120 @@ +# DFlash streaming speculative decoding pipeline for gpt-oss-20b — MULTI-NODE. +# +# Same streaming transport / dispatch as hf_streaming_eagle3_multi_node.yaml: task_1 +# splits N nodes into K serve replicas + (N-K) DDP trainers via SERVE_NODES; hidden +# states move serve -> trainer over NIXL RDMA. DFlash just consumes a different set of +# captured layers and trains a block-diffusion draft instead of an autoregressive one. +# See common/eagle3/train_eagle_streaming.sh for dispatch, rendezvous, and sharding. +# +# gpt-oss-20b notes (vs the dense Qwen3-8B example): +# - MoE (32 experts, 4 active/token) shipped MXFP4-quantized. The serve already +# passes --no-enable-flashinfer-autotune (added for NVFP4/MXFP4 MoE: the autotuner +# re-tunes on the first serving step and can stall a worker past vLLM's +# execute-model timeout, killing EngineCore). H100 has no FP4 HW, so vLLM +# dequant/emulates MXFP4 on Hopper. +# - 24 hidden layers, 5 draft layers -> build_target_layer_ids(24,5)=[1,6,11,16,21] +# (the draft's fc input) plus the final layer for self-logit distillation. vLLM's +# capture ids are those +1 -> [2,7,12,17,22], plus final layer 24. +# - dflash_mask_token_id: gpt-oss has no tokenizer mask token (would fall back to +# None), so set it explicitly to a reserved/unused vocab slot, <|reserved_200013|>. +# +# 3-step pipeline: +# task_0: Build input conversations (jsonl) +# task_1: Streaming train — 2 serve nodes (2 GPU, TP=2) + 2 trainer nodes (2 GPU) +# task_2: vLLM smoke test with DFlash speculative decoding +# +# Usage: +# uv run launch.py --yaml examples/openai/gpt-oss-20b/hf_streaming_dflash_multi_node.yaml --yes + +job_name: gpt-oss-20b_DFlash_streaming_multi_node +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: /hf-local/openai/gpt-oss-20b + + # Step 1: Build input conversations + 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: 1 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 + + # Step 2: Streaming DFlash training — 2 serve replicas (TP=2) + 2 trainer nodes (2 GPU each). + # DFlash extracts 5 target layers (build_target_layer_ids(24,5)=[1,6,11,16,21], the + # draft's fc input) plus the final layer for self-logit distillation. vLLM's capture + # ids are those +1 -> [2,7,12,17,22], plus final layer 24. + task_1: + script: common/eagle3/train_eagle_streaming.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dflash.yaml + - model.model_name_or_path=<> + - model.use_fake_base_for_offline=true + - data.mode=streaming + - data.data_path=/scratchspace/data/train.jsonl + - training.output_dir=/scratchspace/dflash + - training.training_seq_len=4096 + - training.disable_tqdm=true + # Supervise only the assistant content, not the prompt. The serve captures hidden + # states over the full corpus conversation (it does not generate), so the prompt + + # gpt-oss Harmony boilerplate is present; training over the full sequence + # (answer_only_loss=false) lets the draft trivially memorize that boilerplate -> + # inflated train metrics that don't translate to real acceptance. The training-only + # chat template adds {% generation %} tags so the tokenizer emits the assistant mask + # (byte-identical tokenization to the stock template otherwise). + - training.answer_only_loss=true + - data.chat_template=examples/openai/gpt-oss-20b/chat_template_train.jinja + - training.num_train_epochs=1 + - training.max_steps=2000 + # dflash.yaml sets report_to=tensorboard, which hard-fails if tensorboard + # isn't in the serve container; the streaming trainer doesn't need it. + - training.report_to=none + - dflash.dflash_block_size=16 + - dflash.dflash_num_anchors=512 + - dflash.dflash_loss_decay_factor=7 + # gpt-oss has no tokenizer mask token; use a reserved vocab slot (<|reserved_200013|>). + - dflash.dflash_mask_token_id=200013 + - dflash.dflash_architecture_config.num_hidden_layers=5 + environment: + - HF_MODEL_CKPT: <> + # No spaces: nemo_run emits `export FOO=value` unquoted. + - EAGLE_CAPTURE_IDS: "[2,7,12,17,22,24]" + - SERVE_TP: "2" + # extract_hidden_states packs all 6 capture layers into one KV spec (per-token 6*hidden_size*2B=34560); block_size must make the attn page 2*block_size*(num_kv_heads/SERVE_TP)*head_dim*2 >= that, so at TP=2 (4 kv_heads/rank) need >=64. + - SERVE_EXTRA_ARGS: "--block-size=64" + # K serve replica nodes (Slurm nodes 0..K-1); the rest are trainers. + - SERVE_NODES: "2" + # Per-rank in-flight fetches; keep low so the cold serve isn't flooded past its execute-model timeout. + - STREAMING_NUM_WORKERS: "4" + # DFlash uses a custom modeling file; export must trust remote code. + - EXPORT_EXTRA_ARGS: "--trust_remote_code" + slurm_config: + _factory_: "slurm_factory" + nodes: 4 + ntasks_per_node: 1 + gpus_per_node: 2 + container: vllm/vllm-openai:latest + + # Step 3: vLLM smoke test (DFlash, uses exported checkpoint from training) + task_2: + script: common/specdec/vllm_smoke_test.sh + environment: + - HF_MODEL_CKPT: <> + - DRAFT_MODEL: /scratchspace/export + - SPEC_METHOD: "dflash" + - NUM_SPEC_TOKENS: "7" + - MIN_ACCEPTANCE_LENGTH: "1.2" + slurm_config: + _factory_: "slurm_factory" + container: "vllm/vllm-openai:nightly" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 diff --git a/tools/launcher/examples/openai/gpt-oss-20b/hf_streaming_eagle3_multi_node.yaml b/tools/launcher/examples/openai/gpt-oss-20b/hf_streaming_eagle3_multi_node.yaml new file mode 100644 index 00000000000..7ad003d8fd8 --- /dev/null +++ b/tools/launcher/examples/openai/gpt-oss-20b/hf_streaming_eagle3_multi_node.yaml @@ -0,0 +1,114 @@ +# EAGLE3 streaming speculative decoding pipeline for gpt-oss-20b — MULTI-NODE. +# +# task_1 splits N nodes into K serve replicas + (N-K) DDP trainers via SERVE_NODES; +# see common/eagle3/train_eagle_streaming.sh for dispatch, rendezvous, and sharding. +# +# gpt-oss-20b notes (vs the dense Qwen3-8B example): +# - MoE (32 experts, 4 active/token) shipped MXFP4-quantized. The serve already +# passes --no-enable-flashinfer-autotune (added for NVFP4/MXFP4 MoE: the autotuner +# re-tunes on the first serving step and can stall a worker past vLLM's +# execute-model timeout, killing EngineCore). H100 has no FP4 HW, so vLLM +# dequant/emulates MXFP4 on Hopper. +# - 24 hidden layers -> default_eagle_aux_layer_ids(24)=[1,11,20]; vLLM capture ids +# are each +1 plus the final layer 24 -> EAGLE_CAPTURE_IDS="[2,12,21,24]". +# - gpt_oss is native in recent transformers/vLLM (no --trust-remote-code needed). +# +# 3-step pipeline: +# task_0: Build input conversations (jsonl) +# task_1: Streaming train — 2 serve nodes (2 GPU, TP=2) + 2 trainer nodes (2 GPU) +# task_2: Benchmark — evaluate speculative decoding speedup via VLLM +# +# Usage: +# uv run launch.py --yaml examples/openai/gpt-oss-20b/hf_streaming_eagle3_multi_node.yaml --yes + +job_name: gpt-oss-20b_EAGLE3_streaming_multi_node +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: /hf-local/openai/gpt-oss-20b + + # Step 1: Build input conversations + 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: 1 + container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 + + # Step 2: Streaming EAGLE3 training — 2 serve replicas (TP=2) + 2 trainer nodes (2 GPU each). + # Capture ids: default_eagle_aux_layer_ids(24)=[1,11,20] +1, plus final layer 24. + task_1: + script: common/eagle3/train_eagle_streaming.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/eagle3.yaml + - model.model_name_or_path=<> + - model.use_fake_base_for_offline=true + - data.mode=streaming + - data.data_path=/scratchspace/data/train.jsonl + # Supervise only the assistant content, not the prompt. gpt-oss's Harmony chat + # template is mostly fixed structural/boilerplate tokens (system prompt, channel + # markers); training over the full sequence (answer_only_loss=false) lets the draft + # trivially memorize that boilerplate -> inflated train_acc that does NOT translate + # to real spec-decode acceptance. The training-only chat template adds {% generation %} + # tags around the assistant content so the tokenizer can emit the assistant mask + # (it tokenizes byte-identically to the stock template otherwise). + - training.answer_only_loss=true + - data.chat_template=examples/openai/gpt-oss-20b/chat_template_train.jinja + - training.output_dir=/scratchspace/eagle3 + - training.training_seq_len=4096 + - training.disable_tqdm=true + - training.ar_validate_steps=500000 + - training.num_train_epochs=1 + - training.max_steps=2000 + - eagle.eagle_use_torch_compile=false + environment: + - HF_MODEL_CKPT: <> + # No spaces: nemo_run emits `export FOO=value` unquoted. + - EAGLE_CAPTURE_IDS: "[2,12,21,24]" + - SERVE_TP: "2" + # K serve replica nodes (Slurm nodes 0..K-1); the rest are trainers. + - SERVE_NODES: "2" + # Per-rank in-flight fetches; keep low so the cold MXFP4-MoE serve isn't flooded past its execute-model timeout (kills EngineCore). + - STREAMING_NUM_WORKERS: "4" + slurm_config: + _factory_: "slurm_factory" + nodes: 4 + ntasks_per_node: 1 + gpus_per_node: 2 + container: vllm/vllm-openai:latest + + # Step 3: Benchmark speculative decoding (VLLM backend) + task_2: + script: common/specdec_bench/quick_check.sh + args: + - --draft_model_dir /scratchspace/export + - --draft_length 3 + - --output_length 4096 + - --engine VLLM + - --tp_size 1 + - --ep_size 1 + - --speculative_algorithm EAGLE3 + - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl + - --concurrency 32 + # gpt-oss emits Harmony channel tags (<|channel|>analysis...<|channel|>final...) + # in raw output. The MT-bench multi-turn loop re-feeds each answer as the next + # turn's assistant message; gpt-oss's chat template rejects content containing + # channel tags. --postprocess gptoss extracts just the final-channel text before + # re-feeding. (Affects only chat history, not the acceptance-length measurement.) + - --postprocess gptoss + environment: + - HF_MODEL_CKPT: <> + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 1 + container: vllm/vllm-openai:latest