Add Domino online training path

configs/qwen3-8b-domino.json

@@ -0,0 +1,50 @@
+{
+  "architectures": [
+    "DFlashDraftModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "dflash.DFlashDraftModel"
+  },
+  "block_size": 16,
+  "bos_token_id": 151643,
+  "dflash_config": {
+    "mask_token_id": 151669,
+    "target_layer_ids": [1, 9, 17, 25, 33],
+    "projector_type": "domino",
+    "pure_draft_prefix_len": 1,
+    "emb_dim": 256,
+    "gru_hidden_dim": 1024,
+    "shift_label": true
+  },
+  "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
+}

examples/run_qwen3_8b_domino_online.sh

@@ -0,0 +1,47 @@
+#!/bin/bash
+set -euo pipefail
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+ROOT_DIR=$(dirname "$SCRIPT_DIR")
+
+export CUDA_HOME=/usr/local/cuda
+export CONDA_PREFIX=${CONDA_PREFIX:-$HOME/.conda/envs/specforge}
+
+export PATH=$CUDA_HOME/bin:$CONDA_PREFIX/bin:$PATH
+export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$CONDA_PREFIX/lib:$LD_LIBRARY_PATH
+
+export NCCL_DEBUG=WARN
+export SPECFORGE_DATA_NUM_PROC=32
+
+NUM_GPUS=${1:-8}
+ATTENTION_BACKEND=${2:-flex_attention}
+TARGET_MODEL_PATH=${TARGET_MODEL_PATH:-/path/to/Qwen3-8B}
+TRAIN_DATA_PATH=${TRAIN_DATA_PATH:-/path/to/sharegpt_train.jsonl}
+
+torchrun \
+    --standalone \
+    --nproc_per_node $NUM_GPUS \
+     $ROOT_DIR/scripts/train_domino.py \
+    --target-model-path $TARGET_MODEL_PATH \
+    --draft-config-path $ROOT_DIR/configs/qwen3-8b-domino.json \
+    --train-data-path $TRAIN_DATA_PATH \
+    --output-dir $ROOT_DIR/outputs/same_data/qwen3-8b-domino_sharegpt \
+    --num-epochs 6 \
+    --batch-size 2 \
+    --learning-rate 6e-4 \
+    --warmup-ratio 0.04 \
+    --max-grad-norm 1.0 \
+    --max-length 3072 \
+    --chat-template qwen \
+    --attention-backend $ATTENTION_BACKEND \
+    --num-anchors 256 \
+    --loss-decay-gamma 7.0 \
+    --log-interval 50 \
+    --save-interval 2000 \
+    --report-to wandb \
+    --wandb-project specforge-qwen3-8b-domino \
+    --target-model-backend sglang \
+    --block-size 16 \
+    --lambda-base-start 1.0 \
+    --lambda-base-decay-ratio 1.0 \
+    --wandb-name qwen3-8b-domino_sharegpt