Skip to content

Commit 394037d

Browse files
dimapihtardimapihtar
authored
add ckpt tests (#2496)
Signed-off-by: dimapihtar <dpihtar@gmail.com>
1 parent dd0a1d9 commit 394037d

10 files changed

Lines changed: 692 additions & 0 deletions

‎.github/workflows/cicd-main.yml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -380,6 +380,12 @@ jobs:
380380
# - script: L2_Launch_quantization_export
381381
- script: L2_Launch_recipes_llama_cuda_graphs
382382
- script: L2_Launch_utils
383+
- script: L2_Launch_ckpts_mbridge_to_mlm_llama32_1b
384+
- script: L2_Launch_ckpts_mlm_to_mbridge_llama32_1b
385+
- script: L2_Launch_ckpts_mbridge_to_mlm_qwen3_4b
386+
- script: L2_Launch_ckpts_mlm_to_mbridge_qwen3_4b
387+
- script: L2_Launch_ckpts_mbridge_to_mlm_nemotronh_4b
388+
- script: L2_Launch_ckpts_mlm_to_mbridge_nemotronh_4b
383389
needs: [pre-flight, cicd-unit-tests]
384390
runs-on: ${{ needs.pre-flight.outputs.runner_prefix }}-gpu-x2
385391
if: |

‎tests/functional_tests/L2_Launch_ckpts_mbridge_to_mlm_llama32_1b.sh‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
#!/bin/bash
2+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
set -xeuo pipefail # Exit immediately if a command exits with a non-zero status
17+
18+
export CUDA_VISIBLE_DEVICES="0,1"
19+
20+
# Run recipe functional tests on 2 GPUs
21+
# This script tests recipe configurations with their default settings to ensure
22+
# they can run basic training without crashes
23+
uv run python -m torch.distributed.run --nproc_per_node=2 --nnodes=1 -m coverage run --data-file=/opt/Megatron-Bridge/.coverage --source=/opt/Megatron-Bridge/ --parallel-mode -m pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/llama32_1b/test_llama32_1b_ckpt.py::TestLlama32Ckpt::test_llama32_1B_ckpt_mbridge
24+
coverage combine -q
25+
26+
pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/llama32_1b/test_llama32_1b_ckpt.py::TestLlama32Ckpt::test_llama32_1B_ckpt_core
27+
28+
pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/llama32_1b/test_llama32_1b_ckpt.py::TestLlama32Ckpt::test_remove_artifacts

‎tests/functional_tests/L2_Launch_ckpts_mbridge_to_mlm_nemotronh_4b.sh‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
#!/bin/bash
2+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
set -xeuo pipefail # Exit immediately if a command exits with a non-zero status
17+
18+
export CUDA_VISIBLE_DEVICES="0,1"
19+
20+
# Run recipe functional tests on 2 GPUs
21+
# This script tests recipe configurations with their default settings to ensure
22+
# they can run basic training without crashes
23+
uv run python -m torch.distributed.run --nproc_per_node=2 --nnodes=1 -m coverage run --data-file=/opt/Megatron-Bridge/.coverage --source=/opt/Megatron-Bridge/ --parallel-mode -m pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/nemotronh_4b/test_nemotronh_4b_ckpt.py::TestNemotronhCkpt::test_nemotronh_4b_ckpt_mbridge
24+
coverage combine -q
25+
26+
pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/nemotronh_4b/test_nemotronh_4b_ckpt.py::TestNemotronhCkpt::test_nemotronh_4b_ckpt_mcore
27+
28+
pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/nemotronh_4b/test_nemotronh_4b_ckpt.py::TestNemotronhCkpt::test_remove_artifacts

‎tests/functional_tests/L2_Launch_ckpts_mbridge_to_mlm_qwen3_4b.sh‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
#!/bin/bash
2+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
set -xeuo pipefail # Exit immediately if a command exits with a non-zero status
17+
18+
export CUDA_VISIBLE_DEVICES="0,1"
19+
20+
# Run recipe functional tests on 2 GPUs
21+
# This script tests recipe configurations with their default settings to ensure
22+
# they can run basic training without crashes
23+
uv run python -m torch.distributed.run --nproc_per_node=2 --nnodes=1 -m coverage run --data-file=/opt/Megatron-Bridge/.coverage --source=/opt/Megatron-Bridge/ --parallel-mode -m pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py::TestQwen3Ckpt::test_qwen3_4b_ckpt_mbridge
24+
coverage combine -q
25+
26+
pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py::TestQwen3Ckpt::test_qwen3_4b_ckpt_mcore
27+
28+
pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py::TestQwen3Ckpt::test_remove_artifacts

‎tests/functional_tests/L2_Launch_ckpts_mlm_to_mbridge_llama32_1b.sh‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
#!/bin/bash
2+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
set -xeuo pipefail # Exit immediately if a command exits with a non-zero status
17+
18+
export CUDA_VISIBLE_DEVICES="0,1"
19+
20+
# Run recipe functional tests on 2 GPUs
21+
# This script tests recipe configurations with their default settings to ensure
22+
# they can run basic training without crashes
23+
pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/llama32_1b/test_llama32_1b_ckpt.py::TestLlama32Ckpt::test_llama32_1B_ckpt_core
24+
25+
uv run python -m torch.distributed.run --nproc_per_node=2 --nnodes=1 -m coverage run --data-file=/opt/Megatron-Bridge/.coverage --source=/opt/Megatron-Bridge/ --parallel-mode -m pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/llama32_1b/test_llama32_1b_ckpt.py::TestLlama32Ckpt::test_llama32_1B_ckpt_mbridge
26+
coverage combine -q
27+
28+
pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/llama32_1b/test_llama32_1b_ckpt.py::TestLlama32Ckpt::test_remove_artifacts

‎tests/functional_tests/L2_Launch_ckpts_mlm_to_mbridge_nemotronh_4b.sh‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
#!/bin/bash
2+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
set -xeuo pipefail # Exit immediately if a command exits with a non-zero status
17+
18+
export CUDA_VISIBLE_DEVICES="0,1"
19+
20+
# Run recipe functional tests on 2 GPUs
21+
# This script tests recipe configurations with their default settings to ensure
22+
# they can run basic training without crashes
23+
pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/nemotronh_4b/test_nemotronh_4b_ckpt.py::TestNemotronhCkpt::test_nemotronh_4b_ckpt_mcore
24+
25+
uv run python -m torch.distributed.run --nproc_per_node=2 --nnodes=1 -m coverage run --data-file=/opt/Megatron-Bridge/.coverage --source=/opt/Megatron-Bridge/ --parallel-mode -m pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/nemotronh_4b/test_nemotronh_4b_ckpt.py::TestNemotronhCkpt::test_nemotronh_4b_ckpt_mbridge
26+
coverage combine -q
27+
28+
29+
pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/nemotronh_4b/test_nemotronh_4b_ckpt.py::TestNemotronhCkpt::test_remove_artifacts

‎tests/functional_tests/L2_Launch_ckpts_mlm_to_mbridge_qwen3_4b.sh‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
#!/bin/bash
2+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
set -xeuo pipefail # Exit immediately if a command exits with a non-zero status
17+
18+
export CUDA_VISIBLE_DEVICES="0,1"
19+
20+
# Run recipe functional tests on 2 GPUs
21+
# This script tests recipe configurations with their default settings to ensure
22+
# they can run basic training without crashes
23+
pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py::TestQwen3Ckpt::test_qwen3_4b_ckpt_mcore
24+
25+
uv run python -m torch.distributed.run --nproc_per_node=2 --nnodes=1 -m coverage run --data-file=/opt/Megatron-Bridge/.coverage --source=/opt/Megatron-Bridge/ --parallel-mode -m pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py::TestQwen3Ckpt::test_qwen3_4b_ckpt_mbridge
26+
coverage combine -q
27+
28+
pytest -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA tests/functional_tests/ckpts/qwen3_4b/test_qwen3_4b_ckpt.py::TestQwen3Ckpt::test_remove_artifacts

‎tests/functional_tests/ckpts/llama32_1b/test_llama32_1b_ckpt.py‎

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Functional smoke tests for LLaMA checkpointing."""
16+
17+
import os
18+
import shutil
19+
import sys
20+
21+
import pytest
22+
from torch.distributed.run import main as torchrun_main
23+
24+
from megatron.bridge.recipes.llama import llama32_1b_pretrain_config
25+
from megatron.bridge.training.gpt_step import forward_step
26+
from megatron.bridge.training.pretrain import pretrain
27+
28+
29+
BASE_DIR = "/workspace/test_ckpts/llama32_1b"
30+
MBRIDGE_CKPT = f"{BASE_DIR}/mbridge"
31+
MCORE_CKPT = f"{BASE_DIR}/mcore"
32+
TB_DIR = f"{BASE_DIR}/tb"
33+
34+
35+
class TestLlama32Ckpt:
36+
"""Test class for LLama checkpoint functional tests."""
37+
38+
@pytest.mark.run_only_on("GPU")
39+
def test_llama32_1B_ckpt_mbridge(self):
40+
"""Functional test for LLama MBridge checkpoint."""
41+
42+
config = llama32_1b_pretrain_config()
43+
44+
config.checkpoint.save = MBRIDGE_CKPT
45+
config.checkpoint.load = MCORE_CKPT if os.path.exists(MCORE_CKPT) else None
46+
config.checkpoint.load_optim = False
47+
48+
config.model.seq_length = 8192
49+
50+
config.train.train_iters = 10 if config.checkpoint.load else 5
51+
config.train.eval_iters = 5
52+
config.train.save_interval = 5
53+
config.train.global_batch_size = 8
54+
config.train.micro_batch_size = 1
55+
56+
config.scheduler.lr_warmup_iters = 2
57+
58+
config.logger.log_interval = 1
59+
60+
pretrain(config=config, forward_step_func=forward_step)
61+
62+
@pytest.mark.run_only_on("GPU")
63+
def test_llama32_1B_ckpt_core(self, monkeypatch):
64+
"""Functional test for LLama MCore checkpoint."""
65+
66+
load_dir = MBRIDGE_CKPT if os.path.exists(MBRIDGE_CKPT) else None
67+
train_iters = 10 if load_dir else 5
68+
69+
# Set environment variables
70+
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1")
71+
monkeypatch.setenv("CUDA_DEVICE_MAX_CONNECTIONS", "1")
72+
73+
# Set MLM script
74+
monkeypatch.setattr(
75+
sys,
76+
"argv",
77+
[
78+
"torchrun",
79+
"--nproc-per-node=2",
80+
"/opt/Megatron-Bridge/3rdparty/Megatron-LM/pretrain_gpt.py",
81+
"--load",
82+
"/workspace/test_ckpts/llama32_1b_mbridge",
83+
"--save",
84+
"/workspace/test_ckpts/llama32_1b_mcore",
85+
"--init-method-std",
86+
"0.014",
87+
"--disable-bias-linear",
88+
"--use-rope-scaling",
89+
"--swiglu",
90+
"--use-rotary-position-embeddings",
91+
"--num-layers",
92+
"16",
93+
"--hidden-size",
94+
"2048",
95+
"--num-attention-heads",
96+
"32",
97+
"--ffn-hidden-size",
98+
"8192",
99+
"--kv-channels",
100+
"64",
101+
"--group-query-attention",
102+
"--position-embedding-type",
103+
"rope",
104+
"--attention-backend",
105+
"fused",
106+
"--num-query-groups",
107+
"8",
108+
"--normalization",
109+
"RMSNorm",
110+
"--attention-dropout",
111+
"0.0",
112+
"--hidden-dropout",
113+
"0.0",
114+
"--tensor-model-parallel-size",
115+
"1",
116+
"--pipeline-model-parallel-size",
117+
"1",
118+
"--seq-length",
119+
"8192",
120+
"--max-position-embeddings",
121+
"8192",
122+
"--micro-batch-size",
123+
"1",
124+
"--global-batch-size",
125+
"8",
126+
"--mock-data",
127+
"--tokenizer-type",
128+
"NullTokenizer",
129+
"--vocab-size",
130+
"131072",
131+
"--train-iters",
132+
f"{train_iters}",
133+
"--save-interval",
134+
"5",
135+
"--eval-interval",
136+
"5",
137+
"--eval-iters",
138+
"5",
139+
"--load",
140+
load_dir,
141+
"--save",
142+
MCORE_CKPT,
143+
"--ckpt-format",
144+
"torch_dist",
145+
"--log-progress",
146+
"--bf16",
147+
"--lr",
148+
"4.5e-4",
149+
"--min-lr",
150+
"4.5e-5",
151+
"--num-workers",
152+
"2",
153+
"--tensorboard-dir",
154+
TB_DIR,
155+
"--log-interval",
156+
"1",
157+
"--log-throughput",
158+
"--no-load-optim",
159+
],
160+
)
161+
162+
# Run MLM script
163+
torchrun_main()
164+
165+
def test_remove_artifacts(self):
166+
"""Removes model artifacts"""
167+
shutil.rmtree(BASE_DIR)
168+
169+
assert not os.path.exists(BASE_DIR)

0 commit comments

Comments
 (0)