Add eight GPU SWIFT and Megatron training scripts

This commit is contained in:
Codex
2026-06-25 20:34:57 +08:00
parent 4f9e0a30a9
commit d2cafeb8a0
5 changed files with 200 additions and 4 deletions

View File

@@ -12,6 +12,7 @@
- `scripts/train_qwen35_9b_full.sh`: Qwen3.5-9B bf16 full SFT。 - `scripts/train_qwen35_9b_full.sh`: Qwen3.5-9B bf16 full SFT。
- `scripts/train_qwen36_27b_lora.sh`: Qwen3.6-27B rank=32 LoRA。 - `scripts/train_qwen36_27b_lora.sh`: Qwen3.6-27B rank=32 LoRA。
- `scripts/train_qwen36_27b_full.sh`: Qwen3.6-27B bf16 full SFT。 - `scripts/train_qwen36_27b_full.sh`: Qwen3.6-27B bf16 full SFT。
- `scripts/train_qwen36_27b_megatron_full.sh`: 基于 Megatron-SWIFT/MCore-Bridge 的 Qwen3.6-27B full SFT 入口。
- `scripts/run_all_experiments.sh`: 按 LoRA 9B -> full 9B -> LoRA 27B -> full 27B 的顺序执行完整实验。 - `scripts/run_all_experiments.sh`: 按 LoRA 9B -> full 9B -> LoRA 27B -> full 27B 的顺序执行完整实验。
- `runs/`: TensorBoard 日志目录。 - `runs/`: TensorBoard 日志目录。
- `logs/`: 训练 stdout/stderr 和实际命令记录。 - `logs/`: 训练 stdout/stderr 和实际命令记录。
@@ -165,6 +166,9 @@ export HF_TOKEN=<optional-token>
- full SFT: `learning_rate=1e-5` - full SFT: `learning_rate=1e-5`
- LoRA: `learning_rate=5e-5` - LoRA: `learning_rate=5e-5`
- 默认每卡训练 batch size 为 `1` - 默认每卡训练 batch size 为 `1`
- 默认 `NPROC_PER_NODE=8`
- 默认 `CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7`
- 默认 `deepspeed=zero2`
B300/g0049 当前是 8 张 B300每卡约 275GB 显存。由于本实验默认上下文长度是 `262144`,预设 batch size 保守取 `1`,防止 27B/full 或长样本直接 OOM。吞吐测试时可以从脚本外部调大建议先从 LoRA/9B 开始试。 B300/g0049 当前是 8 张 B300每卡约 275GB 显存。由于本实验默认上下文长度是 `262144`,预设 batch size 保守取 `1`,防止 27B/full 或长样本直接 OOM。吞吐测试时可以从脚本外部调大建议先从 LoRA/9B 开始试。
@@ -189,6 +193,9 @@ export QWEN36_27B_FULL_BF16_PER_DEVICE_BATCH_SIZE=1
# scheduler/warmup # scheduler/warmup
export WARMUP_RATIO=0.1 export WARMUP_RATIO=0.1
export LR_SCHEDULER_TYPE=cosine export LR_SCHEDULER_TYPE=cosine
# 短跑 throughput/debug不跑完整 epoch
export MAX_STEPS=10
``` ```
命令: 命令:
@@ -200,6 +207,44 @@ export LR_SCHEDULER_TYPE=cosine
./scripts/train_qwen36_27b_full.sh ./scripts/train_qwen36_27b_full.sh
``` ```
## Megatron-SWIFT
仓库也提供了基于 Megatron-SWIFT/MCore-Bridge 的 full SFT 入口,参考官方 quick-start 的 `megatron sft` 写法:
```bash
./scripts/train_qwen36_27b_megatron_full.sh
```
默认参数:
- `NPROC_PER_NODE=8`
- `CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7`
- `MEGATRON_MODEL=Qwen/Qwen3.6-27B`
- `TENSOR_MODEL_PARALLEL_SIZE=4`
- `PIPELINE_MODEL_PARALLEL_SIZE=1`
- `CONTEXT_PARALLEL_SIZE=1`
- `MICRO_BATCH_SIZE=1`
- `GLOBAL_BATCH_SIZE=8`
- `MAX_LENGTH=262144`
- `LR=1e-5`
- `MIN_LR=1e-6`
- `LR_WARMUP_FRACTION=0.1`
Megatron 多机/共享盘训练要求 dataset cache 一致。脚本默认把 `MODELSCOPE_CACHE`、Megatron 输出和日志放在共享路径:
```text
/mnt/beegfs/workspace/ti_coding_agent_probe/
```
如果换机器或换共享盘,需要覆盖:
```bash
export BEEGFS_ROOT=/mnt/beegfs/workspace/ti_coding_agent_probe
export MODELSCOPE_CACHE=$BEEGFS_ROOT/modelscope_cache
export MEGATRON_OUTPUT_ROOT=$BEEGFS_ROOT/megatron_outputs
export MEGATRON_LOG_ROOT=$BEEGFS_ROOT/megatron_logs
```
## 一键完整实验 ## 一键完整实验
确认 GPU 空闲后执行: 确认 GPU 空闲后执行:

View File

@@ -103,6 +103,7 @@ The default experiment uses:
- explicit cosine LR scheduler via `--lr_scheduler_type cosine` - explicit cosine LR scheduler via `--lr_scheduler_type cosine`
- `max_length=262144` - `max_length=262144`
- conservative per-device train batch size `1` - conservative per-device train batch size `1`
- default `NPROC_PER_NODE=8`, `CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7`, and `DEEPSPEED=zero2`
- checkpoint save every 1000 steps - checkpoint save every 1000 steps
- validation every 1000 steps - validation every 1000 steps
- TensorBoard logging under `runs/` - TensorBoard logging under `runs/`
@@ -116,6 +117,7 @@ export LORA_PER_DEVICE_BATCH_SIZE=2
export FULL_PER_DEVICE_BATCH_SIZE=1 export FULL_PER_DEVICE_BATCH_SIZE=1
export QWEN35_9B_LORA_R32_PER_DEVICE_BATCH_SIZE=2 export QWEN35_9B_LORA_R32_PER_DEVICE_BATCH_SIZE=2
export QWEN36_27B_FULL_BF16_PER_DEVICE_BATCH_SIZE=1 export QWEN36_27B_FULL_BF16_PER_DEVICE_BATCH_SIZE=1
export MAX_STEPS=10
``` ```
Run-specific variables have the highest precedence, then global `PER_DEVICE_BATCH_SIZE` / `GRAD_ACCUM_STEPS`, then train-type defaults, then the safe default of 1. Run-specific variables have the highest precedence, then global `PER_DEVICE_BATCH_SIZE` / `GRAD_ACCUM_STEPS`, then train-type defaults, then the safe default of 1.
@@ -128,3 +130,19 @@ export QWEN36_27B_MODEL_ID=<hf-id>
export QWEN35_9B_MODEL_PATH=<local-path> export QWEN35_9B_MODEL_PATH=<local-path>
export QWEN36_27B_MODEL_PATH=<local-path> export QWEN36_27B_MODEL_PATH=<local-path>
``` ```
## Megatron-SWIFT
Use `scripts/train_qwen36_27b_megatron_full.sh` for Megatron-SWIFT/MCore-Bridge full SFT. It follows the official `megatron sft` quick-start style and defaults to:
- `NPROC_PER_NODE=8`
- `CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7`
- `MEGATRON_MODEL=Qwen/Qwen3.6-27B`
- `TENSOR_MODEL_PARALLEL_SIZE=4`
- `MICRO_BATCH_SIZE=1`
- `GLOBAL_BATCH_SIZE=8`
- `MAX_LENGTH=262144`
- `LR=1e-5`
- `LR_WARMUP_FRACTION=0.1`
For multi-node or shared-disk runs, keep `MODELSCOPE_CACHE` on shared storage. The default is `/mnt/beegfs/workspace/ti_coding_agent_probe/modelscope_cache`.

116
scripts/megatron_train_common.sh Executable file
View File

@@ -0,0 +1,116 @@
#!/usr/bin/env bash
set -euo pipefail
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "${ROOT_DIR}"
export http_proxy="${http_proxy:-http://100.72.0.101:8888}"
export https_proxy="${https_proxy:-http://100.72.0.101:8888}"
export HTTP_PROXY="${HTTP_PROXY:-${http_proxy}}"
export HTTPS_PROXY="${HTTPS_PROXY:-${https_proxy}}"
export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}"
export TOKENIZERS_PARALLELISM="${TOKENIZERS_PARALLELISM:-false}"
export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
if [[ -f .venv/bin/activate ]]; then
source .venv/bin/activate
elif [[ "${DRY_RUN:-0}" != "1" ]]; then
echo "Missing .venv. Run ./scripts/setup_env.sh first." >&2
exit 2
fi
BEEGFS_ROOT="${BEEGFS_ROOT:-/mnt/beegfs/workspace/ti_coding_agent_probe}"
export MODELSCOPE_CACHE="${MODELSCOPE_CACHE:-${BEEGFS_ROOT}/modelscope_cache}"
MEGATRON_OUTPUT_ROOT="${MEGATRON_OUTPUT_ROOT:-${BEEGFS_ROOT}/megatron_outputs}"
MEGATRON_LOG_ROOT="${MEGATRON_LOG_ROOT:-${BEEGFS_ROOT}/megatron_logs}"
mkdir -p "${MODELSCOPE_CACHE}" "${MEGATRON_OUTPUT_ROOT}" "${MEGATRON_LOG_ROOT}" logs
TRAIN_DATASET="${TRAIN_DATASET:-${TRAIN_JSONL:-data/processed/training_probe/train.jsonl}}"
VAL_DATASET="${VAL_DATASET:-${VAL_JSONL:-data/processed/training_probe/validation.jsonl}}"
MEGATRON_MODEL="${MEGATRON_MODEL:-${QWEN36_27B_MODEL_ID:-Qwen/Qwen3.6-27B}}"
MEGATRON_RUN_NAME="${MEGATRON_RUN_NAME:-qwen36_27b_megatron_full}"
MEGATRON_OUTPUT_DIR="${MEGATRON_OUTPUT_DIR:-${MEGATRON_OUTPUT_ROOT}/${MEGATRON_RUN_NAME}}"
NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
TENSOR_MODEL_PARALLEL_SIZE="${TENSOR_MODEL_PARALLEL_SIZE:-4}"
PIPELINE_MODEL_PARALLEL_SIZE="${PIPELINE_MODEL_PARALLEL_SIZE:-1}"
CONTEXT_PARALLEL_SIZE="${CONTEXT_PARALLEL_SIZE:-1}"
MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-1}"
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-8}"
MAX_LENGTH="${MAX_LENGTH:-262144}"
LR="${LR:-1e-5}"
MIN_LR="${MIN_LR:-1e-6}"
LR_WARMUP_FRACTION="${LR_WARMUP_FRACTION:-0.1}"
NUM_EPOCHS="${NUM_EPOCHS:-1}"
SAVE_STEPS="${SAVE_STEPS:-1000}"
EVAL_STEPS="${EVAL_STEPS:-1000}"
DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS:-4}"
DATASET_NUM_PROC="${DATASET_NUM_PROC:-4}"
RECOMPUTE_GRANULARITY="${RECOMPUTE_GRANULARITY:-full}"
RECOMPUTE_METHOD="${RECOMPUTE_METHOD:-uniform}"
RECOMPUTE_NUM_LAYERS="${RECOMPUTE_NUM_LAYERS:-1}"
CROSS_ENTROPY_LOSS_FUSION="${CROSS_ENTROPY_LOSS_FUSION:-true}"
GRADIENT_ACCUMULATION_FUSION="${GRADIENT_ACCUMULATION_FUSION:-false}"
SAVE_SAFETENSORS="${SAVE_SAFETENSORS:-true}"
SEQUENCE_PARALLEL="${SEQUENCE_PARALLEL:-true}"
require_file() {
if [[ ! -f "$1" ]]; then
echo "Missing required file: $1" >&2
exit 2
fi
}
run_megatron_sft() {
require_file "${TRAIN_DATASET}"
require_file "${VAL_DATASET}"
mkdir -p "${MEGATRON_OUTPUT_DIR}"
local log_file="${MEGATRON_LOG_ROOT}/${MEGATRON_RUN_NAME}.log"
local cmd=(
megatron sft
--model "${MEGATRON_MODEL}"
--save_safetensors "${SAVE_SAFETENSORS}"
--dataset "${TRAIN_DATASET}"
--val_dataset "${VAL_DATASET}"
--tensor_model_parallel_size "${TENSOR_MODEL_PARALLEL_SIZE}"
--pipeline_model_parallel_size "${PIPELINE_MODEL_PARALLEL_SIZE}"
--context_parallel_size "${CONTEXT_PARALLEL_SIZE}"
--sequence_parallel "${SEQUENCE_PARALLEL}"
--micro_batch_size "${MICRO_BATCH_SIZE}"
--global_batch_size "${GLOBAL_BATCH_SIZE}"
--recompute_granularity "${RECOMPUTE_GRANULARITY}"
--recompute_method "${RECOMPUTE_METHOD}"
--recompute_num_layers "${RECOMPUTE_NUM_LAYERS}"
--finetune true
--cross_entropy_loss_fusion "${CROSS_ENTROPY_LOSS_FUSION}"
--gradient_accumulation_fusion "${GRADIENT_ACCUMULATION_FUSION}"
--lr "${LR}"
--lr_warmup_fraction "${LR_WARMUP_FRACTION}"
--min_lr "${MIN_LR}"
--num_train_epochs "${NUM_EPOCHS}"
--output_dir "${MEGATRON_OUTPUT_DIR}"
--save_steps "${SAVE_STEPS}"
--eval_steps "${EVAL_STEPS}"
--max_length "${MAX_LENGTH}"
--dataloader_num_workers "${DATALOADER_NUM_WORKERS}"
--dataset_num_proc "${DATASET_NUM_PROC}"
--no_save_optim true
--no_save_rng true
--model_author swift
--model_name ti-coding-agent
)
printf 'CUDA_VISIBLE_DEVICES=%q NPROC_PER_NODE=%q MODELSCOPE_CACHE=%q ' \
"${CUDA_VISIBLE_DEVICES}" "${NPROC_PER_NODE}" "${MODELSCOPE_CACHE}" | tee "${log_file}.cmd"
printf '%q ' "${cmd[@]}" | tee -a "${log_file}.cmd"
echo | tee -a "${log_file}.cmd"
if [[ "${DRY_RUN:-0}" == "1" ]]; then
return 0
fi
CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" \
NPROC_PER_NODE="${NPROC_PER_NODE}" \
MODELSCOPE_CACHE="${MODELSCOPE_CACHE}" \
"${cmd[@]}" 2>&1 | tee "${log_file}"
}

View File

@@ -10,6 +10,7 @@ export HTTP_PROXY="${HTTP_PROXY:-${http_proxy}}"
export HTTPS_PROXY="${HTTPS_PROXY:-${https_proxy}}" export HTTPS_PROXY="${HTTPS_PROXY:-${https_proxy}}"
export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}" export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}"
export TOKENIZERS_PARALLELISM="${TOKENIZERS_PARALLELISM:-false}" export TOKENIZERS_PARALLELISM="${TOKENIZERS_PARALLELISM:-false}"
export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
if [[ -f .venv/bin/activate ]]; then if [[ -f .venv/bin/activate ]]; then
source .venv/bin/activate source .venv/bin/activate
@@ -26,6 +27,9 @@ SAVE_STEPS="${SAVE_STEPS:-1000}"
EVAL_STEPS="${EVAL_STEPS:-1000}" EVAL_STEPS="${EVAL_STEPS:-1000}"
LOGGING_STEPS="${LOGGING_STEPS:-1}" LOGGING_STEPS="${LOGGING_STEPS:-1}"
NUM_EPOCHS="${NUM_EPOCHS:-1}" NUM_EPOCHS="${NUM_EPOCHS:-1}"
NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
DEEPSPEED="${DEEPSPEED:-zero2}"
WARMUP_RATIO="${WARMUP_RATIO:-0.1}" WARMUP_RATIO="${WARMUP_RATIO:-0.1}"
LR_SCHEDULER_TYPE="${LR_SCHEDULER_TYPE:-cosine}" LR_SCHEDULER_TYPE="${LR_SCHEDULER_TYPE:-cosine}"
LORA_RANK="${LORA_RANK:-32}" LORA_RANK="${LORA_RANK:-32}"
@@ -93,7 +97,7 @@ run_swift_train() {
--model "${model_path}" --model "${model_path}"
--dataset "${TRAIN_JSONL}" --dataset "${TRAIN_JSONL}"
--val_dataset "${VAL_JSONL}" --val_dataset "${VAL_JSONL}"
--train_type "${train_type}" --tuner_type "${train_type}"
--torch_dtype bfloat16 --torch_dtype bfloat16
--num_train_epochs "${NUM_EPOCHS}" --num_train_epochs "${NUM_EPOCHS}"
--per_device_train_batch_size "${per_device_batch_size}" --per_device_train_batch_size "${per_device_batch_size}"
@@ -111,16 +115,21 @@ run_swift_train() {
--output_dir "${output_dir}" --output_dir "${output_dir}"
--save_total_limit "${SAVE_TOTAL_LIMIT:-3}" --save_total_limit "${SAVE_TOTAL_LIMIT:-3}"
--dataloader_num_workers "${DATALOADER_NUM_WORKERS:-4}" --dataloader_num_workers "${DATALOADER_NUM_WORKERS:-4}"
--deepspeed "${DEEPSPEED}"
) )
if [[ "${train_type}" == "lora" ]]; then if [[ "${train_type}" == "lora" ]]; then
cmd+=(--lora_rank "${LORA_RANK}") cmd+=(--lora_rank "${LORA_RANK}")
fi fi
if [[ -n "${MAX_STEPS:-}" ]]; then
cmd+=(--max_steps "${MAX_STEPS}")
fi
printf '%q ' "${cmd[@]}" | tee "${log_file}.cmd" printf 'CUDA_VISIBLE_DEVICES=%q NPROC_PER_NODE=%q ' "${CUDA_VISIBLE_DEVICES}" "${NPROC_PER_NODE}" | tee "${log_file}.cmd"
echo printf '%q ' "${cmd[@]}" | tee -a "${log_file}.cmd"
echo | tee -a "${log_file}.cmd"
if [[ "${DRY_RUN:-0}" == "1" ]]; then if [[ "${DRY_RUN:-0}" == "1" ]]; then
return 0 return 0
fi fi
"${cmd[@]}" 2>&1 | tee "${log_file}" CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" NPROC_PER_NODE="${NPROC_PER_NODE}" "${cmd[@]}" 2>&1 | tee "${log_file}"
} }

View File

@@ -0,0 +1,8 @@
#!/usr/bin/env bash
set -euo pipefail
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/megatron_train_common.sh"
MEGATRON_MODEL="${MEGATRON_MODEL:-${QWEN36_27B_MODEL_ID:-Qwen/Qwen3.6-27B}}"
MEGATRON_RUN_NAME="${MEGATRON_RUN_NAME:-qwen36_27b_megatron_full}"
run_megatron_sft