Skip to content

Commit d52b6ad

Browse files
jsw-zorrojsw-zorro
committed
feat(examples): add Qwen3-8B Megatron math RL recipe
Add examples/math/qwen3-8b-megatron-delta — the FSDP qwen3-8b-m2po-delta recipe with the trainer engine switched to the Megatron backend (backend: megatron, tensor_parallel_size: 4). Identical data, algorithm, and weight-transfer path, so it doubles as a clean FSDP-vs-Megatron A/B. End-to-end on 8x H200 (4 RaaS + 4 trainer TP=4, delta TCP): 100 training steps, weight_transfer/delta_sparsity ~0.92 (HF-space delta), and a rising task reward (first-20 mean 0.55 -> later windows ~0.61-0.63).
1 parent a48dada commit d52b6ad

7 files changed

Lines changed: 469 additions & 0 deletions

File tree

‎examples/math/qwen3-8b-megatron-delta/README.md‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
# Qwen3-8B Math RL — Megatron backend, delta TCP weight transfer
2+
3+
Same math RL recipe as [`qwen3-8b-m2po-delta`](../qwen3-8b-m2po-delta) (M2PO,
4+
DeepScaleR data, ctx 16k, lr 5e-6, sparse delta weight sync) but the trainer
5+
uses the **Megatron-LM backend** instead of FSDP. The only difference is the
6+
`trainer_base.engine` block:
7+
8+
```yaml
9+
engine:
10+
backend: megatron
11+
data_parallel_size: 1
12+
tensor_parallel_size: 4
13+
pipeline_parallel_size: 1
14+
```
15+
16+
This makes it a clean FSDP-vs-Megatron A/B: identical data, algorithm, and
17+
weight-transfer path, so reward curves should track each other.
18+
19+
## How weight sync works (Megatron)
20+
21+
The trainer reconstructs the global model from Megatron's TP/PP/EP/VPP
22+
layout into HuggingFace-named tensors (via `export_hf_named_params`,
23+
backed by mbridge) and streams them into the CPU transfer buffer. Because
24+
the buffer holds HF-layout bytes, the sparse **delta** is computed in HF
25+
space and the RaaS receive path is identical to FSDP. See
26+
[`docs/en/architecture/megatron-weight-sync.md`](../../../docs/en/architecture/megatron-weight-sync.md).
27+
28+
## GPU layout (8 GPUs, single node)
29+
30+
| Component | GPUs | Parallelism |
31+
|-----------|------|-------------|
32+
| RaaS (SGLang, model0) | 0,1,2,3 | DP=4 |
33+
| Trainer model0 (Megatron) | 4,5,6,7 | TP=4 |
34+
35+
## Run
36+
37+
```bash
38+
bash examples/math/qwen3-8b-megatron-delta/scripts/run_qwen3-8b-megatron-delta.sh
39+
```
40+
41+
Or launch the three components separately (terminals 1/2/3):
42+
43+
```bash
44+
bash examples/math/qwen3-8b-megatron-delta/scripts/1_astraflow.sh
45+
bash examples/math/qwen3-8b-megatron-delta/scripts/2_raas.sh
46+
bash examples/math/qwen3-8b-megatron-delta/scripts/3_trainer_model0.sh
47+
```
48+
49+
## Scaling to PP / MoE
50+
51+
For pipeline or expert parallelism (and MoE models), set the corresponding
52+
sizes in the `engine` block, e.g. `pipeline_parallel_size: 2` or
53+
`expert_parallel_size: 2`. The backend auto-selects Megatron when `pp>1` or
54+
`ep>1`. Ensure `data_parallel_size * tensor_parallel_size *
55+
pipeline_parallel_size` equals the number of trainer GPUs.

‎examples/math/qwen3-8b-megatron-delta/scripts/1_astraflow.sh‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
#!/bin/bash
2+
set -euo pipefail
3+
# [1/3] Launch AstraFlow HTTP service
4+
#
5+
# Usage (terminal 1):
6+
# bash examples/math/qwen3-8b-m2po-delta/scripts/1_astraflow.sh
7+
8+
export CUDA_VISIBLE_DEVICES=""
9+
10+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
11+
REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
12+
cd "${REPO_ROOT}"
13+
export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"
14+
15+
YAML_DIR="${SCRIPT_DIR}/yaml"
16+
export EXPERIMENT_CONFIG="${EXPERIMENT_CONFIG:-${YAML_DIR}/experiment.yaml}"
17+
source "${REPO_ROOT}/examples/_common/utils.sh"
18+
# Export EXP_NAME and TRIAL_NAME from the experiment YAML.
19+
astraflow_load_experiment_env
20+
21+
export ASTRAFLOW_HOST="${ASTRAFLOW_HOST:-0.0.0.0}"
22+
export ASTRAFLOW_PORT="${ASTRAFLOW_PORT:-8000}"
23+
24+
# NCCL / PYTORCH / WANDB tweaks + LOG_DIR. Defined in examples/_common/utils.sh.
25+
astraflow_setup_env
26+
27+
echo "=== AstraFlow HTTP Service ==="
28+
echo "Experiment config : ${EXPERIMENT_CONFIG}"
29+
echo "Port : ${ASTRAFLOW_PORT}"
30+
echo "==============================="
31+
32+
python3 -u -m astraflow \
33+
--config "${EXPERIMENT_CONFIG}" \
34+
--port "${ASTRAFLOW_PORT}" \
35+
--host "${ASTRAFLOW_HOST}" \
36+
2>&1 | tee "${LOG_DIR}/astraflow.log"

‎examples/math/qwen3-8b-megatron-delta/scripts/2_raas.sh‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
#!/bin/bash
2+
set -euo pipefail
3+
# [2/3] Launch RaaS inference server (SGLang + TCP receiver)
4+
#
5+
# Usage (terminal 2, after AstraFlow is ready):
6+
# bash examples/math/qwen3-8b-m2po-delta/scripts/2_raas.sh
7+
8+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
9+
REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
10+
cd "${REPO_ROOT}"
11+
export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"
12+
13+
YAML_DIR="${SCRIPT_DIR}/yaml"
14+
export EXPERIMENT_CONFIG="${EXPERIMENT_CONFIG:-${YAML_DIR}/experiment.yaml}"
15+
export RAAS_CONFIG="${RAAS_CONFIG:-${YAML_DIR}/raas.yaml}"
16+
source "${REPO_ROOT}/examples/_common/utils.sh"
17+
# Export EXP_NAME and TRIAL_NAME from the experiment YAML.
18+
astraflow_load_experiment_env
19+
20+
export CUDA_VISIBLE_DEVICES="${SERVICE_CUDA_VISIBLE_DEVICES:-0,1,2,3}"
21+
export RAAS_HOST="${RAAS_HOST:-0.0.0.0}"
22+
export RAAS_PORT="${RAAS_PORT:-19190}"
23+
export ASTRAFLOW_PORT="${ASTRAFLOW_PORT:-8000}"
24+
export ASTRAFLOW_URL="${ASTRAFLOW_URL:-http://127.0.0.1:${ASTRAFLOW_PORT}}"
25+
26+
# NCCL / PYTORCH / WANDB tweaks + LOG_DIR. Defined in examples/_common/utils.sh.
27+
astraflow_setup_env
28+
29+
echo "=== RaaS Inference Server (SGLang + TCP receiver) ==="
30+
echo "Experiment config : ${EXPERIMENT_CONFIG}"
31+
echo "RaaS config : ${RAAS_CONFIG}"
32+
echo "GPUs : ${CUDA_VISIBLE_DEVICES}"
33+
echo "Port : ${RAAS_PORT}"
34+
echo "AstraFlow URL : ${ASTRAFLOW_URL}"
35+
echo "======================================================="
36+
37+
python3 -u -m astraflow.raas.server \
38+
--host "${RAAS_HOST}" \
39+
--port "${RAAS_PORT}" \
40+
--config "${EXPERIMENT_CONFIG}" \
41+
--config "${RAAS_CONFIG}" \
42+
--engine-id "${ENGINE_ID:-default}" \
43+
--astraflow-url "${ASTRAFLOW_URL}" \
44+
2>&1 | tee "${LOG_DIR}/raas.log"

‎examples/math/qwen3-8b-megatron-delta/scripts/3_trainer_model0.sh‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
#!/bin/bash
2+
set -euo pipefail
3+
# [3/3] Launch Trainer for model0 (TCP, sender_agent on local_rank 0)
4+
#
5+
# Usage (terminal 3, after AstraFlow and RaaS are ready):
6+
# bash examples/math/qwen3-8b-m2po-delta/scripts/3_trainer_model0.sh
7+
8+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
9+
REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
10+
cd "${REPO_ROOT}"
11+
export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"
12+
13+
YAML_DIR="${SCRIPT_DIR}/yaml"
14+
export EXPERIMENT_CONFIG="${EXPERIMENT_CONFIG:-${YAML_DIR}/experiment.yaml}"
15+
source "${REPO_ROOT}/examples/_common/utils.sh"
16+
# Export EXP_NAME and TRIAL_NAME from the experiment YAML.
17+
astraflow_load_experiment_env
18+
19+
export CUDA_VISIBLE_DEVICES="${TRAINER_MODEL0_GPUS:-4,5,6,7}"
20+
TRAINER0_NPROC="$(echo "${CUDA_VISIBLE_DEVICES}" | awk -F',' '{print NF}')"
21+
22+
export RAAS_PORT="${RAAS_PORT:-19190}"
23+
export ASTRAFLOW_PORT="${ASTRAFLOW_PORT:-8000}"
24+
export ASTRAFLOW_URL="http://127.0.0.1:${ASTRAFLOW_PORT}"
25+
export ASTRAFLOW_RAAS_URL="http://127.0.0.1:${RAAS_PORT}"
26+
27+
# sender_agent (in trainer) listens on this HTTP port
28+
export WEIGHT_TRANSFER_HTTP_PORT="${WEIGHT_TRANSFER_HTTP_PORT_MODEL0:-19861}"
29+
30+
# NCCL / PYTORCH / WANDB tweaks + LOG_DIR. Defined in examples/_common/utils.sh.
31+
astraflow_setup_env
32+
33+
echo "=== Trainer model0 (TCP) ==="
34+
echo "Experiment config : ${EXPERIMENT_CONFIG}"
35+
echo "GPUs : ${CUDA_VISIBLE_DEVICES} (Megatron TP${TRAINER0_NPROC})"
36+
echo "AstraFlow : ${ASTRAFLOW_URL}"
37+
echo "RaaS : ${ASTRAFLOW_RAAS_URL}"
38+
echo "Sender HTTP : ${WEIGHT_TRANSFER_HTTP_PORT}"
39+
echo "WANDB mode : ${WANDB_MODE:-online}"
40+
echo "=========================================="
41+
42+
torchrun --nnodes 1 --nproc-per-node "${TRAINER0_NPROC}" \
43+
--master-addr "${MASTER_ADDR:-127.0.0.1}" --master-port "${MASTER_PORT_MODEL0:-29541}" \
44+
examples/launch_trainer.py \
45+
--config "${EXPERIMENT_CONFIG}" \
46+
--trainer trainer_model0 \
47+
"$@" 2>&1 | tee "${LOG_DIR}/trainer_model0.log"

‎examples/math/qwen3-8b-megatron-delta/scripts/run_qwen3-8b-megatron-delta.sh‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
#!/bin/bash
2+
set -euo pipefail
3+
# All-in-one launcher for AstraFlow v2 math training (Qwen3-8B, M2PO, Megatron TP4, TCP).
4+
#
5+
# Launches 3 processes:
6+
# 1. AstraFlow HTTP service (CPU-only)
7+
# 2. RaaS inference server (SGLang, SERVICE_CUDA_VISIBLE_DEVICES)
8+
# 3. Trainer model0 (math, TRAINER_MODEL0_GPUS)
9+
#
10+
# Usage:
11+
# bash examples/math/qwen3-8b-megatron-delta/scripts/run_qwen3-8b-megatron-delta.sh
12+
13+
# =============================================================================
14+
# Part 1: Load env and settings
15+
# =============================================================================
16+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
17+
REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
18+
cd "${REPO_ROOT}"
19+
export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}"
20+
21+
YAML_DIR="${SCRIPT_DIR}/yaml"
22+
export EXPERIMENT_CONFIG="${EXPERIMENT_CONFIG:-${YAML_DIR}/experiment.yaml}"
23+
export RAAS_CONFIG="${RAAS_CONFIG:-${YAML_DIR}/raas.yaml}"
24+
source "${REPO_ROOT}/examples/_common/utils.sh"
25+
# Export EXP_NAME and TRIAL_NAME from the experiment YAML.
26+
# Defined in examples/_common/utils.sh.
27+
astraflow_load_experiment_env
28+
29+
# =============================================================================
30+
# Part 2: Set up env
31+
# =============================================================================
32+
# GPU assignments (default: 4 GPUs for inference, 4 for training)
33+
export SERVICE_CUDA_VISIBLE_DEVICES="${SERVICE_CUDA_VISIBLE_DEVICES:-0,1,2,3}"
34+
export TRAINER_MODEL0_GPUS="${TRAINER_MODEL0_GPUS:-4,5,6,7}"
35+
# Ports / URLs (each component gets its own port)
36+
export RAAS_HOST="${RAAS_HOST:-0.0.0.0}"
37+
export RAAS_PORT="${RAAS_PORT:-19190}"
38+
export ASTRAFLOW_HOST="${ASTRAFLOW_HOST:-0.0.0.0}"
39+
export ASTRAFLOW_PORT="${ASTRAFLOW_PORT:-8000}"
40+
export ASTRAFLOW_URL="http://127.0.0.1:${ASTRAFLOW_PORT}"
41+
export WEIGHT_TRANSFER_HTTP_PORT_MODEL0="${WEIGHT_TRANSFER_HTTP_PORT_MODEL0:-19861}"
42+
43+
TRAINER0_NPROC="$(echo "${TRAINER_MODEL0_GPUS}" | awk -F',' '{print NF}')"
44+
45+
# NCCL / PYTORCH / WANDB tweaks + LOG_DIR.
46+
# Defined in examples/_common/utils.sh.
47+
astraflow_setup_env
48+
49+
# =============================================================================
50+
# Part 3: Print info and clean up
51+
# =============================================================================
52+
echo "=== AstraFlow v2 (Qwen3-8B, math, M2PO, ctx16k, TCP delta) ==="
53+
echo "Experiment config : ${EXPERIMENT_CONFIG}"
54+
echo "RaaS config : ${RAAS_CONFIG}"
55+
echo "RaaS GPUs : ${SERVICE_CUDA_VISIBLE_DEVICES}"
56+
echo "Trainer model0 GPUs : ${TRAINER_MODEL0_GPUS} (Megatron TP${TRAINER0_NPROC})"
57+
echo "RaaS port : ${RAAS_PORT}"
58+
echo "AstraFlow port : ${ASTRAFLOW_PORT}"
59+
echo "Sender HTTP model0 : ${WEIGHT_TRANSFER_HTTP_PORT_MODEL0}"
60+
echo "WANDB mode : ${WANDB_MODE:-online}"
61+
echo "=========================================================="
62+
63+
trap astraflow_cleanup_trap EXIT INT TERM
64+
65+
# Kill leftover processes and shared memory from prior runs.
66+
# Defined in examples/_common/utils.sh.
67+
astraflow_kill_stale
68+
69+
# =============================================================================
70+
# Part 4: Launch training
71+
# =============================================================================
72+
echo "[1/3] Starting AstraFlow HTTP service..."
73+
CUDA_VISIBLE_DEVICES="" \
74+
python3 -u -m astraflow \
75+
--config "${EXPERIMENT_CONFIG}" \
76+
--port "${ASTRAFLOW_PORT}" \
77+
--host "${ASTRAFLOW_HOST}" \
78+
2>&1 | tee "${LOG_DIR}/astraflow.log" &
79+
sleep 5
80+
81+
echo "[2/3] Starting RaaS inference server (SGLang + TCP receiver)..."
82+
CUDA_VISIBLE_DEVICES="${SERVICE_CUDA_VISIBLE_DEVICES}" \
83+
python3 -u -m astraflow.raas.server \
84+
--host "${RAAS_HOST}" \
85+
--port "${RAAS_PORT}" \
86+
--config "${EXPERIMENT_CONFIG}" \
87+
--config "${RAAS_CONFIG}" \
88+
--engine-id "${ENGINE_ID:-default}" \
89+
--astraflow-url "${ASTRAFLOW_URL}" \
90+
2>&1 | tee "${LOG_DIR}/raas.log" &
91+
sleep 15
92+
93+
export ASTRAFLOW_RAAS_URL="http://127.0.0.1:${RAAS_PORT}"
94+
95+
echo "[3/3] Starting trainer model0..."
96+
CUDA_VISIBLE_DEVICES="${TRAINER_MODEL0_GPUS}" \
97+
WEIGHT_TRANSFER_HTTP_PORT="${WEIGHT_TRANSFER_HTTP_PORT_MODEL0}" \
98+
torchrun --nnodes 1 --nproc-per-node "${TRAINER0_NPROC}" \
99+
--master-addr "${MASTER_ADDR:-127.0.0.1}" --master-port "${MASTER_PORT_MODEL0:-29541}" \
100+
examples/launch_trainer.py \
101+
--config "${EXPERIMENT_CONFIG}" \
102+
--trainer trainer_model0 \
103+
"$@" \
104+
2>&1 | tee "${LOG_DIR}/trainer_model0.log"

0 commit comments

Comments
 (0)