From d2cafeb8a0062368b00d088c29dfb4f5749b5e2c Mon Sep 17 00:00:00 2001 From: Codex Date: Thu, 25 Jun 2026 20:34:57 +0800 Subject: [PATCH] Add eight GPU SWIFT and Megatron training scripts --- README.md | 45 +++++++++ SKILL.md | 18 ++++ scripts/megatron_train_common.sh | 116 ++++++++++++++++++++++ scripts/swift_train_common.sh | 17 +++- scripts/train_qwen36_27b_megatron_full.sh | 8 ++ 5 files changed, 200 insertions(+), 4 deletions(-) create mode 100755 scripts/megatron_train_common.sh create mode 100755 scripts/train_qwen36_27b_megatron_full.sh diff --git a/README.md b/README.md index bd0bec2..f3e5ca0 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,7 @@ - `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_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 的顺序执行完整实验。 - `runs/`: TensorBoard 日志目录。 - `logs/`: 训练 stdout/stderr 和实际命令记录。 @@ -165,6 +166,9 @@ export HF_TOKEN= - full SFT: `learning_rate=1e-5` - LoRA: `learning_rate=5e-5` - 默认每卡训练 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 开始试。 @@ -189,6 +193,9 @@ export QWEN36_27B_FULL_BF16_PER_DEVICE_BATCH_SIZE=1 # scheduler/warmup export WARMUP_RATIO=0.1 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 ``` +## 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 空闲后执行： diff --git a/SKILL.md b/SKILL.md index f28605d..a0df865 100644 --- a/SKILL.md +++ b/SKILL.md @@ -103,6 +103,7 @@ The default experiment uses: - explicit cosine LR scheduler via `--lr_scheduler_type cosine` - `max_length=262144` - 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 - validation every 1000 steps - TensorBoard logging under `runs/` @@ -116,6 +117,7 @@ export LORA_PER_DEVICE_BATCH_SIZE=2 export FULL_PER_DEVICE_BATCH_SIZE=1 export QWEN35_9B_LORA_R32_PER_DEVICE_BATCH_SIZE=2 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. @@ -128,3 +130,19 @@ export QWEN36_27B_MODEL_ID= export QWEN35_9B_MODEL_PATH= export QWEN36_27B_MODEL_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`. diff --git a/scripts/megatron_train_common.sh b/scripts/megatron_train_common.sh new file mode 100755 index 0000000..4d1fd51 --- /dev/null +++ b/scripts/megatron_train_common.sh @@ -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}" +} diff --git a/scripts/swift_train_common.sh b/scripts/swift_train_common.sh index cda66d4..e176825 100755 --- a/scripts/swift_train_common.sh +++ b/scripts/swift_train_common.sh @@ -10,6 +10,7 @@ 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 @@ -26,6 +27,9 @@ SAVE_STEPS="${SAVE_STEPS:-1000}" EVAL_STEPS="${EVAL_STEPS:-1000}" LOGGING_STEPS="${LOGGING_STEPS:-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}" LR_SCHEDULER_TYPE="${LR_SCHEDULER_TYPE:-cosine}" LORA_RANK="${LORA_RANK:-32}" @@ -93,7 +97,7 @@ run_swift_train() { --model "${model_path}" --dataset "${TRAIN_JSONL}" --val_dataset "${VAL_JSONL}" - --train_type "${train_type}" + --tuner_type "${train_type}" --torch_dtype bfloat16 --num_train_epochs "${NUM_EPOCHS}" --per_device_train_batch_size "${per_device_batch_size}" @@ -111,16 +115,21 @@ run_swift_train() { --output_dir "${output_dir}" --save_total_limit "${SAVE_TOTAL_LIMIT:-3}" --dataloader_num_workers "${DATALOADER_NUM_WORKERS:-4}" + --deepspeed "${DEEPSPEED}" ) if [[ "${train_type}" == "lora" ]]; then cmd+=(--lora_rank "${LORA_RANK}") fi + if [[ -n "${MAX_STEPS:-}" ]]; then + cmd+=(--max_steps "${MAX_STEPS}") + fi - printf '%q ' "${cmd[@]}" | tee "${log_file}.cmd" - echo + printf 'CUDA_VISIBLE_DEVICES=%q NPROC_PER_NODE=%q ' "${CUDA_VISIBLE_DEVICES}" "${NPROC_PER_NODE}" | 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 - "${cmd[@]}" 2>&1 | tee "${log_file}" + CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" NPROC_PER_NODE="${NPROC_PER_NODE}" "${cmd[@]}" 2>&1 | tee "${log_file}" } diff --git a/scripts/train_qwen36_27b_megatron_full.sh b/scripts/train_qwen36_27b_megatron_full.sh new file mode 100755 index 0000000..f864064 --- /dev/null +++ b/scripts/train_qwen36_27b_megatron_full.sh @@ -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 +