diff --git a/configs/qwen3.5-4b-dspark.json b/configs/qwen3.5-4b-dspark.json new file mode 100644 index 000000000..065741df9 --- /dev/null +++ b/configs/qwen3.5-4b-dspark.json @@ -0,0 +1,53 @@ +{ + "architectures": [ + "DSparkDraftModel" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": { + "AutoModel": "dspark.DSparkDraftModel" + }, + "block_size": 16, + "bos_token_id": 248043, + "dflash_config": { + "attention_mode": "gqa", + "mask_token_id": 248070, + "target_layer_ids": [1, 8, 15, 22, 29], + "projector_type": "dspark", + "markov_rank": 256, + "markov_head_type": "vanilla", + "confidence_head_alpha": 1.0, + "enable_confidence_head": true, + "confidence_head_with_markov": true + }, + "dtype": "bfloat16", + "eos_token_id": 248044, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 2560, + "initializer_range": 0.02, + "intermediate_size": 9728, + "layer_types": [ + "full_attention", + "full_attention", + "full_attention", + "full_attention", + "full_attention" + ], + "max_position_embeddings": 262144, + "max_window_layers": 5, + "model_type": "qwen3", + "num_attention_heads": 32, + "num_hidden_layers": 5, + "num_key_value_heads": 8, + "num_target_layers": 32, + "pad_token_id": 248044, + "rms_norm_eps": 1e-06, + "rope_scaling": null, + "rope_theta": 10000000, + "sliding_window": null, + "tie_word_embeddings": true, + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 248320 +} diff --git a/examples/configs/online/disaggregated/managed-local/qwen3-4b-dspark-disaggregated-npu.yaml b/examples/configs/online/disaggregated/managed-local/qwen3-4b-dspark-disaggregated-npu.yaml new file mode 100644 index 000000000..717db51c0 --- /dev/null +++ b/examples/configs/online/disaggregated/managed-local/qwen3-4b-dspark-disaggregated-npu.yaml @@ -0,0 +1,76 @@ +# Single-node managed full stack for Ascend NPU: one `specforge train` command +# owns Mooncake, one patched SGLang capture server (device 0), and a 14-rank +# trainer (devices 2-15). Device ordinals are injected via +# ASCEND_RT_VISIBLE_DEVICES on Ascend hosts. +# DSpark's L1 / confidence losses additionally stream the target's final +# hidden state (target_last_hidden_states) through the same capture sink. +# Qwen3-4B is a text target: no embedding_key override and the `qwen` chat +# template (mirrors the external qwen3-4b-dspark recipe). +model: + target_model_path: "Qwen/Qwen3-4B" + draft_model_config: "configs/qwen3-4b-dspark.json" + target_backend: sglang + torch_dtype: "bfloat16" + sglang_attention_backend: ascend +data: + train_data_path: "./cache/dataset/train_regen.jsonl" + max_length: 1024 + chat_template: "qwen" + build_dataset_num_proc: 32 + cache_dir: "./cache" +training: + strategy: "dspark" + num_epochs: 10 + batch_size: 1 + accumulation_steps: 4 + learning_rate: 0.0006 + warmup_ratio: 0.04 + max_grad_norm: 1 + attention_backend: "sdpa" + # Validated DSpark recipe on Ascend (colocated run); raise toward 512 when + # memory allows, mirroring the GPU disaggregated default. + num_anchors: 186 + loss_decay_gamma: 7 + dspark_ce_loss_alpha: 0.1 + dspark_l1_loss_alpha: 0.9 + dspark_confidence_head_alpha: 1.0 + save_interval: 3000 + log_interval: 50 + dist_timeout: 30 + seed: 42 +tracking: + report_to: "tensorboard" +run_id: "qwen3-4b-dspark-npu-managed" +output_dir: "./outputs/qwen3-4b-dspark-npu-managed" + +deployment: + mode: disaggregated + trainer: + nnodes: 1 + nproc_per_node: 14 + disaggregated: + control_dir: outputs/qwen3-4b-dspark-npu-managed/control + consumer_state_dir: outputs/qwen3-4b-dspark-npu-managed/consumer-state + backend: mooncake + managed_local: + # 16-card A3 layout: capture server owns device 0, trainer owns 2-15. + # On 8-card hosts use devices 1-7 with trainer.nproc_per_node: 7. + trainer_cuda_visible_devices: + ["2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "14", "15"] + mooncake: + # Host segment mounted by the capture-server sink for feature storage. + # DSpark additionally streams target_last_hidden_states (~+20% over + # DFlash); 16 GiB fits a 64GB-card host next to the target model. + global_segment_size_bytes: 17179869184 + # Ascend get_into does register_buffer -> transfer -> unregister per + # tensor; under capture/trainer startup contention this does not fit + # the upstream managed-local 500ms lease and fails with LEASE_EXPIRED + # (-707). Use Mooncake's stock 5000ms; the ~19.5s teardown drain + # window still covers it. If the local mooncake_master build predates + # the --default_kv_lease_ttl flag, set this to null instead (omits the + # flag, same 5000ms server default). + default_kv_lease_ttl_ms: 5000 + capture_servers: + - port: 30000 + cuda_visible_devices: ["0"] + tp_size: 1 diff --git a/examples/configs/online/disaggregated/managed-local/qwen3.5-4b-dspark-disaggregated-npu.yaml b/examples/configs/online/disaggregated/managed-local/qwen3.5-4b-dspark-disaggregated-npu.yaml new file mode 100644 index 000000000..000c77dfd --- /dev/null +++ b/examples/configs/online/disaggregated/managed-local/qwen3.5-4b-dspark-disaggregated-npu.yaml @@ -0,0 +1,76 @@ +# Single-node managed full stack for Ascend NPU: one `specforge train` command +# owns Mooncake, one patched SGLang capture server (device 0), and a 14-rank +# trainer (devices 2-15). Device ordinals are injected via +# ASCEND_RT_VISIBLE_DEVICES on Ascend hosts. +# DSpark's L1 / confidence losses additionally stream the target's final +# hidden state (target_last_hidden_states) through the same capture sink. +model: + target_model_path: "Qwen/Qwen3.5-4B" + draft_model_config: "configs/qwen3.5-4b-dspark.json" + target_backend: sglang + trust_remote_code: true + embedding_key: "model.language_model.embed_tokens.weight" + torch_dtype: "bfloat16" + sglang_attention_backend: ascend +data: + train_data_path: "./cache/dataset/train_regen.jsonl" + max_length: 1024 + chat_template: "qwen3.5" + build_dataset_num_proc: 32 + cache_dir: "./cache" +training: + strategy: "dspark" + num_epochs: 10 + batch_size: 1 + accumulation_steps: 4 + learning_rate: 0.0006 + warmup_ratio: 0.04 + max_grad_norm: 1 + attention_backend: "sdpa" + # Validated DSpark recipe on Ascend (colocated run); raise toward 512 when + # memory allows, mirroring the GPU disaggregated default. + num_anchors: 186 + loss_decay_gamma: 7 + dspark_ce_loss_alpha: 0.1 + dspark_l1_loss_alpha: 0.9 + dspark_confidence_head_alpha: 1.0 + save_interval: 3000 + log_interval: 50 + dist_timeout: 30 + seed: 42 +tracking: + report_to: "tensorboard" +run_id: "qwen3.5-4b-dspark-npu-managed" +output_dir: "./outputs/qwen3.5-4b-dspark-npu-managed" + +deployment: + mode: disaggregated + trainer: + nnodes: 1 + nproc_per_node: 14 + disaggregated: + control_dir: outputs/qwen3.5-4b-dspark-npu-managed/control + consumer_state_dir: outputs/qwen3.5-4b-dspark-npu-managed/consumer-state + backend: mooncake + managed_local: + # 16-card A3 layout: capture server owns device 0, trainer owns 2-15. + # On 8-card hosts use devices 1-7 with trainer.nproc_per_node: 7. + trainer_cuda_visible_devices: + ["2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "14", "15"] + mooncake: + # Host segment mounted by the capture-server sink for feature storage. + # DSpark additionally streams target_last_hidden_states (~+20% over + # DFlash); 16 GiB fits a 64GB-card host next to the target model. + global_segment_size_bytes: 17179869184 + # Ascend get_into does register_buffer -> transfer -> unregister per + # tensor; under capture/trainer startup contention this does not fit + # the upstream managed-local 500ms lease and fails with LEASE_EXPIRED + # (-707). Use Mooncake's stock 5000ms; the ~19.5s teardown drain + # window still covers it. If the local mooncake_master build predates + # the --default_kv_lease_ttl flag, set this to null instead (omits the + # flag, same 5000ms server default). + default_kv_lease_ttl_ms: 5000 + capture_servers: + - port: 30000 + cuda_visible_devices: ["0"] + tp_size: 1 diff --git a/tests/test_runtime/test_package_architecture.py b/tests/test_runtime/test_package_architecture.py index c775c1760..962368a0f 100644 --- a/tests/test_runtime/test_package_architecture.py +++ b/tests/test_runtime/test_package_architecture.py @@ -716,6 +716,7 @@ def test_dspark_configs_are_qwen3_gqa_only(self): "kimi-k3-dspark.json", "qwen3-4b-dspark.json", "qwen3-8b-dspark.json", + "qwen3.5-4b-dspark.json", "qwen3.6-27b-dspark.json", }, )