diff --git a/configs/qwen3.5-4b-dta.json b/configs/qwen3.5-4b-dta.json new file mode 100644 index 000000000..24b7b39d4 --- /dev/null +++ b/configs/qwen3.5-4b-dta.json @@ -0,0 +1,49 @@ +{ + "architectures": [ + "DFlashDraftModel" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "auto_map": { + "AutoModel": "dflash.DFlashDraftModel" + }, + "block_size": 16, + "bos_token_id": 248043, + "dflash_config": { + "mask_token_id": 248070, + "target_layer_ids": [1, 8, 15, 22, 29], + "training_mode": "vp_drafter", + "prefix_weight_base": 0.9 + }, + "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, + "transformers_version": "4.57.1", + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 248320 + } diff --git a/examples/run_qwen3.5_4b_dta_online_npu.sh b/examples/run_qwen3.5_4b_dta_online_npu.sh new file mode 100755 index 000000000..a4d8fc6a7 --- /dev/null +++ b/examples/run_qwen3.5_4b_dta_online_npu.sh @@ -0,0 +1,48 @@ +#!/bin/bash +# VP-Drafter (D2SD) training for Qwen3.5-4B on Ascend NPU +# Backend: HF + SDPA attention + HCCL distributed +# +# Required environment variables (override before invoking the script): +# TARGET_MODEL_PATH Path to Qwen3.5-4B weights +# TRAIN_DATA_PATH Path to the training jsonl + +set -eu + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +ROOT_DIR=$(dirname "$SCRIPT_DIR") + +: "${TARGET_MODEL_PATH:?Set TARGET_MODEL_PATH to the Qwen3.5-4B weights directory}" +: "${TRAIN_DATA_PATH:?Set TRAIN_DATA_PATH to the training jsonl file}" + +NPU_DEVICES=${1:-0,1,2,3,4,5,6,7} +NUM_DEVICES=$(echo "$NPU_DEVICES" | tr ',' '\n' | wc -l) + +export ASCEND_RT_VISIBLE_DEVICES=$NPU_DEVICES +export PYTORCH_NPU_ALLOC_CONF=max_split_size_mb:32 + +torchrun \ + --standalone \ + --nproc_per_node "$NUM_DEVICES" \ + "$ROOT_DIR/scripts/train_dflash.py" \ + --target-model-path "$TARGET_MODEL_PATH" \ + --target-model-backend hf \ + --draft-config-path "$ROOT_DIR/configs/qwen3.5-4b-dta.json" \ + --train-data-path "$TRAIN_DATA_PATH" \ + --output-dir "$ROOT_DIR/outputs/qwen3.5-4b-dta-npu" \ + --num-epochs 6 \ + --batch-size 1 \ + --accumulation-steps 4 \ + --learning-rate 6e-4 \ + --warmup-ratio 0.04 \ + --max-grad-norm 1.0 \ + --max-length 1024 \ + --chat-template qwen3.5 \ + --attention-backend sdpa \ + --num-anchors 16 \ + --loss-decay-gamma 7.0 \ + --log-interval 50 \ + --save-interval 3000 \ + --report-to tensorboard \ + --block-size 16 \ + --embedding-key model.language_model.embed_tokens.weight \ + --trust-remote-code