From dca112012f54f06f19b50d0b0a47eaa605bc9b57 Mon Sep 17 00:00:00 2001 From: Yi Sun Date: Sun, 19 Jul 2026 02:56:53 +0000 Subject: [PATCH 1/4] Expand DSpark target support and training scalability --- configs/glm-5.2-dspark.json | 46 ++ configs/inkling-dspark.json | 46 ++ configs/qwen3-4b-dspark.json | 1 + configs/qwen3-8b-dspark.json | 45 ++ docs/basic_usage/disaggregated_training.md | 7 +- docs/benchmarks/benchmark.md | 124 ++- examples/configs/README.md | 17 +- .../configs/glm-5.2-dspark-disaggregated.yaml | 41 + .../configs/inkling-dspark-disaggregated.yaml | 43 + .../qwen3-4b-dspark-disaggregated.yaml | 2 + .../qwen3-8b-dspark-disaggregated.yaml | 40 + .../run_qwen3_8b_dflash_disagg_2node.sh | 1 - patches/sglang/v0.5.14/spec-capture.patch | 80 +- .../gates/run_disaggregated_overfit_gate.sh | 1 - .../algorithms/common/dflash_family_model.py | 739 +++++++++++------- specforge/algorithms/model_providers.py | 1 + specforge/benchmarks/__init__.py | 1 + specforge/benchmarks/sglang.py | 220 ++++++ specforge/cli.py | 40 +- specforge/config/schema.py | 26 + specforge/core/chunking.py | 104 +++ specforge/data/parse.py | 46 +- specforge/data/preprocessing.py | 4 +- specforge/data/template.py | 31 + .../inference/adapters/server_capture.py | 18 +- specforge/inference/sglang_patch_inventory.md | 11 +- specforge/launch_plan.py | 23 +- specforge/modeling/draft/dspark.py | 47 +- specforge/modeling/target/target_utils.py | 120 ++- specforge/optimizer.py | 95 ++- specforge/training/assembly.py | 43 +- specforge/training/capture_contract.py | 13 +- specforge/training/checkpoint.py | 66 +- specforge/training/controller.py | 91 ++- specforge/training/model_loading.py | 11 +- specforge/training/strategies/base.py | 10 +- .../test_benchmarks/test_sglang_benchmark.py | 100 +++ tests/test_config/test_launch_topology.py | 59 +- tests/test_config/test_schema.py | 24 + .../test_unified_feature_reachability.py | 2 +- tests/test_data/test_inkling_template.py | 66 ++ tests/test_modeling/test_domino_model.py | 118 +++ tests/test_modeling/test_draft_registry.py | 45 ++ .../test_target/test_target_utils_loading.py | 87 ++- .../test_bf16_optimizer_clip_grad_norm.py | 148 +++- tests/test_runtime/test_checkpoint_manager.py | 50 ++ tests/test_runtime/test_launch_plan.py | 26 +- .../test_runtime/test_package_architecture.py | 33 + tests/test_runtime/test_server_capture.py | 8 +- .../test_runtime/test_server_capture_gate.py | 12 +- tests/test_runtime/test_trainer.py | 59 +- tests/test_utils/test_chunking.py | 67 ++ tests/test_utils/test_dflash_losses.py | 249 +++++- 53 files changed, 2885 insertions(+), 522 deletions(-) create mode 100644 configs/glm-5.2-dspark.json create mode 100644 configs/inkling-dspark.json create mode 100644 configs/qwen3-8b-dspark.json create mode 100644 examples/configs/glm-5.2-dspark-disaggregated.yaml create mode 100644 examples/configs/inkling-dspark-disaggregated.yaml create mode 100644 examples/configs/qwen3-8b-dspark-disaggregated.yaml create mode 100644 specforge/benchmarks/__init__.py create mode 100644 specforge/benchmarks/sglang.py create mode 100644 specforge/core/chunking.py create mode 100644 tests/test_benchmarks/test_sglang_benchmark.py create mode 100644 tests/test_data/test_inkling_template.py create mode 100644 tests/test_utils/test_chunking.py diff --git a/configs/glm-5.2-dspark.json b/configs/glm-5.2-dspark.json new file mode 100644 index 000000000..28d66955d --- /dev/null +++ b/configs/glm-5.2-dspark.json @@ -0,0 +1,46 @@ +{ + "architectures": ["DSparkDraftModel"], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": {"AutoModel": "dspark.DSparkDraftModel"}, + "block_size": 7, + "bos_token_id": null, + "dflash_config": { + "attention_mode": "gqa", + "confidence_head_alpha": 1.0, + "confidence_head_with_markov": true, + "enable_confidence_head": true, + "markov_head_type": "vanilla", + "markov_rank": 256, + "mask_token_id": 154856, + "projector_type": "dspark", + "target_layer_ids": [1, 19, 38, 57, 76] + }, + "dtype": "bfloat16", + "eos_token_id": [154820, 154827, 154829], + "head_dim": 64, + "hidden_act": "silu", + "hidden_size": 6144, + "initializer_range": 0.02, + "intermediate_size": 12288, + "layer_types": [ + "full_attention", "full_attention", "full_attention", + "full_attention", "full_attention" + ], + "max_position_embeddings": 1048576, + "max_window_layers": 5, + "model_type": "qwen3", + "num_attention_heads": 64, + "num_hidden_layers": 5, + "num_key_value_heads": 16, + "num_target_layers": 78, + "pad_token_id": 154820, + "rms_norm_eps": 1e-05, + "rope_scaling": null, + "rope_theta": 8000000, + "sliding_window": null, + "tie_word_embeddings": false, + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 154880 +} diff --git a/configs/inkling-dspark.json b/configs/inkling-dspark.json new file mode 100644 index 000000000..fd9a55bff --- /dev/null +++ b/configs/inkling-dspark.json @@ -0,0 +1,46 @@ +{ + "architectures": ["DSparkDraftModel"], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": {"AutoModel": "dspark.DSparkDraftModel"}, + "block_size": 7, + "bos_token_id": null, + "dflash_config": { + "attention_mode": "gqa", + "confidence_head_alpha": 1.0, + "confidence_head_with_markov": true, + "enable_confidence_head": true, + "markov_head_type": "vanilla", + "markov_rank": 256, + "mask_token_id": 200064, + "projector_type": "dspark", + "target_layer_ids": [5, 17, 35, 47, 59] + }, + "dtype": "bfloat16", + "eos_token_id": 200006, + "head_dim": 64, + "hidden_act": "silu", + "hidden_size": 6144, + "initializer_range": 0.02, + "intermediate_size": 12288, + "layer_types": [ + "full_attention", "full_attention", "full_attention", + "full_attention", "full_attention" + ], + "max_position_embeddings": 1048576, + "max_window_layers": 5, + "model_type": "qwen3", + "num_attention_heads": 64, + "num_hidden_layers": 5, + "num_key_value_heads": 16, + "num_target_layers": 66, + "pad_token_id": 200006, + "rms_norm_eps": 1e-05, + "rope_scaling": null, + "rope_theta": 8000000, + "sliding_window": null, + "tie_word_embeddings": false, + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 201024 +} diff --git a/configs/qwen3-4b-dspark.json b/configs/qwen3-4b-dspark.json index f47281114..c8a966555 100644 --- a/configs/qwen3-4b-dspark.json +++ b/configs/qwen3-4b-dspark.json @@ -10,6 +10,7 @@ "block_size": 7, "bos_token_id": 151643, "dflash_config": { + "attention_mode": "gqa", "mask_token_id": 151669, "target_layer_ids": [1, 9, 17, 25, 33], "projector_type": "dspark", diff --git a/configs/qwen3-8b-dspark.json b/configs/qwen3-8b-dspark.json new file mode 100644 index 000000000..5857de00a --- /dev/null +++ b/configs/qwen3-8b-dspark.json @@ -0,0 +1,45 @@ +{ + "architectures": ["DSparkDraftModel"], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": {"AutoModel": "dspark.DSparkDraftModel"}, + "block_size": 7, + "bos_token_id": 151643, + "dflash_config": { + "attention_mode": "gqa", + "confidence_head_alpha": 1.0, + "confidence_head_with_markov": true, + "enable_confidence_head": true, + "markov_head_type": "vanilla", + "markov_rank": 256, + "mask_token_id": 151669, + "projector_type": "dspark", + "target_layer_ids": [1, 9, 17, 25, 33] + }, + "dtype": "bfloat16", + "eos_token_id": 151645, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 12288, + "layer_types": [ + "full_attention", "full_attention", "full_attention", + "full_attention", "full_attention" + ], + "max_position_embeddings": 40960, + "max_window_layers": 5, + "model_type": "qwen3", + "num_attention_heads": 32, + "num_hidden_layers": 5, + "num_key_value_heads": 8, + "num_target_layers": 36, + "rms_norm_eps": 1e-06, + "rope_scaling": null, + "rope_theta": 1000000, + "sliding_window": null, + "tie_word_embeddings": false, + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 151936 +} diff --git a/docs/basic_usage/disaggregated_training.md b/docs/basic_usage/disaggregated_training.md index ac15214f3..f445fbca3 100644 --- a/docs/basic_usage/disaggregated_training.md +++ b/docs/basic_usage/disaggregated_training.md @@ -318,9 +318,10 @@ The online producer sends prompts to the URLs in with the model, capture method, and auxiliary layer ids matching the draft config. DFlash, Domino, and DSpark use the DFlash capture contract; EAGLE3 and P-EAGLE use the EAGLE3 capture contract. Capture rejects chunked prefill and -radix-cache paths that can truncate the captured sequence. Online capture is -text-only: VLM training, including Qwen2.5-VL, is not supported. Online -evaluation is also not supported. +gives every request attempt a unique radix-cache namespace so cached prefixes +cannot truncate the captured sequence. Online capture is text-only: VLM +training, including Qwen2.5-VL, is not supported. Online evaluation is also not +supported. The repository's strict e2e gate remains a full local test-stack orchestrator: diff --git a/docs/benchmarks/benchmark.md b/docs/benchmarks/benchmark.md index 8783e46d9..b4c8ec371 100644 --- a/docs/benchmarks/benchmark.md +++ b/docs/benchmarks/benchmark.md @@ -1,87 +1,83 @@ -# Benchmarking speculative decoding +# Benchmarking inference serving -The repository keeps the model-quality and serving-performance benchmark suite -under `benchmarks/`. It is independent from training: train and export a draft -through the unified `specforge` CLI, then benchmark that exported artifact -against an SGLang server. +Use `benchmarks/bench_eagle3.py` to compare EAGLE3 serving configurations and +dataset quality. Use `specforge benchmark sglang` to measure any existing +SGLang deployment without assuming a particular speculative algorithm. -## Run server and benchmarks together +| Runner | Server lifecycle | Measurements | Output | +| --- | --- | --- | --- | +| `python benchmarks/bench_eagle3.py` | Launches SGLang per configuration or uses an existing server | Latency, output throughput, acceptance length, and dataset accuracy when available | Timestamped JSON under `--output-dir` | +| `specforge benchmark sglang` | Uses an existing SGLang server | Aggregate output throughput, acceptance length, and verification count when reported by SGLang | Console and optional `--output-json` | -From the repository root: +## EAGLE3 benchmark matrix + +### Launch SGLang for each configuration ```bash python benchmarks/bench_eagle3.py \ - --model-path meta-llama/Llama-3.1-8B-Instruct \ - --speculative-draft-model-path /path/to/exported-draft \ - --port 30000 \ - --trust-remote-code \ - --mem-fraction-static 0.8 \ - --tp-size 1 \ - --attention-backend fa3 \ - --config-list 1,0,0,0 1,3,1,4 \ - --benchmark-list mtbench gsm8k:5 ceval:5:accountant \ - --dtype bfloat16 + --model-path meta-llama/Llama-3.1-8B-Instruct \ + --speculative-draft-model-path /path/to/exported-draft \ + --port 30000 \ + --trust-remote-code \ + --mem-fraction-static 0.8 \ + --tp-size 1 \ + --attention-backend fa3 \ + --config-list 1,0,0,0 1,3,1,4 \ + --benchmark-list mtbench gsm8k:5 ceval:5:accountant \ + --dtype bfloat16 ``` -Each `--config-list` entry is -`batch-size,num-steps,topk,num-draft-tokens`. Benchmark selectors use -`name[:num-prompts[:subset,...]]`. Available datasets include AIME, C-Eval, -FinanceQA, GPQA, GSM8K, HumanEval, LiveCodeBench, MATH-500, MBPP, MMLU, -MMStar, MT-Bench, and SimpleQA. - -## Benchmark an existing server +- `--config-list` uses `batch-size,num-steps,topk,num-draft-tokens`; `1,0,0,0` is a target-only baseline. +- `--benchmark-list` uses `name[:num-prompts[:subset,...]]`. +- Supported datasets are AIME, C-Eval, FinanceQA, GPQA, GSM8K, HumanEval, LiveCodeBench, MATH-500, MBPP, MMLU, MMStar, MT-Bench, and SimpleQA. +- The runner starts a fresh server for each configuration, runs every requested dataset, flushes the cache between datasets, and writes results under `--output-dir`. -Start SGLang separately, then add `--skip-launch-server`: +### Use an existing SGLang server ```bash python benchmarks/bench_eagle3.py \ - --model-path meta-llama/Llama-3.1-8B-Instruct \ - --port 30000 \ - --config-list 1,3,1,4 \ - --benchmark-list mtbench:5 gsm8k:5 humaneval:5 math500:5 \ - --skip-launch-server + --model-path meta-llama/Llama-3.1-8B-Instruct \ + --port 30000 \ + --config-list 1,3,1,4 \ + --benchmark-list mtbench:5 gsm8k:5 humaneval:5 math500:5 \ + --skip-launch-server ``` -Results are written as timestamped JSON under `--output-dir`. The standalone -GPU microbenchmarks `bench_domino_mfu.py`, -`specforge/benchmarks/benchmark_flex_attention.py`, and -`specforge/benchmarks/benchmark_loss.py` cover trainer MFU, attention, and loss -kernel behavior respectively. - -HumanEval and MBPP execute model-generated Python while scoring. Run those two -benchmarks only in an isolated container or devbox with no credentials or -production data mounted. - -## Retained validation artifacts +With `--skip-launch-server`, the runner does not change the server's speculative settings. The first `--config-list` entry supplies only the request batch size. -The repository keeps the existing training-equivalence and convergence plots -as provenance for the runtime cutover. They are reference results, not a -substitute for rerunning the benchmark suite on the current commit. +## General SGLang benchmark -### DataFlow cutover equivalence +The existing-server runner supports target-only and speculative deployments on +GSM8K, MATH-500, HumanEval, MBPP, and MT-Bench. -These 100-step traces compare the former trainer (`main`) with colocated and -disaggregated DataFlow execution. Keeping them next to the benchmark suite -makes the behavioral evidence available after the duplicate trainers are -removed. - -![Qwen2.5-0.5B DFlash legacy, colocated, and disaggregated loss equivalence](../../examples/assets/qwen2.5-0.5b-dflash-vs-main.png) - -![Qwen2.5-0.5B EAGLE3 legacy, colocated, and disaggregated loss equivalence](../../examples/assets/qwen2.5-0.5b-eagle3-vs-main.png) +```bash +specforge benchmark sglang \ + --model Qwen/Qwen3-8B \ + --dataset gsm8k \ + --base-url http://127.0.0.1:30000 \ + --num-prompts 1024 \ + --concurrency 16 \ + --output-json ./qwen3-8b-gsm8k.json +``` -The [Qwen2.5-7B EAGLE3 offline parity record](eagle3-disaggregated-parity.md) -preserves the corresponding two-node metric comparison and its current -rank-dispatch command. +- Start SGLang with the target-only or speculative configuration you want to measure; this command does not launch or reconfigure the server. +- The runner flushes the server cache, runs one concurrency-sized warmup batch, excludes warmup from the measurement, and then sends `--num-prompts` requests. +- The report includes output-token throughput and includes average acceptance length and speculative verification count when SGLang returns those fields. -### Qwen3.6-27B DFlash training curves +## Comparing results -The retained curves cover both the eight-H200 colocated run and the two-GPU -online-disaggregated producer/consumer run. +Measure target-only and speculative decoding with the same target revision, +tokenizer and chat template, prompts, sampling parameters, output length, +hardware, tensor parallelism, and concurrency. For the EAGLE3 matrix, include a +zero-step configuration in one run. With the general SGLang runner, benchmark +matched target-only and speculative servers separately and compute speedup from +their throughput results. -The separate [Domino disaggregated performance findings](domino-disaggregated-performance.md) -preserve the measured one-server + DP7 tuning study, its MFU analysis, and the -canonical YAML form of the relevant controls. +Do not compare absolute throughput across the matrix and existing-server runners +because their batching and request scheduling differ. -![Qwen3.6-27B DFlash colocated training loss and draft-token accuracy](../../examples/assets/qwen36-27b-dflash-nemotron-6ep.png) +## Safety -![Qwen3.6-27B DFlash online-disaggregated training loss and draft-token accuracy](../../examples/disagg/assets/qwen36-27b-dflash-nemotron-disagg.png) +The EAGLE3 HumanEval and MBPP scorers execute model-generated Python. Run them +only in an isolated environment without credentials or production data. The +general SGLang runner measures decoding and does not execute generated code. diff --git a/examples/configs/README.md b/examples/configs/README.md index 5d5618749..84088be7e 100644 --- a/examples/configs/README.md +++ b/examples/configs/README.md @@ -140,6 +140,7 @@ should make their training strategy and topology explicit. | `model.torch_dtype` | `bfloat16` | `bfloat16`, `float16`, or `float32`. | | `model.cache_dir` | `null` | Model/tokenizer download cache. This is distinct from `data.cache_dir`. | | `model.mask_token_id` | `null` | DFlash-family/P-EAGLE mask token override. Otherwise it resolves from the draft config and then the tokenizer. | +| `model.tokenizer_pad_token_id` | `null` | Explicit non-negative tokenizer pad ID. Use it for released tokenizers that omit padding metadata. | | `model.sglang_attention_backend` | `flashinfer` | SGLang attention implementation for an in-process or managed capture server. | | `model.sglang_mem_fraction_static` | `0.4` | SGLang static-memory fraction in `(0, 1]`; inherited by managed capture servers unless they override it. | | `model.sglang_context_length` | `null` | Positive explicit context limit. Managed capture requires at least `data.max_length + 7`; omitting it derives that value. | @@ -151,6 +152,16 @@ should make their training strategy and topology explicit. | `model.sglang_ep_size` | `1` | SGLang expert-parallel size; it must divide and not exceed every managed capture server's `tp_size`. | | `model.sglang_max_running_requests` | `null` | Positive SGLang request-concurrency limit. | | `model.sglang_max_total_tokens` | `null` | Positive SGLang token-pool limit. | +| `model.sglang_dp_size` | `null` | Optional SGLang data-parallel size. | +| `model.sglang_moe_a2a_backend` | `null` | Optional SGLang MoE all-to-all backend name. | +| `model.sglang_moe_runner_backend` | `null` | Optional SGLang MoE runner backend name. | +| `model.sglang_page_size` | `null` | Optional positive SGLang KV-cache page size. | +| `model.sglang_quantization` | `null` | Optional SGLang target quantization mode. | +| `model.sglang_fp4_gemm_runner_backend` | `null` | Optional SGLang FP4 GEMM runner backend. | +| `model.sglang_mamba_radix_cache_strategy` | `null` | Optional hybrid Mamba/radix cache strategy. | +| `model.sglang_max_mamba_cache_size` | `null` | Optional positive Mamba cache size. | +| `model.sglang_swa_full_tokens_ratio` | `null` | Optional SGLang sliding-window full-token ratio in `(0, 1]`. | +| `model.sglang_mamba_full_memory_ratio` | `null` | Optional SGLang Mamba full-memory ratio in `(0, 1]`. | ### `data`: choose exactly one training source @@ -189,9 +200,11 @@ Common fields: | `training.total_steps` | `null` | Positive optimizer/loss schedule horizon; it does not itself stop an online stream. A finite online disaggregated run may omit both fields: the producer publishes the exact horizon derived from prepared prompts, epochs, DP size, batch size, and accumulation. | | `training.batch_size` | `1` | Per-rank microbatch size. P-EAGLE and USP require 1. | | `training.accumulation_steps` | `1` | Positive microbatches per optimizer update. | +| `training.fsdp_sharding` | `SHARD_GRAD_OP` | Trainer FSDP mode: `SHARD_GRAD_OP`, `FULL_SHARD`, or `NO_SHARD`. | | `training.learning_rate` | `1e-4` | Positive peak learning rate. | | `training.warmup_ratio` | `0.015` | Fraction in `[0, 1]` used for scheduler warmup. | | `training.max_grad_norm` | `0.5` | Positive gradient-clipping norm. | +| `training.optimizer_cpu_offload` | `false` | Keep the optimizer's FP32 master parameters and Adam state on CPU. | | `training.attention_backend` | `flex_attention` | `eager`, `sdpa`, `flex_attention`, `fa`, or `usp`; the selected strategy must support it. | | `training.tp_size` | `1` | Online disaggregated consumers must keep it at 1; configure target TP on capture servers. Offline non-USP ranks consume disjoint data. | | `training.sp_ulysses_size` | `1` | Ulysses sequence-parallel factor for offline EAGLE3 USP. | @@ -212,8 +225,8 @@ Strategy-specific fields should be written only when tuning that objective: | Strategy | Fields and defaults | | --- | --- | | EAGLE3 | `training.ttt_length` (`7`), `training.lk_loss_type` (`null`; `lambda` or `alpha`), `training.kl_scale` (`1.0`), `training.kl_decay` (`1.0`) | -| DFlash / Domino / D-PACE | `training.num_anchors` (`512`), `training.loss_decay_gamma` (`null`), `training.loss_type` (`dflash`), `training.dpace_alpha` (`0.5`), `training.lambda_base_start` (`1.0`), `training.lambda_base_decay_ratio` (`0.5`) | -| DSpark | `training.dspark_ce_loss_alpha` (`0.1`), `training.dspark_l1_loss_alpha` (`0.9`), `training.dspark_confidence_head_alpha` (`1.0`) | +| DFlash / Domino / D-PACE | `training.num_anchors` (`512`), `training.loss_decay_gamma` (`null`), `training.objective_chunk_blocks` (`128`; `0` materializes all objective logits), `training.loss_type` (`dflash`), `training.dpace_alpha` (`0.5`), `training.lambda_base_start` (`1.0`), `training.lambda_base_decay_ratio` (`0.5`) | +| DSpark | Token-pooled objective with valid-first-target anchors and distributed ratio telemetry. Configure the shared `training.num_anchors` (`512`), `training.loss_decay_gamma` (`null`; production recipes use `4.0`), and `training.objective_chunk_blocks` (`128`; `0` materializes all objective logits), plus `training.dspark_ce_loss_alpha` (`0.1`), `training.dspark_l1_loss_alpha` (`0.9`), and `training.dspark_confidence_head_alpha` (`1.0`). | | P-EAGLE | `training.num_depths` (`8`), `training.down_sample_ratio` (`0.8`), `training.down_sample_ratio_min` (`0.2`), `training.norm_before_residual` (`null`) | New recipes must not write the loader-only migration fields diff --git a/examples/configs/glm-5.2-dspark-disaggregated.yaml b/examples/configs/glm-5.2-dspark-disaggregated.yaml new file mode 100644 index 000000000..7cfca7e73 --- /dev/null +++ b/examples/configs/glm-5.2-dspark-disaggregated.yaml @@ -0,0 +1,41 @@ +model: + target_model_path: zai-org/GLM-5.2-FP8 + draft_model_config: configs/glm-5.2-dspark.json + target_backend: sglang + trust_remote_code: true +data: + train_data_path: ./cache/dataset/glm52_dspark_train.jsonl + max_length: 4096 + chat_template: glm-5.2 + cache_dir: cache + build_dataset_num_proc: 64 +training: + strategy: dspark + num_epochs: 10 + # Portable one-rank equivalent of the source recipe's global batch 512. + batch_size: 1 + accumulation_steps: 512 + learning_rate: 0.0006 + warmup_ratio: 0.04 + max_grad_norm: 1.0 + num_anchors: 512 + loss_decay_gamma: 4.0 + objective_chunk_blocks: 128 + # Optimizer-step equivalent of about 500 source microsteps at microbatch 16. + save_interval: 16 + dist_timeout: 30 + seed: 42 +run_id: glm-5.2-dspark-disaggregated +output_dir: outputs/glm-5.2-dspark-disaggregated +deployment: + mode: disaggregated + trainer: + nnodes: 1 + nproc_per_node: 1 + disaggregated: + control_dir: outputs/glm-5.2-dspark-disaggregated/control + backend: mooncake + server_urls: [http://127.0.0.1:30000] + mooncake_metadata_server: http://127.0.0.1:35880/metadata + mooncake_master_server_addr: 127.0.0.1:35551 + mooncake_protocol: tcp diff --git a/examples/configs/inkling-dspark-disaggregated.yaml b/examples/configs/inkling-dspark-disaggregated.yaml new file mode 100644 index 000000000..afadeac74 --- /dev/null +++ b/examples/configs/inkling-dspark-disaggregated.yaml @@ -0,0 +1,43 @@ +model: + target_model_path: thinkingmachines/Inkling + draft_model_config: configs/inkling-dspark.json + target_backend: sglang + tokenizer_pad_token_id: 200006 + embedding_key: model.llm.embed.weight + lm_head_key: model.llm.unembed.weight +data: + train_data_path: ./cache/dataset/inkling_dspark_train.jsonl + max_length: 4096 + chat_template: inkling-thinking + cache_dir: cache + build_dataset_num_proc: 64 +training: + strategy: dspark + num_epochs: 10 + # Portable one-rank equivalent of the source recipe's global batch 512. + batch_size: 1 + accumulation_steps: 512 + learning_rate: 0.0006 + warmup_ratio: 0.04 + max_grad_norm: 1.0 + num_anchors: 512 + loss_decay_gamma: 4.0 + objective_chunk_blocks: 128 + # Optimizer-step equivalent of about 250 source microsteps at microbatch 16. + save_interval: 8 + dist_timeout: 30 + seed: 42 +run_id: inkling-dspark-disaggregated +output_dir: outputs/inkling-dspark-disaggregated +deployment: + mode: disaggregated + trainer: + nnodes: 1 + nproc_per_node: 1 + disaggregated: + control_dir: outputs/inkling-dspark-disaggregated/control + backend: mooncake + server_urls: [http://127.0.0.1:30000] + mooncake_metadata_server: http://127.0.0.1:35880/metadata + mooncake_master_server_addr: 127.0.0.1:35551 + mooncake_protocol: tcp diff --git a/examples/configs/qwen3-4b-dspark-disaggregated.yaml b/examples/configs/qwen3-4b-dspark-disaggregated.yaml index 7cfe46091..962d9c41f 100644 --- a/examples/configs/qwen3-4b-dspark-disaggregated.yaml +++ b/examples/configs/qwen3-4b-dspark-disaggregated.yaml @@ -17,6 +17,8 @@ training: warmup_ratio: 0.04 max_grad_norm: 1.0 num_anchors: 512 + loss_decay_gamma: 4.0 + objective_chunk_blocks: 128 save_interval: 1000 dist_timeout: 30 seed: 42 diff --git a/examples/configs/qwen3-8b-dspark-disaggregated.yaml b/examples/configs/qwen3-8b-dspark-disaggregated.yaml new file mode 100644 index 000000000..0d7c54f03 --- /dev/null +++ b/examples/configs/qwen3-8b-dspark-disaggregated.yaml @@ -0,0 +1,40 @@ +model: + target_model_path: Qwen/Qwen3-8B + draft_model_config: configs/qwen3-8b-dspark.json + target_backend: sglang +data: + train_data_path: ./cache/dataset/perfectblend_raw_900k.jsonl + max_length: 4096 + chat_template: qwen + cache_dir: cache + build_dataset_num_proc: 64 +training: + strategy: dspark + num_epochs: 10 + # Portable one-rank equivalent of the source recipe's global batch 512. + batch_size: 1 + accumulation_steps: 512 + learning_rate: 0.0006 + warmup_ratio: 0.04 + max_grad_norm: 1.0 + num_anchors: 512 + loss_decay_gamma: 4.0 + objective_chunk_blocks: 128 + # Optimizer-step equivalent of 2,000 source microsteps at global microbatch 32. + save_interval: 125 + dist_timeout: 30 + seed: 42 +run_id: qwen3-8b-dspark-disaggregated +output_dir: outputs/qwen3-8b-dspark-disaggregated +deployment: + mode: disaggregated + trainer: + nnodes: 1 + nproc_per_node: 1 + disaggregated: + control_dir: outputs/qwen3-8b-dspark-disaggregated/control + backend: mooncake + server_urls: [http://127.0.0.1:30000] + mooncake_metadata_server: http://127.0.0.1:35880/metadata + mooncake_master_server_addr: 127.0.0.1:35551 + mooncake_protocol: tcp diff --git a/examples/disagg/run_qwen3_8b_dflash_disagg_2node.sh b/examples/disagg/run_qwen3_8b_dflash_disagg_2node.sh index adb93f0f3..60531035c 100755 --- a/examples/disagg/run_qwen3_8b_dflash_disagg_2node.sh +++ b/examples/disagg/run_qwen3_8b_dflash_disagg_2node.sh @@ -236,7 +236,6 @@ run_inference_node() { --tp-size "$SERVER_TP" \ --mem-fraction-static "$SERVER_MEM_FRACTION" \ --chunked-prefill-size -1 \ - --disable-radix-cache \ --enable-spec-capture \ --spec-capture-method dflash \ --spec-capture-aux-layer-ids "${capture_layers[@]}" \ diff --git a/patches/sglang/v0.5.14/spec-capture.patch b/patches/sglang/v0.5.14/spec-capture.patch index 15bcfa45d..33a674213 100644 --- a/patches/sglang/v0.5.14/spec-capture.patch +++ b/patches/sglang/v0.5.14/spec-capture.patch @@ -53,17 +53,14 @@ diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/manager index 951f35495..2359c31b7 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py -@@ -284,6 +284,10 @@ class GenerateReqInput(BaseReq): - # Batch-level: List[List[int]] (one per request). After __getitem__: List[int]. - multi_item_delimiter_indices: Optional[Union[List[List[int]], List[int]]] = None - +@@ -283,3 +283,7 @@ class GenerateReqInput(BaseReq): + # Spec-training capture sink instructions (see spec_capture_sink.py). + # Batch-level: List[Optional[dict]]; per-request after __getitem__. + spec_capture: Optional[Union[List[Optional[Dict]], Dict]] = None + - def contains_mm_input(self) -> bool: - return ( - has_valid_data(self.image_data) + # Pre-computed delimiter indices for multi-item scoring. + # Batch-level: List[List[int]] (one per request). After __getitem__: List[int]. + multi_item_delimiter_indices: Optional[Union[List[List[int]], List[int]]] = None @@ -743,6 +747,11 @@ class GenerateReqInput(BaseReq): if self.multi_item_delimiter_indices is not None else None @@ -76,48 +73,31 @@ index 951f35495..2359c31b7 100644 ) cache[i] = sub return sub -@@ -842,6 +851,9 @@ class TokenizedGenerateReqInput(BaseReq): - # Pre-computed delimiter indices for multi-item scoring - multi_item_delimiter_indices: Optional[List[int]] = None - +@@ -842,2 +851,5 @@ class TokenizedGenerateReqInput(BaseReq): + # Spec-training capture sink instructions (see GenerateReqInput.spec_capture) + spec_capture: Optional[Dict] = None + - # For observability - time_stats: Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] = None - -@@ -1179,6 +1191,9 @@ class BatchTokenIDOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin): - # The trainer step id. Used to know which step's weights are used for sampling. - token_steps: List[List[int]] = None - + # Pre-computed delimiter indices for multi-item scoring + multi_item_delimiter_indices: Optional[List[int]] = None +@@ -1179,1 +1191,4 @@ class BatchTokenIDOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin): + # Spec-training capture: one result dict per request (see spec_capture_sink). + spec_capture: Optional[List[Any]] = None + - # Customized info - customized_info: Optional[Dict[str, List[Any]]] = None - # Detailed breakdown of cached tokens by source (device/host/storage) -@@ -1247,6 +1262,9 @@ class BatchStrOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin): # The trainer step id. Used to know which step's weights are used for sampling. - token_steps: List[List[int]] = None - +@@ -1247,1 +1262,4 @@ class BatchStrOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin): + # Spec-training capture: one result dict per request (see spec_capture_sink). + spec_capture: Optional[List[Any]] = None + - # Customized info - customized_info: Optional[Dict[str, List[Any]]] = None - # Detailed breakdown of cached tokens by source (device/host/storage) + # The trainer step id. Used to know which step's weights are used for sampling. diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index f1dc81d17..a93a466b7 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py -@@ -708,6 +708,7 @@ class Req(ReqDllmMixin): +@@ -708,3 +708,4 @@ class Req(ReqDllmMixin): ] = None, return_pooled_hidden_states: bool = False, - multi_item_delimiter_indices: Optional[List[int]] = None, + spec_capture: Optional[Dict[str, Any]] = None, - ): - # Input and output info - self.rid = rid + multi_item_delimiter_indices: Optional[List[int]] = None, @@ -771,7 +772,13 @@ class Req(ReqDllmMixin): } self.sampling_params = sampling_params @@ -137,7 +117,7 @@ diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/manager index abba37441..3018b8247 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py -@@ -576,6 +576,25 @@ class Scheduler( +@@ -576,6 +576,19 @@ class Scheduler( self.init_batch_result_processor() @@ -150,12 +130,6 @@ index abba37441..3018b8247 100644 + "(single-pass prefill) so captured hidden states cover the " + "whole sequence" + ) -+ if not server_args.disable_radix_cache: -+ raise ValueError( -+ "--enable-spec-capture requires --disable-radix-cache: a " -+ "prefix-cache hit skips prefilling (and capturing) the " -+ "cached tokens, silently truncating the stored hidden states" -+ ) + from sglang.srt import spec_capture_sink + + spec_capture_sink.maybe_init_sink(server_args) @@ -163,7 +137,7 @@ index abba37441..3018b8247 100644 self.is_initializing = False def init_zbal_on_npu(self): -@@ -2054,6 +2067,7 @@ class Scheduler( +@@ -2054,6 +2061,7 @@ class Scheduler( dllm_config=self.dllm_config, time_stats=recv_req.time_stats, multi_item_delimiter_indices=recv_req.multi_item_delimiter_indices, @@ -261,14 +235,11 @@ diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py index f95c59f6f..9e64326a1 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py -@@ -273,6 +273,7 @@ class _GenerationStreamAccumulator: - spec_correct_drafts_histogram: list = field(default_factory=list) +@@ -274,3 +274,4 @@ class _GenerationStreamAccumulator: retraction_counts: list = field(default_factory=list) output_hidden_states: Optional[list] = None + spec_capture: list = field(default_factory=list) routed_experts: Optional[list] = None - indexer_topk: Optional[list] = None - customized_info: dict = field(default_factory=dict) @@ -482,6 +483,8 @@ class _GenerationStreamAccumulator: self.output_hidden_states.append(hs) else: @@ -278,14 +249,11 @@ index f95c59f6f..9e64326a1 100644 if self.return_routed_experts: self.routed_experts.append( req.routed_experts if req.return_routed_experts else None -@@ -540,6 +543,7 @@ class _GenerationStreamAccumulator: - output_token_ids_logprobs_idx=self.output_token_ids_logprobs_idx, +@@ -541,3 +544,4 @@ class _GenerationStreamAccumulator: output_token_entropy_val=None, output_hidden_states=self.output_hidden_states, + spec_capture=self.spec_capture or None, routed_experts=self.routed_experts, - indexer_topk=self.indexer_topk, - customized_info=self.customized_info, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index bf932611a..f0e821ec4 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py @@ -313,10 +281,7 @@ diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/sr index 1cff5c983..a5935fa3c 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py -@@ -471,6 +471,19 @@ class ModelRunner(ModelRunnerKVCacheMixin): - # if there is no aux layer, set to None - self.eagle_aux_hidden_state_layer_ids = None - +@@ -518,3 +518,19 @@ class ModelRunner(ModelRunnerKVCacheMixin): + if server_args.enable_spec_capture and not self.is_draft_worker: + # Aux capture without a draft worker, routed to the strategy's own + # capture method (they wire different submodules — e.g. VL models @@ -324,15 +289,18 @@ index 1cff5c983..a5935fa3c 100644 + if getattr(server_args, "spec_capture_method", "eagle3") == "dflash": + self.dflash_use_aux_hidden_state = True + self.dflash_target_layer_ids = server_args.spec_capture_aux_layer_ids ++ self.dflash_family_use_aux_hidden_state = True ++ self.dflash_family_target_layer_ids = ( ++ server_args.spec_capture_aux_layer_ids ++ ) + else: + self.eagle_use_aux_hidden_state = True + self.eagle_aux_hidden_state_layer_ids = ( + server_args.spec_capture_aux_layer_ids + ) -+ - if self.spec_algorithm.is_dflash() and not self.is_draft_worker: - from sglang.srt.speculative.dflash_utils import parse_dflash_draft_config - + # Apply the rank zero filter to logger + if server_args.show_time_cost: + enable_show_time_cost() diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index c7162c16d..d515b3c69 100644 --- a/python/sglang/srt/server_args.py diff --git a/scripts/gates/run_disaggregated_overfit_gate.sh b/scripts/gates/run_disaggregated_overfit_gate.sh index b449e5085..71d9b8897 100755 --- a/scripts/gates/run_disaggregated_overfit_gate.sh +++ b/scripts/gates/run_disaggregated_overfit_gate.sh @@ -259,7 +259,6 @@ CAPTURE_COMMAND=( --context-length "$MAX_LENGTH" --mem-fraction-static "$CAPTURE_MEM_FRACTION" --chunked-prefill-size -1 - --disable-radix-cache --enable-spec-capture --spec-capture-method "$CAPTURE_METHOD" --spec-capture-aux-layer-ids "${CAPTURE_LAYER_ARGS[@]}" diff --git a/specforge/algorithms/common/dflash_family_model.py b/specforge/algorithms/common/dflash_family_model.py index 20d6f0b7e..48771e224 100644 --- a/specforge/algorithms/common/dflash_family_model.py +++ b/specforge/algorithms/common/dflash_family_model.py @@ -7,6 +7,7 @@ import torch.nn as nn import torch.nn.functional as F +from specforge.core.chunking import checkpointed_chunk_reduce from specforge.modeling.draft.dflash import DFlashDraftModel try: @@ -130,6 +131,7 @@ def __init__( attention_backend: str = "flex_attention", num_anchors: int = 512, loss_decay_gamma: Optional[float] = None, + objective_chunk_blocks: int = 128, loss_type: str = "dflash", dpace_alpha: float = 0.5, ): @@ -140,6 +142,8 @@ def __init__( ) if not 0.0 <= dpace_alpha <= 1.0: raise ValueError(f"dpace_alpha must be in [0, 1], got {dpace_alpha}") + if objective_chunk_blocks < 0: + raise ValueError("objective_chunk_blocks must be >= 0") self.draft_model = draft_model self.lm_head = target_lm_head @@ -149,6 +153,7 @@ def __init__( self.attention_backend = attention_backend self.num_anchors = num_anchors self.loss_decay_gamma = loss_decay_gamma + self.objective_chunk_blocks = int(objective_chunk_blocks) self.loss_type = loss_type self.dpace_alpha = dpace_alpha @@ -307,6 +312,59 @@ def _forward_draft_blocks( ) return anchor_positions, block_keep_mask, output_hidden + def _dflash_objective_chunk_terms( + self, + hidden: torch.Tensor, + target_ids: torch.Tensor, + weight_mask: torch.Tensor, + ) -> Tuple[torch.Tensor, ...]: + """Return additive DFlash/D-PACE loss and accuracy terms.""" + + batch_size, num_blocks, block_size, hidden_size = hidden.shape + logits = self.lm_head( + hidden.reshape(batch_size, num_blocks * block_size, hidden_size) + ).reshape(batch_size, num_blocks, block_size, -1) + neg_log_q = F.cross_entropy( + logits.reshape(-1, logits.shape[-1]), + target_ids.reshape(-1), + reduction="none", + ).reshape_as(target_ids) + + if self.loss_type == "dflash": + loss_weights = weight_mask + if self.loss_decay_gamma is not None and self.loss_decay_gamma > 0: + positions = torch.arange( + self.block_size, + device=hidden.device, + ).view(1, 1, -1) + decay_weights = torch.exp( + -(positions - 1).clamp(min=0).float() / self.loss_decay_gamma + ) + loss_weights = loss_weights * decay_weights + loss_num = (neg_log_q * loss_weights).sum() + loss_den = loss_weights.sum() + elif self.loss_type in _DPACE_LOSS_TYPES: + with torch.no_grad(): + target_probability = torch.exp(-neg_log_q) + dpace_weights = self._dpace_weight( + target_probability, + weight_mask, + weight_mask > 0, + self.loss_type, + ) + loss_num = (neg_log_q * weight_mask * dpace_weights).sum() + loss_den = loss_num.new_zeros(()) + else: # defensive: __init__ validates the configured loss type. + raise ValueError(f"unknown loss_type {self.loss_type!r}") + + with torch.no_grad(): + predicted_ids = logits.argmax(dim=-1) + correct_num = ( + ((predicted_ids == target_ids) & (weight_mask > 0.5)).sum().float() + ) + accuracy_den = weight_mask.sum() + return loss_num, loss_den, correct_num, accuracy_den + def forward( self, input_ids: torch.Tensor, @@ -328,8 +386,6 @@ def forward( loss_mask=loss_mask, ) - logits = self.lm_head(output_hidden) - # --- Labels: same-position prediction (position k predicts token anchor+k) --- label_offsets = torch.arange(0, self.block_size, device=device).view(1, 1, -1) label_indices = anchor_positions.unsqueeze(-1) + label_offsets @@ -358,48 +414,25 @@ def forward( ) weight_mask = weight_mask * original_loss_mask_gathered - binary_eval_mask = weight_mask.view(-1) - - # --- Cross entropy --- - flat_logits = logits.view(-1, logits.size(-1)) - flat_targets = target_ids.view(-1) - - loss_per_token = F.cross_entropy(flat_logits, flat_targets, reduction="none") - + hidden_4d = output_hidden.reshape( + bsz, + anchor_positions.shape[1], + self.block_size, + -1, + ) + loss_num, loss_den, correct_num, accuracy_denom = checkpointed_chunk_reduce( + self._dflash_objective_chunk_terms, + hidden_4d, + target_ids, + weight_mask, + chunk_size=self.objective_chunk_blocks, + dim=1, + ) if self.loss_type == "dflash": - # Preserve the existing DFlash weighted-mean behavior. - loss_weights = weight_mask - if self.loss_decay_gamma is not None and self.loss_decay_gamma > 0: - k = torch.arange(self.block_size, device=device).view(1, 1, -1) - decay_weights = torch.exp( - -(k - 1).clamp(min=0).float() / self.loss_decay_gamma - ) - loss_weights = loss_weights * decay_weights - - flat_weights = loss_weights.view(-1) - valid_token_count = flat_weights.sum() + 1e-6 - loss = (loss_per_token * flat_weights).sum() / valid_token_count - elif self.loss_type in _DPACE_LOSS_TYPES: - neg_log_q = loss_per_token.view_as(target_ids) - with torch.no_grad(): - q = torch.exp(-neg_log_q) - dpace_weights = self._dpace_weight( - q, - weight_mask, - weight_mask > 0, - self.loss_type, - ) - loss_weights = weight_mask * dpace_weights - loss = (neg_log_q * loss_weights).sum() / float(bsz) + loss = loss_num / (loss_den + 1e-6) else: - raise ValueError(f"unknown loss_type {self.loss_type!r}") - - # --- Accuracy --- - with torch.no_grad(): - pred_ids = torch.argmax(flat_logits, dim=-1) - correct = (pred_ids == flat_targets) & (binary_eval_mask > 0.5) - accuracy_denom = binary_eval_mask.sum() - accuracy = correct.sum().float() / (accuracy_denom + 1e-6) + loss = loss_num / float(bsz) + accuracy = correct_num / (accuracy_denom + 1e-6) return loss, accuracy, {"accuracy_denom": accuracy_denom.detach()} @@ -417,6 +450,7 @@ def __init__( attention_backend: str = "flex_attention", num_anchors: int = 512, loss_decay_gamma: Optional[float] = None, + objective_chunk_blocks: int = 128, shift_label: bool = False, ): super().__init__( @@ -428,6 +462,7 @@ def __init__( attention_backend=attention_backend, num_anchors=num_anchors, loss_decay_gamma=loss_decay_gamma, + objective_chunk_blocks=objective_chunk_blocks, loss_type="dflash", ) self.shift_label = shift_label @@ -508,86 +543,77 @@ def _apply_domino_head( prev_token_embeddings=head_token_embeddings, ) - def _compute_extra_metrics( - self, - pred_ids: torch.Tensor, - flat_base_logits: torch.Tensor, - flat_targets: torch.Tensor, - binary_eval_mask: torch.Tensor, - actual_token_count: torch.Tensor, - target_ids: torch.Tensor, - eval_weight_mask: torch.Tensor, - final_loss: torch.Tensor, - base_loss: torch.Tensor, - lambda_base: float, - ) -> Dict[str, torch.Tensor]: - bsz, n, bs = target_ids.shape - - base_pred_ids = torch.argmax(flat_base_logits, dim=-1) - base_correct = (base_pred_ids == flat_targets) & (binary_eval_mask > 0.5) - base_accuracy = base_correct.sum().float() / actual_token_count - - valid_mask_4d = (eval_weight_mask > 0).bool() - pred_accept_len = compute_accept_len( - pred_ids.view(bsz, n, bs), target_ids, valid_mask_4d - ) - base_accept_len = compute_accept_len( - base_pred_ids.view(bsz, n, bs), target_ids, valid_mask_4d - ) - - valid_block_mask = valid_mask_4d.any(dim=2) - num_valid_blocks = valid_block_mask.sum().float() + 1e-6 - avg_accept_len = ( - (pred_accept_len + 1.0) * valid_block_mask.float() - ).sum() / num_valid_blocks - base_avg_accept_len = ( - (base_accept_len + 1.0) * valid_block_mask.float() - ).sum() / num_valid_blocks - - return { - "final_loss": final_loss.detach(), - "base_loss": base_loss.detach(), - "base_accuracy": base_accuracy.detach(), - "accept_len": avg_accept_len.detach(), - "base_accept_len": base_avg_accept_len.detach(), - "lambda_base": torch.tensor(lambda_base, device=final_loss.device), - } - - def _compute_weighted_losses( + def _domino_objective_chunk_terms( self, - final_logits: torch.Tensor, - base_logits: torch.Tensor, + hidden: torch.Tensor, + prev_ids: torch.Tensor, target_ids: torch.Tensor, weight_mask: torch.Tensor, - lambda_base: float, - ) -> Tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - ]: - flat_logits = final_logits.reshape(-1, final_logits.size(-1)) - flat_base_logits = base_logits.reshape(-1, base_logits.size(-1)) - flat_targets = target_ids.reshape(-1) - flat_weights = weight_mask.reshape(-1) - - valid_token_count = flat_weights.sum() + 1e-6 + eval_weight_mask: torch.Tensor, + ) -> Tuple[torch.Tensor, ...]: + """Return additive Domino loss and telemetry terms for one block slice.""" - final_loss_per_token = F.cross_entropy( - flat_logits, flat_targets, reduction="none" + batch_size, num_blocks, block_size, hidden_size = hidden.shape + base_logits = self.lm_head( + hidden.reshape(batch_size, num_blocks * block_size, hidden_size) + ).reshape(batch_size, num_blocks, block_size, -1) + final_logits = self._apply_domino_head( + base_logits4d=base_logits, + hidden4d=hidden, + prev_ids=prev_ids, + target_ids=target_ids, ) - final_loss = (final_loss_per_token * flat_weights).sum() / valid_token_count + final_ce = F.cross_entropy( + final_logits.reshape(-1, final_logits.shape[-1]), + target_ids.reshape(-1), + reduction="none", + ).reshape_as(target_ids) + base_ce = F.cross_entropy( + base_logits.reshape(-1, base_logits.shape[-1]), + target_ids.reshape(-1), + reduction="none", + ).reshape_as(target_ids) + final_num = (final_ce * weight_mask).sum() + base_num = (base_ce * weight_mask).sum() + loss_den = weight_mask.sum() - base_loss_per_token = F.cross_entropy( - flat_base_logits, flat_targets, reduction="none" + with torch.no_grad(): + predicted_ids = final_logits.argmax(dim=-1) + base_predicted_ids = base_logits.argmax(dim=-1) + binary_accuracy_mask = eval_weight_mask > 0.5 + correct_num = ( + ((predicted_ids == target_ids) & binary_accuracy_mask).sum().float() + ) + base_correct_num = ( + ((base_predicted_ids == target_ids) & binary_accuracy_mask) + .sum() + .float() + ) + accuracy_den = eval_weight_mask.sum() + + valid_mask = eval_weight_mask > 0 + accepted = compute_accept_len(predicted_ids, target_ids, valid_mask) + base_accepted = compute_accept_len( + base_predicted_ids, + target_ids, + valid_mask, + ) + valid_blocks = valid_mask.any(dim=-1).float() + accept_num = ((accepted + 1.0) * valid_blocks).sum() + base_accept_num = ((base_accepted + 1.0) * valid_blocks).sum() + accept_den = valid_blocks.sum() + + return ( + final_num, + base_num, + loss_den, + correct_num, + base_correct_num, + accuracy_den, + accept_num, + base_accept_num, + accept_den, ) - base_loss = (base_loss_per_token * flat_weights).sum() / valid_token_count - - loss = (1.0 - lambda_base) * final_loss + lambda_base * base_loss - - return loss, final_loss, base_loss, flat_logits, flat_base_logits, flat_targets def forward( self, @@ -625,21 +651,12 @@ def forward( ) bsz, n, bs = target_ids.shape - base_logits = self.lm_head(output_hidden) hidden4d, prev_ids = self._build_domino_head_inputs( input_ids=input_ids, anchor_positions=anchor_positions, target_ids=target_ids, output_hidden=output_hidden, ) - base_logits4d = base_logits.reshape(bsz, n, bs, -1) - final_logits = self._apply_domino_head( - base_logits4d=base_logits4d, - hidden4d=hidden4d, - prev_ids=prev_ids, - target_ids=target_ids, - ).reshape(bsz, n * bs, -1) - weight_mask = ( block_keep_mask.unsqueeze(-1).expand(-1, -1, self.block_size).float() ) @@ -657,7 +674,6 @@ def forward( weight_mask = weight_mask * original_loss_mask_gathered eval_weight_mask = weight_mask.clone() - binary_eval_mask = weight_mask.view(-1) if self.loss_decay_gamma is not None and self.loss_decay_gamma > 0: k = torch.arange(self.block_size, device=device).view(1, 1, -1) @@ -667,36 +683,41 @@ def forward( ) weight_mask = weight_mask * decay_weights - loss, final_loss, base_loss, flat_logits, flat_base_logits, flat_targets = ( - self._compute_weighted_losses( - final_logits=final_logits, - base_logits=base_logits, - target_ids=target_ids, - weight_mask=weight_mask, - lambda_base=lambda_base, - ) + ( + final_num, + base_num, + loss_den, + correct_num, + base_correct_num, + accuracy_denom, + accept_num, + base_accept_num, + accept_den, + ) = checkpointed_chunk_reduce( + self._domino_objective_chunk_terms, + hidden4d, + prev_ids, + target_ids, + weight_mask, + eval_weight_mask, + chunk_size=self.objective_chunk_blocks, + dim=1, ) - with torch.no_grad(): - pred_ids = torch.argmax(flat_logits, dim=-1) - correct = (pred_ids == flat_targets) & (binary_eval_mask > 0.5) - accuracy_denom = binary_eval_mask.sum() - actual_token_count = accuracy_denom + 1e-6 - accuracy = correct.sum().float() / actual_token_count - - metrics = self._compute_extra_metrics( - pred_ids=pred_ids, - flat_base_logits=flat_base_logits, - flat_targets=flat_targets, - binary_eval_mask=binary_eval_mask, - actual_token_count=actual_token_count, - target_ids=target_ids, - eval_weight_mask=eval_weight_mask, - final_loss=final_loss, - base_loss=base_loss, - lambda_base=lambda_base, - ) - metrics["accuracy_denom"] = accuracy_denom.detach() + valid_token_count = loss_den + 1e-6 + final_loss = final_num / valid_token_count + base_loss = base_num / valid_token_count + loss = (1.0 - lambda_base) * final_loss + lambda_base * base_loss + accuracy = correct_num / (accuracy_denom + 1e-6) + metrics = { + "final_loss": final_loss.detach(), + "base_loss": base_loss.detach(), + "base_accuracy": (base_correct_num / (accuracy_denom + 1e-6)).detach(), + "accept_len": (accept_num / (accept_den + 1e-6)).detach(), + "base_accept_len": (base_accept_num / (accept_den + 1e-6)).detach(), + "lambda_base": torch.tensor(lambda_base, device=loss.device), + "accuracy_denom": accuracy_denom.detach(), + } return loss, accuracy, metrics @@ -717,6 +738,7 @@ def __init__( dspark_ce_loss_alpha: float = 0.1, dspark_l1_loss_alpha: float = 0.9, dspark_confidence_head_alpha: float = 1.0, + objective_chunk_blocks: int = 128, ): super().__init__( draft_model=draft_model, @@ -727,6 +749,7 @@ def __init__( attention_backend=attention_backend, num_anchors=num_anchors, loss_decay_gamma=loss_decay_gamma, + objective_chunk_blocks=objective_chunk_blocks, loss_type="dflash", ) if dspark_ce_loss_alpha < 0: @@ -735,7 +758,6 @@ def __init__( raise ValueError("dspark_l1_loss_alpha must be >= 0") if dspark_confidence_head_alpha < 0: raise ValueError("dspark_confidence_head_alpha must be >= 0") - self.loss_type = "dspark" self.dspark_ce_loss_alpha = float(dspark_ce_loss_alpha) self.dspark_l1_loss_alpha = float(dspark_l1_loss_alpha) @@ -756,41 +778,37 @@ def _build_anchor_candidate_mask( def _sample_anchor_positions( self, seq_len: int, loss_mask: torch.Tensor, device: torch.device ) -> Tuple[torch.Tensor, torch.Tensor]: - """Sample fixed-width DSpark anchors; invalid slots are masked out.""" - valid = self._build_anchor_candidate_mask(seq_len, loss_mask) - bsz = loss_mask.shape[0] - num_candidates = valid.shape[1] - max_n = int(self.num_anchors) - if num_candidates == 0: - anchors = torch.zeros(bsz, max_n, dtype=torch.long, device=device) - keep_mask = torch.zeros(bsz, max_n, dtype=torch.bool, device=device) - return anchors, keep_mask + """Sample only anchors with a valid first target, without dummy width.""" + valid = self._build_anchor_candidate_mask(seq_len, loss_mask) + if valid.shape[1] == 0: + raise ValueError("DSpark needs sequences with at least two tokens") valid_counts = valid.sum(dim=1) - indices = ( - torch.arange(num_candidates, device=device).unsqueeze(0).expand(bsz, -1) - ) - masked_indices = torch.where( - valid, - indices, - torch.full_like(indices, seq_len + 1), + width = min(self.num_anchors, int(valid_counts.max().item())) + if width <= 0: + raise ValueError( + "DSpark found no valid anchor with two consecutive loss tokens" + ) + indices = torch.arange(valid.shape[1], device=device).expand( + loss_mask.shape[0], -1 ) - random_vals = torch.rand(bsz, num_candidates, device=device) - random_vals = torch.where(valid, random_vals, torch.full_like(random_vals, 2.0)) - _, sorted_idx = random_vals.sort(dim=1) - gathered = torch.gather(masked_indices, 1, sorted_idx) - if num_candidates < max_n: - pad = torch.full( - (bsz, max_n - num_candidates), - seq_len + 1, - dtype=gathered.dtype, - device=device, + random_values = torch.rand(valid.shape, device=device) + random_values.masked_fill_(~valid, 2.0) + order = random_values.argsort(dim=1) + candidates = torch.gather(indices, 1, order)[:, :width] + keep_mask = torch.arange(width, device=device).unsqueeze( + 0 + ) < valid_counts.clamp(max=width).unsqueeze(1) + anchors = ( + torch.where( + keep_mask, + candidates, + torch.full_like(candidates, valid.shape[1]), ) - gathered = torch.cat([gathered, pad], dim=1) - anchors = gathered[:, :max_n].sort(dim=1).values - keep_mask = torch.arange(max_n, device=device).unsqueeze(0) < ( - valid_counts.unsqueeze(1).clamp(max=max_n) + .sort(dim=1) + .values ) + keep_mask = anchors < valid.shape[1] anchors = torch.where(keep_mask, anchors, torch.zeros_like(anchors)) return anchors, keep_mask @@ -843,103 +861,265 @@ def _dspark_loss_weight_mask( loss_weight_mask = loss_weight_mask * decay_weights return loss_weight_mask - def _aligned_target_logits( + def _aligned_target_hidden( self, - target_last_hidden_states: Optional[torch.Tensor], + target_last_hidden_states: torch.Tensor, safe_label_indices: torch.Tensor, - ) -> Optional[torch.Tensor]: - if target_last_hidden_states is None: - return None + ) -> torch.Tensor: + """Gather the target state that predicts each DSpark label token.""" + target_pred_indices = (safe_label_indices - 1).clamp(min=0) - aligned_target_hidden = torch.gather( - target_last_hidden_states.unsqueeze(1).expand( - -1, - safe_label_indices.size(1), - -1, - -1, - ), - 2, - target_pred_indices.unsqueeze(-1).expand( - -1, - -1, - -1, - target_last_hidden_states.size(-1), - ), + batch_size = target_last_hidden_states.shape[0] + hidden_size = target_last_hidden_states.shape[-1] + gather_indices = target_pred_indices.reshape(batch_size, -1, 1).expand( + -1, -1, hidden_size ) - return self.lm_head(aligned_target_hidden) + return torch.gather( + target_last_hidden_states, + 1, + gather_indices, + ).reshape(*safe_label_indices.shape, hidden_size) - def _compute_dspark_loss( + def _dspark_objective_chunk_terms( self, - *, - draft_logits: torch.Tensor, + hidden: torch.Tensor, + prev_token_ids: torch.Tensor, target_ids: torch.Tensor, + loss_weights: torch.Tensor, eval_mask: torch.Tensor, - confidence_pred: Optional[torch.Tensor], - aligned_target_logits: Optional[torch.Tensor], - ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: - vocab_size = draft_logits.size(-1) - loss_weight_mask = self._dspark_loss_weight_mask(eval_mask) - flat_logits = draft_logits.reshape(-1, vocab_size) - flat_targets = target_ids.reshape(-1) - flat_weights = loss_weight_mask.reshape(-1) - - loss_per_token = F.cross_entropy(flat_logits, flat_targets, reduction="none") - ce_loss_den = flat_weights.sum() - ce_loss = (loss_per_token * flat_weights).sum() / (ce_loss_den + 1e-6) - - l1_loss = ce_loss.new_zeros(()) - accept_rate_3d = None - needs_target_distribution = ( - self.dspark_l1_loss_alpha > 0 or confidence_pred is not None - ) - if aligned_target_logits is not None and needs_target_distribution: - draft_probs = torch.softmax(draft_logits.float(), dim=-1) - target_probs = torch.softmax(aligned_target_logits.float(), dim=-1) - l1_dist = (draft_probs - target_probs).abs().sum(dim=-1) - accept_rate_3d = 1.0 - 0.5 * l1_dist - accept_rate_3d = accept_rate_3d.clamp_(0.0, 1.0) - if self.dspark_l1_loss_alpha > 0: - l1_loss = (l1_dist * loss_weight_mask).sum() / (ce_loss_den + 1e-6) - elif self.dspark_l1_loss_alpha > 0 or self.dspark_confidence_head_alpha > 0: - raise ValueError( - "DSpark L1/confidence loss requires target_last_hidden_states. " - "Use the disaggregated DSpark server-capture path so the " - "consumer receives target_last_hidden_states." + aligned_target_hidden: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, ...]: + """Return additive loss and telemetry numerators for one block slice.""" + + batch_size, num_blocks, block_size, hidden_size = hidden.shape + base_logits = self.lm_head( + hidden.reshape(batch_size, num_blocks * block_size, hidden_size) + ).reshape(batch_size, num_blocks, block_size, -1) + draft_logits = self.draft_model.apply_logits_head( + base_logits, + prev_token_ids=prev_token_ids, + hidden_states=hidden, + ) + vocab_size = draft_logits.shape[-1] + cross_entropy = F.cross_entropy( + draft_logits.reshape(-1, vocab_size), + target_ids.reshape(-1), + reduction="none", + ).reshape_as(target_ids) + ce_num = (cross_entropy * loss_weights).sum() + + zero = ce_num.new_zeros(()) + l1_num = zero + confidence_num = zero + confidence_error_num = zero + teacher_agreement_num = zero + teacher_top1_num = zero + draft_top1_num = zero + tau_num = zero + tau_den = zero + accept_probability = None + + draft_probabilities = None + teacher_ids = None + if aligned_target_hidden is not None: + with torch.no_grad(): + target_logits = self.lm_head( + aligned_target_hidden.reshape( + batch_size, + num_blocks * block_size, + hidden_size, + ) + ).reshape_as(draft_logits) + target_probabilities = torch.softmax(target_logits.float(), dim=-1) + teacher_ids = target_logits.argmax(dim=-1) + draft_probabilities = torch.softmax(draft_logits.float(), dim=-1) + l1_per_token = ( + (draft_probabilities - target_probabilities).abs().sum(dim=-1) ) + accept_probability = (1.0 - 0.5 * l1_per_token).clamp(0.0, 1.0) + if self.dspark_l1_loss_alpha > 0: + l1_num = (l1_per_token * loss_weights).sum() - confidence_loss = ce_loss.new_zeros(()) - confidence_abs_error = ce_loss.new_zeros(()) - if confidence_pred is not None: - if accept_rate_3d is None: + confidence_pred = self.draft_model.predict_confidence( + hidden, + prev_token_ids=prev_token_ids, + ) + if confidence_pred is not None and self.dspark_confidence_head_alpha > 0: + if accept_probability is None: raise ValueError( - "DSpark confidence head requires aligned target logits." + "DSpark confidence loss requires target_last_hidden_states" ) - confidence_errors = F.binary_cross_entropy_with_logits( + confidence_per_token = F.binary_cross_entropy_with_logits( confidence_pred.float(), - accept_rate_3d.detach(), + accept_probability.detach(), reduction="none", ) - confidence_loss = (confidence_errors * loss_weight_mask).sum() / ( - ce_loss_den + 1e-6 + confidence_num = (confidence_per_token * loss_weights).sum() + confidence_error_num = ( + (confidence_pred.float().sigmoid() - accept_probability).abs() + * loss_weights + ).sum() + + with torch.no_grad(): + predicted_ids = draft_logits.argmax(dim=-1) + correct = ((predicted_ids == target_ids) & eval_mask).float() + correct_num = correct.sum() + eval_den = eval_mask.float().sum() + ce_position_num = (cross_entropy.detach() * eval_mask).sum(dim=(0, 1)) + correct_position_num = correct.sum(dim=(0, 1)) + position_den = eval_mask.float().sum(dim=(0, 1)) + if aligned_target_hidden is not None: + assert draft_probabilities is not None and teacher_ids is not None + teacher_agreement_num = ( + (predicted_ids == teacher_ids).float() * eval_mask + ).sum() + teacher_top1_num = ( + target_probabilities.max(dim=-1).values * eval_mask + ).sum() + draft_top1_num = ( + draft_probabilities.max(dim=-1).values * eval_mask + ).sum() + valid_blocks = eval_mask.any(dim=-1).float() + accepted_expectation = ( + accept_probability.detach() * eval_mask + ).cumprod(dim=-1).sum(dim=-1) + 1.0 + tau_num = (accepted_expectation * valid_blocks).sum() + tau_den = valid_blocks.sum() + + return ( + ce_num, + l1_num, + confidence_num, + confidence_error_num, + correct_num, + eval_den, + ce_position_num, + correct_position_num, + position_den, + teacher_agreement_num, + teacher_top1_num, + draft_top1_num, + tau_num, + tau_den, + ) + + def _compute_dspark_loss( + self, + *, + output_hidden: torch.Tensor, + target_ids: torch.Tensor, + eval_mask: torch.Tensor, + prev_token_ids: torch.Tensor, + safe_label_indices: torch.Tensor, + target_last_hidden_states: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, Dict[str, object]]: + """Token-pooled DSpark objective with bounded vocab-logit memory.""" + + batch_size, num_blocks, block_size = target_ids.shape + hidden_4d = output_hidden.reshape( + batch_size, + num_blocks, + block_size, + -1, + ) + loss_weights = self._dspark_loss_weight_mask(eval_mask) + local_loss_den = loss_weights.sum() + need_target = self.dspark_l1_loss_alpha > 0 or ( + self.dspark_confidence_head_alpha > 0 + and getattr(self.draft_model, "confidence_head", None) is not None + ) + aligned_target_hidden = None + if need_target: + if target_last_hidden_states is None: + raise ValueError( + "DSpark L1/confidence loss requires target_last_hidden_states" + ) + aligned_target_hidden = self._aligned_target_hidden( + target_last_hidden_states, + safe_label_indices, ) - with torch.no_grad(): - confidence_abs_error = ( - (confidence_pred.float().sigmoid() - accept_rate_3d).abs() - * loss_weight_mask - ).sum() / (ce_loss_den + 1e-6) + totals = checkpointed_chunk_reduce( + self._dspark_objective_chunk_terms, + hidden_4d, + prev_token_ids, + target_ids, + loss_weights, + eval_mask, + aligned_target_hidden, + chunk_size=self.objective_chunk_blocks, + dim=1, + ) + + ( + ce_num, + l1_num, + confidence_num, + confidence_error_num, + correct_num, + eval_den, + ce_position_num, + correct_position_num, + position_den, + teacher_agreement_num, + teacher_top1_num, + draft_top1_num, + tau_num, + tau_den, + ) = totals + + global_loss_den = local_loss_den.detach().clone() + world_size = 1 + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + world_size = dist.get_world_size() + dist.all_reduce(global_loss_den, op=dist.ReduceOp.SUM) + if float(global_loss_den) <= 0: + raise ValueError("DSpark objective has no supervised target tokens") loss = ( - self.dspark_ce_loss_alpha * ce_loss - + self.dspark_l1_loss_alpha * l1_loss - + self.dspark_confidence_head_alpha * confidence_loss + world_size + * ( + self.dspark_ce_loss_alpha * ce_num + + self.dspark_l1_loss_alpha * l1_num + + self.dspark_confidence_head_alpha * confidence_num + ) + / global_loss_den ) - metrics = { - "ce_loss": ce_loss.detach(), - "l1_loss": l1_loss.detach(), - "confidence_loss": confidence_loss.detach(), - "confidence_abs_error": confidence_abs_error.detach(), + + ratio_metrics = { + "acc": (correct_num, eval_den), + "ce_loss": (ce_num.detach(), local_loss_den.detach()), + "l1_loss": (l1_num.detach(), local_loss_den.detach()), + "confidence_loss": ( + confidence_num.detach(), + local_loss_den.detach(), + ), + "confidence_abs_error": ( + confidence_error_num.detach(), + local_loss_den.detach(), + ), + "ce_position": (ce_position_num, position_den), + "accuracy_position": (correct_position_num, position_den), } - return loss, metrics + if aligned_target_hidden is not None: + ratio_metrics.update( + { + "teacher_agreement": (teacher_agreement_num, eval_den), + "teacher_top1_prob": (teacher_top1_num, eval_den), + "draft_top1_prob": (draft_top1_num, eval_den), + "tau_probabilistic": (tau_num, tau_den), + } + ) + metrics: Dict[str, object] = { + "ratio_metrics": { + name: (numerator.detach(), denominator.detach()) + for name, (numerator, denominator) in ratio_metrics.items() + }, + "accuracy_denom": eval_den.detach(), + } + accuracy = correct_num / eval_den.clamp_min(1.0) + return loss, {"accuracy": accuracy.detach(), **metrics} def forward( self, @@ -947,22 +1127,18 @@ def forward( hidden_states: torch.Tensor, loss_mask: torch.Tensor, target_last_hidden_states: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor]]: + ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, object]]: """Parallel DSpark training forward pass.""" if self.attention_backend == "flex_attention" and not FLEX_ATTENTION_AVAILABLE: raise ValueError( "flex_attention is not available on this device; use sdpa/eager." ) - bsz = input_ids.shape[0] anchor_positions, block_keep_mask, output_hidden = self._forward_draft_blocks( input_ids=input_ids, hidden_states=hidden_states, loss_mask=loss_mask, ) - logits = self.lm_head(output_hidden) - num_blocks = anchor_positions.size(1) - output_hidden_4d = output_hidden.reshape(bsz, num_blocks, self.block_size, -1) ( target_ids, eval_mask, @@ -978,34 +1154,13 @@ def forward( [anchor_token_ids.unsqueeze(-1), target_ids[:, :, :-1]], dim=-1, ) - draft_logits = logits.reshape(bsz, num_blocks, self.block_size, -1) - draft_logits = self.draft_model.apply_logits_head( - draft_logits, - prev_token_ids=prev_token_ids, - hidden_states=output_hidden_4d, - ) - confidence_pred = self.draft_model.predict_confidence( - output_hidden_4d, - prev_token_ids=prev_token_ids, - ) - aligned_target_logits = self._aligned_target_logits( - target_last_hidden_states, - safe_label_indices, - ) loss, metrics = self._compute_dspark_loss( - draft_logits=draft_logits, + output_hidden=output_hidden, target_ids=target_ids, eval_mask=eval_mask, - confidence_pred=confidence_pred, - aligned_target_logits=aligned_target_logits, + prev_token_ids=prev_token_ids, + safe_label_indices=safe_label_indices, + target_last_hidden_states=target_last_hidden_states, ) - flat_logits = draft_logits.reshape(-1, draft_logits.size(-1)) - flat_targets = target_ids.reshape(-1) - binary_eval_mask = eval_mask.reshape(-1) - with torch.no_grad(): - pred_ids = torch.argmax(flat_logits, dim=-1) - correct = (pred_ids == flat_targets) & binary_eval_mask - accuracy_denom = binary_eval_mask.to(torch.float32).sum() - accuracy = correct.sum().float() / (accuracy_denom + 1e-6) - metrics["accuracy_denom"] = accuracy_denom.detach() + accuracy = metrics.pop("accuracy") return loss, accuracy, metrics diff --git a/specforge/algorithms/model_providers.py b/specforge/algorithms/model_providers.py index e8bd66ecc..d5385eb35 100644 --- a/specforge/algorithms/model_providers.py +++ b/specforge/algorithms/model_providers.py @@ -320,6 +320,7 @@ def _build_dflash_family_model( "attention_backend": cfg.training.attention_backend, "num_anchors": cfg.training.num_anchors, "loss_decay_gamma": cfg.training.loss_decay_gamma, + "objective_chunk_blocks": cfg.training.objective_chunk_blocks, } model = model_factory(common).to(device=_device(), dtype=_torch_dtype(cfg)) return AlgorithmModelParts( diff --git a/specforge/benchmarks/__init__.py b/specforge/benchmarks/__init__.py new file mode 100644 index 000000000..821b93413 --- /dev/null +++ b/specforge/benchmarks/__init__.py @@ -0,0 +1 @@ +"""Runtime benchmarks shipped with SpecForge.""" diff --git a/specforge/benchmarks/sglang.py b/specforge/benchmarks/sglang.py new file mode 100644 index 000000000..de4725b8b --- /dev/null +++ b/specforge/benchmarks/sglang.py @@ -0,0 +1,220 @@ +"""Benchmark a running SGLang server. + +Request scheduling is adapted from z-lab/dflash's MIT benchmark: +https://github.com/z-lab/dflash/blob/main/dflash/benchmark.py + +The runner is speculative-algorithm agnostic and consumes SGLang's optional +speculative-decoding metadata when the server returns it. +""" + +from __future__ import annotations + +import json +import random +import statistics +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import asdict, dataclass +from typing import Any, Optional + +DATASETS: dict[str, dict[str, Any]] = { + "gsm8k": { + "load_args": ("openai/gsm8k", "main"), + "load_kwargs": {"split": "test"}, + "format": lambda row: ( + f"{row['question']}\nPlease reason step by step, and put your final " + "answer within \\boxed{}." + ), + }, + "math500": { + "load_args": ("HuggingFaceH4/MATH-500",), + "load_kwargs": {"split": "test"}, + "format": lambda row: ( + f"{row['problem']}\nPlease reason step by step, and put your final " + "answer within \\boxed{}." + ), + }, + "humaneval": { + "load_args": ("openai/openai_humaneval",), + "load_kwargs": {"split": "test"}, + "format": lambda row: ( + "Write a solution to the following problem and make sure that it " + f"passes the tests:\n```python\n{row['prompt']}\n```" + ), + }, + "mbpp": { + "load_args": ("google-research-datasets/mbpp", "sanitized"), + "load_kwargs": {"split": "test"}, + "format": lambda row: row["prompt"], + }, + "mt-bench": { + "load_args": ("HuggingFaceH4/mt_bench_prompts",), + "load_kwargs": {"split": "train"}, + "format": lambda row: row["prompt"], + "multi_turn": True, + }, +} + + +@dataclass(frozen=True) +class BenchmarkResult: + backend: str + dataset: str + samples: int + output_tokens: int + latency_seconds: float + throughput_tokens_per_second: float + average_acceptance_length: Optional[float] = None + spec_verify_count: Optional[int] = None + + +def _load_prompts(name: str, max_samples: Optional[int]) -> list[list[str]]: + from datasets import load_dataset + + if max_samples is not None and max_samples <= 0: + raise ValueError("--max-samples must be positive") + descriptor = DATASETS[name] + dataset = load_dataset( + *descriptor["load_args"], + **descriptor["load_kwargs"], + ) + prompts: list[list[str]] = [] + for row in dataset: + formatted = descriptor["format"](row) + turns = list(formatted) if descriptor.get("multi_turn") else [formatted] + prompts.append(turns) + if not prompts: + raise ValueError(f"dataset {name!r} did not contain any prompts") + if max_samples is not None and len(prompts) > max_samples: + random.Random(42).shuffle(prompts) + prompts = prompts[:max_samples] + return prompts + + +def _apply_chat_template(tokenizer, messages, enable_thinking: bool) -> str: + kwargs = { + "tokenize": False, + "add_generation_prompt": True, + "enable_thinking": enable_thinking, + } + try: + return tokenizer.apply_chat_template(messages, **kwargs) + except TypeError: + kwargs.pop("enable_thinking") + return tokenizer.apply_chat_template(messages, **kwargs) + + +def _send_sglang(args, prompt: str) -> dict[str, Any]: + import requests + + response = requests.post( + args.base_url.rstrip("/") + "/generate", + json={ + "text": prompt, + "sampling_params": { + "temperature": args.temperature, + "top_p": args.top_p, + "top_k": args.top_k, + "max_new_tokens": args.max_new_tokens, + }, + }, + timeout=args.timeout_seconds, + ) + response.raise_for_status() + payload = response.json() + return payload[0] if isinstance(payload, list) else payload + + +def _run_sglang(args) -> BenchmarkResult: + import requests + from transformers import AutoTokenizer + + if args.concurrency <= 0: + raise ValueError("--concurrency must be positive") + if args.num_prompts <= 0: + raise ValueError("--num-prompts must be positive") + tokenizer = AutoTokenizer.from_pretrained( + args.model, + trust_remote_code=args.trust_remote_code, + ) + dataset = _load_prompts(args.dataset, args.max_samples) + prompt_count = args.num_prompts + warmup_count = args.concurrency + prompts = [ + _apply_chat_template( + tokenizer, + [{"role": "user", "content": dataset[index % len(dataset)][0]}], + args.enable_thinking, + ) + for index in range(prompt_count + warmup_count) + ] + + try: + requests.get( + args.base_url.rstrip("/") + "/flush_cache", + timeout=min(args.timeout_seconds, 60), + ).raise_for_status() + except requests.RequestException: + print("Warning: /flush_cache failed. Continuing.") + + with ThreadPoolExecutor(max_workers=warmup_count) as executor: + list( + executor.map( + lambda prompt: _send_sglang(args, prompt), prompts[:warmup_count] + ) + ) + prompts = prompts[warmup_count:] + + total_tokens = 0 + verify_count = 0 + acceptance_lengths: list[float] = [] + start = time.perf_counter() + with ThreadPoolExecutor(max_workers=args.concurrency) as executor: + futures = [executor.submit(_send_sglang, args, prompt) for prompt in prompts] + for future in as_completed(futures): + output = future.result() + metadata = output.get("meta_info", {}) or {} + total_tokens += int(metadata.get("completion_tokens", 0)) + verify_count += int(metadata.get("spec_verify_ct", 0)) + if metadata.get("spec_accept_length") is not None: + try: + acceptance_lengths.append(float(metadata["spec_accept_length"])) + except (TypeError, ValueError): + pass + elapsed = time.perf_counter() - start + return BenchmarkResult( + backend="sglang", + dataset=args.dataset, + samples=prompt_count, + output_tokens=total_tokens, + latency_seconds=elapsed, + throughput_tokens_per_second=total_tokens / max(elapsed, 1e-12), + average_acceptance_length=( + statistics.fmean(acceptance_lengths) if acceptance_lengths else None + ), + spec_verify_count=verify_count or None, + ) + + +def _print_result(result: BenchmarkResult) -> None: + print(f"Backend: {result.backend}") + print(f"Dataset: {result.dataset} ({result.samples} completed prompts/turns)") + print(f"Output throughput: {result.throughput_tokens_per_second:.2f} tok/s") + if result.average_acceptance_length is not None: + print(f"Average acceptance length: {result.average_acceptance_length:.3f}") + if result.spec_verify_count is not None: + print(f"Speculative verify count: {result.spec_verify_count}") + + +def run(args) -> int: + random.seed(42) + result = _run_sglang(args) + _print_result(result) + if args.output_json: + with open(args.output_json, "w", encoding="utf-8") as output_file: + json.dump(asdict(result), output_file, indent=2, sort_keys=True) + output_file.write("\n") + return 0 + + +__all__ = ["BenchmarkResult", "DATASETS", "run"] diff --git a/specforge/cli.py b/specforge/cli.py index 53b7d8a37..a122a7834 100644 --- a/specforge/cli.py +++ b/specforge/cli.py @@ -210,6 +210,40 @@ def main(argv: Optional[List[str]] = None) -> int: help="target model path supplying a frozen embedding for HF export", ) export.add_argument("--embedding-key", default="model.embed_tokens.weight") + benchmark = sub.add_parser( + "benchmark", + help="benchmark a running inference server", + ) + benchmark_subcommands = benchmark.add_subparsers( + dest="benchmark_command", + required=True, + ) + sglang_benchmark = benchmark_subcommands.add_parser( + "sglang", + help="benchmark a running SGLang server", + description=( + "Measure throughput and optional speculative-decoding telemetry from " + "a running SGLang server." + ), + ) + sglang_benchmark.add_argument("--model", required=True) + sglang_benchmark.add_argument( + "--dataset", + choices=("gsm8k", "math500", "humaneval", "mbpp", "mt-bench"), + required=True, + ) + sglang_benchmark.add_argument("--max-new-tokens", type=int, default=2048) + sglang_benchmark.add_argument("--temperature", type=float, default=0.0) + sglang_benchmark.add_argument("--top-p", type=float, default=1.0) + sglang_benchmark.add_argument("--top-k", type=int, default=1) + sglang_benchmark.add_argument("--max-samples", type=int) + sglang_benchmark.add_argument("--num-prompts", type=int, default=1024) + sglang_benchmark.add_argument("--concurrency", type=int, default=1) + sglang_benchmark.add_argument("--base-url", default="http://127.0.0.1:30000") + sglang_benchmark.add_argument("--timeout-seconds", type=int, default=3600) + sglang_benchmark.add_argument("--enable-thinking", action="store_true") + sglang_benchmark.add_argument("--trust-remote-code", action="store_true") + sglang_benchmark.add_argument("--output-json") args = parser.parse_args(argv) if args.command == "train": @@ -239,7 +273,11 @@ def main(argv: Optional[List[str]] = None) -> int: return 128 + received.signum return 0 return run_commands(plan) - elif args.to == "hf": + if args.command == "benchmark": + from specforge.benchmarks.sglang import run + + return run(args) + if args.to == "hf": from specforge.export.to_hf import export_to_hf export_to_hf( diff --git a/specforge/config/schema.py b/specforge/config/schema.py index fde4af74a..553681159 100644 --- a/specforge/config/schema.py +++ b/specforge/config/schema.py @@ -72,6 +72,9 @@ class ModelConfig(StrictConfigModel): #: DFlash-family and P-EAGLE mask token. ``None`` resolves it from the #: draft config, then the target tokenizer. mask_token_id: Optional[int] = None + #: Explicit tokenizer padding ID for released targets whose tokenizer + #: metadata omits it or whose padding row is outside the unpadded vocab. + tokenizer_pad_token_id: Optional[int] = Field(default=None, ge=0) #: SGLang target-engine tuning. Ignored by hf/custom backends. sglang_attention_backend: str = "flashinfer" sglang_mem_fraction_static: float = Field(default=0.4, gt=0.0, le=1.0) @@ -84,6 +87,24 @@ class ModelConfig(StrictConfigModel): sglang_ep_size: int = Field(default=1, gt=0) sglang_max_running_requests: Optional[int] = Field(default=None, gt=0) sglang_max_total_tokens: Optional[int] = Field(default=None, gt=0) + sglang_dp_size: Optional[int] = Field(default=None, gt=0) + sglang_moe_a2a_backend: Optional[str] = None + sglang_moe_runner_backend: Optional[str] = None + sglang_page_size: Optional[int] = Field(default=None, gt=0) + sglang_quantization: Optional[str] = None + sglang_fp4_gemm_runner_backend: Optional[str] = None + sglang_mamba_radix_cache_strategy: Optional[str] = None + sglang_max_mamba_cache_size: Optional[int] = Field(default=None, gt=0) + sglang_swa_full_tokens_ratio: Optional[float] = Field( + default=None, + gt=0.0, + le=1.0, + ) + sglang_mamba_full_memory_ratio: Optional[float] = Field( + default=None, + gt=0.0, + le=1.0, + ) @model_validator(mode="after") def _validate_input_modality(self): @@ -466,6 +487,9 @@ class TrainingConfig(StrictConfigModel): learning_rate: float = Field(default=1e-4, gt=0.0) warmup_ratio: float = Field(default=0.015, ge=0.0, le=1.0) max_grad_norm: float = Field(default=0.5, gt=0.0) + #: Keep FP32 Adam masters and moments on CPU while the trainable draft + #: remains on the accelerator. + optimizer_cpu_offload: bool = False ttt_length: int = Field(default=7, gt=0) attention_backend: Literal["eager", "sdpa", "flex_attention", "fa", "usp"] = ( "flex_attention" @@ -483,6 +507,8 @@ class TrainingConfig(StrictConfigModel): #: DFlash-family objective/model knobs. num_anchors: int = Field(default=512, gt=0) loss_decay_gamma: Optional[float] = None + #: Anchor blocks per objective slice (0 = materialize all objective logits). + objective_chunk_blocks: int = Field(default=128, ge=0) loss_type: Literal[ "dflash", "dpace", diff --git a/specforge/core/chunking.py b/specforge/core/chunking.py new file mode 100644 index 000000000..031f2b07f --- /dev/null +++ b/specforge/core/chunking.py @@ -0,0 +1,104 @@ +# coding=utf-8 +# Licensed under the Apache License, Version 2.0 (the "License"). +"""Reusable chunked reductions for memory-bounded training objectives.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Optional + +import torch + +ChunkTerms = tuple[torch.Tensor, ...] + + +def checkpointed_chunk_reduce( + function: Callable[..., ChunkTerms], + *aligned_tensors: Optional[torch.Tensor], + chunk_size: int, + dim: int = 0, +) -> ChunkTerms: + """Sum additive terms over aligned tensor slices. + + ``chunk_size=0`` evaluates the full dimension once without checkpointing. + Positive chunk sizes bound intermediates and use non-reentrant activation + checkpointing when gradients are enabled and an input requires gradients. + ``None`` arguments are forwarded unchanged, which keeps optional objective + inputs aligned with the tensors that are sliced. + """ + + if chunk_size < 0: + raise ValueError(f"chunk_size must be >= 0, got {chunk_size}") + + tensors = tuple(tensor for tensor in aligned_tensors if tensor is not None) + if not tensors: + raise ValueError("chunked reduction requires at least one tensor") + + first = tensors[0] + normalized_dim = dim if dim >= 0 else first.ndim + dim + if normalized_dim < 0 or normalized_dim >= first.ndim: + raise ValueError(f"dim {dim} is invalid for a {first.ndim}D tensor") + length = first.shape[normalized_dim] + if length == 0: + raise ValueError("chunked reduction received an empty dimension") + + for tensor in tensors[1:]: + tensor_dim = dim if dim >= 0 else tensor.ndim + dim + if tensor_dim < 0 or tensor_dim >= tensor.ndim: + raise ValueError(f"dim {dim} is invalid for a {tensor.ndim}D tensor") + if tensor.shape[tensor_dim] != length: + raise ValueError( + "chunked reduction inputs must be aligned: " + f"expected dimension length {length}, got {tensor.shape[tensor_dim]}" + ) + + effective_chunk_size = chunk_size or length + totals: Optional[ChunkTerms] = None + for start in range(0, length, effective_chunk_size): + width = min(effective_chunk_size, length - start) + chunk_args = tuple( + ( + tensor.narrow( + dim if dim >= 0 else tensor.ndim + dim, + start, + width, + ) + if tensor is not None + else None + ) + for tensor in aligned_tensors + ) + should_checkpoint = ( + chunk_size > 0 + and torch.is_grad_enabled() + and any( + tensor is not None and tensor.requires_grad for tensor in chunk_args + ) + ) + if should_checkpoint: + from torch.utils.checkpoint import checkpoint + + chunk_terms = checkpoint( + function, + *chunk_args, + use_reentrant=False, + ) + else: + chunk_terms = function(*chunk_args) + + if not isinstance(chunk_terms, tuple) or not all( + isinstance(term, torch.Tensor) for term in chunk_terms + ): + raise TypeError("chunk function must return a tuple of tensors") + if totals is None: + totals = chunk_terms + continue + if len(totals) != len(chunk_terms): + raise ValueError("chunk function returned a different number of terms") + totals = tuple(left + right for left, right in zip(totals, chunk_terms)) + + assert totals is not None + return totals + + +__all__ = ["checkpointed_chunk_reduce"] diff --git a/specforge/data/parse.py b/specforge/data/parse.py index 1217bedfe..70b41fdb9 100644 --- a/specforge/data/parse.py +++ b/specforge/data/parse.py @@ -9,7 +9,7 @@ from .template import ChatTemplate -__all__ = ["GeneralParser", "HarmonyParser", "ThinkingParser"] +__all__ = ["GeneralParser", "GLMParser", "HarmonyParser", "ThinkingParser"] class Parser(ABC): @@ -153,6 +153,29 @@ def set_assistant_pattern(self, chat_template: ChatTemplate): + re.escape("] USER:") + "|$))" ) + elif chat_template.assistant_pattern_type == "inkling": + terminators = "|".join( + re.escape(token) + for token in ( + "<|message_user|>", + "<|message_tool|>", + "<|message_system|>", + ) + ) + self.assistant_pattern = ( + re.escape(self.assistant_message_separator) + + r"([\s\S]*?(?:" + + terminators + + "|$))" + ) + elif chat_template.assistant_pattern_type == "glm": + self.assistant_pattern = ( + re.escape(self.assistant_message_separator) + + r"(?:)?" + + r"([\s\S]*?(?:" + + re.escape(self.chat_template.end_of_turn_token) + + "|$))" + ) else: self.assistant_pattern = ( re.escape(self.assistant_message_separator) @@ -242,7 +265,7 @@ def parse( parts.append(f"{assistant_header}{msg['content']}{end_of_turn}") conversation = "".join(parts) - if not self.tokenizer.pad_token_id: + if self.tokenizer.pad_token_id is None: self.tokenizer.pad_token_id = self.tokenizer.unk_token_id # get input_ids @@ -386,7 +409,7 @@ def parse( ) conversation = prompt_text - if not self.tokenizer.pad_token_id: + if self.tokenizer.pad_token_id is None: self.tokenizer.pad_token_id = self.tokenizer.unk_token_id encoding = self.tokenizer( @@ -444,7 +467,14 @@ def __init__( chat_template: ChatTemplate, ): super().__init__(tokenizer, chat_template) - self.standard_keys = {"role", "content", "tool_calls", "reasoning_content"} + self.standard_keys = { + "role", + "content", + "tool_calls", + "reasoning_content", + "name", + "tool_call_id", + } def apply_chat_template(self, messages, tool, **kwargs) -> str: """Apply chat template to all messages, handling reasoning_content and tool_calls.""" @@ -475,3 +505,11 @@ def parse( return super().parse( conversation, max_length, preformatted, train_only_last_turn, tool, **kwargs ) + + +class GLMParser(GeneralParser): + """Render GLM-5.2's hybrid-thinking template consistently for training.""" + + def apply_chat_template(self, messages, tool, **kwargs) -> str: + kwargs.setdefault("enable_thinking", False) + return super().apply_chat_template(messages, tool, **kwargs) diff --git a/specforge/data/preprocessing.py b/specforge/data/preprocessing.py index 1fae8e130..6f5cf5338 100644 --- a/specforge/data/preprocessing.py +++ b/specforge/data/preprocessing.py @@ -36,7 +36,7 @@ from transformers import PreTrainedTokenizer from ..distributed import get_draft_sp_group, get_sp_ring_group -from .parse import GeneralParser, HarmonyParser, ThinkingParser +from .parse import GeneralParser, GLMParser, HarmonyParser, ThinkingParser from .template import TEMPLATE_REGISTRY, ChatTemplate # define a type called conversation @@ -140,6 +140,8 @@ def preprocess_conversations( parser = GeneralParser(tokenizer, chat_template) elif chat_template.parser_type == "thinking": parser = ThinkingParser(tokenizer, chat_template) + elif chat_template.parser_type == "glm": + parser = GLMParser(tokenizer, chat_template) elif chat_template.parser_type == "openai-harmony": parser = HarmonyParser(tokenizer, chat_template) else: diff --git a/specforge/data/template.py b/specforge/data/template.py index 747becd79..cfb409829 100644 --- a/specforge/data/template.py +++ b/specforge/data/template.py @@ -294,6 +294,19 @@ def get_all_template_names(self) -> List[str]: ), ) +TEMPLATE_REGISTRY.register( + name="glm-5.2", + template=ChatTemplate( + assistant_header="<|assistant|>", + user_header="<|user|>", + system_prompt=None, + end_of_turn_token="<|user|>", + parser_type="glm", + assistant_pattern_type="glm", + ignore_token=["<|user|>"], + ), +) + TEMPLATE_REGISTRY.register( name="gemma", template=ChatTemplate( @@ -337,3 +350,21 @@ def get_all_template_names(self) -> List[str]: enable_thinking=True, ), ) + +TEMPLATE_REGISTRY.register( + name="inkling-thinking", + template=ChatTemplate( + assistant_header="<|message_model|>", + user_header="<|message_user|>", + system_prompt=None, + end_of_turn_token="<|message_user|>", + parser_type="thinking", + assistant_pattern_type="inkling", + enable_thinking=True, + ignore_token=[ + "<|message_user|>", + "<|message_tool|>", + "<|message_system|>", + ], + ), +) diff --git a/specforge/inference/adapters/server_capture.py b/specforge/inference/adapters/server_capture.py index 4cb3fc192..f30a4d292 100644 --- a/specforge/inference/adapters/server_capture.py +++ b/specforge/inference/adapters/server_capture.py @@ -25,6 +25,7 @@ from __future__ import annotations import logging +import uuid from dataclasses import dataclass from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Union @@ -186,7 +187,7 @@ def _request_inputs(self, tasks: List[PromptTask]) -> Dict[str, Any]: raise TypeError( "ServerInputAdapter.build_request_inputs must return a mapping" ) - reserved = {"sampling_params", "spec_capture"} + reserved = {"extra_key", "sampling_params", "spec_capture"} conflicts = sorted(reserved & set(request_inputs)) if conflicts: raise ValueError( @@ -317,6 +318,11 @@ def produce_refs( retryable, mirroring ``generate_features`` semantics. """ body = self._request_inputs(tasks) + # A fresh cache namespace forces a full prefill on every attempt while + # leaving SGLang free to use the radix data structure required by hybrid + # targets. Retries need a new key because the prior attempt may have + # populated its namespace before the response was lost. + body["extra_key"] = [uuid.uuid4().hex for _ in tasks] body["sampling_params"] = {"temperature": 0.0, "max_new_tokens": 1} capture_payloads = [self._spec_capture_payload(t) for t in tasks] body["spec_capture"] = capture_payloads @@ -394,9 +400,9 @@ def produce_refs( f"{expected_identity}" ) ref = self._ref_from_result(task, result, capture) - # A capture shorter than the prompt is corrupt (classic cause: a - # radix-cache prefix hit skips prefilling — and capturing — the - # cached tokens; the patched scheduler refuses that config). + # A capture shorter than the prompt is corrupt: it means cache + # isolation or another full-prefill invariant failed despite the + # request's fresh namespace. expected_len = len(task.payload["input_ids"]) short = { name: spec.shape @@ -411,8 +417,8 @@ def produce_refs( task_id=task.task_id, reason=( f"server_capture: captured seq len != prompt len " - f"{expected_len} for {short} — was the server " - f"started without --disable-radix-cache?" + f"{expected_len} for {short}; the capture request " + "did not execute a complete prefill" ), retryable=False, ) diff --git a/specforge/inference/sglang_patch_inventory.md b/specforge/inference/sglang_patch_inventory.md index 370f2d8db..8ce6f2402 100644 --- a/specforge/inference/sglang_patch_inventory.md +++ b/specforge/inference/sglang_patch_inventory.md @@ -1,7 +1,8 @@ # SGLang patch inventory and supported version -SpecForge pins `sglang==0.5.14`. There are two deliberately separate SGLang -integration surfaces. +SpecForge pins `sglang==0.5.14`. The online patch is also kept compatible with +SGLang's public `inkling-support` layout. There are two deliberately separate +SGLang integration surfaces. ## Online: external spec-capture server @@ -21,6 +22,12 @@ providers map generic server artifacts (`aux`, `last_hidden`, passthrough inputs) to training feature names. No trainer or producer process imports SGLang model-runner internals or loads a target model. +The same patch is dry-run validated against the v0.5.14 tag and the public +`inkling-support` branch. Capture requests carry a unique `extra_key`, so every +training sample executes a full prefill even when radix cache support is +present. Capture launch configs therefore leave radix cache enabled, including +for hybrid targets that require the unified radix tree. + Apply the patch with `scripts/apply_sglang_spec_capture_patch.sh`. The server-capture unit and GPU gates must pass before updating the SGLang pin. diff --git a/specforge/launch_plan.py b/specforge/launch_plan.py index 98f852f49..299714151 100644 --- a/specforge/launch_plan.py +++ b/specforge/launch_plan.py @@ -434,7 +434,6 @@ def _managed_local_services( str(cfg.model.sglang_ep_size), "--chunked-prefill-size", "-1", - "--disable-radix-cache", "--enable-spec-capture", "--spec-capture-method", contract.method, @@ -458,6 +457,28 @@ def _managed_local_services( for value, flag in ( (cfg.model.sglang_max_running_requests, "--max-running-requests"), (cfg.model.sglang_max_total_tokens, "--max-total-tokens"), + (cfg.model.sglang_dp_size, "--dp-size"), + (cfg.model.sglang_moe_a2a_backend, "--moe-a2a-backend"), + (cfg.model.sglang_moe_runner_backend, "--moe-runner-backend"), + (cfg.model.sglang_page_size, "--page-size"), + (cfg.model.sglang_quantization, "--quantization"), + ( + cfg.model.sglang_fp4_gemm_runner_backend, + "--fp4-gemm-runner-backend", + ), + ( + cfg.model.sglang_mamba_radix_cache_strategy, + "--mamba-radix-cache-strategy", + ), + (cfg.model.sglang_max_mamba_cache_size, "--max-mamba-cache-size"), + ( + cfg.model.sglang_swa_full_tokens_ratio, + "--swa-full-tokens-ratio", + ), + ( + cfg.model.sglang_mamba_full_memory_ratio, + "--mamba-full-memory-ratio", + ), ): if value is not None: argv.extend((flag, str(value))) diff --git a/specforge/modeling/draft/dspark.py b/specforge/modeling/draft/dspark.py index c65e45929..1c5853d8b 100644 --- a/specforge/modeling/draft/dspark.py +++ b/specforge/modeling/draft/dspark.py @@ -286,7 +286,32 @@ class DSparkDraftModel(DFlashDraftModel): expected_projector_type = "dspark" def __init__(self, config) -> None: - dflash_config = getattr(config, "dflash_config", None) or {} + dflash_config = dict(getattr(config, "dflash_config", None) or {}) + num_heads = int(config.num_attention_heads) + num_kv_heads = int(config.num_key_value_heads) + if num_heads % num_kv_heads: + raise ValueError( + "DSpark requires num_key_value_heads to divide " + f"num_attention_heads, got {num_kv_heads} and {num_heads}" + ) + attention_mode = str(dflash_config.get("attention_mode", "gqa")).lower() + if attention_mode not in {"gqa", "mha"}: + raise ValueError( + "DSpark dflash_config.attention_mode must be 'gqa' or 'mha', " + f"got {attention_mode!r}" + ) + if attention_mode == "gqa" and num_kv_heads >= num_heads: + raise ValueError( + "DSpark defaults to GQA and requires num_key_value_heads < " + "num_attention_heads; set dflash_config.attention_mode='mha' " + "to opt into equal query/KV head counts" + ) + if attention_mode == "mha" and num_kv_heads != num_heads: + raise ValueError( + "DSpark MHA opt-in requires num_key_value_heads == " + f"num_attention_heads, got {num_kv_heads} and {num_heads}" + ) + dflash_config["attention_mode"] = attention_mode projector_type = dflash_config.get("projector_type") if projector_type is None: dflash_config["projector_type"] = self.expected_projector_type @@ -295,8 +320,28 @@ def __init__(self, config) -> None: raise ValueError( "DSparkDraftModel requires " "dflash_config.projector_type='dspark'." ) + config.dflash_config = dflash_config super().__init__(config) + def _sample_draft_tokens( + self, + target: nn.Module, + draft_hidden: torch.Tensor, + block_output_ids: torch.LongTensor, + ) -> torch.LongTensor: + """Generate a block with the same Markov correction used in training.""" + + proposal_hidden = draft_hidden[:, -self.block_size + 1 :, :] + base_logits = target.lm_head(proposal_hidden) + if self.markov_head is None: + return _sample(base_logits) + sampled_tokens, _ = self.markov_head.sample_block_tokens( + base_logits, + first_prev_token_ids=block_output_ids[:, 0], + hidden_states=proposal_hidden, + ) + return sampled_tokens + def _init_draft_head(self, config, dflash_config: dict) -> None: self.markov_head = build_markov_head(config, dflash_config) confidence_alpha = float(dflash_config.get("confidence_head_alpha", 0.0) or 0.0) diff --git a/specforge/modeling/target/target_utils.py b/specforge/modeling/target/target_utils.py index 6f78cf098..c8ca1afcb 100644 --- a/specforge/modeling/target/target_utils.py +++ b/specforge/modeling/target/target_utils.py @@ -6,11 +6,74 @@ import torch import torch.nn as nn -from huggingface_hub import snapshot_download +from huggingface_hub import hf_hub_download, snapshot_download from safetensors import safe_open from transformers import AutoConfig +class _RawConfigShim: + """Attribute view for released checkpoints with an unregistered model type.""" + + def __init__(self, data: dict): + object.__setattr__(self, "_data", data) + + def __getattr__(self, name): + try: + value = self._data[name] + except KeyError: + raise AttributeError(name) from None + return _RawConfigShim(value) if isinstance(value, dict) else value + + def to_dict(self) -> dict: + return dict(self._data) + + +def load_target_config( + model_path: str, + *, + cache_dir: Optional[str] = None, + trust_remote_code: bool = False, +): + """Load a target config, falling back to its public raw ``config.json``.""" + + try: + return AutoConfig.from_pretrained( + model_path, + cache_dir=cache_dir, + trust_remote_code=trust_remote_code, + ) + except (ValueError, KeyError, OSError) as auto_error: + if os.path.isdir(model_path): + config_path = os.path.join(model_path, "config.json") + elif os.path.isfile(model_path): + config_path = model_path + else: + try: + config_path = hf_hub_download( + repo_id=model_path, + filename="config.json", + cache_dir=cache_dir, + ) + except Exception: + raise auto_error + try: + with open(config_path, encoding="utf-8") as config_file: + return _RawConfigShim(json.load(config_file)) + except (OSError, ValueError): + raise auto_error + + +def target_text_config(config): + return getattr(config, "text_config", config) + + +def target_vocab_size(config) -> int: + text_config = target_text_config(config) + return int( + getattr(text_config, "padded_vocab_size", None) or text_config.vocab_size + ) + + class TargetEmbeddingsAndHead(nn.Module): """ Efficiently loads only the embedding layer and lm_head from a pretrained model. @@ -20,23 +83,15 @@ class TargetEmbeddingsAndHead(nn.Module): def __init__(self, config): super().__init__() self.config = config - # Support for MLLMs with separate text_config - if hasattr(config, "text_config"): - self.embed_tokens = nn.Embedding( - config.text_config.vocab_size, - config.text_config.hidden_size, - padding_idx=config.text_config.pad_token_id, - ) - self.lm_head = nn.Linear( - config.text_config.hidden_size, - config.text_config.vocab_size, - bias=False, - ) - else: - self.embed_tokens = nn.Embedding( - config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id - ) - self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + text_config = target_text_config(config) + vocab_size = target_vocab_size(text_config) + hidden_size = int(text_config.hidden_size) + self.embed_tokens = nn.Embedding( + vocab_size, + hidden_size, + padding_idx=getattr(text_config, "pad_token_id", None), + ) + self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False) @classmethod def from_pretrained( @@ -51,8 +106,10 @@ def from_pretrained( ) -> "TargetEmbeddingsAndHead": # 1. Load Config - config = AutoConfig.from_pretrained( - model_path, cache_dir=cache_dir, trust_remote_code=trust_remote_code + config = load_target_config( + model_path, + cache_dir=cache_dir, + trust_remote_code=trust_remote_code, ) instance = cls(config) @@ -79,6 +136,21 @@ def from_pretrained( # 4. Load Weights instance._load_weights(local_model_path, embed_key, lm_head_key, tie_weights) + text_config = target_text_config(config) + mup_multiplier = getattr( + text_config, + "logits_mup_width_multiplier", + getattr(config, "logits_mup_width_multiplier", None), + ) + if mup_multiplier: + if tie_weights: + raise RuntimeError( + "cannot fold logits_mup_width_multiplier into a tied " + "embedding/LM head" + ) + instance.lm_head.weight.data.div_(float(mup_multiplier)) + instance.lm_head_mup_folded = float(mup_multiplier) + # 5. Move to Device & Freeze instance.to(device=device, dtype=dtype) instance.eval() @@ -201,3 +273,11 @@ def _load_file_content( loaded_keys.add(k) return loaded_keys + + +__all__ = [ + "TargetEmbeddingsAndHead", + "load_target_config", + "target_text_config", + "target_vocab_size", +] diff --git a/specforge/optimizer.py b/specforge/optimizer.py index 2318ebcce..3d5e1aabe 100644 --- a/specforge/optimizer.py +++ b/specforge/optimizer.py @@ -21,13 +21,20 @@ def __init__( max_grad_norm=0.5, total_steps=800_000, warmup_ratio=0.015, + offload_master=False, ): # defaults copied from EAGLE traineagle3 ds_config.json self.model = model self.model_params = [p for p in model.parameters() if p.requires_grad] self.max_grad_norm = max_grad_norm + self.offload_master = bool(offload_master) self.fp32_params = [ - p.detach().clone().to(torch.float32) for p in self.model_params + ( + p.detach().to(device="cpu", dtype=torch.float32).clone() + if self.offload_master + else p.detach().clone().to(torch.float32) + ) + for p in self.model_params ] for mp in self.fp32_params: mp.requires_grad = True @@ -53,15 +60,13 @@ def configure_grad_norm_reduction( self._grad_norm_process_group = process_group self._reduce_grad_norm_across_ranks = enabled - def _clip_grad_norm(self): - """Clip all FSDP shards with one global L2-norm coefficient.""" - grads = [mp.grad for mp in self.fp32_params if mp.grad is not None] - if grads: - total_norm_sq = torch.stack([grad.square().sum() for grad in grads]).sum() - else: - device = self.fp32_params[0].device if self.fp32_params else "cpu" - total_norm_sq = torch.zeros((), dtype=torch.float32, device=device) + def _reduce_grad_norm(self, total_norm_sq): + """All-reduce the squared L2 norm across shard ranks and derive the + clip coefficient. + ``total_norm_sq`` must already live on a device the process group can + reduce (e.g. CUDA for NCCL). Returns ``(total_norm, clip_coef)``. + """ if ( self._reduce_grad_norm_across_ranks and dist.is_available() @@ -72,27 +77,82 @@ def _clip_grad_norm(self): op=dist.ReduceOp.SUM, group=self._grad_norm_process_group, ) - total_norm = total_norm_sq.sqrt() clip_coef = torch.clamp(self.max_grad_norm / (total_norm + 1e-6), max=1.0) + return total_norm, clip_coef + + def _grad_norm_and_clip_coefficient(self): + """Compute the global grad norm from the model params on their own + device, where NCCL can reduce it safely, without materialising master + gradients first.""" + grads = [p.grad.detach() for p in self.model_params if p.grad is not None] + if grads: + total_norm_sq = torch.stack( + [grad.float().square().sum() for grad in grads] + ).sum() + else: + device = self.model_params[0].device if self.model_params else "cpu" + total_norm_sq = torch.zeros((), dtype=torch.float32, device=device) + return self._reduce_grad_norm(total_norm_sq) + + def _clip_grad_norm(self): + """Clip already-populated FP32 master gradients in place. + + Convenience entry point for optimizer tests and custom loops. When + masters are CPU-offloaded, only the scalar norm is moved to the model + device so a NCCL process group can still participate in the reduction. + """ + grads = [master.grad for master in self.fp32_params if master.grad is not None] + if grads: + local_norm_sq = torch.stack( + [grad.float().square().sum() for grad in grads] + ).sum() + else: + master_device = self.fp32_params[0].device if self.fp32_params else "cpu" + local_norm_sq = torch.zeros((), dtype=torch.float32, device=master_device) + + reduction_device = ( + self.model_params[0].device if self.model_params else local_norm_sq.device + ) + total_norm, clip_coef = self._reduce_grad_norm( + local_norm_sq.to(reduction_device) + ) for grad in grads: - grad.mul_(clip_coef) + coefficient = ( + clip_coef + if clip_coef.device == grad.device + else float(clip_coef.item()) + ) + grad.mul_(coefficient) return total_norm def step(self): + grad_norm, clip_coefficient = self._grad_norm_and_clip_coefficient() + cpu_clip_coefficient = ( + float(clip_coefficient.item()) if self.offload_master else None + ) with torch.no_grad(): for p, mp in zip(self.model_params, self.fp32_params): - mp.grad = ( - p.grad.detach().to(torch.float32) if p.grad is not None else None + if p.grad is None: + mp.grad = None + continue + master_grad = p.grad.detach().to( + device=mp.device, + dtype=torch.float32, + ) + master_grad.mul_( + cpu_clip_coefficient + if cpu_clip_coefficient is not None + else clip_coefficient ) - grad_norm = self._clip_grad_norm() + mp.grad = master_grad self.last_grad_norm = grad_norm.detach() self.optimizer.step() self.optimizer.zero_grad() self.scheduler.step() with torch.no_grad(): for p, mp in zip(self.model_params, self.fp32_params): - p.data.copy_(mp.data.to(p.dtype)) + p.data.copy_(mp.data.to(device=p.device, dtype=p.dtype)) p.grad = None return self.last_grad_norm @@ -109,6 +169,9 @@ def load_state_dict(self, state_dict): f"{saved_max_grad_norm} but this run has " f"max_grad_norm={self.max_grad_norm}" ) + # offload_master is a pure device-placement choice: restored fp32 + # masters and Adam moments are relocated to the current master device, + # so toggling it on resume is safe and intentionally not gated here. self.optimizer.load_state_dict(state_dict["optimizer_state_dict"]) print_on_rank0("Successfully loaded optimizer state_dict.") self.scheduler.load_state_dict(state_dict["scheduler_state_dict"]) @@ -135,7 +198,7 @@ def load_state_dict(self, state_dict): ) with torch.no_grad(): for p, mp in zip(self.model_params, self.fp32_params): - mp.data.copy_(p.detach().to(torch.float32)) + mp.data.copy_(p.detach().to(device=mp.device, dtype=mp.dtype)) def state_dict(self): return { diff --git a/specforge/training/assembly.py b/specforge/training/assembly.py index 889fd843c..57044f3a5 100644 --- a/specforge/training/assembly.py +++ b/specforge/training/assembly.py @@ -137,11 +137,24 @@ def _load_text_tokenizer(cfg: Config): ) from transformers import AutoTokenizer - return AutoTokenizer.from_pretrained( + tokenizer = AutoTokenizer.from_pretrained( cfg.model.target_model_path, cache_dir=cfg.model.cache_dir, trust_remote_code=cfg.model.trust_remote_code, ) + if cfg.model.tokenizer_pad_token_id is not None: + tokenizer.pad_token_id = cfg.model.tokenizer_pad_token_id + elif tokenizer.pad_token_id is None: + fallback_id = tokenizer.eos_token_id + if fallback_id is None: + fallback_id = tokenizer.unk_token_id + if fallback_id is None: + raise ValueError( + "target tokenizer has no pad, EOS, or unknown token ID; set " + "model.tokenizer_pad_token_id explicitly" + ) + tokenizer.pad_token_id = fallback_id + return tokenizer def _load_input_tools( @@ -163,20 +176,27 @@ def _load_input_tools( def build_model_bundle(cfg: Config, *, algorithm: AlgorithmRegistration) -> ModelBundle: """Build the method-specific composite model and frozen target pieces.""" import torch - from transformers import AutoConfig + + from specforge.modeling.target.target_utils import ( + load_target_config, + target_text_config, + ) + from specforge.modeling.target.target_utils import ( + target_vocab_size as resolve_target_vocab_size, + ) provider = algorithm.providers.model draft_config, draft_model = _load_draft(cfg, algorithm) needs_input_tools = provider.needs_input_tools(cfg, draft_model) input_tools = _load_input_tools(cfg, algorithm) if needs_input_tools else None - target_config = AutoConfig.from_pretrained( + target_config = load_target_config( cfg.model.target_model_path, cache_dir=cfg.model.cache_dir, trust_remote_code=cfg.model.trust_remote_code, ) - text_config = _target_text_config(target_config) + text_config = target_text_config(target_config) target_hidden_size = int(text_config.hidden_size) - target_vocab_size = int(text_config.vocab_size) + target_vocab_size = resolve_target_vocab_size(target_config) draft_vocab_size = int( getattr(draft_config, "draft_vocab_size", draft_config.vocab_size) ) @@ -246,6 +266,7 @@ def __call__(self, draft_module): max_grad_norm=t.max_grad_norm, warmup_ratio=t.warmup_ratio, total_steps=self.total_steps, + offload_master=t.optimizer_cpu_offload, ) @@ -294,10 +315,18 @@ def _close_configured_logger(logger) -> None: def _prompt_cache_key(cfg: Config, *, path: Optional[str] = None) -> str: - import json + source_path = path or cfg.data.prompts_path or cfg.data.train_data_path + content_hash = None + if source_path and os.path.isfile(source_path): + source_hasher = hashlib.sha256() + with open(source_path, "rb") as source_file: + for chunk in iter(lambda: source_file.read(8 * 1024 * 1024), b""): + source_hasher.update(chunk) + content_hash = source_hasher.hexdigest() identity = { - "path": path or cfg.data.prompts_path or cfg.data.train_data_path, + "path": source_path, + "content_hash": content_hash, "max_length": cfg.data.max_length, "chat_template": cfg.data.chat_template, "is_preformatted": cfg.data.is_preformatted, diff --git a/specforge/training/capture_contract.py b/specforge/training/capture_contract.py index dcfd69c64..feca7d43e 100644 --- a/specforge/training/capture_contract.py +++ b/specforge/training/capture_contract.py @@ -24,18 +24,21 @@ def resolve_server_capture_contract( algorithm: AlgorithmRegistration, ) -> ServerCaptureContract: """Resolve engine flags and feature dimensions from canonical model config.""" - from transformers import AutoConfig - + from specforge.modeling.target.target_utils import ( + load_target_config, + target_text_config, + target_vocab_size, + ) from specforge.training.model_loading import draft_config_dict streaming = algorithm.providers.server_streaming_for(cfg.model.input_modality) - target_cfg = AutoConfig.from_pretrained( + target_cfg = load_target_config( cfg.model.target_model_path, cache_dir=cfg.model.cache_dir, trust_remote_code=cfg.model.trust_remote_code, ) - target_cfg = getattr(target_cfg, "text_config", target_cfg) + target_cfg = target_text_config(target_cfg) model_provider = algorithm.providers.model draft_cfg = draft_config_dict(cfg, provider=model_provider.draft_config) layers = model_provider.resolve_capture_layers(cfg, draft_cfg, target_cfg) @@ -58,7 +61,7 @@ def resolve_server_capture_contract( method=streaming.capture_method, aux_layer_ids=tuple(layers), target_hidden_size=int(target_cfg.hidden_size), - target_vocab_size=int(target_cfg.vocab_size), + target_vocab_size=target_vocab_size(target_cfg), draft_vocab_size=int( draft_cfg.get("draft_vocab_size") or draft_cfg["vocab_size"] ), diff --git a/specforge/training/checkpoint.py b/specforge/training/checkpoint.py index 23790172e..88e8396c6 100644 --- a/specforge/training/checkpoint.py +++ b/specforge/training/checkpoint.py @@ -235,15 +235,67 @@ def read_resume_state( Accepts a checkpoint dir, its ``training_state.pt``, a ``file://`` URI, or an output/run root containing exactly one checkpoint lineage. """ - path = cls.resolve_resume_dir(path) - state = torch.load( - os.path.join(path, STATE_FILE), - map_location=map_location, - weights_only=False, - ) initialized = torch.distributed.is_initialized() - rank = torch.distributed.get_rank() if initialized else 0 + original_path = path + local_error = None + state = None + try: + path = cls.resolve_resume_dir(path) + state = torch.load( + os.path.join(path, STATE_FILE), + map_location=map_location, + weights_only=False, + ) + except Exception as exc: + local_error = exc + world = torch.distributed.get_world_size() if initialized else 1 + if initialized and world > 1: + descriptor = { + "error": ( + None + if local_error is None + else f"{type(local_error).__name__}: {local_error}" + ), + "checkpoint": ( + None if local_error is not None else os.path.basename(path) + ), + "run_id": None if state is None else state.get("run_id"), + "global_step": None if state is None else state.get("global_step"), + "strategy": None if state is None else state.get("strategy"), + } + descriptors = [None] * world + torch.distributed.all_gather_object(descriptors, descriptor) + failures = [ + (rank, item["error"]) + for rank, item in enumerate(descriptors) + if item["error"] is not None + ] + if failures: + raise RuntimeError( + f"resume checkpoint {original_path!r} is not readable on " + "every rank: " + + "; ".join(f"rank {rank}: {error}" for rank, error in failures) + ) + identities = { + ( + item["checkpoint"], + item["run_id"], + item["global_step"], + item["strategy"], + ) + for item in descriptors + } + if len(identities) != 1: + raise ValueError( + f"resume checkpoint {original_path!r} resolves to different " + f"training states across ranks: {descriptors}" + ) + if local_error is not None: + raise local_error + assert state is not None + + rank = torch.distributed.get_rank() if initialized else 0 saved_world = state.get("world_size") rank_path = os.path.join(path, cls._rank_file(rank)) if os.path.exists(rank_path) and (saved_world is None or saved_world == world): diff --git a/specforge/training/controller.py b/specforge/training/controller.py index 36237a5df..de6fcb29e 100644 --- a/specforge/training/controller.py +++ b/specforge/training/controller.py @@ -67,7 +67,9 @@ def _scalar(x: Any) -> float: return float(x) -def _dp_mean_scalars(values: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: +def _dp_mean_scalars( + values: Dict[str, torch.Tensor], *, process_group: Any = None +) -> Dict[str, torch.Tensor]: """Average scalar metrics across DP ranks with one collective. Uses the established DFlash metric convention (DP mean): @@ -82,16 +84,84 @@ def _dp_mean_scalars(values: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor] if not values or not (dist.is_available() and dist.is_initialized()): return values - world = dist.get_world_size() + world = ( + dist.get_world_size() + if process_group is None + else dist.get_world_size(group=process_group) + ) if world <= 1: return values names = list(values) packed = torch.stack([values[name].detach().float().reshape(()) for name in names]) - dist.all_reduce(packed) + if process_group is None: + dist.all_reduce(packed) + else: + dist.all_reduce(packed, group=process_group) packed /= world return {name: packed[index] for index, name in enumerate(names)} +def _reduce_ratio_metrics( + values: Dict[str, Any], + *, + device: torch.device, + process_group: Any, + reduce: bool, +) -> Dict[str, float]: + """Form telemetry ratios only after summing their numerators and counts.""" + + if not values: + return {} + normalized = [] + for name in sorted(values): + pair = values[name] + if not isinstance(pair, (tuple, list)) or len(pair) != 2: + raise TypeError( + f"ratio metric {name!r} must be a (numerator, denominator) pair" + ) + numerator = torch.as_tensor(pair[0]).detach().float().flatten().to(device) + denominator = torch.as_tensor(pair[1]).detach().float().flatten().to(device) + if numerator.shape != denominator.shape: + raise ValueError( + f"ratio metric {name!r} shape mismatch: " + f"{tuple(numerator.shape)} vs {tuple(denominator.shape)}" + ) + normalized.append((name, numerator, denominator)) + + packed = torch.cat( + [ + tensor + for _, numerator, denominator in normalized + for tensor in (numerator, denominator) + ] + ) + if reduce: + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + if process_group is None: + dist.all_reduce(packed) + else: + dist.all_reduce(packed, group=process_group) + + output: Dict[str, float] = {} + cursor = 0 + for name, numerator, _denominator in normalized: + width = numerator.numel() + summed_numerator = packed[cursor : cursor + width] + cursor += width + summed_denominator = packed[cursor : cursor + width] + cursor += width + ratios = summed_numerator / summed_denominator.clamp_min(1e-12) + if width == 1: + output[name] = float(ratios.item()) + else: + output.update( + {f"{name}_{index}": float(value) for index, value in enumerate(ratios)} + ) + return output + + _EAGLE3_STRUCTURED_METRIC_KEYS = frozenset( { "acces", @@ -281,10 +351,18 @@ def _result(self, out: StepOutput, grad_norm, stepped: bool) -> StepResult: # scalar diagnostics are DP-averaged in a single collective at optimizer # boundaries; non-boundary results stay rank-local. metrics: Dict[str, Any] = dict(structured or {}) + metrics.update( + _reduce_ratio_metrics( + out.ratio_metrics, + device=metric_device, + process_group=process_group, + reduce=stepped, + ) + ) scalar_metrics: Dict[str, torch.Tensor] = {} if "loss" not in metrics: scalar_metrics["loss"] = out.loss - if "accuracy" in out.metrics: + if "accuracy" in out.metrics and "acc" not in metrics: accuracy = out.metrics["accuracy"] if isinstance(accuracy, torch.Tensor): scalar_metrics["acc"] = ( @@ -310,7 +388,10 @@ def _result(self, out: StepOutput, grad_norm, stepped: bool) -> StepResult: continue scalar_metrics[key] = scalar if stepped: - scalar_metrics = _dp_mean_scalars(scalar_metrics) + scalar_metrics = _dp_mean_scalars( + scalar_metrics, + process_group=process_group, + ) metrics.update({key: _scalar(value) for key, value in scalar_metrics.items()}) gn = _scalar(grad_norm) if grad_norm is not None else None if gn is not None: diff --git a/specforge/training/model_loading.py b/specforge/training/model_loading.py index 3d0345624..75fb32f82 100644 --- a/specforge/training/model_loading.py +++ b/specforge/training/model_loading.py @@ -217,19 +217,24 @@ def _generate_draft_config( "or a pretrained model.draft_checkpoint_path containing config.json" ) - from transformers import AutoConfig + from specforge.modeling.target.target_utils import ( + load_target_config, + target_text_config, + target_vocab_size, + ) - target_config = AutoConfig.from_pretrained( + target_config = load_target_config( cfg.model.target_model_path, cache_dir=cfg.model.cache_dir, trust_remote_code=cfg.model.trust_remote_code, ) - target_config = _target_text_config(target_config) + target_config = target_text_config(target_config) payload: Dict[str, Any] = {} for field in _TARGET_ARCHITECTURE_FIELDS: value = getattr(target_config, field, None) if value is not None: payload[field] = _serializable_config_value(value) + payload["vocab_size"] = target_vocab_size(target_config) required = [name for name in ("vocab_size", "hidden_size") if name not in payload] if required: diff --git a/specforge/training/strategies/base.py b/specforge/training/strategies/base.py index 3d53d524a..60114b743 100644 --- a/specforge/training/strategies/base.py +++ b/specforge/training/strategies/base.py @@ -17,7 +17,7 @@ from __future__ import annotations import abc -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any, Dict, Optional, Tuple import torch @@ -33,6 +33,7 @@ class StepOutput: loss: torch.Tensor metrics: Dict[str, Any] + ratio_metrics: Dict[str, Tuple[Any, Any]] = field(default_factory=dict) @dataclass(frozen=True) @@ -487,6 +488,7 @@ def forward_loss( metrics = { "accuracy": accuracy.detach(), } + ratio_metrics = model_metrics.get("ratio_metrics", {}) for name in ( "accuracy_denom", "ce_loss", @@ -496,7 +498,11 @@ def forward_loss( ): if name in model_metrics: metrics[name] = model_metrics[name] - return StepOutput(loss=loss, metrics=metrics) + return StepOutput( + loss=loss, + metrics=metrics, + ratio_metrics=ratio_metrics, + ) def checkpoint_state_filter(self, state_dict: Dict[str, Any]) -> Dict[str, Any]: return { diff --git a/tests/test_benchmarks/test_sglang_benchmark.py b/tests/test_benchmarks/test_sglang_benchmark.py new file mode 100644 index 000000000..a924690fd --- /dev/null +++ b/tests/test_benchmarks/test_sglang_benchmark.py @@ -0,0 +1,100 @@ +import unittest +from contextlib import redirect_stdout +from io import StringIO +from types import SimpleNamespace +from unittest import mock + +from specforge.benchmarks import sglang + + +class SGLangBenchmarkTest(unittest.TestCase): + def test_mt_bench_prompts_preserve_turns(self): + rows = [{"prompt": ["first", "second"]}] + with mock.patch("datasets.load_dataset", return_value=rows): + prompts = sglang._load_prompts("mt-bench", max_samples=None) + self.assertEqual(prompts, [["first", "second"]]) + + def test_prompt_loader_rejects_empty_inputs(self): + with mock.patch("datasets.load_dataset", return_value=[]): + with self.assertRaisesRegex(ValueError, "did not contain any prompts"): + sglang._load_prompts("gsm8k", max_samples=None) + with self.assertRaisesRegex(ValueError, "--max-samples must be positive"): + sglang._load_prompts("gsm8k", max_samples=0) + + def test_sglang_path_excludes_warmup_from_reported_totals(self): + args = SimpleNamespace( + model="thinkingmachines/Inkling", + trust_remote_code=False, + dataset="gsm8k", + max_samples=None, + num_prompts=3, + concurrency=2, + enable_thinking=False, + base_url="http://127.0.0.1:30000", + timeout_seconds=30, + ) + response = { + "meta_info": { + "completion_tokens": 2, + "spec_verify_ct": 1, + "spec_accept_length": 3.0, + } + } + flush_response = mock.Mock() + flush_response.raise_for_status.return_value = None + with ( + mock.patch( + "transformers.AutoTokenizer.from_pretrained", + return_value=SimpleNamespace(), + ), + mock.patch.object(sglang, "_load_prompts", return_value=[["prompt"]]), + mock.patch.object(sglang, "_apply_chat_template", return_value="rendered"), + mock.patch.object(sglang, "_send_sglang", return_value=response) as send, + mock.patch("requests.get", return_value=flush_response) as flush, + ): + result = sglang._run_sglang(args) + + self.assertEqual(send.call_count, args.num_prompts + args.concurrency) + self.assertEqual(result.samples, 3) + self.assertEqual(result.output_tokens, 6) + self.assertEqual(result.spec_verify_count, 3) + self.assertEqual(result.average_acceptance_length, 3.0) + flush.assert_called_once_with( + "http://127.0.0.1:30000/flush_cache", + timeout=30, + ) + + def test_shared_cli_dispatches_sglang_benchmark(self): + from specforge.cli import main + + with mock.patch.object(sglang, "run", return_value=7) as run: + status = main( + [ + "benchmark", + "sglang", + "--model", + "thinkingmachines/Inkling", + "--dataset", + "gsm8k", + ] + ) + + self.assertEqual(status, 7) + self.assertEqual(run.call_args.args[0].model, "thinkingmachines/Inkling") + + def test_cli_help_describes_the_backend_not_an_algorithm(self): + from specforge.cli import main + + output = StringIO() + with redirect_stdout(output), self.assertRaises(SystemExit) as exited: + main(["benchmark", "sglang", "--help"]) + + self.assertEqual(exited.exception.code, 0) + help_text = " ".join(output.getvalue().split()) + self.assertIn("a running SGLang server", help_text) + self.assertNotIn("DSpark", help_text) + self.assertNotIn("DFlash", help_text) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_config/test_launch_topology.py b/tests/test_config/test_launch_topology.py index ace91c73e..8af4174f5 100644 --- a/tests/test_config/test_launch_topology.py +++ b/tests/test_config/test_launch_topology.py @@ -15,9 +15,11 @@ "deepseek-v3-671b-eagle3-offline.yaml": 8, "deepseek-v3-671b-eagle3-online.yaml": 8, "gemma3-1b-eagle3-online.yaml": 1, + "glm-5.2-dspark-disaggregated.yaml": 1, "gpt-oss-120b-eagle3-online.yaml": 8, "gpt-oss-20b-eagle3-online.yaml": 8, "lfm2.5-1.2b-instruct-dflash-online.yaml": 8, + "inkling-dspark-disaggregated.yaml": 1, "ling-flash-2.0-eagle3-offline.yaml": 8, "ling-flash-2.0-eagle3-online.yaml": 8, "llama3.1-8b-eagle3-offline.yaml": 1, @@ -45,6 +47,7 @@ "qwen3-8b-domino-multiserver-disaggregated.yaml": 2, "qwen3-8b-domino-online.yaml": 8, "qwen3-8b-dpace-online.yaml": 8, + "qwen3-8b-dspark-disaggregated.yaml": 1, "qwen3-8b-dta-online.yaml": 8, "qwen3-8b-eagle3-offline-disaggregated.yaml": 1, "qwen3-8b-eagle3-offline.yaml": 1, @@ -75,6 +78,18 @@ } EXPECTED_DISAGGREGATED = { + "glm-5.2-dspark-disaggregated.yaml": { + "control_dir": "outputs/glm-5.2-dspark-disaggregated/control", + "backend": "mooncake", + "server_urls": ["http://127.0.0.1:30000"], + **LOCAL_MOONCAKE_ENDPOINTS, + }, + "inkling-dspark-disaggregated.yaml": { + "control_dir": "outputs/inkling-dspark-disaggregated/control", + "backend": "mooncake", + "server_urls": ["http://127.0.0.1:30000"], + **LOCAL_MOONCAKE_ENDPOINTS, + }, "qwen2.5-7b-eagle3-offline-disaggregated.yaml": { "control_dir": ("outputs/qwen2.5-7b-eagle3-offline-disaggregated/control"), "backend": "shared_dir", @@ -93,6 +108,12 @@ "server_urls": ["http://127.0.0.1:30000"], **LOCAL_MOONCAKE_ENDPOINTS, }, + "qwen3-8b-dspark-disaggregated.yaml": { + "control_dir": "outputs/qwen3-8b-dspark-disaggregated/control", + "backend": "mooncake", + "server_urls": ["http://127.0.0.1:30000"], + **LOCAL_MOONCAKE_ENDPOINTS, + }, "qwen3-8b-dflash-1server-dp7-disaggregated.yaml": { "control_dir": ("outputs/qwen3-8b-dflash-1server-dp7-disaggregated/control"), "backend": "mooncake", @@ -266,7 +287,7 @@ def _recipes() -> dict[str, Path]: class ExampleLaunchTopologyTest(unittest.TestCase): def test_every_recipe_has_the_explicit_golden_topology(self): recipes = _recipes() - self.assertEqual(len(EXPECTED_NPROC_PER_NODE), 55) + self.assertEqual(len(EXPECTED_NPROC_PER_NODE), 58) self.assertEqual(set(recipes), set(EXPECTED_NPROC_PER_NODE)) for filename, nproc_per_node in EXPECTED_NPROC_PER_NODE.items(): @@ -334,6 +355,42 @@ def test_every_recipe_keeps_trainer_tp_at_one(self): self.assertEqual(config.training.sp_ulysses_size, 1) self.assertEqual(config.training.sp_ring_size, 1) + def test_migrated_dspark_recipes_match_source_training_contract(self): + expected_save_intervals = { + "qwen3-8b-dspark-disaggregated.yaml": 125, + "glm-5.2-dspark-disaggregated.yaml": 16, + "inkling-dspark-disaggregated.yaml": 8, + } + for filename, save_interval in expected_save_intervals.items(): + with self.subTest(config=filename): + config = Config.from_file(str(EXAMPLE_CONFIG_DIR / filename)) + topology = config.deployment.trainer + global_batch = ( + topology.nnodes + * topology.nproc_per_node + * config.training.batch_size + * config.training.accumulation_steps + ) + self.assertEqual(global_batch, 512) + self.assertEqual(config.training.num_epochs, 10) + self.assertIsNone(config.training.max_steps) + self.assertAlmostEqual(config.training.learning_rate, 6e-4) + self.assertEqual(config.training.num_anchors, 512) + self.assertEqual(config.training.loss_decay_gamma, 4.0) + self.assertEqual(config.training.objective_chunk_blocks, 128) + self.assertEqual(config.training.save_interval, save_interval) + + qwen = Config.from_file( + str(EXAMPLE_CONFIG_DIR / "qwen3-8b-dspark-disaggregated.yaml") + ) + self.assertEqual(qwen.data.chat_template, "qwen") + + qwen4b = Config.from_file( + str(EXAMPLE_CONFIG_DIR / "qwen3-4b-dspark-disaggregated.yaml") + ) + self.assertEqual(qwen4b.training.loss_decay_gamma, 4.0) + self.assertEqual(qwen4b.training.objective_chunk_blocks, 128) + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/tests/test_config/test_schema.py b/tests/test_config/test_schema.py index a2a86b379..b37ace421 100644 --- a/tests/test_config/test_schema.py +++ b/tests/test_config/test_schema.py @@ -303,6 +303,30 @@ def test_unknown_backend_rejected(self): with self.assertRaises(ValidationError): Config.model_validate(bad) + def test_removed_dspark_objective_selector_is_rejected(self): + payload = _online_payload("dspark") + payload["training"]["dspark_objective_mode"] = "legacy" + with self.assertRaisesRegex(ValidationError, "dspark_objective_mode"): + Config.model_validate(payload) + + def test_objective_chunk_blocks_is_shared_and_typed(self): + for strategy in ("dflash", "domino", "dspark"): + payload = _online_payload(strategy) + payload["training"]["objective_chunk_blocks"] = 0 + with self.subTest(strategy=strategy): + config = Config.model_validate(payload) + self.assertEqual(config.training.objective_chunk_blocks, 0) + + payload = _online_payload("dflash") + payload["training"]["objective_chunk_blocks"] = -1 + with self.assertRaisesRegex(ValidationError, "objective_chunk_blocks"): + Config.model_validate(payload) + + payload = _online_payload("dspark") + payload["training"]["dspark_objective_chunk_blocks"] = 1 + with self.assertRaisesRegex(ValidationError, "dspark_objective_chunk_blocks"): + Config.model_validate(payload) + def test_unknown_fields_and_unsupported_modes_fail_early(self): with self.assertRaises(ValidationError): Config.model_validate( diff --git a/tests/test_config/test_unified_feature_reachability.py b/tests/test_config/test_unified_feature_reachability.py index 3bce6ba14..baf031339 100644 --- a/tests/test_config/test_unified_feature_reachability.py +++ b/tests/test_config/test_unified_feature_reachability.py @@ -151,7 +151,7 @@ def test_all_example_configs_validate_through_the_typed_entry(self): for path in EXAMPLE_CONFIG_DIR.glob("*.yaml") if not path.name.startswith(".") ) - self.assertEqual(len(paths), 55) + self.assertEqual(len(paths), 58) resolved_runs = { path.name: resolve_run(Config.from_file(str(path))) for path in paths diff --git a/tests/test_data/test_inkling_template.py b/tests/test_data/test_inkling_template.py new file mode 100644 index 000000000..c7b76f2d5 --- /dev/null +++ b/tests/test_data/test_inkling_template.py @@ -0,0 +1,66 @@ +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace + +from specforge.data.parse import ThinkingParser +from specforge.data.template import TEMPLATE_REGISTRY +from specforge.training.assembly import _prompt_cache_key + + +def _cache_config(path: str): + return SimpleNamespace( + data=SimpleNamespace( + prompts_path="", + train_data_path=path, + max_length=4096, + chat_template="inkling-thinking", + is_preformatted=False, + train_only_last_turn=False, + max_prompts=None, + ), + model=SimpleNamespace( + target_model_path="thinkingmachines/Inkling", + draft_model_config="configs/inkling-dspark.json", + draft_checkpoint_path=None, + draft_num_hidden_layers=None, + draft_block_size=None, + input_modality="text", + ), + training=SimpleNamespace(strategy="dspark"), + ) + + +class TestInklingTemplate(unittest.TestCase): + def test_tool_response_identity_survives_sanitization(self): + parser = ThinkingParser( + SimpleNamespace(), + TEMPLATE_REGISTRY.get("inkling-thinking"), + ) + cleaned = parser._sanitize_message( + { + "role": "tool", + "content": "sunny", + "name": "weather", + "tool_call_id": "call-7", + "private": "discard", + } + ) + self.assertEqual(cleaned["name"], "weather") + self.assertEqual(cleaned["tool_call_id"], "call-7") + self.assertNotIn("private", cleaned) + + def test_prompt_cache_key_tracks_source_content(self): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp, "prompts.jsonl") + path.write_text('{"text":"first"}\n', encoding="utf-8") + config = _cache_config(str(path)) + first = _prompt_cache_key(config) + path.write_text('{"text":"second"}\n', encoding="utf-8") + second = _prompt_cache_key(config) + + self.assertNotEqual(first, second) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_modeling/test_domino_model.py b/tests/test_modeling/test_domino_model.py index 970f9d9e3..c80d917e5 100644 --- a/tests/test_modeling/test_domino_model.py +++ b/tests/test_modeling/test_domino_model.py @@ -1,6 +1,7 @@ # coding=utf-8 """Focused CPU regressions for Domino's auxiliary logits head.""" +import copy import unittest from types import MethodType from unittest.mock import patch @@ -102,6 +103,123 @@ def fixed_draft_blocks(self, input_ids, hidden_states, loss_mask): 0.0, ) + def test_chunked_objective_matches_full_loss_metrics_and_gradients(self): + from specforge.algorithms.common.dflash_family_model import OnlineDominoModel + + hidden_size, vocab_size, block_size = 4, 7, 4 + for shift_label in (False, True): + with self.subTest(shift_label=shift_label): + torch.manual_seed(19) + draft = _bare_domino( + hidden_size=hidden_size, + gru_hidden_size=3, + embedding_size=2, + vocab_size=vocab_size, + block_size=block_size, + ) + draft.shift_label = shift_label + target_head = nn.Linear(hidden_size, vocab_size, bias=False) + target_embedding = nn.Embedding(vocab_size, hidden_size) + target_head.requires_grad_(False) + target_embedding.requires_grad_(False) + full = OnlineDominoModel( + draft_model=draft, + target_lm_head=target_head, + target_embed_tokens=target_embedding, + mask_token_id=0, + block_size=block_size, + attention_backend="sdpa", + num_anchors=2, + objective_chunk_blocks=0, + shift_label=shift_label, + ) + chunked = copy.deepcopy(full) + chunked.objective_chunk_blocks = 1 + + anchors = torch.tensor([[0, 3]]) + keep_mask = torch.ones(1, 2, dtype=torch.bool) + full_hidden = torch.randn( + 1, + 2 * block_size, + hidden_size, + requires_grad=True, + ) + chunked_hidden = full_hidden.detach().clone().requires_grad_() + + def fixed_blocks(output_hidden): + def _forward(self, input_ids, hidden_states, loss_mask): + del self, input_ids, hidden_states, loss_mask + return anchors, keep_mask, output_hidden + + return _forward + + full._forward_draft_blocks = MethodType( + fixed_blocks(full_hidden), + full, + ) + chunked._forward_draft_blocks = MethodType( + fixed_blocks(chunked_hidden), + chunked, + ) + input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 0, 1]]) + loss_mask = torch.tensor([[1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0]]) + hidden_states = torch.zeros(1, input_ids.shape[1], hidden_size) + + full_loss, full_accuracy, full_metrics = full( + input_ids=input_ids, + hidden_states=hidden_states, + loss_mask=loss_mask, + lambda_base=0.35, + ) + chunked_loss, chunked_accuracy, chunked_metrics = chunked( + input_ids=input_ids, + hidden_states=hidden_states, + loss_mask=loss_mask, + lambda_base=0.35, + ) + + torch.testing.assert_close( + chunked_loss, + full_loss, + rtol=1e-6, + atol=1e-7, + ) + torch.testing.assert_close(chunked_accuracy, full_accuracy) + self.assertEqual(chunked_metrics.keys(), full_metrics.keys()) + for name in full_metrics: + torch.testing.assert_close( + chunked_metrics[name], + full_metrics[name], + rtol=1e-6, + atol=1e-7, + ) + + full_loss.backward() + chunked_loss.backward() + torch.testing.assert_close( + chunked_hidden.grad, + full_hidden.grad, + rtol=1e-6, + atol=1e-7, + ) + full_parameters = dict(full.draft_model.named_parameters()) + chunked_parameters = dict(chunked.draft_model.named_parameters()) + self.assertEqual(full_parameters.keys(), chunked_parameters.keys()) + for name, parameter in full_parameters.items(): + chunked_parameter = chunked_parameters[name] + self.assertEqual( + chunked_parameter.grad is None, + parameter.grad is None, + name, + ) + if parameter.grad is not None: + torch.testing.assert_close( + chunked_parameter.grad, + parameter.grad, + rtol=1e-6, + atol=1e-7, + ) + def test_npu_bf16_gru_gradients_reach_registered_weights(self): torch.manual_seed(7) model = _bare_domino().to(dtype=torch.bfloat16) diff --git a/tests/test_modeling/test_draft_registry.py b/tests/test_modeling/test_draft_registry.py index fe97555a1..7dea35200 100644 --- a/tests/test_modeling/test_draft_registry.py +++ b/tests/test_modeling/test_draft_registry.py @@ -5,11 +5,15 @@ auto loaders resolve it from that registry. """ +import copy import json import os import tempfile import unittest +from types import SimpleNamespace +from unittest import mock +import torch from transformers import LlamaConfig, Qwen3Config from specforge.modeling.auto import AutoDraftModel, AutoDraftModelConfig @@ -154,6 +158,47 @@ def test_from_config_builds_dspark_as_explicit_draft(self): self.assertIsNotNone(model.markov_head) self.assertIsNotNone(model.confidence_head) + def test_dspark_defaults_to_gqa_and_requires_explicit_mha(self): + implicit_mha = copy.deepcopy(TINY_DSPARK) + implicit_mha["num_key_value_heads"] = implicit_mha["num_attention_heads"] + path = _write(implicit_mha) + self.addCleanup(os.unlink, path) + config = AutoDraftModelConfig.from_file(path) + with self.assertRaisesRegex(ValueError, "defaults to GQA"): + AutoDraftModel.from_config(config) + + explicit_mha = copy.deepcopy(implicit_mha) + explicit_mha["dflash_config"]["attention_mode"] = "mha" + path = _write(explicit_mha) + self.addCleanup(os.unlink, path) + config = AutoDraftModelConfig.from_file(path) + model = AutoDraftModel.from_config(config) + self.assertEqual(model.config.dflash_config["attention_mode"], "mha") + + def test_dspark_generation_uses_the_markov_head(self): + path = _write(TINY_DSPARK) + self.addCleanup(os.unlink, path) + model = AutoDraftModel.from_config(AutoDraftModelConfig.from_file(path)) + expected = torch.tensor([[11, 12, 13]]) + target = SimpleNamespace( + lm_head=mock.Mock(return_value=torch.zeros(1, model.block_size - 1, 256)) + ) + draft_hidden = torch.zeros(1, model.block_size, model.config.hidden_size) + block_ids = torch.tensor([[7, 0, 0, 0]]) + with mock.patch.object( + model.markov_head, + "sample_block_tokens", + return_value=(expected, None), + ) as sampler: + actual = model._sample_draft_tokens(target, draft_hidden, block_ids) + + torch.testing.assert_close(actual, expected) + self.assertEqual(sampler.call_args.kwargs["first_prev_token_ids"].item(), 7) + self.assertEqual( + tuple(sampler.call_args.kwargs["hidden_states"].shape), + (1, model.block_size - 1, model.config.hidden_size), + ) + def test_from_config_builds_domino_as_explicit_draft(self): path = _write(TINY_DOMINO) self.addCleanup(os.unlink, path) diff --git a/tests/test_modeling/test_target/test_target_utils_loading.py b/tests/test_modeling/test_target/test_target_utils_loading.py index d907cc74e..61dd1717c 100644 --- a/tests/test_modeling/test_target/test_target_utils_loading.py +++ b/tests/test_modeling/test_target/test_target_utils_loading.py @@ -7,7 +7,10 @@ import torch -from specforge.modeling.target.target_utils import TargetEmbeddingsAndHead +from specforge.modeling.target.target_utils import ( + TargetEmbeddingsAndHead, + load_target_config, +) class _FakeSafeOpen: @@ -162,6 +165,88 @@ def test_shape_mismatch_fails_closed(self): tmp, self.embed_key, self.head_key, tie_weights=False ) + def test_nested_text_config_uses_padded_vocab(self): + config = SimpleNamespace( + text_config=SimpleNamespace( + vocab_size=4, + padded_vocab_size=6, + hidden_size=3, + pad_token_id=5, + ) + ) + module = TargetEmbeddingsAndHead(config) + self.assertEqual(tuple(module.embed_tokens.weight.shape), (6, 3)) + self.assertEqual(tuple(module.lm_head.weight.shape), (6, 3)) + self.assertEqual(module.embed_tokens.padding_idx, 5) + + def test_public_raw_config_fallback_preserves_nested_fields(self): + with tempfile.TemporaryDirectory() as tmp: + Path(tmp, "config.json").write_text( + json.dumps( + { + "model_type": "not-yet-registered", + "text_config": { + "vocab_size": 4, + "padded_vocab_size": 6, + "hidden_size": 3, + }, + } + ), + encoding="utf-8", + ) + with patch( + "specforge.modeling.target.target_utils.AutoConfig.from_pretrained", + side_effect=ValueError("unknown model type"), + ): + config = load_target_config(tmp) + + self.assertEqual(config.text_config.hidden_size, 3) + self.assertEqual(config.text_config.padded_vocab_size, 6) + + def test_logits_mup_multiplier_is_folded_into_frozen_head(self): + config = SimpleNamespace( + tie_word_embeddings=False, + text_config=SimpleNamespace( + vocab_size=4, + padded_vocab_size=6, + hidden_size=3, + pad_token_id=None, + logits_mup_width_multiplier=4.0, + ), + ) + + def load_weights(instance, *_args, **_kwargs): + instance.embed_tokens.weight.data.fill_(2.0) + instance.lm_head.weight.data.fill_(8.0) + return {self.embed_key, self.head_key} + + with ( + tempfile.TemporaryDirectory() as tmp, + patch( + "specforge.modeling.target.target_utils.load_target_config", + return_value=config, + ), + patch.object(TargetEmbeddingsAndHead, "_load_weights", load_weights), + ): + module = TargetEmbeddingsAndHead.from_pretrained( + tmp, + embed_key=self.embed_key, + lm_head_key=self.head_key, + device="cpu", + dtype=torch.float32, + ) + + torch.testing.assert_close( + module.embed_tokens.weight, + torch.full_like(module.embed_tokens.weight, 2.0), + ) + torch.testing.assert_close( + module.lm_head.weight, + torch.full_like(module.lm_head.weight, 2.0), + ) + self.assertEqual(module.lm_head_mup_folded, 4.0) + self.assertFalse(module.lm_head.weight.requires_grad) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_optimizer/test_bf16_optimizer_clip_grad_norm.py b/tests/test_optimizer/test_bf16_optimizer_clip_grad_norm.py index 40e4abec7..377ae7687 100644 --- a/tests/test_optimizer/test_bf16_optimizer_clip_grad_norm.py +++ b/tests/test_optimizer/test_bf16_optimizer_clip_grad_norm.py @@ -10,10 +10,10 @@ from specforge.training.backend import FSDPTrainingBackend, ParallelConfig -def _make_optimizer(seed=0): +def _make_optimizer(seed=0, **kwargs): torch.manual_seed(seed) model = torch.nn.Linear(8, 8, bias=False) - return model, BF16Optimizer(model, lr=1e-3, max_grad_norm=0.5) + return model, BF16Optimizer(model, lr=1e-3, max_grad_norm=0.5, **kwargs) class TestClipGradNormSingleProcess(unittest.TestCase): @@ -71,6 +71,114 @@ def configure_grad_norm_reduction(self, **kwargs): backend.set_optimizer(replicated) self.assertFalse(replicated.config["enabled"]) + def test_cpu_offload_matches_resident_optimizer_update(self): + resident_model, resident = _make_optimizer(seed=7, offload_master=False) + offload_model, offload = _make_optimizer(seed=7, offload_master=True) + grad = torch.linspace(-0.2, 0.3, 64).reshape(8, 8) + resident_model.weight.grad = grad.clone() + offload_model.weight.grad = grad.clone() + + resident_norm = resident.step() + offload_norm = offload.step() + + torch.testing.assert_close(offload_norm, resident_norm) + torch.testing.assert_close(offload_model.weight, resident_model.weight) + self.assertTrue( + all(param.device.type == "cpu" for param in offload.fp32_params) + ) + self.assertTrue( + all( + tensor.device.type == "cpu" + for state in offload.optimizer.state.values() + for tensor in state.values() + if isinstance(tensor, torch.Tensor) + ) + ) + + def test_resume_allows_cpu_offload_mode_change(self): + model, resident = _make_optimizer(seed=3, offload_master=False) + for param in model.parameters(): + param.grad = torch.full_like(param, 0.05) + resident.step() + checkpoint = resident.state_dict() + + # Toggling CPU offload on resume is a pure placement change and is + # allowed: masters and Adam moments land on the current master device. + _offload_model, offload = _make_optimizer(seed=99, offload_master=True) + # load_state_dict logs via rank0, which needs a process group. + created_pg = False + if dist.is_available() and not dist.is_initialized(): + store = dist.FileStore( + os.path.join(tempfile.mkdtemp(prefix="opt_pg_"), "store"), 1 + ) + dist.init_process_group("gloo", store=store, rank=0, world_size=1) + created_pg = True + try: + offload.load_state_dict(checkpoint) + finally: + if created_pg: + dist.destroy_process_group() + + for restored, saved in zip(offload.fp32_params, checkpoint["fp32_params"]): + self.assertEqual(restored.device.type, "cpu") + torch.testing.assert_close(restored.detach(), saved) + self.assertTrue( + all( + tensor.device.type == "cpu" + for state in offload.optimizer.state.values() + for tensor in state.values() + if isinstance(tensor, torch.Tensor) + ) + ) + + @unittest.skipUnless( + torch.cuda.is_available(), "requires CUDA to exercise offload transfers" + ) + def test_cpu_offload_gpu_model_matches_resident_update(self): + torch.manual_seed(7) + resident_model = torch.nn.Linear(8, 8, bias=False).cuda() + torch.manual_seed(7) + offload_model = torch.nn.Linear(8, 8, bias=False).cuda() + resident = BF16Optimizer( + resident_model, lr=1e-3, max_grad_norm=0.5, offload_master=False + ) + offload = BF16Optimizer( + offload_model, lr=1e-3, max_grad_norm=0.5, offload_master=True + ) + grad = torch.linspace(-0.2, 0.3, 64).reshape(8, 8).cuda() + resident_model.weight.grad = grad.clone() + offload_model.weight.grad = grad.clone() + + resident_norm = resident.step() + offload_norm = offload.step() + + # Norm is reduced on the model device (CUDA) in both cases. + self.assertEqual(resident_norm.device.type, "cuda") + self.assertEqual(offload_norm.device.type, "cuda") + torch.testing.assert_close(offload_norm, resident_norm) + + # The trainable draft stays on the accelerator and is updated in place. + self.assertEqual(offload_model.weight.device.type, "cuda") + torch.testing.assert_close( + offload_model.weight, resident_model.weight, atol=1e-5, rtol=1e-4 + ) + + # Resident masters/Adam state live on CUDA; offloaded ones live on CPU. + self.assertTrue( + all(param.device.type == "cuda" for param in resident.fp32_params) + ) + self.assertTrue( + all(param.device.type == "cpu" for param in offload.fp32_params) + ) + self.assertTrue( + all( + tensor.device.type == "cpu" + for state in offload.optimizer.state.values() + for tensor in state.values() + if isinstance(tensor, torch.Tensor) + ) + ) + def _distributed_worker(rank, world_size, init_file, results): dist.init_process_group( @@ -103,7 +211,43 @@ def _distributed_worker(rank, world_size, init_file, results): dist.destroy_process_group() +def _distributed_step_worker(rank, world_size, init_file, results): + dist.init_process_group( + backend="gloo", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + ) + try: + model, optimizer = _make_optimizer() + grad_value = 1.0 if rank == 0 else 2.0 + for param in model.parameters(): + param.grad = torch.full_like(param, grad_value) + # Drive the production path: step() reduces the norm on the model device. + norm = optimizer.step() + results[rank] = norm.item() + finally: + dist.destroy_process_group() + + class TestClipGradNormDistributed(unittest.TestCase): + def test_step_reduces_norm_across_ranks(self): + world_size = 2 + with tempfile.TemporaryDirectory() as tmpdir: + init_file = os.path.join(tmpdir, "init") + manager = mp.Manager() + results = manager.dict() + mp.spawn( + _distributed_step_worker, + args=(world_size, init_file, results), + nprocs=world_size, + join=True, + ) + + global_norm = (64 * 1.0**2 + 64 * 2.0**2) ** 0.5 + for rank in range(world_size): + self.assertAlmostEqual(results[rank], global_norm, places=4) + def test_disjoint_shards_use_same_global_clip_coefficient(self): world_size = 2 with tempfile.TemporaryDirectory() as tmpdir: diff --git a/tests/test_runtime/test_checkpoint_manager.py b/tests/test_runtime/test_checkpoint_manager.py index 8b6b5ecbd..b05d00ab6 100644 --- a/tests/test_runtime/test_checkpoint_manager.py +++ b/tests/test_runtime/test_checkpoint_manager.py @@ -12,6 +12,7 @@ import tempfile import unittest from datetime import timedelta +from unittest import mock import torch @@ -358,6 +359,55 @@ def test_world_size_mismatch(self): st = CheckpointManager.read_resume_state(ckpt, require_full_state=False) self.assertEqual(st["backend"], {}) + def test_resume_fails_world_wide_when_any_rank_cannot_read(self): + from specforge.training.checkpoint import CheckpointManager + + out = tempfile.mkdtemp(prefix="ckpt_world_read_") + ckpt = _mgr(out).save( + _state(4, world_size=2, run_id="run", strategy="dspark"), + 4, + rank_state={"optimizer": {}, "rng": {}}, + ) + + def gather(descriptors, local): + descriptors[0] = local + descriptors[1] = { + **local, + "error": "FileNotFoundError: rank-local checkpoint missing", + } + + with ( + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.get_world_size", return_value=2), + mock.patch("torch.distributed.get_rank", return_value=0), + mock.patch("torch.distributed.all_gather_object", side_effect=gather), + self.assertRaisesRegex(RuntimeError, "not readable on every rank"), + ): + CheckpointManager.read_resume_state(ckpt) + + def test_resume_rejects_different_checkpoint_identity_across_ranks(self): + from specforge.training.checkpoint import CheckpointManager + + out = tempfile.mkdtemp(prefix="ckpt_world_identity_") + ckpt = _mgr(out).save( + _state(4, world_size=2, run_id="run", strategy="dspark"), + 4, + rank_state={"optimizer": {}, "rng": {}}, + ) + + def gather(descriptors, local): + descriptors[0] = local + descriptors[1] = {**local, "global_step": 3} + + with ( + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.get_world_size", return_value=2), + mock.patch("torch.distributed.get_rank", return_value=0), + mock.patch("torch.distributed.all_gather_object", side_effect=gather), + self.assertRaisesRegex(ValueError, "different training states"), + ): + CheckpointManager.read_resume_state(ckpt) + def _dist_worker(rank, world, port, out_dir, results_dir): import torch.distributed as dist diff --git a/tests/test_runtime/test_launch_plan.py b/tests/test_runtime/test_launch_plan.py index a819dd966..92a9176ee 100644 --- a/tests/test_runtime/test_launch_plan.py +++ b/tests/test_runtime/test_launch_plan.py @@ -539,7 +539,7 @@ def test_managed_local_rejects_external_and_nonlocal_modes(self): with self.assertRaisesRegex(ValidationError, message): Config.model_validate(raw) - def test_managed_local_accepts_minimum_explicit_capture_context(self): + def test_managed_local_accepts_minimum_context_and_leaves_radix_enabled(self): with tempfile.TemporaryDirectory() as root: cfg = _managed_config(os.path.join(root, "attempt")) raw = cfg.model_dump() @@ -554,6 +554,7 @@ def test_managed_local_accepts_minimum_explicit_capture_context(self): argv = plan.services[1].command.argv self.assertEqual(argv[argv.index("--context-length") + 1], "135") + self.assertNotIn("--disable-radix-cache", argv) def test_managed_local_plan_owns_mooncake_and_multiple_capture_servers(self): servers = [ @@ -584,6 +585,16 @@ def test_managed_local_plan_owns_mooncake_and_multiple_capture_servers(self): sglang_enable_torch_compile=True, sglang_max_running_requests=64, sglang_max_total_tokens=8192, + sglang_dp_size=2, + sglang_moe_a2a_backend="deepep", + sglang_moe_runner_backend="triton", + sglang_page_size=64, + sglang_quantization="fp8", + sglang_fp4_gemm_runner_backend="cutlass", + sglang_mamba_radix_cache_strategy="lru", + sglang_max_mamba_cache_size=1024, + sglang_swa_full_tokens_ratio=0.5, + sglang_mamba_full_memory_ratio=0.25, ) cfg = Config.model_validate(raw) with mock.patch( @@ -647,6 +658,19 @@ def test_managed_local_plan_owns_mooncake_and_multiple_capture_servers(self): self.assertIn(flag, argv) self.assertEqual(argv[argv.index("--max-running-requests") + 1], "64") self.assertEqual(argv[argv.index("--max-total-tokens") + 1], "8192") + for flag, expected in ( + ("--dp-size", "2"), + ("--moe-a2a-backend", "deepep"), + ("--moe-runner-backend", "triton"), + ("--page-size", "64"), + ("--quantization", "fp8"), + ("--fp4-gemm-runner-backend", "cutlass"), + ("--mamba-radix-cache-strategy", "lru"), + ("--max-mamba-cache-size", "1024"), + ("--swa-full-tokens-ratio", "0.5"), + ("--mamba-full-memory-ratio", "0.25"), + ): + self.assertEqual(argv[argv.index(flag) + 1], expected) self.assertEqual(argv[argv.index("--context-length") + 1], "2055") producer, consumer = plan.commands expected_urls = "http://127.0.0.1:30000,http://127.0.0.1:30001" diff --git a/tests/test_runtime/test_package_architecture.py b/tests/test_runtime/test_package_architecture.py index 3ee02f4ac..3960ad777 100644 --- a/tests/test_runtime/test_package_architecture.py +++ b/tests/test_runtime/test_package_architecture.py @@ -3,6 +3,7 @@ from __future__ import annotations import ast +import json import re import tokenize import tomllib @@ -255,8 +256,11 @@ } CANONICAL_DRAFT_CONFIGS = { + "glm-5.2-dspark.json", + "inkling-dspark.json", "llama3-8B-eagle3.json", "qwen3-4b-dspark.json", + "qwen3-8b-dspark.json", "qwen3-8b-dflash.json", "qwen3-8b-domino.json", "qwen3-8b-eagle3.json", @@ -697,6 +701,35 @@ def test_only_canonical_draft_configs_are_checked_in(self): present = {path.name for path in (REPO_ROOT / "configs").glob("*.json")} self.assertTrue(CANONICAL_DRAFT_CONFIGS.issubset(present)) + def test_dspark_configs_are_qwen3_gqa_only(self): + dspark_configs = {} + for path in sorted((REPO_ROOT / "configs").glob("*.json")): + payload = json.loads(path.read_text(encoding="utf-8")) + if "DSparkDraftModel" in payload.get("architectures", []): + dspark_configs[path.name] = payload + + self.assertEqual( + set(dspark_configs), + { + "glm-5.2-dspark.json", + "inkling-dspark.json", + "qwen3-4b-dspark.json", + "qwen3-8b-dspark.json", + }, + ) + for name, payload in dspark_configs.items(): + with self.subTest(config=name): + self.assertEqual(payload["model_type"], "qwen3") + self.assertEqual(payload["dflash_config"]["attention_mode"], "gqa") + self.assertLess( + payload["num_key_value_heads"], + payload["num_attention_heads"], + ) + self.assertEqual( + payload["num_attention_heads"] % payload["num_key_value_heads"], + 0, + ) + def test_examples_and_scripts_do_not_bypass_the_cli(self): direct_imports = [] for root in (REPO_ROOT / "scripts", REPO_ROOT / "examples"): diff --git a/tests/test_runtime/test_server_capture.py b/tests/test_runtime/test_server_capture.py index 68e6ad62c..3b6a41635 100644 --- a/tests/test_runtime/test_server_capture.py +++ b/tests/test_runtime/test_server_capture.py @@ -326,10 +326,13 @@ def recording_server(url, json_body, timeout): request["sampling_params"], ) self.assertEqual(2, len(request["spec_capture"])) + self.assertEqual(2, len(request["extra_key"])) + self.assertEqual(2, len(set(request["extra_key"]))) + self.assertTrue(all(len(value) == 32 for value in request["extra_key"])) def test_generic_adapter_cannot_override_runtime_owned_request_fields(self): task = _task(0, 4) - for reserved_field in ("sampling_params", "spec_capture"): + for reserved_field in ("extra_key", "sampling_params", "spec_capture"): with self.subTest(field=reserved_field): input_adapter = _GenericRequestInputAdapter( { @@ -558,9 +561,11 @@ def test_retry_reuses_generation_after_a_lost_response(self): ) original_post = adapter.post_fn requests = [] + cache_keys = [] def lose_first_response(url, json_body, timeout): requests.append(json_body["spec_capture"][0]) + cache_keys.append(json_body["extra_key"][0]) response = original_post(url, json_body, timeout) if len(requests) == 1: raise ConnectionError("response lost after server write") @@ -584,6 +589,7 @@ def lose_first_response(url, json_body, timeout): self.assertEqual(1, ref.metadata["generation"]) self.assertEqual([False, True], [request["replace"] for request in requests]) + self.assertNotEqual(cache_keys[0], cache_keys[1]) self.assertEqual(0, store.discard_external_attempts()) self.assertTrue(all("/g1/" in key for key in backend._d)) self.assertFalse(any("/g2/" in key for key in backend._d)) diff --git a/tests/test_runtime/test_server_capture_gate.py b/tests/test_runtime/test_server_capture_gate.py index 81598adf2..0d264fbab 100644 --- a/tests/test_runtime/test_server_capture_gate.py +++ b/tests/test_runtime/test_server_capture_gate.py @@ -14,6 +14,9 @@ agnostic), within the documented bf16 tolerance; - strategy-agnosticism: the same server serves eagle3 (aux + last_hidden) and dflash (aux only) requests, named per strategy by the client schema. +- cache isolation: radix cache stays enabled and sequential identical prompts + still capture every token because each request attempt has a fresh + ``extra_key`` namespace. The PR workflow runs this gate explicitly on its GPU runner. Local runs need a GPU, sglang patched with @@ -142,7 +145,6 @@ def setUpClass(cls): "0.3", "--chunked-prefill-size", "-1", - "--disable-radix-cache", "--enable-spec-capture", "--spec-capture-aux-layer-ids", *[str(i) for i in AUX_LAYER_IDS], @@ -278,7 +280,7 @@ def test_eagle3_zero_copy_end_to_end(self): from specforge.inference.capture import CaptureConfig from specforge.runtime.contracts import SampleRef - rows = [[5, 6, 7, 8, 9, 10], [11, 12, 13, 14]] + rows = [[5, 6, 7, 8, 9, 10], [5, 6, 7, 8, 9, 10]] store = self._store("gate-eagle3") adapter = SGLangServerCaptureAdapter( f"http://localhost:{PORT}", @@ -299,7 +301,11 @@ def test_eagle3_zero_copy_end_to_end(self): target_repr="hidden_state", target_hidden_size=H, ) - refs = adapter.produce_refs(self._tasks(rows), capture=contract) + # Submit sequentially so the second identical prompt would hit the + # first request's radix entry if the adapter did not isolate attempts. + refs = [] + for task in self._tasks(rows): + refs.extend(adapter.produce_refs([task], capture=contract)) for ref in refs: self.assertIsInstance(ref, SampleRef, f"expected a ref, got failure: {ref}") diff --git a/tests/test_runtime/test_trainer.py b/tests/test_runtime/test_trainer.py index eee4dba39..efc8ebcfc 100644 --- a/tests/test_runtime/test_trainer.py +++ b/tests/test_runtime/test_trainer.py @@ -18,7 +18,12 @@ from specforge.runtime.data_plane.sample_ref_queue import SampleRefQueue from specforge.training.backend import TrainingBackend from specforge.training.checkpoint import CheckpointManager -from specforge.training.controller import Checkpoint, TrainerController, TrainerCore +from specforge.training.controller import ( + Checkpoint, + TrainerController, + TrainerCore, + _reduce_ratio_metrics, +) from specforge.training.strategies.base import DraftTrainStrategy, StepOutput @@ -167,6 +172,58 @@ def test_strategy_scalar_metrics_are_preserved(self): self.assertEqual(result.metrics["lambda_base"], 0.75) self.assertNotIn("non_scalar_debug", result.metrics) + def test_ratio_metrics_override_mean_of_means_accuracy(self): + strat = FakeStrategy() + core = TrainerCore(strat, FakeBackend(strat.model), accumulation_steps=1) + result = core._result( + StepOutput( + loss=torch.tensor(2.0), + metrics={"accuracy": torch.tensor(0.99)}, + ratio_metrics={ + "acc": (torch.tensor(2.0), torch.tensor(4.0)), + "ce_position": ( + torch.tensor([1.0, 3.0]), + torch.tensor([2.0, 6.0]), + ), + }, + ), + grad_norm=None, + stepped=False, + ) + + self.assertEqual(result.metrics["acc"], 0.5) + self.assertEqual(result.metrics["ce_position_0"], 0.5) + self.assertEqual(result.metrics["ce_position_1"], 0.5) + + def test_ratio_metrics_sum_numerators_and_denominators_before_dividing(self): + remote = torch.tensor([8.0, 6.0, 3.0, 1.0, 2.0, 2.0]) + + def all_reduce(packed, *, group): + self.assertEqual(group, "dp") + packed.add_(remote) + + with ( + mock.patch("torch.distributed.is_available", return_value=True), + mock.patch("torch.distributed.is_initialized", return_value=True), + mock.patch("torch.distributed.all_reduce", side_effect=all_reduce), + ): + metrics = _reduce_ratio_metrics( + { + "acc": (torch.tensor(2.0), torch.tensor(4.0)), + "ce_position": ( + torch.tensor([1.0, 3.0]), + torch.tensor([2.0, 6.0]), + ), + }, + device=torch.device("cpu"), + process_group="dp", + reduce=True, + ) + + self.assertEqual(metrics["acc"], 1.0) + self.assertEqual(metrics["ce_position_0"], 1.0) + self.assertEqual(metrics["ce_position_1"], 0.5) + @staticmethod def _eagle_output( *, diff --git a/tests/test_utils/test_chunking.py b/tests/test_utils/test_chunking.py new file mode 100644 index 000000000..70b6cd58e --- /dev/null +++ b/tests/test_utils/test_chunking.py @@ -0,0 +1,67 @@ +# coding=utf-8 +"""Tests for reusable memory-bounded objective reductions.""" + +import unittest + +import torch + +from specforge.core.chunking import checkpointed_chunk_reduce + + +class CheckpointedChunkReduceTest(unittest.TestCase): + def test_sums_additive_terms_and_preserves_gradients(self): + values = torch.arange(12, dtype=torch.double).reshape(2, 6) + values.requires_grad_() + weights = torch.linspace(0.5, 1.5, 12, dtype=torch.double).reshape(2, 6) + + def terms(value_chunk, weight_chunk, optional): + self.assertIsNone(optional) + return ( + (value_chunk * weight_chunk).sum(), + (value_chunk.square() * weight_chunk).sum(), + ) + + linear, quadratic = checkpointed_chunk_reduce( + terms, + values, + weights, + None, + chunk_size=2, + dim=1, + ) + + torch.testing.assert_close(linear, (values * weights).sum()) + torch.testing.assert_close(quadratic, (values.square() * weights).sum()) + quadratic.backward() + torch.testing.assert_close(values.grad, 2 * values.detach() * weights) + + def test_zero_chunk_size_runs_one_full_slice(self): + calls = [] + values = torch.arange(5) + + def terms(chunk): + calls.append(tuple(chunk.shape)) + return (chunk.sum(),) + + (total,) = checkpointed_chunk_reduce( + terms, + values, + chunk_size=0, + ) + + self.assertEqual(calls, [(5,)]) + self.assertEqual(total.item(), 10) + + def test_rejects_misaligned_inputs(self): + with self.assertRaisesRegex(ValueError, "must be aligned"): + checkpointed_chunk_reduce( + lambda left, right: (left.sum() + right.sum(),), + torch.zeros(2, 3), + torch.zeros(2, 4), + chunk_size=1, + dim=1, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_utils/test_dflash_losses.py b/tests/test_utils/test_dflash_losses.py index a978e9aaa..3a0457044 100644 --- a/tests/test_utils/test_dflash_losses.py +++ b/tests/test_utils/test_dflash_losses.py @@ -7,6 +7,7 @@ output, and LM head output are made deterministic. """ +import copy import importlib.util import sys import types @@ -106,6 +107,26 @@ def predict_confidence(self, hidden_states, prev_token_ids=None): return self.confidence_logits.to(device=hidden_states.device) +class _LearnableDSparkDraft(_FixedDraft): + def __init__(self, hidden_size: int): + super().__init__(hidden_size) + self.signal = nn.Parameter(torch.linspace(-0.2, 0.2, hidden_size)) + self.confidence_head = nn.Identity() + + def forward(self, position_ids, noise_embedding, target_hidden, attention_mask): + del position_ids, target_hidden, attention_mask + bsz, draft_len = noise_embedding.shape[:2] + return self.signal.view(1, 1, -1).expand(bsz, draft_len, -1) + + def apply_logits_head(self, base_logits, prev_token_ids, hidden_states): + del prev_token_ids, hidden_states + return base_logits + + def predict_confidence(self, hidden_states, prev_token_ids=None): + del prev_token_ids + return hidden_states[..., 0] + + class _FixedHead(nn.Module): def __init__(self, logits: torch.Tensor): super().__init__() @@ -120,9 +141,12 @@ def __init__(self, draft_logits: torch.Tensor, target_logits: torch.Tensor): super().__init__() self.register_buffer("draft_logits", draft_logits) self.register_buffer("target_logits", target_logits) + self._calls = 0 def forward(self, hidden_states): - if hidden_states.ndim == 4: + use_target = self._calls % 2 == 1 + self._calls += 1 + if use_target: return self.target_logits.to(device=hidden_states.device) bsz, n_blocks, block_size, vocab_size = self.draft_logits.shape return self.draft_logits.reshape(bsz, n_blocks * block_size, vocab_size).to( @@ -394,6 +418,76 @@ def test_alpha_changes_dpace_loss(self): high_alpha = self._forward_loss(loss_type="dpace", dpace_alpha=0.9) self.assertNotAlmostEqual(low_alpha.item(), high_alpha.item(), places=8) + def test_dflash_family_chunking_matches_full_loss_metrics_and_gradients(self): + for loss_type in ( + "dflash", + "dpace", + "dpace-cumulative-confidence-only", + "dpace-continuation-value-only", + ): + with self.subTest(loss_type=loss_type): + torch.manual_seed(77) + head = nn.Linear(4, self.logits.shape[-1], bias=False).double() + full = _make_model( + self.logits, + self.anchors, + self.keep_mask, + draft_model=_LearnableDSparkDraft(4).double(), + lm_head=head, + loss_type=loss_type, + loss_decay_gamma=3.0, + objective_chunk_blocks=0, + ) + chunked = _make_model( + self.logits, + self.anchors, + self.keep_mask, + draft_model=_LearnableDSparkDraft(4).double(), + lm_head=copy.deepcopy(head), + loss_type=loss_type, + loss_decay_gamma=3.0, + objective_chunk_blocks=1, + ) + chunked.load_state_dict(full.state_dict()) + + full_loss, full_accuracy, full_metrics = full( + input_ids=self.input_ids, + hidden_states=self.hidden_states, + loss_mask=self.loss_mask, + ) + chunked_loss, chunked_accuracy, chunked_metrics = chunked( + input_ids=self.input_ids, + hidden_states=self.hidden_states, + loss_mask=self.loss_mask, + ) + + torch.testing.assert_close( + chunked_loss, + full_loss, + rtol=1e-6, + atol=1e-8, + ) + torch.testing.assert_close(chunked_accuracy, full_accuracy) + torch.testing.assert_close( + chunked_metrics["accuracy_denom"], + full_metrics["accuracy_denom"], + ) + + full_loss.backward() + chunked_loss.backward() + torch.testing.assert_close( + chunked.draft_model.signal.grad, + full.draft_model.signal.grad, + rtol=1e-6, + atol=1e-7, + ) + torch.testing.assert_close( + chunked.lm_head.weight.grad, + full.lm_head.weight.grad, + rtol=1e-6, + atol=1e-7, + ) + def test_invalid_loss_type_rejected(self): with self.assertRaisesRegex(ValueError, "loss_type"): _make_model( @@ -443,8 +537,12 @@ def test_dspark_ce_only_uses_next_token_labels_and_contiguous_mask(self): ) neg_log_q = _neg_log_q(self.logits, targets) weights = eval_mask.double() - want = (neg_log_q * weights).sum() / (weights.sum() + 1e-6) - torch.testing.assert_close(loss, want, rtol=0, atol=1e-6) + want = (neg_log_q * weights).sum() / weights.sum() + torch.testing.assert_close(loss, want, rtol=0, atol=1e-10) + acc_num, acc_den = _metrics["ratio_metrics"]["acc"] + torch.testing.assert_close(acc_den, weights.sum().float()) + torch.testing.assert_close(accuracy, acc_num / acc_den) + self.assertIn("ce_position", _metrics["ratio_metrics"]) def test_dspark_ce_only_skips_target_distribution(self): target_logits = torch.randn_like(self.logits) @@ -468,7 +566,7 @@ def test_dspark_ce_only_skips_target_distribution(self): self.assertTrue(torch.isfinite(loss)) - def test_dspark_l1_and_confidence_match_reference(self): + def test_dspark_l1_and_confidence_match_token_pooled_objective(self): torch.manual_seed(321) target_logits = torch.randn_like(self.logits) confidence_logits = torch.randn( @@ -504,33 +602,38 @@ def test_dspark_l1_and_confidence_match_reference(self): self.logits.shape[2], ) weights = eval_mask.double() - ce = (_neg_log_q(self.logits, targets) * weights).sum() / (weights.sum() + 1e-6) + denominator = weights.sum() + ce = (_neg_log_q(self.logits, targets) * weights).sum() / denominator draft_probs = torch.softmax(self.logits.float(), dim=-1) target_probs = torch.softmax(target_logits.float(), dim=-1) l1_dist = (draft_probs - target_probs).abs().sum(dim=-1).double() - l1 = (l1_dist * weights).sum() / (weights.sum() + 1e-6) + l1 = (l1_dist * weights).sum() / denominator accept_rate = (1.0 - 0.5 * l1_dist).clamp(0.0, 1.0) conf = F.binary_cross_entropy_with_logits( confidence_logits.float(), accept_rate.float(), reduction="none", ).double() - conf = (conf * weights).sum() / (weights.sum() + 1e-6) + conf = (conf * weights).sum() / denominator want = 0.1 * ce + 0.9 * l1 + conf torch.testing.assert_close(loss, want, rtol=0, atol=1e-6) - torch.testing.assert_close(metrics["ce_loss"], ce, rtol=0, atol=1e-6) + ratio_metrics = metrics["ratio_metrics"] + ce_num, ce_den = ratio_metrics["ce_loss"] + l1_num, l1_den = ratio_metrics["l1_loss"] + confidence_num, confidence_den = ratio_metrics["confidence_loss"] + torch.testing.assert_close(ce_num / ce_den, ce, rtol=0, atol=1e-6) torch.testing.assert_close( - metrics["l1_loss"], l1, rtol=0, atol=1e-6, check_dtype=False + l1_num / l1_den, l1, rtol=0, atol=1e-6, check_dtype=False ) torch.testing.assert_close( - metrics["confidence_loss"], + confidence_num / confidence_den, conf, rtol=0, atol=1e-6, check_dtype=False, ) - def test_dspark_requires_target_logits_for_l1_or_confidence(self): + def test_dspark_requires_target_hidden_states_for_l1_or_confidence(self): model = _make_dspark_model( self.logits, self.anchors, @@ -546,6 +649,130 @@ def test_dspark_requires_target_logits_for_l1_or_confidence(self): loss_mask=self.loss_mask, ) + def test_dspark_chunking_matches_full_loss_and_gradient(self): + torch.manual_seed(99) + head = nn.Linear(4, self.logits.shape[-1], bias=False).double() + full = _make_dspark_model( + self.logits, + self.anchors, + self.keep_mask, + draft_model=_LearnableDSparkDraft(4).double(), + lm_head=head, + dspark_ce_loss_alpha=0.1, + dspark_l1_loss_alpha=0.9, + dspark_confidence_head_alpha=1.0, + objective_chunk_blocks=0, + ) + chunked = _make_dspark_model( + self.logits, + self.anchors, + self.keep_mask, + draft_model=_LearnableDSparkDraft(4).double(), + lm_head=copy.deepcopy(head), + dspark_ce_loss_alpha=0.1, + dspark_l1_loss_alpha=0.9, + dspark_confidence_head_alpha=1.0, + objective_chunk_blocks=1, + ) + chunked.load_state_dict(full.state_dict()) + target_hidden = torch.randn_like(self.hidden_states) + + full_loss, _full_acc, full_metrics = full( + input_ids=self.input_ids, + hidden_states=self.hidden_states, + loss_mask=self.loss_mask, + target_last_hidden_states=target_hidden, + ) + chunked_loss, _chunked_acc, chunked_metrics = chunked( + input_ids=self.input_ids, + hidden_states=self.hidden_states, + loss_mask=self.loss_mask, + target_last_hidden_states=target_hidden, + ) + torch.testing.assert_close(chunked_loss, full_loss, rtol=1e-6, atol=1e-8) + for name, full_pair in full_metrics["ratio_metrics"].items(): + chunked_pair = chunked_metrics["ratio_metrics"][name] + torch.testing.assert_close(chunked_pair[0], full_pair[0]) + torch.testing.assert_close(chunked_pair[1], full_pair[1]) + + full_loss.backward() + chunked_loss.backward() + torch.testing.assert_close( + chunked.draft_model.signal.grad, + full.draft_model.signal.grad, + rtol=1e-6, + atol=1e-7, + ) + + def test_dspark_ce_only_chunking_recomputes(self): + torch.manual_seed(101) + head = nn.Linear(4, self.logits.shape[-1], bias=False).double() + full = _make_dspark_model( + self.logits, + self.anchors, + self.keep_mask, + draft_model=_LearnableDSparkDraft(4).double(), + lm_head=head, + dspark_ce_loss_alpha=1.0, + dspark_l1_loss_alpha=0.0, + dspark_confidence_head_alpha=0.0, + objective_chunk_blocks=0, + ) + chunked = _make_dspark_model( + self.logits, + self.anchors, + self.keep_mask, + draft_model=_LearnableDSparkDraft(4).double(), + lm_head=copy.deepcopy(head), + dspark_ce_loss_alpha=1.0, + dspark_l1_loss_alpha=0.0, + dspark_confidence_head_alpha=0.0, + objective_chunk_blocks=1, + ) + chunked.load_state_dict(full.state_dict()) + + full_loss, _full_acc, _full_metrics = full( + input_ids=self.input_ids, + hidden_states=self.hidden_states, + loss_mask=self.loss_mask, + ) + chunked_loss, _chunked_acc, _chunked_metrics = chunked( + input_ids=self.input_ids, + hidden_states=self.hidden_states, + loss_mask=self.loss_mask, + ) + torch.testing.assert_close(chunked_loss, full_loss, rtol=1e-6, atol=1e-8) + full_loss.backward() + chunked_loss.backward() + torch.testing.assert_close( + chunked.draft_model.signal.grad, + full.draft_model.signal.grad, + rtol=1e-6, + atol=1e-7, + ) + + def test_dspark_sampler_keeps_sparse_high_index_anchor(self): + model = _make_dspark_model( + self.logits, + self.anchors, + self.keep_mask, + ) + model.num_anchors = 2 + sparse_mask = torch.tensor( + [ + [0.0, 0.0, 0.0, 0.0, 1.0, 1.0], + [1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + ] + ) + anchors, keep = OnlineDSparkModel._sample_anchor_positions( + model, + seq_len=6, + loss_mask=sparse_mask, + device=sparse_mask.device, + ) + self.assertEqual(anchors[0, 0].item(), 4) + self.assertEqual(keep[0].tolist(), [True, False]) + if __name__ == "__main__": unittest.main() From 70b433ec5623e25c6a15276bbbdb8f6d84244d18 Mon Sep 17 00:00:00 2001 From: Yi Sun Date: Sun, 19 Jul 2026 23:56:43 +0000 Subject: [PATCH 2/4] refactor: flatten benchmark CLI and derive SGLang argv from config SGLang is the only supported benchmark target, so the `sglang` sub-subcommand added no discrimination. Move its arguments and help directly onto `specforge benchmark`; update the docs and dispatch test. Also replace the hand-maintained flag lists in `_managed_local_services` with a `_sglang_argv` helper that maps every `sglang_*` ModelConfig field to its `--foo-bar` flag, taking server-level and fallback values via overrides. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/benchmarks/benchmark.md | 6 +- specforge/cli.py | 36 +++---- specforge/launch_plan.py | 93 +++++++++---------- .../test_benchmarks/test_sglang_benchmark.py | 3 +- 4 files changed, 63 insertions(+), 75 deletions(-) diff --git a/docs/benchmarks/benchmark.md b/docs/benchmarks/benchmark.md index b4c8ec371..fe856b4ed 100644 --- a/docs/benchmarks/benchmark.md +++ b/docs/benchmarks/benchmark.md @@ -1,13 +1,13 @@ # Benchmarking inference serving Use `benchmarks/bench_eagle3.py` to compare EAGLE3 serving configurations and -dataset quality. Use `specforge benchmark sglang` to measure any existing +dataset quality. Use `specforge benchmark` to measure any existing SGLang deployment without assuming a particular speculative algorithm. | Runner | Server lifecycle | Measurements | Output | | --- | --- | --- | --- | | `python benchmarks/bench_eagle3.py` | Launches SGLang per configuration or uses an existing server | Latency, output throughput, acceptance length, and dataset accuracy when available | Timestamped JSON under `--output-dir` | -| `specforge benchmark sglang` | Uses an existing SGLang server | Aggregate output throughput, acceptance length, and verification count when reported by SGLang | Console and optional `--output-json` | +| `specforge benchmark` | Uses an existing SGLang server | Aggregate output throughput, acceptance length, and verification count when reported by SGLang | Console and optional `--output-json` | ## EAGLE3 benchmark matrix @@ -51,7 +51,7 @@ The existing-server runner supports target-only and speculative deployments on GSM8K, MATH-500, HumanEval, MBPP, and MT-Bench. ```bash -specforge benchmark sglang \ +specforge benchmark \ --model Qwen/Qwen3-8B \ --dataset gsm8k \ --base-url http://127.0.0.1:30000 \ diff --git a/specforge/cli.py b/specforge/cli.py index a122a7834..bd057dd53 100644 --- a/specforge/cli.py +++ b/specforge/cli.py @@ -212,38 +212,30 @@ def main(argv: Optional[List[str]] = None) -> int: export.add_argument("--embedding-key", default="model.embed_tokens.weight") benchmark = sub.add_parser( "benchmark", - help="benchmark a running inference server", - ) - benchmark_subcommands = benchmark.add_subparsers( - dest="benchmark_command", - required=True, - ) - sglang_benchmark = benchmark_subcommands.add_parser( - "sglang", help="benchmark a running SGLang server", description=( "Measure throughput and optional speculative-decoding telemetry from " "a running SGLang server." ), ) - sglang_benchmark.add_argument("--model", required=True) - sglang_benchmark.add_argument( + benchmark.add_argument("--model", required=True) + benchmark.add_argument( "--dataset", choices=("gsm8k", "math500", "humaneval", "mbpp", "mt-bench"), required=True, ) - sglang_benchmark.add_argument("--max-new-tokens", type=int, default=2048) - sglang_benchmark.add_argument("--temperature", type=float, default=0.0) - sglang_benchmark.add_argument("--top-p", type=float, default=1.0) - sglang_benchmark.add_argument("--top-k", type=int, default=1) - sglang_benchmark.add_argument("--max-samples", type=int) - sglang_benchmark.add_argument("--num-prompts", type=int, default=1024) - sglang_benchmark.add_argument("--concurrency", type=int, default=1) - sglang_benchmark.add_argument("--base-url", default="http://127.0.0.1:30000") - sglang_benchmark.add_argument("--timeout-seconds", type=int, default=3600) - sglang_benchmark.add_argument("--enable-thinking", action="store_true") - sglang_benchmark.add_argument("--trust-remote-code", action="store_true") - sglang_benchmark.add_argument("--output-json") + benchmark.add_argument("--max-new-tokens", type=int, default=2048) + benchmark.add_argument("--temperature", type=float, default=0.0) + benchmark.add_argument("--top-p", type=float, default=1.0) + benchmark.add_argument("--top-k", type=int, default=1) + benchmark.add_argument("--max-samples", type=int) + benchmark.add_argument("--num-prompts", type=int, default=1024) + benchmark.add_argument("--concurrency", type=int, default=1) + benchmark.add_argument("--base-url", default="http://127.0.0.1:30000") + benchmark.add_argument("--timeout-seconds", type=int, default=3600) + benchmark.add_argument("--enable-thinking", action="store_true") + benchmark.add_argument("--trust-remote-code", action="store_true") + benchmark.add_argument("--output-json") args = parser.parse_args(argv) if args.command == "train": diff --git a/specforge/launch_plan.py b/specforge/launch_plan.py index 299714151..eace81056 100644 --- a/specforge/launch_plan.py +++ b/specforge/launch_plan.py @@ -19,7 +19,7 @@ from urllib import request as urllib_request from urllib.parse import urlsplit, urlunsplit -from specforge.config import SGLANG_CAPTURE_CONTEXT_HEADROOM, Config +from specforge.config import SGLANG_CAPTURE_CONTEXT_HEADROOM, Config, ModelConfig if TYPE_CHECKING: from specforge.algorithms.registry import AlgorithmRegistration @@ -359,6 +359,33 @@ def _managed_local_environment(cfg: Config) -> dict[str, str]: return values +def _sglang_argv( + model: ModelConfig, + *, + overrides: Optional[Mapping[str, object]] = None, +) -> list[str]: + """Derive SGLang CLI args from all ``sglang_*`` fields on *model*. + + Flag names follow the naming convention ``sglang_foo_bar`` -> ``--foo-bar``. + Fields whose values require non-trivial resolution (server-level overrides, + fallback computations) are passed via *overrides* keyed by the original + field name; every other ``sglang_*`` field is read directly from *model*. + """ + resolved = overrides or {} + argv: list[str] = [] + for name in ModelConfig.model_fields: + if not name.startswith("sglang_"): + continue + value = resolved[name] if name in resolved else getattr(model, name) + flag = "--" + name.removeprefix("sglang_").replace("_", "-") + if isinstance(value, bool): + if value: + argv.append(flag) + elif value is not None: + argv.extend((flag, str(value))) + return argv + + def _managed_local_services( cfg: Config, *, @@ -422,16 +449,6 @@ def _managed_local_services( "--skip-tokenizer-init", "--tp-size", str(server.tp_size), - "--context-length", - str(capture_context_length), - "--mem-fraction-static", - str( - server.mem_fraction_static - if server.mem_fraction_static is not None - else cfg.model.sglang_mem_fraction_static - ), - "--ep-size", - str(cfg.model.sglang_ep_size), "--chunked-prefill-size", "-1", "--enable-spec-capture", @@ -443,45 +460,25 @@ def _managed_local_services( "127.0.0.1", "--port", str(server.port), - "--attention-backend", - server.attention_backend or cfg.model.sglang_attention_backend, ] ) - for enabled, flag in ( - (cfg.model.sglang_enable_nccl_nvls, "--enable-nccl-nvls"), - (cfg.model.sglang_enable_symm_mem, "--enable-symm-mem"), - (cfg.model.sglang_enable_torch_compile, "--enable-torch-compile"), - ): - if enabled: - argv.append(flag) - for value, flag in ( - (cfg.model.sglang_max_running_requests, "--max-running-requests"), - (cfg.model.sglang_max_total_tokens, "--max-total-tokens"), - (cfg.model.sglang_dp_size, "--dp-size"), - (cfg.model.sglang_moe_a2a_backend, "--moe-a2a-backend"), - (cfg.model.sglang_moe_runner_backend, "--moe-runner-backend"), - (cfg.model.sglang_page_size, "--page-size"), - (cfg.model.sglang_quantization, "--quantization"), - ( - cfg.model.sglang_fp4_gemm_runner_backend, - "--fp4-gemm-runner-backend", - ), - ( - cfg.model.sglang_mamba_radix_cache_strategy, - "--mamba-radix-cache-strategy", - ), - (cfg.model.sglang_max_mamba_cache_size, "--max-mamba-cache-size"), - ( - cfg.model.sglang_swa_full_tokens_ratio, - "--swa-full-tokens-ratio", - ), - ( - cfg.model.sglang_mamba_full_memory_ratio, - "--mamba-full-memory-ratio", - ), - ): - if value is not None: - argv.extend((flag, str(value))) + argv.extend( + _sglang_argv( + cfg.model, + overrides={ + "sglang_context_length": capture_context_length, + "sglang_mem_fraction_static": ( + server.mem_fraction_static + if server.mem_fraction_static is not None + else cfg.model.sglang_mem_fraction_static + ), + "sglang_attention_backend": ( + server.attention_backend + or cfg.model.sglang_attention_backend + ), + }, + ) + ) service_env = { **shared_env, "CUDA_VISIBLE_DEVICES": ",".join(server.cuda_visible_devices), diff --git a/tests/test_benchmarks/test_sglang_benchmark.py b/tests/test_benchmarks/test_sglang_benchmark.py index a924690fd..d22ef109c 100644 --- a/tests/test_benchmarks/test_sglang_benchmark.py +++ b/tests/test_benchmarks/test_sglang_benchmark.py @@ -71,7 +71,6 @@ def test_shared_cli_dispatches_sglang_benchmark(self): status = main( [ "benchmark", - "sglang", "--model", "thinkingmachines/Inkling", "--dataset", @@ -87,7 +86,7 @@ def test_cli_help_describes_the_backend_not_an_algorithm(self): output = StringIO() with redirect_stdout(output), self.assertRaises(SystemExit) as exited: - main(["benchmark", "sglang", "--help"]) + main(["benchmark", "--help"]) self.assertEqual(exited.exception.code, 0) help_text = " ".join(output.getvalue().split()) From 213aae253acf8154e21bf94006aaa2a7d508f9fd Mon Sep 17 00:00:00 2001 From: wjp666666 <1969554248@qq.com> Date: Tue, 21 Jul 2026 12:46:06 +0000 Subject: [PATCH 3/4] fix unittest fix unittest and lint --- specforge/algorithms/common/dflash_family_model.py | 3 ++- specforge/launch_plan.py | 3 +-- tests/test_config/test_launch_topology.py | 2 +- tests/test_config/test_unified_feature_reachability.py | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/specforge/algorithms/common/dflash_family_model.py b/specforge/algorithms/common/dflash_family_model.py index 48771e224..85b530641 100644 --- a/specforge/algorithms/common/dflash_family_model.py +++ b/specforge/algorithms/common/dflash_family_model.py @@ -1074,7 +1074,8 @@ def _compute_dspark_loss( if dist.is_available() and dist.is_initialized(): world_size = dist.get_world_size() - dist.all_reduce(global_loss_den, op=dist.ReduceOp.SUM) + if world_size > 1: + dist.all_reduce(global_loss_den, op=dist.ReduceOp.SUM) if float(global_loss_den) <= 0: raise ValueError("DSpark objective has no supervised target tokens") loss = ( diff --git a/specforge/launch_plan.py b/specforge/launch_plan.py index eace81056..8fbacff78 100644 --- a/specforge/launch_plan.py +++ b/specforge/launch_plan.py @@ -473,8 +473,7 @@ def _managed_local_services( else cfg.model.sglang_mem_fraction_static ), "sglang_attention_backend": ( - server.attention_backend - or cfg.model.sglang_attention_backend + server.attention_backend or cfg.model.sglang_attention_backend ), }, ) diff --git a/tests/test_config/test_launch_topology.py b/tests/test_config/test_launch_topology.py index b9636c0f5..b47c64b77 100644 --- a/tests/test_config/test_launch_topology.py +++ b/tests/test_config/test_launch_topology.py @@ -291,7 +291,7 @@ def _recipes() -> dict[str, Path]: class ExampleLaunchTopologyTest(unittest.TestCase): def test_every_recipe_has_the_explicit_golden_topology(self): recipes = _recipes() - self.assertEqual(len(EXPECTED_NPROC_PER_NODE), 59) + self.assertEqual(len(EXPECTED_NPROC_PER_NODE), 62) self.assertEqual(set(recipes), set(EXPECTED_NPROC_PER_NODE)) for filename, nproc_per_node in EXPECTED_NPROC_PER_NODE.items(): diff --git a/tests/test_config/test_unified_feature_reachability.py b/tests/test_config/test_unified_feature_reachability.py index 26e2cd60a..a4dd1a1d9 100644 --- a/tests/test_config/test_unified_feature_reachability.py +++ b/tests/test_config/test_unified_feature_reachability.py @@ -146,7 +146,7 @@ def test_all_example_configs_validate_through_the_typed_entry(self): for path in EXAMPLE_CONFIG_DIR.glob("*.yaml") if not path.name.startswith(".") ) - self.assertEqual(len(paths), 59) + self.assertEqual(len(paths), 62) resolved_runs = { path.name: resolve_run(Config.from_file(str(path))) for path in paths From d73f9ae12646f85ed55ace5bef0c413458ea9f0a Mon Sep 17 00:00:00 2001 From: wjp666666 <1969554248@qq.com> Date: Tue, 21 Jul 2026 16:04:33 +0000 Subject: [PATCH 4/4] fix eos_token more than 1 --- specforge/training/assembly.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/specforge/training/assembly.py b/specforge/training/assembly.py index 57044f3a5..5fa2e3091 100644 --- a/specforge/training/assembly.py +++ b/specforge/training/assembly.py @@ -146,6 +146,9 @@ def _load_text_tokenizer(cfg: Config): tokenizer.pad_token_id = cfg.model.tokenizer_pad_token_id elif tokenizer.pad_token_id is None: fallback_id = tokenizer.eos_token_id + if isinstance(fallback_id, (list, tuple)): + fallback_id = fallback_id[0] if fallback_id else None + if fallback_id is None: fallback_id = tokenizer.unk_token_id if fallback_id is None: