diff --git a/config/dspark/dspark_qwen3_27b.py b/config/dspark/dspark_qwen3_27b.py new file mode 100644 index 00000000..6a9a2891 --- /dev/null +++ b/config/dspark/dspark_qwen3_27b.py @@ -0,0 +1,76 @@ +import os + +from deepspec.trainer import Qwen3DSparkTrainer +from deepspec.utils.constant import BASE_CKPT_DIR, BASE_TB_DIR, QWEN_36_27B + + +project_name = "deepspec" +exp_name = "dspark_block7_qwen3_27b" +seed = 42 + +model = dict( + target_model_name_or_path=QWEN_36_27B, + block_size=7, + num_draft_layers=5, + target_layer_ids=[1, 16, 31, 46, 61], + mask_token_id=151669, + num_anchors=512, + + # markov head + markov_rank=256, + markov_head_type="vanilla", + + # confidence head + confidence_head_alpha=1.0, + confidence_head_with_markov=True, + + # loss + loss_decay_gamma=4.0, + ce_loss_alpha=0.1, + l1_loss_alpha=0.9, +) + +train = dict( + trainer_cls=Qwen3DSparkTrainer, + lr=6.0e-4, + warmup_ratio=0.04, + weight_decay=0.0, + precision="bf16", + local_batch_size=1, + global_batch_size=512, + num_train_epochs=10, + max_train_steps=None, + max_grad_norm=1.0, + sharding_strategy="no_shard", + torch_compile=True, +) + +logging = dict( + logging_steps=10, + checkpointing_steps=3000, +) + +data = dict( + target_cache_path=None, + chat_template="qwen", + max_length=4096, + num_workers=4, +) + + +def finalize_cfg(cfg): + logging_cfg = dict(cfg["logging"]) + project_name = str(cfg["project_name"]) + exp_name = str(cfg["exp_name"]) + logging_cfg["checkpoint_dir"] = os.path.join( + BASE_CKPT_DIR, + project_name, + exp_name, + ) + logging_cfg["tensorboard_dir"] = os.path.join( + BASE_TB_DIR, + project_name, + exp_name, + ) + cfg["logging"] = logging_cfg + return cfg diff --git a/deepspec/utils/constant/public.py b/deepspec/utils/constant/public.py index a57ea33a..66cafeaf 100644 --- a/deepspec/utils/constant/public.py +++ b/deepspec/utils/constant/public.py @@ -7,6 +7,7 @@ QWEN_3_4B = "Qwen/Qwen3-4B" QWEN_3_8B = "Qwen/Qwen3-8B" QWEN_3_14B = "Qwen/Qwen3-14B" +QWEN_36_27B = "Qwen/Qwen3.6-27B" GEMMA_4_12B = "google/gemma-4-12B-it" BASE_TB_DIR = os.path.expanduser("~/tensorboard") BASE_CKPT_DIR = os.path.expanduser("~/checkpoints") diff --git a/docs/train_dspark_qwen36_27b.md b/docs/train_dspark_qwen36_27b.md new file mode 100644 index 00000000..508dc729 --- /dev/null +++ b/docs/train_dspark_qwen36_27b.md @@ -0,0 +1,131 @@ +# Training a DSpark Drafter for Qwen3.6-27B + +The pipeline has 3 stages: **data preparation** → **target cache** → **training**. Below are the commands adapted for `Qwen/Qwen3.6-27B`. + +## Prerequisites + +```bash +pip install -r requirements.txt +pip install "sglang[all]" # for data generation server +``` + +## Step 1: Download and Split Prompt Data + +```bash +python scripts/data/download_and_split.py \ + --dataset-name mlabonne/open-perfectblend \ + --test-size 0.05 \ + --train-output-path train_datasets/perfectblend_train.jsonl \ + --test-output-dir eval_datasets \ + --skip-existing +``` + +## Step 2: Regenerate Assistant Answers with Qwen3.6-27B + +The 27B model needs tensor parallelism (TP). Each SGLang worker below uses 2 GPUs (`--tp 2`), so 8 GPUs gives 4 workers. + +### Terminal 1 — Launch SGLang servers + +```bash +#!/usr/bin/env bash +model_path=Qwen/Qwen3.6-27B +num_workers=4 +start_port=30000 +log_dir=logs/sglang_qwen36_27b + +mkdir -p "${log_dir}" + +for ((i = 0; i < num_workers; i++)); do + port=$((start_port + i)) + gpu_start=$((i * 2)) + gpu_end=$((gpu_start + 1)) + log_file="${log_dir}/worker_${i}_port_${port}.log" + + echo "Starting worker ${i} on GPUs ${gpu_start},${gpu_end} port=${port}" + CUDA_VISIBLE_DEVICES=${gpu_start},${gpu_end} sglang serve \ + --model-path "${model_path}" \ + --host 0.0.0.0 \ + --port "${port}" \ + --tp 2 \ + --dtype bfloat16 \ + --mem-fraction-static 0.9 \ + > "${log_file}" 2>&1 & +done +wait +``` + +### Terminal 2 — Regenerate data + +```bash +python scripts/data/generate_train_data.py \ + --model Qwen/Qwen3.6-27B \ + --server-address \ + 127.0.0.1:30000 \ + 127.0.0.1:30001 \ + 127.0.0.1:30002 \ + 127.0.0.1:30003 \ + --concurrency 16 \ + --temperature 0.7 \ + --top-p 0.8 \ + --top-k 20 \ + --min-p 0 \ + --max-tokens 4096 \ + --disable-thinking \ + --resume \ + --input-file-path train_datasets/perfectblend_train.jsonl \ + --output-file-path train_datasets/qwen36_27b/perfectblend_train_regen.jsonl +``` + +Stop SGLang before Step 3 if sharing GPUs. + +## Step 3: Prepare Target Cache + +This runs the 27B target model on all training data and saves intermediate hidden states. + +> **Storage warning:** With 5 target layers and hidden_size=3456 on the full dataset, the cache can be very large (hundreds of TB). Reduce the training set or use fewer `target_layer_ids` if storage is limited. + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python scripts/data/prepare_target_cache.py \ + --config config/dspark/dspark_qwen3_27b.py \ + --train-data-path train_datasets/qwen36_27b/perfectblend_train_regen.jsonl \ + --output-dir ${HOME}/.cache/deepspec/qwen36_27b_target_cache \ + --local-batch-size 8 +``` + +Use `--local-batch-size 4` or lower if you hit OOM — the 27B model is larger. + +## Step 4: Train the DSpark Drafter + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python train.py \ + --config config/dspark/dspark_qwen3_27b.py \ + --opts "data.target_cache_path=${HOME}/.cache/deepspec/qwen36_27b_target_cache" +``` + +For multi-node training, set `MASTER_ADDR`, `MASTER_PORT`, `RANK`, and `WORLD_SIZE` environment variables. You may also want to switch sharding: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python train.py \ + --config config/dspark/dspark_qwen3_27b.py \ + --opts "data.target_cache_path=${HOME}/.cache/deepspec/qwen36_27b_target_cache" \ + --opts "train.sharding_strategy=full_shard" +``` + +## Step 5: Evaluate + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 python eval.py \ + --target_name_or_path Qwen/Qwen3.6-27B \ + --draft_name_or_path ${HOME}/checkpoints/deepspec/dspark_block7_qwen3_27b/step_latest +``` + +## Key Differences from Smaller Qwen3 Models + +| Parameter | Qwen3-4B/8B | Qwen3-14B | Qwen3.6-27B | +|-----------|-------------|-----------|-------------| +| `num_hidden_layers` | 36 | 40 | **64** | +| `target_layer_ids` | [1,9,17,25,33] | [1,10,19,28,37] | **[1,16,31,46,61]** | +| SGLang TP | 1 | 1 | **2** (needs 2 GPUs per server) | +| `local_batch_size` (cache) | 16 | 16 | **8** (lower due to memory) | + +The draft model itself remains small (5 layers) regardless of target size — only the target cache preparation and serving require more GPU memory.