Ant dev vs main

ant_pipeline/inference/run_serve_eagle3_qwen3.sh

@@ -0,0 +1,67 @@
+set -ex
+
+USE_EAGLE3=true # false or true
+ENGINE_TYPE=trt # trt or sglang
+
+# for sglang
+speculative_num_steps=5
+speculative_eagle_topk=8
+speculative_num_draft_tokens=4
+
+# for trt
+max_draft_len=4
+
+# base_model=/mnt/modelops/models/Qwen3-32B
+# eagle_model=/mnt/modelops/models/AngelSlim/Qwen3-32B_eagle3
+
+# base_model=/mnt/modelops/models/Qwen3-14B
+# eagle_model=/mnt/modelops/models/AngelSlim/Qwen3-14B_eagle3
+
+base_model=/mnt/modelops/models/Qwen3-30B-A3B/
+eagle_model=/mnt/modelops/train/eagle3/output/qwen3-30B-A3b-eagle3_nnodes_8/epoch_9
+
+
+if [ ${ENGINE_TYPE} == "sglang" ];then
+  command=(
+    python3 -m sglang.launch_server \
+        --model ${base_model}  \
+        --host 127.0.0.1 \
+        --port 9122 \
+        --mem-fraction 0.85 \
+        --cuda-graph-max-bs 64 \
+        --max-running-requests 64 \
+        --chunked-prefill-size 8192 \
+        --tp-size 8
+  )
+
+  if $USE_EAGLE3; then
+    command+=( 
+        --speculative-algorithm EAGLE3 \
+        --speculative-draft-model-path ${eagle_model} \
+        --speculative-num-steps ${speculative_num_steps} \
+        --speculative-eagle-topk ${speculative_eagle_topk} \
+        --speculative-num-draft-tokens ${speculative_num_draft_tokens} 
+    )
+  fi
+else
+  command=(
+      python3 /mnt/modelops/460695/opensource/Ant-TensorRT-LLM/tensorrt_llm/commands/trtllm-serve.py \
+      --model_path ${base_model} \
+      --port 9122 \
+      --host 127.0.0.1 \
+      --backend pytorch \
+      --max_batch_size 16 \
+      --max_num_tokens 8192 \
+  )
+
+  if $USE_EAGLE3; then
+    command+=(--spec_algo eagle3 --draft_model_path ${eagle_model} --max_draft_len ${max_draft_len})
+  fi
+
+fi
+
+"${command[@]}"
+
+
+
+

ant_pipeline/inference/run_sglang_benchmark.sh

@@ -0,0 +1,21 @@
+set -ex
+
+# pip install sglang[all]
+target_model_path=/mnt/modelops/models/Qwen3-30B-A3B/
+draft_model_path=/mnt/modelops/train/eagle3/output/qwen3-30B-A3b-eagle3_nnodes_8/epoch_9
+
+python3 -m sglang.launch_server \
+    --model $target_model_path  \
+    --speculative-algorithm EAGLE3 \
+    --speculative-draft-model-path $draft_model_path \
+    --speculative-num-steps 3 \
+    --speculative-eagle-topk 1 \
+    --speculative-num-draft-tokens 4 \
+    --mem-fraction-static 0.85 \
+    --cuda-graph-max-bs 32 \
+    --tp 2 \
+    --context-length 8192 \
+    --trust-remote-code \
+    --host 127.0.0.1 \
+    --port 9122 \
+    --dtype bfloat16

ant_pipeline/train/prepare_datasets.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+set -ex
+
+script_dir=$(cd "$(dirname "$0")" && pwd)
+
+# 数据集列表
+datasets=(
+    "ultrachat:/mnt/modelops/datasets/HuggingFaceH4/ultrachat_200k/"
+)
+
+# 循环处理每个数据集
+for entry in "${datasets[@]}"; do
+    dataset_name="${entry%%:*}"   # 冒号前的部分
+    data_path="${entry#*:}"       # 冒号后的部分
+
+    output_path="/mnt/modelops/datasets/specforge_postprocess_${dataset_name}"
+
+    rm -rf "${output_path}"
+    mkdir -p "${output_path}"
+
+    python3 "${script_dir}/../../scripts/prepare_data.py" \
+        --dataset "${dataset_name}" \
+        --output_path "${output_path}" \
+        --data-path "${data_path}"
+done

ant_pipeline/train/prepare_env.sh

@@ -0,0 +1,22 @@
+set -ex
+
+script_dir=$(cd "$(dirname "$0")" && pwd)
+
+pip install --upgrade pip setuptools wheel
+pip install --trusted-host pypi.org --trusted-host files.pythonhosted.org --use-pep517 ninja
+pip install jinja2==3.1.2
+
+pushd $script_dir/../../
+# pip install -r requirements.txt # 默认环境已装好 
+pip install -e .
+popd
+
+# 全量编译 安装 flash-attn
+# MAX_JOBS=$(nproc) pip install flash-attn --no-build-isolation --no-cache-dir --force-reinstall # 默认环境已装好 
+
+# 全量编译 安装 deepspeed
+# MAX_JOBS=$(nproc) \
+# TORCH_CUDA_ARCH_LIST="9.0" \
+# DS_BUILD_OP=1 \
+# TORCH_EXTENSIONS_DIR=~/.cache/torch_extensions \
+# pip install deepspeed --force-reinstall --global-option="build_ext" --global-option="-j192"

ant_pipeline/train/train_by_deepspeed.sh

@@ -0,0 +1,83 @@
+#!/bin/bash
+set -ex
+
+# 获取路径
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)
+ROOT_DIR=$(dirname "$SCRIPT_DIR")/../
+
+pushd ${ROOT_DIR}
+pip install -e .
+popd
+
+# 设置环境变量
+export TORCH_CUDA_ARCH_LIST="9.0"
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+export TORCH_EXTENSIONS_DIR=~/.cache/torch_extensions
+export MAX_JOBS=$(nproc)
+export DS_BUILD_FUSED_ADAM=1  #只编译FUSED_ADAM 优化器
+
+# 获取训练配置
+model_idx=${1}
+nnodes=${2}
+
+if [ $nnodes -eq 1 ]; then
+    MASTER_ADDR=localhost
+    MASTER_PORT=$(( ( RANDOM % 10000 )  + 20000 ))
+    RANK=0
+else 
+    export OMP_NUM_THREADS=8
+    export NCCL_NET=IB
+    export NCCL_IB_DISABLE=0
+    export NCCL_SOCKET_IFNAME=^docker0,lo,eth0
+    export NCCL_IB_GID_INDEX=3
+    export NCCL_DEBUG=0
+fi
+
+config_file=${SCRIPT_DIR}/train_configs.json
+config=$(jq -r ".\"$model_idx\"" "$config_file")
+
+# 读取参数
+target_model_path=$(jq -r '.target_model_path' <<< "$config")
+draft_model_config=$(jq -r '.draft_model_config' <<< "$config")
+batch_size=$(jq -r '.batch_size' <<< "$config")
+zero_stage=$(jq -r '.zero_stage' <<< "$config")
+
+# 设置相关路径参数
+train_name=$(basename "$draft_model_config" .json)
+save_path=/mnt/modelops/train/eagle3/
+log_path=/mnt/modelops/train/eagle3/logs
+output_path=/mnt/modelops/train/eagle3/output
+mkdir -p ${log_path}
+mkdir -p ${output_path}
+
+output_dir=${output_path}/${train_name}_nnodes_${nnodes}
+train_log_path=${log_path}/${train_name}_nnodes_${nnodes}_rank${RANK}_train.log
+nvidia_smi_log_path=${log_path}/${train_name}_nnodes_${nnodes}_rank${RANK}_nvidia_smi.log
+
+# 添加GPU 监控
+nvidia-smi --query-gpu=timestamp,index,utilization.gpu,memory.used,memory.total --format=csv -l 1 > ${nvidia_smi_log_path} &
+
+# 正式开启训练
+torchrun \
+    --nnodes=$nnodes \
+    --nproc_per_node=8 \
+    --node_rank=$RANK \
+    --master_addr=$MASTER_ADDR \
+    --master_port=$MASTER_PORT \
+    $ROOT_DIR/scripts/train_eagle3_online_deepspeed.py \
+    --target-model-path ${target_model_path} \
+    --draft-model-config ${ROOT_DIR}/configs/${draft_model_config} \
+    --train-data-path /mnt/modelops/datasets/specforge_postprocess_ultrachat/ultrachat.jsonl \
+    --output-dir $output_dir \
+    --num-epochs 10 \
+    --batch-size ${batch_size} \
+    --learning-rate 1e-4 \
+    --max-length 2048 \
+    --chat-template qwen \
+    --cache-dir ${save_path}/cache/ \
+    --embedding-key model.embed_tokens.weight \
+    --ttt-length 7 \
+    --zero-stage ${zero_stage} \
+    2>&1 | tee $train_log_path
+
+echo "Training completed successfully. Output saved to $output_dir"

ant_pipeline/train/train_configs.json

@@ -0,0 +1,65 @@
+{
+    "1": 
+    {
+      "target_model_path": "/mnt/modelops/models/Qwen3-32B",
+      "draft_model_config": "qwen3-32b-eagle3.json",
+      "batch_size": 1,
+      "zero_stage": 1
+    },
+    "2": 
+    {
+      "target_model_path": "/mnt/modelops/models/Qwen3-8B",
+      "draft_model_config": "qwen3-8b-eagle3.json",
+      "batch_size": 6,
+      "zero_stage": 1
+    },
+    "3": 
+    {
+      "target_model_path": "/mnt/modelops/models/Qwen3-14B",
+      "draft_model_config": "qwen3-14b-eagle3.json",
+      "batch_size": 4,
+      "zero_stage": 1
+    },
+    "4": 
+    {
+      "target_model_path": "/mnt/modelops/models/Qwen2.5-72B",
+      "draft_model_config": "qwen2.5-72b-eagle3.json",
+      "batch_size": 2,
+      "zero_stage": 3
+    },
+    "5": 
+    {
+      "target_model_path": "/mnt/modelops/models/Qwen3-30B-A3B",
+      "draft_model_config": "qwen3-30B-A3b-eagle3.json",
+      "batch_size": 1,
+      "zero_stage": 1
+    },
+    "6": 
+    {
+      "target_model_path": "/mnt/modelops/models/Qwen3-235B-A22B-Instruct-2507",
+      "draft_model_config": "qwen3-235B-A22B-eagle3.json",
+      "batch_size": 1,
+      "zero_stage": 3
+    },
+    "7": 
+    {
+      "target_model_path": "/mnt/modelops/487922/deepseek-ai__DeepSeek-V2-Lite-Chat",
+      "draft_model_config": "deepseek-v2-lite-chat-eagle3.json",
+      "batch_size": 4,
+      "zero_stage": 1
+    },
+    "8": 
+    {
+      "target_model_path": "/mnt/modelops/models/QwQ-32B",
+      "draft_model_config": "qwq-32b-eagle3.json",
+      "batch_size": 2,
+      "zero_stage": 1
+    },
+    "9": 
+    {
+      "target_model_path": "/mnt/modelops/models/Qwen3-4B/",
+      "draft_model_config": "qwen3-4b-eagle3.json",
+      "batch_size": 12,
+      "zero_stage": 1
+    }
+}

configs/deepseek-v2-lite-chat-eagle3.json

@@ -0,0 +1,39 @@
+{
+  "architectures": [
+    "LlamaForCausalLMEagle3"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 100000,
+  "eos_token_id": 100001,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 10944,
+  "max_position_embeddings": 163840,
+  "max_window_layers": 64,
+  "model_type": "llama",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 1,
+  "num_key_value_heads": 16,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": {
+    "beta_fast": 32.0,
+    "beta_slow": 1.0,
+    "factor": 40.0,
+    "mscale": 0.707,
+    "mscale_all_dim": 0.707,
+    "original_max_position_embeddings": 4096,
+    "type": "yarn"
+  },
+  "rope_theta": 10000,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.33.1",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 102400,
+  "draft_vocab_size": 32000
+}

configs/qwen2.5-72b-eagle3.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "LlamaForCausalLMEagle3"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "hidden_act": "silu",
+  "hidden_size": 8192,
+  "initializer_range": 0.02,
+  "intermediate_size": 29568,
+  "max_position_embeddings": 32768,
+  "max_window_layers": 70,
+  "model_type": "llama",
+  "num_attention_heads": 64,
+  "num_hidden_layers": 1,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000.0,
+  "sliding_window": 131072,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.43.1",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 152064,
+  "draft_vocab_size": 32000
+}