Add eight GPU SWIFT and Megatron training scripts

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=<optional-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 空闲后执行：

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=<hf-id>
 export QWEN35_9B_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`.

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}"
+}

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}"
 }

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
+