Skip to content

Commit cdc7b79

Browse files
author
lixianduo
committed
Add upgrade skills for Megatron-LM-FL upstream sync
1 parent 1df9eda commit cdc7b79

14 files changed

Lines changed: 3320 additions & 0 deletions

‎.claude/skills‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
../skills

‎skills/flagscale-train-upstream-sync/SKILL.md‎

Lines changed: 250 additions & 0 deletions
Large diffs are not rendered by default.

‎skills/flagscale-train-upstream-sync/phases/01-setup-and-branch.md‎

Lines changed: 305 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,305 @@
1+
# Phase 1: Setup & Branch Structure
2+
3+
Stages 1-2 of the FlagScale training upstream sync workflow.
4+
5+
---
6+
7+
### Stage 1: Environment & Repo Setup
8+
9+
Run the **Repo Detection Preamble** below to set `MG_FL_DIR` and `FLAGSCALE_DIR`. Both repos
10+
are required for this skill. If either is missing, the preamble offers to clone it.
11+
12+
#### Step 1: Locate both repos
13+
14+
```bash
15+
# --- Detect Megatron-LM-FL ---
16+
MG_FL_DIR=""
17+
for candidate in \
18+
"$(pwd)" \
19+
"$(pwd)/Megatron-LM-FL" \
20+
"$(pwd)/../Megatron-LM-FL" \
21+
"$(dirname $(pwd))/Megatron-LM-FL"; do
22+
if [ -d "$candidate/megatron" ] && [ -f "$candidate/setup.py" -o -f "$candidate/pyproject.toml" ]; then
23+
MG_FL_DIR=$(cd "$candidate" && pwd)
24+
break
25+
fi
26+
done
27+
28+
if [ -z "$MG_FL_DIR" ]; then
29+
echo "Megatron-LM-FL not found locally. Cloning..."
30+
git clone https://github.com/flagos-ai/Megatron-LM-FL.git
31+
MG_FL_DIR=$(cd Megatron-LM-FL && pwd)
32+
fi
33+
echo "Megatron-LM-FL: $MG_FL_DIR"
34+
35+
# --- Detect FlagScale ---
36+
# FlagScale must live as a sibling of Megatron-LM-FL in the same workspace directory.
37+
WORKSPACE_DIR="$(dirname "$MG_FL_DIR")"
38+
FLAGSCALE_DIR=""
39+
for candidate in \
40+
"$WORKSPACE_DIR/FlagScale" \
41+
"$WORKSPACE_DIR/flagscale"; do
42+
if [ -d "$candidate/flagscale" ] && [ "$(cd "$candidate" && pwd)" != "$MG_FL_DIR" ]; then
43+
FLAGSCALE_DIR=$(cd "$candidate" && pwd)
44+
break
45+
fi
46+
done
47+
48+
if [ -z "$FLAGSCALE_DIR" ]; then
49+
echo "FlagScale not found in $WORKSPACE_DIR. Cloning..."
50+
git clone https://github.com/flagos-ai/FlagScale.git "$WORKSPACE_DIR/FlagScale"
51+
FLAGSCALE_DIR="$WORKSPACE_DIR/FlagScale"
52+
fi
53+
echo "FlagScale: $FLAGSCALE_DIR"
54+
55+
# --- Verify Megatron-LM-FL is installed ---
56+
cd "$MG_FL_DIR"
57+
python3 -c "
58+
import megatron.core
59+
import megatron.plugin
60+
from megatron.plugin.platform import cur_platform
61+
print('megatron.core OK')
62+
print('megatron.plugin OK')
63+
print('cur_platform:', cur_platform)
64+
" || {
65+
echo "megatron.core/plugin not importable. Installing Megatron-LM-FL..."
66+
pip install -e . --no-build-isolation
67+
python3 -c "
68+
import megatron.core
69+
import megatron.plugin
70+
print('megatron.core OK')
71+
print('megatron.plugin OK')
72+
" || {
73+
echo "ERROR: megatron.core/plugin still not importable after install."
74+
echo "Check that pyproject.toml includes both megatron.core and megatron.plugin."
75+
echo "Fix Megatron-LM-FL first."
76+
exit 1
77+
}
78+
}
79+
80+
echo "Both repos ready."
81+
```
82+
83+
#### Step 2: Add upstream remote and fetch the target version
84+
85+
```bash
86+
cd "$FLAGSCALE_DIR"
87+
git remote -v | grep mg-upstream || \
88+
git remote add mg-upstream https://github.com/NVIDIA/Megatron-LM.git
89+
git fetch mg-upstream --tags
90+
```
91+
92+
#### Step 3: Confirm versions and branch names
93+
94+
The base and target versions must match the Megatron-LM-FL sync. Check Megatron-LM-FL's
95+
`SYNC_POINT.md` or branch history to identify them:
96+
97+
```bash
98+
cd "$MG_FL_DIR"
99+
cat SYNC_POINT.md 2>/dev/null || echo "No SYNC_POINT.md found"
100+
git log --oneline --all | grep -i "base\|core_v0" | head -10
101+
```
102+
103+
**Before any branch operations, ask the user for six parameters:**
104+
105+
1. **Base upstream version** — the upstream release the fork is currently based on
106+
2. **Base upstream commit** — (optional) specific commit SHA on the base version to use;
107+
if empty, use the tag/branch tip
108+
3. **Target upstream version** — the upstream release to sync to
109+
4. **Target upstream commit** — (optional) specific commit SHA on the target version to use;
110+
if empty, use the tag/branch tip
111+
5. **FlagScale dev branch name** — the FlagScale branch for the upgraded training code
112+
6. **Megatron-LM-FL dev branch name** — the Megatron-LM-FL dev branch (from the companion skill)
113+
114+
**Prompt the user:**
115+
> Please specify the following:
116+
> 1. Base upstream version the fork is currently based on (e.g. `core_v0.15.0rc7`)
117+
> 2. Base upstream commit (leave empty for tag/branch tip, e.g. `abc1234`)
118+
> 3. Target upstream version to sync to (e.g. `core_v0.16.1`)
119+
> 4. Target upstream commit (leave empty for tag/branch tip, e.g. `def5678`)
120+
> 5. FlagScale dev branch name (e.g. `dev-train-0.16`)
121+
> 6. Megatron-LM-FL dev branch name (e.g. `dev-0.16.1`)
122+
123+
Store the answers and use them throughout **all** subsequent stages:
124+
125+
```bash
126+
cd "$FLAGSCALE_DIR"
127+
BASE_VERSION="<user answer 1>" # e.g. core_v0.15.0rc7
128+
BASE_VERSION_COMMIT="<user answer 2>" # e.g. abc1234 (empty = tag/branch tip)
129+
TARGET_VERSION="<user answer 3>" # e.g. core_v0.16.1
130+
TARGET_VERSION_COMMIT="<user answer 4>" # e.g. def5678 (empty = tag/branch tip)
131+
FS_DEV_BRANCH="<user answer 5>" # e.g. dev-train-0.16
132+
MG_DEV_BRANCH="<user answer 6>" # e.g. dev-0.16.1
133+
134+
echo "Base version (old upstream): $BASE_VERSION"
135+
echo "Base commit: ${BASE_VERSION_COMMIT:-tip}"
136+
echo "Target version (new upstream): $TARGET_VERSION"
137+
echo "Target commit: ${TARGET_VERSION_COMMIT:-tip}"
138+
echo "FlagScale dev branch: $FS_DEV_BRANCH"
139+
echo "Megatron-LM-FL dev branch: $MG_DEV_BRANCH"
140+
141+
# Verify both tags/commits exist
142+
if [ -n "$BASE_VERSION_COMMIT" ]; then
143+
git rev-parse "$BASE_VERSION_COMMIT" > /dev/null 2>&1 || {
144+
echo "ERROR: commit $BASE_VERSION_COMMIT not found"; exit 1
145+
}
146+
else
147+
git rev-parse mg-upstream/$BASE_VERSION > /dev/null 2>&1 || {
148+
echo "ERROR: mg-upstream/$BASE_VERSION not found"; exit 1
149+
}
150+
fi
151+
if [ -n "$TARGET_VERSION_COMMIT" ]; then
152+
git rev-parse "$TARGET_VERSION_COMMIT" > /dev/null 2>&1 || {
153+
echo "ERROR: commit $TARGET_VERSION_COMMIT not found"; exit 1
154+
}
155+
else
156+
git rev-parse mg-upstream/$TARGET_VERSION > /dev/null 2>&1 || {
157+
echo "ERROR: mg-upstream/$TARGET_VERSION not found"; exit 1
158+
}
159+
fi
160+
```
161+
162+
#### Step 4: Reset main to origin/main
163+
164+
If local main has stale commits from previous upgrade attempts, reset it:
165+
166+
```bash
167+
cd "$FLAGSCALE_DIR"
168+
169+
# Check if local main has diverged from origin/main
170+
LOCAL_MAIN=$(git rev-parse main)
171+
ORIGIN_MAIN=$(git rev-parse origin/main)
172+
173+
if [ "$LOCAL_MAIN" != "$ORIGIN_MAIN" ]; then
174+
echo "Local main ($LOCAL_MAIN) differs from origin/main ($ORIGIN_MAIN)"
175+
echo "Commits on local main not in origin/main:"
176+
git log --oneline origin/main..main
177+
echo ""
178+
echo "Resetting local main to origin/main..."
179+
git checkout main
180+
git reset --hard origin/main
181+
echo "main reset to origin/main: $(git rev-parse --short HEAD)"
182+
else
183+
echo "main is already at origin/main: $(git rev-parse --short HEAD)"
184+
fi
185+
```
186+
187+
---
188+
189+
### Stage 2: Create Branch Structure for Three-Way Merge
190+
191+
This follows the same pattern as Megatron-LM-FL's `base`/`dev`/`main` branches. The three
192+
branches represent:
193+
- `base-train` — FlagScale's tree with training/legacy replaced by old upstream content
194+
- `main` — FlagScale's current tree (with FlagScale's customizations)
195+
- `dev-train` — FlagScale's tree with training/legacy replaced by new upstream content
196+
197+
The diff `base-train..main` captures FlagScale's customizations. Applying that diff to
198+
`dev-train` produces the upgraded result.
199+
200+
#### Step 1: Create `base-train` branch
201+
202+
Start from `main` (FlagScale's current state), then replace training/legacy content with the
203+
old upstream version. This creates a synthetic commit where the only difference from `main` is
204+
that training/legacy matches the old upstream.
205+
206+
```bash
207+
cd "$FLAGSCALE_DIR"
208+
209+
git branch -D base-train 2>/dev/null || true
210+
git checkout -b base-train main
211+
212+
# Resolve the base ref (commit SHA if provided, otherwise tag/branch)
213+
BASE_REF="${BASE_VERSION_COMMIT:-mg-upstream/$BASE_VERSION}"
214+
215+
# Replace training/ content with old upstream version
216+
# First, remove existing training and legacy dirs
217+
rm -rf flagscale/train/megatron/training/
218+
rm -rf flagscale/train/megatron/legacy/
219+
220+
# Extract old upstream training/ into FlagScale's path
221+
mkdir -p flagscale/train/megatron/training
222+
git archive $BASE_REF -- megatron/training/ | \
223+
tar -x --strip-components=2 -C flagscale/train/megatron/training/
224+
225+
# Extract old upstream legacy/ into FlagScale's path (if it exists)
226+
if git ls-tree $BASE_REF megatron/legacy/ > /dev/null 2>&1; then
227+
mkdir -p flagscale/train/megatron/legacy
228+
git archive $BASE_REF -- megatron/legacy/ | \
229+
tar -x --strip-components=2 -C flagscale/train/megatron/legacy/
230+
fi
231+
232+
git add -A
233+
git commit -m "base-train: replace training/legacy with upstream $BASE_VERSION content"
234+
echo "base-train branch created (ref: $BASE_REF)"
235+
```
236+
237+
#### Step 2: Create `dev-train` branch
238+
239+
Same approach, but with the new upstream version:
240+
241+
```bash
242+
cd "$FLAGSCALE_DIR"
243+
244+
git branch -D dev-train 2>/dev/null || true
245+
git checkout -b dev-train main
246+
247+
# Resolve the target ref (commit SHA if provided, otherwise tag/branch)
248+
TARGET_REF="${TARGET_VERSION_COMMIT:-mg-upstream/$TARGET_VERSION}"
249+
250+
# Replace training/ content with new upstream version
251+
rm -rf flagscale/train/megatron/training/
252+
rm -rf flagscale/train/megatron/legacy/
253+
254+
# Extract new upstream training/ into FlagScale's path
255+
mkdir -p flagscale/train/megatron/training
256+
git archive $TARGET_REF -- megatron/training/ | \
257+
tar -x --strip-components=2 -C flagscale/train/megatron/training/
258+
259+
# Extract new upstream legacy/ into FlagScale's path (if it exists)
260+
if git ls-tree $TARGET_REF megatron/legacy/ > /dev/null 2>&1; then
261+
mkdir -p flagscale/train/megatron/legacy
262+
git archive $TARGET_REF -- megatron/legacy/ | \
263+
tar -x --strip-components=2 -C flagscale/train/megatron/legacy/
264+
fi
265+
266+
git add -A
267+
git commit -m "dev-train: replace training/legacy with upstream $TARGET_VERSION content"
268+
echo "dev-train branch created (ref: $TARGET_REF)"
269+
```
270+
271+
#### Step 3: Verify the branch structure
272+
273+
```bash
274+
cd "$FLAGSCALE_DIR"
275+
276+
echo "=== Branch structure ==="
277+
echo "base-train: training/legacy = upstream $BASE_VERSION"
278+
echo "main: training/legacy = FlagScale's current version (with customizations)"
279+
echo "dev-train: training/legacy = upstream $TARGET_VERSION"
280+
281+
echo ""
282+
echo "=== FlagScale customizations (diff base-train..main) ==="
283+
git diff --stat base-train..main -- flagscale/train/megatron/training/ flagscale/train/megatron/legacy/
284+
285+
echo ""
286+
echo "=== Upstream changes (diff base-train..dev-train) ==="
287+
git diff --stat base-train..dev-train -- flagscale/train/megatron/training/ flagscale/train/megatron/legacy/
288+
```
289+
290+
Return to dev-train for the merge work:
291+
292+
```bash
293+
git checkout dev-train
294+
```
295+
296+
#### Checkpoint: commit
297+
298+
No pre-commit needed here — the branches are synthetic reference points.
299+
300+
```bash
301+
echo "Branch structure ready. Proceeding to Stage 3."
302+
```
303+
304+
---
305+

0 commit comments

Comments
 (0)