Add eight GPU SWIFT and Megatron training scripts

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
+