diff --git a/.gemini/config.yaml b/.gemini/config.yaml index a348aa0899..43eda2ef30 100644 --- a/.gemini/config.yaml +++ b/.gemini/config.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Mainly adopted from https://github.com/vllm-project/vllm/blob/main/.gemini/config.yaml. # https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github have_fun: false # Just review the code diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index affe1e2a42..8db49aeecd 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Global reviewers for the entire repository * @FlagOpen/flagscale-reviewers diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index 1f45d8c5cb..d9282379d1 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -1,3 +1,19 @@ + + ### PR Category diff --git a/.github/configs/ascend.yml b/.github/configs/ascend.yml index db071e052d..bce2d5bdaa 100644 --- a/.github/configs/ascend.yml +++ b/.github/configs/ascend.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Ascend NPU Hardware Configuration # This file defines CI/CD settings for Huawei Ascend NPU-based testing # Test configurations are defined in tests/test_utils/config/platforms/ascend.yaml diff --git a/.github/configs/cuda.yml b/.github/configs/cuda.yml index 03b363157b..680b249e09 100644 --- a/.github/configs/cuda.yml +++ b/.github/configs/cuda.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # CUDA Hardware Configuration # This file defines CI/CD settings for CUDA-based testing # Test configurations are defined in tests/test_utils/config/platforms/cuda.yaml diff --git a/.github/configs/metax.yml b/.github/configs/metax.yml index d7dce35ae2..862d4ae81b 100644 --- a/.github/configs/metax.yml +++ b/.github/configs/metax.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # MetaX C550 Hardware Configuration # This file defines CI/CD settings for MetaX C550 GPU (MACA) testing # Test configurations are defined in tests/test_utils/config/platforms/metax.yaml diff --git a/.github/configs/template.yml b/.github/configs/template.yml index afbe2758f5..e4fed8d49f 100644 --- a/.github/configs/template.yml +++ b/.github/configs/template.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Template for adding new hardware # IMPORTANT: Fill in all fields below for your new hardware diff --git a/.github/workflows/all_tests.yml b/.github/workflows/all_tests.yml index 447dab2761..561d07b410 100644 --- a/.github/workflows/all_tests.yml +++ b/.github/workflows/all_tests.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + name: all_tests on: diff --git a/.github/workflows/all_tests_common.yml b/.github/workflows/all_tests_common.yml index 3db4af4820..7a4fba0884 100644 --- a/.github/workflows/all_tests_common.yml +++ b/.github/workflows/all_tests_common.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + name: Common All Tests on: diff --git a/.github/workflows/build_image_cuda.yml b/.github/workflows/build_image_cuda.yml index 1aae881c39..e4ed6edb0d 100644 --- a/.github/workflows/build_image_cuda.yml +++ b/.github/workflows/build_image_cuda.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + name: Build Docker Images - CUDA on: diff --git a/.github/workflows/claude.yml b/.github/workflows/claude.yml index 78d43004e8..188a88b883 100644 --- a/.github/workflows/claude.yml +++ b/.github/workflows/claude.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + name: Claude Code on: diff --git a/.github/workflows/format_check.yml b/.github/workflows/format_check.yml index 06b575535b..d4ca0701e1 100644 --- a/.github/workflows/format_check.yml +++ b/.github/workflows/format_check.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + name: format_check on: diff --git a/.github/workflows/functional_tests_benchmark.yml b/.github/workflows/functional_tests_benchmark.yml index 341efd49a2..281753bb56 100644 --- a/.github/workflows/functional_tests_benchmark.yml +++ b/.github/workflows/functional_tests_benchmark.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + name: Common Functional Tests - Benchmark on: diff --git a/.github/workflows/functional_tests_cli.yml b/.github/workflows/functional_tests_cli.yml index 70c3415e29..957be8a1c2 100644 --- a/.github/workflows/functional_tests_cli.yml +++ b/.github/workflows/functional_tests_cli.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + name: Common Functional Tests - CLI on: diff --git a/.github/workflows/functional_tests_hetero_train.yml b/.github/workflows/functional_tests_hetero_train.yml index b59feda3b5..2b64fe8c36 100644 --- a/.github/workflows/functional_tests_hetero_train.yml +++ b/.github/workflows/functional_tests_hetero_train.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + name: Common Functional Tests - Heterogeneous Training on: diff --git a/.github/workflows/functional_tests_inference.yml b/.github/workflows/functional_tests_inference.yml index d2b53ac8d4..bd783c8708 100644 --- a/.github/workflows/functional_tests_inference.yml +++ b/.github/workflows/functional_tests_inference.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + name: Common Functional Tests - Inference # Inference task workflow - currently minimal tests diff --git a/.github/workflows/functional_tests_megatron_fl_trigger.yml b/.github/workflows/functional_tests_megatron_fl_trigger.yml index 1e4717ad91..350e8bf54d 100644 --- a/.github/workflows/functional_tests_megatron_fl_trigger.yml +++ b/.github/workflows/functional_tests_megatron_fl_trigger.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + name: Functional Tests - Training (Megatron-LM-FL) # Triggered by Megatron-LM-FL Repo cicd workflow diff --git a/.github/workflows/functional_tests_rl.yml b/.github/workflows/functional_tests_rl.yml index 325d92f280..e175b0298e 100644 --- a/.github/workflows/functional_tests_rl.yml +++ b/.github/workflows/functional_tests_rl.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + name: Common Functional Tests - RL # RL task workflow - currently minimal tests diff --git a/.github/workflows/functional_tests_serve.yml b/.github/workflows/functional_tests_serve.yml index 73d992e3d2..05aac229c9 100644 --- a/.github/workflows/functional_tests_serve.yml +++ b/.github/workflows/functional_tests_serve.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + name: Common Functional Tests - Serve # Serve task workflow - currently minimal tests diff --git a/.github/workflows/functional_tests_train.yml b/.github/workflows/functional_tests_train.yml index 38ff486f95..a44a75e614 100644 --- a/.github/workflows/functional_tests_train.yml +++ b/.github/workflows/functional_tests_train.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + name: Common Functional Tests - Training on: diff --git a/.github/workflows/push_image_harbor.yml b/.github/workflows/push_image_harbor.yml index 62b57ab165..42abec7fe7 100644 --- a/.github/workflows/push_image_harbor.yml +++ b/.github/workflows/push_image_harbor.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + name: Push Images to Harbor on: diff --git a/.github/workflows/test_failure_analysis.yml b/.github/workflows/test_failure_analysis.yml index d784877041..748d37e6cb 100644 --- a/.github/workflows/test_failure_analysis.yml +++ b/.github/workflows/test_failure_analysis.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + name: Failure Analysis and Auto-Retry Flaky Tests # This example demonstrates using structured outputs to detect flaky test failures diff --git a/.github/workflows/unit_tests_common.yml b/.github/workflows/unit_tests_common.yml index 0c8627da97..3f41194399 100644 --- a/.github/workflows/unit_tests_common.yml +++ b/.github/workflows/unit_tests_common.yml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + name: Common Unit Tests on: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 338f85ec49..4e31ecc36d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + exclude: | (?x)^( flagscale/train/megatron/.*| diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md index 03e0b807e6..23cfbe4332 100644 --- a/CODE_OF_CONDUCT.md +++ b/CODE_OF_CONDUCT.md @@ -1,3 +1,19 @@ + + # Code of Conduct ## Our Pledge diff --git a/MAINTAINERS.md b/MAINTAINERS.md index 184254b857..b3e07230b0 100644 --- a/MAINTAINERS.md +++ b/MAINTAINERS.md @@ -1,3 +1,19 @@ + + # Maintainers | GitHub ID | Email Address | diff --git a/README_cn.md b/README_cn.md index 9f8d33cfb1..c50e1b9963 100644 --- a/README_cn.md +++ b/README_cn.md @@ -1,3 +1,19 @@ + + [github+banner-20260130](https://flagos.io/) [中文版|[English](./README.md)] diff --git a/__init__.py b/__init__.py index e69de29bb2..a491ed2726 100644 --- a/__init__.py +++ b/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/docker/build.sh b/docker/build.sh index ed8abbd6a4..f7089c6e06 100755 --- a/docker/build.sh +++ b/docker/build.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # FlagScale Docker Build Script # # NOTE: This script is experimental and requires further testing. diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 06e04bcfff..bcdc9bab95 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -1,3 +1,19 @@ + + ## Change History - **[2026/03]** Released [v1.0.0](https://github.com/flagos-ai/FlagScale/releases/tag/v1.0.0): diff --git a/docs/ascend.md b/docs/ascend.md index 86329f6f8e..02707ea227 100644 --- a/docs/ascend.md +++ b/docs/ascend.md @@ -1,3 +1,19 @@ + + # Ascend Guide ## Quick Start diff --git a/docs/getting-started.md b/docs/getting-started.md index 55b07aed25..8d4b811997 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -1,3 +1,19 @@ + + # Getting Started ## Overview diff --git a/docs/robo.md b/docs/robo.md index dcff0d1222..7aa82bd187 100644 --- a/docs/robo.md +++ b/docs/robo.md @@ -1,3 +1,19 @@ + + ## FlagOS-Robo Overview 🤖 FlagOS-Robo is built upon the unified and open-source AI system software stack, [FlagOS](https://flagos.io), which supports various AI chips. diff --git a/examples/aquila/__init__.py b/examples/aquila/__init__.py index e69de29bb2..a491ed2726 100644 --- a/examples/aquila/__init__.py +++ b/examples/aquila/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/examples/aquila/conf/hetero_auto_tuner.yaml b/examples/aquila/conf/hetero_auto_tuner.yaml index 951e2d2fb4..0fff73a982 100644 --- a/examples/aquila/conf/hetero_auto_tuner.yaml +++ b/examples/aquila/conf/hetero_auto_tuner.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - train: 1_8b - _self_ diff --git a/examples/aquila/conf/inference.yaml b/examples/aquila/conf/inference.yaml index 0cb7b8a620..e349426508 100644 --- a/examples/aquila/conf/inference.yaml +++ b/examples/aquila/conf/inference.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: 7b diff --git a/examples/aquila/conf/inference/7b.yaml b/examples/aquila/conf/inference/7b.yaml index 8cbb5ac302..7b1b5d9168 100644 --- a/examples/aquila/conf/inference/7b.yaml +++ b/examples/aquila/conf/inference/7b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: model: BAAI/Aquila-7B tokenizer: BAAI/Aquila-7B diff --git a/examples/aquila/conf/train.yaml b/examples/aquila/conf/train.yaml index 2ef8726787..9cd072b35b 100644 --- a/examples/aquila/conf/train.yaml +++ b/examples/aquila/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - train: 7b - _self_ diff --git a/examples/aquila/conf/train/1_8b.yaml b/examples/aquila/conf/train/1_8b.yaml index 21f8b90e5a..35a2931c42 100644 --- a/examples/aquila/conf/train/1_8b.yaml +++ b/examples/aquila/conf/train/1_8b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: #skip_iters_range: [10, 20] distributed_backend: flagcx diff --git a/examples/aquila/conf/train/34b.yaml b/examples/aquila/conf/train/34b.yaml index b882ef87e2..11e5e38a60 100644 --- a/examples/aquila/conf/train/34b.yaml +++ b/examples/aquila/conf/train/34b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 4 pipeline_model_parallel_size: 4 diff --git a/examples/aquila/conf/train/3b.yaml b/examples/aquila/conf/train/3b.yaml index 29d1735928..6569630cfa 100644 --- a/examples/aquila/conf/train/3b.yaml +++ b/examples/aquila/conf/train/3b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: distributed_backend: nccl reset_position_ids: True diff --git a/examples/aquila/conf/train/70b.yaml b/examples/aquila/conf/train/70b.yaml index d9a184c852..105c249f78 100644 --- a/examples/aquila/conf/train/70b.yaml +++ b/examples/aquila/conf/train/70b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 8 pipeline_model_parallel_size: 4 diff --git a/examples/aquila/conf/train/7b.yaml b/examples/aquila/conf/train/7b.yaml index bd3cb53dec..159a53392c 100644 --- a/examples/aquila/conf/train/7b.yaml +++ b/examples/aquila/conf/train/7b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 1 pipeline_model_parallel_size: 1 diff --git a/examples/aquila/conf/train/hetero.yaml b/examples/aquila/conf/train/hetero.yaml index 89e481564b..1514372642 100644 --- a/examples/aquila/conf/train/hetero.yaml +++ b/examples/aquila/conf/train/hetero.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 2 # pp stages diff --git a/examples/aquila/conf/train_auto_tuner.yaml b/examples/aquila/conf/train_auto_tuner.yaml index a5ecdc9f98..c6ccc0d8e3 100644 --- a/examples/aquila/conf/train_auto_tuner.yaml +++ b/examples/aquila/conf/train_auto_tuner.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - train: 7b - _self_ diff --git a/examples/aquila/conf/train_cloud.yaml b/examples/aquila/conf/train_cloud.yaml index 6a0f657d62..73927cd251 100644 --- a/examples/aquila/conf/train_cloud.yaml +++ b/examples/aquila/conf/train_cloud.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - train: 7b - _self_ diff --git a/examples/aquila/conf/train_hetero.yaml b/examples/aquila/conf/train_hetero.yaml index 5393dff14c..c7d6a17cef 100644 --- a/examples/aquila/conf/train_hetero.yaml +++ b/examples/aquila/conf/train_hetero.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - train: hetero - _self_ diff --git a/examples/aquila/tokenizer/merges.txt b/examples/aquila/tokenizer/merges.txt index 8d41af9ecf..97215499ef 100644 --- a/examples/aquila/tokenizer/merges.txt +++ b/examples/aquila/tokenizer/merges.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + #version: 0.2 - Trained by `huggingface/tokenizers` Ġ Ġ ä ¸ diff --git a/examples/aquila/tokenizer/special_tokens.txt b/examples/aquila/tokenizer/special_tokens.txt index df9c9338c7..c818a0b2f6 100644 --- a/examples/aquila/tokenizer/special_tokens.txt +++ b/examples/aquila/tokenizer/special_tokens.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + <|startofpiece|> <|endofpiece|> <|LDWANG|> diff --git a/examples/aquila/utils/__init__.py b/examples/aquila/utils/__init__.py index e69de29bb2..a491ed2726 100644 --- a/examples/aquila/utils/__init__.py +++ b/examples/aquila/utils/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/examples/aquila/utils/convo_dataset.py b/examples/aquila/utils/convo_dataset.py index 9db833d4f9..c5e04a0e14 100644 --- a/examples/aquila/utils/convo_dataset.py +++ b/examples/aquila/utils/convo_dataset.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """GPT style dataset.""" import copy diff --git a/examples/aquila/utils/convo_prompt.py b/examples/aquila/utils/convo_prompt.py index 92a8dbe3e1..a20bb08883 100644 --- a/examples/aquila/utils/convo_prompt.py +++ b/examples/aquila/utils/convo_prompt.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import dataclasses from enum import Enum, auto from typing import Any diff --git a/examples/aquila/utils/cyg_conversation.py b/examples/aquila/utils/cyg_conversation.py index 91eab70b8f..a61bd0edda 100644 --- a/examples/aquila/utils/cyg_conversation.py +++ b/examples/aquila/utils/cyg_conversation.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import dataclasses from enum import Enum, auto from typing import Any diff --git a/examples/cloud/conf/serve.yaml b/examples/cloud/conf/serve.yaml index 8807df24d4..a23b8ec32b 100644 --- a/examples/cloud/conf/serve.yaml +++ b/examples/cloud/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: cloud_model diff --git a/examples/cloud/conf/serve/cloud_model.yaml b/examples/cloud/conf/serve/cloud_model.yaml index 09a02328d5..732c38b038 100644 --- a/examples/cloud/conf/serve/cloud_model.yaml +++ b/examples/cloud/conf/serve/cloud_model.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine: vllm engine_args: diff --git a/examples/deepseek_r1/conf/hostfile.txt b/examples/deepseek_r1/conf/hostfile.txt index 0d8b1e05f4..de1a4088bf 100644 --- a/examples/deepseek_r1/conf/hostfile.txt +++ b/examples/deepseek_r1/conf/hostfile.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ip slots type=xxx[optional] # master node x.x.x.x slots=8 type=gpu diff --git a/examples/deepseek_r1/conf/serve.yaml b/examples/deepseek_r1/conf/serve.yaml index 6f08e9d701..b8a0c4a434 100644 --- a/examples/deepseek_r1/conf/serve.yaml +++ b/examples/deepseek_r1/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 671b diff --git a/examples/deepseek_r1/conf/serve/671b.yaml b/examples/deepseek_r1/conf/serve/671b.yaml index 7235e59731..1c1fa478f9 100644 --- a/examples/deepseek_r1/conf/serve/671b.yaml +++ b/examples/deepseek_r1/conf/serve/671b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /models/deepseek_r1 # path of weight of deepseek r1 diff --git a/examples/deepseek_r1_distill_qwen/conf/serve.yaml b/examples/deepseek_r1_distill_qwen/conf/serve.yaml index 021c3ce47d..0db5b39c32 100644 --- a/examples/deepseek_r1_distill_qwen/conf/serve.yaml +++ b/examples/deepseek_r1_distill_qwen/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 32b diff --git a/examples/deepseek_r1_distill_qwen/conf/serve/32b.yaml b/examples/deepseek_r1_distill_qwen/conf/serve/32b.yaml index 06dd3e25b9..d5fe287035 100644 --- a/examples/deepseek_r1_distill_qwen/conf/serve/32b.yaml +++ b/examples/deepseek_r1_distill_qwen/conf/serve/32b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /share/DeepSeek-R1-Distill-Qwen-32B # should be customized diff --git a/examples/deepseek_v3/conf/finetune.yaml b/examples/deepseek_v3/conf/finetune.yaml index fc5e0ddc2d..f3d079ac59 100644 --- a/examples/deepseek_v3/conf/finetune.yaml +++ b/examples/deepseek_v3/conf/finetune.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # DeepSeek 16B, 2.4A Model defaults: - _self_ diff --git a/examples/deepseek_v3/conf/serve.yaml b/examples/deepseek_v3/conf/serve.yaml index f4554f3b67..f6434bfc2d 100644 --- a/examples/deepseek_v3/conf/serve.yaml +++ b/examples/deepseek_v3/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 671b diff --git a/examples/deepseek_v3/conf/serve/671b.yaml b/examples/deepseek_v3/conf/serve/671b.yaml index 3def3b41c9..9f783a83d3 100644 --- a/examples/deepseek_v3/conf/serve/671b.yaml +++ b/examples/deepseek_v3/conf/serve/671b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /models/deepseek_v3/model # path of weight of deepseek v3 diff --git a/examples/deepseek_v3/conf/serve_auto_tuner.yaml b/examples/deepseek_v3/conf/serve_auto_tuner.yaml index 5f627e28aa..404a8dfdfa 100644 --- a/examples/deepseek_v3/conf/serve_auto_tuner.yaml +++ b/examples/deepseek_v3/conf/serve_auto_tuner.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 671b diff --git a/examples/deepseek_v3/conf/train.yaml b/examples/deepseek_v3/conf/train.yaml index 2dda3fca1f..3d044af28f 100644 --- a/examples/deepseek_v3/conf/train.yaml +++ b/examples/deepseek_v3/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # DeepSeek 16B, 2.4A Model defaults: - _self_ diff --git a/examples/deepseek_v3/conf/train/11b_a2b.yaml b/examples/deepseek_v3/conf/train/11b_a2b.yaml index fc3699ab25..e8b58b1e1b 100644 --- a/examples/deepseek_v3/conf/train/11b_a2b.yaml +++ b/examples/deepseek_v3/conf/train/11b_a2b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: no_shared_fs: ${experiment.runner.no_shared_fs} num_workers: 2 diff --git a/examples/deepseek_v3/conf/train/16b_a3b.yaml b/examples/deepseek_v3/conf/train/16b_a3b.yaml index ea10ee7a3e..114c868378 100644 --- a/examples/deepseek_v3/conf/train/16b_a3b.yaml +++ b/examples/deepseek_v3/conf/train/16b_a3b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: no_shared_fs: ${experiment.runner.no_shared_fs} num_workers: 2 diff --git a/examples/deepseek_v3/conf/train/16b_a3b_finetune.yaml b/examples/deepseek_v3/conf/train/16b_a3b_finetune.yaml index 518727434a..99c1ba56f7 100644 --- a/examples/deepseek_v3/conf/train/16b_a3b_finetune.yaml +++ b/examples/deepseek_v3/conf/train/16b_a3b_finetune.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: no_shared_fs: ${experiment.runner.no_shared_fs} num_workers: 2 diff --git a/examples/deepseek_v3/conf/train/engram.yaml b/examples/deepseek_v3/conf/train/engram.yaml index 1026e12eeb..677b2ec18e 100644 --- a/examples/deepseek_v3/conf/train/engram.yaml +++ b/examples/deepseek_v3/conf/train/engram.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # DeepSeek Engram 27B system: no_shared_fs: ${experiment.runner.no_shared_fs} diff --git a/examples/deepseek_v3/conf/train/hetero.yaml b/examples/deepseek_v3/conf/train/hetero.yaml index 24fdac5a9b..35a5fe4efa 100644 --- a/examples/deepseek_v3/conf/train/hetero.yaml +++ b/examples/deepseek_v3/conf/train/hetero.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: no_shared_fs: ${experiment.runner.no_shared_fs} num_workers: 16 diff --git a/examples/deepseek_v3/conf/train_11b_a2b.yaml b/examples/deepseek_v3/conf/train_11b_a2b.yaml index e1c39fc148..e81e13b9e4 100644 --- a/examples/deepseek_v3/conf/train_11b_a2b.yaml +++ b/examples/deepseek_v3/conf/train_11b_a2b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # DeepSeek light model with 16 layers, 11b a2b # For DualPipeV Performance Testing defaults: diff --git a/examples/deepseek_v3/conf/train_auto_tuner.yaml b/examples/deepseek_v3/conf/train_auto_tuner.yaml index ed4caeb3a2..6845531ce7 100644 --- a/examples/deepseek_v3/conf/train_auto_tuner.yaml +++ b/examples/deepseek_v3/conf/train_auto_tuner.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # DeepSeek 16B, 2.4A Model defaults: - _self_ diff --git a/examples/deepseek_v3/conf/train_engram.yaml b/examples/deepseek_v3/conf/train_engram.yaml index 30a4f9fa32..d035ae0e3c 100644 --- a/examples/deepseek_v3/conf/train_engram.yaml +++ b/examples/deepseek_v3/conf/train_engram.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: engram diff --git a/examples/deepseek_v3/conf/train_hetero.yaml b/examples/deepseek_v3/conf/train_hetero.yaml index e5d2271977..f4ac7a23c2 100644 --- a/examples/deepseek_v3/conf/train_hetero.yaml +++ b/examples/deepseek_v3/conf/train_hetero.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # DeepSeek 16B, 2.4A Model defaults: - _self_ diff --git a/examples/emu3.5/conf/image_generation.yaml b/examples/emu3.5/conf/image_generation.yaml index d0428f10f0..07f68541c6 100644 --- a/examples/emu3.5/conf/image_generation.yaml +++ b/examples/emu3.5/conf/image_generation.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: t2i # x2i or t2i diff --git a/examples/emu3.5/conf/inference/t2i.yaml b/examples/emu3.5/conf/inference/t2i.yaml index 1c44dfdea7..96a74b39cc 100644 --- a/examples/emu3.5/conf/inference/t2i.yaml +++ b/examples/emu3.5/conf/inference/t2i.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: save: ${experiment.exp_dir} vq_model: ${experiment.vq_model} diff --git a/examples/emu3.5/conf/inference/visual_guidance.yaml b/examples/emu3.5/conf/inference/visual_guidance.yaml index bdba0d74e6..81f1262792 100644 --- a/examples/emu3.5/conf/inference/visual_guidance.yaml +++ b/examples/emu3.5/conf/inference/visual_guidance.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: save: ${experiment.exp_dir} vq_model: ${experiment.vq_model} diff --git a/examples/emu3.5/conf/inference/visual_narrative.yaml b/examples/emu3.5/conf/inference/visual_narrative.yaml index 62cb75f03a..7897dd9a91 100644 --- a/examples/emu3.5/conf/inference/visual_narrative.yaml +++ b/examples/emu3.5/conf/inference/visual_narrative.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: save: ${experiment.exp_dir} vq_model: ${experiment.vq_model} diff --git a/examples/emu3.5/conf/inference/x2i.yaml b/examples/emu3.5/conf/inference/x2i.yaml index ea0c5092b3..1acbaccdc1 100644 --- a/examples/emu3.5/conf/inference/x2i.yaml +++ b/examples/emu3.5/conf/inference/x2i.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: save: ${experiment.exp_dir} vq_model: ${experiment.vq_model} diff --git a/examples/emu3.5/conf/interleaved_generation.yaml b/examples/emu3.5/conf/interleaved_generation.yaml index 7ad6765818..950b532b7d 100644 --- a/examples/emu3.5/conf/interleaved_generation.yaml +++ b/examples/emu3.5/conf/interleaved_generation.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: visual_guidance # visual_narrative or visual_guidance diff --git a/examples/emu3.5/conf/serve/emu3p5.py b/examples/emu3.5/conf/serve/emu3p5.py index 538541a6a6..56e1923abb 100644 --- a/examples/emu3.5/conf/serve/emu3p5.py +++ b/examples/emu3.5/conf/serve/emu3p5.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Modified from FlagSale/flagscale/inference/inference_emu3p5.py import os diff --git a/examples/emu3.5/conf/serve/emu3p5.yaml b/examples/emu3.5/conf/serve/emu3p5.yaml index 3022d0e131..43c63d5646 100644 --- a/examples/emu3.5/conf/serve/emu3p5.yaml +++ b/examples/emu3.5/conf/serve/emu3p5.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: A module: examples/emu3.5/conf/serve/emu3p5.py name: Emu3_5InferencePipeline diff --git a/examples/emu3.5/conf/serve_emu3p5.yaml b/examples/emu3.5/conf/serve_emu3p5.yaml index bacb0cd2bf..9f2f5d416a 100644 --- a/examples/emu3.5/conf/serve_emu3p5.yaml +++ b/examples/emu3.5/conf/serve_emu3p5.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: emu3p5 diff --git a/examples/emu3/conf/chat.yaml b/examples/emu3/conf/chat.yaml index 61710c6e2d..17b090b02f 100644 --- a/examples/emu3/conf/chat.yaml +++ b/examples/emu3/conf/chat.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: i2t diff --git a/examples/emu3/conf/compress.yaml b/examples/emu3/conf/compress.yaml index b7d41cbe4d..b814495a65 100644 --- a/examples/emu3/conf/compress.yaml +++ b/examples/emu3/conf/compress.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - compress: fp8_dynamic diff --git a/examples/emu3/conf/compress/fp8_dynamic.yaml b/examples/emu3/conf/compress/fp8_dynamic.yaml index 309156969a..d10e7ab609 100644 --- a/examples/emu3/conf/compress/fp8_dynamic.yaml +++ b/examples/emu3/conf/compress/fp8_dynamic.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - model - _self_ diff --git a/examples/emu3/conf/compress/model.yaml b/examples/emu3/conf/compress/model.yaml index 7e21b35f47..221b921993 100644 --- a/examples/emu3/conf/compress/model.yaml +++ b/examples/emu3/conf/compress/model.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 1 pipeline_model_parallel_size: 1 diff --git a/examples/emu3/conf/compress/w4a16.yaml b/examples/emu3/conf/compress/w4a16.yaml index a403886fb5..3650ce22c1 100644 --- a/examples/emu3/conf/compress/w4a16.yaml +++ b/examples/emu3/conf/compress/w4a16.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - model - _self_ diff --git a/examples/emu3/conf/gen.yaml b/examples/emu3/conf/gen.yaml index 1a29a83937..e238d0c531 100644 --- a/examples/emu3/conf/gen.yaml +++ b/examples/emu3/conf/gen.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: t2i diff --git a/examples/emu3/conf/inference/i2t.yaml b/examples/emu3/conf/inference/i2t.yaml index 9f1b3ab05e..d7080aef45 100644 --- a/examples/emu3/conf/inference/i2t.yaml +++ b/examples/emu3/conf/inference/i2t.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: model: ${experiment.model} vq_model: ${experiment.vq_model} diff --git a/examples/emu3/conf/inference/t2i.yaml b/examples/emu3/conf/inference/t2i.yaml index 58086d4c62..5351cce7f5 100644 --- a/examples/emu3/conf/inference/t2i.yaml +++ b/examples/emu3/conf/inference/t2i.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: model: ${experiment.model} vq_model: ${experiment.vq_model} diff --git a/examples/ernie45/conf/hostfile.txt b/examples/ernie45/conf/hostfile.txt index 0d8b1e05f4..de1a4088bf 100644 --- a/examples/ernie45/conf/hostfile.txt +++ b/examples/ernie45/conf/hostfile.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ip slots type=xxx[optional] # master node x.x.x.x slots=8 type=gpu diff --git a/examples/ernie45/conf/serve.yaml b/examples/ernie45/conf/serve.yaml index d33413fa66..d4c33913ef 100644 --- a/examples/ernie45/conf/serve.yaml +++ b/examples/ernie45/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 300b diff --git a/examples/ernie45/conf/serve/300b.yaml b/examples/ernie45/conf/serve/300b.yaml index 93ab11afdd..f609e9590e 100644 --- a/examples/ernie45/conf/serve/300b.yaml +++ b/examples/ernie45/conf/serve/300b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: PaddlePaddle/ERNIE-4.5-300B-A47B-PT # add Post-trained models diff --git a/examples/flux_1_dev/conf/inference.yaml b/examples/flux_1_dev/conf/inference.yaml index ad9215a251..8131b0b468 100644 --- a/examples/flux_1_dev/conf/inference.yaml +++ b/examples/flux_1_dev/conf/inference.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: t2i diff --git a/examples/flux_1_dev/conf/inference/t2i.yaml b/examples/flux_1_dev/conf/inference/t2i.yaml index 6aeab70604..2d71c3321b 100644 --- a/examples/flux_1_dev/conf/inference/t2i.yaml +++ b/examples/flux_1_dev/conf/inference/t2i.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + engine: model: ${experiment.model} loader: diffusers diff --git a/examples/gr00t_n1_5/conf/serve.yaml b/examples/gr00t_n1_5/conf/serve.yaml index d240a35d4c..2ed66f22cb 100644 --- a/examples/gr00t_n1_5/conf/serve.yaml +++ b/examples/gr00t_n1_5/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: gr00t_n1_5 diff --git a/examples/gr00t_n1_5/conf/serve/gr00t_n1_5.yaml b/examples/gr00t_n1_5/conf/serve/gr00t_n1_5.yaml index 45279c2e35..5a9df44f26 100644 --- a/examples/gr00t_n1_5/conf/serve/gr00t_n1_5.yaml +++ b/examples/gr00t_n1_5/conf/serve/gr00t_n1_5.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model # Not in use engine_args: host: 0.0.0.0 diff --git a/examples/gr00t_n1_5/conf/train.yaml b/examples/gr00t_n1_5/conf/train.yaml index 89ef63badb..1c019483c0 100644 --- a/examples/gr00t_n1_5/conf/train.yaml +++ b/examples/gr00t_n1_5/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: gr00t_n1_5 diff --git a/examples/gr00t_n1_5/conf/train/gr00t_n1_5.yaml b/examples/gr00t_n1_5/conf/train/gr00t_n1_5.yaml index 0cfd4cc9d8..9bae66440e 100644 --- a/examples/gr00t_n1_5/conf/train/gr00t_n1_5.yaml +++ b/examples/gr00t_n1_5/conf/train/gr00t_n1_5.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: batch_size: 32 train_steps: 10000 diff --git a/examples/grok2/conf/serve.yaml b/examples/grok2/conf/serve.yaml index 556e64c9c3..0bbbd901ec 100644 --- a/examples/grok2/conf/serve.yaml +++ b/examples/grok2/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 270b diff --git a/examples/grok2/conf/serve/270b.yaml b/examples/grok2/conf/serve/270b.yaml index 3bdc23c461..6e3a92cb4a 100644 --- a/examples/grok2/conf/serve/270b.yaml +++ b/examples/grok2/conf/serve/270b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /share/grok-2 diff --git a/examples/kimi_k2/conf/hostfile.txt b/examples/kimi_k2/conf/hostfile.txt index 0d8b1e05f4..de1a4088bf 100644 --- a/examples/kimi_k2/conf/hostfile.txt +++ b/examples/kimi_k2/conf/hostfile.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ip slots type=xxx[optional] # master node x.x.x.x slots=8 type=gpu diff --git a/examples/kimi_k2/conf/serve.yaml b/examples/kimi_k2/conf/serve.yaml index ac050787d1..2ed31467f6 100644 --- a/examples/kimi_k2/conf/serve.yaml +++ b/examples/kimi_k2/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 1t diff --git a/examples/kimi_k2/conf/serve/1t.yaml b/examples/kimi_k2/conf/serve/1t.yaml index ff80115504..75e2354eb5 100644 --- a/examples/kimi_k2/conf/serve/1t.yaml +++ b/examples/kimi_k2/conf/serve/1t.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /models/Kimi-K2-Instruct # add Post-trained models diff --git a/examples/llama2/conf/train.yaml b/examples/llama2/conf/train.yaml index 553ac3d7f5..a62aa1c1c7 100644 --- a/examples/llama2/conf/train.yaml +++ b/examples/llama2/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - train: 7b - _self_ diff --git a/examples/llama2/conf/train/7b.yaml b/examples/llama2/conf/train/7b.yaml index bee99d65cf..ddae25641d 100644 --- a/examples/llama2/conf/train/7b.yaml +++ b/examples/llama2/conf/train/7b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 1 pipeline_model_parallel_size: 1 diff --git a/examples/llama3/conf/train.yaml b/examples/llama3/conf/train.yaml index 7cf4426346..1b42275dff 100644 --- a/examples/llama3/conf/train.yaml +++ b/examples/llama3/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - train: 70b - _self_ diff --git a/examples/llama3/conf/train/70b.yaml b/examples/llama3/conf/train/70b.yaml index 8e4a06be51..44f962cc2a 100644 --- a/examples/llama3/conf/train/70b.yaml +++ b/examples/llama3/conf/train/70b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 8 pipeline_model_parallel_size: 4 diff --git a/examples/llama3/conf/train/70b_finetune.yaml b/examples/llama3/conf/train/70b_finetune.yaml index aca7f56b10..848cf505c2 100644 --- a/examples/llama3/conf/train/70b_finetune.yaml +++ b/examples/llama3/conf/train/70b_finetune.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 8 pipeline_model_parallel_size: 4 diff --git a/examples/llama3/conf/train/8b.yaml b/examples/llama3/conf/train/8b.yaml index 41e3c1d826..214e86d4ae 100644 --- a/examples/llama3/conf/train/8b.yaml +++ b/examples/llama3/conf/train/8b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 4 pipeline_model_parallel_size: 2 diff --git a/examples/llama3/conf/train/8b_hetero.yaml b/examples/llama3/conf/train/8b_hetero.yaml index 1724e94895..daf49cae91 100644 --- a/examples/llama3/conf/train/8b_hetero.yaml +++ b/examples/llama3/conf/train/8b_hetero.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 4 pipeline_model_parallel_size: 2 diff --git a/examples/llama3/conf/train_hetero.yaml b/examples/llama3/conf/train_hetero.yaml index 9e8abdd603..8c19a330a0 100644 --- a/examples/llama3/conf/train_hetero.yaml +++ b/examples/llama3/conf/train_hetero.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - train: 8b_hetero - _self_ diff --git a/examples/llava1_5/conf/train.yaml b/examples/llava1_5/conf/train.yaml index 9fc4a40ce8..0b3375f720 100644 --- a/examples/llava1_5/conf/train.yaml +++ b/examples/llava1_5/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - train: 7b - _self_ diff --git a/examples/llava1_5/conf/train/7b.yaml b/examples/llava1_5/conf/train/7b.yaml index fc13c683f5..c4b74aaa4a 100644 --- a/examples/llava1_5/conf/train/7b.yaml +++ b/examples/llava1_5/conf/train/7b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 1 pipeline_model_parallel_size: 1 diff --git a/examples/llava_onevision/conf/compress.yaml b/examples/llava_onevision/conf/compress.yaml index 42819009ae..734ed91d60 100644 --- a/examples/llava_onevision/conf/compress.yaml +++ b/examples/llava_onevision/conf/compress.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - compress: w4a16 diff --git a/examples/llava_onevision/conf/compress/fp8_dynamic.yaml b/examples/llava_onevision/conf/compress/fp8_dynamic.yaml index 5ae15de6a9..3066b7b13e 100644 --- a/examples/llava_onevision/conf/compress/fp8_dynamic.yaml +++ b/examples/llava_onevision/conf/compress/fp8_dynamic.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - model - _self_ diff --git a/examples/llava_onevision/conf/compress/model.yaml b/examples/llava_onevision/conf/compress/model.yaml index 877494e539..672f4815d5 100644 --- a/examples/llava_onevision/conf/compress/model.yaml +++ b/examples/llava_onevision/conf/compress/model.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 1 pipeline_model_parallel_size: 1 diff --git a/examples/llava_onevision/conf/compress/w4a16.yaml b/examples/llava_onevision/conf/compress/w4a16.yaml index fde0e51925..b406875b73 100644 --- a/examples/llava_onevision/conf/compress/w4a16.yaml +++ b/examples/llava_onevision/conf/compress/w4a16.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - model - _self_ diff --git a/examples/llava_onevision/conf/hostfile.txt b/examples/llava_onevision/conf/hostfile.txt index 0d8b1e05f4..de1a4088bf 100644 --- a/examples/llava_onevision/conf/hostfile.txt +++ b/examples/llava_onevision/conf/hostfile.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ip slots type=xxx[optional] # master node x.x.x.x slots=8 type=gpu diff --git a/examples/llava_onevision/conf/serve.yaml b/examples/llava_onevision/conf/serve.yaml index 40622dd46a..824e14e032 100644 --- a/examples/llava_onevision/conf/serve.yaml +++ b/examples/llava_onevision/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 7b diff --git a/examples/llava_onevision/conf/serve/7b.yaml b/examples/llava_onevision/conf/serve/7b.yaml index 65112005a2..d4161efc58 100644 --- a/examples/llava_onevision/conf/serve/7b.yaml +++ b/examples/llava_onevision/conf/serve/7b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /models/llava diff --git a/examples/llava_onevision/conf/train.yaml b/examples/llava_onevision/conf/train.yaml index bb710c96aa..4824f3a2fb 100644 --- a/examples/llava_onevision/conf/train.yaml +++ b/examples/llava_onevision/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - train: 1_5b - _self_ diff --git a/examples/llava_onevision/conf/train/1_5b.yaml b/examples/llava_onevision/conf/train/1_5b.yaml index e43dc1d063..cb8789ecfc 100644 --- a/examples/llava_onevision/conf/train/1_5b.yaml +++ b/examples/llava_onevision/conf/train/1_5b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 1 pipeline_model_parallel_size: 1 diff --git a/examples/llava_onevision/conf/train/7b.yaml b/examples/llava_onevision/conf/train/7b.yaml index a49fe940cd..f685a9ad44 100644 --- a/examples/llava_onevision/conf/train/7b.yaml +++ b/examples/llava_onevision/conf/train/7b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 2 pipeline_model_parallel_size: 1 diff --git a/examples/minicpm_o_2.6/conf/serve.yaml b/examples/minicpm_o_2.6/conf/serve.yaml index c05bea03d8..a973d6d8bf 100644 --- a/examples/minicpm_o_2.6/conf/serve.yaml +++ b/examples/minicpm_o_2.6/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 7b diff --git a/examples/minicpm_o_2.6/conf/serve/7b.yaml b/examples/minicpm_o_2.6/conf/serve/7b.yaml index 6d5fa06acc..9aec76ebb2 100644 --- a/examples/minicpm_o_2.6/conf/serve/7b.yaml +++ b/examples/minicpm_o_2.6/conf/serve/7b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /models/minicpm_o_2.6 # path of weight of minicpm_o_2.6 diff --git a/examples/minicpm_v_4/conf/serve.yaml b/examples/minicpm_v_4/conf/serve.yaml index 9e4d8210c6..90248dac9c 100644 --- a/examples/minicpm_v_4/conf/serve.yaml +++ b/examples/minicpm_v_4/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 4b diff --git a/examples/minicpm_v_4/conf/serve/4b.yaml b/examples/minicpm_v_4/conf/serve/4b.yaml index 1cd4812186..1d2e444b0d 100644 --- a/examples/minicpm_v_4/conf/serve/4b.yaml +++ b/examples/minicpm_v_4/conf/serve/4b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /share/models/MiniCPM-V-4 diff --git a/examples/mixtral/conf/train.yaml b/examples/mixtral/conf/train.yaml index 5ec96cf551..501ee66f7a 100644 --- a/examples/mixtral/conf/train.yaml +++ b/examples/mixtral/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: 8x7b diff --git a/examples/mixtral/conf/train/8x7b.yaml b/examples/mixtral/conf/train/8x7b.yaml index 9332c5e980..2bdbe7253d 100644 --- a/examples/mixtral/conf/train/8x7b.yaml +++ b/examples/mixtral/conf/train/8x7b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 2 pipeline_model_parallel_size: 4 diff --git a/examples/openjourney/conf/inference.yaml b/examples/openjourney/conf/inference.yaml index 1846e77934..3bb3886a02 100644 --- a/examples/openjourney/conf/inference.yaml +++ b/examples/openjourney/conf/inference.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: t2i diff --git a/examples/openjourney/conf/inference/t2i.yaml b/examples/openjourney/conf/inference/t2i.yaml index bdaf44a0a0..9c8667c8e5 100644 --- a/examples/openjourney/conf/inference/t2i.yaml +++ b/examples/openjourney/conf/inference/t2i.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + engine: model: ${experiment.model} loader: diffusers diff --git a/examples/pi0/client_pi0.py b/examples/pi0/client_pi0.py index 4074ad3839..a48b339208 100644 --- a/examples/pi0/client_pi0.py +++ b/examples/pi0/client_pi0.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import base64 import json diff --git a/examples/pi0/conf/inference.yaml b/examples/pi0/conf/inference.yaml index 61337c5957..f7afdd20eb 100644 --- a/examples/pi0/conf/inference.yaml +++ b/examples/pi0/conf/inference.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: pi0 diff --git a/examples/pi0/conf/inference/pi0.yaml b/examples/pi0/conf/inference/pi0.yaml index eba8605925..270dbaf772 100644 --- a/examples/pi0/conf/inference/pi0.yaml +++ b/examples/pi0/conf/inference/pi0.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + engine: # Model variant: "pi0" or "pi0.5" model_variant: "pi0" diff --git a/examples/pi0/conf/serve.yaml b/examples/pi0/conf/serve.yaml index a3e775821a..d2d3d9cb38 100644 --- a/examples/pi0/conf/serve.yaml +++ b/examples/pi0/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: pi0 diff --git a/examples/pi0/conf/serve/pi0.yaml b/examples/pi0/conf/serve/pi0.yaml index 5f29f03442..ec4f36ecc1 100644 --- a/examples/pi0/conf/serve/pi0.yaml +++ b/examples/pi0/conf/serve/pi0.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: host: 0.0.0.0 diff --git a/examples/pi0/conf/train.yaml b/examples/pi0/conf/train.yaml index 859bc50631..9cd55f9242 100644 --- a/examples/pi0/conf/train.yaml +++ b/examples/pi0/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: pi0 diff --git a/examples/pi0/conf/train/pi0.yaml b/examples/pi0/conf/train/pi0.yaml index 29b8e2b51a..de1ebf894e 100644 --- a/examples/pi0/conf/train/pi0.yaml +++ b/examples/pi0/conf/train/pi0.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: batch_size: 1 train_steps: 100000 diff --git a/examples/pi0/download.py b/examples/pi0/download.py index 48391c1e72..d3a38e797c 100755 --- a/examples/pi0/download.py +++ b/examples/pi0/download.py @@ -1,4 +1,19 @@ #!/usr/bin/env python3 + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """ Download models or datasets from HuggingFace Hub or ModelScope to a user-defined folder. diff --git a/examples/pi0/dump_dataset_inputs.py b/examples/pi0/dump_dataset_inputs.py index 1cb9c61acb..202f6bd5a7 100644 --- a/examples/pi0/dump_dataset_inputs.py +++ b/examples/pi0/dump_dataset_inputs.py @@ -1,4 +1,19 @@ #!/usr/bin/env python + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """ Extract inference inputs (images, state, task) from a LeRobotDataset. diff --git a/examples/pi0_5/conf/inference.yaml b/examples/pi0_5/conf/inference.yaml index c95eba6603..20bb253cf9 100644 --- a/examples/pi0_5/conf/inference.yaml +++ b/examples/pi0_5/conf/inference.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: pi0_5 diff --git a/examples/pi0_5/conf/inference/pi0_5.yaml b/examples/pi0_5/conf/inference/pi0_5.yaml index 975b3f0f2c..b3e07b4f9d 100644 --- a/examples/pi0_5/conf/inference/pi0_5.yaml +++ b/examples/pi0_5/conf/inference/pi0_5.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + engine: # Model variant: "pi0" or "pi0.5" model_variant: "pi0.5" diff --git a/examples/pi0_5/conf/serve.yaml b/examples/pi0_5/conf/serve.yaml index 7d8b361e12..30899fec3a 100644 --- a/examples/pi0_5/conf/serve.yaml +++ b/examples/pi0_5/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: pi0_5 diff --git a/examples/pi0_5/conf/serve/pi0_5.yaml b/examples/pi0_5/conf/serve/pi0_5.yaml index 4eb6bff94f..ee7de717b7 100644 --- a/examples/pi0_5/conf/serve/pi0_5.yaml +++ b/examples/pi0_5/conf/serve/pi0_5.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: host: 0.0.0.0 diff --git a/examples/pi0_5/conf/train.yaml b/examples/pi0_5/conf/train.yaml index 60e88123b4..12cbae4a4f 100644 --- a/examples/pi0_5/conf/train.yaml +++ b/examples/pi0_5/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: pi0_5 diff --git a/examples/pi0_5/conf/train/pi0_5.yaml b/examples/pi0_5/conf/train/pi0_5.yaml index a33721daa3..1e07c7a9c3 100644 --- a/examples/pi0_5/conf/train/pi0_5.yaml +++ b/examples/pi0_5/conf/train/pi0_5.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: batch_size: 16 train_steps: 30000 diff --git a/examples/qwen2/conf/rl.yaml b/examples/qwen2/conf/rl.yaml index be4180d070..fca6147207 100644 --- a/examples/qwen2/conf/rl.yaml +++ b/examples/qwen2/conf/rl.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - rl: 7b - _self_ diff --git a/examples/qwen2/conf/rl/7b.yaml b/examples/qwen2/conf/rl/7b.yaml index 17bb67633c..14384bf94e 100644 --- a/examples/qwen2/conf/rl/7b.yaml +++ b/examples/qwen2/conf/rl/7b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + #default verl config path config-path: config diff --git a/examples/qwen2_5/__init__.py b/examples/qwen2_5/__init__.py index e69de29bb2..a491ed2726 100644 --- a/examples/qwen2_5/__init__.py +++ b/examples/qwen2_5/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/examples/qwen2_5/conf/hostfile.txt b/examples/qwen2_5/conf/hostfile.txt index 0d8b1e05f4..de1a4088bf 100644 --- a/examples/qwen2_5/conf/hostfile.txt +++ b/examples/qwen2_5/conf/hostfile.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ip slots type=xxx[optional] # master node x.x.x.x slots=8 type=gpu diff --git a/examples/qwen2_5/conf/hostfile_nodes_envs.txt b/examples/qwen2_5/conf/hostfile_nodes_envs.txt index f460d34eff..391b0aaa1b 100644 --- a/examples/qwen2_5/conf/hostfile_nodes_envs.txt +++ b/examples/qwen2_5/conf/hostfile_nodes_envs.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ip slots type=xxx[optional] # master node x.x.x.x slots=4 type=gpu diff --git a/examples/qwen2_5/conf/rl.yaml b/examples/qwen2_5/conf/rl.yaml index 64d080ce4b..884810a12e 100644 --- a/examples/qwen2_5/conf/rl.yaml +++ b/examples/qwen2_5/conf/rl.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - rl: 0_5b - _self_ diff --git a/examples/qwen2_5/conf/rl/0_5b.yaml b/examples/qwen2_5/conf/rl/0_5b.yaml index 89df024601..d214b5681e 100644 --- a/examples/qwen2_5/conf/rl/0_5b.yaml +++ b/examples/qwen2_5/conf/rl/0_5b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + #default verl config path config-path: config diff --git a/examples/qwen2_5/conf/serve.yaml b/examples/qwen2_5/conf/serve.yaml index a8c96978c8..1bd1ab51b2 100644 --- a/examples/qwen2_5/conf/serve.yaml +++ b/examples/qwen2_5/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 7b diff --git a/examples/qwen2_5/conf/serve/72b.yaml b/examples/qwen2_5/conf/serve/72b.yaml index d4025afc7d..42f71481fd 100644 --- a/examples/qwen2_5/conf/serve/72b.yaml +++ b/examples/qwen2_5/conf/serve/72b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /models/Qwen2.5-72B-Instruct diff --git a/examples/qwen2_5/conf/serve/72b_nodes_envs.yaml b/examples/qwen2_5/conf/serve/72b_nodes_envs.yaml index 84c3a07ce8..82b7512aee 100644 --- a/examples/qwen2_5/conf/serve/72b_nodes_envs.yaml +++ b/examples/qwen2_5/conf/serve/72b_nodes_envs.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /models/Qwen2.5-72B-Instruct diff --git a/examples/qwen2_5/conf/serve/7b.yaml b/examples/qwen2_5/conf/serve/7b.yaml index 2c5f5b6770..f97ee0281f 100644 --- a/examples/qwen2_5/conf/serve/7b.yaml +++ b/examples/qwen2_5/conf/serve/7b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /models/Qwen2.5-7B-Instruct diff --git a/examples/qwen2_5/conf/serve/7b_multiple_instance.yaml b/examples/qwen2_5/conf/serve/7b_multiple_instance.yaml index 26bbf75d9b..dce3a93b37 100644 --- a/examples/qwen2_5/conf/serve/7b_multiple_instance.yaml +++ b/examples/qwen2_5/conf/serve/7b_multiple_instance.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /models/Qwen2.5-7B-Instruct diff --git a/examples/qwen2_5/conf/serve/multiple_models.yaml b/examples/qwen2_5/conf/serve/multiple_models.yaml index 2403d6b961..b2e8aacbff 100644 --- a/examples/qwen2_5/conf/serve/multiple_models.yaml +++ b/examples/qwen2_5/conf/serve/multiple_models.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: A module: examples/qwen2_5/models.py name: Model1Logic diff --git a/examples/qwen2_5/conf/serve_auto_tuner.yaml b/examples/qwen2_5/conf/serve_auto_tuner.yaml index e1697a304c..3515852800 100644 --- a/examples/qwen2_5/conf/serve_auto_tuner.yaml +++ b/examples/qwen2_5/conf/serve_auto_tuner.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 7b diff --git a/examples/qwen2_5/conf/serve_disagg_xpyd.yaml b/examples/qwen2_5/conf/serve_disagg_xpyd.yaml index 063dcf0a24..c574b7ae0a 100644 --- a/examples/qwen2_5/conf/serve_disagg_xpyd.yaml +++ b/examples/qwen2_5/conf/serve_disagg_xpyd.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 7b diff --git a/examples/qwen2_5/conf/serve_multiple_instance.yaml b/examples/qwen2_5/conf/serve_multiple_instance.yaml index 1562edf29e..7cca45b654 100644 --- a/examples/qwen2_5/conf/serve_multiple_instance.yaml +++ b/examples/qwen2_5/conf/serve_multiple_instance.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 7b_multiple_instance diff --git a/examples/qwen2_5/conf/serve_multiple_models.yaml b/examples/qwen2_5/conf/serve_multiple_models.yaml index ec6e3adabd..31cceb49c2 100644 --- a/examples/qwen2_5/conf/serve_multiple_models.yaml +++ b/examples/qwen2_5/conf/serve_multiple_models.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: multiple_models diff --git a/examples/qwen2_5/conf/serve_nodes_envs.yaml b/examples/qwen2_5/conf/serve_nodes_envs.yaml index 5bb9da2451..00747d229b 100644 --- a/examples/qwen2_5/conf/serve_nodes_envs.yaml +++ b/examples/qwen2_5/conf/serve_nodes_envs.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 72b_nodes_envs diff --git a/examples/qwen2_5/conf/serve_ssh.yaml b/examples/qwen2_5/conf/serve_ssh.yaml index 416c53c573..6d88751c96 100644 --- a/examples/qwen2_5/conf/serve_ssh.yaml +++ b/examples/qwen2_5/conf/serve_ssh.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 7b diff --git a/examples/qwen2_5/conf/train.yaml b/examples/qwen2_5/conf/train.yaml index 41522e97a8..8bfaedc74d 100644 --- a/examples/qwen2_5/conf/train.yaml +++ b/examples/qwen2_5/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: 1_5b diff --git a/examples/qwen2_5/conf/train/10b.yaml b/examples/qwen2_5/conf/train/10b.yaml index d9dd8a0a14..c36213360b 100644 --- a/examples/qwen2_5/conf/train/10b.yaml +++ b/examples/qwen2_5/conf/train/10b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 1 pipeline_model_parallel_size: 2 diff --git a/examples/qwen2_5/conf/train/1_5b.yaml b/examples/qwen2_5/conf/train/1_5b.yaml index b24b53be33..9879ef0cca 100644 --- a/examples/qwen2_5/conf/train/1_5b.yaml +++ b/examples/qwen2_5/conf/train/1_5b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 1 pipeline_model_parallel_size: 1 diff --git a/examples/qwen2_5/models.py b/examples/qwen2_5/models.py index f9633e29a0..353bc290d7 100644 --- a/examples/qwen2_5/models.py +++ b/examples/qwen2_5/models.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from vllm import LLM, SamplingParams diff --git a/examples/qwen2_5_vl/conf/serve.yaml b/examples/qwen2_5_vl/conf/serve.yaml index 90bd6c3691..0c74560bee 100644 --- a/examples/qwen2_5_vl/conf/serve.yaml +++ b/examples/qwen2_5_vl/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 32b_instruct diff --git a/examples/qwen2_5_vl/conf/serve/32b_instruct.yaml b/examples/qwen2_5_vl/conf/serve/32b_instruct.yaml index 5143e92eb5..8427b970e4 100644 --- a/examples/qwen2_5_vl/conf/serve/32b_instruct.yaml +++ b/examples/qwen2_5_vl/conf/serve/32b_instruct.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /share/Qwen2.5-VL-32B-Instruct # should be customized diff --git a/examples/qwen2_5_vl/conf/train.yaml b/examples/qwen2_5_vl/conf/train.yaml index 99330d0c34..b48691e1ab 100644 --- a/examples/qwen2_5_vl/conf/train.yaml +++ b/examples/qwen2_5_vl/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - train: 7b - _self_ diff --git a/examples/qwen2_5_vl/conf/train/32b.yaml b/examples/qwen2_5_vl/conf/train/32b.yaml index 642f7e9ac1..fe740e8016 100644 --- a/examples/qwen2_5_vl/conf/train/32b.yaml +++ b/examples/qwen2_5_vl/conf/train/32b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: num_workers: 1 calculate_per_token_loss: true diff --git a/examples/qwen2_5_vl/conf/train/7b.yaml b/examples/qwen2_5_vl/conf/train/7b.yaml index f3f62e27da..f849b01076 100644 --- a/examples/qwen2_5_vl/conf/train/7b.yaml +++ b/examples/qwen2_5_vl/conf/train/7b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: num_workers: 1 calculate_per_token_loss: true diff --git a/examples/qwen3/conf/hostfile.txt b/examples/qwen3/conf/hostfile.txt index 28c4c38aad..a8ba9e8141 100644 --- a/examples/qwen3/conf/hostfile.txt +++ b/examples/qwen3/conf/hostfile.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ip slots type=xxx[optional] # master node x.x.x.x slots=2 type=gpu diff --git a/examples/qwen3/conf/inference/4b.yaml b/examples/qwen3/conf/inference/4b.yaml index 1dab361957..35209c3052 100644 --- a/examples/qwen3/conf/inference/4b.yaml +++ b/examples/qwen3/conf/inference/4b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: model: Qwen/Qwen3-4B trust_remote_code: true diff --git a/examples/qwen3/conf/inference_fl.yaml b/examples/qwen3/conf/inference_fl.yaml index 66131e1e56..2438b697d1 100644 --- a/examples/qwen3/conf/inference_fl.yaml +++ b/examples/qwen3/conf/inference_fl.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: 4b diff --git a/examples/qwen3/conf/rl.yaml b/examples/qwen3/conf/rl.yaml index a1eef41a5f..81c5ec2279 100644 --- a/examples/qwen3/conf/rl.yaml +++ b/examples/qwen3/conf/rl.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - rl: 0_6b - _self_ diff --git a/examples/qwen3/conf/rl/0_6b.yaml b/examples/qwen3/conf/rl/0_6b.yaml index 9cec28cc88..9a7617d986 100644 --- a/examples/qwen3/conf/rl/0_6b.yaml +++ b/examples/qwen3/conf/rl/0_6b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + #default verl config path config-path: config diff --git a/examples/qwen3/conf/rl/0_6b_megatron.yaml b/examples/qwen3/conf/rl/0_6b_megatron.yaml index 8e88e17fa0..18d29d6bdc 100644 --- a/examples/qwen3/conf/rl/0_6b_megatron.yaml +++ b/examples/qwen3/conf/rl/0_6b_megatron.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + #default verl config path config-path: config diff --git a/examples/qwen3/conf/serve.yaml b/examples/qwen3/conf/serve.yaml index f0b4656f31..a7b31ac130 100644 --- a/examples/qwen3/conf/serve.yaml +++ b/examples/qwen3/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 8b diff --git a/examples/qwen3/conf/serve/0_6b.yaml b/examples/qwen3/conf/serve/0_6b.yaml index d233ad5a4a..8be03ed860 100644 --- a/examples/qwen3/conf/serve/0_6b.yaml +++ b/examples/qwen3/conf/serve/0_6b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /tmp/models/Qwen3-0.6B/ diff --git a/examples/qwen3/conf/serve/8b.yaml b/examples/qwen3/conf/serve/8b.yaml index 360d49b92f..28b5720f22 100644 --- a/examples/qwen3/conf/serve/8b.yaml +++ b/examples/qwen3/conf/serve/8b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: ${oc.env:QWEN3_PATH} diff --git a/examples/qwen3/conf/serve/sglang_qwen3_8b.yaml b/examples/qwen3/conf/serve/sglang_qwen3_8b.yaml index 2f3871a7e7..314320c676 100644 --- a/examples/qwen3/conf/serve/sglang_qwen3_8b.yaml +++ b/examples/qwen3/conf/serve/sglang_qwen3_8b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /models/Qwen3-8B/ diff --git a/examples/qwen3/conf/serve_atmb.yaml b/examples/qwen3/conf/serve_atmb.yaml index b64fa95513..72e594eb76 100644 --- a/examples/qwen3/conf/serve_atmb.yaml +++ b/examples/qwen3/conf/serve_atmb.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # atmb: Auto Tuning for Multi-Backend defaults: - _self_ diff --git a/examples/qwen3/conf/serve_sglang_multiple_nodes.yaml b/examples/qwen3/conf/serve_sglang_multiple_nodes.yaml index 0d32b50746..0d3c840d05 100644 --- a/examples/qwen3/conf/serve_sglang_multiple_nodes.yaml +++ b/examples/qwen3/conf/serve_sglang_multiple_nodes.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # atmb: Auto Tuning for Multi-Backend defaults: - _self_ diff --git a/examples/qwen3/conf/train.yaml b/examples/qwen3/conf/train.yaml index 5a173a4e59..5bea356dbb 100644 --- a/examples/qwen3/conf/train.yaml +++ b/examples/qwen3/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ # - train: 30b_a3b diff --git a/examples/qwen3/conf/train/0_6b.yaml b/examples/qwen3/conf/train/0_6b.yaml index edc664a2af..62b630f2d2 100644 --- a/examples/qwen3/conf/train/0_6b.yaml +++ b/examples/qwen3/conf/train/0_6b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: no_shared_fs: ${experiment.runner.no_shared_fs} num_workers: 16 diff --git a/examples/qwen3/conf/train/10b.yaml b/examples/qwen3/conf/train/10b.yaml index 2a268c07a4..ca84c822c5 100644 --- a/examples/qwen3/conf/train/10b.yaml +++ b/examples/qwen3/conf/train/10b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: no_shared_fs: ${experiment.runner.no_shared_fs} num_workers: 16 diff --git a/examples/qwen3/conf/train/10b_muon.yaml b/examples/qwen3/conf/train/10b_muon.yaml index e64d80b210..8f45ed8856 100644 --- a/examples/qwen3/conf/train/10b_muon.yaml +++ b/examples/qwen3/conf/train/10b_muon.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: no_shared_fs: ${experiment.runner.no_shared_fs} num_workers: 16 diff --git a/examples/qwen3/conf/train/14b.yaml b/examples/qwen3/conf/train/14b.yaml index eb3638e999..04a3aa1905 100644 --- a/examples/qwen3/conf/train/14b.yaml +++ b/examples/qwen3/conf/train/14b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: no_shared_fs: ${experiment.runner.no_shared_fs} num_workers: 16 diff --git a/examples/qwen3/conf/train/235b_a22b.yaml b/examples/qwen3/conf/train/235b_a22b.yaml index c559c57907..14913e1547 100644 --- a/examples/qwen3/conf/train/235b_a22b.yaml +++ b/examples/qwen3/conf/train/235b_a22b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: no_shared_fs: ${experiment.runner.no_shared_fs} num_workers: 16 diff --git a/examples/qwen3/conf/train/30b_a3b.yaml b/examples/qwen3/conf/train/30b_a3b.yaml index f1934d68d0..8700cc4464 100644 --- a/examples/qwen3/conf/train/30b_a3b.yaml +++ b/examples/qwen3/conf/train/30b_a3b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: no_shared_fs: ${experiment.runner.no_shared_fs} num_workers: 2 diff --git a/examples/qwen3/conf/train/32b.yaml b/examples/qwen3/conf/train/32b.yaml index 52812809a0..3512597e3f 100644 --- a/examples/qwen3/conf/train/32b.yaml +++ b/examples/qwen3/conf/train/32b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: no_shared_fs: ${experiment.runner.no_shared_fs} num_workers: 2 diff --git a/examples/qwen3/conf/train/engram.yaml b/examples/qwen3/conf/train/engram.yaml index 9b04379c62..6c52b32831 100644 --- a/examples/qwen3/conf/train/engram.yaml +++ b/examples/qwen3/conf/train/engram.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: no_shared_fs: ${experiment.runner.no_shared_fs} num_workers: 2 diff --git a/examples/qwen3/conf/train/te_fl.yaml b/examples/qwen3/conf/train/te_fl.yaml index 3cd60605a8..efe738b679 100644 --- a/examples/qwen3/conf/train/te_fl.yaml +++ b/examples/qwen3/conf/train/te_fl.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: no_shared_fs: ${experiment.runner.no_shared_fs} num_workers: 2 diff --git a/examples/qwen3/conf/train_auto_tuner.yaml b/examples/qwen3/conf/train_auto_tuner.yaml index feccbe57e2..26620426f6 100644 --- a/examples/qwen3/conf/train_auto_tuner.yaml +++ b/examples/qwen3/conf/train_auto_tuner.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ # - train: 30b_a3b diff --git a/examples/qwen3/conf/train_engram.yaml b/examples/qwen3/conf/train_engram.yaml index a9c5b93d7b..ea420007aa 100644 --- a/examples/qwen3/conf/train_engram.yaml +++ b/examples/qwen3/conf/train_engram.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: engram diff --git a/examples/qwen3/conf/train_te_fl.yaml b/examples/qwen3/conf/train_te_fl.yaml index bfcb2c9996..b2f47f3bc0 100644 --- a/examples/qwen3/conf/train_te_fl.yaml +++ b/examples/qwen3/conf/train_te_fl.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: te_fl diff --git a/examples/qwen35/conf/train.yaml b/examples/qwen35/conf/train.yaml index dc815bf228..af25e7182e 100644 --- a/examples/qwen35/conf/train.yaml +++ b/examples/qwen35/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: 4b diff --git a/examples/qwen35/conf/train/0.8b.yaml b/examples/qwen35/conf/train/0.8b.yaml index 435fcc8b18..acc2384b5b 100644 --- a/examples/qwen35/conf/train/0.8b.yaml +++ b/examples/qwen35/conf/train/0.8b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: vision_ration: 0.1 num_workers: 1 diff --git a/examples/qwen35/conf/train/27b.yaml b/examples/qwen35/conf/train/27b.yaml index 7ac30de5f6..d595c834c3 100644 --- a/examples/qwen35/conf/train/27b.yaml +++ b/examples/qwen35/conf/train/27b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: vision_ration: 0.1 num_workers: 1 diff --git a/examples/qwen35/conf/train/35b_a3b.yaml b/examples/qwen35/conf/train/35b_a3b.yaml index 1fb30cd8a6..01c6a60111 100644 --- a/examples/qwen35/conf/train/35b_a3b.yaml +++ b/examples/qwen35/conf/train/35b_a3b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: vision_ration: 0.1 num_workers: 1 diff --git a/examples/qwen35/conf/train/4b.yaml b/examples/qwen35/conf/train/4b.yaml index 8cb3f6b91f..996524374c 100644 --- a/examples/qwen35/conf/train/4b.yaml +++ b/examples/qwen35/conf/train/4b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: vision_ration: 0.1 num_workers: 1 diff --git a/examples/qwen3_next/conf/inference/80b_a3b.yaml b/examples/qwen3_next/conf/inference/80b_a3b.yaml index 8a95042c82..8b8e46e51f 100644 --- a/examples/qwen3_next/conf/inference/80b_a3b.yaml +++ b/examples/qwen3_next/conf/inference/80b_a3b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: model: Qwen/Qwen3-Next-80B-A3B-Instruct trust_remote_code: true diff --git a/examples/qwen3_next/conf/inference_fl.yaml b/examples/qwen3_next/conf/inference_fl.yaml index a81e36b982..397b675876 100644 --- a/examples/qwen3_next/conf/inference_fl.yaml +++ b/examples/qwen3_next/conf/inference_fl.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: 80b_a3b diff --git a/examples/qwen3_o/conf/serve.yaml b/examples/qwen3_o/conf/serve.yaml index cb15b31a2d..bf5959391f 100644 --- a/examples/qwen3_o/conf/serve.yaml +++ b/examples/qwen3_o/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 30b diff --git a/examples/qwen3_o/conf/serve/30b.yaml b/examples/qwen3_o/conf/serve/30b.yaml index 7cf8c29ec6..fe9b0368f8 100644 --- a/examples/qwen3_o/conf/serve/30b.yaml +++ b/examples/qwen3_o/conf/serve/30b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /share/Qwen3-Omni-30B-A3B-Instruct diff --git a/examples/qwen3_vl/conf/train.yaml b/examples/qwen3_vl/conf/train.yaml index 3027fe121e..a4472e0de8 100644 --- a/examples/qwen3_vl/conf/train.yaml +++ b/examples/qwen3_vl/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - train: 8b - _self_ diff --git a/examples/qwen3_vl/conf/train/32b.yaml b/examples/qwen3_vl/conf/train/32b.yaml index e6ffb3dac4..57e3c646e9 100644 --- a/examples/qwen3_vl/conf/train/32b.yaml +++ b/examples/qwen3_vl/conf/train/32b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: num_workers: 1 calculate_per_token_loss: true diff --git a/examples/qwen3_vl/conf/train/4b.yaml b/examples/qwen3_vl/conf/train/4b.yaml index e9d7677389..0abd79dffb 100644 --- a/examples/qwen3_vl/conf/train/4b.yaml +++ b/examples/qwen3_vl/conf/train/4b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: num_workers: 1 calculate_per_token_loss: true diff --git a/examples/qwen3_vl/conf/train/8b.yaml b/examples/qwen3_vl/conf/train/8b.yaml index b0bd812182..e61b73854d 100644 --- a/examples/qwen3_vl/conf/train/8b.yaml +++ b/examples/qwen3_vl/conf/train/8b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: # enable vision recompute especially when memory is limited # vision_recompute_activations: True diff --git a/examples/qwen_gr00t/conf/inference.yaml b/examples/qwen_gr00t/conf/inference.yaml index 1a988df98f..84f78683ae 100644 --- a/examples/qwen_gr00t/conf/inference.yaml +++ b/examples/qwen_gr00t/conf/inference.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: qwen_gr00t diff --git a/examples/qwen_gr00t/conf/inference/qwen_gr00t.yaml b/examples/qwen_gr00t/conf/inference/qwen_gr00t.yaml index e7169843f4..ac45f92af4 100644 --- a/examples/qwen_gr00t/conf/inference/qwen_gr00t.yaml +++ b/examples/qwen_gr00t/conf/inference/qwen_gr00t.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + engine: model_variant: "QwenGr00t" model: /workspace/models/qwen_gr00t_train/checkpoints/last diff --git a/examples/qwen_gr00t/conf/serve.yaml b/examples/qwen_gr00t/conf/serve.yaml index 18f961dd7d..6958010c31 100644 --- a/examples/qwen_gr00t/conf/serve.yaml +++ b/examples/qwen_gr00t/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: qwen_gr00t diff --git a/examples/qwen_gr00t/conf/serve/qwen_gr00t.yaml b/examples/qwen_gr00t/conf/serve/qwen_gr00t.yaml index c0dec5eaf7..db17f303c4 100644 --- a/examples/qwen_gr00t/conf/serve/qwen_gr00t.yaml +++ b/examples/qwen_gr00t/conf/serve/qwen_gr00t.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model # Not in use engine_args: host: 0.0.0.0 diff --git a/examples/qwen_gr00t/conf/train.yaml b/examples/qwen_gr00t/conf/train.yaml index 1b98e3c9b6..6a8655ce0c 100644 --- a/examples/qwen_gr00t/conf/train.yaml +++ b/examples/qwen_gr00t/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: qwen_gr00t diff --git a/examples/qwen_gr00t/conf/train/qwen_gr00t.yaml b/examples/qwen_gr00t/conf/train/qwen_gr00t.yaml index 04072a5e28..20f387d3be 100644 --- a/examples/qwen_gr00t/conf/train/qwen_gr00t.yaml +++ b/examples/qwen_gr00t/conf/train/qwen_gr00t.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: batch_size: 8 train_steps: 30000 diff --git a/examples/qwen_image/conf/inference.yaml b/examples/qwen_image/conf/inference.yaml index 32adc7a7fe..0220d4cd39 100644 --- a/examples/qwen_image/conf/inference.yaml +++ b/examples/qwen_image/conf/inference.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: t2i diff --git a/examples/qwen_image/conf/inference/t2i.yaml b/examples/qwen_image/conf/inference/t2i.yaml index 7cbd927e32..32214e3ed1 100644 --- a/examples/qwen_image/conf/inference/t2i.yaml +++ b/examples/qwen_image/conf/inference/t2i.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + engine: model: ${experiment.model} loader: diffusers diff --git a/examples/qwq/conf/serve.yaml b/examples/qwq/conf/serve.yaml index c24d1baec9..93aff8870e 100644 --- a/examples/qwq/conf/serve.yaml +++ b/examples/qwq/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 32b diff --git a/examples/qwq/conf/serve/32b.yaml b/examples/qwq/conf/serve/32b.yaml index d8e3b38f71..1c6f170f6d 100644 --- a/examples/qwq/conf/serve/32b.yaml +++ b/examples/qwq/conf/serve/32b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /share/QwQ-32B # should be customized diff --git a/examples/qwq/conf/train.yaml b/examples/qwq/conf/train.yaml index 54662b4a56..8a88633cf8 100644 --- a/examples/qwq/conf/train.yaml +++ b/examples/qwq/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: 32b diff --git a/examples/qwq/conf/train/32b.yaml b/examples/qwq/conf/train/32b.yaml index 7a6c0dd670..014eea3b67 100644 --- a/examples/qwq/conf/train/32b.yaml +++ b/examples/qwq/conf/train/32b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: no_shared_fs: ${experiment.runner.no_shared_fs} num_workers: 2 diff --git a/examples/robobrain/conf/compress.yaml b/examples/robobrain/conf/compress.yaml index 6a6bfa68a5..0ed8050a24 100644 --- a/examples/robobrain/conf/compress.yaml +++ b/examples/robobrain/conf/compress.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - compress: w8a16 diff --git a/examples/robobrain/conf/compress/model.yaml b/examples/robobrain/conf/compress/model.yaml index 835162e63c..10623a0a98 100644 --- a/examples/robobrain/conf/compress/model.yaml +++ b/examples/robobrain/conf/compress/model.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 1 pipeline_model_parallel_size: 1 diff --git a/examples/robobrain/conf/compress/w8a16.yaml b/examples/robobrain/conf/compress/w8a16.yaml index f194fb47c9..e924e147f9 100644 --- a/examples/robobrain/conf/compress/w8a16.yaml +++ b/examples/robobrain/conf/compress/w8a16.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - model - _self_ diff --git a/examples/robobrain/conf/serve.yaml b/examples/robobrain/conf/serve.yaml index a9d87d3a06..b96cabd450 100644 --- a/examples/robobrain/conf/serve.yaml +++ b/examples/robobrain/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 7b diff --git a/examples/robobrain/conf/serve/7b.yaml b/examples/robobrain/conf/serve/7b.yaml index 8505a31cff..70e2fc31cd 100644 --- a/examples/robobrain/conf/serve/7b.yaml +++ b/examples/robobrain/conf/serve/7b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /models/robobrain diff --git a/examples/robobrain2/conf/inference.yaml b/examples/robobrain2/conf/inference.yaml index 9ed1457b68..c186352d02 100644 --- a/examples/robobrain2/conf/inference.yaml +++ b/examples/robobrain2/conf/inference.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: 3b diff --git a/examples/robobrain2/conf/inference/32b.yaml b/examples/robobrain2/conf/inference/32b.yaml index 697676cf09..b82accdeca 100644 --- a/examples/robobrain2/conf/inference/32b.yaml +++ b/examples/robobrain2/conf/inference/32b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: # path to checkpoint model: /tmp/models/BAAI/RoboBrain2.0-7B diff --git a/examples/robobrain2/conf/inference/3b.yaml b/examples/robobrain2/conf/inference/3b.yaml index e81cbe7a2e..70c08f1129 100644 --- a/examples/robobrain2/conf/inference/3b.yaml +++ b/examples/robobrain2/conf/inference/3b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: # path to checkpoint model: /nfs/hcr/models/BAAI/RoboBrain2.0-3B diff --git a/examples/robobrain2/conf/inference/7b.yaml b/examples/robobrain2/conf/inference/7b.yaml index 5568ad95e4..f1acbe30f5 100644 --- a/examples/robobrain2/conf/inference/7b.yaml +++ b/examples/robobrain2/conf/inference/7b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: # path to checkpoint model: /tmp/models/BAAI/RoboBrain2.0-7B diff --git a/examples/robobrain2/conf/serve.yaml b/examples/robobrain2/conf/serve.yaml index 0387dec5f0..3394f19e08 100644 --- a/examples/robobrain2/conf/serve.yaml +++ b/examples/robobrain2/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 3b diff --git a/examples/robobrain2/conf/serve/32b.yaml b/examples/robobrain2/conf/serve/32b.yaml index 22d461d932..2160896a3c 100644 --- a/examples/robobrain2/conf/serve/32b.yaml +++ b/examples/robobrain2/conf/serve/32b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /share/project/lms/robobrain2_32b # should be customized diff --git a/examples/robobrain2/conf/serve/3b.yaml b/examples/robobrain2/conf/serve/3b.yaml index d378a87c8e..fdf3bfba3c 100644 --- a/examples/robobrain2/conf/serve/3b.yaml +++ b/examples/robobrain2/conf/serve/3b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /tmp/models/BAAI/RoboBrain2.0-3B # should be customized diff --git a/examples/robobrain2/conf/serve/7b.yaml b/examples/robobrain2/conf/serve/7b.yaml index ff96872566..9960f77f67 100644 --- a/examples/robobrain2/conf/serve/7b.yaml +++ b/examples/robobrain2/conf/serve/7b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /share/project/lms/robobrain2_7b # should be customized diff --git a/examples/robobrain2_5/conf/inference.yaml b/examples/robobrain2_5/conf/inference.yaml index 21f64ecf9e..520598a6d7 100644 --- a/examples/robobrain2_5/conf/inference.yaml +++ b/examples/robobrain2_5/conf/inference.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: 8b diff --git a/examples/robobrain2_5/conf/inference/32b.yaml b/examples/robobrain2_5/conf/inference/32b.yaml index 697676cf09..b82accdeca 100644 --- a/examples/robobrain2_5/conf/inference/32b.yaml +++ b/examples/robobrain2_5/conf/inference/32b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: # path to checkpoint model: /tmp/models/BAAI/RoboBrain2.0-7B diff --git a/examples/robobrain2_5/conf/inference/4b.yaml b/examples/robobrain2_5/conf/inference/4b.yaml index 5a0d7d1cd7..9a92912a73 100644 --- a/examples/robobrain2_5/conf/inference/4b.yaml +++ b/examples/robobrain2_5/conf/inference/4b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: # path to checkpoint model: /tmp/models/Qwen/Qwen3-VL-4B-Instruct diff --git a/examples/robobrain2_5/conf/inference/8b.yaml b/examples/robobrain2_5/conf/inference/8b.yaml index 7f109959ba..e6178e3d2e 100644 --- a/examples/robobrain2_5/conf/inference/8b.yaml +++ b/examples/robobrain2_5/conf/inference/8b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: # path to checkpoint model: /tmp/models/Qwen/Qwen3-VL-8B-Instruct diff --git a/examples/robobrain2_5/conf/serve.yaml b/examples/robobrain2_5/conf/serve.yaml index cbaaacdd21..628205f244 100644 --- a/examples/robobrain2_5/conf/serve.yaml +++ b/examples/robobrain2_5/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 8b diff --git a/examples/robobrain2_5/conf/serve/32b.yaml b/examples/robobrain2_5/conf/serve/32b.yaml index 941907b7dd..257e89181f 100644 --- a/examples/robobrain2_5/conf/serve/32b.yaml +++ b/examples/robobrain2_5/conf/serve/32b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /tmp/models/Qwen/Qwen3-VL-32B-Instruct # should be customized diff --git a/examples/robobrain2_5/conf/serve/4b.yaml b/examples/robobrain2_5/conf/serve/4b.yaml index c7d93b91d9..343947c815 100644 --- a/examples/robobrain2_5/conf/serve/4b.yaml +++ b/examples/robobrain2_5/conf/serve/4b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /tmp/models/Qwen/Qwen3-VL-4B-Instruct # should be customized diff --git a/examples/robobrain2_5/conf/serve/8b.yaml b/examples/robobrain2_5/conf/serve/8b.yaml index 92718cd631..562a89f7ce 100644 --- a/examples/robobrain2_5/conf/serve/8b.yaml +++ b/examples/robobrain2_5/conf/serve/8b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /tmp/models/Qwen/Qwen3-VL-8B-Instruct # should be customized diff --git a/examples/robobrain_x0/client_agilex.py b/examples/robobrain_x0/client_agilex.py index 89e4af313b..84d5b3d8ab 100644 --- a/examples/robobrain_x0/client_agilex.py +++ b/examples/robobrain_x0/client_agilex.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Adopted from FlagOpen/RoboBrain-X0 (https://github.com/FlagOpen/RoboBrain-X0/blob/main/agilex/client_agilex.py) import argparse import base64 diff --git a/examples/robobrain_x0/conf/serve.yaml b/examples/robobrain_x0/conf/serve.yaml index f93cef68ff..86b089b809 100644 --- a/examples/robobrain_x0/conf/serve.yaml +++ b/examples/robobrain_x0/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: robobrain_x0 diff --git a/examples/robobrain_x0/conf/serve/robobrain_x0.yaml b/examples/robobrain_x0/conf/serve/robobrain_x0.yaml index 95603101fa..b1230a6721 100644 --- a/examples/robobrain_x0/conf/serve/robobrain_x0.yaml +++ b/examples/robobrain_x0/conf/serve/robobrain_x0.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: /models/BAAI/RoboBrain-X0-Preview diff --git a/examples/robobrain_x0/conf/train.yaml b/examples/robobrain_x0/conf/train.yaml index b8ac208925..f7f9ec6274 100644 --- a/examples/robobrain_x0/conf/train.yaml +++ b/examples/robobrain_x0/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - train: robobrain_x0 - _self_ diff --git a/examples/robobrain_x0/conf/train/robobrain_x0.yaml b/examples/robobrain_x0/conf/train/robobrain_x0.yaml index 94127dde1f..b08500a26a 100644 --- a/examples/robobrain_x0/conf/train/robobrain_x0.yaml +++ b/examples/robobrain_x0/conf/train/robobrain_x0.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: vision_recompute_layer_steps: 16 vision_ration: 0.1 diff --git a/examples/robobrain_x0/conf/train/robobrain_x0_lerobot.yaml b/examples/robobrain_x0/conf/train/robobrain_x0_lerobot.yaml index 2d9306a1e4..7f0871e474 100644 --- a/examples/robobrain_x0/conf/train/robobrain_x0_lerobot.yaml +++ b/examples/robobrain_x0/conf/train/robobrain_x0_lerobot.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # RoboBrain-X0 training configuration with LeRobotDataset format # This configuration uses the LeRobot dataset format instead of WebDataset/Energon # Reference: https://github.com/huggingface/lerobot diff --git a/examples/robobrain_x0_5/client_libero.py b/examples/robobrain_x0_5/client_libero.py index 8b1bf07872..c438e52e24 100644 --- a/examples/robobrain_x0_5/client_libero.py +++ b/examples/robobrain_x0_5/client_libero.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Adopted from FlagOpen/RoboBrain-X0 (https://github.com/FlagOpen/RoboBrain-X0/blob/main/agilex/client_agilex.py) import argparse import base64 diff --git a/examples/robobrain_x0_5/conf/serve.yaml b/examples/robobrain_x0_5/conf/serve.yaml index a1762bae3a..9bba689115 100644 --- a/examples/robobrain_x0_5/conf/serve.yaml +++ b/examples/robobrain_x0_5/conf/serve.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: libero_qwengroot diff --git a/examples/robobrain_x0_5/conf/serve/libero_qwengroot.yaml b/examples/robobrain_x0_5/conf/serve/libero_qwengroot.yaml index a98000972b..5192dcedcb 100644 --- a/examples/robobrain_x0_5/conf/serve/libero_qwengroot.yaml +++ b/examples/robobrain_x0_5/conf/serve/libero_qwengroot.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine_args: model: outputs/libero_qwengroot/checkpoints/ckpt_out diff --git a/examples/robobrain_x0_5/conf/train.yaml b/examples/robobrain_x0_5/conf/train.yaml index fd50551e59..6f71ff4cac 100644 --- a/examples/robobrain_x0_5/conf/train.yaml +++ b/examples/robobrain_x0_5/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: libero_qwengroot diff --git a/examples/robobrain_x0_5/conf/train/libero_qwengroot.yaml b/examples/robobrain_x0_5/conf/train/libero_qwengroot.yaml index 4527dc8796..20d821a180 100644 --- a/examples/robobrain_x0_5/conf/train/libero_qwengroot.yaml +++ b/examples/robobrain_x0_5/conf/train/libero_qwengroot.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + seed: 42 trackers: [jsonl, wandb] wandb_entity: jinhuiye diff --git a/examples/rwkv/conf/train.yaml b/examples/rwkv/conf/train.yaml index 760eed172a..c4d5b1b19a 100644 --- a/examples/rwkv/conf/train.yaml +++ b/examples/rwkv/conf/train.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - train: 3b - _self_ diff --git a/examples/rwkv/conf/train/0_1b.yaml b/examples/rwkv/conf/train/0_1b.yaml index ba40537c52..1f89fd9a18 100644 --- a/examples/rwkv/conf/train/0_1b.yaml +++ b/examples/rwkv/conf/train/0_1b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: distributed_backend: nccl reset_position_ids: False diff --git a/examples/rwkv/conf/train/0_4b.yaml b/examples/rwkv/conf/train/0_4b.yaml index 458ed75efe..481e646510 100644 --- a/examples/rwkv/conf/train/0_4b.yaml +++ b/examples/rwkv/conf/train/0_4b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: distributed_backend: nccl reset_position_ids: False diff --git a/examples/rwkv/conf/train/1_5b.yaml b/examples/rwkv/conf/train/1_5b.yaml index a239c797c8..bdfac6487b 100644 --- a/examples/rwkv/conf/train/1_5b.yaml +++ b/examples/rwkv/conf/train/1_5b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: distributed_backend: nccl reset_position_ids: False diff --git a/examples/rwkv/conf/train/3b.yaml b/examples/rwkv/conf/train/3b.yaml index 701b253934..800eb69d3c 100644 --- a/examples/rwkv/conf/train/3b.yaml +++ b/examples/rwkv/conf/train/3b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: distributed_backend: nccl reset_position_ids: False diff --git a/examples/rwkv/conf/train/7b.yaml b/examples/rwkv/conf/train/7b.yaml index 965fe354bf..487db664d1 100644 --- a/examples/rwkv/conf/train/7b.yaml +++ b/examples/rwkv/conf/train/7b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: distributed_backend: nccl reset_position_ids: False diff --git a/examples/rwkv/tokenizer/__init__.py b/examples/rwkv/tokenizer/__init__.py index e69de29bb2..a491ed2726 100644 --- a/examples/rwkv/tokenizer/__init__.py +++ b/examples/rwkv/tokenizer/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/examples/wan2_1/conf/inference.yaml b/examples/wan2_1/conf/inference.yaml index 6b1782fdc6..e99b3db56b 100644 --- a/examples/wan2_1/conf/inference.yaml +++ b/examples/wan2_1/conf/inference.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: 1.3b diff --git a/examples/wan2_1/conf/inference/1.3b.yaml b/examples/wan2_1/conf/inference/1.3b.yaml index 12e6f78e43..4d92eb328c 100644 --- a/examples/wan2_1/conf/inference/1.3b.yaml +++ b/examples/wan2_1/conf/inference/1.3b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + engine: model: Wan-AI/Wan2.1-T2V-1.3B-Diffusers loader: diffusers diff --git a/examples/wan2_1/conf/inference/14b.yaml b/examples/wan2_1/conf/inference/14b.yaml index 20507595dc..15a5498069 100644 --- a/examples/wan2_1/conf/inference/14b.yaml +++ b/examples/wan2_1/conf/inference/14b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + engine: model: Wan-AI/Wan2.1-T2V-14B-Diffusers loader: diffusers diff --git a/flagscale/__init__.py b/flagscale/__init__.py index 85301627f7..325d90ad55 100644 --- a/flagscale/__init__.py +++ b/flagscale/__init__.py @@ -1,3 +1,18 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + def _get_version() -> str: """Get version from importlib.metadata or parse pyproject.toml as fallback.""" try: diff --git a/flagscale/agent/__init__.py b/flagscale/agent/__init__.py index e69de29bb2..a491ed2726 100644 --- a/flagscale/agent/__init__.py +++ b/flagscale/agent/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flagscale/agent/collaboration/__init__.py b/flagscale/agent/collaboration/__init__.py index 973b753435..14206a00f0 100644 --- a/flagscale/agent/collaboration/__init__.py +++ b/flagscale/agent/collaboration/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from .collaborator import Collaborator __all__ = ["Collaborator"] diff --git a/flagscale/agent/collaboration/collaborator.py b/flagscale/agent/collaboration/collaborator.py index 1bec67e9cc..c65a503ef3 100644 --- a/flagscale/agent/collaboration/collaborator.py +++ b/flagscale/agent/collaboration/collaborator.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import datetime import json import threading diff --git a/flagscale/agent/tool_match/__init__.py b/flagscale/agent/tool_match/__init__.py index b3a7cac389..6886942a7b 100644 --- a/flagscale/agent/tool_match/__init__.py +++ b/flagscale/agent/tool_match/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Tool matching module for intelligent tool selection""" from .tool_matcher import ToolMatcher diff --git a/flagscale/agent/tool_match/tool_matcher.py b/flagscale/agent/tool_match/tool_matcher.py index 150de34d61..0ef4768f65 100644 --- a/flagscale/agent/tool_match/tool_matcher.py +++ b/flagscale/agent/tool_match/tool_matcher.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Intelligent tool matcher using semantic embeddings""" import logging diff --git a/flagscale/agent/tool_match/tool_registry.py b/flagscale/agent/tool_match/tool_registry.py index cd913dc657..ec0d813513 100644 --- a/flagscale/agent/tool_match/tool_registry.py +++ b/flagscale/agent/tool_match/tool_registry.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Tool registry for managing available tools""" from typing import Any diff --git a/flagscale/cli.py b/flagscale/cli.py index eeaa2f925b..075e635816 100644 --- a/flagscale/cli.py +++ b/flagscale/cli.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import subprocess import sys from enum import Enum diff --git a/flagscale/compress/__init__.py b/flagscale/compress/__init__.py index e69de29bb2..a491ed2726 100644 --- a/flagscale/compress/__init__.py +++ b/flagscale/compress/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flagscale/compress/adapter.py b/flagscale/compress/adapter.py index 28f2a2d239..ce6a34bb0f 100644 --- a/flagscale/compress/adapter.py +++ b/flagscale/compress/adapter.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import torch from compressed_tensors.quantization import ( QuantizationConfig, diff --git a/flagscale/compress/algo/__init__.py b/flagscale/compress/algo/__init__.py index e69de29bb2..a491ed2726 100644 --- a/flagscale/compress/algo/__init__.py +++ b/flagscale/compress/algo/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flagscale/compress/algo/algo_base.py b/flagscale/compress/algo/algo_base.py index d21415df01..20c3c8547f 100644 --- a/flagscale/compress/algo/algo_base.py +++ b/flagscale/compress/algo/algo_base.py @@ -1,3 +1,18 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + class BaseALGO: def __init__(self, name): self.name = name diff --git a/flagscale/compress/combined_algo.py b/flagscale/compress/combined_algo.py index c72143b141..07724d7f61 100644 --- a/flagscale/compress/combined_algo.py +++ b/flagscale/compress/combined_algo.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import copy diff --git a/flagscale/compress/compressor.py b/flagscale/compress/compressor.py index a500e462d8..be72367b2d 100644 --- a/flagscale/compress/compressor.py +++ b/flagscale/compress/compressor.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import os import shutil diff --git a/flagscale/compress/compressor_emu3.py b/flagscale/compress/compressor_emu3.py index e921f487a9..b3fc999cfd 100644 --- a/flagscale/compress/compressor_emu3.py +++ b/flagscale/compress/compressor_emu3.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import random diff --git a/flagscale/compress/compressor_llava_ov.py b/flagscale/compress/compressor_llava_ov.py index e34ebdabcf..8b508ef9c4 100644 --- a/flagscale/compress/compressor_llava_ov.py +++ b/flagscale/compress/compressor_llava_ov.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import ast import copy diff --git a/flagscale/eval/__init__.py b/flagscale/eval/__init__.py index e69de29bb2..a491ed2726 100644 --- a/flagscale/eval/__init__.py +++ b/flagscale/eval/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flagscale/eval/robo.py b/flagscale/eval/robo.py index e649d8d168..cecf974a22 100644 --- a/flagscale/eval/robo.py +++ b/flagscale/eval/robo.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import hashlib import hmac diff --git a/flagscale/inference/__init__.py b/flagscale/inference/__init__.py index e69de29bb2..a491ed2726 100644 --- a/flagscale/inference/__init__.py +++ b/flagscale/inference/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flagscale/inference/arguments.py b/flagscale/inference/arguments.py index b838725f0e..e63e8f5f8a 100644 --- a/flagscale/inference/arguments.py +++ b/flagscale/inference/arguments.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse from omegaconf import OmegaConf diff --git a/flagscale/inference/core/__init__.py b/flagscale/inference/core/__init__.py index e69de29bb2..a491ed2726 100644 --- a/flagscale/inference/core/__init__.py +++ b/flagscale/inference/core/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flagscale/inference/core/diffusion/taylorseer_transformation.py b/flagscale/inference/core/diffusion/taylorseer_transformation.py index 012bbc7bb6..3885dd27e1 100644 --- a/flagscale/inference/core/diffusion/taylorseer_transformation.py +++ b/flagscale/inference/core/diffusion/taylorseer_transformation.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Modified from # https://github.com/vipshop/cache-dit/blob/v0.3.0/src/cache_dit/cache_factory/cache_contexts/taylorseer.py diff --git a/flagscale/inference/core/diffusion/timestep_tracker_transformation.py b/flagscale/inference/core/diffusion/timestep_tracker_transformation.py index 28e6a065c9..3bd6c7d418 100644 --- a/flagscale/inference/core/diffusion/timestep_tracker_transformation.py +++ b/flagscale/inference/core/diffusion/timestep_tracker_transformation.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import inspect import torch diff --git a/flagscale/inference/core/inference_diffusion.py b/flagscale/inference/core/inference_diffusion.py index c6eae09d1b..ea9c1e08e2 100644 --- a/flagscale/inference/core/inference_diffusion.py +++ b/flagscale/inference/core/inference_diffusion.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse from omegaconf import DictConfig, ListConfig, OmegaConf diff --git a/flagscale/inference/core/inference_engine.py b/flagscale/inference/core/inference_engine.py index a12aca91f7..cc49b51876 100644 --- a/flagscale/inference/core/inference_engine.py +++ b/flagscale/inference/core/inference_engine.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import dataclasses import importlib import os diff --git a/flagscale/inference/emu_utils/__init__.py b/flagscale/inference/emu_utils/__init__.py index c170221141..5920582364 100644 --- a/flagscale/inference/emu_utils/__init__.py +++ b/flagscale/inference/emu_utils/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from .processing_emu3 import Emu3Processor from .processing_emu3p5 import Emu3p5Processor diff --git a/flagscale/inference/emu_utils/processing_emu3.py b/flagscale/inference/emu_utils/processing_emu3.py index 7ce124bcb5..6759c56016 100644 --- a/flagscale/inference/emu_utils/processing_emu3.py +++ b/flagscale/inference/emu_utils/processing_emu3.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Copy from: # https://github.com/baaivision/Emu3/blob/main/emu3/mllm/processing_emu3.py # https://github.com/baaivision/Emu3/blob/main/emu3/mllm/utils_emu3.py diff --git a/flagscale/inference/emu_utils/processing_emu3p5.py b/flagscale/inference/emu_utils/processing_emu3p5.py index 1cf19c6134..d3abb615c6 100644 --- a/flagscale/inference/emu_utils/processing_emu3p5.py +++ b/flagscale/inference/emu_utils/processing_emu3p5.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os.path as osp import numpy as np diff --git a/flagscale/inference/emu_utils/prompt_case.py b/flagscale/inference/emu_utils/prompt_case.py index 3d16fb7b7b..c8271f9e15 100644 --- a/flagscale/inference/emu_utils/prompt_case.py +++ b/flagscale/inference/emu_utils/prompt_case.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + T2I_CASE = [ { "prompt": """A lively comic-style illustration depicting two humorous cartoon dogs interacting near a freshly dug backyard hole surrounded by scattered dirt, garden tools, blooming flowers, and a wooden fence background. At the upper-left side, Dog One stands nervously near the messy hole, ears down and eyes wide open with an expression of concern. Its speech bubble is an oval shape, outlined neatly with smooth, slightly rounded corners, positioned clearly above Dog One's head. Inside, clearly readable playful handwritten-style text emphasizes the dog's worried tone, saying, "You sure the humans won't notice this giant hole here?". Toward the lower-right side, Dog Two sits calmly and confidently with a cheerful, carefree expression, wagging its tail gently. Its speech bubble is rectangular with softly rounded edges, placed slightly overlapping with Dog One's speech bubble to guide the reader naturally downward diagonally across the frame. Dog Two's friendly, humorous response appears in a whimsical italicized comic font, clearly stating, "Relax! We'll just blame it on the neighbor's cat again!". Each speech bubble creates the playful and engaging backyard scene.""", diff --git a/flagscale/inference/inference_emu3.py b/flagscale/inference/inference_emu3.py index 66918a5924..6e1ea0aa53 100644 --- a/flagscale/inference/inference_emu3.py +++ b/flagscale/inference/inference_emu3.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import torch from PIL import Image from transformers import AutoImageProcessor, AutoModel, AutoTokenizer diff --git a/flagscale/inference/inference_emu3p5.py b/flagscale/inference/inference_emu3p5.py index da580d1c74..8a4681e052 100644 --- a/flagscale/inference/inference_emu3p5.py +++ b/flagscale/inference/inference_emu3p5.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os from omegaconf import DictConfig diff --git a/flagscale/inference/inference_llm.py b/flagscale/inference/inference_llm.py index 5b08631a98..d06900f468 100644 --- a/flagscale/inference/inference_llm.py +++ b/flagscale/inference/inference_llm.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from vllm import LLM from vllm.sampling_params import SamplingParams diff --git a/flagscale/inference/inference_pi.py b/flagscale/inference/inference_pi.py index 084bdc0d2f..a48bcea587 100644 --- a/flagscale/inference/inference_pi.py +++ b/flagscale/inference/inference_pi.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import json diff --git a/flagscale/inference/inference_qwen_gr00t.py b/flagscale/inference/inference_qwen_gr00t.py index 144ea4d1f2..0db5343ca3 100644 --- a/flagscale/inference/inference_qwen_gr00t.py +++ b/flagscale/inference/inference_qwen_gr00t.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import numpy as np diff --git a/flagscale/inference/inference_robobrain2.py b/flagscale/inference/inference_robobrain2.py index 039faa965c..60dc255196 100644 --- a/flagscale/inference/inference_robobrain2.py +++ b/flagscale/inference/inference_robobrain2.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from PIL import Image from transformers import AutoTokenizer diff --git a/flagscale/inference/runtime_context.py b/flagscale/inference/runtime_context.py index 5d13ee86c2..fde935b87b 100644 --- a/flagscale/inference/runtime_context.py +++ b/flagscale/inference/runtime_context.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import contextlib import contextvars import itertools diff --git a/flagscale/inference/utils.py b/flagscale/inference/utils.py index 05f1854100..f7d596c159 100644 --- a/flagscale/inference/utils.py +++ b/flagscale/inference/utils.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from typing import Any import torch diff --git a/flagscale/logger.py b/flagscale/logger.py index 61dcf26a1c..8aec57cc75 100644 --- a/flagscale/logger.py +++ b/flagscale/logger.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import logging import sys diff --git a/flagscale/models/__init__.py b/flagscale/models/__init__.py index e69de29bb2..a491ed2726 100644 --- a/flagscale/models/__init__.py +++ b/flagscale/models/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flagscale/models/configs/__init__.py b/flagscale/models/configs/__init__.py index e69de29bb2..a491ed2726 100644 --- a/flagscale/models/configs/__init__.py +++ b/flagscale/models/configs/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flagscale/models/megatron/__init__.py b/flagscale/models/megatron/__init__.py index e69de29bb2..a491ed2726 100644 --- a/flagscale/models/megatron/__init__.py +++ b/flagscale/models/megatron/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flagscale/models/megatron/deepseek_v4/deepseek_builder.py b/flagscale/models/megatron/deepseek_v4/deepseek_builder.py index b75d00f4cc..d05506926a 100644 --- a/flagscale/models/megatron/deepseek_v4/deepseek_builder.py +++ b/flagscale/models/megatron/deepseek_v4/deepseek_builder.py @@ -15,7 +15,7 @@ get_gpt_mtp_block_spec, get_gpt_layer_with_transformer_engine_spec, get_gpt_layer_local_spec, - get_gpt_layer_with_inference_spec + get_gpt_layer_with_inference_spec, ) from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add from megatron.core.transformer.spec_utils import ModuleSpec @@ -33,14 +33,15 @@ from megatron.core.transformer.enums import LayerType from megatron.training.utils import get_args from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( - get_transformer_block_with_experimental_attention_variant_spec, + get_transformer_block_with_experimental_attention_variant_spec, _get_backend_spec_provider, get_dsv4_hybrid_module_spec_for_backend, _get_moe_module_spec, - get_moe_layer_pattern + get_moe_layer_pattern, ) from megatron.core.transformer.hyper_connection import HyperConnectionModule from megatron.core.transformer.engram import EngramModule + try: import transformer_engine as te # pylint: disable=unused-import @@ -80,7 +81,6 @@ from .deepseek_model import DeepSeekModel - def get_deepseek_layer_spec( use_te: bool, config: TransformerConfig, @@ -106,7 +106,7 @@ def get_deepseek_layer_spec( else backend.layer_norm(rms_norm=rms_norm, for_qk=False) ) if build_engram: - engram_module = EngramModule + engram_module = EngramModule else: engram_module = None submodules = DeepSeekTransformerLayerSubmodules( @@ -118,7 +118,7 @@ def get_deepseek_layer_spec( mlp=moe_layer_spec, mlp_bda=get_bias_dropout_add, mlp_hyper_connection=HyperConnectionModule, - engram=ModuleSpec(module=engram_module) + engram=ModuleSpec(module=engram_module), ) return ModuleSpec(module=DeepSeekTransformerLayer, submodules=submodules) @@ -149,7 +149,6 @@ def get_deepseek_decoder_block_spec( build_engram=False, ) - # Create the layer specs for the model. layer_specs = [] for layer_number in range(config.num_layers): @@ -157,7 +156,9 @@ def get_deepseek_decoder_block_spec( is_engram_layer = True else: is_engram_layer = False - layer_specs.append(moe_deepseek_engram_layer_spec if is_engram_layer else moe_deepseek_layer_spec) + layer_specs.append( + moe_deepseek_engram_layer_spec if is_engram_layer else moe_deepseek_layer_spec + ) # Slice the layer specs to only include the layers that are built in this pipeline stage. # Note: MCore layer_number starts at 1 @@ -194,9 +195,11 @@ def get_deepseek_decoder_block_spec( return block_spec -def deepseek_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_collection=None): +def deepseek_builder( + args, pre_process, post_process, vp_stage=None, config=None, pg_collection=None +): """Drop-in replacement builder compatible with model_provider(...).""" - print_rank_0('building DeepSeek model (engram and mhc file) ...') + print_rank_0("building DeepSeek model (engram and mhc file) ...") if config is None: if args.yaml_cfg is not None: @@ -204,7 +207,6 @@ def deepseek_builder(args, pre_process, post_process, vp_stage=None, config=None else: config = core_transformer_config_from_args(args) - if args.use_legacy_models: raise NotImplementedError("Legacy GPT models do not support deepseek module insertion.") else: @@ -215,20 +217,24 @@ def deepseek_builder(args, pre_process, post_process, vp_stage=None, config=None if args.heterogeneous_layers_config_path is not None: assert not (config.transformer_impl == "inference_optimized") - raise NotImplementedError("Using heterogeneous layers is not supported with deepseek builder.") + raise NotImplementedError( + "Using heterogeneous layers is not supported with deepseek builder." + ) transformer_layer_spec = get_deepseek_decoder_block_spec( config=config, use_transformer_engine=use_te, normalization=args.normalization, qk_l2_norm=args.qk_l2_norm, vp_stage=vp_stage, - use_moe=True + use_moe=True, ) mtp_block_spec = None if args.mtp_num_layers is not None: assert not (config.transformer_impl == "inference_optimized") - transformer_layer_spec_for_mtp = get_deepseek_layer_spec(use_te, config, build_engram=False) + transformer_layer_spec_for_mtp = get_deepseek_layer_spec( + use_te, config, build_engram=False + ) mtp_block_spec = get_gpt_mtp_block_spec( config, transformer_layer_spec_for_mtp, diff --git a/flagscale/models/megatron/deepseek_v4/deepseek_model.py b/flagscale/models/megatron/deepseek_v4/deepseek_model.py index 2f3448c4e3..bd215f9b1a 100644 --- a/flagscale/models/megatron/deepseek_v4/deepseek_model.py +++ b/flagscale/models/megatron/deepseek_v4/deepseek_model.py @@ -1,237 +1,258 @@ -from typing import Optional - -import torch - -## megatron-core -from megatron.core.models.gpt import GPTModel -from megatron.core.packed_seq_params import PackedSeqParams -from megatron.core.utils import deprecate_inference_params -from megatron.core.inference.contexts import BaseInferenceContext -from megatron.core.transformer.engram import get_or_create_hash_mapping - - -class LazyHashInputIds: - """ - Lazy wrapper for hash input IDs that computes asynchronously and - synchronizes only when accessed. This allows hash computation to overlap - with preprocessing and early decoder layers. - """ - - def __init__(self, hash_mapping, input_ids, hash_stream=None): - self.hash_mapping = hash_mapping - self.input_ids = input_ids - self.hash_stream = hash_stream - self._result = None - self._is_async_pending = False - # Async - if self.hash_stream is not None: - # self.hash_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(self.hash_stream): - self._result = self.hash_mapping.hash(self.input_ids) - self._is_async_pending = True - # record result to use across stream - self._record_current_stream() - - def _record_current_stream(self): - """Helper to record current stream on all result tensors""" - if self._result is None: - return - current_stream = torch.cuda.current_stream() - if isinstance(self._result, dict): - for t in self._result.values(): - if isinstance(t, torch.Tensor): - t.record_stream(current_stream) - elif isinstance(self._result, torch.Tensor): - self._result.record_stream(current_stream) - - def __getitem__(self, key): - # Case 1: Async compute -> wait - if self._is_async_pending: - torch.cuda.current_stream().wait_stream(self.hash_stream) - self._is_async_pending = False # Async finish - self._record_current_stream() - - # Case 2: Sync but no compute -> start compute - elif self._result is None: - self._result = self.hash_mapping.hash(self.input_ids) - - # Case 3: Async or sync compute is finished. - # print(f"[rank{torch.distributed.get_rank()}]: LazyHashInputIds result = {self._result}") - return self._result[key] - - def get(self, key, default=None): - """Get hash result with default value.""" - try: - return self[key] - except KeyError: - return default - - -class DeepSeekModel(GPTModel): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - if self.config.use_engram: - self.engram_hash = get_or_create_hash_mapping( - engram_vocab_size=self.config.engram_vocab_size, - max_ngram_size=self.config.max_ngram_size, - n_embed_per_ngram=self.config.n_embed_per_ngram, - n_head_per_ngram=self.config.n_head_per_ngram, - layer_ids=self.config.engram_layer_ids, - tokenizer_name_or_path=self.config.engram_tokenizer_name_or_path, - pad_id=self.config.engram_pad_id, - seed=self.config.engram_seed, - ) - else: - self.engram_hash = None - - # Optional: Create a separate CUDA stream for hash computation - # This allows overlapping hash computation with preprocessing - self._hash_stream = None - if torch.cuda.is_available(): - self._hash_stream = torch.cuda.Stream() - - def forward( - self, - input_ids: torch.Tensor, - position_ids: torch.Tensor, - attention_mask: torch.Tensor, - decoder_input: torch.Tensor = None, - labels: torch.Tensor = None, - inference_context: BaseInferenceContext = None, - packed_seq_params: PackedSeqParams = None, - extra_block_kwargs: dict = None, - runtime_gather_output: bool | None = None, - *, - inference_params: BaseInferenceContext | None = None, - loss_mask: torch.Tensor | None = None, - padding_mask: Optional[torch.Tensor] = None, - is_spec_decode: Optional[bool] = None, - ) -> torch.Tensor: - assert input_ids is not None, "Input ids can not be None for EngramModel" - inference_context = deprecate_inference_params(inference_context, inference_params) - - # Create lazy hash input IDs that computes asynchronously - # The computation will start immediately in a separate stream (if available) - # but will only synchronize when actually accessed in the decoder - if self.config.use_engram: - engram_hash_input_ids = LazyHashInputIds( - hash_mapping=self.engram_hash, - input_ids=input_ids, - hash_stream=self._hash_stream, - ) - else: - engram_hash_input_ids = None - - # Preprocessing can run in parallel with hash computation - preproc_output = self._preprocess( - input_ids=input_ids, - position_ids=position_ids, - decoder_input=decoder_input, - inference_context=inference_context, - packed_seq_params=packed_seq_params, - padding_mask=padding_mask, - ) - - (decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset, padding_mask) = ( - preproc_output[:6] - ) - - rotary_pos_cos_sin = preproc_output[6] if len(preproc_output) == 7 else None - - # Pass input_ids to decoder for hash-based MoE routing - decoder_extra_block_kwargs = extra_block_kwargs or {} - if self.config.moe_n_hash_layers > 0 and input_ids is not None: - decoder_extra_block_kwargs['input_ids'] = input_ids - if self.config.use_engram: - decoder_extra_block_kwargs['engram_hash_input_ids'] = engram_hash_input_ids - - # Run decoder - decoder_output = self.decoder( - hidden_states=decoder_input, - attention_mask=attention_mask, - inference_context=inference_context, - rotary_pos_emb=rotary_pos_emb, - rotary_pos_cos=rotary_pos_cos, - rotary_pos_sin=rotary_pos_sin, - rotary_pos_cos_sin=rotary_pos_cos_sin, - packed_seq_params=packed_seq_params, - sequence_len_offset=sequence_len_offset, - **decoder_extra_block_kwargs, - ) - # torch.cuda.nvtx.range_pop() - # When mHC + MTP, the decoder returns (contracted, multi-stream). - # MTP needs multi-stream; lm_head needs contracted. - if isinstance(decoder_output, tuple): - hidden_states, mhc_multistream = decoder_output - else: - hidden_states = decoder_output - mhc_multistream = None - - return self._postprocess( - hidden_states=hidden_states, - input_ids=input_ids, - position_ids=position_ids, - labels=labels, - rotary_pos_emb=rotary_pos_emb, - rotary_pos_cos=rotary_pos_cos, - rotary_pos_sin=rotary_pos_sin, - mtp_in_postprocess=self.mtp_process, - loss_mask=loss_mask, - decoder_input=decoder_input, - attention_mask=attention_mask, - inference_params=inference_params, - packed_seq_params=packed_seq_params, - sequence_len_offset=sequence_len_offset, - runtime_gather_output=runtime_gather_output, - extra_block_kwargs=extra_block_kwargs, - inference_context=inference_context, - is_spec_decode=is_spec_decode, - mhc_multistream=mhc_multistream, - ) - - def build_schedule_plan( - self, - input_ids: torch.Tensor, - position_ids: torch.Tensor, - attention_mask: torch.Tensor, - decoder_input: torch.Tensor = None, - labels: torch.Tensor = None, - inference_context: BaseInferenceContext = None, - packed_seq_params: PackedSeqParams = None, - extra_block_kwargs: dict = None, - runtime_gather_output: Optional[bool] = None, - inference_params: Optional[BaseInferenceContext] = None, - loss_mask: Optional[torch.Tensor] = None, - ): - """ - Adaptation of overlap_moe_expert_parallel_comm. - """ - # Precompute the engram_hash_iput_ids, it will be used to create a TransformerChunkSchedulePlan. - engram_hash_input_ids = LazyHashInputIds( - hash_mapping=self.engram_hash, - input_ids=input_ids, - hash_stream=self._hash_stream, - ) - if extra_block_kwargs is None: - extra_block_kwargs = { - "engram_hash_input_ids": engram_hash_input_ids, - } - return super().build_schedule_plan( - input_ids, - position_ids, - attention_mask, - decoder_input, - labels=labels, - loss_mask=loss_mask, - extra_block_kwargs=extra_block_kwargs - ) - - def sharded_state_dict( - self, prefix: str = "", sharded_offsets: tuple = (), metadata: dict | None = None - ): - # Engram makes the layers non-homogeneous, so force the sharded-state-dict path - # to use layer-specific keys. - if metadata is None: - metadata = {} - metadata["non_homogeneous_layers"] = True - return super().sharded_state_dict(prefix=prefix, sharded_offsets=sharded_offsets, metadata=metadata) +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional + +import torch + +## megatron-core +from megatron.core.models.gpt import GPTModel +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.utils import deprecate_inference_params +from megatron.core.inference.contexts import BaseInferenceContext +from megatron.core.transformer.engram import get_or_create_hash_mapping + + +class LazyHashInputIds: + """ + Lazy wrapper for hash input IDs that computes asynchronously and + synchronizes only when accessed. This allows hash computation to overlap + with preprocessing and early decoder layers. + """ + + def __init__(self, hash_mapping, input_ids, hash_stream=None): + self.hash_mapping = hash_mapping + self.input_ids = input_ids + self.hash_stream = hash_stream + self._result = None + self._is_async_pending = False + # Async + if self.hash_stream is not None: + # self.hash_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self.hash_stream): + self._result = self.hash_mapping.hash(self.input_ids) + self._is_async_pending = True + # record result to use across stream + self._record_current_stream() + + def _record_current_stream(self): + """Helper to record current stream on all result tensors""" + if self._result is None: + return + current_stream = torch.cuda.current_stream() + if isinstance(self._result, dict): + for t in self._result.values(): + if isinstance(t, torch.Tensor): + t.record_stream(current_stream) + elif isinstance(self._result, torch.Tensor): + self._result.record_stream(current_stream) + + def __getitem__(self, key): + # Case 1: Async compute -> wait + if self._is_async_pending: + torch.cuda.current_stream().wait_stream(self.hash_stream) + self._is_async_pending = False # Async finish + self._record_current_stream() + + # Case 2: Sync but no compute -> start compute + elif self._result is None: + self._result = self.hash_mapping.hash(self.input_ids) + + # Case 3: Async or sync compute is finished. + # print(f"[rank{torch.distributed.get_rank()}]: LazyHashInputIds result = {self._result}") + return self._result[key] + + def get(self, key, default=None): + """Get hash result with default value.""" + try: + return self[key] + except KeyError: + return default + + +class DeepSeekModel(GPTModel): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if self.config.use_engram: + self.engram_hash = get_or_create_hash_mapping( + engram_vocab_size=self.config.engram_vocab_size, + max_ngram_size=self.config.max_ngram_size, + n_embed_per_ngram=self.config.n_embed_per_ngram, + n_head_per_ngram=self.config.n_head_per_ngram, + layer_ids=self.config.engram_layer_ids, + tokenizer_name_or_path=self.config.engram_tokenizer_name_or_path, + pad_id=self.config.engram_pad_id, + seed=self.config.engram_seed, + ) + else: + self.engram_hash = None + + # Optional: Create a separate CUDA stream for hash computation + # This allows overlapping hash computation with preprocessing + self._hash_stream = None + if torch.cuda.is_available(): + self._hash_stream = torch.cuda.Stream() + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: torch.Tensor, + decoder_input: torch.Tensor = None, + labels: torch.Tensor = None, + inference_context: BaseInferenceContext = None, + packed_seq_params: PackedSeqParams = None, + extra_block_kwargs: dict = None, + runtime_gather_output: bool | None = None, + *, + inference_params: BaseInferenceContext | None = None, + loss_mask: torch.Tensor | None = None, + padding_mask: Optional[torch.Tensor] = None, + is_spec_decode: Optional[bool] = None, + ) -> torch.Tensor: + assert input_ids is not None, "Input ids can not be None for EngramModel" + inference_context = deprecate_inference_params(inference_context, inference_params) + + # Create lazy hash input IDs that computes asynchronously + # The computation will start immediately in a separate stream (if available) + # but will only synchronize when actually accessed in the decoder + if self.config.use_engram: + engram_hash_input_ids = LazyHashInputIds( + hash_mapping=self.engram_hash, + input_ids=input_ids, + hash_stream=self._hash_stream, + ) + else: + engram_hash_input_ids = None + + # Preprocessing can run in parallel with hash computation + preproc_output = self._preprocess( + input_ids=input_ids, + position_ids=position_ids, + decoder_input=decoder_input, + inference_context=inference_context, + packed_seq_params=packed_seq_params, + padding_mask=padding_mask, + ) + + ( + decoder_input, + rotary_pos_emb, + rotary_pos_cos, + rotary_pos_sin, + sequence_len_offset, + padding_mask, + ) = preproc_output[:6] + + rotary_pos_cos_sin = preproc_output[6] if len(preproc_output) == 7 else None + + # Pass input_ids to decoder for hash-based MoE routing + decoder_extra_block_kwargs = extra_block_kwargs or {} + if self.config.moe_n_hash_layers > 0 and input_ids is not None: + decoder_extra_block_kwargs["input_ids"] = input_ids + if self.config.use_engram: + decoder_extra_block_kwargs["engram_hash_input_ids"] = engram_hash_input_ids + + # Run decoder + decoder_output = self.decoder( + hidden_states=decoder_input, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + rotary_pos_cos_sin=rotary_pos_cos_sin, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + **decoder_extra_block_kwargs, + ) + # torch.cuda.nvtx.range_pop() + # When mHC + MTP, the decoder returns (contracted, multi-stream). + # MTP needs multi-stream; lm_head needs contracted. + if isinstance(decoder_output, tuple): + hidden_states, mhc_multistream = decoder_output + else: + hidden_states = decoder_output + mhc_multistream = None + + return self._postprocess( + hidden_states=hidden_states, + input_ids=input_ids, + position_ids=position_ids, + labels=labels, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + mtp_in_postprocess=self.mtp_process, + loss_mask=loss_mask, + decoder_input=decoder_input, + attention_mask=attention_mask, + inference_params=inference_params, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + runtime_gather_output=runtime_gather_output, + extra_block_kwargs=extra_block_kwargs, + inference_context=inference_context, + is_spec_decode=is_spec_decode, + mhc_multistream=mhc_multistream, + ) + + def build_schedule_plan( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: torch.Tensor, + decoder_input: torch.Tensor = None, + labels: torch.Tensor = None, + inference_context: BaseInferenceContext = None, + packed_seq_params: PackedSeqParams = None, + extra_block_kwargs: dict = None, + runtime_gather_output: Optional[bool] = None, + inference_params: Optional[BaseInferenceContext] = None, + loss_mask: Optional[torch.Tensor] = None, + ): + """ + Adaptation of overlap_moe_expert_parallel_comm. + """ + # Precompute the engram_hash_iput_ids, it will be used to create a TransformerChunkSchedulePlan. + engram_hash_input_ids = LazyHashInputIds( + hash_mapping=self.engram_hash, + input_ids=input_ids, + hash_stream=self._hash_stream, + ) + if extra_block_kwargs is None: + extra_block_kwargs = { + "engram_hash_input_ids": engram_hash_input_ids, + } + return super().build_schedule_plan( + input_ids, + position_ids, + attention_mask, + decoder_input, + labels=labels, + loss_mask=loss_mask, + extra_block_kwargs=extra_block_kwargs, + ) + + def sharded_state_dict( + self, prefix: str = "", sharded_offsets: tuple = (), metadata: dict | None = None + ): + # Engram makes the layers non-homogeneous, so force the sharded-state-dict path + # to use layer-specific keys. + if metadata is None: + metadata = {} + metadata["non_homogeneous_layers"] = True + return super().sharded_state_dict( + prefix=prefix, sharded_offsets=sharded_offsets, metadata=metadata + ) diff --git a/flagscale/models/megatron/deepseek_v4/deepseek_transformer_block.py b/flagscale/models/megatron/deepseek_v4/deepseek_transformer_block.py index eb6ac204b7..a2f52bf1ee 100644 --- a/flagscale/models/megatron/deepseek_v4/deepseek_transformer_block.py +++ b/flagscale/models/megatron/deepseek_v4/deepseek_transformer_block.py @@ -1,375 +1,399 @@ -## built-in -from contextlib import nullcontext -from typing import TYPE_CHECKING, List, Optional, Set, Tuple, Union, cast - -import torch -from torch import Tensor - -from megatron.core import tensor_parallel -from megatron.core.enums import Fp8Recipe -from megatron.core.fp4_utils import get_fp4_context -from megatron.core.fp8_utils import get_fp8_context -from megatron.core.inference.contexts import BaseInferenceContext -from megatron.core.packed_seq_params import PackedSeqParams -if TYPE_CHECKING: - from megatron.core.tensor_parallel.random import CheckpointManager -else: - CheckpointManager = None - -# megatron-core -from megatron.core.transformer.transformer_block import TransformerBlock, apply_module, get_transformer_layer_offset -from megatron.core.utils import ( - WrappedTensor, - deprecate_inference_params, - get_pg_rank, - make_viewless_tensor, -) -from .hyper_connection import HyperConnectionModule - - -class DeepSeekTransformerBlock(TransformerBlock): - ## TODO: Finish at next version. These codes are copied from Megatron-LM dev branch, megatron/core/transformer/transformer_block.py - def _build_mhc_recompute_layer_plan( - self, use_mhc_recompute: bool - ) -> Tuple[List[Optional[CheckpointManager]], List[bool]]: - """Pre-build per-layer MHC recompute managers and block-end markers.""" - num_layers = len(self.layers) - layer_managers: List[Optional[CheckpointManager]] = [None] * num_layers - is_recompute_block_end: List[bool] = [False] * num_layers - - if not use_mhc_recompute or num_layers == 0: - return layer_managers, is_recompute_block_end - - mhc_recompute_layer_num = self.config.mhc_recompute_layer_num - mhc_manager = CheckpointManager() - - for l_no in range(num_layers): - is_last_in_transformer_block = l_no == num_layers - 1 - is_last_in_recompute_block = is_last_in_transformer_block - if mhc_recompute_layer_num is not None: - is_last_in_recompute_block = is_last_in_transformer_block or ( - (l_no + 1) % mhc_recompute_layer_num == 0 - ) - - layer_managers[l_no] = mhc_manager - is_recompute_block_end[l_no] = is_last_in_recompute_block - - if is_last_in_recompute_block and not is_last_in_transformer_block: - mhc_manager = CheckpointManager() - - return layer_managers, is_recompute_block_end - - @staticmethod - def _finalize_mhc_recompute_layer(mhc_manager, hidden_states, is_last_in_recompute_block): - """ - This function is only implemented in megatron dev branch, not any release version. - And if not using mhc recompute, this is no needed. So we just keep the function signature and leave it empty for now. - """ - pass - # @staticmethod - # def _finalize_mhc_recompute_layer( - # mhc_manager: Optional[CheckpointManager], - # hidden_states: Tensor, - # is_last_in_recompute_block: bool, - # ) -> None: - # """Finalize MHC recompute state for the current layer when block ends.""" - # if mhc_manager is not None and is_last_in_recompute_block: - # mhc_manager.discard_all_outputs_and_register_unified_recompute(hidden_states) - - def forward( - self, - hidden_states: Union[Tensor, WrappedTensor], - attention_mask: Optional[Tensor], - context: Optional[Tensor] = None, - context_mask: Optional[Tensor] = None, - rotary_pos_emb: Optional[Tensor] = None, - rotary_pos_cos: Optional[Tensor] = None, - rotary_pos_sin: Optional[Tensor] = None, - rotary_pos_cos_sin: Optional[Tensor] = None, - attention_bias: Optional[Tensor] = None, - inference_context: Optional[BaseInferenceContext] = None, - packed_seq_params: Optional[PackedSeqParams] = None, - sequence_len_offset: Optional[Tensor] = None, - padding_mask: Optional[Tensor] = None, - extract_layer_indices: Optional[Set[int]] = None, - *, - inference_params: Optional[BaseInferenceContext] = None, - dynamic_inference_decode_only: Optional[bool] = None, - input_ids: Optional[Tensor] = None, # Used for hash-based MoE Router - engram_hash_input_ids: Optional[Tensor] = None, # This is the input ids used for hashing in engram, which can be different from the input_ids used for attention mask and positional embedding. Only used when engram is enabled and at least one of the layers in this block has engram hashing. - ): - """ - Perform the forward pass through the transformer block. - - This method handles the core computation of the transformer, including - self-attention, optional cross-attention, and feed-forward operations. - - Args: - hidden_states (Union[Tensor, WrappedTensor]): Input tensor of shape [s, b, h] - where s is the sequence length, b is the batch size, and h is the hidden size. - Can be passed as a WrappedTensor during inference to avoid an obsolete - reference in the calling function. - attention_mask (Tensor): Boolean tensor of shape [1, 1, s, s] for masking - self-attention. - context (Tensor, optional): Context tensor for cross-attention. - context_mask (Tensor, optional): Mask for cross-attention context - rotary_pos_emb (Tensor, optional): Rotary positional embeddings. - rotary_pos_cos (Optional[Tensor]): Rotary embedding cosine. - rotary_pos_sin (Optional[Tensor]): Rotary embedding sine. - rotary_pos_cos_sin (Optional[Tensor]): Combined rotary embedding cosine and sine. - Currently used exclusively for inference with dynamic batching and flashinfer RoPE. - attention_bias (Tensor): Bias tensor for Q * K.T of shape in shape broadcastable - to [b, num_head, sq, skv], e.g. [1, 1, sq, skv]. - Used as an alternative to apply attention mask for TE cuDNN attention. - inference_context (BaseInferenceContext, optional): Parameters for inference-time - optimizations. - packed_seq_params (PackedSeqParams, optional): Parameters for packed sequence - processing. - extract_layer_indices (Set[int], optional): A set of global - layer indices (0-based across all pipeline stages) from - which to extract intermediate hidden states. If - non-empty, the forward pass will collect hidden_states - after each specified layer. - dynamic_inference_decode_only: Optional[bool]: If true, indicates that the current - inference context is for decode-only. This args is only used to uniquely - identify decode and non-decode cuda graph runners in the cuda graph manager. - - Returns: - Union[Tensor, Tuple[Tensor, List[Tensor]]]: - - If extract_layer_indices is None or empty: Returns the output hidden states tensor - of shape [s, b, h]. - - If extract_layer_indices is non-empty: Returns a tuple - of (hidden_states, intermediate_hidden_states) where - intermediate_hidden_states is a list of tensors - corresponding to hidden states after each layer in - extract_layer_indices. - """ - - ########## FlagScale Begin ########## - # for refined recompute - self.current_microbatch = -1 - if len(self.layers) > 0: - if hasattr(self.layers[0], 'current_microbatch'): - self.current_microbatch = self.layers[0].current_microbatch - saved_recompute_granularity = self.config.recompute_granularity - ########## FlagScale End ########## - - inference_context = deprecate_inference_params(inference_context, inference_params) - # Remove 'dynamic_inference_decode_only' from kwargs if present - # this is only used to uniquely identify decode and non-decode cuda graph - # runners in the cuda graph manager - - # Initialize feature collection (consistent with FastGen's Wan implementation) - if extract_layer_indices is None: - extract_layer_indices = set() - intermediate_hidden_states: List[Tensor] = [] - - # Calculate the global layer offset for this pipeline stage - # This is needed to convert local layer indices to global indices for feature extraction - pp_group = self.pg_collection.pp if hasattr(self.pg_collection, 'pp') else None - layer_offset = get_transformer_layer_offset( - self.config, self.vp_stage, get_pg_rank(pp_group) - ) - - # Delete the obsolete reference to the initial input tensor if necessary - if isinstance(hidden_states, WrappedTensor): - hidden_states = hidden_states.unwrap() - - if not self.pre_process: - # See set_input_tensor() - hidden_states = self.input_tensor - - # Viewless tensor. - # - We only need to create a viewless tensor in the case of micro batch - # size (mbs) == 1, since in this case, 'hidden_states.transpose()' - # above creates a view tensor, and '.contiguous()' is a pass-through. - # For mbs >= 2, '.contiguous()' creates a new tensor, eliminating - # the need to make it viewless. - # - # However, we don't explicitly check mbs == 1 here because - # make_viewless_tensor() has negligible overhead when its input - # is already viewless. - # - # - For the 'else' case above, calling make_viewless_tensor() here is - # likely redundant, since p2p_communication.py (likely originator) - # already creates viewless tensors. That said, make_viewless_tensor() - # is called here to be future-proof and corner-case-proof. - hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) - - # Expand hidden states for hyper connections at the start of the block - # Only expand at the first PP stage; subsequent stages receive n-stream from previous stage - if self.config.enable_hyper_connections and self.pre_process: - hidden_states = HyperConnectionModule.input_expand( - hidden_states, self.config.num_residual_streams - ) # [s, b, C] -> [s, b, n*C] - - if self.config.sequence_parallel: - rng_context = tensor_parallel.get_cuda_rng_tracker().fork() - else: - rng_context = nullcontext() - - # If fp8_recipe is delayed, wrap the entire pass with get_fp8_context(), - # otherwise do nothing extra at the outer level - # if we are using other fp8 recipes, then the context manager enter&exit are free - # we can wrap fp8_context within the for loop over layers, so that we can fine-grained - # control which layer will be fp8 or bf16 - # For FP4: NVFP4BlockScaling doesn't have delayed scaling, always uses inner context - if self.config.fp8: - use_outer_quantization_context = self.config.fp8_recipe == Fp8Recipe.delayed - use_inner_quantization_context = self.config.fp8_recipe != Fp8Recipe.delayed - outer_quantization_context = ( - get_fp8_context(self.config) if use_outer_quantization_context else nullcontext() - ) - elif self.config.fp4: - use_outer_quantization_context = False - use_inner_quantization_context = True - outer_quantization_context = nullcontext() - else: - # No quantization - use_outer_quantization_context = False - use_inner_quantization_context = False - outer_quantization_context = nullcontext() - - # Determine if MHC recompute should be used - # Only enable when: training mode AND hyper connections AND 'mhc' in recompute_modules - use_mhc_recompute = ( - self.training - and self.config.enable_hyper_connections - and self.config.recompute_granularity == 'selective' - and "mhc" in self.config.recompute_modules - ) - mhc_layer_managers, mhc_is_last_in_recompute_block = self._build_mhc_recompute_layer_plan( - use_mhc_recompute - ) - - with rng_context, outer_quantization_context: - # Forward pass. - if self.config.recompute_granularity == 'full' and self.training: - checkpointed_result = self._checkpointed_forward( - hidden_states=hidden_states, - attention_mask=attention_mask, - context=context, - context_mask=context_mask, - rotary_pos_emb=rotary_pos_emb, - attention_bias=attention_bias, - packed_seq_params=packed_seq_params, - use_inner_quantization_context=use_inner_quantization_context, - padding_mask=padding_mask, - extract_layer_indices=extract_layer_indices, - layer_offset=layer_offset, - input_ids=input_ids - ) - # Handle return value from _checkpointed_forward - if len(extract_layer_indices) > 0: - # (hidden_states, intermediate_hidden_states) tuple - hidden_states, intermediate_hidden_states = checkpointed_result - else: - # No intermediate_hidden_states requested: just hidden_states - hidden_states = checkpointed_result - else: - for l_no, layer in enumerate(self.layers): - # Get appropriate inner quantization context - if use_inner_quantization_context: - if self.config.fp8: - inner_quantization_context = get_fp8_context( - self.config, layer.layer_number - 1 - ) - elif self.config.fp4: - inner_quantization_context = get_fp4_context( - self.config, layer.layer_number - 1 - ) - else: - inner_quantization_context = nullcontext() - else: - inner_quantization_context = nullcontext() - - mhc_manager = mhc_layer_managers[l_no] - if mhc_manager is not None: - mhc_manager.is_last_layer_in_recompute_block = ( - mhc_is_last_in_recompute_block[l_no] - ) - - with self.offload_context, inner_quantization_context: - #### FlagScale Begin #### - # Pre-compute embeddings for the next DeepSeekTransformerLayer if engram exists, to overlap with current layer's computation - if l_no < len(self.layers) - 1: - next_layer = self.layers[l_no + 1] - if getattr(next_layer, "is_engram_layer", False): - next_layer.pre_compute_embedding(engram_hash_input_ids) - #### FlagScale End #### - hidden_states, context = layer( - hidden_states=hidden_states, - attention_mask=attention_mask, - context=context, - context_mask=context_mask, - rotary_pos_emb=rotary_pos_emb, - rotary_pos_cos=rotary_pos_cos, - rotary_pos_sin=rotary_pos_sin, - rotary_pos_cos_sin=rotary_pos_cos_sin, - attention_bias=attention_bias, - inference_context=inference_context, - packed_seq_params=packed_seq_params, - sequence_len_offset=sequence_len_offset, - padding_mask=padding_mask, - mhc_recompute_manager=mhc_manager, - input_ids=input_ids - ) - - self._finalize_mhc_recompute_layer( - mhc_manager=mhc_manager, - hidden_states=hidden_states, - is_last_in_recompute_block=mhc_is_last_in_recompute_block[l_no], - ) - - if ( - torch.is_grad_enabled() - and self.config.cpu_offloading - and self.group_prefetch_offload_commit_async is not None - ): - hidden_states = self.group_prefetch_offload_commit_async(hidden_states) - - # Extract intermediate embeddings using global layer index - if (l_no + layer_offset) in extract_layer_indices: - intermediate_hidden_states.append(hidden_states) - - # Only contract if the final layer norm is in this stage - if self.config.enable_hyper_connections and self.has_final_layernorm_in_this_stage(): - hidden_states = HyperConnectionModule.output_contract( - hidden_states, self.config.num_residual_streams - ) # [s, b, n*C] -> [s, b, C] - - # Final layer norm. - if self.final_layernorm is not None: - hidden_states = apply_module(self.final_layernorm)(cast(Tensor, hidden_states)) - # TENorm produces a "viewed" tensor. This will result in schedule.py's - # deallocate_output_tensor() throwing an error, so a viewless tensor is - # created to prevent this. - hidden_states = make_viewless_tensor( - inp=hidden_states, requires_grad=True, keep_graph=True - ) - - # If this TransformerBlock is empty, input and output hidden states will be the same node - # on the computational graph and will lead to unexpected errors in pipeline schedules. - if not self.pre_process and len(self.layers) == 0 and not self.final_layernorm: - hidden_states = hidden_states.clone() - ########## FlagScale Begin ########## - self.config.recompute_granularity = saved_recompute_granularity - ########## FlagScale End ########## - - if len(extract_layer_indices) > 0: - return hidden_states, intermediate_hidden_states - - return hidden_states - - def sharded_state_dict( - self, prefix: str = "", sharded_offsets: tuple = (), metadata: dict | None = None - ): - # Engram let the layers be non-homogeneous, so we need to set the flag in metadata to let the sharded state dict logic know. - # This is usefule when all layer are same, the TransformerBlock will be homogeneous, it generate sharded_state_dict will same keys for all layer and need all layers have the same structure. - # The layer has engram module does not fit this assumption. - # If the flag is set to True, the sharded_state_dict will use layer_number to generate different keys for different layer, which is same to models has dense layer leading and moe layer following. - # Actually, engram really causes the layers to be non-homogeneous. - if metadata is None: - metadata = {} - metadata["non_homogeneous_layers"] = True - return super().sharded_state_dict(prefix=prefix, sharded_offsets=sharded_offsets, metadata=metadata) +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +## built-in +from contextlib import nullcontext +from typing import TYPE_CHECKING, List, Optional, Set, Tuple, Union, cast + +import torch +from torch import Tensor + +from megatron.core import tensor_parallel +from megatron.core.enums import Fp8Recipe +from megatron.core.fp4_utils import get_fp4_context +from megatron.core.fp8_utils import get_fp8_context +from megatron.core.inference.contexts import BaseInferenceContext +from megatron.core.packed_seq_params import PackedSeqParams + +if TYPE_CHECKING: + from megatron.core.tensor_parallel.random import CheckpointManager +else: + CheckpointManager = None + +# megatron-core +from megatron.core.transformer.transformer_block import ( + TransformerBlock, + apply_module, + get_transformer_layer_offset, +) +from megatron.core.utils import ( + WrappedTensor, + deprecate_inference_params, + get_pg_rank, + make_viewless_tensor, +) +from .hyper_connection import HyperConnectionModule + + +class DeepSeekTransformerBlock(TransformerBlock): + ## TODO: Finish at next version. These codes are copied from Megatron-LM dev branch, megatron/core/transformer/transformer_block.py + def _build_mhc_recompute_layer_plan( + self, use_mhc_recompute: bool + ) -> Tuple[List[Optional[CheckpointManager]], List[bool]]: + """Pre-build per-layer MHC recompute managers and block-end markers.""" + num_layers = len(self.layers) + layer_managers: List[Optional[CheckpointManager]] = [None] * num_layers + is_recompute_block_end: List[bool] = [False] * num_layers + + if not use_mhc_recompute or num_layers == 0: + return layer_managers, is_recompute_block_end + + mhc_recompute_layer_num = self.config.mhc_recompute_layer_num + mhc_manager = CheckpointManager() + + for l_no in range(num_layers): + is_last_in_transformer_block = l_no == num_layers - 1 + is_last_in_recompute_block = is_last_in_transformer_block + if mhc_recompute_layer_num is not None: + is_last_in_recompute_block = is_last_in_transformer_block or ( + (l_no + 1) % mhc_recompute_layer_num == 0 + ) + + layer_managers[l_no] = mhc_manager + is_recompute_block_end[l_no] = is_last_in_recompute_block + + if is_last_in_recompute_block and not is_last_in_transformer_block: + mhc_manager = CheckpointManager() + + return layer_managers, is_recompute_block_end + + @staticmethod + def _finalize_mhc_recompute_layer(mhc_manager, hidden_states, is_last_in_recompute_block): + """ + This function is only implemented in megatron dev branch, not any release version. + And if not using mhc recompute, this is no needed. So we just keep the function signature and leave it empty for now. + """ + pass + + # @staticmethod + # def _finalize_mhc_recompute_layer( + # mhc_manager: Optional[CheckpointManager], + # hidden_states: Tensor, + # is_last_in_recompute_block: bool, + # ) -> None: + # """Finalize MHC recompute state for the current layer when block ends.""" + # if mhc_manager is not None and is_last_in_recompute_block: + # mhc_manager.discard_all_outputs_and_register_unified_recompute(hidden_states) + + def forward( + self, + hidden_states: Union[Tensor, WrappedTensor], + attention_mask: Optional[Tensor], + context: Optional[Tensor] = None, + context_mask: Optional[Tensor] = None, + rotary_pos_emb: Optional[Tensor] = None, + rotary_pos_cos: Optional[Tensor] = None, + rotary_pos_sin: Optional[Tensor] = None, + rotary_pos_cos_sin: Optional[Tensor] = None, + attention_bias: Optional[Tensor] = None, + inference_context: Optional[BaseInferenceContext] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + sequence_len_offset: Optional[Tensor] = None, + padding_mask: Optional[Tensor] = None, + extract_layer_indices: Optional[Set[int]] = None, + *, + inference_params: Optional[BaseInferenceContext] = None, + dynamic_inference_decode_only: Optional[bool] = None, + input_ids: Optional[Tensor] = None, # Used for hash-based MoE Router + engram_hash_input_ids: Optional[ + Tensor + ] = None, # This is the input ids used for hashing in engram, which can be different from the input_ids used for attention mask and positional embedding. Only used when engram is enabled and at least one of the layers in this block has engram hashing. + ): + """ + Perform the forward pass through the transformer block. + + This method handles the core computation of the transformer, including + self-attention, optional cross-attention, and feed-forward operations. + + Args: + hidden_states (Union[Tensor, WrappedTensor]): Input tensor of shape [s, b, h] + where s is the sequence length, b is the batch size, and h is the hidden size. + Can be passed as a WrappedTensor during inference to avoid an obsolete + reference in the calling function. + attention_mask (Tensor): Boolean tensor of shape [1, 1, s, s] for masking + self-attention. + context (Tensor, optional): Context tensor for cross-attention. + context_mask (Tensor, optional): Mask for cross-attention context + rotary_pos_emb (Tensor, optional): Rotary positional embeddings. + rotary_pos_cos (Optional[Tensor]): Rotary embedding cosine. + rotary_pos_sin (Optional[Tensor]): Rotary embedding sine. + rotary_pos_cos_sin (Optional[Tensor]): Combined rotary embedding cosine and sine. + Currently used exclusively for inference with dynamic batching and flashinfer RoPE. + attention_bias (Tensor): Bias tensor for Q * K.T of shape in shape broadcastable + to [b, num_head, sq, skv], e.g. [1, 1, sq, skv]. + Used as an alternative to apply attention mask for TE cuDNN attention. + inference_context (BaseInferenceContext, optional): Parameters for inference-time + optimizations. + packed_seq_params (PackedSeqParams, optional): Parameters for packed sequence + processing. + extract_layer_indices (Set[int], optional): A set of global + layer indices (0-based across all pipeline stages) from + which to extract intermediate hidden states. If + non-empty, the forward pass will collect hidden_states + after each specified layer. + dynamic_inference_decode_only: Optional[bool]: If true, indicates that the current + inference context is for decode-only. This args is only used to uniquely + identify decode and non-decode cuda graph runners in the cuda graph manager. + + Returns: + Union[Tensor, Tuple[Tensor, List[Tensor]]]: + - If extract_layer_indices is None or empty: Returns the output hidden states tensor + of shape [s, b, h]. + - If extract_layer_indices is non-empty: Returns a tuple + of (hidden_states, intermediate_hidden_states) where + intermediate_hidden_states is a list of tensors + corresponding to hidden states after each layer in + extract_layer_indices. + """ + + ########## FlagScale Begin ########## + # for refined recompute + self.current_microbatch = -1 + if len(self.layers) > 0: + if hasattr(self.layers[0], "current_microbatch"): + self.current_microbatch = self.layers[0].current_microbatch + saved_recompute_granularity = self.config.recompute_granularity + ########## FlagScale End ########## + + inference_context = deprecate_inference_params(inference_context, inference_params) + # Remove 'dynamic_inference_decode_only' from kwargs if present + # this is only used to uniquely identify decode and non-decode cuda graph + # runners in the cuda graph manager + + # Initialize feature collection (consistent with FastGen's Wan implementation) + if extract_layer_indices is None: + extract_layer_indices = set() + intermediate_hidden_states: List[Tensor] = [] + + # Calculate the global layer offset for this pipeline stage + # This is needed to convert local layer indices to global indices for feature extraction + pp_group = self.pg_collection.pp if hasattr(self.pg_collection, "pp") else None + layer_offset = get_transformer_layer_offset( + self.config, self.vp_stage, get_pg_rank(pp_group) + ) + + # Delete the obsolete reference to the initial input tensor if necessary + if isinstance(hidden_states, WrappedTensor): + hidden_states = hidden_states.unwrap() + + if not self.pre_process: + # See set_input_tensor() + hidden_states = self.input_tensor + + # Viewless tensor. + # - We only need to create a viewless tensor in the case of micro batch + # size (mbs) == 1, since in this case, 'hidden_states.transpose()' + # above creates a view tensor, and '.contiguous()' is a pass-through. + # For mbs >= 2, '.contiguous()' creates a new tensor, eliminating + # the need to make it viewless. + # + # However, we don't explicitly check mbs == 1 here because + # make_viewless_tensor() has negligible overhead when its input + # is already viewless. + # + # - For the 'else' case above, calling make_viewless_tensor() here is + # likely redundant, since p2p_communication.py (likely originator) + # already creates viewless tensors. That said, make_viewless_tensor() + # is called here to be future-proof and corner-case-proof. + hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + + # Expand hidden states for hyper connections at the start of the block + # Only expand at the first PP stage; subsequent stages receive n-stream from previous stage + if self.config.enable_hyper_connections and self.pre_process: + hidden_states = HyperConnectionModule.input_expand( + hidden_states, self.config.num_residual_streams + ) # [s, b, C] -> [s, b, n*C] + + if self.config.sequence_parallel: + rng_context = tensor_parallel.get_cuda_rng_tracker().fork() + else: + rng_context = nullcontext() + + # If fp8_recipe is delayed, wrap the entire pass with get_fp8_context(), + # otherwise do nothing extra at the outer level + # if we are using other fp8 recipes, then the context manager enter&exit are free + # we can wrap fp8_context within the for loop over layers, so that we can fine-grained + # control which layer will be fp8 or bf16 + # For FP4: NVFP4BlockScaling doesn't have delayed scaling, always uses inner context + if self.config.fp8: + use_outer_quantization_context = self.config.fp8_recipe == Fp8Recipe.delayed + use_inner_quantization_context = self.config.fp8_recipe != Fp8Recipe.delayed + outer_quantization_context = ( + get_fp8_context(self.config) if use_outer_quantization_context else nullcontext() + ) + elif self.config.fp4: + use_outer_quantization_context = False + use_inner_quantization_context = True + outer_quantization_context = nullcontext() + else: + # No quantization + use_outer_quantization_context = False + use_inner_quantization_context = False + outer_quantization_context = nullcontext() + + # Determine if MHC recompute should be used + # Only enable when: training mode AND hyper connections AND 'mhc' in recompute_modules + use_mhc_recompute = ( + self.training + and self.config.enable_hyper_connections + and self.config.recompute_granularity == "selective" + and "mhc" in self.config.recompute_modules + ) + mhc_layer_managers, mhc_is_last_in_recompute_block = self._build_mhc_recompute_layer_plan( + use_mhc_recompute + ) + + with rng_context, outer_quantization_context: + # Forward pass. + if self.config.recompute_granularity == "full" and self.training: + checkpointed_result = self._checkpointed_forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + context=context, + context_mask=context_mask, + rotary_pos_emb=rotary_pos_emb, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + use_inner_quantization_context=use_inner_quantization_context, + padding_mask=padding_mask, + extract_layer_indices=extract_layer_indices, + layer_offset=layer_offset, + input_ids=input_ids, + ) + # Handle return value from _checkpointed_forward + if len(extract_layer_indices) > 0: + # (hidden_states, intermediate_hidden_states) tuple + hidden_states, intermediate_hidden_states = checkpointed_result + else: + # No intermediate_hidden_states requested: just hidden_states + hidden_states = checkpointed_result + else: + for l_no, layer in enumerate(self.layers): + # Get appropriate inner quantization context + if use_inner_quantization_context: + if self.config.fp8: + inner_quantization_context = get_fp8_context( + self.config, layer.layer_number - 1 + ) + elif self.config.fp4: + inner_quantization_context = get_fp4_context( + self.config, layer.layer_number - 1 + ) + else: + inner_quantization_context = nullcontext() + else: + inner_quantization_context = nullcontext() + + mhc_manager = mhc_layer_managers[l_no] + if mhc_manager is not None: + mhc_manager.is_last_layer_in_recompute_block = ( + mhc_is_last_in_recompute_block[l_no] + ) + + with self.offload_context, inner_quantization_context: + #### FlagScale Begin #### + # Pre-compute embeddings for the next DeepSeekTransformerLayer if engram exists, to overlap with current layer's computation + if l_no < len(self.layers) - 1: + next_layer = self.layers[l_no + 1] + if getattr(next_layer, "is_engram_layer", False): + next_layer.pre_compute_embedding(engram_hash_input_ids) + #### FlagScale End #### + hidden_states, context = layer( + hidden_states=hidden_states, + attention_mask=attention_mask, + context=context, + context_mask=context_mask, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + rotary_pos_cos_sin=rotary_pos_cos_sin, + attention_bias=attention_bias, + inference_context=inference_context, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + padding_mask=padding_mask, + mhc_recompute_manager=mhc_manager, + input_ids=input_ids, + ) + + self._finalize_mhc_recompute_layer( + mhc_manager=mhc_manager, + hidden_states=hidden_states, + is_last_in_recompute_block=mhc_is_last_in_recompute_block[l_no], + ) + + if ( + torch.is_grad_enabled() + and self.config.cpu_offloading + and self.group_prefetch_offload_commit_async is not None + ): + hidden_states = self.group_prefetch_offload_commit_async(hidden_states) + + # Extract intermediate embeddings using global layer index + if (l_no + layer_offset) in extract_layer_indices: + intermediate_hidden_states.append(hidden_states) + + # Only contract if the final layer norm is in this stage + if self.config.enable_hyper_connections and self.has_final_layernorm_in_this_stage(): + hidden_states = HyperConnectionModule.output_contract( + hidden_states, self.config.num_residual_streams + ) # [s, b, n*C] -> [s, b, C] + + # Final layer norm. + if self.final_layernorm is not None: + hidden_states = apply_module(self.final_layernorm)(cast(Tensor, hidden_states)) + # TENorm produces a "viewed" tensor. This will result in schedule.py's + # deallocate_output_tensor() throwing an error, so a viewless tensor is + # created to prevent this. + hidden_states = make_viewless_tensor( + inp=hidden_states, requires_grad=True, keep_graph=True + ) + + # If this TransformerBlock is empty, input and output hidden states will be the same node + # on the computational graph and will lead to unexpected errors in pipeline schedules. + if not self.pre_process and len(self.layers) == 0 and not self.final_layernorm: + hidden_states = hidden_states.clone() + ########## FlagScale Begin ########## + self.config.recompute_granularity = saved_recompute_granularity + ########## FlagScale End ########## + + if len(extract_layer_indices) > 0: + return hidden_states, intermediate_hidden_states + + return hidden_states + + def sharded_state_dict( + self, prefix: str = "", sharded_offsets: tuple = (), metadata: dict | None = None + ): + # Engram let the layers be non-homogeneous, so we need to set the flag in metadata to let the sharded state dict logic know. + # This is usefule when all layer are same, the TransformerBlock will be homogeneous, it generate sharded_state_dict will same keys for all layer and need all layers have the same structure. + # The layer has engram module does not fit this assumption. + # If the flag is set to True, the sharded_state_dict will use layer_number to generate different keys for different layer, which is same to models has dense layer leading and moe layer following. + # Actually, engram really causes the layers to be non-homogeneous. + if metadata is None: + metadata = {} + metadata["non_homogeneous_layers"] = True + return super().sharded_state_dict( + prefix=prefix, sharded_offsets=sharded_offsets, metadata=metadata + ) diff --git a/flagscale/models/megatron/deepseek_v4/deepseek_transformer_layer.py b/flagscale/models/megatron/deepseek_v4/deepseek_transformer_layer.py index a4ac775fea..1a9f60e978 100644 --- a/flagscale/models/megatron/deepseek_v4/deepseek_transformer_layer.py +++ b/flagscale/models/megatron/deepseek_v4/deepseek_transformer_layer.py @@ -52,7 +52,10 @@ def __init__(self, config, submodules, *args, **kwargs): self.engram = None self._deepseek_engram_hash_input_ids = None self._mhc_recompute_manager = None - if self.config.engram_layer_ids is not None and self.layer_number - 1 in self.config.engram_layer_ids: + if ( + self.config.engram_layer_ids is not None + and self.layer_number - 1 in self.config.engram_layer_ids + ): self.is_engram_layer = True else: self.is_engram_layer = False @@ -102,7 +105,7 @@ def _forward_attention( sequence_len_offset: Optional[Tensor] = None, padding_mask: Optional[Tensor] = None, input_ids: Optional[Tensor] = None, - mhc_recompute_manager = None, + mhc_recompute_manager=None, *, inference_params: Optional[Any] = None, ): @@ -131,7 +134,9 @@ def _forward_attention( ) def pre_compute_embedding(self, engram_hash_input_ids): - if isinstance(self.engram, IdentityOp) or (self.layer_number not in self.config.engram_layer_ids): + if isinstance(self.engram, IdentityOp) or ( + self.layer_number not in self.config.engram_layer_ids + ): return hash_input_ids = engram_hash_input_ids[self.layer_number - 1] self.engram.pre_compute_embedding(hash_input_ids) diff --git a/flagscale/models/megatron/engram/engram_builder.py b/flagscale/models/megatron/engram/engram_builder.py index 2e541a58d3..192e044f18 100644 --- a/flagscale/models/megatron/engram/engram_builder.py +++ b/flagscale/models/megatron/engram/engram_builder.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ruff: noqa: TC001 # ruff: noqa: F401 ## built-in @@ -6,34 +20,42 @@ from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add from megatron.core.models.backends import BackendSpecProvider + ## megatron-core from megatron.core.models.gpt.gpt_layer_specs import ( - get_gpt_decoder_block_spec, get_gpt_layer_local_spec, + get_gpt_decoder_block_spec, + get_gpt_layer_local_spec, get_gpt_layer_with_inference_spec, - get_gpt_layer_with_transformer_engine_spec, get_gpt_mtp_block_spec, - get_mlp_module_spec_for_backend) + get_gpt_layer_with_transformer_engine_spec, + get_gpt_mtp_block_spec, + get_mlp_module_spec_for_backend, +) from megatron.core.models.gpt.gpt_model import GPTModel -from megatron.core.transformer.attention import (SelfAttention, - SelfAttentionSubmodules) +from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules from megatron.core.transformer.enums import AttnMaskType, LayerType from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.multi_latent_attention import ( - MLASelfAttention, MLASelfAttentionSubmodules) + MLASelfAttention, + MLASelfAttentionSubmodules, +) from megatron.core.transformer.spec_utils import ModuleSpec from megatron.core.transformer.torch_norm import L2Norm from megatron.core.transformer.transformer_block import ( - TransformerBlockSubmodules, get_num_layers_to_build) + TransformerBlockSubmodules, + get_num_layers_to_build, +) from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import ( - TransformerLayerSubmodules, get_transformer_layer_offset) + TransformerLayerSubmodules, + get_transformer_layer_offset, +) from megatron.training import get_args, print_rank_0 from megatron.training.arguments import core_transformer_config_from_args try: import transformer_engine as te # pylint: disable=unused-import from megatron.core.extensions.transformer_engine import TENorm - from megatron.core.extensions.transformer_engine_spec_provider import \ - TESpecProvider + from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider HAVE_TE = True except ImportError: @@ -126,13 +148,9 @@ def get_engram_transformer_layer_spec( assert HAVE_KITCHEN backend: BackendSpecProvider = KitchenSpecProvider(fallback=TESpecProvider()) if use_te_op_fuser: - raise AssertionError( - "use_te_op_fuser not compatible with using kitchen in mlp." - ) + raise AssertionError("use_te_op_fuser not compatible with using kitchen in mlp.") if use_te_activation_func: - raise AssertionError( - "use_te_activation_func not compatible with using kitchen." - ) + raise AssertionError("use_te_activation_func not compatible with using kitchen.") else: backend = TESpecProvider() @@ -194,14 +212,10 @@ def get_engram_transformer_layer_spec( core_attention=backend.core_attention(), linear_proj=backend.row_parallel_linear(), q_layernorm=( - L2Norm - if qk_l2_norm - else (qk_norm if qk_layernorm else IdentityOp) + L2Norm if qk_l2_norm else (qk_norm if qk_layernorm else IdentityOp) ), k_layernorm=( - L2Norm - if qk_l2_norm - else (qk_norm if qk_layernorm else IdentityOp) + L2Norm if qk_l2_norm else (qk_norm if qk_layernorm else IdentityOp) ), ), ), @@ -279,8 +293,7 @@ def get_engram_decoder_block_spec( if use_moe: if isinstance(config.moe_layer_freq, int): moe_layer_pattern = [ - 1 if (i % config.moe_layer_freq == 0) else 0 - for i in range(config.num_layers) + 1 if (i % config.moe_layer_freq == 0) else 0 for i in range(config.num_layers) ] elif isinstance(config.moe_layer_freq, list): moe_layer_pattern = config.moe_layer_freq @@ -301,9 +314,7 @@ def get_engram_decoder_block_spec( for layer_number in range(config.num_layers): is_engram_layer = True if layer_number in config.engram_layer_ids else False if moe_layer_pattern[layer_number] == 1: - layer_specs.append( - moe_engram_layer_spec if is_engram_layer else moe_orig_layer_spec - ) + layer_specs.append(moe_engram_layer_spec if is_engram_layer else moe_orig_layer_spec) elif moe_layer_pattern[layer_number] == 0: layer_specs.append( dense_engram_layer_spec if is_engram_layer else dense_orig_layer_spec @@ -346,9 +357,7 @@ def get_engram_decoder_block_spec( return block_spec -def engram_builder( - args, pre_process, post_process, vp_stage=None, config=None, pg_collection=None -): +def engram_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_collection=None): print_rank_0("building Engram model ...") config = core_transformer_config_from_args(args) diff --git a/flagscale/models/megatron/engram/engram_model.py b/flagscale/models/megatron/engram/engram_model.py index 6e473ff61f..7cd0b0dfc8 100644 --- a/flagscale/models/megatron/engram/engram_model.py +++ b/flagscale/models/megatron/engram/engram_model.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ruff: noqa: RUF013 ## built-in from typing import Optional @@ -33,7 +47,7 @@ def __init__(self, hash_mapping, input_ids, hash_stream=None): self.input_ids = input_ids self.hash_stream = hash_stream self._result = None - self._is_async_pending = False + self._is_async_pending = False # Async if self.hash_stream is not None: # self.hash_stream.wait_stream(cur_platform.current_stream()) @@ -61,11 +75,11 @@ def __getitem__(self, key): cur_platform.current_stream().wait_stream(self.hash_stream) self._is_async_pending = False # Async finish self._record_current_stream() - + # Case 2: Sync but no compute -> start compute elif self._result is None: self._result = self.hash_mapping.hash(self.input_ids) - + # Case 3: Async or sync compute is finished. # print(f"[rank{torch.distributed.get_rank()}]: LazyHashInputIds result = {self._result}") return self._result[key] @@ -226,6 +240,5 @@ def build_schedule_plan( decoder_input, labels=labels, loss_mask=loss_mask, - extra_block_kwargs=extra_block_kwargs + extra_block_kwargs=extra_block_kwargs, ) - diff --git a/flagscale/models/megatron/engram/engram_transformer_layer.py b/flagscale/models/megatron/engram/engram_transformer_layer.py index e6d4f884c7..a782bfeb20 100644 --- a/flagscale/models/megatron/engram/engram_transformer_layer.py +++ b/flagscale/models/megatron/engram/engram_transformer_layer.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ruff: noqa: RUF013 # ruff: noqa: E711 @@ -23,7 +37,7 @@ deprecate_inference_params, make_viewless_tensor, ) -from megatron.core.transformer.engram import EngramModule +from megatron.core.transformer.engram import EngramModule class EngramTransformerLayer(TransformerLayer): @@ -77,14 +91,16 @@ def forward( sequence_len_offset=sequence_len_offset, inference_params=inference_params, ) - + def pre_compute_embedding(self, hash_input_ids: Tensor): self.engram.pre_compute_embedding(hash_input_ids) def sharded_state_dict( self, prefix: str = "", sharded_offsets: tuple = (), metadata: dict | None = None ): - return super().sharded_state_dict(prefix=prefix, sharded_offsets=sharded_offsets, metadata=metadata) + return super().sharded_state_dict( + prefix=prefix, sharded_offsets=sharded_offsets, metadata=metadata + ) class EngramTransformerBlock(TransformerBlock): @@ -340,7 +356,7 @@ def forward( hidden_states = hidden_states.clone() return hidden_states - + def sharded_state_dict( self, prefix: str = "", sharded_offsets: tuple = (), metadata: dict | None = None ): @@ -350,4 +366,6 @@ def sharded_state_dict( # If the flag is set to True, the sharded_state_dict will use layer_number to generate different keys for different layer, which is same to models has dense layer leading and moe layer following. # Actually, engram really causes the layers to be non-homogeneous. metadata["non_homogeneous_layers"] = True - return super().sharded_state_dict(prefix=prefix, sharded_offsets=sharded_offsets, metadata=metadata) + return super().sharded_state_dict( + prefix=prefix, sharded_offsets=sharded_offsets, metadata=metadata + ) diff --git a/flagscale/models/megatron/llava_onevision/__init__.py b/flagscale/models/megatron/llava_onevision/__init__.py index e69de29bb2..a491ed2726 100644 --- a/flagscale/models/megatron/llava_onevision/__init__.py +++ b/flagscale/models/megatron/llava_onevision/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flagscale/models/megatron/llava_onevision/clip_vit_model.py b/flagscale/models/megatron/llava_onevision/clip_vit_model.py index 3cdcefa815..ffb6797a36 100644 --- a/flagscale/models/megatron/llava_onevision/clip_vit_model.py +++ b/flagscale/models/megatron/llava_onevision/clip_vit_model.py @@ -61,9 +61,7 @@ def __init__( self.add_class_token = add_class_token self.class_token_len = class_token_len - self.seq_length = self.num_patches + ( - self.class_token_len if self.add_class_token else 0 - ) + self.seq_length = self.num_patches + (self.class_token_len if self.add_class_token else 0) # siglip same as LLaVA-NeXT args = get_args() @@ -86,9 +84,7 @@ def __init__( ) self.position_ids = torch.arange(self.seq_length).expand(1, -1).cuda() - self.position_embeddings = torch.nn.Embedding( - self.seq_length, self.visual_hidden_size - ) + self.position_embeddings = torch.nn.Embedding(self.seq_length, self.visual_hidden_size) self.add_class_token = add_class_token if self.add_class_token: @@ -179,17 +175,13 @@ def forward( return x -def get_num_image_embeddings( - img_h, img_w, patch_dim, disable_vision_class_token, class_token_len -): +def get_num_image_embeddings(img_h, img_w, patch_dim, disable_vision_class_token, class_token_len): """Get the number of image embeddings per image tile.""" add_class_token = not disable_vision_class_token num_patches_per_dim_h = img_h // patch_dim num_patches_per_dim_w = img_w // patch_dim num_patches = num_patches_per_dim_h * num_patches_per_dim_w - num_image_embeddings_per_tile = num_patches + ( - class_token_len if add_class_token else 0 - ) + num_image_embeddings_per_tile = num_patches + (class_token_len if add_class_token else 0) return num_image_embeddings_per_tile diff --git a/flagscale/models/megatron/llava_onevision/config.py b/flagscale/models/megatron/llava_onevision/config.py index 7de522cae8..d395c73797 100644 --- a/flagscale/models/megatron/llava_onevision/config.py +++ b/flagscale/models/megatron/llava_onevision/config.py @@ -1,5 +1,6 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. -""" Configuration for the vision tower, the llm tower and the projector. """ +"""Configuration for the vision tower, the llm tower and the projector.""" + import torch from megatron.core.activations import quick_gelu, squared_relu, fast_gelu @@ -16,9 +17,7 @@ def get_language_model_config(config): # Add qwen2_7b config if config.language_model_type == "qwen2_7b": config.activation_func = torch.nn.functional.silu - config.add_bias_linear = ( - False # linear_qkv has bias but linear_proj and MLP has no bias - ) + config.add_bias_linear = False # linear_qkv has bias but linear_proj and MLP has no bias config.add_qkv_bias = True config.bias_activation_fusion = False config.gated_linear_unit = True @@ -33,9 +32,7 @@ def get_language_model_config(config): # Add qwen2_1.5b config elif config.language_model_type == "qwen2_1.5b": config.activation_func = torch.nn.functional.silu - config.add_bias_linear = ( - False # linear_qkv has bias but linear_proj and MLP has no bias - ) + config.add_bias_linear = False # linear_qkv has bias but linear_proj and MLP has no bias config.add_qkv_bias = True config.bias_activation_fusion = False config.gated_linear_unit = True diff --git a/flagscale/models/megatron/llava_onevision/dataloader_provider.py b/flagscale/models/megatron/llava_onevision/dataloader_provider.py index 7c069e75ef..9e81e07d23 100644 --- a/flagscale/models/megatron/llava_onevision/dataloader_provider.py +++ b/flagscale/models/megatron/llava_onevision/dataloader_provider.py @@ -2,7 +2,10 @@ import os import torch -from flagscale.models.megatron.llava_onevision.dataset_helpers import print_error_handler, AnyResTaskEncoder +from flagscale.models.megatron.llava_onevision.dataset_helpers import ( + print_error_handler, + AnyResTaskEncoder, +) from megatron.core import mpu from megatron.energon import ( @@ -100,15 +103,13 @@ def train_valid_test_dataloaders_provider(train_val_test_num_samples): ) if os.path.exists(data_save_name): try: - dataset_state_dict = torch.load(data_save_name, map_location="cpu", weights_only=False) - train_dataloader.restore_state_rank( - dataset_state_dict["dataloader_state_dict"] + dataset_state_dict = torch.load( + data_save_name, map_location="cpu", weights_only=False ) + train_dataloader.restore_state_rank(dataset_state_dict["dataloader_state_dict"]) print_rank_0(f"restored dataset state from {data_save_name}") except Exception as e: - print_rank_0( - "loading dataloader checkpoint failed. Skipping. " + str(e) - ) + print_rank_0("loading dataloader checkpoint failed. Skipping. " + str(e)) if args.training_dataset_only: return ( EnergonDataloader(train_dataloader), diff --git a/flagscale/models/megatron/llava_onevision/dataset_helpers.py b/flagscale/models/megatron/llava_onevision/dataset_helpers.py index 4348a0e86a..00062dec58 100644 --- a/flagscale/models/megatron/llava_onevision/dataset_helpers.py +++ b/flagscale/models/megatron/llava_onevision/dataset_helpers.py @@ -9,11 +9,7 @@ import numpy as np import torch -from megatron.energon import ( - Batch, - DefaultTaskEncoder, - InterleavedSample -) +from megatron.energon import Batch, DefaultTaskEncoder, InterleavedSample from megatron.training import get_args @@ -37,6 +33,7 @@ class AnyResTaskSample: image_sizes: List[torch.Tensor] modalities: List[torch.Tensor] + # Typing for the resulting batch data after encode_batch() @dataclass class AnyResTaskBatch(Batch): @@ -52,16 +49,16 @@ class AnyResTaskBatch(Batch): modalities: torch.Tensor -class AnyResTaskEncoder(DefaultTaskEncoder[InterleavedSample, InterleavedSample, AnyResTaskBatch, dict]): +class AnyResTaskEncoder( + DefaultTaskEncoder[InterleavedSample, InterleavedSample, AnyResTaskBatch, dict] +): """ A task encoder for anyres. This encoder is just a wrapper around data that has already been made. Production data can be referenced to LLaVA-NeXT and can be input into vision tower. """ - def __init__( - self - ): + def __init__(self): # Specify the batch_type for default batching (batching is performed here "manually" by # overwriting the `batch` method) super().__init__() @@ -70,7 +67,9 @@ def __init__( def encode_sample(self, sample: InterleavedSample): if not isinstance(sample, InterleavedSample): - raise ValueError(f"This encoder only supports InterleavedSample, but got {type(sample)}.") + raise ValueError( + f"This encoder only supports InterleavedSample, but got {type(sample)}." + ) yield self.encode_interleaved(sample) def encode_interleaved(self, sample: InterleavedSample): @@ -82,7 +81,9 @@ def encode_interleaved(self, sample: InterleavedSample): elif len(sample.sequence) == 5: input_ids, labels, images, image_sizes, modalities = sample.sequence else: - assert ValueError("The sequence must have 4 or 5 elements, but got {len(sample.sequence)}.") + assert ValueError( + "The sequence must have 4 or 5 elements, but got {len(sample.sequence)}." + ) # process modalities to tensor modalities_list = [] @@ -107,7 +108,7 @@ def encode_interleaved(self, sample: InterleavedSample): labels_shape=torch.tensor(labels.shape), images=images, image_sizes=image_sizes, - modalities=modalities + modalities=modalities, ) def batch(self, samples: List[AnyResTaskSample]) -> AnyResTaskBatch: @@ -117,10 +118,23 @@ def batch(self, samples: List[AnyResTaskSample]) -> AnyResTaskBatch: labels_shape = torch.stack([s.labels_shape for s in samples], dim=0) # Double loop images = torch.cat([image.flatten() for s in samples for image in s.images], dim=0) - split_image_sizes = torch.stack([torch.tensor(image.shape) for s in samples for image in s.images], dim=0) + split_image_sizes = torch.stack( + [torch.tensor(image.shape) for s in samples for image in s.images], dim=0 + ) # Adapt video data by decord - image_sizes = torch.stack([image_sizes if len(image_sizes.shape) == 1 else torch.tensor((1, image_sizes.item())) for s in samples for image_sizes in s.image_sizes], dim=0) - modalities = torch.stack([modalities for s in samples for modalities in s.modalities], dim=0) + image_sizes = torch.stack( + [ + image_sizes + if len(image_sizes.shape) == 1 + else torch.tensor((1, image_sizes.item())) + for s in samples + for image_sizes in s.image_sizes + ], + dim=0, + ) + modalities = torch.stack( + [modalities for s in samples for modalities in s.modalities], dim=0 + ) batch = AnyResTaskBatch( __keys__=[s.__key__ for s in samples], @@ -132,7 +146,7 @@ def batch(self, samples: List[AnyResTaskSample]) -> AnyResTaskBatch: images=images, image_sizes=image_sizes, split_image_sizes=split_image_sizes, - modalities=modalities + modalities=modalities, ) return batch diff --git a/flagscale/models/megatron/llava_onevision/layer_specs.py b/flagscale/models/megatron/llava_onevision/layer_specs.py index b56e0b07e1..46be2e0d1b 100644 --- a/flagscale/models/megatron/llava_onevision/layer_specs.py +++ b/flagscale/models/megatron/llava_onevision/layer_specs.py @@ -36,7 +36,7 @@ from megatron.core.transformer.torch_layer_norm import WrappedTorchLayerNorm - warnings.warn(f'Apex is not installed. Falling back to Torch LayerNorm') + warnings.warn(f"Apex is not installed. Falling back to Torch LayerNorm") LNImpl = WrappedTorchLayerNorm diff --git a/flagscale/models/megatron/llava_onevision/llava_onevision_model.py b/flagscale/models/megatron/llava_onevision/llava_onevision_model.py index 5ef82617fd..442540abe2 100644 --- a/flagscale/models/megatron/llava_onevision/llava_onevision_model.py +++ b/flagscale/models/megatron/llava_onevision/llava_onevision_model.py @@ -85,9 +85,7 @@ def __init__( super().__init__(config=language_transformer_config) if has_config_logger_enabled(language_transformer_config): - log_config_to_disk( - language_transformer_config, locals(), prefix=type(self).__name__ - ) + log_config_to_disk(language_transformer_config, locals(), prefix=type(self).__name__) logging.getLogger(__name__).warning( "LLaVA OneVision model is under active development. " @@ -108,17 +106,13 @@ def __init__( args = self.args # Init image_newline if "unpad" in args.mm_patch_merge_type: - embed_std = 1 / torch.sqrt( - torch.tensor(args.hidden_size, dtype=torch.bfloat16) - ) + embed_std = 1 / torch.sqrt(torch.tensor(args.hidden_size, dtype=torch.bfloat16)) self.image_newline = torch.nn.Parameter( torch.randn(args.hidden_size, dtype=torch.bfloat16) * embed_std ) # Add share_embeddings_and_output_weights to the language model. - self.share_embeddings_and_output_weights = ( - not args.untie_embeddings_and_output_weights - ) + self.share_embeddings_and_output_weights = not args.untie_embeddings_and_output_weights if self.add_decoder: self.language_model = GPTModel( @@ -268,9 +262,7 @@ def prepare_inputs_labels_for_multimodal( if vision_tower is None or images is None or input_ids.shape[1] == 1: # [BugFix]: comment out the embed_tokens # input_ids = self.embed_tokens(input_ids) - loss_mask = torch.where( - labels == IGNORE_INDEX, torch.tensor(0), torch.tensor(1) - ) + loss_mask = torch.where(labels == IGNORE_INDEX, torch.tensor(0), torch.tensor(1)) return input_ids, position_ids, attention_mask, labels, loss_mask if isinstance(modalities, str): @@ -335,9 +327,9 @@ def prepare_inputs_labels_for_multimodal( image_aspect_ratio = args.image_aspect_ratio if image_aspect_ratio != "square": - assert ( - "anyres" in image_aspect_ratio - ), f"Unexpected image_aspect_ratio: {image_aspect_ratio}" + assert "anyres" in image_aspect_ratio, ( + f"Unexpected image_aspect_ratio: {image_aspect_ratio}" + ) mm_newline_position = args.mm_newline_position assert mm_newline_position in [ @@ -374,9 +366,7 @@ def prepare_inputs_labels_for_multimodal( image_feature = torch.cat( ( image_feature, - self.image_newline[None].to( - image_feature.device - ), + self.image_newline[None].to(image_feature.device), ), dim=0, ) @@ -392,9 +382,7 @@ def prepare_inputs_labels_for_multimodal( base_image_feature = image_feature[0] # Patch iamge features image_feature = image_feature[1:] - assert ( - args.img_h == args.img_w - ), "Only support square image size." + assert args.img_h == args.img_w, "Only support square image size." height = width = args.img_h // args.patch_dim assert height * width == base_image_feature.shape[0] @@ -403,33 +391,24 @@ def prepare_inputs_labels_for_multimodal( r"anyres_max_(\d+)", image_aspect_ratio ) if matched_anyres_max_num_patches: - max_num_patches = int( - matched_anyres_max_num_patches.group(1) - ) + max_num_patches = int(matched_anyres_max_num_patches.group(1)) - if ( - image_aspect_ratio == "anyres" - or "anyres_max" in image_aspect_ratio - ): + if image_aspect_ratio == "anyres" or "anyres_max" in image_aspect_ratio: vision_tower_image_size = args.img_h - assert ( - args.image_grid_pinpoints is not None - ), "image_grid_pinpoints must be provided." - num_patch_width, num_patch_height = ( - get_anyres_image_grid_shape( - image_sizes[image_idx], - args.image_grid_pinpoints, - vision_tower_image_size, - ) + assert args.image_grid_pinpoints is not None, ( + "image_grid_pinpoints must be provided." + ) + num_patch_width, num_patch_height = get_anyres_image_grid_shape( + image_sizes[image_idx], + args.image_grid_pinpoints, + vision_tower_image_size, ) image_feature = image_feature.view( num_patch_height, num_patch_width, height, width, -1 ) if "maxpool2x2" in mm_patch_merge_type: - image_feature = image_feature.permute( - 4, 0, 2, 1, 3 - ).contiguous() + image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous() image_feature = image_feature.flatten(1, 2).flatten(2, 3) image_feature = nn.functional.max_pool2d(image_feature, 2) image_feature = image_feature.flatten(1, 2).transpose(0, 1) @@ -439,13 +418,9 @@ def prepare_inputs_labels_for_multimodal( and matched_anyres_max_num_patches ): unit = image_feature.shape[2] - image_feature = image_feature.permute( - 4, 0, 2, 1, 3 - ).contiguous() + image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous() image_feature = image_feature.flatten(1, 2).flatten(2, 3) - image_feature = unpad_image( - image_feature, image_sizes[image_idx] - ) + image_feature = unpad_image(image_feature, image_sizes[image_idx]) c, h, w = image_feature.shape times = math.sqrt(h * w / (max_num_patches * unit**2)) if times > 1.1: @@ -466,13 +441,9 @@ def prepare_inputs_labels_for_multimodal( ) image_feature = image_feature.flatten(1, 2).transpose(0, 1) elif "unpad" in mm_patch_merge_type: - image_feature = image_feature.permute( - 4, 0, 2, 1, 3 - ).contiguous() + image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous() image_feature = image_feature.flatten(1, 2).flatten(2, 3) - image_feature = unpad_image( - image_feature, image_sizes[image_idx] - ) + image_feature = unpad_image(image_feature, image_sizes[image_idx]) image_feature = torch.cat( ( image_feature, @@ -484,16 +455,12 @@ def prepare_inputs_labels_for_multimodal( ) image_feature = image_feature.flatten(1, 2).transpose(0, 1) else: - image_feature = image_feature.permute( - 0, 2, 1, 3, 4 - ).contiguous() + image_feature = image_feature.permute(0, 2, 1, 3, 4).contiguous() image_feature = image_feature.flatten(0, 3) if "nobase" in mm_patch_merge_type: pass else: - image_feature = torch.cat( - (base_image_feature, image_feature), dim=0 - ) + image_feature = torch.cat((base_image_feature, image_feature), dim=0) new_image_features.append(image_feature) else: # single image operations image_feature = image_feature[0] @@ -504,9 +471,7 @@ def prepare_inputs_labels_for_multimodal( new_image_features.append(image_feature) image_features = new_image_features else: - raise ValueError( - f"Unexpected mm_patch_merge_type: {args.mm_patch_merge_type}" - ) + raise ValueError(f"Unexpected mm_patch_merge_type: {args.mm_patch_merge_type}") else: image_features = self.encode_images(images) @@ -544,9 +509,7 @@ def prepare_inputs_labels_for_multimodal( if num_images == 0: cur_image_features = image_features[cur_image_idx] cur_input_embeds_1 = self.embed_tokens(cur_input_ids) - cur_input_embeds = torch.cat( - [cur_input_embeds_1, cur_image_features[0:0]], dim=0 - ) + cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features[0:0]], dim=0) new_input_embeds.append(cur_input_embeds) new_labels.append(labels[batch_idx]) cur_image_idx += 1 @@ -562,9 +525,7 @@ def prepare_inputs_labels_for_multimodal( cur_labels_noim = [] for i in range(len(image_token_indices) - 1): cur_input_ids_noim.append( - cur_input_ids[ - image_token_indices[i] + 1 : image_token_indices[i + 1] - ] + cur_input_ids[image_token_indices[i] + 1 : image_token_indices[i + 1]] ) cur_labels_noim.append( cur_labels[image_token_indices[i] + 1 : image_token_indices[i + 1]] @@ -604,13 +565,9 @@ def prepare_inputs_labels_for_multimodal( tokenizer_model_max_length = args.max_position_embeddings new_input_embeds = [ - x[:tokenizer_model_max_length] - for x, modality in zip(new_input_embeds, modalities) - ] - new_labels = [ - x[:tokenizer_model_max_length] - for x, modality in zip(new_labels, modalities) + x[:tokenizer_model_max_length] for x, modality in zip(new_input_embeds, modalities) ] + new_labels = [x[:tokenizer_model_max_length] for x, modality in zip(new_labels, modalities)] # Combine them max_len = max(x.shape[0] for x in new_input_embeds) batch_size = len(new_input_embeds) @@ -632,9 +589,7 @@ def prepare_inputs_labels_for_multimodal( ) tokenizer = get_tokenizer() - for i, (cur_new_embed, cur_new_labels) in enumerate( - zip(new_input_embeds, new_labels) - ): + for i, (cur_new_embed, cur_new_labels) in enumerate(zip(new_input_embeds, new_labels)): cur_len = cur_new_embed.shape[0] if tokenizer.padding_side == "left": new_input_embeds_padded.append( @@ -684,9 +639,7 @@ def prepare_inputs_labels_for_multimodal( loss_mask = None else: new_labels = new_labels_padded - loss_mask = torch.where( - new_labels == IGNORE_INDEX, torch.tensor(0), torch.tensor(1) - ) + loss_mask = torch.where(new_labels == IGNORE_INDEX, torch.tensor(0), torch.tensor(1)) if _attention_mask is None: attention_mask = None @@ -722,7 +675,6 @@ def forward( position_ids: Optional[torch.LongTensor] = None, image_token_index: Optional[int] = IMAGE_TOKEN_INDEX, ) -> torch.Tensor: - input_embeds, position_ids, attention_mask, labels, loss_mask = ( self.prepare_inputs_labels_for_multimodal( input_ids, @@ -776,9 +728,7 @@ def get_2dPool(self, image_feature, stride=2): ) else: - raise ValueError( - f"Unexpected mm_spatial_pool_mode: {args.mm_spatial_pool_mode}" - ) + raise ValueError(f"Unexpected mm_spatial_pool_mode: {args.mm_spatial_pool_mode}") image_feature = image_feature.permute(0, 2, 3, 1) image_feature = image_feature.view(num_frames, -1, num_dim) return image_feature @@ -912,8 +862,9 @@ def select_best_resolution(original_size, possible_resolutions): for width, height in possible_resolutions: # Calculate the downscaled size to keep the aspect ratio scale = min(width / original_width, height / original_height) - downscaled_width, downscaled_height = int(original_width * scale), int( - original_height * scale + downscaled_width, downscaled_height = ( + int(original_width * scale), + int(original_height * scale), ) # Calculate effective and wasted resolutions diff --git a/flagscale/models/megatron/llava_onevision/mlp_converter.py b/flagscale/models/megatron/llava_onevision/mlp_converter.py index fa68fa4363..a492347c20 100644 --- a/flagscale/models/megatron/llava_onevision/mlp_converter.py +++ b/flagscale/models/megatron/llava_onevision/mlp_converter.py @@ -13,7 +13,6 @@ def convert(input_path, output_path, tensor_parallel_size): new_state_dicts = [{"model": dict()} for _ in range(tensor_parallel_size)] for name, tensor in state_dict.items(): - # Map parameter names to ones used in megatron. new_name = "" new_tensor = tensor @@ -73,5 +72,3 @@ def convert(input_path, output_path, tensor_parallel_size): convert(args.input, args.output, args.tensor_parallel_size) print("done.") - - diff --git a/flagscale/models/megatron/qwen2_5_vl/__init__.py b/flagscale/models/megatron/qwen2_5_vl/__init__.py index e69de29bb2..a491ed2726 100644 --- a/flagscale/models/megatron/qwen2_5_vl/__init__.py +++ b/flagscale/models/megatron/qwen2_5_vl/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flagscale/models/megatron/qwen2_5_vl/language_module.py b/flagscale/models/megatron/qwen2_5_vl/language_module.py index 81768fdd10..d0aabcc4c1 100644 --- a/flagscale/models/megatron/qwen2_5_vl/language_module.py +++ b/flagscale/models/megatron/qwen2_5_vl/language_module.py @@ -33,13 +33,21 @@ def __init__( config: TransformerConfig, vocab_size: int, max_sequence_length: int, - position_embedding_type: Literal['learned_absolute', 'rope', 'none'] = 'learned_absolute', + position_embedding_type: Literal["learned_absolute", "rope", "none"] = "learned_absolute", num_tokentypes: int = 0, - scatter_to_sequence_parallel: bool = False, # chage default to False + scatter_to_sequence_parallel: bool = False, # chage default to False ): - assert scatter_to_sequence_parallel == False, "QwenVLLanguageModelEmbedding does not support scatter_to_sequence_parallel" - super().__init__(config, vocab_size, max_sequence_length, position_embedding_type, num_tokentypes, scatter_to_sequence_parallel) - + assert scatter_to_sequence_parallel == False, ( + "QwenVLLanguageModelEmbedding does not support scatter_to_sequence_parallel" + ) + super().__init__( + config, + vocab_size, + max_sequence_length, + position_embedding_type, + num_tokentypes, + scatter_to_sequence_parallel, + ) def forward( self, @@ -49,7 +57,7 @@ def forward( image_input_mask: Tensor = None, video_input_mask: Tensor = None, image_embeds: Tensor = None, - video_embeds: Tensor = None + video_embeds: Tensor = None, ) -> Tensor: """Forward pass of the embedding module. @@ -89,9 +97,13 @@ def forward( if not self.reduce_scatter_embeddings: embeddings = embeddings.clone() if image_embeds is not None: - embeddings[image_input_mask] = image_embeds.to(embeddings.device, embeddings.dtype) + embeddings[image_input_mask] = image_embeds.to( + embeddings.device, embeddings.dtype + ) if video_embeds is not None: - embeddings[video_input_mask] = video_embeds.to(embeddings.device, embeddings.dtype) + embeddings[video_input_mask] = video_embeds.to( + embeddings.device, embeddings.dtype + ) embeddings = tensor_parallel.scatter_to_sequence_parallel_region(embeddings) # `scatter_to_sequence_parallel_region` returns a view, which prevents # the original tensor from being garbage collected. Clone to facilitate GC. @@ -142,8 +154,8 @@ def __init__( parallel_output: bool = True, share_embeddings_and_output_weights: bool = False, position_embedding_type: Literal[ - 'learned_absolute', 'rope', 'mrope', 'none' - ] = 'learned_absolute', + "learned_absolute", "rope", "mrope", "none" + ] = "learned_absolute", rotary_percent: float = 1.0, rotary_base: int = 10000, rope_scaling: bool = False, @@ -152,20 +164,25 @@ def __init__( seq_len_interpolation_factor: Optional[float] = None, mtp_block_spec: Optional[ModuleSpec] = None, ) -> None: - super().__init__(config=config, transformer_layer_spec=transformer_layer_spec, - vocab_size=vocab_size, max_sequence_length=max_sequence_length, - pre_process=pre_process, post_process=post_process, - fp16_lm_cross_entropy=fp16_lm_cross_entropy, - parallel_output=parallel_output, - share_embeddings_and_output_weights=share_embeddings_and_output_weights, - position_embedding_type=position_embedding_type, - rotary_percent=rotary_percent, - rotary_base=rotary_base, - rope_scaling=rope_scaling, - rope_scaling_factor=rope_scaling_factor, - scatter_embedding_sequence_parallel=scatter_embedding_sequence_parallel, - seq_len_interpolation_factor=seq_len_interpolation_factor, - mtp_block_spec=mtp_block_spec) + super().__init__( + config=config, + transformer_layer_spec=transformer_layer_spec, + vocab_size=vocab_size, + max_sequence_length=max_sequence_length, + pre_process=pre_process, + post_process=post_process, + fp16_lm_cross_entropy=fp16_lm_cross_entropy, + parallel_output=parallel_output, + share_embeddings_and_output_weights=share_embeddings_and_output_weights, + position_embedding_type=position_embedding_type, + rotary_percent=rotary_percent, + rotary_base=rotary_base, + rope_scaling=rope_scaling, + rope_scaling_factor=rope_scaling_factor, + scatter_embedding_sequence_parallel=scatter_embedding_sequence_parallel, + seq_len_interpolation_factor=seq_len_interpolation_factor, + mtp_block_spec=mtp_block_spec, + ) if self.pre_process: self.embedding = QwenVLLanguageModelEmbedding( config=self.config, diff --git a/flagscale/models/megatron/qwen2_5_vl/layer_specs.py b/flagscale/models/megatron/qwen2_5_vl/layer_specs.py index b74cb4ff57..a89fcef470 100644 --- a/flagscale/models/megatron/qwen2_5_vl/layer_specs.py +++ b/flagscale/models/megatron/qwen2_5_vl/layer_specs.py @@ -20,7 +20,7 @@ TELayerNormColumnParallelLinear, TENorm, TERowParallelLinear, - TEColumnParallelLinear + TEColumnParallelLinear, ) from megatron.core.transformer.enums import AttnMaskType @@ -30,17 +30,14 @@ from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules from megatron.core.transformer.mlp import MLP, MLPSubmodules -from megatron.core.transformer.attention import (SelfAttentionSubmodules, SelfAttention) +from megatron.core.transformer.attention import SelfAttentionSubmodules, SelfAttention from .vision_attention import SelfAttentionVision + # Use this spec to use lower level Transformer Engine modules (required for fp8 training) -def get_gpt_layer_with_transformer_engine_spec( - qk_layernorm: bool = False -) -> ModuleSpec: - mlp = get_mlp_module_spec( - use_te=True, num_experts=None, moe_grouped_gemm=False - ) +def get_gpt_layer_with_transformer_engine_spec(qk_layernorm: bool = False) -> ModuleSpec: + mlp = get_mlp_module_spec(use_te=True, num_experts=None, moe_grouped_gemm=False) return ModuleSpec( module=TransformerLayer, submodules=TransformerLayerSubmodules( @@ -62,10 +59,9 @@ def get_gpt_layer_with_transformer_engine_spec( ), ) -def get_qwen2vl_vision_model_spec( - is_vit=False -) -> ModuleSpec: - attn_mask_type = AttnMaskType.no_mask # THD --> causal_pad + +def get_qwen2vl_vision_model_spec(is_vit=False) -> ModuleSpec: + attn_mask_type = AttnMaskType.no_mask # THD --> causal_pad mlp = ModuleSpec( module=MLP, @@ -98,7 +94,10 @@ def get_qwen2vl_vision_model_spec( # Helper function to get module spec for MLP/MoE def get_mlp_module_spec( - use_te: bool = True, num_experts: int = None, moe_grouped_gemm: bool = False, add_norm: bool = True + use_te: bool = True, + num_experts: int = None, + moe_grouped_gemm: bool = False, + add_norm: bool = True, ) -> ModuleSpec: if num_experts is None: # Dense MLP w/ or w/o TE modules. diff --git a/flagscale/models/megatron/qwen2_5_vl/qwen2_5_vl_model.py b/flagscale/models/megatron/qwen2_5_vl/qwen2_5_vl_model.py index ca348997cb..f1b9a18331 100644 --- a/flagscale/models/megatron/qwen2_5_vl/qwen2_5_vl_model.py +++ b/flagscale/models/megatron/qwen2_5_vl/qwen2_5_vl_model.py @@ -17,6 +17,7 @@ from flagscale.models.megatron.qwen2_5_vl.vit_model import Qwen2_5VisionModel from flagscale.models.megatron.qwen2_5_vl.language_module import QwenVLLanguageModel + # Note: This is under development and may be missing features. class Qwen2_5VLModel(MegatronModule): """Qwen2.5VL multi-modal model. @@ -59,10 +60,9 @@ def __init__( vision_projection_config: TransformerConfig, vision_projection_layer_spec: ModuleSpec, vision_projection_type: str = "mlp", - allow_missing_vision_projection_checkpoint: bool = False, parallel_output: bool = True, - language_position_embedding_type: str = 'rope', + language_position_embedding_type: str = "rope", language_rotary_percent: float = 1.0, pre_process: bool = True, post_process: bool = True, @@ -70,7 +70,7 @@ def __init__( add_decoder: bool = True, language_rotary_base: int = 10000, fp16_lm_cross_entropy: bool = False, - language_share_embeddings_and_output_weights: bool=False + language_share_embeddings_and_output_weights: bool = False, ) -> None: super().__init__(config=language_transformer_config) @@ -88,7 +88,9 @@ def __init__( self.vision_projection = None self.language_model = None - self.square_merge_size = vision_projection_config.ffn_hidden_size // vision_transformer_config.hidden_size + self.square_merge_size = ( + vision_projection_config.ffn_hidden_size // vision_transformer_config.hidden_size + ) # This attribute is needed to check if an all-reduce is required # on the word embeddings inside `finalize_model_grads._allreduce_word_embedding_grads`. @@ -101,7 +103,7 @@ def __init__( vision_projection_layer_spec, projection_type=vision_projection_type, pre_process=True, - post_process=True + post_process=True, ) self.language_model = QwenVLLanguageModel( @@ -116,7 +118,7 @@ def __init__( post_process=self.post_process, rotary_base=language_rotary_base, fp16_lm_cross_entropy=fp16_lm_cross_entropy, - share_embeddings_and_output_weights=language_share_embeddings_and_output_weights + share_embeddings_and_output_weights=language_share_embeddings_and_output_weights, ) self.share_embeddings_and_output_weights = ( self.language_model.share_embeddings_and_output_weights @@ -134,7 +136,7 @@ def set_input_tensor(self, input_tensor) -> None: # gives us non-lists or None if not isinstance(input_tensor, list): input_tensor = [input_tensor] - assert len(input_tensor) == 1, 'input_tensor should only be length 1 for Qwen2VL' + assert len(input_tensor) == 1, "input_tensor should only be length 1 for Qwen2VL" if self.pre_process: self.encoder_hidden_state = input_tensor[0] @@ -174,7 +176,6 @@ def forward( video_start_index: int = -1, image_input_mask: torch.Tensor = None, video_input_mask: torch.Tensor = None, - attention_mask: torch.Tensor = None, labels: torch.Tensor = None, inference_params: InferenceParams = None, @@ -216,8 +217,8 @@ def forward( elif self.config.fp16: vision_data = vision_data.to(torch.float16) vision_embeds = self.vision_model( - vision_data=vision_data, # If None, vision model should use intermediate outputs (EPP > 1) - grid_thw=vision_grid_thw # should provided in each EPP stage + vision_data=vision_data, # If None, vision model should use intermediate outputs (EPP > 1) + grid_thw=vision_grid_thw, # should provided in each EPP stage ) # If running inference, the language model KV cache will be updated for image token positions. @@ -231,8 +232,8 @@ def forward( # If running inference, we can skip image token computation if they were computed already earlier for this sample. if use_inference_kv_cache: language_embeddings: torch.Tensor = self.language_model.embedding( - input_ids=input_ids, - position_ids=None # NOTE: disable + input_ids=input_ids, + position_ids=None, # NOTE: disable ) # [text_seq_len, b, h_language] # NOTE: why not cat here? is it the combined embeddings useless? combined_embeddings = language_embeddings @@ -247,35 +248,37 @@ def forward( image_embeds = vision_embeds[:video_start_index] video_embeds = vision_embeds[video_start_index:] else: - raise ValueError(f"Expect video token start index in range [0, {vision_embeds.shape[0]}], but got {video_start_index}") + raise ValueError( + f"Expect video token start index in range [0, {vision_embeds.shape[0]}], but got {video_start_index}" + ) if image_embeds is not None: - image_input_mask = image_input_mask.T # shape [seqlen, mbs] + image_input_mask = image_input_mask.T # shape [seqlen, mbs] if video_embeds is not None: video_input_mask = video_input_mask.T combined_embeddings = self.language_model.embedding( input_ids=input_ids, - position_ids=None, # NOTE: disable + position_ids=None, # NOTE: disable image_input_mask=image_input_mask, video_input_mask=video_input_mask, image_embeds=image_embeds, - video_embeds=video_embeds + video_embeds=video_embeds, ) # [text_seq_len, b, h_language] else: combined_embeddings = self.language_model.embedding( input_ids=input_ids, - position_ids=None # NOTE: disable + position_ids=None, # NOTE: disable ) # [text_seq_len, b, h_language] else: combined_embeddings = None output = self.language_model( input_ids=None, - position_ids=position_ids, # None in encoder - attention_mask=attention_mask, # None in encoder - decoder_input=combined_embeddings, # only not None in the first decoder PP stage - labels=labels, # only not None in the last decoder PP stage - inference_params=inference_params, # currently always None - packed_seq_params=packed_seq_params, # currently always None + position_ids=position_ids, # None in encoder + attention_mask=attention_mask, # None in encoder + decoder_input=combined_embeddings, # only not None in the first decoder PP stage + labels=labels, # only not None in the last decoder PP stage + inference_params=inference_params, # currently always None + packed_seq_params=packed_seq_params, # currently always None **(extra_block_kwargs or {}), ) return output diff --git a/flagscale/models/megatron/qwen2_5_vl/tensor_parallel.py b/flagscale/models/megatron/qwen2_5_vl/tensor_parallel.py index bfb0222f82..b64c2330d0 100644 --- a/flagscale/models/megatron/qwen2_5_vl/tensor_parallel.py +++ b/flagscale/models/megatron/qwen2_5_vl/tensor_parallel.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Adopted from https://github.com/alibaba/Pai-Megatron-Patch/blob/8949a6647cbf6b39837ad3dd911fa4aa0726895b/megatron_patch/tensor_parallel.py # NOTE: we slightly modify this file to support zero-size tensor, such as [0, 128] when we don't have video data @@ -19,10 +33,10 @@ def _check_data_types(keys, data, target_dtype): """Check that all the keys have the same target data type.""" for key in keys: - assert ( - data[key].dtype == target_dtype - ), '{} has data type {} which ' 'is different than {}'.format( - key, data[key].dtype, target_dtype + assert data[key].dtype == target_dtype, ( + "{} has data type {} which is different than {}".format( + key, data[key].dtype, target_dtype + ) ) @@ -36,14 +50,14 @@ def _build_key_size_numel_dictionaries(keys, data): if get_tensor_model_parallel_rank() == 0: offset = 0 for key in keys: - assert data[key].dim() < max_dim, 'you should increase MAX_DATA_DIM' + assert data[key].dim() < max_dim, "you should increase MAX_DATA_DIM" size = data[key].size() for i, s in enumerate(size): sizes[i + offset] = s offset += max_dim # Move to GPU and broadcast. - sizes_cuda = torch.tensor(sizes, dtype=torch.long, device='cuda') + sizes_cuda = torch.tensor(sizes, dtype=torch.long, device="cuda") torch.distributed.broadcast( sizes_cuda, get_tensor_model_parallel_src_rank(), group=get_tensor_model_parallel_group() ) @@ -93,7 +107,9 @@ def broadcast_data(keys, data, datatype): # Flatten the data associated with the keys flatten_data = torch.cat([data[key].contiguous().view(-1) for key in keys], dim=0).cuda() else: - flatten_data = torch.empty(total_numel, device=cur_platform.current_device(), dtype=datatype) + flatten_data = torch.empty( + total_numel, device=cur_platform.current_device(), dtype=datatype + ) # Broadcast torch.distributed.broadcast( diff --git a/flagscale/models/megatron/qwen2_5_vl/transformer_config.py b/flagscale/models/megatron/qwen2_5_vl/transformer_config.py index c138dbefe8..164892daba 100644 --- a/flagscale/models/megatron/qwen2_5_vl/transformer_config.py +++ b/flagscale/models/megatron/qwen2_5_vl/transformer_config.py @@ -23,31 +23,33 @@ def get_vision_model_config(args, config): # mlp: hidden_size -> intermediate_size -> embed_dim, silu # NOTE: here we provide a workaround to solve the wrong layer amount when VPP of decoder is on - if config.num_layers in[28, 36]: - config.ffn_hidden_size = 3420 # 7B 72B + if config.num_layers in [28, 36]: + config.ffn_hidden_size = 3420 # 7B 72B else: - config.ffn_hidden_size = 3456 # 32B + config.ffn_hidden_size = 3456 # 32B if parallel_state.get_virtual_pipeline_model_parallel_world_size() is not None: - config.num_layers = 32 * parallel_state.get_virtual_pipeline_model_parallel_world_size() # depth + config.num_layers = ( + 32 * parallel_state.get_virtual_pipeline_model_parallel_world_size() + ) # depth else: - config.num_layers = 32 # depth - config.num_attention_heads = 16 # num_heads - config.add_bias_linear = True # all nn.Linear has bias (MLP, attn) - config.add_qkv_bias = True # qkv_proj in attn has bias - config.hidden_size = 1280 # hidden_size + config.num_layers = 32 # depth + config.num_attention_heads = 16 # num_heads + config.add_bias_linear = True # all nn.Linear has bias (MLP, attn) + config.add_qkv_bias = True # qkv_proj in attn has bias + config.hidden_size = 1280 # hidden_size config.hidden_dropout = 0.0 config.attention_dropout = 0.0 # config.gated_linear_unit = False # no gated # config.activation_func = quick_gelu # hidden_act config.kv_channels = config.hidden_size // config.num_attention_heads - config.num_query_groups = config.num_attention_heads # no GQA - config.layernorm_zero_centered_gamma = False # False - config.apply_query_key_layer_scaling = False # factor=math.sqrt(head_dim) - config.bias_activation_fusion = False # no swiglu, set false - config.bias_dropout_fusion = False # no dropout, set false - config.attention_softmax_in_fp32 = True # use True + config.num_query_groups = config.num_attention_heads # no GQA + config.layernorm_zero_centered_gamma = False # False + config.apply_query_key_layer_scaling = False # factor=math.sqrt(head_dim) + config.bias_activation_fusion = False # no swiglu, set false + config.bias_dropout_fusion = False # no dropout, set false + config.attention_softmax_in_fp32 = True # use True # config.normalization = 'LayerNorm' # use RMSNorm config.seq_length = args.seq_length @@ -68,9 +70,9 @@ def get_vision_model_config(args, config): config.num_layers_in_first_pipeline_stage = None config.num_layers_in_last_pipeline_stage = None if args.vision_recompute_layer_steps != 0: - config.recompute_method="uniform" - config.recompute_granularity="full" - config.recompute_num_layers=args.vision_recompute_layer_steps # 16 for 32B + config.recompute_method = "uniform" + config.recompute_granularity = "full" + config.recompute_num_layers = args.vision_recompute_layer_steps # 16 for 32B return config @@ -85,7 +87,7 @@ def get_vision_projection_config(config, embed_dim, spatial_merge_size): config.gated_linear_unit = False config.bias_activation_fusion = False config.add_bias_linear = True - config.ffn_hidden_size = embed_dim * (spatial_merge_size ** 2) + config.ffn_hidden_size = embed_dim * (spatial_merge_size**2) config.activation_func = torch.nn.functional.gelu config.tp_comm_overlap = False config.sequence_parallel = False diff --git a/flagscale/models/megatron/qwen2_5_vl/vision_attention.py b/flagscale/models/megatron/qwen2_5_vl/vision_attention.py index d2431acb35..f962e93c51 100644 --- a/flagscale/models/megatron/qwen2_5_vl/vision_attention.py +++ b/flagscale/models/megatron/qwen2_5_vl/vision_attention.py @@ -418,7 +418,7 @@ def forward( sequence_len_offset: Optional[int] = None, *, inference_params: Optional[BaseInferenceContext] = None, - rotary_pos_cos_sin = None, + rotary_pos_cos_sin=None, ) -> Tuple[Tensor, Tensor]: """ Perform a forward pass through the attention module. @@ -446,9 +446,9 @@ def forward( inference_context = deprecate_inference_params(inference_context, inference_params) if inference_context and inference_context.is_dynamic_batching(): - assert is_fa_min_version( - "2.7.3" - ), "flash attn verion v2.7.3 and above is required for dynamic batching." + assert is_fa_min_version("2.7.3"), ( + "flash attn verion v2.7.3 and above is required for dynamic batching." + ) # hidden_states: [sq, b, h] if self.config.flash_decode and not self.training and inference_context is not None: @@ -608,9 +608,9 @@ def forward( kv_lengths, block_table, ) - core_attn_out = rearrange(core_attn_out, 's b h d -> s b (h d)') + core_attn_out = rearrange(core_attn_out, "s b h d -> s b (h d)") - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + if packed_seq_params is not None and packed_seq_params.qkv_format == "thd": # reshape to same output shape as unpacked case # (t, np, hn) -> (t, b=1, h=np*hn) # t is the pack size = sum (sq_i) @@ -641,7 +641,6 @@ def __init__( attn_mask_type=AttnMaskType.padding, cp_comm_type: str = None, pg_collection: ProcessGroupCollection = None, - ): super().__init__( config=config, @@ -663,7 +662,7 @@ def __init__( bias=self.config.add_bias_linear or self.config.add_qkv_bias, skip_bias_add=False, is_expert=False, - tp_comm_buffer_name='qkv', + tp_comm_buffer_name="qkv", ) if submodules.q_layernorm is not None: @@ -785,12 +784,10 @@ def get_query_key_value_tensors(self, hidden_states, key_value_states=None): ] if SplitAlongDim is not None: - # [sq, b, ng, (np/ng + 2) * hn] # --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn] (query, key, value) = SplitAlongDim(mixed_qkv, 3, split_arg_list) else: - # [sq, b, ng, (np/ng + 2) * hn] # --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn] (query, key, value) = torch.split(mixed_qkv, split_arg_list, dim=3) diff --git a/flagscale/models/megatron/qwen2_5_vl/vision_transformer_block.py b/flagscale/models/megatron/qwen2_5_vl/vision_transformer_block.py index 3ba04864ae..0fb5be4395 100644 --- a/flagscale/models/megatron/qwen2_5_vl/vision_transformer_block.py +++ b/flagscale/models/megatron/qwen2_5_vl/vision_transformer_block.py @@ -126,7 +126,7 @@ def checkpoint_handler(forward_func): rotary_pos_emb, ) - if self.config.recompute_method == 'uniform': + if self.config.recompute_method == "uniform": # Uniformly divide the total number of Transformer layers and checkpoint # the input activation of each divided chunk. # A method to further reduce memory usage reducing checkpoints. @@ -138,7 +138,7 @@ def checkpoint_handler(forward_func): layer_idx += self.config.recompute_num_layers - elif self.config.recompute_method == 'block': + elif self.config.recompute_method == "block": # Checkpoint the input activation of only a set number of individual # Transformer layers and skip the rest. # A method fully use the device memory removing redundant re-computation. @@ -257,7 +257,7 @@ def forward( with rng_context, outer_fp8_context: # Forward pass. - if self.config.recompute_granularity == 'full' and self.training: + if self.config.recompute_granularity == "full" and self.training: hidden_states = self._checkpointed_forward( hidden_states=hidden_states, attention_mask=attention_mask, diff --git a/flagscale/models/megatron/qwen2_5_vl/vit_model.py b/flagscale/models/megatron/qwen2_5_vl/vit_model.py index deeb01e608..2b35288b0c 100644 --- a/flagscale/models/megatron/qwen2_5_vl/vit_model.py +++ b/flagscale/models/megatron/qwen2_5_vl/vit_model.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Mainly adopted from https://github.com/alibaba/Pai-Megatron-Patch/blob/8949a6647cbf6b39837ad3dd911fa4aa0726895b/megatron_patch/model/qwen2_5_vl/visionmodel.py. from typing import Optional @@ -49,7 +63,9 @@ def __init__( flat_dim = in_channels * temporal_patch_size * patch_size * patch_size self.proj = nn.Linear(flat_dim, embed_dim, bias=bias) else: - self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=self.kernel_size, stride=self.stride, bias=bias) + self.proj = nn.Conv3d( + in_channels, embed_dim, kernel_size=self.kernel_size, stride=self.stride, bias=bias + ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: """ @@ -67,8 +83,9 @@ def enable_linear(self): # and https://github.com/pytorch/pytorch/issues/166122 # and https://github.com/huggingface/transformers/pull/45041 # By default, we use CUDNN's convolution ops with optimization. - return self.kernel_size == self.stride and \ - version.parse(torch.__version__) > version.parse('2.9.0') + return self.kernel_size == self.stride and version.parse(torch.__version__) > version.parse( + "2.9.0" + ) def _forward_matmul(self, hidden_states): target_dtype = self.proj.weight.dtype @@ -83,11 +100,11 @@ def _forward_conv(self, hidden_states): hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim) return hidden_states + # copied from https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py class VisionRotaryEmbedding(nn.Module): - """ + """ """ - """ def __init__(self, dim: int, theta: float = 10000.0) -> None: super().__init__() # NOTE(lizhiyu): print inv_freq to check it. @@ -100,6 +117,7 @@ def forward(self, seqlen: int) -> torch.Tensor: freqs = torch.outer(seq, self.inv_freq) return freqs + # reference from https://github.com/huggingface/transformers/blob/0ad3710d4767d4ac7ee95f33f8554373e59efade/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py#L243 class Qwen2_5VisionModel(VisionModule): """Qwen2.5 ViT vision model. @@ -122,9 +140,8 @@ def __init__( projection_config: TransformerConfig, projection_layer_spec: ModuleSpec, projection_type: str = "mlp", - pre_process: bool = True, - post_process: bool = False + post_process: bool = False, ) -> None: super().__init__(config=transformer_config) @@ -164,7 +181,7 @@ def __init__( spec=transformer_layer_spec, pre_process=self.pre_process, post_process=self.post_process, - post_layer_norm=True + post_layer_norm=True, ) self.merge_hidden_size = projection_config.ffn_hidden_size @@ -175,7 +192,7 @@ def __init__( projection_config, projection_layer_spec, projection_type, - projection_config.ffn_hidden_size + projection_config.ffn_hidden_size, ) else: self.projection = None @@ -188,7 +205,7 @@ def set_input_tensor(self, input_tensor: torch.Tensor) -> None: Args: input_tensor (Tensor): Sets the input tensor for the model. """ - if self.pre_process: # always True + if self.pre_process: # always True self.input_tensor = input_tensor else: raise NotImplementedError() @@ -223,9 +240,9 @@ def rot_pos_emb(self, grid_thw): return rotary_pos_emb def get_window_index(self, grid_thw): - ''' + """ grid_thw: (tiles, 3) -> - ''' + """ window_index: list = [] cu_window_seqlens: list = [0] window_index_id = 0 @@ -234,13 +251,17 @@ def get_window_index(self, grid_thw): for grid_t, grid_h, grid_w in grid_thw: llm_grid_h, llm_grid_w = ( - grid_h // self.spatial_merge_size, # 224 // 2 = 112 + grid_h // self.spatial_merge_size, # 224 // 2 = 112 grid_w // self.spatial_merge_size, ) - index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w) - pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size # vit_merger_window_size = 4 + index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape( + grid_t, llm_grid_h, llm_grid_w + ) + pad_h = ( + vit_merger_window_size - llm_grid_h % vit_merger_window_size + ) # vit_merger_window_size = 4 pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size - num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size # 向上取整 + num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size # 向上取整 num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size index_padded = F.pad(index, (0, pad_w, 0, pad_h), "constant", -100) index_padded = index_padded.reshape( @@ -293,7 +314,7 @@ def forward( assert inference_params is None # Rotary positional embeddings (embedding is None for PP intermediate devices) - #vision_data (t, 3) --> (t, embed_dim) + # vision_data (t, 3) --> (t, embed_dim) vision_data = self.patch_embed(vision_data) # window_index: [tiles, num_windows] cu_window_seqlens: [tiles * num_windows] window_index, cu_window_seqlens = self.get_window_index(grid_thw) @@ -305,19 +326,23 @@ def forward( cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens) seq_len, _ = vision_data.size() - vision_data = vision_data.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + vision_data = vision_data.reshape( + seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1 + ) vision_data = vision_data[window_index, :, :] vision_data = vision_data.reshape(seq_len, 1, -1) rotary_pos_emb = self.rot_pos_emb(grid_thw) - rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + rotary_pos_emb = rotary_pos_emb.reshape( + seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1 + ) rotary_pos_emb = rotary_pos_emb[window_index, :, :] rotary_pos_emb = rotary_pos_emb.reshape(seq_len, 1, 1, -1).repeat(1, 1, 1, 2) hidden_states = self.decoder( - hidden_states = vision_data, - attention_mask = None, - inference_params = inference_params, + hidden_states=vision_data, + attention_mask=None, + inference_params=inference_params, rotary_pos_emb=rotary_pos_emb, packed_seq_params=self.build_packed_seq_params(None, cu_window_seqlens), packed_seq_params_full=self.build_packed_seq_params(grid_thw), @@ -339,14 +364,14 @@ def build_packed_seq_params( seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]) cu_seqlens = seqlens.cumsum(dim=0) cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0).int() - else: # the step of cu_seqlens is window_size, not sampel seq_length + else: # the step of cu_seqlens is window_size, not sampel seq_length seqlens = cu_seqlens[1:] - cu_seqlens[:-1] max_seqlen_q = seqlens.max() return PackedSeqParams( cu_seqlens_q=cu_seqlens, cu_seqlens_kv=cu_seqlens, - qkv_format='thd', + qkv_format="thd", max_seqlen_q=max_seqlen_q, - max_seqlen_kv=max_seqlen_q + max_seqlen_kv=max_seqlen_q, ) diff --git a/flagscale/models/megatron/qwen3_vl/__init__.py b/flagscale/models/megatron/qwen3_vl/__init__.py index e69de29bb2..a491ed2726 100644 --- a/flagscale/models/megatron/qwen3_vl/__init__.py +++ b/flagscale/models/megatron/qwen3_vl/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flagscale/models/megatron/qwen3_vl/language_model.py b/flagscale/models/megatron/qwen3_vl/language_model.py index e7ca8a5da4..7d3013f9a8 100644 --- a/flagscale/models/megatron/qwen3_vl/language_model.py +++ b/flagscale/models/megatron/qwen3_vl/language_model.py @@ -32,6 +32,7 @@ from flagscale.models.megatron.qwen2_5_vl.language_module import QwenVLLanguageModelEmbedding + ######### New Impl for Qwen3-VL Language Model ######### # reference from https://github.com/huggingface/transformers/blob/d08b98b965176ea9cf8c8e8b24995c955b7e2ec9/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py#L278 class Qwen3VLLanguageRotaryEmbedding(torch.nn.Module): @@ -95,8 +96,8 @@ def apply_interleaved_mrope(self, freqs, mrope_section): """ freqs_t = freqs[0] # just overwrite the first dimension T for dim, offset in enumerate((1, 2), start=1): # H, W - length = mrope_section[dim] * 3 # 20 * 3 = 60 - idx = slice(offset, length, 3) # 1, 4, 7, ..., 59 for H; 2, 5, 8, ..., 59 for W + length = mrope_section[dim] * 3 # 20 * 3 = 60 + idx = slice(offset, length, 3) # 1, 4, 7, ..., 59 for H; 2, 5, 8, ..., 59 for W freqs_t[..., idx] = freqs[dim, ..., idx] return freqs_t @@ -111,13 +112,17 @@ def forward(self, position_ids: torch.Tensor, mrope_section: List[int]) -> Tenso Returns: Tensor: Embeddings after applying RoPE. """ - seq = position_ids.to(device=self.inv_freq.device, dtype=torch.float32) # shape (3, bs, seq_length) + seq = position_ids.to( + device=self.inv_freq.device, dtype=torch.float32 + ) # shape (3, bs, seq_length) if self.seq_len_interpolation_factor is not None: seq *= 1 / self.seq_len_interpolation_factor # shape (3, bs, dim, 1) - inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, seq.shape[1], -1, 1) + inv_freq_expanded = ( + self.inv_freq[None, None, :, None].float().expand(3, seq.shape[1], -1, 1) + ) # shape (3, bs, 1, seq_length) seq_expanded = seq[:, :, None, :].float() # shape (3, bs, seq_length, dim) @@ -143,6 +148,7 @@ def forward(self, position_ids: torch.Tensor, mrope_section: List[int]) -> Tenso emb = get_pos_emb_on_this_cp_rank(emb, 0, self.cp_group) return emb + class Qwen3VLLanguageModule(GPTModel): """Qwen3-VL Language Module. @@ -174,8 +180,8 @@ def __init__( parallel_output: bool = True, share_embeddings_and_output_weights: bool = False, position_embedding_type: Literal[ - 'learned_absolute', 'rope', 'mrope', 'none' - ] = 'learned_absolute', + "learned_absolute", "rope", "mrope", "none" + ] = "learned_absolute", rotary_percent: float = 1.0, rotary_base: int = 10000, rope_scaling: bool = False, @@ -201,7 +207,7 @@ def __init__( self.share_embeddings_and_output_weights = share_embeddings_and_output_weights self.vp_stage = vp_stage - if hasattr(self.config, 'position_embedding_type'): + if hasattr(self.config, "position_embedding_type"): self.position_embedding_type = self.config.position_embedding_type else: self.position_embedding_type = position_embedding_type @@ -214,7 +220,7 @@ def __init__( self.max_position_embeddings = max_sequence_length self.rotary_percent = rotary_percent - if hasattr(self.config, 'rotary_base'): + if hasattr(self.config, "rotary_base"): self.rotary_base = self.config.rotary_base else: self.rotary_base = rotary_base @@ -229,7 +235,7 @@ def __init__( max_sequence_length=self.max_sequence_length, position_embedding_type=position_embedding_type, ) - if self.position_embedding_type == 'mrope' and not self.config.multi_latent_attention: + if self.position_embedding_type == "mrope" and not self.config.multi_latent_attention: self.rotary_pos_emb = Qwen3VLLanguageRotaryEmbedding( kv_channels=self.config.kv_channels, rotary_percent=rotary_percent, @@ -238,10 +244,9 @@ def __init__( rotary_base=rotary_base, ) self.mrope_section = self.config.mrope_section - assert ( - self.mrope_section is not None - ), "mrope require mrope_section setting, but we got None from TransformerConfig" - + assert self.mrope_section is not None, ( + "mrope require mrope_section setting, but we got None from TransformerConfig" + ) # Cache for RoPE tensors which do not change between iterations. self.rotary_pos_emb_cache = {} @@ -263,7 +268,6 @@ def __init__( # Output if self.post_process: - if self.config.defer_embedding_wgrad_compute: # The embedding activation buffer preserves a reference to the input activations # of the final embedding projection layer GEMM. It will hold the activations for @@ -299,23 +303,31 @@ def __init__( if has_config_logger_enabled(self.config): log_config_to_disk( - self.config, self.state_dict(), prefix=f'{type(self).__name__}_init_ckpt' + self.config, self.state_dict(), prefix=f"{type(self).__name__}_init_ckpt" ) for name, module in self.named_modules(): - if hasattr(module, 'finish_init'): + if hasattr(module, "finish_init"): quant_config = get_quant_config_or_none(name, self.config.quant_recipe) module.finish_init(quant_config) - def forward(self, input_ids, position_ids, attention_mask, - decoder_input = None, labels = None, inference_context = None, - packed_seq_params = None, extra_block_kwargs = None, - runtime_gather_output = None, - # args for deepstack - visual_pos_masks: Optional[torch.Tensor] = None, - deepstack_visual_embeds: Optional[list[torch.Tensor]] = None, - *, inference_params = None, - loss_mask = None): - + def forward( + self, + input_ids, + position_ids, + attention_mask, + decoder_input=None, + labels=None, + inference_context=None, + packed_seq_params=None, + extra_block_kwargs=None, + runtime_gather_output=None, + # args for deepstack + visual_pos_masks: Optional[torch.Tensor] = None, + deepstack_visual_embeds: Optional[list[torch.Tensor]] = None, + *, + inference_params=None, + loss_mask=None, + ): inference_context = deprecate_inference_params(inference_context, inference_params) decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset = ( @@ -338,8 +350,8 @@ def forward(self, input_ids, position_ids, attention_mask, rotary_pos_sin=rotary_pos_sin, packed_seq_params=packed_seq_params, sequence_len_offset=sequence_len_offset, - visual_pos_masks = visual_pos_masks, - deepstack_visual_embeds = deepstack_visual_embeds, + visual_pos_masks=visual_pos_masks, + deepstack_visual_embeds=deepstack_visual_embeds, **(extra_block_kwargs or {}), ) diff --git a/flagscale/models/megatron/qwen3_vl/language_transformer_block.py b/flagscale/models/megatron/qwen3_vl/language_transformer_block.py index 0a86705c43..7a6b60cc3a 100755 --- a/flagscale/models/megatron/qwen3_vl/language_transformer_block.py +++ b/flagscale/models/megatron/qwen3_vl/language_transformer_block.py @@ -23,7 +23,10 @@ get_pg_rank, make_viewless_tensor, ) -from megatron.core.tensor_parallel.mappings import (gather_from_sequence_parallel_region, scatter_to_sequence_parallel_region) +from megatron.core.tensor_parallel.mappings import ( + gather_from_sequence_parallel_region, + scatter_to_sequence_parallel_region, +) try: import transformer_engine.pytorch as te # pylint: disable=unused-import @@ -61,7 +64,6 @@ class LanguageTransformerBlock(TransformerBlock): - def _checkpointed_forward( self, hidden_states: Tensor, @@ -142,7 +144,7 @@ def checkpoint_handler(forward_func): rotary_pos_emb, ) - if self.config.recompute_method == 'uniform': + if self.config.recompute_method == "uniform": # Uniformly divide the total number of Transformer layers and checkpoint # the input activation of each divided chunk. # A method to further reduce memory usage reducing checkpoints. @@ -154,9 +156,11 @@ def checkpoint_handler(forward_func): # layer_idx += self.config.recompute_num_layers # NOTE: Assume that this is first pipeline stage that has at least three layers. - # The other stages will just pass None for visual_pos_masks and deepstack_visual_embeds. + # The other stages will just pass None for visual_pos_masks and deepstack_visual_embeds. if visual_pos_masks is not None and deepstack_visual_embeds is not None: - assert len(self.layers) >= len(deepstack_visual_embeds), f"First pipeline stage should have at least {len(self.layers)} layers for deepstack." + assert len(self.layers) >= len(deepstack_visual_embeds), ( + f"First pipeline stage should have at least {len(self.layers)} layers for deepstack." + ) if layer_idx < len(deepstack_visual_embeds): hidden_states = self._deepstack_process( hidden_states, @@ -166,7 +170,7 @@ def checkpoint_handler(forward_func): layer_idx += self.config.recompute_num_layers - elif self.config.recompute_method == 'block': + elif self.config.recompute_method == "block": # Checkpoint the input activation of only a set number of individual # Transformer layers and skip the rest. # A method fully use the device memory removing redundant re-computation. @@ -192,8 +196,8 @@ def checkpoint_handler(forward_func): return hidden_states - """Transformer class.""" + def forward( self, hidden_states: Union[Tensor, WrappedTensor], @@ -298,9 +302,13 @@ def forward( with rng_context, outer_quantization_context: # Forward pass. - if self.config.recompute_granularity == 'full' and self.training: - assert self.config.recompute_method == 'uniform' and self.config.recompute_num_layers == 1, \ + if self.config.recompute_granularity == "full" and self.training: + assert ( + self.config.recompute_method == "uniform" + and self.config.recompute_num_layers == 1 + ), ( f"Only uniform recompute with recompute_num_layers=1 is supported for full recompute in Qwen3-VL." + ) hidden_states = self._checkpointed_forward( hidden_states=hidden_states, attention_mask=attention_mask, @@ -310,8 +318,8 @@ def forward( attention_bias=attention_bias, packed_seq_params=packed_seq_params, use_inner_quantization_context=use_inner_quantization_context, - visual_pos_masks = visual_pos_masks, - deepstack_visual_embeds = deepstack_visual_embeds, + visual_pos_masks=visual_pos_masks, + deepstack_visual_embeds=deepstack_visual_embeds, ) else: for l_no, layer in enumerate(self.layers): @@ -348,7 +356,9 @@ def forward( # NOTE: Assume that this is first pipeline stage that has at least three layers. # The other stages will just pass None for visual_pos_masks and deepstack_visual_embeds. if visual_pos_masks is not None and deepstack_visual_embeds is not None: - assert len(self.layers) >= len(deepstack_visual_embeds), f"First pipeline stage should have at least {len(self.layers)} layers for deepstack." + assert len(self.layers) >= len(deepstack_visual_embeds), ( + f"First pipeline stage should have at least {len(self.layers)} layers for deepstack." + ) if l_no < len(deepstack_visual_embeds): hidden_states = self._deepstack_process( hidden_states, @@ -379,9 +389,11 @@ def forward( return hidden_states - def _deepstack_process( - self, hidden_states: torch.Tensor, visual_pos_masks: torch.Tensor, visual_embeds: torch.Tensor + self, + hidden_states: torch.Tensor, + visual_pos_masks: torch.Tensor, + visual_embeds: torch.Tensor, ): visual_pos_masks = visual_pos_masks.to(hidden_states.device) visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype) diff --git a/flagscale/models/megatron/qwen3_vl/layer_specs.py b/flagscale/models/megatron/qwen3_vl/layer_specs.py index 03da284f81..637d8ddaf2 100644 --- a/flagscale/models/megatron/qwen3_vl/layer_specs.py +++ b/flagscale/models/megatron/qwen3_vl/layer_specs.py @@ -25,7 +25,7 @@ TERowParallelLinear, TEColumnParallelLinear, TEColumnParallelGroupedLinear, - TERowParallelGroupedLinear + TERowParallelGroupedLinear, ) from megatron.core.transformer.enums import AttnMaskType @@ -45,11 +45,13 @@ from megatron.core.transformer.moe.experts import TEGroupedMLP from flagscale.models.megatron.qwen3_vl.vision_attention import VisionSelfAttention + + # Use this spec to use lower level Transformer Engine modules (required for fp8 training) def get_gpt_layer_with_transformer_engine_spec( num_experts: Optional[int] = None, moe_grouped_gemm: Optional[bool] = False, - qk_layernorm: bool = False + qk_layernorm: bool = False, ) -> ModuleSpec: mlp = get_mlp_module_spec( use_te=True, num_experts=num_experts, moe_grouped_gemm=moe_grouped_gemm @@ -75,10 +77,9 @@ def get_gpt_layer_with_transformer_engine_spec( ), ) -def get_qwen3vl_vision_model_spec( - is_vit=True -) -> ModuleSpec: - attn_mask_type = AttnMaskType.no_mask # THD full attention + +def get_qwen3vl_vision_model_spec(is_vit=True) -> ModuleSpec: + attn_mask_type = AttnMaskType.no_mask # THD full attention mlp = ModuleSpec( module=MLP, @@ -111,7 +112,10 @@ def get_qwen3vl_vision_model_spec( # Helper function to get module spec for MLP/MoE def get_mlp_module_spec( - use_te: bool = True, num_experts: int = None, moe_grouped_gemm: bool = False, add_norm: bool = True + use_te: bool = True, + num_experts: int = None, + moe_grouped_gemm: bool = False, + add_norm: bool = True, ) -> ModuleSpec: assert use_te, "Only Transformer Engine backend is supported" if num_experts is None: @@ -141,8 +145,9 @@ def get_mlp_module_spec( experts=ModuleSpec( module=TEGroupedMLP, submodules=MLPSubmodules( - linear_fc1=TEColumnParallelGroupedLinear, linear_fc2=TERowParallelGroupedLinear - ) + linear_fc1=TEColumnParallelGroupedLinear, + linear_fc2=TERowParallelGroupedLinear, + ), ), - ) + ), ) diff --git a/flagscale/models/megatron/qwen3_vl/model.py b/flagscale/models/megatron/qwen3_vl/model.py index 0392163fdc..59d16e380b 100644 --- a/flagscale/models/megatron/qwen3_vl/model.py +++ b/flagscale/models/megatron/qwen3_vl/model.py @@ -19,6 +19,7 @@ from .vision_model import Qwen3VisionModel from .language_model import Qwen3VLLanguageModule + # Note: This is under development and may be missing features. class Qwen3VLModel(MegatronModule): """Qwen3VL multi-modal model. @@ -61,9 +62,8 @@ def __init__( vision_projection_config: TransformerConfig, vision_projection_layer_spec: ModuleSpec, vision_projection_type: str = "mlp", - parallel_output: bool = True, - language_position_embedding_type: str = 'mrope', + language_position_embedding_type: str = "mrope", language_rotary_percent: float = 1.0, pre_process: bool = True, post_process: bool = True, @@ -71,8 +71,8 @@ def __init__( add_decoder: bool = True, language_rotary_base: int = 10000, fp16_lm_cross_entropy: bool = False, - language_share_embeddings_and_output_weights: bool=False, - vp_stage: int=None + language_share_embeddings_and_output_weights: bool = False, + vp_stage: int = None, ) -> None: super().__init__(config=language_transformer_config) @@ -90,7 +90,9 @@ def __init__( self.vision_projection = None self.language_model = None - self.square_merge_size = vision_projection_config.ffn_hidden_size // vision_transformer_config.hidden_size + self.square_merge_size = ( + vision_projection_config.ffn_hidden_size // vision_transformer_config.hidden_size + ) # This attribute is needed to check if an all-reduce is required # on the word embeddings inside `finalize_model_grads._allreduce_word_embedding_grads`. @@ -103,7 +105,7 @@ def __init__( vision_projection_layer_spec, projection_type=vision_projection_type, pre_process=True, - post_process=True + post_process=True, ) self.language_model = Qwen3VLLanguageModule( @@ -139,7 +141,7 @@ def set_input_tensor(self, input_tensor) -> None: # gives us non-lists or None if not isinstance(input_tensor, list): input_tensor = [input_tensor] - assert len(input_tensor) == 1, 'input_tensor should only be length 1 for Qwen2VL' + assert len(input_tensor) == 1, "input_tensor should only be length 1 for Qwen2VL" if self.pre_process: self.encoder_hidden_state = input_tensor[0] @@ -177,10 +179,8 @@ def forward( vision_data: torch.Tensor = None, vision_grid_thw: torch.Tensor = None, video_start_index: int = -1, - image_input_mask: torch.Tensor = None, video_input_mask: torch.Tensor = None, - attention_mask: torch.Tensor = None, labels: torch.Tensor = None, inference_params: InferenceParams = None, @@ -217,8 +217,8 @@ def forward( deepstack_feature_lists = None if vision_grid_thw.shape[0] > 0: vision_embeds, deepstack_feature_lists = self.vision_model( - vision_data=vision_data, # If None, vision model should use intermediate outputs (EPP > 1) - grid_thw=vision_grid_thw # should provided in each EPP stage + vision_data=vision_data, # If None, vision model should use intermediate outputs (EPP > 1) + grid_thw=vision_grid_thw, # should provided in each EPP stage ) # If running inference, the language model KV cache will be updated for image token positions. @@ -233,13 +233,13 @@ def forward( if use_inference_kv_cache: language_embeddings: torch.Tensor = self.language_model.embedding( input_ids=input_ids, - position_ids=None # NOTE: disable + position_ids=None, # NOTE: disable ) # [text_seq_len, b, h_language] # NOTE: why not cat here? is it the combined embeddings useless? combined_embeddings = language_embeddings elif vision_embeds is not None: if image_input_mask is not None: - image_input_mask = image_input_mask.T # shape [seqlen, mbs] + image_input_mask = image_input_mask.T # shape [seqlen, mbs] if video_input_mask is not None: video_input_mask = video_input_mask.T @@ -256,20 +256,22 @@ def forward( video_embeds = vision_embeds[video_start_index:] visual_pos_masks = torch.logical_or(image_input_mask, video_input_mask) else: - raise ValueError(f"Expect video token start index in range [0, {vision_embeds.shape[0]}], but got {video_start_index}") + raise ValueError( + f"Expect video token start index in range [0, {vision_embeds.shape[0]}], but got {video_start_index}" + ) combined_embeddings = self.language_model.embedding( input_ids=input_ids, - position_ids=None, # NOTE: disable + position_ids=None, # NOTE: disable image_input_mask=image_input_mask, video_input_mask=video_input_mask, image_embeds=image_embeds, - video_embeds=video_embeds + video_embeds=video_embeds, ) # [text_seq_len, b, h_language] else: combined_embeddings = self.language_model.embedding( input_ids=input_ids, - position_ids=None # NOTE: disable + position_ids=None, # NOTE: disable ) # [text_seq_len, b, h_language] visual_pos_masks = None deepstack_feature_lists = None @@ -281,12 +283,12 @@ def forward( output = self.language_model( input_ids=None, - position_ids=position_ids, # None in encoder - attention_mask=attention_mask, # None in encoder - decoder_input=combined_embeddings, # only not None in the first decoder PP stage - labels=labels, # only not None in the last decoder PP stage - inference_params=inference_params, # currently always None - packed_seq_params=packed_seq_params, # currently always None + position_ids=position_ids, # None in encoder + attention_mask=attention_mask, # None in encoder + decoder_input=combined_embeddings, # only not None in the first decoder PP stage + labels=labels, # only not None in the last decoder PP stage + inference_params=inference_params, # currently always None + packed_seq_params=packed_seq_params, # currently always None visual_pos_masks=visual_pos_masks, deepstack_visual_embeds=deepstack_feature_lists, **(extra_block_kwargs or {}), @@ -373,24 +375,47 @@ def get_rope_index( text_len = ed - st st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) # t_index is always 0 because llm_grid_t is always 1 (we use timestamps to encode the temporal information for videos) - t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) + t_index = ( + torch.arange(llm_grid_t) + .view(-1, 1) + .expand(-1, llm_grid_h * llm_grid_w) + .flatten() + ) + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + text_len + st_idx + ) st = ed + llm_grid_t * llm_grid_h * llm_grid_w if st < len(input_tokens): st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 text_len = len(input_tokens) - st - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) - mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) + mrope_position_deltas = torch.tensor( + mrope_position_deltas, device=input_ids.device + ).unsqueeze(1) return position_ids, mrope_position_deltas else: if attention_mask is not None: diff --git a/flagscale/models/megatron/qwen3_vl/transformer_config.py b/flagscale/models/megatron/qwen3_vl/transformer_config.py index 86342ef213..ce99df3e1d 100644 --- a/flagscale/models/megatron/qwen3_vl/transformer_config.py +++ b/flagscale/models/megatron/qwen3_vl/transformer_config.py @@ -1,4 +1,3 @@ - # Copyright (c) 2025, BAAI. All rights reserved. # # Adopted from https://github.com/alibaba/Pai-Megatron-Patch/blob/main/megatron_patch/model/qwen3_vl/transformer_config.py @@ -27,8 +26,7 @@ @dataclass class Qwen3VLTransformerConfig(TransformerConfig): - - transformer_impl: str = 'transformer_engine' + transformer_impl: str = "transformer_engine" rotary_base: int = None rotary_scaling_factor: int = None max_position_embeddings: int = None @@ -47,7 +45,6 @@ class Qwen3VLTransformerConfig(TransformerConfig): # add_qkv_bias = True - def get_vision_model_config(args, config): # Given a Transformer Config from decoder, build vision encoder config # diff: out_hidden_size & intermediate_size @@ -56,35 +53,37 @@ def get_vision_model_config(args, config): assert parallel_state.get_virtual_pipeline_model_parallel_world_size() is None, "NotSupported" if args.num_layers == 36 and args.hidden_size == 2560: # 4B - config.num_layers = 24 # depth - config.hidden_size = 1024 # hidden_size + config.num_layers = 24 # depth + config.hidden_size = 1024 # hidden_size config.ffn_hidden_size = 4096 - config.deepstack_visual_indexes=[5, 11, 17] + config.deepstack_visual_indexes = [5, 11, 17] else: # 8B, 32B, 30BA3B, 235BA22B - config.num_layers = 27 # depth - config.hidden_size = 1152 # hidden_size + config.num_layers = 27 # depth + config.hidden_size = 1152 # hidden_size config.ffn_hidden_size = 4304 - config.deepstack_visual_indexes=[8, 16, 24] + config.deepstack_visual_indexes = [8, 16, 24] - config.num_attention_heads = 16 # num_heads - config.add_bias_linear = True # all nn.Linear has bias (MLP, attn) - config.add_qkv_bias = True # qkv_proj in attn has bias + config.num_attention_heads = 16 # num_heads + config.add_bias_linear = True # all nn.Linear has bias (MLP, attn) + config.add_qkv_bias = True # qkv_proj in attn has bias config.hidden_dropout = 0.0 config.attention_dropout = 0.0 - config.gated_linear_unit = False # no gated + config.gated_linear_unit = False # no gated # NOTE(lizhyu): different from Qwen2.5-VL, use GELU here vs "gelu_pytorch_tanh" in huggingface implementation. - config.activation_func = functools.partial(torch.nn.functional.gelu, approximate="tanh") # Using GELU activation not SiLU + config.activation_func = functools.partial( + torch.nn.functional.gelu, approximate="tanh" + ) # Using GELU activation not SiLU config.kv_channels = config.hidden_size // config.num_attention_heads - config.num_query_groups = config.num_attention_heads # no GQA - config.layernorm_zero_centered_gamma = False # False - config.apply_query_key_layer_scaling = False # factor=math.sqrt(head_dim) - config.bias_activation_fusion = False # no swiglu, set false - config.bias_dropout_fusion = False # no dropout, set false - config.attention_softmax_in_fp32 = True # use True - config.normalization = 'LayerNorm' # use LayerNorm + config.num_query_groups = config.num_attention_heads # no GQA + config.layernorm_zero_centered_gamma = False # False + config.apply_query_key_layer_scaling = False # factor=math.sqrt(head_dim) + config.bias_activation_fusion = False # no swiglu, set false + config.bias_dropout_fusion = False # no dropout, set false + config.attention_softmax_in_fp32 = True # use True + config.normalization = "LayerNorm" # use LayerNorm config.seq_length = args.seq_length config.tp_comm_overlap = False @@ -103,8 +102,8 @@ def get_vision_model_config(args, config): # Reset recompute settings for vision encoder, otherwise it may use the language's settings. if args.vision_recompute_activations: - config.recompute_granularity = 'full' - config.recompute_method = 'uniform' + config.recompute_granularity = "full" + config.recompute_method = "uniform" config.recompute_num_layers = 1 return config @@ -122,8 +121,8 @@ def get_vision_projection_config(config, embed_dim, spatial_merge_size): config.gated_linear_unit = False config.bias_activation_fusion = False config.add_bias_linear = True - config.ffn_hidden_size = embed_dim * (spatial_merge_size ** 2) - config.activation_func = torch.nn.GELU() # Using GELU activation not SiLU + config.ffn_hidden_size = embed_dim * (spatial_merge_size**2) + config.activation_func = torch.nn.GELU() # Using GELU activation not SiLU config.tp_comm_overlap = False config.sequence_parallel = False return config diff --git a/flagscale/models/megatron/qwen3_vl/vision_attention.py b/flagscale/models/megatron/qwen3_vl/vision_attention.py index 6b9cf812f0..5b139d1406 100644 --- a/flagscale/models/megatron/qwen3_vl/vision_attention.py +++ b/flagscale/models/megatron/qwen3_vl/vision_attention.py @@ -121,7 +121,7 @@ def __init__( bias=self.config.add_bias_linear or self.config.add_qkv_bias, skip_bias_add=False, is_expert=False, - tp_comm_buffer_name='qkv', + tp_comm_buffer_name="qkv", tp_group=self.pg_collection.tp, ) @@ -135,9 +135,9 @@ def __init__( ) else: tp_world_size = get_tensor_model_parallel_world_size() - assert ( - tp_world_size <= 1 - ), "TP world size must be less than 1 for qk_layernorm_hidden_dim" + assert tp_world_size <= 1, ( + "TP world size must be less than 1 for qk_layernorm_hidden_dim" + ) # nums_head_cur_rank = divide(self.config.num_attention_heads, tp_world_size) self.q_layernorm = build_module( submodules.q_layernorm, @@ -158,9 +158,9 @@ def __init__( ) else: tp_world_size = get_tensor_model_parallel_world_size() - assert ( - tp_world_size <= 1 - ), "TP world size must be less than 1 for qk_layernorm_hidden_dim" + assert tp_world_size <= 1, ( + "TP world size must be less than 1 for qk_layernorm_hidden_dim" + ) # nums_head_cur_rank = divide(self.config.num_attention_heads, tp_world_size) self.k_layernorm = build_module( submodules.k_layernorm, @@ -270,12 +270,10 @@ def get_query_key_value_tensors(self, hidden_states, key_value_states=None): ] if SplitAlongDim is not None: - # [sq, b, ng, (np/ng + 2) * hn] # --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn] (query, key, value) = SplitAlongDim(mixed_qkv, 3, split_arg_list) else: - # [sq, b, ng, (np/ng + 2) * hn] # --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn] (query, key, value) = torch.split(mixed_qkv, split_arg_list, dim=3) @@ -357,9 +355,9 @@ def forward( inference_context = deprecate_inference_params(inference_context, inference_params) if inference_context and inference_context.is_dynamic_batching(): - assert HAVE_FA3 or is_fa_min_version( - "2.7.3" - ), "flash attn verion v2.7.3 and above is required for dynamic batching." + assert HAVE_FA3 or is_fa_min_version("2.7.3"), ( + "flash attn verion v2.7.3 and above is required for dynamic batching." + ) # hidden_states: [sq, b, h] if self.config.flash_decode and not self.training and inference_context is not None: @@ -535,9 +533,9 @@ def forward( kv_lengths, block_table, ) - core_attn_out = rearrange(core_attn_out, 's b h d -> s b (h d)') + core_attn_out = rearrange(core_attn_out, "s b h d -> s b (h d)") - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + if packed_seq_params is not None and packed_seq_params.qkv_format == "thd": # reshape to same output shape as unpacked case # (t, np, hn) -> (t, b=1, h=np*hn) # t is the pack size = sum (sq_i) diff --git a/flagscale/models/megatron/qwen3_vl/vision_model.py b/flagscale/models/megatron/qwen3_vl/vision_model.py index 9811c6004a..52292ab58b 100644 --- a/flagscale/models/megatron/qwen3_vl/vision_model.py +++ b/flagscale/models/megatron/qwen3_vl/vision_model.py @@ -19,6 +19,7 @@ from flagscale.models.megatron.qwen2_5_vl.vit_model import PatchEmbed from flagscale.models.megatron.qwen3_vl.vision_transformer_block import VisionTransformerBlock + # reference from https://github.com/huggingface/transformers/blob/d08b98b965176ea9cf8c8e8b24995c955b7e2ec9/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py#L79 # this is different from RotaryEmbedding in megatron.core.transformer.rotary_embedding. # Don't cat the even and odd parts here. @@ -59,9 +60,8 @@ def __init__( projection_config: TransformerConfig, projection_layer_spec: ModuleSpec, projection_type: str = "mlp", - pre_process: bool = True, - post_process: bool = False + post_process: bool = False, ) -> None: super().__init__(config=transformer_config) @@ -83,10 +83,12 @@ def __init__( temporal_patch_size=temporal_patch_size, in_channels=in_channels, embed_dim=embed_dim, - bias=True + bias=True, ) - self.pos_embed = nn.Embedding(transformer_config.num_position_embeddings, transformer_config.hidden_size) + self.pos_embed = nn.Embedding( + transformer_config.num_position_embeddings, transformer_config.hidden_size + ) self.num_grid_per_side = int(transformer_config.num_position_embeddings**0.5) head_dim = embed_dim // num_heads @@ -105,7 +107,6 @@ def __init__( pre_process=self.pre_process, post_process=self.post_process, post_layer_norm=True, - # NOTE: for deepstack projection_config=projection_config, projection_layer_spec=projection_layer_spec, @@ -121,7 +122,7 @@ def __init__( projection_config, projection_layer_spec, projection_type, - projection_config.ffn_hidden_size + projection_config.ffn_hidden_size, ) self.input_tensor = None @@ -138,7 +139,7 @@ def set_input_tensor(self, input_tensor: Tensor) -> None: if not isinstance(input_tensor, list): input_tensor = [input_tensor] - assert len(input_tensor) == 1, 'input_tensor should only be length 1 for gpt/bert' + assert len(input_tensor) == 1, "input_tensor should only be length 1 for gpt/bert" self.decoder.set_input_tensor(input_tensor[0]) def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor: @@ -270,15 +271,19 @@ def forward( vision_data = vision_data + pos_embeds seq_len, _ = vision_data.size() - vision_data = vision_data.reshape(seq_len, 1, -1) # change to (seq_len, batch_size=1, hidden_size) + vision_data = vision_data.reshape( + seq_len, 1, -1 + ) # change to (seq_len, batch_size=1, hidden_size) rotary_pos_emb = self.rot_pos_emb(grid_thw) - rotary_pos_emb = rotary_pos_emb.reshape(seq_len, 1, 1, -1).repeat(1, 1, 1, 2) # change to (seq_len, batch_size=1, 1, dim), and repeat for even adn odd + rotary_pos_emb = rotary_pos_emb.reshape(seq_len, 1, 1, -1).repeat( + 1, 1, 1, 2 + ) # change to (seq_len, batch_size=1, 1, dim), and repeat for even adn odd hidden_states, deepstack_feature_lists = self.decoder( - hidden_states = vision_data, - attention_mask = None, - inference_params = inference_params, + hidden_states=vision_data, + attention_mask=None, + inference_params=inference_params, rotary_pos_emb=rotary_pos_emb, packed_seq_params=self.build_packed_seq_params(grid_thw), **(extra_block_kwargs or {}), @@ -291,7 +296,6 @@ def forward( def build_packed_seq_params( self, grid_thw: Optional[torch.Tensor], - ) -> PackedSeqParams: # NOTE: each frame is a sequence (rather than each grid) seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]) @@ -301,7 +305,7 @@ def build_packed_seq_params( return PackedSeqParams( cu_seqlens_q=cu_seqlens, cu_seqlens_kv=cu_seqlens, - qkv_format='thd', + qkv_format="thd", max_seqlen_q=max_seqlen_q, - max_seqlen_kv=max_seqlen_q + max_seqlen_kv=max_seqlen_q, ) diff --git a/flagscale/models/megatron/qwen3_vl/vision_rope_utils.py b/flagscale/models/megatron/qwen3_vl/vision_rope_utils.py index 346cbce5aa..a321c80516 100644 --- a/flagscale/models/megatron/qwen3_vl/vision_rope_utils.py +++ b/flagscale/models/megatron/qwen3_vl/vision_rope_utils.py @@ -36,12 +36,12 @@ __all__ = [ - 'apply_rotary_pos_emb', - 'apply_rotary_emb_flash', - 'apply_rotary_pos_emb_with_cos_sin', - 'fused_apply_rotary_pos_emb', - 'fused_apply_rotary_pos_emb_thd', - 'get_pos_emb_on_this_cp_rank', + "apply_rotary_pos_emb", + "apply_rotary_emb_flash", + "apply_rotary_pos_emb_with_cos_sin", + "fused_apply_rotary_pos_emb", + "fused_apply_rotary_pos_emb_thd", + "get_pos_emb_on_this_cp_rank", ] @@ -88,6 +88,7 @@ def _rotate_half(x: Tensor, rotary_interleaved: bool) -> Tensor: x_new = torch.stack((-x2, x1), dim=-1) return x_new.view(x_new.shape[0], x_new.shape[1], x_new.shape[2], -1) + # NOTE(lizhiyu): The difference from the original _apply_rotary_pos_emb_bshd is that we use float not half. def _apply_rotary_pos_emb_bshd( t: Tensor, diff --git a/flagscale/models/megatron/qwen3_vl/vision_transformer_block.py b/flagscale/models/megatron/qwen3_vl/vision_transformer_block.py index 39f3c60f61..0d67dc289d 100755 --- a/flagscale/models/megatron/qwen3_vl/vision_transformer_block.py +++ b/flagscale/models/megatron/qwen3_vl/vision_transformer_block.py @@ -24,6 +24,7 @@ ) from megatron.core.models.vision.multimodal_projector import MultimodalProjector from megatron.core.transformer.spec_utils import ModuleSpec, build_module + try: import transformer_engine.pytorch as te # pylint: disable=unused-import @@ -65,19 +66,31 @@ class VisionTransformerBlock(TransformerBlock): 1. It adds deepstack merger and norm layers 2. It returns hidden_states and deepstack features in the forward function """ - def __init__(self, config, spec, - post_layer_norm = True, pre_process = True, post_process = True, - pg_collection = None, vp_stage = None, - projection_config = None, # Note: DeepStack - projection_layer_spec = None, # Note: DeepStack - projection_type = 'mlp', # Note: DeepStack): - ): - super().__init__(config, spec, post_layer_norm, pre_process, post_process, pg_collection, vp_stage) + + def __init__( + self, + config, + spec, + post_layer_norm=True, + pre_process=True, + post_process=True, + pg_collection=None, + vp_stage=None, + projection_config=None, # Note: DeepStack + projection_layer_spec=None, # Note: DeepStack + projection_type="mlp", # Note: DeepStack): + ): + super().__init__( + config, spec, post_layer_norm, pre_process, post_process, pg_collection, vp_stage + ) if self.final_layernorm != None: # NOTE(lizhiyu): replace final layernorm with TENorm if using TE self.final_layernorm = None - self.final_layernorm = torch.nn.LayerNorm(normalized_shape=self.config.hidden_size, eps=self.config.layernorm_epsilon,) + self.final_layernorm = torch.nn.LayerNorm( + normalized_shape=self.config.hidden_size, + eps=self.config.layernorm_epsilon, + ) # NOTE: DeepStack self.deepstack_visual_indexes = self.config.deepstack_visual_indexes @@ -87,7 +100,7 @@ def __init__(self, config, spec, projection_config, projection_layer_spec, projection_type, - projection_config.ffn_hidden_size + projection_config.ffn_hidden_size, ) for _ in range(len(self.config.deepstack_visual_indexes)) ] @@ -102,7 +115,10 @@ def __init__(self, config, spec, # eps=self.config.layernorm_epsilon, # ), # NOTE(lizhiyu): replace with torch LayerNorm - torch.nn.LayerNorm(normalized_shape = projection_config.ffn_hidden_size, eps = self.config.layernorm_epsilon,) + torch.nn.LayerNorm( + normalized_shape=projection_config.ffn_hidden_size, + eps=self.config.layernorm_epsilon, + ) for _ in range(len(self.config.deepstack_visual_indexes)) ] ) @@ -185,7 +201,7 @@ def checkpoint_handler(forward_func): rotary_pos_emb, ) - if self.config.recompute_method == 'uniform': + if self.config.recompute_method == "uniform": # Uniformly divide the total number of Transformer layers and checkpoint # the input activation of each divided chunk. # A method to further reduce memory usage reducing checkpoints. @@ -216,7 +232,7 @@ def checkpoint_handler(forward_func): layer_idx += self.config.recompute_num_layers - elif self.config.recompute_method == 'block': + elif self.config.recompute_method == "block": # Checkpoint the input activation of only a set number of individual # Transformer layers and skip the rest. # A method fully use the device memory removing redundant re-computation. @@ -243,6 +259,7 @@ def checkpoint_handler(forward_func): return hidden_states """Transformer class.""" + def forward( self, hidden_states: Union[Tensor, WrappedTensor], @@ -344,9 +361,13 @@ def forward( with rng_context, outer_quantization_context: # Forward pass. - if self.config.recompute_granularity == 'full' and self.training: - assert self.config.recompute_method == 'uniform' and self.config.recompute_num_layers == 1, \ + if self.config.recompute_granularity == "full" and self.training: + assert ( + self.config.recompute_method == "uniform" + and self.config.recompute_num_layers == 1 + ), ( f"Only uniform recompute with recompute_num_layers=1 is supported for full recompute in Qwen3-VL." + ) hidden_states = self._checkpointed_forward( hidden_states=hidden_states, attention_mask=attention_mask, @@ -356,7 +377,7 @@ def forward( attention_bias=attention_bias, packed_seq_params=packed_seq_params, use_inner_quantization_context=use_inner_quantization_context, - deepstack_feature_lists = deepstack_feature_lists, # Note: DeepStack + deepstack_feature_lists=deepstack_feature_lists, # Note: DeepStack ) else: for l_no, layer in enumerate(self.layers): diff --git a/flagscale/models/megatron/rwkv/__init__.py b/flagscale/models/megatron/rwkv/__init__.py index b60aca583a..d99a777e73 100644 --- a/flagscale/models/megatron/rwkv/__init__.py +++ b/flagscale/models/megatron/rwkv/__init__.py @@ -1 +1,15 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from .rwkv_model import RWKVModel diff --git a/flagscale/models/megatron/rwkv/rwkv_model.py b/flagscale/models/megatron/rwkv/rwkv_model.py index 2952e0f7b7..795f1f46c1 100644 --- a/flagscale/models/megatron/rwkv/rwkv_model.py +++ b/flagscale/models/megatron/rwkv/rwkv_model.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # This code is modified based on the RWKV GitHub repository: # https://github.com/BlinkDL/RWKV-LM @@ -22,13 +36,19 @@ try: import torch._dynamo as _dynamo except Exception: + class _Dummy: - def disable(self, f): return f - def is_compiling(self): return False + def disable(self, f): + return f + + def is_compiling(self): + return False + _dynamo = _Dummy() _raw_token_shift = token_shift + @_dynamo.disable def _safe_token_shift(*args, **kwargs): out = _raw_token_shift(*args, **kwargs) @@ -37,11 +57,14 @@ def _safe_token_shift(*args, **kwargs): return main return out + token_shift = _safe_token_shift + def __nop(ob): return ob + # ROCm detection (robust) ROCm_flag = hasattr(torch.version, "hip") and (torch.version.hip is not None) @@ -76,7 +99,10 @@ def __nop(ob): ] load( name="wind_backstepping_hip", - sources=["megatron/core/models/rwkv/cuda/wkv7_hip.hip", "megatron/core/models/rwkv/cuda/wkv7_op.hip"], + sources=[ + "megatron/core/models/rwkv/cuda/wkv7_hip.hip", + "megatron/core/models/rwkv/cuda/wkv7_op.hip", + ], is_python_module=False, verbose=True, extra_cuda_cflags=flags, @@ -93,7 +119,10 @@ def __nop(ob): ] load( name="wind_backstepping", - sources=["megatron/core/models/rwkv/cuda/wkv7_cuda.cu", "megatron/core/models/rwkv/cuda/wkv7_op.cpp"], + sources=[ + "megatron/core/models/rwkv/cuda/wkv7_cuda.cu", + "megatron/core/models/rwkv/cuda/wkv7_op.cpp", + ], is_python_module=False, verbose=True, extra_cuda_cflags=flags, @@ -106,8 +135,7 @@ def forward(ctx, w, q, k, v, z, b): assert T % CHUNK_LEN == 0 assert all(i.dtype in [torch.bfloat16, torch.float16] for i in [w, q, k, v, z, b]) w, q, k, v, z, b, s, sa = ctx.saved_tensors - dw, dq, dk, dv, dz, db = [torch.empty_like(x) for x in [ - w, q, k, v, z, b]] + dw, dq, dk, dv, dz, db = [torch.empty_like(x) for x in [w, q, k, v, z, b]] torch.ops.wind_backstepping.backward( w, q, k, v, z, b, dy, s, sa, dw, dq, dk, dv, dz, db ) @@ -171,10 +199,10 @@ def __init__(self, args, layer_id): def ortho_init(x, scale): shape = x.shape if len(shape) == 2: - gain = (math.sqrt(shape[0] / shape[1]) if shape[0] > shape[1] else 1) + gain = math.sqrt(shape[0] / shape[1]) if shape[0] > shape[1] else 1 nn.init.orthogonal_(x, gain=gain * scale) elif len(shape) == 3: - gain = (math.sqrt(shape[1] / shape[2]) if shape[1] > shape[2] else 1) + gain = math.sqrt(shape[1] / shape[2]) if shape[1] > shape[2] else 1 for i in range(shape[0]): nn.init.orthogonal_(x[i], gain=gain * scale) else: @@ -236,9 +264,9 @@ def ortho_init(x, scale): def forward(self, x, v_first): B, T, C = x.size() xx = token_shift(x) - xr, xw, xk, xv, xa, xg = fused_addcmul_rwkv7(x, xx, self.x_r, - self.x_w, self.x_k, self.x_v, - self.x_a, self.x_g) + xr, xw, xk, xv, xa, xg = fused_addcmul_rwkv7( + x, xx, self.x_r, self.x_w, self.x_k, self.x_v, self.x_a, self.x_g + ) r = self.receptance(xr) # soft-clamp to (-inf, -0.5) @@ -257,8 +285,7 @@ def forward(self, x, v_first): a = torch.sigmoid(self.a0 + (xa @ self.a1) @ self.a2) g = torch.sigmoid(xg @ self.g1) @ self.g2 - kk = F.normalize((k * self.k_k).view(B, T, self.n_head, -1), - dim=-1, p=2.0).view(B, T, C) + kk = F.normalize((k * self.k_k).view(B, T, self.n_head, -1), dim=-1, p=2.0).view(B, T, C) k = fused_k_rwkv7(k, a, self.k_a) x = RUN_CUDA_RWKV7g(r, w, k, v, -kk, kk * a) @@ -290,9 +317,7 @@ def __init__(self, args, layer_id): self.key = nn.Linear(args.n_embd, args.n_embd * 4, bias=False) self.value = nn.Linear(args.n_embd * 4, args.n_embd, bias=False) - self.key.weight.data.uniform_( - -0.5 / (args.n_embd**0.5), 0.5 / (args.n_embd**0.5) - ) + self.key.weight.data.uniform_(-0.5 / (args.n_embd**0.5), 0.5 / (args.n_embd**0.5)) self.value.weight.data.zero_() @CompileFunction @@ -351,9 +376,9 @@ def __init__( *, pre_process: bool = True, post_process: bool = True, - parallel_output: bool = True, # keep for future TP; no effect with TP=1 + parallel_output: bool = True, # keep for future TP; no effect with TP=1 use_grad_checkpoint: bool = False, # PyTorch checkpoint at block granularity - dtype: torch.dtype | None = None, # e.g., torch.bfloat16 to match --bf16 + dtype: torch.dtype | None = None, # e.g., torch.bfloat16 to match --bf16 ): super().__init__() self.vocab_size = vocab_size @@ -371,14 +396,15 @@ def __init__( num_attention_heads=1, use_cpu_initialization=True, bf16=True, - fp16=False + fp16=False, ) # Embedding / Blocks / Head - align with your original RWKV stack self.emb = nn.Embedding(vocab_size, hidden_size) # Use _ArgsShim to adapt expected args for Block - self.blocks = nn.ModuleList([Block(_ArgsShim(hidden_size, n_layer), i) - for i in range(n_layer)]) + self.blocks = nn.ModuleList( + [Block(_ArgsShim(hidden_size, n_layer), i) for i in range(n_layer)] + ) self.ln_out = nn.LayerNorm(hidden_size) self.head = nn.Linear(hidden_size, vocab_size, bias=False) @@ -420,10 +446,7 @@ def forward( # 4) Training branch: return per-token cross entropy for Megatron to aggregate if labels is not None: # Use transpose for CE: [B, V, T] vs labels [B, T] - per_tok = F.cross_entropy( - logits.transpose(1, 2), labels, - reduction="none" - ) # [B, T] + per_tok = F.cross_entropy(logits.transpose(1, 2), labels, reduction="none") # [B, T] # Do NOT reduce here. Optionally pre-mask per-token loss for clarity. if loss_mask is not None: @@ -446,24 +469,22 @@ def set_input_tensor(self, input_tensor: Tensor) -> None: # gives us non-lists or None self.input_tensor = input_tensor - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): + def state_dict_for_save_checkpoint(self, prefix="", keep_vars=False): """Customized state dict for checkpoint saving.""" state_dict_ = {} # Save embeddings - state_dict_[self._word_embeddings_for_head_key] = self.emb.state_dict( - keep_vars=keep_vars - ) + state_dict_[self._word_embeddings_for_head_key] = self.emb.state_dict(keep_vars=keep_vars) # Save each block blocks_state = {} for i, blk in enumerate(self.blocks): - blocks_state[f'block_{i}'] = blk.state_dict() - state_dict_['blocks'] = blocks_state + blocks_state[f"block_{i}"] = blk.state_dict() + state_dict_["blocks"] = blocks_state # Save output layer - state_dict_['ln_out'] = self.ln_out.state_dict(keep_vars=keep_vars) - state_dict_['head'] = self.head.state_dict(keep_vars=keep_vars) + state_dict_["ln_out"] = self.ln_out.state_dict(keep_vars=keep_vars) + state_dict_["head"] = self.head.state_dict(keep_vars=keep_vars) return state_dict_ @@ -474,17 +495,17 @@ def load_state_dict(self, state_dict, strict=True): self.emb.load_state_dict(state_dict[self._word_embeddings_for_head_key], strict=strict) # Load blocks - if 'blocks' in state_dict: + if "blocks" in state_dict: for i, blk in enumerate(self.blocks): - blk_key = f'block_{i}' - if blk_key in state_dict['blocks']: - blk.load_state_dict(state_dict['blocks'][blk_key], strict=strict) + blk_key = f"block_{i}" + if blk_key in state_dict["blocks"]: + blk.load_state_dict(state_dict["blocks"][blk_key], strict=strict) # Load output layer - if 'ln_out' in state_dict: - self.ln_out.load_state_dict(state_dict['ln_out'], strict=strict) - if 'head' in state_dict: - self.head.load_state_dict(state_dict['head'], strict=strict) + if "ln_out" in state_dict: + self.ln_out.load_state_dict(state_dict["ln_out"], strict=strict) + if "head" in state_dict: + self.head.load_state_dict(state_dict["head"], strict=strict) class _ArgsShim: @@ -492,6 +513,7 @@ class _ArgsShim: Minimal adapter to provide attributes that Block(...) expects from `args`. Extend this class if Block accesses additional fields in your repo. """ + def __init__(self, n_embd: int, n_layer: int): self.n_embd = n_embd self.n_layer = n_layer diff --git a/flagscale/models/robobrain_x/action_encoder.py b/flagscale/models/robobrain_x/action_encoder.py index 1e07a86366..e4efe11b47 100644 --- a/flagscale/models/robobrain_x/action_encoder.py +++ b/flagscale/models/robobrain_x/action_encoder.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Mainly adopted from starVLA/starVLA: # https://github.com/starVLA/starVLA/blob/starVLA/starVLA/model/modules/action_model/flow_matching_head/action_encoder.py import torch diff --git a/flagscale/models/robobrain_x/cross_attention_dit.py b/flagscale/models/robobrain_x/cross_attention_dit.py index acba75ac50..da36b990c8 100644 --- a/flagscale/models/robobrain_x/cross_attention_dit.py +++ b/flagscale/models/robobrain_x/cross_attention_dit.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Mainly adopted from starVLA/starVLA: # https://github.com/starVLA/starVLA/blob/starVLA/starVLA/model/modules/action_model/flow_matching_head/cross_attention_dit.py diff --git a/flagscale/models/utils/__init__.py b/flagscale/models/utils/__init__.py index e69de29bb2..a491ed2726 100644 --- a/flagscale/models/utils/__init__.py +++ b/flagscale/models/utils/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flagscale/models/vla/__init__.py b/flagscale/models/vla/__init__.py index a951c6536a..a86610b402 100644 --- a/flagscale/models/vla/__init__.py +++ b/flagscale/models/vla/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from .base_policy import TrainablePolicy from .gr00t_n1_5 import Gr00tN15 from .protocols import ActionModel, VLMBackbone diff --git a/flagscale/models/vla/action_model/__init__.py b/flagscale/models/vla/action_model/__init__.py index 3a113e505a..5845e1aa37 100644 --- a/flagscale/models/vla/action_model/__init__.py +++ b/flagscale/models/vla/action_model/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from .gr00t_action_header import GR00TActionHead __all__ = ["GR00TActionHead"] diff --git a/flagscale/models/vla/base_policy.py b/flagscale/models/vla/base_policy.py index 2d2eb02656..ae79b4fa64 100644 --- a/flagscale/models/vla/base_policy.py +++ b/flagscale/models/vla/base_policy.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from __future__ import annotations import inspect diff --git a/flagscale/models/vla/gr00t_n1_5/__init__.py b/flagscale/models/vla/gr00t_n1_5/__init__.py index 276a8c1052..271924cf6e 100644 --- a/flagscale/models/vla/gr00t_n1_5/__init__.py +++ b/flagscale/models/vla/gr00t_n1_5/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from flagscale.models.vla.gr00t_n1_5.configuration_gr00t_n1_5 import Gr00tN15Config from flagscale.models.vla.gr00t_n1_5.modeling_gr00t_n1_5 import Gr00tN15 diff --git a/flagscale/models/vla/gr00t_n1_5/configuration_gr00t_n1_5.py b/flagscale/models/vla/gr00t_n1_5/configuration_gr00t_n1_5.py index 5baf75f63b..d943e0d83e 100644 --- a/flagscale/models/vla/gr00t_n1_5/configuration_gr00t_n1_5.py +++ b/flagscale/models/vla/gr00t_n1_5/configuration_gr00t_n1_5.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from __future__ import annotations from dataclasses import dataclass, field diff --git a/flagscale/models/vla/gr00t_n1_5/eagle2/__init__.py b/flagscale/models/vla/gr00t_n1_5/eagle2/__init__.py index e69de29bb2..a491ed2726 100644 --- a/flagscale/models/vla/gr00t_n1_5/eagle2/__init__.py +++ b/flagscale/models/vla/gr00t_n1_5/eagle2/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flagscale/models/vla/protocols.py b/flagscale/models/vla/protocols.py index 413f5fa2a9..7d1ab4c994 100644 --- a/flagscale/models/vla/protocols.py +++ b/flagscale/models/vla/protocols.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from typing import Protocol from torch import Tensor diff --git a/flagscale/models/vla/qwen_gr00t/__init__.py b/flagscale/models/vla/qwen_gr00t/__init__.py index 67476a3b7f..d5b6188de1 100644 --- a/flagscale/models/vla/qwen_gr00t/__init__.py +++ b/flagscale/models/vla/qwen_gr00t/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from .configuration_qwen_gr00t import QwenGr00tConfig from .modeling_qwen_gr00t import QwenGr00t diff --git a/flagscale/models/vla/qwen_gr00t/configuration_qwen_gr00t.py b/flagscale/models/vla/qwen_gr00t/configuration_qwen_gr00t.py index 44333a8059..c403a12af1 100644 --- a/flagscale/models/vla/qwen_gr00t/configuration_qwen_gr00t.py +++ b/flagscale/models/vla/qwen_gr00t/configuration_qwen_gr00t.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from __future__ import annotations from dataclasses import dataclass, field diff --git a/flagscale/models/vla/qwen_gr00t/task_encoder_qwen_gr00t.py b/flagscale/models/vla/qwen_gr00t/task_encoder_qwen_gr00t.py index c3596547f5..f239f79f6e 100644 --- a/flagscale/models/vla/qwen_gr00t/task_encoder_qwen_gr00t.py +++ b/flagscale/models/vla/qwen_gr00t/task_encoder_qwen_gr00t.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import json import logging diff --git a/flagscale/models/vla/registry.py b/flagscale/models/vla/registry.py index 36cf0ae56a..f2287ffd4e 100644 --- a/flagscale/models/vla/registry.py +++ b/flagscale/models/vla/registry.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + VLM_REGISTRY: dict[str, type] = {} ACTION_MODEL_REGISTRY: dict[str, type] = {} diff --git a/flagscale/models/vla/utils.py b/flagscale/models/vla/utils.py index b0ee5fdd12..1db439b983 100644 --- a/flagscale/models/vla/utils.py +++ b/flagscale/models/vla/utils.py @@ -1,3 +1,18 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + def get_vlm_config(vlm_config) -> dict: """ Extract common fields from any VLM config, handling structural differences. diff --git a/flagscale/models/vla/vlm/__init__.py b/flagscale/models/vla/vlm/__init__.py index 4b736223f7..b867ea4199 100644 --- a/flagscale/models/vla/vlm/__init__.py +++ b/flagscale/models/vla/vlm/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from .qwenvl_backbone import Qwen3VLBackbone, Qwen25VLBackbone, QwenVLBackbone __all__ = ["QwenVLBackbone", "Qwen25VLBackbone", "Qwen3VLBackbone"] diff --git a/flagscale/models/vla/vlm/qwenvl_backbone.py b/flagscale/models/vla/vlm/qwenvl_backbone.py index 7c78dca626..9ed1b11935 100644 --- a/flagscale/models/vla/vlm/qwenvl_backbone.py +++ b/flagscale/models/vla/vlm/qwenvl_backbone.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Mainly adopted from: # https://github.com/starVLA/starVLA/blob/3f7feefbc5fc25890ad3a7d262b8a0aea1339aa7/starVLA/model/modules/vlm/QWen3.py diff --git a/flagscale/patches_utils.py b/flagscale/patches_utils.py index c39f51e358..9029ae5c36 100644 --- a/flagscale/patches_utils.py +++ b/flagscale/patches_utils.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # this file is used for adding tools func to processing patches diff --git a/flagscale/platforms/__init__.py b/flagscale/platforms/__init__.py index 9d84d8161e..2639ddf377 100644 --- a/flagscale/platforms/__init__.py +++ b/flagscale/platforms/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from .platform_base import PlatformBase from .platform_manager import get_platform, set_platform from .platform_register import register_platforms diff --git a/flagscale/platforms/platform_base.py b/flagscale/platforms/platform_base.py index 682716e298..a23915c1e6 100644 --- a/flagscale/platforms/platform_base.py +++ b/flagscale/platforms/platform_base.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import abc from abc import ABC diff --git a/flagscale/platforms/platform_cuda.py b/flagscale/platforms/platform_cuda.py index a395ba9470..c09f355873 100644 --- a/flagscale/platforms/platform_cuda.py +++ b/flagscale/platforms/platform_cuda.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import torch from .platform_base import PlatformBase diff --git a/flagscale/platforms/platform_manager.py b/flagscale/platforms/platform_manager.py index 9fe8a2c04b..2633d10944 100644 --- a/flagscale/platforms/platform_manager.py +++ b/flagscale/platforms/platform_manager.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os from .platform_register import PLATFORMS diff --git a/flagscale/platforms/platform_musa.py b/flagscale/platforms/platform_musa.py index ed3fb53be6..97c64f6ba7 100644 --- a/flagscale/platforms/platform_musa.py +++ b/flagscale/platforms/platform_musa.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import torch from .platform_base import PlatformBase diff --git a/flagscale/platforms/platform_npu.py b/flagscale/platforms/platform_npu.py index b83c80f53e..24bcdee1ad 100644 --- a/flagscale/platforms/platform_npu.py +++ b/flagscale/platforms/platform_npu.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import torch from .platform_base import PlatformBase diff --git a/flagscale/run.py b/flagscale/run.py index 5f135b3066..0f83a9ae97 100644 --- a/flagscale/run.py +++ b/flagscale/run.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import sys import warnings diff --git a/flagscale/runner/__init__.py b/flagscale/runner/__init__.py index e69de29bb2..a491ed2726 100644 --- a/flagscale/runner/__init__.py +++ b/flagscale/runner/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flagscale/runner/auto_tuner/__init__.py b/flagscale/runner/auto_tuner/__init__.py index aa94e73db5..e6038242ab 100644 --- a/flagscale/runner/auto_tuner/__init__.py +++ b/flagscale/runner/auto_tuner/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from flagscale.runner.auto_tuner.tuner import AutoTunerBase from flagscale.runner.auto_tuner.tuner_serve import ServeAutoTuner from flagscale.runner.auto_tuner.tuner_train import TrainAutoTuner diff --git a/flagscale/runner/auto_tuner/generate.py b/flagscale/runner/auto_tuner/generate.py index 1736446617..ee24eb92c9 100644 --- a/flagscale/runner/auto_tuner/generate.py +++ b/flagscale/runner/auto_tuner/generate.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import copy import os diff --git a/flagscale/runner/auto_tuner/hetero/__init__.py b/flagscale/runner/auto_tuner/hetero/__init__.py index 067fd0ce6a..6f6732578f 100644 --- a/flagscale/runner/auto_tuner/hetero/__init__.py +++ b/flagscale/runner/auto_tuner/hetero/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from .hetero_generator import HeteroGenerator from .hetero_pruner import HeteroPruner from .hetero_recorder import HeteroRecorder diff --git a/flagscale/runner/auto_tuner/hetero/hetero_generator.py b/flagscale/runner/auto_tuner/hetero/hetero_generator.py index 03ff3f4af6..e51a7468c9 100644 --- a/flagscale/runner/auto_tuner/hetero/hetero_generator.py +++ b/flagscale/runner/auto_tuner/hetero/hetero_generator.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from flagscale.runner.auto_tuner.generate import Generator diff --git a/flagscale/runner/auto_tuner/hetero/hetero_pruner.py b/flagscale/runner/auto_tuner/hetero/hetero_pruner.py index 94fe7ad93b..7c82ff25e1 100644 --- a/flagscale/runner/auto_tuner/hetero/hetero_pruner.py +++ b/flagscale/runner/auto_tuner/hetero/hetero_pruner.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import csv import json import logging diff --git a/flagscale/runner/auto_tuner/hetero/hetero_recorder.py b/flagscale/runner/auto_tuner/hetero/hetero_recorder.py index be396c7021..b5e1cfa496 100644 --- a/flagscale/runner/auto_tuner/hetero/hetero_recorder.py +++ b/flagscale/runner/auto_tuner/hetero/hetero_recorder.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import csv import json import os diff --git a/flagscale/runner/auto_tuner/hetero/hetero_searcher.py b/flagscale/runner/auto_tuner/hetero/hetero_searcher.py index 719dfdca43..ed5decc779 100644 --- a/flagscale/runner/auto_tuner/hetero/hetero_searcher.py +++ b/flagscale/runner/auto_tuner/hetero/hetero_searcher.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import copy import itertools import logging diff --git a/flagscale/runner/auto_tuner/hetero/hetero_theoretical_memory.py b/flagscale/runner/auto_tuner/hetero/hetero_theoretical_memory.py index 796bb3cc90..349c392e25 100644 --- a/flagscale/runner/auto_tuner/hetero/hetero_theoretical_memory.py +++ b/flagscale/runner/auto_tuner/hetero/hetero_theoretical_memory.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import logging import math import os diff --git a/flagscale/runner/auto_tuner/memory_model.py b/flagscale/runner/auto_tuner/memory_model.py index bd48ab4c3f..a30ea12d98 100644 --- a/flagscale/runner/auto_tuner/memory_model.py +++ b/flagscale/runner/auto_tuner/memory_model.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from flagscale.runner.auto_tuner.hetero.hetero_theoretical_memory import ( hetero_report_theoretical_memory, ) diff --git a/flagscale/runner/auto_tuner/platform.py b/flagscale/runner/auto_tuner/platform.py index 01273b0658..23f8f824a7 100644 --- a/flagscale/runner/auto_tuner/platform.py +++ b/flagscale/runner/auto_tuner/platform.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os diff --git a/flagscale/runner/auto_tuner/prune/history.py b/flagscale/runner/auto_tuner/prune/history.py index cd9157ec13..f584ae2192 100644 --- a/flagscale/runner/auto_tuner/prune/history.py +++ b/flagscale/runner/auto_tuner/prune/history.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import logging from flagscale.runner.auto_tuner.utils import beside, compare_by_recompute diff --git a/flagscale/runner/auto_tuner/prune/memory.py b/flagscale/runner/auto_tuner/prune/memory.py index b912174141..6c23d18f3f 100644 --- a/flagscale/runner/auto_tuner/prune/memory.py +++ b/flagscale/runner/auto_tuner/prune/memory.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import logging logger = logging.getLogger("FlagScale-AutoTuner") diff --git a/flagscale/runner/auto_tuner/prune/pruner.py b/flagscale/runner/auto_tuner/prune/pruner.py index 60c1cf82af..a378277713 100644 --- a/flagscale/runner/auto_tuner/prune/pruner.py +++ b/flagscale/runner/auto_tuner/prune/pruner.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from flagscale.runner.auto_tuner.prune.history import _HISTORY_BASED_PRUNE_FUNC from flagscale.runner.auto_tuner.prune.memory import ( prune_by_memory_model, diff --git a/flagscale/runner/auto_tuner/record/recorder.py b/flagscale/runner/auto_tuner/record/recorder.py index 876c3c2b91..b28a00f3f4 100644 --- a/flagscale/runner/auto_tuner/record/recorder.py +++ b/flagscale/runner/auto_tuner/record/recorder.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import json import logging import os diff --git a/flagscale/runner/auto_tuner/search/algorithm.py b/flagscale/runner/auto_tuner/search/algorithm.py index eca538e370..5706a3eafb 100644 --- a/flagscale/runner/auto_tuner/search/algorithm.py +++ b/flagscale/runner/auto_tuner/search/algorithm.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from abc import ABC, abstractmethod from flagscale.runner.auto_tuner.utils import ( diff --git a/flagscale/runner/auto_tuner/search/searcher.py b/flagscale/runner/auto_tuner/search/searcher.py index 2d67bc6827..5a7ec79e3b 100644 --- a/flagscale/runner/auto_tuner/search/searcher.py +++ b/flagscale/runner/auto_tuner/search/searcher.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import copy import itertools import json diff --git a/flagscale/runner/auto_tuner/tuner.py b/flagscale/runner/auto_tuner/tuner.py index 1cca2c6695..d8bd035c2c 100644 --- a/flagscale/runner/auto_tuner/tuner.py +++ b/flagscale/runner/auto_tuner/tuner.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import time from abc import ABC, abstractmethod diff --git a/flagscale/runner/auto_tuner/tuner_serve.py b/flagscale/runner/auto_tuner/tuner_serve.py index 05bd3f881f..ee7855491d 100644 --- a/flagscale/runner/auto_tuner/tuner_serve.py +++ b/flagscale/runner/auto_tuner/tuner_serve.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import copy import datetime import logging diff --git a/flagscale/runner/auto_tuner/tuner_train.py b/flagscale/runner/auto_tuner/tuner_train.py index 29a81a226e..1fd0bc67c4 100644 --- a/flagscale/runner/auto_tuner/tuner_train.py +++ b/flagscale/runner/auto_tuner/tuner_train.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import copy import datetime import json diff --git a/flagscale/runner/auto_tuner/utils.py b/flagscale/runner/auto_tuner/utils.py index 42b99cfd18..21429df796 100644 --- a/flagscale/runner/auto_tuner/utils.py +++ b/flagscale/runner/auto_tuner/utils.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import sys from types import SimpleNamespace diff --git a/flagscale/runner/autotuner_factory.py b/flagscale/runner/autotuner_factory.py index 60bdaa7a43..14f359c2db 100644 --- a/flagscale/runner/autotuner_factory.py +++ b/flagscale/runner/autotuner_factory.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from typing import TypeVar from flagscale.runner.auto_tuner import AutoTunerBase, ServeAutoTuner, TrainAutoTuner diff --git a/flagscale/runner/backend/__init__.py b/flagscale/runner/backend/__init__.py index 8fac5bc222..e3f04c4ef2 100644 --- a/flagscale/runner/backend/__init__.py +++ b/flagscale/runner/backend/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from flagscale.runner.backend.backend_base import BackendBase from flagscale.runner.backend.backend_llama_cpp import LlamaCppBackend from flagscale.runner.backend.backend_megatron import MegatronBackend diff --git a/flagscale/runner/backend/backend_base.py b/flagscale/runner/backend/backend_base.py index a8a455bb68..f852956a50 100644 --- a/flagscale/runner/backend/backend_base.py +++ b/flagscale/runner/backend/backend_base.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from abc import ABC, abstractmethod from omegaconf import DictConfig diff --git a/flagscale/runner/backend/backend_llama_cpp.py b/flagscale/runner/backend/backend_llama_cpp.py index f891b061dc..cc635b3e0d 100644 --- a/flagscale/runner/backend/backend_llama_cpp.py +++ b/flagscale/runner/backend/backend_llama_cpp.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import json import os diff --git a/flagscale/runner/backend/backend_megatron.py b/flagscale/runner/backend/backend_megatron.py index 7139e224aa..9c3946db06 100644 --- a/flagscale/runner/backend/backend_megatron.py +++ b/flagscale/runner/backend/backend_megatron.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os from datetime import datetime diff --git a/flagscale/runner/backend/backend_native_compress.py b/flagscale/runner/backend/backend_native_compress.py index d0cbab6bdd..5cb05608f1 100644 --- a/flagscale/runner/backend/backend_native_compress.py +++ b/flagscale/runner/backend/backend_native_compress.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os from datetime import datetime diff --git a/flagscale/runner/backend/backend_native_serve.py b/flagscale/runner/backend/backend_native_serve.py index 8c76aaf106..2ec1e7f5b8 100644 --- a/flagscale/runner/backend/backend_native_serve.py +++ b/flagscale/runner/backend/backend_native_serve.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import json import os diff --git a/flagscale/runner/backend/backend_native_train.py b/flagscale/runner/backend/backend_native_train.py index 155b2b9336..47fed89ced 100644 --- a/flagscale/runner/backend/backend_native_train.py +++ b/flagscale/runner/backend/backend_native_train.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os from datetime import datetime diff --git a/flagscale/runner/backend/backend_sglang.py b/flagscale/runner/backend/backend_sglang.py index 649da06611..ee796832e0 100644 --- a/flagscale/runner/backend/backend_sglang.py +++ b/flagscale/runner/backend/backend_sglang.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import copy import json import os diff --git a/flagscale/runner/backend/backend_verl.py b/flagscale/runner/backend/backend_verl.py index cf8a8d91fc..7ad8fc8a31 100644 --- a/flagscale/runner/backend/backend_verl.py +++ b/flagscale/runner/backend/backend_verl.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os from omegaconf import DictConfig, OmegaConf diff --git a/flagscale/runner/backend/backend_vllm.py b/flagscale/runner/backend/backend_vllm.py index 73ccf129cc..7376fb560e 100644 --- a/flagscale/runner/backend/backend_vllm.py +++ b/flagscale/runner/backend/backend_vllm.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import collections import json import os diff --git a/flagscale/runner/elastic/__init__.py b/flagscale/runner/elastic/__init__.py index e69de29bb2..a491ed2726 100644 --- a/flagscale/runner/elastic/__init__.py +++ b/flagscale/runner/elastic/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/flagscale/runner/elastic/diagnostic.py b/flagscale/runner/elastic/diagnostic.py index 17193df01c..caed34a702 100644 --- a/flagscale/runner/elastic/diagnostic.py +++ b/flagscale/runner/elastic/diagnostic.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import time diff --git a/flagscale/runner/elastic/gpu_health_check.py b/flagscale/runner/elastic/gpu_health_check.py index 73f5ccc052..698d59219e 100644 --- a/flagscale/runner/elastic/gpu_health_check.py +++ b/flagscale/runner/elastic/gpu_health_check.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """ Complete GPU Health Check Implementation diff --git a/flagscale/runner/elastic/log_collector.py b/flagscale/runner/elastic/log_collector.py index 0aba472189..9ceda2e3cb 100644 --- a/flagscale/runner/elastic/log_collector.py +++ b/flagscale/runner/elastic/log_collector.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import glob import os import shlex diff --git a/flagscale/runner/elastic/monitor_launcher.py b/flagscale/runner/elastic/monitor_launcher.py index 32d0f02942..5a349992dd 100644 --- a/flagscale/runner/elastic/monitor_launcher.py +++ b/flagscale/runner/elastic/monitor_launcher.py @@ -1,4 +1,19 @@ #!/usr/bi/env python3 + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """ An independent monitoring service launcher for starting the monitoring service during the training process, including log collection and diagnostic report generation diff --git a/flagscale/runner/elastic/monitor_service.py b/flagscale/runner/elastic/monitor_service.py index b19f26585b..141e806a58 100644 --- a/flagscale/runner/elastic/monitor_service.py +++ b/flagscale/runner/elastic/monitor_service.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import signal import subprocess diff --git a/flagscale/runner/elastic/simulated_fault.py b/flagscale/runner/elastic/simulated_fault.py index 9eab834edd..4b8b19b093 100644 --- a/flagscale/runner/elastic/simulated_fault.py +++ b/flagscale/runner/elastic/simulated_fault.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import random import time diff --git a/flagscale/runner/launcher/__init__.py b/flagscale/runner/launcher/__init__.py index 1597c1b948..35e7570479 100644 --- a/flagscale/runner/launcher/__init__.py +++ b/flagscale/runner/launcher/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from flagscale.runner.launcher.launcher_base import LauncherBase from flagscale.runner.launcher.launcher_cloud import CloudLauncher from flagscale.runner.launcher.launcher_ssh import SshLauncher diff --git a/flagscale/runner/launcher/launcher_base.py b/flagscale/runner/launcher/launcher_base.py index 6cb8cde8f1..0445f81ad7 100644 --- a/flagscale/runner/launcher/launcher_base.py +++ b/flagscale/runner/launcher/launcher_base.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from abc import ABC, abstractmethod diff --git a/flagscale/runner/launcher/launcher_cloud.py b/flagscale/runner/launcher/launcher_cloud.py index 9794f6893d..05a8dc47b0 100644 --- a/flagscale/runner/launcher/launcher_cloud.py +++ b/flagscale/runner/launcher/launcher_cloud.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import shlex import subprocess diff --git a/flagscale/runner/launcher/launcher_ssh.py b/flagscale/runner/launcher/launcher_ssh.py index f238cdc8c2..8b952d0dea 100644 --- a/flagscale/runner/launcher/launcher_ssh.py +++ b/flagscale/runner/launcher/launcher_ssh.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import asyncio import multiprocessing import os diff --git a/flagscale/runner/runner_base.py b/flagscale/runner/runner_base.py index 6f80625bbb..ac569b4fb4 100644 --- a/flagscale/runner/runner_base.py +++ b/flagscale/runner/runner_base.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from abc import ABC from omegaconf import DictConfig diff --git a/flagscale/runner/runner_base_legacy.py b/flagscale/runner/runner_base_legacy.py index faf38d4c88..75b1b51ab9 100644 --- a/flagscale/runner/runner_base_legacy.py +++ b/flagscale/runner/runner_base_legacy.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from abc import ABC, abstractmethod from enum import Enum diff --git a/flagscale/runner/runner_factory.py b/flagscale/runner/runner_factory.py index ea120b5b60..cab86bde7e 100644 --- a/flagscale/runner/runner_factory.py +++ b/flagscale/runner/runner_factory.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from typing import TypeVar from flagscale.runner.backend import ( diff --git a/flagscale/runner/runner_inference.py b/flagscale/runner/runner_inference.py index ec9cc4efd6..e7c97e9a10 100644 --- a/flagscale/runner/runner_inference.py +++ b/flagscale/runner/runner_inference.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import shlex diff --git a/flagscale/runner/runner_serve.py b/flagscale/runner/runner_serve.py index b6c3f179de..7625760ccc 100644 --- a/flagscale/runner/runner_serve.py +++ b/flagscale/runner/runner_serve.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import asyncio import collections import contextlib diff --git a/flagscale/runner/runner_train.py b/flagscale/runner/runner_train.py index 9eeb0e9bd5..6d8693a196 100644 --- a/flagscale/runner/runner_train.py +++ b/flagscale/runner/runner_train.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import multiprocessing import os import shlex diff --git a/flagscale/runner/straggler/README_STRAGGLER.md b/flagscale/runner/straggler/README_STRAGGLER.md index 076f93fb36..b83945ca80 100644 --- a/flagscale/runner/straggler/README_STRAGGLER.md +++ b/flagscale/runner/straggler/README_STRAGGLER.md @@ -1,3 +1,19 @@ + + # Straggler on Metax C550 ## Scope diff --git a/flagscale/runner/straggler/__init__.py b/flagscale/runner/straggler/__init__.py index c1e227b057..3f63e13b32 100644 --- a/flagscale/runner/straggler/__init__.py +++ b/flagscale/runner/straggler/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """FlagScale straggler detection utilities.""" from .comm import CommProfiler, CommStatsCollector, GlooCommHook, NCCLCommHook diff --git a/flagscale/runner/straggler/comm.py b/flagscale/runner/straggler/comm.py index 873df7d1c1..dba50a3f8d 100644 --- a/flagscale/runner/straggler/comm.py +++ b/flagscale/runner/straggler/comm.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Communication monitoring helpers for straggler analysis.""" import time diff --git a/flagscale/runner/straggler/config.py b/flagscale/runner/straggler/config.py index 4cd8365a58..321cfa9c11 100644 --- a/flagscale/runner/straggler/config.py +++ b/flagscale/runner/straggler/config.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Configuration types for FlagScale straggler detection.""" from dataclasses import dataclass, field diff --git a/flagscale/runner/straggler/detector.py b/flagscale/runner/straggler/detector.py index 33aa7d0b37..82774e0a79 100644 --- a/flagscale/runner/straggler/detector.py +++ b/flagscale/runner/straggler/detector.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Main detector implementation for FlagScale straggler analysis.""" import json diff --git a/flagscale/runner/straggler/healthcheck.py b/flagscale/runner/straggler/healthcheck.py index c621403f05..4525f82e24 100644 --- a/flagscale/runner/straggler/healthcheck.py +++ b/flagscale/runner/straggler/healthcheck.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Network health helpers for straggler workflows.""" import socket diff --git a/flagscale/runner/straggler/report.py b/flagscale/runner/straggler/report.py index 8f5272488a..bf1728ec9b 100644 --- a/flagscale/runner/straggler/report.py +++ b/flagscale/runner/straggler/report.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Serialization and text formatting for straggler reports.""" from typing import Any diff --git a/flagscale/runner/straggler/section.py b/flagscale/runner/straggler/section.py index 0bc9573973..a1429eed39 100644 --- a/flagscale/runner/straggler/section.py +++ b/flagscale/runner/straggler/section.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Section-level timing utilities for straggler detection.""" import time diff --git a/flagscale/runner/utils.py b/flagscale/runner/utils.py index 083f2ca612..3c60474773 100644 --- a/flagscale/runner/utils.py +++ b/flagscale/runner/utils.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import asyncio import collections import json diff --git a/flagscale/serve/__init__.py b/flagscale/serve/__init__.py index 646bfce3c0..566d35bcf0 100644 --- a/flagscale/serve/__init__.py +++ b/flagscale/serve/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from .utils import load_args, task_config __all__ = ["task_config", "load_args"] diff --git a/flagscale/serve/args_mapping/common_args.yaml b/flagscale/serve/args_mapping/common_args.yaml index 8f8a0e0535..3213a5b3b6 100644 --- a/flagscale/serve/args_mapping/common_args.yaml +++ b/flagscale/serve/args_mapping/common_args.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Common serve args for all backends common_args: # Server binding configuration diff --git a/flagscale/serve/args_mapping/gen_template_conf.py b/flagscale/serve/args_mapping/gen_template_conf.py index 5c5d55630b..899b757afe 100644 --- a/flagscale/serve/args_mapping/gen_template_conf.py +++ b/flagscale/serve/args_mapping/gen_template_conf.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # It creates a mapping conf of argument names to placeholder values and saves # the configuration to a file named "template_conf.yaml". import yaml diff --git a/flagscale/serve/args_mapping/gen_template_funcs.py b/flagscale/serve/args_mapping/gen_template_funcs.py index 83dbced78c..bd2ecfe6aa 100644 --- a/flagscale/serve/args_mapping/gen_template_funcs.py +++ b/flagscale/serve/args_mapping/gen_template_funcs.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # It creates a py file named "template_funcs.py" with functions that are # responsible for converting the key-value to the desired format import argparse diff --git a/flagscale/serve/args_mapping/mapping.py b/flagscale/serve/args_mapping/mapping.py index 05b6f41836..cb43469a25 100644 --- a/flagscale/serve/args_mapping/mapping.py +++ b/flagscale/serve/args_mapping/mapping.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import importlib from pathlib import Path diff --git a/flagscale/serve/args_mapping/mapping.yaml b/flagscale/serve/args_mapping/mapping.yaml index 6f804ed0c3..3d46b29a3e 100644 --- a/flagscale/serve/args_mapping/mapping.yaml +++ b/flagscale/serve/args_mapping/mapping.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llama_cpp: # Serve args(keys) mapping from vllm to llama_cpp key_mapping: diff --git a/flagscale/serve/args_mapping/mapping_funcs/llama_cpp.py b/flagscale/serve/args_mapping/mapping_funcs/llama_cpp.py index e8b6b36eef..d9618b7bf0 100644 --- a/flagscale/serve/args_mapping/mapping_funcs/llama_cpp.py +++ b/flagscale/serve/args_mapping/mapping_funcs/llama_cpp.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import json import os diff --git a/flagscale/serve/arguments.py b/flagscale/serve/arguments.py index 06c864a740..a6faf98a1b 100644 --- a/flagscale/serve/arguments.py +++ b/flagscale/serve/arguments.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import yaml diff --git a/flagscale/serve/core.py b/flagscale/serve/core.py index 1784b8b1e8..690f23c461 100644 --- a/flagscale/serve/core.py +++ b/flagscale/serve/core.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import ray diff --git a/flagscale/serve/dag_utils.py b/flagscale/serve/dag_utils.py index 50ea99c394..c2d0f3fc9a 100644 --- a/flagscale/serve/dag_utils.py +++ b/flagscale/serve/dag_utils.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import socket from flagscale.logger import logger diff --git a/flagscale/serve/data_process/action_chunk_to_fast_token.py b/flagscale/serve/data_process/action_chunk_to_fast_token.py index 1969babb04..ced0445cc5 100644 --- a/flagscale/serve/data_process/action_chunk_to_fast_token.py +++ b/flagscale/serve/data_process/action_chunk_to_fast_token.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Adopted from FlagOpen/RoboBrain-X0 (https://github.com/FlagOpen/RoboBrain-X0/blob/main/data_process/action_token/action_chunk_to_fast_token.py) diff --git a/flagscale/serve/engine.py b/flagscale/serve/engine.py index 0247670ce1..73e47f4403 100644 --- a/flagscale/serve/engine.py +++ b/flagscale/serve/engine.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import ast import importlib import importlib.util diff --git a/flagscale/serve/managers/pynvml_manager.py b/flagscale/serve/managers/pynvml_manager.py index 5a3dd88be3..e0b39103ed 100644 --- a/flagscale/serve/managers/pynvml_manager.py +++ b/flagscale/serve/managers/pynvml_manager.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import pynvml diff --git a/flagscale/serve/metric/__init__.py b/flagscale/serve/metric/__init__.py index 03dfd0db57..aa1a741b5e 100644 --- a/flagscale/serve/metric/__init__.py +++ b/flagscale/serve/metric/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from .serve_metric import calculate_metrics __all__ = ["calculate_metrics"] diff --git a/flagscale/serve/metric/serve_metric.py b/flagscale/serve/metric/serve_metric.py index 3ee2481b2f..6f84751e6f 100644 --- a/flagscale/serve/metric/serve_metric.py +++ b/flagscale/serve/metric/serve_metric.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import warnings from dataclasses import dataclass diff --git a/flagscale/serve/msgpack_numpy.py b/flagscale/serve/msgpack_numpy.py index b8dc50379b..859d291240 100644 --- a/flagscale/serve/msgpack_numpy.py +++ b/flagscale/serve/msgpack_numpy.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Copied from: # https://github.com/starVLA/starVLA/blob/3f7feefbc5fc25890ad3a7d262b8a0aea1339aa7/deployment/model_server/tools/msgpack_numpy.py diff --git a/flagscale/serve/processor/__init__.py b/flagscale/serve/processor/__init__.py index 6e9f116b76..b2a553c76a 100644 --- a/flagscale/serve/processor/__init__.py +++ b/flagscale/serve/processor/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from .image_layout_processor import ImageLayoutProcessorStep from .image_resize_processor import ImageResizeProcessorStep diff --git a/flagscale/serve/processor/image_layout_processor.py b/flagscale/serve/processor/image_layout_processor.py index d22dce0c4d..8484566ef1 100644 --- a/flagscale/serve/processor/image_layout_processor.py +++ b/flagscale/serve/processor/image_layout_processor.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from dataclasses import dataclass from typing import Any, Literal diff --git a/flagscale/serve/processor/image_resize_processor.py b/flagscale/serve/processor/image_resize_processor.py index 7de7457087..960a8796fa 100644 --- a/flagscale/serve/processor/image_resize_processor.py +++ b/flagscale/serve/processor/image_resize_processor.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from dataclasses import dataclass, field from typing import Any diff --git a/flagscale/serve/run_fs_serve_vllm.py b/flagscale/serve/run_fs_serve_vllm.py index bda11b4ee2..350ba08da4 100644 --- a/flagscale/serve/run_fs_serve_vllm.py +++ b/flagscale/serve/run_fs_serve_vllm.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import logging from ray import serve diff --git a/flagscale/serve/run_inference_engine.py b/flagscale/serve/run_inference_engine.py index 87a53b3969..44eaa0c9aa 100644 --- a/flagscale/serve/run_inference_engine.py +++ b/flagscale/serve/run_inference_engine.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import logging as logger import os import socket diff --git a/flagscale/serve/run_serve.py b/flagscale/serve/run_serve.py index b517efe414..85c7326a96 100644 --- a/flagscale/serve/run_serve.py +++ b/flagscale/serve/run_serve.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from flagscale.serve.arguments import parse_config from flagscale.serve.engine import ServeEngine diff --git a/flagscale/serve/run_serve_robobrain_x0.5_libero.py b/flagscale/serve/run_serve_robobrain_x0.5_libero.py index 64189df130..542fbaa31f 100644 --- a/flagscale/serve/run_serve_robobrain_x0.5_libero.py +++ b/flagscale/serve/run_serve_robobrain_x0.5_libero.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Adopted from FlagOpen/RoboBrain-X0 (https://github.com/FlagOpen/RoboBrain-X0/blob/main/agilex/server_agilex.py) import base64 diff --git a/flagscale/serve/run_serve_robobrain_x0_agilex.py b/flagscale/serve/run_serve_robobrain_x0_agilex.py index c185c8b7de..962c6a465a 100644 --- a/flagscale/serve/run_serve_robobrain_x0_agilex.py +++ b/flagscale/serve/run_serve_robobrain_x0_agilex.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """ Adopted from FlagOpen/RoboBrain-X0 (https://github.com/FlagOpen/RoboBrain-X0/blob/main/agilex/server_agilex.py) diff --git a/flagscale/serve/run_serve_vla.py b/flagscale/serve/run_serve_vla.py index 5083c2f4f3..6eb9274aaa 100644 --- a/flagscale/serve/run_serve_vla.py +++ b/flagscale/serve/run_serve_vla.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Mainly adopted from: # https://github.com/starVLA/starVLA/blob/3f7feefbc5fc25890ad3a7d262b8a0aea1339aa7/deployment/model_server/server_policy.py diff --git a/flagscale/serve/utils.py b/flagscale/serve/utils.py index 18830e95c0..02718b8c73 100644 --- a/flagscale/serve/utils.py +++ b/flagscale/serve/utils.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import threading from functools import wraps diff --git a/flagscale/serve/websocket_policy_server.py b/flagscale/serve/websocket_policy_server.py index 60034e790b..c4c15defb5 100644 --- a/flagscale/serve/websocket_policy_server.py +++ b/flagscale/serve/websocket_policy_server.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Mainly adopted from: # https://github.com/starVLA/starVLA/blob/3f7feefbc5fc25890ad3a7d262b8a0aea1339aa7/deployment/model_server/tools/websocket_policy_server.py diff --git a/flagscale/train/datasets/lerobot_mixture_dataset.py b/flagscale/train/datasets/lerobot_mixture_dataset.py index 32414fa526..68f2df904c 100644 --- a/flagscale/train/datasets/lerobot_mixture_dataset.py +++ b/flagscale/train/datasets/lerobot_mixture_dataset.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import hashlib import random from collections.abc import Sequence diff --git a/flagscale/train/megatron/__init__.py b/flagscale/train/megatron/__init__.py index c917f8f6d1..19a492b86d 100644 --- a/flagscale/train/megatron/__init__.py +++ b/flagscale/train/megatron/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Namespace package configuration # This allows megatron.core and megatron.plugin to be imported from the installed # megatron-core package, while other modules (training, rl, post_training, legacy) @@ -6,12 +20,14 @@ # Make this a namespace package to allow imports from multiple locations try: import pkgutil + __path__ = pkgutil.extend_path(__path__, __name__) except (AttributeError, NameError): # If __path__ doesn't exist yet, create it import os + __path__ = [os.path.dirname(__file__)] # The installed megatron-core package will have megatron.core and megatron.plugin # They will be automatically available through the namespace package mechanism -# No need to explicitly import them here \ No newline at end of file +# No need to explicitly import them here diff --git a/flagscale/train/megatron/gpt_builders.py b/flagscale/train/megatron/gpt_builders.py index e3a1d56e7c..08b487474f 100644 --- a/flagscale/train/megatron/gpt_builders.py +++ b/flagscale/train/megatron/gpt_builders.py @@ -25,14 +25,23 @@ # NOTE: Loading `megatron.legacy.model` earlier fails due to circular import -def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_collection=None, dualpipev_stage=None): - print_rank_0('building GPT model ...') +def gpt_builder( + args, + pre_process, + post_process, + vp_stage=None, + config=None, + pg_collection=None, + dualpipev_stage=None, +): + print_rank_0("building GPT model ...") if config is None: if args.yaml_cfg is not None: config = core_transformer_config_from_yaml(args, "language_model") else: ######### FlagScale Begin ######### from megatron.plugin.hetero.parallel_context import get_parallel_context + para_ctx = get_parallel_context() if para_ctx is not None: config = para_ctx.get_transformer_config() @@ -80,7 +89,7 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_ if args.mtp_num_layers is not None: assert not (config.transformer_impl == "inference_optimized") if ( - hasattr(transformer_layer_spec, 'layer_specs') + hasattr(transformer_layer_spec, "layer_specs") and len(transformer_layer_spec.layer_specs) == 0 ): # Get the decoder layer spec explicitly if no decoder layer in the last stage, @@ -89,7 +98,11 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_ else: # Define the decoder block spec decoder_layer_specs = get_gpt_decoder_layer_specs( - config, use_transformer_engine=use_te, normalization=args.normalization, qk_l2_norm=args.qk_l2_norm, vp_stage=vp_stage + config, + use_transformer_engine=use_te, + normalization=args.normalization, + qk_l2_norm=args.qk_l2_norm, + vp_stage=vp_stage, ) transformer_layer_spec_for_mtp = decoder_layer_specs[-1] # Use spec of the last layer in decoder block as spec of the transformer layer in MTP diff --git a/flagscale/train/megatron/legacy/fp16_deprecated/loss_scaler.py b/flagscale/train/megatron/legacy/fp16_deprecated/loss_scaler.py index cb64aa9289..dd77ca194d 100755 --- a/flagscale/train/megatron/legacy/fp16_deprecated/loss_scaler.py +++ b/flagscale/train/megatron/legacy/fp16_deprecated/loss_scaler.py @@ -1,26 +1,29 @@ # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. - -"""For backward compatibility, we need the class definitions to deserialize.""" - -class LossScaler: - def __init__(self, scale=1): - self.cur_scale = scale - -class DynamicLossScaler: - def __init__(self, - init_scale=2**32, - scale_factor=2., - scale_window=1000, - min_scale=1, - delayed_shift=1, - consecutive_hysteresis=False): - self.cur_scale = init_scale - self.cur_iter = 0 - self.last_overflow_iter = -1 - self.scale_factor = scale_factor - self.scale_window = scale_window - self.min_scale = min_scale - self.delayed_shift = delayed_shift - self.cur_hysteresis = delayed_shift - self.consecutive_hysteresis = consecutive_hysteresis - + +"""For backward compatibility, we need the class definitions to deserialize.""" + + +class LossScaler: + def __init__(self, scale=1): + self.cur_scale = scale + + +class DynamicLossScaler: + def __init__( + self, + init_scale=2**32, + scale_factor=2.0, + scale_window=1000, + min_scale=1, + delayed_shift=1, + consecutive_hysteresis=False, + ): + self.cur_scale = init_scale + self.cur_iter = 0 + self.last_overflow_iter = -1 + self.scale_factor = scale_factor + self.scale_window = scale_window + self.min_scale = min_scale + self.delayed_shift = delayed_shift + self.cur_hysteresis = delayed_shift + self.consecutive_hysteresis = consecutive_hysteresis diff --git a/flagscale/train/megatron/legacy/model/bert_model.py b/flagscale/train/megatron/legacy/model/bert_model.py index 0b233a8e55..42f22918d3 100644 --- a/flagscale/train/megatron/legacy/model/bert_model.py +++ b/flagscale/train/megatron/legacy/model/bert_model.py @@ -17,8 +17,10 @@ from .module import MegatronModule from megatron.plugin.platform import get_platform + cur_platform = get_platform() + def bert_extended_attention_mask(attention_mask): # We create a 3D attention mask from a 2D tensor mask. # [b, 1, s] @@ -31,15 +33,15 @@ def bert_extended_attention_mask(attention_mask): extended_attention_mask = attention_mask_bss.unsqueeze(1) # Convert attention mask to binary: - extended_attention_mask = (extended_attention_mask < 0.5) + extended_attention_mask = extended_attention_mask < 0.5 return extended_attention_mask + def bert_position_ids(token_ids): # Create position ids seq_length = token_ids.size(1) - position_ids = torch.arange(seq_length, dtype=torch.long, - device=token_ids.device) + position_ids = torch.arange(seq_length, dtype=torch.long, device=token_ids.device) position_ids = position_ids.unsqueeze(0).expand_as(token_ids) return position_ids @@ -63,8 +65,8 @@ def __init__(self, mpu_vocab_size, config, parallel_output): self.parallel_output = parallel_output self.dense = get_linear_layer(config.hidden_size, config.hidden_size, config.init_method) - setattr(self.dense.weight, 'sequence_parallel', config.sequence_parallel) - setattr(self.dense.bias, 'sequence_parallel', config.sequence_parallel) + setattr(self.dense.weight, "sequence_parallel", config.sequence_parallel) + setattr(self.dense.bias, "sequence_parallel", config.sequence_parallel) self.norm = get_norm(config) self.gelu = torch.nn.functional.gelu @@ -77,10 +79,9 @@ def forward(self, hidden_states, word_embeddings_weight): hidden_states = self.dense(hidden_states) hidden_states = self.gelu(hidden_states) hidden_states = self.norm(hidden_states) - output = parallel_lm_logits(hidden_states, - word_embeddings_weight, - self.parallel_output, - bias=self.bias) + output = parallel_lm_logits( + hidden_states, word_embeddings_weight, self.parallel_output, bias=self.bias + ) return output def load_state_dict(self, state_dict, strict=True): @@ -95,14 +96,11 @@ def load_state_dict(self, state_dict, strict=True): super().load_state_dict(state_dict_, strict) -def post_language_model_processing(lm_output, pooled_output, - lm_head, binary_head, - lm_labels, - logit_weights, - fp16_lm_cross_entropy): +def post_language_model_processing( + lm_output, pooled_output, lm_head, binary_head, lm_labels, logit_weights, fp16_lm_cross_entropy +): # Output. - lm_logits = lm_head( - lm_output, logit_weights) + lm_logits = lm_head(lm_output, logit_weights) binary_logits = None if binary_head is not None: @@ -110,32 +108,33 @@ def post_language_model_processing(lm_output, pooled_output, if lm_labels is None: # [s b h] => [b s h] - return lm_logits.transpose(0,1).contiguous(), binary_logits + return lm_logits.transpose(0, 1).contiguous(), binary_logits else: # [b s] => [s b] - lm_labels = lm_labels.transpose(0,1).contiguous() + lm_labels = lm_labels.transpose(0, 1).contiguous() # lm_logits : [s, b, h] and lm_labels: [s, b] if fp16_lm_cross_entropy: assert lm_logits.dtype == torch.half lm_loss = tensor_parallel.vocab_parallel_cross_entropy(lm_logits, lm_labels) else: - lm_loss = tensor_parallel.vocab_parallel_cross_entropy(lm_logits.float(), - lm_labels) + lm_loss = tensor_parallel.vocab_parallel_cross_entropy(lm_logits.float(), lm_labels) # [s, b] => [b s] - lm_loss = lm_loss.transpose(0,1).contiguous() + lm_loss = lm_loss.transpose(0, 1).contiguous() return lm_loss, binary_logits class BertModel(MegatronModule): """Bert Language model.""" - def __init__(self, - config, - num_tokentypes=2, - add_binary_head=True, - parallel_output=True, - pre_process=True, - post_process=True): + def __init__( + self, + config, + num_tokentypes=2, + add_binary_head=True, + parallel_output=True, + pre_process=True, + post_process=True, + ): super().__init__(config=config) args = get_args() @@ -158,34 +157,38 @@ def __init__(self, add_pooler=self.add_binary_head, encoder_attn_mask_type=AttnMaskType.padding, pre_process=self.pre_process, - post_process=self.post_process) + post_process=self.post_process, + ) self.initialize_word_embeddings() if self.post_process: - self.lm_head = BertLMHead(self.shared_embedding_or_output_weight().size(0), config, parallel_output) - self._lm_head_key = 'lm_head' + self.lm_head = BertLMHead( + self.shared_embedding_or_output_weight().size(0), config, parallel_output + ) + self._lm_head_key = "lm_head" self.binary_head = None if self.add_binary_head: - self.binary_head = get_linear_layer(config.hidden_size, 2, - config.init_method) - self._binary_head_key = 'binary_head' + self.binary_head = get_linear_layer(config.hidden_size, 2, config.init_method) + self._binary_head_key = "binary_head" def set_input_tensor(self, input_tensor): """See megatron.legacy.model.transformer.set_input_tensor()""" self.language_model.set_input_tensor(input_tensor) - def forward(self, bert_model_input, attention_mask, - tokentype_ids=None, lm_labels=None, inference_context=None): - + def forward( + self, + bert_model_input, + attention_mask, + tokentype_ids=None, + lm_labels=None, + inference_context=None, + ): extended_attention_mask = bert_extended_attention_mask(attention_mask) input_ids = bert_model_input position_ids = bert_position_ids(input_ids) lm_output = self.language_model( - input_ids, - position_ids, - extended_attention_mask, - tokentype_ids=tokentype_ids + input_ids, position_ids, extended_attention_mask, tokentype_ids=tokentype_ids ) if self.post_process and self.add_binary_head: @@ -193,7 +196,6 @@ def forward(self, bert_model_input, attention_mask, # Return pooled output (e.g., when computing Bert embeddings). if self.return_embeddings: - # Sum attention mask. embeddings = torch.transpose(lm_output, 0, 1) masks = torch.sum(attention_mask, dim=1) @@ -202,9 +204,10 @@ def forward(self, bert_model_input, attention_mask, output = torch.zeros( size=(embeddings.shape[0], embeddings.shape[2]), dtype=torch.float32, - device=cur_platform.current_device()) + device=cur_platform.current_device(), + ) for i, (embedding, mask) in enumerate(zip(embeddings, masks)): - output[i, :] = torch.mean(embedding[1: mask - 1], dim=0) + output[i, :] = torch.mean(embedding[1 : mask - 1], dim=0) return output @@ -212,48 +215,51 @@ def forward(self, bert_model_input, attention_mask, pooled_output = None if self.post_process: - return post_language_model_processing(lm_output, pooled_output, - self.lm_head, self.binary_head, - lm_labels, - self.shared_embedding_or_output_weight(), - self.fp16_lm_cross_entropy) + return post_language_model_processing( + lm_output, + pooled_output, + self.lm_head, + self.binary_head, + lm_labels, + self.shared_embedding_or_output_weight(), + self.fp16_lm_cross_entropy, + ) else: return lm_output - - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): + def state_dict_for_save_checkpoint(self, prefix="", keep_vars=False): """For easy load when model is combined with other heads, add an extra key.""" state_dict_ = {} - state_dict_[self._language_model_key] \ - = self.language_model.state_dict_for_save_checkpoint(prefix=prefix, - keep_vars=keep_vars) + state_dict_[self._language_model_key] = self.language_model.state_dict_for_save_checkpoint( + prefix=prefix, keep_vars=keep_vars + ) if self.post_process: - state_dict_[self._lm_head_key] \ - = self.lm_head.state_dict_for_save_checkpoint(prefix=prefix, - keep_vars=keep_vars) + state_dict_[self._lm_head_key] = self.lm_head.state_dict_for_save_checkpoint( + prefix=prefix, keep_vars=keep_vars + ) if self.post_process and self.add_binary_head: - state_dict_[self._binary_head_key] \ - = self.binary_head.state_dict(prefix=prefix, keep_vars=keep_vars) + state_dict_[self._binary_head_key] = self.binary_head.state_dict( + prefix=prefix, keep_vars=keep_vars + ) # Save word_embeddings. if self.post_process and not self.pre_process: - state_dict_[self._word_embeddings_for_head_key] \ - = self.word_embeddings.state_dict(prefix=prefix, keep_vars=keep_vars) + state_dict_[self._word_embeddings_for_head_key] = self.word_embeddings.state_dict( + prefix=prefix, keep_vars=keep_vars + ) return state_dict_ def load_state_dict(self, state_dict, strict=True): """Customized load.""" - self.language_model.load_state_dict( - state_dict[self._language_model_key], strict=strict) + self.language_model.load_state_dict(state_dict[self._language_model_key], strict=strict) if self.post_process: - self.lm_head.load_state_dict( - state_dict[self._lm_head_key], strict=strict) + self.lm_head.load_state_dict(state_dict[self._lm_head_key], strict=strict) if self.post_process and self.add_binary_head: - self.binary_head.load_state_dict( - state_dict[self._binary_head_key], strict=strict) + self.binary_head.load_state_dict(state_dict[self._binary_head_key], strict=strict) # Load word_embeddings. if self.post_process and not self.pre_process: self.word_embeddings.load_state_dict( - state_dict[self._word_embeddings_for_head_key], strict=strict) + state_dict[self._word_embeddings_for_head_key], strict=strict + ) diff --git a/flagscale/train/megatron/legacy/model/biencoder_model.py b/flagscale/train/megatron/legacy/model/biencoder_model.py index df787686b4..a4d5183e13 100644 --- a/flagscale/train/megatron/legacy/model/biencoder_model.py +++ b/flagscale/train/megatron/legacy/model/biencoder_model.py @@ -16,36 +16,42 @@ from megatron.legacy.model.utils import scaled_init_method_normal from .module import MegatronModule -def get_model_provider(only_query_model=False, only_context_model=False, - biencoder_shared_query_context_model=False): +def get_model_provider( + only_query_model=False, only_context_model=False, biencoder_shared_query_context_model=False +): def model_provider(pre_process=True, post_process=True): """Build the model.""" - print_rank_0('building Bienoder model ...') - model = biencoder_model_provider(only_query_model=only_query_model, - only_context_model = only_context_model, - biencoder_shared_query_context_model = \ - biencoder_shared_query_context_model, - pre_process=pre_process, post_process=post_process) + print_rank_0("building Bienoder model ...") + model = biencoder_model_provider( + only_query_model=only_query_model, + only_context_model=only_context_model, + biencoder_shared_query_context_model=biencoder_shared_query_context_model, + pre_process=pre_process, + post_process=post_process, + ) return model return model_provider -def biencoder_model_provider(only_query_model=False, - only_context_model=False, - biencoder_shared_query_context_model=False, - pre_process=True, - post_process=True): +def biencoder_model_provider( + only_query_model=False, + only_context_model=False, + biencoder_shared_query_context_model=False, + pre_process=True, + post_process=True, +): """Build the model.""" - assert mpu.get_tensor_model_parallel_world_size() == 1 and \ - mpu.get_pipeline_model_parallel_world_size() == 1, \ - "Model parallel size > 1 not supported for ICT" + assert ( + mpu.get_tensor_model_parallel_world_size() == 1 + and mpu.get_pipeline_model_parallel_world_size() == 1 + ), "Model parallel size > 1 not supported for ICT" - print_rank_0('building BiEncoderModel...') + print_rank_0("building BiEncoderModel...") # simpler to just keep using 2 tokentypes since # the LM we initialize with has 2 tokentypes @@ -54,10 +60,10 @@ def biencoder_model_provider(only_query_model=False, parallel_output=False, only_query_model=only_query_model, only_context_model=only_context_model, - biencoder_shared_query_context_model=\ - biencoder_shared_query_context_model, + biencoder_shared_query_context_model=biencoder_shared_query_context_model, pre_process=pre_process, - post_process=post_process) + post_process=post_process, + ) return model @@ -65,14 +71,16 @@ def biencoder_model_provider(only_query_model=False, class BiEncoderModel(MegatronModule): """Bert-based module for Biencoder model.""" - def __init__(self, - num_tokentypes=1, - parallel_output=True, - only_query_model=False, - only_context_model=False, - biencoder_shared_query_context_model=False, - pre_process=True, - post_process=True): + def __init__( + self, + num_tokentypes=1, + parallel_output=True, + only_query_model=False, + only_context_model=False, + biencoder_shared_query_context_model=False, + pre_process=True, + post_process=True, + ): super(BiEncoderModel, self).__init__() args = get_args() @@ -80,10 +88,10 @@ def __init__(self, num_tokentypes=num_tokentypes, parallel_output=parallel_output, pre_process=pre_process, - post_process=post_process) + post_process=post_process, + ) - self.biencoder_shared_query_context_model = \ - biencoder_shared_query_context_model + self.biencoder_shared_query_context_model = biencoder_shared_query_context_model assert not (only_context_model and only_query_model) self.use_context_model = not only_query_model self.use_query_model = not only_context_model @@ -91,18 +99,18 @@ def __init__(self, if self.biencoder_shared_query_context_model: self.model = PretrainedBertModel(**bert_kwargs) - self._model_key = 'shared_model' + self._model_key = "shared_model" self.query_model, self.context_model = self.model, self.model else: if self.use_query_model: # this model embeds (pseudo-)queries - Embed_input in the paper self.query_model = PretrainedBertModel(**bert_kwargs) - self._query_key = 'query_model' + self._query_key = "query_model" if self.use_context_model: # this model embeds evidence blocks - Embed_doc in the paper self.context_model = PretrainedBertModel(**bert_kwargs) - self._context_key = 'context_model' + self._context_key = "context_model" def set_input_tensor(self, input_tensor): """See megatron.legacy.model.transformer.set_input_tensor()""" @@ -111,23 +119,28 @@ def set_input_tensor(self, input_tensor): # self.language_model.set_input_tensor(input_tensor) return - def forward(self, query_tokens, query_attention_mask, query_types, - context_tokens, context_attention_mask, context_types): + def forward( + self, + query_tokens, + query_attention_mask, + query_types, + context_tokens, + context_attention_mask, + context_types, + ): """Run a forward pass for each of the models and return the respective embeddings.""" if self.use_query_model: - query_logits = self.embed_text(self.query_model, - query_tokens, - query_attention_mask, - query_types) + query_logits = self.embed_text( + self.query_model, query_tokens, query_attention_mask, query_types + ) else: raise ValueError("Cannot embed query without the query model.") if self.use_context_model: - context_logits = self.embed_text(self.context_model, - context_tokens, - context_attention_mask, - context_types) + context_logits = self.embed_text( + self.context_model, context_tokens, context_attention_mask, context_types + ) else: raise ValueError("Cannot embed block without the block model.") return query_logits, context_logits @@ -135,28 +148,26 @@ def forward(self, query_tokens, query_attention_mask, query_types, @staticmethod def embed_text(model, tokens, attention_mask, token_types): """Embed a batch of tokens using the model""" - logits = model(tokens, - attention_mask, - token_types) + logits = model(tokens, attention_mask, token_types) return logits - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): + def state_dict_for_save_checkpoint(self, prefix="", keep_vars=False): """Save dict with state dicts of each of the models.""" state_dict_ = {} if self.biencoder_shared_query_context_model: - state_dict_[self._model_key] = \ - self.model.state_dict_for_save_checkpoint( - prefix=prefix, keep_vars=keep_vars) + state_dict_[self._model_key] = self.model.state_dict_for_save_checkpoint( + prefix=prefix, keep_vars=keep_vars + ) else: if self.use_query_model: - state_dict_[self._query_key] = \ - self.query_model.state_dict_for_save_checkpoint( - prefix=prefix, keep_vars=keep_vars) + state_dict_[self._query_key] = self.query_model.state_dict_for_save_checkpoint( + prefix=prefix, keep_vars=keep_vars + ) if self.use_context_model: - state_dict_[self._context_key] = \ - self.context_model.state_dict_for_save_checkpoint( - prefix=prefix, keep_vars=keep_vars) + state_dict_[self._context_key] = self.context_model.state_dict_for_save_checkpoint( + prefix=prefix, keep_vars=keep_vars + ) return state_dict_ @@ -164,18 +175,15 @@ def load_state_dict(self, state_dict, strict=True): """Load the state dicts of each of the models""" if self.biencoder_shared_query_context_model: print_rank_0("Loading shared query-context model") - self.model.load_state_dict(state_dict[self._model_key], \ - strict=strict) + self.model.load_state_dict(state_dict[self._model_key], strict=strict) else: if self.use_query_model: print_rank_0("Loading query model") - self.query_model.load_state_dict( \ - state_dict[self._query_key], strict=strict) + self.query_model.load_state_dict(state_dict[self._query_key], strict=strict) if self.use_context_model: print_rank_0("Loading context model") - self.context_model.load_state_dict( \ - state_dict[self._context_key], strict=strict) + self.context_model.load_state_dict(state_dict[self._context_key], strict=strict) def init_state_dict_from_bert(self): """Initialize the state from a pretrained BERT model @@ -189,37 +197,41 @@ def init_state_dict_from_bert(self): tracker_filename = get_checkpoint_tracker_filename(args.bert_load) if not os.path.isfile(tracker_filename): raise FileNotFoundError("Could not find BERT checkpoint") - with open(tracker_filename, 'r') as f: + with open(tracker_filename, "r") as f: iteration = int(f.read().strip()) assert iteration > 0 checkpoint_name = get_checkpoint_name(args.bert_load, iteration, False) if mpu.get_data_parallel_rank() == 0: - print('global rank {} is loading BERT checkpoint {}'.format( - torch.distributed.get_rank(), checkpoint_name)) + print( + "global rank {} is loading BERT checkpoint {}".format( + torch.distributed.get_rank(), checkpoint_name + ) + ) # Load the checkpoint. try: - state_dict = torch.load(checkpoint_name, map_location='cpu') + state_dict = torch.load(checkpoint_name, map_location="cpu") except ModuleNotFoundError: from megatron.legacy.fp16_deprecated import loss_scaler + # For backward compatibility. - print_rank_0(' > deserializing using the old code structure ...') - sys.modules['fp16.loss_scaler'] = sys.modules[ - 'megatron.fp16_deprecated.loss_scaler'] - sys.modules['megatron.fp16.loss_scaler'] = sys.modules[ - 'megatron.fp16_deprecated.loss_scaler'] - state_dict = torch.load(checkpoint_name, map_location='cpu') - sys.modules.pop('fp16.loss_scaler', None) - sys.modules.pop('megatron.fp16.loss_scaler', None) + print_rank_0(" > deserializing using the old code structure ...") + sys.modules["fp16.loss_scaler"] = sys.modules["megatron.fp16_deprecated.loss_scaler"] + sys.modules["megatron.fp16.loss_scaler"] = sys.modules[ + "megatron.fp16_deprecated.loss_scaler" + ] + state_dict = torch.load(checkpoint_name, map_location="cpu") + sys.modules.pop("fp16.loss_scaler", None) + sys.modules.pop("megatron.fp16.loss_scaler", None) except Exception: - print_rank_0('could not load the BERT checkpoint') + print_rank_0("could not load the BERT checkpoint") sys.exit() - checkpoint_version = state_dict.get('checkpoint_version', 0) + checkpoint_version = state_dict.get("checkpoint_version", 0) # load the LM state dict into each model - model_dict = state_dict['model']['language_model'] + model_dict = state_dict["model"]["language_model"] if self.biencoder_shared_query_context_model: self.model.language_model.load_state_dict(model_dict) @@ -229,17 +241,15 @@ def init_state_dict_from_bert(self): self.query_model.language_model.load_state_dict(model_dict) # give each model the same ict_head to begin with as well if self.biencoder_projection_dim > 0: - query_proj_state_dict = \ - self.state_dict_for_save_checkpoint()\ - [self._query_key]['projection_enc'] + query_proj_state_dict = self.state_dict_for_save_checkpoint()[self._query_key][ + "projection_enc" + ] fix_query_key_value_ordering(self.query_model, checkpoint_version) if self.use_context_model: self.context_model.language_model.load_state_dict(model_dict) - if self.query_model is not None and \ - self.biencoder_projection_dim > 0: - self.context_model.projection_enc.load_state_dict\ - (query_proj_state_dict) + if self.query_model is not None and self.biencoder_projection_dim > 0: + self.context_model.projection_enc.load_state_dict(query_proj_state_dict) fix_query_key_value_ordering(self.context_model, checkpoint_version) @@ -247,8 +257,7 @@ class PretrainedBertModel(MegatronModule): """BERT-based encoder for queries or contexts used for learned information retrieval.""" - def __init__(self, num_tokentypes=2, - parallel_output=True, pre_process=True, post_process=True): + def __init__(self, num_tokentypes=2, parallel_output=True, pre_process=True, post_process=True): super(PretrainedBertModel, self).__init__() args = get_args() @@ -259,8 +268,7 @@ def __init__(self, num_tokentypes=2, self.pre_process = pre_process self.post_process = post_process init_method = init_method_normal(args.init_method_std) - scaled_init_method = scaled_init_method_normal( - args.init_method_std, args.num_layers) + scaled_init_method = scaled_init_method_normal(args.init_method_std, args.num_layers) self.language_model, self._language_model_key = get_language_model( num_tokentypes=num_tokentypes, @@ -269,23 +277,23 @@ def __init__(self, num_tokentypes=2, init_method=init_method, scaled_init_method=scaled_init_method, pre_process=self.pre_process, - post_process=self.post_process) + post_process=self.post_process, + ) if args.biencoder_projection_dim > 0: - self.projection_enc = get_linear_layer(args.hidden_size, - args.biencoder_projection_dim, - init_method) - self._projection_enc_key = 'projection_enc' + self.projection_enc = get_linear_layer( + args.hidden_size, args.biencoder_projection_dim, init_method + ) + self._projection_enc_key = "projection_enc" def forward(self, input_ids, attention_mask, tokentype_ids=None): extended_attention_mask = attention_mask.unsqueeze(1) - #extended_attention_mask = bert_extended_attention_mask(attention_mask) + # extended_attention_mask = bert_extended_attention_mask(attention_mask) position_ids = bert_position_ids(input_ids) - lm_output = self.language_model(input_ids, - position_ids, - extended_attention_mask, - tokentype_ids=tokentype_ids) + lm_output = self.language_model( + input_ids, position_ids, extended_attention_mask, tokentype_ids=tokentype_ids + ) # This mask will be used in average-pooling and max-pooling pool_mask = (input_ids == self.pad_id).unsqueeze(2) @@ -301,29 +309,27 @@ def forward(self, input_ids, attention_mask, tokentype_ids=None): return pooled_output - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): + def state_dict_for_save_checkpoint(self, prefix="", keep_vars=False): """For easy load when model is combined with other heads, add an extra key.""" state_dict_ = {} - state_dict_[self._language_model_key] \ - = self.language_model.state_dict_for_save_checkpoint( - prefix=prefix, keep_vars=keep_vars) + state_dict_[self._language_model_key] = self.language_model.state_dict_for_save_checkpoint( + prefix=prefix, keep_vars=keep_vars + ) if self.biencoder_projection_dim > 0: - state_dict_[self._projection_enc_key] = \ - self.projection_enc.state_dict(prefix=prefix, - keep_vars=keep_vars) + state_dict_[self._projection_enc_key] = self.projection_enc.state_dict( + prefix=prefix, keep_vars=keep_vars + ) return state_dict_ def load_state_dict(self, state_dict, strict=True): """Customized load.""" print_rank_0("loading pretrained weights") - self.language_model.load_state_dict( - state_dict[self._language_model_key], strict=strict) + self.language_model.load_state_dict(state_dict[self._language_model_key], strict=strict) if self.biencoder_projection_dim > 0: print_rank_0("loading projection head weights") - self.projection_enc.load_state_dict( - state_dict[self._projection_enc_key], strict=strict) + self.projection_enc.load_state_dict(state_dict[self._projection_enc_key], strict=strict) diff --git a/flagscale/train/megatron/legacy/model/classification.py b/flagscale/train/megatron/legacy/model/classification.py index c9fe165280..41e688e624 100644 --- a/flagscale/train/megatron/legacy/model/classification.py +++ b/flagscale/train/megatron/legacy/model/classification.py @@ -15,13 +15,7 @@ class Classification(MegatronModule): - - def __init__(self, - config, - num_classes, - num_tokentypes=2, - pre_process=True, - post_process=True): + def __init__(self, config, num_classes, num_tokentypes=2, pre_process=True, post_process=True): super().__init__(config=config, share_embeddings_and_output_weights=False) args = get_args() @@ -35,31 +29,28 @@ def __init__(self, add_pooler=True, encoder_attn_mask_type=AttnMaskType.padding, pre_process=self.pre_process, - post_process=self.post_process) + post_process=self.post_process, + ) # Multi-choice head. if self.post_process: self.classification_dropout = torch.nn.Dropout(args.hidden_dropout) - self.classification_head = get_linear_layer(args.hidden_size, - self.num_classes, - config.init_method) - self._classification_head_key = 'classification_head' + self.classification_head = get_linear_layer( + args.hidden_size, self.num_classes, config.init_method + ) + self._classification_head_key = "classification_head" def set_input_tensor(self, input_tensor): """See megatron.legacy.model.transformer.set_input_tensor()""" self.language_model.set_input_tensor(input_tensor) def forward(self, model_input, attention_mask, tokentype_ids=None): - extended_attention_mask = bert_extended_attention_mask(attention_mask) input_ids = model_input position_ids = bert_position_ids(input_ids) lm_output = self.language_model( - input_ids, - position_ids, - extended_attention_mask, - tokentype_ids=tokentype_ids + input_ids, position_ids, extended_attention_mask, tokentype_ids=tokentype_ids ) if self.post_process: @@ -73,29 +64,31 @@ def forward(self, model_input, attention_mask, tokentype_ids=None): return classification_logits return lm_output - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): + def state_dict_for_save_checkpoint(self, prefix="", keep_vars=False): """For easy load when model is combined with other heads, add an extra key.""" state_dict_ = {} - state_dict_[self._language_model_key] \ - = self.language_model.state_dict_for_save_checkpoint(prefix=prefix, - keep_vars=keep_vars) + state_dict_[self._language_model_key] = self.language_model.state_dict_for_save_checkpoint( + prefix=prefix, keep_vars=keep_vars + ) if self.post_process: - state_dict_[self._classification_head_key] \ - = self.classification_head.state_dict(prefix=prefix, keep_vars=keep_vars) + state_dict_[self._classification_head_key] = self.classification_head.state_dict( + prefix=prefix, keep_vars=keep_vars + ) return state_dict_ def load_state_dict(self, state_dict, strict=True): """Customized load.""" - self.language_model.load_state_dict( - state_dict[self._language_model_key], strict=strict) + self.language_model.load_state_dict(state_dict[self._language_model_key], strict=strict) if self.post_process: if self._classification_head_key in state_dict: self.classification_head.load_state_dict( - state_dict[self._classification_head_key], strict=strict) + state_dict[self._classification_head_key], strict=strict + ) else: - print_rank_last('***WARNING*** could not find {} in the checkpoint, ' - 'initializing to random'.format( - self._classification_head_key)) + print_rank_last( + "***WARNING*** could not find {} in the checkpoint, " + "initializing to random".format(self._classification_head_key) + ) diff --git a/flagscale/train/megatron/legacy/model/enums.py b/flagscale/train/megatron/legacy/model/enums.py index bab179d1a0..78d83c0000 100644 --- a/flagscale/train/megatron/legacy/model/enums.py +++ b/flagscale/train/megatron/legacy/model/enums.py @@ -2,17 +2,21 @@ import enum + class LayerType(enum.Enum): encoder = 1 decoder = 2 - + + class AttnType(enum.Enum): self_attn = 1 cross_attn = 2 + class AttnMaskType(enum.Enum): padding = 1 causal = 2 + # For backward compatibility with old model checkpoints from megatron.core.enums import ModelType diff --git a/flagscale/train/megatron/legacy/model/fused_bias_gelu.py b/flagscale/train/megatron/legacy/model/fused_bias_gelu.py index e00e63148b..ccd12f0d5b 100644 --- a/flagscale/train/megatron/legacy/model/fused_bias_gelu.py +++ b/flagscale/train/megatron/legacy/model/fused_bias_gelu.py @@ -12,10 +12,12 @@ # actual gelu is: # x * 0.5 * (1.0 + torch.erf(x * 0.70710678)) + @jit_fuser def bias_gelu(bias, y): x = bias + y - return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))) + return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))) + # gradient of tanh approximation of gelu # gradient of actual gelu is: @@ -25,8 +27,11 @@ def bias_gelu_back(g, bias, y): x = bias + y tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243 - ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out) - return ff*g + ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * ( + 1 + tanh_out + ) + return ff * g + class GeLUFunction(torch.autograd.Function): @staticmethod @@ -41,4 +46,5 @@ def backward(ctx, grad_output): tmp = bias_gelu_back(grad_output, bias, input) return tmp, tmp + bias_gelu_impl = GeLUFunction.apply diff --git a/flagscale/train/megatron/legacy/model/fused_layer_norm.py b/flagscale/train/megatron/legacy/model/fused_layer_norm.py index 5c35483874..207181b85e 100644 --- a/flagscale/train/megatron/legacy/model/fused_layer_norm.py +++ b/flagscale/train/megatron/legacy/model/fused_layer_norm.py @@ -1,8 +1,8 @@ # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. """This code is copied fron NVIDIA apex: - https://github.com/NVIDIA/apex - with some changes. """ + https://github.com/NVIDIA/apex +with some changes.""" import inspect import numbers @@ -15,6 +15,7 @@ try: from apex.contrib.layer_norm.layer_norm import FastLayerNormFN + HAVE_PERSIST_LAYER_NORM = True except ImportError: HAVE_PERSIST_LAYER_NORM = False @@ -29,11 +30,14 @@ class MixedFusedLayerNorm(torch.nn.Module): - - def __init__(self, normalized_shape, eps=1e-5, - no_persist_layer_norm=True, - sequence_parallel=False, - apply_layernorm_1p=False): + def __init__( + self, + normalized_shape, + eps=1e-5, + no_persist_layer_norm=True, + sequence_parallel=False, + apply_layernorm_1p=False, + ): super(MixedFusedLayerNorm, self).__init__() self.apply_layernorm_1p = apply_layernorm_1p @@ -44,11 +48,33 @@ def __init__(self, normalized_shape, eps=1e-5, # List of hiddens sizes supported in the persistent layer norm kernel # If the hidden size is not supported, fall back to the non-persistent # kernel. - persist_ln_hidden_sizes = [1024, 1536, 2048, 2304, 3072, 3840, 4096, - 5120, 6144, 8192, 10240, 12288, 12800, 15360, 16384, 18432, 20480, - 24576, 25600, 30720, 32768, 40960, 49152, 65536] - if normalized_shape not in persist_ln_hidden_sizes or \ - not HAVE_PERSIST_LAYER_NORM: + persist_ln_hidden_sizes = [ + 1024, + 1536, + 2048, + 2304, + 3072, + 3840, + 4096, + 5120, + 6144, + 8192, + 10240, + 12288, + 12800, + 15360, + 16384, + 18432, + 20480, + 24576, + 25600, + 30720, + 32768, + 40960, + 49152, + 65536, + ] + if normalized_shape not in persist_ln_hidden_sizes or not HAVE_PERSIST_LAYER_NORM: no_persist_layer_norm = True if isinstance(normalized_shape, numbers.Integral): @@ -62,38 +88,38 @@ def __init__(self, normalized_shape, eps=1e-5, self.sequence_parallel = sequence_parallel # set sequence parallelism flag on weight and bias parameters - setattr(self.weight, 'sequence_parallel', self.sequence_parallel) - setattr(self.bias, 'sequence_parallel', self.sequence_parallel) - - - def reset_parameters(self): + setattr(self.weight, "sequence_parallel", self.sequence_parallel) + setattr(self.bias, "sequence_parallel", self.sequence_parallel) - if self.apply_layernorm_1p: - init.zeros_(self.weight) - init.zeros_(self.bias) - else: - init.ones_(self.weight) - init.zeros_(self.bias) - - def forward(self, input): - - weight = self.weight + 1 if self.apply_layernorm_1p else self.weight - - if self.no_persist_layer_norm: - assert fused_layer_norm_affine is not None, \ - "fused_layer_norm_affine is not available, please install apex from https://github.com/NVIDIA/apex" - return fused_layer_norm_affine(input, weight, self.bias, self.normalized_shape, eps=self.eps) - else: - if 'memory_efficient' in inspect.getfullargspec(FastLayerNormFN.forward).args: - output = FastLayerNormFN.apply(input, weight, self.bias, self.eps, False) + def reset_parameters(self): + if self.apply_layernorm_1p: + init.zeros_(self.weight) + init.zeros_(self.bias) + else: + init.ones_(self.weight) + init.zeros_(self.bias) + + def forward(self, input): + weight = self.weight + 1 if self.apply_layernorm_1p else self.weight + + if self.no_persist_layer_norm: + assert fused_layer_norm_affine is not None, ( + "fused_layer_norm_affine is not available, please install apex from https://github.com/NVIDIA/apex" + ) + return fused_layer_norm_affine( + input, weight, self.bias, self.normalized_shape, eps=self.eps + ) else: - output = FastLayerNormFN.apply(input, weight, self.bias, self.eps) - # Apex's fast layer norm function outputs a 'view' tensor (i.e., has - # a populated '_base' field). This will result in schedule.py's - # deallocate_output_tensor() throwing an error, so a viewless tensor is - # created to prevent this. - output = make_viewless_tensor(inp = output, - requires_grad = input.requires_grad, - keep_graph = True) - - return output + if "memory_efficient" in inspect.getfullargspec(FastLayerNormFN.forward).args: + output = FastLayerNormFN.apply(input, weight, self.bias, self.eps, False) + else: + output = FastLayerNormFN.apply(input, weight, self.bias, self.eps) + # Apex's fast layer norm function outputs a 'view' tensor (i.e., has + # a populated '_base' field). This will result in schedule.py's + # deallocate_output_tensor() throwing an error, so a viewless tensor is + # created to prevent this. + output = make_viewless_tensor( + inp=output, requires_grad=input.requires_grad, keep_graph=True + ) + + return output diff --git a/flagscale/train/megatron/legacy/model/fused_softmax.py b/flagscale/train/megatron/legacy/model/fused_softmax.py index 58f900bddd..1fb2face87 100644 --- a/flagscale/train/megatron/legacy/model/fused_softmax.py +++ b/flagscale/train/megatron/legacy/model/fused_softmax.py @@ -19,12 +19,10 @@ def forward(ctx, inputs, scale): try: import scaled_upper_triang_masked_softmax_cuda except (ImportError, ModuleNotFoundError): - print(f'Please install Apex to use fused_softmax') + print(f"Please install Apex to use fused_softmax") scale_t = torch.tensor([scale]) - softmax_results = scaled_upper_triang_masked_softmax_cuda.forward( - inputs, scale_t[0] - ) + softmax_results = scaled_upper_triang_masked_softmax_cuda.forward(inputs, scale_t[0]) ctx.save_for_backward(softmax_results, scale_t) return softmax_results @@ -34,7 +32,7 @@ def backward(ctx, output_grads): try: import scaled_upper_triang_masked_softmax_cuda except (ImportError, ModuleNotFoundError): - print(f'Please install Apex to use fused_softmax') + print(f"Please install Apex to use fused_softmax") softmax_results, scale_t = ctx.saved_tensors input_grads = scaled_upper_triang_masked_softmax_cuda.backward( @@ -57,7 +55,7 @@ def forward(ctx, inputs, mask, scale): try: import scaled_masked_softmax_cuda except (ImportError, ModuleNotFoundError): - print(f'Please install Apex to use fused_softmax') + print(f"Please install Apex to use fused_softmax") scale_t = torch.tensor([scale]) @@ -70,13 +68,11 @@ def backward(ctx, output_grads): try: import scaled_masked_softmax_cuda except (ImportError, ModuleNotFoundError): - print(f'Please install Apex to use fused_softmax') + print(f"Please install Apex to use fused_softmax") softmax_results, scale_t = ctx.saved_tensors - input_grads = scaled_masked_softmax_cuda.backward( - output_grads, softmax_results, scale_t[0] - ) + input_grads = scaled_masked_softmax_cuda.backward(output_grads, softmax_results, scale_t[0]) return input_grads, None, None @@ -92,13 +88,11 @@ def forward(ctx, inputs, scale): try: import scaled_softmax_cuda except (ImportError, ModuleNotFoundError): - print(f'Please install Apex to use fused_softmax') + print(f"Please install Apex to use fused_softmax") scale_t = torch.tensor([scale]) - softmax_results = scaled_softmax_cuda.forward( - inputs, scale_t[0] - ) + softmax_results = scaled_softmax_cuda.forward(inputs, scale_t[0]) ctx.save_for_backward(softmax_results, scale_t) return softmax_results @@ -107,13 +101,11 @@ def backward(ctx, output_grads): try: import scaled_softmax_cudaa except (ImportError, ModuleNotFoundError): - print(f'Please install Apex to use fused_softmax') + print(f"Please install Apex to use fused_softmax") softmax_results, scale_t = ctx.saved_tensors - input_grads = scaled_softmax_cuda.backward( - output_grads, softmax_results, scale_t[0] - ) + input_grads = scaled_softmax_cuda.backward(output_grads, softmax_results, scale_t[0]) return input_grads, None, None @@ -144,9 +136,9 @@ def __init__( super(FusedScaleMaskSoftmax, self).__init__() self.input_in_fp16 = input_in_fp16 self.input_in_bf16 = input_in_bf16 - assert not ( - self.input_in_fp16 and self.input_in_bf16 - ), "both fp16 and bf16 flags cannot be active at the same time." + assert not (self.input_in_fp16 and self.input_in_bf16), ( + "both fp16 and bf16 flags cannot be active at the same time." + ) self.input_in_float16 = self.input_in_fp16 or self.input_in_bf16 self.attn_mask_type = attn_mask_type self.scaled_masked_softmax_fusion = scaled_masked_softmax_fusion @@ -154,9 +146,7 @@ def __init__( self.softmax_in_fp32 = softmax_in_fp32 self.scale = scale - assert ( - self.scale is None or softmax_in_fp32 - ), "softmax should be in fp32 when scaled" + assert self.scale is None or softmax_in_fp32, "softmax should be in fp32 when scaled" def forward(self, input, mask): # [b, np, sq, sk] @@ -229,6 +219,6 @@ def get_batch_per_block(sq, sk, b, np): try: import scaled_masked_softmax_cuda except (ImportError, ModuleNotFoundError): - print(f'Please install Apex to use fused_softmax') + print(f"Please install Apex to use fused_softmax") return scaled_masked_softmax_cuda.get_batch_per_block(sq, sk, b, np) diff --git a/flagscale/train/megatron/legacy/model/gpt_model.py b/flagscale/train/megatron/legacy/model/gpt_model.py index 66fd0979c4..9b25994bfd 100644 --- a/flagscale/train/megatron/legacy/model/gpt_model.py +++ b/flagscale/train/megatron/legacy/model/gpt_model.py @@ -15,44 +15,40 @@ from .module import MegatronModule -def post_language_model_processing(lm_output, labels, logit_weights, - parallel_output, - fp16_lm_cross_entropy): - +def post_language_model_processing( + lm_output, labels, logit_weights, parallel_output, fp16_lm_cross_entropy +): # Output. Format [s b h] - output = parallel_lm_logits( - lm_output, - logit_weights, - parallel_output) + output = parallel_lm_logits(lm_output, logit_weights, parallel_output) if labels is None: # [s b h] => [b s h] - return output.transpose(0,1).contiguous() + return output.transpose(0, 1).contiguous() else: # [b s] => [s b] - labels = labels.transpose(0,1).contiguous() + labels = labels.transpose(0, 1).contiguous() if fp16_lm_cross_entropy: assert output.dtype == torch.half loss = tensor_parallel.vocab_parallel_cross_entropy(output, labels) else: loss = tensor_parallel.vocab_parallel_cross_entropy(output.float(), labels) - + # [s b] => [b, s] - loss = loss.transpose(0,1).contiguous() + loss = loss.transpose(0, 1).contiguous() return loss class GPTModel(MegatronModule): """GPT-2 Language model.""" - def __init__(self, - config, - num_tokentypes=0, - parallel_output=True, - pre_process=True, - post_process=True): + def __init__( + self, config, num_tokentypes=0, parallel_output=True, pre_process=True, post_process=True + ): args = get_args() - super().__init__(config=config, share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights) + super().__init__( + config=config, + share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, + ) self.parallel_output = parallel_output self.pre_process = pre_process @@ -66,8 +62,9 @@ def __init__(self, add_pooler=False, encoder_attn_mask_type=AttnMaskType.causal, pre_process=self.pre_process, - post_process=self.post_process) - + post_process=self.post_process, + ) + if not args.untie_embeddings_and_output_weights: self.initialize_word_embeddings() @@ -75,46 +72,64 @@ def set_input_tensor(self, input_tensor): """See megatron.legacy.model.transformer.set_input_tensor()""" self.language_model.set_input_tensor(input_tensor) - def forward(self, input_ids, position_ids, attention_mask, - labels=None, tokentype_ids=None, inference_context=None, *, inference_params=None): - + def forward( + self, + input_ids, + position_ids, + attention_mask, + labels=None, + tokentype_ids=None, + inference_context=None, + *, + inference_params=None, + ): inference_context = deprecate_inference_params(inference_context, inference_params) lm_output = self.language_model( - input_ids, - position_ids, - attention_mask, - inference_context=inference_context) + input_ids, position_ids, attention_mask, inference_context=inference_context + ) if self.post_process: return post_language_model_processing( - lm_output, labels, - self.language_model.output_layer.weight if self.untie_embeddings_and_output_weights else self.shared_embedding_or_output_weight(), + lm_output, + labels, + self.language_model.output_layer.weight + if self.untie_embeddings_and_output_weights + else self.shared_embedding_or_output_weight(), self.parallel_output, - self.fp16_lm_cross_entropy) + self.fp16_lm_cross_entropy, + ) else: return lm_output - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): - + def state_dict_for_save_checkpoint(self, prefix="", keep_vars=False): state_dict_ = {} - state_dict_[self._language_model_key] \ - = self.language_model.state_dict_for_save_checkpoint( - prefix=prefix, keep_vars=keep_vars) + state_dict_[self._language_model_key] = self.language_model.state_dict_for_save_checkpoint( + prefix=prefix, keep_vars=keep_vars + ) # Save word_embeddings. - if self.post_process and not self.pre_process and not self.untie_embeddings_and_output_weights: - state_dict_[self._word_embeddings_for_head_key] \ - = self.word_embeddings.state_dict(prefix=prefix, - keep_vars=keep_vars) + if ( + self.post_process + and not self.pre_process + and not self.untie_embeddings_and_output_weights + ): + state_dict_[self._word_embeddings_for_head_key] = self.word_embeddings.state_dict( + prefix=prefix, keep_vars=keep_vars + ) return state_dict_ def load_state_dict(self, state_dict, strict=True): """Customized load.""" # Load word_embeddings. - if self.post_process and not self.pre_process and not self.untie_embeddings_and_output_weights: + if ( + self.post_process + and not self.pre_process + and not self.untie_embeddings_and_output_weights + ): self.word_embeddings.load_state_dict( - state_dict[self._word_embeddings_for_head_key], strict=strict) + state_dict[self._word_embeddings_for_head_key], strict=strict + ) if self._language_model_key in state_dict: state_dict = state_dict[self._language_model_key] self.language_model.load_state_dict(state_dict, strict=strict) diff --git a/flagscale/train/megatron/legacy/model/language_model.py b/flagscale/train/megatron/legacy/model/language_model.py index 383230edb7..8fa365fa1a 100644 --- a/flagscale/train/megatron/legacy/model/language_model.py +++ b/flagscale/train/megatron/legacy/model/language_model.py @@ -83,7 +83,7 @@ def get_language_model( post_process=post_process, ) # key used for checkpoints. - language_model_key = 'language_model' + language_model_key = "language_model" return language_model, language_model_key @@ -159,13 +159,13 @@ def __init__( self.word_embeddings = tensor_parallel.VocabParallelEmbedding( vocab_size, self.hidden_size, config=config, init_method=config.init_method ) - self._word_embeddings_key = 'word_embeddings' + self._word_embeddings_key = "word_embeddings" # Position embedding (serial). - self.add_position_embedding = args.position_embedding_type == 'learned_absolute' + self.add_position_embedding = args.position_embedding_type == "learned_absolute" if self.add_position_embedding: self.position_embeddings = torch.nn.Embedding(max_sequence_length, self.hidden_size) - self._position_embeddings_key = 'position_embeddings' + self._position_embeddings_key = "position_embeddings" # Initialize the position embeddings. if args.perform_initialization: self.init_method(self.position_embeddings.weight) @@ -174,7 +174,7 @@ def __init__( # Add this as an optional field that can be added through # method call so we can load a pretrain model without # token types and add them as needed. - self._tokentype_embeddings_key = 'tokentype_embeddings' + self._tokentype_embeddings_key = "tokentype_embeddings" if self.num_tokentypes > 0: self.tokentype_embeddings = torch.nn.Embedding(self.num_tokentypes, self.hidden_size) # Initialize the token-type embeddings. @@ -206,9 +206,9 @@ def add_tokentype_embeddings(self, num_tokentypes): This allows us to load the model normally and then add this embedding. """ if self.tokentype_embeddings is not None: - raise Exception('tokentype embeddings is already initialized') + raise Exception("tokentype embeddings is already initialized") if torch.distributed.get_rank() == 0: - print('adding embedding for {} tokentypes'.format(num_tokentypes), flush=True) + print("adding embedding for {} tokentypes".format(num_tokentypes), flush=True) self.num_tokentypes = num_tokentypes self.tokentype_embeddings = torch.nn.Embedding(num_tokentypes, self.hidden_size) # Initialize the token-type embeddings. @@ -252,7 +252,7 @@ def forward(self, input_ids, position_ids, tokentype_ids=None): return embeddings - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): + def state_dict_for_save_checkpoint(self, prefix="", keep_vars=False): """For easy load.""" state_dict_ = {} @@ -280,8 +280,8 @@ def load_state_dict(self, state_dict, strict=True): # for backward compatibility. state_dict_ = {} for key in state_dict.keys(): - if 'word_embeddings' in key: - state_dict_[key.split('word_embeddings.')[1]] = state_dict[key] + if "word_embeddings" in key: + state_dict_[key.split("word_embeddings.")[1]] = state_dict[key] self.word_embeddings.load_state_dict(state_dict_, strict=strict) # Position embedding. @@ -292,8 +292,8 @@ def load_state_dict(self, state_dict, strict=True): # for backward compatibility. state_dict_ = {} for key in state_dict.keys(): - if 'position_embeddings' in key: - state_dict_[key.split('position_embeddings.')[1]] = state_dict[key] + if "position_embeddings" in key: + state_dict_[key.split("position_embeddings.")[1]] = state_dict[key] self.position_embeddings.load_state_dict(state_dict_, strict=strict) # Tokentype embedding. @@ -304,14 +304,14 @@ def load_state_dict(self, state_dict, strict=True): else: # for backward compatibility. for key in state_dict.keys(): - if 'tokentype_embeddings' in key: - state_dict_[key.split('tokentype_embeddings.')[1]] = state_dict[key] + if "tokentype_embeddings" in key: + state_dict_[key.split("tokentype_embeddings.")[1]] = state_dict[key] if len(state_dict_.keys()) > 0: self.tokentype_embeddings.load_state_dict(state_dict_, strict=strict) else: print( - '***WARNING*** expected tokentype embeddings in the ' - 'checkpoint but could not find it', + "***WARNING*** expected tokentype embeddings in the " + "checkpoint but could not find it", flush=True, ) @@ -372,10 +372,10 @@ def __init__( config, self.num_tokentypes, ) - self._embedding_key = 'embedding' + self._embedding_key = "embedding" # Rotary positional embeddings - self.use_rotary_position_embeddings = args.position_embedding_type == 'rope' + self.use_rotary_position_embeddings = args.position_embedding_type == "rope" if self.use_rotary_position_embeddings: self.seq_length = args.seq_length rotary_dim = ( @@ -403,7 +403,7 @@ def __init__( pre_process=self.pre_process, post_process=self.post_process, ) - self._encoder_key = 'encoder' + self._encoder_key = "encoder" else: self.encoder = None @@ -418,7 +418,7 @@ def __init__( pre_process=self.pre_process, post_process=self.post_process, ) - self._decoder_key = 'decoder' + self._decoder_key = "decoder" else: self.decoder = None @@ -426,7 +426,7 @@ def __init__( # Pooler. if self.add_pooler: self.pooler = Pooler(self.hidden_size, self.init_method) - self._pooler_key = 'pooler' + self._pooler_key = "pooler" if self.untie_embeddings_and_output_weights: self.output_layer = tensor_parallel.ColumnParallelLinear( @@ -436,7 +436,7 @@ def __init__( init_method=self.init_method, bias=False, ) # Setting bias to False always to keep it consistent with embedding tying that also does not have a bias. - self._output_layer_key = 'output_layer' + self._output_layer_key = "output_layer" def set_input_tensor(self, input_tensor): """See megatron.legacy.model.transformer.set_input_tensor()""" @@ -447,14 +447,14 @@ def set_input_tensor(self, input_tensor): input_tensor = [input_tensor] if self.add_encoder and self.add_decoder: - assert ( - len(input_tensor) == 1 - ), 'input_tensor should only be length 1 for stage with both encoder and decoder' + assert len(input_tensor) == 1, ( + "input_tensor should only be length 1 for stage with both encoder and decoder" + ) self.encoder.set_input_tensor(input_tensor[0]) elif self.add_encoder: - assert ( - len(input_tensor) == 1 - ), 'input_tensor should only be length 1 for stage with only encoder' + assert len(input_tensor) == 1, ( + "input_tensor should only be length 1 for stage with only encoder" + ) self.encoder.set_input_tensor(input_tensor[0]) elif self.add_decoder: if len(input_tensor) == 2: @@ -464,9 +464,9 @@ def set_input_tensor(self, input_tensor): self.decoder.set_input_tensor(None) self.encoder_hidden_state = input_tensor[0] else: - raise Exception('input_tensor must have either length 1 or 2') + raise Exception("input_tensor must have either length 1 or 2") else: - raise Exception('Stage must have at least either encoder or decoder') + raise Exception("Stage must have at least either encoder or decoder") def forward( self, @@ -485,7 +485,6 @@ def forward( *, inference_params: Optional[BaseInferenceContext] = None, ): - inference_context = deprecate_inference_params(inference_context, inference_params) # Encoder embedding. @@ -552,7 +551,7 @@ def forward( else: return decoder_output, encoder_output - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): + def state_dict_for_save_checkpoint(self, prefix="", keep_vars=False): """For easy load.""" state_dict_ = {} @@ -592,7 +591,7 @@ def load_state_dict(self, state_dict, strict=True): # for backward compatibility. state_dict_ = {} for key in state_dict.keys(): - if '_embeddings' in key: + if "_embeddings" in key: state_dict_[key] = state_dict[key] self.embedding.load_state_dict(state_dict_, strict=strict) @@ -601,19 +600,19 @@ def load_state_dict(self, state_dict, strict=True): if self._encoder_key in state_dict: state_dict_ = state_dict[self._encoder_key] # For backward compatibility. - elif 'transformer' in state_dict: - state_dict_ = state_dict['transformer'] + elif "transformer" in state_dict: + state_dict_ = state_dict["transformer"] else: # For backward compatibility. state_dict_ = {} for key in state_dict.keys(): - if 'transformer.' in key: - state_dict_[key.split('transformer.')[1]] = state_dict[key] + if "transformer." in key: + state_dict_[key.split("transformer.")[1]] = state_dict[key] # For backward compatibility. state_dict_self_attention = {} for key in state_dict_.keys(): - if '.attention.' in key: + if ".attention." in key: state_dict_self_attention[key.replace(".attention.", ".self_attention.")] = ( state_dict_[key] ) @@ -626,14 +625,14 @@ def load_state_dict(self, state_dict, strict=True): # Pooler. if self.post_process: if self.add_pooler: - assert 'pooler' in state_dict, 'could not find data for pooler in the checkpoint' + assert "pooler" in state_dict, "could not find data for pooler in the checkpoint" self.pooler.load_state_dict(state_dict[self._pooler_key], strict=strict) if self.untie_embeddings_and_output_weights: - assert ( - 'output_layer' in state_dict - ), 'could not find data for output_layer in the checkpoint' + assert "output_layer" in state_dict, ( + "could not find data for output_layer in the checkpoint" + ) self.output_layer.load_state_dict(state_dict[self._output_layer_key], strict=strict) # Decoder. if self.add_decoder: - assert 'decoder' in state_dict, 'could not find data for pooler in the checkpoint' + assert "decoder" in state_dict, "could not find data for pooler in the checkpoint" self.decoder.load_state_dict(state_dict[self._decoder_key], strict=strict) diff --git a/flagscale/train/megatron/legacy/model/module.py b/flagscale/train/megatron/legacy/model/module.py index 5e37b73dc5..22a80fef1a 100644 --- a/flagscale/train/megatron/legacy/model/module.py +++ b/flagscale/train/megatron/legacy/model/module.py @@ -11,6 +11,7 @@ from megatron.plugin.platform import get_platform + cur_platform = get_platform() _FLOAT_TYPES = (torch.FloatTensor, cur_platform.FloatTensor) @@ -27,7 +28,7 @@ def __init__(self, config=None, share_embeddings_and_output_weights=True): self.config = config self.share_embeddings_and_output_weights = share_embeddings_and_output_weights - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): + def state_dict_for_save_checkpoint(self, prefix="", keep_vars=False): """Use this function to override the state dict for saving checkpoints.""" return self.state_dict(prefix=prefix, keep_vars=keep_vars) @@ -38,8 +39,8 @@ def shared_embedding_or_output_weight(self): else: if not self.share_embeddings_and_output_weights: raise Exception( - 'shared_embedding_or_output_weight() called for last ' - 'stage, but share_embeddings_and_output_weights is false' + "shared_embedding_or_output_weight() called for last " + "stage, but share_embeddings_and_output_weights is false" ) return self.word_embeddings.weight @@ -47,8 +48,8 @@ def initialize_word_embeddings(self): args = get_args() if not self.share_embeddings_and_output_weights: raise Exception( - 'initialize_word_embeddings() was called but ' - 'share_embeddings_and_output_weights is false' + "initialize_word_embeddings() was called but " + "share_embeddings_and_output_weights is false" ) # This function just initializes the word embeddings in the final stage @@ -82,7 +83,7 @@ def initialize_word_embeddings(self): # update is the same on both stages. if mpu.is_pipeline_last_stage(ignore_virtual=False) and not self.pre_process: assert not mpu.is_pipeline_first_stage(ignore_virtual=False) - self._word_embeddings_for_head_key = 'word_embeddings_for_head' + self._word_embeddings_for_head_key = "word_embeddings_for_head" # set word_embeddings weights to 0 here, then copy first # stage's weights using all_reduce below. self.word_embeddings = tensor_parallel.VocabParallelEmbedding( diff --git a/flagscale/train/megatron/legacy/model/multiple_choice.py b/flagscale/train/megatron/legacy/model/multiple_choice.py index bec0548c40..2c0a74f756 100644 --- a/flagscale/train/megatron/legacy/model/multiple_choice.py +++ b/flagscale/train/megatron/legacy/model/multiple_choice.py @@ -15,12 +15,7 @@ class MultipleChoice(MegatronModule): - - def __init__(self, - config, - num_tokentypes=2, - pre_process=True, - post_process=True): + def __init__(self, config, num_tokentypes=2, pre_process=True, post_process=True): super(MultipleChoice, self).__init__(share_embeddings_and_output_weights=False) args = get_args() @@ -33,21 +28,20 @@ def __init__(self, add_pooler=True, encoder_attn_mask_type=AttnMaskType.padding, pre_process=self.pre_process, - post_process=self.post_process) + post_process=self.post_process, + ) # Multi-choice head. if self.post_process: self.multichoice_dropout = torch.nn.Dropout(args.hidden_dropout) - self.multichoice_head = get_linear_layer(args.hidden_size, 1, - init_method) - self._multichoice_head_key = 'multichoice_head' + self.multichoice_head = get_linear_layer(args.hidden_size, 1, init_method) + self._multichoice_head_key = "multichoice_head" def set_input_tensor(self, input_tensor): """See megatron.legacy.model.transformer.set_input_tensor()""" self.language_model.set_input_tensor(input_tensor) def forward(self, model_input, attention_mask, tokentype_ids=None): - # [batch, choices, sequence] --> [batch * choices, sequence] --> # transformer --> [batch, choices] --> softmax @@ -68,10 +62,7 @@ def forward(self, model_input, attention_mask, tokentype_ids=None): position_ids = bert_position_ids(input_ids) lm_output = self.language_model( - input_ids, - position_ids, - extended_attention_mask, - tokentype_ids=tokentype_ids + input_ids, position_ids, extended_attention_mask, tokentype_ids=tokentype_ids ) if self.post_process: _, pooled_output = lm_output @@ -84,29 +75,31 @@ def forward(self, model_input, attention_mask, tokentype_ids=None): return multichoice_logits return lm_output - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): + def state_dict_for_save_checkpoint(self, prefix="", keep_vars=False): """For easy load when model is combined with other heads, add an extra key.""" state_dict_ = {} - state_dict_[self._language_model_key] \ - = self.language_model.state_dict_for_save_checkpoint(prefix=prefix, - keep_vars=keep_vars) + state_dict_[self._language_model_key] = self.language_model.state_dict_for_save_checkpoint( + prefix=prefix, keep_vars=keep_vars + ) if self.post_process: - state_dict_[self._multichoice_head_key] \ - = self.multichoice_head.state_dict(prefix=prefix, keep_vars=keep_vars) + state_dict_[self._multichoice_head_key] = self.multichoice_head.state_dict( + prefix=prefix, keep_vars=keep_vars + ) return state_dict_ def load_state_dict(self, state_dict, strict=True): """Customized load.""" - self.language_model.load_state_dict( - state_dict[self._language_model_key], strict=strict) + self.language_model.load_state_dict(state_dict[self._language_model_key], strict=strict) if self.post_process: if self._multichoice_head_key in state_dict: self.multichoice_head.load_state_dict( - state_dict[self._multichoice_head_key], strict=strict) + state_dict[self._multichoice_head_key], strict=strict + ) else: - print_rank_last('***WARNING*** could not find {} in the checkpoint, ' - 'initializing to random'.format( - self._multichoice_head_key)) + print_rank_last( + "***WARNING*** could not find {} in the checkpoint, " + "initializing to random".format(self._multichoice_head_key) + ) diff --git a/flagscale/train/megatron/legacy/model/realm_model.py b/flagscale/train/megatron/legacy/model/realm_model.py index 838da96126..691bedab07 100644 --- a/flagscale/train/megatron/legacy/model/realm_model.py +++ b/flagscale/train/megatron/legacy/model/realm_model.py @@ -15,17 +15,22 @@ from megatron.legacy.model.bert_model import bert_extended_attention_mask, bert_position_ids from megatron.plugin.platform import get_platform + cur_platform = get_platform() + def general_ict_model_provider(only_query_model=False, only_block_model=False): """Build the model.""" args = get_args() - assert args.ict_head_size is not None, \ + assert args.ict_head_size is not None, ( "Need to specify --ict-head-size to provide an ICTBertModel" - assert mpu.get_tensor_model_parallel_world_size() == 1 and mpu.get_pipeline_model_parallel_world_size() == 1, \ - "Model parallel size > 1 not supported for ICT" + ) + assert ( + mpu.get_tensor_model_parallel_world_size() == 1 + and mpu.get_pipeline_model_parallel_world_size() == 1 + ), "Model parallel size > 1 not supported for ICT" - print_rank_0('building ICTBertModel...') + print_rank_0("building ICTBertModel...") # simpler to just keep using 2 tokentypes since the LM we initialize with has 2 tokentypes model = ICTBertModel( @@ -33,24 +38,28 @@ def general_ict_model_provider(only_query_model=False, only_block_model=False): num_tokentypes=2, parallel_output=True, only_query_model=only_query_model, - only_block_model=only_block_model) + only_block_model=only_block_model, + ) return model class ICTBertModel(MegatronModule): """Bert-based module for Inverse Cloze task.""" - def __init__(self, - ict_head_size, - num_tokentypes=1, - parallel_output=True, - only_query_model=False, - only_block_model=False): + + def __init__( + self, + ict_head_size, + num_tokentypes=1, + parallel_output=True, + only_query_model=False, + only_block_model=False, + ): super(ICTBertModel, self).__init__() bert_kwargs = dict( ict_head_size=ict_head_size, num_tokentypes=num_tokentypes, - parallel_output=parallel_output + parallel_output=parallel_output, ) assert not (only_block_model and only_query_model) self.use_block_model = not only_query_model @@ -59,12 +68,12 @@ def __init__(self, if self.use_query_model: # this model embeds (pseudo-)queries - Embed_input in the paper self.query_model = IREncoderBertModel(**bert_kwargs) - self._query_key = 'question_model' + self._query_key = "question_model" if self.use_block_model: # this model embeds evidence blocks - Embed_doc in the paper self.block_model = IREncoderBertModel(**bert_kwargs) - self._block_key = 'context_model' + self._block_key = "context_model" def forward(self, query_tokens, query_attention_mask, block_tokens, block_attention_mask): """Run a forward pass for each of the models and return the respective embeddings.""" @@ -76,7 +85,9 @@ def embed_query(self, query_tokens, query_attention_mask): """Embed a batch of tokens using the query model""" if self.use_query_model: query_types = cur_platform.LongTensor(*query_tokens.shape).fill_(0) - query_ict_logits, _ = self.query_model.forward(query_tokens, query_attention_mask, query_types) + query_ict_logits, _ = self.query_model.forward( + query_tokens, query_attention_mask, query_types + ) return query_ict_logits else: raise ValueError("Cannot embed query without query model.") @@ -85,23 +96,25 @@ def embed_block(self, block_tokens, block_attention_mask): """Embed a batch of tokens using the block model""" if self.use_block_model: block_types = cur_platform.LongTensor(*block_tokens.shape).fill_(0) - block_ict_logits, _ = self.block_model.forward(block_tokens, block_attention_mask, block_types) + block_ict_logits, _ = self.block_model.forward( + block_tokens, block_attention_mask, block_types + ) return block_ict_logits else: raise ValueError("Cannot embed block without block model.") - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): + def state_dict_for_save_checkpoint(self, prefix="", keep_vars=False): """Save dict with state dicts of each of the models.""" state_dict_ = {} if self.use_query_model: - state_dict_[self._query_key] \ - = self.query_model.state_dict_for_save_checkpoint( - prefix=prefix, keep_vars=keep_vars) + state_dict_[self._query_key] = self.query_model.state_dict_for_save_checkpoint( + prefix=prefix, keep_vars=keep_vars + ) if self.use_block_model: - state_dict_[self._block_key] \ - = self.block_model.state_dict_for_save_checkpoint( - prefix=prefix, keep_vars=keep_vars) + state_dict_[self._block_key] = self.block_model.state_dict_for_save_checkpoint( + prefix=prefix, keep_vars=keep_vars + ) return state_dict_ @@ -109,13 +122,11 @@ def load_state_dict(self, state_dict, strict=True): """Load the state dicts of each of the models""" if self.use_query_model: print("Loading ICT query model", flush=True) - self.query_model.load_state_dict( - state_dict[self._query_key], strict=strict) + self.query_model.load_state_dict(state_dict[self._query_key], strict=strict) if self.use_block_model: print("Loading ICT block model", flush=True) - self.block_model.load_state_dict( - state_dict[self._block_key], strict=strict) + self.block_model.load_state_dict(state_dict[self._block_key], strict=strict) def init_state_dict_from_bert(self): """Initialize the state from a pretrained BERT model on iteration zero of ICT pretraining""" @@ -123,32 +134,38 @@ def init_state_dict_from_bert(self): tracker_filename = get_checkpoint_tracker_filename(args.bert_load) if not os.path.isfile(tracker_filename): raise FileNotFoundError("Could not find BERT load for ICT") - with open(tracker_filename, 'r') as f: + with open(tracker_filename, "r") as f: iteration = int(f.read().strip()) assert iteration > 0 checkpoint_name = get_checkpoint_name(args.bert_load, iteration, False) if mpu.get_data_parallel_rank() == 0: - print('global rank {} is loading checkpoint {}'.format( - torch.distributed.get_rank(), checkpoint_name)) + print( + "global rank {} is loading checkpoint {}".format( + torch.distributed.get_rank(), checkpoint_name + ) + ) try: - state_dict = torch.load(checkpoint_name, map_location='cpu') + state_dict = torch.load(checkpoint_name, map_location="cpu") except Exception: raise ValueError("Could not load checkpoint") # load the LM state dict into each model - model_dict = state_dict['model']['language_model'] + model_dict = state_dict["model"]["language_model"] self.query_model.language_model.load_state_dict(model_dict) self.block_model.language_model.load_state_dict(model_dict) # give each model the same ict_head to begin with as well - query_ict_head_state_dict = self.state_dict_for_save_checkpoint()[self._query_key]['ict_head'] + query_ict_head_state_dict = self.state_dict_for_save_checkpoint()[self._query_key][ + "ict_head" + ] self.block_model.ict_head.load_state_dict(query_ict_head_state_dict) class IREncoderBertModel(MegatronModule): """BERT-based encoder for queries or blocks used for learned information retrieval.""" + def __init__(self, ict_head_size, num_tokentypes=2, parallel_output=True): super(IREncoderBertModel, self).__init__() args = get_args() @@ -156,52 +173,47 @@ def __init__(self, ict_head_size, num_tokentypes=2, parallel_output=True): self.ict_head_size = ict_head_size self.parallel_output = parallel_output init_method = init_method_normal(args.init_method_std) - scaled_init_method = scaled_init_method_normal(args.init_method_std, - args.num_layers) + scaled_init_method = scaled_init_method_normal(args.init_method_std, args.num_layers) self.language_model, self._language_model_key = get_language_model( num_tokentypes=num_tokentypes, add_pooler=True, encoder_attn_mask_type=AttnMaskType.padding, init_method=init_method, - scaled_init_method=scaled_init_method) + scaled_init_method=scaled_init_method, + ) self.ict_head = get_linear_layer(args.hidden_size, ict_head_size, init_method) - self._ict_head_key = 'ict_head' + self._ict_head_key = "ict_head" def forward(self, input_ids, attention_mask, tokentype_ids=None): extended_attention_mask = bert_extended_attention_mask( - attention_mask, next(self.language_model.parameters()).dtype) + attention_mask, next(self.language_model.parameters()).dtype + ) position_ids = bert_position_ids(input_ids) lm_output, pooled_output = self.language_model( - input_ids, - position_ids, - extended_attention_mask, - tokentype_ids=tokentype_ids) + input_ids, position_ids, extended_attention_mask, tokentype_ids=tokentype_ids + ) # Output. ict_logits = self.ict_head(pooled_output) return ict_logits, None - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): + def state_dict_for_save_checkpoint(self, prefix="", keep_vars=False): """For easy load when model is combined with other heads, add an extra key.""" state_dict_ = {} - state_dict_[self._language_model_key] \ - = self.language_model.state_dict_for_save_checkpoint(prefix=prefix, - keep_vars=keep_vars) - state_dict_[self._ict_head_key] \ - = self.ict_head.state_dict(prefix=prefix, - keep_vars=keep_vars) + state_dict_[self._language_model_key] = self.language_model.state_dict_for_save_checkpoint( + prefix=prefix, keep_vars=keep_vars + ) + state_dict_[self._ict_head_key] = self.ict_head.state_dict( + prefix=prefix, keep_vars=keep_vars + ) return state_dict_ def load_state_dict(self, state_dict, strict=True): """Customized load.""" - self.language_model.load_state_dict( - state_dict[self._language_model_key], strict=strict) - self.ict_head.load_state_dict( - state_dict[self._ict_head_key], strict=strict) - - + self.language_model.load_state_dict(state_dict[self._language_model_key], strict=strict) + self.ict_head.load_state_dict(state_dict[self._ict_head_key], strict=strict) diff --git a/flagscale/train/megatron/legacy/model/rms_norm.py b/flagscale/train/megatron/legacy/model/rms_norm.py index 21ba00c600..6b33cdfa60 100644 --- a/flagscale/train/megatron/legacy/model/rms_norm.py +++ b/flagscale/train/megatron/legacy/model/rms_norm.py @@ -3,13 +3,11 @@ import torch from torch import nn -class RMSNorm(torch.nn.Module): - def __init__(self, - dim: int, - eps: float = 1e-6, - sequence_parallel: bool = False, - config: dict = None): +class RMSNorm(torch.nn.Module): + def __init__( + self, dim: int, eps: float = 1e-6, sequence_parallel: bool = False, config: dict = None + ): """RMS Normaliation module Args: @@ -22,7 +20,7 @@ def __init__(self, self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) - setattr(self.weight, 'sequence_parallel', sequence_parallel) + setattr(self.weight, "sequence_parallel", sequence_parallel) def _norm(self, x): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) diff --git a/flagscale/train/megatron/legacy/model/t5_model.py b/flagscale/train/megatron/legacy/model/t5_model.py index 1662188334..2ae1d3fc97 100644 --- a/flagscale/train/megatron/legacy/model/t5_model.py +++ b/flagscale/train/megatron/legacy/model/t5_model.py @@ -9,15 +9,11 @@ from megatron.legacy.model.enums import AttnMaskType from megatron.legacy.model.language_model import parallel_lm_logits, get_language_model from megatron.legacy.model import LayerNorm -from megatron.legacy.model.utils import ( - openai_gelu, - get_linear_layer -) +from megatron.legacy.model.utils import openai_gelu, get_linear_layer from .module import MegatronModule def t5_extended_attention_mask(attention_mask_list): - def attn_mask_postprocess(attn_mask): # [b, 1, s, s] extended_attention_mask = attn_mask.unsqueeze(1) @@ -29,8 +25,7 @@ def attn_mask_postprocess(attn_mask): def t5_position_ids(token_ids): # Create position ids seq_length = token_ids.size(1) - position_ids = torch.arange(seq_length, dtype=torch.long, - device=token_ids.device) + position_ids = torch.arange(seq_length, dtype=torch.long, device=token_ids.device) position_ids = position_ids.unsqueeze(0).expand_as(token_ids) return position_ids @@ -54,24 +49,25 @@ def __init__(self, mpu_vocab_size, parallel_output): self.parallel_output = parallel_output def forward(self, hidden_states, word_embeddings_weight): - output = parallel_lm_logits(hidden_states, - word_embeddings_weight, - self.parallel_output, - bias=self.bias) + output = parallel_lm_logits( + hidden_states, word_embeddings_weight, self.parallel_output, bias=self.bias + ) return output class T5Model(MegatronModule): """T5 Language model.""" - def __init__(self, - config, - num_tokentypes=0, - parallel_output=True, - pre_process=True, - post_process=True, - add_encoder=True, - add_decoder=True): + def __init__( + self, + config, + num_tokentypes=0, + parallel_output=True, + pre_process=True, + post_process=True, + add_encoder=True, + add_decoder=True, + ): super().__init__(config=config) args = get_args() @@ -90,7 +86,8 @@ def __init__(self, add_decoder=add_decoder, encoder_attn_mask_type=AttnMaskType.padding, pre_process=self.pre_process, - post_process=self.post_process) + post_process=self.post_process, + ) self.initialize_word_embeddings() @@ -101,9 +98,9 @@ def __init__(self, if self.post_process and self.add_decoder: self.lm_head = T5LMHead( - self.shared_embedding_or_output_weight().size(0), - parallel_output) - self._lm_head_key = 'lm_head' + self.shared_embedding_or_output_weight().size(0), parallel_output + ) + self._lm_head_key = "lm_head" # Tells schedules.py that this model has a skip connection between the encoder's output and the decoder # (and hence both the encoder and decoder's tensors are required for correct backprop). @@ -113,47 +110,59 @@ def set_input_tensor(self, input_tensor): """See megatron.legacy.model.transformer.set_input_tensor()""" self.language_model.set_input_tensor(input_tensor) - def forward(self, encoder_input_ids, decoder_input_ids, encoder_attn_mask, - decoder_attn_mask, encoder_decoder_attn_mask, - tokentype_ids=None, lm_labels=None, enc_hidden_states=None): - + def forward( + self, + encoder_input_ids, + decoder_input_ids, + encoder_attn_mask, + decoder_attn_mask, + encoder_decoder_attn_mask, + tokentype_ids=None, + lm_labels=None, + enc_hidden_states=None, + ): # Converting the attention masks to proper parameter settings - encoder_attn_mask, decoder_attn_mask, encoder_decoder_attn_mask = t5_extended_attention_mask( - [encoder_attn_mask, decoder_attn_mask, encoder_decoder_attn_mask]) + encoder_attn_mask, decoder_attn_mask, encoder_decoder_attn_mask = ( + t5_extended_attention_mask( + [encoder_attn_mask, decoder_attn_mask, encoder_decoder_attn_mask] + ) + ) encoder_position_ids = t5_position_ids(encoder_input_ids) decoder_position_ids = t5_position_ids(decoder_input_ids) - lm_output = self.language_model(encoder_input_ids, - encoder_position_ids, - encoder_attn_mask, - decoder_input_ids, - decoder_position_ids, - decoder_attn_mask, - encoder_decoder_attn_mask, - tokentype_ids=tokentype_ids, - enc_hidden_states=enc_hidden_states) + lm_output = self.language_model( + encoder_input_ids, + encoder_position_ids, + encoder_attn_mask, + decoder_input_ids, + decoder_position_ids, + decoder_attn_mask, + encoder_decoder_attn_mask, + tokentype_ids=tokentype_ids, + enc_hidden_states=enc_hidden_states, + ) if self.post_process and self.add_decoder: decoder_output, encoder_output = lm_output # Output. [s, b, h] - lm_logits = self.lm_head(decoder_output, - self.shared_embedding_or_output_weight()) + lm_logits = self.lm_head(decoder_output, self.shared_embedding_or_output_weight()) if lm_labels is None: # [s b h] => [b s h] - return lm_logits.transpose(0,1).contiguous() + return lm_logits.transpose(0, 1).contiguous() else: # [b s] => [s b] - lm_labels = lm_labels.transpose(0,1).contiguous() + lm_labels = lm_labels.transpose(0, 1).contiguous() if self.fp16_lm_cross_entropy: assert lm_logits.dtype == torch.half lm_loss = tensor_parallel.vocab_parallel_cross_entropy(lm_logits, lm_labels) else: - lm_loss = tensor_parallel.vocab_parallel_cross_entropy(lm_logits.float(), - lm_labels) + lm_loss = tensor_parallel.vocab_parallel_cross_entropy( + lm_logits.float(), lm_labels + ) # [s b] => [b s] - lm_loss = lm_loss.transpose(0,1).contiguous() + lm_loss = lm_loss.transpose(0, 1).contiguous() return lm_loss elif self.add_decoder and not self.add_encoder: decoder_output, encoder_output = lm_output @@ -162,34 +171,33 @@ def forward(self, encoder_input_ids, decoder_input_ids, encoder_attn_mask, encoder_output = lm_output return encoder_output - def state_dict_for_save_checkpoint(self, prefix='', keep_vars=False): + def state_dict_for_save_checkpoint(self, prefix="", keep_vars=False): """For easy load when model is combined with other heads, add an extra key.""" state_dict_ = {} - state_dict_[self._language_model_key] \ - = self.language_model.state_dict_for_save_checkpoint(prefix=prefix, - keep_vars=keep_vars) + state_dict_[self._language_model_key] = self.language_model.state_dict_for_save_checkpoint( + prefix=prefix, keep_vars=keep_vars + ) if self.post_process and self.add_decoder: - state_dict_[self._lm_head_key] \ - = self.lm_head.state_dict_for_save_checkpoint(prefix=prefix, - keep_vars=keep_vars) - # Save word_embeddings. + state_dict_[self._lm_head_key] = self.lm_head.state_dict_for_save_checkpoint( + prefix=prefix, keep_vars=keep_vars + ) + # Save word_embeddings. if self.post_process and not self.pre_process and self.add_decoder: - state_dict_[self._word_embeddings_for_head_key] \ - = self.word_embeddings.state_dict(prefix=prefix, - keep_vars=keep_vars) + state_dict_[self._word_embeddings_for_head_key] = self.word_embeddings.state_dict( + prefix=prefix, keep_vars=keep_vars + ) return state_dict_ def load_state_dict(self, state_dict, strict=True): """Customized load.""" - self.language_model.load_state_dict( - state_dict[self._language_model_key], strict=strict) + self.language_model.load_state_dict(state_dict[self._language_model_key], strict=strict) if self.post_process and self.add_decoder: - self.lm_head.load_state_dict(state_dict[self._lm_head_key], - strict=strict) + self.lm_head.load_state_dict(state_dict[self._lm_head_key], strict=strict) # Load word embeddings. if self.post_process and not self.pre_process and self.add_decoder: self.word_embeddings.load_state_dict( - state_dict[self._word_embeddings_for_head_key], strict=strict) + state_dict[self._word_embeddings_for_head_key], strict=strict + ) diff --git a/flagscale/train/megatron/legacy/model/transformer.py b/flagscale/train/megatron/legacy/model/transformer.py index fbe29b4014..d8e71935a3 100644 --- a/flagscale/train/megatron/legacy/model/transformer.py +++ b/flagscale/train/megatron/legacy/model/transformer.py @@ -1,6 +1,7 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. """Transformer.""" + import math import os from contextlib import nullcontext @@ -74,30 +75,34 @@ """ from megatron.plugin.platform import get_platform + cur_platform = get_platform() + class DropPath(MegatronModule): """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). """ - def __init__(self, drop_prob=0.): + def __init__(self, drop_prob=0.0): super(DropPath, self).__init__() self.drop_prob = drop_prob def forward(self, hidden_state): - if self.drop_prob == 0. or not self.training: + if self.drop_prob == 0.0 or not self.training: return hidden_state keep_prob = 1 - self.drop_prob # work with diff dim tensors, not just 2D ConvNets # hidden_state: [s, b, h] shape = (1,) + (hidden_state.shape[1],) + (1,) * (hidden_state.ndim - 2) - random_tensor = keep_prob + \ - torch.rand(shape, dtype=hidden_state.dtype, device=hidden_state.device) + random_tensor = keep_prob + torch.rand( + shape, dtype=hidden_state.dtype, device=hidden_state.device + ) random_tensor.floor_() # binarize output = hidden_state.div(keep_prob) * random_tensor return output + class ParallelMLP(MegatronModule): """MLP. @@ -137,13 +142,17 @@ def __init__(self, config, is_expert=False): elif args.onnx_safe: self.activation_func = erf_gelu elif args.swiglu: + def swiglu(x): x = torch.chunk(x, 2, dim=-1) return F.silu(x[0]) * x[1] + self.activation_func = swiglu elif args.squared_relu: + def squared_relu(x): return torch.pow(F.relu(x), 2) + self.activation_func = squared_relu else: self.bias_gelu_fusion = args.bias_gelu_fusion @@ -162,7 +171,6 @@ def squared_relu(x): ) def forward(self, hidden_states): - # [s, b, 4hp] intermediate_parallel, bias_parallel = self.dense_h_to_4h(hidden_states) @@ -179,6 +187,7 @@ def forward(self, hidden_states): output, output_bias = self.dense_4h_to_h(intermediate_parallel) return output, output_bias + def sinkhorn(cost, tol=0.0001): cost = torch.exp(cost) d0 = torch.ones(cost.size(0), device=cost.device, dtype=cost.dtype) @@ -188,11 +197,11 @@ def sinkhorn(cost, tol=0.0001): error = 1e9 d1_old = d1 while error > tol: - d0 = (1/d0.size(0))*1/(torch.sum(d1*cost,1) + eps) - d1 = (1/d1.size(0))*1/(torch.sum(d0.unsqueeze(1)*cost,0)+eps) - error = torch.mean(torch.abs(d1_old-d1)) + d0 = (1 / d0.size(0)) * 1 / (torch.sum(d1 * cost, 1) + eps) + d1 = (1 / d1.size(0)) * 1 / (torch.sum(d0.unsqueeze(1) * cost, 0) + eps) + error = torch.mean(torch.abs(d1_old - d1)) d1_old = d1 - return d1*cost*d0.unsqueeze(1) + return d1 * cost * d0.unsqueeze(1) def get_router_linear_layer(config): @@ -200,7 +209,7 @@ def get_router_linear_layer(config): router = torch.nn.Linear(args.hidden_size, args.num_experts, bias=False) with get_cuda_rng_tracker().fork(get_data_parallel_rng_tracker_name()): config.init_method(router.weight) - setattr(router.weight, 'sequence_parallel',config.sequence_parallel) + setattr(router.weight, "sequence_parallel", config.sequence_parallel) return router @@ -208,6 +217,7 @@ class SwitchMLP(MegatronModule): """ Routes input to one of N MLP "experts" """ + def __init__(self, config): super(SwitchMLP, self).__init__() args = get_args() @@ -219,7 +229,9 @@ def __init__(self, config): assert args.num_experts % self.expert_parallel_size == 0 self.num_local_experts = args.num_experts // self.expert_parallel_size local_expert_indices_offset = mpu.get_expert_model_parallel_rank() * self.num_local_experts - self.local_expert_indices = [local_expert_indices_offset + i for i in range(self.num_local_experts)] + self.local_expert_indices = [ + local_expert_indices_offset + i for i in range(self.num_local_experts) + ] self.local_experts = torch.nn.ModuleList() for i in range(self.num_local_experts): @@ -228,7 +240,7 @@ def __init__(self, config): self.tp_ep_group = get_expert_tensor_and_model_parallel_group() def gather_indices(self, local_indices): - """ Gather tensors and concatinate along the first dimension.""" + """Gather tensors and concatinate along the first dimension.""" world_size = self.tp_ep_group.size() # Bypass the function if we are using only 1 GPU. if world_size == 1: @@ -238,8 +250,9 @@ def gather_indices(self, local_indices): dim_size[0] = dim_size[0] * world_size # TODO pre allocate memory - output = torch.empty(dim_size, dtype=local_indices.dtype, - device=cur_platform.current_device()) + output = torch.empty( + dim_size, dtype=local_indices.dtype, device=cur_platform.current_device() + ) torch.distributed._all_gather_base( output, local_indices.contiguous(), group=self.tp_ep_group ) @@ -273,8 +286,9 @@ def forward(self, hidden_states): # Converting [s, b, h] to [s*b, h]. # Each vector could be routed differently if self.sequence_parallel or (self.expert_parallel_size > 1): - global_hidden_states = \ - gather_from_sequence_parallel_region(hidden_states, group=self.tp_ep_group) + global_hidden_states = gather_from_sequence_parallel_region( + hidden_states, group=self.tp_ep_group + ) global_indices = self.gather_indices(max_ind) else: global_hidden_states = hidden_states @@ -295,21 +309,22 @@ def forward(self, hidden_states): output_bias_total[local_indices, :] = output_bias if self.sequence_parallel or (self.expert_parallel_size > 1): - output_total = \ - reduce_scatter_to_sequence_parallel_region(output_total, group=self.tp_ep_group) + output_total = reduce_scatter_to_sequence_parallel_region( + output_total, group=self.tp_ep_group + ) if self.add_bias: - output_bias_total = \ - reduce_scatter_to_sequence_parallel_region(output_bias_total, group=self.tp_ep_group) + output_bias_total = reduce_scatter_to_sequence_parallel_region( + output_bias_total, group=self.tp_ep_group + ) # bias is duplicated across tensor parallelism ranks; # reduce scatter reduces bias across tensor parallel_ranks - output_bias_total = \ - output_bias_total/mpu.get_tensor_model_parallel_world_size() + output_bias_total = output_bias_total / mpu.get_tensor_model_parallel_world_size() - output_total = output_total*max_prob + output_total = output_total * max_prob output_total = output_total.view(s, b, h) if self.add_bias: - output_bias_total = output_bias_total*max_prob + output_bias_total = output_bias_total * max_prob output_bias_total = output_bias_total.view(s, b, h) else: output_bias_total = None @@ -318,9 +333,7 @@ def forward(self, hidden_states): class CoreAttention(MegatronModule): - - def __init__(self, layer_number, config, - attn_mask_type=AttnMaskType.padding): + def __init__(self, layer_number, config, attn_mask_type=AttnMaskType.padding): super(CoreAttention, self).__init__() self.fp16 = config.fp16 self.bf16 = config.bf16 @@ -337,12 +350,13 @@ def __init__(self, layer_number, config, # Per attention head and per partition values. world_size = mpu.get_tensor_model_parallel_world_size() - self.hidden_size_per_partition = core.utils.divide(projection_size, - world_size) + self.hidden_size_per_partition = core.utils.divide(projection_size, world_size) self.hidden_size_per_attention_head = core.utils.divide( - projection_size, config.num_attention_heads) + projection_size, config.num_attention_heads + ) self.num_attention_heads_per_partition = core.utils.divide( - config.num_attention_heads, world_size) + config.num_attention_heads, world_size + ) coeff = None self.norm_factor = math.sqrt(self.hidden_size_per_attention_head) @@ -351,49 +365,53 @@ def __init__(self, layer_number, config, self.norm_factor *= coeff self.scale_mask_softmax = FusedScaleMaskSoftmax( - self.fp16, self.bf16, + self.fp16, + self.bf16, self.attn_mask_type, config.masked_softmax_fusion, attention_mask_func, self.attention_softmax_in_fp32, - coeff) + coeff, + ) # Dropout. Note that for a single iteration, this layer will generate # different outputs on different number of parallel partitions but # on average it should not be partition dependent. self.attention_dropout = torch.nn.Dropout(config.attention_dropout) - def forward(self, query_layer, key_layer, - value_layer, attention_mask): - + def forward(self, query_layer, key_layer, value_layer, attention_mask): # =================================== # Raw attention scores. [b, np, s, s] # =================================== # [b, np, sq, sk] - output_size = (query_layer.size(1), - query_layer.size(2), - query_layer.size(0), - key_layer.size(0)) + output_size = ( + query_layer.size(1), + query_layer.size(2), + query_layer.size(0), + key_layer.size(0), + ) # [sq, b, np, hn] -> [sq, b * np, hn] - query_layer = query_layer.reshape(output_size[2], - output_size[0] * output_size[1], -1) + query_layer = query_layer.reshape(output_size[2], output_size[0] * output_size[1], -1) # [sk, b, np, hn] -> [sk, b * np, hn] - key_layer = key_layer.view(output_size[3], - output_size[0] * output_size[1], -1) + key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1) # preallocting input tensor: [b * np, sq, sk] matmul_input_buffer = mpu.get_global_memory_buffer().get_tensor( - (output_size[0]*output_size[1], output_size[2], output_size[3]), - query_layer.dtype, "mpu") + (output_size[0] * output_size[1], output_size[2], output_size[3]), + query_layer.dtype, + "mpu", + ) # Raw attention scores. [b * np, sq, sk] matmul_result = torch.baddbmm( matmul_input_buffer, - query_layer.transpose(0, 1), # [b * np, sq, hn] + query_layer.transpose(0, 1), # [b * np, sq, hn] key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk] - beta=0.0, alpha=(1.0/self.norm_factor)) + beta=0.0, + alpha=(1.0 / self.norm_factor), + ) # change view to [b, np, sq, sk] attention_scores = matmul_result.view(*output_size) @@ -403,8 +421,7 @@ def forward(self, query_layer, key_layer, # =========================== # attention scores and attention mask [b, np, sq, sk] - attention_probs = self.scale_mask_softmax(attention_scores, - attention_mask) + attention_probs = self.scale_mask_softmax(attention_scores, attention_mask) # This is actually dropping out entire tokens to attend to, which might # seem a bit unusual, but is taken from the original Transformer paper. @@ -422,18 +439,18 @@ def forward(self, query_layer, key_layer, # [sk, b, np, hn] --> [b, np, sq, hn] # context layer shape: [b, np, sq, hn] - output_size = (value_layer.size(1), - value_layer.size(2), - query_layer.size(0), - value_layer.size(3)) + output_size = ( + value_layer.size(1), + value_layer.size(2), + query_layer.size(0), + value_layer.size(3), + ) # change view [sk, b * np, hn] - value_layer = value_layer.view(value_layer.size(0), - output_size[0] * output_size[1], -1) + value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1) # change view [b * np, sq, sk] - attention_probs = attention_probs.view(output_size[0] * output_size[1], - output_size[2], -1) + attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1) # matmul: [b * np, sq, hn] context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1)) @@ -445,8 +462,7 @@ def forward(self, query_layer, key_layer, context_layer = context_layer.permute(2, 0, 1, 3).contiguous() # [sq, b, np, hn] --> [sq, b, hp] - new_context_layer_shape = context_layer.size()[:-2] + \ - (self.hidden_size_per_partition,) + new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,) context_layer = context_layer.view(*new_context_layer_shape) return context_layer @@ -462,12 +478,15 @@ class FlashSelfAttention(torch.nn.Module): attention_dropout: The dropout rate to apply to the attention (default: 0.0) """ - def __init__(self, causal=False, softmax_scale=None, attention_dropout=0.0, - device=None, dtype=None): + + def __init__( + self, causal=False, softmax_scale=None, attention_dropout=0.0, device=None, dtype=None + ): super().__init__() - assert flash_attn_unpadded_func is not None, ('Please install FlashAttention first, ' - 'e.g., with pip install flash-attn') - assert rearrange is not None, 'Please install einops first, e.g., with pip install einops' + assert flash_attn_unpadded_func is not None, ( + "Please install FlashAttention first, e.g., with pip install flash-attn" + ) + assert rearrange is not None, "Please install einops first, e.g., with pip install einops" self.causal = causal self.softmax_scale = softmax_scale self.dropout_p = attention_dropout @@ -479,15 +498,16 @@ def forward(self, q, k, v): q, k, v: The tensor containing the query, key, and value. (B, S, H, D) """ - assert all((i.dtype in [torch.float16, torch.bfloat16] for i in (q,k,v))) - assert all((i.is_cuda for i in (q,k,v))) + assert all((i.dtype in [torch.float16, torch.bfloat16] for i in (q, k, v))) + assert all((i.is_cuda for i in (q, k, v))) batch_size, seqlen_q = q.shape[0], q.shape[1] seqlen_k = k.shape[1] - q, k, v = [rearrange(x, 'b s ... -> (b s) ...') for x in [q, k, v]] - cu_seqlens_q = torch.arange(0, (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32, - device=q.device) + q, k, v = [rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v]] + cu_seqlens_q = torch.arange( + 0, (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32, device=q.device + ) if self.training: # during training q,k,v always have same seqlen @@ -500,17 +520,25 @@ def forward(self, q, k, v): # turn off FA causal mask after first inference autoregressive iteration # only on first autoregressive step q,k,v have same seqlen is_causal = seqlen_q == seqlen_k - cu_seqlens_k = torch.arange(0, (batch_size + 1) * seqlen_k, step=seqlen_k, dtype=torch.int32, - device=q.device) + cu_seqlens_k = torch.arange( + 0, (batch_size + 1) * seqlen_k, step=seqlen_k, dtype=torch.int32, device=q.device + ) dropout_p = 0 output = flash_attn_unpadded_func( - q, k, v, cu_seqlens_q, cu_seqlens_k, seqlen_q, seqlen_k, + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqlen_q, + seqlen_k, dropout_p, - softmax_scale=self.softmax_scale, causal=is_causal + softmax_scale=self.softmax_scale, + causal=is_causal, ) - output = rearrange(output, '(b s) ... -> b s ...', b=batch_size) + output = rearrange(output, "(b s) ... -> b s ...", b=batch_size) return output @@ -521,9 +549,13 @@ class ParallelAttention(MegatronModule): and returns output of the same size. """ - def __init__(self, config, layer_number, - attention_type=AttnType.self_attn, - attn_mask_type=AttnMaskType.padding): + def __init__( + self, + config, + layer_number, + attention_type=AttnType.self_attn, + attn_mask_type=AttnMaskType.padding, + ): super(ParallelAttention, self).__init__() args = get_args() self.layer_number = max(1, layer_number) @@ -541,33 +573,43 @@ def __init__(self, config, layer_number, else: kv_projection_size = args.kv_channels * args.num_attention_heads - self.use_flash_attn = args.use_flash_attn \ - and attention_type == AttnType.self_attn \ + self.use_flash_attn = ( + args.use_flash_attn + and attention_type == AttnType.self_attn and self.attn_mask_type == AttnMaskType.causal + ) if self.use_flash_attn: if flash_attn_unpadded_func is None: - raise ImportError('FlashAttention is not installed, please install with ' - 'pip install flash-attn') - assert attention_type == AttnType.self_attn, ('FlashAttention code path only supports ' - 'self-attention for now') - assert self.attn_mask_type == AttnMaskType.causal, ('FlashAttention code path only ' - 'supports causal mask for now') + raise ImportError( + "FlashAttention is not installed, please install with pip install flash-attn" + ) + assert attention_type == AttnType.self_attn, ( + "FlashAttention code path only supports self-attention for now" + ) + assert self.attn_mask_type == AttnMaskType.causal, ( + "FlashAttention code path only supports causal mask for now" + ) if rearrange is None: - raise ImportError('einops is not installed, please install with pip install einops') + raise ImportError("einops is not installed, please install with pip install einops") # Per attention head and per partition values. world_size = mpu.get_tensor_model_parallel_world_size() self.hidden_size_per_attention_head = core.utils.divide( - query_projection_size, config.num_attention_heads) + query_projection_size, config.num_attention_heads + ) self.num_attention_heads_per_partition = core.utils.divide( - config.num_attention_heads, world_size) + config.num_attention_heads, world_size + ) if self.group_query_attention: if args.num_query_groups % world_size != 0: - raise NotImplementedError('Currently the num_query_groups should be ' - 'a multiple of the tensor parallel size') + raise NotImplementedError( + "Currently the num_query_groups should be " + "a multiple of the tensor parallel size" + ) self.num_query_groups_per_partition = core.utils.divide( - args.num_query_groups, world_size) + args.num_query_groups, world_size + ) else: self.num_query_groups_per_partition = self.num_attention_heads_per_partition @@ -579,12 +621,15 @@ def __init__(self, config, layer_number, config=config, init_method=config.init_method, bias=args.add_bias_linear or args.add_qkv_bias, - gather_output=False) + gather_output=False, + ) else: assert attention_type == AttnType.cross_attn if self.group_query_attention: - raise NotImplementedError("Grouped query attention not implemented for cross-attention.") + raise NotImplementedError( + "Grouped query attention not implemented for cross-attention." + ) assert query_projection_size == kv_projection_size self.query = tensor_parallel.ColumnParallelLinear( @@ -593,7 +638,8 @@ def __init__(self, config, layer_number, config=config, init_method=config.init_method, bias=config.add_bias_linear, - gather_output=False) + gather_output=False, + ) self.key_value = tensor_parallel.ColumnParallelLinear( config.hidden_size, @@ -601,13 +647,12 @@ def __init__(self, config, layer_number, config=config, init_method=config.init_method, bias=config.add_bias_linear, - gather_output=False) + gather_output=False, + ) - self.core_attention = CoreAttention(self.layer_number, config, - self.attn_mask_type) + self.core_attention = CoreAttention(self.layer_number, config, self.attn_mask_type) self.checkpoint_core_attention = ( - config.recompute_granularity == 'selective' - and "core_attn" in config.recompute_modules + config.recompute_granularity == "selective" and "core_attn" in config.recompute_modules ) if self.use_flash_attn: @@ -623,28 +668,34 @@ def __init__(self, config, layer_number, init_method=config.output_layer_init_method, bias=args.add_bias_linear, input_is_parallel=True, - skip_bias_add=True) + skip_bias_add=True, + ) - def _checkpointed_attention_forward(self, query_layer, key_layer, - value_layer, attention_mask, - rotary_pos_emb=None): + def _checkpointed_attention_forward( + self, query_layer, key_layer, value_layer, attention_mask, rotary_pos_emb=None + ): """Forward method with activation checkpointing.""" + def custom_forward(*inputs): query_layer = inputs[0] key_layer = inputs[1] value_layer = inputs[2] attention_mask = inputs[3] - output_ = self.core_attention(query_layer, key_layer, - value_layer, attention_mask) + output_ = self.core_attention(query_layer, key_layer, value_layer, attention_mask) return output_ - q_pos_emb, k_pos_emb = (None, None) if rotary_pos_emb is None \ - else rotary_pos_emb + q_pos_emb, k_pos_emb = (None, None) if rotary_pos_emb is None else rotary_pos_emb hidden_states = tensor_parallel.checkpoint( custom_forward, - False, query_layer, key_layer, value_layer, attention_mask, - q_pos_emb, k_pos_emb) + False, + query_layer, + key_layer, + value_layer, + attention_mask, + q_pos_emb, + k_pos_emb, + ) return hidden_states @@ -655,11 +706,19 @@ def _allocate_memory(self, inference_max_sequence_len, batch_size, num_attention num_attention_heads, self.hidden_size_per_attention_head, dtype=self.params_dtype, - device=cur_platform.current_device()) + device=cur_platform.current_device(), + ) - def forward(self, hidden_states, attention_mask, - encoder_output=None, inference_context=None, - rotary_pos_emb=None, *, inference_params=None): + def forward( + self, + hidden_states, + attention_mask, + encoder_output=None, + inference_context=None, + rotary_pos_emb=None, + *, + inference_params=None, + ): # hidden_states: [sq, b, h] inference_context = deprecate_inference_params(inference_context, inference_params) @@ -673,24 +732,26 @@ def forward(self, hidden_states, attention_mask, inf_max_seq_len = inference_context.max_sequence_length inf_max_batch_size = inference_context.max_batch_size inference_key_memory = self._allocate_memory( - inf_max_seq_len, inf_max_batch_size, - self.num_query_groups_per_partition) + inf_max_seq_len, inf_max_batch_size, self.num_query_groups_per_partition + ) inference_value_memory = self._allocate_memory( - inf_max_seq_len, inf_max_batch_size, - self.num_query_groups_per_partition) + inf_max_seq_len, inf_max_batch_size, self.num_query_groups_per_partition + ) inference_context.key_value_memory_dict[self.layer_number] = ( - inference_key_memory, inference_value_memory) + inference_key_memory, + inference_value_memory, + ) is_first_step = True else: - inference_key_memory, inference_value_memory = \ + inference_key_memory, inference_value_memory = ( inference_context.key_value_memory_dict[self.layer_number] + ) # ===================== # Query, Key, and Value # ===================== if self.attention_type == AttnType.self_attn: - # Attention heads [sq, b, h] --> [sq, b, ng * (np/ng + 2) * hn)] mixed_x_layer, _ = self.query_key_value(hidden_states) @@ -698,49 +759,58 @@ def forward(self, hidden_states, attention_mask, new_tensor_shape = mixed_x_layer.size()[:-1] + ( self.num_query_groups_per_partition, ( - (self.num_attention_heads_per_partition // self.num_query_groups_per_partition + 2) + ( + self.num_attention_heads_per_partition + // self.num_query_groups_per_partition + + 2 + ) * self.hidden_size_per_attention_head ), ) mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) # [sq, b, ng, (np/ng + 2) * hn] --> [sq, b, ng, np/ng * hn], [sq, b, ng, hn], [sq, b, ng, hn] - (query_layer, - key_layer, - value_layer) = torch.split( + (query_layer, key_layer, value_layer) = torch.split( mixed_x_layer, [ ( - self.num_attention_heads_per_partition // self.num_query_groups_per_partition + self.num_attention_heads_per_partition + // self.num_query_groups_per_partition * self.hidden_size_per_attention_head ), self.hidden_size_per_attention_head, - self.hidden_size_per_attention_head + self.hidden_size_per_attention_head, ], - dim=3) + dim=3, + ) # [sq, b, ng, np/ng * hn] -> [sq, b, np, hn] - - query_layer = query_layer.view(query_layer.size(0), query_layer.size(1), -1, self.hidden_size_per_attention_head) + query_layer = query_layer.view( + query_layer.size(0), query_layer.size(1), -1, self.hidden_size_per_attention_head + ) else: # Attention heads [sk, b, h] --> [sk, b, (np * 2 * hn)] mixed_kv_layer, _ = self.key_value(encoder_output) # [sk, b, (np * 2 * hn)] --> [sk, b, np, 2 * hn] - new_tensor_shape = mixed_kv_layer.size()[:-1] + \ - (self.num_attention_heads_per_partition, - 2 * self.hidden_size_per_attention_head) + new_tensor_shape = mixed_kv_layer.size()[:-1] + ( + self.num_attention_heads_per_partition, + 2 * self.hidden_size_per_attention_head, + ) mixed_kv_layer = mixed_kv_layer.view(*new_tensor_shape) # [sk, b, np, 2 * hn] --> 2 [sk, b, np, hn] - (key_layer, - value_layer) = tensor_parallel.split_tensor_along_last_dim(mixed_kv_layer, 2) + (key_layer, value_layer) = tensor_parallel.split_tensor_along_last_dim( + mixed_kv_layer, 2 + ) # Attention head [sq, b, h] --> [sq, b, hp] query_layer, _ = self.query(hidden_states) # [sq, b, hp] --> [sq, b, np, hn] - new_tensor_shape = query_layer.size()[:-1] + \ - (self.num_attention_heads_per_partition, - self.hidden_size_per_attention_head) + new_tensor_shape = query_layer.size()[:-1] + ( + self.num_attention_heads_per_partition, + self.hidden_size_per_attention_head, + ) query_layer = query_layer.view(*new_tensor_shape) # ================================== @@ -752,7 +822,7 @@ def forward(self, hidden_states, attention_mask, if isinstance(rotary_pos_emb, tuple): rotary_pos_emb = rotary_pos_emb else: - rotary_pos_emb = ((rotary_pos_emb,) * 2) + rotary_pos_emb = (rotary_pos_emb,) * 2 if inference_context: batch_start = inference_context.batch_size_offset @@ -760,18 +830,19 @@ def forward(self, hidden_states, attention_mask, assert batch_end <= inference_key_memory.size(1) sequence_start = inference_context.sequence_len_offset sequence_end = sequence_start + key_layer.size(0) - assert sequence_end <= inference_key_memory.size(0), ("Current sequence length is " - "longer than expected maximum sequence length! Increase inference_max_seq_length.") + assert sequence_end <= inference_key_memory.size(0), ( + "Current sequence length is " + "longer than expected maximum sequence length! Increase inference_max_seq_length." + ) # Copy key and values. - inference_key_memory[sequence_start:sequence_end, - batch_start:batch_end, ...] = key_layer - inference_value_memory[sequence_start:sequence_end, - batch_start:batch_end, ...] = value_layer - key_layer = inference_key_memory[ - :sequence_end, batch_start:batch_end, ...] - value_layer = inference_value_memory[ - :sequence_end, batch_start:batch_end, ...] - + inference_key_memory[sequence_start:sequence_end, batch_start:batch_end, ...] = ( + key_layer + ) + inference_value_memory[sequence_start:sequence_end, batch_start:batch_end, ...] = ( + value_layer + ) + key_layer = inference_key_memory[:sequence_end, batch_start:batch_end, ...] + value_layer = inference_value_memory[:sequence_end, batch_start:batch_end, ...] # adjust the key rotary positional embedding if rotary_pos_emb is not None: @@ -800,19 +871,17 @@ def forward(self, hidden_states, attention_mask, # expand the key_layer and value_layer [sk, b, ng, hn] -> [sk, b, np, hn] if self.num_attention_heads_per_partition // self.num_query_groups_per_partition > 1: key_layer = key_layer.repeat_interleave( - self.num_attention_heads_per_partition // self.num_query_groups_per_partition, - dim = 2 + self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim=2 ) value_layer = value_layer.repeat_interleave( - self.num_attention_heads_per_partition // self.num_query_groups_per_partition, - dim = 2 + self.num_attention_heads_per_partition // self.num_query_groups_per_partition, dim=2 ) # apply relative positional encoding (rotary embedding) if rotary_pos_emb is not None: q_pos_emb, k_pos_emb = rotary_pos_emb - query_layer = apply_rotary_pos_emb(query_layer, q_pos_emb,self.config) - key_layer = apply_rotary_pos_emb(key_layer, k_pos_emb,self.config) + query_layer = apply_rotary_pos_emb(query_layer, q_pos_emb, self.config) + key_layer = apply_rotary_pos_emb(key_layer, k_pos_emb, self.config) # TODO, can apply positional embedding to value_layer so it has # absolute positional embedding. # otherwise, only relative positional embedding takes effect @@ -821,19 +890,23 @@ def forward(self, hidden_states, attention_mask, if not self.use_flash_attn: if self.checkpoint_core_attention: context_layer = self._checkpointed_attention_forward( - query_layer, key_layer, value_layer, attention_mask) + query_layer, key_layer, value_layer, attention_mask + ) else: context_layer = self.core_attention( - query_layer, key_layer, value_layer, attention_mask) + query_layer, key_layer, value_layer, attention_mask + ) else: - q, k, v = [rearrange(x, 's b ... -> b s ...').contiguous() - for x in (query_layer, key_layer, value_layer)] + q, k, v = [ + rearrange(x, "s b ... -> b s ...").contiguous() + for x in (query_layer, key_layer, value_layer) + ] if not self.sequence_parallel: with tensor_parallel.get_cuda_rng_tracker().fork(): context_layer = self.core_attention_flash(q, k, v) else: context_layer = self.core_attention_flash(q, k, v) - context_layer = rearrange(context_layer, 'b s h d -> s b (h d)').contiguous() + context_layer = rearrange(context_layer, "b s h d -> s b (h d)").contiguous() # ================= # Output. [sq, b, h] @@ -856,22 +929,21 @@ def bias_dropout_add(x, bias, residual, prob, training): def get_bias_dropout_add(training): def _bias_dropout_add(x, bias, residual, prob): return bias_dropout_add(x, bias, residual, prob, training) + return _bias_dropout_add @jit_fuser -def bias_dropout_add_fused_train(x: torch.Tensor, - bias: Optional[torch.Tensor], - residual: torch.Tensor, - prob: float) -> torch.Tensor: +def bias_dropout_add_fused_train( + x: torch.Tensor, bias: Optional[torch.Tensor], residual: torch.Tensor, prob: float +) -> torch.Tensor: return bias_dropout_add(x, bias, residual, prob, True) @jit_fuser -def bias_dropout_add_fused_inference(x: torch.Tensor, - bias: Optional[torch.Tensor], - residual: torch.Tensor, - prob: float) -> torch.Tensor: +def bias_dropout_add_fused_inference( + x: torch.Tensor, bias: Optional[torch.Tensor], residual: torch.Tensor, prob: float +) -> torch.Tensor: return bias_dropout_add(x, bias, residual, prob, False) @@ -882,18 +954,21 @@ class ParallelTransformerLayer(MegatronModule): output of the same size. """ - def __init__(self, config, - layer_number, layer_type=LayerType.encoder, - self_attn_mask_type=AttnMaskType.padding, - drop_path_rate=0.): + def __init__( + self, + config, + layer_number, + layer_type=LayerType.encoder, + self_attn_mask_type=AttnMaskType.padding, + drop_path_rate=0.0, + ): args = get_args() super(ParallelTransformerLayer, self).__init__() self.layer_number = layer_number self.layer_type = layer_type - self.apply_residual_connection_post_norm \ - = config.apply_residual_connection_post_layernorm + self.apply_residual_connection_post_norm = config.apply_residual_connection_post_layernorm self.bf16 = config.bf16 self.fp32_residual_connection = config.fp32_residual_connection @@ -906,7 +981,8 @@ def __init__(self, config, config, layer_number, attention_type=AttnType.self_attn, - attn_mask_type=self_attn_mask_type) + attn_mask_type=self_attn_mask_type, + ) self.hidden_dropout = config.hidden_dropout self.bias_dropout_fusion = config.bias_dropout_fusion self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0.0 else None @@ -921,25 +997,20 @@ def __init__(self, config, self.mlp = ParallelMLP(config) # Set bias+dropout+add fusion grad_enable execution handler. - TORCH_MAJOR = int(torch.__version__.split('.')[0]) - TORCH_MINOR = int(torch.__version__.split('.')[1]) + TORCH_MAJOR = int(torch.__version__.split(".")[0]) + TORCH_MINOR = int(torch.__version__.split(".")[1]) use_nvfuser = TORCH_MAJOR > 1 or (TORCH_MAJOR == 1 and TORCH_MINOR >= 10) - self.bias_dropout_add_exec_handler = \ - nullcontext if use_nvfuser else torch.enable_grad + self.bias_dropout_add_exec_handler = nullcontext if use_nvfuser else torch.enable_grad - def default_decoder_cross_attention(self, - encoder_output, - enc_dec_attn_mask, - norm_input, - norm_output, - bias_dropout_add_func): - '''Cross attention for a standard encoder-decoder model.''' + def default_decoder_cross_attention( + self, encoder_output, enc_dec_attn_mask, norm_input, norm_output, bias_dropout_add_func + ): + """Cross attention for a standard encoder-decoder model.""" # Attention. - attention_output, attention_bias = \ - self.inter_attention(norm_output, - enc_dec_attn_mask, - encoder_output=encoder_output) + attention_output, attention_bias = self.inter_attention( + norm_output, enc_dec_attn_mask, encoder_output=encoder_output + ) # Residual connection. if self.apply_residual_connection_post_norm: @@ -953,23 +1024,25 @@ def default_decoder_cross_attention(self, # Bias-dropout-add. with self.bias_dropout_add_exec_handler(): norm_input = bias_dropout_add_func( - attention_output, - attention_bias, - residual, - self.hidden_dropout) + attention_output, attention_bias, residual, self.hidden_dropout + ) # Normalize. norm_output = self.post_inter_attention_norm(norm_input) return norm_input, norm_output - def forward(self, hidden_states, attention_mask, - encoder_output=None, enc_dec_attn_mask=None, - inference_context=None, - rotary_pos_emb=None, - *, - inference_params=None): - + def forward( + self, + hidden_states, + attention_mask, + encoder_output=None, + enc_dec_attn_mask=None, + inference_context=None, + rotary_pos_emb=None, + *, + inference_params=None, + ): inference_context = deprecate_inference_params(inference_context, inference_params) # hidden_states: [s, b, h] @@ -978,12 +1051,12 @@ def forward(self, hidden_states, attention_mask, norm_output = self.input_norm(hidden_states) # Self attention. - attention_output, attention_bias = \ - self.self_attention( - norm_output, - attention_mask, - inference_context=inference_context, - rotary_pos_emb=rotary_pos_emb) + attention_output, attention_bias = self.self_attention( + norm_output, + attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + ) # Residual connection. if self.apply_residual_connection_post_norm: @@ -1008,14 +1081,12 @@ def forward(self, hidden_states, attention_mask, attention_bias = attention_bias.expand_as(residual) with self.bias_dropout_add_exec_handler(): norm_input = bias_dropout_add_func( - attention_output, - attention_bias, - residual, - self.hidden_dropout) + attention_output, attention_bias, residual, self.hidden_dropout + ) else: - out = torch.nn.functional.dropout(attention_output + attention_bias, - p=self.hidden_dropout, - training=self.training) + out = torch.nn.functional.dropout( + attention_output + attention_bias, p=self.hidden_dropout, training=self.training + ) norm_input = residual + self.drop_path(out) # Layer norm post the self attention. @@ -1025,16 +1096,11 @@ def forward(self, hidden_states, attention_mask, if self.layer_type == LayerType.encoder: pass elif self.layer_type == LayerType.decoder: - norm_input, norm_output = \ - self.default_decoder_cross_attention( - encoder_output, - enc_dec_attn_mask, - norm_input, - norm_output, - bias_dropout_add_func) + norm_input, norm_output = self.default_decoder_cross_attention( + encoder_output, enc_dec_attn_mask, norm_input, norm_output, bias_dropout_add_func + ) else: - raise Exception("Unsupported layer type, '%s'." % - self.layer_type.name) + raise Exception("Unsupported layer type, '%s'." % self.layer_type.name) # MLP. mlp_output, mlp_bias = self.mlp(norm_output) @@ -1049,11 +1115,7 @@ def forward(self, hidden_states, attention_mask, if mlp_bias is not None: mlp_bias = mlp_bias.expand_as(residual) with self.bias_dropout_add_exec_handler(): - output = bias_dropout_add_func( - mlp_output, - mlp_bias, - residual, - self.hidden_dropout) + output = bias_dropout_add_func(mlp_output, mlp_bias, residual, self.hidden_dropout) # Jit compiled function creates 'view' tensor. This tensor # potentially gets saved in the MPU checkpoint function context, @@ -1061,16 +1123,16 @@ def forward(self, hidden_states, attention_mask, # won't result in memory savings (like the data loader, or # p2p_communication), it serves to document the origin of this # 'view' tensor. - output = core.utils.make_viewless_tensor(inp = output, - requires_grad = output.requires_grad, - keep_graph = True) + output = core.utils.make_viewless_tensor( + inp=output, requires_grad=output.requires_grad, keep_graph=True + ) else: if mlp_bias is not None: mlp_output = mlp_output + mlp_bias - out = torch.nn.functional.dropout(mlp_output, - p=self.hidden_dropout, - training=self.training) + out = torch.nn.functional.dropout( + mlp_output, p=self.hidden_dropout, training=self.training + ) output = residual + self.drop_path(out) return output @@ -1096,9 +1158,16 @@ def __init__(self, layer_number): super().__init__() self.layer_number = layer_number - def forward(self, hidden_states, attention_mask, - encoder_output=None, enc_dec_attn_mask=None, - inference_context=None, *, inference_params=None): + def forward( + self, + hidden_states, + attention_mask, + encoder_output=None, + enc_dec_attn_mask=None, + inference_context=None, + *, + inference_params=None, + ): return hidden_states.clone() @@ -1106,8 +1175,9 @@ def _get_num_layers(args, model_type, is_decoder=False): """Compute the number of transformer layers resident on the current rank.""" if mpu.get_pipeline_model_parallel_world_size() > 1: assert args.num_layers == args.encoder_num_layers - assert args.num_layers % args.transformer_pipeline_model_parallel_size == 0, \ - 'num_layers must be divisible by transformer_pipeline_model_parallel_size' + assert args.num_layers % args.transformer_pipeline_model_parallel_size == 0, ( + "num_layers must be divisible by transformer_pipeline_model_parallel_size" + ) # When a standalone embedding stage is used, all transformer layers # are divided among pipeline rank >= 1, while on pipeline rank 0, @@ -1116,8 +1186,8 @@ def _get_num_layers(args, model_type, is_decoder=False): num_layers = ( 0 if args.account_for_embedding_in_pipeline_split - and mpu.get_pipeline_model_parallel_rank() == 0 else - args.num_layers // args.transformer_pipeline_model_parallel_size + and mpu.get_pipeline_model_parallel_rank() == 0 + else args.num_layers // args.transformer_pipeline_model_parallel_size ) else: if not is_decoder: @@ -1130,13 +1200,17 @@ def _get_num_layers(args, model_type, is_decoder=False): class ParallelTransformer(MegatronModule): """Transformer class.""" - def __init__(self, config, - model_type, layer_type=LayerType.encoder, - self_attn_mask_type=AttnMaskType.padding, - post_norm=True, - pre_process=True, - post_process=True, - drop_path_rate=0.0): + def __init__( + self, + config, + model_type, + layer_type=LayerType.encoder, + self_attn_mask_type=AttnMaskType.padding, + post_norm=True, + pre_process=True, + post_process=True, + drop_path_rate=0.0, + ): super(ParallelTransformer, self).__init__() args = get_args() @@ -1155,8 +1229,9 @@ def __init__(self, config, self.recompute_granularity = config.recompute_granularity self.recompute_method = config.recompute_method self.recompute_num_layers = config.recompute_num_layers - self.distribute_saved_activations = \ + self.distribute_saved_activations = ( config.distribute_saved_activations and not config.sequence_parallel + ) self.sequence_parallel = config.sequence_parallel @@ -1164,7 +1239,7 @@ def __init__(self, config, self.transformer_engine_v_0_10 = False self.transformer_engine_v_0_11 = False self.transformer_engine_v_0_8 = False - if self.transformer_impl == 'transformer_engine': + if self.transformer_impl == "transformer_engine": global transformer_engine import transformer_engine @@ -1175,15 +1250,17 @@ def __init__(self, config, if core.utils.is_te_min_version("0.11.0"): self.transformer_engine_v_0_11 = True - assert not args.squared_relu, ("TransformerEngine does not support squared " - "relu activation.") + assert not args.squared_relu, ( + "TransformerEngine does not support squared relu activation." + ) self.use_fp8 = args.fp8 is not None self.fp8_recipe = None self.fp8_group = None if self.use_fp8: - assert args.transformer_impl == 'transformer_engine', \ - 'transformer-engine required for fp8 training and inference' + assert args.transformer_impl == "transformer_engine", ( + "transformer-engine required for fp8 training and inference" + ) self.fp8_group = mpu.get_amax_reduction_group(tp_only_amax_red=config.tp_only_amax_red) if args.fp8 == "e4m3": fp8_format = transformer_engine.common.recipe.Format.E4M3 @@ -1203,40 +1280,45 @@ def __init__(self, config, self.num_microbatches_in_previous_step = -1 self.microbatch_count = 0 self.checkpoint_core_attention = ( - config.recompute_granularity == 'selective' - and "core_attn" in config.recompute_modules + config.recompute_granularity == "selective" and "core_attn" in config.recompute_modules ) # Number of layers. - self.num_layers = _get_num_layers(args, model_type, - layer_type==LayerType.decoder) + self.num_layers = _get_num_layers(args, model_type, layer_type == LayerType.decoder) self.drop_path_rates = [ - rate.item() for rate in - torch.linspace(0, self.drop_path_rate, config.num_layers)] + rate.item() for rate in torch.linspace(0, self.drop_path_rate, config.num_layers) + ] def build_layer(layer_number): - if args.transformer_impl == 'local': + if args.transformer_impl == "local": return ParallelTransformerLayer( config, layer_number, layer_type=layer_type, self_attn_mask_type=self_attn_mask_type, - drop_path_rate=self.drop_path_rates[layer_number - 1]) + drop_path_rate=self.drop_path_rates[layer_number - 1], + ) else: # This argument is only available from TE v0.10 onwards. extra_transformer_engine_kwargs = {} if self.transformer_engine_v_0_8: extra_transformer_engine_kwargs["bias"] = args.add_bias_linear if self.transformer_engine_v_0_10: - extra_transformer_engine_kwargs["activation"] = "swiglu" if args.swiglu else "gelu" + extra_transformer_engine_kwargs["activation"] = ( + "swiglu" if args.swiglu else "gelu" + ) if self.transformer_engine_v_0_11: extra_transformer_engine_kwargs["normalization"] = args.normalization - assert config.attention_softmax_in_fp32, "TransformerEngine only supports softmax compute in FP32." + assert config.attention_softmax_in_fp32, ( + "TransformerEngine only supports softmax compute in FP32." + ) assert ( - (bool(int(os.getenv("NVTE_APPLY_QK_LAYER_SCALING", "0"))) and args.fp16) == config.apply_query_key_layer_scaling - ), ("Unsupported config for apply_query_key_layer_scaling in TransformerEngine. If --apply-query-key-layer-scaling is " - "provided, set env-var NVTE_APPLY_QK_LAYER_SCALING=1 and you must be using fp16.") + bool(int(os.getenv("NVTE_APPLY_QK_LAYER_SCALING", "0"))) and args.fp16 + ) == config.apply_query_key_layer_scaling, ( + "Unsupported config for apply_query_key_layer_scaling in TransformerEngine. If --apply-query-key-layer-scaling is " + "provided, set env-var NVTE_APPLY_QK_LAYER_SCALING=1 and you must be using fp16." + ) return transformer_engine.pytorch.TransformerLayer( config.hidden_size, config.ffn_hidden_size, @@ -1249,7 +1331,9 @@ def build_layer(layer_number): layer_number=layer_number, kv_channels=config.kv_channels, self_attn_mask_type=self_attn_mask_type.name, - tp_group=mpu.get_tensor_model_parallel_group() if mpu.is_initialized() else None, + tp_group=mpu.get_tensor_model_parallel_group() + if mpu.is_initialized() + else None, tp_size=mpu.get_tensor_model_parallel_world_size(), get_rng_state_tracker=get_cuda_rng_tracker if get_cuda_rng_tracker().is_initialized() @@ -1265,12 +1349,13 @@ def build_layer(layer_number): drop_path_rate=self.drop_path_rates[layer_number - 1], set_parallel_mode=True, fuse_qkv_params=True, - **extra_transformer_engine_kwargs) + **extra_transformer_engine_kwargs, + ) if config.virtual_pipeline_model_parallel_size is not None: - assert config.num_layers % config.virtual_pipeline_model_parallel_size == 0, \ - 'num_layers_per_stage must be divisible by ' \ - 'virtual_pipeline_model_parallel_size' + assert config.num_layers % config.virtual_pipeline_model_parallel_size == 0, ( + "num_layers_per_stage must be divisible by virtual_pipeline_model_parallel_size" + ) # Number of layers in each model chunk is the number of layers in the stage, # divided by the number of model chunks in a stage. self.num_layers = self.num_layers // config.virtual_pipeline_model_parallel_size @@ -1283,8 +1368,8 @@ def build_layer(layer_number): # Stage 0: [0, 1] [4, 5] # Stage 1: [2, 3] [6, 7] offset = mpu.get_virtual_pipeline_model_parallel_rank() * ( - config.num_layers // config.virtual_pipeline_model_parallel_size) + \ - (mpu.get_pipeline_model_parallel_rank() * self.num_layers) + config.num_layers // config.virtual_pipeline_model_parallel_size + ) + (mpu.get_pipeline_model_parallel_rank() * self.num_layers) else: # Each stage gets a contiguous set of layers. offset = mpu.get_pipeline_model_parallel_rank() * self.num_layers @@ -1299,10 +1384,11 @@ def build_layer(layer_number): # this, we assign a 'no-op' layer on these ranks, which will # disconnect the input tensor from the output tensor. self.num_layers = 1 - self.layers = torch.nn.ModuleList([ NoopTransformerLayer(1) ]) + self.layers = torch.nn.ModuleList([NoopTransformerLayer(1)]) else: self.layers = torch.nn.ModuleList( - [build_layer(i + 1 + offset) for i in range(self.num_layers)]) + [build_layer(i + 1 + offset) for i in range(self.num_layers)] + ) if self.post_process and self.post_norm: # Final layer norm before output. @@ -1311,10 +1397,17 @@ def build_layer(layer_number): def _get_layer(self, layer_number): return self.layers[layer_number] - def _checkpointed_forward(self, hidden_states, attention_mask, - encoder_output, enc_dec_attn_mask, - rotary_pos_emb, is_first_microbatch): + def _checkpointed_forward( + self, + hidden_states, + attention_mask, + encoder_output, + enc_dec_attn_mask, + rotary_pos_emb, + is_first_microbatch, + ): """Forward method with activation checkpointing.""" + def custom(start, end): def custom_forward(*args, **kwargs): x_, *args = args @@ -1322,69 +1415,103 @@ def custom_forward(*args, **kwargs): layer = self._get_layer(index) x_ = layer(x_, *args, **kwargs) return x_ + return custom_forward te_forward_kwargs = {} - if self.transformer_impl == 'transformer_engine': - te_forward_kwargs['is_first_microbatch'] = is_first_microbatch + if self.transformer_impl == "transformer_engine": + te_forward_kwargs["is_first_microbatch"] = is_first_microbatch if self.transformer_engine_v_0_10: - te_forward_kwargs['rotary_pos_emb'] = rotary_pos_emb + te_forward_kwargs["rotary_pos_emb"] = rotary_pos_emb - if self.recompute_method == 'uniform': + if self.recompute_method == "uniform": # Uniformly divide the total number of Transformer layers and # checkpoint the input activation of each divided chunk. # A method to further reduce memory usage reducing checkpoints. l = 0 while l < self.num_layers: - if self.transformer_impl == 'transformer_engine': + if self.transformer_impl == "transformer_engine": hidden_states = transformer_engine.pytorch.checkpoint( custom(l, l + self.recompute_num_layers), self.distribute_saved_activations, tensor_parallel.get_cuda_rng_tracker, mpu.get_tensor_model_parallel_group(), - hidden_states, attention_mask, encoder_output, - enc_dec_attn_mask, **te_forward_kwargs) + hidden_states, + attention_mask, + encoder_output, + enc_dec_attn_mask, + **te_forward_kwargs, + ) else: hidden_states = tensor_parallel.checkpoint( custom(l, l + self.recompute_num_layers), self.distribute_saved_activations, - hidden_states, attention_mask, - encoder_output, enc_dec_attn_mask, - None, None, None, None, rotary_pos_emb) + hidden_states, + attention_mask, + encoder_output, + enc_dec_attn_mask, + None, + None, + None, + None, + rotary_pos_emb, + ) l += self.recompute_num_layers - elif self.recompute_method == 'block': + elif self.recompute_method == "block": # Checkpoint the input activation of only a set number of individual # Transformer layers and skip the rest. # A method fully use the device memory removing redundant re-computation. for l in range(self.num_layers): if l < self.recompute_num_layers: - if self.transformer_impl == 'transformer_engine': + if self.transformer_impl == "transformer_engine": hidden_states = transformer_engine.pytorch.checkpoint( custom(l, l + 1), self.distribute_saved_activations, tensor_parallel.get_cuda_rng_tracker, mpu.get_tensor_model_parallel_group(), - hidden_states, attention_mask, encoder_output, - enc_dec_attn_mask, **te_forward_kwargs) + hidden_states, + attention_mask, + encoder_output, + enc_dec_attn_mask, + **te_forward_kwargs, + ) else: hidden_states = tensor_parallel.checkpoint( custom(l, l + 1), self.distribute_saved_activations, - hidden_states, attention_mask, - encoder_output, enc_dec_attn_mask, - None, None, None, None, rotary_pos_emb) + hidden_states, + attention_mask, + encoder_output, + enc_dec_attn_mask, + None, + None, + None, + None, + rotary_pos_emb, + ) else: - if self.transformer_impl == 'transformer_engine': + if self.transformer_impl == "transformer_engine": hidden_states = custom(l, l + 1)( - hidden_states, attention_mask, encoder_output, - enc_dec_attn_mask, **te_forward_kwargs) + hidden_states, + attention_mask, + encoder_output, + enc_dec_attn_mask, + **te_forward_kwargs, + ) else: hidden_states = custom(l, l + 1)( - hidden_states, attention_mask, - encoder_output, enc_dec_attn_mask, - None, None, None, None, rotary_pos_emb) + hidden_states, + attention_mask, + encoder_output, + enc_dec_attn_mask, + None, + None, + None, + None, + rotary_pos_emb, + ) else: raise ValueError("Invalid activation recompute method.") @@ -1400,20 +1527,26 @@ def set_input_tensor(self, input_tensor): forward_step_func""" self.input_tensor = input_tensor - def forward(self, hidden_states, attention_mask, - encoder_output=None, enc_dec_attn_mask=None, - inference_context=None, - rotary_pos_emb=None, - *, - inference_params=None): + def forward( + self, + hidden_states, + attention_mask, + encoder_output=None, + enc_dec_attn_mask=None, + inference_context=None, + rotary_pos_emb=None, + *, + inference_params=None, + ): # hidden_states: [s, b, h] inference_context = deprecate_inference_params(inference_context, inference_params) # Checks. if inference_context: - assert self.recompute_granularity is None, \ - 'inference does not work with activation checkpointing' + assert self.recompute_granularity is None, ( + "inference does not work with activation checkpointing" + ) if not self.pre_process: # See set_input_tensor() @@ -1450,50 +1583,51 @@ def forward(self, hidden_states, attention_mask, with rng_context: # The fp8_autocast context manager is a no-op when enabled=True # The if...else serves to short circuit name resolution for fp8_autocast - with transformer_engine.pytorch.fp8_autocast( - enabled=self.use_fp8, - fp8_recipe=self.fp8_recipe, - fp8_group=self.fp8_group - ) if self.use_fp8 else nullcontext(): + with ( + transformer_engine.pytorch.fp8_autocast( + enabled=self.use_fp8, fp8_recipe=self.fp8_recipe, fp8_group=self.fp8_group + ) + if self.use_fp8 + else nullcontext() + ): # Determine if the current iteration is first microbatch if self.num_microbatches_in_previous_step != get_num_microbatches(): - self.microbatch_count = 0 # Reset count on new batch size rampup interval + self.microbatch_count = 0 # Reset count on new batch size rampup interval self.num_microbatches_in_previous_step = get_num_microbatches() is_first_microbatch = self.microbatch_count % get_num_microbatches() == 0 # Forward pass. - if self.recompute_granularity == 'full': - hidden_states = self._checkpointed_forward(hidden_states, - attention_mask, - encoder_output, - enc_dec_attn_mask, - rotary_pos_emb, - is_first_microbatch) + if self.recompute_granularity == "full": + hidden_states = self._checkpointed_forward( + hidden_states, + attention_mask, + encoder_output, + enc_dec_attn_mask, + rotary_pos_emb, + is_first_microbatch, + ) else: forward_kwargs = { - 'encoder_output': encoder_output, - 'enc_dec_attn_mask': enc_dec_attn_mask, + "encoder_output": encoder_output, + "enc_dec_attn_mask": enc_dec_attn_mask, } - if self.transformer_impl == 'local': - forward_kwargs['inference_context'] = inference_context + if self.transformer_impl == "local": + forward_kwargs["inference_context"] = inference_context else: - forward_kwargs['inference_params'] = inference_context + forward_kwargs["inference_params"] = inference_context - if self.transformer_impl == 'transformer_engine': - forward_kwargs['is_first_microbatch'] = is_first_microbatch - forward_kwargs['checkpoint_core_attention'] = self.checkpoint_core_attention + if self.transformer_impl == "transformer_engine": + forward_kwargs["is_first_microbatch"] = is_first_microbatch + forward_kwargs["checkpoint_core_attention"] = self.checkpoint_core_attention if self.transformer_engine_v_0_10: - forward_kwargs['rotary_pos_emb'] = rotary_pos_emb + forward_kwargs["rotary_pos_emb"] = rotary_pos_emb else: - forward_kwargs['rotary_pos_emb'] = rotary_pos_emb + forward_kwargs["rotary_pos_emb"] = rotary_pos_emb for index in range(self.num_layers): layer = self._get_layer(index) - hidden_states = layer( - hidden_states, - attention_mask, - **forward_kwargs) + hidden_states = layer(hidden_states, attention_mask, **forward_kwargs) # Skip counter update for eval and activation checkpointing if torch.is_grad_enabled() and self.training: diff --git a/flagscale/train/megatron/legacy/model/utils.py b/flagscale/train/megatron/legacy/model/utils.py index 1ea05bb96f..31579f1fef 100644 --- a/flagscale/train/megatron/legacy/model/utils.py +++ b/flagscale/train/megatron/legacy/model/utils.py @@ -10,8 +10,10 @@ from megatron.legacy.model import LayerNorm, RMSNorm from megatron.core.jit import jit_fuser + def init_method_normal(sigma): """Init method based on N(0, sigma).""" + def init_(tensor): return torch.nn.init.normal_(tensor, mean=0.0, std=sigma) @@ -46,17 +48,19 @@ def get_linear_layer(rows, columns, init_method): @jit_fuser def gelu_impl(x): """OpenAI's gelu implementation.""" - return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x * + return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x))) + - (1.0 + 0.044715 * x * x))) def openai_gelu(x): return gelu_impl(x) -#This is actually Python equivalent of torch.nn.functional.gelu(), also with type hints for ONNX exporter +# This is actually Python equivalent of torch.nn.functional.gelu(), also with type hints for ONNX exporter @jit_fuser def erf_gelu(x): - return x * 0.5 * (torch.erf(x / 1.41421).to(dtype=x.dtype)+torch.ones_like(x).to(dtype=x.dtype)) + return ( + x * 0.5 * (torch.erf(x / 1.41421).to(dtype=x.dtype) + torch.ones_like(x).to(dtype=x.dtype)) + ) def get_norm(config): @@ -67,13 +71,18 @@ def get_norm(config): eps=config.layernorm_epsilon, no_persist_layer_norm=not config.persist_layer_norm, sequence_parallel=config.sequence_parallel, - apply_layernorm_1p=args.apply_layernorm_1p) + apply_layernorm_1p=args.apply_layernorm_1p, + ) elif args.normalization == "RMSNorm": if "apply_layernorm_1p" in args and args.apply_layernorm_1p: - raise NotImplementedError('RMSNorm does not currently support the layernorm_1p formulation.') + raise NotImplementedError( + "RMSNorm does not currently support the layernorm_1p formulation." + ) - return RMSNorm(dim=config.hidden_size, - eps=config.layernorm_epsilon, - sequence_parallel=config.sequence_parallel) + return RMSNorm( + dim=config.hidden_size, + eps=config.layernorm_epsilon, + sequence_parallel=config.sequence_parallel, + ) else: raise Exception(f"unsupported norm type '{args.normalization}'.") diff --git a/flagscale/train/megatron/legacy/model/vision/classification.py b/flagscale/train/megatron/legacy/model/vision/classification.py index f9419c71de..da180d6741 100644 --- a/flagscale/train/megatron/legacy/model/vision/classification.py +++ b/flagscale/train/megatron/legacy/model/vision/classification.py @@ -10,11 +10,11 @@ from megatron.legacy.model.vision.mit_backbone import mit_b3_avg from megatron.legacy.model.module import MegatronModule + class VitClassificationModel(MegatronModule): """Vision Transformer Model.""" - def __init__(self, config, num_classes, finetune=False, - pre_process=True, post_process=True): + def __init__(self, config, num_classes, finetune=False, pre_process=True, post_process=True): super(VitClassificationModel, self).__init__() args = get_args() self.config = config @@ -28,7 +28,7 @@ def __init__(self, config, num_classes, finetune=False, config=config, pre_process=self.pre_process, post_process=self.post_process, - single_token_output=True + single_token_output=True, ) if self.post_process: @@ -36,9 +36,7 @@ def __init__(self, config, num_classes, finetune=False, self.head = VitMlpHead(config, self.hidden_size, self.num_classes) else: self.head = get_linear_layer( - self.hidden_size, - self.num_classes, - torch.nn.init.zeros_ + self.hidden_size, self.num_classes, torch.nn.init.zeros_ ) def set_input_tensor(self, input_tensor): @@ -57,8 +55,7 @@ def forward(self, input): class MitClassificationModel(MegatronModule): """Mix vision Transformer Model.""" - def __init__(self, num_classes, - pre_process=True, post_process=True): + def __init__(self, num_classes, pre_process=True, post_process=True): super(MitClassificationModel, self).__init__() args = get_args() @@ -71,7 +68,7 @@ def __init__(self, num_classes, def _init_weights(self, m): if isinstance(m, torch.nn.Linear): - trunc_normal_(m.weight, std=.02) + trunc_normal_(m.weight, std=0.02) if isinstance(m, torch.nn.Linear) and m.bias is not None: torch.nn.init.constant_(m.bias, 0) diff --git a/flagscale/train/megatron/legacy/model/vision/dino.py b/flagscale/train/megatron/legacy/model/vision/dino.py index 20ca2100f6..cf4eec42e0 100644 --- a/flagscale/train/megatron/legacy/model/vision/dino.py +++ b/flagscale/train/megatron/legacy/model/vision/dino.py @@ -21,9 +21,17 @@ class DINOLoss(torch.nn.Module): - def __init__(self, out_dim, ncrops, warmup_teacher_temp, teacher_temp, - warmup_teacher_temp_epochs, nepochs, student_temp=0.1, - center_momentum=0.9): + def __init__( + self, + out_dim, + ncrops, + warmup_teacher_temp, + teacher_temp, + warmup_teacher_temp_epochs, + nepochs, + student_temp=0.1, + center_momentum=0.9, + ): super().__init__() self.student_temp = student_temp self.center_momentum = center_momentum @@ -31,11 +39,12 @@ def __init__(self, out_dim, ncrops, warmup_teacher_temp, teacher_temp, self.register_buffer("center", torch.zeros(1, out_dim)) # we apply a warm up for the teacher temperature because # a too high temperature makes the training instable at the beginning - self.teacher_temp_schedule = np.concatenate(( - np.linspace(warmup_teacher_temp, - teacher_temp, warmup_teacher_temp_epochs), - np.ones(nepochs - warmup_teacher_temp_epochs) * teacher_temp - )) + self.teacher_temp_schedule = np.concatenate( + ( + np.linspace(warmup_teacher_temp, teacher_temp, warmup_teacher_temp_epochs), + np.ones(nepochs - warmup_teacher_temp_epochs) * teacher_temp, + ) + ) self.teacher_temp = teacher_temp def forward(self, student_output, teacher_output, iteration): @@ -79,6 +88,7 @@ def update_center(self, teacher_output): batch_center = batch_center / (len(teacher_output) * torch.distributed.get_world_size()) self.center = self.center * self.center_momentum + batch_center * (1 - self.center_momentum) + class DINOHead(torch.nn.Module): def __init__(self, in_dim, out_dim, norm_last_layer=True, nlayers=3): super().__init__() @@ -97,14 +107,16 @@ def __init__(self, in_dim, out_dim, norm_last_layer=True, nlayers=3): layers.append(torch.nn.Linear(hidden_dim, bottleneck_dim)) self.mlp = torch.nn.Sequential(*layers) self.apply(self._init_weights) - self.last_layer = torch.nn.utils.weight_norm(torch.nn.Linear(bottleneck_dim, out_dim, bias=False)) + self.last_layer = torch.nn.utils.weight_norm( + torch.nn.Linear(bottleneck_dim, out_dim, bias=False) + ) self.last_layer.weight_g.data.fill_(1) if norm_last_layer: self.last_layer.weight_g.requires_grad = False def _init_weights(self, m): if isinstance(m, torch.nn.Linear): - trunc_normal_(m.weight, std=.02) + trunc_normal_(m.weight, std=0.02) if isinstance(m, torch.nn.Linear) and m.bias is not None: torch.nn.init.constant_(m.bias, 0) @@ -116,7 +128,6 @@ def forward(self, x): class MultiCropWrapper(MegatronModule): - """ Perform forward pass separately on each resolution input. The inputs corresponding to a single resolution are clubbed and single @@ -125,10 +136,11 @@ class MultiCropWrapper(MegatronModule): concatenate all the output features and run the head forward on these concatenated features. """ + def __init__(self, backbone, head): super(MultiCropWrapper, self).__init__() # disable layers dedicated to ImageNet labels classification - #backbone.fc, backbone.head = torch.nn.Identity(), torch.nn.Identity() + # backbone.fc, backbone.head = torch.nn.Identity(), torch.nn.Identity() self.backbone = backbone self.head = head @@ -136,14 +148,17 @@ def forward(self, x): # convert to list if not isinstance(x, list): x = [x] - idx_crops = torch.cumsum(torch.unique_consecutive( - torch.tensor([inp.shape[-1] for inp in x]), - return_counts=True, - )[1], 0) + idx_crops = torch.cumsum( + torch.unique_consecutive( + torch.tensor([inp.shape[-1] for inp in x]), + return_counts=True, + )[1], + 0, + ) start_idx = 0 for end_idx in idx_crops: - _out = self.backbone(torch.cat(x[start_idx: end_idx])) + _out = self.backbone(torch.cat(x[start_idx:end_idx])) if start_idx == 0: output = _out else: @@ -156,17 +171,18 @@ def forward(self, x): return output -def cosine_scheduler(base_value, final_value, epochs, niter_per_ep, - warmup_epochs=0, start_warmup_value=0): +def cosine_scheduler( + base_value, final_value, epochs, niter_per_ep, warmup_epochs=0, start_warmup_value=0 +): warmup_schedule = np.array([]) warmup_iters = warmup_epochs * niter_per_ep if warmup_epochs > 0: - warmup_schedule = \ - np.linspace(start_warmup_value, base_value, warmup_iters) + warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters) iters = np.arange(epochs * niter_per_ep - warmup_iters) - schedule = final_value + 0.5 * (base_value - final_value) \ - * (1 + np.cos(np.pi * iters / len(iters))) + schedule = final_value + 0.5 * (base_value - final_value) * ( + 1 + np.cos(np.pi * iters / len(iters)) + ) schedule = np.concatenate((warmup_schedule, schedule)) assert len(schedule) == epochs * niter_per_ep @@ -176,43 +192,43 @@ def cosine_scheduler(base_value, final_value, epochs, niter_per_ep, def get_student_backbone_and_num_features(config, pre_process=True, post_process=True): args = get_args() - if args.vision_backbone_type == 'vit': - student = VitBackbone(config, - pre_process=pre_process, - post_process=post_process, - drop_path_rate=0.1, - single_token_output=True) + if args.vision_backbone_type == "vit": + student = VitBackbone( + config, + pre_process=pre_process, + post_process=post_process, + drop_path_rate=0.1, + single_token_output=True, + ) num_features = args.hidden_size - elif args.vision_backbone_type == 'mit': + elif args.vision_backbone_type == "mit": student = mit_b5_avg(drop_path_rate=0.1) num_features = 512 - elif args.vision_backbone_type == 'swin': + elif args.vision_backbone_type == "swin": student = get_swin() num_features = student.num_features else: - raise Exception('{} vision backbone is not supported.'.format( - args.vision_backbone_type)) + raise Exception("{} vision backbone is not supported.".format(args.vision_backbone_type)) return student, num_features + def get_teacher_backbone_and_num_features(config, pre_process=True, post_process=True): args = get_args() - if args.vision_backbone_type == 'vit': - teacher = VitBackbone(config, - pre_process=pre_process, - post_process=post_process, - single_token_output=True) + if args.vision_backbone_type == "vit": + teacher = VitBackbone( + config, pre_process=pre_process, post_process=post_process, single_token_output=True + ) num_features = args.hidden_size - elif args.vision_backbone_type == 'mit': + elif args.vision_backbone_type == "mit": teacher = mit_b5_avg(drop_path_rate=0.0) num_features = 512 - elif args.vision_backbone_type == 'swin': + elif args.vision_backbone_type == "swin": teacher = get_swin(is_teacher=True) num_features = teacher.num_features else: - raise Exception('{} vision backbone is not supported.'.format( - args.vision_backbone_type)) + raise Exception("{} vision backbone is not supported.".format(args.vision_backbone_type)) return teacher, num_features @@ -236,27 +252,23 @@ def __init__(self, config, pre_process=True, post_process=True): self.post_process = post_process self.momentum_teacher = 0.996 - student_backbone, num_features = \ - get_student_backbone_and_num_features(config, pre_process, post_process) + student_backbone, num_features = get_student_backbone_and_num_features( + config, pre_process, post_process + ) self.student = MultiCropWrapper( student_backbone, - DINOHead(num_features, self.out_dim, - norm_last_layer=args.dino_norm_last_layer) + DINOHead(num_features, self.out_dim, norm_last_layer=args.dino_norm_last_layer), ) self.momentum_schedule = cosine_scheduler( - self.momentum_teacher, 1, - args.train_iters // args.iter_per_epoch, - args.iter_per_epoch + self.momentum_teacher, 1, args.train_iters // args.iter_per_epoch, args.iter_per_epoch ) - teacher_backbone, num_features = \ - get_teacher_backbone_and_num_features(config, pre_process, post_process) - self.teacher = MultiCropWrapper( - teacher_backbone, - DINOHead(num_features, self.out_dim) + teacher_backbone, num_features = get_teacher_backbone_and_num_features( + config, pre_process, post_process ) + self.teacher = MultiCropWrapper(teacher_backbone, DINOHead(num_features, self.out_dim)) self.teacher.load_state_dict(self.student.state_dict()) for p in self.teacher.parameters(): @@ -288,4 +300,3 @@ def update_momentum(self, iteration): m = self.momentum_schedule[iteration] for param_q, param_k in zip(self.student.parameters(), self.teacher.parameters()): param_k.data.mul_(m).add_((1 - m) * param_q.detach().data) - diff --git a/flagscale/train/megatron/legacy/model/vision/esvit_swin_backbone.py b/flagscale/train/megatron/legacy/model/vision/esvit_swin_backbone.py index c7c67b40c6..07862f8ee2 100644 --- a/flagscale/train/megatron/legacy/model/vision/esvit_swin_backbone.py +++ b/flagscale/train/megatron/legacy/model/vision/esvit_swin_backbone.py @@ -23,8 +23,9 @@ class Mlp(nn.Module): - def __init__(self, in_features, hidden_features=None, - out_features=None, act_layer=nn.GELU, drop=0.): + def __init__( + self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0 + ): super(Mlp, self).__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features @@ -85,18 +86,27 @@ class WindowAttention(nn.Module): proj_drop (float, optional): Dropout ratio of output. Default: 0.0 """ - def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.): - + def __init__( + self, + dim, + window_size, + num_heads, + qkv_bias=True, + qk_scale=None, + attn_drop=0.0, + proj_drop=0.0, + ): super(WindowAttention, self).__init__() self.dim = dim self.window_size = window_size # Wh, Ww self.num_heads = num_heads head_dim = dim // num_heads - self.scale = qk_scale or head_dim ** -0.5 + self.scale = qk_scale or head_dim**-0.5 # define a parameter table of relative position bias self.relative_position_bias_table = nn.Parameter( - torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads) + ) # 2*Wh-1 * 2*Ww-1, nH # get pair-wise relative position index for each token inside the window coords_h = torch.arange(self.window_size[0]) @@ -116,7 +126,7 @@ def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, at self.proj = nn.Linear(dim, dim) self.proj_drop = nn.Dropout(proj_drop) - trunc_normal_(self.relative_position_bias_table, std=.02) + trunc_normal_(self.relative_position_bias_table, std=0.02) self.softmax = nn.Softmax(dim=-1) def forward(self, x, mask=None): @@ -126,20 +136,31 @@ def forward(self, x, mask=None): mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None """ B_, N, C = x.shape - qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + qkv = ( + self.qkv(x) + .reshape(B_, N, 3, self.num_heads, C // self.num_heads) + .permute(2, 0, 3, 1, 4) + ) q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) q = q * self.scale - attn = (q @ k.transpose(-2, -1)) - - relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( - self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH - relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + attn = q @ k.transpose(-2, -1) + + relative_position_bias = self.relative_position_bias_table[ + self.relative_position_index.view(-1) + ].view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1 + ) # Wh*Ww,Wh*Ww,nH + relative_position_bias = relative_position_bias.permute( + 2, 0, 1 + ).contiguous() # nH, Wh*Ww, Wh*Ww attn = attn + relative_position_bias.unsqueeze(0) if mask is not None: nW = mask.shape[0] - attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0).type(attn.type()) + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze( + 0 + ).type(attn.type()) attn = attn.view(-1, self.num_heads, N, N) attn = self.softmax(attn) else: @@ -154,7 +175,7 @@ def forward(self, x, mask=None): return x, attn_out def extra_repr(self) -> str: - return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}' + return f"dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}" def flops(self, N): # calculate flops for 1 window with token length of N @@ -194,9 +215,22 @@ class SwinTransformerBlock(nn.Module): norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm """ - def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0, - mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., - act_layer=nn.GELU, norm_layer=nn.LayerNorm): + def __init__( + self, + dim, + input_resolution, + num_heads, + window_size=7, + shift_size=0, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + ): super().__init__() self.dim = dim self.input_resolution = input_resolution @@ -212,46 +246,60 @@ def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0 self.norm1 = norm_layer(dim) self.attn = WindowAttention( - dim, window_size=(self.window_size, self.window_size), num_heads=num_heads, - qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) - - self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + dim, + window_size=(self.window_size, self.window_size), + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=drop, + ) + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() self.norm2 = norm_layer(dim) mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + self.mlp = Mlp( + in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop + ) self.H = input_resolution[0] self.W = input_resolution[1] self.attn_mask_dict = {} - def create_attn_mask(self, H, W): # calculate attention mask for SW-MSA Hp = int(np.ceil(H / self.window_size)) * self.window_size Wp = int(np.ceil(W / self.window_size)) * self.window_size img_mask = torch.zeros((1, Hp, Wp, 1)) # 1 Hp Wp 1 - h_slices = (slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None)) - w_slices = (slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None)) + h_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) + w_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) cnt = 0 for h in h_slices: for w in w_slices: img_mask[:, h, w, :] = cnt cnt += 1 - mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = window_partition( + img_mask, self.window_size + ) # nW, window_size, window_size, 1 mask_windows = mask_windows.view(-1, self.window_size * self.window_size) attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) - attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill( + attn_mask == 0, float(0.0) + ) return attn_mask - def forward(self, x): B, L, C = x.shape H = int(sqrt(L)) @@ -283,8 +331,12 @@ def forward(self, x): attn_mask = None # partition windows - x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C - x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + x_windows = window_partition( + shifted_x, self.window_size + ) # nW*B, window_size, window_size, C + x_windows = x_windows.view( + -1, self.window_size * self.window_size, C + ) # nW*B, window_size*window_size, C # W-MSA/SW-MSA attn_windows, attn = self.attn(x_windows, attn_mask) # nW*B, window_size*window_size, C @@ -311,8 +363,10 @@ def forward(self, x): return x, attn def extra_repr(self) -> str: - return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \ - f"window_size={self.window_size}, shift_size={self.shift_size} mlp_ratio={self.mlp_ratio}" + return ( + f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " + f"window_size={self.window_size}, shift_size={self.shift_size} mlp_ratio={self.mlp_ratio}" + ) def flops(self): flops = 0 @@ -345,7 +399,7 @@ def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm): self.norm = norm_layer(4 * dim) def forward(self, x): - """ Forward function. + """Forward function. Args: x: Input feature, tensor size (B, H*W, C). H, W: Spatial resolution of the input feature. @@ -373,7 +427,6 @@ def forward(self, x): return x - def extra_repr(self) -> str: return f"input_resolution={self.input_resolution}, dim={self.dim}" @@ -402,25 +455,46 @@ class BasicLayer(nn.Module): downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None """ - def __init__(self, dim, input_resolution, depth, num_heads, window_size, - mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., - drop_path=0., norm_layer=nn.LayerNorm, downsample=None): - + def __init__( + self, + dim, + input_resolution, + depth, + num_heads, + window_size, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + norm_layer=nn.LayerNorm, + downsample=None, + ): super().__init__() self.dim = dim self.input_resolution = input_resolution self.depth = depth - self.blocks = nn.ModuleList([ - SwinTransformerBlock(dim=dim, input_resolution=input_resolution, - num_heads=num_heads, window_size=window_size, - shift_size=0 if (i % 2 == 0) else window_size // 2, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop, attn_drop=attn_drop, - drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, - norm_layer=norm_layer) - for i in range(depth)]) + self.blocks = nn.ModuleList( + [ + SwinTransformerBlock( + dim=dim, + input_resolution=input_resolution, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer, + ) + for i in range(depth) + ] + ) if downsample is not None: self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer) else: @@ -451,7 +525,6 @@ def forward_with_attention(self, x): x = self.downsample(x) return x, attns - def extra_repr(self) -> str: return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}" @@ -465,8 +538,7 @@ def flops(self): class PatchEmbed(nn.Module): - """ Image to Patch Embedding - """ + """Image to Patch Embedding""" def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None): super().__init__() @@ -495,7 +567,6 @@ def forward(self, x): x = self.norm(x) return x - def flops(self): Ho, Wo = self.patches_resolution flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1]) @@ -503,8 +574,9 @@ def flops(self): flops += Ho * Wo * self.embed_dim return flops + class SwinTransformer(nn.Module): - r""" Swin Transformer + r"""Swin Transformer A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - https://arxiv.org/pdf/2103.14030 Args: @@ -527,11 +599,27 @@ class SwinTransformer(nn.Module): patch_norm (bool): If True, add normalization after patch embedding. """ - def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, - embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], - window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None, - drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1, - norm_layer=nn.LayerNorm, ape=False, patch_norm=True, **kwargs): + def __init__( + self, + img_size=224, + patch_size=4, + in_chans=3, + num_classes=1000, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.1, + norm_layer=nn.LayerNorm, + ape=False, + patch_norm=True, + **kwargs, + ): super().__init__() self.num_classes = num_classes @@ -543,33 +631,45 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, self.mlp_ratio = mlp_ratio self.patch_embed = PatchEmbed( - img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, - norm_layer=norm_layer if self.patch_norm else None) + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None, + ) num_patches = self.patch_embed.num_patches patches_resolution = self.patch_embed.patches_resolution self.patches_resolution = patches_resolution if self.ape: self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) - trunc_normal_(self.absolute_pos_embed, std=.02) + trunc_normal_(self.absolute_pos_embed, std=0.02) self.pos_drop = nn.Dropout(p=drop_rate) - dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, sum(depths)) + ] # stochastic depth decay rule self.layers = nn.ModuleList() for i_layer in range(self.num_layers): - layer = BasicLayer(dim=int(embed_dim * 2 ** i_layer), - input_resolution=(patches_resolution[0] // (2 ** i_layer), - patches_resolution[1] // (2 ** i_layer)), - depth=depths[i_layer], - num_heads=num_heads[i_layer], - window_size=window_size, - mlp_ratio=self.mlp_ratio, - qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, - drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], - norm_layer=norm_layer, - downsample=PatchMerging if (i_layer < self.num_layers - 1) else None) + layer = BasicLayer( + dim=int(embed_dim * 2**i_layer), + input_resolution=( + patches_resolution[0] // (2**i_layer), + patches_resolution[1] // (2**i_layer), + ), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=self.mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])], + norm_layer=norm_layer, + downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + ) self.layers.append(layer) self.norm = norm_layer(self.num_features) @@ -579,7 +679,7 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, def _init_weights(self, m): if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=.02) + trunc_normal_(m.weight, std=0.02) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): @@ -588,12 +688,12 @@ def _init_weights(self, m): @torch.jit.ignore def no_weight_decay(self): - return {'absolute_pos_embed'} + return {"absolute_pos_embed"} @torch.jit.ignore def no_weight_decay_keywords(self): # todo: to be implemented - return {'relative_position_bias_table'} + return {"relative_position_bias_table"} def forward(self, x): x = self.patch_embed(x) @@ -610,7 +710,6 @@ def forward(self, x): return x - def forward_feature_maps(self, x): x = self.patch_embed(x) if self.ape: @@ -626,23 +725,20 @@ def forward_feature_maps(self, x): return x, x_grid - def forward_selfattention(self, x, n=1): # n=1 return the last layer attn map; otherwise return attn maps in all layers - x = self.patch_embed(x) if self.ape: x = x + self.absolute_pos_embed x = self.pos_drop(x) - if n==1: + if n == 1: return self.forward_last_selfattention(x) else: return self.forward_all_selfattention(x) def forward_last_selfattention(self, x): - for i, layer in enumerate(self.layers): if i < len(self.layers) - 1: x = layer(x) @@ -659,9 +755,7 @@ def forward_all_selfattention(self, x): return attn_out - def forward_return_n_last_blocks(self, x, n=1, return_patch_avgpool=False, depth=[]): - num_blks = sum(depth) start_idx = num_blks - n @@ -673,7 +767,6 @@ def forward_return_n_last_blocks(self, x, n=1, return_patch_avgpool=False, depth start_blk = start_idx - sum_cur sum_cur = sum_cur_new - x = self.patch_embed(x) if self.ape: x = x + self.absolute_pos_embed @@ -688,20 +781,17 @@ def forward_return_n_last_blocks(self, x, n=1, return_patch_avgpool=False, depth if i >= start_stage: for x_ in fea[start_blk:]: - - if i == len(self.layers)-1: # use the norm in the last stage + if i == len(self.layers) - 1: # use the norm in the last stage x_ = self.norm(x_) - x_avg = torch.flatten(self.avgpool(x_.transpose(1, 2)), 1) # B C - # print(f'Stage {i}, x_avg {x_avg.shape}') + x_avg = torch.flatten(self.avgpool(x_.transpose(1, 2)), 1) # B C + # print(f'Stage {i}, x_avg {x_avg.shape}') output.append(x_avg) start_blk = 0 return torch.cat(output, dim=-1) - - def flops(self): flops = 0 flops += self.patch_embed.flops() @@ -709,33 +799,35 @@ def flops(self): flops += layer.flops() if dist.get_rank() == 0: print(f"GFLOPs layer_{i}: {layer.flops() / 1e9}") - flops += self.num_features * self.patches_resolution[0] * self.patches_resolution[1] // (2 ** self.num_layers) + flops += ( + self.num_features + * self.patches_resolution[0] + * self.patches_resolution[1] + // (2**self.num_layers) + ) flops += self.num_features * self.num_classes return flops - def init_weights(self, pretrained='', pretrained_layers=[], verbose=True): + def init_weights(self, pretrained="", pretrained_layers=[], verbose=True): if os.path.isfile(pretrained): - pretrained_dict = torch.load(pretrained, map_location='cpu') - logging.info(f'=> loading pretrained model {pretrained}') + pretrained_dict = torch.load(pretrained, map_location="cpu") + logging.info(f"=> loading pretrained model {pretrained}") model_dict = self.state_dict() - pretrained_dict = { - k: v for k, v in pretrained_dict.items() - if k in model_dict.keys() - } + pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict.keys()} need_init_state_dict = {} for k, v in pretrained_dict.items(): need_init = ( - k.split('.')[0] in pretrained_layers - or pretrained_layers[0] == '*' - or 'relative_position_index' not in k - or 'attn_mask' not in k + k.split(".")[0] in pretrained_layers + or pretrained_layers[0] == "*" + or "relative_position_index" not in k + or "attn_mask" not in k ) if need_init: if verbose: - logging.info(f'=> init {k} from {pretrained}') + logging.info(f"=> init {k} from {pretrained}") - if 'relative_position_bias_table' in k and v.size() != model_dict[k].size(): + if "relative_position_bias_table" in k and v.size() != model_dict[k].size(): relative_position_bias_table_pretrained = v relative_position_bias_table_current = model_dict[k] L1, nH1 = relative_position_bias_table_pretrained.size() @@ -745,18 +837,26 @@ def init_weights(self, pretrained='', pretrained_layers=[], verbose=True): else: if L1 != L2: logging.info( - '=> load_pretrained: resized variant: {} to {}' - .format((L1, nH1), (L2, nH2)) + "=> load_pretrained: resized variant: {} to {}".format( + (L1, nH1), (L2, nH2) + ) + ) + S1 = int(L1**0.5) + S2 = int(L2**0.5) + relative_position_bias_table_pretrained_resized = ( + torch.nn.functional.interpolate( + relative_position_bias_table_pretrained.permute(1, 0).view( + 1, nH1, S1, S1 + ), + size=(S2, S2), + mode="bicubic", + ) ) - S1 = int(L1 ** 0.5) - S2 = int(L2 ** 0.5) - relative_position_bias_table_pretrained_resized = torch.nn.functional.interpolate( - relative_position_bias_table_pretrained.permute(1, 0).view(1, nH1, S1, S1), - size=(S2, S2), - mode='bicubic') - v = relative_position_bias_table_pretrained_resized.view(nH2, L2).permute(1, 0) - - if 'absolute_pos_embed' in k and v.size() != model_dict[k].size(): + v = relative_position_bias_table_pretrained_resized.view( + nH2, L2 + ).permute(1, 0) + + if "absolute_pos_embed" in k and v.size() != model_dict[k].size(): absolute_pos_embed_pretrained = v absolute_pos_embed_current = model_dict[k] _, L1, C1 = absolute_pos_embed_pretrained.size() @@ -766,16 +866,26 @@ def init_weights(self, pretrained='', pretrained_layers=[], verbose=True): else: if L1 != L2: logging.info( - '=> load_pretrained: resized variant: {} to {}' - .format((1, L1, C1), (1, L2, C2)) + "=> load_pretrained: resized variant: {} to {}".format( + (1, L1, C1), (1, L2, C2) + ) + ) + S1 = int(L1**0.5) + S2 = int(L2**0.5) + absolute_pos_embed_pretrained = ( + absolute_pos_embed_pretrained.reshape(-1, S1, S1, C1) ) - S1 = int(L1 ** 0.5) - S2 = int(L2 ** 0.5) - absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.reshape(-1, S1, S1, C1) - absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.permute(0, 3, 1, 2) - absolute_pos_embed_pretrained_resized = torch.nn.functional.interpolate( - absolute_pos_embed_pretrained, size=(S2, S2), mode='bicubic') - v = absolute_pos_embed_pretrained_resized.permute(0, 2, 3, 1).flatten(1, 2) + absolute_pos_embed_pretrained = ( + absolute_pos_embed_pretrained.permute(0, 3, 1, 2) + ) + absolute_pos_embed_pretrained_resized = ( + torch.nn.functional.interpolate( + absolute_pos_embed_pretrained, size=(S2, S2), mode="bicubic" + ) + ) + v = absolute_pos_embed_pretrained_resized.permute( + 0, 2, 3, 1 + ).flatten(1, 2) need_init_state_dict[k] = v self.load_state_dict(need_init_state_dict, strict=False) @@ -783,27 +893,21 @@ def init_weights(self, pretrained='', pretrained_layers=[], verbose=True): def freeze_pretrained_layers(self, frozen_layers=[]): for name, module in self.named_modules(): if ( - name.split('.')[0] in frozen_layers - or '.'.join(name.split('.')[0:2]) in frozen_layers - or (len(frozen_layers) > 0 and frozen_layers[0] == '*') + name.split(".")[0] in frozen_layers + or ".".join(name.split(".")[0:2]) in frozen_layers + or (len(frozen_layers) > 0 and frozen_layers[0] == "*") ): for _name, param in module.named_parameters(): param.requires_grad = False - logging.info( - '=> set param {} requires grad to False' - .format(name) - ) + logging.info("=> set param {} requires grad to False".format(name)) for name, param in self.named_parameters(): if ( - name.split('.')[0] in frozen_layers - or (len(frozen_layers) > 0 and frozen_layers[0] == '*') - and param.requires_grad is True + name.split(".")[0] in frozen_layers + or (len(frozen_layers) > 0 and frozen_layers[0] == "*") + and param.requires_grad is True ): param.requires_grad = False - logging.info( - '=> set param {} requires grad to False' - .format(name) - ) + logging.info("=> set param {} requires grad to False".format(name)) return self @@ -815,7 +919,7 @@ def get_swin(is_teacher=False): depths = [2, 2, 6, 2] num_heads = [3, 6, 12, 24] drop_path_rate = 0.1 - elif args.swin_backbone_type == 'h3': + elif args.swin_backbone_type == "h3": embed_dim = 384 depths = [2, 2, 18, 2] num_heads = [6, 12, 24, 48] @@ -846,4 +950,3 @@ def get_swin(is_teacher=False): ) return swin - diff --git a/flagscale/train/megatron/legacy/model/vision/inpainting.py b/flagscale/train/megatron/legacy/model/vision/inpainting.py index f71f5e3209..4f81c73bd1 100644 --- a/flagscale/train/megatron/legacy/model/vision/inpainting.py +++ b/flagscale/train/megatron/legacy/model/vision/inpainting.py @@ -17,7 +17,6 @@ class VitInpaintingModel(MegatronModule): - def __init__(self, config, pre_process=True, post_process=True): super(VitInpaintingModel, self).__init__() args = get_args() @@ -40,29 +39,26 @@ def __init__(self, config, pre_process=True, post_process=True): if self.post_process: self.linear_decoder = get_linear_layer( - self.hidden_size, - self.backbone.flatten_dim, - torch.nn.init.zeros_ + self.hidden_size, self.backbone.flatten_dim, torch.nn.init.zeros_ ) def set_input_tensor(self, input_tensor): self.backbone.set_input_tensor(input_tensor) def forward(self, input): - hidden_states = self.backbone(input) if not self.post_process: return hidden_states decoded_output = self.linear_decoder(hidden_states) output = einops.rearrange( - decoded_output, - "b (h w) (p1 p2 c) -> b c (h p1) (w p2)", - p1=self.patch_dim, - p2=self.patch_dim, - h=self.img_h//self.patch_dim, - w=self.img_w//self.patch_dim, - ) + decoded_output, + "b (h w) (p1 p2 c) -> b c (h p1) (w p2)", + p1=self.patch_dim, + p2=self.patch_dim, + h=self.img_h // self.patch_dim, + w=self.img_w // self.patch_dim, + ) return output @@ -71,6 +67,7 @@ class MLP(torch.nn.Module): """ Linear Embedding """ + def __init__(self, input_dim=2048, embed_dim=768): super().__init__() self.proj = torch.nn.Linear(input_dim, embed_dim) @@ -106,7 +103,9 @@ def __init__(self, pre_process=True, post_process=True): self.linear_c2 = MLP(input_dim=c2_in_channels, embed_dim=self.embedding_dim) self.linear_c1 = MLP(input_dim=c1_in_channels, embed_dim=self.embedding_dim) - self.conv_fuse = torch.nn.Conv2d(self.embedding_dim*4, self.embedding_dim, 1, 1, bias=False) + self.conv_fuse = torch.nn.Conv2d( + self.embedding_dim * 4, self.embedding_dim, 1, 1, bias=False + ) self.norm = apex.parallel.SyncBatchNorm(self.embedding_dim) self.dropout = torch.nn.Dropout2d(0.1) @@ -121,13 +120,13 @@ def forward(self, input): n, _, h, w = c4.shape _c4 = self.linear_c4(c4).permute(0, 2, 1).reshape(n, -1, c4.shape[2], c4.shape[3]) - _c4 = resize(_c4, size=c1.size()[2:], mode='bilinear', align_corners=False) + _c4 = resize(_c4, size=c1.size()[2:], mode="bilinear", align_corners=False) _c3 = self.linear_c3(c3).permute(0, 2, 1).reshape(n, -1, c3.shape[2], c3.shape[3]) - _c3 = resize(_c3, size=c1.size()[2:], mode='bilinear', align_corners=False) + _c3 = resize(_c3, size=c1.size()[2:], mode="bilinear", align_corners=False) _c2 = self.linear_c2(c2).permute(0, 2, 1).reshape(n, -1, c2.shape[2], c2.shape[3]) - _c2 = resize(_c2, size=c1.size()[2:], mode='bilinear', align_corners=False) + _c2 = resize(_c2, size=c1.size()[2:], mode="bilinear", align_corners=False) _c1 = self.linear_c1(c1).permute(0, 2, 1).reshape(n, -1, c1.shape[2], c1.shape[3]) @@ -145,8 +144,8 @@ def forward(self, input): "b (c p1 p2) h w -> b c (h p1) (w p2)", p1=self.patch_dim, p2=self.patch_dim, - h=self.img_h//self.patch_dim, - w=self.img_w//self.patch_dim, + h=self.img_h // self.patch_dim, + w=self.img_w // self.patch_dim, ) return output diff --git a/flagscale/train/megatron/legacy/model/vision/knn_monitor.py b/flagscale/train/megatron/legacy/model/vision/knn_monitor.py index 53e7495085..3ecf6c5082 100644 --- a/flagscale/train/megatron/legacy/model/vision/knn_monitor.py +++ b/flagscale/train/megatron/legacy/model/vision/knn_monitor.py @@ -9,8 +9,10 @@ _FEATURE_BANK = None from megatron.plugin.platform import get_platform + cur_platform = get_platform() + def build_data_loader(dataset, drop_last=True, shuffle=False): """Data loader. Note that batch-size is the local (per GPU) batch-size.""" # Sampler. @@ -20,8 +22,7 @@ def build_data_loader(dataset, drop_last=True, shuffle=False): world_size = mpu.get_data_parallel_world_size() rank = mpu.get_data_parallel_rank() sampler = torch.utils.data.distributed.DistributedSampler( - dataset, num_replicas=world_size, rank=rank, - drop_last=drop_last, shuffle=shuffle + dataset, num_replicas=world_size, rank=rank, drop_last=drop_last, shuffle=shuffle ) # Data loader. Note that batch size is the per GPU batch size. @@ -46,11 +47,11 @@ def compute_feature_bank(model): train_ds = ImageFolder( root=args.data_path[0], transform=ClassificationTransform((args.img_h, args.img_w), train=False), - data_per_class_fraction=1.0 + data_per_class_fraction=1.0, ) classes = len(train_ds.classes) dataloader = build_data_loader(train_ds) - + for m in model: m.eval() @@ -62,7 +63,7 @@ def compute_feature_bank(model): feature = F.normalize(teacher_feature.float(), dim=1) feature_bank.append(feature) feature_label.append(labels) - + for m in model: m.train() @@ -70,20 +71,17 @@ def compute_feature_bank(model): feature_bank = torch.cat(feature_bank, dim=0).contiguous() feature_label = torch.cat(feature_label, dim=0).contiguous() - feature_banks = [torch.zeros_like(feature_bank) - for i in range(mpu.get_data_parallel_world_size())] - torch.distributed.all_gather(feature_banks, - feature_bank, - group=mpu.get_data_parallel_group()) + feature_banks = [ + torch.zeros_like(feature_bank) for i in range(mpu.get_data_parallel_world_size()) + ] + torch.distributed.all_gather(feature_banks, feature_bank, group=mpu.get_data_parallel_group()) - assert torch.all(torch.eq(feature_banks[mpu.get_data_parallel_rank()], - feature_bank)) + assert torch.all(torch.eq(feature_banks[mpu.get_data_parallel_rank()], feature_bank)) - feature_labels = [torch.zeros_like(feature_label) - for i in range(mpu.get_data_parallel_world_size())] - torch.distributed.all_gather(feature_labels, - feature_label, - group=mpu.get_data_parallel_group()) + feature_labels = [ + torch.zeros_like(feature_label) for i in range(mpu.get_data_parallel_world_size()) + ] + torch.distributed.all_gather(feature_labels, feature_label, group=mpu.get_data_parallel_group()) # [D, N] feature_banks = torch.cat(feature_banks, dim=0).t().contiguous() @@ -110,23 +108,17 @@ def knn_predict(feature, feature_bank, feature_labels, classes, knn_k, knn_t): # [B, K] sim_weight, sim_indices = sim_matrix.topk(k=knn_k, dim=-1) # [B, K] - sim_labels = torch.gather(feature_labels.expand(feature.size(0), -1), - dim=-1, - index=sim_indices) + sim_labels = torch.gather(feature_labels.expand(feature.size(0), -1), dim=-1, index=sim_indices) sim_weight = (sim_weight / knn_t).exp() # counts for each class - one_hot_label = torch.zeros(feature.size(0) * knn_k, - classes, - device=sim_labels.device) + one_hot_label = torch.zeros(feature.size(0) * knn_k, classes, device=sim_labels.device) # [B*K, C] - one_hot_label = one_hot_label.scatter(dim=-1, - index=sim_labels.view(-1, 1), - value=1.0) + one_hot_label = one_hot_label.scatter(dim=-1, index=sim_labels.view(-1, 1), value=1.0) # weighted score ---> [B, C] pred_scores = torch.sum( - one_hot_label.view(feature.size(0), -1, classes) * sim_weight.unsqueeze(dim=-1), - dim=1) + one_hot_label.view(feature.size(0), -1, classes) * sim_weight.unsqueeze(dim=-1), dim=1 + ) pred_labels = pred_scores.argsort(dim=-1, descending=True) return pred_labels diff --git a/flagscale/train/megatron/legacy/model/vision/mit_backbone.py b/flagscale/train/megatron/legacy/model/vision/mit_backbone.py index 3ca2303c30..b8627e0fad 100644 --- a/flagscale/train/megatron/legacy/model/vision/mit_backbone.py +++ b/flagscale/train/megatron/legacy/model/vision/mit_backbone.py @@ -11,12 +11,9 @@ class Mlp(nn.Module): - def __init__(self, - in_features, - hidden_features=None, - out_features=None, - act_layer=nn.GELU, - drop=0.): + def __init__( + self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0 + ): super().__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features @@ -30,7 +27,7 @@ def __init__(self, def _init_weights(self, m): if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=.02) + trunc_normal_(m.weight, std=0.02) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): @@ -54,21 +51,23 @@ def forward(self, x, H, W): class Attention(nn.Module): - def __init__(self, - dim, - num_heads=8, - qkv_bias=False, - qk_scale=None, - attn_drop=0., - proj_drop=0., - sr_ratio=1): + def __init__( + self, + dim, + num_heads=8, + qkv_bias=False, + qk_scale=None, + attn_drop=0.0, + proj_drop=0.0, + sr_ratio=1, + ): super().__init__() assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}." self.dim = dim self.num_heads = num_heads head_dim = dim // num_heads - self.scale = qk_scale or head_dim ** -0.5 + self.scale = qk_scale or head_dim**-0.5 self.q = nn.Linear(dim, dim, bias=qkv_bias) self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias) @@ -85,7 +84,7 @@ def __init__(self, def _init_weights(self, m): if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=.02) + trunc_normal_(m.weight, std=0.02) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): @@ -106,9 +105,17 @@ def forward(self, x, H, W): x_ = x.permute(0, 2, 1).reshape(B, C, H, W) x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1) x_ = self.norm(x_) - kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + kv = ( + self.kv(x_) + .reshape(B, -1, 2, self.num_heads, C // self.num_heads) + .permute(2, 0, 3, 1, 4) + ) else: - kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + kv = ( + self.kv(x) + .reshape(B, -1, 2, self.num_heads, C // self.num_heads) + .permute(2, 0, 3, 1, 4) + ) k, v = kv[0], kv[1] attn = (q @ k.transpose(-2, -1)) * self.scale @@ -123,26 +130,44 @@ def forward(self, x, H, W): class Block(nn.Module): - - def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., - drop_path=0., act_layer=nn.GELU, norm_layer=LayerNorm, sr_ratio=1): + def __init__( + self, + dim, + num_heads, + mlp_ratio=4.0, + qkv_bias=False, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + norm_layer=LayerNorm, + sr_ratio=1, + ): super().__init__() self.norm1 = norm_layer(dim) self.attn = Attention( dim, - num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, - attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio) + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=drop, + sr_ratio=sr_ratio, + ) # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here - self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() self.norm2 = norm_layer(dim) mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + self.mlp = Mlp( + in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop + ) self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=.02) + trunc_normal_(m.weight, std=0.02) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): @@ -163,23 +188,27 @@ def forward(self, x, H, W): class OverlapPatchEmbed(nn.Module): - """ Image to Patch Embedding - """ + """Image to Patch Embedding""" def __init__(self, img_size=224, patch_size=7, stride=4, in_chans=3, embed_dim=768): super().__init__() img_size = (img_size, img_size) patch_size = (patch_size, patch_size) - self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride, - padding=(patch_size[0] // 2, patch_size[1] // 2)) + self.proj = nn.Conv2d( + in_chans, + embed_dim, + kernel_size=patch_size, + stride=stride, + padding=(patch_size[0] // 2, patch_size[1] // 2), + ) self.norm = LayerNorm(embed_dim) self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=.02) + trunc_normal_(m.weight, std=0.02) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): @@ -202,64 +231,145 @@ def forward(self, x): class MixVisionTransformer(nn.Module): - def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dims=[64, 128, 256, 512], - num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0., - attn_drop_rate=0., drop_path_rate=0., norm_layer=LayerNorm, - depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1], output_avg=False): + def __init__( + self, + img_size=224, + patch_size=16, + in_chans=3, + num_classes=1000, + embed_dims=[64, 128, 256, 512], + num_heads=[1, 2, 4, 8], + mlp_ratios=[4, 4, 4, 4], + qkv_bias=False, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.0, + norm_layer=LayerNorm, + depths=[3, 4, 6, 3], + sr_ratios=[8, 4, 2, 1], + output_avg=False, + ): super().__init__() self.num_classes = num_classes self.depths = depths self.output_avg = output_avg # patch_embed - self.patch_embed1 = OverlapPatchEmbed(img_size=img_size, patch_size=7, stride=4, in_chans=in_chans, - embed_dim=embed_dims[0]) - self.patch_embed2 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_chans=embed_dims[0], - embed_dim=embed_dims[1]) - self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_chans=embed_dims[1], - embed_dim=embed_dims[2]) - self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 16, patch_size=3, stride=2, in_chans=embed_dims[2], - embed_dim=embed_dims[3]) + self.patch_embed1 = OverlapPatchEmbed( + img_size=img_size, patch_size=7, stride=4, in_chans=in_chans, embed_dim=embed_dims[0] + ) + self.patch_embed2 = OverlapPatchEmbed( + img_size=img_size // 4, + patch_size=3, + stride=2, + in_chans=embed_dims[0], + embed_dim=embed_dims[1], + ) + self.patch_embed3 = OverlapPatchEmbed( + img_size=img_size // 8, + patch_size=3, + stride=2, + in_chans=embed_dims[1], + embed_dim=embed_dims[2], + ) + self.patch_embed4 = OverlapPatchEmbed( + img_size=img_size // 16, + patch_size=3, + stride=2, + in_chans=embed_dims[2], + embed_dim=embed_dims[3], + ) # transformer encoder - dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, sum(depths)) + ] # stochastic depth decay rule cur = 0 - self.block1 = nn.ModuleList([Block( - dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=mlp_ratios[0], qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, - sr_ratio=sr_ratios[0]) - for i in range(depths[0])]) + self.block1 = nn.ModuleList( + [ + Block( + dim=embed_dims[0], + num_heads=num_heads[0], + mlp_ratio=mlp_ratios[0], + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[cur + i], + norm_layer=norm_layer, + sr_ratio=sr_ratios[0], + ) + for i in range(depths[0]) + ] + ) self.norm1 = norm_layer(embed_dims[0]) cur += depths[0] - self.block2 = nn.ModuleList([Block( - dim=embed_dims[1], num_heads=num_heads[1], mlp_ratio=mlp_ratios[1], qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, - sr_ratio=sr_ratios[1]) - for i in range(depths[1])]) + self.block2 = nn.ModuleList( + [ + Block( + dim=embed_dims[1], + num_heads=num_heads[1], + mlp_ratio=mlp_ratios[1], + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[cur + i], + norm_layer=norm_layer, + sr_ratio=sr_ratios[1], + ) + for i in range(depths[1]) + ] + ) self.norm2 = norm_layer(embed_dims[1]) cur += depths[1] - self.block3 = nn.ModuleList([Block( - dim=embed_dims[2], num_heads=num_heads[2], mlp_ratio=mlp_ratios[2], qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, - sr_ratio=sr_ratios[2]) - for i in range(depths[2])]) + self.block3 = nn.ModuleList( + [ + Block( + dim=embed_dims[2], + num_heads=num_heads[2], + mlp_ratio=mlp_ratios[2], + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[cur + i], + norm_layer=norm_layer, + sr_ratio=sr_ratios[2], + ) + for i in range(depths[2]) + ] + ) self.norm3 = norm_layer(embed_dims[2]) cur += depths[2] - self.block4 = nn.ModuleList([Block( - dim=embed_dims[3], num_heads=num_heads[3], mlp_ratio=mlp_ratios[3], qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer, - sr_ratio=sr_ratios[3]) - for i in range(depths[3])]) + self.block4 = nn.ModuleList( + [ + Block( + dim=embed_dims[3], + num_heads=num_heads[3], + mlp_ratio=mlp_ratios[3], + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[cur + i], + norm_layer=norm_layer, + sr_ratio=sr_ratios[3], + ) + for i in range(depths[3]) + ] + ) self.norm4 = norm_layer(embed_dims[3]) self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=.02) + trunc_normal_(m.weight, std=0.02) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): @@ -334,7 +444,7 @@ def forward_features(self, x): def forward(self, x): x = self.forward_features(x) - + if self.output_avg: x = x[3].mean(dim=1) @@ -354,62 +464,132 @@ def forward(self, x, H, W): return x + class mit_b0(MixVisionTransformer): def __init__(self, **kwargs): super(mit_b0, self).__init__( - patch_size=4, embed_dims=[32, 64, 160, 256], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4], - qkv_bias=True, norm_layer=partial(LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1], - drop_rate=0.0, drop_path_rate=0.1) + patch_size=4, + embed_dims=[32, 64, 160, 256], + num_heads=[1, 2, 5, 8], + mlp_ratios=[4, 4, 4, 4], + qkv_bias=True, + norm_layer=partial(LayerNorm, eps=1e-6), + depths=[2, 2, 2, 2], + sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, + drop_path_rate=0.1, + ) class mit_b1(MixVisionTransformer): def __init__(self, **kwargs): super(mit_b1, self).__init__( - patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4], - qkv_bias=True, norm_layer=partial(LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1], - drop_rate=0.0, drop_path_rate=0.1) + patch_size=4, + embed_dims=[64, 128, 320, 512], + num_heads=[1, 2, 5, 8], + mlp_ratios=[4, 4, 4, 4], + qkv_bias=True, + norm_layer=partial(LayerNorm, eps=1e-6), + depths=[2, 2, 2, 2], + sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, + drop_path_rate=0.1, + ) class mit_b2(MixVisionTransformer): def __init__(self, **kwargs): super(mit_b2, self).__init__( - patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4], - qkv_bias=True, norm_layer=partial(LayerNorm, eps=1e-6), depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1], - drop_rate=0.0, drop_path_rate=0.1) + patch_size=4, + embed_dims=[64, 128, 320, 512], + num_heads=[1, 2, 5, 8], + mlp_ratios=[4, 4, 4, 4], + qkv_bias=True, + norm_layer=partial(LayerNorm, eps=1e-6), + depths=[3, 4, 6, 3], + sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, + drop_path_rate=0.1, + ) + - class mit_b3(MixVisionTransformer): def __init__(self, **kwargs): super(mit_b3, self).__init__( - patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4], - qkv_bias=True, norm_layer=partial(LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1], - drop_rate=0.0, drop_path_rate=0.1) + patch_size=4, + embed_dims=[64, 128, 320, 512], + num_heads=[1, 2, 5, 8], + mlp_ratios=[4, 4, 4, 4], + qkv_bias=True, + norm_layer=partial(LayerNorm, eps=1e-6), + depths=[3, 4, 18, 3], + sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, + drop_path_rate=0.1, + ) + class mit_b3_avg(MixVisionTransformer): def __init__(self, drop_path_rate=0.1, **kwargs): super(mit_b3_avg, self).__init__( - patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4], - qkv_bias=True, norm_layer=partial(LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1], - drop_rate=0.0, drop_path_rate=drop_path_rate, output_avg=True) + patch_size=4, + embed_dims=[64, 128, 320, 512], + num_heads=[1, 2, 5, 8], + mlp_ratios=[4, 4, 4, 4], + qkv_bias=True, + norm_layer=partial(LayerNorm, eps=1e-6), + depths=[3, 4, 18, 3], + sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, + drop_path_rate=drop_path_rate, + output_avg=True, + ) + class mit_b4(MixVisionTransformer): def __init__(self, **kwargs): super(mit_b4, self).__init__( - patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4], - qkv_bias=True, norm_layer=partial(LayerNorm, eps=1e-6), depths=[3, 8, 27, 3], sr_ratios=[8, 4, 2, 1], - drop_rate=0.0, drop_path_rate=0.1) + patch_size=4, + embed_dims=[64, 128, 320, 512], + num_heads=[1, 2, 5, 8], + mlp_ratios=[4, 4, 4, 4], + qkv_bias=True, + norm_layer=partial(LayerNorm, eps=1e-6), + depths=[3, 8, 27, 3], + sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, + drop_path_rate=0.1, + ) + class mit_b5(MixVisionTransformer): def __init__(self, **kwargs): super(mit_b5, self).__init__( - patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4], - qkv_bias=True, norm_layer=partial(LayerNorm, eps=1e-6), depths=[3, 6, 40, 3], sr_ratios=[8, 4, 2, 1], - drop_rate=0.0, drop_path_rate=0.1) + patch_size=4, + embed_dims=[64, 128, 320, 512], + num_heads=[1, 2, 5, 8], + mlp_ratios=[4, 4, 4, 4], + qkv_bias=True, + norm_layer=partial(LayerNorm, eps=1e-6), + depths=[3, 6, 40, 3], + sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, + drop_path_rate=0.1, + ) + class mit_b5_avg(MixVisionTransformer): def __init__(self, drop_path_rate=0.1, **kwargs): super(mit_b5_avg, self).__init__( - patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4], - qkv_bias=True, norm_layer=partial(LayerNorm, eps=1e-6), depths=[3, 6, 40, 3], sr_ratios=[8, 4, 2, 1], - drop_rate=0.0, drop_path_rate=drop_path_rate, output_avg=True) - + patch_size=4, + embed_dims=[64, 128, 320, 512], + num_heads=[1, 2, 5, 8], + mlp_ratios=[4, 4, 4, 4], + qkv_bias=True, + norm_layer=partial(LayerNorm, eps=1e-6), + depths=[3, 6, 40, 3], + sr_ratios=[8, 4, 2, 1], + drop_rate=0.0, + drop_path_rate=drop_path_rate, + output_avg=True, + ) diff --git a/flagscale/train/megatron/legacy/model/vision/swin_backbone.py b/flagscale/train/megatron/legacy/model/vision/swin_backbone.py index 231802c8f2..724bbce6df 100644 --- a/flagscale/train/megatron/legacy/model/vision/swin_backbone.py +++ b/flagscale/train/megatron/legacy/model/vision/swin_backbone.py @@ -17,8 +17,9 @@ class Mlp(nn.Module): - def __init__(self, in_features, hidden_features=None, - out_features=None, act_layer=nn.GELU, drop=0.): + def __init__( + self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0 + ): super().__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features @@ -69,7 +70,7 @@ def window_reverse(windows, window_size, H, W): class WindowAttention(nn.Module): - r""" Window based multi-head self attention (W-MSA) module with relative position bias. + r"""Window based multi-head self attention (W-MSA) module with relative position bias. It supports both of shifted and non-shifted window. Args: @@ -82,18 +83,27 @@ class WindowAttention(nn.Module): proj_drop (float, optional): Dropout ratio of output. Default: 0.0 """ - def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.): - + def __init__( + self, + dim, + window_size, + num_heads, + qkv_bias=True, + qk_scale=None, + attn_drop=0.0, + proj_drop=0.0, + ): super().__init__() self.dim = dim self.window_size = window_size # Wh, Ww self.num_heads = num_heads head_dim = dim // num_heads - self.scale = qk_scale or head_dim ** -0.5 + self.scale = qk_scale or head_dim**-0.5 # define a parameter table of relative position bias self.relative_position_bias_table = nn.Parameter( - torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads) + ) # 2*Wh-1 * 2*Ww-1, nH # get pair-wise relative position index for each token inside the window coords_h = torch.arange(self.window_size[0]) @@ -113,7 +123,7 @@ def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, at self.proj = nn.Linear(dim, dim) self.proj_drop = nn.Dropout(proj_drop) - trunc_normal_(self.relative_position_bias_table, std=.02) + trunc_normal_(self.relative_position_bias_table, std=0.02) self.softmax = nn.Softmax(dim=-1) def forward(self, x, mask=None): @@ -123,15 +133,24 @@ def forward(self, x, mask=None): mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None """ B_, N, C = x.shape - qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + qkv = ( + self.qkv(x) + .reshape(B_, N, 3, self.num_heads, C // self.num_heads) + .permute(2, 0, 3, 1, 4) + ) q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) q = q * self.scale - attn = (q @ k.transpose(-2, -1)) - - relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( - self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH - relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + attn = q @ k.transpose(-2, -1) + + relative_position_bias = self.relative_position_bias_table[ + self.relative_position_index.view(-1) + ].view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1 + ) # Wh*Ww,Wh*Ww,nH + relative_position_bias = relative_position_bias.permute( + 2, 0, 1 + ).contiguous() # nH, Wh*Ww, Wh*Ww attn = attn + relative_position_bias.unsqueeze(0) if mask is not None: @@ -150,7 +169,7 @@ def forward(self, x, mask=None): return x def extra_repr(self) -> str: - return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}' + return f"dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}" def flops(self, N): # calculate flops for 1 window with token length of N @@ -167,7 +186,7 @@ def flops(self, N): class SwinTransformerBlock(nn.Module): - r""" Swin Transformer Block. + r"""Swin Transformer Block. Args: dim (int): Number of input channels. @@ -185,9 +204,22 @@ class SwinTransformerBlock(nn.Module): norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm """ - def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0, - mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., - act_layer=nn.GELU, norm_layer=nn.LayerNorm): + def __init__( + self, + dim, + input_resolution, + num_heads, + window_size=7, + shift_size=0, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + act_layer=nn.GELU, + norm_layer=nn.LayerNorm, + ): super().__init__() self.dim = dim self.input_resolution = input_resolution @@ -203,18 +235,26 @@ def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0 self.norm1 = norm_layer(dim) self.attn = WindowAttention( - dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, - qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) - - self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + dim, + window_size=to_2tuple(self.window_size), + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=drop, + ) + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() self.norm2 = norm_layer(dim) mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + self.mlp = Mlp( + in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop + ) self.H = input_resolution[0] self.W = input_resolution[1] - self.attn_mask_dict = {} + self.attn_mask_dict = {} def create_attn_mask(self, H, W): # calculate attention mask for SW-MSA @@ -222,26 +262,33 @@ def create_attn_mask(self, H, W): Hp = int(np.ceil(H / self.window_size)) * self.window_size Wp = int(np.ceil(W / self.window_size)) * self.window_size img_mask = torch.zeros((1, Hp, Wp, 1)) # 1 Hp Wp 1 - h_slices = (slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None)) - w_slices = (slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None)) + h_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) + w_slices = ( + slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None), + ) cnt = 0 for h in h_slices: for w in w_slices: img_mask[:, h, w, :] = cnt cnt += 1 - mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = window_partition( + img_mask, self.window_size + ) # nW, window_size, window_size, 1 mask_windows = mask_windows.view(-1, self.window_size * self.window_size) attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) - attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill( + attn_mask == 0, float(0.0) + ) return attn_mask - def forward(self, x): B, L, C = x.shape H = int(sqrt(L)) @@ -258,8 +305,12 @@ def forward(self, x): shifted_x = x # partition windows - x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C - x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + x_windows = window_partition( + shifted_x, self.window_size + ) # nW*B, window_size, window_size, C + x_windows = x_windows.view( + -1, self.window_size * self.window_size, C + ) # nW*B, window_size*window_size, C # W-MSA/SW-MSA attn_windows = self.attn(x_windows, mask=self.attn_mask) # nW*B, window_size*window_size, C @@ -282,8 +333,10 @@ def forward(self, x): return x def extra_repr(self) -> str: - return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \ - f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" + return ( + f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " + f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" + ) def flops(self): flops = 0 @@ -301,7 +354,7 @@ def flops(self): class PatchMerging(nn.Module): - r""" Patch Merging Layer. + r"""Patch Merging Layer. Args: input_resolution (tuple[int]): Resolution of input feature. @@ -350,7 +403,7 @@ def flops(self): class BasicLayer(nn.Module): - """ A basic Swin Transformer layer for one stage. + """A basic Swin Transformer layer for one stage. Args: dim (int): Number of input channels. @@ -369,10 +422,23 @@ class BasicLayer(nn.Module): use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. """ - def __init__(self, dim, input_resolution, depth, num_heads, window_size, - mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., - drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False): - + def __init__( + self, + dim, + input_resolution, + depth, + num_heads, + window_size, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop=0.0, + attn_drop=0.0, + drop_path=0.0, + norm_layer=nn.LayerNorm, + downsample=None, + use_checkpoint=False, + ): super().__init__() self.dim = dim self.input_resolution = input_resolution @@ -380,16 +446,25 @@ def __init__(self, dim, input_resolution, depth, num_heads, window_size, self.use_checkpoint = use_checkpoint # build blocks - self.blocks = nn.ModuleList([ - SwinTransformerBlock(dim=dim, input_resolution=input_resolution, - num_heads=num_heads, window_size=window_size, - shift_size=0 if (i % 2 == 0) else window_size // 2, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop, attn_drop=attn_drop, - drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, - norm_layer=norm_layer) - for i in range(depth)]) + self.blocks = nn.ModuleList( + [ + SwinTransformerBlock( + dim=dim, + input_resolution=input_resolution, + num_heads=num_heads, + window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop, + attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer, + ) + for i in range(depth) + ] + ) # patch merging layer if downsample is not None: @@ -421,7 +496,7 @@ def flops(self): class PatchEmbed(nn.Module): - r""" Image to Patch Embedding + r"""Image to Patch Embedding Args: img_size (int): Image size. Default: 224. @@ -453,8 +528,9 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_la def forward(self, x): B, C, H, W = x.shape # FIXME look at relaxing size constraints - assert H == self.img_size[0] and W == self.img_size[1], \ + assert H == self.img_size[0] and W == self.img_size[1], ( f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." + ) x = self.proj(x).flatten(2).transpose(1, 2) # B Ph*Pw C if self.norm is not None: x = self.norm(x) @@ -469,7 +545,7 @@ def flops(self): class SwinTransformer(nn.Module): - r""" Swin Transformer + r"""Swin Transformer A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - https://arxiv.org/pdf/2103.14030 @@ -493,12 +569,28 @@ class SwinTransformer(nn.Module): use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False """ - def __init__(self, img_size=224, patch_size=4, in_chans=3, - embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], - window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None, - drop_rate=0., attn_drop_rate=0., drop_path_rate=0.3, - norm_layer=partial(nn.LayerNorm, eps=1e-6), ape=False, patch_norm=True, - use_checkpoint=False, output_avg=False, **kwargs): + def __init__( + self, + img_size=224, + patch_size=4, + in_chans=3, + embed_dim=96, + depths=[2, 2, 6, 2], + num_heads=[3, 6, 12, 24], + window_size=7, + mlp_ratio=4.0, + qkv_bias=True, + qk_scale=None, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.3, + norm_layer=partial(nn.LayerNorm, eps=1e-6), + ape=False, + patch_norm=True, + use_checkpoint=False, + output_avg=False, + **kwargs, + ): super().__init__() self.num_layers = len(depths) @@ -510,11 +602,15 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, self.img_size = to_2tuple(img_size) self.patch_size = to_2tuple(patch_size) self.output_avg = output_avg - + # split image into non-overlapping patches self.patch_embed = PatchEmbed( - img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, - norm_layer=norm_layer if self.patch_norm else None) + img_size=img_size, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None, + ) num_patches = self.patch_embed.num_patches patches_resolution = self.patch_embed.patches_resolution self.patches_resolution = patches_resolution @@ -522,36 +618,44 @@ def __init__(self, img_size=224, patch_size=4, in_chans=3, # absolute position embedding if self.ape: self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) - trunc_normal_(self.absolute_pos_embed, std=.02) + trunc_normal_(self.absolute_pos_embed, std=0.02) self.pos_drop = nn.Dropout(p=drop_rate) # stochastic depth - dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + dpr = [ + x.item() for x in torch.linspace(0, drop_path_rate, sum(depths)) + ] # stochastic depth decay rule # build layers self.layers = nn.ModuleList() for i_layer in range(self.num_layers): - layer = BasicLayer(dim=int(embed_dim * 2 ** i_layer), - input_resolution=(patches_resolution[0] // (2 ** i_layer), - patches_resolution[1] // (2 ** i_layer)), - depth=depths[i_layer], - num_heads=num_heads[i_layer], - window_size=window_size, - mlp_ratio=self.mlp_ratio, - qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, - drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], - norm_layer=norm_layer, - downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, - use_checkpoint=use_checkpoint) + layer = BasicLayer( + dim=int(embed_dim * 2**i_layer), + input_resolution=( + patches_resolution[0] // (2**i_layer), + patches_resolution[1] // (2**i_layer), + ), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=self.mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])], + norm_layer=norm_layer, + downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, + use_checkpoint=use_checkpoint, + ) self.layers.append(layer) self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=.02) + trunc_normal_(m.weight, std=0.02) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): @@ -560,11 +664,11 @@ def _init_weights(self, m): @torch.jit.ignore def no_weight_decay(self): - return {'absolute_pos_embed'} + return {"absolute_pos_embed"} @torch.jit.ignore def no_weight_decay_keywords(self): - return {'relative_position_bias_table'} + return {"relative_position_bias_table"} def forward(self, x): x = self.patch_embed(x) @@ -584,7 +688,7 @@ def forward(self, x): px = px.permute(0, 2, 1).contiguous() px = px.reshape(b, c, h, w) # is this a fair assumption ?? i think it's baked into the architecture - h, w = h//2, w//2 + h, w = h // 2, w // 2 outs.append(px) if self.output_avg: @@ -597,7 +701,12 @@ def flops(self): flops += self.patch_embed.flops() for i, layer in enumerate(self.layers): flops += layer.flops() - flops += self.num_features * self.patches_resolution[0] * self.patches_resolution[1] // (2 ** self.num_layers) + flops += ( + self.num_features + * self.patches_resolution[0] + * self.patches_resolution[1] + // (2**self.num_layers) + ) flops += self.num_features * self.num_classes return flops @@ -610,7 +719,10 @@ def get_swin(drop_path_rate=0.3, output_avg=False): depths = [2, 2, 18, 2] num_heads = [4, 8, 16, 32] swin = SwinTransformer( - img_size=(args.img_h, args.img_w,), + img_size=( + args.img_h, + args.img_w, + ), in_chans=3, patch_size=args.patch_dim, embed_dim=embed_dim, @@ -622,4 +734,3 @@ def get_swin(drop_path_rate=0.3, output_avg=False): ) return swin - diff --git a/flagscale/train/megatron/legacy/model/vision/utils.py b/flagscale/train/megatron/legacy/model/vision/utils.py index 6d29a877f1..63f9598b3e 100644 --- a/flagscale/train/megatron/legacy/model/vision/utils.py +++ b/flagscale/train/megatron/legacy/model/vision/utils.py @@ -4,25 +4,23 @@ import torch.nn.functional as F -def resize(input, - size=None, - scale_factor=None, - mode='nearest', - align_corners=None, - warning=True): +def resize(input, size=None, scale_factor=None, mode="nearest", align_corners=None, warning=True): if warning: if size is not None and align_corners: input_h, input_w = tuple(int(x) for x in input.shape[2:]) output_h, output_w = tuple(int(x) for x in size) if output_h > input_h or output_w > output_h: - if ((output_h > 1 and output_w > 1 and input_h > 1 - and input_w > 1) and (output_h - 1) % (input_h - 1) - and (output_w - 1) % (input_w - 1)): + if ( + (output_h > 1 and output_w > 1 and input_h > 1 and input_w > 1) + and (output_h - 1) % (input_h - 1) + and (output_w - 1) % (input_w - 1) + ): warnings.warn( - f'When align_corners={align_corners}, ' - 'the output would more aligned if ' - f'input size {(input_h, input_w)} is `x+1` and ' - f'out size {(output_h, output_w)} is `nx+1`') + f"When align_corners={align_corners}, " + "the output would more aligned if " + f"input size {(input_h, input_w)} is `x+1` and " + f"out size {(output_h, output_w)} is `nx+1`" + ) if isinstance(size, torch.Size): size = tuple(int(x) for x in size) return F.interpolate(input, size, scale_factor, mode, align_corners) diff --git a/flagscale/train/megatron/legacy/model/vision/vit_backbone.py b/flagscale/train/megatron/legacy/model/vision/vit_backbone.py index 7bd24b8877..9d236cff5f 100644 --- a/flagscale/train/megatron/legacy/model/vision/vit_backbone.py +++ b/flagscale/train/megatron/legacy/model/vision/vit_backbone.py @@ -19,8 +19,10 @@ CLASS_TOKEN_LENGTH = 8 from megatron.plugin.platform import get_platform + cur_platform = get_platform() + class VitMlpHead(MegatronModule): """Pooler layer. @@ -51,9 +53,9 @@ def forward(self, hidden_states): def isPerfectSquare(x): - if(x >= 0): + if x >= 0: sr = math.sqrt(x) - return (int(sr) * int(sr) == x) + return int(sr) * int(sr) == x return False @@ -66,7 +68,6 @@ def twod_interpolate_position_embeddings_hook( unexpected_keys, error_msgs, ): - args = get_args() num_patches_per_dim_h = args.img_h // args.patch_dim num_patches_per_dim_w = args.img_w // args.patch_dim @@ -80,9 +81,11 @@ def twod_interpolate_position_embeddings_hook( input_param = state_dict[key] input_seq_len = input_param.shape[0] - assert(isPerfectSquare(input_seq_len) or isPerfectSquare(input_seq_len - CLASS_TOKEN_LENGTH)) + assert isPerfectSquare(input_seq_len) or isPerfectSquare(input_seq_len - CLASS_TOKEN_LENGTH) input_has_class_token = not isPerfectSquare(input_seq_len) - num_tok_input = input_seq_len - CLASS_TOKEN_LENGTH if input_has_class_token else input_seq_len + num_tok_input = ( + input_seq_len - CLASS_TOKEN_LENGTH if input_has_class_token else input_seq_len + ) num_tok_output = num_patches output_has_class_token = args.class_token_present @@ -97,14 +100,11 @@ def twod_interpolate_position_embeddings_hook( assert input_param.shape[1] == hidden_size if num_tok_input != num_tok_output: - gs_input = int(math.sqrt(num_tok_input)) gs_new = (num_patches_per_dim_h, num_patches_per_dim_w) input_param_grid = input_param_grid.transpose(0, 1).contiguous() - input_param_grid = input_param_grid.reshape( - (1, -1, gs_input, gs_input) - ) + input_param_grid = input_param_grid.reshape((1, -1, gs_input, gs_input)) input_param_grid = input_param_grid.float() scale_factor = (gs_new[0] / gs_input, gs_new[1] / gs_input) @@ -119,10 +119,7 @@ def twod_interpolate_position_embeddings_hook( assert input_param_grid.shape[1] == hidden_size input_param = input_param_grid - assert ( - input_param.shape[0] == num_tok_output - and input_param.shape[1] == hidden_size - ) + assert input_param.shape[0] == num_tok_output and input_param.shape[1] == hidden_size if output_has_class_token: input_param = torch.cat((input_param_tok, input_param), dim=0) @@ -133,14 +130,16 @@ def twod_interpolate_position_embeddings_hook( class VitBackbone(MegatronModule): """Vision Transformer Model.""" - def __init__(self, - config, - pre_process=True, - post_process=True, - class_token=True, - single_token_output=False, - post_layer_norm=True, - drop_path_rate=0.0): + def __init__( + self, + config, + pre_process=True, + post_process=True, + class_token=True, + single_token_output=False, + post_layer_norm=True, + drop_path_rate=0.0, + ): super(VitBackbone, self).__init__(share_embeddings_and_output_weights=False) args = get_args() self.config = config @@ -176,20 +175,16 @@ def __init__(self, torch.randn(1, CLASS_TOKEN_LENGTH, self.hidden_size) ) torch.nn.init.zeros_(self.cls_token) - self.position_ids = torch.arange(self.seq_length).expand(1, -1).to(cur_platform.device()) + self.position_ids = ( + torch.arange(self.seq_length).expand(1, -1).to(cur_platform.device()) + ) # Linear encoder - self.linear_encoder = torch.nn.Linear( - self.flatten_dim, self.hidden_size - ) + self.linear_encoder = torch.nn.Linear(self.flatten_dim, self.hidden_size) # embedding - self.position_embeddings = torch.nn.Embedding( - self.seq_length, self.hidden_size - ) - init_method_normal(args.init_method_std)( - self.position_embeddings.weight - ) + self.position_embeddings = torch.nn.Embedding(self.seq_length, self.hidden_size) + init_method_normal(args.init_method_std)(self.position_embeddings.weight) args.class_token_present = self.class_token self.position_embeddings._register_load_state_dict_pre_hook( @@ -205,7 +200,7 @@ def __init__(self, pre_process=self.pre_process, post_process=self.post_process, post_layer_norm=self.post_layer_norm, - drop_path_rate=self.drop_path_rate + drop_path_rate=self.drop_path_rate, ) def set_input_tensor(self, input_tensor): @@ -213,7 +208,6 @@ def set_input_tensor(self, input_tensor): self.transformer.set_input_tensor(input_tensor) def forward(self, input): - if self.pre_process: rearranged_input = einops.rearrange( input, @@ -230,8 +224,9 @@ def forward(self, input): cls_tokens = self.cls_token.expand(encoder_output.shape[0], -1, -1) concatenated_tokens = torch.cat((cls_tokens, encoder_output), dim=1) - token_embeddings = concatenated_tokens + \ - self.position_embeddings(self.position_ids[:, :concatenated_tokens.shape[1]]) + token_embeddings = concatenated_tokens + self.position_embeddings( + self.position_ids[:, : concatenated_tokens.shape[1]] + ) # [b, s, h] => [s, b, h] token_embeddings = token_embeddings.transpose(0, 1).contiguous() hidden_states = self.embedding_dropout(token_embeddings) @@ -248,4 +243,3 @@ def forward(self, input): hidden_states = hidden_states.transpose(0, 1).contiguous() return hidden_states - diff --git a/flagscale/train/megatron/model_provider.py b/flagscale/train/megatron/model_provider.py index 4b58cacfa5..09d2fed9fe 100644 --- a/flagscale/train/megatron/model_provider.py +++ b/flagscale/train/megatron/model_provider.py @@ -12,6 +12,7 @@ try: from megatron.post_training.model_builder import modelopt_gpt_mamba_builder + has_nvidia_modelopt = True except ImportError: has_nvidia_modelopt = False @@ -22,7 +23,13 @@ def model_provider( - model_builder: Callable, pre_process=True, post_process=True, vp_stage: Optional[int] = None, config=None, pg_collection=None, dualpipev_stage: Optional[int] = None + model_builder: Callable, + pre_process=True, + post_process=True, + vp_stage: Optional[int] = None, + config=None, + pg_collection=None, + dualpipev_stage: Optional[int] = None, ) -> Union[GPTModel, megatron.legacy.model.GPTModel, MambaModel]: """Builds the model. @@ -49,21 +56,31 @@ def model_provider( def oom_observer(device, alloc, device_alloc, device_free): # snapshot right after an OOM happened - print('saving allocated state during OOM') + print("saving allocated state during OOM") filename = f"oom_rank-{torch.distributed.get_rank()}_{args.memory_snapshot_path}" torch.cuda.memory._dump_snapshot(filename) torch._C._cuda_attach_out_of_memory_observer(oom_observer) - if has_nvidia_modelopt and getattr(args, 'modelopt_enabled', False): + if has_nvidia_modelopt and getattr(args, "modelopt_enabled", False): # [ModelOpt]: Use custom builder + spec when modelopt is enabled model_builder = modelopt_gpt_mamba_builder assert not args.use_dualpipev if dualpipev_stage is not None: - return model_builder(args, pre_process, post_process, vp_stage, config=config, pg_collection=pg_collection, dualpipev_stage=dualpipev_stage) - return model_builder(args, pre_process, post_process, vp_stage, config=config, pg_collection=pg_collection) + return model_builder( + args, + pre_process, + post_process, + vp_stage, + config=config, + pg_collection=pg_collection, + dualpipev_stage=dualpipev_stage, + ) + return model_builder( + args, pre_process, post_process, vp_stage, config=config, pg_collection=pg_collection + ) def count_parameters_in_layer(model, layer_name): diff --git a/flagscale/train/megatron/plugin_flagscale/__init__.py b/flagscale/train/megatron/plugin_flagscale/__init__.py index 5ddaa1975c..7ba6cd45ae 100644 --- a/flagscale/train/megatron/plugin_flagscale/__init__.py +++ b/flagscale/train/megatron/plugin_flagscale/__init__.py @@ -1,4 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Load the centralized override registry so that all lazy mappings are # available before any @overridable function is called. import megatron.plugin_flagscale.override_registry # noqa: F401 - diff --git a/flagscale/train/megatron/plugin_flagscale/dist_signal_handler.py b/flagscale/train/megatron/plugin_flagscale/dist_signal_handler.py index a47af2e067..50b8ede7f0 100644 --- a/flagscale/train/megatron/plugin_flagscale/dist_signal_handler.py +++ b/flagscale/train/megatron/plugin_flagscale/dist_signal_handler.py @@ -4,6 +4,7 @@ from megatron.plugin.platform import get_platform from functools import wraps + cur_platform = get_platform() logger = logging.getLogger(__name__) @@ -12,13 +13,13 @@ def get_device(local_rank=None): logger.info("FlagScale Plugin Decorator: Unsing overrided dist_singal_handler.get_device") backend = torch.distributed.get_backend() - if backend == 'nccl': + if backend == "nccl": if local_rank is None: device = torch.device(cur_platform.device_name()) else: - device = torch.device(f'{cur_platform.device_name()}:{local_rank}') - elif backend == 'gloo': - device = torch.device('cpu') + device = torch.device(f"{cur_platform.device_name()}:{local_rank}") + elif backend == "gloo": + device = torch.device("cpu") else: raise RuntimeError - return device \ No newline at end of file + return device diff --git a/flagscale/train/megatron/plugin_flagscale/npu_plugin.py b/flagscale/train/megatron/plugin_flagscale/npu_plugin.py index ed53ac989e..8025fac761 100644 --- a/flagscale/train/megatron/plugin_flagscale/npu_plugin.py +++ b/flagscale/train/megatron/plugin_flagscale/npu_plugin.py @@ -1,37 +1,56 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import time import torch import torch_npu + def get_device_arch_version(): return 8 + def _compile_dependencies(): if torch.distributed.get_rank() == 0: start_time = time.time() - print('> compiling dataset index builder ...') + print("> compiling dataset index builder ...") from megatron.core.datasets.utils import compile_helpers compile_helpers() print( - '>>> done with dataset index builder. Compilation time: {:.3f} seconds'.format(time.time() - start_time), + ">>> done with dataset index builder. Compilation time: {:.3f} seconds".format( + time.time() - start_time + ), flush=True, ) + def get_device(local_rank=None): backend = torch.distributed.get_backend() - if backend == 'hccl': + if backend == "hccl": if local_rank is None: - device = torch.device('cuda') + device = torch.device("cuda") else: - device = torch.device(f'cuda:{local_rank}') - elif backend == 'nccl': + device = torch.device(f"cuda:{local_rank}") + elif backend == "nccl": if local_rank is None: device = torch.device(cur_platform.device_name()) else: - device = torch.device(f'{cur_platform.device_name()}:{local_rank}') - elif backend == 'gloo': - device = torch.device('cpu') + device = torch.device(f"{cur_platform.device_name()}:{local_rank}") + elif backend == "gloo": + device = torch.device("cpu") else: raise RuntimeError(f"Unsupported distributed backend: {backend}") - return device \ No newline at end of file + return device diff --git a/flagscale/train/megatron/plugin_flagscale/override_registry.py b/flagscale/train/megatron/plugin_flagscale/override_registry.py index 4707ff4d9c..db77c4e6c3 100644 --- a/flagscale/train/megatron/plugin_flagscale/override_registry.py +++ b/flagscale/train/megatron/plugin_flagscale/override_registry.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """ Centralized override registry for FlagScale training """ diff --git a/flagscale/train/megatron/plugin_flagscale/test_override.py b/flagscale/train/megatron/plugin_flagscale/test_override.py index 20deb1ed1b..2f21ccee20 100644 --- a/flagscale/train/megatron/plugin_flagscale/test_override.py +++ b/flagscale/train/megatron/plugin_flagscale/test_override.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """ Test override of get_device in plugin_flagscale. @@ -8,15 +22,14 @@ Or standalone: python megatron/plugin_flagscale/test_override.py """ + import sys import os import unittest from unittest.mock import patch, MagicMock # Ensure both Megatron-LM-FL and FlagScale's megatron paths are importable -MEGATRON_LM_FL_ROOT = os.path.abspath( - os.path.join(os.path.dirname(__file__), "..", "..", "..") -) +MEGATRON_LM_FL_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) if MEGATRON_LM_FL_ROOT not in sys.path: sys.path.insert(0, MEGATRON_LM_FL_ROOT) @@ -61,14 +74,16 @@ def test_no_override_uses_original(self): @overridable def get_device(local_rank=None): import torch + # Simulate original logic if local_rank is None: return torch.device(mock_platform.device_name()) else: - return torch.device(f'{mock_platform.device_name()}:{local_rank}') + return torch.device(f"{mock_platform.device_name()}:{local_rank}") with patch("torch.distributed.get_backend", return_value="nccl"): import torch + device = get_device() self.assertEqual(device, torch.device("cuda")) @@ -83,6 +98,7 @@ def test_override_replaces_get_device(self): @overridable def get_device(local_rank=None): import torch + # Original: hardcoded cuda if local_rank is None: return torch.device("cuda") @@ -96,11 +112,11 @@ def get_device_override(local_rank=None): if local_rank is None: return mock_platform.device_name() else: - return f'{mock_platform.device_name()}:{local_rank}' + return f"{mock_platform.device_name()}:{local_rank}" # The method_key for module-level function is "module_basename.func_name" # Since get_device is defined in this test, module is __main__ or test file - module_parts = get_device.__module__.split('.') + module_parts = get_device.__module__.split(".") module_name = module_parts[-1] register_override_method(f"{module_name}.get_device", get_device_override) @@ -134,7 +150,7 @@ def get_device_plugin(local_rank=None): if local_rank is None: return mock_platform.device_name() else: - return f'{mock_platform.device_name()}:{local_rank}' + return f"{mock_platform.device_name()}:{local_rank}" plugin_mod.get_device = get_device_plugin sys.modules["megatron.plugin_flagscale._test_get_device"] = plugin_mod @@ -168,10 +184,10 @@ def test_gloo_backend_returns_cpu(self): def get_device_plugin(local_rank=None): backend = torch.distributed.get_backend() - if backend == 'nccl': + if backend == "nccl": return torch.device("cuda") - elif backend == 'gloo': - return torch.device('cpu') + elif backend == "gloo": + return torch.device("cpu") else: raise RuntimeError @@ -179,7 +195,7 @@ def get_device_plugin(local_rank=None): def get_device(local_rank=None): return torch.device("cuda") # dummy original - module_parts = get_device.__module__.split('.') + module_parts = get_device.__module__.split(".") module_name = module_parts[-1] register_override_method(f"{module_name}.get_device", get_device_plugin) @@ -193,10 +209,10 @@ def test_unsupported_backend_raises(self): def get_device_plugin(local_rank=None): backend = torch.distributed.get_backend() - if backend == 'nccl': + if backend == "nccl": return torch.device("cuda") - elif backend == 'gloo': - return torch.device('cpu') + elif backend == "gloo": + return torch.device("cpu") else: raise RuntimeError(f"Unsupported backend: {backend}") @@ -204,7 +220,7 @@ def get_device_plugin(local_rank=None): def get_device(local_rank=None): return torch.device("cuda") - module_parts = get_device.__module__.split('.') + module_parts = get_device.__module__.split(".") module_name = module_parts[-1] register_override_method(f"{module_name}.get_device", get_device_plugin) diff --git a/flagscale/train/megatron/train_deepseek_v4.py b/flagscale/train/megatron/train_deepseek_v4.py index 41b5af980a..85ba72cb72 100644 --- a/flagscale/train/megatron/train_deepseek_v4.py +++ b/flagscale/train/megatron/train_deepseek_v4.py @@ -43,7 +43,6 @@ def get_batch_on_this_tp_rank(data_iterator): - args = get_args() def _broadcast(item): @@ -55,54 +54,52 @@ def _broadcast(item): ) if mpu.get_tensor_model_parallel_rank() == 0: - assert data_iterator is not None data = next(data_iterator) batch = { - 'tokens': data["tokens"].cuda(non_blocking=True), - 'labels': data["labels"].cuda(non_blocking=True), - 'loss_mask': data["loss_mask"].cuda(non_blocking=True), - 'attention_mask': ( + "tokens": data["tokens"].cuda(non_blocking=True), + "labels": data["labels"].cuda(non_blocking=True), + "loss_mask": data["loss_mask"].cuda(non_blocking=True), + "attention_mask": ( None if "attention_mask" not in data else data["attention_mask"].cuda(non_blocking=True) ), - 'position_ids': data["position_ids"].cuda(non_blocking=True), + "position_ids": data["position_ids"].cuda(non_blocking=True), } if args.pipeline_model_parallel_size == 1: - _broadcast(batch['tokens']) - _broadcast(batch['labels']) - _broadcast(batch['loss_mask']) - _broadcast(batch['attention_mask']) - _broadcast(batch['position_ids']) + _broadcast(batch["tokens"]) + _broadcast(batch["labels"]) + _broadcast(batch["loss_mask"]) + _broadcast(batch["attention_mask"]) + _broadcast(batch["position_ids"]) elif mpu.is_pipeline_first_stage(): - _broadcast(batch['tokens']) - _broadcast(batch['attention_mask']) - _broadcast(batch['position_ids']) + _broadcast(batch["tokens"]) + _broadcast(batch["attention_mask"]) + _broadcast(batch["position_ids"]) ######### FlagScale Begin ######## - if mpu.get_dualpipev_pipeline_model_parallel_world_size() is not None: - _broadcast(batch['loss_mask']) - _broadcast(batch['labels']) + if mpu.get_dualpipev_pipeline_model_parallel_world_size() is not None: + _broadcast(batch["loss_mask"]) + _broadcast(batch["labels"]) ######### FlagScale End ######## elif mpu.is_pipeline_last_stage(): # Multi-Token Prediction (MTP) layers need tokens and position_ids to calculate embedding. # Currently the Multi-Token Prediction (MTP) layers is fixed on the last stage, so we need # to broadcast tokens and position_ids to all of the tensor parallel ranks on the last stage. - _broadcast(batch['tokens']) + _broadcast(batch["tokens"]) if args.mtp_num_layers is not None: - _broadcast(batch['position_ids']) - _broadcast(batch['labels']) - _broadcast(batch['loss_mask']) - _broadcast(batch['attention_mask']) + _broadcast(batch["position_ids"]) + _broadcast(batch["labels"]) + _broadcast(batch["loss_mask"]) + _broadcast(batch["attention_mask"]) else: - _broadcast(batch['tokens']) + _broadcast(batch["tokens"]) else: - tokens = torch.empty( (args.micro_batch_size, args.seq_length), dtype=torch.int64, @@ -169,11 +166,11 @@ def _broadcast(item): _broadcast(tokens) batch = { - 'tokens': tokens, - 'labels': labels, - 'loss_mask': loss_mask, - 'attention_mask': attention_mask, - 'position_ids': position_ids, + "tokens": tokens, + "labels": labels, + "loss_mask": loss_mask, + "attention_mask": attention_mask, + "position_ids": position_ids, } return batch @@ -257,7 +254,7 @@ def loss_func( num_tokens = loss_mask.sum().clone().detach().to(torch.int) reporting_loss = torch.cat([loss.clone().detach().view(1), num_tokens.view(1)]) - return (loss, num_tokens, {'lm loss': reporting_loss}) + return (loss, num_tokens, {"lm loss": reporting_loss}) def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = False): @@ -272,21 +269,22 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa timers = get_timers() # Get the batch. - timers('batch-generator', log_level=2).start() + timers("batch-generator", log_level=2).start() global stimer with stimer(bdata=True): vp_stage = get_attr_wrapped_model(model, "vp_stage") tokens, labels, loss_mask, attention_mask, position_ids = get_batch(data_iterator, vp_stage) # print(f"in train_engram, rank: {torch.distributed.get_rank()}, {tokens=}") - timers('batch-generator').stop() + timers("batch-generator").stop() with stimer: if args.use_legacy_models: output_tensor = model(tokens, position_ids, attention_mask, labels=labels) else: if return_schedule_plan: - assert args.overlap_moe_expert_parallel_comm, \ + assert args.overlap_moe_expert_parallel_comm, ( "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan" + ) schedule_plan = model.build_schedule_plan( tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask ) @@ -362,7 +360,10 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None print_rank_0("> building train, validation, and test datasets for GPT ...") train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder( - dataset_type, train_val_test_num_samples, partial(is_dataset_built_on_rank, vp_stage=vp_stage), config + dataset_type, + train_val_test_num_samples, + partial(is_dataset_built_on_rank, vp_stage=vp_stage), + config, ).build() print_rank_0("> finished creating GPT datasets ...") @@ -371,22 +372,21 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None if __name__ == "__main__": - # Temporary for transition to core datasets train_valid_test_datasets_provider.is_distributed = True # Optionally enable inprocess restart on pretrain pretrain, store = inprocess_restart.maybe_wrap_for_inprocess_restart(pretrain) - extra_valid_datasets_provider.is_distributed = True ######## FlagScale ######## + extra_valid_datasets_provider.is_distributed = True ######## FlagScale ######## pretrain( train_valid_test_datasets_provider, partial(model_provider, deepseek_builder), ModelType.encoder_or_decoder, forward_step, - args_defaults={'tokenizer_type': 'GPT2BPETokenizer'}, + args_defaults={"tokenizer_type": "GPT2BPETokenizer"}, extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None, store=store, - extra_valid_dataset_provider=extra_valid_datasets_provider + extra_valid_dataset_provider=extra_valid_datasets_provider, ) diff --git a/flagscale/train/megatron/train_engram.py b/flagscale/train/megatron/train_engram.py index 62bc1034e5..0559598126 100644 --- a/flagscale/train/megatron/train_engram.py +++ b/flagscale/train/megatron/train_engram.py @@ -37,6 +37,7 @@ from megatron.training.training import pretrain from megatron.plugin.hetero.parallel_context import get_parallel_context from megatron.plugin.platform import get_platform + cur_platform = get_platform() # engram @@ -46,7 +47,6 @@ def get_batch_on_this_tp_rank(data_iterator): - args = get_args() def _broadcast(item): @@ -58,54 +58,54 @@ def _broadcast(item): ) if mpu.get_tensor_model_parallel_rank() == 0: - assert data_iterator is not None data = next(data_iterator) batch = { - 'tokens': data["tokens"].to(device=cur_platform.device_name(), non_blocking=True), - 'labels': data["labels"].to(device=cur_platform.device_name(), non_blocking=True), - 'loss_mask': data["loss_mask"].to(device=cur_platform.device_name(), non_blocking=True), - 'attention_mask': ( + "tokens": data["tokens"].to(device=cur_platform.device_name(), non_blocking=True), + "labels": data["labels"].to(device=cur_platform.device_name(), non_blocking=True), + "loss_mask": data["loss_mask"].to(device=cur_platform.device_name(), non_blocking=True), + "attention_mask": ( None if "attention_mask" not in data else data["attention_mask"].to(device=cur_platform.device_name(), non_blocking=True) ), - 'position_ids': data["position_ids"].to(device=cur_platform.device_name(), non_blocking=True), + "position_ids": data["position_ids"].to( + device=cur_platform.device_name(), non_blocking=True + ), } if args.pipeline_model_parallel_size == 1: - _broadcast(batch['tokens']) - _broadcast(batch['labels']) - _broadcast(batch['loss_mask']) - _broadcast(batch['attention_mask']) - _broadcast(batch['position_ids']) + _broadcast(batch["tokens"]) + _broadcast(batch["labels"]) + _broadcast(batch["loss_mask"]) + _broadcast(batch["attention_mask"]) + _broadcast(batch["position_ids"]) elif mpu.is_pipeline_first_stage(): - _broadcast(batch['tokens']) - _broadcast(batch['attention_mask']) - _broadcast(batch['position_ids']) + _broadcast(batch["tokens"]) + _broadcast(batch["attention_mask"]) + _broadcast(batch["position_ids"]) ######### FlagScale Begin ######## - if mpu.get_dualpipev_pipeline_model_parallel_world_size() is not None: - _broadcast(batch['loss_mask']) - _broadcast(batch['labels']) + if mpu.get_dualpipev_pipeline_model_parallel_world_size() is not None: + _broadcast(batch["loss_mask"]) + _broadcast(batch["labels"]) ######### FlagScale End ######## elif mpu.is_pipeline_last_stage(): # Multi-Token Prediction (MTP) layers need tokens and position_ids to calculate embedding. # Currently the Multi-Token Prediction (MTP) layers is fixed on the last stage, so we need # to broadcast tokens and position_ids to all of the tensor parallel ranks on the last stage. - _broadcast(batch['tokens']) + _broadcast(batch["tokens"]) if args.mtp_num_layers is not None: - _broadcast(batch['position_ids']) - _broadcast(batch['labels']) - _broadcast(batch['loss_mask']) - _broadcast(batch['attention_mask']) + _broadcast(batch["position_ids"]) + _broadcast(batch["labels"]) + _broadcast(batch["loss_mask"]) + _broadcast(batch["attention_mask"]) else: - _broadcast(batch['tokens']) + _broadcast(batch["tokens"]) else: - tokens = torch.empty( (args.micro_batch_size, args.seq_length), dtype=torch.int64, @@ -172,11 +172,11 @@ def _broadcast(item): _broadcast(tokens) batch = { - 'tokens': tokens, - 'labels': labels, - 'loss_mask': loss_mask, - 'attention_mask': attention_mask, - 'position_ids': position_ids, + "tokens": tokens, + "labels": labels, + "loss_mask": loss_mask, + "attention_mask": attention_mask, + "position_ids": position_ids, } return batch @@ -219,7 +219,7 @@ def loss_func( """ args = get_args() - if has_nvidia_modelopt and getattr(args, 'modelopt_enabled', False): # [ModelOpt] + if has_nvidia_modelopt and getattr(args, "modelopt_enabled", False): # [ModelOpt] return loss_func_modelopt(loss_mask, output_tensor, model=model) losses = output_tensor.view(-1).float() @@ -260,7 +260,7 @@ def loss_func( num_tokens = loss_mask.sum().clone().detach().to(torch.int) reporting_loss = torch.cat([loss.clone().detach().view(1), num_tokens.view(1)]) - return (loss, num_tokens, {'lm loss': reporting_loss}) + return (loss, num_tokens, {"lm loss": reporting_loss}) def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = False): @@ -275,21 +275,22 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa timers = get_timers() # Get the batch. - timers('batch-generator', log_level=2).start() + timers("batch-generator", log_level=2).start() global stimer with stimer(bdata=True): vp_stage = get_attr_wrapped_model(model, "vp_stage") tokens, labels, loss_mask, attention_mask, position_ids = get_batch(data_iterator, vp_stage) # print(f"in train_engram, rank: {torch.distributed.get_rank()}, {tokens=}") - timers('batch-generator').stop() + timers("batch-generator").stop() with stimer: if args.use_legacy_models: output_tensor = model(tokens, position_ids, attention_mask, labels=labels) else: if return_schedule_plan: - assert args.overlap_moe_expert_parallel_comm, \ + assert args.overlap_moe_expert_parallel_comm, ( "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan" + ) schedule_plan = model.build_schedule_plan( tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask ) @@ -365,7 +366,10 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None print_rank_0("> building train, validation, and test datasets for GPT ...") train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder( - dataset_type, train_val_test_num_samples, partial(is_dataset_built_on_rank, vp_stage=vp_stage), config + dataset_type, + train_val_test_num_samples, + partial(is_dataset_built_on_rank, vp_stage=vp_stage), + config, ).build() print_rank_0("> finished creating GPT datasets ...") @@ -374,22 +378,21 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None if __name__ == "__main__": - # Temporary for transition to core datasets train_valid_test_datasets_provider.is_distributed = True # Optionally enable inprocess restart on pretrain pretrain, store = inprocess_restart.maybe_wrap_for_inprocess_restart(pretrain) - extra_valid_datasets_provider.is_distributed = True ######## FlagScale ######## + extra_valid_datasets_provider.is_distributed = True ######## FlagScale ######## pretrain( train_valid_test_datasets_provider, partial(model_provider, engram_builder), ModelType.encoder_or_decoder, forward_step, - args_defaults={'tokenizer_type': 'GPT2BPETokenizer'}, + args_defaults={"tokenizer_type": "GPT2BPETokenizer"}, extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None, store=store, - extra_valid_dataset_provider=extra_valid_datasets_provider + extra_valid_dataset_provider=extra_valid_datasets_provider, ) diff --git a/flagscale/train/megatron/train_gpt.py b/flagscale/train/megatron/train_gpt.py index aaf4b9d4d7..dc185b7609 100644 --- a/flagscale/train/megatron/train_gpt.py +++ b/flagscale/train/megatron/train_gpt.py @@ -4,6 +4,7 @@ # Capture the true program start time BEFORE any heavy imports. import time + _PROGRAM_START_TIME = time.time() import json @@ -11,7 +12,8 @@ # Suppress warnings on all ranks but rank 0. import os import warnings -rank = int(os.environ.get('RANK', 0)) + +rank = int(os.environ.get("RANK", 0)) if rank != 0: warnings.filterwarnings("ignore", category=UserWarning) warnings.filterwarnings("ignore", category=FutureWarning) @@ -29,7 +31,12 @@ from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.models.gpt import GPTModel from megatron.core.rerun_state_machine import get_rerun_state_machine -from megatron.core.utils import get_attr_wrapped_model, get_thd_batch_on_this_cp_rank, get_batch_on_this_hybrid_cp_rank, StragglerDetector +from megatron.core.utils import ( + get_attr_wrapped_model, + get_thd_batch_on_this_cp_rank, + get_batch_on_this_hybrid_cp_rank, + StragglerDetector, +) from megatron.training import ( get_args, get_timers, @@ -76,39 +83,49 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): config = core_transformer_config_from_args(args) # TODO: this is pretty hacky, find a better way is_packed_sequence = get_args().sft # SFT always uses packed sequence - if not is_first_or_last_pipeline_stage(vp_stage) and not is_packed_sequence and ( - (not mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage))): + if ( + not is_first_or_last_pipeline_stage(vp_stage) + and not is_packed_sequence + and (not mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage)) + ): return None, None, None, None, None, None # get batches based on the TP rank you are on batch = get_batch_on_this_tp_rank( data_iterator, - mtp_on_this_rank=mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage) - ) + mtp_on_this_rank=mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage), + ) - cu_seqlens = batch.pop('cu_seqlens', None) - cu_seqlens_padded = batch.pop('cu_seqlens_padded', None) - max_seqlen = batch.pop('max_seqlen', None) - local_cp_size = batch.pop('local_cp_size', None) + cu_seqlens = batch.pop("cu_seqlens", None) + cu_seqlens_padded = batch.pop("cu_seqlens_padded", None) + max_seqlen = batch.pop("max_seqlen", None) + local_cp_size = batch.pop("local_cp_size", None) if local_cp_size is not None: local_cp_size = int(local_cp_size.item()) if cu_seqlens is not None: - assert ( - cu_seqlens.dim() == 2 and cu_seqlens.shape[0] == 1 - ), "micro-batch-size must be 1 for packing" + assert cu_seqlens.dim() == 2 and cu_seqlens.shape[0] == 1, ( + "micro-batch-size must be 1 for packing" + ) cu_seqlens = cu_seqlens[0] assert max_seqlen.dim() == 1 # For middle pipeline stages with packed sequences, only cu_seqlens and # max_seqlen are needed (for attention masking); skip the full batch. if not is_first_or_last_pipeline_stage(vp_stage) and is_packed_sequence: - return None, None, None, None, None, PackedSeqParams( - cu_seqlens_q=cu_seqlens, - cu_seqlens_kv=cu_seqlens, - max_seqlen_q=int(max_seqlen[0].item()), - max_seqlen_kv=int(max_seqlen[0].item()), - qkv_format='thd', + return ( + None, + None, + None, + None, + None, + PackedSeqParams( + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=int(max_seqlen[0].item()), + max_seqlen_kv=int(max_seqlen[0].item()), + qkv_format="thd", + ), ) if cu_seqlens is None and local_cp_size is None: @@ -116,8 +133,10 @@ def get_batch(data_iterator, vp_stage: Optional[int] = None): batch = get_batch_on_this_cp_rank(batch) # The implementation of this function is in MCore packed_seq_params = None elif local_cp_size is None: # Packed THD format - batch, packed_seq_params = get_thd_batch_on_this_cp_rank(batch, cu_seqlens, cu_seqlens_padded, max_seqlen) - else: # Hybrid CP format + batch, packed_seq_params = get_thd_batch_on_this_cp_rank( + batch, cu_seqlens, cu_seqlens_padded, max_seqlen + ) + else: # Hybrid CP format batch, packed_seq_params = get_batch_on_this_hybrid_cp_rank(batch, local_cp_size) return (*batch.values(), packed_seq_params) @@ -145,7 +164,7 @@ def loss_func( """ args = get_args() - if has_nvidia_modelopt and getattr(args, 'modelopt_enabled', False): # [ModelOpt] + if has_nvidia_modelopt and getattr(args, "modelopt_enabled", False): # [ModelOpt] loss, num_tokens, report = loss_func_modelopt(loss_mask, output_tensor, model=model) else: losses = output_tensor.view(-1).float() @@ -153,7 +172,7 @@ def loss_func( loss = torch.sum(losses * loss_mask) num_tokens = loss_mask.sum().clone().detach().to(torch.int) - report = {'lm loss': torch.cat([loss.clone().detach().view(1), num_tokens.view(1)])} + report = {"lm loss": torch.cat([loss.clone().detach().view(1), num_tokens.view(1)])} # Check individual rank losses are not NaN prior to DP all-reduce. rerun_state_machine = get_rerun_state_machine() @@ -201,27 +220,35 @@ def forward_step(data_iterator, model: GPTModel, return_schedule_plan: bool = Fa timers = get_timers() # Get the batch. - timers('batch-generator', log_level=2).start() + timers("batch-generator", log_level=2).start() global stimer with stimer(bdata=True): vp_stage = get_attr_wrapped_model(model, "vp_stage") - tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params = get_batch(data_iterator, vp_stage) - timers('batch-generator').stop() + tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params = get_batch( + data_iterator, vp_stage + ) + timers("batch-generator").stop() with stimer: if args.use_legacy_models: output_tensor = model(tokens, position_ids, attention_mask, labels=labels) else: if return_schedule_plan: - assert args.overlap_moe_expert_parallel_comm, \ + assert args.overlap_moe_expert_parallel_comm, ( "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan" + ) schedule_plan = model.build_schedule_plan( tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask ) return schedule_plan, partial(loss_func, loss_mask, model=model) else: output_tensor = model( - tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask, packed_seq_params=packed_seq_params + tokens, + position_ids, + attention_mask, + labels=labels, + loss_mask=loss_mask, + packed_seq_params=packed_seq_params, ) # [ModelOpt]: model is needed to access ModelOpt distillation losses @@ -235,9 +262,8 @@ def is_dataset_built_on_rank(vp_stage=None, is_packed_sequence=False): return False elif is_packed_sequence: return True - return ( - is_first_or_last_pipeline_stage(vp_stage) - or mtp_on_this_rank(config, ignore_virtual=False, vp_stage=vp_stage) + return is_first_or_last_pipeline_stage(vp_stage) or mtp_on_this_rank( + config, ignore_virtual=False, vp_stage=vp_stage ) @@ -278,7 +304,7 @@ def core_gpt_dataset_config_from_args(args): "defer_npy_index_mmap": args.dataloader_defer_npy_index_mmap, "context_parallel_size": args.context_parallel_size, "data_parallel_size": args.data_parallel_size, - "sequence_parallel_size": args.tensor_model_parallel_size*args.sequence_parallel, + "sequence_parallel_size": args.tensor_model_parallel_size * args.sequence_parallel, "hybrid_context_parallel": args.hybrid_context_parallel, } @@ -323,7 +349,6 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None if config is None: config = core_gpt_dataset_config_from_args(args) - is_packed_sequence = False if args.sft: dataset_type = SFTDataset @@ -338,7 +363,9 @@ def train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=None print_rank_0("> building train, validation, and test datasets for GPT ...") - is_dataset_built = partial(is_dataset_built_on_rank, vp_stage=vp_stage, is_packed_sequence=is_packed_sequence) + is_dataset_built = partial( + is_dataset_built_on_rank, vp_stage=vp_stage, is_packed_sequence=is_packed_sequence + ) train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder( dataset_type, train_val_test_num_samples, is_dataset_built, config ).build() @@ -385,7 +412,7 @@ def get_embedding_ranks(pp_ranks: List[int]): partial(model_provider, gpt_builder), ModelType.encoder_or_decoder, forward_step, - args_defaults={'tokenizer_type': 'GPT2BPETokenizer'}, + args_defaults={"tokenizer_type": "GPT2BPETokenizer"}, extra_args_provider=add_modelopt_args if has_nvidia_modelopt else None, store=store, get_embedding_ranks=get_embedding_ranks, diff --git a/flagscale/train/megatron/train_llava.py b/flagscale/train/megatron/train_llava.py index 2a1619d35f..cac797d9ed 100644 --- a/flagscale/train/megatron/train_llava.py +++ b/flagscale/train/megatron/train_llava.py @@ -1,5 +1,6 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. """Pretrain or SFT multimodal.""" + from copy import deepcopy from functools import partial import os @@ -8,14 +9,19 @@ import torch -sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), - os.path.pardir, os.path.pardir))) +sys.path.append( + os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir)) +) from megatron.training import get_args, get_timers, get_tokenizer, print_rank_0 from megatron.training.arguments import core_transformer_config_from_args from megatron.core import mpu, tensor_parallel from megatron.core.enums import ModelType -from examples.multimodal.config import get_language_model_config, get_vision_model_config, get_vision_projection_config +from examples.multimodal.config import ( + get_language_model_config, + get_vision_model_config, + get_vision_projection_config, +) from megatron.core.models.multimodal.llava_model import LLaVAModel from examples.multimodal.layer_specs import get_layer_spec, get_mlp_module_spec, get_layer_spec_te from megatron.training.utils import average_losses_across_data_parallel_group @@ -23,11 +29,13 @@ from megatron.training.training import pretrain from megatron.plugin.platform import get_platform + cur_platform = get_platform() + def model_provider( - pre_process=True, post_process=True, add_encoder=True, add_decoder=True, - parallel_output=True) -> LLaVAModel: + pre_process=True, post_process=True, add_encoder=True, add_decoder=True, parallel_output=True +) -> LLaVAModel: """Builds the model. Args: @@ -46,7 +54,7 @@ def model_provider( use_te = args.use_te - print_rank_0('building a multimodal model ...') + print_rank_0("building a multimodal model ...") num_image_tokens = get_image_token_count() @@ -54,11 +62,15 @@ def model_provider( args.decoder_seq_length = args.seq_length + num_image_tokens args.seq_length = num_image_tokens if torch.distributed.get_rank() == 0: - warnings.warn("Changed decoder_seq_length to num_image_tokens ({num_image_tokens}) + user-specified seq_length ({old_seq_length}).") + warnings.warn( + "Changed decoder_seq_length to num_image_tokens ({num_image_tokens}) + user-specified seq_length ({old_seq_length})." + ) if args.decoder_seq_length > args.max_position_embeddings: args.max_position_embeddings = args.decoder_seq_length - warnings.warn("Expanded max_position_embeddings to {args.max_position_embeddings} to accommodate the full sequence of vit output + llm output.") + warnings.warn( + "Expanded max_position_embeddings to {args.max_position_embeddings} to accommodate the full sequence of vit output + llm output." + ) base_config = core_transformer_config_from_args(get_args()) base_config.language_model_type = args.language_model_type @@ -72,7 +84,9 @@ def model_provider( language_transformer_layer_spec = get_layer_spec(is_vit=False) vision_config = deepcopy(base_config) - vision_config = get_vision_model_config(vision_config, apply_query_key_layer_scaling=args.apply_query_key_layer_scaling) + vision_config = get_vision_model_config( + vision_config, apply_query_key_layer_scaling=args.apply_query_key_layer_scaling + ) if use_te: vision_transformer_layer_spec = get_layer_spec_te(is_vit=True) @@ -80,15 +94,23 @@ def model_provider( vision_transformer_layer_spec = get_layer_spec(is_vit=True) vision_projection_config = deepcopy(base_config) - vision_projection_config = get_vision_projection_config(vision_projection_config, language_config.hidden_size) + vision_projection_config = get_vision_projection_config( + vision_projection_config, language_config.hidden_size + ) if args.encoder_pipeline_model_parallel_size > 0: - assert args.encoder_pipeline_model_parallel_size == 1, "ViT can only live on 1 pipeline stage." + assert args.encoder_pipeline_model_parallel_size == 1, ( + "ViT can only live on 1 pipeline stage." + ) vision_config.pipeline_model_parallel_size = args.encoder_pipeline_model_parallel_size - vision_projection_config.pipeline_model_parallel_size = args.encoder_pipeline_model_parallel_size + vision_projection_config.pipeline_model_parallel_size = ( + args.encoder_pipeline_model_parallel_size + ) if args.encoder_tensor_model_parallel_size > 0: vision_config.tensor_model_parallel_size = args.encoder_tensor_model_parallel_size - vision_projection_config.tensor_model_parallel_size = args.encoder_tensor_model_parallel_size + vision_projection_config.tensor_model_parallel_size = ( + args.encoder_tensor_model_parallel_size + ) vision_projection_layer_spec = get_mlp_module_spec(use_te=use_te).submodules @@ -117,7 +139,11 @@ def model_provider( language_rotary_base=args.rotary_base, ) - model.freeze(freeze_language_model=args.freeze_LM, freeze_vision_model=args.freeze_ViT, freeze_vision_projection=False) + model.freeze( + freeze_language_model=args.freeze_LM, + freeze_vision_model=args.freeze_ViT, + freeze_vision_projection=False, + ) return model @@ -148,28 +174,30 @@ def get_batch(data_iterator): tokens_ = data_text.long() - img_raw = data_img['img'].reshape(-1, 3, args.img_h, args.img_w) + img_raw = data_img["img"].reshape(-1, 3, args.img_h, args.img_w) cur_platform.range_push("index tokens") tokenizer = get_tokenizer() text_length = args.decoder_seq_length - args.seq_length tokens = tokens_[:, :text_length].contiguous() - labels = tokens_[:, 1:text_length+1].contiguous() + labels = tokens_[:, 1 : text_length + 1].contiguous() assert tokens.shape == labels.shape, f"tokens: {tokens.shape} != labels: {labels.shape}" cur_platform.range_pop() cur_platform.range_push("get_ltor_masks_and_position_ids") - if hasattr(tokenizer, 'eod'): + if hasattr(tokenizer, "eod"): eod_token = tokenizer.eod - elif hasattr(tokenizer, 'eos_id'): + elif hasattr(tokenizer, "eos_id"): eod_token = tokenizer.eos_id - attention_mask, loss_mask, position_ids = \ - get_ltor_masks_and_position_ids(tokens, eod_token, - args.reset_position_ids, - args.reset_attention_mask, - args.eod_mask_loss, - question_length=prompt_len) + attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids( + tokens, + eod_token, + args.reset_position_ids, + args.reset_attention_mask, + args.eod_mask_loss, + question_length=prompt_len, + ) cur_platform.range_pop() return tokens, labels, loss_mask, attention_mask, position_ids, img_raw @@ -188,13 +216,15 @@ def get_image_token_count(): return num_image_tokens -def get_ltor_masks_and_position_ids(data, - eod_token, - reset_position_ids, - reset_attention_mask, - eod_mask_loss, - question_length=None, - weights=None): +def get_ltor_masks_and_position_ids( + data, + eod_token, + reset_position_ids, + reset_attention_mask, + eod_mask_loss, + question_length=None, + weights=None, +): """Build masks and position id for left to right model.""" # Extract batch size and sequence length. @@ -205,9 +235,9 @@ def get_ltor_masks_and_position_ids(data, att_mask_batch = micro_batch_size else: att_mask_batch = 1 - attention_mask = torch.tril(torch.ones( - (att_mask_batch, seq_length, seq_length), device=data.device)).view( - att_mask_batch, 1, seq_length, seq_length) + attention_mask = torch.tril( + torch.ones((att_mask_batch, seq_length, seq_length), device=data.device) + ).view(att_mask_batch, 1, seq_length, seq_length) # Loss mask. loss_mask = torch.ones(data.size(), dtype=torch.float, device=data.device) @@ -215,22 +245,19 @@ def get_ltor_masks_and_position_ids(data, loss_mask[data == eod_token] = 0.0 # Position ids. - position_ids = torch.arange(seq_length, dtype=torch.long, - device=data.device) + position_ids = torch.arange(seq_length, dtype=torch.long, device=data.device) position_ids = position_ids.unsqueeze(0).expand_as(data) # We need to clone as the ids will be modifed based on batch index. if reset_position_ids: position_ids = position_ids.clone() - if question_length is not None: for b in range(micro_batch_size): - loss_mask[b, :max(0, question_length[b].item())] = 0.0 + loss_mask[b, : max(0, question_length[b].item())] = 0.0 if reset_position_ids or reset_attention_mask: # Loop through the batches: for b in range(micro_batch_size): - # Find indecies where EOD token is. eod_index = position_ids[b, data[b] == eod_token] # Detach indecies from positions if going to modify positions. @@ -243,14 +270,14 @@ def get_ltor_masks_and_position_ids(data, i = eod_index[j] # Mask attention loss. if reset_attention_mask: - attention_mask[b, 0, (i + 1):, :(i + 1)] = 0 + attention_mask[b, 0, (i + 1) :, : (i + 1)] = 0 # Reset positions. if reset_position_ids: - position_ids[b, (i + 1):] -= (i + 1 - prev_index) + position_ids[b, (i + 1) :] -= i + 1 - prev_index prev_index = i + 1 # Convert attention mask to binary: - attention_mask = (attention_mask < 0.5) + attention_mask = attention_mask < 0.5 if weights is not None: loss_mask = loss_mask * weights @@ -261,15 +288,14 @@ def loss_func(loss_mask, output_tensor): losses = output_tensor.float() if loss_mask is not None: loss_mask = loss_mask.view(-1).float() - loss = torch.sum(losses.view(-1) * loss_mask) / max( 1,loss_mask.sum() ) + loss = torch.sum(losses.view(-1) * loss_mask) / max(1, loss_mask.sum()) else: loss = torch.mean(losses) # Reduce loss for logging. averaged_loss = average_losses_across_data_parallel_group([loss]) - return loss, {'lm loss': averaged_loss[0]} - + return loss, {"lm loss": averaged_loss[0]} def forward_step(data_iterator, model: LLaVAModel): @@ -287,31 +313,42 @@ def forward_step(data_iterator, model: LLaVAModel): timers = get_timers() # Get the batch. - timers('batch-generator', log_level=2).start() + timers("batch-generator", log_level=2).start() tokens, labels, loss_mask, attention_mask, position_ids, images = get_batch(data_iterator) - timers('batch-generator').stop() + timers("batch-generator").stop() - output_tensor, loss_mask = model(images, tokens, position_ids, attention_mask, labels, loss_mask) + output_tensor, loss_mask = model( + images, tokens, position_ids, attention_mask, labels, loss_mask + ) return output_tensor, partial(loss_func, loss_mask) + def add_multimodal_extra_args(parser): """Extra arguments.""" - group = parser.add_argument_group(title='multimodal arguments') - group.add_argument('--valid-path', nargs='*', default=None, - help='Path to the training dataset. Accepted format:' - '1) a single data path, 2) multiple datasets in the' - 'form: dataset1-weight dataset1-path dataset2-weight ' - 'dataset2-path ...') - group.add_argument('--dataset-config', type=str, default=None) + group = parser.add_argument_group(title="multimodal arguments") + group.add_argument( + "--valid-path", + nargs="*", + default=None, + help="Path to the training dataset. Accepted format:" + "1) a single data path, 2) multiple datasets in the" + "form: dataset1-weight dataset1-path dataset2-weight " + "dataset2-path ...", + ) + group.add_argument("--dataset-config", type=str, default=None) group.add_argument("--prompt-path", type=str, default=None) - group.add_argument('--freeze-LM', action='store_true', default=False) - group.add_argument('--freeze-ViT', action='store_true', default=False) - group.add_argument('--language-model-type', type=str, required=True) + group.add_argument("--freeze-LM", action="store_true", default=False) + group.add_argument("--freeze-ViT", action="store_true", default=False) + group.add_argument("--language-model-type", type=str, required=True) group.add_argument("--disable-vision-class-token", action="store_true", default=False) - group.add_argument("--allow-missing-vision-projection-checkpoint", action="store_true", default=False) + group.add_argument( + "--allow-missing-vision-projection-checkpoint", action="store_true", default=False + ) group.add_argument("--use-te", action="store_true", default=False) - group.add_argument("--dataloader-save", type=str, default=None, help="Energon dataloader state save path") + group.add_argument( + "--dataloader-save", type=str, default=None, help="Energon dataloader state save path" + ) return parser @@ -357,7 +394,7 @@ def llava_position_embedding_ranks(pp_ranks): model_provider, ModelType.encoder_and_decoder, forward_step, - args_defaults={'tokenizer_type': 'GPT2BPETokenizer'}, + args_defaults={"tokenizer_type": "GPT2BPETokenizer"}, extra_args_provider=add_multimodal_extra_args, get_embedding_ranks=llava_embedding_ranks, get_position_embedding_ranks=llava_position_embedding_ranks, diff --git a/flagscale/train/megatron/train_llava_onevision.py b/flagscale/train/megatron/train_llava_onevision.py index eb1d2d372b..1b5de336b2 100644 --- a/flagscale/train/megatron/train_llava_onevision.py +++ b/flagscale/train/megatron/train_llava_onevision.py @@ -1,5 +1,6 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. """Pretrain or SFT LLaVA-NeXT model.""" + from copy import deepcopy from functools import partial import os @@ -9,9 +10,7 @@ import torch sys.path.append( - os.path.abspath( - os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir) - ) + os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir)) ) from megatron.training import get_args, get_timers, get_tokenizer, print_rank_0 @@ -40,6 +39,7 @@ from megatron.training.training import pretrain from megatron.plugin.platform import get_platform + cur_platform = get_platform() @@ -103,9 +103,7 @@ def model_provider( if use_te: vision_transformer_layer_spec = get_layer_spec_te(is_vit=True) else: - vision_transformer_layer_spec = get_layer_spec( - is_vit=True, normalization="LayerNorm" - ) + vision_transformer_layer_spec = get_layer_spec(is_vit=True, normalization="LayerNorm") vision_projection_config = deepcopy(base_config) vision_projection_config = get_vision_projection_config( @@ -113,19 +111,15 @@ def model_provider( ) if args.encoder_pipeline_model_parallel_size > 0: - assert ( - args.encoder_pipeline_model_parallel_size == 1 - ), "ViT can only live on 1 pipeline stage." - vision_config.pipeline_model_parallel_size = ( - args.encoder_pipeline_model_parallel_size + assert args.encoder_pipeline_model_parallel_size == 1, ( + "ViT can only live on 1 pipeline stage." ) + vision_config.pipeline_model_parallel_size = args.encoder_pipeline_model_parallel_size vision_projection_config.pipeline_model_parallel_size = ( args.encoder_pipeline_model_parallel_size ) if args.encoder_tensor_model_parallel_size > 0: - vision_config.tensor_model_parallel_size = ( - args.encoder_tensor_model_parallel_size - ) + vision_config.tensor_model_parallel_size = args.encoder_tensor_model_parallel_size vision_projection_config.tensor_model_parallel_size = ( args.encoder_tensor_model_parallel_size ) @@ -184,26 +178,20 @@ def get_batch(data_iterator): else: data = None - input_ids = tensor_parallel.broadcast_data(["input_ids"], data, torch.int64)[ - "input_ids" + input_ids = tensor_parallel.broadcast_data(["input_ids"], data, torch.int64)["input_ids"] + input_ids_shape = tensor_parallel.broadcast_data(["input_ids_shape"], data, torch.int64)[ + "input_ids_shape" ] - input_ids_shape = tensor_parallel.broadcast_data( - ["input_ids_shape"], data, torch.int64 - )["input_ids_shape"] labels = tensor_parallel.broadcast_data(["labels"], data, torch.int64)["labels"] labels_shape = tensor_parallel.broadcast_data(["labels_shape"], data, torch.int64)[ "labels_shape" ] images = tensor_parallel.broadcast_data(["images"], data, torch.float32)["images"] - split_image_sizes = tensor_parallel.broadcast_data( - ["split_image_sizes"], data, torch.int64 - )["split_image_sizes"] - image_sizes = tensor_parallel.broadcast_data(["image_sizes"], data, torch.int64)[ - "image_sizes" - ] - modalities = tensor_parallel.broadcast_data(["modalities"], data, torch.int64)[ - "modalities" + split_image_sizes = tensor_parallel.broadcast_data(["split_image_sizes"], data, torch.int64)[ + "split_image_sizes" ] + image_sizes = tensor_parallel.broadcast_data(["image_sizes"], data, torch.int64)["image_sizes"] + modalities = tensor_parallel.broadcast_data(["modalities"], data, torch.int64)["modalities"] # Convert to list # input_ids to list @@ -211,9 +199,7 @@ def get_batch(data_iterator): start_idx = 0 for shape in input_ids_shape: num_elements = torch.prod(shape).item() - sub_tensor = input_ids[start_idx : start_idx + num_elements].reshape( - shape.tolist() - ) + sub_tensor = input_ids[start_idx : start_idx + num_elements].reshape(shape.tolist()) input_ids_list.append(sub_tensor) start_idx += num_elements assert start_idx == input_ids.numel() @@ -224,9 +210,7 @@ def get_batch(data_iterator): start_idx = 0 for shape in labels_shape: num_elements = torch.prod(shape).item() - sub_tensor = labels[start_idx : start_idx + num_elements].reshape( - shape.tolist() - ) + sub_tensor = labels[start_idx : start_idx + num_elements].reshape(shape.tolist()) labels_list.append(sub_tensor) start_idx += num_elements assert start_idx == labels.numel() @@ -237,9 +221,7 @@ def get_batch(data_iterator): start_idx = 0 for shape in split_image_sizes: num_elements = torch.prod(shape).item() - sub_tensor = images[start_idx : start_idx + num_elements].reshape( - shape.tolist() - ) + sub_tensor = images[start_idx : start_idx + num_elements].reshape(shape.tolist()) images_list.append(sub_tensor) start_idx += num_elements assert start_idx == images.numel() @@ -284,9 +266,7 @@ def get_batch(data_iterator): padding_value=tokenizer.pad_token_id, tokenizer=tokenizer, ) - labels = pad_sequence( - labels, batch_first=True, padding_value=IGNORE_INDEX, tokenizer=tokenizer - ) + labels = pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX, tokenizer=tokenizer) # Attention mask same as LLaVA-NeXT attention_mask = input_ids.ne(tokenizer.pad_token_id) cur_platform.range_pop() @@ -325,9 +305,7 @@ def loss_func(labels: torch.Tensor, loss_mask: torch.Tensor, logits: torch.Tenso shift_logits = logits[:-1, :, :].contiguous() shift_labels = labels[1:, ...].contiguous() - losses = tensor_parallel.vocab_parallel_cross_entropy( - shift_logits.float(), shift_labels - ) + losses = tensor_parallel.vocab_parallel_cross_entropy(shift_logits.float(), shift_labels) losses = losses.transpose(0, 1).contiguous().float() if loss_mask is not None: loss_mask = loss_mask[..., 1:].contiguous() @@ -357,13 +335,9 @@ def forward_step(data_iterator, model: LLaVAOneVisionModel): # Get the batch. timers("batch-generator", log_level=2).start() - input_ids, labels, attention_mask, images, image_sizes, modalities = get_batch( - data_iterator - ) + input_ids, labels, attention_mask, images, image_sizes, modalities = get_batch(data_iterator) if "text" in modalities and ("image" in modalities or "video" in modalities): - raise ValueError( - "Both text and other modalities are present in the same batch." - ) + raise ValueError("Both text and other modalities are present in the same batch.") timers("batch-generator").stop() output_tensor, labels, loss_mask = model( @@ -390,9 +364,7 @@ def add_multimodal_extra_args(parser): group.add_argument("--freeze-LM", action="store_true", default=False) group.add_argument("--freeze-ViT", action="store_true", default=False) group.add_argument("--language-model-type", type=str, required=True) - group.add_argument( - "--disable-vision-class-token", action="store_true", default=False - ) + group.add_argument("--disable-vision-class-token", action="store_true", default=False) group.add_argument( "--allow-missing-vision-projection-checkpoint", action="store_true", @@ -441,15 +413,11 @@ def add_multimodal_extra_args(parser): group.add_argument( "--pos-skipping-range", type=int, default=4096, help="Position skipping range" ) - group.add_argument( - "--add-faster-video", default=False, help="Whetehr add fatser video token" - ) + group.add_argument("--add-faster-video", default=False, help="Whetehr add fatser video token") group.add_argument( "--mm-spatial-pool-mode", type=str, default="bilinear", help="Spatial pool mode" ) - group.add_argument( - "--mm-newline-position", type=str, default="grid", help="Newline position." - ) + group.add_argument("--mm-newline-position", type=str, default="grid", help="Newline position.") return parser diff --git a/flagscale/train/megatron/train_qwen2_5_vl.py b/flagscale/train/megatron/train_qwen2_5_vl.py index fdc26b3b7c..e3005051ef 100644 --- a/flagscale/train/megatron/train_qwen2_5_vl.py +++ b/flagscale/train/megatron/train_qwen2_5_vl.py @@ -73,6 +73,7 @@ from megatron.training.training import pretrain from megatron.plugin.platform import get_platform + cur_platform = get_platform() stimer = StragglerDetector() @@ -134,7 +135,7 @@ def model_provider( args.rotary_seq_len_interpolation_factor is not None or args.rotary_seq_len_interpolation_factor != 1 ): - print_rank_0('Multimodal RoPE currently not support RoPE interpolation, set to None...') + print_rank_0("Multimodal RoPE currently not support RoPE interpolation, set to None...") args.rotary_seq_len_interpolation_factor = None vision_config = get_vision_model_config(args, deepcopy(config)) @@ -162,7 +163,7 @@ def model_provider( drop_vision_class_token=False, # NOTE: no class token to drop? vision_projection_config=vision_projector_config, vision_projection_layer_spec=vision_projector_spec, - vision_projection_type='mlp', + vision_projection_type="mlp", allow_missing_vision_projection_checkpoint=args.allow_missing_vision_projection_checkpoint, language_position_embedding_type=args.position_embedding_type, language_rotary_percent=args.rotary_percent, @@ -461,8 +462,8 @@ def get_batch(data_iterator): # shape: n_video_samples video_thw_grids = broadcast_data(["video_thw_grids"], data, torch.long)["video_thw_grids"] # shape: n_video_samples - second_per_grid_ts = broadcast_data(['second_per_grid_ts'], data, torch.float32)[ - 'second_per_grid_ts' + second_per_grid_ts = broadcast_data(["second_per_grid_ts"], data, torch.float32)[ + "second_per_grid_ts" ] image_input_mask = broadcast_data(["image_input_mask"], data, torch.bool)["image_input_mask"] @@ -522,7 +523,7 @@ def loss_func( """ args = get_args() - if has_nvidia_modelopt and getattr(args, 'modelopt_enabled', False): # [ModelOpt] + if has_nvidia_modelopt and getattr(args, "modelopt_enabled", False): # [ModelOpt] return loss_func_modelopt(loss_mask, output_tensor, model=model) losses = output_tensor.view(-1).float() @@ -563,7 +564,7 @@ def loss_func( num_tokens = loss_mask.sum().clone().detach().to(torch.int) reporting_loss = torch.cat([loss.clone().detach().view(1), num_tokens.view(1)]) - return (loss, num_tokens, {'lm loss': reporting_loss}) + return (loss, num_tokens, {"lm loss": reporting_loss}) def forward_step(data_iterator, model: Qwen2_5VLModel): @@ -577,7 +578,7 @@ def forward_step(data_iterator, model: Qwen2_5VLModel): timers = get_timers() # Get the batch. - timers('batch-generator', log_level=2).start() + timers("batch-generator", log_level=2).start() global stimer with stimer(bdata=True): ( @@ -593,7 +594,7 @@ def forward_step(data_iterator, model: Qwen2_5VLModel): image_input_mask, video_input_mask, ) = get_batch(data_iterator) - timers('batch-generator').stop() + timers("batch-generator").stop() vision_data = torch.cat([imgs, videos], dim=0) vision_grid = torch.cat([image_thw_grids, video_thw_grids], dim=0) with stimer: @@ -879,7 +880,7 @@ def add_multimodal_extra_args(parser): model_provider, ModelType.encoder_or_decoder, forward_step, - args_defaults={'tokenizer_type': 'Qwen2VLTokenizer'}, + args_defaults={"tokenizer_type": "Qwen2VLTokenizer"}, extra_args_provider=add_multimodal_extra_args, process_non_loss_data_func=write_online_eval_to_tensorboard, non_loss_data_func=run_online_eval, diff --git a/flagscale/train/megatron/train_qwen35.py b/flagscale/train/megatron/train_qwen35.py index 0a38f8f55d..9f5faab6f2 100644 --- a/flagscale/train/megatron/train_qwen35.py +++ b/flagscale/train/megatron/train_qwen35.py @@ -47,10 +47,12 @@ has_nvidia_modelopt = False from megatron.training.training import pretrain + stimer = StragglerDetector() # Qwen2.5-VL data handling from megatron.core.num_microbatches_calculator import get_num_microbatches + torch._dynamo.config.suppress_errors = True from megatron.core.parallel_state import ( get_tensor_model_parallel_rank, @@ -81,11 +83,12 @@ from flagscale.models.megatron.qwen35.layer_specs import ( get_qwen35_language_model_spec, get_qwen35_mtp_block_spec, - get_mlp_module_spec + get_mlp_module_spec, ) from flagscale.models.megatron.qwen3_vl.layer_specs import get_qwen3vl_vision_model_spec from megatron.plugin.platform import get_platform + cur_platform = get_platform() from tools.datasets.qwenvl.data.dataset_helpers import TaskEncoder, print_error_handler @@ -104,13 +107,18 @@ def model_provider( config = core_transformer_config_from_args(args, Qwen35TransformerConfig) # Qwen3.5 uses zero-centered gamma for RMSNorm; override if needed # (core_transformer_config_from_args may be affected by apply_layernorm_1p) - config.layernorm_zero_centered_gamma = getattr(args, 'layernorm_zero_centered_gamma', True) + config.layernorm_zero_centered_gamma = getattr(args, "layernorm_zero_centered_gamma", True) use_te = args.transformer_impl == "transformer_engine" if not use_te: raise NotImplementedError("Qwen3.5 model is only implemented with TransformerEngine!") - if args.rotary_seq_len_interpolation_factor is not None or args.rotary_seq_len_interpolation_factor != 1: - print_rank_0('Multimodal RoPE currently does not support RoPE interpolation, set to None...') + if ( + args.rotary_seq_len_interpolation_factor is not None + or args.rotary_seq_len_interpolation_factor != 1 + ): + print_rank_0( + "Multimodal RoPE currently does not support RoPE interpolation, set to None..." + ) args.rotary_seq_len_interpolation_factor = None # Vision configs (identical encoder to Qwen3-VL) @@ -142,22 +150,18 @@ def model_provider( language_transformer_layer_spec=language_layer_spec, language_vocab_size=args.padded_vocab_size, language_max_sequence_length=args.max_position_embeddings, - vision_transformer_config=vision_config, vision_transformer_layer_spec=vision_model_spec, vision_projection_config=vision_projector_config, vision_projection_layer_spec=vision_projector_spec, - vision_projection_type='mlp', - + vision_projection_type="mlp", language_position_embedding_type=args.position_embedding_type, language_rotary_percent=args.rotary_percent, language_rotary_base=args.rotary_base, - pre_process=pre_process, post_process=post_process, add_decoder=add_decoder, add_encoder=add_encoder, - fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, parallel_output=True, language_share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, @@ -204,9 +208,7 @@ def get_ltor_masks_and_position_ids( return attention_mask, loss_mask, position_ids -def get_batch( - data_iterator, model: Qwen35Model = None -) -> Tuple: +def get_batch(data_iterator, model: Qwen35Model = None) -> Tuple: """Generate a batch.""" imgs = None tokens = None @@ -239,7 +241,9 @@ def get_batch( cur_platform.empty_cache() video_thw_grids = broadcast_data(["video_thw_grids"], data, torch.long)["video_thw_grids"] - second_per_grid_ts = broadcast_data(['second_per_grid_ts'], data, torch.float32)['second_per_grid_ts'] + second_per_grid_ts = broadcast_data(["second_per_grid_ts"], data, torch.float32)[ + "second_per_grid_ts" + ] image_input_mask = broadcast_data(["image_input_mask"], data, torch.bool)["image_input_mask"] video_input_mask = broadcast_data(["video_input_mask"], data, torch.bool)["video_input_mask"] cur_platform.range_pop() @@ -332,7 +336,7 @@ def loss_func( num_tokens = loss_mask.sum().clone().detach().to(torch.int) reporting_loss = torch.cat([loss.clone().detach().view(1), num_tokens.view(1)]) - return (loss, num_tokens, {'lm loss': reporting_loss}) + return (loss, num_tokens, {"lm loss": reporting_loss}) def forward_step(data_iterator, model: Qwen35Model): @@ -340,7 +344,7 @@ def forward_step(data_iterator, model: Qwen35Model): args = get_args() timers = get_timers() - timers('batch-generator', log_level=2).start() + timers("batch-generator", log_level=2).start() global stimer with stimer(bdata=True): ( @@ -356,7 +360,7 @@ def forward_step(data_iterator, model: Qwen35Model): image_input_mask, video_input_mask, ) = get_batch(data_iterator, model=unwrap_model(model)) - timers('batch-generator').stop() + timers("batch-generator").stop() vision_data = torch.cat([imgs, videos], dim=0) vision_grid = torch.cat([image_thw_grids, video_thw_grids], dim=0) @@ -472,7 +476,9 @@ def train_valid_test_dataloaders_provider(train_val_test_num_samples): ) if os.path.exists(data_save_name): try: - dataset_state_dict = torch.load(data_save_name, map_location="cpu", weights_only=False) + dataset_state_dict = torch.load( + data_save_name, map_location="cpu", weights_only=False + ) train_dataloader.restore_state_rank(dataset_state_dict["dataloader_state_dict"]) print_rank_0(f"restored dataset state from {data_save_name}") except Exception as e: @@ -492,6 +498,7 @@ def train_valid_test_dataloaders_provider(train_val_test_num_samples): class EnergonDataloader: """Wrapper for Megatron Energon dataloader.""" + def __init__(self, dataloader): self._dataloader = dataloader self._iter = iter(cyclic_iter(dataloader)) @@ -530,7 +537,9 @@ def add_qwen35_extra_args(parser): group.add_argument("--image-max-pixels", type=int, default=768 * 768) group.add_argument("--image-min-pixels", type=int, default=32 * 32) group.add_argument("--vision-recompute-activations", action="store_true", default=False) - group.add_argument("--no-use-system-prompt", dest="use_system_prompt", action="store_false", default=True) + group.add_argument( + "--no-use-system-prompt", dest="use_system_prompt", action="store_false", default=True + ) group.add_argument( "--convert-checkpoint-from-megatron-to-transformers", action="store_true", @@ -563,7 +572,7 @@ def add_qwen35_extra_args(parser): model_provider, ModelType.encoder_or_decoder, forward_step, - args_defaults={'tokenizer_type': 'Qwen2VLTokenizer'}, + args_defaults={"tokenizer_type": "Qwen2VLTokenizer"}, extra_args_provider=add_qwen35_extra_args, process_non_loss_data_func=write_online_eval_to_tensorboard, non_loss_data_func=run_online_eval, diff --git a/flagscale/train/megatron/train_qwen3_vl.py b/flagscale/train/megatron/train_qwen3_vl.py index 62066a218b..84e1f910cd 100644 --- a/flagscale/train/megatron/train_qwen3_vl.py +++ b/flagscale/train/megatron/train_qwen3_vl.py @@ -30,7 +30,7 @@ # torch.serialization.add_safe_globals([Namespace]) from megatron.core import parallel_state -from megatron.training.checkpointing import get_checkpoint_name # for dataloder +from megatron.training.checkpointing import get_checkpoint_name # for dataloder from megatron.core.enums import ModelType @@ -54,12 +54,18 @@ has_nvidia_modelopt = False from megatron.training.training import pretrain + stimer = StragglerDetector() #### especially for qwen2.5-vl #### from megatron.core.num_microbatches_calculator import get_num_microbatches + torch._dynamo.config.suppress_errors = True -from megatron.core.parallel_state import get_tensor_model_parallel_rank, get_pipeline_model_parallel_world_size, get_pipeline_model_parallel_rank +from megatron.core.parallel_state import ( + get_tensor_model_parallel_rank, + get_pipeline_model_parallel_world_size, + get_pipeline_model_parallel_rank, +) from megatron.energon import ( LimitDataset, RepeatDataset, @@ -75,22 +81,28 @@ from flagscale.models.megatron.qwen2_5_vl.tensor_parallel import broadcast_data -from flagscale.models.megatron.qwen3_vl.layer_specs import (get_gpt_layer_with_transformer_engine_spec, - get_qwen3vl_vision_model_spec, - get_mlp_module_spec) +from flagscale.models.megatron.qwen3_vl.layer_specs import ( + get_gpt_layer_with_transformer_engine_spec, + get_qwen3vl_vision_model_spec, + get_mlp_module_spec, +) from flagscale.models.megatron.qwen3_vl.model import Qwen3VLModel from flagscale.models.megatron.qwen3_vl.transformer_config import ( Qwen3VLTransformerConfig, get_vision_model_config, - get_vision_projection_config + get_vision_projection_config, ) from megatron.plugin.platform import get_platform + cur_platform = get_platform() from tools.datasets.qwenvl.data.dataset_helpers import TaskEncoder, print_error_handler + #### especially for qwen2.5-vl #### -IGNORE_IDX=-100 +IGNORE_IDX = -100 + + def model_provider( pre_process=True, post_process=True, add_encoder=True, add_decoder=True ) -> Union[Qwen3VLModel]: @@ -103,21 +115,27 @@ def model_provider( if not use_te: raise NotImplementedError("The Qwen3-VL model is only implemented with TransformerEngine!") - if args.rotary_seq_len_interpolation_factor is not None or args.rotary_seq_len_interpolation_factor != 1: - print_rank_0('Multimodal RoPE currently not support RoPE interpolation, set to None...') + if ( + args.rotary_seq_len_interpolation_factor is not None + or args.rotary_seq_len_interpolation_factor != 1 + ): + print_rank_0("Multimodal RoPE currently not support RoPE interpolation, set to None...") args.rotary_seq_len_interpolation_factor = None vision_config = get_vision_model_config(args, deepcopy(config)) vision_config.pipeline_model_parallel_size = 1 vision_config.first_pipeline_num_layers = None - vision_projector_config = get_vision_projection_config(deepcopy(config), vision_config.hidden_size, args.spatial_merge_size) + vision_projector_config = get_vision_projection_config( + deepcopy(config), vision_config.hidden_size, args.spatial_merge_size + ) print_rank_0("building Qwen3-VL model in TE...") # Layer Specs of vit, llm and projector transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec( num_experts=args.num_experts, moe_grouped_gemm=args.moe_grouped_gemm, - qk_layernorm=args.qk_layernorm) + qk_layernorm=args.qk_layernorm, + ) vision_model_spec = get_qwen3vl_vision_model_spec() vision_projector_spec = get_mlp_module_spec(add_norm=False).submodules if args.enable_variable_seq_lengths: @@ -128,25 +146,20 @@ def model_provider( language_transformer_layer_spec=transformer_layer_spec, language_vocab_size=args.padded_vocab_size, language_max_sequence_length=args.max_position_embeddings, - vision_transformer_config=vision_config, vision_transformer_layer_spec=vision_model_spec, # drop_vision_class_token=False, # no use - vision_projection_config=vision_projector_config, vision_projection_layer_spec=vision_projector_spec, - vision_projection_type='mlp', + vision_projection_type="mlp", # allow_missing_vision_projection_checkpoint= args.allow_missing_vision_projection_checkpoint, - language_position_embedding_type=args.position_embedding_type, language_rotary_percent=args.rotary_percent, language_rotary_base=args.rotary_base, - pre_process=pre_process, post_process=post_process, add_decoder=add_decoder, add_encoder=add_encoder, - fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, parallel_output=True, language_share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, @@ -155,28 +168,29 @@ def model_provider( model.freeze( freeze_language_model=args.freeze_LM, freeze_vision_model=args.freeze_ViT, - freeze_vision_projection=False + freeze_vision_projection=False, ) return model + def get_ltor_masks_and_position_ids( - input_ids, - image_thw_grids, - video_thw_grids, - target, - pad_token, - second_per_grid_ts, - ignore_index=None, - model: Qwen3VLModel = None - ): + input_ids, + image_thw_grids, + video_thw_grids, + target, + pad_token, + second_per_grid_ts, + ignore_index=None, + model: Qwen3VLModel = None, +): """Build masks and position id for left to right model.""" # Position ids. [3 X bs X seqlen] position_ids, _ = model.get_rope_index( input_ids=input_ids, image_grid_thw=image_thw_grids, video_grid_thw=video_thw_grids, - attention_mask=input_ids != pad_token + attention_mask=input_ids != pad_token, ) # Loss mask. @@ -190,6 +204,7 @@ def get_ltor_masks_and_position_ids( return attention_mask, loss_mask, position_ids + def get_batch(data_iterator, model: Qwen3VLModel = None) -> Tuple: """Generate a batch""" imgs = None @@ -207,15 +222,16 @@ def get_batch(data_iterator, model: Qwen3VLModel = None) -> Tuple: pad_token_id = IGNORE_IDX # while (data["target"] == pad_token_id).all() or (data["target"].shape[-1] < 986 or data["target"].shape[-1] > 1000): # for debug while (data["target"] == pad_token_id).all(): - logging.getLogger(__name__).warning("The current data is invalid because the target is all pad_token_id! Get next data to avoid fail, but it's better to check the data!") + logging.getLogger(__name__).warning( + "The current data is invalid because the target is all pad_token_id! Get next data to avoid fail, but it's better to check the data!" + ) data = next(data_iterator) else: data = None + data_text = broadcast_data(["text"], data, torch.int64)["text"] - data_text = broadcast_data(["text"], data, torch.int64)["text"] - - target = broadcast_data(["target"], data, torch.int64)["target"] + target = broadcast_data(["target"], data, torch.int64)["target"] # shape: num_tiles x c x h x w imgs = broadcast_data(["imgs"], data, torch.float32)["imgs"] @@ -231,8 +247,9 @@ def get_batch(data_iterator, model: Qwen3VLModel = None) -> Tuple: # shape: n_video_samples video_thw_grids = broadcast_data(["video_thw_grids"], data, torch.long)["video_thw_grids"] # shape: n_video_samples - second_per_grid_ts = broadcast_data(['second_per_grid_ts'], data, torch.float32)['second_per_grid_ts'] - + second_per_grid_ts = broadcast_data(["second_per_grid_ts"], data, torch.float32)[ + "second_per_grid_ts" + ] image_input_mask = broadcast_data(["image_input_mask"], data, torch.bool)["image_input_mask"] video_input_mask = broadcast_data(["video_input_mask"], data, torch.bool)["video_input_mask"] @@ -250,7 +267,14 @@ def get_batch(data_iterator, model: Qwen3VLModel = None) -> Tuple: # NOTE: no sequence packing in LLM inputs cur_platform.range_push("get_ltor_masks_and_position_ids") attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids( - tokens, image_thw_grids, video_thw_grids, labels, pad_token=tokenizer.pad_token_id, second_per_grid_ts=second_per_grid_ts, ignore_index=IGNORE_IDX, model=model, + tokens, + image_thw_grids, + video_thw_grids, + labels, + pad_token=tokenizer.pad_token_id, + second_per_grid_ts=second_per_grid_ts, + ignore_index=IGNORE_IDX, + model=model, ) cur_platform.range_pop() @@ -265,9 +289,10 @@ def get_batch(data_iterator, model: Qwen3VLModel = None) -> Tuple: image_thw_grids, video_thw_grids, image_input_mask, - video_input_mask + video_input_mask, ) + # define spiky loss as a loss that's 10x the max loss observed SPIKY_LOSS_FACTOR = 10 @@ -290,7 +315,7 @@ def loss_func( """ args = get_args() - if has_nvidia_modelopt and getattr(args, 'modelopt_enabled', False): # [ModelOpt] + if has_nvidia_modelopt and getattr(args, "modelopt_enabled", False): # [ModelOpt] return loss_func_modelopt(loss_mask, output_tensor, model=model) losses = output_tensor.view(-1).float() @@ -331,7 +356,7 @@ def loss_func( num_tokens = loss_mask.sum().clone().detach().to(torch.int) reporting_loss = torch.cat([loss.clone().detach().view(1), num_tokens.view(1)]) - return (loss, num_tokens, {'lm loss': reporting_loss}) + return (loss, num_tokens, {"lm loss": reporting_loss}) def forward_step(data_iterator, model: Qwen3VLModel): @@ -345,7 +370,7 @@ def forward_step(data_iterator, model: Qwen3VLModel): timers = get_timers() # Get the batch. - timers('batch-generator', log_level=2).start() + timers("batch-generator", log_level=2).start() global stimer with stimer(bdata=True): ( @@ -359,31 +384,33 @@ def forward_step(data_iterator, model: Qwen3VLModel): image_thw_grids, video_thw_grids, image_input_mask, - video_input_mask + video_input_mask, ) = get_batch(data_iterator, model=unwrap_model(model)) - timers('batch-generator').stop() + timers("batch-generator").stop() vision_data = torch.cat([imgs, videos], dim=0) vision_grid = torch.cat([image_thw_grids, video_thw_grids], dim=0) with stimer: output_tensor = model( - input_ids = tokens, - position_ids = position_ids, - vision_data = vision_data, - vision_grid_thw = vision_grid, - video_start_index = image_input_mask.sum().cpu().item(), - image_input_mask = image_input_mask, - video_input_mask = video_input_mask, - attention_mask = attention_mask, - labels = labels + input_ids=tokens, + position_ids=position_ids, + vision_data=vision_data, + vision_grid_thw=vision_grid, + video_start_index=image_input_mask.sum().cpu().item(), + image_input_mask=image_input_mask, + video_input_mask=video_input_mask, + attention_mask=attention_mask, + labels=labels, ) return output_tensor, partial(loss_func, loss_mask, model=model) + def run_online_eval(model): """Run an evaluation benchmark during training.""" # Do nothing. return [] + def write_online_eval_to_tensorboard(data, iteration, writer): """Write online evaluation data to Tensorboard.""" if not writer: @@ -393,6 +420,7 @@ def write_online_eval_to_tensorboard(data, iteration, writer): for k, v in item.items(): writer.add_scalar(k, v, iteration) + def datasets_provider(worker_config=None): """Create multimodal train, validation and test datasets.""" args = get_args() @@ -403,8 +431,8 @@ def datasets_provider(worker_config=None): task_encoder=TaskEncoder(), worker_config=worker_config, virtual_epoch_length=0, - max_samples_per_sequence=args.max_samples_per_sequence, # sequential shuffle in a tar - shuffle_buffer_size=args.shuffle_buffer_size, # shuffle in a sequential + max_samples_per_sequence=args.max_samples_per_sequence, # sequential shuffle in a tar + shuffle_buffer_size=args.shuffle_buffer_size, # shuffle in a sequential handler=print_error_handler, repeat=True, image_decode="pil", @@ -435,24 +463,26 @@ def datasets_provider(worker_config=None): return train_dataset, val_datasets_without_source_datasets, None + def is_first_or_last_stage(pp_size, transformer_pipeline_model_parallel_size): """Check if the current pipeline parallel stage is the first or last stage.""" - if pp_size == 1: # No pipeline parallelism. + if pp_size == 1: # No pipeline parallelism. return True is_valid_rank = False pp_rank = get_pipeline_model_parallel_rank() if transformer_pipeline_model_parallel_size == 0: # No separate pipeline stage for the vision model. Run the dataloader on the first and last pipeline stage. - is_valid_rank = pp_rank in (0, pp_size-1) + is_valid_rank = pp_rank in (0, pp_size - 1) elif transformer_pipeline_model_parallel_size == 1: # Separate pipeline stage for the vision model. Run the dataloader on the first vision and LM stage and last LM stage. - is_valid_rank = pp_rank in (0, 1, pp_size-1) + is_valid_rank = pp_rank in (0, 1, pp_size - 1) else: raise NotImplementedError("encoder-pipeline-model-parallel-size > 1 is not supported yet") return is_valid_rank + def is_dataloader_rank(transformer_pipeline_model_parallel_size): """Check if we should have the dataloader on this tensor and pipeline parallel rank.""" # Run dataloader only on the first tensor parallel rank (will be broadcasted to others). @@ -464,6 +494,7 @@ def is_dataloader_rank(transformer_pipeline_model_parallel_size): return is_first_rank + def train_valid_test_dataloaders_provider(train_val_test_num_samples): """Build multimodal train, validation and test dataloaders.""" args = get_args() @@ -495,12 +526,14 @@ def train_valid_test_dataloaders_provider(train_val_test_num_samples): data_save_name = get_checkpoint_name( args.dataloader_save, args.iteration, - pipeline_rank=0, # Only the first pipeline parallel rank stores the dataloader checkpoint. + pipeline_rank=0, # Only the first pipeline parallel rank stores the dataloader checkpoint. basename=f"train_dataloader_dprank{dp_rank:03d}.pt", ) if os.path.exists(data_save_name): try: - dataset_state_dict = torch.load(data_save_name, map_location="cpu", weights_only=False) + dataset_state_dict = torch.load( + data_save_name, map_location="cpu", weights_only=False + ) train_dataloader.restore_state_rank(dataset_state_dict["dataloader_state_dict"]) print_rank_0(f"restored dataset state from {data_save_name}") except Exception as e: @@ -513,12 +546,14 @@ def train_valid_test_dataloaders_provider(train_val_test_num_samples): ] else: valid_dataloader = EnergonDataloader(None) - test_dataloader = None # NOTE: no test + test_dataloader = None # NOTE: no test return EnergonDataloader(train_dataloader), valid_dataloader, EnergonDataloader(test_dataloader) + class EnergonDataloader: """A wrapper to use Megatron Energon dataloader with the Megatron-LM training loop.""" + def __init__(self, dataloader): self._dataloader = dataloader self._iter = iter(cyclic_iter(dataloader)) @@ -542,7 +577,12 @@ def cyclic_iter(iter): def add_multimodal_extra_args(parser): """Extra arguments.""" group = parser.add_argument_group(title="multimodal arguments") - group.add_argument("--disable-vision-class-token", action="store_true", default=False, help="Disable vision class token") + group.add_argument( + "--disable-vision-class-token", + action="store_true", + default=False, + help="Disable vision class token", + ) group.add_argument( "--dataloader-save", type=str, default=None, help="Energon dataloader state save path" ) @@ -553,20 +593,59 @@ def add_multimodal_extra_args(parser): group.add_argument("--temporal-patch-size", type=int, default=2) group.add_argument("--patch-size", type=int, default=16) group.add_argument("--max-padding-length", type=int, default=2048) - group.add_argument("--enable-variable-seq-lengths", action="store_true", default=False, help="Enable variable sequence lengths") - group.add_argument("--vision-root", type=str, default = None, help="The vision dirctory root path.") - group.add_argument("--max-samples-per-sequence", type=int, default=2**31-1, help="max sequencial seqence samples in a slice") - group.add_argument("--shuffle-buffer-size", type=int, default=0, help="the buffer size to shuffle the samples in a seqence") + group.add_argument( + "--enable-variable-seq-lengths", + action="store_true", + default=False, + help="Enable variable sequence lengths", + ) + group.add_argument( + "--vision-root", type=str, default=None, help="The vision dirctory root path." + ) + group.add_argument( + "--max-samples-per-sequence", + type=int, + default=2**31 - 1, + help="max sequencial seqence samples in a slice", + ) + group.add_argument( + "--shuffle-buffer-size", + type=int, + default=0, + help="the buffer size to shuffle the samples in a seqence", + ) # learning rate - group.add_argument("--vision-ration", type=float, default=0.1, help="the learning rate ration of vision(inlude merger) compared with llm") - group.add_argument("--image-max-pixels", type=int, default=768*768, help="the maximum pixels of a single image") - group.add_argument("--image-min-pixels", type=int, default=32*32, help="the minimum pixels of a single image") + group.add_argument( + "--vision-ration", + type=float, + default=0.1, + help="the learning rate ration of vision(inlude merger) compared with llm", + ) + group.add_argument( + "--image-max-pixels", + type=int, + default=768 * 768, + help="the maximum pixels of a single image", + ) + group.add_argument( + "--image-min-pixels", type=int, default=32 * 32, help="the minimum pixels of a single image" + ) # vision model recompute - group.add_argument("--vision-recompute-activations", action="store_true", default=False, help="Recompute vision model activations") + group.add_argument( + "--vision-recompute-activations", + action="store_true", + default=False, + help="Recompute vision model activations", + ) # data processing - group.add_argument("--no-use-system-prompt", dest="use_system_prompt", action="store_false", default=True, help="Don't use system prompt") - + group.add_argument( + "--no-use-system-prompt", + dest="use_system_prompt", + action="store_false", + default=True, + help="Don't use system prompt", + ) # just for checkpoint conversion group.add_argument( @@ -577,15 +656,21 @@ def add_multimodal_extra_args(parser): "If False, convert a Transformers checkpoint to a Megatron checkpoint." ), ) - group.add_argument("--freeze-LM", action="store_true", default=False, help="Freeze the language model") - group.add_argument("--freeze-ViT", action="store_true", default=False, help="Freeze the vision model") + group.add_argument( + "--freeze-LM", action="store_true", default=False, help="Freeze the language model" + ) + group.add_argument( + "--freeze-ViT", action="store_true", default=False, help="Freeze the vision model" + ) group.add_argument( "--allow-missing-vision-projection-checkpoint", action="store_true", default=False, help="Allow missing vision projection checkpoint", ) - group.add_argument("--use-te", action="store_true", default=False, help="Use transformer engine") + group.add_argument( + "--use-te", action="store_true", default=False, help="Use transformer engine" + ) return parser @@ -597,7 +682,7 @@ def add_multimodal_extra_args(parser): model_provider, ModelType.encoder_or_decoder, forward_step, - args_defaults={'tokenizer_type': 'Qwen2VLTokenizer'}, + args_defaults={"tokenizer_type": "Qwen2VLTokenizer"}, extra_args_provider=add_multimodal_extra_args, process_non_loss_data_func=write_online_eval_to_tensorboard, non_loss_data_func=run_online_eval, diff --git a/flagscale/train/megatron/train_robobrain_x0.5_qwengroot.py b/flagscale/train/megatron/train_robobrain_x0.5_qwengroot.py index 87fc9323ef..79cf0f616a 100644 --- a/flagscale/train/megatron/train_robobrain_x0.5_qwengroot.py +++ b/flagscale/train/megatron/train_robobrain_x0.5_qwengroot.py @@ -29,6 +29,7 @@ from flagscale.models.robobrain_x.qwen_groot import Qwen_GR00T from megatron.plugin.platform import get_platform + cur_platform = get_platform() # Sane Defaults @@ -173,7 +174,11 @@ def main(cfg) -> None: while not done: batch = next(data_iter) - qwen_inputs, state, actions = batch.get("qwen_inputs"), batch.get("state"), batch.get("actions") + qwen_inputs, state, actions = ( + batch.get("qwen_inputs"), + batch.get("state"), + batch.get("actions"), + ) if not qwen_inputs or not actions: continue for i in qwen_inputs: diff --git a/flagscale/train/megatron/train_robobrain_x0.py b/flagscale/train/megatron/train_robobrain_x0.py index 9473bb4265..a6d0094307 100644 --- a/flagscale/train/megatron/train_robobrain_x0.py +++ b/flagscale/train/megatron/train_robobrain_x0.py @@ -122,8 +122,10 @@ LAST_LARGE_IMG = False from megatron.plugin.platform import get_platform + cur_platform = get_platform() + def model_provider( pre_process=True, post_process=True, add_encoder=True, add_decoder=True ) -> Union[Qwen2_5VLModel]: @@ -562,7 +564,7 @@ def loss_func( """ args = get_args() - if has_nvidia_modelopt and getattr(args, 'modelopt_enabled', False): # [ModelOpt] + if has_nvidia_modelopt and getattr(args, "modelopt_enabled", False): # [ModelOpt] return loss_func_modelopt(loss_mask, output_tensor, model=model) losses = output_tensor.view(-1).float() diff --git a/flagscale/train/megatron/train_rwkv.py b/flagscale/train/megatron/train_rwkv.py index 4f60058447..f78c6db5f3 100644 --- a/flagscale/train/megatron/train_rwkv.py +++ b/flagscale/train/megatron/train_rwkv.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os from functools import partial @@ -106,7 +120,7 @@ def loss_func( """ args = get_args() - if has_nvidia_modelopt and getattr(args, 'modelopt_enabled', False): # [ModelOpt] + if has_nvidia_modelopt and getattr(args, "modelopt_enabled", False): # [ModelOpt] return loss_func_modelopt(loss_mask, output_tensor, model=model) # Reshape output tensor to 1D for element-wise multiplication. @@ -148,7 +162,7 @@ def loss_func( num_tokens = loss_mask.sum().clone().detach().to(torch.int) reporting_loss = torch.cat([loss.clone().detach().view(1), num_tokens.view(1)]) - return (loss, num_tokens, {'lm loss': reporting_loss}) + return (loss, num_tokens, {"lm loss": reporting_loss}) def forward_step(data_iterator, model: RWKVModel, return_schedule_plan: bool = False): @@ -163,11 +177,11 @@ def forward_step(data_iterator, model: RWKVModel, return_schedule_plan: bool = F timers = get_timers() # Get the batch. - timers('batch-generator', log_level=2).start() + timers("batch-generator", log_level=2).start() global stimer with stimer(bdata=True): tokens, labels, loss_mask, attention_mask, position_ids = get_batch(data_iterator) - timers('batch-generator').stop() + timers("batch-generator").stop() # Safety check for pipeline parallel intermediate stages if tokens is None: @@ -179,9 +193,9 @@ def forward_step(data_iterator, model: RWKVModel, return_schedule_plan: bool = F output_tensor = model(tokens, labels=labels) else: if return_schedule_plan: - assert ( - args.overlap_moe_expert_parallel_comm - ), "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan" + assert args.overlap_moe_expert_parallel_comm, ( + "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan" + ) schedule_plan = model.build_schedule_plan( tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask ) @@ -256,7 +270,6 @@ def is_dataset_built_on_rank(): if __name__ == "__main__": - # Temporary for transition to core datasets train_valid_test_datasets_provider.is_distributed = True @@ -268,6 +281,6 @@ def is_dataset_built_on_rank(): model_provider, ModelType.encoder_or_decoder, forward_step, - args_defaults={'tokenizer_type': 'RWKVTokenizer'}, + args_defaults={"tokenizer_type": "RWKVTokenizer"}, store=store, ) diff --git a/flagscale/train/megatron/training/__init__.py b/flagscale/train/megatron/training/__init__.py index 3546dfd576..de84f97983 100644 --- a/flagscale/train/megatron/training/__init__.py +++ b/flagscale/train/megatron/training/__init__.py @@ -10,9 +10,7 @@ from .global_vars import get_one_logger from .global_vars import get_adlr_autoresume from .global_vars import get_timers -from .initialize import initialize_megatron +from .initialize import initialize_megatron from .training import pretrain, get_model, get_train_valid_test_num_samples, set_startup_timestamps -from .utils import (print_rank_0, - is_last_rank, - print_rank_last) +from .utils import print_rank_0, is_last_rank, print_rank_last diff --git a/flagscale/train/megatron/training/argument_utils.py b/flagscale/train/megatron/training/argument_utils.py index b9f7c7b22d..91c648bd86 100644 --- a/flagscale/train/megatron/training/argument_utils.py +++ b/flagscale/train/megatron/training/argument_utils.py @@ -14,17 +14,20 @@ # TODO: support arg renames + class TypeInferenceError(Exception): """Custom exception type to be conditionally handled by ArgumentGroupFactory.""" + pass + class ArgumentGroupFactory: """Utility that adds an argument group to an ArgumentParser based on the attributes of a dataclass. This utility uses dataclass metadata including type annotations and docstrings to automatically infer the type, default, and other argparse keyword arguments. - You can override or supplement the automatically inferred argparse kwargs for any + You can override or supplement the automatically inferred argparse kwargs for any dataclass field by providing an "argparse_meta" key in the field's metadata dict. The value should be a dict of kwargs that will be passed to ArgumentParser.add_argument(). These metadata kwargs take precedence over the automatically inferred values. @@ -53,13 +56,13 @@ class YourConfig: that require some customized or additional handling. Args: - src_cfg_class: The source dataclass type (not instance) whose fields will be - converted into command-line arguments. Each field's type annotation determines - the argument type, default values become argument defaults, and field-level + src_cfg_class: The source dataclass type (not instance) whose fields will be + converted into command-line arguments. Each field's type annotation determines + the argument type, default values become argument defaults, and field-level docstrings are extracted to populate argument help text. - exclude: Optional list of attribute names from `src_cfg_class` to exclude from + exclude: Optional list of attribute names from `src_cfg_class` to exclude from argument generation. Useful for omitting internal fields, computed properties, - or attributes that should be configured through other means. If None, all + or attributes that should be configured through other means. If None, all dataclass fields will be converted to command-line arguments. Default: None. """ @@ -73,12 +76,12 @@ def _format_arg_name(self, config_attr_name: str, prefix: Optional[str] = None) Args: config_attr_name: dataclass attribute name - prefix: prefix string to add to the dataclass attribute name. e.g. 'no' for bool + prefix: prefix string to add to the dataclass attribute name. e.g. 'no' for bool settings that are default True. A hyphen is added after the prefix. Default: None """ arg_name = config_attr_name if prefix: - arg_name = prefix + '_' + arg_name + arg_name = prefix + "_" + arg_name arg_name = "--" + arg_name.replace("_", "-") return arg_name @@ -88,6 +91,7 @@ def _get_enum_kwargs(self, config_type: enum.EnumMeta) -> dict[str, Any]: With these settings, the user must provide a valid enum value, e.g. 'flash', for `AttnBackend.flash`. """ + def enum_type_handler(cli_arg): return config_type[cli_arg] @@ -111,7 +115,9 @@ def _extract_type(self, config_type: type) -> dict[str, Any]: if origin in [types.UnionType, typing.Union]: # Handle Optional and Union - if type_tuple[1] == type(None): # Optional type. First element is value inside Optional[] + if type_tuple[1] == type( + None + ): # Optional type. First element is value inside Optional[] return self._extract_type(type_tuple[0]) else: raise TypeInferenceError(f"Unions not supported by argparse: {config_type}") @@ -122,17 +128,20 @@ def _extract_type(self, config_type: type) -> dict[str, Any]: kwargs["nargs"] = "+" return kwargs else: - raise TypeInferenceError(f"Multi-type lists not supported by argparse: {config_type}") + raise TypeInferenceError( + f"Multi-type lists not supported by argparse: {config_type}" + ) elif origin is typing.Literal: choices_types = [type(choice) for choice in type_tuple] - assert all([t == choices_types[0] for t in choices_types]), "Type of each choice in a Literal type should all be the same." + assert all([t == choices_types[0] for t in choices_types]), ( + "Type of each choice in a Literal type should all be the same." + ) kwargs = {"type": choices_types[0], "choices": type_tuple} return kwargs else: raise TypeInferenceError(f"Unsupported type: {config_type}") - def _build_argparse_kwargs_from_field(self, attribute: Field) -> dict[str, Any]: """Assemble kwargs for add_argument(). @@ -142,7 +151,9 @@ def _build_argparse_kwargs_from_field(self, attribute: Field) -> dict[str, Any]: argparse_kwargs = {} argparse_kwargs["arg_names"] = [self._format_arg_name(attribute.name)] argparse_kwargs["dest"] = attribute.name - argparse_kwargs["help"] = self.field_docstrings[attribute.name] if attribute.name in self.field_docstrings else "" + argparse_kwargs["help"] = ( + self.field_docstrings[attribute.name] if attribute.name in self.field_docstrings else "" + ) # dataclasses specifies that both should not be set if isinstance(attribute.default, type(dataclasses.MISSING)): @@ -156,7 +167,6 @@ def _build_argparse_kwargs_from_field(self, attribute: Field) -> dict[str, Any]: # save metadata here, but update at the end so the metadata has highest precedence attr_argparse_meta = attribute.metadata["argparse_meta"] - # if we cannot infer the argparse type, all of this logic may fail. we try to defer # to the developer-specified metadata if present try: @@ -164,12 +174,17 @@ def _build_argparse_kwargs_from_field(self, attribute: Field) -> dict[str, Any]: # use store_true or store_false action for enable/disable flags, which doesn't accept a 'type' if argparse_kwargs["type"] == bool: - argparse_kwargs["action"] = "store_true" if attribute.default == False else "store_false" + argparse_kwargs["action"] = ( + "store_true" if attribute.default == False else "store_false" + ) argparse_kwargs.pop("type") # add '--no-*' and '--disable-*' prefix if this is a store_false argument if argparse_kwargs["action"] == "store_false": - argparse_kwargs["arg_names"] = [self._format_arg_name(attribute.name, prefix="no"), self._format_arg_name(attribute.name, prefix="disable")] + argparse_kwargs["arg_names"] = [ + self._format_arg_name(attribute.name, prefix="no"), + self._format_arg_name(attribute.name, prefix="disable"), + ] except TypeInferenceError as e: if attr_argparse_meta is not None: print( @@ -181,7 +196,7 @@ def _build_argparse_kwargs_from_field(self, attribute: Field) -> dict[str, Any]: else: raise e - # metadata provided by field takes precedence + # metadata provided by field takes precedence if attr_argparse_meta is not None: argparse_kwargs.update(attr_argparse_meta) @@ -231,12 +246,16 @@ def _get_field_docstrings(self, src_cfg_class: type) -> dict[str, str]: if a_cond and b_cond: # These should be guaranteed by typechecks above, but assert just in case - assert isinstance(a.target.id, str), "Dataclass attribute not in the expected format. Name is not a string." - assert isinstance(b.value.value, str), "Dataclass attribute docstring is not a string." + assert isinstance(a.target.id, str), ( + "Dataclass attribute not in the expected format. Name is not a string." + ) + assert isinstance(b.value.value, str), ( + "Dataclass attribute docstring is not a string." + ) # Formatting docstring = inspect.cleandoc(b.value.value) - docstring = ' '.join(docstring.split()) + docstring = " ".join(docstring.split()) field_docstrings[a.target.id] = docstring diff --git a/flagscale/train/megatron/training/arguments.py b/flagscale/train/megatron/training/arguments.py index 2105958a79..139dd5bc2f 100644 --- a/flagscale/train/megatron/training/arguments.py +++ b/flagscale/train/megatron/training/arguments.py @@ -48,10 +48,12 @@ from megatron.training.arguments_fs import add_flagscale_arguments from megatron.plugin.platform import get_platform + cur_platform = get_platform() + def add_megatron_arguments(parser: argparse.ArgumentParser): - """"Add Megatron-LM arguments to the given parser.""" + """ "Add Megatron-LM arguments to the given parser.""" # Standard arguments. parser = _add_network_size_args(parser) @@ -89,10 +91,10 @@ def add_megatron_arguments(parser: argparse.ArgumentParser): return parser + def parse_args(extra_args_provider=None, ignore_unknown_args=False): """Parse all arguments.""" - parser = argparse.ArgumentParser(description='Megatron-LM Arguments', - allow_abbrev=False) + parser = argparse.ArgumentParser(description="Megatron-LM Arguments", allow_abbrev=False) parser = add_megatron_arguments(parser) parser = add_flagscale_arguments(parser) @@ -110,20 +112,21 @@ def parse_args(extra_args_provider=None, ignore_unknown_args=False): # Experimental yaml if args.yaml_cfg is not None: from .yaml_arguments import load_yaml - assert args.yaml_cfg and not args.use_legacy_models, \ + + assert args.yaml_cfg and not args.use_legacy_models, ( "Yaml config is not supported with legacy models." + ) args = load_yaml(args.yaml_cfg) - # Args from environment - args.rank = int(os.getenv('RANK', '0')) - args.world_size = int(os.getenv("WORLD_SIZE", '1')) + args.rank = int(os.getenv("RANK", "0")) + args.world_size = int(os.getenv("WORLD_SIZE", "1")) # Args to disable MSC if not args.enable_msc: MultiStorageClientFeature.disable() assert MultiStorageClientFeature.is_enabled() is False - warn_rank_0('The MSC feature is disabled.') + warn_rank_0("The MSC feature is disabled.") return args @@ -157,17 +160,22 @@ def validate_model_config_args_from_heterogeneous_config(args): args.heterogeneous_layers_config_path ).read_text() - hf_config_dict = types.SimpleNamespace(**json.loads(args.heterogeneous_layers_config_encoded_json)) + hf_config_dict = types.SimpleNamespace( + **json.loads(args.heterogeneous_layers_config_encoded_json) + ) assert hf_config_dict.hidden_act == "silu", ( f"hidden_act in heterogeneous config is {hf_config_dict.hidden_act}, should be silu" ) n_kv_heads_in_group = [ - config["attention"]["n_heads_in_group"] for config in hf_config_dict.block_configs + config["attention"]["n_heads_in_group"] + for config in hf_config_dict.block_configs if config["attention"]["n_heads_in_group"] is not None ] - assert all(num == n_kv_heads_in_group[0] for num in n_kv_heads_in_group), "num query head must be consistent across all layers" + assert all(num == n_kv_heads_in_group[0] for num in n_kv_heads_in_group), ( + "num query head must be consistent across all layers" + ) args_to_validate = { "swiglu": True, @@ -193,26 +201,28 @@ def validate_model_config_args_from_heterogeneous_config(args): incompatible_args[key] = (provided_value, value) if incompatible_args: - incompatible_args_str = ', '.join([ - f"{k}: {provided_value} (provided) != {value} (expected)" - for k, (provided_value, value) in incompatible_args.items() - ]) - raise ValueError( - f"Arguments differ from heterogeneous config: {incompatible_args_str}" + incompatible_args_str = ", ".join( + [ + f"{k}: {provided_value} (provided) != {value} (expected)" + for k, (provided_value, value) in incompatible_args.items() + ] ) + raise ValueError(f"Arguments differ from heterogeneous config: {incompatible_args_str}") + def _eval_pattern(pattern): - """ Validate and evaluate a string containing a Python list expression """ + """Validate and evaluate a string containing a Python list expression""" assert isinstance(pattern, str) # validate input, only allow comma, digits, [, ], (, ), +, and * - if bool(re.compile(r'[^,\d\[\]\(\)\+\*]').search(pattern)): + if bool(re.compile(r"[^,\d\[\]\(\)\+\*]").search(pattern)): raise ValueError(f"Invalid pattern: {pattern}") return eval(pattern) + def no_rope_freq_type(x): - """ Controls which layers to skip performing Rotary Position Embedding. + """Controls which layers to skip performing Rotary Position Embedding. - An integer N: Represents a 1:N ratio, meaning RoPE is skipped every N-1 layers. - A string "N": Same as above, but provided as a string - A string containing a Python list expression that defines a custom pattern, e.g.: @@ -227,13 +237,14 @@ def no_rope_freq_type(x): if x is None or isinstance(x, int): return x assert isinstance(x, str) - if '[' in x: + if "[" in x: # it's a custom pattern return _eval_pattern(x) else: # it's a single int but in str return int(x) + def moe_freq_type(x): """Frequency between MoE layers and Dense layers. @@ -252,13 +263,14 @@ def moe_freq_type(x): if isinstance(x, int): return x assert isinstance(x, str) - if '[' in x: + if "[" in x: # it's a custom pattern return _eval_pattern(x) else: # it's a single int but in str return int(x) + def la_freq_type(x): """Frequency between LA (linear attention) layers and SDPA (scaled dot-product attention) layers. @@ -277,7 +289,7 @@ def la_freq_type(x): if x is None or isinstance(x, int): return x assert isinstance(x, str) - if '[' in x: + if "[" in x: # it's a custom pattern return _eval_pattern(x) else: @@ -302,11 +314,14 @@ def dsv4_hybrid_attn_type(x): if x is None or isinstance(x, list): return x assert isinstance(x, str) - if '[' in x: + if "[" in x: # it's a custom pattern return _eval_pattern(x) else: - raise ValueError(f"Invalid dsv4_hybrid_attn_type: {x}. Must be a list of integers or a string containing a Python list expression.") + raise ValueError( + f"Invalid dsv4_hybrid_attn_type: {x}. Must be a list of integers or a string containing a Python list expression." + ) + def tuple_type(x): """ @@ -318,27 +333,30 @@ def tuple_type(x): if x is None or isinstance(x, tuple): return x assert isinstance(x, str) - return tuple(int(i) for i in x.strip('()').split(',')) + return tuple(int(i) for i in x.strip("()").split(",")) -def validate_args(args, defaults={}): +def validate_args(args, defaults={}): # Temporary - assert args.non_persistent_ckpt_type in ['global', 'local', None], \ - 'Currently only global and local checkpoints are supported' - if args.non_persistent_ckpt_type == 'local': + assert args.non_persistent_ckpt_type in ["global", "local", None], ( + "Currently only global and local checkpoints are supported" + ) + if args.non_persistent_ckpt_type == "local": try: - from nvidia_resiliency_ext.checkpointing.local.ckpt_managers.local_manager import \ - LocalCheckpointManager + from nvidia_resiliency_ext.checkpointing.local.ckpt_managers.local_manager import ( + LocalCheckpointManager, + ) except ModuleNotFoundError as e: - raise RuntimeError('nvidia_resiliency_ext is required for local checkpointing') from e + raise RuntimeError("nvidia_resiliency_ext is required for local checkpointing") from e # validate model config args from heterogeneous config (if provided). validate_model_config_args_from_heterogeneous_config(args) # Set args.use_dist_ckpt from args.ckpt_format. if args.use_legacy_models: - assert args.ckpt_format == "torch", \ + assert args.ckpt_format == "torch", ( "legacy model format only supports the 'torch' checkpoint format." + ) update_use_dist_ckpt(args) # FlagScale begin @@ -347,9 +365,12 @@ def validate_args(args, defaults={}): multiple_of = defaults.get("multiple_of", None) hidden_dim_multiplier = defaults.get("hidden_dim_multiplier", None) - if not enable_hetero: - total_model_size = args.tensor_model_parallel_size * args.pipeline_model_parallel_size * args.context_parallel_size + total_model_size = ( + args.tensor_model_parallel_size + * args.pipeline_model_parallel_size + * args.context_parallel_size + ) # Total model size. assert args.world_size % total_model_size == 0, ( @@ -357,12 +378,18 @@ def validate_args(args, defaults={}): ) if args.attention_backend == AttnBackend.local: - assert args.spec[0] == 'local' , '--attention-backend local is only supported with --spec local' + assert args.spec[0] == "local", ( + "--attention-backend local is only supported with --spec local" + ) # Pipeline model parallel size. args.transformer_pipeline_model_parallel_size = args.pipeline_model_parallel_size - total_model_size = args.tensor_model_parallel_size * args.pipeline_model_parallel_size * args.context_parallel_size + total_model_size = ( + args.tensor_model_parallel_size + * args.pipeline_model_parallel_size + * args.context_parallel_size + ) if not enable_hetero: args.data_parallel_size = args.world_size // total_model_size @@ -417,8 +444,8 @@ def validate_args(args, defaults={}): # Offload mode requires CG persistence: CG recapture runs dummy forward # passes that corrupt the preserved KV data. - assert ( - (not args.rl_kv_cache_management_mode == "offload") or (args.rl_persist_cuda_graphs) + assert (not args.rl_kv_cache_management_mode == "offload") or ( + args.rl_persist_cuda_graphs ), "--rl-kv-cache-management-mode=offload requires --rl-persist-cuda-graphs" # There's no need to manually offload the KV cache with UVM. @@ -427,7 +454,7 @@ def validate_args(args, defaults={}): and args.rl_kv_cache_management_mode == "offload" ), "--rl-kv-cache-management-mode=offload is incompatible with UVM" # We currently cannot recapture CGs in offload mode. - assert not( + assert not ( not args.rl_persist_cuda_graphs and args.rl_kv_cache_management_mode == "offload" ), "Cannot recapture CUDA graphs while offloading KV cache." @@ -445,116 +472,140 @@ def validate_args(args, defaults={}): ) # Resolve deprecated --rl-parallel-generation-tasks -> --rl-num-parallel-generations. - assert args.rl_num_parallel_generations is None \ - or args.rl_parallel_generation_tasks is None, \ - "Cannot specify both --rl-num-parallel-generations and " \ - "--rl-parallel-generation-tasks. Use --rl-num-parallel-generations " \ + assert ( + args.rl_num_parallel_generations is None or args.rl_parallel_generation_tasks is None + ), ( + "Cannot specify both --rl-num-parallel-generations and " + "--rl-parallel-generation-tasks. Use --rl-num-parallel-generations " "(--rl-parallel-generation-tasks is deprecated)." + ) if args.rl_parallel_generation_tasks is not None: print_rank_0( "WARNING: --rl-parallel-generation-tasks is deprecated, " - "use --rl-num-parallel-generations instead.") + "use --rl-num-parallel-generations instead." + ) args.rl_num_parallel_generations = ( - args.rl_parallel_generation_tasks * args.grpo_group_size) + args.rl_parallel_generation_tasks * args.grpo_group_size + ) # Resolve --rl-num-parallel-generations / --rl-num-parallel-generation-batches. - assert args.rl_num_parallel_generations is None \ - or args.rl_num_parallel_generation_batches is None, \ - "--rl-num-parallel-generations and --rl-num-parallel-generation-batches " \ + assert ( + args.rl_num_parallel_generations is None + or args.rl_num_parallel_generation_batches is None + ), ( + "--rl-num-parallel-generations and --rl-num-parallel-generation-batches " "are mutually exclusive." + ) if args.rl_num_parallel_generations is not None: - assert args.rl_partial_rollouts, \ + assert args.rl_partial_rollouts, ( "--rl-num-parallel-generations requires --rl-partial-rollouts." - assert args.rl_num_parallel_generations % args.grpo_group_size == 0, \ - f"--rl-num-parallel-generations ({args.rl_num_parallel_generations}) " \ + ) + assert args.rl_num_parallel_generations % args.grpo_group_size == 0, ( + f"--rl-num-parallel-generations ({args.rl_num_parallel_generations}) " f"must be divisible by --grpo-group-size ({args.grpo_group_size})." + ) args.rl_parallel_generation_tasks = ( - args.rl_num_parallel_generations // args.grpo_group_size) + args.rl_num_parallel_generations // args.grpo_group_size + ) if args.rl_generation_batch_size is None: args.rl_generation_batch_size = 1 elif args.rl_num_parallel_generation_batches is not None: - assert args.rl_partial_rollouts, \ + assert args.rl_partial_rollouts, ( "--rl-num-parallel-generation-batches requires --rl-partial-rollouts." + ) if args.rl_generation_batch_size is None: args.rl_generation_batch_size = args.grpo_prompts_per_step args.rl_parallel_generation_tasks = ( - args.rl_num_parallel_generation_batches * args.rl_generation_batch_size) + args.rl_num_parallel_generation_batches * args.rl_generation_batch_size + ) else: if args.rl_generation_batch_size is None: args.rl_generation_batch_size = 1 args.rl_parallel_generation_tasks = 512 # Derive enforce_order after all resolution is complete. - args.rl_enforce_generation_order = (args.rl_generation_batch_size > 1) + args.rl_enforce_generation_order = args.rl_generation_batch_size > 1 args.grpo_samples_per_iteration = args.grpo_prompts_per_step * args.grpo_group_size if args.rl_use_sequence_packing: - assert args.micro_batch_size == 1, \ + assert args.micro_batch_size == 1, ( "micro_batch_size must be 1 when using sequence packing. To increase compute per micro batch increase the sequence length." + ) - print_rank_0('using world size: {}, data-parallel size: {}, ' - 'context-parallel size: {}, ' - 'hierarchical context-parallel sizes: {}, ' - 'tensor-model-parallel size: {}, ' - 'pipeline-model-parallel size: {}'.format( - args.world_size, args.data_parallel_size, - args.context_parallel_size, - args.hierarchical_context_parallel_sizes, - args.tensor_model_parallel_size, - args.pipeline_model_parallel_size)) + print_rank_0( + "using world size: {}, data-parallel size: {}, " + "context-parallel size: {}, " + "hierarchical context-parallel sizes: {}, " + "tensor-model-parallel size: {}, " + "pipeline-model-parallel size: {}".format( + args.world_size, + args.data_parallel_size, + args.context_parallel_size, + args.hierarchical_context_parallel_sizes, + args.tensor_model_parallel_size, + args.pipeline_model_parallel_size, + ) + ) # Checks. if args.hierarchical_context_parallel_sizes: from numpy import prod + assert args.context_parallel_size == prod(args.hierarchical_context_parallel_sizes) if "a2a+p2p" in args.cp_comm_type: - assert args.hierarchical_context_parallel_sizes is not None, \ - "--hierarchical-context-parallel-sizes must be set when a2a+p2p is used in cp comm" + assert args.hierarchical_context_parallel_sizes is not None, ( + "--hierarchical-context-parallel-sizes must be set when a2a+p2p is used in cp comm" + ) if args.expert_tensor_parallel_size is None: if not enable_hetero: args.expert_tensor_parallel_size = args.tensor_model_parallel_size # Deprecated arguments. - assert args.batch_size is None, '--batch-size argument is no longer ' \ - 'valid, use --micro-batch-size instead' + assert args.batch_size is None, ( + "--batch-size argument is no longer valid, use --micro-batch-size instead" + ) del args.batch_size - assert args.warmup is None, '--warmup argument is no longer valid, use ' \ - '--lr-warmup-fraction instead' + assert args.warmup is None, ( + "--warmup argument is no longer valid, use --lr-warmup-fraction instead" + ) del args.warmup - assert args.model_parallel_size is None, '--model-parallel-size is no ' \ - 'longer valid, use --tensor-model-parallel-size instead' + assert args.model_parallel_size is None, ( + "--model-parallel-size is no longer valid, use --tensor-model-parallel-size instead" + ) del args.model_parallel_size if args.checkpoint_activations: - print_rank_0('--checkpoint-activations is no longer valid, use --recompute-activations, ' - 'or, for more control, --recompute-granularity and --recompute-method.') + print_rank_0( + "--checkpoint-activations is no longer valid, use --recompute-activations, " + "or, for more control, --recompute-granularity and --recompute-method." + ) exit() del args.checkpoint_activations if args.recompute_activations: - args.recompute_granularity = 'selective' + args.recompute_granularity = "selective" del args.recompute_activations if args.enable_cuda_graph or args.external_cuda_graph: - assert ( - args.cuda_graph_impl == "none" - ), "Do not use --enable-cuda-graph or --external-cuda-graph with --cuda-graph-impl." - assert ( - not args.enable_cuda_graph or not args.external_cuda_graph - ), "--enable-cuda-graph and --external-cuda-graph cannot be enabled at the same time." + assert args.cuda_graph_impl == "none", ( + "Do not use --enable-cuda-graph or --external-cuda-graph with --cuda-graph-impl." + ) + assert not args.enable_cuda_graph or not args.external_cuda_graph, ( + "--enable-cuda-graph and --external-cuda-graph cannot be enabled at the same time." + ) if args.enable_cuda_graph: print_rank_0( - '--enable-cuda-graph is deprecated, use --cuda-graph-impl=local instead.', args.rank + "--enable-cuda-graph is deprecated, use --cuda-graph-impl=local instead.", args.rank ) args.cuda_graph_impl = "local" del args.enable_cuda_graph if args.external_cuda_graph: print_rank_0( - '--external-cuda-graph is deprecated, use --cuda-graph-impl=transformer_engine instead.', + "--external-cuda-graph is deprecated, use --cuda-graph-impl=transformer_engine instead.", args.rank, ) args.cuda_graph_impl = "transformer_engine" @@ -566,65 +617,84 @@ def validate_args(args, defaults={}): # arguments that are passed to the program. We check this by # ensuring the arg is set to None. if getattr(args, key, None) is not None: - warn_rank_0('Overriding default arguments for {key}:{v} ' - 'with {key}:{v2}'.format(key=key, v=defaults[key], - v2=getattr(args, key))) + warn_rank_0( + "Overriding default arguments for {key}:{v} with {key}:{v2}".format( + key=key, v=defaults[key], v2=getattr(args, key) + ) + ) else: setattr(args, key, defaults[key]) if args.data_path is not None and args.split is None: - legacy_default_split_value = '969, 30, 1' - warn_rank_0('Please specify --split when using --data-path. Using legacy default value ' - f'of "{legacy_default_split_value}"') + legacy_default_split_value = "969, 30, 1" + warn_rank_0( + "Please specify --split when using --data-path. Using legacy default value " + f'of "{legacy_default_split_value}"' + ) args.split = legacy_default_split_value use_data_path = (args.data_path is not None) or (args.data_args_path is not None) if use_data_path: # Exactly one of the two has to be None if we use it. assert (args.data_path is None) or (args.data_args_path is None) - use_per_split_data_path = any( - elt is not None - for elt in [args.train_data_path, args.valid_data_path, args.test_data_path]) or \ - args.per_split_data_args_path is not None + use_per_split_data_path = ( + any( + elt is not None + for elt in [args.train_data_path, args.valid_data_path, args.test_data_path] + ) + or args.per_split_data_args_path is not None + ) if use_per_split_data_path: - # Exactly one of the two has to be None if we use it. - assert any(elt is not None - for elt in [args.train_data_path, args.valid_data_path, args.test_data_path]) is False or \ - args.per_split_data_args_path is None + # Exactly one of the two has to be None if we use it. + assert ( + any( + elt is not None + for elt in [args.train_data_path, args.valid_data_path, args.test_data_path] + ) + is False + or args.per_split_data_args_path is None + ) if args.phase_transition_iterations: args.phase_transition_iterations = sorted( int(x.strip()) for x in args.phase_transition_iterations.split(",") ) - assert args.rampup_batch_size is None, "multi-phase training does not support batch size ramp-up" + assert args.rampup_batch_size is None, ( + "multi-phase training does not support batch size ramp-up" + ) # Batch size. assert args.micro_batch_size is not None assert args.micro_batch_size > 0 if args.global_batch_size is None: args.global_batch_size = args.micro_batch_size * args.data_parallel_size - print_rank_0('setting global batch size to {}'.format(args.global_batch_size)) + print_rank_0("setting global batch size to {}".format(args.global_batch_size)) assert args.global_batch_size > 0 if args.perform_rl_step: num_generated_samples_per_inference_iteration = ( - args.grpo_samples_per_iteration * args.grpo_iterations) + args.grpo_samples_per_iteration * args.grpo_iterations + ) # Ensure that the number of prompts we collect is a multiple of the global batch size. # TODO: Make this account for batch size rampup? - assert num_generated_samples_per_inference_iteration % args.global_batch_size == 0, \ + assert num_generated_samples_per_inference_iteration % args.global_batch_size == 0, ( f"grpo_group_size * grpo_prompts_per_step * grpo_iterations should be divisible by global_batch_size" + ) # For now only exit/checkpoint on iterations where we generate data. We don't currently # have a way to checkpoint the generated data. num_training_iterations_per_inference_iteration = ( - num_generated_samples_per_inference_iteration // args.global_batch_size) + num_generated_samples_per_inference_iteration // args.global_batch_size + ) if args.exit_interval is not None: - assert args.exit_interval % num_training_iterations_per_inference_iteration == 0, \ + assert args.exit_interval % num_training_iterations_per_inference_iteration == 0, ( f"exit_interval should be divisible by number of global batches per inference iteration." + ) if args.save_interval is not None: - assert args.save_interval % num_training_iterations_per_inference_iteration == 0, \ + assert args.save_interval % num_training_iterations_per_inference_iteration == 0, ( f"save_interval should be divisible by number of global batches per inference iteration." + ) # === Hybrid layer pattern: deprecation handling and validation === @@ -632,8 +702,8 @@ def validate_args(args, defaults={}): used_hybrid_override_pattern = False if args.hybrid_override_pattern is not None: assert args.hybrid_layer_pattern is None, ( - '--hybrid-override-pattern and --hybrid-layer-pattern cannot both be specified. ' - '--hybrid-override-pattern is deprecated; use --hybrid-layer-pattern instead.' + "--hybrid-override-pattern and --hybrid-layer-pattern cannot both be specified. " + "--hybrid-override-pattern is deprecated; use --hybrid-layer-pattern instead." ) warn_rank_0( "--hybrid-override-pattern is deprecated. Use --hybrid-layer-pattern instead.", @@ -652,9 +722,12 @@ def validate_args(args, defaults={}): ) from megatron.core.ssm.mamba_hybrid_layer_allocation import ( - Symbols, parse_hybrid_pattern, get_hybrid_total_layer_count, + Symbols, + parse_hybrid_pattern, + get_hybrid_total_layer_count, get_hybrid_total_pipeline_segment_count, ) + sep = Symbols.MTP_SEPARATOR # Backward compat: convert legacy mtp_hybrid_override_pattern to unified format @@ -664,9 +737,11 @@ def validate_args(args, defaults={}): and args.mtp_num_layers > 0 and (args.hybrid_layer_pattern is None or sep not in args.hybrid_layer_pattern) ): - main_pattern = args.hybrid_layer_pattern or '' + main_pattern = args.hybrid_layer_pattern or "" mtp_pattern = args.mtp_hybrid_override_pattern - args.hybrid_layer_pattern = main_pattern + sep + sep.join([mtp_pattern] * args.mtp_num_layers) + args.hybrid_layer_pattern = ( + main_pattern + sep + sep.join([mtp_pattern] * args.mtp_num_layers) + ) args.mtp_hybrid_override_pattern = None print_rank_0(f"Converted legacy MTP pattern to unified: {args.hybrid_layer_pattern}") @@ -676,8 +751,8 @@ def validate_args(args, defaults={}): num_layers_in_pattern = get_hybrid_total_layer_count(args.hybrid_layer_pattern) if args.num_layers is not None and args.num_layers != num_layers_in_pattern: warn_rank_0( - f'--hybrid-layer-pattern is set; ignoring --num-layers ({args.num_layers}) and ' - f'using the layer count derived from the pattern ({num_layers_in_pattern}).', + f"--hybrid-layer-pattern is set; ignoring --num-layers ({args.num_layers}) and " + f"using the layer count derived from the pattern ({num_layers_in_pattern}).", args.rank, ) args.num_layers = num_layers_in_pattern @@ -688,39 +763,39 @@ def validate_args(args, defaults={}): has_pipes = Symbols.PIPE in args.hybrid_layer_pattern.split(sep)[0] if has_pipes: assert args.decoder_first_pipeline_num_layers is None, ( - 'If --hybrid-layer-pattern contains pipe separators, ' - '--decoder-first-pipeline-num-layers should not be specified ' - 'as the pipeline layout is explicitly defined.' + "If --hybrid-layer-pattern contains pipe separators, " + "--decoder-first-pipeline-num-layers should not be specified " + "as the pipeline layout is explicitly defined." ) assert args.decoder_last_pipeline_num_layers is None, ( - 'If --hybrid-layer-pattern contains pipe separators, ' - '--decoder-last-pipeline-num-layers should not be specified ' - 'as the pipeline layout is explicitly defined.' + "If --hybrid-layer-pattern contains pipe separators, " + "--decoder-last-pipeline-num-layers should not be specified " + "as the pipeline layout is explicitly defined." ) assert args.num_layers_per_virtual_pipeline_stage is None, ( - '--num-layers-per-virtual-pipeline-stage should not be used with ' - '--hybrid-layer-pattern. To specify virtual pipelining, describe a number of ' - 'pipeline segments in --hybrid-layer-pattern that is a multiple of ' - '--pipeline-model-parallel-size greater than 1' + "--num-layers-per-virtual-pipeline-stage should not be used with " + "--hybrid-layer-pattern. To specify virtual pipelining, describe a number of " + "pipeline segments in --hybrid-layer-pattern that is a multiple of " + "--pipeline-model-parallel-size greater than 1" ) assert args.num_virtual_stages_per_pipeline_rank is None, ( - '--num-virtual-stages-per-pipeline-rank should not be used with ' - '--hybrid-layer-pattern. Virtual pipeline stages are derived from the ' - 'number of | segments in the pattern.' + "--num-virtual-stages-per-pipeline-rank should not be used with " + "--hybrid-layer-pattern. Virtual pipeline stages are derived from the " + "number of | segments in the pattern." ) assert args.pipeline_model_parallel_layout is None, ( - '--pipeline-model-parallel-layout should not be used with --hybrid-layer-pattern. ' - 'Pipeline stage layout is defined by | separators in the pattern.' + "--pipeline-model-parallel-layout should not be used with --hybrid-layer-pattern. " + "Pipeline stage layout is defined by | separators in the pattern." ) assert not args.account_for_embedding_in_pipeline_split, ( - '--account-for-embedding-in-pipeline-split should not be used with ' - '--hybrid-layer-pattern. Pipeline stage layout is defined by | separators ' - 'in the pattern.' + "--account-for-embedding-in-pipeline-split should not be used with " + "--hybrid-layer-pattern. Pipeline stage layout is defined by | separators " + "in the pattern." ) assert not args.account_for_loss_in_pipeline_split, ( - '--account-for-loss-in-pipeline-split should not be used with ' - '--hybrid-layer-pattern. Pipeline stage layout is defined by | separators ' - 'in the pattern.' + "--account-for-loss-in-pipeline-split should not be used with " + "--hybrid-layer-pattern. Pipeline stage layout is defined by | separators " + "in the pattern." ) # Derive VPP from pipe segments in the pattern @@ -733,10 +808,10 @@ def validate_args(args, defaults={}): args.virtual_pipeline_model_parallel_size = None else: assert hybrid_pipeline_segments % args.transformer_pipeline_model_parallel_size == 0, ( - 'The number of hybrid pipeline segments described by --hybrid-layer-pattern must ' - 'be evenly divisible by --pipeline-model-parallel-size. ' - f'Got {hybrid_pipeline_segments} segments and ' - f'{args.transformer_pipeline_model_parallel_size} pipeline parallel size.' + "The number of hybrid pipeline segments described by --hybrid-layer-pattern must " + "be evenly divisible by --pipeline-model-parallel-size. " + f"Got {hybrid_pipeline_segments} segments and " + f"{args.transformer_pipeline_model_parallel_size} pipeline parallel size." ) if hybrid_pipeline_segments > args.transformer_pipeline_model_parallel_size: # Must be set here in order to assign virtual parallel ranks in @@ -760,13 +835,15 @@ def validate_args(args, defaults={}): f"MTP depth count ({inferred_mtp_num_layers}) in pattern " f"'{args.hybrid_layer_pattern}'. " f"Using the inferred value ({inferred_mtp_num_layers}).", - args.rank + args.rank, ) args.mtp_num_layers = inferred_mtp_num_layers # MTP validation if args.mtp_num_layers: - assert not args.use_legacy_models, "The legacy Megatron models does not support Multi-Token Prediction (MTP)." + assert not args.use_legacy_models, ( + "The legacy Megatron models does not support Multi-Token Prediction (MTP)." + ) assert args.position_embedding_type in ("rope", "none", "mrope"), ( f"Multi-Token Prediction (MTP) is not supported with {args.position_embedding_type} position embedding type." + f"The supported position embedding types are rope, none, and mrope." @@ -775,14 +852,16 @@ def validate_args(args, defaults={}): # Validate MTP args for hybrid vs non-hybrid models if args.hybrid_layer_pattern is not None: # Mamba/hybrid model MTP validation - if args.mtp_num_layers and not (args.hybrid_layer_pattern and sep in args.hybrid_layer_pattern): + if args.mtp_num_layers and not ( + args.hybrid_layer_pattern and sep in args.hybrid_layer_pattern + ): # Hybrid model wants MTP but no unified pattern - check for legacy args if args.mtp_hybrid_override_pattern is None: warn_rank_0( "Hybrid model with --mtp-num-layers but no MTP pattern. " "Use unified --hybrid-layer-pattern with '/' separator (e.g., 'M*M*/MM/MM') " "or legacy --mtp-hybrid-override-pattern for old checkpoints.", - args.rank + args.rank, ) else: # Non-hybrid (GPT) model MTP validation @@ -791,7 +870,7 @@ def validate_args(args, defaults={}): "--mtp-hybrid-override-pattern is for Mamba/hybrid models only. " "For GPT models, MTP replicates the main transformer layer structure. " "This argument will be ignored.", - args.rank + args.rank, ) # === End of hybrid layer pattern: deprecation handling and validation === @@ -801,18 +880,20 @@ def validate_args(args, defaults={}): + int(args.num_virtual_stages_per_pipeline_rank is not None) + int(args.pipeline_model_parallel_layout is not None) ) <= 1, ( - 'No more than one of the following arguments can be set at the same time: ' - '--num-layers-per-virtual-pipeline-stage, --num-virtual-stages-per-pipeline-rank,' - '--pipeline-model-parallel-layout. ' - f'{args.num_layers_per_virtual_pipeline_stage=}, ' - f'{args.num_virtual_stages_per_pipeline_rank=}, ' - f'{args.pipeline_model_parallel_layout=}.' + "No more than one of the following arguments can be set at the same time: " + "--num-layers-per-virtual-pipeline-stage, --num-virtual-stages-per-pipeline-rank," + "--pipeline-model-parallel-layout. " + f"{args.num_layers_per_virtual_pipeline_stage=}, " + f"{args.num_virtual_stages_per_pipeline_rank=}, " + f"{args.pipeline_model_parallel_layout=}." ) if args.pipeline_model_parallel_layout is not None: # Parse the input flattened layout to a list and get the vpp size. # We will validate the layout more carefully in the TransformerConfig constructor. - num_stages = PipelineParallelLayerLayout.get_num_stages_from_str(args.pipeline_model_parallel_layout) + num_stages = PipelineParallelLayerLayout.get_num_stages_from_str( + args.pipeline_model_parallel_layout + ) assert num_stages % args.pipeline_model_parallel_size == 0, ( f"The length of pipeline_model_parallel_layout must be divisible" f" by pipeline_model_parallel_size ({num_stages=}," @@ -821,11 +902,20 @@ def validate_args(args, defaults={}): args.virtual_pipeline_model_parallel_size = num_stages // args.pipeline_model_parallel_size if args.virtual_pipeline_model_parallel_size == 1: args.virtual_pipeline_model_parallel_size = None - elif args.num_layers_per_virtual_pipeline_stage is not None or args.num_virtual_stages_per_pipeline_rank is not None: - assert enable_hetero is False, 'num_layers_per_virtual_pipeline_stage is not supported with heterogeneous parallelism for now' + elif ( + args.num_layers_per_virtual_pipeline_stage is not None + or args.num_virtual_stages_per_pipeline_rank is not None + ): + assert enable_hetero is False, ( + "num_layers_per_virtual_pipeline_stage is not supported with heterogeneous parallelism for now" + ) if args.num_virtual_stages_per_pipeline_rank is None: - assert args.decoder_first_pipeline_num_layers is None and args.decoder_last_pipeline_num_layers is None, \ - 'please use --num-virtual-stages-per-pipeline-rank to specify virtual pipeline parallel degree when enable uneven pipeline parallelism' + assert ( + args.decoder_first_pipeline_num_layers is None + and args.decoder_last_pipeline_num_layers is None + ), ( + "please use --num-virtual-stages-per-pipeline-rank to specify virtual pipeline parallel degree when enable uneven pipeline parallelism" + ) if args.num_layers is not None: num_layers = args.num_layers else: @@ -837,14 +927,21 @@ def validate_args(args, defaults={}): if args.account_for_loss_in_pipeline_split: num_layers += 1 - assert num_layers % args.transformer_pipeline_model_parallel_size == 0, \ - 'number of layers of the model must be divisible pipeline model parallel size' - num_layers_per_pipeline_stage = num_layers // args.transformer_pipeline_model_parallel_size + assert num_layers % args.transformer_pipeline_model_parallel_size == 0, ( + "number of layers of the model must be divisible pipeline model parallel size" + ) + num_layers_per_pipeline_stage = ( + num_layers // args.transformer_pipeline_model_parallel_size + ) - assert num_layers_per_pipeline_stage % args.num_layers_per_virtual_pipeline_stage == 0, \ - 'number of layers per pipeline stage must be divisible number of layers per virtual pipeline stage' - args.virtual_pipeline_model_parallel_size = num_layers_per_pipeline_stage // \ - args.num_layers_per_virtual_pipeline_stage + assert ( + num_layers_per_pipeline_stage % args.num_layers_per_virtual_pipeline_stage == 0 + ), ( + "number of layers per pipeline stage must be divisible number of layers per virtual pipeline stage" + ) + args.virtual_pipeline_model_parallel_size = ( + num_layers_per_pipeline_stage // args.num_layers_per_virtual_pipeline_stage + ) else: args.virtual_pipeline_model_parallel_size = args.num_virtual_stages_per_pipeline_rank if args.virtual_pipeline_model_parallel_size == 1: @@ -854,7 +951,10 @@ def validate_args(args, defaults={}): if args.hybrid_layer_pattern is None: args.virtual_pipeline_model_parallel_size = None - if args.decoder_first_pipeline_num_layers is None and args.decoder_last_pipeline_num_layers is None: + if ( + args.decoder_first_pipeline_num_layers is None + and args.decoder_last_pipeline_num_layers is None + ): # Divisibility check not applicable for T5 models which specify encoder_num_layers # and decoder_num_layers, or for hybrid models using --hybrid-layer-pattern. if args.num_layers is not None and args.hybrid_layer_pattern is None: @@ -867,91 +967,113 @@ def validate_args(args, defaults={}): num_layers += 1 if enable_hetero is False: - assert num_layers % args.transformer_pipeline_model_parallel_size == 0, \ - 'Number of layers should be divisible by the pipeline-model-parallel size' - + assert num_layers % args.transformer_pipeline_model_parallel_size == 0, ( + "Number of layers should be divisible by the pipeline-model-parallel size" + ) if args.virtual_pipeline_model_parallel_size is not None: if args.overlap_p2p_comm: - assert args.pipeline_model_parallel_size > 1, \ - 'When interleaved schedule is used, pipeline-model-parallel size '\ - 'should be greater than 1' + assert args.pipeline_model_parallel_size > 1, ( + "When interleaved schedule is used, pipeline-model-parallel size " + "should be greater than 1" + ) else: - assert args.pipeline_model_parallel_size > 2, \ - 'When interleaved schedule is used and p2p communication overlap is disabled, '\ - 'pipeline-model-parallel size should be greater than 2 to avoid having multiple '\ - 'p2p sends and recvs between same 2 ranks per communication batch' + assert args.pipeline_model_parallel_size > 2, ( + "When interleaved schedule is used and p2p communication overlap is disabled, " + "pipeline-model-parallel size should be greater than 2 to avoid having multiple " + "p2p sends and recvs between same 2 ranks per communication batch" + ) else: # Overlap P2P communication is disabled if not using the interleaved schedule. args.overlap_p2p_comm = False args.align_param_gather = False # Only print warning if PP size > 1. if args.rank == 0 and args.pipeline_model_parallel_size > 1: - print('WARNING: Setting args.overlap_p2p_comm and args.align_param_gather to False ' - 'since non-interleaved schedule does not support overlapping p2p communication ' - 'and aligned param AG') + print( + "WARNING: Setting args.overlap_p2p_comm and args.align_param_gather to False " + "since non-interleaved schedule does not support overlapping p2p communication " + "and aligned param AG" + ) print_rank_0( f"Number of virtual stages per pipeline stage: {args.virtual_pipeline_model_parallel_size}" ) if args.overlap_param_gather: - assert args.use_distributed_optimizer or args.use_megatron_fsdp \ - or args.optimizer == 'dist_muon', \ - '--overlap-param-gather only supported with distributed optimizer, megatron fsdp, or dist_muon' - assert args.overlap_grad_reduce, \ - 'Must use --overlap-param-gather with --overlap-grad-reduce' - assert not args.use_legacy_models, \ - '--overlap-param-gather only supported with MCore models' + assert ( + args.use_distributed_optimizer + or args.use_megatron_fsdp + or args.optimizer == "dist_muon" + ), ( + "--overlap-param-gather only supported with distributed optimizer, megatron fsdp, or dist_muon" + ) + assert args.overlap_grad_reduce, ( + "Must use --overlap-param-gather with --overlap-grad-reduce" + ) + assert not args.use_legacy_models, "--overlap-param-gather only supported with MCore models" if args.use_torch_fsdp2: - assert is_torch_min_version("2.4.0"), \ - 'FSDP2 requires PyTorch >= 2.4.0 with FSDP 2 support.' - assert args.pipeline_model_parallel_size == 1, \ - '--use-torch-fsdp2 is not supported with pipeline parallelism' - assert args.expert_model_parallel_size == 1, \ - '--use-torch-fsdp2 is not supported with expert parallelism' - assert not args.use_distributed_optimizer, \ + assert is_torch_min_version("2.4.0"), "FSDP2 requires PyTorch >= 2.4.0 with FSDP 2 support." + assert args.pipeline_model_parallel_size == 1, ( + "--use-torch-fsdp2 is not supported with pipeline parallelism" + ) + assert args.expert_model_parallel_size == 1, ( + "--use-torch-fsdp2 is not supported with expert parallelism" + ) + assert not args.use_distributed_optimizer, ( "--use-torch-fsdp2 is not supported with MCore's distributed optimizer" - assert not args.gradient_accumulation_fusion, \ - '--use-torch-fsdp2 is not supported with gradient accumulation fusion' - assert args.ckpt_format in ('torch_dist', 'torch_dcp'), \ - '--use-torch-fsdp2 requires --ckpt-format torch_dist or torch_dcp' - assert args.untie_embeddings_and_output_weights, \ - '--use-torch-fsdp2 requires --untie-embeddings-and-output-weights' - assert not args.fp16, \ - '--use-torch-fsdp2 not supported with fp16 yet' - assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1", \ - 'FSDP always requires CUDA_DEVICE_MAX_CONNECTIONS value large than one' + ) + assert not args.gradient_accumulation_fusion, ( + "--use-torch-fsdp2 is not supported with gradient accumulation fusion" + ) + assert args.ckpt_format in ("torch_dist", "torch_dcp"), ( + "--use-torch-fsdp2 requires --ckpt-format torch_dist or torch_dcp" + ) + assert args.untie_embeddings_and_output_weights, ( + "--use-torch-fsdp2 requires --untie-embeddings-and-output-weights" + ) + assert not args.fp16, "--use-torch-fsdp2 not supported with fp16 yet" + assert os.environ.get("CUDA_DEVICE_MAX_CONNECTIONS") != "1", ( + "FSDP always requires CUDA_DEVICE_MAX_CONNECTIONS value large than one" + ) if args.fp8_param_gather and is_te_min_version("2.0.0"): args.fp8_param_gather = False warn_rank_0( - 'FSDP2 FP8 param gather is not supported yet in TE 2.0, will fallback to bf16' - 'all_gather instead, turning off fp8_param_gather', + "FSDP2 FP8 param gather is not supported yet in TE 2.0, will fallback to bf16" + "all_gather instead, turning off fp8_param_gather", args.rank, ) if args.fp4_param_gather and not is_te_min_version("2.7.0.dev0"): args.fp4_param_gather = False warn_rank_0( - 'FSDP2 FP4 param gather is not supported yet in TE 2.0, will fallback to bf16' - 'all_gather instead, turning off fp4_param_gather', + "FSDP2 FP4 param gather is not supported yet in TE 2.0, will fallback to bf16" + "all_gather instead, turning off fp4_param_gather", args.rank, ) if args.overlap_param_gather_with_optimizer_step: - assert args.use_distributed_optimizer, \ - '--overlap-param-gather-with-optimizer-step only supported with distributed optimizer' - assert args.overlap_param_gather, \ - 'Must use --overlap-param-gather-with-optimizer-step with --overlap-param-gather' - assert args.virtual_pipeline_model_parallel_size is not None, \ - '--overlap-param-gather-with-optimizer-step only supported with interleaved pipeline parallelism' - assert not args.use_dist_ckpt, \ - '--overlap-param-gather-with-optimizer-step not supported with distributed checkpointing yet' + assert args.use_distributed_optimizer, ( + "--overlap-param-gather-with-optimizer-step only supported with distributed optimizer" + ) + assert args.overlap_param_gather, ( + "Must use --overlap-param-gather-with-optimizer-step with --overlap-param-gather" + ) + assert args.virtual_pipeline_model_parallel_size is not None, ( + "--overlap-param-gather-with-optimizer-step only supported with interleaved pipeline parallelism" + ) + assert not args.use_dist_ckpt, ( + "--overlap-param-gather-with-optimizer-step not supported with distributed checkpointing yet" + ) # Map string data-type to torch.dtype. dtype_map = { - 'fp32': torch.float32, 'bf16': torch.bfloat16, 'fp16': torch.float16, 'fp8': torch.uint8, 'auto': None, None: None, + "fp32": torch.float32, + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp8": torch.uint8, + "auto": None, + None: None, } map_dtype = lambda d: d if isinstance(d, torch.dtype) else dtype_map[d] @@ -969,12 +1091,20 @@ def validate_args(args, defaults={}): args.megatron_fsdp_grad_comm_dtype = torch.bfloat16 if args.fp8_param_gather: - assert args.use_distributed_optimizer or args.use_torch_fsdp2 or args.use_megatron_fsdp or not torch.is_grad_enabled(), \ - '--fp8-param-gather only supported with distributed optimizer, torch fsdp2, megatron fsdp, or inference mode' + assert ( + args.use_distributed_optimizer + or args.use_torch_fsdp2 + or args.use_megatron_fsdp + or not torch.is_grad_enabled() + ), ( + "--fp8-param-gather only supported with distributed optimizer, torch fsdp2, megatron fsdp, or inference mode" + ) # FP4 and FP8 are mutually exclusive if args.fp4 and args.fp8: - raise ValueError("--fp4-format and --fp8-format cannot be used simultaneously. Please choose one.") + raise ValueError( + "--fp4-format and --fp8-format cannot be used simultaneously. Please choose one." + ) # FP4 param requires FP4 mode if args.fp4_param_gather and not args.fp4: @@ -982,11 +1112,13 @@ def validate_args(args, defaults={}): # FP4 requires TE >= 2.7.0.dev0 if args.fp4 and not is_te_min_version("2.7.0.dev0"): - raise ValueError("--fp4-format requires Transformer Engine >= 2.7.0.dev0 for NVFP4BlockScaling support.") + raise ValueError( + "--fp4-format requires Transformer Engine >= 2.7.0.dev0 for NVFP4BlockScaling support." + ) if ( - args.fp8_recipe == 'mxfp8' - and args.transformer_impl == 'inference_optimized' + args.fp8_recipe == "mxfp8" + and args.transformer_impl == "inference_optimized" and not is_flashinfer_min_version("0.6.4") ): raise ValueError("MXFP8 with inference optimized layers requires FlashInfer >= 0.6.4") @@ -999,22 +1131,27 @@ def validate_args(args, defaults={}): if args.data_parallel_sharding_strategy in ["optim_grads_params", "optim_grads"]: warn_rank_0( - 'Please make sure your TransformerEngine support FSDP + gradient accumulation fusion', + "Please make sure your TransformerEngine support FSDP + gradient accumulation fusion", args.rank, ) if args.data_parallel_sharding_strategy == "optim_grads_params": - assert args.check_weight_hash_across_dp_replicas_interval is None, \ - 'check_weight_hash_across_dp_replicas_interval is not supported with optim_grads_params' + assert args.check_weight_hash_across_dp_replicas_interval is None, ( + "check_weight_hash_across_dp_replicas_interval is not supported with optim_grads_params" + ) - assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1", \ - 'FSDP always requires CUDA_DEVICE_MAX_CONNECTIONS value large than one' + assert os.environ.get("CUDA_DEVICE_MAX_CONNECTIONS") != "1", ( + "FSDP always requires CUDA_DEVICE_MAX_CONNECTIONS value large than one" + ) - assert args.ckpt_format == "fsdp_dtensor", \ + assert args.ckpt_format == "fsdp_dtensor", ( "Megatron FSDP only supports fsdp_dtensor checkpoint format" - + ) + if args.fsdp_manual_registration: - assert args.use_megatron_fsdp, "FSDP manual registration is only supported with Megatron FSDP" + assert args.use_megatron_fsdp, ( + "FSDP manual registration is only supported with Megatron FSDP" + ) assert args.nccl_ub, "FSDP manual registration is only supported with nccl-ub option" if args.use_megatron_fsdp: @@ -1029,43 +1166,56 @@ def validate_args(args, defaults={}): # where NaNs in grads / loss are signal to the loss scaler. if not args.loss_scale: args.check_for_nan_in_loss_and_grad = False - warn_rank_0('Setting args.check_for_nan_in_loss_and_grad to False since ' - 'dynamic loss scaling is being used') + warn_rank_0( + "Setting args.check_for_nan_in_loss_and_grad to False since " + "dynamic loss scaling is being used" + ) if args.bf16: assert not args.fp16 args.params_dtype = torch.bfloat16 # bfloat16 requires gradient accumulation and all-reduce to # be done in fp32. if args.accumulate_allreduce_grads_in_fp32: - assert args.main_grads_dtype == torch.float32, \ + assert args.main_grads_dtype == torch.float32, ( "--main-grads-dtype can only be fp32 when --accumulate-allreduce-grads-in-fp32 is set" + ) if args.grad_reduce_in_bf16: args.accumulate_allreduce_grads_in_fp32 = False elif not args.accumulate_allreduce_grads_in_fp32 and args.main_grads_dtype == torch.float32: args.accumulate_allreduce_grads_in_fp32 = True - print_rank_0('accumulate and all-reduce gradients in fp32 for bfloat16 data type.') + print_rank_0("accumulate and all-reduce gradients in fp32 for bfloat16 data type.") if args.cuda_graph_impl == "local" and CudaGraphScope.full_iteration in args.cuda_graph_scope: - assert not args.check_for_nan_in_loss_and_grad, \ - "--no-check-for-nan-in-loss-and-grad should be set with --cuda-graph-scope=full_iteration for training. Note: If you are trying to use full_iteration CUDA graphs for inference, please use --cuda-graph-scope full_iteration_inference instead" - - if args.cuda_graph_impl == "local" and CudaGraphScope.full_iteration_inference in args.cuda_graph_scope: + assert not args.check_for_nan_in_loss_and_grad, ( + "--no-check-for-nan-in-loss-and-grad should be set with --cuda-graph-scope=full_iteration for training. Note: If you are trying to use full_iteration CUDA graphs for inference, please use --cuda-graph-scope full_iteration_inference instead" + ) + + if ( + args.cuda_graph_impl == "local" + and CudaGraphScope.full_iteration_inference in args.cuda_graph_scope + ): if args.fp8 is not None: - assert args.transformer_impl == "inference_optimized", \ - "fp8 with full_iteration_inference CUDA graphs is only supported with " \ + assert args.transformer_impl == "inference_optimized", ( + "fp8 with full_iteration_inference CUDA graphs is only supported with " "--transformer-impl=inference_optimized" - assert args.fp8_recipe == "mxfp8", \ + ) + assert args.fp8_recipe == "mxfp8", ( "Only --fp8-recipe=mxfp8 is supported with full_iteration_inference CUDA graphs" + ) - if args.cuda_graph_impl == 'local': - assert args.inference_dynamic_batching_num_cuda_graphs > 0 or args.inference_dynamic_batching_num_cuda_graphs == -1, \ - 'inference_dynamic_batching_num_cuda_graphs should be a positive integer or -1' \ - '-1 means that we will automatically determine the number of CUDA graphs to capture based on the `max_requests` value.' + if args.cuda_graph_impl == "local": + assert ( + args.inference_dynamic_batching_num_cuda_graphs > 0 + or args.inference_dynamic_batching_num_cuda_graphs == -1 + ), ( + "inference_dynamic_batching_num_cuda_graphs should be a positive integer or -1" + "-1 means that we will automatically determine the number of CUDA graphs to capture based on the `max_requests` value." + ) - print_rank_0('using {} for parameters ...'.format(args.params_dtype)) + print_rank_0("using {} for parameters ...".format(args.params_dtype)) if args.dataloader_type is None: - args.dataloader_type = 'single' + args.dataloader_type = "single" # data assert args.num_dataset_builder_threads > 0 @@ -1089,45 +1239,42 @@ def validate_args(args, defaults={}): if args.train_iters and not args.skip_train: # If we use iteration-based training, make sure the # sample-based options are off. - assert args.train_samples is None, \ - 'expected iteration-based training' - assert args.lr_decay_samples is None, \ - 'expected iteration-based learning rate decay' - assert args.lr_warmup_samples == 0, \ - 'expected iteration-based learning rate warmup' - assert args.rampup_batch_size is None, \ - 'expected no batch-size rampup for iteration-based training' + assert args.train_samples is None, "expected iteration-based training" + assert args.lr_decay_samples is None, "expected iteration-based learning rate decay" + assert args.lr_warmup_samples == 0, "expected iteration-based learning rate warmup" + assert args.rampup_batch_size is None, ( + "expected no batch-size rampup for iteration-based training" + ) if args.lr_warmup_fraction is not None: - assert args.lr_warmup_iters == 0, \ - 'can only specify one of lr-warmup-fraction and lr-warmup-iters' + assert args.lr_warmup_iters == 0, ( + "can only specify one of lr-warmup-fraction and lr-warmup-iters" + ) # Sample-based training. if args.train_samples and not args.skip_train: # If we use sample-based training, make sure the # iteration-based options are off. - assert args.train_iters is None, \ - 'expected sample-based training' - assert args.lr_decay_iters is None, \ - 'expected sample-based learning rate decay' - assert args.lr_warmup_iters == 0, \ - 'expected sample-based learnig rate warmup' + assert args.train_iters is None, "expected sample-based training" + assert args.lr_decay_iters is None, "expected sample-based learning rate decay" + assert args.lr_warmup_iters == 0, "expected sample-based learnig rate warmup" if args.lr_warmup_fraction is not None: - assert args.lr_warmup_samples == 0, \ - 'can only specify one of lr-warmup-fraction ' \ - 'and lr-warmup-samples' + assert args.lr_warmup_samples == 0, ( + "can only specify one of lr-warmup-fraction and lr-warmup-samples" + ) if args.num_layers is not None: - assert args.encoder_num_layers is None, \ - 'cannot have both num-layers and encoder-num-layers specified' + assert args.encoder_num_layers is None, ( + "cannot have both num-layers and encoder-num-layers specified" + ) args.encoder_num_layers = args.num_layers else: - assert args.encoder_num_layers is not None, \ - 'either num-layers or encoder-num-layers should be specified' + assert args.encoder_num_layers is not None, ( + "either num-layers or encoder-num-layers should be specified" + ) args.num_layers = args.encoder_num_layers # Check required arguments. - required_args = ['num_layers', 'hidden_size', 'num_attention_heads', - 'max_position_embeddings'] + required_args = ["num_layers", "hidden_size", "num_attention_heads", "max_position_embeddings"] for req_arg in required_args: _check_arg_is_not_none(args, req_arg) @@ -1138,11 +1285,11 @@ def validate_args(args, defaults={}): if multiple_of is not None: hidden_dim = int(4 * args.hidden_size * 2 / 3) if hidden_dim_multiplier is not None: - assert hidden_dim_multiplier > 0, \ - 'multiplier for hidden dim should be greater than zero' + assert hidden_dim_multiplier > 0, ( + "multiplier for hidden dim should be greater than zero" + ) hidden_dim = int(hidden_dim * hidden_dim_multiplier) - args.ffn_hidden_size = multiple_of * \ - ((hidden_dim + multiple_of - 1) // multiple_of) + args.ffn_hidden_size = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) else: # reduce the dimnesion for MLP since projections happens on # two linear layers. this keeps the number of paramters in @@ -1158,9 +1305,10 @@ def validate_args(args, defaults={}): args.kv_channels = args.hidden_size // args.num_attention_heads if args.seq_length is not None and args.context_parallel_size > 1: - assert args.seq_length % (args.context_parallel_size * 2) == 0, \ - 'seq-length should be a multiple of 2 * context-parallel-size ' \ - 'if context-parallel-size > 1.' + assert args.seq_length % (args.context_parallel_size * 2) == 0, ( + "seq-length should be a multiple of 2 * context-parallel-size " + "if context-parallel-size > 1." + ) if args.seq_length is not None: assert args.encoder_seq_length is None @@ -1170,9 +1318,10 @@ def validate_args(args, defaults={}): args.seq_length = args.encoder_seq_length if args.seq_length is not None: - assert args.max_position_embeddings >= args.seq_length, \ - f"max_position_embeddings ({args.max_position_embeddings}) must be greater than " \ + assert args.max_position_embeddings >= args.seq_length, ( + f"max_position_embeddings ({args.max_position_embeddings}) must be greater than " f"or equal to seq_length ({args.seq_length})." + ) if args.decoder_seq_length is not None: assert args.max_position_embeddings >= args.decoder_seq_length if args.lr is not None: @@ -1187,19 +1336,20 @@ def validate_args(args, defaults={}): assert args.log_memory_interval % args.log_interval == 0 # Mixed precision checks. if args.fp16_lm_cross_entropy: - assert args.fp16, 'lm cross entropy in fp16 only support in fp16 mode.' + assert args.fp16, "lm cross entropy in fp16 only support in fp16 mode." if args.fp32_residual_connection: - assert args.fp16 or args.bf16, \ - 'residual connection in fp32 only supported when using fp16 or bf16.' + assert args.fp16 or args.bf16, ( + "residual connection in fp32 only supported when using fp16 or bf16." + ) if args.moe_grouped_gemm: - assert args.bf16, 'Currently GroupedGEMM for MoE only supports bf16 dtype.' + assert args.bf16, "Currently GroupedGEMM for MoE only supports bf16 dtype." dc = cur_platform.get_device_capability() assert dc[0] >= 8, "Unsupported compute capability for GroupedGEMM kernels." if args.no_weight_decay_cond_type is not None: print_rank_0( - 'WARNING: --no-weight-decay-cond-type is deprecated. Please use --apply-wd-to-qk-layernorm instead.', + "WARNING: --no-weight-decay-cond-type is deprecated. Please use --apply-wd-to-qk-layernorm instead.", args.rank, ) if args.no_weight_decay_cond_type == "apply_wd_to_qk_layernorm": @@ -1208,7 +1358,7 @@ def validate_args(args, defaults={}): raise ValueError(f"Invalid no_weight_decay_cond_type: {args.no_weight_decay_cond_type}") args.no_weight_decay_cond_type = None - if args.weight_decay_incr_style == 'constant': + if args.weight_decay_incr_style == "constant": assert args.start_weight_decay is None assert args.end_weight_decay is None args.start_weight_decay = args.weight_decay @@ -1220,30 +1370,33 @@ def validate_args(args, defaults={}): # Persistent fused layer norm. if not is_torch_min_version("1.11.0a0"): args.no_persist_layer_norm = True - print_rank_0('Persistent fused layer norm kernel is supported from ' - 'pytorch v1.11 (nvidia pytorch container paired with v1.11). ' - 'Defaulting to no_persist_layer_norm=True') + print_rank_0( + "Persistent fused layer norm kernel is supported from " + "pytorch v1.11 (nvidia pytorch container paired with v1.11). " + "Defaulting to no_persist_layer_norm=True" + ) # Activation recomputing. if args.distribute_saved_activations: - assert args.tensor_model_parallel_size > 1, 'can distribute ' \ - 'recomputed activations only across tensor model ' \ - 'parallel groups' - assert args.recompute_granularity == 'full', \ - 'distributed recompute activations is only '\ - 'application to full recompute granularity' - assert args.recompute_method is not None, \ - 'for distributed recompute activations to work you '\ - 'need to use a recompute method ' - assert is_torch_min_version("1.10.0a0"), \ - 'distributed recompute activations are supported for pytorch ' \ - 'v1.10 and above (Nvidia Pytorch container >= 21.07). Current ' \ - f'pytorch version is v{get_torch_version()}.' - - if args.recompute_granularity == 'selective': - assert args.recompute_method is None, \ - 'recompute method is not yet supported for ' \ - 'selective recomputing granularity' + assert args.tensor_model_parallel_size > 1, ( + "can distribute recomputed activations only across tensor model parallel groups" + ) + assert args.recompute_granularity == "full", ( + "distributed recompute activations is only application to full recompute granularity" + ) + assert args.recompute_method is not None, ( + "for distributed recompute activations to work you need to use a recompute method " + ) + assert is_torch_min_version("1.10.0a0"), ( + "distributed recompute activations are supported for pytorch " + "v1.10 and above (Nvidia Pytorch container >= 21.07). Current " + f"pytorch version is v{get_torch_version()}." + ) + + if args.recompute_granularity == "selective": + assert args.recompute_method is None, ( + "recompute method is not yet supported for selective recomputing granularity" + ) # disable sequence parallelism when tp=1 # to avoid change in numerics when @@ -1257,19 +1410,32 @@ def validate_args(args, defaults={}): args.sequence_parallel = False if args.tp_comm_overlap: - assert args.sequence_parallel == True, 'Tensor parallel communication/GEMM overlap can happen only when sequence parallelism is enabled' + assert args.sequence_parallel == True, ( + "Tensor parallel communication/GEMM overlap can happen only when sequence parallelism is enabled" + ) if args.hybrid_context_parallel: - assert not args.pipeline_model_parallel_size > 1, 'Hybrid context parallelism not supported with pipeline parallelism' - assert not args.enable_cuda_graph, 'Hybrid context parallelism not supported with CUDA Graph' - assert not args.use_megatron_fsdp, 'Hybrid context parallelism not supported with Megatron FSDP' - assert args.dataloader_type == 'single', 'Hybrid context parallelism only supported with single dataloader type' - assert args.calculate_per_token_loss, 'Hybrid context parallelism must be used with --calculate-per-token-loss' + assert not args.pipeline_model_parallel_size > 1, ( + "Hybrid context parallelism not supported with pipeline parallelism" + ) + assert not args.enable_cuda_graph, ( + "Hybrid context parallelism not supported with CUDA Graph" + ) + assert not args.use_megatron_fsdp, ( + "Hybrid context parallelism not supported with Megatron FSDP" + ) + assert args.dataloader_type == "single", ( + "Hybrid context parallelism only supported with single dataloader type" + ) + assert args.calculate_per_token_loss, ( + "Hybrid context parallelism must be used with --calculate-per-token-loss" + ) # disable async_tensor_model_parallel_allreduce when # model parallel memory optimization is enabled - if (args.tensor_model_parallel_size > 1 or args.context_parallel_size > 1) \ - and get_device_arch_version() < 10: + if ( + args.tensor_model_parallel_size > 1 or args.context_parallel_size > 1 + ) and get_device_arch_version() < 10: # CUDA_DEVICE_MAX_CONNECTIONS requirement no longer exists since the Blackwell architecture if args.use_torch_fsdp2 or args.use_megatron_fsdp: fsdp_impl = "Torch-FSDP2" if args.use_torch_fsdp2 else "Megatron-FSDP" @@ -1292,18 +1458,19 @@ def validate_args(args, defaults={}): args.rank, ) else: - assert os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') == "1", \ - "Using tensor model parallelism or context parallelism require setting the environment variable " \ + assert os.environ.get("CUDA_DEVICE_MAX_CONNECTIONS") == "1", ( + "Using tensor model parallelism or context parallelism require setting the environment variable " "CUDA_DEVICE_MAX_CONNECTIONS to 1" + ) # Setting FSDP communication groups for high priority streams for Blackwell and later architectures # Assigning high priority to communication streams ensures that communication kernels are scheduled # with higher priority, minimizing the exposed communication when it is overlapped with other computation kernels. if args.use_torch_fsdp2 or args.use_megatron_fsdp and get_device_arch_version() >= 10: - if 'dp_cp' not in args.high_priority_stream_groups: - args.high_priority_stream_groups.append('dp_cp') - if args.expert_model_parallel_size > 1 and 'ep_dp' not in args.high_priority_stream_groups: - args.high_priority_stream_groups.append('ep_dp') + if "dp_cp" not in args.high_priority_stream_groups: + args.high_priority_stream_groups.append("dp_cp") + if args.expert_model_parallel_size > 1 and "ep_dp" not in args.high_priority_stream_groups: + args.high_priority_stream_groups.append("ep_dp") # Disable bias gelu fusion if we are disabling bias altogether if not args.add_bias_linear: @@ -1314,45 +1481,49 @@ def validate_args(args, defaults={}): args.add_qkv_bias = True if args.qk_clip: - assert is_te_min_version("2.9.0"), \ - '--qk-clip is only supported with TE >= 2.9.0.' - assert 0.0 < args.qk_clip_alpha < 1.0, \ - '--qk-clip-alpha must be between 0.0 and 1.0 when using --qk-clip.' - assert args.qk_clip_threshold > 0, \ - '--qk-clip-threshold must be greater than 0 when using --qk-clip.' + assert is_te_min_version("2.9.0"), "--qk-clip is only supported with TE >= 2.9.0." + assert 0.0 < args.qk_clip_alpha < 1.0, ( + "--qk-clip-alpha must be between 0.0 and 1.0 when using --qk-clip." + ) + assert args.qk_clip_threshold > 0, ( + "--qk-clip-threshold must be greater than 0 when using --qk-clip." + ) # decoupled log max attention logit check if args.log_max_attention_logit: - assert is_te_min_version("2.9.0"), \ - '--log-max-attention-logit is only supported with TE >= 2.9.0.' + assert is_te_min_version("2.9.0"), ( + "--log-max-attention-logit is only supported with TE >= 2.9.0." + ) if args.decoupled_lr is not None or args.decoupled_min_lr is not None: - assert not args.use_legacy_models, \ - '--decoupled-lr and --decoupled-min-lr is not supported in legacy models.' + assert not args.use_legacy_models, ( + "--decoupled-lr and --decoupled-min-lr is not supported in legacy models." + ) # Legacy RoPE arguments if args.use_rotary_position_embeddings: - args.position_embedding_type = 'rope' + args.position_embedding_type = "rope" if args.rotary_interleaved and args.use_legacy_models: - raise RuntimeError('--rotary-interleaved is not supported in legacy models.') - if args.position_embedding_type != 'rope': + raise RuntimeError("--rotary-interleaved is not supported in legacy models.") + if args.position_embedding_type != "rope": args.apply_rope_fusion = False # Would just need to add 'NoPE' as a position_embedding_type to support this, but for now # don't allow it to keep things simple - if not args.add_position_embedding and args.position_embedding_type != 'rope': - raise RuntimeError('--no-position-embedding is deprecated, use --position-embedding-type') + if not args.add_position_embedding and args.position_embedding_type != "rope": + raise RuntimeError("--no-position-embedding is deprecated, use --position-embedding-type") # Relative position embeddings arguments - if args.position_embedding_type == 'relative': - assert ( - args.transformer_impl == "transformer_engine" - ), 'Local transformer implementation currently does not support attention bias-based position embeddings.' + if args.position_embedding_type == "relative": + assert args.transformer_impl == "transformer_engine", ( + "Local transformer implementation currently does not support attention bias-based position embeddings." + ) # MultiModal rotary embeddings arguments if args.position_embedding_type == "mrope": - assert args.mrope_section is not None, \ - '--mrope-section should be set when using --position-embedding-type mrope.' + assert args.mrope_section is not None, ( + "--mrope-section should be set when using --position-embedding-type mrope." + ) # MoE Spec check if args.num_experts == 0: @@ -1366,38 +1537,50 @@ def validate_args(args, defaults={}): assert not args.use_legacy_models, "Context parallelism is not supported in legacy models." # Expert parallelism check - if args.expert_model_parallel_size > 1: - assert args.num_experts is not None, "num_experts must be non None to use expert model parallelism" - assert args.num_experts % args.expert_model_parallel_size == 0, \ + if args.expert_model_parallel_size > 1: + assert args.num_experts is not None, ( + "num_experts must be non None to use expert model parallelism" + ) + assert args.num_experts % args.expert_model_parallel_size == 0, ( "Number of experts should be a multiple of expert model parallel_size." + ) # MoE router check - if isinstance(args.moe_router_load_balancing_type, list) and len(args.moe_router_load_balancing_type) == 1: + if ( + isinstance(args.moe_router_load_balancing_type, list) + and len(args.moe_router_load_balancing_type) == 1 + ): args.moe_router_load_balancing_type = args.moe_router_load_balancing_type[0] if isinstance(args.moe_aux_loss_coeff, list) and len(args.moe_aux_loss_coeff) == 1: args.moe_aux_loss_coeff = args.moe_aux_loss_coeff[0] # Distributed checkpointing checks if args.use_dist_ckpt and args.use_legacy_models: - raise RuntimeError('--use-dist-ckpt is not supported in legacy models.') + raise RuntimeError("--use-dist-ckpt is not supported in legacy models.") # torch_dcp (torch.distributed.checkpoint) checkpointing format checks. if args.ckpt_format == "torch_dcp": assert args.use_torch_fsdp2, "--ckpt-format torch_dcp is only tested with FSDP." - assert args.tensor_model_parallel_size <= 1, \ + assert args.tensor_model_parallel_size <= 1, ( "--ckpt-format torch_dcp is not tested with megatron tensor parallelism." - assert args.pipeline_model_parallel_size <= 1, \ + ) + assert args.pipeline_model_parallel_size <= 1, ( "--ckpt-format torch_dcp is not tested with megatron pipeline parallelism." + ) # fsdp_dtensor checkpointing format checks. if args.ckpt_format == "fsdp_dtensor": - assert args.use_megatron_fsdp, "--ckpt-format fsdp_dtensor is only tested with Megatron FSDP." + assert args.use_megatron_fsdp, ( + "--ckpt-format fsdp_dtensor is only tested with Megatron FSDP." + ) # Data blend checks - assert args.mock_data + \ - bool(args.data_path) + \ - any([args.train_data_path, args.valid_data_path, args.test_data_path]) \ - <= 1, "A single data source must be provided in training mode, else None" + assert ( + args.mock_data + + bool(args.data_path) + + any([args.train_data_path, args.valid_data_path, args.test_data_path]) + <= 1 + ), "A single data source must be provided in training mode, else None" if args.fim_data: extra_tokens = [ @@ -1410,16 +1593,21 @@ def validate_args(args, defaults={}): assert not args.mock_data, "Mock dataset is not supported with FIM dataset." assert args.fim_rate, "--fim-rate should be specified." assert args.fim_spm_rate, "--fim-spm-rate should be specified." - assert all(token is not None for token in extra_tokens), "FIM extra tokens should be specified." + assert all(token is not None for token in extra_tokens), ( + "FIM extra tokens should be specified." + ) # Deterministic mode if args.deterministic_mode: assert not args.use_flash_attn, "Flash attention can not be used in deterministic mode." - assert not args.cross_entropy_loss_fusion, "Cross Entropy Fusion is currently not deterministic." + assert not args.cross_entropy_loss_fusion, ( + "Cross Entropy Fusion is currently not deterministic." + ) all_reduce_choices = ["Tree", "Ring", "CollnetDirect", "CollnetChain", "^NVLS"] - assert os.getenv("NCCL_ALGO", -1) != -1 and os.getenv("NCCL_ALGO") in all_reduce_choices, \ + assert os.getenv("NCCL_ALGO", -1) != -1 and os.getenv("NCCL_ALGO") in all_reduce_choices, ( f"NCCL_ALGO must be one of {all_reduce_choices}." + ) torch.use_deterministic_algorithms(True) @@ -1443,37 +1631,48 @@ def validate_args(args, defaults={}): assert args.use_dist_ckpt if args.dist_ckpt_optim_fully_reshardable: - assert not args.distrib_optim_fully_reshardable_mem_efficient, \ - '--distrib-optim-fully-reshardable-mem-efficient requires -enable-gloo-process-groups' + assert not args.distrib_optim_fully_reshardable_mem_efficient, ( + "--distrib-optim-fully-reshardable-mem-efficient requires -enable-gloo-process-groups" + ) if args.fake_process_group: - assert args.moe_token_dispatcher_type != "flex", "Fake process group is not supported with flex token dispatcher." + assert args.moe_token_dispatcher_type != "flex", ( + "Fake process group is not supported with flex token dispatcher." + ) # Disable nan check for fake process group args.check_for_nan_in_loss_and_grad = False - warn_rank_0('check_for_nan_in_loss_and_grad is set to False for fake process group.') + warn_rank_0("check_for_nan_in_loss_and_grad is set to False for fake process group.") # Disable gloo process groups for fake process group args.use_gloo_process_groups = False - warn_rank_0('use_gloo_process_groups is set to False for fake process group.') + warn_rank_0("use_gloo_process_groups is set to False for fake process group.") # Checkpointing if args.ckpt_fully_parallel_save_deprecated and args.rank == 0: - print('--ckpt-fully-parallel-save flag is deprecated and has no effect.' - ' Use --no-ckpt-fully-parallel-save to disable parallel save.') + print( + "--ckpt-fully-parallel-save flag is deprecated and has no effect." + " Use --no-ckpt-fully-parallel-save to disable parallel save." + ) if ( args.use_dist_ckpt and not args.ckpt_fully_parallel_save and args.use_distributed_optimizer and args.rank == 0 ): - print('Warning: With non-parallel ckpt save and DistributedOptimizer,' - ' it will be impossible to resume training with different parallelism.' - ' Consider removing flag --no-ckpt-fully-parallel-save.') + print( + "Warning: With non-parallel ckpt save and DistributedOptimizer," + " it will be impossible to resume training with different parallelism." + " Consider removing flag --no-ckpt-fully-parallel-save." + ) if args.use_dist_ckpt_deprecated and args.rank == 0: - print('--use-dist-ckpt is deprecated and has no effect.' - ' Use --ckpt-format to select the checkpoint format.') + print( + "--use-dist-ckpt is deprecated and has no effect." + " Use --ckpt-format to select the checkpoint format." + ) if args.dist_ckpt_format_deprecated and args.rank == 0: - print('--dist-ckpt-format is deprecated and has no effect.' - ' Use --ckpt-format to select the checkpoint format.') + print( + "--dist-ckpt-format is deprecated and has no effect." + " Use --ckpt-format to select the checkpoint format." + ) if args.use_dist_ckpt and args.ckpt_fully_parallel_load: if args.ckpt_fully_parallel_load_exchange_algo != "broadcast": @@ -1483,64 +1682,85 @@ def validate_args(args, defaults={}): ) if args.load_main_params_from_ckpt: - assert args.no_load_optim, '--load-main-params-from-ckpt must be used with --no-load-optim.' + assert args.no_load_optim, "--load-main-params-from-ckpt must be used with --no-load-optim." if args.use_dist_ckpt and args.async_save: if not args.use_persistent_ckpt_worker: warn_rank_0( - '--async-save is not supported without --use-persistent-ckpt-worker. ' - 'Disabling --async-save.' + "--async-save is not supported without --use-persistent-ckpt-worker. " + "Disabling --async-save." ) args.async_save = False # Inference args if args.inference_batch_times_seqlen_threshold > -1: - assert args.pipeline_model_parallel_size > 1, \ + assert args.pipeline_model_parallel_size > 1, ( "--inference-batch-times-seqlen-threshold requires setting --pipeline-model-parallel-size > 1." - assert ( - args.cuda_graph_impl == "none" - ), "Pipeline-parallel microbatched inference is incompatible with CUDA graphs" + ) + assert args.cuda_graph_impl == "none", ( + "Pipeline-parallel microbatched inference is incompatible with CUDA graphs" + ) if args.inference_dynamic_batching: assert args.inference_dynamic_batching_buffer_size_gb is not None - assert args.inference_dynamic_batching_block_size % 256 == 0, "block size should be a multiple of 256" + assert args.inference_dynamic_batching_block_size % 256 == 0, ( + "block size should be a multiple of 256" + ) - if args.cuda_graph_impl == "local" and args.expert_model_parallel_size > 1 and args.transformer_impl != "inference_optimized": - assert args.moe_pad_experts_for_cuda_graph_inference, \ - "--moe-pad-experts-for-cuda-graph-inference must be set when using CUDA graphs with expert parallelism" + if ( + args.cuda_graph_impl == "local" + and args.expert_model_parallel_size > 1 + and args.transformer_impl != "inference_optimized" + ): + assert args.moe_pad_experts_for_cuda_graph_inference, ( + "--moe-pad-experts-for-cuda-graph-inference must be set when using CUDA graphs with expert parallelism" + ) # MoE upcycling check if args.moe_use_upcycling: assert args.save is not None, "When using upcycling, the --save option must be specified." if not args.no_load_optim: args.no_load_optim = True - warn_rank_0('enabling --no-load-optim for upcycling.') + warn_rank_0("enabling --no-load-optim for upcycling.") if not args.no_load_rng: args.no_load_rng = True - warn_rank_0('enabling --no-load-rng for upcycling.') + warn_rank_0("enabling --no-load-rng for upcycling.") # --skip-train checks. # In RL inference-only mode, --no-load-optim is user-controlled: it determines whether the # optimizer is created (needed for --rl-offload-optimizer-during-inference) or skipped entirely. if args.skip_train and not args.perform_rl_step and not args.no_load_optim: args.no_load_optim = True - warn_rank_0('enabling --no-load-optim when skipping training.') - if args.skip_train and args.perform_rl_step and args.no_load_optim and args.rl_offload_optimizer_during_inference: - assert False, \ - '--no-load-optim with --skip-train --perform-rl-step skips the optimizer; ' \ - '--rl-offload-optimizer-during-inference is incompatible (no optimizer to offload).' + warn_rank_0("enabling --no-load-optim when skipping training.") + if ( + args.skip_train + and args.perform_rl_step + and args.no_load_optim + and args.rl_offload_optimizer_during_inference + ): + assert False, ( + "--no-load-optim with --skip-train --perform-rl-step skips the optimizer; " + "--rl-offload-optimizer-during-inference is incompatible (no optimizer to offload)." + ) # Muon optimizer check - if 'muon' in args.optimizer: - - if args.optimizer == 'muon': - assert not args.overlap_grad_reduce, "Muon optimizer does not support overlap grad reduce. Use dist_muon instead." - assert not args.overlap_param_gather, "Muon optimizer does not support overlap param gather. Use dist_muon instead." + if "muon" in args.optimizer: + if args.optimizer == "muon": + assert not args.overlap_grad_reduce, ( + "Muon optimizer does not support overlap grad reduce. Use dist_muon instead." + ) + assert not args.overlap_param_gather, ( + "Muon optimizer does not support overlap param gather. Use dist_muon instead." + ) - assert not args.use_distributed_optimizer, "Muon optimizer does not support distributed optimizer for now." + assert not args.use_distributed_optimizer, ( + "Muon optimizer does not support distributed optimizer for now." + ) assert not args.use_torch_fsdp2, "Muon optimizer does not support Torch-FSDP2 for now." assert not args.use_megatron_fsdp, "Muon optimizer does not support Megatron-FSDP for now." - assert args.ckpt_format in ["torch", "torch_dist"], "Muon optimizer supports torch and torch_dist checkpoint format." + assert args.ckpt_format in ["torch", "torch_dist"], ( + "Muon optimizer supports torch and torch_dist checkpoint format." + ) # Optimizer CPU offload check if args.optimizer_cpu_offload: @@ -1554,17 +1774,24 @@ def validate_args(args, defaults={}): ) if args.non_persistent_ckpt_type == "local": - assert args.non_persistent_local_ckpt_dir is not None, "Tried to use local checkpointing without specifying --local-ckpt-dir!" + assert args.non_persistent_local_ckpt_dir is not None, ( + "Tried to use local checkpointing without specifying --local-ckpt-dir!" + ) if args.replication: - assert args.replication_jump is not None, "--replication requires the value of --replication-jump!" - assert args.non_persistent_ckpt_type == "local", f"--replication requires args.non_persistent_ckpt_type == 'local', but got: {args.non_persistent_ckpt_type}" + assert args.replication_jump is not None, ( + "--replication requires the value of --replication-jump!" + ) + assert args.non_persistent_ckpt_type == "local", ( + f"--replication requires args.non_persistent_ckpt_type == 'local', but got: {args.non_persistent_ckpt_type}" + ) elif args.replication_jump: warn_rank_0("--replication-jump was specified despite not using replication. Ignoring.") args.replication_jump = None if args.delay_wgrad_compute: - assert args.transformer_impl == 'transformer_engine', \ + assert args.transformer_impl == "transformer_engine", ( "Delaying wgrad compute is only supported with transformer_engine implementation" + ) if args.overlap_grad_reduce: assert is_te_min_version("2.8.0"), ( "overlap_grad_reduce is only supported with TE >= 2.8.0 when enabling delay_wgrad_compute" @@ -1575,18 +1802,18 @@ def validate_args(args, defaults={}): and not args.moe_shared_expert_overlap ) if wgrad_in_graph_scope: - assert is_te_min_version( - "2.12.0" - ), "CUDA graph with delay_wgrad_compute requires TE version >= 2.12.0." + assert is_te_min_version("2.12.0"), ( + "CUDA graph with delay_wgrad_compute requires TE version >= 2.12.0." + ) assert args.gradient_accumulation_fusion, ( - 'CUDA graph with delay_wgrad_compute requires gradient_accumulation_fusion ' - 'to be enabled. This is because the default gradient accumulation does not ' - 'use static memory addresses, which breaks CUDA graph requirements.' + "CUDA graph with delay_wgrad_compute requires gradient_accumulation_fusion " + "to be enabled. This is because the default gradient accumulation does not " + "use static memory addresses, which breaks CUDA graph requirements." ) if CudaGraphScope.attn in args.cuda_graph_scope: - assert ( - not args.add_bias_linear and not args.add_qkv_bias - ), "CUDA graph with delay_wgrad_compute doesn't support attn bias for now." + assert not args.add_bias_linear and not args.add_qkv_bias, ( + "CUDA graph with delay_wgrad_compute doesn't support attn bias for now." + ) if not args.gradient_accumulation_fusion: assert is_te_min_version("2.7.0"), ( @@ -1595,14 +1822,18 @@ def validate_args(args, defaults={}): ) if args.fine_grained_activation_offloading: - assert args.transformer_impl == 'transformer_engine', \ + assert args.transformer_impl == "transformer_engine", ( "Fine-grained activation offloading is only supported with transformer_engine implementation" + ) if is_te_min_version("2.10.0"): - assert os.getenv("NVTE_CPU_OFFLOAD_V1", "0") == "1", \ + assert os.getenv("NVTE_CPU_OFFLOAD_V1", "0") == "1", ( "For fine-grained activation offloading with TE >= 2.10.0, NVTE_CPU_OFFLOAD_V1 should be set to 1 to avoid offloading weights." + ) if args.mtp_num_layers: - assert not args.use_legacy_models, "The legacy Megatron models does not support Multi-Token Prediction (MTP)." + assert not args.use_legacy_models, ( + "The legacy Megatron models does not support Multi-Token Prediction (MTP)." + ) # MTP is compatible with position embedding types that use position_ids. supported_position_types = ["learned_absolute", "rope", "mrope", "none"] assert args.position_embedding_type in supported_position_types, ( @@ -1636,20 +1867,30 @@ def validate_args(args, defaults={}): assert args.cuda_graph_scope == ["full"], "full scope cannot be used with other scopes." args.cuda_graph_scope = [] warn_rank_0( - 'full scope is deprecated. Use empty cuda_graph_scope to capture the whole layer.' + "full scope is deprecated. Use empty cuda_graph_scope to capture the whole layer." ) - + if args.multi_latent_attention: - assert not args.group_query_attention, "Group query attention is mutually exclusive with multi latent attention." - + assert not args.group_query_attention, ( + "Group query attention is mutually exclusive with multi latent attention." + ) + if args.mla_down_proj_fusion: - assert args.multi_latent_attention, "--mla-down-proj-fusion requires --multi-latent-attention" + assert args.multi_latent_attention, ( + "--mla-down-proj-fusion requires --multi-latent-attention" + ) # MoE latent projections if args.moe_latent_size is not None: - assert args.moe_latent_size > 0, "MoE latent projection dimension has to be greater than zero." - assert args.num_experts is not None, "MoE latent projections are applicable only for MoE models." - assert not args.use_legacy_models, "MoE latent projections are only supported for mcore models." + assert args.moe_latent_size > 0, ( + "MoE latent projection dimension has to be greater than zero." + ) + assert args.num_experts is not None, ( + "MoE latent projections are applicable only for MoE models." + ) + assert not args.use_legacy_models, ( + "MoE latent projections are only supported for mcore models." + ) if args.tiktoken_special_tokens and not args.tokenizer_special_tokens: warn_rank_0( @@ -1657,7 +1898,7 @@ def validate_args(args, defaults={}): "Use --tokenizer-special-tokens instead." ) args.tokenizer_special_tokens = args.tiktoken_special_tokens - + if args.tokenizer_hf_use_fast: warn_rank_0( "--tokenizer-hf-use-fast argument is deprecated and will be removed soon. " @@ -1681,23 +1922,23 @@ def validate_args(args, defaults={}): def _print_args(title, args): """Print arguments.""" from megatron.training.utils import is_rank0 + if is_rank0(): - print(f'------------------------ {title} ------------------------', flush=True) + print(f"------------------------ {title} ------------------------", flush=True) str_list = [] for arg in vars(args): - dots = '.' * (48 - len(arg)) - str_list.append(' {} {} {}'.format(arg, dots, getattr(args, arg))) + dots = "." * (48 - len(arg)) + str_list.append(" {} {} {}".format(arg, dots, getattr(args, arg))) for arg in sorted(str_list, key=lambda x: x.lower()): print(arg, flush=True) - print(f'-------------------- end of {title} ---------------------', flush=True) + print(f"-------------------- end of {title} ---------------------", flush=True) def _check_arg_is_not_none(args, arg): - assert getattr(args, arg) is not None, '{} argument is None'.format(arg) + assert getattr(args, arg) is not None, "{} argument is None".format(arg) def core_transformer_config_from_args(args, config_class=None): - # Config class. config_class = config_class or TransformerConfig @@ -1705,10 +1946,14 @@ def core_transformer_config_from_args(args, config_class=None): if config_class == TransformerConfig: config_class = MLATransformerConfig else: - assert issubclass(config_class, MLATransformerConfig), "config_class should be a subclass of MLATransformerConfig when using multi_latent_attention." - + assert issubclass(config_class, MLATransformerConfig), ( + "config_class should be a subclass of MLATransformerConfig when using multi_latent_attention." + ) + if args.heterogeneous_layers_config_path is not None: - assert not args.multi_latent_attention, "Multi latent attention with heterogeneous layers is not supported." + assert not args.multi_latent_attention, ( + "Multi latent attention with heterogeneous layers is not supported." + ) config_class = HeterogeneousTransformerConfig # Translate args to core transformer configuration @@ -1716,306 +1961,436 @@ def core_transformer_config_from_args(args, config_class=None): for f in dataclasses.fields(config_class): if hasattr(args, f.name): kw_args[f.name] = getattr(args, f.name) - kw_args['persist_layer_norm'] = not args.no_persist_layer_norm - kw_args['deallocate_pipeline_outputs'] = True - kw_args['pipeline_dtype'] = args.params_dtype - kw_args['batch_p2p_comm'] = not args.overlap_p2p_comm - kw_args['num_moe_experts'] = args.num_experts - kw_args['actual_vocab_size'] = args.padded_vocab_size - kw_args['rotary_interleaved'] = args.rotary_interleaved - kw_args['num_layers_in_first_pipeline_stage']= args.decoder_first_pipeline_num_layers - kw_args['num_layers_in_last_pipeline_stage']= args.decoder_last_pipeline_num_layers - kw_args['fp8_param'] = args.fp8_param_gather - kw_args['fp4_param'] = args.fp4_param_gather + kw_args["persist_layer_norm"] = not args.no_persist_layer_norm + kw_args["deallocate_pipeline_outputs"] = True + kw_args["pipeline_dtype"] = args.params_dtype + kw_args["batch_p2p_comm"] = not args.overlap_p2p_comm + kw_args["num_moe_experts"] = args.num_experts + kw_args["actual_vocab_size"] = args.padded_vocab_size + kw_args["rotary_interleaved"] = args.rotary_interleaved + kw_args["num_layers_in_first_pipeline_stage"] = args.decoder_first_pipeline_num_layers + kw_args["num_layers_in_last_pipeline_stage"] = args.decoder_last_pipeline_num_layers + kw_args["fp8_param"] = args.fp8_param_gather + kw_args["fp4_param"] = args.fp4_param_gather if args.swiglu: - kw_args['activation_func'] = F.silu - kw_args['gated_linear_unit'] = True - kw_args['bias_activation_fusion'] = args.bias_swiglu_fusion + kw_args["activation_func"] = F.silu + kw_args["gated_linear_unit"] = True + kw_args["bias_activation_fusion"] = args.bias_swiglu_fusion else: - kw_args['bias_activation_fusion'] = args.bias_gelu_fusion + kw_args["bias_activation_fusion"] = args.bias_gelu_fusion if args.squared_relu: assert not args.swiglu - kw_args['activation_func'] = squared_relu + kw_args["activation_func"] = squared_relu elif args.quick_geglu: assert not args.swiglu - kw_args['gated_linear_unit'] = True - kw_args['activation_func'] = quick_gelu + kw_args["gated_linear_unit"] = True + kw_args["activation_func"] = quick_gelu if args.init_method_xavier_uniform: - kw_args['init_method'] = torch.nn.init.xavier_uniform_ - kw_args['scaled_init_method'] = torch.nn.init.xavier_uniform_ + kw_args["init_method"] = torch.nn.init.xavier_uniform_ + kw_args["scaled_init_method"] = torch.nn.init.xavier_uniform_ if args.group_query_attention: - kw_args['num_query_groups'] = args.num_query_groups + kw_args["num_query_groups"] = args.num_query_groups else: - kw_args['num_query_groups'] = None - kw_args['config_logger_dir'] = args.config_logger_dir + kw_args["num_query_groups"] = None + kw_args["config_logger_dir"] = args.config_logger_dir if args.rope_type is None: # Pop 'rope_type' to let the config class use the default value. - kw_args.pop('rope_type', None) + kw_args.pop("rope_type", None) else: - assert (args.multi_latent_attention or args.rope_type == 'rope'), ( + assert args.multi_latent_attention or args.rope_type == "rope", ( f'Common attention only support rope_type="rope", but got {args.rope_type}.' ) if len(args.cp_comm_type) == 1: - kw_args['cp_comm_type'] = args.cp_comm_type[0] + kw_args["cp_comm_type"] = args.cp_comm_type[0] if args.hybrid_layer_pattern is not None: - kw_args['is_hybrid_model'] = True + kw_args["is_hybrid_model"] = True - kw_args['inference_sampling_seed'] = args.seed + kw_args["inference_sampling_seed"] = args.seed # handle quantization config # NOTE: Kitchen arguments are only added to the namespace when # Kitchen library is available. if hasattr(args, "kitchen_config_file") and args.kitchen_config_file is not None: - kw_args['use_kitchen'] = True - kw_args['quant_recipe'] = load_quantization_recipe(args.kitchen_config_file) - elif hasattr(args, 'kitchen_recipe_number') and args.kitchen_recipe_number is not None: - kw_args['use_kitchen'] = True - kw_args['quant_recipe'] = kitchen_quantization_recipe_config(args.kitchen_recipe_number) + kw_args["use_kitchen"] = True + kw_args["quant_recipe"] = load_quantization_recipe(args.kitchen_config_file) + elif hasattr(args, "kitchen_recipe_number") and args.kitchen_recipe_number is not None: + kw_args["use_kitchen"] = True + kw_args["quant_recipe"] = kitchen_quantization_recipe_config(args.kitchen_recipe_number) - kw_args['moe_latent_size'] = args.moe_latent_size + kw_args["moe_latent_size"] = args.moe_latent_size if args.te_precision_config_file: - assert not 'quant_recipe' in kw_args, "Quantization recipe already configured." + assert not "quant_recipe" in kw_args, "Quantization recipe already configured." # TODO(kwyss): Prohibit fp8_params or fp4_params with this flexibility - kw_args['quant_recipe'] = load_quantization_recipe(args.te_precision_config_file) + kw_args["quant_recipe"] = load_quantization_recipe(args.te_precision_config_file) if hasattr(args, "use_kitchen_attention"): - kw_args['use_kitchen_attention'] = args.use_kitchen_attention + kw_args["use_kitchen_attention"] = args.use_kitchen_attention if hasattr(args, "kitchen_attention_backend"): - kw_args['kitchen_attention_backend'] = args.kitchen_attention_backend + kw_args["kitchen_attention_backend"] = args.kitchen_attention_backend # Return config. return config_class(**kw_args) def _add_transformer_engine_args(parser): - group = parser.add_argument_group(title='Transformer-Engine') + group = parser.add_argument_group(title="Transformer-Engine") # FP4 related arguments - group.add_argument('--fp4-param-gather', action='store_true', - help='Keep the compute param in fp4 (do not use any other intermediate ' - 'dtype) and perform the param all-gather in fp4.') + group.add_argument( + "--fp4-param-gather", + action="store_true", + help="Keep the compute param in fp4 (do not use any other intermediate " + "dtype) and perform the param all-gather in fp4.", + ) # FP8 related arguments - group.add_argument('--fp8-param-gather', action='store_true', - help='Keep the compute param in fp8 (do not use any other intermediate ' - 'dtype) and perform the param all-gather in fp8.') + group.add_argument( + "--fp8-param-gather", + action="store_true", + help="Keep the compute param in fp8 (do not use any other intermediate " + "dtype) and perform the param all-gather in fp8.", + ) # TE precision config file - group.add_argument('--te-precision-config-file', default=None, - help='Configuration file to select per-module precision overrides. ' - 'See TransformerEngineMixedPrecision.md') + group.add_argument( + "--te-precision-config-file", + default=None, + help="Configuration file to select per-module precision overrides. " + "See TransformerEngineMixedPrecision.md", + ) return parser -def _add_inference_args(parser): - group = parser.add_argument_group(title='inference') - - group.add_argument('--inference-batch-times-seqlen-threshold', - type=int, default=-1, - help='If (batch-size * sequence-length) is smaller than this threshold' - 'then batches will not be split up for pipelining.' - 'Requires setting --pipeline-model-parallel-size > 1.' - 'Setting this to -1 indicates that batch pipelining is not used.') - group.add_argument('--max-tokens-to-oom', - type=int, default=12000, - help='Maximum number of tokens during inference' - 'tokens here is # in prompt + # to generate' - 'Allows us to throw an error before OOM crashes server') - group.add_argument('--output-bert-embeddings', action='store_true', - help='Output Bert embeddings (via mean pooling) from ' - 'model, rather than its binary head output or entire ' - 'hidden batch.') - group.add_argument('--bert-embedder-type', default="megatron", - choices=["megatron", "huggingface"], - help='Select either Megatron or Huggingface as the ' - 'Bert embedder.') - group.add_argument('--cuda-graph-scope', nargs='+', type=lambda scope: CudaGraphScope[scope] if scope != "full" else scope, default=[], - help='Determines the CUDA graphs capturing scope. ' - 'choices: "attn", "mlp", "moe", "moe_router", "moe_preprocess", "mamba", "full_iteration". ' - '"attn": captures operations in TransformerLayer._forward_attention(). ' - '"mlp": captures operations in TransformerLayer._forward_mlp() for a dense layer. ' - '"moe": captures operations in TransformerLayer._forward_mlp() for a MoE layer. ' - '"moe_router": captures operations in TransformerLayer._forward_mlp() up to MoELayer.router(), ' - 'including the shared experts if they are not overlapped with EP comm. ' - '"moe_preprocess": captures operations in MoELayer.preprocess(). Must be used together with "moe_router". ' - '"mamba": captures the mamba layer. ' - '"full_iteration": captures a whole training iteration. ' - '"full_iteration_inference": captures a whole inference iteration. ' - 'full_iteration and full_iteration_inference scopes are only supported with --cuda-graph-impl=local, other scopes are only supported with --cuda-graph-impl=transformer_engine. ' - 'If not specified, the default scope is to capture the whole Transformer layer. ' - 'For backward compatibility, we still allow passing "full" to specify capturing the whole layer, and convert it to an empty list.') - group.add_argument('--use-legacy-static-engine', action='store_true', default=False, - help='Use legacy static engine. (Current static engine uses dynamic engine under the hood)', - dest='use_legacy_static_engine') - group.add_argument('--inference-max-requests', type=int, default=8, - help='Maximum number of requests for inference.', - dest='inference_max_requests') - group.add_argument('--inference-max-seq-length', type=int, default=2560, - help='Maximum sequence length expected for inference (prefill + decode).', - dest='inference_max_seq_length') - group.add_argument('--inference-dynamic-batching', - action='store_true', default=False, - help='Enable dynamic batching mode.') - group.add_argument('--inference-dynamic-batching-buffer-size-gb', - type=float, default=40., - help='Amount of on-GPU memory allocated for the KV cache. ' - 'The total amount of memory allocated for the KV cache ' - '(CPU + GPU memory) depends on the value set for the ' - 'unified virtual memory (UVM) level (via ' - '`--inference-dynamic-batching-unified-memory-level`).' - 'If the UVM level is 0, then only GPU memory is used and ' - 'the total memory equals `buffer_size_gb`. If the UVM ' - 'level is 1, then additional memory is utilized on the ' - 'CPU and the total memory equals `buffer_size_gb + ' - 'paused_buffer_size_gb`.') - group.add_argument('--inference-dynamic-batching-paused-buffer-size-gb', - type=float, default=None, - help='Amount of memory reserved for paused requests in ' - 'the dynamic inference context. Active requests are ' - 'paused when there are not enough active blocks available ' - 'to continue generating a request.') - group.add_argument('--inference-dynamic-batching-mamba-memory-ratio', type=float, default=None, - help='Percentage of memory buffer to allocate for Mamba states. ' - 'If not specified, allocates Mamba state tensors for each KV cache block. ' - 'Only used for hybrid models.') - group.add_argument('--inference-dynamic-batching-block-size', - type=int, default=256, - help='KV cache block size. ' - 'It should be a multiple of 256') - group.add_argument('--inference-dynamic-batching-max-requests', - type=int, default=None, - help='Override the inference context\'s `max_requests`. ' - 'By default, `max_requests` is set to the number of ' - 'blocks in the context\'s memory buffer.') - group.add_argument('--inference-dynamic-batching-max-tokens', - type=int, default=None, - help='Override the inference context\'s default `max_tokens`.') - group.add_argument('--inference-dynamic-batching-num-cuda-graphs', - type=int, default=16, - help='Maximum number of cuda graphs to capture, where the ' - 'cuda graph batch sizes range from 1 to `max_requests`. ' - '(See `dynamic_context.py` for details on how ' - '`max_requests` is computed). Due to rounding, the actual ' - 'number of cuda graphs may not equal this argument.' - 'The user can also pass -1, in which case we automatically determine the number of graphs ' \ - 'to capture based on the `max_requests`.') - group.add_argument('--inference-dynamic-batching-track-paused-request-events', - action='store_true', - help='Track paused request ids by adding \'paused\' events ' - 'to each request\'s event history. This has a very minor ' - 'impact on latency.') - group.add_argument('--inference-dynamic-batching-track-generated-token-events', - action='store_true', - help='Track per-token events with timestamps for each generated token. ' - 'When enabled, each generated token creates a GENERATED_TOKEN event ' - 'with a timestamp, useful for per-token latency analysis.') - group.add_argument('--decode-only-cuda-graphs', - action='store_true', default=False, - help='Only use cuda graphs for decode-only steps, not prefill and mixed steps.') - group.add_argument('--inference-dynamic-batching-unified-memory-level', - type=int, default=0, choices=[0, 1], - help='Set unified memory usage within the dynamic ' - 'inference context. The levels are: 0) no unified memory, ' - '1) allocate `memory_buffer` in unified memory. ' - 'Eventually, additional levels will be included to ' - 'control other tensors within the context.') - group.add_argument('--enable-chunked-prefill', dest='enable_chunked_prefill', - action='store_true', default=False, - help="Enable chunked prefill (disabled by default)") - group.add_argument('--num-speculative-tokens', type=int, default=0, - help='Number of speculative tokens generated during decode') - group.add_argument('--inference-dynamic-batching-prefix-caching', - dest='inference_dynamic_batching_enable_prefix_caching', - action=argparse.BooleanOptionalAction, - default=False, - help='Enable/disable prefix caching for dynamic batching inference. ' - 'When disabled, KV cache blocks cannot be shared between ' - 'requests with identical prompt prefixes.') - group.add_argument('--inference-dynamic-batching-prefix-caching-eviction-policy', - type=str, default='ref_zero', - choices=['ref_zero', 'lru'], - dest='inference_dynamic_batching_prefix_caching_eviction_policy', - help='Eviction policy for prefix caching blocks. ' - '"ref_zero" (default) immediately returns blocks to the ' - 'free pool when ref_count hits 0. "lru" keeps blocks ' - 'cached and evicts via LRU only when space is needed.') - group.add_argument('--inference-dynamic-batching-prefix-caching-coordinator-policy', - type=str, default='first_prefix_block', - choices=['longest_prefix', 'first_prefix_block', 'round_robin'], - dest='inference_dynamic_batching_prefix_caching_coordinator_policy', - help='Coordinator routing policy for prefix caching. ' - '"first_prefix_block" (default) routes based on the first ' - 'block hash only. "longest_prefix" routes to the rank with ' - 'the longest matching prefix. "round_robin" ignores prefix ' - 'affinity and cycles through ranks.') - group.add_argument('--inference-dynamic-batching-prefix-caching-routing-alpha', - type=float, default=0.5, - dest='inference_dynamic_batching_prefix_caching_routing_alpha', - help='Weight for prefix-aware routing score: ' - 'score = alpha * match + (1 - alpha) * normalized_load. ' - 'Higher alpha favors prefix cache hits; lower alpha ' - 'favors load balance. Default: 0.5.') - group.add_argument('--inference-dynamic-batching-prefix-caching-mamba-gb', - type=float, default=None, - dest='inference_dynamic_batching_prefix_caching_mamba_gb', - help='GPU memory budget (in GB) for the Mamba state cache ' - 'used by prefix caching on hybrid models. When set, Mamba ' - 'states at block boundaries are cached for reuse.') - group.add_argument('--inference-dynamic-batching-cuda-graph-max-tokens', - type=int, default=16384, - help='Maximum number of tokens to capture in a cuda graph.') - group.add_argument('--inference-dynamic-batching-cuda-graph-mixed-prefill-count', - type=int, default=16, - help='Number of mixed prefill requests to capture in a cuda graph.') - group.add_argument('--inference-logging-step-interval', type=int, default=0, - help='Step interval for logging inference metrics. ' - 'Default to 0 to disable inference logging.') - group.add_argument('--inference-text-gen-server-logging', action=argparse.BooleanOptionalAction, - required=False, default=False, - help='Enable per-request logging in the inference text generation server.') - group.add_argument('--inference-wandb-logging', action=argparse.BooleanOptionalAction, - required=False, default=False, help='Enable inference wandb logging.') - group.add_argument("--inference-coordinator-port", type=int, - help="This port will be used to setup the inference coordinator on node-0") - group.add_argument('--mamba-inference-conv-states-dtype', type=str, - choices=['bf16', 'fp16', 'fp32'], default='bf16', - help='Dtype for the Mamba inference conv states tensor') - group.add_argument('--mamba-inference-ssm-states-dtype', type=str, - choices=['bf16', 'fp16', 'fp32'], default='bf16', - help='Dtype for the Mamba inference SSM states tensor') - group.add_argument('--inference-use-synchronous-zmq-collectives', action=argparse.BooleanOptionalAction, - required=False, default=False, help='Use synchronous ZMQ collectives for inference. Helps in reducing performance variability for MoEs.') - return parser +def _add_inference_args(parser): + group = parser.add_argument_group(title="inference") -def _add_network_size_args(parser): - exclude = [ - # cannot provide callables over CLI - "timers", - "finalize_model_grads_func", - "grad_scale_func", - "no_sync_func", - "grad_sync_func", - "param_sync_func", - "_cpu_offloading_context", - "init_method", - "output_layer_init_method", - "embedding_init_method", - "activation_func", - # types affect docstring - "pipeline_model_parallel_layout", - "window_size", - "window_attn_skip_freq", - "no_rope_freq", - "moe_layer_freq", - "linear_attention_freq", - "csa_compress_ratios", - "moe_router_load_balancing_type", - "moe_aux_loss_coeff", - "cp_comm_type", - "cuda_graph_scope", - # no CLI argument exists for these - "virtual_pipeline_model_parallel_size", - "params_dtype", + group.add_argument( + "--inference-batch-times-seqlen-threshold", + type=int, + default=-1, + help="If (batch-size * sequence-length) is smaller than this threshold" + "then batches will not be split up for pipelining." + "Requires setting --pipeline-model-parallel-size > 1." + "Setting this to -1 indicates that batch pipelining is not used.", + ) + group.add_argument( + "--max-tokens-to-oom", + type=int, + default=12000, + help="Maximum number of tokens during inference" + "tokens here is # in prompt + # to generate" + "Allows us to throw an error before OOM crashes server", + ) + group.add_argument( + "--output-bert-embeddings", + action="store_true", + help="Output Bert embeddings (via mean pooling) from " + "model, rather than its binary head output or entire " + "hidden batch.", + ) + group.add_argument( + "--bert-embedder-type", + default="megatron", + choices=["megatron", "huggingface"], + help="Select either Megatron or Huggingface as the Bert embedder.", + ) + group.add_argument( + "--cuda-graph-scope", + nargs="+", + type=lambda scope: CudaGraphScope[scope] if scope != "full" else scope, + default=[], + help="Determines the CUDA graphs capturing scope. " + 'choices: "attn", "mlp", "moe", "moe_router", "moe_preprocess", "mamba", "full_iteration". ' + '"attn": captures operations in TransformerLayer._forward_attention(). ' + '"mlp": captures operations in TransformerLayer._forward_mlp() for a dense layer. ' + '"moe": captures operations in TransformerLayer._forward_mlp() for a MoE layer. ' + '"moe_router": captures operations in TransformerLayer._forward_mlp() up to MoELayer.router(), ' + "including the shared experts if they are not overlapped with EP comm. " + '"moe_preprocess": captures operations in MoELayer.preprocess(). Must be used together with "moe_router". ' + '"mamba": captures the mamba layer. ' + '"full_iteration": captures a whole training iteration. ' + '"full_iteration_inference": captures a whole inference iteration. ' + "full_iteration and full_iteration_inference scopes are only supported with --cuda-graph-impl=local, other scopes are only supported with --cuda-graph-impl=transformer_engine. " + "If not specified, the default scope is to capture the whole Transformer layer. " + 'For backward compatibility, we still allow passing "full" to specify capturing the whole layer, and convert it to an empty list.', + ) + group.add_argument( + "--use-legacy-static-engine", + action="store_true", + default=False, + help="Use legacy static engine. (Current static engine uses dynamic engine under the hood)", + dest="use_legacy_static_engine", + ) + group.add_argument( + "--inference-max-requests", + type=int, + default=8, + help="Maximum number of requests for inference.", + dest="inference_max_requests", + ) + group.add_argument( + "--inference-max-seq-length", + type=int, + default=2560, + help="Maximum sequence length expected for inference (prefill + decode).", + dest="inference_max_seq_length", + ) + group.add_argument( + "--inference-dynamic-batching", + action="store_true", + default=False, + help="Enable dynamic batching mode.", + ) + group.add_argument( + "--inference-dynamic-batching-buffer-size-gb", + type=float, + default=40.0, + help="Amount of on-GPU memory allocated for the KV cache. " + "The total amount of memory allocated for the KV cache " + "(CPU + GPU memory) depends on the value set for the " + "unified virtual memory (UVM) level (via " + "`--inference-dynamic-batching-unified-memory-level`)." + "If the UVM level is 0, then only GPU memory is used and " + "the total memory equals `buffer_size_gb`. If the UVM " + "level is 1, then additional memory is utilized on the " + "CPU and the total memory equals `buffer_size_gb + " + "paused_buffer_size_gb`.", + ) + group.add_argument( + "--inference-dynamic-batching-paused-buffer-size-gb", + type=float, + default=None, + help="Amount of memory reserved for paused requests in " + "the dynamic inference context. Active requests are " + "paused when there are not enough active blocks available " + "to continue generating a request.", + ) + group.add_argument( + "--inference-dynamic-batching-mamba-memory-ratio", + type=float, + default=None, + help="Percentage of memory buffer to allocate for Mamba states. " + "If not specified, allocates Mamba state tensors for each KV cache block. " + "Only used for hybrid models.", + ) + group.add_argument( + "--inference-dynamic-batching-block-size", + type=int, + default=256, + help="KV cache block size. It should be a multiple of 256", + ) + group.add_argument( + "--inference-dynamic-batching-max-requests", + type=int, + default=None, + help="Override the inference context's `max_requests`. " + "By default, `max_requests` is set to the number of " + "blocks in the context's memory buffer.", + ) + group.add_argument( + "--inference-dynamic-batching-max-tokens", + type=int, + default=None, + help="Override the inference context's default `max_tokens`.", + ) + group.add_argument( + "--inference-dynamic-batching-num-cuda-graphs", + type=int, + default=16, + help="Maximum number of cuda graphs to capture, where the " + "cuda graph batch sizes range from 1 to `max_requests`. " + "(See `dynamic_context.py` for details on how " + "`max_requests` is computed). Due to rounding, the actual " + "number of cuda graphs may not equal this argument." + "The user can also pass -1, in which case we automatically determine the number of graphs " + "to capture based on the `max_requests`.", + ) + group.add_argument( + "--inference-dynamic-batching-track-paused-request-events", + action="store_true", + help="Track paused request ids by adding 'paused' events " + "to each request's event history. This has a very minor " + "impact on latency.", + ) + group.add_argument( + "--inference-dynamic-batching-track-generated-token-events", + action="store_true", + help="Track per-token events with timestamps for each generated token. " + "When enabled, each generated token creates a GENERATED_TOKEN event " + "with a timestamp, useful for per-token latency analysis.", + ) + group.add_argument( + "--decode-only-cuda-graphs", + action="store_true", + default=False, + help="Only use cuda graphs for decode-only steps, not prefill and mixed steps.", + ) + group.add_argument( + "--inference-dynamic-batching-unified-memory-level", + type=int, + default=0, + choices=[0, 1], + help="Set unified memory usage within the dynamic " + "inference context. The levels are: 0) no unified memory, " + "1) allocate `memory_buffer` in unified memory. " + "Eventually, additional levels will be included to " + "control other tensors within the context.", + ) + group.add_argument( + "--enable-chunked-prefill", + dest="enable_chunked_prefill", + action="store_true", + default=False, + help="Enable chunked prefill (disabled by default)", + ) + group.add_argument( + "--num-speculative-tokens", + type=int, + default=0, + help="Number of speculative tokens generated during decode", + ) + group.add_argument( + "--inference-dynamic-batching-prefix-caching", + dest="inference_dynamic_batching_enable_prefix_caching", + action=argparse.BooleanOptionalAction, + default=False, + help="Enable/disable prefix caching for dynamic batching inference. " + "When disabled, KV cache blocks cannot be shared between " + "requests with identical prompt prefixes.", + ) + group.add_argument( + "--inference-dynamic-batching-prefix-caching-eviction-policy", + type=str, + default="ref_zero", + choices=["ref_zero", "lru"], + dest="inference_dynamic_batching_prefix_caching_eviction_policy", + help="Eviction policy for prefix caching blocks. " + '"ref_zero" (default) immediately returns blocks to the ' + 'free pool when ref_count hits 0. "lru" keeps blocks ' + "cached and evicts via LRU only when space is needed.", + ) + group.add_argument( + "--inference-dynamic-batching-prefix-caching-coordinator-policy", + type=str, + default="first_prefix_block", + choices=["longest_prefix", "first_prefix_block", "round_robin"], + dest="inference_dynamic_batching_prefix_caching_coordinator_policy", + help="Coordinator routing policy for prefix caching. " + '"first_prefix_block" (default) routes based on the first ' + 'block hash only. "longest_prefix" routes to the rank with ' + 'the longest matching prefix. "round_robin" ignores prefix ' + "affinity and cycles through ranks.", + ) + group.add_argument( + "--inference-dynamic-batching-prefix-caching-routing-alpha", + type=float, + default=0.5, + dest="inference_dynamic_batching_prefix_caching_routing_alpha", + help="Weight for prefix-aware routing score: " + "score = alpha * match + (1 - alpha) * normalized_load. " + "Higher alpha favors prefix cache hits; lower alpha " + "favors load balance. Default: 0.5.", + ) + group.add_argument( + "--inference-dynamic-batching-prefix-caching-mamba-gb", + type=float, + default=None, + dest="inference_dynamic_batching_prefix_caching_mamba_gb", + help="GPU memory budget (in GB) for the Mamba state cache " + "used by prefix caching on hybrid models. When set, Mamba " + "states at block boundaries are cached for reuse.", + ) + group.add_argument( + "--inference-dynamic-batching-cuda-graph-max-tokens", + type=int, + default=16384, + help="Maximum number of tokens to capture in a cuda graph.", + ) + group.add_argument( + "--inference-dynamic-batching-cuda-graph-mixed-prefill-count", + type=int, + default=16, + help="Number of mixed prefill requests to capture in a cuda graph.", + ) + group.add_argument( + "--inference-logging-step-interval", + type=int, + default=0, + help="Step interval for logging inference metrics. " + "Default to 0 to disable inference logging.", + ) + group.add_argument( + "--inference-text-gen-server-logging", + action=argparse.BooleanOptionalAction, + required=False, + default=False, + help="Enable per-request logging in the inference text generation server.", + ) + group.add_argument( + "--inference-wandb-logging", + action=argparse.BooleanOptionalAction, + required=False, + default=False, + help="Enable inference wandb logging.", + ) + group.add_argument( + "--inference-coordinator-port", + type=int, + help="This port will be used to setup the inference coordinator on node-0", + ) + group.add_argument( + "--mamba-inference-conv-states-dtype", + type=str, + choices=["bf16", "fp16", "fp32"], + default="bf16", + help="Dtype for the Mamba inference conv states tensor", + ) + group.add_argument( + "--mamba-inference-ssm-states-dtype", + type=str, + choices=["bf16", "fp16", "fp32"], + default="bf16", + help="Dtype for the Mamba inference SSM states tensor", + ) + group.add_argument( + "--inference-use-synchronous-zmq-collectives", + action=argparse.BooleanOptionalAction, + required=False, + default=False, + help="Use synchronous ZMQ collectives for inference. Helps in reducing performance variability for MoEs.", + ) + return parser + + +def _add_network_size_args(parser): + exclude = [ + # cannot provide callables over CLI + "timers", + "finalize_model_grads_func", + "grad_scale_func", + "no_sync_func", + "grad_sync_func", + "param_sync_func", + "_cpu_offloading_context", + "init_method", + "output_layer_init_method", + "embedding_init_method", + "activation_func", + # types affect docstring + "pipeline_model_parallel_layout", + "window_size", + "window_attn_skip_freq", + "no_rope_freq", + "moe_layer_freq", + "linear_attention_freq", + "csa_compress_ratios", + "moe_router_load_balancing_type", + "moe_aux_loss_coeff", + "cp_comm_type", + "cuda_graph_scope", + # no CLI argument exists for these + "virtual_pipeline_model_parallel_size", + "params_dtype", "enable_autocast", "autocast_dtype", "num_microbatches_with_partial_activation_checkpoints", @@ -2096,83 +2471,160 @@ def _add_network_size_args(parser): transformer_factory = ArgumentGroupFactory(TransformerConfig, exclude=exclude) transformer_group = transformer_factory.build_group(parser, "transformer configuration") - group = parser.add_argument_group(title='network size') - - group.add_argument('--encoder-num-layers', type=int, default=None, - help='Number of encoder transformer layers.') - group.add_argument('--decoder-num-layers', type=int, default=None, - help='Number of decoder transformer layers.') - group.add_argument('--group-query-attention', action='store_true', - help='Use group-query attention.') - group.add_argument('--window-size', type=tuple_type, default=None, - help='Window size for window attention. If not provided, ' - 'window attention will be disabled.') - group.add_argument('--window-attn-skip-freq', type=moe_freq_type, default=None, - help='Frequency of layers to skip window attention. Accepts either: ' - '- An integer N: Represents a (N-1):1 ratio, meaning one full attention layer ' - 'after (N-1) SWA layers. ' - '- A string containing a Python list expression that defines a custom pattern, ' - 'e.g.: "[1,1,1,0]*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' - 'where 1 indicates SWA and 0 indicates full attention. ') - group.add_argument('--max-position-embeddings', type=int, default=None, - help='Maximum number of position embeddings to use. ' - 'This is the size of position embedding.') - group.add_argument('--position-embedding-type', type=str, default='learned_absolute', - choices=['learned_absolute', 'rope', 'mrope', 'relative', 'none'], - help='Position embedding type.') - group.add_argument('--relative-attention-num-buckets', type=int, default=32, - help='Number of buckets for relative position embeddings.') - group.add_argument('--relative-attention-max-distance', type=int, default=128, - help='Maximum distance for relative position embeddings calculation.') - group.add_argument('--use-rotary-position-embeddings', action='store_true', - help='Use rotary positional embeddings or not. ' - 'Deprecated: use --position-embedding-type') - group.add_argument('--rotary-base', type=int, default=10000, - help='Base to use for rotary positional embeddings, default 10000') - group.add_argument('--rotary-percent', type=float, default=1.0, - help='Percent of rotary dimension to use, default 100%%') - group.add_argument('--rotary-seq-len-interpolation-factor', type=int, default=None, - help='Sequence length interpolation factor for rotary embeddings.') - group.add_argument('--use-rope-scaling', action='store_true', - help='Apply rope scaling as used in llama3.x') - group.add_argument('--rope-scaling-factor', type=float, default=8.0, - help='Rope scaling factor in llama3.x models') - group.add_argument('--no-rope-freq', type=no_rope_freq_type, default=None, - help='Controls which layers to skip performing Rotary Position Embedding. Accepts either: ' - '- An integer N: Represents a 1:N ratio, meaning RoPE is skipped every N-1 layers. ' - '- A string containing a Python list expression that defines a custom pattern, e.g.: ' - '"([0]*3+[1]*1)*3" evaluates to [0,0,0,1,0,0,0,1,0,0,0,1] ' - 'where 1 indicates no-rope layer. This patten is equivalent to --no-rope-freq=4.' - 'By default this is disabled and set to None, indicating RoPE will be performed' - 'on every layer.' - ) - group.add_argument('--no-position-embedding', - action='store_false', - help='Disable position embedding. Deprecated: use --position-embedding-type', - dest='add_position_embedding') - group.add_argument('--make-vocab-size-divisible-by', type=int, default=128, - help='Pad the vocab size to be divisible by this value.' - 'This is added for computational efficieny reasons.') - group.add_argument('--openai-gelu', action='store_true', - help='Use OpenAIs GeLU implementation. This option' - 'should not be used unless for backward compatibility' - 'reasons.') - group.add_argument('--squared-relu', action='store_true', - help='Use squared relu activation instead of default gelu') - group.add_argument('--swiglu', action='store_true', - help='Use gated linear units and SiLU activation instead of default gelu') - group.add_argument('--quick-geglu', action='store_true', - help='Use quick geglu activation instead of default gelu') - group.add_argument('--onnx-safe', type=bool, required=False, - help='Use workarounds for known problems with ' - 'Torch ONNX exporter') - group.add_argument('--bert-no-binary-head', action='store_false', - help='Disable BERT binary head.', - dest='bert_binary_head') - group.add_argument('--untie-embeddings-and-output-weights', action='store_true', - help='Untie embeddings and output weights.') + group = parser.add_argument_group(title="network size") + + group.add_argument( + "--encoder-num-layers", type=int, default=None, help="Number of encoder transformer layers." + ) + group.add_argument( + "--decoder-num-layers", type=int, default=None, help="Number of decoder transformer layers." + ) + group.add_argument( + "--group-query-attention", action="store_true", help="Use group-query attention." + ) + group.add_argument( + "--window-size", + type=tuple_type, + default=None, + help="Window size for window attention. If not provided, " + "window attention will be disabled.", + ) + group.add_argument( + "--window-attn-skip-freq", + type=moe_freq_type, + default=None, + help="Frequency of layers to skip window attention. Accepts either: " + "- An integer N: Represents a (N-1):1 ratio, meaning one full attention layer " + "after (N-1) SWA layers. " + "- A string containing a Python list expression that defines a custom pattern, " + 'e.g.: "[1,1,1,0]*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' + "where 1 indicates SWA and 0 indicates full attention. ", + ) + group.add_argument( + "--max-position-embeddings", + type=int, + default=None, + help="Maximum number of position embeddings to use. " + "This is the size of position embedding.", + ) + group.add_argument( + "--position-embedding-type", + type=str, + default="learned_absolute", + choices=["learned_absolute", "rope", "mrope", "relative", "none"], + help="Position embedding type.", + ) + group.add_argument( + "--relative-attention-num-buckets", + type=int, + default=32, + help="Number of buckets for relative position embeddings.", + ) + group.add_argument( + "--relative-attention-max-distance", + type=int, + default=128, + help="Maximum distance for relative position embeddings calculation.", + ) + group.add_argument( + "--use-rotary-position-embeddings", + action="store_true", + help="Use rotary positional embeddings or not. Deprecated: use --position-embedding-type", + ) + group.add_argument( + "--rotary-base", + type=int, + default=10000, + help="Base to use for rotary positional embeddings, default 10000", + ) + group.add_argument( + "--rotary-percent", + type=float, + default=1.0, + help="Percent of rotary dimension to use, default 100%%", + ) + group.add_argument( + "--rotary-seq-len-interpolation-factor", + type=int, + default=None, + help="Sequence length interpolation factor for rotary embeddings.", + ) + group.add_argument( + "--use-rope-scaling", action="store_true", help="Apply rope scaling as used in llama3.x" + ) + group.add_argument( + "--rope-scaling-factor", + type=float, + default=8.0, + help="Rope scaling factor in llama3.x models", + ) + group.add_argument( + "--no-rope-freq", + type=no_rope_freq_type, + default=None, + help="Controls which layers to skip performing Rotary Position Embedding. Accepts either: " + "- An integer N: Represents a 1:N ratio, meaning RoPE is skipped every N-1 layers. " + "- A string containing a Python list expression that defines a custom pattern, e.g.: " + '"([0]*3+[1]*1)*3" evaluates to [0,0,0,1,0,0,0,1,0,0,0,1] ' + "where 1 indicates no-rope layer. This patten is equivalent to --no-rope-freq=4." + "By default this is disabled and set to None, indicating RoPE will be performed" + "on every layer.", + ) + group.add_argument( + "--no-position-embedding", + action="store_false", + help="Disable position embedding. Deprecated: use --position-embedding-type", + dest="add_position_embedding", + ) + group.add_argument( + "--make-vocab-size-divisible-by", + type=int, + default=128, + help="Pad the vocab size to be divisible by this value." + "This is added for computational efficieny reasons.", + ) + group.add_argument( + "--openai-gelu", + action="store_true", + help="Use OpenAIs GeLU implementation. This option" + "should not be used unless for backward compatibility" + "reasons.", + ) + group.add_argument( + "--squared-relu", + action="store_true", + help="Use squared relu activation instead of default gelu", + ) + group.add_argument( + "--swiglu", + action="store_true", + help="Use gated linear units and SiLU activation instead of default gelu", + ) + group.add_argument( + "--quick-geglu", + action="store_true", + help="Use quick geglu activation instead of default gelu", + ) + group.add_argument( + "--onnx-safe", + type=bool, + required=False, + help="Use workarounds for known problems with Torch ONNX exporter", + ) + group.add_argument( + "--bert-no-binary-head", + action="store_false", + help="Disable BERT binary head.", + dest="bert_binary_head", + ) + group.add_argument( + "--untie-embeddings-and-output-weights", + action="store_true", + help="Untie embeddings and output weights.", + ) return parser + def _add_straggler_detector_args(parser): from megatron.training.config import StragglerDetectionConfig @@ -2181,330 +2633,616 @@ def _add_straggler_detector_args(parser): return parser + def _add_workload_inspector_server_args(parser): - group = parser.add_argument_group(title='workload inspector') - group.add_argument('--run-workload-inspector-server', action='store_true', - help='If set, enables workload inspector server for on-demand profiling.') + group = parser.add_argument_group(title="workload inspector") + group.add_argument( + "--run-workload-inspector-server", + action="store_true", + help="If set, enables workload inspector server for on-demand profiling.", + ) return parser -def _add_inprocess_restart_args(parser): - group = parser.add_argument_group(title='In-process restart') - - group.add_argument('--inprocess-restart', action='store_true', - help='Enables in-process restart.') - - group.add_argument('--inprocess-max-iterations', default=None, type=int, - help='Maximum number of in-process restart iterations.') - group.add_argument('--inprocess-monitor-thread-interval', default=1.0, type=float, - help='Monitoring interval (in seconds) for the monitoring thread.') - group.add_argument('--inprocess-monitor-process-interval', default=1.0, type=float, - help='Monitoring interval (in seconds) for the monitoring process.') - group.add_argument('--inprocess-progress-watchdog-interval', default=1.0, type=float, - help='Interval (in seconds) for automatic progress watchdog timestamp ' - 'updates.') - group.add_argument('--inprocess-heartbeat-interval', default=30, type=float, - help='Monitoring interval (in seconds) for detecting unresponsive ranks.') - - group.add_argument('--inprocess-soft-timeout', default=60, type=float, - help='Soft progress timeout (in seconds).') - group.add_argument('--inprocess-hard-timeout', default=90, type=float, - help='Hard progress timeout (in seconds).') - group.add_argument('--inprocess-heartbeat-timeout', default=60, type=float, - help='Timeout (in seconds) for a missing rank detection heartbeat.') - - group.add_argument('--inprocess-barrier-timeout', default=120, type=float, - help='Timeout (in seconds) for internal distributed barrier') - group.add_argument('--inprocess-completion-timeout', default=120, type=float, - help='Timeout (in seconds) for barrier on completion on all ranks') - - group.add_argument('--inprocess-last-call-wait', default=1, type=float, - help='Time interval (in seconds) for other ranks to report concurrent ' - 'terminal failures.') - group.add_argument('--inprocess-termination-grace-time', default=1, type=float, - help='Interval (in seconds) between SIGTERM and SIGKILL issued on hard ' - 'timeout') - - group.add_argument('--inprocess-granularity', default='node', type=str, - choices=['node', 'rank'], - help='Granularity for in-process restart.') - group.add_argument('--inprocess-active-world-size', - default=int(os.getenv('WORLD_SIZE', '1')), type=int, - help='The number of ranks initially executing the workload. ' - 'The remaining ranks from the allocation are set aside ' - 'as warm reserve.') - group.add_argument('--inprocess-empty-cuda-cache', action='store_true', - help='Release all unoccupied cached GPU memory on every in-process restart.') - return parser -def _add_one_logger_args(parser): - group = parser.add_argument_group(title='one logger') - group.add_argument('--no-one-logger', action='store_false', - help='If set, disable using one_logger to track E2E metrics' - 'Note that one_logger is an internal tool and not ' - 'available externally. For installation, please go to ' - 'https://confluence.nvidia.com/display/MLWFO/Package+Repositories' - 'for more details', - dest='enable_one_logger') - group.add_argument('--one-logger-project', type=str, default='megatron-lm', - help='The one-logger project name. Will ignore if ' - '--no-one-logger is set') - group.add_argument('--one-logger-run-name', type=str, default=None, - help='The one-logger run name displayed. Will ignore if ' - '--no-one-logger is set') - group.add_argument('--one-logger-async', action='store_true', - help='If set, forces one_logger to use async mode.') - group.add_argument('--app-tag-run-name', type=str, default=None, - help='Jobs belonging to same training run, suppose to ' - 'have the same name. It will be used to track progress of ' - 'a training done over multiple different jobs') - group.add_argument('--app-tag-run-version', type=str, default='0.0.0', - help='The version of the training of which current job is ' - 'part of. It will be used to track the changes in the ' - 'application side which might change the performance ' - 'baseline') - return parser +def _add_inprocess_restart_args(parser): + group = parser.add_argument_group(title="In-process restart") + group.add_argument( + "--inprocess-restart", action="store_true", help="Enables in-process restart." + ) -def _add_ft_package_args(parser): - group = parser.add_argument_group(title='ft_package') - group.add_argument('--enable-ft-package', action='store_true', - help='If set, Fault Tolerance package is enabled. ' - 'Note: This feature is for Nvidia internal use only.') - group.add_argument('--calc-ft-timeouts', action='store_true', - help='If set, FT package will try to automatically compute the timeouts. ' - 'Note: This feature is for Nvidia internal use only.') - group.add_argument('--ft-num-warmup-iters', type=int, default=5, - help='Number of warmup iterations before monitoring step section and ' - 'out-of-section timeouts. The first N iterations are excluded from ' - 'timeout monitoring as they can be significantly slower than steady-state. ' - 'Default: 5. Note: This feature is for Nvidia internal use only.') - return parser + group.add_argument( + "--inprocess-max-iterations", + default=None, + type=int, + help="Maximum number of in-process restart iterations.", + ) + group.add_argument( + "--inprocess-monitor-thread-interval", + default=1.0, + type=float, + help="Monitoring interval (in seconds) for the monitoring thread.", + ) + group.add_argument( + "--inprocess-monitor-process-interval", + default=1.0, + type=float, + help="Monitoring interval (in seconds) for the monitoring process.", + ) + group.add_argument( + "--inprocess-progress-watchdog-interval", + default=1.0, + type=float, + help="Interval (in seconds) for automatic progress watchdog timestamp updates.", + ) + group.add_argument( + "--inprocess-heartbeat-interval", + default=30, + type=float, + help="Monitoring interval (in seconds) for detecting unresponsive ranks.", + ) + group.add_argument( + "--inprocess-soft-timeout", + default=60, + type=float, + help="Soft progress timeout (in seconds).", + ) + group.add_argument( + "--inprocess-hard-timeout", + default=90, + type=float, + help="Hard progress timeout (in seconds).", + ) + group.add_argument( + "--inprocess-heartbeat-timeout", + default=60, + type=float, + help="Timeout (in seconds) for a missing rank detection heartbeat.", + ) -def _add_logging_args(parser): - from megatron.training.config import LoggerConfig + group.add_argument( + "--inprocess-barrier-timeout", + default=120, + type=float, + help="Timeout (in seconds) for internal distributed barrier", + ) + group.add_argument( + "--inprocess-completion-timeout", + default=120, + type=float, + help="Timeout (in seconds) for barrier on completion on all ranks", + ) - log_factory = ArgumentGroupFactory(LoggerConfig, exclude = ["log_throughput_to_tensorboard", "throughput_window_size", "memory_keys", "log_l2_norm_grad_to_tensorboard", "log_runtime_to_tensorboard", "runtime_time_unit", "filter_warnings", "modules_to_filter", "set_level_for_all_loggers", "save_config_filepath"]) - group = log_factory.build_group(parser, title="logging") + group.add_argument( + "--inprocess-last-call-wait", + default=1, + type=float, + help="Time interval (in seconds) for other ranks to report concurrent terminal failures.", + ) + group.add_argument( + "--inprocess-termination-grace-time", + default=1, + type=float, + help="Interval (in seconds) between SIGTERM and SIGKILL issued on hard timeout", + ) + group.add_argument( + "--inprocess-granularity", + default="node", + type=str, + choices=["node", "rank"], + help="Granularity for in-process restart.", + ) + group.add_argument( + "--inprocess-active-world-size", + default=int(os.getenv("WORLD_SIZE", "1")), + type=int, + help="The number of ranks initially executing the workload. " + "The remaining ranks from the allocation are set aside " + "as warm reserve.", + ) + group.add_argument( + "--inprocess-empty-cuda-cache", + action="store_true", + help="Release all unoccupied cached GPU memory on every in-process restart.", + ) return parser -def _add_regularization_args(parser): - group = parser.add_argument_group(title='regularization') - - group.add_argument('--weight-decay', type=float, default=0.01, - help='Weight decay coefficient for L2 regularization.') - group.add_argument('--apply-wd-to-qk-layernorm', action='store_true', - help='Apply weight decay to qk layernorm as a special case.') - group.add_argument('--clip-grad', type=float, default=1.0, - help='Gradient clipping based on global L2 norm.') - group.add_argument('--adam-beta1', type=float, default=0.9, - help='First coefficient for computing running averages ' - 'of gradient and its square') - group.add_argument('--adam-beta2', type=float, default=0.999, - help='Second coefficient for computing running averages ' - 'of gradient and its square') - group.add_argument('--adam-eps', type=float, default=1e-08, - help='Term added to the denominator to improve' - 'numerical stability') - group.add_argument('--sgd-momentum', type=float, default=0.9, - help='Momentum factor for sgd') - group.add_argument('--muon-momentum', type=float, default=0.9, - help='Momentum factor for Muon optimizer') - group.add_argument('--muon-no-split-qkv', action='store_false', default=True, - dest='muon_split_qkv', - help='Whether to split QKV parameters for Muon optimizer') - group.add_argument('--muon-use-nesterov', action='store_true', - help='Whether to use Nesterov-style momentum in the internal SGD') - group.add_argument('--muon-scale-mode', type=str, default='spectral', - choices=['spectral', 'unit_rms_norm', 'shape_scaling'], - help='Scale mode for Muon optimizer. With MuP, set ' - '--muon-scale-mode unit_rms_norm to use unit_rms_norm scaling, ' - 'or set --muon-scale-mode spectral to keep spectral scaling.') - group.add_argument('--muon-fp32-matmul-prec', type=str, default='medium', - choices=['low', 'medium', 'high'], - help='FP32 matmul precision for Newton-Schulz iteration') - group.add_argument('--muon-coefficient-type', type=str, default='quintic', - help='Newton-Schulz coefficient type for the Muon optimizer. ' - 'Valid types are discovered from the installed emerging_optimizers ' - 'package (e.g. simple, quintic, polar_express, aol). ' - 'Validated at optimizer creation time.') - group.add_argument('--muon-num-ns-steps', type=int, default=5, - help='Number of Newton-Schulz steps for Muon optimizer') - group.add_argument('--muon-tp-mode', type=str, default='blockwise', - choices=['blockwise', 'duplicated', 'distributed'], - help='How to perform NS calculation for tensor model parallel weights') - group.add_argument('--muon-extra-scale-factor', type=float, default=1.0, - help='Additional scale factor for the muon update') - group.add_argument('--muon-scalar-optimizer', type=str, default='adam', - choices=['adam', 'lion'], - help='Optimizer for scalar parameters (embeddings, biases, norms) ' - 'when using muon. Defaults to adam.') - group.add_argument('--lion-beta1', type=float, default=0.95, - help='First beta coefficient for Lion optimizer ' - '(used in sign update). Default: 0.95.') - group.add_argument('--lion-beta2', type=float, default=0.98, - help='Second beta coefficient for Lion optimizer ' - '(used in momentum EMA update). Default: 0.98.') - - group.add_argument('--no-weight-decay-cond-type', type=str, choices=['apply_wd_to_qk_layernorm'], - help='Type of no weight decay condition. Choices: ' - 'None (default): apply weight decay to 1D weights and biases.' - '"apply_wd_to_qk_layernorm": additionally apply weight decay to ' - 'qk layernorm as a special case.' - 'DEPRECATED. Please use --apply-wd-to-qk-layernorm instead. ') +def _add_one_logger_args(parser): + group = parser.add_argument_group(title="one logger") + group.add_argument( + "--no-one-logger", + action="store_false", + help="If set, disable using one_logger to track E2E metrics" + "Note that one_logger is an internal tool and not " + "available externally. For installation, please go to " + "https://confluence.nvidia.com/display/MLWFO/Package+Repositories" + "for more details", + dest="enable_one_logger", + ) + group.add_argument( + "--one-logger-project", + type=str, + default="megatron-lm", + help="The one-logger project name. Will ignore if --no-one-logger is set", + ) + group.add_argument( + "--one-logger-run-name", + type=str, + default=None, + help="The one-logger run name displayed. Will ignore if --no-one-logger is set", + ) + group.add_argument( + "--one-logger-async", + action="store_true", + help="If set, forces one_logger to use async mode.", + ) + group.add_argument( + "--app-tag-run-name", + type=str, + default=None, + help="Jobs belonging to same training run, suppose to " + "have the same name. It will be used to track progress of " + "a training done over multiple different jobs", + ) + group.add_argument( + "--app-tag-run-version", + type=str, + default="0.0.0", + help="The version of the training of which current job is " + "part of. It will be used to track the changes in the " + "application side which might change the performance " + "baseline", + ) + return parser + + +def _add_ft_package_args(parser): + group = parser.add_argument_group(title="ft_package") + group.add_argument( + "--enable-ft-package", + action="store_true", + help="If set, Fault Tolerance package is enabled. " + "Note: This feature is for Nvidia internal use only.", + ) + group.add_argument( + "--calc-ft-timeouts", + action="store_true", + help="If set, FT package will try to automatically compute the timeouts. " + "Note: This feature is for Nvidia internal use only.", + ) + group.add_argument( + "--ft-num-warmup-iters", + type=int, + default=5, + help="Number of warmup iterations before monitoring step section and " + "out-of-section timeouts. The first N iterations are excluded from " + "timeout monitoring as they can be significantly slower than steady-state. " + "Default: 5. Note: This feature is for Nvidia internal use only.", + ) + return parser + + +def _add_logging_args(parser): + from megatron.training.config import LoggerConfig + + log_factory = ArgumentGroupFactory( + LoggerConfig, + exclude=[ + "log_throughput_to_tensorboard", + "throughput_window_size", + "memory_keys", + "log_l2_norm_grad_to_tensorboard", + "log_runtime_to_tensorboard", + "runtime_time_unit", + "filter_warnings", + "modules_to_filter", + "set_level_for_all_loggers", + "save_config_filepath", + ], + ) + group = log_factory.build_group(parser, title="logging") + + return parser + + +def _add_regularization_args(parser): + group = parser.add_argument_group(title="regularization") + + group.add_argument( + "--weight-decay", + type=float, + default=0.01, + help="Weight decay coefficient for L2 regularization.", + ) + group.add_argument( + "--apply-wd-to-qk-layernorm", + action="store_true", + help="Apply weight decay to qk layernorm as a special case.", + ) + group.add_argument( + "--clip-grad", type=float, default=1.0, help="Gradient clipping based on global L2 norm." + ) + group.add_argument( + "--adam-beta1", + type=float, + default=0.9, + help="First coefficient for computing running averages of gradient and its square", + ) + group.add_argument( + "--adam-beta2", + type=float, + default=0.999, + help="Second coefficient for computing running averages of gradient and its square", + ) + group.add_argument( + "--adam-eps", + type=float, + default=1e-08, + help="Term added to the denominator to improvenumerical stability", + ) + group.add_argument("--sgd-momentum", type=float, default=0.9, help="Momentum factor for sgd") + group.add_argument( + "--muon-momentum", type=float, default=0.9, help="Momentum factor for Muon optimizer" + ) + group.add_argument( + "--muon-no-split-qkv", + action="store_false", + default=True, + dest="muon_split_qkv", + help="Whether to split QKV parameters for Muon optimizer", + ) + group.add_argument( + "--muon-use-nesterov", + action="store_true", + help="Whether to use Nesterov-style momentum in the internal SGD", + ) + group.add_argument( + "--muon-scale-mode", + type=str, + default="spectral", + choices=["spectral", "unit_rms_norm", "shape_scaling"], + help="Scale mode for Muon optimizer. With MuP, set " + "--muon-scale-mode unit_rms_norm to use unit_rms_norm scaling, " + "or set --muon-scale-mode spectral to keep spectral scaling.", + ) + group.add_argument( + "--muon-fp32-matmul-prec", + type=str, + default="medium", + choices=["low", "medium", "high"], + help="FP32 matmul precision for Newton-Schulz iteration", + ) + group.add_argument( + "--muon-coefficient-type", + type=str, + default="quintic", + help="Newton-Schulz coefficient type for the Muon optimizer. " + "Valid types are discovered from the installed emerging_optimizers " + "package (e.g. simple, quintic, polar_express, aol). " + "Validated at optimizer creation time.", + ) + group.add_argument( + "--muon-num-ns-steps", + type=int, + default=5, + help="Number of Newton-Schulz steps for Muon optimizer", + ) + group.add_argument( + "--muon-tp-mode", + type=str, + default="blockwise", + choices=["blockwise", "duplicated", "distributed"], + help="How to perform NS calculation for tensor model parallel weights", + ) + group.add_argument( + "--muon-extra-scale-factor", + type=float, + default=1.0, + help="Additional scale factor for the muon update", + ) + group.add_argument( + "--muon-scalar-optimizer", + type=str, + default="adam", + choices=["adam", "lion"], + help="Optimizer for scalar parameters (embeddings, biases, norms) " + "when using muon. Defaults to adam.", + ) + group.add_argument( + "--lion-beta1", + type=float, + default=0.95, + help="First beta coefficient for Lion optimizer (used in sign update). Default: 0.95.", + ) + group.add_argument( + "--lion-beta2", + type=float, + default=0.98, + help="Second beta coefficient for Lion optimizer " + "(used in momentum EMA update). Default: 0.98.", + ) + + group.add_argument( + "--no-weight-decay-cond-type", + type=str, + choices=["apply_wd_to_qk_layernorm"], + help="Type of no weight decay condition. Choices: " + "None (default): apply weight decay to 1D weights and biases." + '"apply_wd_to_qk_layernorm": additionally apply weight decay to ' + "qk layernorm as a special case." + "DEPRECATED. Please use --apply-wd-to-qk-layernorm instead. ", + ) return parser def _add_rl_args(parser): - group = parser.add_argument_group(title='rl') - group.add_argument('--perform-rl-step', action='store_true', - help="Use the RL training step.") - group.add_argument('--rl-prompts-per-eval', type=int, default=32, - help='Number of prompts to evaluate for for each RL task.' - 'This evaluation can be very expensive when using environments' - 'that evaluate pass@k so we default to a lower number.') + group = parser.add_argument_group(title="rl") + group.add_argument("--perform-rl-step", action="store_true", help="Use the RL training step.") + group.add_argument( + "--rl-prompts-per-eval", + type=int, + default=32, + help="Number of prompts to evaluate for for each RL task." + "This evaluation can be very expensive when using environments" + "that evaluate pass@k so we default to a lower number.", + ) # TODO(rkirby): allow for "complete" evaluation when --rl-prompts-per-eval is set to -1 - group.add_argument('--grpo-prompts-per-step', type=int, default=32, - help="Number of GRPO groups (G in the paper).") - group.add_argument('--grpo-group-size', type=int, default=2, - help="Number of samples per a GRPO group.") - group.add_argument('--rl-num-parallel-generations', type=int, default=None, - help='Number of rollouts being generated by the inference engine simultaneously. ' - 'Internally divided by grpo_group_size. ' - 'Requires --rl-partial-rollouts. ' - 'Mutually exclusive with --rl-num-parallel-generation-batches.') - group.add_argument('--rl-num-parallel-generation-batches', type=int, default=None, - help='Number of generation batches in flight. ' - 'Set to L+1 to allow for L steps of staleness between the inference and training policies. ' - 'Each batch contains grpo_prompts_per_step groups by default. ' - 'Requires --rl-partial-rollouts. ' - 'Mutually exclusive with --rl-num-parallel-generations.') - group.add_argument('--rl-generation-batch-size', type=int, default=None, - help='Override the number of groups per generation batch. ' - 'Defaults to grpo_prompts_per_step when ' - '--rl-num-parallel-generation-batches is set.') - group.add_argument('--grpo-iterations', type=int, default=2, - help="Number of iterations per a GRPO implementation.") + group.add_argument( + "--grpo-prompts-per-step", + type=int, + default=32, + help="Number of GRPO groups (G in the paper).", + ) + group.add_argument( + "--grpo-group-size", type=int, default=2, help="Number of samples per a GRPO group." + ) + group.add_argument( + "--rl-num-parallel-generations", + type=int, + default=None, + help="Number of rollouts being generated by the inference engine simultaneously. " + "Internally divided by grpo_group_size. " + "Requires --rl-partial-rollouts. " + "Mutually exclusive with --rl-num-parallel-generation-batches.", + ) + group.add_argument( + "--rl-num-parallel-generation-batches", + type=int, + default=None, + help="Number of generation batches in flight. " + "Set to L+1 to allow for L steps of staleness between the inference and training policies. " + "Each batch contains grpo_prompts_per_step groups by default. " + "Requires --rl-partial-rollouts. " + "Mutually exclusive with --rl-num-parallel-generations.", + ) + group.add_argument( + "--rl-generation-batch-size", + type=int, + default=None, + help="Override the number of groups per generation batch. " + "Defaults to grpo_prompts_per_step when " + "--rl-num-parallel-generation-batches is set.", + ) + group.add_argument( + "--grpo-iterations", + type=int, + default=2, + help="Number of iterations per a GRPO implementation.", + ) # As in DAPO, we keep upper/lower eps different. # To have a vanilla GRPO, set them to be the same. - group.add_argument('--grpo-clamp-eps-lower', type=float, default=0.01, - help="Lower GRPO clipping bound.") - group.add_argument('--grpo-clamp-eps-upper', type=float, default=0.01, - help="Upper GRPO clipping bound. In vanilla implementation, equals to the lower one.") - group.add_argument('--grpo-kl-beta', type=float, default=0.001, - help="KL term weight in the GRPO loss.") - group.add_argument('--grpo-entropy-term-weight', type=float, default=0.0, - help="Entropy term weight in GRPO loss.") - group.add_argument('--grpo-filter-groups-with-same-reward', action='store_true', - help="Filter groups with same reward.") - group.add_argument('--langrl-env-config', type=str, default=None, - help="Path to YAML config file for RL environment configuration.") - group.add_argument('--rl-default-temperature', type=float, default=1.0, - help="Default temperature for model inference.") - group.add_argument('--rl-default-top-p', type=float, default=0, - help="Default top-p for model inference.") - group.add_argument('--rl-default-top-k', type=int, default=-1, - help="Default top-k for model inference.") - group.add_argument('--rl-offload-optimizer-during-inference', action='store_true', - help='Offload optimizer state to CPU during inference/rollout to save GPU memory') - group.add_argument('--rl-kv-cache-management-mode', type=str, default='persist', - choices=['persist', 'offload', 'recompute'], - help='KV cache management mode during RL training: ' - 'persist: leave KV cache in GPU memory (default), ' - 'offload: offload KV cache to CPU during training, ' - 'recompute: deallocate KV cache and recompute from scratch each cycle') - group.add_argument('--rl-persist-cuda-graphs', action=argparse.BooleanOptionalAction, type=bool, default=False, - help='Persist CUDA graphs when the inference engine is suspended. ' - 'If False, CUDA graphs are deleted on suspend and re-captured on resume.') - group.add_argument('--rl-partial-rollouts', action=argparse.BooleanOptionalAction, default=False, - help='Allow inference to continue generating rollouts while training updates ' - 'the policy weights. This enables off-policy training where rollouts may ' - 'be generated with a stale version of the policy. Use ' - '--rl-num-parallel-generations or --rl-num-parallel-generation-batches ' - 'to control the degree of staleness.') - group.add_argument('--rl-inference-logprobs-is-correction', action=argparse.BooleanOptionalAction, type=bool, default=False, - help='If set, use inference logprobs in importance sampling correction of the loss.') - group.add_argument('--rl-importance-sampling-truncation-coef', type=float, default=None, - help="If --inference-logprobs-is-correction is on and this coefficient is set, apply truncation for the IS correction at GRPO loss.") - group.add_argument('--rl-use-sequence-packing', action=argparse.BooleanOptionalAction, type=bool, default=False, - help='Enable sequence packing') - group.add_argument('--rl-sequence-packing-max-sequences-per-bin', type=int, default=50, - help='Maximum number of sequences that can be packed into a single bin. ') - group.add_argument('--rl-sequence-packing-algo', type=str, default='fifo', - choices=['fifo', 'round-robin'], - help='Algorithm for distributing packed bins across ranks. ' - 'fifo: first-in-first-out sequential distribution, ' - 'round-robin: distribute bins cyclically across ranks for better load balancing') - group.add_argument('--rl-training-cuda-graphs', action=argparse.BooleanOptionalAction, type=bool, - default=False, - help='If set, do not toggle CUDA graphs on/off between inference and training phases.') - group.add_argument('--rl-inference-tensor-model-parallel-size', type=int, default=None, - help='Degree of tensor model parallelism for inference for RL.') - group.add_argument( - '--rl-inference-pipeline-model-parallel-size', - type=int, - default=None, - help='Degree of pipeline model parallelism for inference for RL.', - ) - group.add_argument( - '--rl-inference-expert-model-parallel-size', - type=int, - default=None, - help='Degree of expert model parallelism for inference for RL.', - ) - group.add_argument( - '--rl-inference-expert-tensor-model-parallel-size', - type=int, - default=None, - help='Degree of expert tensor model parallelism for inference for RL. ' - 'For MoE models, this controls the TP size for expert layers specifically. ' - 'Defaults to training expert_tensor_parallel_size if not specified.', - ) - group.add_argument( - '--rl-inference-model-unified-memory-level', + group.add_argument( + "--grpo-clamp-eps-lower", type=float, default=0.01, help="Lower GRPO clipping bound." + ) + group.add_argument( + "--grpo-clamp-eps-upper", + type=float, + default=0.01, + help="Upper GRPO clipping bound. In vanilla implementation, equals to the lower one.", + ) + group.add_argument( + "--grpo-kl-beta", type=float, default=0.001, help="KL term weight in the GRPO loss." + ) + group.add_argument( + "--grpo-entropy-term-weight", + type=float, + default=0.0, + help="Entropy term weight in GRPO loss.", + ) + group.add_argument( + "--grpo-filter-groups-with-same-reward", + action="store_true", + help="Filter groups with same reward.", + ) + group.add_argument( + "--langrl-env-config", + type=str, + default=None, + help="Path to YAML config file for RL environment configuration.", + ) + group.add_argument( + "--rl-default-temperature", + type=float, + default=1.0, + help="Default temperature for model inference.", + ) + group.add_argument( + "--rl-default-top-p", type=float, default=0, help="Default top-p for model inference." + ) + group.add_argument( + "--rl-default-top-k", type=int, default=-1, help="Default top-k for model inference." + ) + group.add_argument( + "--rl-offload-optimizer-during-inference", + action="store_true", + help="Offload optimizer state to CPU during inference/rollout to save GPU memory", + ) + group.add_argument( + "--rl-kv-cache-management-mode", + type=str, + default="persist", + choices=["persist", "offload", "recompute"], + help="KV cache management mode during RL training: " + "persist: leave KV cache in GPU memory (default), " + "offload: offload KV cache to CPU during training, " + "recompute: deallocate KV cache and recompute from scratch each cycle", + ) + group.add_argument( + "--rl-persist-cuda-graphs", + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help="Persist CUDA graphs when the inference engine is suspended. " + "If False, CUDA graphs are deleted on suspend and re-captured on resume.", + ) + group.add_argument( + "--rl-partial-rollouts", + action=argparse.BooleanOptionalAction, + default=False, + help="Allow inference to continue generating rollouts while training updates " + "the policy weights. This enables off-policy training where rollouts may " + "be generated with a stale version of the policy. Use " + "--rl-num-parallel-generations or --rl-num-parallel-generation-batches " + "to control the degree of staleness.", + ) + group.add_argument( + "--rl-inference-logprobs-is-correction", + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help="If set, use inference logprobs in importance sampling correction of the loss.", + ) + group.add_argument( + "--rl-importance-sampling-truncation-coef", + type=float, + default=None, + help="If --inference-logprobs-is-correction is on and this coefficient is set, apply truncation for the IS correction at GRPO loss.", + ) + group.add_argument( + "--rl-use-sequence-packing", + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help="Enable sequence packing", + ) + group.add_argument( + "--rl-sequence-packing-max-sequences-per-bin", + type=int, + default=50, + help="Maximum number of sequences that can be packed into a single bin. ", + ) + group.add_argument( + "--rl-sequence-packing-algo", + type=str, + default="fifo", + choices=["fifo", "round-robin"], + help="Algorithm for distributing packed bins across ranks. " + "fifo: first-in-first-out sequential distribution, " + "round-robin: distribute bins cyclically across ranks for better load balancing", + ) + group.add_argument( + "--rl-training-cuda-graphs", + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help="If set, do not toggle CUDA graphs on/off between inference and training phases.", + ) + group.add_argument( + "--rl-inference-tensor-model-parallel-size", + type=int, + default=None, + help="Degree of tensor model parallelism for inference for RL.", + ) + group.add_argument( + "--rl-inference-pipeline-model-parallel-size", + type=int, + default=None, + help="Degree of pipeline model parallelism for inference for RL.", + ) + group.add_argument( + "--rl-inference-expert-model-parallel-size", + type=int, + default=None, + help="Degree of expert model parallelism for inference for RL.", + ) + group.add_argument( + "--rl-inference-expert-tensor-model-parallel-size", + type=int, + default=None, + help="Degree of expert tensor model parallelism for inference for RL. " + "For MoE models, this controls the TP size for expert layers specifically. " + "Defaults to training expert_tensor_parallel_size if not specified.", + ) + group.add_argument( + "--rl-inference-model-unified-memory-level", type=int, default=0, choices=[0, 1], help=( - 'Allocate the separate RL inference model parameters from a unified virtual memory (UVM) ' - 'CUDA mempool. Level 0 disables UVM (default). Level 1 enables UVM allocation so the ' - 'inference model weights can be prefetched to CPU when idle while keeping CUDA-graph-safe ' - 'device pointers.' + "Allocate the separate RL inference model parameters from a unified virtual memory (UVM) " + "CUDA mempool. Level 0 disables UVM (default). Level 1 enables UVM allocation so the " + "inference model weights can be prefetched to CPU when idle while keeping CUDA-graph-safe " + "device pointers." ), ) group.add_argument( - '--rl-offload-inference-model-weights-when-idle', + "--rl-offload-inference-model-weights-when-idle", action=argparse.BooleanOptionalAction, required=False, default=False, help=( - 'When using a separate RL inference model, offload its weights to CPU when not doing rollout ' - 'inference, and restore to GPU right before inference. Works with two backends: ' - '1) UVM (when --rl-inference-model-unified-memory-level=1), or ' - '2) torch_memory_saver (when UVM is not enabled; requires torch_memory_saver to be installed).' + "When using a separate RL inference model, offload its weights to CPU when not doing rollout " + "inference, and restore to GPU right before inference. Works with two backends: " + "1) UVM (when --rl-inference-model-unified-memory-level=1), or " + "2) torch_memory_saver (when UVM is not enabled; requires torch_memory_saver to be installed)." ), ) - group.add_argument('--refit-method', type=str, default='gloo', - choices=['nccl', 'gloo', 'nvshmem'], - help=('Method to refit the model weights between training and inference models during RL. ' - 'nccl: use NCCLCopyService to refit using NCCL; ' - 'gloo: use GlooCopyService over CPU; ' - 'nvshmem: use NVSHMEMCopyService to refit using the NVSHMEM.')) - group.add_argument('--rl-verify-model-weights-swap', action=argparse.BooleanOptionalAction, default=False, - help='If set, verify that the model weights were correctly transferred by comparing forward pass outputs on' - 'the first swap of model weights.') - - group.add_argument('--rl-parallel-generation-tasks', type=int, default=None, - help='Deprecated: use --rl-num-parallel-generations instead.') - group.add_argument('--rl-skip-bos-token', action=argparse.BooleanOptionalAction, type=bool, default=False, - help='Skip BOS token at the beginning of the sequences. Default is False.') - group.add_argument('--rl-inference-parsers', nargs='*', default=[], - help='List of response parsers to enable for RL inference ' - '(e.g. --rl-inference-parsers deepseek-r1-reasoning qwen3-coder-tool).') + group.add_argument( + "--refit-method", + type=str, + default="gloo", + choices=["nccl", "gloo", "nvshmem"], + help=( + "Method to refit the model weights between training and inference models during RL. " + "nccl: use NCCLCopyService to refit using NCCL; " + "gloo: use GlooCopyService over CPU; " + "nvshmem: use NVSHMEMCopyService to refit using the NVSHMEM." + ), + ) + group.add_argument( + "--rl-verify-model-weights-swap", + action=argparse.BooleanOptionalAction, + default=False, + help="If set, verify that the model weights were correctly transferred by comparing forward pass outputs on" + "the first swap of model weights.", + ) + + group.add_argument( + "--rl-parallel-generation-tasks", + type=int, + default=None, + help="Deprecated: use --rl-num-parallel-generations instead.", + ) + group.add_argument( + "--rl-skip-bos-token", + action=argparse.BooleanOptionalAction, + type=bool, + default=False, + help="Skip BOS token at the beginning of the sequences. Default is False.", + ) + group.add_argument( + "--rl-inference-parsers", + nargs="*", + default=[], + help="List of response parsers to enable for RL inference " + "(e.g. --rl-inference-parsers deepseek-r1-reasoning qwen3-coder-tool).", + ) return parser + def _add_training_args(parser): from megatron.training.config import TrainingConfig from megatron.training.config import ProfilingConfig @@ -2515,94 +3253,178 @@ def _add_training_args(parser): train_factory = ArgumentGroupFactory(TrainingConfig) group = train_factory.build_group(parser, "training") - group.add_argument('--batch-size', type=int, default=None, - help='Old batch size parameter, do not use. ' - 'Use --micro-batch-size instead') - group.add_argument('--recompute-activations', action='store_true', - help='recompute activation to allow for training ' - 'with larger models, sequences, and batch sizes.') - group.add_argument('--no-check-for-nan-in-loss-and-grad', action='store_false', - help='Check for NaNs in loss and grad', - dest='check_for_nan_in_loss_and_grad') - group.add_argument('--check-for-large-grads', action='store_true', - help='Check for unexpectedly large grads', - dest='check_for_large_grads') - group.add_argument('--result-rejected-tracker-filename', type=str, default=None, - help='Optional name of file tracking `result_rejected` events.') - group.add_argument('--tp-comm-overlap-cfg', type=str, default=None, - help='Config file when tp_comm_overlap is enabled.') + group.add_argument( + "--batch-size", + type=int, + default=None, + help="Old batch size parameter, do not use. Use --micro-batch-size instead", + ) + group.add_argument( + "--recompute-activations", + action="store_true", + help="recompute activation to allow for training " + "with larger models, sequences, and batch sizes.", + ) + group.add_argument( + "--no-check-for-nan-in-loss-and-grad", + action="store_false", + help="Check for NaNs in loss and grad", + dest="check_for_nan_in_loss_and_grad", + ) + group.add_argument( + "--check-for-large-grads", + action="store_true", + help="Check for unexpectedly large grads", + dest="check_for_large_grads", + ) + group.add_argument( + "--result-rejected-tracker-filename", + type=str, + default=None, + help="Optional name of file tracking `result_rejected` events.", + ) + group.add_argument( + "--tp-comm-overlap-cfg", + type=str, + default=None, + help="Config file when tp_comm_overlap is enabled.", + ) # deprecated - group.add_argument('--checkpoint-activations', action='store_true', - help='Checkpoint activation to allow for training ' - 'with larger models, sequences, and batch sizes.') - group.add_argument('--no-masked-softmax-fusion', - action='store_false', - help='Disable fusion of query_key_value scaling, ' - 'masking, and softmax.', - dest='masked_softmax_fusion') - group.add_argument('--no-bias-gelu-fusion', action='store_false', - help='Disable bias and gelu fusion.', - dest='bias_gelu_fusion') - group.add_argument('--no-bias-swiglu-fusion', action='store_false', - help='Disable bias and swiglu fusion, the fusion is ' - 'available only when using megatron-core.', - dest='bias_swiglu_fusion') - group.add_argument('--no-bias-dropout-fusion', action='store_false', - help='Disable bias and dropout fusion.', - dest='bias_dropout_fusion') - group.add_argument('--no-rope-fusion', action='store_false', - help='Disable rope fusion, the fusion is available ' - 'only when using megatron-core.', - dest='apply_rope_fusion') - group.add_argument('--rope-type', type=str, default=None, - choices=['rope', 'yarn'], - help='Type of rope to use. Note that MLA takes yarn by default, ' - 'and common attention takes rope by default.') - group.add_argument('--use-flash-attn', action='store_true', - help='use FlashAttention implementation of attention. ' - 'https://arxiv.org/abs/2205.14135') - group.add_argument('--optimizer', type=str, default='adam', - choices=['adam', 'sgd', 'muon', 'dist_muon', 'lion'], - help='Optimizer function') - group.add_argument('--optimizer-cpu-offload', action='store_true', - help='Offload optimizer state to CPU') - group.add_argument('--optimizer-cuda-graph', action='store_true', - help='Enable CUDA graph for optimizer step') - group.add_argument('--optimizer-offload-fraction', type=float, default=1.0, - help='Ratio of optimizer state to offload to CPU') - group.add_argument('--use-torch-optimizer-for-cpu-offload', action='store_true', - help="Use torch.optim.Optimizer instead of Megatron's optimizer in optimizer cpu offload mode.") - group.add_argument('--overlap-cpu-optimizer-d2h-h2d', action='store_true', default=False, - help='Overlap CPU optimizer step, gradients D2H and updated parameters H2D.') - group.add_argument('--dump-param-to-param-group-map', type=str, default=None, - help="Path to a file containing parameter-to-parameter-group mapping. " - "Provide a JSON file that specifies which parameters belong to which " - "parameter group for global coordination.") - group.add_argument('--no-pin-cpu-grads', action='store_false', dest='pin_cpu_grads', - help='Disable pinning of CPU memory for gradients.') - group.add_argument('--no-pin-cpu-params', action='store_false', dest='pin_cpu_params', - help='Disable pinning of CPU memory for parameters.') - group.add_argument('--dataloader-type', type=str, default=None, - choices=['single', 'cyclic', 'external'], - help='Single pass vs multiple pass data loader') - group.add_argument('--no-persist-layer-norm', action='store_true', - help='Disable using persistent fused layer norm kernel. ' - 'This kernel supports only a set of hidden sizes. Please ' - 'check persist_ln_hidden_sizes if your hidden ' - 'size is supported.') - group.add_argument('--no-gradient-accumulation-fusion', - action='store_false', - help='Disable fusing gradient accumulation to weight ' - 'gradient computation of linear layers', - dest='gradient_accumulation_fusion') - group.add_argument('--use-mcore-models', action='store_true', - dest='deprecated_use_mcore_models', - help='DEPRECATED. Use the implementation from megatron core.' - 'Now ignored and mcore models are the default, use ' - '--use-legacy-models to not use core models.') - group.add_argument('--use-legacy-models', action='store_true', - help='Use the legacy Megatron models, not Megatron-Core models.') + group.add_argument( + "--checkpoint-activations", + action="store_true", + help="Checkpoint activation to allow for training " + "with larger models, sequences, and batch sizes.", + ) + group.add_argument( + "--no-masked-softmax-fusion", + action="store_false", + help="Disable fusion of query_key_value scaling, masking, and softmax.", + dest="masked_softmax_fusion", + ) + group.add_argument( + "--no-bias-gelu-fusion", + action="store_false", + help="Disable bias and gelu fusion.", + dest="bias_gelu_fusion", + ) + group.add_argument( + "--no-bias-swiglu-fusion", + action="store_false", + help="Disable bias and swiglu fusion, the fusion is " + "available only when using megatron-core.", + dest="bias_swiglu_fusion", + ) + group.add_argument( + "--no-bias-dropout-fusion", + action="store_false", + help="Disable bias and dropout fusion.", + dest="bias_dropout_fusion", + ) + group.add_argument( + "--no-rope-fusion", + action="store_false", + help="Disable rope fusion, the fusion is available only when using megatron-core.", + dest="apply_rope_fusion", + ) + group.add_argument( + "--rope-type", + type=str, + default=None, + choices=["rope", "yarn"], + help="Type of rope to use. Note that MLA takes yarn by default, " + "and common attention takes rope by default.", + ) + group.add_argument( + "--use-flash-attn", + action="store_true", + help="use FlashAttention implementation of attention. https://arxiv.org/abs/2205.14135", + ) + group.add_argument( + "--optimizer", + type=str, + default="adam", + choices=["adam", "sgd", "muon", "dist_muon", "lion"], + help="Optimizer function", + ) + group.add_argument( + "--optimizer-cpu-offload", action="store_true", help="Offload optimizer state to CPU" + ) + group.add_argument( + "--optimizer-cuda-graph", action="store_true", help="Enable CUDA graph for optimizer step" + ) + group.add_argument( + "--optimizer-offload-fraction", + type=float, + default=1.0, + help="Ratio of optimizer state to offload to CPU", + ) + group.add_argument( + "--use-torch-optimizer-for-cpu-offload", + action="store_true", + help="Use torch.optim.Optimizer instead of Megatron's optimizer in optimizer cpu offload mode.", + ) + group.add_argument( + "--overlap-cpu-optimizer-d2h-h2d", + action="store_true", + default=False, + help="Overlap CPU optimizer step, gradients D2H and updated parameters H2D.", + ) + group.add_argument( + "--dump-param-to-param-group-map", + type=str, + default=None, + help="Path to a file containing parameter-to-parameter-group mapping. " + "Provide a JSON file that specifies which parameters belong to which " + "parameter group for global coordination.", + ) + group.add_argument( + "--no-pin-cpu-grads", + action="store_false", + dest="pin_cpu_grads", + help="Disable pinning of CPU memory for gradients.", + ) + group.add_argument( + "--no-pin-cpu-params", + action="store_false", + dest="pin_cpu_params", + help="Disable pinning of CPU memory for parameters.", + ) + group.add_argument( + "--dataloader-type", + type=str, + default=None, + choices=["single", "cyclic", "external"], + help="Single pass vs multiple pass data loader", + ) + group.add_argument( + "--no-persist-layer-norm", + action="store_true", + help="Disable using persistent fused layer norm kernel. " + "This kernel supports only a set of hidden sizes. Please " + "check persist_ln_hidden_sizes if your hidden " + "size is supported.", + ) + group.add_argument( + "--no-gradient-accumulation-fusion", + action="store_false", + help="Disable fusing gradient accumulation to weight gradient computation of linear layers", + dest="gradient_accumulation_fusion", + ) + group.add_argument( + "--use-mcore-models", + action="store_true", + dest="deprecated_use_mcore_models", + help="DEPRECATED. Use the implementation from megatron core." + "Now ignored and mcore models are the default, use " + "--use-legacy-models to not use core models.", + ) + group.add_argument( + "--use-legacy-models", + action="store_true", + help="Use the legacy Megatron models, not Megatron-Core models.", + ) return parser @@ -2622,8 +3444,11 @@ def _add_initialization_args(parser): rng_factory = ArgumentGroupFactory(RNGConfig) group = rng_factory.build_group(parser, "RNG and initialization") - group.add_argument('--init-method-xavier-uniform', action='store_true', - help='Enable Xavier uniform parameter initialization') + group.add_argument( + "--init-method-xavier-uniform", + action="store_true", + help="Enable Xavier uniform parameter initialization", + ) return parser @@ -2634,21 +3459,39 @@ def _add_learning_rate_args(parser): sched_factory = ArgumentGroupFactory(SchedulerConfig, exclude=["no_weight_decay_cond_type"]) group = sched_factory.build_group(parser, title="learning rate and weight decay") - group.add_argument('--lr', type=float, default=None, - help='Initial learning rate. Depending on decay style ' - 'and initial warmup, the learning rate at each ' - 'iteration would be different.') - group.add_argument('--warmup', type=int, default=None, - help='Old lr warmup argument, do not use. Use one of the' - '--lr-warmup-* arguments above') - group.add_argument('--min-lr', type=float, default=0.0, - help='Minimum value for learning rate. The scheduler' - 'clip values below this threshold.') - group.add_argument('--decoupled-lr', type=float, default=None, - help='Separate learning rate for the input and output layer') - group.add_argument('--decoupled-min-lr', type=float, default=None, - help='Minimum value for learning rate for the input and output layer. The scheduler' - 'clip values below this threshold') + group.add_argument( + "--lr", + type=float, + default=None, + help="Initial learning rate. Depending on decay style " + "and initial warmup, the learning rate at each " + "iteration would be different.", + ) + group.add_argument( + "--warmup", + type=int, + default=None, + help="Old lr warmup argument, do not use. Use one of the--lr-warmup-* arguments above", + ) + group.add_argument( + "--min-lr", + type=float, + default=0.0, + help="Minimum value for learning rate. The schedulerclip values below this threshold.", + ) + group.add_argument( + "--decoupled-lr", + type=float, + default=None, + help="Separate learning rate for the input and output layer", + ) + group.add_argument( + "--decoupled-min-lr", + type=float, + default=None, + help="Minimum value for learning rate for the input and output layer. The scheduler" + "clip values below this threshold", + ) return parser @@ -2656,60 +3499,120 @@ def _add_learning_rate_args(parser): def _add_checkpointing_args(parser): from megatron.training.config import CheckpointConfig - ckpt_factory = ArgumentGroupFactory(CheckpointConfig, exclude=["most_recent_k", "save_tokenizer_assets", "save_optim", "save_rng", "load_optim", "load_rng"]) + ckpt_factory = ArgumentGroupFactory( + CheckpointConfig, + exclude=[ + "most_recent_k", + "save_tokenizer_assets", + "save_optim", + "save_rng", + "load_optim", + "load_rng", + ], + ) group = ckpt_factory.build_group(parser, "checkpointing") - group.add_argument('--no-save-optim', action='store_true', default=None, - help='Do not save current optimizer.') - group.add_argument('--no-save-rng', action='store_true', default=None, - help='Do not save current rng state.') - group.add_argument('--no-load-optim', action='store_true', default=None, - help='Do not load optimizer when loading checkpoint.') - group.add_argument('--no-load-rng', action='store_true', default=None, - help='Do not load rng state when loading checkpoint.') - group.add_argument('--use-dist-ckpt', action='store_true', - dest='use_dist_ckpt_deprecated', - help='Deprecated: see --ckpt-format.') - group.add_argument('--dist-ckpt-format', - dest='dist_ckpt_format_deprecated', - help='Deprecated: see --ckpt-format.') - group.add_argument('--dist-ckpt-workers', type=int, default=1, - help='Number of workers for distributed checkpointing. ' - 'Only used for async save. ' - 'If set to 1, the checkpointing is performed in a single process.') - group.add_argument('--ckpt-fully-parallel-save', action='store_true', - dest='ckpt_fully_parallel_save_deprecated', - help='Deprecated: see --no-ckpt-fully-parallel-save.') + group.add_argument( + "--no-save-optim", action="store_true", default=None, help="Do not save current optimizer." + ) + group.add_argument( + "--no-save-rng", action="store_true", default=None, help="Do not save current rng state." + ) + group.add_argument( + "--no-load-optim", + action="store_true", + default=None, + help="Do not load optimizer when loading checkpoint.", + ) + group.add_argument( + "--no-load-rng", + action="store_true", + default=None, + help="Do not load rng state when loading checkpoint.", + ) + group.add_argument( + "--use-dist-ckpt", + action="store_true", + dest="use_dist_ckpt_deprecated", + help="Deprecated: see --ckpt-format.", + ) + group.add_argument( + "--dist-ckpt-format", + dest="dist_ckpt_format_deprecated", + help="Deprecated: see --ckpt-format.", + ) + group.add_argument( + "--dist-ckpt-workers", + type=int, + default=1, + help="Number of workers for distributed checkpointing. " + "Only used for async save. " + "If set to 1, the checkpointing is performed in a single process.", + ) + group.add_argument( + "--ckpt-fully-parallel-save", + action="store_true", + dest="ckpt_fully_parallel_save_deprecated", + help="Deprecated: see --no-ckpt-fully-parallel-save.", + ) return parser def _add_mixed_precision_args(parser): - group = parser.add_argument_group(title='mixed precision') - - group.add_argument('--grad-reduce-in-bf16', action='store_true', - help='Reduce gradients in bfloat16.') - group.add_argument('--loss-scale', type=float, default=None, - help='Static loss scaling, positive power of 2 ' - 'values can improve fp16 convergence. If None, dynamic' - 'loss scaling is used.') - group.add_argument('--initial-loss-scale', type=float, default=2**32, - help='Initial loss-scale for dynamic loss scaling.') - group.add_argument('--min-loss-scale', type=float, default=1.0, - help='Minimum loss scale for dynamic loss scaling.') - group.add_argument('--loss-scale-window', type=float, default=1000, - help='Window over which to raise/lower dynamic scale.') - group.add_argument('--hysteresis', type=int, default=2, - help='hysteresis for dynamic loss scaling') - group.add_argument('--attention-softmax-in-fp32', action='store_true', - help='Run attention masking and softmax in fp32.') - group.add_argument('--accumulate-allreduce-grads-in-fp32', - action='store_true', - help='Gradient accumulation and all-reduce in fp32.') - group.add_argument('--fp16-lm-cross-entropy', action='store_true', - help='Move the cross entropy unreduced loss calculation' - 'for lm head to fp16.') - group.add_argument('--reuse-grad-buf-for-mxfp8-param-ag', action='store_true', - help='If True, reuse the grad buffer for MXFP8 parameter all-gather.') + group = parser.add_argument_group(title="mixed precision") + + group.add_argument( + "--grad-reduce-in-bf16", action="store_true", help="Reduce gradients in bfloat16." + ) + group.add_argument( + "--loss-scale", + type=float, + default=None, + help="Static loss scaling, positive power of 2 " + "values can improve fp16 convergence. If None, dynamic" + "loss scaling is used.", + ) + group.add_argument( + "--initial-loss-scale", + type=float, + default=2**32, + help="Initial loss-scale for dynamic loss scaling.", + ) + group.add_argument( + "--min-loss-scale", + type=float, + default=1.0, + help="Minimum loss scale for dynamic loss scaling.", + ) + group.add_argument( + "--loss-scale-window", + type=float, + default=1000, + help="Window over which to raise/lower dynamic scale.", + ) + group.add_argument( + "--hysteresis", type=int, default=2, help="hysteresis for dynamic loss scaling" + ) + group.add_argument( + "--attention-softmax-in-fp32", + action="store_true", + help="Run attention masking and softmax in fp32.", + ) + group.add_argument( + "--accumulate-allreduce-grads-in-fp32", + action="store_true", + help="Gradient accumulation and all-reduce in fp32.", + ) + group.add_argument( + "--fp16-lm-cross-entropy", + action="store_true", + help="Move the cross entropy unreduced loss calculationfor lm head to fp16.", + ) + group.add_argument( + "--reuse-grad-buf-for-mxfp8-param-ag", + action="store_true", + help="If True, reuse the grad buffer for MXFP8 parameter all-gather.", + ) return parser @@ -2720,120 +3623,251 @@ def _add_distributed_args(parser): dist_init_factory = ArgumentGroupFactory(DistributedInitConfig) group = dist_init_factory.build_group(parser, "distributed init") - group.add_argument('--decoder-first-pipeline-num-layers', - type=int, default=None, - help=('The number of transformer layers on the first pipeline stage of the decoder. ' - 'Default None is even split of transformer layers across all pipeline stages')) - group.add_argument('--decoder-last-pipeline-num-layers', - type=int, default=None, - help=('The number of transformer layers on the last pipeline stage of the decoder. ' - 'Default None is even split of transformer layers across all pipeline stages')) - group.add_argument('--pipeline-model-parallel-layout', - type=str, default=None, - help=('A string that describes a custom pipeline model parallel layout. ' - 'e.g., "E|(t|)*3,m|m||L". E, L, t, m denotes embedding, loss, transformer ' - 'decoder layer, and mtp layer, respectively. Stages are split by "|". ' - 'Replicated stages or layers can be described with multiplication. ' - 'Commas can be used cosmetically. ' - 'Default None is not using this argument to set the layout.')) - group.add_argument('--model-parallel-size', type=int, default=None, - help='Old model parallel argument, do not use. Use ' - '--tensor-model-parallel-size instead.') - group.add_argument('--num-layers-per-virtual-pipeline-stage', type=int, default=None, - help='Number of layers per virtual pipeline stage') - group.add_argument('--num-virtual-stages-per-pipeline-rank', type=int, default=None, - help='Number of virtual pipeline stages per pipeline parallelism rank') - group.add_argument('--no-overlap-p2p-communication', action='store_false', - help='overlap pipeline parallel communication with forward and backward chunks in 1F1B', - dest='overlap_p2p_comm') - group.add_argument('--overlap-grad-reduce', action='store_true', - default=False, help='If set, overlap DDP grad reduce.') - group.add_argument('--ddp-num-buckets', type=int, default=None, - help='Number of buckets for data-parallel communication') - group.add_argument('--ddp-bucket-size', type=int, default=None, - help='Bucket size for data-parallel communication') - group.add_argument('--ddp-pad-buckets-for-high-nccl-busbw', action='store_true', - default=False, help='If set, make sure the bucket size is divisible by a large power ' - 'of 2 (2^16) to ensure NCCL collectives have high bus bandwidth at large DP counts, ' - 'since NCCL message size (which for ring algorithms is bucket_size / dp_size) ' - 'apparently needs to be divisible by a power of 2 for high busbw.') - group.add_argument('--ddp-reduce-scatter-with-fp32-accumulation', action='store_true', - default=False, help='If set, use a reduce-scatter implementation which sends lower-precision ' - 'values over the wire (using an all-to-all to keep total communication overhead in line ' - 'with the standard ring implementation) but performs accumulation locally in FP32.') - group.add_argument('--ddp-param-name-patterns-for-fp32-local-accumulation', nargs='+', default=[], - help='List of param_name patterns (in Python\'s fnmatch format) to match against ' - 'to do local gradient accumulation in FP32. The special pattern \'all\' matches ' - 'every parameter.') - group.add_argument('--ddp-average-in-collective', action='store_true', - default=False, help='If set, average directly in data-parallel communication collective.') - group.add_argument('--overlap-param-gather', action='store_true', - default=False, help='If set, overlap param all-gather in distributed optimizer.') - group.add_argument('--overlap-param-gather-with-optimizer-step', action='store_true', - default=False, help='If set, overlap param all-gather of first bucket with optimizer step.') - group.add_argument('--no-align-param-gather', action='store_false', - help='If not set, all PP stages will launch param all-gathers simultaneously. ' - 'Otherwise, each PP stage will independently launch as needed.', - dest='align_param_gather') - group.add_argument('--no-scatter-gather-tensors-in-pipeline', action='store_false', - help='If not set, use scatter/gather to optimize communication of tensors in pipeline.', - dest='scatter_gather_tensors_in_pipeline') - group.add_argument('--use-distributed-optimizer', action='store_true', - help='Use distributed optimizer.') - group.add_argument('--use-nccl-ub', action='store_true', dest='nccl_ub', - help='Use the userbuffer registration for DP/FSDP communication buffers.' - 'This option will reduce GPU SM usage for the DP/FSDP communication,' - 'which is improving the performance of the overlapped computation.') - group.add_argument('--disable-symmetric-registration', action='store_true', dest='disable_symmetric_registration', - default=False, help='Disable symmetric (window) registration for NCCL userbuffer registration.' - 'This option will force to use conventional (local) userbuffer registration when use-nccl-ub is set.') - group.add_argument('--fsdp-manual-registration', action='store_true', dest='fsdp_manual_registration', - default=False, help='Manually register the FSDP communication buffers to NCCL user buffer.' - 'This option is only effective when use-megatron-fsdp and use-nccl-ub is set.') - group.add_argument('--create-all-gather-group', action='store_true', - help='Create a separate process group for all-gather operations ' - 'to overlap reduce-scatter and all-gather operations.') - group.add_argument('--data-parallel-sharding-strategy', type=str, default='no_shard', - choices=['no_shard', 'optim', 'optim_grads', 'optim_grads_params'], - help='Sharding strategy of data parallelism.') - group.add_argument('--outer-dp-sharding-strategy', type=str, default='no_shard', - choices=['no_shard', 'optim'], - help='Sharding strategy for outer data parallel group in Hybrid Sharded Data Parallel (HSDP) mode. ' - 'Valid values are "no_shard" (DP Replication) and "optim" (Optimizer State Hybrid Sharding). ' - 'The "optim" option is only supported when --data-parallel-sharding-strategy is "optim_grads_params". ' - 'This option is only effective when Hybrid FSDP is enabled (i.e., when dp_outer_dim is not None). ' - 'Default: "no_shard".') - group.add_argument('--no-gradient-reduce-div-fusion', action='store_false', dest='gradient_reduce_div_fusion', - help='If not set, fuse the division in gradient reduce.') - group.add_argument('--fsdp-double-buffer', action='store_true', - help="Enable double buffering for temporary memory needed for Megatron FSDP communications. " - "Double-buffering the communication memory improves memory management efficiency by " - "reusing previously allocated buffers, rather than creating new buffers for each FSDP communication. " - "This is required for user buffer registration and is enabled by default when using NCCL user buffers.") - group.add_argument('--suggested-communication-unit-size', type=int, default=None, - help='Specifies the number of elements to communicate at once during FSDP (Fully Sharded Data Parallel) operations. ' - 'This flag also affects FSDP all-gather prefetch behavior. Setting a larger value increases the communication buffer size, ' - 'while a smaller value disables prefetching and may degrade performance. Adjust this value based on your system\'s memory ' - 'and performance requirements.') - group.add_argument('--keep-fp8-transpose-cache', action='store_true', - help='If set, keep the fp8 transpose cache when using Megatron FSDP.') - group.add_argument('--enable-full-sharding-in-hsdp', action='store_true', - help='If set, enable full sharding in megatron-fsdp Hybrid Sharded Data Parallel (HSDP) mode.') - group.add_argument('--num-distributed-optimizer-instances', type=int, default=1, - help='Number of Distributed Optimizer copies across Data Parallel domain.') - group.add_argument('--torch-fsdp2-no-reshard-after-forward', action='store_false', dest='torch_fsdp2_reshard_after_forward', - help='Whether to reshard weights after forward pass when using PyTorch FSDP2. ' - 'Set to enable FSDP ZeRO-2.') - group.add_argument('--cp-comm-type', nargs='+', type=str, default=["p2p"], - help='Inter-gpu communication type for context parallelism: ' - 'p2p, a2a, allgather or a2a+p2p. If a single string is provided, ' - 'all layers will share the same communication type. Users can also ' - 'specify separated types for each layer like ' - '--cp-comm-type p2p p2p a2a a2a a2a+p2p a2a+p2p') - group.add_argument('--fake-process-group', action='store_true', default=False, - help='If set, initialize with fake distributed process group and all distributed communication operations will be skipped. \ + group.add_argument( + "--decoder-first-pipeline-num-layers", + type=int, + default=None, + help=( + "The number of transformer layers on the first pipeline stage of the decoder. " + "Default None is even split of transformer layers across all pipeline stages" + ), + ) + group.add_argument( + "--decoder-last-pipeline-num-layers", + type=int, + default=None, + help=( + "The number of transformer layers on the last pipeline stage of the decoder. " + "Default None is even split of transformer layers across all pipeline stages" + ), + ) + group.add_argument( + "--pipeline-model-parallel-layout", + type=str, + default=None, + help=( + "A string that describes a custom pipeline model parallel layout. " + 'e.g., "E|(t|)*3,m|m||L". E, L, t, m denotes embedding, loss, transformer ' + 'decoder layer, and mtp layer, respectively. Stages are split by "|". ' + "Replicated stages or layers can be described with multiplication. " + "Commas can be used cosmetically. " + "Default None is not using this argument to set the layout." + ), + ) + group.add_argument( + "--model-parallel-size", + type=int, + default=None, + help="Old model parallel argument, do not use. Use --tensor-model-parallel-size instead.", + ) + group.add_argument( + "--num-layers-per-virtual-pipeline-stage", + type=int, + default=None, + help="Number of layers per virtual pipeline stage", + ) + group.add_argument( + "--num-virtual-stages-per-pipeline-rank", + type=int, + default=None, + help="Number of virtual pipeline stages per pipeline parallelism rank", + ) + group.add_argument( + "--no-overlap-p2p-communication", + action="store_false", + help="overlap pipeline parallel communication with forward and backward chunks in 1F1B", + dest="overlap_p2p_comm", + ) + group.add_argument( + "--overlap-grad-reduce", + action="store_true", + default=False, + help="If set, overlap DDP grad reduce.", + ) + group.add_argument( + "--ddp-num-buckets", + type=int, + default=None, + help="Number of buckets for data-parallel communication", + ) + group.add_argument( + "--ddp-bucket-size", + type=int, + default=None, + help="Bucket size for data-parallel communication", + ) + group.add_argument( + "--ddp-pad-buckets-for-high-nccl-busbw", + action="store_true", + default=False, + help="If set, make sure the bucket size is divisible by a large power " + "of 2 (2^16) to ensure NCCL collectives have high bus bandwidth at large DP counts, " + "since NCCL message size (which for ring algorithms is bucket_size / dp_size) " + "apparently needs to be divisible by a power of 2 for high busbw.", + ) + group.add_argument( + "--ddp-reduce-scatter-with-fp32-accumulation", + action="store_true", + default=False, + help="If set, use a reduce-scatter implementation which sends lower-precision " + "values over the wire (using an all-to-all to keep total communication overhead in line " + "with the standard ring implementation) but performs accumulation locally in FP32.", + ) + group.add_argument( + "--ddp-param-name-patterns-for-fp32-local-accumulation", + nargs="+", + default=[], + help="List of param_name patterns (in Python's fnmatch format) to match against " + "to do local gradient accumulation in FP32. The special pattern 'all' matches " + "every parameter.", + ) + group.add_argument( + "--ddp-average-in-collective", + action="store_true", + default=False, + help="If set, average directly in data-parallel communication collective.", + ) + group.add_argument( + "--overlap-param-gather", + action="store_true", + default=False, + help="If set, overlap param all-gather in distributed optimizer.", + ) + group.add_argument( + "--overlap-param-gather-with-optimizer-step", + action="store_true", + default=False, + help="If set, overlap param all-gather of first bucket with optimizer step.", + ) + group.add_argument( + "--no-align-param-gather", + action="store_false", + help="If not set, all PP stages will launch param all-gathers simultaneously. " + "Otherwise, each PP stage will independently launch as needed.", + dest="align_param_gather", + ) + group.add_argument( + "--no-scatter-gather-tensors-in-pipeline", + action="store_false", + help="If not set, use scatter/gather to optimize communication of tensors in pipeline.", + dest="scatter_gather_tensors_in_pipeline", + ) + group.add_argument( + "--use-distributed-optimizer", action="store_true", help="Use distributed optimizer." + ) + group.add_argument( + "--use-nccl-ub", + action="store_true", + dest="nccl_ub", + help="Use the userbuffer registration for DP/FSDP communication buffers." + "This option will reduce GPU SM usage for the DP/FSDP communication," + "which is improving the performance of the overlapped computation.", + ) + group.add_argument( + "--disable-symmetric-registration", + action="store_true", + dest="disable_symmetric_registration", + default=False, + help="Disable symmetric (window) registration for NCCL userbuffer registration." + "This option will force to use conventional (local) userbuffer registration when use-nccl-ub is set.", + ) + group.add_argument( + "--fsdp-manual-registration", + action="store_true", + dest="fsdp_manual_registration", + default=False, + help="Manually register the FSDP communication buffers to NCCL user buffer." + "This option is only effective when use-megatron-fsdp and use-nccl-ub is set.", + ) + group.add_argument( + "--create-all-gather-group", + action="store_true", + help="Create a separate process group for all-gather operations " + "to overlap reduce-scatter and all-gather operations.", + ) + group.add_argument( + "--data-parallel-sharding-strategy", + type=str, + default="no_shard", + choices=["no_shard", "optim", "optim_grads", "optim_grads_params"], + help="Sharding strategy of data parallelism.", + ) + group.add_argument( + "--outer-dp-sharding-strategy", + type=str, + default="no_shard", + choices=["no_shard", "optim"], + help="Sharding strategy for outer data parallel group in Hybrid Sharded Data Parallel (HSDP) mode. " + 'Valid values are "no_shard" (DP Replication) and "optim" (Optimizer State Hybrid Sharding). ' + 'The "optim" option is only supported when --data-parallel-sharding-strategy is "optim_grads_params". ' + "This option is only effective when Hybrid FSDP is enabled (i.e., when dp_outer_dim is not None). " + 'Default: "no_shard".', + ) + group.add_argument( + "--no-gradient-reduce-div-fusion", + action="store_false", + dest="gradient_reduce_div_fusion", + help="If not set, fuse the division in gradient reduce.", + ) + group.add_argument( + "--fsdp-double-buffer", + action="store_true", + help="Enable double buffering for temporary memory needed for Megatron FSDP communications. " + "Double-buffering the communication memory improves memory management efficiency by " + "reusing previously allocated buffers, rather than creating new buffers for each FSDP communication. " + "This is required for user buffer registration and is enabled by default when using NCCL user buffers.", + ) + group.add_argument( + "--suggested-communication-unit-size", + type=int, + default=None, + help="Specifies the number of elements to communicate at once during FSDP (Fully Sharded Data Parallel) operations. " + "This flag also affects FSDP all-gather prefetch behavior. Setting a larger value increases the communication buffer size, " + "while a smaller value disables prefetching and may degrade performance. Adjust this value based on your system's memory " + "and performance requirements.", + ) + group.add_argument( + "--keep-fp8-transpose-cache", + action="store_true", + help="If set, keep the fp8 transpose cache when using Megatron FSDP.", + ) + group.add_argument( + "--enable-full-sharding-in-hsdp", + action="store_true", + help="If set, enable full sharding in megatron-fsdp Hybrid Sharded Data Parallel (HSDP) mode.", + ) + group.add_argument( + "--num-distributed-optimizer-instances", + type=int, + default=1, + help="Number of Distributed Optimizer copies across Data Parallel domain.", + ) + group.add_argument( + "--torch-fsdp2-no-reshard-after-forward", + action="store_false", + dest="torch_fsdp2_reshard_after_forward", + help="Whether to reshard weights after forward pass when using PyTorch FSDP2. " + "Set to enable FSDP ZeRO-2.", + ) + group.add_argument( + "--cp-comm-type", + nargs="+", + type=str, + default=["p2p"], + help="Inter-gpu communication type for context parallelism: " + "p2p, a2a, allgather or a2a+p2p. If a single string is provided, " + "all layers will share the same communication type. Users can also " + "specify separated types for each layer like " + "--cp-comm-type p2p p2p a2a a2a a2a+p2p a2a+p2p", + ) + group.add_argument( + "--fake-process-group", + action="store_true", + default=False, + help="If set, initialize with fake distributed process group and all distributed communication operations will be skipped. \ This is quite useful for profiling memory usage of distributed training with just one GPU. \ - Setting WORLD_SIZE and RANK to the specific values for target distribtued scale.') + Setting WORLD_SIZE and RANK to the specific values for target distribtued scale.", + ) return parser @@ -2847,261 +3881,471 @@ def _add_validation_args(parser): def _add_tokenizer_args(parser): - group = parser.add_argument_group(title='tokenizer') - group.add_argument('--vocab-size', type=int, default=None, - help='Size of vocab before EOD or padding.') - group.add_argument('--padded-vocab-size', type=int, default=None, - help='Vocabulary size of the model (padded to be divisible by ' - 'tensor model parallel size). If not provided, it will be ' - 'automatically calculated from vocab-size.') - group.add_argument('--vocab-file', type=str, default=None, - help='Path to the vocab file.') - group.add_argument('--merge-file', type=str, default=None, - help='Path to the BPE merge file.') - group.add_argument('--vocab-extra-ids', type=int, default=0, - help='Number of additional vocabulary tokens. ' - 'They are used for span masking in the T5 model') - group.add_argument('--tokenizer-type', type=str, - default=None, - choices=['BertWordPieceLowerCase', - 'BertWordPieceCase', - 'GPT2BPETokenizer', - 'SentencePieceTokenizer', - 'GPTSentencePieceTokenizer', - 'HuggingFaceTokenizer', - 'Llama2Tokenizer', - 'TikTokenizer', - 'MultimodalTokenizer', - 'NullTokenizer', - 'NullMultimodalTokenizer', - 'SFTTokenizer', - 'AquilaTokenizerFS', - 'HFTokenizerFS', - 'HFTokenizersTokenizerFS', - 'Llama3TokenizerFS', - 'QwenTokenizerFS', - 'Qwen2TokenizerFS', - 'Qwen2VLTokenizer'], - help='What type of tokenizer to use.') - group.add_argument('--tokenizer-model', type=str, default=None, - help='Sentencepiece tokenizer model.') - group.add_argument('--tokenizer-metadata', type=str, default=None, - help='Path to tokenizer metadata in json format.') - group.add_argument('--tokenizer-special-tokens', type=str, nargs='+', default=None, - help='List of special tokens. For TikTokenizer needs to have ' - '["", "", "", "", "", "", ""]') - group.add_argument('--tiktoken-pattern', type=str, default=None, - help='Which tiktoken pattern to use. Options: [v1, v2]') - group.add_argument('--tiktoken-num-special-tokens', type=int, default=1000, - help='Number of special tokens in tiktoken tokenizer') - group.add_argument('--tiktoken-special-tokens', type=str, nargs='+', default=None, - help='List of tiktoken special tokens, needs to have ' - '["", "", "", "", "", "", ""]') - group.add_argument('--tokenizer-sentencepiece-legacy', action='store_true', default=False, - help='SentencePiece tokenizer wrapper legacy behavior. Allows special tokens usage.') - group.add_argument('--tokenizer-hf-use-fast', action='store_true', default=True, - help='Whether to use fast HuggingFace tokenizer.') - group.add_argument('--tokenizer-hf-include-special-tokens', action='store_true', default=True, - help='Converting text to ids will include special for HuggingFace tokenizer.') - group.add_argument('--tokenizer-hf-no-use-fast', action='store_true', default=False, - help='Whether to use fast HuggingFace tokenizer.') - group.add_argument('--tokenizer-hf-no-include-special-tokens', action='store_true', default=False, - help='Converting text to ids will not include special for HuggingFace tokenizer.') - group.add_argument("--trust-remote-code", action="store_true", default=False, - help='Whether or not to allow PreTrainedTokenizer to execute remote code') + group = parser.add_argument_group(title="tokenizer") + group.add_argument( + "--vocab-size", type=int, default=None, help="Size of vocab before EOD or padding." + ) + group.add_argument( + "--padded-vocab-size", + type=int, + default=None, + help="Vocabulary size of the model (padded to be divisible by " + "tensor model parallel size). If not provided, it will be " + "automatically calculated from vocab-size.", + ) + group.add_argument("--vocab-file", type=str, default=None, help="Path to the vocab file.") + group.add_argument("--merge-file", type=str, default=None, help="Path to the BPE merge file.") + group.add_argument( + "--vocab-extra-ids", + type=int, + default=0, + help="Number of additional vocabulary tokens. " + "They are used for span masking in the T5 model", + ) + group.add_argument( + "--tokenizer-type", + type=str, + default=None, + choices=[ + "BertWordPieceLowerCase", + "BertWordPieceCase", + "GPT2BPETokenizer", + "SentencePieceTokenizer", + "GPTSentencePieceTokenizer", + "HuggingFaceTokenizer", + "Llama2Tokenizer", + "TikTokenizer", + "MultimodalTokenizer", + "NullTokenizer", + "NullMultimodalTokenizer", + "SFTTokenizer", + "AquilaTokenizerFS", + "HFTokenizerFS", + "HFTokenizersTokenizerFS", + "Llama3TokenizerFS", + "QwenTokenizerFS", + "Qwen2TokenizerFS", + "Qwen2VLTokenizer", + ], + help="What type of tokenizer to use.", + ) + group.add_argument( + "--tokenizer-model", type=str, default=None, help="Sentencepiece tokenizer model." + ) + group.add_argument( + "--tokenizer-metadata", + type=str, + default=None, + help="Path to tokenizer metadata in json format.", + ) + group.add_argument( + "--tokenizer-special-tokens", + type=str, + nargs="+", + default=None, + help="List of special tokens. For TikTokenizer needs to have " + '["", "", "", "", "", "", ""]', + ) + group.add_argument( + "--tiktoken-pattern", + type=str, + default=None, + help="Which tiktoken pattern to use. Options: [v1, v2]", + ) + group.add_argument( + "--tiktoken-num-special-tokens", + type=int, + default=1000, + help="Number of special tokens in tiktoken tokenizer", + ) + group.add_argument( + "--tiktoken-special-tokens", + type=str, + nargs="+", + default=None, + help="List of tiktoken special tokens, needs to have " + '["", "", "", "", "", "", ""]', + ) + group.add_argument( + "--tokenizer-sentencepiece-legacy", + action="store_true", + default=False, + help="SentencePiece tokenizer wrapper legacy behavior. Allows special tokens usage.", + ) + group.add_argument( + "--tokenizer-hf-use-fast", + action="store_true", + default=True, + help="Whether to use fast HuggingFace tokenizer.", + ) + group.add_argument( + "--tokenizer-hf-include-special-tokens", + action="store_true", + default=True, + help="Converting text to ids will include special for HuggingFace tokenizer.", + ) + group.add_argument( + "--tokenizer-hf-no-use-fast", + action="store_true", + default=False, + help="Whether to use fast HuggingFace tokenizer.", + ) + group.add_argument( + "--tokenizer-hf-no-include-special-tokens", + action="store_true", + default=False, + help="Converting text to ids will not include special for HuggingFace tokenizer.", + ) + group.add_argument( + "--trust-remote-code", + action="store_true", + default=False, + help="Whether or not to allow PreTrainedTokenizer to execute remote code", + ) return parser def _add_data_args(parser): - group = parser.add_argument_group(title='data and dataloader') - - group.add_argument('--data-path', nargs='*', default=None, - help='The weight and prefix list for a set of train, validation, and test' - 'datasets which split according to --split. The accepted formats are: ' - '(1) a single prefix, ' - '(2) a list of weight prefix pairs e.g. weight1 prefix1 weight2 prefix2, ' - '(3) a list of prefixes e.g. prefix1 prefix2. ' - 'For (3), weights are inferred from the lengths of the contributing datasets. ' - 'This argument is exclusive to the other independent --*-data-path arguments.') - group.add_argument('--phase-transition-iterations', type=str, default=None, - help='Comma-separated list of iterations where phase ' - 'transitions occur. Requires fixed global batch size across phases. ' - 'Does not support batch size ramp-up.') - group.add_argument('--split', type=str, default=None, - help='Comma-separated list of proportions for training,' - ' validation, and test split. For example the split ' - '`90,5,5` will use 90%% of data for training, 5%% for ' - 'validation and 5%% for test.') - group.add_argument('--train-data-path', nargs='*', default=None, - help='The weight and prefix list for an independent train dataset. ' - 'Follows the same pattern rules as --data-path.') - group.add_argument('--valid-data-path', nargs='*', default=None, - help='The weight and prefix list for an independent validation dataset. ' - 'Follows the same pattern rules as --data-path.') - group.add_argument('--test-data-path', nargs='*', default=None, - help='The weight and prefix list for an independent test dataset. ' - 'Follows the same pattern rules as --data-path.') - group.add_argument('--data-args-path', type=str, default=None, - help='Path to data-args. Instead of feeding `--data-path` ' - 'with weighted dataset, we pass in a file path from which ' - 'we read that argument. This is useful when the list of data is ' - 'too big.') - group.add_argument('--per-split-data-args-path', type=str, default=None, - help='Path to per-split-data-args. Instead of feeding ' - '`--(train|valid|test)-data-path` with weighted dataset, ' - 'we pass in a file path from which we read those arguments. ' - 'This is useful when the list of data is too big. Format is a ' - 'json file with `train`, `valid, `test` keys') - group.add_argument('--per-dataset-sequences-path', default=None, - help='Path to a json file with the sequences per dataset. Check the tools/build_sequences_per_dataset.py script to build this file.') - group.add_argument('--dataloader-fast-cache-load', action='store_true', - help='Option to use the fast cache loading path when building the datasets. Requires all the dataset caches to be built and stored in --data-cache-path.') - group.add_argument('--dataloader-defer-npy-index-mmap', action='store_true', - help='Defer the mmap of the dataset indexes (.npy files) until the first access. Requires all the dataset caches to be built and stored in --data-cache-path.') - group.add_argument('--data-cache-path', default=None, - help='Path to a directory to hold cached index files.') - group.add_argument('--no-mmap-bin-files', action='store_false', - help='Disable mmap-ing of .bin files.', - dest='mmap_bin_files') - group.add_argument('--mock-data', action='store_true', - help='Skip data loading and validation and opt for artificial ' - 'generation of mock data when an implementation is available.') - group.add_argument('--seq-length', type=int, default=None, - help='Maximum sequence length to process.') - group.add_argument('--encoder-seq-length', type=int, default=None, - help='Maximum encoder sequence length to process.' - 'This should be exclusive of --seq-length') - group.add_argument('--decoder-seq-length', type=int, default=None, - help="Maximum decoder sequence length to process.") - group.add_argument('--sample-rate', type=float, default=1.0, - help='sample rate for training data. Supposed to be 0 ' - ' < sample_rate < 1') - group.add_argument('--mask-prob', type=float, default=0.15, - help='Probability of replacing a token with mask.') - group.add_argument('--short-seq-prob', type=float, default=0.1, - help='Probability of producing a short sequence.') - group.add_argument('--num-workers', type=int, default=2, - help="Dataloader number of workers.") - group.add_argument('--reset-position-ids', action='store_true', - help='Reset posistion ids after end-of-document token.') - group.add_argument('--reset-attention-mask', action='store_true', - help='Reset self attention mask after ' - 'end-of-document token.') - group.add_argument('--eod-mask-loss', action='store_true', - help='Mask loss for the end of document tokens.') - group.add_argument('--no-create-attention-mask-in-dataloader', action='store_false', - help='If set, do not create attention_masks in dataloader.', - dest='create_attention_mask_in_dataloader') - group.add_argument('--num-dataset-builder-threads', type=int, default=1, - help='Number of parallel threads per rank for dataset builder') - group.add_argument('--object-storage-cache-path', type=str, default=None, - help='Path to cache index files when using s3 or msc dataloader') - group.add_argument('--mid-level-dataset-surplus', type=float, default=0.005, - help='The sample surplus to build for the mid-level datasets(s)') - group.add_argument('--allow-ambiguous-pad-tokens', action='store_true', - help='Whether to prevent pad tokens already present in the dataset ' - 'from being masked out when the pad token incorrectly shares the same id ' - 'with other special tokens in the tokenizer. Note that this argument has ' - 'no effect when the tokenizer correctly provides a unique id for the pad. ' - 'Masking out such ambiguous pad tokens results in training instability. ' - 'Such a scenario is best resolved by fixing the tokenizer; leaving this ' - 'option as False provides a workaround. ' - 'When left to the default of False, any token ids that collide with the ' - 'pad token id - as provided by the tokenizer - will not be masked out of ' - 'the loss calculation: it cannot be determined whether they are truly pad. ' - 'If instead this argument is set, the training flow will treat all tokens ' - 'that share the same id as the pad token as true pad tokens, potentially ' - 'causing severe training instability.') - group.add_argument('--fim-data', action='store_true', help='Whether to use the FIM dataset.') - group.add_argument('--fim-rate', type=float, default=0.5, - help='Probability to convert a training sample into a FIM format.') - group.add_argument('--fim-spm-rate', type=float, default=0.5, - help='Probability that the a FIM sample uses the SPM format over the PSM format.') - group.add_argument('--fim-split-sample', type=str, default=None, - help='String around which to split the sample for FIM.') - group.add_argument('--fim-fragment-rate', type=float, default=None, - help='Rate of FIM on each fragment when --fim-split-sample is not None.') - group.add_argument('--fim-no-prefix', type=str, default=None, - help='Do not apply FIM to fragments that start with this prefix') - group.add_argument('--fim-prefix-token', type=str, default='', - help='FIM prefix token') - group.add_argument('--fim-middle-token', type=str, default='', - help='FIM middle token') - group.add_argument('--fim-suffix-token', type=str, default='', - help='FIM suffix token') - group.add_argument('--fim-pad-token', type=str, default='', - help='FIM PAD token') - group.add_argument('--fim-eod-token', type=str, default='<|endoftext|>', - help='FIM EOD token') + group = parser.add_argument_group(title="data and dataloader") + + group.add_argument( + "--data-path", + nargs="*", + default=None, + help="The weight and prefix list for a set of train, validation, and test" + "datasets which split according to --split. The accepted formats are: " + "(1) a single prefix, " + "(2) a list of weight prefix pairs e.g. weight1 prefix1 weight2 prefix2, " + "(3) a list of prefixes e.g. prefix1 prefix2. " + "For (3), weights are inferred from the lengths of the contributing datasets. " + "This argument is exclusive to the other independent --*-data-path arguments.", + ) + group.add_argument( + "--phase-transition-iterations", + type=str, + default=None, + help="Comma-separated list of iterations where phase " + "transitions occur. Requires fixed global batch size across phases. " + "Does not support batch size ramp-up.", + ) + group.add_argument( + "--split", + type=str, + default=None, + help="Comma-separated list of proportions for training," + " validation, and test split. For example the split " + "`90,5,5` will use 90%% of data for training, 5%% for " + "validation and 5%% for test.", + ) + group.add_argument( + "--train-data-path", + nargs="*", + default=None, + help="The weight and prefix list for an independent train dataset. " + "Follows the same pattern rules as --data-path.", + ) + group.add_argument( + "--valid-data-path", + nargs="*", + default=None, + help="The weight and prefix list for an independent validation dataset. " + "Follows the same pattern rules as --data-path.", + ) + group.add_argument( + "--test-data-path", + nargs="*", + default=None, + help="The weight and prefix list for an independent test dataset. " + "Follows the same pattern rules as --data-path.", + ) + group.add_argument( + "--data-args-path", + type=str, + default=None, + help="Path to data-args. Instead of feeding `--data-path` " + "with weighted dataset, we pass in a file path from which " + "we read that argument. This is useful when the list of data is " + "too big.", + ) + group.add_argument( + "--per-split-data-args-path", + type=str, + default=None, + help="Path to per-split-data-args. Instead of feeding " + "`--(train|valid|test)-data-path` with weighted dataset, " + "we pass in a file path from which we read those arguments. " + "This is useful when the list of data is too big. Format is a " + "json file with `train`, `valid, `test` keys", + ) + group.add_argument( + "--per-dataset-sequences-path", + default=None, + help="Path to a json file with the sequences per dataset. Check the tools/build_sequences_per_dataset.py script to build this file.", + ) + group.add_argument( + "--dataloader-fast-cache-load", + action="store_true", + help="Option to use the fast cache loading path when building the datasets. Requires all the dataset caches to be built and stored in --data-cache-path.", + ) + group.add_argument( + "--dataloader-defer-npy-index-mmap", + action="store_true", + help="Defer the mmap of the dataset indexes (.npy files) until the first access. Requires all the dataset caches to be built and stored in --data-cache-path.", + ) + group.add_argument( + "--data-cache-path", default=None, help="Path to a directory to hold cached index files." + ) + group.add_argument( + "--no-mmap-bin-files", + action="store_false", + help="Disable mmap-ing of .bin files.", + dest="mmap_bin_files", + ) + group.add_argument( + "--mock-data", + action="store_true", + help="Skip data loading and validation and opt for artificial " + "generation of mock data when an implementation is available.", + ) + group.add_argument( + "--seq-length", type=int, default=None, help="Maximum sequence length to process." + ) + group.add_argument( + "--encoder-seq-length", + type=int, + default=None, + help="Maximum encoder sequence length to process.This should be exclusive of --seq-length", + ) + group.add_argument( + "--decoder-seq-length", + type=int, + default=None, + help="Maximum decoder sequence length to process.", + ) + group.add_argument( + "--sample-rate", + type=float, + default=1.0, + help="sample rate for training data. Supposed to be 0 < sample_rate < 1", + ) + group.add_argument( + "--mask-prob", type=float, default=0.15, help="Probability of replacing a token with mask." + ) + group.add_argument( + "--short-seq-prob", + type=float, + default=0.1, + help="Probability of producing a short sequence.", + ) + group.add_argument("--num-workers", type=int, default=2, help="Dataloader number of workers.") + group.add_argument( + "--reset-position-ids", + action="store_true", + help="Reset posistion ids after end-of-document token.", + ) + group.add_argument( + "--reset-attention-mask", + action="store_true", + help="Reset self attention mask after end-of-document token.", + ) + group.add_argument( + "--eod-mask-loss", action="store_true", help="Mask loss for the end of document tokens." + ) + group.add_argument( + "--no-create-attention-mask-in-dataloader", + action="store_false", + help="If set, do not create attention_masks in dataloader.", + dest="create_attention_mask_in_dataloader", + ) + group.add_argument( + "--num-dataset-builder-threads", + type=int, + default=1, + help="Number of parallel threads per rank for dataset builder", + ) + group.add_argument( + "--object-storage-cache-path", + type=str, + default=None, + help="Path to cache index files when using s3 or msc dataloader", + ) + group.add_argument( + "--mid-level-dataset-surplus", + type=float, + default=0.005, + help="The sample surplus to build for the mid-level datasets(s)", + ) + group.add_argument( + "--allow-ambiguous-pad-tokens", + action="store_true", + help="Whether to prevent pad tokens already present in the dataset " + "from being masked out when the pad token incorrectly shares the same id " + "with other special tokens in the tokenizer. Note that this argument has " + "no effect when the tokenizer correctly provides a unique id for the pad. " + "Masking out such ambiguous pad tokens results in training instability. " + "Such a scenario is best resolved by fixing the tokenizer; leaving this " + "option as False provides a workaround. " + "When left to the default of False, any token ids that collide with the " + "pad token id - as provided by the tokenizer - will not be masked out of " + "the loss calculation: it cannot be determined whether they are truly pad. " + "If instead this argument is set, the training flow will treat all tokens " + "that share the same id as the pad token as true pad tokens, potentially " + "causing severe training instability.", + ) + group.add_argument("--fim-data", action="store_true", help="Whether to use the FIM dataset.") + group.add_argument( + "--fim-rate", + type=float, + default=0.5, + help="Probability to convert a training sample into a FIM format.", + ) + group.add_argument( + "--fim-spm-rate", + type=float, + default=0.5, + help="Probability that the a FIM sample uses the SPM format over the PSM format.", + ) + group.add_argument( + "--fim-split-sample", + type=str, + default=None, + help="String around which to split the sample for FIM.", + ) + group.add_argument( + "--fim-fragment-rate", + type=float, + default=None, + help="Rate of FIM on each fragment when --fim-split-sample is not None.", + ) + group.add_argument( + "--fim-no-prefix", + type=str, + default=None, + help="Do not apply FIM to fragments that start with this prefix", + ) + group.add_argument( + "--fim-prefix-token", type=str, default="", help="FIM prefix token" + ) + group.add_argument( + "--fim-middle-token", type=str, default="", help="FIM middle token" + ) + group.add_argument( + "--fim-suffix-token", type=str, default="", help="FIM suffix token" + ) + group.add_argument("--fim-pad-token", type=str, default="", help="FIM PAD token") + group.add_argument("--fim-eod-token", type=str, default="<|endoftext|>", help="FIM EOD token") return parser def _add_autoresume_args(parser): - group = parser.add_argument_group(title='autoresume') + group = parser.add_argument_group(title="autoresume") - group.add_argument('--adlr-autoresume', action='store_true', - help='Enable autoresume on adlr cluster.') - group.add_argument('--adlr-autoresume-interval', type=int, default=1000, - help='Intervals over which check for autoresume' - 'termination signal') + group.add_argument( + "--adlr-autoresume", action="store_true", help="Enable autoresume on adlr cluster." + ) + group.add_argument( + "--adlr-autoresume-interval", + type=int, + default=1000, + help="Intervals over which check for autoresumetermination signal", + ) return parser def _add_biencoder_args(parser): - group = parser.add_argument_group(title='biencoder') + group = parser.add_argument_group(title="biencoder") # network size - group.add_argument('--ict-head-size', type=int, default=None, - help='Size of block embeddings to be used in ICT and ' - 'REALM (paper default: 128)') - group.add_argument('--biencoder-projection-dim', type=int, default=0, - help='Size of projection head used in biencoder (paper' - ' default: 128)') - group.add_argument('--biencoder-shared-query-context-model', action='store_true', - help='Whether to share the parameters of the query ' - 'and context models or not') + group.add_argument( + "--ict-head-size", + type=int, + default=None, + help="Size of block embeddings to be used in ICT and REALM (paper default: 128)", + ) + group.add_argument( + "--biencoder-projection-dim", + type=int, + default=0, + help="Size of projection head used in biencoder (paper default: 128)", + ) + group.add_argument( + "--biencoder-shared-query-context-model", + action="store_true", + help="Whether to share the parameters of the query and context models or not", + ) # checkpointing - group.add_argument('--ict-load', type=str, default=None, - help='Directory containing an ICTBertModel checkpoint') - group.add_argument('--bert-load', type=str, default=None, - help='Directory containing an BertModel checkpoint ' - '(needed to start ICT and REALM)') + group.add_argument( + "--ict-load", type=str, default=None, help="Directory containing an ICTBertModel checkpoint" + ) + group.add_argument( + "--bert-load", + type=str, + default=None, + help="Directory containing an BertModel checkpoint (needed to start ICT and REALM)", + ) # data - group.add_argument('--titles-data-path', type=str, default=None, - help='Path to titles dataset used for ICT') - group.add_argument('--query-in-block-prob', type=float, default=0.1, - help='Probability of keeping query in block for ' - 'ICT dataset') - group.add_argument('--use-one-sent-docs', action='store_true', - help='Whether to use one sentence documents in ICT') - group.add_argument('--evidence-data-path', type=str, default=None, - help='Path to Wikipedia Evidence frm DPR paper') + group.add_argument( + "--titles-data-path", type=str, default=None, help="Path to titles dataset used for ICT" + ) + group.add_argument( + "--query-in-block-prob", + type=float, + default=0.1, + help="Probability of keeping query in block for ICT dataset", + ) + group.add_argument( + "--use-one-sent-docs", + action="store_true", + help="Whether to use one sentence documents in ICT", + ) + group.add_argument( + "--evidence-data-path", + type=str, + default=None, + help="Path to Wikipedia Evidence frm DPR paper", + ) # training - group.add_argument('--retriever-report-topk-accuracies', nargs='+', type=int, - default=[], help="Which top-k accuracies to report " - "(e.g. '1 5 20')") - group.add_argument('--retriever-score-scaling', action='store_true', - help='Whether to scale retriever scores by inverse ' - 'square root of hidden size') + group.add_argument( + "--retriever-report-topk-accuracies", + nargs="+", + type=int, + default=[], + help="Which top-k accuracies to report (e.g. '1 5 20')", + ) + group.add_argument( + "--retriever-score-scaling", + action="store_true", + help="Whether to scale retriever scores by inverse square root of hidden size", + ) # faiss index - group.add_argument('--block-data-path', type=str, default=None, - help='Where to save/load BlockData to/from') - group.add_argument('--embedding-path', type=str, default=None, - help='Where to save/load Open-Retrieval Embedding' - ' data to/from') + group.add_argument( + "--block-data-path", type=str, default=None, help="Where to save/load BlockData to/from" + ) + group.add_argument( + "--embedding-path", + type=str, + default=None, + help="Where to save/load Open-Retrieval Embedding data to/from", + ) # indexer - group.add_argument('--indexer-batch-size', type=int, default=128, - help='How large of batches to use when doing indexing ' - 'jobs') - group.add_argument('--indexer-log-interval', type=int, default=1000, - help='After how many batches should the indexer ' - 'report progress') + group.add_argument( + "--indexer-batch-size", + type=int, + default=128, + help="How large of batches to use when doing indexing jobs", + ) + group.add_argument( + "--indexer-log-interval", + type=int, + default=1000, + help="After how many batches should the indexer report progress", + ) return parser @@ -3109,157 +4353,257 @@ def _add_vision_args(parser): group = parser.add_argument_group(title="vision") # general vision arguements - group.add_argument('--num-classes', type=int, default=1000, - help='num of classes in vision classificaiton task') - group.add_argument('--img-h', type=int, default=224, - help='Image height for vision classification task') - group.add_argument('--img-w', type=int, default=224, - help='Image height for vision classification task') - group.add_argument('--num-channels', type=int, default=3, - help='Number of channels in input image data') - group.add_argument('--patch-dim', type=int, default=16, - help='patch dimension') - group.add_argument('--classes-fraction', type=float, default=1.0, - help='training with fraction of classes.') - group.add_argument('--data-per-class-fraction', type=float, default=1.0, - help='training with fraction of data per class.') - group.add_argument('--no-data-sharding', action='store_false', - help='Disable data sharding.', - dest='data_sharding') - group.add_argument('--head-lr-mult', type=float, default=1.0, - help='learning rate multiplier for head during finetuning') + group.add_argument( + "--num-classes", type=int, default=1000, help="num of classes in vision classificaiton task" + ) + group.add_argument( + "--img-h", type=int, default=224, help="Image height for vision classification task" + ) + group.add_argument( + "--img-w", type=int, default=224, help="Image height for vision classification task" + ) + group.add_argument( + "--num-channels", type=int, default=3, help="Number of channels in input image data" + ) + group.add_argument("--patch-dim", type=int, default=16, help="patch dimension") + group.add_argument( + "--classes-fraction", type=float, default=1.0, help="training with fraction of classes." + ) + group.add_argument( + "--data-per-class-fraction", + type=float, + default=1.0, + help="training with fraction of data per class.", + ) + group.add_argument( + "--no-data-sharding", + action="store_false", + help="Disable data sharding.", + dest="data_sharding", + ) + group.add_argument( + "--head-lr-mult", + type=float, + default=1.0, + help="learning rate multiplier for head during finetuning", + ) # pretraining type and backbone selection` - group.add_argument('--vision-pretraining', action='store_true', - help='flag to indicate vision pretraining') - group.add_argument('--vision-pretraining-type', type=str, default='classify', - choices=['classify', 'inpaint', 'dino'], - help='pretraining objectives') - group.add_argument('--vision-backbone-type', type=str, default='vit', - choices=['vit', 'mit', 'swin'], - help='backbone types types') - group.add_argument('--swin-backbone-type', type=str, default='tiny', - choices=['tiny', 'base', 'h3'], - help='pretraining objectives') + group.add_argument( + "--vision-pretraining", action="store_true", help="flag to indicate vision pretraining" + ) + group.add_argument( + "--vision-pretraining-type", + type=str, + default="classify", + choices=["classify", "inpaint", "dino"], + help="pretraining objectives", + ) + group.add_argument( + "--vision-backbone-type", + type=str, + default="vit", + choices=["vit", "mit", "swin"], + help="backbone types types", + ) + group.add_argument( + "--swin-backbone-type", + type=str, + default="tiny", + choices=["tiny", "base", "h3"], + help="pretraining objectives", + ) # inpainting arguments - group.add_argument('--mask-type', type=str, default='random', - choices=['random', 'row'], - help='mask types') - group.add_argument('--mask-factor', type=float, default=1.0, - help='mask size scaling parameter') + group.add_argument( + "--mask-type", type=str, default="random", choices=["random", "row"], help="mask types" + ) + group.add_argument("--mask-factor", type=float, default=1.0, help="mask size scaling parameter") # dino arguments - group.add_argument('--iter-per-epoch', type=int, default=1250, - help='iterations per epoch') - group.add_argument('--dino-local-img-size', type=int, default=96, - help='Image size for vision classification task') - group.add_argument('--dino-local-crops-number', type=int, default=10, - help='Number of local crops') - group.add_argument('--dino-head-hidden-size', type=int, default=2048, - help='Hidden dimension size in dino head') - group.add_argument('--dino-bottleneck-size', type=int, default=256, - help='Bottle neck dimension in dino head ') - group.add_argument('--dino-freeze-last-layer', type=float, default=1, - help='Freezing last layer weights') - group.add_argument('--dino-norm-last-layer', action='store_true', - help='Disable Norm in last layer.') - group.add_argument('--dino-warmup-teacher-temp', type=float, default=0.04, - help='warump teacher temperature') - group.add_argument('--dino-teacher-temp', type=float, default=0.07, - help='teacher temperature') - group.add_argument('--dino-warmup-teacher-temp-epochs', type=int, default=30, - help='warmup teacher temperaure epochs') + group.add_argument("--iter-per-epoch", type=int, default=1250, help="iterations per epoch") + group.add_argument( + "--dino-local-img-size", + type=int, + default=96, + help="Image size for vision classification task", + ) + group.add_argument( + "--dino-local-crops-number", type=int, default=10, help="Number of local crops" + ) + group.add_argument( + "--dino-head-hidden-size", type=int, default=2048, help="Hidden dimension size in dino head" + ) + group.add_argument( + "--dino-bottleneck-size", type=int, default=256, help="Bottle neck dimension in dino head " + ) + group.add_argument( + "--dino-freeze-last-layer", type=float, default=1, help="Freezing last layer weights" + ) + group.add_argument( + "--dino-norm-last-layer", action="store_true", help="Disable Norm in last layer." + ) + group.add_argument( + "--dino-warmup-teacher-temp", type=float, default=0.04, help="warump teacher temperature" + ) + group.add_argument("--dino-teacher-temp", type=float, default=0.07, help="teacher temperature") + group.add_argument( + "--dino-warmup-teacher-temp-epochs", + type=int, + default=30, + help="warmup teacher temperaure epochs", + ) return parser + def _add_moe_args(parser): group = parser.add_argument_group(title="moe") # General arguments - group.add_argument('--num-experts', type=int, default=None, - help='Number of Experts in MoE (None means no MoE)') - group.add_argument('--moe-layer-freq', type=moe_freq_type, default=1, - help='Frequency between MoE layers and Dense layers. Accepts either: ' - '- An integer N: Represents a 1:N ratio, meaning one expert layer for every N-1 dense layers ' - '- A string containing a Python list expression that defines a custom pattern, e.g.: ' - '"([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' - 'where 1 indicates an expert layer and 0 indicates a dense layer. ' - 'Examples: "([0]+[1]*23)": 1 dense layer followed by 23 expert layers, ' - '"([1]*3+[0]*2)*2": Three expert layers followed by two dense layers, repeated twice.') - group.add_argument('--moe-use-upcycling', action='store_true', - help='Load a checkpoint of a dense model, convert it into an MoE model, and save the converted model to the path specified by --save. ' - 'Upcycling is implemented on the top of distributed checkpointing, so it supports parallel modes different from the dense model.') + group.add_argument( + "--num-experts", type=int, default=None, help="Number of Experts in MoE (None means no MoE)" + ) + group.add_argument( + "--moe-layer-freq", + type=moe_freq_type, + default=1, + help="Frequency between MoE layers and Dense layers. Accepts either: " + "- An integer N: Represents a 1:N ratio, meaning one expert layer for every N-1 dense layers " + "- A string containing a Python list expression that defines a custom pattern, e.g.: " + '"([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' + "where 1 indicates an expert layer and 0 indicates a dense layer. " + 'Examples: "([0]+[1]*23)": 1 dense layer followed by 23 expert layers, ' + '"([1]*3+[0]*2)*2": Three expert layers followed by two dense layers, repeated twice.', + ) + group.add_argument( + "--moe-use-upcycling", + action="store_true", + help="Load a checkpoint of a dense model, convert it into an MoE model, and save the converted model to the path specified by --save. " + "Upcycling is implemented on the top of distributed checkpointing, so it supports parallel modes different from the dense model.", + ) # Router arguments - group.add_argument('--moe-router-load-balancing-type', nargs='+', type=str, - choices=['aux_loss', 'seq_aux_loss', 'global_aux_loss', 'sinkhorn', 'none'], - default='aux_loss', - help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer; "seq_aux_loss" corresponds to the load balancing loss used in DeepSeekV2, which computes the loss for each individual sample; "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss".') - group.add_argument('--moe-aux-loss-coeff', type=float, nargs='+', default=0.0, - help='Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.') + group.add_argument( + "--moe-router-load-balancing-type", + nargs="+", + type=str, + choices=["aux_loss", "seq_aux_loss", "global_aux_loss", "sinkhorn", "none"], + default="aux_loss", + help='Determines the load balancing strategy for the router. "aux_loss" corresponds to the load balancing loss used in GShard and SwitchTransformer; "seq_aux_loss" corresponds to the load balancing loss used in DeepSeekV2, which computes the loss for each individual sample; "sinkhorn" corresponds to the balancing algorithm used in S-BASE, and "none" implies no load balancing. The default is "aux_loss".', + ) + group.add_argument( + "--moe-aux-loss-coeff", + type=float, + nargs="+", + default=0.0, + help="Scaling coefficient for the aux loss: a starting value of 1e-2 is recommended.", + ) # Token dispatcher arguments # MoE communication overlap arguments - group.add_argument('--moe-upcycling-granularity', type=int, default=1, - help='This param sepecifics how many times smaller is the expert hidden size compared with the original dense FFN hidden size. ' - 'For using granular upcycling strategy, please set this param as a positive integer. If this param is set to 1, it means using the default upcycling strategy.') + group.add_argument( + "--moe-upcycling-granularity", + type=int, + default=1, + help="This param sepecifics how many times smaller is the expert hidden size compared with the original dense FFN hidden size. " + "For using granular upcycling strategy, please set this param as a positive integer. If this param is set to 1, it means using the default upcycling strategy.", + ) return parser + def _add_mla_args(parser): group = parser.add_argument_group(title="mla") - group.add_argument('--q-lora-rank', type=int, default=None, - help="Rank of Query tensor's low rank representation.") - group.add_argument('--kv-lora-rank', type=int, default=32, - help="Rank of Key and Value tensors' low rank representation.") - group.add_argument('--qk-head-dim', type=int, default=128, - help="Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim") - group.add_argument('--qk-pos-emb-head-dim', type=int, default=64, - help="Dimension of the position embedding in the QK projection.") - group.add_argument('--v-head-dim', type=int, default=128, - help="Dimension of the head in the V projection.") - group.add_argument('--rotary-scaling-factor', type=float, default=1.0, - help="Rotary scaling factor for the rotary embeddings.") - group.add_argument('--mscale', type=float, default=1.0, - help="Mscale for YaRN RoPE in multi-latent attention.") - group.add_argument('--mscale-all-dim', type=float, default=0.0, - help="Mscale all dimensions for YaRN RoPE in multi-latent attention.") - group.add_argument('--cache-mla-latents', action='store_true', default=False, - help="If set caches the mla down projected latents with mla flash decode.") - group.add_argument( - '--o-groups', + group.add_argument( + "--q-lora-rank", + type=int, + default=None, + help="Rank of Query tensor's low rank representation.", + ) + group.add_argument( + "--kv-lora-rank", + type=int, + default=32, + help="Rank of Key and Value tensors' low rank representation.", + ) + group.add_argument( + "--qk-head-dim", + type=int, + default=128, + help="Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim", + ) + group.add_argument( + "--qk-pos-emb-head-dim", + type=int, + default=64, + help="Dimension of the position embedding in the QK projection.", + ) + group.add_argument( + "--v-head-dim", type=int, default=128, help="Dimension of the head in the V projection." + ) + group.add_argument( + "--rotary-scaling-factor", + type=float, + default=1.0, + help="Rotary scaling factor for the rotary embeddings.", + ) + group.add_argument( + "--mscale", type=float, default=1.0, help="Mscale for YaRN RoPE in multi-latent attention." + ) + group.add_argument( + "--mscale-all-dim", + type=float, + default=0.0, + help="Mscale all dimensions for YaRN RoPE in multi-latent attention.", + ) + group.add_argument( + "--cache-mla-latents", + action="store_true", + default=False, + help="If set caches the mla down projected latents with mla flash decode.", + ) + group.add_argument( + "--o-groups", type=int, default=8, - help="Number of groups for grouped output (wo_a). 0 = single linear." + help="Number of groups for grouped output (wo_a). 0 = single linear.", ) group.add_argument( - '--o-lora-rank', + "--o-lora-rank", type=int, default=1024, - help="Low-rank dimension per group for grouped output (wo_a). Used when o-groups > 0." + help="Low-rank dimension per group for grouped output (wo_a). Used when o-groups > 0.", ) group.add_argument( - '--mla-down-proj-fusion', - action='store_true', + "--mla-down-proj-fusion", + action="store_true", default=False, help="Enable fused q/kv down-projection and fused input layernorm when backend supports. " - "Otherwise fall back to the unfused MLA.", + "Otherwise fall back to the unfused MLA.", ) return parser + def _add_experimental_attention_variant_args(parser): group = parser.add_argument_group(title="experimental_attention_variant") # Linear attention - group.add_argument('--linear-attention-freq', type=la_freq_type, default=None, - help='Frequency between LA (linear attention) layers and' - ' SDPA (scaled dot-product attention) layers. Accepts either: ' - '- An integer N: Represents a (N-1):N ratio, meaning (N-1) LA layers for every 1 SDPA layer ' - '- A string containing a Python list expression that defines a custom pattern, e.g.: ' - '"([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' - 'where 1 indicates an LA layer and 0 indicates a SDPA layer. ' - 'Examples: "([0]+[1]*23)": 1 SDPA layer followed by 23 LA layers, ' - '"([1]*3+[0]*2)*2": Three LA layers followed by two SDPA layers, repeated twice.') + group.add_argument( + "--linear-attention-freq", + type=la_freq_type, + default=None, + help="Frequency between LA (linear attention) layers and" + " SDPA (scaled dot-product attention) layers. Accepts either: " + "- An integer N: Represents a (N-1):N ratio, meaning (N-1) LA layers for every 1 SDPA layer " + "- A string containing a Python list expression that defines a custom pattern, e.g.: " + '"([1]*3+[0]*1)*3" evaluates to [1,1,1,0,1,1,1,0,1,1,1,0] ' + "where 1 indicates an LA layer and 0 indicates a SDPA layer. " + 'Examples: "([0]+[1]*23)": 1 SDPA layer followed by 23 LA layers, ' + '"([1]*3+[0]*2)*2": Three LA layers followed by two SDPA layers, repeated twice.', + ) # DeepSeek V4 hybrid attention - group.add_argument('--csa-compress-ratios', type=dsv4_hybrid_attn_type, default=None, - help="""Frequency between HCA (Heavy compress attention) layers and CSA (Compress sparse attention) layers. + group.add_argument( + "--csa-compress-ratios", + type=dsv4_hybrid_attn_type, + default=None, + help="""Frequency between HCA (Heavy compress attention) layers and CSA (Compress sparse attention) layers. Accepts either: - A list of integers: Represents the compress ratios for each layer, where 0 indicates SDPA, 128 indicates HCA, 4 indicate CSA with the given compress ratio. Now only support such values. - A string containing a Python list expression that defines a custom pattern, e.g.: @@ -3270,16 +4614,18 @@ def _add_experimental_attention_variant_args(parser): Examples: "([128]+[4]*23)": 1 HCA layer followed by 23 CSA layers "([128]*3+[4]*2)*2": Three HCA layers followed by two CSA layers, repeated twice. - """) + """, + ) group.add_argument( - '--no-dsa-kernel-fusion', - action='store_false', - help='Disable fused DSA sparse-attention kernels (FlashMLA + cuDNN DSA) ' - 'and fall back to unfused PyTorch implementations.', - dest='apply_dsa_kernel_fusion', + "--no-dsa-kernel-fusion", + action="store_false", + help="Disable fused DSA sparse-attention kernels (FlashMLA + cuDNN DSA) " + "and fall back to unfused PyTorch implementations.", + dest="apply_dsa_kernel_fusion", ) return parser + def _add_heterogeneous_args(parser): """ Heterogeneous models refer to transformer architectures where individual layers can differ @@ -3319,86 +4665,147 @@ def _add_heterogeneous_args(parser): ] """ group = parser.add_argument_group(title="heterogeneous architecture") - group.add_argument('--heterogeneous-layers-config-path', type=str, default=None, - help='Path to json file containing heterogeneous model configuration. ' - 'Use the format of the HuggingFace config files in llama nemotron ' - 'models, e.g. https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json.') - group.add_argument('--heterogeneous-layers-config-encoded-json', type=str, default=None, - help='This is encoded json string of the heterogeneous model configuration. Used to keep the content ' - 'of the heterogeneous model specification in args when the model is loaded from a checkpoint. ' - 'Use the format of the HuggingFace config files in llama nemotron ' - 'models, e.g. https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json.') + group.add_argument( + "--heterogeneous-layers-config-path", + type=str, + default=None, + help="Path to json file containing heterogeneous model configuration. " + "Use the format of the HuggingFace config files in llama nemotron " + "models, e.g. https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json.", + ) + group.add_argument( + "--heterogeneous-layers-config-encoded-json", + type=str, + default=None, + help="This is encoded json string of the heterogeneous model configuration. Used to keep the content " + "of the heterogeneous model specification in args when the model is loaded from a checkpoint. " + "Use the format of the HuggingFace config files in llama nemotron " + "models, e.g. https://huggingface.co/nvidia/Llama-3_3-Nemotron-Super-49B-v1/resolve/main/config.json.", + ) return parser + def _add_experimental_args(parser): - group = parser.add_argument_group(title='experimental') - - group.add_argument('--enable-experimental', action='store_true', - help='Enable experimental features.') - group.add_argument('--spec', type=str, default=None, nargs='*', - help='Specify the pair ' - 'that returns a spec to customize a model, transformer ' - 'block, or transformer layer, depending on the use case.' - 'To use local spec specify local as the argument.' - 'For more details, see the model class, ' - '`transformer_block.py`, or `transformer_layer.py`') - group.add_argument('--hybrid-layer-pattern', type=str, default=None, - help='Specify a hybrid layer pattern using M (mamba), G (gdn), ' - '* (attention), - (mlp), E (moe). Use | to define pipeline stage ' - 'boundaries for flexible virtual pipeline parallel (fVPP). Use / to ' - 'separate MTP patterns. Example: "M-M-|M-M*-|M-M-|M-M*-" or "M-M-|M-M*-/MM/MM". ' - 'When this flag is used, it is the sole indicator that a hybrid model ' - 'is being run.') - group.add_argument('--hybrid-override-pattern', type=str, default=None, - help='Deprecated. Use --hybrid-layer-pattern instead. ' - 'If specified, its value will be forwarded to --hybrid-layer-pattern.') - group.add_argument('--yaml-cfg', type=str, default=None, - help = 'Config file to add additional arguments') + group = parser.add_argument_group(title="experimental") + + group.add_argument( + "--enable-experimental", action="store_true", help="Enable experimental features." + ) + group.add_argument( + "--spec", + type=str, + default=None, + nargs="*", + help="Specify the pair " + "that returns a spec to customize a model, transformer " + "block, or transformer layer, depending on the use case." + "To use local spec specify local as the argument." + "For more details, see the model class, " + "`transformer_block.py`, or `transformer_layer.py`", + ) + group.add_argument( + "--hybrid-layer-pattern", + type=str, + default=None, + help="Specify a hybrid layer pattern using M (mamba), G (gdn), " + "* (attention), - (mlp), E (moe). Use | to define pipeline stage " + "boundaries for flexible virtual pipeline parallel (fVPP). Use / to " + 'separate MTP patterns. Example: "M-M-|M-M*-|M-M-|M-M*-" or "M-M-|M-M*-/MM/MM". ' + "When this flag is used, it is the sole indicator that a hybrid model " + "is being run.", + ) + group.add_argument( + "--hybrid-override-pattern", + type=str, + default=None, + help="Deprecated. Use --hybrid-layer-pattern instead. " + "If specified, its value will be forwarded to --hybrid-layer-pattern.", + ) + group.add_argument( + "--yaml-cfg", type=str, default=None, help="Config file to add additional arguments" + ) # Args of precision-aware optimizer. - group.add_argument('--use-precision-aware-optimizer', action='store_true', - help='Use the precision-aware optimizer in TransformerEngine, which allows ' - 'setting the main params and optimizer states to lower precision, such as ' - 'fp16, bf16 and fp8.') - group.add_argument('--main-grads-dtype', default='fp32', choices=['fp32', 'bf16'], - help='Dtype of main grads when enabling precision-aware-optimizer.') - group.add_argument('--main-params-dtype', default='fp32', choices=['fp32', 'fp16'], - help='Dtype of main params when enabling precision-aware-optimizer.') - group.add_argument('--exp-avg-dtype', default='fp32', choices=['fp32', 'fp16', 'bf16', 'fp8'], - help='Dtype of exp_avg (1st moment in adam optimizer) when enabling ' - 'precision-aware-optimizer. This dtype is used for storing the ' - 'optimizer state in memory during training but does not affect ' - 'the precision in the kernel computation.') - group.add_argument('--exp-avg-sq-dtype', default='fp32', choices=['fp32', 'fp16', 'bf16', 'fp8'], - help='Dtype of exp_avg_sq (2nd moment in adam optimizer) when enabling ' - 'precision-aware-optimizer. This dtype is used for storing the ' - 'optimizer state in memory during training but does not affect ' - 'the precision in the kernel computation.') + group.add_argument( + "--use-precision-aware-optimizer", + action="store_true", + help="Use the precision-aware optimizer in TransformerEngine, which allows " + "setting the main params and optimizer states to lower precision, such as " + "fp16, bf16 and fp8.", + ) + group.add_argument( + "--main-grads-dtype", + default="fp32", + choices=["fp32", "bf16"], + help="Dtype of main grads when enabling precision-aware-optimizer.", + ) + group.add_argument( + "--main-params-dtype", + default="fp32", + choices=["fp32", "fp16"], + help="Dtype of main params when enabling precision-aware-optimizer.", + ) + group.add_argument( + "--exp-avg-dtype", + default="fp32", + choices=["fp32", "fp16", "bf16", "fp8"], + help="Dtype of exp_avg (1st moment in adam optimizer) when enabling " + "precision-aware-optimizer. This dtype is used for storing the " + "optimizer state in memory during training but does not affect " + "the precision in the kernel computation.", + ) + group.add_argument( + "--exp-avg-sq-dtype", + default="fp32", + choices=["fp32", "fp16", "bf16", "fp8"], + help="Dtype of exp_avg_sq (2nd moment in adam optimizer) when enabling " + "precision-aware-optimizer. This dtype is used for storing the " + "optimizer state in memory during training but does not affect " + "the precision in the kernel computation.", + ) # Megatron-FSDP Arguments - group.add_argument('--megatron-fsdp-main-params-dtype', default='fp32', choices=['fp32', 'bf16', 'fp16', 'auto'], - help="Data type for the main weight buffer utilized for distributed optimization " - "and quantization with Megatron-FSDP. If 'auto', then the native model parameter " - "data-type will be used for the main weight data-type.") - group.add_argument('--megatron-fsdp-main-grads-dtype', default='auto', choices=['fp32', 'bf16', 'fp16', 'auto'], - help="Data type for the main gradient buffer utilized for distributed optimization " - "with Megatron-FSDP. If 'auto', then the native model gradient data-type will " - "be used for the main gradient / accumulation data-type.") - group.add_argument("--megatron-fsdp-grad-comm-dtype", default='auto', choices=['fp32', 'fp16', 'bf16', 'auto'], - help="When using Megatron-FSDP, this controls the data-type used when communicating " - "model gradients during FSDP. If 'auto', then the main gradient data-type will " - "be used for the gradient communication / reduction data-type. When using NCCL " - "v2.27+, reduction is always computed in FP32 if using NCCL Symmetric kernels.") - + group.add_argument( + "--megatron-fsdp-main-params-dtype", + default="fp32", + choices=["fp32", "bf16", "fp16", "auto"], + help="Data type for the main weight buffer utilized for distributed optimization " + "and quantization with Megatron-FSDP. If 'auto', then the native model parameter " + "data-type will be used for the main weight data-type.", + ) + group.add_argument( + "--megatron-fsdp-main-grads-dtype", + default="auto", + choices=["fp32", "bf16", "fp16", "auto"], + help="Data type for the main gradient buffer utilized for distributed optimization " + "with Megatron-FSDP. If 'auto', then the native model gradient data-type will " + "be used for the main gradient / accumulation data-type.", + ) + group.add_argument( + "--megatron-fsdp-grad-comm-dtype", + default="auto", + choices=["fp32", "fp16", "bf16", "auto"], + help="When using Megatron-FSDP, this controls the data-type used when communicating " + "model gradients during FSDP. If 'auto', then the main gradient data-type will " + "be used for the gradient communication / reduction data-type. When using NCCL " + "v2.27+, reduction is always computed in FP32 if using NCCL Symmetric kernels.", + ) + return parser def _add_msc_args(parser): group = parser.add_argument_group(title="msc") - group.add_argument('--disable-msc', default=True, action='store_false', dest='enable_msc', - help='Disable the usage of Multi-Storage Client (MSC) in Megatron Core.') + group.add_argument( + "--disable-msc", + default=True, + action="store_false", + dest="enable_msc", + help="Disable the usage of Multi-Storage Client (MSC) in Megatron Core.", + ) return parser + def _add_kitchen_quantization_arguments(parser: argparse.ArgumentParser): """Add quant-specific arguments to the main parser @@ -3414,14 +4821,14 @@ def _add_kitchen_quantization_arguments(parser: argparse.ArgumentParser): group = parser.add_argument_group(title="kitchen") recipe_or_config_group = group.add_mutually_exclusive_group(required=False) recipe_or_config_group.add_argument( - '--kitchen-config-file', + "--kitchen-config-file", type=str, default=None, help="Use the config .yaml file at the specified location to " "configure kitchen quantization.", ) recipe_or_config_group.add_argument( - '--kitchen-recipe-number', + "--kitchen-recipe-number", type=int, default=None, help="Use a default kitchen recipe for all linear layers as defined by QAT_PARAMS index. " @@ -3429,9 +4836,14 @@ def _add_kitchen_quantization_arguments(parser: argparse.ArgumentParser): ) return parser + def _add_sft_args(parser): - group = parser.add_argument_group(title='sft') - group.add_argument('--sft', action="store_true", help='Megatron SFT training') - group.add_argument('--sft-tokenizer-prompt-format', type=str, default="nemotron-h-aligned", - help='SFT prompt format.') + group = parser.add_argument_group(title="sft") + group.add_argument("--sft", action="store_true", help="Megatron SFT training") + group.add_argument( + "--sft-tokenizer-prompt-format", + type=str, + default="nemotron-h-aligned", + help="SFT prompt format.", + ) return parser diff --git a/flagscale/train/megatron/training/arguments_fs.py b/flagscale/train/megatron/training/arguments_fs.py index 22a86fa171..8bca6eb7ee 100644 --- a/flagscale/train/megatron/training/arguments_fs.py +++ b/flagscale/train/megatron/training/arguments_fs.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import ast import itertools import types @@ -37,18 +51,15 @@ def _initialize_distributed(self): device_count = cur_platform.device_count() if torch.distributed.is_initialized(): - if args.rank == 0: print( - "torch distributed is already initialized, " - "skipping initialization ...", + "torch distributed is already initialized, skipping initialization ...", flush=True, ) args.rank = torch.distributed.get_rank() args.world_size = torch.distributed.get_world_size() else: - if args.rank == 0: print("> initializing torch distributed ...", flush=True) # Manually set the device ids. @@ -86,9 +97,9 @@ def pre_validate_args(self): self._build_rank_mapper() if self.args.hetero_process_meshes is not None: - assert ( - len(self.args.hetero_process_meshes) % 5 == 0 - ), f"length of hetero_process_meshes {self.args.hetero_process_meshes} should be divisible by 5, the format should be tp0, cp0, dp0, pp0, tp1, cp1, dp1, pp1, ..." + assert len(self.args.hetero_process_meshes) % 5 == 0, ( + f"length of hetero_process_meshes {self.args.hetero_process_meshes} should be divisible by 5, the format should be tp0, cp0, dp0, pp0, tp1, cp1, dp1, pp1, ..." + ) hetero_process_meshes_tp = self.args.hetero_process_meshes[0::5] hetero_process_meshes_cp = self.args.hetero_process_meshes[1::5] hetero_process_meshes_ep = self.args.hetero_process_meshes[2::5] @@ -98,53 +109,54 @@ def pre_validate_args(self): # Check if tensor parallel sizes are inconsistent across meshes # NOTE: If TP degrees differ, sequence parallelism must be enabled if len(set(hetero_process_meshes_tp)) > 1: - assert ( - self.args.sequence_parallel - ), f"Sequence parallelism must be enabled (`sequence_parallel=True`) when tensor parallelism degrees differ across heterogeneous meshes. Found TP degrees: {hetero_process_meshes_tp}" + assert self.args.sequence_parallel, ( + f"Sequence parallelism must be enabled (`sequence_parallel=True`) when tensor parallelism degrees differ across heterogeneous meshes. Found TP degrees: {hetero_process_meshes_tp}" + ) # Expert tensor parallel size if self.expert_tensor_parallel_size_per_process_mesh is not None: assert len(self.expert_tensor_parallel_size_per_process_mesh) == len( hetero_process_meshes_tp - ), f"length of expert_tensor_parallel_size_per_process_mesh {len(self.expert_tensor_parallel_size_per_process_mesh)} should be equal to length of hetero_process_meshes_tp {len(hetero_process_meshes_tp)}" + ), ( + f"length of expert_tensor_parallel_size_per_process_mesh {len(self.expert_tensor_parallel_size_per_process_mesh)} should be equal to length of hetero_process_meshes_tp {len(hetero_process_meshes_tp)}" + ) # Data parallel size # NOTE: Use the first data parallel size as the global data parallel size to loader data self.args.data_parallel_size = hetero_process_meshes_dp[0] assert all( - self.args.data_parallel_size * self.args.micro_batch_size % hetero_dp - == 0 + self.args.data_parallel_size * self.args.micro_batch_size % hetero_dp == 0 for hetero_dp in hetero_process_meshes_dp - ), f"data_parallel_size * micro_batch_size {self.args.data_parallel_size * self.args.micro_batch_size} should be divisible by all hetero_process_meshes_dp {hetero_process_meshes_dp}!" + ), ( + f"data_parallel_size * micro_batch_size {self.args.data_parallel_size * self.args.micro_batch_size} should be divisible by all hetero_process_meshes_dp {hetero_process_meshes_dp}!" + ) # NOTE: Only support cp and ep size to be the same assert all( - hetero_cp == hetero_process_meshes_cp[0] - for hetero_cp in hetero_process_meshes_cp + hetero_cp == hetero_process_meshes_cp[0] for hetero_cp in hetero_process_meshes_cp ), f"all hetero_process_meshes_cp {hetero_process_meshes_cp} should be the same!" # Note: Ep size should all be 1 or all be not 1 assert all(1 == hetero_ep for hetero_ep in hetero_process_meshes_ep) or any( 1 != hetero_ep for hetero_ep in hetero_process_meshes_ep - ), f"all hetero_process_meshes_ep {hetero_process_meshes_ep} should be the 1 or none of hetero_process_meshes_ep is not 1!" + ), ( + f"all hetero_process_meshes_ep {hetero_process_meshes_ep} should be the 1 or none of hetero_process_meshes_ep is not 1!" + ) # Pipeline model parallel size - assert self.args.pipeline_model_parallel_size == sum( - hetero_process_meshes_pp - ), f"origin pipeline_model_parallel_size {self.args.pipeline_model_parallel_size} should match sum of hetero_process_meshes_pp {hetero_process_meshes_pp}!" - assert ( - self.args.standalone_embedding_stage == False - ), "standalone not supported with process_meshes set!" + assert self.args.pipeline_model_parallel_size == sum(hetero_process_meshes_pp), ( + f"origin pipeline_model_parallel_size {self.args.pipeline_model_parallel_size} should match sum of hetero_process_meshes_pp {hetero_process_meshes_pp}!" + ) + assert self.args.standalone_embedding_stage == False, ( + "standalone not supported with process_meshes set!" + ) self.args.transformer_pipeline_model_parallel_size = ( self.args.pipeline_model_parallel_size ) # if untie_embeddings_and_output_weights is False, the first and last stage should have the same tp degree - if ( - self.args.untie_embeddings_and_output_weights == False - or self.args.mtp_num_layers - ): - assert ( - hetero_process_meshes_tp[0] == hetero_process_meshes_tp[-1] - ), f"if untie_embeddings_and_output_weights is False or mtp_num_layers is not 0, the first and last stage should have the same tp degree!" + if self.args.untie_embeddings_and_output_weights == False or self.args.mtp_num_layers: + assert hetero_process_meshes_tp[0] == hetero_process_meshes_tp[-1], ( + f"if untie_embeddings_and_output_weights is False or mtp_num_layers is not 0, the first and last stage should have the same tp degree!" + ) if ( hetero_process_meshes_dp[0] != hetero_process_meshes_dp[-1] @@ -152,8 +164,7 @@ def pre_validate_args(self): ): assert ( hetero_process_meshes_dp[0] % hetero_process_meshes_dp[-1] == 0 - or hetero_process_meshes_dp[-1] % hetero_process_meshes_dp[0] - == 0 + or hetero_process_meshes_dp[-1] % hetero_process_meshes_dp[0] == 0 ), ( f"if untie_embeddings_and_output_weights is False and hetero_process_meshes_dp[0] and hetero_process_meshes_dp[-1] are different, " "the hetero_process_meshes_dp[0] should be divisible by hetero_process_meshes_dp[-1] or hetero_process_meshes_dp[-1] should be divisible by hetero_process_meshes_dp[0] currently!" @@ -165,26 +176,27 @@ def pre_validate_args(self): # Virtual parallel size. if self.args.enable_hetero: - assert ( - self.args.num_layers_per_virtual_pipeline_stage == None - ), "virtual pipeline not support now!" + assert self.args.num_layers_per_virtual_pipeline_stage == None, ( + "virtual pipeline not support now!" + ) # Model layer splits if self.args.hetero_pipeline_layer_split is None: num_layers_per_pipeline_stage = ( - self.args.num_layers - // self.args.transformer_pipeline_model_parallel_size + self.args.num_layers // self.args.transformer_pipeline_model_parallel_size ) self.args.hetero_pipeline_layer_split = [ num_layers_per_pipeline_stage ] * self.args.pipeline_model_parallel_size else: - assert ( - sum(self.args.hetero_pipeline_layer_split) == self.args.num_layers - ), f"sum of hetero_pipeline_layer_split {self.args.hetero_pipeline_layer_split} should be equal to num_layers {self.args.num_layers}" + assert sum(self.args.hetero_pipeline_layer_split) == self.args.num_layers, ( + f"sum of hetero_pipeline_layer_split {self.args.hetero_pipeline_layer_split} should be equal to num_layers {self.args.num_layers}" + ) assert self.args.pipeline_model_parallel_size == len( self.args.hetero_pipeline_layer_split - ), f"pipeline_model_parallel_size {self.args.pipeline_model_parallel_size} should be equal to the length of hetero_pipeline_layer_split {self.args.hetero_pipeline_layer_split}" + ), ( + f"pipeline_model_parallel_size {self.args.pipeline_model_parallel_size} should be equal to the length of hetero_pipeline_layer_split {self.args.hetero_pipeline_layer_split}" + ) setattr( self.args, "all_pipeline_model_parallel_size", @@ -197,12 +209,12 @@ def pre_validate_args(self): self.args.hetero_process_meshes = hetero_process_meshes # Device types - assert len(hetero_process_meshes) == len( - self.args.hetero_device_types - ), f"length of hetero_process_meshes {len(hetero_process_meshes)} should match length of hetero_device_types {len(self.args.hetero_device_types)}" - assert ( - self.args.hetero_current_device_type in self.args.hetero_device_types - ), f"hetero_current_device_type {self.args.hetero_current_device_type} should be in hetero_device_types {self.args.hetero_device_types}" + assert len(hetero_process_meshes) == len(self.args.hetero_device_types), ( + f"length of hetero_process_meshes {len(hetero_process_meshes)} should match length of hetero_device_types {len(self.args.hetero_device_types)}" + ) + assert self.args.hetero_current_device_type in self.args.hetero_device_types, ( + f"hetero_current_device_type {self.args.hetero_current_device_type} should be in hetero_device_types {self.args.hetero_device_types}" + ) current_process_mesh_idx = 0 accumulated_world_size = 0 @@ -225,10 +237,7 @@ def pre_validate_args(self): self.args.expert_model_parallel_size = ep self.args.data_parallel_size = dp self.args.pipeline_model_parallel_size = pp - if ( - self.args.expert_tensor_parallel_size_per_process_mesh - is not None - ): + if self.args.expert_tensor_parallel_size_per_process_mesh is not None: self.args.expert_tensor_parallel_size = ( self.args.expert_tensor_parallel_size_per_process_mesh[ current_process_mesh_idx @@ -245,9 +254,13 @@ def pre_validate_args(self): current_process_mesh_idx += 1 # DeepSeek-V4 Temporary if self.args.enable_hyper_connections: - assert not self.args.overlap_moe_expert_parallel_comm, "Hyper-connection is not supported with overlap_moe_expert_parallel_comm yet!" + assert not self.args.overlap_moe_expert_parallel_comm, ( + "Hyper-connection is not supported with overlap_moe_expert_parallel_comm yet!" + ) if self.args.experimental_attention_variant == "dsv4_hybrid": - assert self.args.context_parallel_size == 1, "Context parallelism is not supported with dsv4_hybrid attention variant yet!" + assert self.args.context_parallel_size == 1, ( + "Context parallelism is not supported with dsv4_hybrid attention variant yet!" + ) def post_validate_args(self): """Post-validate the arguments after Megatron function `validate_args`.""" @@ -259,9 +272,9 @@ def _parse_recompute_refined_config(recom_config, recom_config_name): """Parse refined recompute configuration.""" if recom_config is None: return None - assert isinstance( - recom_config, list - ), f"[{recom_config_name}] recompute configuration, is not list." + assert isinstance(recom_config, list), ( + f"[{recom_config_name}] recompute configuration, is not list." + ) recom_config = [ast.literal_eval(item) for item in recom_config] parsed_pp_size = 0 parsed_pp_chunk_config = [] @@ -282,24 +295,23 @@ def _parse_recompute_refined_config(recom_config, recom_config_name): f"for [{recom_config_name}] refined recompute " f"configuration, the sum [{len(cur_pp_stage_per_mc)}] of n0, n1, ... of sub-list should be equal to nums_micro_batch [{args.global_batch_size // (args.micro_batch_size * args.data_parallel_size)}]." ) - if ( - "method" in recom_config_name - or "granularity" in recom_config_name - ): - assert all( - val == 0 or val == 1 for val in cur_pp_stage_per_mc - ), f"the config-flag of {recom_config_name} must be 0 or 1" + if "method" in recom_config_name or "granularity" in recom_config_name: + assert all(val == 0 or val == 1 for val in cur_pp_stage_per_mc), ( + f"the config-flag of {recom_config_name} must be 0 or 1" + ) parsed_pp_chunk_config.append(cur_pp_stage_per_mc) if args.virtual_pipeline_model_parallel_size != None: assert ( parsed_pp_size == args.all_pipeline_model_parallel_size * args.virtual_pipeline_model_parallel_size - ), "for refined recompute configuration, the sum of axis 0 should be equal to pipeline-model-parallel-size * args.virtual_pipeline_model_parallel_size." + ), ( + "for refined recompute configuration, the sum of axis 0 should be equal to pipeline-model-parallel-size * args.virtual_pipeline_model_parallel_size." + ) else: - assert ( - parsed_pp_size == args.all_pipeline_model_parallel_size - ), "for refined recompute configuration, the sum of axis 0 should be equal to pipeline-model-parallel-size." + assert parsed_pp_size == args.all_pipeline_model_parallel_size, ( + "for refined recompute configuration, the sum of axis 0 should be equal to pipeline-model-parallel-size." + ) return parsed_pp_chunk_config if args.recompute_granularity_per_stage_micro_batch != None: @@ -312,23 +324,17 @@ def _parse_recompute_refined_config(recom_config, recom_config_name): "need to use a recompute method " ) - args.recompute_granularity_per_stage_micro_batch = ( - _parse_recompute_refined_config( - args.recompute_granularity_per_stage_micro_batch, - "recompute_granularity_per_stage_micro_batch", - ) + args.recompute_granularity_per_stage_micro_batch = _parse_recompute_refined_config( + args.recompute_granularity_per_stage_micro_batch, + "recompute_granularity_per_stage_micro_batch", ) - args.recompute_method_per_stage_micro_batch = ( - _parse_recompute_refined_config( - args.recompute_method_per_stage_micro_batch, - "recompute_method_per_stage_micro_batch", - ) + args.recompute_method_per_stage_micro_batch = _parse_recompute_refined_config( + args.recompute_method_per_stage_micro_batch, + "recompute_method_per_stage_micro_batch", ) - args.recompute_num_layers_per_stage_micro_batch = ( - _parse_recompute_refined_config( - args.recompute_num_layers_per_stage_micro_batch, - "recompute_num_layers_per_stage_micro_batch", - ) + args.recompute_num_layers_per_stage_micro_batch = _parse_recompute_refined_config( + args.recompute_num_layers_per_stage_micro_batch, + "recompute_num_layers_per_stage_micro_batch", ) # DualPipeV related @@ -342,18 +348,14 @@ def _parse_recompute_refined_config(recom_config, recom_config_name): middle_stage_layers = args.num_layers num_middle_stages = args.pipeline_model_parallel_size if args.decoder_first_pipeline_num_layers is not None: - middle_stage_layers = ( - middle_stage_layers - args.decoder_first_pipeline_num_layers - ) + middle_stage_layers = middle_stage_layers - args.decoder_first_pipeline_num_layers num_middle_stages = num_middle_stages - 1 assert args.decoder_first_pipeline_num_layers % 2 == 0, ( "The first pipeline stage must contain an even number of Transformer layers, " "so that DualPipeV can split it into two model chunks." ) if args.decoder_last_pipeline_num_layers is not None: - middle_stage_layers = ( - middle_stage_layers - args.decoder_last_pipeline_num_layers - ) + middle_stage_layers = middle_stage_layers - args.decoder_last_pipeline_num_layers num_middle_stages = num_middle_stages - 1 assert args.decoder_last_pipeline_num_layers % 2 == 0, ( "The last pipeline stage must contain an even number of Transformer layers, " @@ -361,60 +363,53 @@ def _parse_recompute_refined_config(recom_config, recom_config_name): ) if num_middle_stages > 0: assert middle_stage_layers > 0, "Layers can not be empty" - assert ( - middle_stage_layers % num_middle_stages == 0 - ), "Layers must be even split" + assert middle_stage_layers % num_middle_stages == 0, "Layers must be even split" num_layers_in_middle_stages = middle_stage_layers // num_middle_stages assert num_layers_in_middle_stages % 2 == 0, ( "The middle pipeline stage must contain an even number of Transformer layers, " "so that DualPipeV can split it into two model chunks." ) - assert ( - args.moe_shared_expert_overlap is False - ), " DualPipeV does not support simultaneous use with moe_shared_expert_overlap currently." + assert args.moe_shared_expert_overlap is False, ( + " DualPipeV does not support simultaneous use with moe_shared_expert_overlap currently." + ) if args.moe_fb_overlap: - assert ( - args.overlap_grad_reduce is False - and args.overlap_param_gather is False - ), ( + assert args.overlap_grad_reduce is False and args.overlap_param_gather is False, ( " DualPipeV configured with moe_fb_overlap is incompatible with either overlap_grad_reduce or overlap_param_gather. " " When moe_fb_overlap is enabled, DualPipeV activates the DW-split mechanism provided by Transformer Engine, " " which causes all param.grad attributes to be None during the backward-for-inputs phase. " " This absence of gradient tensors violates the assumptions of both overlap_grad_reduce and overlap_param_gather, precipitating an assertion failure within DDP." ) - assert ( - args.transformer_impl == 'transformer_engine' - ), 'delay_wgrad_compute is only supported with transformer_engine implementation' + assert args.transformer_impl == "transformer_engine", ( + "delay_wgrad_compute is only supported with transformer_engine implementation" + ) - assert ( - args.untie_embeddings_and_output_weights is True - ), " DualPipeV is not supported with shared embedding and lm head" - assert ( - args.mtp_num_layers is None - ), "DualPipeV is not supported with multi-token-predictor currently" + assert args.untie_embeddings_and_output_weights is True, ( + " DualPipeV is not supported with shared embedding and lm head" + ) + assert args.mtp_num_layers is None, ( + "DualPipeV is not supported with multi-token-predictor currently" + ) if args.peft_type is not None: - assert ( - args.transformer_impl == 'transformer_engine' - ), 'PEFT is only supported with transformer_engine implementation' + assert args.transformer_impl == "transformer_engine", ( + "PEFT is only supported with transformer_engine implementation" + ) if ( args.num_experts is not None and args.moe_shared_expert_intermediate_size is not None ): - assert ( - not args.moe_shared_expert_overlap - ), 'PEFT is incompatible with moe_shared_expert_overlap' + assert not args.moe_shared_expert_overlap, ( + "PEFT is incompatible with moe_shared_expert_overlap" + ) assert args.num_experts is None, "PEFT is not tested with MoE currently" assert ( args.recompute_method is None and args.recompute_granularity is None and args.recompute_num_layers is None ), "PEFT will raise comfilcts with recompute currently" - assert ( - args.ckpt_format == 'torch' - ), "PEFT is only tested with torch format checkpoint" + assert args.ckpt_format == "torch", "PEFT is only tested with torch format checkpoint" # Engram related. if self.args.use_engram: @@ -432,39 +427,40 @@ def _parse_recompute_refined_config(recom_config, recom_config_name): ) self.args.engram_embedding_parallel_size = None elif self.args.engram_embedding_parallel_method == "alltoall": - assert ( - self.args.engram_embedding_parallel_size is not None - ), "embedding parallel size should be specified when using alltoall" + assert self.args.engram_embedding_parallel_size is not None, ( + "embedding parallel size should be specified when using alltoall" + ) else: raise ValueError( f"Invalid embedding parallel method: {self.args.engram_embedding_parallel_method}" ) if self.args.engram_offload_embedding_optimizer_states: - assert ( - self.args.engram_embedding_parallel_method == "alltoall" - ), f"Offloading embedding optimizer states is only supported when using alltoall for engram embedding parallelism, now is {self.args.engram_embedding_parallel_method}." - assert ( - self.args.optimizer_cpu_offload - ), "Offloading embedding optimizer states requires optimizer_cpu_offload to be enabled." + assert self.args.engram_embedding_parallel_method == "alltoall", ( + f"Offloading embedding optimizer states is only supported when using alltoall for engram embedding parallelism, now is {self.args.engram_embedding_parallel_method}." + ) + assert self.args.optimizer_cpu_offload, ( + "Offloading embedding optimizer states requires optimizer_cpu_offload to be enabled." + ) warnings.warn( "Offloading embedding optimizer states will offload all embedding optimizer states to CPU, which may cause slowdown. " "Please make sure this is what you want. This is typically used to save GPU memory when Engram embedding is large while accelerators are limited." "If you do not want to offload all embedding optimizer states to CPU, please disable this and set the --optimizer-offload-fraction to a value less than 1 to offload part of the optimizer states to CPU." "Of course you can set the --optimizer-offload-fraction to offload other params meanwhile enable this to offload all embedding optimizer states to CPU." ) - assert ( - not self.args.use_megatron_fsdp - ), "Megatron FSDP is not supported yet; support is planned for a later version." - assert ( - not self.args.init_model_with_meta_device - ), "Init_model_with_meta_device is not supported yet; support is planned for a later version." - assert ( - self.args.use_distributed_optimizer - ), "When use engram, distributed_optimizer must be enabled, because there is a bug caused by allreduce grad norm in model parallel group when do not use distributed_optimizer. We have not found a pretty solution yet, so disable it temporarily." + assert not self.args.use_megatron_fsdp, ( + "Megatron FSDP is not supported yet; support is planned for a later version." + ) + assert not self.args.init_model_with_meta_device, ( + "Init_model_with_meta_device is not supported yet; support is planned for a later version." + ) + assert self.args.use_distributed_optimizer, ( + "When use engram, distributed_optimizer must be enabled, because there is a bug caused by allreduce grad norm in model parallel group when do not use distributed_optimizer. We have not found a pretty solution yet, so disable it temporarily." + ) assert not ( - args.pipeline_model_parallel_size == 1 - and args.overlap_moe_expert_parallel_comm - ), "When no pipeline and enable overlap_moe_expert_parallel_comm, a bug will occur, it will be fixed in a later version." + args.pipeline_model_parallel_size == 1 and args.overlap_moe_expert_parallel_comm + ), ( + "When no pipeline and enable overlap_moe_expert_parallel_comm, a bug will occur, it will be fixed in a later version." + ) def _add_hetero_args(parser): @@ -610,285 +606,286 @@ def _add_peft_args(parser): def _add_network_size_args(parser): - group = parser.add_argument_group(title='flagscale network size') + group = parser.add_argument_group(title="flagscale network size") group.add_argument( - '--norm-init-weight', + "--norm-init-weight", type=float, default=None, help="Norm weight initialization.", ) group.add_argument( - '--multiple-of', + "--multiple-of", type=int, default=None, - help='Multiplier for setting Feed-Forward Network hidden size when swiglu.', + help="Multiplier for setting Feed-Forward Network hidden size when swiglu.", ) group.add_argument( - '--hidden-dim-multiplier', + "--hidden-dim-multiplier", type=float, default=None, - help='Custom Multiplier for setting Feed-Forward Network hidden dim when swiglu.', + help="Custom Multiplier for setting Feed-Forward Network hidden dim when swiglu.", ) return parser def _add_logging_args(parser): - group = parser.add_argument_group(title='flagscale logging') + group = parser.add_argument_group(title="flagscale logging") group.add_argument( - '--wandb-mode', + "--wandb-mode", type=str, - choices=['online', 'offline', 'disabled'], - default='offline', + choices=["online", "offline", "disabled"], + default="offline", help='Can be "online", "offline" or "disabled". Defaults to "offline".', ) group.add_argument( - '--wandb-api-key', + "--wandb-api-key", type=str, - default='', - help='The wandb API keys and must be provided if using online mode.', + default="", + help="The wandb API keys and must be provided if using online mode.", ) group.add_argument( - '--wandb-log-model', action='store_true', help='If set, write model to wandb.' + "--wandb-log-model", action="store_true", help="If set, write model to wandb." ) group.add_argument( - '--wandb-log-model-interval', + "--wandb-log-model-interval", type=int, default=1000, - help='The interval to save the model to wandb.', + help="The interval to save the model to wandb.", ) return parser def _add_training_args(parser): - group = parser.add_argument_group(title='flagscale training') + group = parser.add_argument_group(title="flagscale training") group.add_argument( - '--recompute-granularity-per-stage-micro-batch', - nargs='*', + "--recompute-granularity-per-stage-micro-batch", + nargs="*", type=str, default=None, - help='used with recompute-granularity=full, setting recompute granularity' - 'of each stage and each micro-batch. This argument must be a two-dimension list, ' - 'the sum of the first item of all the sub-lists should be equal to pipeline-model-parallel-size.' - 'Every sub-list is in the form: n0, flag0, n1, flag1,... except the first item, which is the stage number.' - 'The sum of n0, n1, ... should be equal to nums-micro-batch.' - 'granularity flag: 0 means turning off full recompute, 1 means turning on', + help="used with recompute-granularity=full, setting recompute granularity" + "of each stage and each micro-batch. This argument must be a two-dimension list, " + "the sum of the first item of all the sub-lists should be equal to pipeline-model-parallel-size." + "Every sub-list is in the form: n0, flag0, n1, flag1,... except the first item, which is the stage number." + "The sum of n0, n1, ... should be equal to nums-micro-batch." + "granularity flag: 0 means turning off full recompute, 1 means turning on", ) group.add_argument( - '--recompute-method-per-stage-micro-batch', - nargs='*', + "--recompute-method-per-stage-micro-batch", + nargs="*", type=str, default=None, - help='used with recompute-granularity=full, setting recompute method ' - 'of each stage and each micro-batch. This argument must be a two-dimension list, ' - 'the sum of the first item of all the sub-lists should be equal to pipeline-model-parallel-size.' - 'Every sub-list is in the form: n0, flag0, n1, flag1,... except the first item, which is the stage number.' - 'The sum of n0, n1, ... should be equal to nums-micro-batch.' - 'method: 0 means uniform, 1 means block', + help="used with recompute-granularity=full, setting recompute method " + "of each stage and each micro-batch. This argument must be a two-dimension list, " + "the sum of the first item of all the sub-lists should be equal to pipeline-model-parallel-size." + "Every sub-list is in the form: n0, flag0, n1, flag1,... except the first item, which is the stage number." + "The sum of n0, n1, ... should be equal to nums-micro-batch." + "method: 0 means uniform, 1 means block", ) group.add_argument( - '--recompute-num-layers-per-stage-micro-batch', - nargs='*', + "--recompute-num-layers-per-stage-micro-batch", + nargs="*", type=str, default=None, - help='used with recompute-granularity=full, setting recompute num layers ' - 'of each stage and each micro-batch. This argument must be a two-dimension list, ' - 'Every sub-list is in the form: n0, num_laryers0, n1, num_laryers1,... except the first item, which is the stage number.' - 'The sum of n0, n1, ... should be equal to nums-micro-batch. ', + help="used with recompute-granularity=full, setting recompute num layers " + "of each stage and each micro-batch. This argument must be a two-dimension list, " + "Every sub-list is in the form: n0, num_laryers0, n1, num_laryers1,... except the first item, which is the stage number." + "The sum of n0, n1, ... should be equal to nums-micro-batch. ", ) group.add_argument( - '--skip-samples-range', - nargs='+', + "--skip-samples-range", + nargs="+", type=int, default=None, - help='Range of samples to skip during training.', + help="Range of samples to skip during training.", ) group.add_argument( - '--skip-iters-range', - nargs='+', + "--skip-iters-range", + nargs="+", type=int, default=None, - help='Range of iterations to skip during training.', + help="Range of iterations to skip during training.", ) group.add_argument( - '--use-dualpipev', - action='store_true', - help='Use DualPipeV pipeline schedule method', + "--use-dualpipev", + action="store_true", + help="Use DualPipeV pipeline schedule method", ) group.add_argument( - '--moe-fb-overlap', - action='store_true', - help='DualPipeV overlapping of moe a2a communication and forward/backward computation', + "--moe-fb-overlap", + action="store_true", + help="DualPipeV overlapping of moe a2a communication and forward/backward computation", ) return parser def _add_learning_rate_args(parser): - group = parser.add_argument_group(title='flagscale learning rate') + group = parser.add_argument_group(title="flagscale learning rate") ## stablelm2-scheduler consists of multiple stages group.add_argument( - '--lr-decay-stablelm2-cosine-samples', + "--lr-decay-stablelm2-cosine-samples", type=int, default=0, - help='Samples number of cosine scheduler including warmup samples, used in stablelm2 scheduler.', + help="Samples number of cosine scheduler including warmup samples, used in stablelm2 scheduler.", ) group.add_argument( - '--lr-decay-stablelm2-cosine-max-lr', + "--lr-decay-stablelm2-cosine-max-lr", type=float, default=None, - help='Maximum lr of cosine scheduler, used in stablelm2 scheduler.', + help="Maximum lr of cosine scheduler, used in stablelm2 scheduler.", ) group.add_argument( - '--lr-decay-stablelm2-cosine-period-samples', + "--lr-decay-stablelm2-cosine-period-samples", type=int, default=0, - help='Period of cosine scheduler, used in stablelm2 scheduler.', + help="Period of cosine scheduler, used in stablelm2 scheduler.", ) group.add_argument( - '--lr-decay-stablelm2-rsqrt-samples', + "--lr-decay-stablelm2-rsqrt-samples", type=int, default=0, - help='Samples number of rsqrt scheduler used in stablelm2 scheduler.', + help="Samples number of rsqrt scheduler used in stablelm2 scheduler.", ) group.add_argument( - '--lr-decay-stablelm2-decay-samples', + "--lr-decay-stablelm2-decay-samples", type=int, default=0, - help='Samples number of decay scheduler used in stablelm2 scheduler.', + help="Samples number of decay scheduler used in stablelm2 scheduler.", ) group.add_argument( - '--lr-decay-stablelm2-alpha', + "--lr-decay-stablelm2-alpha", type=float, default=1.0, - help='Numerator used in stablelm2 scheduler.', + help="Numerator used in stablelm2 scheduler.", ) group.add_argument( - '--lr-decay-stablelm2-beta', + "--lr-decay-stablelm2-beta", type=float, default=0.0, - help='Denominator used in stablelm2 scheduler.', + help="Denominator used in stablelm2 scheduler.", ) return parser def _add_checkpointing_args(parser): - group = parser.add_argument_group(title='flagscale checkpointing') + group = parser.add_argument_group(title="flagscale checkpointing") group.add_argument( - '--rampup-save-interval', + "--rampup-save-interval", type=int, default=None, - help='Number of iterations between checkpoint saves.in the ramup phase.', + help="Number of iterations between checkpoint saves.in the ramup phase.", ) group.add_argument( - '--save-when-num-microbatches-change', - action='store_true', - help='Save param name to index maps only', + "--save-when-num-microbatches-change", + action="store_true", + help="Save param name to index maps only", ) return parser def _add_distributed_args(parser): - group = parser.add_argument_group(title='flagscale distributed') + group = parser.add_argument_group(title="flagscale distributed") group.add_argument( - '--standalone-embedding-stage', - action='store_true', + "--standalone-embedding-stage", + action="store_true", default=False, - help='If set, *input* embedding layer ' - 'is placed on its own pipeline stage, without any ' - 'transformer layers. (For T5, this flag currently only ' - 'affects the encoder embedding.)', + help="If set, *input* embedding layer " + "is placed on its own pipeline stage, without any " + "transformer layers. (For T5, this flag currently only " + "affects the encoder embedding.)", ) group.add_argument( - '--use-partial-reduce-for-shared-embedding', - action='store_true', - help='Use partial reduce for shared word embedding.', + "--use-partial-reduce-for-shared-embedding", + action="store_true", + help="Use partial reduce for shared word embedding.", ) group.add_argument( - '--no-shared-fs', - action='store_true', - help='Indicate whether not running on a shared file system.', + "--no-shared-fs", + action="store_true", + help="Indicate whether not running on a shared file system.", ) group.add_argument( - '--use-padded-layerwise-optimizer', - action='store_true', - help='Enable pad when use layer-wise optimizer.' + "--use-padded-layerwise-optimizer", + action="store_true", + help="Enable pad when use layer-wise optimizer.", ) return parser def _add_validation_args(parser): - group = parser.add_argument_group(title='flagscale validation') + group = parser.add_argument_group(title="flagscale validation") group.add_argument( - '--extra-eval-interval', + "--extra-eval-interval", type=int, default=None, - help='Interval between running evaluation on ' 'extra validation sets.', + help="Interval between running evaluation on extra validation sets.", ) return parser def _add_tokenizer_args(parser): - group = parser.add_argument_group(title='flagscale tokenizer') + group = parser.add_argument_group(title="flagscale tokenizer") group.add_argument( - '--special-tokens-file', + "--special-tokens-file", type=str, default=None, - help='Path to the BPE special tokens file.', + help="Path to the BPE special tokens file.", ) group.add_argument( - '--tokenizer-path', + "--tokenizer-path", type=str, default=None, - help='Path to the huggingface tokenizer.', + help="Path to the huggingface tokenizer.", ) return parser def _add_data_args(parser): - group = parser.add_argument_group(title='flagscale data') + group = parser.add_argument_group(title="flagscale data") group.add_argument( - '--extra-valid-data-path', - nargs='*', + "--extra-valid-data-path", + nargs="*", default=None, - help='The weight, prefix list for an independent extra validation dataset. ' - 'The accepted format is a list of weight, prefix and tag, ' - 'e.g. weight1 prefix1 tag1 weight2 prefix2 tag2. ' - 'The weight1 means the number of tokens in the prefix1 dataset. ', + help="The weight, prefix list for an independent extra validation dataset. " + "The accepted format is a list of weight, prefix and tag, " + "e.g. weight1 prefix1 tag1 weight2 prefix2 tag2. " + "The weight1 means the number of tokens in the prefix1 dataset. ", ) group.add_argument( - '--finetune-dataset-type', + "--finetune-dataset-type", type=str, default=None, - choices=['CPT', None], - help='datasets type during finetunning.', + choices=["CPT", None], + help="datasets type during finetunning.", ) group.add_argument( - '--apply-sft-dataset-separated-loss-mask-if-existed', - action='store_true', - help='If set, use sft dataset with separated loss mask files, ' - 'if _loss_mask_document.bin and _loss_mask_document.idx existed.', + "--apply-sft-dataset-separated-loss-mask-if-existed", + action="store_true", + help="If set, use sft dataset with separated loss mask files, " + "if _loss_mask_document.bin and _loss_mask_document.idx existed.", ) return parser def _add_vision_args(parser): - group = parser.add_argument_group(title='flagscale vision') + group = parser.add_argument_group(title="flagscale vision") group.add_argument( - '--qk-layernorm-hidden-dim', - action='store_true', - help='Whether to layer normalize the q and k attention embeddings on hidden dimension rather than head dimension', + "--qk-layernorm-hidden-dim", + action="store_true", + help="Whether to layer normalize the q and k attention embeddings on hidden dimension rather than head dimension", ) return parser + def _add_straggler_args(parser): group = parser.add_argument_group(title="flagscale straggler") @@ -955,6 +952,8 @@ def _add_straggler_args(parser): help="Directory used to save FlagScale straggler reports.", ) return parser + + def _add_perf_monitor_args(parser): group = parser.add_argument_group(title="flagscale perf monitor") @@ -1023,65 +1022,66 @@ def _add_perf_monitor_args(parser): ) return parser + def _add_flagos_args(parser): group = parser.add_argument_group(title="flagscale fl") group.add_argument( - '--mg-fl-prefer', + "--mg-fl-prefer", type=str, - choices=['cuda', 'musa', 'txda'], - default='', - help='Backend selection for megatron fl.', + choices=["cuda", "musa", "txda"], + default="", + help="Backend selection for megatron fl.", ) group.add_argument( - '--te-fl-prefer', + "--te-fl-prefer", type=str, - choices=['flagos', 'vendor', 'reference'], - default='vendor', - help='Backend selection for transformer engine fl.', + choices=["flagos", "vendor", "reference"], + default="vendor", + help="Backend selection for transformer engine fl.", ) group.add_argument( - '--te-fl-per-op', + "--te-fl-per-op", type=str, default=None, - help='Backend selection for custom ops.', + help="Backend selection for custom ops.", ) group.add_argument( - '--te-fl-allow-vendors', + "--te-fl-allow-vendors", type=str, default=None, - help='Allow vendors for transformer engine fl.', + help="Allow vendors for transformer engine fl.", ) group.add_argument( - '--te-fl-deny-vendors', + "--te-fl-deny-vendors", type=str, default=None, - help='Deny vendors for transformer engine fl.', + help="Deny vendors for transformer engine fl.", ) group.add_argument( - '--enable-flag-gems', - action='store_true', - help='Enable flag gems to replace torch ops for distributed training.', + "--enable-flag-gems", + action="store_true", + help="Enable flag gems to replace torch ops for distributed training.", ) group.add_argument( - '--flag-gems-log-path', type=str, default=None, help='Path of flag gems logging' + "--flag-gems-log-path", type=str, default=None, help="Path of flag gems logging" ) group.add_argument( - '--flag-gems-unused', nargs='*', default=None, help='Flag Gems unused ops list' + "--flag-gems-unused", nargs="*", default=None, help="Flag Gems unused ops list" ) return parser def _add_flagscale_specific_args(parser): """Add FlagScale-specific arguments that don't fit in other categories.""" - group = parser.add_argument_group(title='flagscale specific') + group = parser.add_argument_group(title="flagscale specific") # Inference args (not in any upstream dataclass config) group.add_argument( - '--inference-wandb-logging-step-interval', + "--inference-wandb-logging-step-interval", type=int, default=0, - help='Step interval for logging inference metrics to wandb. ' - 'Default to 0 to disable inference wandb logging.', + help="Step interval for logging inference metrics to wandb. " + "Default to 0 to disable inference wandb logging.", ) return parser diff --git a/flagscale/train/megatron/training/async_utils.py b/flagscale/train/megatron/training/async_utils.py index a085a7caf7..338b1d946c 100644 --- a/flagscale/train/megatron/training/async_utils.py +++ b/flagscale/train/megatron/training/async_utils.py @@ -4,6 +4,7 @@ This module provides a singleton instance of AsyncCallsQueue which manages the async checkpoint save calls. """ + import logging import time @@ -17,10 +18,14 @@ try: from nvidia_resiliency_ext.checkpointing.async_ckpt.core import AsyncRequest as NVRxAsyncRequest from nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async import _results_queue - from nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver import save_state_dict_async_finalize + from nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver import ( + save_state_dict_async_finalize, + ) except (ImportError, ModuleNotFoundError): from megatron.core.dist_checkpointing.strategies.filesystem_async import _results_queue - from megatron.core.dist_checkpointing.strategies.state_dict_saver import save_state_dict_async_finalize + from megatron.core.dist_checkpointing.strategies.state_dict_saver import ( + save_state_dict_async_finalize, + ) NVRxAsyncRequest = ABC @@ -45,7 +50,7 @@ def _get_async_calls_queue(): return _async_calls_queue -def init_persistent_async_worker(rank: int, mp_mode: str = 'spawn'): +def init_persistent_async_worker(rank: int, mp_mode: str = "spawn"): global _async_calls_queue args = get_args() async_strategy, async_modules = get_async_strategy(getattr(args, "async_strategy", "nvrx")) @@ -69,9 +74,12 @@ def init_persistent_async_worker(rank: int, mp_mode: str = 'spawn'): **kwargs, ) # initialize ckpt write results queue - get_write_results_queue('fork') + get_write_results_queue("fork") if rank == 0: - print(f"init_persistent_async_worker: rank {rank}, Async Caller Started in {time.time() - time_start} seconds", flush=True) + print( + f"init_persistent_async_worker: rank {rank}, Async Caller Started in {time.time() - time_start} seconds", + flush=True, + ) def schedule_async_save(async_request: AsyncRequest | NVRxAsyncRequest): @@ -98,7 +106,7 @@ def maybe_finalize_async_save(blocking: bool = False, terminate=False): return if blocking and not is_empty_async_queue(): - print_rank_0('Unfinalized async checkpoint saves. Finalizing them synchronously now.') + print_rank_0("Unfinalized async checkpoint saves. Finalizing them synchronously now.") async_calls_queue = _async_calls_queue if async_calls_queue is not None: @@ -107,6 +115,7 @@ def maybe_finalize_async_save(blocking: bool = False, terminate=False): # Clean up finished deletion processes to prevent zombies # Import here to avoid circular dependency from .checkpointing import finalize_deletion_processes + finalize_deletion_processes(blocking=blocking or terminate) if terminate and async_calls_queue is not None: @@ -124,7 +133,7 @@ def is_empty_async_queue() -> bool: def reset_persistent_async_worker(async_strategy): global _async_calls_queue, _results_queue - + if _async_calls_queue is not None: _async_calls_queue.close(abort=True) del _async_calls_queue diff --git a/flagscale/train/megatron/training/checkpointing.py b/flagscale/train/megatron/training/checkpointing.py index eafdb85b53..11a536bc5f 100644 --- a/flagscale/train/megatron/training/checkpointing.py +++ b/flagscale/train/megatron/training/checkpointing.py @@ -24,7 +24,10 @@ from megatron.core import dist_checkpointing, mpu, tensor_parallel from megatron.core.dist_checkpointing.mapping import ShardedObject -from megatron.core.dist_checkpointing.strategies.torch import TorchDistLoadShardedStrategy, TorchDistSaveShardedStrategy +from megatron.core.dist_checkpointing.strategies.torch import ( + TorchDistLoadShardedStrategy, + TorchDistSaveShardedStrategy, +) from megatron.core.dist_checkpointing.strategies.fully_parallel import ( FullyParallelLoadStrategyWrapper, FullyParallelSaveStrategyWrapper, @@ -45,13 +48,16 @@ from .utils import append_to_progress_log, is_last_rank, print_rank_0, unwrap_model try: - from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import preprocess_state_dict_for_uneven_dtensor + from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( + preprocess_state_dict_for_uneven_dtensor, + ) from megatron.core.transformer.fsdp_dtensor_checkpoint import ( print_diff_in_state_dicts, handle_fp8_extra_state_case, handle_swiglu_in_state_dict, handle_experts_in_state_dict, ) + HAVE_MEGATRON_FSDP = True except ImportError: HAVE_MEGATRON_FSDP = False @@ -61,32 +67,38 @@ try: from modelopt.torch.opt.plugins import save_modelopt_state, save_sharded_modelopt_state from megatron.post_training.utils import print_distributed_quant_summary + has_nvidia_modelopt = True except Exception: has_nvidia_modelopt = False try: - from nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async import FileSystemWriterAsync - from nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver import save_state_dict_async_plan + from nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async import ( + FileSystemWriterAsync, + ) + from nvidia_resiliency_ext.checkpointing.async_ckpt.state_dict_saver import ( + save_state_dict_async_plan, + ) HAVE_NVRX = True except (ImportError, ModuleNotFoundError): - HAVE_NVRX = False _CHECKPOINT_VERSION = None _LOADED_ITERATION = None logger = getLogger(__name__) -_NON_PERSISTENT_CKPT_SUBDIR = 'non_persistent' +_NON_PERSISTENT_CKPT_SUBDIR = "non_persistent" from megatron.plugin.platform import get_platform + cur_platform = get_platform() # Track deletion processes to prevent zombies _deletion_processes = [] + def finalize_deletion_processes(blocking=False): """Clean up deletion processes to prevent zombie processes. @@ -104,7 +116,9 @@ def finalize_deletion_processes(blocking=False): finished = [] for proc in _deletion_processes: if not proc.is_alive() or blocking: - logger.debug(f"Joining deletion process {proc.pid} (blocking={blocking}, is_alive={proc.is_alive()})") + logger.debug( + f"Joining deletion process {proc.pid} (blocking={blocking}, is_alive={proc.is_alive()})" + ) proc.join() finished.append(proc) for proc in finished: @@ -114,8 +128,7 @@ def finalize_deletion_processes(blocking=False): def set_checkpoint_version(value): global _CHECKPOINT_VERSION if _CHECKPOINT_VERSION is not None: - assert _CHECKPOINT_VERSION == value, \ - "checkpoint versions do not match" + assert _CHECKPOINT_VERSION == value, "checkpoint versions do not match" _CHECKPOINT_VERSION = value @@ -155,31 +168,32 @@ def _compare(arg_name, old_arg_name=None, default=None): else: checkpoint_value = getattr(checkpoint_args, ckpt_arg_name) args_value = getattr(args, arg_name) - error_message = '{} value from checkpoint ({}) is not equal to the ' \ - 'input argument value ({}).'.format( - arg_name, checkpoint_value, args_value) + error_message = ( + "{} value from checkpoint ({}) is not equal to the input argument value ({}).".format( + arg_name, checkpoint_value, args_value + ) + ) assert checkpoint_value == args_value, error_message - _compare('num_layers') - _compare('hidden_size') - _compare('num_attention_heads') - _compare('add_position_embedding', default=True) + _compare("num_layers") + _compare("hidden_size") + _compare("num_attention_heads") + _compare("add_position_embedding", default=True) if args.vocab_file: - _compare('max_position_embeddings') - _compare('make_vocab_size_divisible_by') + _compare("max_position_embeddings") + _compare("make_vocab_size_divisible_by") if not args.use_dist_ckpt: - _compare('padded_vocab_size') - _compare('tokenizer_type') + _compare("padded_vocab_size") + _compare("tokenizer_type") if args.data_parallel_random_init: - _compare('data_parallel_random_init') + _compare("data_parallel_random_init") if args.phase_transition_iterations: - _compare('global_batch_size') + _compare("global_batch_size") if get_checkpoint_version() < 3.0: - _compare('tensor_model_parallel_size', - old_arg_name='model_parallel_size') + _compare("tensor_model_parallel_size", old_arg_name="model_parallel_size") if get_checkpoint_version() >= 3.0 and not args.use_dist_ckpt: - _compare('tensor_model_parallel_size') - _compare('pipeline_model_parallel_size') + _compare("tensor_model_parallel_size") + _compare("pipeline_model_parallel_size") def isfile(filename) -> bool: @@ -200,29 +214,36 @@ def ensure_directory_exists(filename, check_parent=True): os.makedirs(dirname, exist_ok=True) -def get_checkpoint_name(checkpoints_path, iteration, release=False, - pipeline_parallel=None, - tensor_rank=None, pipeline_rank=None, - expert_parallel=None, expert_rank=None, - return_base_dir=False, basename="model_optim_rng.pt"): +def get_checkpoint_name( + checkpoints_path, + iteration, + release=False, + pipeline_parallel=None, + tensor_rank=None, + pipeline_rank=None, + expert_parallel=None, + expert_rank=None, + return_base_dir=False, + basename="model_optim_rng.pt", +): """Determine the directory name for this rank's checkpoint.""" if release: - directory = 'release' + directory = "release" else: - directory = 'iter_{:07d}'.format(iteration) + directory = "iter_{:07d}".format(iteration) if return_base_dir: common_path = os.path.join(checkpoints_path, directory) return common_path # Use both the tensor and pipeline MP rank. if pipeline_parallel is None: - pipeline_parallel = (mpu.get_pipeline_model_parallel_world_size() > 1) + pipeline_parallel = mpu.get_pipeline_model_parallel_world_size() > 1 if tensor_rank is None: tensor_rank = mpu.get_tensor_model_parallel_rank() if pipeline_rank is None: pipeline_rank = mpu.get_pipeline_model_parallel_rank() if expert_parallel is None: - expert_parallel = (mpu.get_expert_model_parallel_world_size() > 1) + expert_parallel = mpu.get_expert_model_parallel_world_size() > 1 if expert_rank is None: expert_rank = mpu.get_expert_model_parallel_rank() @@ -230,14 +251,14 @@ def get_checkpoint_name(checkpoints_path, iteration, release=False, # optimizer, then the optimizer's path must additionally include the # data parallel rank. if not pipeline_parallel: - common_path = os.path.join(checkpoints_path, directory, - f'mp_rank_{tensor_rank:02d}') + common_path = os.path.join(checkpoints_path, directory, f"mp_rank_{tensor_rank:02d}") else: - common_path = os.path.join(checkpoints_path, directory, - f'mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}') + common_path = os.path.join( + checkpoints_path, directory, f"mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}" + ) if expert_parallel: - common_path = common_path + f'_{expert_rank:03d}' + common_path = common_path + f"_{expert_rank:03d}" return os.path.join(common_path, basename) @@ -246,7 +267,7 @@ def get_load_checkpoint_path_by_args(args, load_arg="load"): """Get the checkpoint path based on the arguments.""" load_dir = getattr(args, load_arg) iteration, release = -1, False - tracker_filename = 'because load directory is not defined' + tracker_filename = "because load directory is not defined" if load_dir is not None: tracker_filename = get_checkpoint_tracker_filename(load_dir) if isfile(tracker_filename): @@ -260,8 +281,7 @@ def get_load_checkpoint_path_by_args(args, load_arg="load"): def get_distributed_optimizer_checkpoint_name(model_checkpoint_name): - return os.path.join(os.path.dirname(model_checkpoint_name), - "distrib_optim.pt") + return os.path.join(os.path.dirname(model_checkpoint_name), "distrib_optim.pt") def find_checkpoint_rank_0(checkpoints_path, iteration, release=False): @@ -274,41 +294,65 @@ def find_checkpoint_rank_0(checkpoints_path, iteration, release=False): """ # Look for checkpoint with no pipelining and no expert parallelism - filename = get_checkpoint_name(checkpoints_path, iteration, release, - pipeline_parallel=False, - tensor_rank=0, pipeline_rank=0, - expert_parallel=False, expert_rank=0) + filename = get_checkpoint_name( + checkpoints_path, + iteration, + release, + pipeline_parallel=False, + tensor_rank=0, + pipeline_rank=0, + expert_parallel=False, + expert_rank=0, + ) if isfile(filename): return filename # Look for checkpoint with no pipelining and expert parallelism - filename = get_checkpoint_name(checkpoints_path, iteration, release, - pipeline_parallel=False, - tensor_rank=0, pipeline_rank=0, - expert_parallel=True, expert_rank=0) + filename = get_checkpoint_name( + checkpoints_path, + iteration, + release, + pipeline_parallel=False, + tensor_rank=0, + pipeline_rank=0, + expert_parallel=True, + expert_rank=0, + ) if isfile(filename): return filename # Look for checkpoint with pipelining and no expert parallelism - filename = get_checkpoint_name(checkpoints_path, iteration, release, - pipeline_parallel=True, - tensor_rank=0, pipeline_rank=0, - expert_parallel=False, expert_rank=0) + filename = get_checkpoint_name( + checkpoints_path, + iteration, + release, + pipeline_parallel=True, + tensor_rank=0, + pipeline_rank=0, + expert_parallel=False, + expert_rank=0, + ) if isfile(filename): return filename # Look for checkpoint with pipelining and expert parallelism - filename = get_checkpoint_name(checkpoints_path, iteration, release, - pipeline_parallel=True, - tensor_rank=0, pipeline_rank=0, - expert_parallel=True, expert_rank=0) + filename = get_checkpoint_name( + checkpoints_path, + iteration, + release, + pipeline_parallel=True, + tensor_rank=0, + pipeline_rank=0, + expert_parallel=True, + expert_rank=0, + ) if isfile(filename): return filename # Look for a distributed checkpoint - filename = get_checkpoint_name(checkpoints_path, iteration, release, - pipeline_parallel=True, - return_base_dir=True) + filename = get_checkpoint_name( + checkpoints_path, iteration, release, pipeline_parallel=True, return_base_dir=True + ) if dist_checkpointing.check_is_distributed_checkpoint(filename): return filename @@ -316,10 +360,9 @@ def find_checkpoint_rank_0(checkpoints_path, iteration, release=False): def get_checkpoint_tracker_filename(checkpoints_path): - """Tracker file rescords the latest chckpoint during training to restart from.""" - return os.path.join(checkpoints_path, 'latest_checkpointed_iteration.txt') + return os.path.join(checkpoints_path, "latest_checkpointed_iteration.txt") def checkpoint_exists(checkpoints_path): @@ -335,25 +378,29 @@ def read_metadata(tracker_filename): iteration = -1 release = False - with open_file(tracker_filename, 'r') as f: + with open_file(tracker_filename, "r") as f: metastring = f.read().strip() try: iteration = int(metastring) except ValueError: - release = metastring == 'release' + release = metastring == "release" if not release: - print_rank_0('ERROR: Invalid metadata file {}. Exiting'.format( - tracker_filename)) + print_rank_0("ERROR: Invalid metadata file {}. Exiting".format(tracker_filename)) sys.exit() else: # Set iteration to 0 for release checkpoints iteration = 0 - assert iteration >= 0 or release, 'error parsing metadata file {}'.format( - tracker_filename) + assert iteration >= 0 or release, "error parsing metadata file {}".format(tracker_filename) # Get the max iteration retrieved across the ranks. if torch.distributed.is_initialized(): - iters_cuda = torch.tensor([iteration], dtype=torch.long, device=cur_platform.device_name() if torch.distributed.get_backend() != 'gloo' else 'cpu') + iters_cuda = torch.tensor( + [iteration], + dtype=torch.long, + device=cur_platform.device_name() + if torch.distributed.get_backend() != "gloo" + else "cpu", + ) torch.distributed.all_reduce(iters_cuda, op=torch.distributed.ReduceOp.MAX) max_iter = iters_cuda[0].item() @@ -362,10 +409,12 @@ def read_metadata(tracker_filename): # iteration across all ranks. if iteration != max_iter: rank = torch.distributed.get_rank() - print('WARNING: on rank {} found iteration {} in the ' - 'metadata while max iteration across the ranks ' - 'is {}, replacing it with max iteration.'.format( - rank, iteration, max_iter), flush=True) + print( + "WARNING: on rank {} found iteration {} in the " + "metadata while max iteration across the ranks " + "is {}, replacing it with max iteration.".format(rank, iteration, max_iter), + flush=True, + ) else: # When loading a checkpoint outside of training (for example, # when editing it), we might not have torch distributed @@ -374,25 +423,31 @@ def read_metadata(tracker_filename): return max_iter, release -def get_rng_state(ckpt_format: str, tp_group: torch.distributed.ProcessGroup, pp_group: torch.distributed.ProcessGroup) -> Union[List[Dict[str, Any]], ShardedObject]: +def get_rng_state( + ckpt_format: str, + tp_group: torch.distributed.ProcessGroup, + pp_group: torch.distributed.ProcessGroup, +) -> Union[List[Dict[str, Any]], ShardedObject]: """Collect rng state across data parallel ranks.""" args = get_args() rng_state = { - 'random_rng_state': random.getstate(), - 'np_rng_state': np.random.get_state(), - 'torch_rng_state': torch.get_rng_state(), - 'cuda_rng_state': cur_platform.get_rng_state(), - 'rng_tracker_states': tensor_parallel.get_cuda_rng_tracker().get_states()} + "random_rng_state": random.getstate(), + "np_rng_state": np.random.get_state(), + "torch_rng_state": torch.get_rng_state(), + "cuda_rng_state": cur_platform.get_rng_state(), + "rng_tracker_states": tensor_parallel.get_cuda_rng_tracker().get_states(), + } rng_state_list = None - if args.data_parallel_random_init and torch.distributed.is_initialized() and \ - mpu.get_data_parallel_world_size() > 1: - rng_state_list = \ - [None for i in range(mpu.get_data_parallel_world_size())] + if ( + args.data_parallel_random_init + and torch.distributed.is_initialized() + and mpu.get_data_parallel_world_size() > 1 + ): + rng_state_list = [None for i in range(mpu.get_data_parallel_world_size())] torch.distributed.all_gather_object( - rng_state_list, - rng_state, - group=mpu.get_data_parallel_group()) + rng_state_list, rng_state, group=mpu.get_data_parallel_group() + ) else: rng_state_list = [rng_state] @@ -401,17 +456,21 @@ def get_rng_state(ckpt_format: str, tp_group: torch.distributed.ProcessGroup, pp pp_size = get_pg_size(pp_group) tp_rank = get_pg_rank(tp_group) tp_size = get_pg_size(tp_group) - rng_state_list = ShardedObject('rng_state', rng_state_list, (pp_size, tp_size), (pp_rank, tp_rank), - replica_id=mpu.get_data_parallel_rank(with_context_parallel=True)) + rng_state_list = ShardedObject( + "rng_state", + rng_state_list, + (pp_size, tp_size), + (pp_rank, tp_rank), + replica_id=mpu.get_data_parallel_rank(with_context_parallel=True), + ) elif ckpt_format == "fsdp_dtensor": pp_rank = mpu.get_pipeline_model_parallel_rank() tp_rank = mpu.get_tensor_model_parallel_rank() - rng_state_list = { - f"({pp_rank}, {tp_rank})": rng_state_list - } + rng_state_list = {f"({pp_rank}, {tp_rank})": rng_state_list} return rng_state_list + class CheckpointType(Enum): LEGACY = auto() LOCAL = auto() @@ -420,7 +479,9 @@ class CheckpointType(Enum): FSDP_DTENSOR = auto() -def _build_sharded_state_dict_metadata(args: Namespace, dp_cp_group: Optional[torch.distributed.ProcessGroup] = None) -> dict: +def _build_sharded_state_dict_metadata( + args: Namespace, dp_cp_group: Optional[torch.distributed.ProcessGroup] = None +) -> dict: """Builds metadata used for sharded_state_dict versioning. The whole content metadata is passed to ``shared_state_dict`` model and optimizer methods @@ -438,21 +499,23 @@ def _build_sharded_state_dict_metadata(args: Namespace, dp_cp_group: Optional[to metadata = {} if args.use_distributed_optimizer and args.ckpt_format == "fsdp_dtensor": - metadata['distrib_optim_sharding_type'] = 'fsdp_dtensor' + metadata["distrib_optim_sharding_type"] = "fsdp_dtensor" if args.use_distributed_optimizer and args.ckpt_format != "fsdp_dtensor": if args.dist_ckpt_optim_fully_reshardable: - metadata['distrib_optim_sharding_type'] = 'fully_reshardable' - metadata['distrib_optim_fully_reshardable_mem_efficient'] = args.distrib_optim_fully_reshardable_mem_efficient + metadata["distrib_optim_sharding_type"] = "fully_reshardable" + metadata["distrib_optim_fully_reshardable_mem_efficient"] = ( + args.distrib_optim_fully_reshardable_mem_efficient + ) else: - metadata['distrib_optim_sharding_type'] = 'dp_reshardable' + metadata["distrib_optim_sharding_type"] = "dp_reshardable" - metadata['singleton_local_shards'] = False - metadata['chained_optim_avoid_prefix'] = True + metadata["singleton_local_shards"] = False + metadata["chained_optim_avoid_prefix"] = True # Add dp_cp_group to metadata. If not provided, fallback to global parallel state. if dp_cp_group is None: dp_cp_group = mpu.get_data_parallel_group(with_context_parallel=True) - metadata['dp_cp_group'] = dp_cp_group + metadata["dp_cp_group"] = dp_cp_group return metadata @@ -463,8 +526,10 @@ def save_grads(save_dir, state_dict, iteration, grad_label): NOTE: wgrads for non-expert layers will be duplicated if using expert parallelism, but this can be handled in postprocessing.""" - print_rank_0(f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saving {grad_label} " - f"from iteration {iteration:7d}") + print_rank_0( + f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saving {grad_label} " + f"from iteration {iteration:7d}" + ) if mpu.get_expert_data_parallel_rank() == 0: # Create saving directory. @@ -486,13 +551,32 @@ def save_grads(save_dir, state_dict, iteration, grad_label): # Convert back to dict (e.g., from collections.defaultdict) for easy loading later. torch.save(dict(state_dict), full_save_path) - print_rank_0(f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saved {grad_label} " - f"from iteration {iteration:7d}") + print_rank_0( + f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saved {grad_label} " + f"from iteration {iteration:7d}" + ) -def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floating_point_operations_so_far, - checkpointing_context=None, pipeline_rank=None, expert_rank=None, tensor_rank=None, pipeline_parallel=None, expert_parallel=None, non_persistent_ckpt=False, - train_data_iterator=None, preprocess_common_state_dict_fn = None, release=False, tp_group: Optional[torch.distributed.ProcessGroup] = None, pp_group: Optional[torch.distributed.ProcessGroup] = None, dp_cp_group: Optional[torch.distributed.ProcessGroup] = None): +def save_checkpoint( + iteration, + model, + optimizer, + opt_param_scheduler, + num_floating_point_operations_so_far, + checkpointing_context=None, + pipeline_rank=None, + expert_rank=None, + tensor_rank=None, + pipeline_parallel=None, + expert_parallel=None, + non_persistent_ckpt=False, + train_data_iterator=None, + preprocess_common_state_dict_fn=None, + release=False, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + pp_group: Optional[torch.distributed.ProcessGroup] = None, + dp_cp_group: Optional[torch.distributed.ProcessGroup] = None, +): """Save a model, optimizer and optionally dataloader checkpoint. Checkpointing context is used to persist some checkpointing state @@ -514,7 +598,9 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati args = get_args() if args.async_save and not is_empty_async_queue(): - print_rank_0('WARNING: Starting a checkpoint save before previous has finished. Consider increasing the checkpoint interval.') + print_rank_0( + "WARNING: Starting a checkpoint save before previous has finished. Consider increasing the checkpoint interval." + ) # Prepare E2E metrics at start of save checkpoint productive_metrics = on_save_checkpoint_start(args.async_save) @@ -530,7 +616,7 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati ckpt_type = CheckpointType.GLOBAL if args.use_dist_ckpt else CheckpointType.LEGACY save_dir = args.save if non_persistent_ckpt: - if args.non_persistent_ckpt_type == 'global': + if args.non_persistent_ckpt_type == "global": ckpt_type = CheckpointType.GLOBAL save_dir = ( args.non_persistent_global_ckpt_dir @@ -541,15 +627,19 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati cleanup_old_non_persistent_checkpoint( save_dir, leave_ckpt_num=1, do_async=args.async_save ) - elif args.non_persistent_ckpt_type == 'local': + elif args.non_persistent_ckpt_type == "local": ckpt_type = CheckpointType.LOCAL - save_dir = checkpointing_context['local_checkpoint_manager'].local_ckpt_dir + save_dir = checkpointing_context["local_checkpoint_manager"].local_ckpt_dir else: - raise NotImplementedError(f"Please use local or global non-persistent checkpoints (got: {args.non_persistent_ckpt_type})") + raise NotImplementedError( + f"Please use local or global non-persistent checkpoints (got: {args.non_persistent_ckpt_type})" + ) - ckpt_format = args.ckpt_format if ckpt_type == CheckpointType.GLOBAL else 'torch' - print_rank_0(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saving checkpoint " - f"at iteration {iteration:7d} to {save_dir} in {ckpt_format} format") + ckpt_format = args.ckpt_format if ckpt_type == CheckpointType.GLOBAL else "torch" + print_rank_0( + f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] saving checkpoint " + f"at iteration {iteration:7d} to {save_dir} in {ckpt_format} format" + ) # Collect rng state across data parallel ranks. if tp_group is None and pp_group is None: @@ -560,16 +650,28 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati # Collect rerun state across all ranks rerun_state_machine = get_rerun_state_machine() rerun_state = rerun_state_machine.state_dict( - data_iterator=train_data_iterator, ckpt_format=args.ckpt_format, + data_iterator=train_data_iterator, + ckpt_format=args.ckpt_format, ) # Checkpoint name. - return_base_dir = (ckpt_type != CheckpointType.LEGACY) - checkpoint_name = get_checkpoint_name(save_dir, iteration, release=release, pipeline_parallel=pipeline_parallel, - tensor_rank=tensor_rank, pipeline_rank=pipeline_rank, expert_parallel=expert_parallel, expert_rank=expert_rank, return_base_dir=return_base_dir) + return_base_dir = ckpt_type != CheckpointType.LEGACY + checkpoint_name = get_checkpoint_name( + save_dir, + iteration, + release=release, + pipeline_parallel=pipeline_parallel, + tensor_rank=tensor_rank, + pipeline_rank=pipeline_rank, + expert_parallel=expert_parallel, + expert_rank=expert_rank, + return_base_dir=return_base_dir, + ) # Save dataloader state if the dataloader supports it (currently only Megatron Energon). - maybe_save_dataloader_state(train_data_iterator, iteration, getattr(args, "dataloader_save", None)) + maybe_save_dataloader_state( + train_data_iterator, iteration, getattr(args, "dataloader_save", None) + ) # Save distributed optimizer's custom parameter state. if ( @@ -578,16 +680,17 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati and optimizer is not None and ckpt_type == CheckpointType.LEGACY ): - optim_checkpoint_name = \ - get_distributed_optimizer_checkpoint_name(checkpoint_name) + optim_checkpoint_name = get_distributed_optimizer_checkpoint_name(checkpoint_name) ensure_directory_exists(optim_checkpoint_name) if not optimizer.is_stub_optimizer: optimizer.save_parameter_state(optim_checkpoint_name) # LayerWiseDistributedOptimizer save optimizer state to file on different ranks - if getattr(args, "optimizer", "adam").startswith("dist_") and args.ckpt_format == 'torch': + if getattr(args, "optimizer", "adam").startswith("dist_") and args.ckpt_format == "torch": dp_rank = mpu.get_data_parallel_rank() - optim_checkpoint_name = os.path.join(os.path.dirname(checkpoint_name), f"layer_wise_optimizer_{dp_rank}.pt") + optim_checkpoint_name = os.path.join( + os.path.dirname(checkpoint_name), f"layer_wise_optimizer_{dp_rank}.pt" + ) ensure_directory_exists(optim_checkpoint_name) if not optimizer.is_stub_optimizer: optimizer.save_state_dict_to_file(optim_checkpoint_name) @@ -595,9 +698,17 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati async_save_request = None if args.async_save: if ckpt_type == CheckpointType.LEGACY: - raise NotImplementedError('Async checkpoint save not implemented for legacy checkpoints') - elif ckpt_type == CheckpointType.GLOBAL and args.ckpt_format not in ['torch_dist', 'torch_dcp', 'fsdp_dtensor']: - raise NotImplementedError(f'Async checkpoint save not implemented for {args.ckpt_format} distributed checkpoint format') + raise NotImplementedError( + "Async checkpoint save not implemented for legacy checkpoints" + ) + elif ckpt_type == CheckpointType.GLOBAL and args.ckpt_format not in [ + "torch_dist", + "torch_dcp", + "fsdp_dtensor", + ]: + raise NotImplementedError( + f"Async checkpoint save not implemented for {args.ckpt_format} distributed checkpoint format" + ) rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 @@ -606,15 +717,19 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati # exactly one rank. Neither dp_rank==0 nor edp_rank==0 alone covers all shards when # the dense and expert parallelism layouts disagree (e.g. TP > EP*ETP); the union # does, with at most one rank per (tp_rank, ep_rank) inside any DP group. - if not torch.distributed.is_initialized() \ - or mpu.get_data_parallel_rank() == 0 \ - or mpu.get_expert_data_parallel_rank() == 0 \ - or ckpt_type != CheckpointType.LEGACY: + if ( + not torch.distributed.is_initialized() + or mpu.get_data_parallel_rank() == 0 + or mpu.get_expert_data_parallel_rank() == 0 + or ckpt_type != CheckpointType.LEGACY + ): if ckpt_type != CheckpointType.LEGACY: sharded_sd_metadata = _build_sharded_state_dict_metadata(args, dp_cp_group=dp_cp_group) if args.use_distributed_optimizer: - print_rank_0(f'Storing distributed optimizer sharded state of type' - f' {sharded_sd_metadata["distrib_optim_sharding_type"]}') + print_rank_0( + f"Storing distributed optimizer sharded state of type" + f" {sharded_sd_metadata['distrib_optim_sharding_type']}" + ) else: sharded_sd_metadata = None state_dict = generate_state_dict( @@ -629,53 +744,69 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati rerun_state=rerun_state, ) - state_dict['num_floating_point_operations_so_far'] = num_floating_point_operations_so_far + state_dict["num_floating_point_operations_so_far"] = num_floating_point_operations_so_far if ckpt_type == CheckpointType.GLOBAL and ckpt_format == "torch_dist": if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: # TODO Handle non-empty directories (e.g., after a crash during saving). ensure_directory_exists(checkpoint_name, check_parent=False) - if checkpointing_context is not None and 'save_strategy' in checkpointing_context: - save_strategy = checkpointing_context['save_strategy'] + if checkpointing_context is not None and "save_strategy" in checkpointing_context: + save_strategy = checkpointing_context["save_strategy"] # Already saved once before - don't need to rerun sharding validation validate_sharding_integrity = not args.ckpt_assume_constant_structure else: validate_sharding_integrity = True save_strategy = TorchDistSaveShardedStrategy() - if args.ckpt_assume_constant_structure and args.ckpt_format == 'torch_dist': + if args.ckpt_assume_constant_structure and args.ckpt_format == "torch_dist": save_strategy.use_cached_ckpt_structure = args.ckpt_assume_constant_structure if args.async_save: save_strategy.thread_count = args.dist_ckpt_workers else: # We don't allow per-rank parallel save for sync save - logger.warning('Per-rank parallel save is not supported for sync save. ' - 'Setting args.dist_ckpt_workers to 1') + logger.warning( + "Per-rank parallel save is not supported for sync save. " + "Setting args.dist_ckpt_workers to 1" + ) save_strategy.thread_count = 1 - if checkpointing_context is not None and 'load_strategy' in checkpointing_context: - cached_global_metadata = getattr(checkpointing_context['load_strategy'], 'cached_global_metadata', None) + if ( + checkpointing_context is not None + and "load_strategy" in checkpointing_context + ): + cached_global_metadata = getattr( + checkpointing_context["load_strategy"], "cached_global_metadata", None + ) if cached_global_metadata is not None: logger.debug("Plugging in the read metadata from the load strategy...") save_strategy.cached_global_metadata = cached_global_metadata else: - logger.debug("Failed to plug in the read metadata from the load strategy...") + logger.debug( + "Failed to plug in the read metadata from the load strategy..." + ) if args.ckpt_fully_parallel_save: - if args.ckpt_fully_parallel_save_process_group == 'dp': + if args.ckpt_fully_parallel_save_process_group == "dp": process_group = mpu.get_data_parallel_group(with_context_parallel=True) - elif args.ckpt_fully_parallel_save_process_group == 'ep_dp': + elif args.ckpt_fully_parallel_save_process_group == "ep_dp": process_group = mpu.get_expert_data_parallel_group() - save_strategy = FullyParallelSaveStrategyWrapper(save_strategy, process_group, - args.ckpt_assume_constant_structure) + save_strategy = FullyParallelSaveStrategyWrapper( + save_strategy, process_group, args.ckpt_assume_constant_structure + ) # Store save strategy for future checkpoint saves if checkpointing_context is not None: - checkpointing_context['save_strategy'] = save_strategy + checkpointing_context["save_strategy"] = save_strategy end_ckpt = time() - logger.debug(f"rank: {rank}, takes {end_ckpt - start_ckpt} to prepare state dict for ckpt ") - async_save_request = dist_checkpointing.save(state_dict, checkpoint_name, save_strategy, - async_sharded_save=args.async_save, - validate_access_integrity=validate_sharding_integrity, - preprocess_common_before_consistancy_check=preprocess_common_state_dict_fn, - content_metadata=_clean_metadata_for_serialization(sharded_sd_metadata), - async_strategy=args.async_strategy) + logger.debug( + f"rank: {rank}, takes {end_ckpt - start_ckpt} to prepare state dict for ckpt " + ) + async_save_request = dist_checkpointing.save( + state_dict, + checkpoint_name, + save_strategy, + async_sharded_save=args.async_save, + validate_access_integrity=validate_sharding_integrity, + preprocess_common_before_consistancy_check=preprocess_common_state_dict_fn, + content_metadata=_clean_metadata_for_serialization(sharded_sd_metadata), + async_strategy=args.async_strategy, + ) # [ModelOpt]: save sharded modelopt_state if has_nvidia_modelopt: save_sharded_modelopt_state(model, checkpoint_name, (args.ckpt_format, 1)) @@ -695,9 +826,16 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati ) save_state_dict_ret = save_state_dict_async_plan( - state_dict, fs_storage_writer, None, coordinator_rank, planner=planner, enable_cache=args.ckpt_assume_constant_structure + state_dict, + fs_storage_writer, + None, + coordinator_rank, + planner=planner, + enable_cache=args.ckpt_assume_constant_structure, + ) + async_save_request = get_save_and_finalize_callbacks( + fs_storage_writer, save_state_dict_ret ) - async_save_request = get_save_and_finalize_callbacks(fs_storage_writer, save_state_dict_ret) else: fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter(checkpoint_name) torch.distributed.checkpoint.save( @@ -708,19 +846,28 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati # [ModelOpt]: Inject modelopt_state into state_dict if has_nvidia_modelopt: if ckpt_type == CheckpointType.LOCAL: - print_rank_0('WARNING: Local checkpointing does not support nvidia_modelopt.') + print_rank_0("WARNING: Local checkpointing does not support nvidia_modelopt.") else: save_modelopt_state(model, state_dict) end_ckpt = time() - logger.debug(f"rank: {rank}, takes {end_ckpt - start_ckpt} to prepare state dict for ckpt ") + logger.debug( + f"rank: {rank}, takes {end_ckpt - start_ckpt} to prepare state dict for ckpt " + ) if ckpt_type == CheckpointType.LOCAL: try: - from megatron.core.dist_checkpointing.tensor_aware_state_dict import MCoreTensorAwareStateDict + from megatron.core.dist_checkpointing.tensor_aware_state_dict import ( + MCoreTensorAwareStateDict, + ) except ModuleNotFoundError: - raise RuntimeError("The 'nvidia_resiliency_ext' module is required for local " - "checkpointing but was not found. Please ensure it is installed.") - if (sharded_sd_metadata or {}).get('distrib_optim_sharding_type') in ['fully_reshardable', 'dp_zero_gather_scatter']: + raise RuntimeError( + "The 'nvidia_resiliency_ext' module is required for local " + "checkpointing but was not found. Please ensure it is installed." + ) + if (sharded_sd_metadata or {}).get("distrib_optim_sharding_type") in [ + "fully_reshardable", + "dp_zero_gather_scatter", + ]: # Note: Currently full reshardabilty is not supported when local checkpoints are used. raise RuntimeError( f"Local checkpointing does not support optimizer sharding type '{sharded_sd_metadata['distrib_optim_sharding_type']}'. " @@ -728,16 +875,21 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati ) algo = args.non_persistent_local_ckpt_algo cached_metadata = None - if args.ckpt_assume_constant_structure and 'local_checkpoint_cache' in checkpointing_context: - cached_metadata = checkpointing_context['local_checkpoint_cache'] + if ( + args.ckpt_assume_constant_structure + and "local_checkpoint_cache" in checkpointing_context + ): + cached_metadata = checkpointing_context["local_checkpoint_cache"] state_dict_for_save, cacheable_metadata = MCoreTensorAwareStateDict.from_state_dict( - state_dict, algo=algo, cached_metadata=cached_metadata, - parallelization_group=mpu.get_data_parallel_group(with_context_parallel=True) + state_dict, + algo=algo, + cached_metadata=cached_metadata, + parallelization_group=mpu.get_data_parallel_group(with_context_parallel=True), ) - async_save_request = checkpointing_context['local_checkpoint_manager'].save( + async_save_request = checkpointing_context["local_checkpoint_manager"].save( state_dict_for_save, iteration, is_async=bool(args.async_save) ) - checkpointing_context['local_checkpoint_cache'] = cacheable_metadata + checkpointing_context["local_checkpoint_cache"] = cacheable_metadata else: assert ckpt_type == CheckpointType.LEGACY # Save. @@ -752,63 +904,98 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati torch.distributed.barrier() # And update the latest iteration - if not torch.distributed.is_initialized() \ - or torch.distributed.get_rank() == 0: + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: tracker_filename = get_checkpoint_tracker_filename(save_dir) if ckpt_type == CheckpointType.LOCAL: + def iter_finalize_fn(): - print_rank_0(f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] successfully " - f"saved local checkpoint from iteration {iteration:7d}") + print_rank_0( + f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] successfully " + f"saved local checkpoint from iteration {iteration:7d}" + ) if args.log_progress and args.async_save: - append_to_progress_log(f'Saved async local checkpoint\tIteration: {iteration}', - barrier=False) + append_to_progress_log( + f"Saved async local checkpoint\tIteration: {iteration}", barrier=False + ) else: + def iter_finalize_fn(): prev_iteration = 0 - save_retain_interval = getattr(args, 'save_retain_interval', None) # For backwards compatibility of tests. + save_retain_interval = getattr( + args, "save_retain_interval", None + ) # For backwards compatibility of tests. if save_retain_interval is not None: - if os.path.exists(tracker_filename): # TODO: Make this work with MSC remote paths? - with open_file(tracker_filename, 'r') as f: + if os.path.exists( + tracker_filename + ): # TODO: Make this work with MSC remote paths? + with open_file(tracker_filename, "r") as f: prev_iteration = int(f.read().strip()) - with open_file(tracker_filename, 'w') as f: + with open_file(tracker_filename, "w") as f: f.write("release" if release else str(iteration)) - tensor_rank_to_print = (tensor_rank if tensor_rank is not None else mpu.get_tensor_model_parallel_rank()) + 1 - pipeline_rank_to_print = (pipeline_rank if pipeline_rank is not None else mpu.get_pipeline_model_parallel_rank()) + 1 - print_rank_0(f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] successfully saved " - f"checkpoint from iteration {int(iteration):7d} to {args.save} " - f"[ t {tensor_rank_to_print}/{mpu.get_tensor_model_parallel_world_size()}, " - f"p {pipeline_rank_to_print}/{mpu.get_pipeline_model_parallel_world_size()} ]") + tensor_rank_to_print = ( + tensor_rank if tensor_rank is not None else mpu.get_tensor_model_parallel_rank() + ) + 1 + pipeline_rank_to_print = ( + pipeline_rank + if pipeline_rank is not None + else mpu.get_pipeline_model_parallel_rank() + ) + 1 + print_rank_0( + f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] successfully saved " + f"checkpoint from iteration {int(iteration):7d} to {args.save} " + f"[ t {tensor_rank_to_print}/{mpu.get_tensor_model_parallel_world_size()}, " + f"p {pipeline_rank_to_print}/{mpu.get_pipeline_model_parallel_world_size()} ]" + ) if args.log_progress and args.async_save: - append_to_progress_log(f'Saved async checkpoint\tIteration: {iteration}', - barrier=False) + append_to_progress_log( + f"Saved async checkpoint\tIteration: {iteration}", barrier=False + ) if save_retain_interval is not None: - if prev_iteration > 0 and prev_iteration != iteration and prev_iteration % save_retain_interval != 0: - checkpoint_name = get_checkpoint_name(args.save, iteration=prev_iteration, - return_base_dir=True) + if ( + prev_iteration > 0 + and prev_iteration != iteration + and prev_iteration % save_retain_interval != 0 + ): + checkpoint_name = get_checkpoint_name( + args.save, iteration=prev_iteration, return_base_dir=True + ) # Don't delete if `checkpoint_name` is a symbolic link. - if os.path.islink(checkpoint_name): # TODO: Make this work with MSC remote paths? - print_rank_0(f' skipping deleting checkpoint from iteration {prev_iteration:7d} ' - f'at {args.save} since it is a symbolic link') + if os.path.islink( + checkpoint_name + ): # TODO: Make this work with MSC remote paths? + print_rank_0( + f" skipping deleting checkpoint from iteration {prev_iteration:7d} " + f"at {args.save} since it is a symbolic link" + ) else: # Asynchronous version of delete_checkpoint(args, iteration_to_delete=prev_iteration). # Use multiprocessing to delete checkpoint in background if args.async_save: # Clean up any finished deletion processes before starting a new one finalize_deletion_processes(blocking=False) - ctx = multiprocessing.get_context('fork') + ctx = multiprocessing.get_context("fork") delete_process = ctx.Process( target=_async_delete_checkpoint_impl, - args=(args.save, prev_iteration, args.log_progress, True, - args.async_ckpt_cpu_priority, args.async_ckpt_io_priority), - daemon=True + args=( + args.save, + prev_iteration, + args.log_progress, + True, + args.async_ckpt_cpu_priority, + args.async_ckpt_io_priority, + ), + daemon=True, ) delete_process.start() # Track the process so we can join it later to prevent zombies _deletion_processes.append(delete_process) else: - th = threading.Thread(target=_async_delete_checkpoint_impl, args=(args.save, prev_iteration, args.log_progress)) + th = threading.Thread( + target=_async_delete_checkpoint_impl, + args=(args.save, prev_iteration, args.log_progress), + ) th.start() if args.async_save: @@ -818,10 +1005,11 @@ def iter_finalize_fn(): iter_finalize_fn() # Additional callback for one_logger (last rank) - if not torch.distributed.is_initialized() \ - or is_last_rank(): + if not torch.distributed.is_initialized() or is_last_rank(): + def onelogger_finalize_fn(): on_save_checkpoint_success(productive_metrics, args.async_save) + if args.async_save: assert async_save_request is not None async_save_request.add_finalize_fn(onelogger_finalize_fn) @@ -829,11 +1017,11 @@ def onelogger_finalize_fn(): onelogger_finalize_fn() # Additional callback for wandb (last rank) - if not torch.distributed.is_initialized() \ - or is_last_rank(): + if not torch.distributed.is_initialized() or is_last_rank(): + def wandb_finalize_fn(): ######### FlagScale Begin ######### - #NOTE(lizhiyu): The tracker file is created by rank 0 but wandb_finalize_fn is called on the last rank. + # NOTE(lizhiyu): The tracker file is created by rank 0 but wandb_finalize_fn is called on the last rank. import time as pytime tracker_file = get_checkpoint_tracker_filename(save_dir) @@ -844,17 +1032,24 @@ def wandb_finalize_fn(): for _ in range(max_retries): if isfile(tracker_file): - with open(tracker_file, 'r') as f: + with open(tracker_file, "r") as f: content = f.read().strip() if content == str(iteration): break # Success - print(f'WandB finalization waiting for the tracker file {tracker_file} to update...') + print( + f"WandB finalization waiting for the tracker file {tracker_file} to update..." + ) pytime.sleep(wait_interval_seconds) else: # This block executes if the loop completes without a `break`. - raise RuntimeError(f"Timed out waiting for tracker file {tracker_file} to be updated for iteration {iteration} after {timeout_seconds} seconds.") + raise RuntimeError( + f"Timed out waiting for tracker file {tracker_file} to be updated for iteration {iteration} after {timeout_seconds} seconds." + ) ######### FlagScale End ######### - wandb_utils.on_save_checkpoint_success(checkpoint_name, get_checkpoint_tracker_filename(save_dir), save_dir, iteration) + wandb_utils.on_save_checkpoint_success( + checkpoint_name, get_checkpoint_tracker_filename(save_dir), save_dir, iteration + ) + if args.async_save: assert async_save_request is not None async_save_request.add_finalize_fn(wandb_finalize_fn) @@ -863,8 +1058,10 @@ def wandb_finalize_fn(): if args.async_save: schedule_async_save(async_save_request) - print_rank_0(f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] scheduled " - f"an async checkpoint save at iteration {iteration:7d} to {save_dir}") + print_rank_0( + f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}] scheduled " + f"an async checkpoint save at iteration {iteration:7d} to {save_dir}" + ) # Wait so everyone is done (not necessary) if torch.distributed.is_initialized(): @@ -875,15 +1072,22 @@ def wandb_finalize_fn(): ft_integration.on_checkpointing_end(is_async_finalization=False) + @_disable_gc() -def _async_delete_checkpoint_impl(save_path, iteration_to_delete, log_progress=False, lower_priority=False, - cpu_priority=None, io_priority=None): +def _async_delete_checkpoint_impl( + save_path, + iteration_to_delete, + log_progress=False, + lower_priority=False, + cpu_priority=None, + io_priority=None, +): """Module-level function for async checkpoint deletion. - + This function can be pickled and executed by the async worker process. Note: This is only called from rank 0, so we use regular print() instead of print_rank_0() since torch.distributed won't be initialized in the async worker process. - + Args: save_path (str): Path to the checkpoints directory iteration_to_delete (int): Iteration number of checkpoint to delete @@ -894,19 +1098,29 @@ def _async_delete_checkpoint_impl(save_path, iteration_to_delete, log_progress=F """ if lower_priority: from megatron.core.dist_checkpointing.strategies.async_utils import _set_process_qos + _set_process_qos(cpu_priority=cpu_priority, io_priority=io_priority) - checkpoint_name = get_checkpoint_name(save_path, iteration=iteration_to_delete, - return_base_dir=True) + checkpoint_name = get_checkpoint_name( + save_path, iteration=iteration_to_delete, return_base_dir=True + ) try: shutil.rmtree(checkpoint_name) # TODO: Make this work with MSC remote paths? - print(f' successfully deleted checkpoint from iteration {iteration_to_delete:7d} ' - f'at {save_path}', flush=True) + print( + f" successfully deleted checkpoint from iteration {iteration_to_delete:7d} " + f"at {save_path}", + flush=True, + ) if log_progress: - append_to_progress_log(f'Deleted checkpoint\tIteration: {iteration_to_delete}', barrier=False) + append_to_progress_log( + f"Deleted checkpoint\tIteration: {iteration_to_delete}", barrier=False + ) except Exception as e: - print(f' encountered exception "{e}" when trying to delete checkpoint from ' - f'iteration {iteration_to_delete:7d} at {save_path}', flush=True) + print( + f' encountered exception "{e}" when trying to delete checkpoint from ' + f"iteration {iteration_to_delete:7d} at {save_path}", + flush=True, + ) # Any exception encountered in checkpoint deletion can be ignored and is not fatal. pass @@ -917,17 +1131,20 @@ def cleanup_old_non_persistent_checkpoint(save_dir, leave_ckpt_num=1, do_async=F save_dir = Path(save_dir) iter_prefix = "iter_" - iter_ckpts = save_dir.rglob(f'{iter_prefix}*') - sorted_iter_ckpts = sorted(iter_ckpts, key=lambda ckpt_name: int(ckpt_name.name[len(iter_prefix):])) + iter_ckpts = save_dir.rglob(f"{iter_prefix}*") + sorted_iter_ckpts = sorted( + iter_ckpts, key=lambda ckpt_name: int(ckpt_name.name[len(iter_prefix) :]) + ) if not sorted_iter_ckpts: return rm_iter_ckpts = sorted_iter_ckpts[:-leave_ckpt_num] - print_rank_0(f'Non-persistent checkpoints scheduled for removal: {rm_iter_ckpts}') - print_rank_0(f'Non-persistent checkpoints to be kept: {sorted_iter_ckpts[-leave_ckpt_num:]}') + print_rank_0(f"Non-persistent checkpoints scheduled for removal: {rm_iter_ckpts}") + print_rank_0(f"Non-persistent checkpoints to be kept: {sorted_iter_ckpts[-leave_ckpt_num:]}") def remove_iter_ckpts(_iter_ckpts): for ckpt in _iter_ckpts: shutil.rmtree(ckpt) + if do_async: threading.Thread(target=remove_iter_ckpts, args=(rm_iter_ckpts,)).start() else: @@ -955,10 +1172,15 @@ def maybe_save_dataloader_state(train_iterator, iteration, dataloader_save_path) # If dataloader doesn't support saving state, raise an error. if not hasattr(train_iterator.iterable, "save_state"): - raise RuntimeError(f"Could not find a save_state for the train_iterator of type {type(train_iterator)}") + raise RuntimeError( + f"Could not find a save_state for the train_iterator of type {type(train_iterator)}" + ) # Save dataloader state for each data parallel rank only once. - first_rank = mpu.is_pipeline_first_stage(ignore_virtual=True) and mpu.get_tensor_model_parallel_rank() == 0 + first_rank = ( + mpu.is_pipeline_first_stage(ignore_virtual=True) + and mpu.get_tensor_model_parallel_rank() == 0 + ) if not first_rank: return @@ -967,8 +1189,7 @@ def maybe_save_dataloader_state(train_iterator, iteration, dataloader_save_path) print(f"saving dataloader checkpoint at iteration {iteration} to {dataloader_save_path}") train_dataloader_state_dict = train_iterator.iterable.save_state() data_state_save_path = get_checkpoint_name( - dataloader_save_path, iteration, - basename=f'train_dataloader_dprank{dp_rank:03d}.pt' + dataloader_save_path, iteration, basename=f"train_dataloader_dprank{dp_rank:03d}.pt" ) torch.distributed.barrier(group=mpu.get_data_parallel_group()) @@ -978,7 +1199,7 @@ def maybe_save_dataloader_state(train_iterator, iteration, dataloader_save_path) torch.distributed.barrier(group=mpu.get_data_parallel_group()) dataloader_save_dict = {} - dataloader_save_dict['dataloader_state_dict'] = train_dataloader_state_dict + dataloader_save_dict["dataloader_state_dict"] = train_dataloader_state_dict torch.save(dataloader_save_dict, data_state_save_path) @@ -993,14 +1214,14 @@ def generate_state_dict( model_sd_kwargs=None, rerun_state=None, ): - """Generate a state dict from given model, optimizer, scheduler, rng state and others. """ + """Generate a state dict from given model, optimizer, scheduler, rng state and others.""" # Arguments, iteration, and model. state_dict = {} - state_dict['args'] = args - state_dict['checkpoint_version'] = 3.0 + state_dict["args"] = args + state_dict["checkpoint_version"] = 3.0 if iteration is not None: - state_dict['iteration'] = iteration + state_dict["iteration"] = iteration for i in range(len(model)): key = "model" @@ -1009,13 +1230,16 @@ def generate_state_dict( if args.ckpt_format == "torch_dist": model_sd = model[i].sharded_state_dict( - **(model_sd_kwargs or { - "metadata": { - "dp_cp_group": mpu.get_data_parallel_group(with_context_parallel=True) + **( + model_sd_kwargs + or { + "metadata": { + "dp_cp_group": mpu.get_data_parallel_group(with_context_parallel=True) + } } - }) + ) ) - else: # torch, torch_dcp, fsdp_dtensor + else: # torch, torch_dcp, fsdp_dtensor model_sd = model[i].state_dict_for_save_checkpoint() state_dict[key] = model_sd @@ -1023,35 +1247,38 @@ def generate_state_dict( # Optimizer stuff. if not args.no_save_optim: if optimizer is not None and not optimizer.is_stub_optimizer: - if args.ckpt_format == "torch_dist": optimizer_sd = optimizer.sharded_state_dict( state_dict, - **(optim_sd_kwargs or { - "metadata": { - "dp_cp_group": mpu.get_data_parallel_group(with_context_parallel=True) + **( + optim_sd_kwargs + or { + "metadata": { + "dp_cp_group": mpu.get_data_parallel_group( + with_context_parallel=True + ) + } } - }) + ), ) elif args.ckpt_format == "fsdp_dtensor": if optim_sd_kwargs is None: optim_sd_kwargs = {} if "metadata" not in optim_sd_kwargs: optim_sd_kwargs["metadata"] = {} - optim_sd_kwargs['metadata'].update(_build_sharded_state_dict_metadata(args)) + optim_sd_kwargs["metadata"].update(_build_sharded_state_dict_metadata(args)) optimizer_sd = optimizer.sharded_state_dict(state_dict, **optim_sd_kwargs) else: optimizer_sd = optimizer.state_dict() - state_dict['optimizer'] = optimizer_sd + state_dict["optimizer"] = optimizer_sd if opt_param_scheduler is not None: - state_dict['opt_param_scheduler'] = \ - opt_param_scheduler.state_dict() + state_dict["opt_param_scheduler"] = opt_param_scheduler.state_dict() # Rerun state if rerun_state: - state_dict['rerun_state_machine'] = rerun_state + state_dict["rerun_state_machine"] = rerun_state # RNG states. if not args.no_save_rng and rng_state: @@ -1071,9 +1298,7 @@ def preprocess_fsdp_dtensor_state_dict(args, raw_state_dict, model): state_dict["model"] = model_state_dict state_dict["optimizer"] = optimizer_state_dict else: - model_state_dict, _ = handle_swiglu_in_state_dict( - model, state_dict["model"], None - ) + model_state_dict, _ = handle_swiglu_in_state_dict(model, state_dict["model"], None) state_dict["model"] = model_state_dict if args.num_experts: state_dict["model"] = handle_experts_in_state_dict(state_dict["model"], args.num_experts) @@ -1086,7 +1311,7 @@ def _transpose_first_dim(t, num_splits, num_splits_first, model): input_shape = t.size() # We use a self_attention module but the values extracted aren't # specific to self attention so should work for cross attention as well - while hasattr(model, 'module'): + while hasattr(model, "module"): model = model.module attention_module = model.language_model.encoder.layers[0].self_attention hidden_size_per_attention_head = attention_module.hidden_size_per_attention_head @@ -1097,9 +1322,11 @@ def _transpose_first_dim(t, num_splits, num_splits_first, model): -->(tranpose) [np, num_splits, hn, h] -->(view) [np * num_splits * hn, h] """ - intermediate_shape = \ - (num_splits, num_attention_heads_per_partition, - hidden_size_per_attention_head) + input_shape[1:] + intermediate_shape = ( + num_splits, + num_attention_heads_per_partition, + hidden_size_per_attention_head, + ) + input_shape[1:] t = t.view(*intermediate_shape) t = t.transpose(0, 1).contiguous() @@ -1109,10 +1336,11 @@ def _transpose_first_dim(t, num_splits, num_splits_first, model): -->(tranpose) [np, num_splits, hn, h] -->(view) [np * num_splits * hn, h] """ - intermediate_shape = \ - (num_attention_heads_per_partition, - hidden_size_per_attention_head, num_splits) +\ - input_shape[1:] + intermediate_shape = ( + num_attention_heads_per_partition, + hidden_size_per_attention_head, + num_splits, + ) + input_shape[1:] t = t.view(*intermediate_shape) t = t.transpose(1, 2).contiguous() @@ -1127,10 +1355,10 @@ def fix_query_key_value_ordering(model, checkpoint_version): """ if checkpoint_version < 2.0: if isinstance(model, list): - assert len(model)==1 + assert len(model) == 1 model = model[0] for name, param in model.named_parameters(): - if name.endswith(('.query_key_value.weight', '.query_key_value.bias')): + if name.endswith((".query_key_value.weight", ".query_key_value.bias")): if checkpoint_version == 0: fixed_param = _transpose_first_dim(param.data, 3, True, model) elif checkpoint_version == 1.0: @@ -1139,7 +1367,7 @@ def fix_query_key_value_ordering(model, checkpoint_version): print_rank_0(f"Invalid checkpoint version {checkpoint_version}.") sys.exit() param.data.copy_(fixed_param) - if name.endswith(('.key_value.weight', '.key_value.bias')): + if name.endswith((".key_value.weight", ".key_value.bias")): if checkpoint_version == 0: fixed_param = _transpose_first_dim(param.data, 2, True, model) elif checkpoint_version == 1.0: @@ -1148,8 +1376,11 @@ def fix_query_key_value_ordering(model, checkpoint_version): print_rank_0(f"Invalid checkpoint version {checkpoint_version}.") sys.exit() param.data.copy_(fixed_param) - print_rank_0(" successfully fixed query-key-values ordering for" - " checkpoint version {}".format(checkpoint_version)) + print_rank_0( + " successfully fixed query-key-values ordering for checkpoint version {}".format( + checkpoint_version + ) + ) def _get_non_persistent_iteration(non_persistent_global_dir, args, checkpointing_context=None): @@ -1160,17 +1391,19 @@ def _get_non_persistent_iteration(non_persistent_global_dir, args, checkpointing if isfile(tracker_filename): iteration, release = read_metadata(tracker_filename) if release: - raise RuntimeError('Non-persistent checkpoint can\'t be a release checkpoint') + raise RuntimeError("Non-persistent checkpoint can't be a release checkpoint") else: iteration = -1 - print_rank_0('WARNING: could not find the metadata file {}'.format(tracker_filename)) - print_rank_0(' will not load any non-persistent checkpoint') + print_rank_0("WARNING: could not find the metadata file {}".format(tracker_filename)) + print_rank_0(" will not load any non-persistent checkpoint") return iteration elif args.non_persistent_ckpt_type == "local": - return checkpointing_context['local_checkpoint_manager'].find_latest() + return checkpointing_context["local_checkpoint_manager"].find_latest() else: - assert False, 'Please use local or global non-persistent checkpoints' \ - f'(got: {args.non_persistent_ckpt_type})' + assert False, ( + "Please use local or global non-persistent checkpoints" + f"(got: {args.non_persistent_ckpt_type})" + ) def _load_non_persistent_base_checkpoint( @@ -1181,37 +1414,44 @@ def _load_non_persistent_base_checkpoint( non_persistent_iteration, checkpointing_context=None, ): - """ Load the base state_dict from a non-persistent distributed checkpoint. + """Load the base state_dict from a non-persistent distributed checkpoint. Depending on the non_persistent_ckpt_type, different logic may be required. """ assert args.non_persistent_ckpt_type is not None if args.non_persistent_ckpt_type == "global": if not rank0: print_rank_0( - f'Loading from a non-persistent checkpoint (non-persistent iter {non_persistent_iteration})' + f"Loading from a non-persistent checkpoint (non-persistent iter {non_persistent_iteration})" ) return _load_global_dist_base_checkpoint( - non_persistent_global_dir, args, rank0, sharded_state_dict, non_persistent_iteration, False, - checkpointing_context=checkpointing_context + non_persistent_global_dir, + args, + rank0, + sharded_state_dict, + non_persistent_iteration, + False, + checkpointing_context=checkpointing_context, ) elif args.non_persistent_ckpt_type == "local": intermediate_state_dict, checkpoint_name = checkpointing_context[ - 'local_checkpoint_manager' + "local_checkpoint_manager" ].load() state_dict = intermediate_state_dict.to_state_dict( sharded_state_dict, algo=args.non_persistent_local_ckpt_algo, - parallelization_group = mpu.get_data_parallel_group(with_context_parallel=True) + parallelization_group=mpu.get_data_parallel_group(with_context_parallel=True), ) return state_dict, checkpoint_name, False, CheckpointType.LOCAL else: - raise NotImplementedError(f"Please use local or global non-persistent checkpoints (got: {args.non_persistent_ckpt_type})") + raise NotImplementedError( + f"Please use local or global non-persistent checkpoints (got: {args.non_persistent_ckpt_type})" + ) def _load_global_dist_base_checkpoint( load_dir, args, rank0, sharded_state_dict, iteration, release, checkpointing_context=None ): - """ Load the base state_dict from the given directory containing the global distributed checkpoint """ + """Load the base state_dict from the given directory containing the global distributed checkpoint""" if rank0: checkpoint_name = find_checkpoint_rank_0(load_dir, iteration, release) state_dict = dist_checkpointing.load_common_state_dict(checkpoint_name) @@ -1223,19 +1463,21 @@ def _load_global_dist_base_checkpoint( args.use_dist_ckpt, ) raise RuntimeError( - 'Detected load from a distributed checkpoint, but neither --use-dist-ckpt nor --auto-detect-ckpt-format is set.' + "Detected load from a distributed checkpoint, but neither --use-dist-ckpt nor --auto-detect-ckpt-format is set." ) checkpoint_name = get_checkpoint_name(load_dir, iteration, release, return_base_dir=True) load_strategy = TorchDistLoadShardedStrategy(cache_metadata=args.ckpt_assume_constant_structure) # NOTE: `args.ckpt_fully_parallel_load` applies to both persistent and non-persistent checkpoints. if args.ckpt_fully_parallel_load: - if args.ckpt_fully_parallel_load_process_group == 'dp': + if args.ckpt_fully_parallel_load_process_group == "dp": process_group = mpu.get_data_parallel_group(with_context_parallel=True) - elif args.ckpt_fully_parallel_load_process_group == 'ep_dp': + elif args.ckpt_fully_parallel_load_process_group == "ep_dp": process_group = mpu.get_expert_data_parallel_group() else: - raise ValueError(f"Invalid load process group: {args.ckpt_fully_parallel_load_process_group}") + raise ValueError( + f"Invalid load process group: {args.ckpt_fully_parallel_load_process_group}" + ) load_strategy = FullyParallelLoadStrategyWrapper( load_strategy, process_group, exchange_algo=args.ckpt_fully_parallel_load_exchange_algo @@ -1284,7 +1526,7 @@ def _load_base_checkpoint( sharded_state_dict=None, checkpointing_context=None, ): - """ Load the base state_dict from the given directory + """Load the base state_dict from the given directory If rank0 is true, just loads rank 0 checkpoint, ignoring arguments. """ @@ -1298,7 +1540,7 @@ def _load_base_checkpoint( non_persistent_global_dir, args, checkpointing_context ) iteration, release = -1, False - tracker_filename = 'because load directory is not defined' + tracker_filename = "because load directory is not defined" if load_dir is not None: tracker_filename = get_checkpoint_tracker_filename(load_dir) if isfile(tracker_filename): @@ -1323,14 +1565,14 @@ def _load_base_checkpoint( checkpointing_context, ) else: - print_rank_0('WARNING: non-persistent checkpoints are older than persistent checkpoint') + print_rank_0("WARNING: non-persistent checkpoints are older than persistent checkpoint") # Otherwise we are dealing with global checkpoints # If no tracker file, return nothing if iteration == -1: if not rank0: - print_rank_0('WARNING: could not find the metadata file {}'.format(tracker_filename)) - print_rank_0(' will not load any checkpoints and will start from random') + print_rank_0("WARNING: could not find the metadata file {}".format(tracker_filename)) + print_rank_0(" will not load any checkpoints and will start from random") # Conditionally exit if checkpoint not found. if args.exit_on_missing_checkpoint: print_rank_0(">> '--exit-on-missing-checkpoint' set ... exiting. <<") @@ -1347,10 +1589,10 @@ def _load_base_checkpoint( if not rank0: dist_infix = "distributed " if ckpt_format == "torch_dist" else "" if release: - print_rank_0(f' loading release {dist_infix}checkpoint from {load_dir}') + print_rank_0(f" loading release {dist_infix}checkpoint from {load_dir}") else: print_rank_0( - f' loading {dist_infix}checkpoint from {load_dir} at iteration {iteration}' + f" loading {dist_infix}checkpoint from {load_dir} at iteration {iteration}" ) ckpt_type = None @@ -1358,7 +1600,13 @@ def _load_base_checkpoint( # Handle global distributed checkpoint if ckpt_format == "torch_dist": return _load_global_dist_base_checkpoint( - load_dir, args, rank0, sharded_state_dict, iteration, release, checkpointing_context=checkpointing_context + load_dir, + args, + rank0, + sharded_state_dict, + iteration, + release, + checkpointing_context=checkpointing_context, ) elif ckpt_format == "torch": ckpt_type = CheckpointType.LEGACY @@ -1366,26 +1614,30 @@ def _load_base_checkpoint( if rank0: checkpoint_name = find_checkpoint_rank_0(load_dir, iteration, release) else: - checkpoint_name = get_checkpoint_name(load_dir, iteration, release, return_base_dir=False) + checkpoint_name = get_checkpoint_name( + load_dir, iteration, release, return_base_dir=False + ) try: - state_dict = torch.load(checkpoint_name, map_location='cpu') + state_dict = torch.load(checkpoint_name, map_location="cpu") except ModuleNotFoundError: from megatron.legacy.fp16_deprecated import loss_scaler # For backward compatibility. if not rank0: - print_rank_0(' > deserializing using the old code structure ...') - sys.modules['fp16.loss_scaler'] = sys.modules['megatron.legacy.fp16_deprecated.loss_scaler'] - sys.modules['megatron.fp16.loss_scaler'] = sys.modules[ - 'megatron.legacy.fp16_deprecated.loss_scaler' + print_rank_0(" > deserializing using the old code structure ...") + sys.modules["fp16.loss_scaler"] = sys.modules[ + "megatron.legacy.fp16_deprecated.loss_scaler" + ] + sys.modules["megatron.fp16.loss_scaler"] = sys.modules[ + "megatron.legacy.fp16_deprecated.loss_scaler" ] - sys.modules['megatron.model'] = sys.modules['megatron.legacy.model'] - state_dict = torch.load(checkpoint_name, map_location='cpu') - sys.modules.pop('fp16.loss_scaler', None) - sys.modules.pop('megatron.fp16.loss_scaler', None) - sys.modules.pop('megatron.model', None) + sys.modules["megatron.model"] = sys.modules["megatron.legacy.model"] + state_dict = torch.load(checkpoint_name, map_location="cpu") + sys.modules.pop("fp16.loss_scaler", None) + sys.modules.pop("megatron.fp16.loss_scaler", None) + sys.modules.pop("megatron.model", None) except Exception as e: - print('could not load the checkpoint') + print("could not load the checkpoint") print(e) sys.exit() elif ckpt_format == "torch_dcp": @@ -1415,18 +1667,21 @@ def _load_base_checkpoint( return {}, checkpoint_name, release, CheckpointType.FSDP_DTENSOR state_dict = sharded_state_dict - raw_optimizer_state_dict = state_dict["optimizer"].copy() if "optimizer" in state_dict else None + raw_optimizer_state_dict = ( + state_dict["optimizer"].copy() if "optimizer" in state_dict else None + ) raw_model_state_dict = state_dict["model"].copy() if "model" in state_dict else None model = state_dict.pop("_model") state_dict = preprocess_fsdp_dtensor_state_dict(args, state_dict, model[0]) ckpt_type = CheckpointType.FSDP_DTENSOR fs_storage_reader = torch.distributed.checkpoint.FileSystemReader(checkpoint_name) - allow_partial_load = not getattr(args, 'strict_fsdp_dtensor_load', False) + allow_partial_load = not getattr(args, "strict_fsdp_dtensor_load", False) if allow_partial_load: state_dict_metadata = fs_storage_reader.read_metadata().state_dict_metadata rank = torch.distributed.get_rank() import time as _time + _time.sleep(rank * 0.001) # Make that logs of different ranks do not overlap print_diff_in_state_dicts(state_dict_metadata, state_dict) @@ -1448,9 +1703,7 @@ def _load_base_checkpoint( return state_dict, checkpoint_name, release, ckpt_type -def load_args_from_checkpoint( - args, load_arg='load', checkpointing_context=None -): +def load_args_from_checkpoint(args, load_arg="load", checkpointing_context=None): """Set required arguments from the checkpoint specified in the arguments. @@ -1466,7 +1719,7 @@ def load_args_from_checkpoint( load_dir = getattr(args, load_arg) if load_dir is None: - print_rank_0('No load directory specified, using provided arguments.') + print_rank_0("No load directory specified, using provided arguments.") return args state_dict, checkpoint_name, release, ckpt_type = _load_base_checkpoint( @@ -1478,39 +1731,41 @@ def load_args_from_checkpoint( # Args. if not state_dict: - print_rank_0('Checkpoint not found to provide arguments, using provided arguments.') + print_rank_0("Checkpoint not found to provide arguments, using provided arguments.") return args - if 'args' not in state_dict: - print_rank_0('Checkpoint provided does not have arguments saved, using provided arguments.') + if "args" not in state_dict: + print_rank_0("Checkpoint provided does not have arguments saved, using provided arguments.") return args - checkpoint_args = state_dict['args'] - checkpoint_version = state_dict.get('checkpoint_version', 0) - args.iteration = state_dict['iteration'] + checkpoint_args = state_dict["args"] + checkpoint_version = state_dict.get("checkpoint_version", 0) + args.iteration = state_dict["iteration"] # One-off conversion for foundation models - if hasattr(checkpoint_args, 'disable_bias_linear'): + if hasattr(checkpoint_args, "disable_bias_linear"): setattr( - checkpoint_args, 'add_bias_linear', not getattr(checkpoint_args, 'disable_bias_linear') + checkpoint_args, "add_bias_linear", not getattr(checkpoint_args, "disable_bias_linear") ) # For backward compatibility. - if hasattr(checkpoint_args, 'apply_layernorm_rms'): - checkpoint_args.normalization = 'RMSNorm' + if hasattr(checkpoint_args, "apply_layernorm_rms"): + checkpoint_args.normalization = "RMSNorm" # Backward compat: old checkpoints have hybrid_override_pattern but not hybrid_layer_pattern - if (getattr(checkpoint_args, 'hybrid_override_pattern', None) is not None - and getattr(checkpoint_args, 'hybrid_layer_pattern', None) is None): + if ( + getattr(checkpoint_args, "hybrid_override_pattern", None) is not None + and getattr(checkpoint_args, "hybrid_layer_pattern", None) is None + ): setattr( - checkpoint_args, 'hybrid_layer_pattern', - getattr(checkpoint_args, 'hybrid_override_pattern'), + checkpoint_args, + "hybrid_layer_pattern", + getattr(checkpoint_args, "hybrid_override_pattern"), ) # num_layers is now derived from hybrid_layer_pattern in validate_args, and should not be # set at the same time as hybrid_layer_pattern. - if hasattr(checkpoint_args, 'num_layers'): - setattr(checkpoint_args, 'num_layers', None) - + if hasattr(checkpoint_args, "num_layers"): + setattr(checkpoint_args, "num_layers", None) def _set_arg(arg_name, old_arg_name=None, force=False): if not force and getattr(args, arg_name, None) is not None: @@ -1528,100 +1783,110 @@ def _set_arg(arg_name, old_arg_name=None, force=False): print_rank_0(f"Checkpoint did not provide arguments {arg_name}") # Model args. - _set_arg('num_layers') - _set_arg('hidden_size') - _set_arg('ffn_hidden_size') - _set_arg('seq_length') - _set_arg('num_attention_heads') - _set_arg('num_query_groups', force=True) - _set_arg('group_query_attention', force=True) - _set_arg('kv_channels') - _set_arg('max_position_embeddings') - _set_arg('position_embedding_type', force=True) - _set_arg('add_position_embedding', force=True) - _set_arg('use_rotary_position_embeddings', force=True) - _set_arg('rotary_base', force=True) - _set_arg('rotary_percent', force=True) - _set_arg('rotary_interleaved', force=True) - _set_arg('add_bias_linear', force=True) - _set_arg('add_qkv_bias', force=True) - _set_arg('squared_relu', force=True) - _set_arg('swiglu', force=True) - _set_arg('multiple_of', force=True) - _set_arg('hidden_dim_multiplier', force=True) - _set_arg('untie_embeddings_and_output_weights', force=True) - _set_arg('apply_layernorm_1p', force=True) - _set_arg('normalization', force=True) - _set_arg('apply_query_key_layer_scaling', force=True) - _set_arg('attention_dropout', force=True) - _set_arg('hidden_dropout', force=True) + _set_arg("num_layers") + _set_arg("hidden_size") + _set_arg("ffn_hidden_size") + _set_arg("seq_length") + _set_arg("num_attention_heads") + _set_arg("num_query_groups", force=True) + _set_arg("group_query_attention", force=True) + _set_arg("kv_channels") + _set_arg("max_position_embeddings") + _set_arg("position_embedding_type", force=True) + _set_arg("add_position_embedding", force=True) + _set_arg("use_rotary_position_embeddings", force=True) + _set_arg("rotary_base", force=True) + _set_arg("rotary_percent", force=True) + _set_arg("rotary_interleaved", force=True) + _set_arg("add_bias_linear", force=True) + _set_arg("add_qkv_bias", force=True) + _set_arg("squared_relu", force=True) + _set_arg("swiglu", force=True) + _set_arg("multiple_of", force=True) + _set_arg("hidden_dim_multiplier", force=True) + _set_arg("untie_embeddings_and_output_weights", force=True) + _set_arg("apply_layernorm_1p", force=True) + _set_arg("normalization", force=True) + _set_arg("apply_query_key_layer_scaling", force=True) + _set_arg("attention_dropout", force=True) + _set_arg("hidden_dropout", force=True) # Legacy MTP pattern for old checkpoints - _set_arg('mtp_hybrid_override_pattern', force=True) - _set_arg('mtp_num_layers', force=True) - _set_arg('mtp_use_repeated_layer', force=True) + _set_arg("mtp_hybrid_override_pattern", force=True) + _set_arg("mtp_num_layers", force=True) + _set_arg("mtp_use_repeated_layer", force=True) - _set_arg('spec', force=True) + _set_arg("spec", force=True) - _set_arg('num_experts', force=True) - _set_arg('mtp_num_layers', force=True) - _set_arg('moe_layer_freq', force=True) - if getattr(checkpoint_args, 'num_experts', None) is not None: - _set_arg('moe_ffn_hidden_size', force=True) + _set_arg("num_experts", force=True) + _set_arg("mtp_num_layers", force=True) + _set_arg("moe_layer_freq", force=True) + if getattr(checkpoint_args, "num_experts", None) is not None: + _set_arg("moe_ffn_hidden_size", force=True) else: - setattr(args, 'moe_ffn_hidden_size', None) - _set_arg('moe_router_topk', force=True) - _set_arg('moe_token_dispatcher_type', force=False) - _set_arg('moe_router_pre_softmax', force=True) - _set_arg('moe_grouped_gemm', force=True) - _set_arg('moe_shared_expert_intermediate_size', force=True) - _set_arg('moe_router_score_function', force=True) - _set_arg('moe_router_enable_expert_bias', force=True) - _set_arg('moe_router_topk_scaling_factor', force=True) + setattr(args, "moe_ffn_hidden_size", None) + _set_arg("moe_router_topk", force=True) + _set_arg("moe_token_dispatcher_type", force=False) + _set_arg("moe_router_pre_softmax", force=True) + _set_arg("moe_grouped_gemm", force=True) + _set_arg("moe_shared_expert_intermediate_size", force=True) + _set_arg("moe_router_score_function", force=True) + _set_arg("moe_router_enable_expert_bias", force=True) + _set_arg("moe_router_topk_scaling_factor", force=True) # Mamba args. - _set_arg('mamba_state_dim', force=True) - _set_arg('mamba_head_dim', force=True) - _set_arg('mamba_num_groups', force=True) - _set_arg('mamba_num_heads', force=True) + _set_arg("mamba_state_dim", force=True) + _set_arg("mamba_head_dim", force=True) + _set_arg("mamba_num_groups", force=True) + _set_arg("mamba_num_heads", force=True) # We need to be able to override hybrid_layer_pattern from the command-line so that different # pipelining can be specified when re-loading a model (e.g. for inference or post-training). - _set_arg('hybrid_layer_pattern') + _set_arg("hybrid_layer_pattern") # Heterogeneous args. - _set_arg('heterogeneous_layers_config_path', force=True) - _set_arg('heterogeneous_layers_config_encoded_json', force=True) + _set_arg("heterogeneous_layers_config_path", force=True) + _set_arg("heterogeneous_layers_config_encoded_json", force=True) # MoE latent projection. - _set_arg('moe_latent_size', force=True) + _set_arg("moe_latent_size", force=True) # Tokenizer args. if args.use_tokenizer_model_from_checkpoint_args: # Using checkpoint version might not always be safe (e.g., if running on different cluster). - _set_arg('tokenizer_model', force=True) - _set_arg('tokenizer_type', force=True) - _set_arg('tiktoken_pattern', force=True) - _set_arg('padded_vocab_size') + _set_arg("tokenizer_model", force=True) + _set_arg("tokenizer_type", force=True) + _set_arg("tiktoken_pattern", force=True) + _set_arg("padded_vocab_size") # Checkpoint args. - _set_arg('ckpt_format') + _set_arg("ckpt_format") # Model parallelism args. if args.use_mp_args_from_checkpoint_args: if checkpoint_version < 3.0: - _set_arg('tensor_model_parallel_size', 'model_parallel_size') + _set_arg("tensor_model_parallel_size", "model_parallel_size") else: - _set_arg('tensor_model_parallel_size', force=True) - _set_arg('pipeline_model_parallel_size', force=True) - _set_arg('virtual_pipeline_model_parallel_size', force=True) - _set_arg('num_layers_per_virtual_pipeline_stage') - _set_arg('expert_model_parallel_size', force=True) + _set_arg("tensor_model_parallel_size", force=True) + _set_arg("pipeline_model_parallel_size", force=True) + _set_arg("virtual_pipeline_model_parallel_size", force=True) + _set_arg("num_layers_per_virtual_pipeline_stage") + _set_arg("expert_model_parallel_size", force=True) return args, checkpoint_args -def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load', strict=True, - checkpointing_context=None, skip_load_to_model_and_opt=False, tp_group: Optional[torch.distributed.ProcessGroup] = None, pp_group: Optional[torch.distributed.ProcessGroup] = None, dp_cp_group: Optional[torch.distributed.ProcessGroup] = None): +def load_checkpoint( + ddp_model, + optimizer, + opt_param_scheduler, + load_arg="load", + strict=True, + checkpointing_context=None, + skip_load_to_model_and_opt=False, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + pp_group: Optional[torch.distributed.ProcessGroup] = None, + dp_cp_group: Optional[torch.distributed.ProcessGroup] = None, +): """Load a model checkpoint and return the iteration. strict (bool): whether to strictly enforce that the keys in :attr:`state_dict` of the checkpoint match the names of @@ -1635,10 +1900,10 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load', load_dir = getattr(args, load_arg) # Finetuning directories - pretrained_dir = getattr(args, 'pretrained_checkpoint', None) + pretrained_dir = getattr(args, "pretrained_checkpoint", None) if pretrained_dir is not None and not checkpoint_exists(load_dir): print_rank_0( - f'Checkpoint file not found in load directory {load_dir} attempting to finetune with checkpoint in {pretrained_dir}' + f"Checkpoint file not found in load directory {load_dir} attempting to finetune with checkpoint in {pretrained_dir}" ) load_dir = pretrained_dir if not checkpoint_exists(load_dir): @@ -1666,7 +1931,7 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load', elif ckpt_type in [CheckpointType.LOCAL, CheckpointType.GLOBAL]: ckpt_format = "torch_dist" elif ckpt_type == None: - pass # Not loaded. + pass # Not loaded. else: raise NotImplementedError(f"checkpoint format {ckpt_format} not supported") @@ -1692,29 +1957,36 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load', args.pipeline_model_parallel_size, ) - ckpt_world_size = getattr(ckpt_args, 'world_size', 0) - run_world_size = getattr(args, 'world_size', 0) - ckpt_dp = getattr(ckpt_args, 'data_parallel_size', 0) - run_dp = getattr(args, 'data_parallel_size', 0) + ckpt_world_size = getattr(ckpt_args, "world_size", 0) + run_world_size = getattr(args, "world_size", 0) + ckpt_dp = getattr(ckpt_args, "data_parallel_size", 0) + run_dp = getattr(args, "data_parallel_size", 0) mismatch_msg = "(TP, PP) mismatch after resume ({} vs {} from checkpoint)".format( run_tp_pp, ckpt_tp_pp ) ########## FlagScale Begin ########## - #Add support for changing parallel strategy from tp/pp to ep for ChainedOptimizer when using dist checkpointing + # Add support for changing parallel strategy from tp/pp to ep for ChainedOptimizer when using dist checkpointing convert_to_ep = ( - getattr(args, 'expert_model_parallel_size', 1) != 1 and - getattr(state_dict['args'], 'expert_model_parallel_size', 1) == 1 + getattr(args, "expert_model_parallel_size", 1) != 1 + and getattr(state_dict["args"], "expert_model_parallel_size", 1) == 1 ) ########## FlagScale End ########## # Determine if RNG state will be loaded - if (ckpt_tp_pp == run_tp_pp and not release and not args.finetune and not args.no_load_rng - and not getattr(ckpt_args, 'no_save_rng', False)): + if ( + ckpt_tp_pp == run_tp_pp + and not release + and not args.finetune + and not args.no_load_rng + and not getattr(ckpt_args, "no_save_rng", False) + ): if tp_group is None and pp_group is None: tp_group = mpu.get_tensor_model_parallel_group() pp_group = mpu.get_pipeline_model_parallel_group() - gen_sd_rng_state = get_rng_state(args.ckpt_format, tp_group, pp_group) # we can load the rng state + gen_sd_rng_state = get_rng_state( + args.ckpt_format, tp_group, pp_group + ) # we can load the rng state else: ignore_rng_state = True gen_sd_rng_state = None @@ -1724,12 +1996,20 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load', if ckpt_type == CheckpointType.LOCAL: sharded_sd_metadata = _build_sharded_state_dict_metadata(args) else: - sharded_sd_metadata = dist_checkpointing.load_content_metadata(preloaded_state_dict=state_dict) - print_rank_0(f'sharded_state_dict metadata loaded from the checkpoint: {sharded_sd_metadata}') + sharded_sd_metadata = dist_checkpointing.load_content_metadata( + preloaded_state_dict=state_dict + ) + print_rank_0( + f"sharded_state_dict metadata loaded from the checkpoint: {sharded_sd_metadata}" + ) # Determine if optimizer state will be loaded - if (not release and not args.finetune and not args.no_load_optim - and not getattr(ckpt_args, 'no_save_optim', False)): + if ( + not release + and not args.finetune + and not args.no_load_optim + and not getattr(ckpt_args, "no_save_optim", False) + ): gen_sd_optim = optimizer gen_sd_opt_param_scheduler = opt_param_scheduler @@ -1739,24 +2019,34 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load', # Can be removed after ending support for MLM optimizer checkpoints with MCore < v0.13 # (for MCore v0.13+ checkpoints `sharded_sd_metadata is not None`) sharded_sd_metadata = { - 'distrib_optim_sharding_type': ('fully_sharded_model_space' - if getattr(ckpt_args, 'ckpt_fully_parallel_save', False) - else 'dp_zero_gather_scatter'), + "distrib_optim_sharding_type": ( + "fully_sharded_model_space" + if getattr(ckpt_args, "ckpt_fully_parallel_save", False) + else "dp_zero_gather_scatter" + ), } if ( ckpt_tp_pp != run_tp_pp - and sharded_sd_metadata['distrib_optim_sharding_type'] + and sharded_sd_metadata["distrib_optim_sharding_type"] not in DistributedOptimizer.checkpoint_fully_reshardable_formats - and convert_to_ep ########## FlagScale Added ########## + and convert_to_ep ########## FlagScale Added ########## ): - raise RuntimeError(f"{mismatch_msg}: not supported for DistributedOptimizer with sharding type" - f" {sharded_sd_metadata['distrib_optim_sharding_type']}." - f" Please use `--ckpt-fully-parallel-save` flag during checkpoint saving.") + raise RuntimeError( + f"{mismatch_msg}: not supported for DistributedOptimizer with sharding type" + f" {sharded_sd_metadata['distrib_optim_sharding_type']}." + f" Please use `--ckpt-fully-parallel-save` flag during checkpoint saving." + ) # Check if fully parallel load is compatible with sharding type - if args.ckpt_fully_parallel_load and sharded_sd_metadata['distrib_optim_sharding_type'] == 'dp_zero_gather_scatter': - raise RuntimeError("Fully parallel load is not supported for dp_zero_gather_scatter checkpoints. " - "Please remove --ckpt-fully-parallel-load flag") + if ( + args.ckpt_fully_parallel_load + and sharded_sd_metadata["distrib_optim_sharding_type"] + == "dp_zero_gather_scatter" + ): + raise RuntimeError( + "Fully parallel load is not supported for dp_zero_gather_scatter checkpoints. " + "Please remove --ckpt-fully-parallel-load flag" + ) else: gen_sd_optim = None gen_sd_opt_param_scheduler = None @@ -1770,7 +2060,9 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load', sharded_sd_metadata = {} sharded_sd_metadata["dp_cp_group"] = dp_cp_group - optim_sd_kwargs = dict(metadata=sharded_sd_metadata, is_loading=True, convert_to_ep=convert_to_ep) ########## FlagScale Added ########## + optim_sd_kwargs = dict( + metadata=sharded_sd_metadata, is_loading=True, convert_to_ep=convert_to_ep + ) ########## FlagScale Added ########## model_sd_kwargs = dict(metadata=sharded_sd_metadata) # Determine if rerun state will be loaded @@ -1781,32 +2073,37 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load', and ckpt_dp == run_dp and not release and not args.finetune - and 'rerun_state_machine' in state_dict + and "rerun_state_machine" in state_dict ): rerun_state_machine = get_rerun_state_machine() - if rerun_state_machine.validate_state_dict(state_dict['rerun_state_machine']): + if rerun_state_machine.validate_state_dict(state_dict["rerun_state_machine"]): gen_sd_rerun_state = rerun_state_machine.state_dict( - data_iterator=None, ckpt_format=ckpt_format, force=True, + data_iterator=None, + ckpt_format=ckpt_format, + force=True, ) ignore_rerun_state = False - if ( - ckpt_world_size != run_world_size - or ckpt_tp_pp != run_tp_pp - or ckpt_dp != run_dp - ): + if ckpt_world_size != run_world_size or ckpt_tp_pp != run_tp_pp or ckpt_dp != run_dp: print_rank_0("Job sharding has changed: Rerun state will be ignored") # [ModelOpt]: Initial loading from non-resume sharded checkpoint to a Distillation Model # will result in key mismatch with loss modules potentially containing parameters, since # it requires generating a state_dict before loading. Here we hide those modules if present. - with contextlib.ExitStack() as stack: # Allows multiple context managers for each model shard + with ( + contextlib.ExitStack() as stack + ): # Allows multiple context managers for each model shard if args.finetune and hasattr(model[0], "hide_loss_modules"): for m in model: stack.enter_context(m.hide_loss_modules()) - load_kwargs['sharded_state_dict'] = generate_state_dict( - args, model, gen_sd_optim, gen_sd_opt_param_scheduler, gen_sd_rng_state, - optim_sd_kwargs=optim_sd_kwargs, model_sd_kwargs=model_sd_kwargs, - rerun_state=gen_sd_rerun_state + load_kwargs["sharded_state_dict"] = generate_state_dict( + args, + model, + gen_sd_optim, + gen_sd_opt_param_scheduler, + gen_sd_rng_state, + optim_sd_kwargs=optim_sd_kwargs, + model_sd_kwargs=model_sd_kwargs, + rerun_state=gen_sd_rerun_state, ) elif args.ckpt_format == "torch_dcp": model_sd = model[0].state_dict() @@ -1839,7 +2136,9 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load', if not args.finetune: if "rerun_state_machine" in state_dict_metadata: gen_sd_rerun_state = get_rerun_state_machine().state_dict( - data_iterator=None, ckpt_format=ckpt_format, force=True, + data_iterator=None, + ckpt_format=ckpt_format, + force=True, ) if not args.no_load_rng: gen_sd_rng_state = get_rng_state(args.ckpt_format, tp_group, pp_group) @@ -1863,8 +2162,7 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load', load_kwargs["sharded_state_dict"] = state_dict state_dict, checkpoint_name, release, ckpt_type = _load_base_checkpoint( - load_dir, args, rank0=False, checkpointing_context=checkpointing_context, - **load_kwargs + load_dir, args, rank0=False, checkpointing_context=checkpointing_context, **load_kwargs ) # Checkpoint not loaded. @@ -1873,7 +2171,7 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load', return 0, 0 # Set checkpoint version. - set_checkpoint_version(state_dict.get('checkpoint_version', 0)) + set_checkpoint_version(state_dict.get("checkpoint_version", 0)) # Convert to regular torch tensor to DTensor. if ckpt_type == CheckpointType.LEGACY and args.ckpt_format == "torch_dcp": @@ -1885,29 +2183,28 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load', iteration = 0 else: try: - iteration = state_dict['iteration'] + iteration = state_dict["iteration"] except KeyError: try: # Backward compatible with older checkpoints - iteration = state_dict['total_iters'] + iteration = state_dict["total_iters"] except KeyError: - print_rank_0('A metadata file exists but unable to load ' - 'iteration from checkpoint {}, exiting'.format(checkpoint_name)) + print_rank_0( + "A metadata file exists but unable to load " + "iteration from checkpoint {}, exiting".format(checkpoint_name) + ) sys.exit() - num_floating_point_operations_so_far = state_dict.get('num_floating_point_operations_so_far', 0) + num_floating_point_operations_so_far = state_dict.get("num_floating_point_operations_so_far", 0) # Check arguments. - if 'args' in state_dict and not args.finetune: - checkpoint_args = state_dict['args'] + if "args" in state_dict and not args.finetune: + checkpoint_args = state_dict["args"] check_checkpoint_args(checkpoint_args) - args.consumed_train_samples = getattr(checkpoint_args, - 'consumed_train_samples', 0) - args.skipped_train_samples = getattr(checkpoint_args, - 'skipped_train_samples', 0) + args.consumed_train_samples = getattr(checkpoint_args, "consumed_train_samples", 0) + args.skipped_train_samples = getattr(checkpoint_args, "skipped_train_samples", 0) update_num_microbatches(consumed_samples=args.consumed_train_samples, verbose=True) - args.consumed_valid_samples = getattr(checkpoint_args, - 'consumed_valid_samples', 0) + args.consumed_valid_samples = getattr(checkpoint_args, "consumed_valid_samples", 0) else: - print_rank_0('could not find arguments in the checkpoint ...') + print_rank_0("could not find arguments in the checkpoint ...") def load_model_state_dict(module, state_dict, strict: bool): """Helper function to load state dict with fallback for missing extra states.""" @@ -1918,62 +2215,79 @@ def load_model_state_dict(module, state_dict, strict: bool): # Fallback support for backward compatibility breaking changes in TransformerEngine load_return = module.load_state_dict(state_dict, strict=False) print(f"load_return: {load_return}") + # Model. if not skip_load_to_model_and_opt: if len(ddp_model) == 1: - load_model_state_dict(ddp_model[0], state_dict['model'], strict) + load_model_state_dict(ddp_model[0], state_dict["model"], strict) else: for i in range(len(ddp_model)): # If there is no corresponding model in the state_dict, it will be ignored. # It means that this is an empty stage. - if 'model%d' % i not in state_dict: + if "model%d" % i not in state_dict: continue - load_model_state_dict(ddp_model[i], state_dict['model%d' % i], strict) + load_model_state_dict(ddp_model[i], state_dict["model%d" % i], strict) # Fix up query/key/value matrix ordering if needed. checkpoint_version = get_checkpoint_version() - print_rank_0(f' checkpoint version {checkpoint_version}') + print_rank_0(f" checkpoint version {checkpoint_version}") fix_query_key_value_ordering(model, checkpoint_version) # Optimizer. if not release and not args.finetune and not args.no_load_optim: try: # Load state dict. - if getattr(args, "optimizer", "adam").startswith("dist_") and args.ckpt_format == 'torch': + if ( + getattr(args, "optimizer", "adam").startswith("dist_") + and args.ckpt_format == "torch" + ): # LayerWiseDistributedOptimizer load optimizer state from file on different ranks dp_rank = mpu.get_data_parallel_rank() - optim_checkpoint_name = os.path.join(os.path.dirname(checkpoint_name), f"layer_wise_optimizer_{dp_rank}.pt") + optim_checkpoint_name = os.path.join( + os.path.dirname(checkpoint_name), f"layer_wise_optimizer_{dp_rank}.pt" + ) optimizer.load_state_dict_from_file(optim_checkpoint_name) - elif not skip_load_to_model_and_opt and optimizer is not None and not optimizer.is_stub_optimizer: - optimizer.load_state_dict(state_dict['optimizer']) + elif ( + not skip_load_to_model_and_opt + and optimizer is not None + and not optimizer.is_stub_optimizer + ): + optimizer.load_state_dict(state_dict["optimizer"]) # Load distributed optimizer's custom parameter state. # For distributed checkpoint it's already loaded in load_state_dict above is_torch_dist = ckpt_format == "torch_dist" - if args.use_distributed_optimizer and not is_torch_dist and ckpt_format not in ["torch_dcp", "fsdp_dtensor"]: + if ( + args.use_distributed_optimizer + and not is_torch_dist + and ckpt_format not in ["torch_dcp", "fsdp_dtensor"] + ): # NOTE: this is a manual read of the tracker file. # This code should not be reached when reading from a non_persistent checkpoint assert not is_torch_dist tracker_filename = get_checkpoint_tracker_filename(load_dir) iteration, release = read_metadata(tracker_filename) - model_checkpoint_name = \ - get_checkpoint_name(load_dir, iteration, release) - optim_checkpoint_name = \ - get_distributed_optimizer_checkpoint_name( - model_checkpoint_name) - optimizer.load_parameter_state(optim_checkpoint_name, - update_legacy_format=args.ckpt_convert_update_legacy_dist_opt_format) + model_checkpoint_name = get_checkpoint_name(load_dir, iteration, release) + optim_checkpoint_name = get_distributed_optimizer_checkpoint_name( + model_checkpoint_name + ) + optimizer.load_parameter_state( + optim_checkpoint_name, + update_legacy_format=args.ckpt_convert_update_legacy_dist_opt_format, + ) # Load scheduler. if opt_param_scheduler is not None: - if 'lr_scheduler' in state_dict: # backward compatbility - opt_param_scheduler.load_state_dict(state_dict['lr_scheduler']) + if "lr_scheduler" in state_dict: # backward compatbility + opt_param_scheduler.load_state_dict(state_dict["lr_scheduler"]) else: - opt_param_scheduler.load_state_dict(state_dict['opt_param_scheduler']) + opt_param_scheduler.load_state_dict(state_dict["opt_param_scheduler"]) except KeyError as e: - print_rank_0('Unable to load optimizer from checkpoint {}. ' - 'Specify --no-load-optim or --finetune to prevent ' - 'attempting to load the optimizer state, ' - 'exiting ...'.format(checkpoint_name)) + print_rank_0( + "Unable to load optimizer from checkpoint {}. " + "Specify --no-load-optim or --finetune to prevent " + "attempting to load the optimizer state, " + "exiting ...".format(checkpoint_name) + ) raise e else: if (args.fp16 or args.bf16) and optimizer is not None: @@ -1985,8 +2299,8 @@ def load_model_state_dict(module, state_dict, strict: bool): # rerun state if not ignore_rerun_state: try: - if 'rerun_state_machine' in state_dict: - get_rerun_state_machine().load_state_dict(state_dict['rerun_state_machine']) + if "rerun_state_machine" in state_dict: + get_rerun_state_machine().load_state_dict(state_dict["rerun_state_machine"]) except Exception as e: print_rank_0(f"Unable to restore RerunMachine from checkpoint: {e}. Skipping.") @@ -1995,87 +2309,92 @@ def load_model_state_dict(module, state_dict, strict: bool): try: cuda_rng_tracker = tensor_parallel.get_cuda_rng_tracker() graph_safe_rng = tensor_parallel.is_graph_safe_cuda_rng_tracker(cuda_rng_tracker) - if 'rng_state' in state_dict: + if "rng_state" in state_dict: if args.ckpt_format == "fsdp_dtensor": # FSDP DTensor checkpoints store rng_state in a different format. tp_rank = mpu.get_tensor_model_parallel_rank() pp_rank = mpu.get_pipeline_model_parallel_rank() - if f"({pp_rank}, {tp_rank})" in state_dict['rng_state']: - rng_state = state_dict['rng_state'][f"({pp_rank}, {tp_rank})"] + if f"({pp_rank}, {tp_rank})" in state_dict["rng_state"]: + rng_state = state_dict["rng_state"][f"({pp_rank}, {tp_rank})"] else: print_rank_0("WARNING: RNG state not found for current TP/PP rank") - rng_state = next(iter(state_dict['rng_state'].values())) + rng_state = next(iter(state_dict["rng_state"].values())) else: - rng_state = state_dict['rng_state'] + rng_state = state_dict["rng_state"] # access rng_state for data parallel rank if args.data_parallel_random_init: rng_state = rng_state[mpu.get_data_parallel_rank()] else: rng_state = rng_state[0] - random.setstate(rng_state['random_rng_state']) - np.random.set_state(rng_state['np_rng_state']) - torch.set_rng_state(rng_state['torch_rng_state']) - cur_platform.set_rng_state(rng_state['cuda_rng_state']) + random.setstate(rng_state["random_rng_state"]) + np.random.set_state(rng_state["np_rng_state"]) + torch.set_rng_state(rng_state["torch_rng_state"]) + cur_platform.set_rng_state(rng_state["cuda_rng_state"]) # Check for empty states array - if not rng_state['rng_tracker_states']: + if not rng_state["rng_tracker_states"]: raise KeyError rng_tracker_states = { k: tensor_parallel.convert_cuda_rng_state(v, to_graphable=graph_safe_rng) - for k, v in rng_state['rng_tracker_states'].items() + for k, v in rng_state["rng_tracker_states"].items() } else: # backward compatability - random.setstate(state_dict['random_rng_state']) - np.random.set_state(state_dict['np_rng_state']) - torch.set_rng_state(state_dict['torch_rng_state']) - cur_platform.set_rng_state(state_dict['cuda_rng_state']) + random.setstate(state_dict["random_rng_state"]) + np.random.set_state(state_dict["np_rng_state"]) + torch.set_rng_state(state_dict["torch_rng_state"]) + cur_platform.set_rng_state(state_dict["cuda_rng_state"]) # Check for empty states array - if not state_dict['rng_tracker_states']: + if not state_dict["rng_tracker_states"]: raise KeyError rng_tracker_states = { k: tensor_parallel.convert_cuda_rng_state(v, to_graphable=graph_safe_rng) - for k, v in state_dict['rng_tracker_states'].items() + for k, v in state_dict["rng_tracker_states"].items() } cuda_rng_tracker.set_states(rng_tracker_states) except KeyError: - print_rank_0('Unable to load rng state from checkpoint {}. ' - 'Specify --no-load-rng or --finetune to prevent ' - 'attempting to load the rng state, ' - 'exiting ...'.format(checkpoint_name)) + print_rank_0( + "Unable to load rng state from checkpoint {}. " + "Specify --no-load-rng or --finetune to prevent " + "attempting to load the rng state, " + "exiting ...".format(checkpoint_name) + ) sys.exit() # Some utilities want to load a checkpoint without distributed being initialized if torch.distributed.is_initialized(): torch.distributed.barrier() - print_rank_0(f' successfully loaded checkpoint from {load_dir} ' - f'[ t {mpu.get_tensor_model_parallel_rank() + 1}/{mpu.get_tensor_model_parallel_world_size()}, ' - f'p {mpu.get_pipeline_model_parallel_rank() + 1}/{mpu.get_pipeline_model_parallel_world_size()} ] ' - f'at iteration {iteration}') - + print_rank_0( + f" successfully loaded checkpoint from {load_dir} " + f"[ t {mpu.get_tensor_model_parallel_rank() + 1}/{mpu.get_tensor_model_parallel_world_size()}, " + f"p {mpu.get_pipeline_model_parallel_rank() + 1}/{mpu.get_pipeline_model_parallel_world_size()} ] " + f"at iteration {iteration}" + ) + if has_nvidia_modelopt: print_distributed_quant_summary(model, msg="After loading checkpoint") - + # Additional callback for wandb (last rank) - if not torch.distributed.is_initialized() \ - or is_last_rank(): + if not torch.distributed.is_initialized() or is_last_rank(): wandb_utils.on_load_checkpoint_success(checkpoint_name, load_dir) cur_platform.empty_cache() if iteration > 0: # Notify FT that a checkpoint was loaded. - is_local_chkpt = (ckpt_type == CheckpointType.LOCAL) + is_local_chkpt = ckpt_type == CheckpointType.LOCAL ft_integration.on_checkpoint_loaded(is_local_chkpt=is_local_chkpt) # Patch checkpoint as needed if required field is not found. if optimizer is not None: log_printed = False for param_group in optimizer.param_groups: - if 'default_config' not in param_group: - param_group['default_config'] = True + if "default_config" not in param_group: + param_group["default_config"] = True if not log_printed: - print_rank_0(">>> Inserting 'default_config' field into optimizer.param_groups...") + print_rank_0( + ">>> Inserting 'default_config' field into optimizer.param_groups..." + ) log_printed = True return iteration, num_floating_point_operations_so_far @@ -2095,8 +2414,9 @@ def _to_dtensor(wrapped_model, model_state_dict): return new_model_sd -def load_biencoder_checkpoint(model, only_query_model=False, - only_context_model=False, custom_load_path=None): +def load_biencoder_checkpoint( + model, only_query_model=False, only_context_model=False, custom_load_path=None +): """ selectively load retrieval models for indexing/retrieving from saved checkpoints @@ -2110,30 +2430,33 @@ def load_biencoder_checkpoint(model, only_query_model=False, tracker_filename = get_checkpoint_tracker_filename(load_path) - with open_file(tracker_filename, 'r') as f: + with open_file(tracker_filename, "r") as f: iteration = int(f.read().strip()) - checkpoint_name = get_checkpoint_name(load_path, iteration, - args.use_distributed_optimizer, - release=False) + checkpoint_name = get_checkpoint_name( + load_path, iteration, args.use_distributed_optimizer, release=False + ) if mpu.get_data_parallel_rank() == 0: - print('global rank {} is loading checkpoint {}'.format( - torch.distributed.get_rank(), checkpoint_name)) + print( + "global rank {} is loading checkpoint {}".format( + torch.distributed.get_rank(), checkpoint_name + ) + ) - state_dict = torch.load(checkpoint_name, map_location='cpu') - ret_state_dict = state_dict['model'] + state_dict = torch.load(checkpoint_name, map_location="cpu") + ret_state_dict = state_dict["model"] if only_query_model: - ret_state_dict.pop('context_model') + ret_state_dict.pop("context_model") if only_context_model: - ret_state_dict.pop('query_model') + ret_state_dict.pop("query_model") assert len(model) == 1 model[0].load_state_dict(ret_state_dict) torch.distributed.barrier() if mpu.get_data_parallel_rank() == 0: - print(' successfully loaded {}'.format(checkpoint_name)) + print(" successfully loaded {}".format(checkpoint_name)) return model diff --git a/flagscale/train/megatron/training/config/common_config.py b/flagscale/train/megatron/training/config/common_config.py index a2c443d1c0..7239fb48ed 100644 --- a/flagscale/train/megatron/training/config/common_config.py +++ b/flagscale/train/megatron/training/config/common_config.py @@ -3,6 +3,7 @@ from typing import Literal import os + @dataclass(kw_only=True) class RNGConfig: """Configuration settings for random number generation.""" @@ -25,7 +26,9 @@ class RNGConfig: class ProfilingConfig: """Configuration settings for profiling the training process.""" - use_nsys_profiler: bool = field(default=False, metadata={"argparse_meta": {"arg_names": ["--profile"], "dest": "profile"}}) + use_nsys_profiler: bool = field( + default=False, metadata={"argparse_meta": {"arg_names": ["--profile"], "dest": "profile"}} + ) """Enable nsys profiling. When using this option, nsys options should be specified in commandline. An example nsys commandline is `nsys profile -s none -t nvtx,cuda -o --force-overwrite true @@ -43,11 +46,11 @@ class ProfilingConfig: pytorch_profiler_collect_shapes: bool = False """Collect tensor shape in pytorch profiler.""" - + pytorch_profiler_collect_callstack: bool = False """Collect callstack in pytorch profiler.""" - - pytorch_profiler_collect_chakra: bool = False + + pytorch_profiler_collect_chakra: bool = False """Collect chakra trace in pytorch profiler.""" profile_ranks: list[int] = field(default_factory=lambda: []) @@ -106,7 +109,9 @@ class DistributedInitConfig: """If set, distributed ranks initialize order is changed from tp-cp-ep-dp-pp to tp-cp-ep-pp-dp. """ - use_gloo_process_groups: bool = field(default=True, metadata={"argparse_meta": {"arg_names": ["--disable-gloo-process-groups"]}}) + use_gloo_process_groups: bool = field( + default=True, metadata={"argparse_meta": {"arg_names": ["--disable-gloo-process-groups"]}} + ) """If enabled, create Gloo process groups for communications.""" use_sharp: bool = False diff --git a/flagscale/train/megatron/training/config/resilience_config.py b/flagscale/train/megatron/training/config/resilience_config.py index dd0bd71652..b284abbe43 100644 --- a/flagscale/train/megatron/training/config/resilience_config.py +++ b/flagscale/train/megatron/training/config/resilience_config.py @@ -2,6 +2,7 @@ from dataclasses import dataclass from typing import Literal + @dataclass(kw_only=True) class RerunStateMachineConfig: """Configuration for the rerun state machine used for result validation or stats.""" @@ -10,7 +11,9 @@ class RerunStateMachineConfig: """Rate at which to inject unexpected results, e.g. 1000 means once every 1000 result validations""" - error_injection_type: Literal["correct_result", "transient_error", "persistent_error"] = "transient_error" + error_injection_type: Literal["correct_result", "transient_error", "persistent_error"] = ( + "transient_error" + ) """Type of error to inject. """ rerun_mode: Literal["disabled", "validate_results", "report_stats"] = "validate_results" @@ -39,4 +42,3 @@ class StragglerDetectionConfig: disable_straggler_on_startup: bool = False """If set, StragglerDetector is disabled on startup.""" - diff --git a/flagscale/train/megatron/training/config/training_config.py b/flagscale/train/megatron/training/config/training_config.py index 027fc7a921..9142e8990d 100644 --- a/flagscale/train/megatron/training/config/training_config.py +++ b/flagscale/train/megatron/training/config/training_config.py @@ -3,6 +3,7 @@ import signal from typing import Literal, Optional + @dataclass(kw_only=True) class TrainingConfig: """Configuration settings related to the training loop.""" @@ -16,7 +17,9 @@ class TrainingConfig: data-parallel-size. If this value is None, then use micro-batch-size * data-parallel-size as the global batch size. This choice will result in 1 for number of micro-batches.""" - rampup_batch_size: list[int] | None = field(default=None, metadata={"argparse_meta": {"nargs": 3}}) + rampup_batch_size: list[int] | None = field( + default=None, metadata={"argparse_meta": {"nargs": 3}} + ) """Batch size ramp up with the following values: , , For example: @@ -163,13 +166,30 @@ class SchedulerConfig: """number of samples to warmup learning rate over. Calculated at runtime from lr_warmup_fraction, lr_warmup_iters, or lr_warmup_samples. """ - - override_opt_param_scheduler: bool = field(default=False, metadata={"argparse_meta": {"arg_names": ["--override-opt_param-scheduler", "--override-opt-param-scheduler"]}}) + + override_opt_param_scheduler: bool = field( + default=False, + metadata={ + "argparse_meta": { + "arg_names": ["--override-opt_param-scheduler", "--override-opt-param-scheduler"] + } + }, + ) """Reset the values of the scheduler (learning rate, warmup iterations, minimum learning rate, maximum number of iterations, and decay style) from input arguments and ignore values from checkpoints. Note that all the above values will be reset.""" - use_checkpoint_opt_param_scheduler: bool = field(default=False, metadata={"argparse_meta": {"arg_names": ["--use-checkpoint-opt_param-scheduler", "--use-checkpoint-opt-param-scheduler"]}}) + use_checkpoint_opt_param_scheduler: bool = field( + default=False, + metadata={ + "argparse_meta": { + "arg_names": [ + "--use-checkpoint-opt_param-scheduler", + "--use-checkpoint-opt-param-scheduler", + ] + } + }, + ) """Use checkpoint to set the values of the scheduler (learning rate, warmup iterations, minimum learning rate, maximum number of iterations, and decay style) from checkpoint and ignore input arguments.""" @@ -285,7 +305,10 @@ class LoggerConfig: runtime_time_unit: str = "hours" """Time unit to use for time logging. """ - barrier_with_L1_time: bool = field(default=True, metadata={"argparse_meta": {"arg_names": ["--no-barrier-with-level-1-timing"]}}) + barrier_with_L1_time: bool = field( + default=True, + metadata={"argparse_meta": {"arg_names": ["--no-barrier-with-level-1-timing"]}}, + ) """If not disabled, use barrier with level 1 time measurements. Note that this is up to the user to make sure calling barrier with their timers will not result in hangs. This can happen if for example the user adds a level 1 timer that is not called by all ranks. @@ -332,7 +355,12 @@ class CheckpointConfig: save: str | None = None """Output directory to save checkpoints to.""" - save_interval: int | None = field(default=None, metadata={"argparse_meta": {"arg_names": ["--save-interval", "--persistent-save-interval"]}}) + save_interval: int | None = field( + default=None, + metadata={ + "argparse_meta": {"arg_names": ["--save-interval", "--persistent-save-interval"]} + }, + ) """Number of iterations between persistent checkpoint saves.""" save_wgrads_interval: int | None = None @@ -436,7 +464,15 @@ class CheckpointConfig: The legacy format was deprecated on Feb 13, 2024. """ - fully_parallel_save: bool = field(default=True, metadata={"argparse_meta": {"arg_names": ["--no-ckpt-fully-parallel-save"], "dest": "ckpt_fully_parallel_save"}}) + fully_parallel_save: bool = field( + default=True, + metadata={ + "argparse_meta": { + "arg_names": ["--no-ckpt-fully-parallel-save"], + "dest": "ckpt_fully_parallel_save", + } + }, + ) """Disable applying full save parallelization across DP for distributed checkpoints. Depending on ckpt format might decrease the number of files in the checkpoint. Makes DistributedOptimizer checkpoint non-reshardable.""" @@ -460,10 +496,20 @@ class CheckpointConfig: async_ckpt_io_priority: Optional[int] = 3 """I/O scheduling class (0-3, 3=idle) for the async checkpoint writer process.""" - fully_parallel_load: bool = field(default=False, metadata={"argparse_meta": {"arg_names": ["--ckpt-fully-parallel-load"], "dest": "ckpt_fully_parallel_load"}}) + fully_parallel_load: bool = field( + default=False, + metadata={ + "argparse_meta": { + "arg_names": ["--ckpt-fully-parallel-load"], + "dest": "ckpt_fully_parallel_load", + } + }, + ) """Apply full load parallelization across DP for distributed checkpoints.""" - ckpt_fully_parallel_load_exchange_algo: Literal["broadcast", "gather_rounds", "gather_object"] = "broadcast" + ckpt_fully_parallel_load_exchange_algo: Literal[ + "broadcast", "gather_rounds", "gather_object" + ] = "broadcast" """Algorithm for fully parallel load of distributed checkpoints. "broadcast"(default): Broadcast the checkpoint from rank 0 to all other ranks. "gather_rounds": Gather the checkpoint from all ranks in rounds. @@ -533,8 +579,9 @@ class CheckpointConfig: def __post_init__(self): from megatron.training.utils import has_nvrx_installed - assert self.async_strategy in ["nvrx", "mcore"], \ + assert self.async_strategy in ["nvrx", "mcore"], ( f"async_strategy {self.async_strategy} is not supported. Available strategies: nvrx, mcore." + ) if self.async_save and self.ckpt_format in ["torch_dcp", "fsdp_dtensor"]: assert has_nvrx_installed(), ( diff --git a/flagscale/train/megatron/training/datasets/concated_indexed_dataset.py b/flagscale/train/megatron/training/datasets/concated_indexed_dataset.py index 88ce04f5ae..cc73825ffa 100644 --- a/flagscale/train/megatron/training/datasets/concated_indexed_dataset.py +++ b/flagscale/train/megatron/training/datasets/concated_indexed_dataset.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from typing import Optional, Tuple, Union import numpy diff --git a/flagscale/train/megatron/training/datasets/data_samplers.py b/flagscale/train/megatron/training/datasets/data_samplers.py index 0e48b3f086..ae17ce77ff 100644 --- a/flagscale/train/megatron/training/datasets/data_samplers.py +++ b/flagscale/train/megatron/training/datasets/data_samplers.py @@ -2,7 +2,6 @@ """Dataloaders.""" - import random import numpy as np @@ -23,9 +22,9 @@ def build_pretraining_data_loader(dataset, consumed_samples): return None args = get_args() - if hasattr(dataset, 'split'): + if hasattr(dataset, "split"): split = dataset.split - elif hasattr(dataset, 'index_split'): + elif hasattr(dataset, "index_split"): split = dataset.index_split else: split = None @@ -38,7 +37,7 @@ def build_pretraining_data_loader(dataset, consumed_samples): data_parallel_rank=mpu.get_data_parallel_rank(), data_parallel_size=mpu.get_data_parallel_world_size(), ) - elif args.dataloader_type == 'single': + elif args.dataloader_type == "single": if args.hybrid_context_parallel: batch_sampler = HybridCPMegatronPretrainingSampler( total_samples=len(dataset), @@ -46,7 +45,8 @@ def build_pretraining_data_loader(dataset, consumed_samples): micro_batch_size=args.micro_batch_size, global_batch_size=args.global_batch_size, data_parallel_rank=mpu.get_data_parallel_rank(), - data_parallel_size=mpu.get_data_parallel_world_size()) + data_parallel_size=mpu.get_data_parallel_world_size(), + ) else: # Megatron sampler batch_sampler = MegatronPretrainingSampler( @@ -54,8 +54,9 @@ def build_pretraining_data_loader(dataset, consumed_samples): consumed_samples=consumed_samples, micro_batch_size=args.micro_batch_size, data_parallel_rank=mpu.get_data_parallel_rank(), - data_parallel_size=mpu.get_data_parallel_world_size()) - elif args.dataloader_type == 'cyclic': + data_parallel_size=mpu.get_data_parallel_world_size(), + ) + elif args.dataloader_type == "cyclic": batch_sampler = MegatronPretrainingRandomSampler( dataset, total_samples=len(dataset), @@ -70,7 +71,7 @@ def build_pretraining_data_loader(dataset, consumed_samples): # torch-compatible dataloader and define samplers, if needed. return dataset else: - raise Exception('{} dataloader type is not supported.'.format(args.dataloader_type)) + raise Exception("{} dataloader type is not supported.".format(args.dataloader_type)) def worker_init_fn(_): import os @@ -91,12 +92,12 @@ def close_nvidia_fds(): if args.exit_signal_handler: DistributedSignalHandler(args.exit_signal).__enter__() - maybe_worker_init_fn = ( - worker_init_fn if args.num_workers > 0 else None - ) + maybe_worker_init_fn = worker_init_fn if args.num_workers > 0 else None # Torch dataloader. if args.hybrid_context_parallel: - extra_kwargs = {"collate_fn": lambda x: x,} + extra_kwargs = { + "collate_fn": lambda x: x, + } else: extra_kwargs = {} return torch.utils.data.DataLoader( @@ -109,6 +110,7 @@ def close_nvidia_fds(): **extra_kwargs, ) + class MegatronPretrainingSampler: """ Sampler for Megatron pretraining dataloaders that divides data samples across @@ -136,16 +138,16 @@ def __init__( self.drop_last = drop_last # Sanity checks. - assert self.total_samples > 0, 'no sample to consume: {}'.format(self.total_samples) - assert ( - self.consumed_samples < self.total_samples - ), 'no samples left to consume: {}, {}'.format(self.consumed_samples, self.total_samples) + assert self.total_samples > 0, "no sample to consume: {}".format(self.total_samples) + assert self.consumed_samples < self.total_samples, ( + "no samples left to consume: {}, {}".format(self.consumed_samples, self.total_samples) + ) assert self.micro_batch_size > 0 assert data_parallel_size > 0 - assert ( - self.data_parallel_rank < data_parallel_size - ), 'data_parallel_rank should be smaller than data size: {}, ' '{}'.format( - self.data_parallel_rank, data_parallel_size + assert self.data_parallel_rank < data_parallel_size, ( + "data_parallel_rank should be smaller than data size: {}, {}".format( + self.data_parallel_rank, data_parallel_size + ) ) def __len__(self): @@ -178,6 +180,7 @@ def __iter__(self): start_idx, end_idx = self.get_start_end_idx() yield batch[start_idx:end_idx] + class HybridCPMegatronPretrainingSampler(MegatronPretrainingSampler): """ Data sampler for hybrid context parallel (Hybrid CP) format. @@ -186,9 +189,24 @@ class HybridCPMegatronPretrainingSampler(MegatronPretrainingSampler): of the entire global batch. """ - def __init__(self, total_samples, consumed_samples, micro_batch_size, global_batch_size, - data_parallel_rank, data_parallel_size, drop_last=True): - super().__init__(total_samples, consumed_samples, micro_batch_size, data_parallel_rank, data_parallel_size, drop_last) + def __init__( + self, + total_samples, + consumed_samples, + micro_batch_size, + global_batch_size, + data_parallel_rank, + data_parallel_size, + drop_last=True, + ): + super().__init__( + total_samples, + consumed_samples, + micro_batch_size, + data_parallel_rank, + data_parallel_size, + drop_last, + ) self.global_batch_size = global_batch_size self.data_parallel_size = data_parallel_size self.num_micro_batches = self.global_batch_size // self.micro_batch_times_data_parallel_size @@ -197,7 +215,11 @@ def __len__(self): return self.total_samples def get_start_end_idx_global_batch(self): - start_idx = [self.data_parallel_rank * self.micro_batch_size + i * self.micro_batch_size * self.data_parallel_size for i in range(self.num_micro_batches)] + start_idx = [ + self.data_parallel_rank * self.micro_batch_size + + i * self.micro_batch_size * self.data_parallel_size + for i in range(self.num_micro_batches) + ] end_idx = [start_idx[i] + self.micro_batch_size for i in range(self.num_micro_batches)] return start_idx, end_idx @@ -210,7 +232,7 @@ def __iter__(self): start_idx, end_idx = self.get_start_end_idx_global_batch() global_batch_idx = [] for i in range(self.num_micro_batches): - global_batch_idx.extend(batch[start_idx[i]:end_idx[i]]) + global_batch_idx.extend(batch[start_idx[i] : end_idx[i]]) yield global_batch_idx batch = [] @@ -219,9 +241,10 @@ def __iter__(self): start_idx, end_idx = self.get_start_end_idx_global_batch() global_batch_idx = [] for i in range(self.num_micro_batches): - global_batch_idx.extend(batch[start_idx[i]:end_idx[i]]) + global_batch_idx.extend(batch[start_idx[i] : end_idx[i]]) yield global_batch_idx + class RandomSeedDataset(Dataset): """ A dataset wrapper that resets the random seed before each sample. @@ -294,13 +317,13 @@ def __init__( self.last_batch_size = self.total_samples % self.micro_batch_times_data_parallel_size # Sanity checks. - assert self.total_samples > 0, 'no sample to consume: {}'.format(self.total_samples) + assert self.total_samples > 0, "no sample to consume: {}".format(self.total_samples) assert self.micro_batch_size > 0 assert data_parallel_size > 0 - assert ( - self.data_parallel_rank < data_parallel_size - ), 'data_parallel_rank should be smaller than data size: {}, ' '{}'.format( - self.data_parallel_rank, data_parallel_size + assert self.data_parallel_rank < data_parallel_size, ( + "data_parallel_rank should be smaller than data size: {}, {}".format( + self.data_parallel_rank, data_parallel_size + ) ) def __len__(self): diff --git a/flagscale/train/megatron/training/datasets/fim_dataset.py b/flagscale/train/megatron/training/datasets/fim_dataset.py index 875f979c91..362430c4bc 100644 --- a/flagscale/train/megatron/training/datasets/fim_dataset.py +++ b/flagscale/train/megatron/training/datasets/fim_dataset.py @@ -248,7 +248,6 @@ def _permute( Maintain the same sample length (if transform creates a few extra tokens, drop them). """ if self.np_rng.binomial(1, fim_rate): # sample bernoulli dist - # Use remove_special_tokens=True so character-level boundaries and re-tokenization # are consistent; otherwise ids_to_text(..., None) keeps special tokens when # include_special_tokens=True, changing contents and breaking e.g. split_sample. diff --git a/flagscale/train/megatron/training/datasets/sft_dataset.py b/flagscale/train/megatron/training/datasets/sft_dataset.py index 9de5d2a52f..e0935b1bdb 100644 --- a/flagscale/train/megatron/training/datasets/sft_dataset.py +++ b/flagscale/train/megatron/training/datasets/sft_dataset.py @@ -36,9 +36,7 @@ def __init__(self, dataset_path: str) -> None: try: from datasets import load_dataset except ImportError: - raise ImportError( - "SFTDataset currently requires datasets library to be installed" - ) + raise ImportError("SFTDataset currently requires datasets library to be installed") self.dataset = load_dataset("json", data_files=dataset_path, split="all") def __len__(self) -> int: @@ -83,13 +81,12 @@ def _split_conversations(self, merged_conversations): split_conversations.append(current) current = [msg] # Then start the new conversation else: - current.append(msg) # Continue accumulating the current conversation + current.append(msg) # Continue accumulating the current conversation if current: # Store any remaining conversation split_conversations.append(current) return split_conversations def __getitem__(self, idx: int) -> Dict[str, Any]: - tokenizer = self.config.tokenizer pack_length = self.config.sequence_length @@ -99,7 +96,7 @@ def __getitem__(self, idx: int) -> Dict[str, Any]: def extend_with_padding(tokens, targets, positions, pad_len): tokens.extend([pad] * pad_len) targets.extend([pad] * pad_len) - positions.extend(range(positions[-1]+1, positions[-1]+1+pad_len)) + positions.extend(range(positions[-1] + 1, positions[-1] + 1 + pad_len)) pack_tokens = [] pack_targets = [] @@ -109,7 +106,6 @@ def extend_with_padding(tokens, targets, positions, pad_len): pad = tokenizer.pad # TODO(duncan): Track number of convs dropped and/or truncated and amount of end-padding for conversation in split_conversations: - tokens, targets = tokenizer.tokenize_conversation( conversation, return_target=True, add_generation_prompt=False ) @@ -117,7 +113,6 @@ def extend_with_padding(tokens, targets, positions, pad_len): tokens_list = tokens.tolist() targets_list = targets.tolist() - pack_tokens.extend(tokens_list) pack_targets.extend(targets_list) @@ -144,7 +139,7 @@ def extend_with_padding(tokens, targets, positions, pad_len): pack_targets = pack_targets[:max_body] pack_tokens.append(pad) pack_targets.append(pad) - pack_positions = pack_positions[:pack_length+1] + pack_positions = pack_positions[: pack_length + 1] # Note len({pack_tokens, pack_targets, pack_positions}) should be pack_length + 1 cu_seqlens[-1] = len(pack_tokens) - 1 break @@ -161,8 +156,8 @@ def extend_with_padding(tokens, targets, positions, pad_len): assert len(pack_positions) == pack_length + 1 # Align and convert to tensors - input_ids = torch.tensor(pack_tokens[:-1], dtype=torch.int64) - labels = torch.tensor(pack_targets[1:], dtype=torch.int64) + input_ids = torch.tensor(pack_tokens[:-1], dtype=torch.int64) + labels = torch.tensor(pack_targets[1:], dtype=torch.int64) position_ids = torch.tensor(pack_positions[:-1], dtype=torch.int64) # Loss mask. @@ -182,11 +177,11 @@ def extend_with_padding(tokens, targets, positions, pad_len): max_seqlen = adjacent_diffs.max() # max_seqlen is a 0-D tensor return { - 'tokens': input_ids, - 'labels': labels, + "tokens": input_ids, + "labels": labels, # 'attention_mask': attention_mask, # PyTorch collate cannot handle NoneType - 'loss_mask': loss_mask, - 'position_ids': position_ids, - 'cu_seqlens': cu_seqlens, - 'max_seqlen': max_seqlen, + "loss_mask": loss_mask, + "position_ids": position_ids, + "cu_seqlens": cu_seqlens, + "max_seqlen": max_seqlen, } diff --git a/flagscale/train/megatron/training/datasets/sft_dataset_fs.py b/flagscale/train/megatron/training/datasets/sft_dataset_fs.py index 40f6404092..4ad017ba1b 100644 --- a/flagscale/train/megatron/training/datasets/sft_dataset_fs.py +++ b/flagscale/train/megatron/training/datasets/sft_dataset_fs.py @@ -86,9 +86,9 @@ def _build_loss_mask_dataset(self) -> None: print(f"> Used Dataset: aux_loss_mask ...") if self.loss_mask_dataset is not None: - assert len(self.dataset) == len( - self.loss_mask_dataset - ), f"Samples are not equal, ({len(self.dataset)} != {len(self.loss_mask_dataset)})" + assert len(self.dataset) == len(self.loss_mask_dataset), ( + f"Samples are not equal, ({len(self.dataset)} != {len(self.loss_mask_dataset)})" + ) def __getitem__(self, idx: Optional[int]) -> Dict[str, torch.Tensor]: """Abstract method implementation diff --git a/flagscale/train/megatron/training/dgrad_logging.py b/flagscale/train/megatron/training/dgrad_logging.py index c046b4709f..bf92689768 100644 --- a/flagscale/train/megatron/training/dgrad_logging.py +++ b/flagscale/train/megatron/training/dgrad_logging.py @@ -24,8 +24,10 @@ def _get_linear_types(): TERowParallelLinear, TELayerNormColumnParallelLinear, ) - types.extend([TELinear, TEColumnParallelLinear, TERowParallelLinear, - TELayerNormColumnParallelLinear]) + + types.extend( + [TELinear, TEColumnParallelLinear, TERowParallelLinear, TELayerNormColumnParallelLinear] + ) except ImportError: pass @@ -35,9 +37,11 @@ def _get_linear_types(): TEColumnParallelGroupedLinear, TERowParallelGroupedLinear, ) + if TEGroupedLinear is not None: - types.extend([TEGroupedLinear, TEColumnParallelGroupedLinear, - TERowParallelGroupedLinear]) + types.extend( + [TEGroupedLinear, TEColumnParallelGroupedLinear, TERowParallelGroupedLinear] + ) except ImportError: pass @@ -49,7 +53,7 @@ def _get_linear_types(): class DataGradLogger: """Captures and saves gradients from all linear layers using backward hooks. - + NOTE: Right now, we only save the dgrads for the last microbatch in a batch on DP replica 0. The code below would need to be extended to save dgrads for all microbatches in a batch.""" @@ -60,6 +64,7 @@ def __init__(self, save_dir: str): def _make_hook(self, model_chunk_name: str, module_name: str): """Create a backward hook for a named module.""" + def hook(_, grad_input, grad_output): for idx, grad in enumerate(grad_output): if grad is not None: @@ -69,6 +74,7 @@ def hook(_, grad_input, grad_output): if grad is not None: grad_name = f"{module_name}/input{idx}" self._dgrads_state_dict[model_chunk_name][grad_name] = grad.detach().cpu() + return hook def save(self, iteration: int): diff --git a/flagscale/train/megatron/training/dist_signal_handler.py b/flagscale/train/megatron/training/dist_signal_handler.py index e9a4aceb17..3308910d47 100644 --- a/flagscale/train/megatron/training/dist_signal_handler.py +++ b/flagscale/train/megatron/training/dist_signal_handler.py @@ -4,13 +4,14 @@ import torch SIGNAL_MAP = { - 'SIGTERM': signal.SIGTERM, - 'SIGINT': signal.SIGINT, - 'SIGUSR1': signal.SIGUSR1, - 'SIGUSR2': signal.SIGUSR2 + "SIGTERM": signal.SIGTERM, + "SIGINT": signal.SIGINT, + "SIGUSR1": signal.SIGUSR1, + "SIGUSR2": signal.SIGUSR2, } from megatron.plugin.platform import get_platform + cur_platform = get_platform() from megatron.plugin.decorators import overridable @@ -27,21 +28,20 @@ def get_world_size(): @overridable def get_device(local_rank=None): backend = torch.distributed.get_backend() - if backend == 'nccl': + if backend == "nccl": if local_rank is None: device = torch.device(cur_platform.device_name()) else: - device = torch.device(f'{cur_platform.device_name()}:{local_rank}') - elif backend == 'gloo': - device = torch.device('cpu') + device = torch.device(f"{cur_platform.device_name()}:{local_rank}") + elif backend == "gloo": + device = torch.device("cpu") else: raise RuntimeError return device def all_gather_item(item, dtype, group=None, async_op=False, local_rank=None): - if not torch.distributed.is_available() or \ - not torch.distributed.is_initialized(): + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): return [item] device = get_device(local_rank) @@ -53,8 +53,7 @@ def all_gather_item(item, dtype, group=None, async_op=False, local_rank=None): tensor = torch.tensor([item], device=device, dtype=dtype) output_tensors = [ - torch.zeros(1, dtype=tensor.dtype, device=tensor.device) - for _ in range(group_size) + torch.zeros(1, dtype=tensor.dtype, device=tensor.device) for _ in range(group_size) ] torch.distributed.all_gather(output_tensors, tensor, group, async_op) output = [elem.item() for elem in output_tensors] @@ -62,13 +61,11 @@ def all_gather_item(item, dtype, group=None, async_op=False, local_rank=None): class DistributedSignalHandler: - def __init__(self, sig: str = 'SIGTERM'): + def __init__(self, sig: str = "SIGTERM"): self.sig = SIGNAL_MAP.get(sig, signal.SIGTERM) def signals_received(self): - all_received = all_gather_item( - self._signal_received, dtype=torch.int32 - ) + all_received = all_gather_item(self._signal_received, dtype=torch.int32) return all_received def __enter__(self): diff --git a/flagscale/train/megatron/training/extra_valid.py b/flagscale/train/megatron/training/extra_valid.py index 31e77ef337..b7366529ec 100644 --- a/flagscale/train/megatron/training/extra_valid.py +++ b/flagscale/train/megatron/training/extra_valid.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import math import torch @@ -9,10 +23,16 @@ from megatron.core.rerun_state_machine import RerunDataIterator from megatron.training.datasets.data_samplers import build_pretraining_data_loader from megatron.training import get_args, get_tokenizer, print_rank_0 -from megatron.training.global_vars import get_tensorboard_writer, get_wandb_writer, get_extra_valid_datasets, set_extra_valid_datasets +from megatron.training.global_vars import ( + get_tensorboard_writer, + get_wandb_writer, + get_extra_valid_datasets, + set_extra_valid_datasets, +) from megatron.training.utils import is_last_rank, print_rank_last from megatron.plugin.platform import get_platform + cur_platform = get_platform() @@ -67,9 +87,9 @@ def extra_valid_datasets_provider(data_path, num_samples): dataset_type, [0, num_samples, 0], is_dataset_built_on_rank, config ).build() - assert ( - extra_train_ds is None and extra_test_ds is None - ), "train_ds and test_ds should be None for extra_valid_ds" + assert extra_train_ds is None and extra_test_ds is None, ( + "train_ds and test_ds should be None for extra_valid_ds" + ) print_rank_0("> finished creating GPT datasets ...") @@ -84,9 +104,9 @@ def build_extra_valid_datasets(build_extra_valid_dataset_provider): if args.extra_valid_data_path is None: return [None] - assert ( - len(args.extra_valid_data_path) % 2 == 0 - ), "extra_valid_data_path format should be a list of weight, prefix and tag." + assert len(args.extra_valid_data_path) % 2 == 0, ( + "extra_valid_data_path format should be a list of weight, prefix and tag." + ) blend = args.extra_valid_data_path raw_num_tokens_per_dataset, raw_prefix_paths_per_dataset = zip( @@ -114,9 +134,9 @@ def build_extra_valid_datasets(build_extra_valid_dataset_provider): args.extra_prefix_paths_list = raw_prefix_paths_per_dataset args.extra_num_samples_list = num_samples_per_dataset - assert len(raw_prefix_paths_per_dataset) == len( - num_samples_per_dataset - ), f"Number of extra_valid data paths {len(raw_prefix_paths_per_dataset)} does not match number of extra_valid data samples {len(num_samples_per_dataset)}" + assert len(raw_prefix_paths_per_dataset) == len(num_samples_per_dataset), ( + f"Number of extra_valid data paths {len(raw_prefix_paths_per_dataset)} does not match number of extra_valid data samples {len(num_samples_per_dataset)}" + ) extra_valid_datasets = [] for path, num_samples in zip(raw_prefix_paths_per_dataset, num_samples_per_dataset): @@ -140,7 +160,6 @@ def build_extra_valid_data_loaders(build_extra_valid_dataset_provider): # Construct the data pipeline if is_distributed or mpu.get_tensor_model_parallel_rank() == 0: - # Build datasets if necessary. if get_extra_valid_datasets() is None: extra_valid_datasets = build_extra_valid_datasets(build_extra_valid_dataset_provider) @@ -156,7 +175,9 @@ def build_extra_valid_data_loaders(build_extra_valid_dataset_provider): # Flags to know if we need to do extra_validation. is_none = map(lambda _: _ is None, extra_valid_dataloaders) do_extra_valid = len(extra_valid_dataloaders) > 0 and not any(is_none) - flags = torch.tensor([int(do_extra_valid)], dtype=torch.long, device=cur_platform.device_name()) + flags = torch.tensor( + [int(do_extra_valid)], dtype=torch.long, device=cur_platform.device_name() + ) else: flags = torch.tensor([0], dtype=torch.long, device=cur_platform.device_name()) diff --git a/flagscale/train/megatron/training/fs_theoretical_memory_usage.py b/flagscale/train/megatron/training/fs_theoretical_memory_usage.py index a152ced3df..ddb2e74878 100644 --- a/flagscale/train/megatron/training/fs_theoretical_memory_usage.py +++ b/flagscale/train/megatron/training/fs_theoretical_memory_usage.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """ Computes theoretical memory footprint for model training referring to megatron. Reference: https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/theoretical_memory_usage.py @@ -354,7 +368,8 @@ def compute_weight_and_optimizer_memory(args, verbose=False): ) sparse_mlp_params_per_ep_rank_noddp = ( - 2 * args.hidden_size + 2 + * args.hidden_size # Shared MoE mlp * (shared_expert_ffn_hidden_size * gated_linear_multiplier) # gate diff --git a/flagscale/train/megatron/training/ft_integration.py b/flagscale/train/megatron/training/ft_integration.py index d89a4f9791..8845551f50 100644 --- a/flagscale/train/megatron/training/ft_integration.py +++ b/flagscale/train/megatron/training/ft_integration.py @@ -64,8 +64,10 @@ _MIN_ITERS_FOR_STEP_TIMEOUT_UPDATE = 16 from megatron.plugin.platform import get_platform + cur_platform = get_platform() + def get_rank_monitor_client() -> Optional[Any]: """Returns the underlying fault tolerance client instance @@ -95,7 +97,7 @@ def setup() -> None: cli = RankMonitorClient() global _GLOBAL_RANK_MONITOR_CLIENT - global_vars._ensure_var_is_not_initialized(_GLOBAL_RANK_MONITOR_CLIENT, 'rank monitor client') + global_vars._ensure_var_is_not_initialized(_GLOBAL_RANK_MONITOR_CLIENT, "rank monitor client") _GLOBAL_RANK_MONITOR_CLIENT = cli global _ft_state_path @@ -283,7 +285,7 @@ def _maybe_update_timeouts(is_closing_ft=False): if is_closing_ft: # with async checkpointing, "checkpointing" section is not updated, # but still we want to see some checkpointing to ensure that is was a complete run - if {'setup', 'step'}.issubset(sections_to_update) and _seen_checkpoints_cnt > 0: + if {"setup", "step"}.issubset(sections_to_update) and _seen_checkpoints_cnt > 0: update_out_of_section = True else: print_rank_0( @@ -306,13 +308,13 @@ def maybe_setup_simulated_fault() -> None: NOTE: This if for FT testing only """ - simulated_fault_desc = os.environ.get('FT_SIM_FAULT_DESC', None) + simulated_fault_desc = os.environ.get("FT_SIM_FAULT_DESC", None) if not simulated_fault_desc: return fault_type: Any # silence mypy rank_to_fail: Any # silence mypy base_delay: Any # silence mypy - fault_type, rank_to_fail, base_delay = simulated_fault_desc.split(';') + fault_type, rank_to_fail, base_delay = simulated_fault_desc.split(";") fault_type = fault_type.strip() rank_to_fail = rank_to_fail.strip() rank_to_fail = int(rank_to_fail) if rank_to_fail else None @@ -338,12 +340,12 @@ def maybe_setup_simulated_fault() -> None: # this rank is not going to simulate a fault, nothing more to do return - if fault_type == 'random': - fault_type = rng.choice(['rank_killed', 'rank_hung']) + if fault_type == "random": + fault_type = rng.choice(["rank_killed", "rank_hung"]) - if fault_type == 'rank_killed': + if fault_type == "rank_killed": target_pid = os.getpid() - elif fault_type == 'rank_hung': + elif fault_type == "rank_hung": target_pid = os.getpid() else: raise Exception(f"Unknown fault type {fault_type} expected one of: rank_killed, rank_hung.") @@ -361,7 +363,7 @@ def __fault_thread(): file=of, flush=True, ) - if fault_type == 'rank_hung': + if fault_type == "rank_hung": os.kill(target_pid, signal.SIGSTOP) else: os.kill(target_pid, signal.SIGKILL) diff --git a/flagscale/train/megatron/training/global_vars.py b/flagscale/train/megatron/training/global_vars.py index 530030c8f3..d4e8bf50df 100644 --- a/flagscale/train/megatron/training/global_vars.py +++ b/flagscale/train/megatron/training/global_vars.py @@ -13,8 +13,11 @@ from megatron.core.config import set_experimental_flag from megatron.core.energy_monitor import EnergyMonitor from megatron.core.jit import disable_jit_fuser -from megatron.core.num_microbatches_calculator import init_num_microbatches_calculator, unset_num_microbatches_calculator -from megatron.training.tokenizer import build_tokenizer ########## FlagScale Add ########## +from megatron.core.num_microbatches_calculator import ( + init_num_microbatches_calculator, + unset_num_microbatches_calculator, +) +from megatron.training.tokenizer import build_tokenizer ########## FlagScale Add ########## from megatron.training.dist_signal_handler import DistributedSignalHandler ########## FlagScale Begin ########## @@ -38,16 +41,15 @@ ########## FlagScale End ########## - def get_args(): """Return arguments.""" - _ensure_var_is_initialized(_GLOBAL_ARGS, 'args') + _ensure_var_is_initialized(_GLOBAL_ARGS, "args") return _GLOBAL_ARGS def get_tokenizer(): """Return tokenizer.""" - _ensure_var_is_initialized(_GLOBAL_TOKENIZER, 'tokenizer') + _ensure_var_is_initialized(_GLOBAL_TOKENIZER, "tokenizer") return _GLOBAL_TOKENIZER @@ -68,6 +70,7 @@ def get_one_logger(): to check if it is initialized.""" return _GLOBAL_ONE_LOGGER + def get_adlr_autoresume(): """ADLR autoresume object. It can be None so no need to check if it is initialized.""" @@ -76,23 +79,24 @@ def get_adlr_autoresume(): def get_timers(): """Return timers.""" - _ensure_var_is_initialized(_GLOBAL_TIMERS, 'timers') + _ensure_var_is_initialized(_GLOBAL_TIMERS, "timers") return _GLOBAL_TIMERS + def get_energy_monitor(): """Return energy monitor.""" - _ensure_var_is_initialized(_GLOBAL_ENERGY_MONITOR, 'energy monitor') + _ensure_var_is_initialized(_GLOBAL_ENERGY_MONITOR, "energy monitor") return _GLOBAL_ENERGY_MONITOR + def get_signal_handler(): - _ensure_var_is_initialized(_GLOBAL_SIGNAL_HANDLER, 'signal handler') + _ensure_var_is_initialized(_GLOBAL_SIGNAL_HANDLER, "signal handler") return _GLOBAL_SIGNAL_HANDLER def _set_signal_handler(exit_signal): - global _GLOBAL_SIGNAL_HANDLER - _ensure_var_is_not_initialized(_GLOBAL_SIGNAL_HANDLER, 'signal handler') + _ensure_var_is_not_initialized(_GLOBAL_SIGNAL_HANDLER, "signal handler") _GLOBAL_SIGNAL_HANDLER = DistributedSignalHandler(exit_signal).__enter__() @@ -107,6 +111,7 @@ def _graceful_shutdown(signum, frame): - Exits the process cleanly """ from megatron.training.utils import print_rank_0 + print_rank_0("\nTermination requested. Performing orderly shutdown.") try: @@ -130,7 +135,7 @@ def set_global_variables(args, build_tokenizer=True): assert args is not None - _ensure_var_is_not_initialized(_GLOBAL_ARGS, 'args') + _ensure_var_is_not_initialized(_GLOBAL_ARGS, "args") set_args(args) init_num_microbatches_calculator( @@ -171,12 +176,13 @@ def set_global_writers(args): assert args is not None - _ensure_var_is_initialized(_GLOBAL_ARGS, 'args') + _ensure_var_is_initialized(_GLOBAL_ARGS, "args") _set_tensorboard_writer(args) _set_one_logger(args) # build wandb writers for all processes in the dp group of the last rank from megatron.core import mpu + mp_groups = mpu.get_model_parallel_group() if not isinstance(mp_groups, list): mp_groups = [mp_groups] @@ -234,7 +240,7 @@ def set_args(args): def _build_tokenizer(args): """Initialize tokenizer.""" global _GLOBAL_TOKENIZER - _ensure_var_is_not_initialized(_GLOBAL_TOKENIZER, 'tokenizer') + _ensure_var_is_not_initialized(_GLOBAL_TOKENIZER, "tokenizer") _GLOBAL_TOKENIZER = build_tokenizer(args) return _GLOBAL_TOKENIZER @@ -248,62 +254,72 @@ def rebuild_tokenizer(args): def _set_tensorboard_writer(args): """Set tensorboard writer.""" global _GLOBAL_TENSORBOARD_WRITER - _ensure_var_is_not_initialized(_GLOBAL_TENSORBOARD_WRITER, - 'tensorboard writer') + _ensure_var_is_not_initialized(_GLOBAL_TENSORBOARD_WRITER, "tensorboard writer") - if hasattr(args, 'tensorboard_dir') and \ - args.tensorboard_dir and args.rank == (args.world_size - 1): + if ( + hasattr(args, "tensorboard_dir") + and args.tensorboard_dir + and args.rank == (args.world_size - 1) + ): try: from torch.utils.tensorboard import SummaryWriter - print('> setting tensorboard ...') + + print("> setting tensorboard ...") _GLOBAL_TENSORBOARD_WRITER = SummaryWriter( - log_dir=args.tensorboard_dir, - max_queue=args.tensorboard_queue_size) + log_dir=args.tensorboard_dir, max_queue=args.tensorboard_queue_size + ) except ModuleNotFoundError: - print('WARNING: TensorBoard writing requested but is not ' - 'available (are you using PyTorch 1.1.0 or later?), ' - 'no TensorBoard logs will be written.', flush=True) + print( + "WARNING: TensorBoard writing requested but is not " + "available (are you using PyTorch 1.1.0 or later?), " + "no TensorBoard logs will be written.", + flush=True, + ) def _set_wandb_writer(args): global _GLOBAL_WANDB_WRITER - _ensure_var_is_not_initialized(_GLOBAL_WANDB_WRITER, - 'wandb writer') - if getattr(args, 'wandb_project', ''): - if args.wandb_exp_name == '': + _ensure_var_is_not_initialized(_GLOBAL_WANDB_WRITER, "wandb writer") + if getattr(args, "wandb_project", ""): + if args.wandb_exp_name == "": raise ValueError("Please specify the wandb experiment name!") import wandb + if args.wandb_save_dir: save_dir = args.wandb_save_dir else: # Defaults to the save dir. - save_dir = os.path.join(args.save, 'wandb') + save_dir = os.path.join(args.save, "wandb") wandb_config = vars(args) - if 'kitchen_config_file' in wandb_config and wandb_config['kitchen_config_file'] is not None: + if ( + "kitchen_config_file" in wandb_config + and wandb_config["kitchen_config_file"] is not None + ): # Log the contents of the config for discovery of what the quantization # settings were. - with open(wandb_config['kitchen_config_file'], "r") as f: - wandb_config['kitchen_config_file_contents'] = f.read() + with open(wandb_config["kitchen_config_file"], "r") as f: + wandb_config["kitchen_config_file_contents"] = f.read() rank = torch.distributed.get_rank() save_dir = os.path.join(save_dir, "rank-{}".format(rank)) wandb_id = f"{args.wandb_exp_name}-rank-{rank}" - name = f'{args.wandb_exp_name}-rank-{rank}' + name = f"{args.wandb_exp_name}-rank-{rank}" group = args.wandb_exp_name wandb_kwargs = { - 'id': wandb_id, - 'dir': save_dir, - 'name': name, - 'group': group, - 'project': args.wandb_project, - 'mode': args.wandb_mode, - 'resume': 'auto', - 'config': wandb_config} + "id": wandb_id, + "dir": save_dir, + "name": name, + "group": group, + "project": args.wandb_project, + "mode": args.wandb_mode, + "resume": "auto", + "config": wandb_config, + } if args.wandb_entity: - wandb_kwargs['entity'] = args.wandb_entity - os.makedirs(wandb_kwargs['dir'], exist_ok=True) - if args.wandb_mode == 'online' or args.wandb_api_key: - assert args.wandb_api_key, 'wandb_api_key is required for online mode' + wandb_kwargs["entity"] = args.wandb_entity + os.makedirs(wandb_kwargs["dir"], exist_ok=True) + if args.wandb_mode == "online" or args.wandb_api_key: + assert args.wandb_api_key, "wandb_api_key is required for online mode" wandb.login(key=args.wandb_api_key) wandb.init(**wandb_kwargs) _GLOBAL_WANDB_WRITER = wandb @@ -311,41 +327,46 @@ def _set_wandb_writer(args): def _set_one_logger(args): global _GLOBAL_ONE_LOGGER - _ensure_var_is_not_initialized(_GLOBAL_ONE_LOGGER, 'one logger') + _ensure_var_is_not_initialized(_GLOBAL_ONE_LOGGER, "one logger") if args.enable_one_logger and args.rank == (args.world_size - 1): - if args.one_logger_async or getattr(args, 'wandb_project', ''): + if args.one_logger_async or getattr(args, "wandb_project", ""): one_logger_async = True else: one_logger_async = False try: from one_logger import OneLogger + config = { - 'project': args.one_logger_project, - 'name': args.one_logger_run_name, - 'async': one_logger_async, + "project": args.one_logger_project, + "name": args.one_logger_run_name, + "async": one_logger_async, } one_logger = OneLogger(config=config) _GLOBAL_ONE_LOGGER = one_logger except Exception: - print('WARNING: one_logger package is required to enable e2e metrics ' - 'tracking. please go to ' - 'https://confluence.nvidia.com/display/MLWFO/Package+Repositories' - ' for details to install it') + print( + "WARNING: one_logger package is required to enable e2e metrics " + "tracking. please go to " + "https://confluence.nvidia.com/display/MLWFO/Package+Repositories" + " for details to install it" + ) + def _set_adlr_autoresume(args): """Initialize ADLR autoresume.""" global _GLOBAL_ADLR_AUTORESUME - _ensure_var_is_not_initialized(_GLOBAL_ADLR_AUTORESUME, 'adlr autoresume') + _ensure_var_is_not_initialized(_GLOBAL_ADLR_AUTORESUME, "adlr autoresume") if args.adlr_autoresume: from megatron.training.utils import print_rank_0 - print_rank_0('enabling autoresume ...') - sys.path.append(os.environ.get('SUBMIT_SCRIPTS', '.')) + + print_rank_0("enabling autoresume ...") + sys.path.append(os.environ.get("SUBMIT_SCRIPTS", ".")) try: from userlib.auto_resume import AutoResume except ImportError: - print_rank_0('ADLR autoresume is not available, exiting ...') + print_rank_0("ADLR autoresume is not available, exiting ...") sys.exit() _GLOBAL_ADLR_AUTORESUME = AutoResume @@ -354,24 +375,26 @@ def _set_adlr_autoresume(args): def _set_timers(args): """Initialize timers.""" global _GLOBAL_TIMERS - _ensure_var_is_not_initialized(_GLOBAL_TIMERS, 'timers') + _ensure_var_is_not_initialized(_GLOBAL_TIMERS, "timers") _GLOBAL_TIMERS = Timers(args.timing_log_level, args.timing_log_option) + def _set_energy_monitor(args): """Initialize energy monitor.""" global _GLOBAL_ENERGY_MONITOR - _ensure_var_is_not_initialized(_GLOBAL_ENERGY_MONITOR, 'energy monitor') + _ensure_var_is_not_initialized(_GLOBAL_ENERGY_MONITOR, "energy monitor") _GLOBAL_ENERGY_MONITOR = EnergyMonitor() def _ensure_var_is_initialized(var, name): """Make sure the input variable is not None.""" - assert var is not None, '{} is not initialized.'.format(name) + assert var is not None, "{} is not initialized.".format(name) def _ensure_var_is_not_initialized(var, name): """Make sure the input variable is not None.""" - assert var is None, '{} is already initialized.'.format(name) + assert var is None, "{} is already initialized.".format(name) + def destroy_global_vars(): global _GLOBAL_ARGS @@ -425,4 +448,6 @@ def set_extra_valid_datasets(extra_valid_datasets): """Set extra_valid datasets.""" "" global _GLOBAL_EXTRA_VALID_DATASETS _GLOBAL_EXTRA_VALID_DATASETS = extra_valid_datasets + + ########## FlagScale End ########## diff --git a/flagscale/train/megatron/training/initialize.py b/flagscale/train/megatron/training/initialize.py index ff2786f28d..962f002f6f 100644 --- a/flagscale/train/megatron/training/initialize.py +++ b/flagscale/train/megatron/training/initialize.py @@ -1,6 +1,7 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. """Megatron initialization.""" + import logging import os import random @@ -22,7 +23,9 @@ RerunMode, initialize_rerun_state_machine, ) -from megatron.core.transformer.custom_layers.batch_invariant_kernels import enable_batch_invariant_mode +from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + enable_batch_invariant_mode, +) from megatron.core.utils import get_te_version, is_te_min_version, is_torch_min_version from megatron.training import get_adlr_autoresume, get_args, get_tensorboard_writer from megatron.training.utils import print_rank_0, warn_rank_0 @@ -42,8 +45,10 @@ logger = logging.getLogger(__name__) from megatron.plugin.platform import get_platform + cur_platform = get_platform() + def initialize_megatron( extra_args_provider=None, args_defaults={}, @@ -80,13 +85,15 @@ def initialize_megatron( args.exit_on_missing_checkpoint = True if args.use_checkpoint_args or args_defaults.get("use_checkpoint_args", False): - assert args.load is not None or args.pretrained_checkpoint is not None, "--use-checkpoint-args requires --load or --pretrained-checkpoint argument" + assert args.load is not None or args.pretrained_checkpoint is not None, ( + "--use-checkpoint-args requires --load or --pretrained-checkpoint argument" + ) assert args.non_persistent_ckpt_type != "local", ( "--use-checkpoint-args is not supported with --non_persistent_ckpt_type=local. " "Two-stage checkpoint loading is not implemented, and all arguments must be defined " "before initializing LocalCheckpointManager." ) - load_args_from_checkpoint(args, load_arg='pretrained_checkpoint') + load_args_from_checkpoint(args, load_arg="pretrained_checkpoint") load_args_from_checkpoint(args) ## FlagScale Begin: Pre Validate Arguments ## @@ -113,15 +120,15 @@ def initialize_megatron( setup_logging() if args.async_save and args.use_persistent_ckpt_worker: - init_persistent_async_worker(args.rank, 'forkserver') + init_persistent_async_worker(args.rank, "forkserver") # init rerun state def state_save_func(): - return {'rng_tracker_states': tensor_parallel.get_cuda_rng_tracker().get_states()} + return {"rng_tracker_states": tensor_parallel.get_cuda_rng_tracker().get_states()} def state_restore_func(state_dict): - if state_dict['rng_tracker_states']: - tensor_parallel.get_cuda_rng_tracker().set_states(state_dict['rng_tracker_states']) + if state_dict["rng_tracker_states"]: + tensor_parallel.get_cuda_rng_tracker().set_states(state_dict["rng_tracker_states"]) args = get_args() initialize_rerun_state_machine( @@ -158,7 +165,6 @@ def finish_mpu_init(): use_cudagraphable_rng=args.cuda_graph_impl != "none", ) - # Setup MoE aux loss scale value. if args.num_experts is not None: from megatron.core.transformer.moe.router import MoEAuxLossAutoScaler @@ -199,9 +205,9 @@ def finish_mpu_init(): # No continuation function return None + @overridable def _compile_dependencies(): - args = get_args() # ========================= @@ -209,6 +215,7 @@ def _compile_dependencies(): # ========================= # TODO: move this to ninja from megatron.plugin.utils import is_built_on_zero_rank + if is_built_on_zero_rank(): start_time = time.time() print("> compiling dataset index builder ...") @@ -216,13 +223,15 @@ def _compile_dependencies(): compile_helpers() print( - ">>> done with dataset index builder. Compilation time: {:.3f} " - "seconds".format(time.time() - start_time), + ">>> done with dataset index builder. Compilation time: {:.3f} seconds".format( + time.time() - start_time + ), flush=True, ) torch.distributed.barrier() + def _initialize_tp_communicators(): """initializing the communicators with user buffers for high-performance tensor-model-parallel communication overlap""" @@ -246,7 +255,7 @@ def _initialize_tp_communicators(): else: ub_cfgs = {} - if getattr(args, 'decoder_tp_comm_overlap', False): + if getattr(args, "decoder_tp_comm_overlap", False): input_shape = [ (args.decoder_seq_length * args.micro_batch_size) // args.context_parallel_size, args.hidden_size, @@ -257,11 +266,16 @@ def _initialize_tp_communicators(): args.hidden_size, ] - if is_te_min_version("2.7.0"): UserBufferQuantizationMode = te_module.base.UserBufferQuantizationMode - quantization_modes = [UserBufferQuantizationMode.FP8 if args.fp8 else UserBufferQuantizationMode.NONE] - if args.fp8 is not None and args.first_last_layers_bf16 and (args.num_layers_at_start_in_bf16 > 0 or args.num_layers_at_end_in_bf16 > 0): + quantization_modes = [ + UserBufferQuantizationMode.FP8 if args.fp8 else UserBufferQuantizationMode.NONE + ] + if ( + args.fp8 is not None + and args.first_last_layers_bf16 + and (args.num_layers_at_start_in_bf16 > 0 or args.num_layers_at_end_in_bf16 > 0) + ): quantization_modes.append(UserBufferQuantizationMode.NONE) # The process group with the target bootstrap backend is created in Transformer Engine. te_module.base.initialize_ub( @@ -281,12 +295,12 @@ def _initialize_tp_communicators(): bootstrap_backend=args.tp_comm_bootstrap_backend, ) else: - if args.tp_comm_bootstrap_backend != 'mpi': + if args.tp_comm_bootstrap_backend != "mpi": warnings.warn( f"Transformer Engine v{get_te_version()} supports only MPI bootstrap backend." ) # Create a MPI process group to help with TP communication overlap bootstrap. - create_group(backend='mpi', group_desc='TP_BOOTSTRAP_GROUP_MPI') + create_group(backend="mpi", group_desc="TP_BOOTSTRAP_GROUP_MPI") te_module.base.initialize_ub( shape=input_shape, @@ -302,13 +316,11 @@ def _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks, s device_count = cur_platform.device_count() if torch.distributed.is_initialized(): - print_rank_0("torch distributed is already initialized, skipping initialization ...") args.rank = torch.distributed.get_rank() args.world_size = torch.distributed.get_world_size() else: - print_rank_0("> initializing torch distributed ...") # Manually set the device ids. if device_count > 0: @@ -330,26 +342,26 @@ def _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks, s # so that the remaining defaults are applied consistently. _fr_path = ( args.flight_recorder_dump_path - or os.environ.get('TORCH_FR_DUMP_TEMP_FILE') - or os.environ.get('TORCH_NCCL_DEBUG_INFO_TEMP_FILE') + or os.environ.get("TORCH_FR_DUMP_TEMP_FILE") + or os.environ.get("TORCH_NCCL_DEBUG_INFO_TEMP_FILE") ) if _fr_path is not None: _fr_dump_prefix = _fr_path if os.path.isdir(_fr_path): - _fr_dump_prefix = os.path.join(_fr_path, '_dump_') + _fr_dump_prefix = os.path.join(_fr_path, "_dump_") warn_rank_0( "Flight recorder: using directory " f"'{_fr_path}' for dump path, appending per-rank prefix " f"'{_fr_dump_prefix}'." ) _fr_env_defaults = { - 'TORCH_FR_DUMP_TEMP_FILE': _fr_dump_prefix, - 'TORCH_NCCL_DEBUG_INFO_TEMP_FILE': _fr_dump_prefix, - 'TORCH_NCCL_TRACE_BUFFER_SIZE': str(args.flight_recorder_trace_buffer_size), - 'TORCH_NCCL_DUMP_ON_TIMEOUT': str(int(args.flight_recorder_dump_on_timeout)), - 'TORCH_INCLUDE_STACK_TRACE': str(int(args.flight_recorder_include_stack_trace)), - 'TORCH_INCLUDE_ONLY_ACTIVE': str(int(args.flight_recorder_include_only_active)), - 'TORCH_NCCL_EXTRA_DUMP_ON_EXEC': str(int(args.flight_recorder_extra_dump_on_exec)), + "TORCH_FR_DUMP_TEMP_FILE": _fr_dump_prefix, + "TORCH_NCCL_DEBUG_INFO_TEMP_FILE": _fr_dump_prefix, + "TORCH_NCCL_TRACE_BUFFER_SIZE": str(args.flight_recorder_trace_buffer_size), + "TORCH_NCCL_DUMP_ON_TIMEOUT": str(int(args.flight_recorder_dump_on_timeout)), + "TORCH_INCLUDE_STACK_TRACE": str(int(args.flight_recorder_include_stack_trace)), + "TORCH_INCLUDE_ONLY_ACTIVE": str(int(args.flight_recorder_include_only_active)), + "TORCH_NCCL_EXTRA_DUMP_ON_EXEC": str(int(args.flight_recorder_extra_dump_on_exec)), } for _var, _default in _fr_env_defaults.items(): if _var in os.environ: @@ -366,26 +378,29 @@ def _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks, s # Call the init process init_process_group_kwargs = { - 'backend': args.distributed_backend, - 'store': store, - 'world_size': args.world_size, - 'rank': args.rank, - 'timeout': timedelta(minutes=args.distributed_timeout_minutes), + "backend": args.distributed_backend, + "store": store, + "world_size": args.world_size, + "rank": args.rank, + "timeout": timedelta(minutes=args.distributed_timeout_minutes), } if args.enable_hetero and args.hetero_use_cpu_communication: # if not all(device_type == args.hetero_device_types[0] for device_type in args.hetero_device_types): # init_process_group_kwargs['backend'] = 'gloo' - init_process_group_kwargs['backend'] = "cpu:gloo" + init_process_group_kwargs["backend"] = "cpu:gloo" # TODO: @aoyulong the init_process_group will be hanging if the device_id is set # if packaging.version.Version(torch.__version__) >= packaging.version.Version("2.3.0"): # init_process_group_kwargs['device_id'] = device_id if args.fake_process_group: - assert is_torch_min_version("2.3.0"), "Fake process group is only supported with PyTorch 2.3.0 and above." + assert is_torch_min_version("2.3.0"), ( + "Fake process group is only supported with PyTorch 2.3.0 and above." + ) from torch.testing._internal.distributed.fake_pg import FakeStore + store = FakeStore() - init_process_group_kwargs['backend'] = 'fake' - init_process_group_kwargs['store'] = store + init_process_group_kwargs["backend"] = "fake" + init_process_group_kwargs["store"] = store torch.distributed.init_process_group(**init_process_group_kwargs) inprocess_restart.maybe_force_nccl_backend_init(device_id) @@ -416,7 +431,7 @@ def _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks, s engram_embedding_parallel_size=args.engram_embedding_parallel_size, distributed_timeout_minutes=args.distributed_timeout_minutes, nccl_communicator_config_path=args.nccl_communicator_config_path, - order='tp-cp-ep-dp-pp' if not args.use_tp_pp_dp_mapping else 'tp-cp-ep-pp-dp', + order="tp-cp-ep-dp-pp" if not args.use_tp_pp_dp_mapping else "tp-cp-ep-pp-dp", get_embedding_ranks=get_embedding_ranks, get_position_embedding_ranks=get_position_embedding_ranks, create_gloo_process_groups=args.use_gloo_process_groups, @@ -514,7 +529,9 @@ def _warmup_jit_function(): # Warmup fused bias+gelu bias = torch.rand( - args.ffn_hidden_size // args.tensor_model_parallel_size, dtype=dtype, device=cur_platform.device_name() + args.ffn_hidden_size // args.tensor_model_parallel_size, + dtype=dtype, + device=cur_platform.device_name(), ) input = torch.rand( ( @@ -551,7 +568,9 @@ def _warmup_jit_function(): dtype=dtype, device=cur_platform.device_name(), ) - bias = torch.rand((args.hidden_size), dtype=dtype, device=cur_platform.device_name()).expand_as(residual) + bias = torch.rand((args.hidden_size), dtype=dtype, device=cur_platform.device_name()).expand_as( + residual + ) dropout_rate = 0.1 # Warmup JIT fusions with the input grad_enable state of both forward # prop and recomputation @@ -577,7 +596,7 @@ def setup_logging() -> None: """ args = get_args() logging_level = None - env_logging_level = os.getenv('MEGATRON_LOGGING_LEVEL', None) + env_logging_level = os.getenv("MEGATRON_LOGGING_LEVEL", None) if env_logging_level is not None: logging_level = int(env_logging_level) if args.logging_level is not None: @@ -585,5 +604,5 @@ def setup_logging() -> None: if logging_level is not None: if is_rank0(): - logger.info(f'Setting logging level to {logging_level}') + logger.info(f"Setting logging level to {logging_level}") logging.getLogger().setLevel(logging_level) diff --git a/flagscale/train/megatron/training/inprocess_restart.py b/flagscale/train/megatron/training/inprocess_restart.py index 17e7bc59d6..a9b249c776 100644 --- a/flagscale/train/megatron/training/inprocess_restart.py +++ b/flagscale/train/megatron/training/inprocess_restart.py @@ -22,24 +22,28 @@ from . import arguments from megatron.plugin.platform import get_platform + cur_platform = get_platform() + def destroy_state(): from . import training + training.destroy_global_state() rerun_state_machine.destroy_rerun_state_machine() + def inprocess_restart(train, args): if inprocess is None: - warnings.warn('In-process restart is not available') + warnings.warn("In-process restart is not available") return train - if 'TORCH_CPP_LOG_LEVEL' not in os.environ or os.environ['TORCH_CPP_LOG_LEVEL'] not in ( - 'error', - 'fatal', + if "TORCH_CPP_LOG_LEVEL" not in os.environ or os.environ["TORCH_CPP_LOG_LEVEL"] not in ( + "error", + "fatal", ): warnings.warn( - 'Set TORCH_CPP_LOG_LEVEL=error to suppress c10d waitForInput timeout warning messages' + "Set TORCH_CPP_LOG_LEVEL=error to suppress c10d waitForInput timeout warning messages" ) # Layers represents a configuration for a layer of branches at a certain @@ -53,7 +57,7 @@ def inprocess_restart(train, args): flag=inprocess.rank_assignment.LayerFlag.RESERVE, ) ] - if args.inprocess_granularity == 'node': + if args.inprocess_granularity == "node": device_count = cur_platform.device_count() layers.append( @@ -84,9 +88,8 @@ def inprocess_restart(train, args): class AbortCheckpoint(inprocess.abort.Abort): def __init__(self, async_strategy): self.async_strategy = async_strategy - def __call__( - self, state: inprocess.state.FrozenState - ) -> inprocess.state.FrozenState: + + def __call__(self, state: inprocess.state.FrozenState) -> inprocess.state.FrozenState: reset_persistent_async_worker(self.async_strategy) return state @@ -101,8 +104,8 @@ def __call__( train = inprocess.Wrapper( store_kwargs={ - 'timeout': timedelta(seconds=300), - 'port': int(os.environ['MASTER_PORT']) + 2, + "timeout": timedelta(seconds=300), + "port": int(os.environ["MASTER_PORT"]) + 2, }, initialize=initialize, abort=abort, @@ -128,17 +131,16 @@ def __call__( def maybe_wrap_for_inprocess_restart(pretrain): - args = arguments.parse_args(ignore_unknown_args=True) if args.inprocess_restart: pretrain = inprocess_restart(pretrain, args) store = torch.distributed.TCPStore( - host_name=os.environ['MASTER_ADDR'], - port=int(os.environ['MASTER_PORT'])+1, - world_size=int(os.getenv('WORLD_SIZE', '1')), - is_master=(int(os.getenv('RANK', '0')) == 0), + host_name=os.environ["MASTER_ADDR"], + port=int(os.environ["MASTER_PORT"]) + 1, + world_size=int(os.getenv("WORLD_SIZE", "1")), + is_master=(int(os.getenv("RANK", "0")) == 0), timeout=timedelta(seconds=300), wait_for_workers=True, use_libuv=True, @@ -150,7 +152,6 @@ def maybe_wrap_for_inprocess_restart(pretrain): def maybe_force_nccl_backend_init(device_id): - args = get_args() # Inprocess uses destroy_process_group to terminate NCCL backend, which diff --git a/flagscale/train/megatron/training/one_logger_utils.py b/flagscale/train/megatron/training/one_logger_utils.py index fd18c569d1..c214b1e0f2 100644 --- a/flagscale/train/megatron/training/one_logger_utils.py +++ b/flagscale/train/megatron/training/one_logger_utils.py @@ -15,9 +15,17 @@ def get_timestamp_in_ms(): return round(time.time() * 1000.0) -def on_train_start(iteration, consumed_train_samples, train_samples, seq_length, - train_iters, save, async_save, log_throughput, - num_floating_point_operations_so_far): +def on_train_start( + iteration, + consumed_train_samples, + train_samples, + seq_length, + train_iters, + save, + async_save, + log_throughput, + num_floating_point_operations_so_far, +): """Function will be called at the start of train function to prepare and track E2E metrics. Args: @@ -38,43 +46,45 @@ def on_train_start(iteration, consumed_train_samples, train_samples, seq_length, with one_logger.get_context_manager(): # Get app train loop start time app_train_loop_start_time = get_timestamp_in_ms() - one_logger.store_set('app_train_loop_start_time', app_train_loop_start_time) + one_logger.store_set("app_train_loop_start_time", app_train_loop_start_time) # Set up initial values in store - one_logger.store_set('iteration_start', iteration) - one_logger.store_set('train_samples_start', consumed_train_samples) + one_logger.store_set("iteration_start", iteration) + one_logger.store_set("train_samples_start", consumed_train_samples) # Init accumulative metric values in one-logger store - one_logger.store_set('train_iterations_time_msecs_total', 0) - one_logger.store_set('tracked_train_iterations', iteration) - one_logger.store_set('validation_iterations_time_msecs_total', 0) - one_logger.store_set('tracked_validation_iterations', 0) - one_logger.store_set('save_checkpoint_count', 0) - one_logger.store_set('save_checkpoint_sync_time_total', 0.0) + one_logger.store_set("train_iterations_time_msecs_total", 0) + one_logger.store_set("tracked_train_iterations", iteration) + one_logger.store_set("validation_iterations_time_msecs_total", 0) + one_logger.store_set("tracked_validation_iterations", 0) + one_logger.store_set("save_checkpoint_count", 0) + one_logger.store_set("save_checkpoint_sync_time_total", 0.0) # Derive train_samples from iters for iteration-based training train_samples_target = train_samples or train_iters * args.global_batch_size train_tokens_target = seq_length * train_samples_target e2e_metrics = { - 'train_samples_start': consumed_train_samples, - 'train_iterations_start': iteration, - 'train_samples_target': train_samples_target, - 'train_iterations_target': train_iters, - 'train_tokens_target': train_tokens_target, - 'app_train_loop_start_time': app_train_loop_start_time, - 'is_save_checkpoint_enabled': save is not None, - 'save_checkpoint_strategy': 'async' if async_save else 'sync', + "train_samples_start": consumed_train_samples, + "train_iterations_start": iteration, + "train_samples_target": train_samples_target, + "train_iterations_target": train_iters, + "train_tokens_target": train_tokens_target, + "app_train_loop_start_time": app_train_loop_start_time, + "is_save_checkpoint_enabled": save is not None, + "save_checkpoint_strategy": "async" if async_save else "sync", } if log_throughput: - e2e_metrics.update({ - 'train_tflop_start': float(num_floating_point_operations_so_far) / (10**12), - }) + e2e_metrics.update( + { + "train_tflop_start": float(num_floating_point_operations_so_far) / (10**12), + } + ) one_logger.log_metrics(e2e_metrics) def _produce_e2e_metrics(log_throughput=False, throughput=None): - """ Generate APP metrics for E2E tracking + """Generate APP metrics for E2E tracking NOTE: always call this function after barrier call Args: @@ -85,17 +95,25 @@ def _produce_e2e_metrics(log_throughput=False, throughput=None): dict: all E2E metrics """ one_logger = get_one_logger() - + if one_logger: with one_logger.get_context_manager(): # Unpack and assign local vars - base_metrics = one_logger.store_get('get_e2e_base_metrics')() - (iteration, train_duration, eval_duration, eval_iterations, - total_flops_since_current_train_start, num_floating_point_operations_so_far, - consumed_train_samples, world_size, seq_length) = base_metrics.values() - - iteration_start = one_logger.store_get('iteration_start') - train_samples_start = one_logger.store_get('train_samples_start') + base_metrics = one_logger.store_get("get_e2e_base_metrics")() + ( + iteration, + train_duration, + eval_duration, + eval_iterations, + total_flops_since_current_train_start, + num_floating_point_operations_so_far, + consumed_train_samples, + world_size, + seq_length, + ) = base_metrics.values() + + iteration_start = one_logger.store_get("iteration_start") + train_samples_start = one_logger.store_get("train_samples_start") train_samples = consumed_train_samples - train_samples_start train_iterations = iteration - iteration_start @@ -105,45 +123,49 @@ def _produce_e2e_metrics(log_throughput=False, throughput=None): else: validation_iterations_time_msecs_avg = None - if not one_logger.store_has_key('first_logged_train_iterations_finish_time'): + if not one_logger.store_has_key("first_logged_train_iterations_finish_time"): one_logger.store_set( - 'first_logged_train_iterations_finish_time', - get_timestamp_in_ms() + "first_logged_train_iterations_finish_time", get_timestamp_in_ms() ) train_tokens = train_samples * seq_length e2e_metrics = { - 'first_logged_train_iterations_finish_time': \ - one_logger.store_get('first_logged_train_iterations_finish_time'), - 'train_iterations_end': iteration, - 'train_samples_end': consumed_train_samples, - 'train_iterations': train_iterations, - 'train_samples': train_samples, - 'train_iterations_time_msecs_avg': train_iterations_time_msecs_avg, - 'validation_iterations_time_total': eval_duration, - 'validation_iterations_time_msecs_avg': validation_iterations_time_msecs_avg, - 'train_tokens': train_tokens, - 'train_iterations_time_total': train_duration, - 'last_logged_train_iterations_finish_time': get_timestamp_in_ms(), + "first_logged_train_iterations_finish_time": one_logger.store_get( + "first_logged_train_iterations_finish_time" + ), + "train_iterations_end": iteration, + "train_samples_end": consumed_train_samples, + "train_iterations": train_iterations, + "train_samples": train_samples, + "train_iterations_time_msecs_avg": train_iterations_time_msecs_avg, + "validation_iterations_time_total": eval_duration, + "validation_iterations_time_msecs_avg": validation_iterations_time_msecs_avg, + "train_tokens": train_tokens, + "train_iterations_time_total": train_duration, + "last_logged_train_iterations_finish_time": get_timestamp_in_ms(), } if log_throughput: if train_duration: - train_throughput_per_gpu = total_flops_since_current_train_start / (train_duration * 10**12 * world_size) + train_throughput_per_gpu = total_flops_since_current_train_start / ( + train_duration * 10**12 * world_size + ) else: train_throughput_per_gpu = 0.0 - train_throughput_per_gpu_max = one_logger.store_get('train_throughput_per_gpu_max') + train_throughput_per_gpu_max = one_logger.store_get("train_throughput_per_gpu_max") if throughput: train_throughput_per_gpu_max = max(throughput, train_throughput_per_gpu_max) - one_logger.store_set('train_throughput_per_gpu_max', train_throughput_per_gpu_max) + one_logger.store_set( + "train_throughput_per_gpu_max", train_throughput_per_gpu_max + ) throughput_metrics = { - 'train_tflop_end': float(num_floating_point_operations_so_far) / (10**12), - 'train_tflop': float(total_flops_since_current_train_start) / (10**12), - 'train_throughput_per_gpu': train_throughput_per_gpu, - 'train_throughput_per_gpu_max': train_throughput_per_gpu_max, + "train_tflop_end": float(num_floating_point_operations_so_far) / (10**12), + "train_tflop": float(total_flops_since_current_train_start) / (10**12), + "train_throughput_per_gpu": train_throughput_per_gpu, + "train_throughput_per_gpu_max": train_throughput_per_gpu_max, } e2e_metrics.update(throughput_metrics) @@ -151,58 +173,70 @@ def _produce_e2e_metrics(log_throughput=False, throughput=None): # Minimal train iteration duration current_train_iterations_time_msecs_total = train_duration * 1000.0 current_train_iteration = iteration - prev_train_iterations_time_msecs_total = one_logger.store_get('train_iterations_time_msecs_total') - tracked_train_iterations = one_logger.store_get('tracked_train_iterations') + prev_train_iterations_time_msecs_total = one_logger.store_get( + "train_iterations_time_msecs_total" + ) + tracked_train_iterations = one_logger.store_get("tracked_train_iterations") if current_train_iteration > tracked_train_iterations: train_iterations_time_msecs = ( - (current_train_iterations_time_msecs_total - prev_train_iterations_time_msecs_total) / - (current_train_iteration - tracked_train_iterations) - ) + current_train_iterations_time_msecs_total + - prev_train_iterations_time_msecs_total + ) / (current_train_iteration - tracked_train_iterations) - if not one_logger.store_has_key('train_iterations_time_msecs_min'): + if not one_logger.store_has_key("train_iterations_time_msecs_min"): train_iterations_time_msecs_min = train_iterations_time_msecs else: train_iterations_time_msecs_min = min( - one_logger.store_get('train_iterations_time_msecs_min'), - train_iterations_time_msecs + one_logger.store_get("train_iterations_time_msecs_min"), + train_iterations_time_msecs, ) - one_logger.store_set('train_iterations_time_msecs_min', train_iterations_time_msecs_min) - one_logger.store_set('train_iterations_time_msecs_total', current_train_iterations_time_msecs_total) - one_logger.store_set('tracked_train_iterations', current_train_iteration) + one_logger.store_set( + "train_iterations_time_msecs_min", train_iterations_time_msecs_min + ) + one_logger.store_set( + "train_iterations_time_msecs_total", current_train_iterations_time_msecs_total + ) + one_logger.store_set("tracked_train_iterations", current_train_iteration) - e2e_metrics.update({ - 'train_iterations_time_msecs_min': train_iterations_time_msecs_min - }) + e2e_metrics.update( + {"train_iterations_time_msecs_min": train_iterations_time_msecs_min} + ) # Minimal validation iteration duration current_validation_iterations_time_msecs_total = eval_duration * 1000.0 current_validation_iteration = eval_iterations - prev_validation_iterations_time_msecs_total = \ - one_logger.store_get('validation_iterations_time_msecs_total') - tracked_validation_iterations = one_logger.store_get('tracked_validation_iterations') + prev_validation_iterations_time_msecs_total = one_logger.store_get( + "validation_iterations_time_msecs_total" + ) + tracked_validation_iterations = one_logger.store_get("tracked_validation_iterations") if current_validation_iteration > tracked_validation_iterations: validation_iterations_time_msecs = ( - (current_validation_iterations_time_msecs_total - prev_validation_iterations_time_msecs_total) / - (current_validation_iteration - tracked_validation_iterations) - ) + current_validation_iterations_time_msecs_total + - prev_validation_iterations_time_msecs_total + ) / (current_validation_iteration - tracked_validation_iterations) # Cache minimal validation iteration duration - if not one_logger.store_has_key('validation_iterations_time_msecs_min'): + if not one_logger.store_has_key("validation_iterations_time_msecs_min"): validation_iterations_time_msecs_min = validation_iterations_time_msecs else: validation_iterations_time_msecs_min = min( - one_logger.store_get('validation_iterations_time_msecs_min'), - validation_iterations_time_msecs + one_logger.store_get("validation_iterations_time_msecs_min"), + validation_iterations_time_msecs, ) - one_logger.store_set('validation_iterations_time_msecs_min', validation_iterations_time_msecs_min) - one_logger.store_set('validation_iterations_time_msecs_total', current_validation_iterations_time_msecs_total) - one_logger.store_set('tracked_validation_iterations', current_validation_iteration) + one_logger.store_set( + "validation_iterations_time_msecs_min", validation_iterations_time_msecs_min + ) + one_logger.store_set( + "validation_iterations_time_msecs_total", + current_validation_iterations_time_msecs_total, + ) + one_logger.store_set("tracked_validation_iterations", current_validation_iteration) - e2e_metrics.update({ - 'validation_iterations_time_msecs_min': validation_iterations_time_msecs_min - }) + e2e_metrics.update( + {"validation_iterations_time_msecs_min": validation_iterations_time_msecs_min} + ) return e2e_metrics @@ -233,69 +267,84 @@ def on_save_checkpoint_start(async_save): dict: productive metrics to be stored to DB after ckpt succeeds """ one_logger = get_one_logger() - + if one_logger: with one_logger.get_context_manager(): # Unpack and assign local vars - base_metrics = one_logger.store_get('get_e2e_base_metrics')() - (iteration, train_duration, eval_duration, eval_iterations, - total_flops_since_current_train_start, num_floating_point_operations_so_far, - consumed_train_samples, world_size, seq_length) = base_metrics.values() - - save_checkpoint_count = one_logger.store_get('save_checkpoint_count') + 1 - one_logger.store_set('save_checkpoint_count', save_checkpoint_count) - one_logger.log_metrics({ - 'train_iterations_save_checkpoint_end': iteration, - 'save_checkpoint_count': save_checkpoint_count, - }) + base_metrics = one_logger.store_get("get_e2e_base_metrics")() + ( + iteration, + train_duration, + eval_duration, + eval_iterations, + total_flops_since_current_train_start, + num_floating_point_operations_so_far, + consumed_train_samples, + world_size, + seq_length, + ) = base_metrics.values() + + save_checkpoint_count = one_logger.store_get("save_checkpoint_count") + 1 + one_logger.store_set("save_checkpoint_count", save_checkpoint_count) + one_logger.log_metrics( + { + "train_iterations_save_checkpoint_end": iteration, + "save_checkpoint_count": save_checkpoint_count, + } + ) productive_metrics = { - 'train_tflop_productive_end': float(num_floating_point_operations_so_far) / (10**12), - 'train_iterations_productive_end': iteration, - 'train_samples_productive_end': consumed_train_samples, - 'train_iterations_time_total_productive': train_duration, - 'validation_iterations_time_total_productive': eval_duration, + "train_tflop_productive_end": float(num_floating_point_operations_so_far) + / (10**12), + "train_iterations_productive_end": iteration, + "train_samples_productive_end": consumed_train_samples, + "train_iterations_time_total_productive": train_duration, + "validation_iterations_time_total_productive": eval_duration, } if async_save: - productive_metrics.update({ - 'save_checkpoint_async_count': save_checkpoint_count, - }) + productive_metrics.update( + { + "save_checkpoint_async_count": save_checkpoint_count, + } + ) return productive_metrics - + def on_pretrain_start(): - """ Function to be called at the start of pretrain function to track E2E meta data - """ + """Function to be called at the start of pretrain function to track E2E meta data""" args = get_args() one_logger = get_one_logger() if one_logger: with one_logger.get_context_manager(): - job_name = os.environ.get('SLURM_JOB_NAME', None) - app_tag_run_name = job_name if not args.app_tag_run_name else args.app_tag_run_name + job_name = os.environ.get("SLURM_JOB_NAME", None) + app_tag_run_name = job_name if not args.app_tag_run_name else args.app_tag_run_name app_tag_run_version = args.app_tag_run_version - one_logger.store_set('app_tag_run_name', app_tag_run_name) - one_logger.store_set('app_tag_run_version', app_tag_run_version) - one_logger.store_set('train_throughput_per_gpu_max', 0.0) - - one_logger.log_metrics({ - 'train_iterations_warmup': 5, - 'data_parallel_size' : args.data_parallel_size, - 'context_parallel_size': args.context_parallel_size, - 'global_batch_size': args.global_batch_size, - 'micro_batch_size': args.micro_batch_size, - 'pipeline_model_parallel_size': args.pipeline_model_parallel_size, - 'tensor_model_parallel_size': args.tensor_model_parallel_size, - 'expert_model_parallel_size' : args.expert_model_parallel_size, - 'world_size': args.world_size, - 'model_seq_length': args.seq_length, - 'app_tag_run_name': app_tag_run_name, - 'app_tag_run_version': app_tag_run_version, - 'is_log_throughput_enabled': args.log_throughput, - 'app_run_type': 'training', - 'summary_data_schema_version': '1.0.0', - 'app_metrics_feature_tags': 'full', - 'one_logger_utils_version': _one_logger_utils_version, - }) + one_logger.store_set("app_tag_run_name", app_tag_run_name) + one_logger.store_set("app_tag_run_version", app_tag_run_version) + one_logger.store_set("train_throughput_per_gpu_max", 0.0) + + one_logger.log_metrics( + { + "train_iterations_warmup": 5, + "data_parallel_size": args.data_parallel_size, + "context_parallel_size": args.context_parallel_size, + "global_batch_size": args.global_batch_size, + "micro_batch_size": args.micro_batch_size, + "pipeline_model_parallel_size": args.pipeline_model_parallel_size, + "tensor_model_parallel_size": args.tensor_model_parallel_size, + "expert_model_parallel_size": args.expert_model_parallel_size, + "world_size": args.world_size, + "model_seq_length": args.seq_length, + "app_tag_run_name": app_tag_run_name, + "app_tag_run_version": app_tag_run_version, + "is_log_throughput_enabled": args.log_throughput, + "app_run_type": "training", + "summary_data_schema_version": "1.0.0", + "app_metrics_feature_tags": "full", + "one_logger_utils_version": _one_logger_utils_version, + } + ) + def track_config_flags(train_iters, skip_train, do_train, do_valid, do_test, dataloader_type): """Track flags about train/validation/test enablement @@ -313,11 +362,14 @@ def track_config_flags(train_iters, skip_train, do_train, do_valid, do_test, dat with one_logger.get_context_manager(): # Track if training is enabled. Can only be done once args.do_train is assigned after dataloader is built. train_enabled = train_iters and (not skip_train) and do_train and train_iters > 0 - one_logger.log_metrics({ - 'is_train_iterations_enabled': train_enabled, - 'is_validation_iterations_enabled': bool(do_valid), - 'is_test_iterations_enabled': bool(do_test), - }) + one_logger.log_metrics( + { + "is_train_iterations_enabled": train_enabled, + "is_validation_iterations_enabled": bool(do_valid), + "is_test_iterations_enabled": bool(do_test), + } + ) + def on_save_checkpoint_success(productive_metrics, async_save): """Function to be called after checkpointing succeeds and checkpoint is persisted for storing productive metrics @@ -331,41 +383,45 @@ def on_save_checkpoint_success(productive_metrics, async_save): if one_logger: with one_logger.get_context_manager(): # Accumulate train_iterations_time_total_productive for current iteration - prod_iteration = productive_metrics['train_iterations_productive_end'] + prod_iteration = productive_metrics["train_iterations_productive_end"] # Log start timestamp of first iteration that was successfully checkpointed - if not one_logger.store_has_key('first_checkpoint_success'): - app_train_loop_start_time = one_logger.store_get('app_train_loop_start_time') - one_logger.store_set('first_checkpoint_success', True) - one_logger.log_metrics({ - 'first_saved_train_iterations_start_time': app_train_loop_start_time - }) + if not one_logger.store_has_key("first_checkpoint_success"): + app_train_loop_start_time = one_logger.store_get("app_train_loop_start_time") + one_logger.store_set("first_checkpoint_success", True) + one_logger.log_metrics( + {"first_saved_train_iterations_start_time": app_train_loop_start_time} + ) # Handle possible out-of-order async checkpoint callbacks need_update = True - if one_logger.store_has_key('iters_prod_max'): - need_update = prod_iteration > one_logger.store_get('iters_prod_max') + if one_logger.store_has_key("iters_prod_max"): + need_update = prod_iteration > one_logger.store_get("iters_prod_max") if need_update: # Update cache - one_logger.store_set('iters_prod_max', prod_iteration) + one_logger.store_set("iters_prod_max", prod_iteration) if async_save: - save_checkpoint_sync_time_total_productive = \ - one_logger.store_pop(f'save_checkpoint_sync_time_total_productive:{prod_iteration}') - last_successful_save_checkpoint_sync_finish_time = \ - one_logger.store_pop(f'save_checkpoint_sync_finish_time:{prod_iteration}') + save_checkpoint_sync_time_total_productive = one_logger.store_pop( + f"save_checkpoint_sync_time_total_productive:{prod_iteration}" + ) + last_successful_save_checkpoint_sync_finish_time = one_logger.store_pop( + f"save_checkpoint_sync_finish_time:{prod_iteration}" + ) # Update productive metrics and log to DB - productive_metrics.update({ - 'save_checkpoint_sync_time_total_productive': save_checkpoint_sync_time_total_productive, - 'last_successful_save_checkpoint_sync_finish_time': last_successful_save_checkpoint_sync_finish_time - }) + productive_metrics.update( + { + "save_checkpoint_sync_time_total_productive": save_checkpoint_sync_time_total_productive, + "last_successful_save_checkpoint_sync_finish_time": last_successful_save_checkpoint_sync_finish_time, + } + ) one_logger.log_metrics(productive_metrics) def on_save_checkpoint_end(save_checkpoint_duration, current_iteration, async_save): """Function to be called after checkpointing ends - + Args: save_checkpoint_duration (float): duration of current save checkpoint process current_iteration (int): current train iteration step number @@ -377,60 +433,68 @@ def on_save_checkpoint_end(save_checkpoint_duration, current_iteration, async_sa save_checkpoint_sync_finish_time = get_timestamp_in_ms() # Track finish timestamp of the sync part of first successful save checkpoint - if (one_logger.store_has_key('first_checkpoint_success') - and not one_logger.store_has_key('first_successful_checkpoint_end')): - one_logger.store_set('first_successful_checkpoint_end', True) - one_logger.log_metrics({ - 'first_successful_save_checkpoint_sync_finish_time': save_checkpoint_sync_finish_time - }) + if one_logger.store_has_key( + "first_checkpoint_success" + ) and not one_logger.store_has_key("first_successful_checkpoint_end"): + one_logger.store_set("first_successful_checkpoint_end", True) + one_logger.log_metrics( + { + "first_successful_save_checkpoint_sync_finish_time": save_checkpoint_sync_finish_time + } + ) - save_checkpoint_sync_count = one_logger.store_get('save_checkpoint_count') + save_checkpoint_sync_count = one_logger.store_get("save_checkpoint_count") # accumulate total sync checkpointing duration - save_checkpoint_sync_time_total = \ - one_logger.store_get('save_checkpoint_sync_time_total') + save_checkpoint_duration - one_logger.store_set('save_checkpoint_sync_time_total', save_checkpoint_sync_time_total) + save_checkpoint_sync_time_total = ( + one_logger.store_get("save_checkpoint_sync_time_total") + save_checkpoint_duration + ) + one_logger.store_set("save_checkpoint_sync_time_total", save_checkpoint_sync_time_total) e2e_metrics = {} if async_save: # Cache total sync checkpointing duration one_logger.store_set( - f'save_checkpoint_sync_time_total_productive:{current_iteration}', - save_checkpoint_sync_time_total + f"save_checkpoint_sync_time_total_productive:{current_iteration}", + save_checkpoint_sync_time_total, ) # Cache finish time for current iteration - one_logger.store_set(f'save_checkpoint_sync_finish_time:{current_iteration}', - save_checkpoint_sync_finish_time) + one_logger.store_set( + f"save_checkpoint_sync_finish_time:{current_iteration}", + save_checkpoint_sync_finish_time, + ) else: - e2e_metrics.update({ - # Track productive total time directly for sync ckpt - 'save_checkpoint_sync_time_total_productive': save_checkpoint_sync_time_total, - 'last_successful_save_checkpoint_sync_finish_time': save_checkpoint_sync_finish_time, - }) + e2e_metrics.update( + { + # Track productive total time directly for sync ckpt + "save_checkpoint_sync_time_total_productive": save_checkpoint_sync_time_total, + "last_successful_save_checkpoint_sync_finish_time": save_checkpoint_sync_finish_time, + } + ) # Tracking min & max value sync checkpointing duration # For the first comparison - if not one_logger.store_has_key('save_checkpoint_sync_time_max'): - one_logger.store_set('save_checkpoint_sync_time_max', save_checkpoint_duration) - if not one_logger.store_has_key('save_checkpoint_sync_time_min'): - one_logger.store_set('save_checkpoint_sync_time_min', save_checkpoint_duration) + if not one_logger.store_has_key("save_checkpoint_sync_time_max"): + one_logger.store_set("save_checkpoint_sync_time_max", save_checkpoint_duration) + if not one_logger.store_has_key("save_checkpoint_sync_time_min"): + one_logger.store_set("save_checkpoint_sync_time_min", save_checkpoint_duration) save_checkpoint_sync_time_max = max( - one_logger.store_get('save_checkpoint_sync_time_max'), - save_checkpoint_duration + one_logger.store_get("save_checkpoint_sync_time_max"), save_checkpoint_duration ) save_checkpoint_sync_time_min = min( - one_logger.store_get('save_checkpoint_sync_time_min'), - save_checkpoint_duration + one_logger.store_get("save_checkpoint_sync_time_min"), save_checkpoint_duration + ) + one_logger.store_set("save_checkpoint_sync_time_max", save_checkpoint_sync_time_max) + one_logger.store_set("save_checkpoint_sync_time_min", save_checkpoint_sync_time_min) + e2e_metrics.update( + { + "save_checkpoint_sync_count": save_checkpoint_sync_count, + "save_checkpoint_sync_time_max": save_checkpoint_sync_time_max, + "save_checkpoint_sync_time_min": save_checkpoint_sync_time_min, + "save_checkpoint_sync_time_total": save_checkpoint_sync_time_total, + } ) - one_logger.store_set('save_checkpoint_sync_time_max', save_checkpoint_sync_time_max) - one_logger.store_set('save_checkpoint_sync_time_min', save_checkpoint_sync_time_min) - e2e_metrics.update({ - 'save_checkpoint_sync_count': save_checkpoint_sync_count, - 'save_checkpoint_sync_time_max': save_checkpoint_sync_time_max, - 'save_checkpoint_sync_time_min': save_checkpoint_sync_time_min, - 'save_checkpoint_sync_time_total': save_checkpoint_sync_time_total, - }) one_logger.log_metrics(e2e_metrics) @@ -446,16 +510,16 @@ def track_app_tag(batch_size, world_size, seq_length): one_logger = get_one_logger() if one_logger: with one_logger.get_context_manager(): - app_tag_run_name = one_logger.store_get('app_tag_run_name') - app_tag_run_version = one_logger.store_get('app_tag_run_version') - current_app_tag = (f'{app_tag_run_name}_{app_tag_run_version}_{batch_size}' - f'_{world_size}_{seq_length}') + app_tag_run_name = one_logger.store_get("app_tag_run_name") + app_tag_run_version = one_logger.store_get("app_tag_run_version") + current_app_tag = ( + f"{app_tag_run_name}_{app_tag_run_version}_{batch_size}_{world_size}_{seq_length}" + ) one_logger.log_app_tag(current_app_tag) def finish(): - """Flush E2E metrics to remote server - """ + """Flush E2E metrics to remote server""" one_logger = get_one_logger() if one_logger: with one_logger.get_context_manager(): diff --git a/flagscale/train/megatron/training/peft/__init__.py b/flagscale/train/megatron/training/peft/__init__.py index cc3626fd75..6ac0622b24 100644 --- a/flagscale/train/megatron/training/peft/__init__.py +++ b/flagscale/train/megatron/training/peft/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """PEFT (Parameter-Efficient Fine-Tuning) module. This module provides implementations of various PEFT methods for fine-tuning @@ -7,5 +21,4 @@ from megatron.training.peft.peft import PEFT, AdapterWrapper from megatron.training.peft.lora import LoRA -__all__ = ['PEFT', 'AdapterWrapper', 'LoRA'] - +__all__ = ["PEFT", "AdapterWrapper", "LoRA"] diff --git a/flagscale/train/megatron/training/peft/lora.py b/flagscale/train/megatron/training/peft/lora.py index 08878e869e..2f7475bda4 100644 --- a/flagscale/train/megatron/training/peft/lora.py +++ b/flagscale/train/megatron/training/peft/lora.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import logging from collections import namedtuple @@ -48,7 +62,7 @@ def forward( return linear_output + adapter_output, bias -class LoRA(PEFT, peft_type='lora'): +class LoRA(PEFT, peft_type="lora"): """ Implements the LoRA (Low-Rank Adaptation) module for parameter-efficient fine-tuning. @@ -110,7 +124,7 @@ def freeze_model(self, model: nn.Module): for name, param in model.named_parameters(): param.requires_grad = False for name, param in model.named_parameters(): - if 'adapter' in name: + if "adapter" in name: param.requires_grad = True def load_state_dict_pre_hooks(self, model: nn.Module): diff --git a/flagscale/train/megatron/training/peft/peft.py b/flagscale/train/megatron/training/peft/peft.py index 1fa424ac42..0185721611 100644 --- a/flagscale/train/megatron/training/peft/peft.py +++ b/flagscale/train/megatron/training/peft/peft.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from abc import ABC, abstractmethod from typing import Any, Callable, Dict, List, Optional, Tuple, Type @@ -57,7 +71,7 @@ def transform(self, module, name=None, prefix=None): def apply_transform(self, model: nn.Module): for full_name, module in model.named_modules(): - prefix, name = full_name.rsplit('.', 1) if '.' in full_name else ('', full_name) + prefix, name = full_name.rsplit(".", 1) if "." in full_name else ("", full_name) replaced_module = self.transform(module, name, prefix) if replaced_module == module: continue @@ -102,9 +116,9 @@ def base_linear_forward(self, x, *args, **kwargs): layernorm_output is different from input x only when linear layer is LayerNormColumnParallelLinear. """ linear_output = self.to_wrap(x, *args, **kwargs) - assert isinstance( - linear_output, tuple - ), f"{self.to_wrap} should return a tuple but instead returns {linear_output}" + assert isinstance(linear_output, tuple), ( + f"{self.to_wrap} should return a tuple but instead returns {linear_output}" + ) """ Four cases for the wrapped module's return values 1. nothing: (out, None) 2. return_bias: (out, bias) @@ -122,7 +136,7 @@ def base_linear_forward(self, x, *args, **kwargs): return linear_output, bias, layernorm_output - def state_dict(self, destination=None, prefix='', keep_vars=False): + def state_dict(self, destination=None, prefix="", keep_vars=False): """Retrieve the state dictionary of the wrapped module and adapter. This method overrides the default state_dict behavior to include both @@ -146,6 +160,6 @@ def state_dict(self, destination=None, prefix='', keep_vars=False): self.to_wrap.state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars) # Store adapter state dict under the "adapter" prefix in the destination dict self.adapter.state_dict( - destination=destination, prefix=f'{prefix}adapter.', keep_vars=keep_vars + destination=destination, prefix=f"{prefix}adapter.", keep_vars=keep_vars ) return destination diff --git a/flagscale/train/megatron/training/peft/utils.py b/flagscale/train/megatron/training/peft/utils.py index 30e163d57d..e8510373c8 100644 --- a/flagscale/train/megatron/training/peft/utils.py +++ b/flagscale/train/megatron/training/peft/utils.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import math import re @@ -106,7 +120,7 @@ def is_expert_linear(fqn): Return whether the current base module is an expert linear module. See ParallelLinearAdapter.is_expert for usage details. """ - return re.match(r'(?!.*shared_).*mlp\..*experts.*\.linear_fc[1-2]$', fqn) is not None + return re.match(r"(?!.*shared_).*mlp\..*experts.*\.linear_fc[1-2]$", fqn) is not None def init_method_normal(sigma): @@ -137,7 +151,6 @@ def init_(tensor): class ParallelLinearAdapter(nn.Module): - def __init__( self, in_features: int, @@ -151,9 +164,8 @@ def __init__( out_init_method: str = "zero", dropout: float = 0.0, alpha: float | None = None, - dropout_position: str = 'post', + dropout_position: str = "post", ): - super().__init__() self.in_features = in_features self.out_features = out_features @@ -220,11 +232,11 @@ def __init__( self.dropout = None def _get_init_fn(self, init_method: str): - if init_method == 'xavier': + if init_method == "xavier": init_fn = nn.init.xavier_normal_ - elif init_method == 'normal': + elif init_method == "normal": init_fn = init_method_normal(0.2) - elif init_method == 'kaiming': + elif init_method == "kaiming": init_fn = init_method_kaiming_uniform(math.sqrt(5)) elif init_method == "zero": init_fn = init_method_const(0.0) @@ -234,13 +246,13 @@ def _get_init_fn(self, init_method: str): def forward(self, x): """ """ - if self.dropout is not None and self.dropout_position == 'pre': + if self.dropout is not None and self.dropout_position == "pre": x = self.dropout(x) x, _ = self.linear_in(x) x, _ = self.linear_out(x) - if self.dropout is not None and self.dropout_position == 'post': + if self.dropout is not None and self.dropout_position == "post": x = self.dropout(x) x = x * (self.alpha / self.dim) @@ -248,7 +260,7 @@ def forward(self, x): return x def sharded_state_dict( - self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[dict] = None + self, prefix: str = "", sharded_offsets: tuple = (), metadata: Optional[dict] = None ): """ Sharded state dict for LoRA adapter. Special treatment is given to the linear_fc1 adapter diff --git a/flagscale/train/megatron/training/spiky_loss.py b/flagscale/train/megatron/training/spiky_loss.py index 809edf259c..9716f4a198 100644 --- a/flagscale/train/megatron/training/spiky_loss.py +++ b/flagscale/train/megatron/training/spiky_loss.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import math import torch diff --git a/flagscale/train/megatron/training/stablelm2_scheduler.py b/flagscale/train/megatron/training/stablelm2_scheduler.py index c8ed2ec6bf..3c7a016704 100644 --- a/flagscale/train/megatron/training/stablelm2_scheduler.py +++ b/flagscale/train/megatron/training/stablelm2_scheduler.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from dataclasses import dataclass diff --git a/flagscale/train/megatron/training/theoretical_memory_usage.py b/flagscale/train/megatron/training/theoretical_memory_usage.py index bdc5b9f577..b0f9673ddc 100644 --- a/flagscale/train/megatron/training/theoretical_memory_usage.py +++ b/flagscale/train/megatron/training/theoretical_memory_usage.py @@ -2,7 +2,6 @@ """Computes theoretical memory footprint for model training.""" - import math from .utils import is_hybrid_model, print_rank_0 @@ -19,7 +18,7 @@ def compute_weight_and_optimizer_memory(args, verbose=False): # MoE. num_experts = 1 if args.num_experts is None else args.num_experts gated_linear_multiplier = 3 / 2 if args.swiglu else 1 - + shared_expert_ffn_hidden_size = ( 0 if args.moe_shared_expert_intermediate_size is None @@ -62,19 +61,29 @@ def compute_weight_and_optimizer_memory(args, verbose=False): if args.multi_latent_attention: assert not args.group_query_attention if args.q_lora_rank is None: - q_term = args.hidden_size * args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim) + q_term = ( + args.hidden_size + * args.num_attention_heads + * (args.qk_head_dim + args.qk_pos_emb_head_dim) + ) else: ## q lora + rope + q norm - q_term = args.q_lora_rank * (args.hidden_size + args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim) + norm_size) - + q_term = args.q_lora_rank * ( + args.hidden_size + + args.num_attention_heads * (args.qk_head_dim + args.qk_pos_emb_head_dim) + + norm_size + ) + self_attn_term = ( q_term - ## kv lora + rope + kv norm + args.kv_lora_rank - * (args.hidden_size + args.num_attention_heads * (args.qk_head_dim + args.v_head_dim) + norm_size) + * ( + args.hidden_size + + args.num_attention_heads * (args.qk_head_dim + args.v_head_dim) + + norm_size + ) + args.hidden_size * args.qk_pos_emb_head_dim - ## o proj + (args.num_attention_heads * args.v_head_dim) * args.hidden_size ) @@ -83,12 +92,11 @@ def compute_weight_and_optimizer_memory(args, verbose=False): 2 * args.hidden_size * args.hidden_size - * ( - # Attention. - ( - (1 + (args.num_query_groups / args.num_attention_heads)) - * query_projection_to_hidden_size_ratio - ) + * + # Attention. + ( + (1 + (args.num_query_groups / args.num_attention_heads)) + * query_projection_to_hidden_size_ratio ) ) @@ -108,7 +116,7 @@ def compute_weight_and_optimizer_memory(args, verbose=False): * args.hidden_size * ( # MoE MLP. - + (moe_ffn_hidden_size * num_experts * gated_linear_multiplier) + +(moe_ffn_hidden_size * num_experts * gated_linear_multiplier) # Shared MoE MLP. + (shared_expert_ffn_hidden_size * gated_linear_multiplier) # Transformer layernorms. @@ -121,7 +129,7 @@ def compute_weight_and_optimizer_memory(args, verbose=False): * args.hidden_size * ( # MoE MLP. - + (moe_ffn_hidden_size * args.moe_router_topk * gated_linear_multiplier) + +(moe_ffn_hidden_size * args.moe_router_topk * gated_linear_multiplier) # Shared MoE MLP. + (shared_expert_ffn_hidden_size * gated_linear_multiplier) # Transformer layernorms. @@ -295,7 +303,12 @@ def compute_activation_memory_without_sp(args, num_microbatches, verbose=False): """Compute activation memory without sequence parallelism""" # 4. Compute per-layer memory - per_layer_memory = args.seq_length * args.micro_batch_size * args.hidden_size * (10 + (24 / args.tensor_model_parallel_size)) + per_layer_memory = ( + args.seq_length + * args.micro_batch_size + * args.hidden_size + * (10 + (24 / args.tensor_model_parallel_size)) + ) if verbose: print( @@ -367,7 +380,9 @@ def compute_activation_memory_without_sp(args, num_microbatches, verbose=False): def report_theoretical_memory(args, num_microbatches=None, verbose=False): if is_hybrid_model(args): - print("Theoretical memory footprints not yet supported for hybrid Mamba-Transformer models.") + print( + "Theoretical memory footprints not yet supported for hybrid Mamba-Transformer models." + ) return weight_and_optimizer_memory = ( @@ -375,7 +390,7 @@ def report_theoretical_memory(args, num_microbatches=None, verbose=False): ) # Choose the appropriate activation memory calculation based on parallelism strategy - if args.sequence_parallel and args.recompute_granularity == 'selective': + if args.sequence_parallel and args.recompute_granularity == "selective": print_rank_0("compute_activation_memory with SP") activation_memory = ( compute_activation_memory(args, num_microbatches=num_microbatches, verbose=verbose) @@ -384,7 +399,9 @@ def report_theoretical_memory(args, num_microbatches=None, verbose=False): else: print_rank_0("compute_activation_memory_without_sp") activation_memory = ( - compute_activation_memory_without_sp(args, num_microbatches=num_microbatches, verbose=verbose) + compute_activation_memory_without_sp( + args, num_microbatches=num_microbatches, verbose=verbose + ) / NUM_BYTES_IN_MEGABYTE ) diff --git a/flagscale/train/megatron/training/tokenizer/__init__.py b/flagscale/train/megatron/training/tokenizer/__init__.py index e2e8ed852b..f2e09cb599 100644 --- a/flagscale/train/megatron/training/tokenizer/__init__.py +++ b/flagscale/train/megatron/training/tokenizer/__init__.py @@ -10,4 +10,4 @@ from .tokenizer import build_tokenizer, register_tokenizer_factory -__all__ = ['build_tokenizer', 'register_tokenizer_factory'] +__all__ = ["build_tokenizer", "register_tokenizer_factory"] diff --git a/flagscale/train/megatron/training/tokenizer/gpt2_tokenization.py b/flagscale/train/megatron/training/tokenizer/gpt2_tokenization.py index 8058554bef..a7ef603041 100644 --- a/flagscale/train/megatron/training/tokenizer/gpt2_tokenization.py +++ b/flagscale/train/megatron/training/tokenizer/gpt2_tokenization.py @@ -15,8 +15,7 @@ """Tokenization classes for OpenAI GPT.""" -from __future__ import (absolute_import, division, print_function, - unicode_literals) +from __future__ import absolute_import, division, print_function, unicode_literals import sys import json @@ -34,6 +33,7 @@ def lru_cache(): return lambda func: func + from collections import OrderedDict logger = logging.getLogger(__name__) @@ -133,18 +133,19 @@ def _cut_text(self, text, offsets): start = end return tokens + PRETRAINED_VOCAB_ARCHIVE_MAP = { - 'gpt2': "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json", + "gpt2": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json", } PRETRAINED_MERGES_ARCHIVE_MAP = { - 'gpt2': "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt", + "gpt2": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt", } PRETRAINED_VOCAB_POSITIONAL_EMBEDDINGS_SIZE_MAP = { - 'gpt2': 1024, + "gpt2": 1024, } -VOCAB_NAME = 'vocab.json' -MERGES_NAME = 'merges.txt' -SPECIAL_TOKENS_NAME = 'special_tokens.txt' +VOCAB_NAME = "vocab.json" +MERGES_NAME = "merges.txt" +SPECIAL_TOKENS_NAME = "special_tokens.txt" @lru_cache() @@ -159,8 +160,11 @@ def bytes_to_unicode(): And avoids mapping to whitespace/control characters the bpe code barfs on. """ _chr = unichr if sys.version_info[0] == 2 else chr - bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + \ - list(range(ord("®"), ord("ÿ") + 1)) + bs = ( + list(range(ord("!"), ord("~") + 1)) + + list(range(ord("¡"), ord("¬") + 1)) + + list(range(ord("®"), ord("ÿ") + 1)) + ) cs = bs[:] n = 0 for b in range(2**8): @@ -190,6 +194,7 @@ class GPT2Tokenizer(object): GPT-2 BPE tokenizer. Peculiarities: - Byte-level BPE """ + @classmethod def from_pretrained(cls, pretrained_model_name_or_path, cache_dir=None, *inputs, **kwargs): """ @@ -211,6 +216,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, cache_dir=None, *inputs, # redirect to the cache, if necessary try: from .file_utils import cached_path + resolved_vocab_file = cached_path(vocab_file, cache_dir=cache_dir) resolved_merges_file = cached_path(merges_file, cache_dir=cache_dir) except EnvironmentError: @@ -219,39 +225,48 @@ def from_pretrained(cls, pretrained_model_name_or_path, cache_dir=None, *inputs, "We assumed '{}' was a path or url but couldn't find files {} and {} " "at this path or url.".format( pretrained_model_name_or_path, - ', '.join(PRETRAINED_VOCAB_ARCHIVE_MAP.keys()), + ", ".join(PRETRAINED_VOCAB_ARCHIVE_MAP.keys()), pretrained_model_name_or_path, - vocab_file, merges_file)) + vocab_file, + merges_file, + ) + ) return None if resolved_vocab_file == vocab_file and resolved_merges_file == merges_file: logger.info("loading vocabulary file {}".format(vocab_file)) logger.info("loading merges file {}".format(merges_file)) else: - logger.info("loading vocabulary file {} from cache at {}".format( - vocab_file, resolved_vocab_file)) - logger.info("loading merges file {} from cache at {}".format( - merges_file, resolved_merges_file)) + logger.info( + "loading vocabulary file {} from cache at {}".format( + vocab_file, resolved_vocab_file + ) + ) + logger.info( + "loading merges file {} from cache at {}".format(merges_file, resolved_merges_file) + ) if pretrained_model_name_or_path in PRETRAINED_VOCAB_POSITIONAL_EMBEDDINGS_SIZE_MAP: # if we're using a pretrained model, ensure the tokenizer wont index sequences longer # than the number of positional embeddings max_len = PRETRAINED_VOCAB_POSITIONAL_EMBEDDINGS_SIZE_MAP[pretrained_model_name_or_path] - kwargs['max_len'] = min(kwargs.get('max_len', int(1e12)), max_len) + kwargs["max_len"] = min(kwargs.get("max_len", int(1e12)), max_len) # Instantiate tokenizer. - if special_tokens_file and 'special_tokens' not in kwargs: - with open(special_tokens_file, encoding='utf-8') as f: - special_tokens = f.read().split('\n')[:-1] + if special_tokens_file and "special_tokens" not in kwargs: + with open(special_tokens_file, encoding="utf-8") as f: + special_tokens = f.read().split("\n")[:-1] else: - special_tokens = kwargs.pop('special_tokens', []) + special_tokens = kwargs.pop("special_tokens", []) tokenizer = cls( resolved_vocab_file, resolved_merges_file, special_tokens=special_tokens, *inputs, - **kwargs) + **kwargs, + ) return tokenizer - def __init__(self, vocab_file, merges_file, errors='replace', - special_tokens=None, max_len=None): + def __init__( + self, vocab_file, merges_file, errors="replace", special_tokens=None, max_len=None + ): self.max_len = max_len if max_len is not None else int(1e12) with open(vocab_file) as f: self.encoder = json.load(f) @@ -259,8 +274,8 @@ def __init__(self, vocab_file, merges_file, errors='replace', self.errors = errors # how to handle errors in decoding self.byte_encoder = bytes_to_unicode() self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} - with open(merges_file, encoding='utf-8') as f: - bpe_data = f.read().split('\n')[1:-1] + with open(merges_file, encoding="utf-8") as f: + bpe_data = f.read().split("\n")[1:-1] bpe_merges = [tuple(merge.split()) for merge in bpe_data] self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges)))) self.cache = {} @@ -268,7 +283,8 @@ def __init__(self, vocab_file, merges_file, errors='replace', # Should haved added re.IGNORECASE so BPE merges can happen for # capitalized versions of contractions self.pat = re.compile( - r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""") + r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""" + ) self.special_tokens = {} self.special_tokens_decoder = {} @@ -278,16 +294,17 @@ def __len__(self): return len(self.encoder) + len(self.special_tokens) def set_special_tokens(self, special_tokens): - """ Add a list of additional tokens to the encoder. - The additional tokens are indexed starting from the last index of the - current vocabulary in the order of the `special_tokens` list. + """Add a list of additional tokens to the encoder. + The additional tokens are indexed starting from the last index of the + current vocabulary in the order of the `special_tokens` list. """ if not special_tokens: self.special_tokens = {} self.special_tokens_decoder = {} return - self.special_tokens = dict((tok, len(self.encoder) + i) - for i, tok in enumerate(special_tokens)) + self.special_tokens = dict( + (tok, len(self.encoder) + i) for i, tok in enumerate(special_tokens) + ) self.special_tokens_decoder = {v: k for k, v in self.special_tokens.items()} logger.info("Special tokens {}".format(self.special_tokens)) @@ -301,7 +318,7 @@ def bpe(self, token): return token while True: - bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float('inf'))) + bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf"))) if bigram not in self.bpe_ranks: break first, second = bigram @@ -328,23 +345,23 @@ def bpe(self, token): break else: pairs = get_pairs(word) - word = ' '.join(word) + word = " ".join(word) self.cache[token] = word return word def tokenize(self, text): - """ Tokenize a string. """ + """Tokenize a string.""" bpe_tokens = [] for token in re.findall(self.pat, text): if sys.version_info[0] == 2: - token = ''.join(self.byte_encoder[ord(b)] for b in token) + token = "".join(self.byte_encoder[ord(b)] for b in token) else: - token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8')) - bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(' ')) + token = "".join(self.byte_encoder[b] for b in token.encode("utf-8")) + bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" ")) return bpe_tokens def convert_tokens_to_ids(self, tokens): - """ Converts a sequence of tokens into ids using the vocab. """ + """Converts a sequence of tokens into ids using the vocab.""" ids = [] if isinstance(tokens, str) or (sys.version_info[0] == 2 and isinstance(tokens, unicode)): if tokens in self.special_tokens: @@ -361,7 +378,8 @@ def convert_tokens_to_ids(self, tokens): "Token indices sequence length is longer than the specified maximum " " sequence length for this OpenAI GPT model ({} > {}). Running this" " sequence through the model will result in indexing errors".format( - len(ids), self.max_len) + len(ids), self.max_len + ) ) return ids @@ -380,8 +398,8 @@ def encode(self, text): return self.convert_tokens_to_ids(self.tokenize(text)) def decode(self, tokens): - text = ''.join([self.decoder[token] for token in tokens]) - text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors=self.errors) + text = "".join([self.decoder[token] for token in tokens]) + text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors) return text def save_vocabulary(self, vocab_path): @@ -393,38 +411,46 @@ def save_vocabulary(self, vocab_path): merge_file = os.path.join(vocab_path, MERGES_NAME) special_tokens_file = os.path.join(vocab_path, SPECIAL_TOKENS_NAME) - with open(vocab_file, 'w', encoding='utf-8') as f: + with open(vocab_file, "w", encoding="utf-8") as f: f.write(json.dumps(self.encoder, ensure_ascii=False)) index = 0 with open(merge_file, "w", encoding="utf-8") as writer: - writer.write(u'#version: 0.2\n') + writer.write("#version: 0.2\n") for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]): if index != token_index: - logger.warning("Saving vocabulary to {}: BPE merge indices are not consecutive." - " Please check that the tokenizer is not corrupted!".format(merge_file)) + logger.warning( + "Saving vocabulary to {}: BPE merge indices are not consecutive." + " Please check that the tokenizer is not corrupted!".format(merge_file) + ) index = token_index - writer.write(' '.join(bpe_tokens) + u'\n') + writer.write(" ".join(bpe_tokens) + "\n") index += 1 index = len(self.encoder) - with open(special_tokens_file, 'w', encoding='utf-8') as writer: + with open(special_tokens_file, "w", encoding="utf-8") as writer: for token, token_index in sorted(self.special_tokens.items(), key=lambda kv: kv[1]): if index != token_index: - logger.warning("Saving special tokens vocabulary to {}: BPE indices are not consecutive." - " Please check that the tokenizer is not corrupted!".format(special_tokens_file)) + logger.warning( + "Saving special tokens vocabulary to {}: BPE indices are not consecutive." + " Please check that the tokenizer is not corrupted!".format( + special_tokens_file + ) + ) index = token_index - writer.write(token + u'\n') + writer.write(token + "\n") index += 1 return vocab_file, merge_file, special_tokens_file class AquilaTokenizer(GPT2Tokenizer): - def __init__(self, vocab_file, merges_file, errors='replace', - special_tokens=None, max_len=None): - super().__init__(vocab_file, merges_file, errors=errors, - special_tokens=special_tokens, max_len=max_len) + def __init__( + self, vocab_file, merges_file, errors="replace", special_tokens=None, max_len=None + ): + super().__init__( + vocab_file, merges_file, errors=errors, special_tokens=special_tokens, max_len=max_len + ) self.tokens_trie = Trie() if len(self.special_tokens) > 0: @@ -436,14 +462,14 @@ def __init__(self, vocab_file, merges_file, errors='replace', self.encoder[v] = k def _tokenize(self, text): - """ Tokenize a string. """ + """Tokenize a string.""" bpe_tokens = [] for token in re.findall(self.pat, text): if sys.version_info[0] == 2: - token = ''.join(self.byte_encoder[ord(b)] for b in token) + token = "".join(self.byte_encoder[ord(b)] for b in token) else: - token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8')) - bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(' ')) + token = "".join(self.byte_encoder[b] for b in token.encode("utf-8")) + bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" ")) return bpe_tokens def tokenize(self, text): @@ -466,6 +492,6 @@ def decode(self, tokens): text.append(self.special_tokens_decoder[token]) else: text.append(self.decoder[token]) - text = ''.join(text) - text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors=self.errors) + text = "".join(text) + text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors) return text diff --git a/flagscale/train/megatron/training/tokenizer/rwkv_tokenization.py b/flagscale/train/megatron/training/tokenizer/rwkv_tokenization.py index c814c55495..4dd22eacc1 100644 --- a/flagscale/train/megatron/training/tokenizer/rwkv_tokenization.py +++ b/flagscale/train/megatron/training/tokenizer/rwkv_tokenization.py @@ -1,5 +1,18 @@ -from __future__ import (absolute_import, division, print_function, - unicode_literals) +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import absolute_import, division, print_function, unicode_literals import sys import json @@ -10,9 +23,11 @@ try: from functools import lru_cache except ImportError: + def lru_cache(): return lambda func: func + logger = logging.getLogger(__name__) @@ -54,6 +69,7 @@ class RWKVTokenizer: RWKV Trie-based tokenizer. Compatible interface with Megatron tokenizer. """ + @classmethod def from_pretrained(cls, tokenizer_path, *inputs, **kwargs): tokenizer = cls(tokenizer_path, *inputs, **kwargs) @@ -64,7 +80,7 @@ def __init__(self, tokenizer_path, special_tokens=None, max_len=None): self.idx2token = {} self.token2idx = {} - files = [f for f in os.listdir(tokenizer_path) if f.endswith('.txt')] + files = [f for f in os.listdir(tokenizer_path) if f.endswith(".txt")] if not files: raise ValueError(f"No .txt vocab files found in {tokenizer_path}") for vocab_file in files: @@ -74,9 +90,11 @@ def __init__(self, tokenizer_path, special_tokens=None, max_len=None): sorted_tokens = [] for line in lines: - idx = int(line[:line.index(' ')]) - token_bytes = eval(line[line.index(' '):line.rindex(' ')]) - token_bytes = token_bytes.encode("utf-8") if isinstance(token_bytes, str) else token_bytes + idx = int(line[: line.index(" ")]) + token_bytes = eval(line[line.index(" ") : line.rindex(" ")]) + token_bytes = ( + token_bytes.encode("utf-8") if isinstance(token_bytes, str) else token_bytes + ) sorted_tokens.append(token_bytes) self.idx2token[idx] = token_bytes self.token2idx[token_bytes] = idx @@ -112,7 +130,7 @@ def encode_bytes(self, src: bytes): return tokens def decode_bytes(self, tokens): - return b''.join([self.idx2token[i] for i in tokens]) + return b"".join([self.idx2token[i] for i in tokens]) def encode(self, text): return self.encode_bytes(text.encode("utf-8")) @@ -121,7 +139,7 @@ def decode(self, tokens): try: return self.decode_bytes(tokens).decode("utf-8") except Exception: - return '\ufffd' + return "\ufffd" def convert_tokens_to_ids(self, tokens): ids = [] @@ -146,5 +164,4 @@ def __len__(self): @property def vocab_size(self): - return len(self.idx2token) + len(getattr(self, 'special_tokens', {})) - + return len(self.idx2token) + len(getattr(self, "special_tokens", {})) diff --git a/flagscale/train/megatron/training/tokenizer/tokenizer.py b/flagscale/train/megatron/training/tokenizer/tokenizer.py index ed6cd6cf61..dbd61899d4 100644 --- a/flagscale/train/megatron/training/tokenizer/tokenizer.py +++ b/flagscale/train/megatron/training/tokenizer/tokenizer.py @@ -45,7 +45,7 @@ def build_tokenizer(args, **kwargs): from megatron.training.utils import print_rank_0 if args.tokenizer_type in _TOKENIZER_FACTORY_REGISTRY: - print_rank_0(f'> building {args.tokenizer_type} tokenizer ...') + print_rank_0(f"> building {args.tokenizer_type} tokenizer ...") tokenizer = _TOKENIZER_FACTORY_REGISTRY[args.tokenizer_type](args, **kwargs) else: tokenizer = _upstream_build_tokenizer(args, **kwargs) @@ -87,12 +87,13 @@ def __init__(self, vocab_file, merge_file, special_tokens_file): super().__init__(path=vocab_file) special_tokens = [] if special_tokens_file: - special_tokens = open(special_tokens_file, encoding='utf-8').read().split('\n')[:-1] - self.tokenizer = AquilaTokenizer(vocab_file, merge_file, errors='replace', - special_tokens=special_tokens, max_len=None) - self.eod_id = self.tokenizer.encoder[''] - self.cls_id = self.tokenizer.encoder['[CLS]'] - self.pad_id = self.tokenizer.encoder['<|endoftext|>'] + special_tokens = open(special_tokens_file, encoding="utf-8").read().split("\n")[:-1] + self.tokenizer = AquilaTokenizer( + vocab_file, merge_file, errors="replace", special_tokens=special_tokens, max_len=None + ) + self.eod_id = self.tokenizer.encoder[""] + self.cls_id = self.tokenizer.encoder["[CLS]"] + self.pad_id = self.tokenizer.encoder["<|endoftext|>"] @property def vocab_size(self): @@ -131,6 +132,7 @@ class _HFTokenizerFS(_FlagScaleTokenizerBase): def __init__(self, tokenizer_path): super().__init__(path=tokenizer_path) from transformers import AutoTokenizer + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True) self.eod_id = self.tokenizer.eos_token_id self.cls_id = self.tokenizer.bos_token_id @@ -183,9 +185,9 @@ class _QwenTokenizerFS(_HFTokenizerFS): def __init__(self, tokenizer_path): super().__init__(tokenizer_path) - self.eod_id = self.tokenizer.encode('<|extra_204|>')[0] - self.cls_id = self.tokenizer.encode('<|extra_203|>')[0] - self.pad_id = self.tokenizer.encode('<|endoftext|>')[0] + self.eod_id = self.tokenizer.encode("<|extra_204|>")[0] + self.cls_id = self.tokenizer.encode("<|extra_203|>")[0] + self.pad_id = self.tokenizer.encode("<|endoftext|>")[0] class _HFTokenizersTokenizerFS(_FlagScaleTokenizerBase): @@ -194,6 +196,7 @@ class _HFTokenizersTokenizerFS(_FlagScaleTokenizerBase): def __init__(self, json_file): super().__init__(path=json_file) from tokenizers import Tokenizer + self.tokenizer = Tokenizer.from_file(json_file) print(f"Vocab size: {self.tokenizer.get_vocab_size()}") self.eod_id = self.tokenizer.token_to_id("<|endoftext|>") @@ -234,9 +237,9 @@ class _Qwen2TokenizerFS(_HFTokenizerFS): def __init__(self, tokenizer_path, args): super().__init__(tokenizer_path) - self.eod_id = self.tokenizer.encode('<|extra_204|>')[0] - self.cls_id = self.tokenizer.encode('<|extra_203|>')[0] - self.pad_id = self.tokenizer.encode('<|endoftext|>')[0] + self.eod_id = self.tokenizer.encode("<|extra_204|>")[0] + self.cls_id = self.tokenizer.encode("<|extra_203|>")[0] + self.pad_id = self.tokenizer.encode("<|endoftext|>")[0] assert args.vocab_size is not None self._vocab_size = args.vocab_size @@ -249,41 +252,59 @@ class _Qwen2VLTokenizer(_FlagScaleTokenizerBase): def __init__(self, tokenizer_path, extra_vocab_size): super().__init__(tokenizer_path) from transformers import AutoTokenizer + self.tokenizer = AutoTokenizer.from_pretrained( tokenizer_path, padding_side="right", use_fast=True, split_special_tokens=False, trust_remote_code=True, - revision = "main", - token = None, + revision="main", + token=None, ) self.extra_vocab_size = extra_vocab_size - self.special_tokens_map = {k:v for k, v in zip(self.tokenizer.all_special_tokens, self.tokenizer.all_special_ids)} - self.image_token = '<|image_pad|>' - self.video_token = '<|video_pad|>' - self.vision_start_token = '<|vision_start|>' - self.vision_end_token = '<|vision_end|>' + self.special_tokens_map = { + k: v for k, v in zip(self.tokenizer.all_special_tokens, self.tokenizer.all_special_ids) + } + self.image_token = "<|image_pad|>" + self.video_token = "<|video_pad|>" + self.vision_start_token = "<|vision_start|>" + self.vision_end_token = "<|vision_end|>" from transformers import AutoProcessor + self.processor = AutoProcessor.from_pretrained( tokenizer_path, - revision = "main", - token = None, + revision="main", + token=None, ) # NOTE: In Qwen2-VL, template in chat_template.json is same within tokenizer_config.json and both can be used. # However, in Qwen 2.5-VL, the two templates are different and only the one in chat_template.json is OK. self.chat_template = self.processor.chat_template - def __call__(self, text, return_tensors=None, - padding=None, max_length=None, truncation=None, add_special_tokens=None): + def __call__( + self, + text, + return_tensors=None, + padding=None, + max_length=None, + truncation=None, + add_special_tokens=None, + ): + return self.tokenizer( + text, + return_tensors=return_tensors, + padding=padding, + max_length=max_length, + truncation=truncation, + add_special_tokens=add_special_tokens, + ) - return self.tokenizer(text, return_tensors=return_tensors, padding=padding, - max_length=max_length, truncation=truncation, add_special_tokens=add_special_tokens) + def apply_chat_template(self, conversations, tokenize: bool = True, **kwargs): + return self.tokenizer.apply_chat_template( + conversations, tokenize=tokenize, chat_template=self.chat_template, **kwargs + ) - def apply_chat_template(self, conversations, tokenize:bool=True, **kwargs): - return self.tokenizer.apply_chat_template(conversations, tokenize=tokenize, chat_template=self.chat_template, **kwargs) - @property def vocab_size(self): return self.tokenizer.vocab_size + self.extra_vocab_size @@ -317,26 +338,27 @@ def pad_token_id(self): @property def eos_token_id(self): return self.tokenizer.eos_token_id - + @property def image_token_id(self): return self.special_tokens_map[self.image_token] - + @property def video_token_id(self): return self.special_tokens_map[self.video_token] - + @property def vision_start_token_id(self): return self.special_tokens_map[self.vision_start_token] - + @property def vision_end_token_id(self): return self.special_tokens_map[self.vision_end_token] - + def encode(self, x): return self.tokenizer.encode(x) + class _RWKVTokenizerFS(_FlagScaleTokenizerBase): """RWKV Trie-based tokenizer.""" @@ -371,6 +393,7 @@ def apply_chat_template(self, *args, **kwargs): # Factory functions and registration # --------------------------------------------------------------------------- + def _build_aquila(args, **kwargs): assert args.vocab_file and args.merge_file and args.special_tokens_file return _AquilaTokenizerFS(args.vocab_file, args.merge_file, args.special_tokens_file) @@ -413,11 +436,11 @@ def _build_rwkv(args, **kwargs): return _RWKVTokenizerFS(args.tokenizer_path) -register_tokenizer_factory('AquilaTokenizerFS', _build_aquila) -register_tokenizer_factory('HFTokenizerFS', _build_hf) -register_tokenizer_factory('Llama3TokenizerFS', _build_llama3) -register_tokenizer_factory('QwenTokenizerFS', _build_qwen) -register_tokenizer_factory('HFTokenizersTokenizerFS', _build_hftokenizers) -register_tokenizer_factory('Qwen2TokenizerFS', _build_qwen2) -register_tokenizer_factory('Qwen2VLTokenizer', _build_qwen2vl) -register_tokenizer_factory('RWKVTokenizer', _build_rwkv) +register_tokenizer_factory("AquilaTokenizerFS", _build_aquila) +register_tokenizer_factory("HFTokenizerFS", _build_hf) +register_tokenizer_factory("Llama3TokenizerFS", _build_llama3) +register_tokenizer_factory("QwenTokenizerFS", _build_qwen) +register_tokenizer_factory("HFTokenizersTokenizerFS", _build_hftokenizers) +register_tokenizer_factory("Qwen2TokenizerFS", _build_qwen2) +register_tokenizer_factory("Qwen2VLTokenizer", _build_qwen2vl) +register_tokenizer_factory("RWKVTokenizer", _build_rwkv) diff --git a/flagscale/train/megatron/training/training.py b/flagscale/train/megatron/training/training.py index d48ddfb2b8..688d00296d 100644 --- a/flagscale/train/megatron/training/training.py +++ b/flagscale/train/megatron/training/training.py @@ -1,15 +1,17 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Pretrain utilities.""" + import time + # The earliest we can measure the start time. _TRAIN_START_TIME = time.time() # Startup timestamps for tracking program initialization phases _STARTUP_TIMESTAMPS = { - 'program_start': None, # Set by entry script before imports - 'main_entry': None, # Set by entry script at start of __main__ - 'pretrain_entry': None, # Set at top of pretrain() + "program_start": None, # Set by entry script before imports + "main_entry": None, # Set by entry script at start of __main__ + "pretrain_entry": None, # Set at top of pretrain() } @@ -26,9 +28,9 @@ def set_startup_timestamps(program_start=None, main_entry=None): global _TRAIN_START_TIME, _STARTUP_TIMESTAMPS if program_start is not None: _TRAIN_START_TIME = program_start - _STARTUP_TIMESTAMPS['program_start'] = program_start + _STARTUP_TIMESTAMPS["program_start"] = program_start if main_entry is not None: - _STARTUP_TIMESTAMPS['main_entry'] = main_entry + _STARTUP_TIMESTAMPS["main_entry"] = main_entry from collections import defaultdict @@ -58,12 +60,13 @@ def set_startup_timestamps(program_start=None, main_entry=None): logging.basicConfig(handlers=[CustomHandler()], level=logging.INFO) from .theoretical_memory_usage import report_theoretical_memory -_LEGACY_TRAIN_START_TIME = time.time() # NOTE(asolergi-nv): Legacy timestamp +_LEGACY_TRAIN_START_TIME = time.time() # NOTE(asolergi-nv): Legacy timestamp import torch try: from megatron.rl import rl_utils + has_rl_utils = True except ImportError: has_rl_utils = False @@ -73,48 +76,48 @@ def set_startup_timestamps(program_start=None, main_entry=None): # as RL_LOGGABLE_TIMER_NAMES instead of being defined here. RL_LOGGABLE_TIMER_NAMES = [ # Top-level RL phases - 'rl/rollout-collection', - 'rl/prepare-data-for-update', + "rl/rollout-collection", + "rl/prepare-data-for-update", # Rollout collection breakdown - 'rl/inference-setup', - 'rl/collect-rollouts', - 'rl/sync-rollouts', - 'rl/suspend-engine', + "rl/inference-setup", + "rl/collect-rollouts", + "rl/sync-rollouts", + "rl/suspend-engine", # Optimizer offload/restore - 'rl/offload-optimizer-before-inference', - 'rl/restore-optimizer-after-inference', - 'rl/offload-kv-cache-after-inference', - 'rl/restore-kv-cache-before-inference', + "rl/offload-optimizer-before-inference", + "rl/restore-optimizer-after-inference", + "rl/offload-kv-cache-after-inference", + "rl/restore-kv-cache-before-inference", # Fine-grained offload/restore breakdown - 'rl/restore/grad-buffers', - 'rl/restore/optimizer-state', - 'rl/restore/wait-for-transfers', - 'rl/offload/grad-buffers', - 'rl/offload/optimizer-state', + "rl/restore/grad-buffers", + "rl/restore/optimizer-state", + "rl/restore/wait-for-transfers", + "rl/offload/grad-buffers", + "rl/offload/optimizer-state", # Weight prefetching - 'rl/prefetch-weights-to-gpu', - 'rl/prefetch-weights-to-cpu', + "rl/prefetch-weights-to-gpu", + "rl/prefetch-weights-to-cpu", # Data preparation - 'rl/compute-group-stats', - 'rl/prepare-advantages', - 'rl/prepare-trajectories', - 'rl/get-ltor-masks', - 'rl/create-dataloader', - 'rl/sequence-packing', - 'rl/align-inference-logprobs', - 'rl/log-wandb-tb', - 'rl/pack-sequences', - 'rl/regather-trajectories', + "rl/compute-group-stats", + "rl/prepare-advantages", + "rl/prepare-trajectories", + "rl/get-ltor-masks", + "rl/create-dataloader", + "rl/sequence-packing", + "rl/align-inference-logprobs", + "rl/log-wandb-tb", + "rl/pack-sequences", + "rl/regather-trajectories", # Logprobs computation - 'rl/compute-logprobs', - 'rl/compute-old-logprobs', - 'rl/compute-ref-logprobs', - 'rl/get-logprobs', - 'rl/forward-pass', - 'rl/log-softmax', + "rl/compute-logprobs", + "rl/compute-old-logprobs", + "rl/compute-ref-logprobs", + "rl/get-logprobs", + "rl/forward-pass", + "rl/log-softmax", # Inference / cuda graphs - 'rl/build-cuda-graphs', - 'rl/wait-for-decode-only', + "rl/build-cuda-graphs", + "rl/wait-for-decode-only", ] try: @@ -164,9 +167,14 @@ def set_startup_timestamps(program_start=None, main_entry=None): from megatron.core.transformer.cuda_graphs import TECudaGraphHelper from megatron.core.transformer.enums import CudaGraphScope from megatron.core.transformer.module import Float16Module -from megatron.core.distributed import DistributedDataParallelConfig, TorchFullyShardedDataParallelConfig +from megatron.core.distributed import ( + DistributedDataParallelConfig, + TorchFullyShardedDataParallelConfig, +) from megatron.core.distributed import DistributedDataParallel as DDP -from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallel as megatron_FSDP +from megatron.core.distributed.fsdp.mcore_fsdp_adapter import ( + FullyShardedDataParallel as megatron_FSDP, +) from megatron.core.optimizer.optimizer import param_group_identifier_keys from megatron.core.optimizer.qk_clip import clip_qk @@ -180,7 +188,13 @@ def set_startup_timestamps(program_start=None, main_entry=None): from megatron.core.distributed import finalize_model_grads from megatron.core.enums import ModelType -from megatron.core.optimizer import get_megatron_optimizer, AdamOptimizerConfig, SGDOptimizerConfig, OptimizerConfig, ParamKey +from megatron.core.optimizer import ( + get_megatron_optimizer, + AdamOptimizerConfig, + SGDOptimizerConfig, + OptimizerConfig, + ParamKey, +) from megatron.core.optimizer.muon import get_megatron_muon_optimizer from megatron.core.rerun_state_machine import ( get_rerun_state_machine, @@ -191,7 +205,11 @@ def set_startup_timestamps(program_start=None, main_entry=None): from megatron.training.initialize import initialize_megatron from megatron.training.initialize import write_args_to_tensorboard from megatron.training.initialize import set_jit_fusion_options -from megatron.training.utils import get_batch_on_this_cp_rank, get_batch_on_this_tp_rank, is_hybrid_model +from megatron.training.utils import ( + get_batch_on_this_cp_rank, + get_batch_on_this_tp_rank, + is_hybrid_model, +) from megatron.training.datasets.data_samplers import build_pretraining_data_loader from megatron.core.datasets.data_schedule import HybridCPDataLoaderWrapper from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler @@ -202,7 +220,7 @@ def set_startup_timestamps(program_start=None, main_entry=None): from megatron.core.parallel_state import ( destroy_global_memory_buffer, destroy_model_parallel, - update_pg_timeout + update_pg_timeout, ) from megatron.core.inference.symmetric_memory import SymmetricMemoryManager from megatron.core.inference.unified_memory import create_unified_mempool @@ -210,6 +228,7 @@ def set_startup_timestamps(program_start=None, main_entry=None): try: from torch_memory_saver import torch_memory_saver + torch_memory_saver.hook_mode = "torch" HAVE_TORCH_MEMORY_SAVER = True except ImportError: @@ -221,7 +240,7 @@ def set_startup_timestamps(program_start=None, main_entry=None): get_current_global_batch_size, get_current_running_global_batch_size, get_num_microbatches, - update_num_microbatches + update_num_microbatches, ) from megatron.training.async_utils import maybe_finalize_async_save @@ -257,7 +276,7 @@ def set_startup_timestamps(program_start=None, main_entry=None): ########## FlagScale Begin ########## from megatron.training.global_vars import get_spiky_loss_detector -from megatron.training.peft import PEFT # Import PEFT from peft module +from megatron.training.peft import PEFT # Import PEFT from peft module from megatron.plugin.hetero.parallel_context import get_parallel_context from flagscale.runner.straggler import ( OptionalSectionContext, @@ -272,6 +291,7 @@ def set_startup_timestamps(program_start=None, main_entry=None): ) from megatron.plugin.platform import get_platform + cur_platform = get_platform() import megatron.plugin_flagscale @@ -358,13 +378,14 @@ def destroy_global_state(): def print_datetime(string, override_timestamp=None): """Note that this call will sync across all ranks. Use override_timestamp if provided; - otherwise use current timestamp.""" + otherwise use current timestamp.""" torch.distributed.barrier() if override_timestamp is None: - time_str = datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f') + time_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f") else: - time_str = datetime.fromtimestamp(override_timestamp).strftime('%Y-%m-%d %H:%M:%S.%f') - print_rank_0(f'[{string}] datetime: {time_str} ') + time_str = datetime.fromtimestamp(override_timestamp).strftime("%Y-%m-%d %H:%M:%S.%f") + print_rank_0(f"[{string}] datetime: {time_str} ") + def num_floating_point_operations(args, batch_size): def mlp_layer_flops(batch_size, seq_len, hidden_size, expansion=4.0, swiglu=False): @@ -372,21 +393,44 @@ def mlp_layer_flops(batch_size, seq_len, hidden_size, expansion=4.0, swiglu=Fals scale_factor = 3.0 / 2.0 if swiglu else 1.0 return 4 * expansion * scale_factor * batch_size * seq_len * hidden_size**2 - def moe_layer_flops(batch_size, seq_len, hidden_size, moe_ffn_hidden_size, - shared_expert_ffn_hidden_size, num_experts_routed_to, - moe_latent_size=None, swiglu=False): + def moe_layer_flops( + batch_size, + seq_len, + hidden_size, + moe_ffn_hidden_size, + shared_expert_ffn_hidden_size, + num_experts_routed_to, + moe_latent_size=None, + swiglu=False, + ): """Calculate FLOPs for an MoE layer.""" scale_factor = 3.0 / 2.0 if swiglu else 1.0 if moe_latent_size is None: - routed_flops = (4 * batch_size * seq_len * hidden_size * - moe_ffn_hidden_size * num_experts_routed_to * scale_factor) + routed_flops = ( + 4 + * batch_size + * seq_len + * hidden_size + * moe_ffn_hidden_size + * num_experts_routed_to + * scale_factor + ) else: # Routed experts run on moe_latent_size. - routed_flops = (4 * batch_size * seq_len * moe_latent_size * - moe_ffn_hidden_size * num_experts_routed_to * scale_factor) + routed_flops = ( + 4 + * batch_size + * seq_len + * moe_latent_size + * moe_ffn_hidden_size + * num_experts_routed_to + * scale_factor + ) # Up proj and down proj. - routed_flops += (4 * batch_size * seq_len * hidden_size * moe_latent_size) - shared_flops = 4 * batch_size * seq_len * hidden_size * shared_expert_ffn_hidden_size * scale_factor + routed_flops += 4 * batch_size * seq_len * hidden_size * moe_latent_size + shared_flops = ( + 4 * batch_size * seq_len * hidden_size * shared_expert_ffn_hidden_size * scale_factor + ) return routed_flops + shared_flops def attn_layer_flops( @@ -404,8 +448,9 @@ def attn_layer_flops( * (hidden_size + (hidden_size * (g / num_heads)) + (seq_len / 2)) ) - def mamba_layer_flops(batch_size, seq_len, hidden_size, state_dim=16, - head_dim=64, num_groups=1, num_heads=128): + def mamba_layer_flops( + batch_size, seq_len, hidden_size, state_dim=16, head_dim=64, num_groups=1, num_heads=128 + ): """Calculate FLOPs for a Mamba layer.""" # Note (rwaleffe): flops estimate for scan should be updated based on new SSD kernels, # but small percent of overall layer flops @@ -426,57 +471,109 @@ def mamba_layer_flops(batch_size, seq_len, hidden_size, state_dim=16, + (2 * batch_size * seq_len * d_in * hidden_size) # out_proj ) - def gdn_layer_flops(batch_size, seq_len, hidden_size, - qk_head_dim=128, v_head_dim=128, - num_qk_heads=16, num_v_heads=32, - conv_kernel_dim=4): + def gdn_layer_flops( + batch_size, + seq_len, + hidden_size, + qk_head_dim=128, + v_head_dim=128, + num_qk_heads=16, + num_v_heads=32, + conv_kernel_dim=4, + ): """Calculate FLOPs for a Gated Delta Net (GDN) layer.""" qk_dim = qk_head_dim * num_qk_heads v_dim = v_head_dim * num_v_heads return ( - 2 * batch_size * seq_len * ( + 2 + * batch_size + * seq_len + * ( # in_proj: hidden_size -> (2*qk_dim + 2*v_dim + 2*num_v_heads) hidden_size * (2 * qk_dim + 2 * v_dim + 2 * num_v_heads) # conv1d + conv_kernel_dim * (2 * qk_dim + v_dim) # gated delta rule: KK^T, VK^T, S(a(I-bKK^T)), and SQ - + num_v_heads * (v_head_dim ** 2) * 4 + + num_v_heads * (v_head_dim**2) * 4 # out_proj: v_dim -> hidden_size + hidden_size * v_dim ) ) - def hybrid_flops(batch_size, seq_len, hidden_size, - num_attn_layers, num_mamba_layers, num_mlp_layers, num_moe_layers, - num_gdn_layers=0, - mamba_state_dim=128, mamba_head_dim=64, - mamba_num_groups=8, mamba_num_heads=128, - num_attn_heads=32, gqa=True, - gqa_groups=8, kv_channels=None, - mlp_expansion=4.0, swiglu=False, - moe_latent_size=None, - moe_ffn_hidden_size=2048, shared_expert_ffn_hidden_size=2048, num_experts_routed_to=1, - gdn_qk_head_dim=128, gdn_v_head_dim=128, - gdn_num_qk_heads=16, gdn_num_v_heads=32, - gdn_conv_kernel_dim=4, - vocab_size=256000, mtp_num_layers=0): + def hybrid_flops( + batch_size, + seq_len, + hidden_size, + num_attn_layers, + num_mamba_layers, + num_mlp_layers, + num_moe_layers, + num_gdn_layers=0, + mamba_state_dim=128, + mamba_head_dim=64, + mamba_num_groups=8, + mamba_num_heads=128, + num_attn_heads=32, + gqa=True, + gqa_groups=8, + kv_channels=None, + mlp_expansion=4.0, + swiglu=False, + moe_latent_size=None, + moe_ffn_hidden_size=2048, + shared_expert_ffn_hidden_size=2048, + num_experts_routed_to=1, + gdn_qk_head_dim=128, + gdn_v_head_dim=128, + gdn_num_qk_heads=16, + gdn_num_v_heads=32, + gdn_conv_kernel_dim=4, + vocab_size=256000, + mtp_num_layers=0, + ): """Calculate total FLOPs for the hybrid model.""" flops_fwd = ( - num_attn_layers * attn_layer_flops(batch_size, seq_len, hidden_size, - num_attn_heads, gqa, gqa_groups, kv_channels) + - num_mlp_layers * mlp_layer_flops(batch_size, seq_len, hidden_size, - mlp_expansion, swiglu) + - num_mamba_layers * mamba_layer_flops(batch_size, seq_len, hidden_size, - mamba_state_dim, mamba_head_dim, - mamba_num_groups, mamba_num_heads) + - num_moe_layers * moe_layer_flops(batch_size, seq_len, hidden_size, moe_ffn_hidden_size, - shared_expert_ffn_hidden_size, num_experts_routed_to, - moe_latent_size, swiglu) + - num_gdn_layers * gdn_layer_flops(batch_size, seq_len, hidden_size, - gdn_qk_head_dim, gdn_v_head_dim, - gdn_num_qk_heads, gdn_num_v_heads, - gdn_conv_kernel_dim) + - (2 * batch_size * seq_len * hidden_size * vocab_size * (1 + mtp_num_layers)) # logits computation + num_attn_layers + * attn_layer_flops( + batch_size, seq_len, hidden_size, num_attn_heads, gqa, gqa_groups, kv_channels + ) + + num_mlp_layers + * mlp_layer_flops(batch_size, seq_len, hidden_size, mlp_expansion, swiglu) + + num_mamba_layers + * mamba_layer_flops( + batch_size, + seq_len, + hidden_size, + mamba_state_dim, + mamba_head_dim, + mamba_num_groups, + mamba_num_heads, + ) + + num_moe_layers + * moe_layer_flops( + batch_size, + seq_len, + hidden_size, + moe_ffn_hidden_size, + shared_expert_ffn_hidden_size, + num_experts_routed_to, + moe_latent_size, + swiglu, + ) + + num_gdn_layers + * gdn_layer_flops( + batch_size, + seq_len, + hidden_size, + gdn_qk_head_dim, + gdn_v_head_dim, + gdn_num_qk_heads, + gdn_num_v_heads, + gdn_conv_kernel_dim, + ) + + ( + 2 * batch_size * seq_len * hidden_size * vocab_size * (1 + mtp_num_layers) + ) # logits computation ) return flops_fwd * 3 @@ -545,7 +642,7 @@ def transformer_flops(): if args.multi_latent_attention: assert not args.group_query_attention - ''' + """ Basic arithmetic let B is batch size, s is seq_len, h is embedding dim, for one self_attnetion block (prenorm is not included) @@ -557,7 +654,7 @@ def transformer_flops(): references https://arxiv.org/abs/2305.10403 https://arxiv.org/abs/2205.05198 - ''' + """ if args.experimental_attention_variant == "dsv4_hybrid": ## DSv4 hybrid MLA projections (per layer, per token). ## In dsv4_hybrid mode, qk_head_dim + qk_pos_emb_head_dim == v_head_dim @@ -566,9 +663,7 @@ def transformer_flops(): ## Full core attention is replaced by sparse attention and is accounted ## for in the dsv4_hybrid branch below. q_term = args.q_lora_rank * ( - args.hidden_size - + args.num_attention_heads * args.v_head_dim - + 1 # q norm + args.hidden_size + args.num_attention_heads * args.v_head_dim + 1 # q norm ) kv_term = args.hidden_size * args.v_head_dim + args.v_head_dim # kv proj + kv norm ## Grouped low-rank output projection: @@ -645,8 +740,7 @@ def transformer_flops(): / 2 # causal mask (only half of the mask is non-zero) * 2 # QK^T and (QK^T)V ## out proj - + query_projection_size - * args.hidden_size + + query_projection_size * args.hidden_size ) ) @@ -655,8 +749,8 @@ def transformer_flops(): if isinstance(args.linear_attention_freq, int): linear_attention_pattern = [ # [1,1,...,1,0,1,1,...,1,0,...] - 0 if ((i + 1) % args.linear_attention_freq == 0) - else 1 for i in range(num_layers) + 0 if ((i + 1) % args.linear_attention_freq == 0) else 1 + for i in range(num_layers) ] elif isinstance(args.linear_attention_freq, list): linear_attention_pattern = args.linear_attention_freq @@ -693,24 +787,18 @@ def transformer_flops(): * fma_expansion_factor * ( ## in proj - args.hidden_size - * (2 * qk_dim + 2 * v_dim + 2 * num_v_heads) + args.hidden_size * (2 * qk_dim + 2 * v_dim + 2 * num_v_heads) ## conv1d - + args.linear_conv_kernel_dim - * (2 * qk_dim + v_dim) + + args.linear_conv_kernel_dim * (2 * qk_dim + v_dim) ## gated delta rule - + num_v_heads - * (v_head_dim ** 2) - * 4 # KK^T, VK^T, S(a(I-bKK^T)), and SQ + + num_v_heads * (v_head_dim**2) * 4 # KK^T, VK^T, S(a(I-bKK^T)), and SQ ## out proj - + args.hidden_size - * v_dim + + args.hidden_size * v_dim ) ) else: raise ValueError( - "Invalid experimental_attention_variant: " - f"{args.experimental_attention_variant}" + f"Invalid experimental_attention_variant: {args.experimental_attention_variant}" ) else: num_linear_attention_layers = 0 @@ -732,8 +820,7 @@ def transformer_flops(): * args.hidden_size * ( # dense layer (deepseek v2, v3 style) - (args.ffn_hidden_size * ffn_expansion_factor) - * num_dense_layers + (args.ffn_hidden_size * ffn_expansion_factor) * num_dense_layers # routed experts + ( (moe_ffn_hidden_size * num_experts_routed_to * ffn_expansion_factor) @@ -751,8 +838,7 @@ def transformer_flops(): ) * num_moe_layers # Shared Experts. - + (shared_expert_ffn_hidden_size * ffn_expansion_factor) - * num_moe_layers + + (shared_expert_ffn_hidden_size * ffn_expansion_factor) * num_moe_layers ) # Self Attention + self_attn_term @@ -782,6 +868,7 @@ def transformer_flops(): from operator import itemgetter from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols, get_hybrid_layer_counts + num_mamba_layers, num_gdn_layers, num_attn_layers, num_mlp_layers, num_moe_layers = ( itemgetter(Symbols.MAMBA, Symbols.GDN, Symbols.ATTENTION, Symbols.MLP, Symbols.MOE)( get_hybrid_layer_counts(args.hybrid_layer_pattern) @@ -812,10 +899,16 @@ def transformer_flops(): mlp_expansion=args.ffn_hidden_size / args.hidden_size, swiglu=args.swiglu, moe_latent_size=args.moe_latent_size, - moe_ffn_hidden_size=(args.moe_ffn_hidden_size if args.moe_ffn_hidden_size is not None - else args.ffn_hidden_size), - shared_expert_ffn_hidden_size=(0 if args.moe_shared_expert_intermediate_size is None - else args.moe_shared_expert_intermediate_size), + moe_ffn_hidden_size=( + args.moe_ffn_hidden_size + if args.moe_ffn_hidden_size is not None + else args.ffn_hidden_size + ), + shared_expert_ffn_hidden_size=( + 0 + if args.moe_shared_expert_intermediate_size is None + else args.moe_shared_expert_intermediate_size + ), num_experts_routed_to=args.moe_router_topk, gdn_qk_head_dim=args.linear_key_head_dim or 128, gdn_v_head_dim=args.linear_value_head_dim or 128, @@ -850,12 +943,12 @@ def get_start_time_from_progress_log(): latest_num_floating_point_operations_uncommitted = None def _get_field(string, type): - return type(string.split(': ')[1]) + return type(string.split(": ")[1]) - with open_file(progress_log_filename, 'r') as f: + with open_file(progress_log_filename, "r") as f: for line in f: line = line.strip() - line_tokens = line.split('\t') + line_tokens = line.split("\t") world_size_in_line = _get_field(line_tokens[2], int) if line_tokens[3] == "Saved checkpoint": latest_num_floating_point_operations = _get_field(line_tokens[7], float) @@ -869,7 +962,9 @@ def _get_field(string, type): # save_checkpoint was called directly (without save_checkpoint_and_time), # which writes "Saved async checkpoint" but not "Saving async checkpoint". if latest_num_floating_point_operations_uncommitted is not None: - latest_num_floating_point_operations = latest_num_floating_point_operations_uncommitted + latest_num_floating_point_operations = ( + latest_num_floating_point_operations_uncommitted + ) latest_num_floating_point_operations_uncommitted = None if world_size_in_line != args.world_size: # Re-start search if we see a different world size. @@ -880,12 +975,14 @@ def _get_field(string, type): if start_time is None: start_time = line_tokens[0] start_num_floating_point_operations = latest_num_floating_point_operations - assert ( - start_time is not None and start_num_floating_point_operations is not None - ), "Should have seen at least one 'Starting job' entry with same world_size" - print_rank_0(f"megatron.training.get_start_time_from_progress_log: " - f"{start_time=}, {start_num_floating_point_operations=}") - return datetime.strptime(start_time, '%Y-%m-%d %H:%M:%S'), start_num_floating_point_operations + assert start_time is not None and start_num_floating_point_operations is not None, ( + "Should have seen at least one 'Starting job' entry with same world_size" + ) + print_rank_0( + f"megatron.training.get_start_time_from_progress_log: " + f"{start_time=}, {start_num_floating_point_operations=}" + ) + return datetime.strptime(start_time, "%Y-%m-%d %H:%M:%S"), start_num_floating_point_operations def preprocess_common_state_dict(common_state_dict): @@ -893,14 +990,15 @@ def preprocess_common_state_dict(common_state_dict): # Convert args key of type namespace to dictionary preprocessed_common_state_dict = copy.deepcopy(common_state_dict) - preprocessed_common_state_dict['args'] = vars(preprocessed_common_state_dict['args']) + preprocessed_common_state_dict["args"] = vars(preprocessed_common_state_dict["args"]) # Remove rank and local rank from state dict if it exists, since they are expected to be different - preprocessed_common_state_dict['args'].pop('local_rank', None) - preprocessed_common_state_dict['args'].pop('rank', None) + preprocessed_common_state_dict["args"].pop("local_rank", None) + preprocessed_common_state_dict["args"].pop("rank", None) if ( - preprocessed_common_state_dict['args']['use_distributed_optimizer'] + preprocessed_common_state_dict["args"]["use_distributed_optimizer"] and "optimizer" in preprocessed_common_state_dict ): + def reorder_inner_param_groups(optimizer_state_dict): # When distributed optimizer loading, source param groups will be reordered, # so we reorder the param groups here to prevent warning. @@ -920,7 +1018,7 @@ def reorder_inner_param_groups(optimizer_state_dict): param_groups.sort(key=key_fn) inner_optimizer["param_groups"] = param_groups - optimizer_state_dict = preprocessed_common_state_dict['optimizer'] + optimizer_state_dict = preprocessed_common_state_dict["optimizer"] if "optimizer" in optimizer_state_dict: # Only 1 optimizer in chained optimizer. reorder_inner_param_groups(optimizer_state_dict) @@ -946,7 +1044,7 @@ def pretrain( non_loss_data_func=None, store=None, inprocess_call_wrapper: Optional[CallWrapper] = None, - extra_valid_dataset_provider=None, ###### FlagScale Add ########## + extra_valid_dataset_provider=None, ###### FlagScale Add ########## ): """Main training program. @@ -992,7 +1090,7 @@ def pretrain( """ # Capture timestamp right at top of pretrain, before initialize_megatron global _STARTUP_TIMESTAMPS - _STARTUP_TIMESTAMPS['pretrain_entry'] = time.time() + _STARTUP_TIMESTAMPS["pretrain_entry"] = time.time() if inprocess_call_wrapper is not None: iteration = inprocess_call_wrapper.iteration @@ -1020,24 +1118,23 @@ def pretrain( timers = get_timers() if args.fine_grained_activation_offloading: - from megatron.core.pipeline_parallel.utils import ( - set_ideal_affinity_for_current_gpu - ) + from megatron.core.pipeline_parallel.utils import set_ideal_affinity_for_current_gpu + set_ideal_affinity_for_current_gpu() ########## FlagScale Begin ########## args = get_args() if args.mg_fl_prefer: - os.environ['MG_FL_PREFER'] = args.mg_fl_prefer + os.environ["MG_FL_PREFER"] = args.mg_fl_prefer # enable flagos:triton / vendor:cuda / reference:torch backend for transformer engine fl if args.te_fl_prefer: - os.environ['TE_FL_PREFER'] = args.te_fl_prefer + os.environ["TE_FL_PREFER"] = args.te_fl_prefer if args.te_fl_per_op: - os.environ['TE_FL_PER_OP'] = args.te_fl_per_op + os.environ["TE_FL_PER_OP"] = args.te_fl_per_op if args.te_fl_allow_vendors: - os.environ['TE_FL_ALLOW_VENDORS'] = args.te_fl_allow_vendors + os.environ["TE_FL_ALLOW_VENDORS"] = args.te_fl_allow_vendors if args.te_fl_deny_vendors: - os.environ['TE_FL_DENY_VENDORS'] = args.te_fl_deny_vendors + os.environ["TE_FL_DENY_VENDORS"] = args.te_fl_deny_vendors # enable flag gems to replace torch ops for distributed training # TODO(lixianduo): fix flag gems re-register error @@ -1046,9 +1143,11 @@ def pretrain( import flag_gems except ImportError: raise RuntimeError("Failed to import 'flag_gems'. Please install flag_gems.") - + try: - flag_gems.enable(record=True, once=True, unused=args.flag_gems_unused, path=args.flag_gems_log_path) + flag_gems.enable( + record=True, once=True, unused=args.flag_gems_unused, path=args.flag_gems_log_path + ) except Exception as e: raise RuntimeError(f"Failed to enable 'flag_gems': {e}.") @@ -1071,14 +1170,16 @@ def pretrain( # Adjust the startup time so it reflects the global minimum. # This will be closer to what scheduler will see (outside of # image ... launches). - program_start = _STARTUP_TIMESTAMPS.get('program_start') - main_entry = _STARTUP_TIMESTAMPS.get('main_entry') - pretrain_entry = _STARTUP_TIMESTAMPS.get('pretrain_entry') + program_start = _STARTUP_TIMESTAMPS.get("program_start") + main_entry = _STARTUP_TIMESTAMPS.get("main_entry") + pretrain_entry = _STARTUP_TIMESTAMPS.get("pretrain_entry") # Initialize program_start_global with a fallback value in case set_startup_timestamps() wasn't called program_start_global = _TRAIN_START_TIME - if _STARTUP_TIMESTAMPS['program_start'] is not None: - program_start_global = torch.tensor([_STARTUP_TIMESTAMPS['program_start']], dtype=torch.double, device=cur_platform.device()) + if _STARTUP_TIMESTAMPS["program_start"] is not None: + program_start_global = torch.tensor( + [_STARTUP_TIMESTAMPS["program_start"]], dtype=torch.double, device=cur_platform.device() + ) torch.distributed.all_reduce(program_start_global, op=torch.distributed.ReduceOp.MIN) program_start_global = program_start_global.item() set_startup_timestamps(program_start=program_start_global) @@ -1086,9 +1187,13 @@ def pretrain( global _LEGACY_TRAIN_START_TIME ########## FlagScale Begin ########## if "cpu:gloo" == torch.distributed.get_backend(): - start_time_tensor = torch.tensor([_LEGACY_TRAIN_START_TIME], dtype=torch.double, device='cpu') + start_time_tensor = torch.tensor( + [_LEGACY_TRAIN_START_TIME], dtype=torch.double, device="cpu" + ) else: - start_time_tensor = torch.tensor([_LEGACY_TRAIN_START_TIME], dtype=torch.double, device=cur_platform.device_name()) + start_time_tensor = torch.tensor( + [_LEGACY_TRAIN_START_TIME], dtype=torch.double, device=cur_platform.device_name() + ) ########## FlagScale End ########## torch.distributed.all_reduce(start_time_tensor, op=torch.distributed.ReduceOp.MIN) _LEGACY_TRAIN_START_TIME = start_time_tensor.item() @@ -1097,32 +1202,47 @@ def pretrain( megatron_init_end = time.time() app_metrics = {} - app_metrics['app_start_time'] = round(program_start_global * 1000.0) - app_metrics['app_model_init_start_time'] = round(program_start_global * 1000.0) + app_metrics["app_start_time"] = round(program_start_global * 1000.0) + app_metrics["app_model_init_start_time"] = round(program_start_global * 1000.0) # Print basic megatron init time (using global min start) # NOTE(asolergi-nv): This is not entirely accurate, but we keep it for backwards compatibility. print_rank_0( - 'time to initialize megatron (seconds): {:.3f}'.format(megatron_init_end - _LEGACY_TRAIN_START_TIME) + "time to initialize megatron (seconds): {:.3f}".format( + megatron_init_end - _LEGACY_TRAIN_START_TIME + ) ) # Note, not entirely accurate as rank 0 might not be the first or last to hit these timestamps - print_datetime('after in-process setup and before initialize_megatron', timestamp_after_inprocess_setup) - print_datetime('after in-job setup and before initialize_megatron', timestamp_after_in_job_setup) + print_datetime( + "after in-process setup and before initialize_megatron", timestamp_after_inprocess_setup + ) + print_datetime( + "after in-job setup and before initialize_megatron", timestamp_after_in_job_setup + ) if program_start is not None and main_entry is not None and pretrain_entry is not None: # Inject startup deltas into timers startup_timers = { - 'startup-program-entry-spread': program_start - program_start_global, # Local program start timestamp vs the global earliest program start timestamp - 'startup-library-setup': main_entry - program_start, # Local library imports - 'startup-program-setup': pretrain_entry - main_entry, # Local __main__ entry to pretrain entry - 'startup-in-process-setup': timestamp_after_inprocess_setup - pretrain_entry, # Local in-process setup - 'startup-in-job-setup': timestamp_after_in_job_setup - timestamp_after_inprocess_setup, # Local in-job setup - 'startup-initialize-megatron': timestamp_after_initialize_megatron - timestamp_after_in_job_setup, # Local initialize megatron - 'startup-set-jit-fusion-options': timestamp_after_set_jit_fusion_options - timestamp_after_initialize_megatron, # Local set JIT fusion options - 'all-reduce-start-timestamps-tensor': megatron_init_end - timestamp_after_set_jit_fusion_options, # 2x All-reduce, first collective call - 'startup-megatron-init-local': megatron_init_end - pretrain_entry, # Local megatron init - 'startup-megatron-init-global': megatron_init_end - program_start_global, # Local megatron init vs the global earliest program start timestamp + "startup-program-entry-spread": program_start + - program_start_global, # Local program start timestamp vs the global earliest program start timestamp + "startup-library-setup": main_entry - program_start, # Local library imports + "startup-program-setup": pretrain_entry + - main_entry, # Local __main__ entry to pretrain entry + "startup-in-process-setup": timestamp_after_inprocess_setup + - pretrain_entry, # Local in-process setup + "startup-in-job-setup": timestamp_after_in_job_setup + - timestamp_after_inprocess_setup, # Local in-job setup + "startup-initialize-megatron": timestamp_after_initialize_megatron + - timestamp_after_in_job_setup, # Local initialize megatron + "startup-set-jit-fusion-options": timestamp_after_set_jit_fusion_options + - timestamp_after_initialize_megatron, # Local set JIT fusion options + "all-reduce-start-timestamps-tensor": megatron_init_end + - timestamp_after_set_jit_fusion_options, # 2x All-reduce, first collective call + "startup-megatron-init-local": megatron_init_end + - pretrain_entry, # Local megatron init + "startup-megatron-init-global": megatron_init_end + - program_start_global, # Local megatron init vs the global earliest program start timestamp } for name, delta in startup_timers.items(): timers(name, log_level=0).set_elapsed(delta) @@ -1130,22 +1250,22 @@ def pretrain( # Print rank 0's absolute timestamps startup_timestamps = { - 'before library-setup': program_start, - 'after library-setup': main_entry, - 'before megatron-init': pretrain_entry, + "before library-setup": program_start, + "after library-setup": main_entry, + "before megatron-init": pretrain_entry, } for name, ts in startup_timestamps.items(): - ts_str = datetime.fromtimestamp(ts).strftime('%Y-%m-%d %H:%M:%S.%f') - print_rank_0(f'[{name}] datetime: {ts_str}') + ts_str = datetime.fromtimestamp(ts).strftime("%Y-%m-%d %H:%M:%S.%f") + print_rank_0(f"[{name}] datetime: {ts_str}") - print_datetime('after megatron is initialized') - app_metrics['app_model_init_finish_time'] = one_logger_utils.get_timestamp_in_ms() + print_datetime("after megatron is initialized") + app_metrics["app_model_init_finish_time"] = one_logger_utils.get_timestamp_in_ms() # Track E2E metrics on pretrain start one_logger_utils.on_pretrain_start() # Context used for persisting some state between checkpoint saves. - if args.non_persistent_ckpt_type == 'local': + if args.non_persistent_ckpt_type == "local": try: from nvidia_resiliency_ext.checkpointing.local.ckpt_managers.local_manager import ( LocalCheckpointManager, @@ -1171,7 +1291,7 @@ def pretrain( repl_strategy = None checkpointing_context = { - 'local_checkpoint_manager': LocalCheckpointManager( + "local_checkpoint_manager": LocalCheckpointManager( args.non_persistent_local_ckpt_dir, repl_strategy=repl_strategy ) } @@ -1179,13 +1299,13 @@ def pretrain( checkpointing_context = {} # Model, optimizer, and learning rate. - timers('model-and-optimizer-setup', log_level=0).start(barrier=True) + timers("model-and-optimizer-setup", log_level=0).start(barrier=True) model, optimizer, opt_param_scheduler = setup_model_and_optimizer( model_provider, model_type, checkpointing_context=checkpointing_context ) - timers('model-and-optimizer-setup').stop() - print_datetime('after model, optimizer, and learning rate ' 'scheduler are built') + timers("model-and-optimizer-setup").stop() + print_datetime("after model, optimizer, and learning rate scheduler are built") config = get_model_config(model[0]) # Build a separate inference model for RL if requested. @@ -1218,7 +1338,9 @@ def pretrain( # Build an isolated inference config so training config remains unchanged inference_config = copy.deepcopy(config) if args.rl_inference_tensor_model_parallel_size is not None: - inference_config.tensor_model_parallel_size = args.rl_inference_tensor_model_parallel_size + inference_config.tensor_model_parallel_size = ( + args.rl_inference_tensor_model_parallel_size + ) if args.rl_inference_pipeline_model_parallel_size is not None: inference_config.pipeline_model_parallel_size = ( args.rl_inference_pipeline_model_parallel_size @@ -1275,19 +1397,23 @@ def pretrain( ) # Data stuff. - app_metrics['app_build_dataiters_start_time'] = one_logger_utils.get_timestamp_in_ms() - timers('train/valid/test-data-iterators-setup', log_level=0).start(barrier=True) + app_metrics["app_build_dataiters_start_time"] = one_logger_utils.get_timestamp_in_ms() + timers("train/valid/test-data-iterators-setup", log_level=0).start(barrier=True) if args.virtual_pipeline_model_parallel_size is not None: train_data_iterator = [] valid_data_iterator = [] test_data_iterator = [] for vp_stage in range(len(model)): - dataset_provider_parameters = inspect.signature(train_valid_test_dataset_provider).parameters - assert "vp_stage" in dataset_provider_parameters, \ + dataset_provider_parameters = inspect.signature( + train_valid_test_dataset_provider + ).parameters + assert "vp_stage" in dataset_provider_parameters, ( "vp_stage must be a kwarg in train_valid_test_dataset_provider when using virtual pipeline parallelism" - vp_stage_train_valid_test_dataset_provider = \ - functools.partial(train_valid_test_dataset_provider, vp_stage=vp_stage) - if getattr(train_valid_test_dataset_provider, 'is_distributed', False): + ) + vp_stage_train_valid_test_dataset_provider = functools.partial( + train_valid_test_dataset_provider, vp_stage=vp_stage + ) + if getattr(train_valid_test_dataset_provider, "is_distributed", False): vp_stage_train_valid_test_dataset_provider.is_distributed = True iterators = build_train_valid_test_data_iterators( vp_stage_train_valid_test_dataset_provider @@ -1300,9 +1426,7 @@ def pretrain( valid_data_iterator = [] test_data_iterator = [] for _ in range(2): - iterators = build_train_valid_test_data_iterators( - train_valid_test_dataset_provider - ) + iterators = build_train_valid_test_data_iterators(train_valid_test_dataset_provider) train_data_iterator.append(iterators[0]) valid_data_iterator.append(iterators[1]) test_data_iterator.append(iterators[2]) @@ -1310,9 +1434,9 @@ def pretrain( train_data_iterator, valid_data_iterator, test_data_iterator = ( build_train_valid_test_data_iterators(train_valid_test_dataset_provider) ) - timers('train/valid/test-data-iterators-setup').stop() - print_datetime('after dataloaders are built') - app_metrics['app_build_dataiters_finish_time'] = one_logger_utils.get_timestamp_in_ms() + timers("train/valid/test-data-iterators-setup").stop() + print_datetime("after dataloaders are built") + app_metrics["app_build_dataiters_finish_time"] = one_logger_utils.get_timestamp_in_ms() # Track if training is enabled. Can only be done once args.do_train is assigned after dataloader is built. one_logger_utils.track_config_flags( @@ -1325,8 +1449,8 @@ def pretrain( ) # Print setup timing. - print_rank_0('done with setup ...') - timers.log(['model-and-optimizer-setup', 'train/valid/test-data-iterators-setup'], barrier=True) + print_rank_0("done with setup ...") + timers.log(["model-and-optimizer-setup", "train/valid/test-data-iterators-setup"], barrier=True) one_logger = get_one_logger() one_logger and one_logger.log_metrics(app_metrics) @@ -1334,13 +1458,13 @@ def pretrain( wandb_writer = get_wandb_writer() if wandb_writer: # Add job name to the wandb config to make it easier to run more singleton dependency jobs. - wandb_writer.config.update({'slurm_job_name': os.getenv("SLURM_JOB_NAME", "N/A")}) + wandb_writer.config.update({"slurm_job_name": os.getenv("SLURM_JOB_NAME", "N/A")}) if not args.skip_train or args.perform_rl_step: if args.skip_train: - print_rank_0('RL inference-only mode (--skip-train --perform-rl-step) ...') + print_rank_0("RL inference-only mode (--skip-train --perform-rl-step) ...") else: - print_rank_0('training ...') + print_rank_0("training ...") iteration = 0 args.curr_iteration = iteration @@ -1360,10 +1484,15 @@ def pretrain( extra_valid_dataset_provider, ) - print_datetime('after training is done') + print_datetime("after training is done") - if not args.auto_tune: ########## FlagScale Add ########## - if not args.skip_train and args.save and iteration != 0 and iteration % args.save_interval != 0: + if not args.auto_tune: ########## FlagScale Add ########## + if ( + not args.skip_train + and args.save + and iteration != 0 + and iteration % args.save_interval != 0 + ): save_checkpoint( iteration, model, @@ -1376,16 +1505,16 @@ def pretrain( ) one_logger and one_logger.log_metrics( - {'app_train_loop_finish_time': one_logger_utils.get_timestamp_in_ms()} + {"app_train_loop_finish_time": one_logger_utils.get_timestamp_in_ms()} ) else: - print_rank_0('skipping training (--skip-train is on) ...') + print_rank_0("skipping training (--skip-train is on) ...") iteration = args.iteration if args.do_valid: - prefix = f'iteration {iteration} on validation set' + prefix = f"iteration {iteration} on validation set" if args.perform_rl_step: rl_eval_model = model rl_training_model = None @@ -1407,15 +1536,20 @@ def pretrain( ) else: evaluate_and_print_results( - prefix, forward_step_func, - valid_data_iterator, model, - iteration, process_non_loss_data_func, config, - verbose=True, write_to_tensorboard=not args.skip_train, - non_loss_data_func=non_loss_data_func + prefix, + forward_step_func, + valid_data_iterator, + model, + iteration, + process_non_loss_data_func, + config, + verbose=True, + write_to_tensorboard=not args.skip_train, + non_loss_data_func=non_loss_data_func, ) if args.do_test: - prefix = f'iteration {iteration} on test set' + prefix = f"iteration {iteration} on test set" evaluate_and_print_results( prefix, forward_step_func, @@ -1431,7 +1565,11 @@ def pretrain( ########## FlagScale Begin ########## if extra_valid_dataset_provider is not None: - from megatron.training.extra_valid import extra_evaluate_and_print_results, build_extra_valid_data_iterators + from megatron.training.extra_valid import ( + extra_evaluate_and_print_results, + build_extra_valid_data_iterators, + ) + # NOTE(zhaoyinglia): Must rebuild the dataloaders for extra validation here, # to guarantee extra validation start from extra_iter=0 every time, # but we don't need to rebuild the datasets. @@ -1444,15 +1582,14 @@ def pretrain( elif args.use_dualpipev: extra_valid_data_iterator = [] for _ in range(2): - extra_iterators = build_extra_valid_data_iterators( - extra_valid_dataset_provider) + extra_iterators = build_extra_valid_data_iterators(extra_valid_dataset_provider) extra_valid_data_iterator.append(extra_iterators) else: - extra_valid_data_iterator = ( - build_extra_valid_data_iterators(extra_valid_dataset_provider) + extra_valid_data_iterator = build_extra_valid_data_iterators( + extra_valid_dataset_provider ) if getattr(args, "do_extra_valid", False): - prefix = f'iteration {iteration} on extra validation set' + prefix = f"iteration {iteration} on extra validation set" for extra_valid_index, extra_valid_data_itr in enumerate(extra_valid_data_iterator): extra_evaluate_and_print_results( extra_valid_index, @@ -1465,7 +1602,7 @@ def pretrain( config, verbose=True, write_to_tensorboard=not args.skip_train, - non_loss_data_func=non_loss_data_func + non_loss_data_func=non_loss_data_func, ) ########## FlagScale End ########## @@ -1478,7 +1615,7 @@ def pretrain( ft_integration.on_checkpointing_end(is_async_finalization=True) one_logger and one_logger.log_metrics( - {'app_finish_time': one_logger_utils.get_timestamp_in_ms()} + {"app_finish_time": one_logger_utils.get_timestamp_in_ms()} ) if args.perform_rl_step: @@ -1489,7 +1626,6 @@ def pretrain( def update_train_iters(args): - # For iteration-based training, we don't need to do anything if args.train_iters: return @@ -1518,10 +1654,16 @@ def update_train_iters(args): iterations += (args.train_samples - consumed_samples) // args.global_batch_size args.train_iters = iterations - print_rank_0(f'setting training iterations to {args.train_iters}') + print_rank_0(f"setting training iterations to {args.train_iters}") -def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap_with_ddp=True, config=None, pg_collection=None): +def get_model( + model_provider_func, + model_type=ModelType.encoder_or_decoder, + wrap_with_ddp=True, + config=None, + pg_collection=None, +): """Build the model.""" args = get_args() args.model_type = model_type @@ -1530,10 +1672,11 @@ def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap if has_nvidia_modelopt: from megatron.post_training.checkpointing import has_modelopt_state + # [ModelOpt]: Check if the checkpoint is a ModelOpt checkpoint and # set a flag to use our model provider if so. if args.load is not None and has_modelopt_state(args.load): - print_rank_0(f'ModelOpt checkpoint detected') + print_rank_0(f"ModelOpt checkpoint detected") args.modelopt_enabled = True elif getattr(args, "export_kd_teacher_load", None): # For distillation ckpts without ModelOpt state @@ -1596,9 +1739,8 @@ def build_model(): model.model_type = model_type return model - if args.init_model_with_meta_device: - with torch.device('meta'): + with torch.device("meta"): model = build_model() else: model = build_model() @@ -1639,8 +1781,7 @@ def build_model(): ) if get_pg_rank(pg_collection.dp) == 0 and get_pg_rank(pg_collection.cp) == 0: print( - ' > number of parameters on (tensor, pipeline) ' - 'model parallel rank ({}, {}): {}'.format( + " > number of parameters on (tensor, pipeline) model parallel rank ({}, {}): {}".format( get_pg_rank(pg_collection.tp), get_pg_rank(pg_collection.pp), num_parameters, @@ -1665,7 +1806,10 @@ def build_model(): # Materialize tensors on meta device (GPU allocation) if not using FSDP2 and not using Megatron FSDP. if args.init_model_with_meta_device and not args.use_torch_fsdp2 and not args.use_megatron_fsdp: - model = [to_empty_if_meta_device(model_module, device=torch.device(cur_platform.device_name())) for model_module in model] + model = [ + to_empty_if_meta_device(model_module, device=torch.device(cur_platform.device_name())) + for model_module in model + ] # Before TE2.x: The model_module.bfloat16()/model_module.half() above will call the inplace # copy of TE's Float8Tensor, which will write an unwanted value (amax calculated @@ -1687,32 +1831,37 @@ def build_model(): if getattr(args, "use_torch_fsdp2", False): reshard_after_forward = getattr(args, "torch_fsdp2_reshard_after_forward", True) - ddp_config = TorchFullyShardedDataParallelConfig(reshard_after_forward=reshard_after_forward) + ddp_config = TorchFullyShardedDataParallelConfig( + reshard_after_forward=reshard_after_forward + ) else: kwargs = {} for f in dataclasses.fields(DistributedDataParallelConfig): if hasattr(args, f.name): kwargs[f.name] = getattr(args, f.name) - kwargs['grad_reduce_in_fp32'] = args.accumulate_allreduce_grads_in_fp32 - kwargs['check_for_nan_in_grad'] = args.check_for_nan_in_loss_and_grad - kwargs['check_for_large_grads'] = args.check_for_large_grads + kwargs["grad_reduce_in_fp32"] = args.accumulate_allreduce_grads_in_fp32 + kwargs["check_for_nan_in_grad"] = args.check_for_nan_in_loss_and_grad + kwargs["check_for_large_grads"] = args.check_for_large_grads if args.ddp_num_buckets is not None: - assert args.ddp_bucket_size is None, \ + assert args.ddp_bucket_size is None, ( "Cannot specify both --ddp-num-buckets and --ddp-bucket-size" - assert args.ddp_num_buckets > 0, \ - "--ddp-num-buckets must be greater than 0" - kwargs['bucket_size'] = num_parameters // args.ddp_num_buckets + ) + assert args.ddp_num_buckets > 0, "--ddp-num-buckets must be greater than 0" + kwargs["bucket_size"] = num_parameters // args.ddp_num_buckets else: - kwargs['bucket_size'] = args.ddp_bucket_size - kwargs['pad_buckets_for_high_nccl_busbw'] = args.ddp_pad_buckets_for_high_nccl_busbw - kwargs['reduce_scatter_with_fp32_accumulation'] = args.ddp_reduce_scatter_with_fp32_accumulation - kwargs['param_name_patterns_for_fp32_local_accumulation'] = \ - tuple(args.ddp_param_name_patterns_for_fp32_local_accumulation) - kwargs['average_in_collective'] = args.ddp_average_in_collective + kwargs["bucket_size"] = args.ddp_bucket_size + kwargs["pad_buckets_for_high_nccl_busbw"] = args.ddp_pad_buckets_for_high_nccl_busbw + kwargs["reduce_scatter_with_fp32_accumulation"] = ( + args.ddp_reduce_scatter_with_fp32_accumulation + ) + kwargs["param_name_patterns_for_fp32_local_accumulation"] = tuple( + args.ddp_param_name_patterns_for_fp32_local_accumulation + ) + kwargs["average_in_collective"] = args.ddp_average_in_collective # Megatron-FSDP arguments. - kwargs['megatron_fsdp_main_params_dtype'] = args.megatron_fsdp_main_params_dtype - kwargs['megatron_fsdp_main_grads_dtype'] = args.megatron_fsdp_main_grads_dtype - kwargs['megatron_fsdp_grad_comm_dtype'] = args.megatron_fsdp_grad_comm_dtype + kwargs["megatron_fsdp_main_params_dtype"] = args.megatron_fsdp_main_params_dtype + kwargs["megatron_fsdp_main_grads_dtype"] = args.megatron_fsdp_main_grads_dtype + kwargs["megatron_fsdp_grad_comm_dtype"] = args.megatron_fsdp_grad_comm_dtype # Initialize DDPConfig. ddp_config = DistributedDataParallelConfig(**kwargs) @@ -1744,7 +1893,8 @@ def build_model(): module=model_chunk, # Turn off bucketing for model_chunk 2 onwards, since communication # for these model chunks is overlapped with compute anyway. - disable_bucketing=(model_chunk_idx > 0) or args.overlap_param_gather_with_optimizer_step, + disable_bucketing=(model_chunk_idx > 0) + or args.overlap_param_gather_with_optimizer_step, ) for (model_chunk_idx, model_chunk) in enumerate(model) ] @@ -1793,21 +1943,23 @@ def get_optimizer_param_scheduler(optimizer): else: lr_warmup_steps = args.lr_warmup_samples else: - raise Exception('either train-iters or train-samples should be provided.') + raise Exception("either train-iters or train-samples should be provided.") ########## FlagScale Begin ########## stablelm2_scheduler_config = None - if args.lr_decay_style == 'stablelm2-scheduler': + if args.lr_decay_style == "stablelm2-scheduler": from megatron.training.stablelm2_scheduler import StableLM2SchedulerConfig + stablelm2_scheduler_config = StableLM2SchedulerConfig( - args.global_batch_size, - args.lr_decay_stablelm2_cosine_samples, - args.lr_decay_stablelm2_rsqrt_samples, - args.lr_decay_stablelm2_alpha, - args.lr_decay_stablelm2_beta, - cosine_max_lr=args.lr_decay_stablelm2_cosine_max_lr, - cosine_period_samples=args.lr_decay_stablelm2_cosine_period_samples, - decay_samples=args.lr_decay_stablelm2_decay_samples) + args.global_batch_size, + args.lr_decay_stablelm2_cosine_samples, + args.lr_decay_stablelm2_rsqrt_samples, + args.lr_decay_stablelm2_alpha, + args.lr_decay_stablelm2_beta, + cosine_max_lr=args.lr_decay_stablelm2_cosine_max_lr, + cosine_period_samples=args.lr_decay_stablelm2_cosine_period_samples, + decay_samples=args.lr_decay_stablelm2_decay_samples, + ) ########## FlagScale End ########## opt_param_scheduler = OptimizerParamScheduler( @@ -1836,7 +1988,7 @@ def get_megatron_optimizer_config(args: Any) -> OptimizerConfig: """Return a Megatron optimizer config object from Megatron's arguments.""" config = None - if args.optimizer == 'adam' or 'muon' in args.optimizer: + if args.optimizer == "adam" or "muon" in args.optimizer: # TODO(deyuf): Muon needs both adam + muon but get() only receive one config # So for now we keep using adam config that's back compat with old way kwargs = {} @@ -1844,7 +1996,7 @@ def get_megatron_optimizer_config(args: Any) -> OptimizerConfig: if hasattr(args, f.name): kwargs[f.name] = getattr(args, f.name) config = AdamOptimizerConfig(**kwargs) - elif args.optimizer == 'sgd': + elif args.optimizer == "sgd": kwargs = {} for f in dataclasses.fields(SGDOptimizerConfig): if hasattr(args, f.name): @@ -1878,7 +2030,9 @@ def setup_model_and_optimizer( model = get_model(model_provider_func, model_type, wrap_with_ddp=wrap_with_ddp) unwrapped_model = unwrap_model(model) - one_logger and one_logger.log_metrics({"app_build_optimzer_start_time": one_logger_utils.get_timestamp_in_ms()}) + one_logger and one_logger.log_metrics( + {"app_build_optimzer_start_time": one_logger_utils.get_timestamp_in_ms()} + ) if skip_optimizer: optimizer, opt_param_scheduler = None, None # In RL inference-only mode, train_iters must still be set despite having no optimizer. @@ -1908,7 +2062,7 @@ def setup_model_and_optimizer( if mup_overrides: config_overrides = {**(config_overrides or {}), **mup_overrides} - if 'muon' not in config.optimizer: + if "muon" not in config.optimizer: # If the user is asking for a non-zero embedding init std, skip weight decay for embeddings # to avoid embeddings from shrinking to zero as recommended in https://arxiv.org/abs/2312.16903 # default_skip_embedding_weight_decay=args.embedding_init_method_std is not None, @@ -1925,11 +2079,13 @@ def setup_model_and_optimizer( model, config_overrides=config_overrides, use_gloo_process_groups=args.use_gloo_process_groups, - layer_wise_distributed_optimizer='dist' in config.optimizer, + layer_wise_distributed_optimizer="dist" in config.optimizer, ) opt_param_scheduler = get_optimizer_param_scheduler(optimizer) - one_logger and one_logger.log_metrics({"app_build_optimzer_finish_time": one_logger_utils.get_timestamp_in_ms()}) + one_logger and one_logger.log_metrics( + {"app_build_optimzer_finish_time": one_logger_utils.get_timestamp_in_ms()} + ) if args.moe_use_upcycling: torch.distributed.barrier() @@ -1963,9 +2119,9 @@ def setup_model_and_optimizer( unwrapped_model, dense_model_for_upcycling, load_kwargs={ - 'model': dense_model_for_upcycling, - 'optimizer': None, - 'opt_param_scheduler': None, + "model": dense_model_for_upcycling, + "optimizer": None, + "opt_param_scheduler": None, }, ) args.iteration = 1 @@ -1976,15 +2132,15 @@ def setup_model_and_optimizer( del dense_model_for_upcycling if (args.fp16 or args.bf16) and optimizer is not None: optimizer.reload_model_params() - print_rank_0(f'Upcycled checkpoint saved to {args.save}') + print_rank_0(f"Upcycled checkpoint saved to {args.save}") if ( args.load is not None or args.pretrained_checkpoint is not None ) and not args.moe_use_upcycling: one_logger and one_logger.log_metrics( - {'load_checkpoint_start_time': one_logger_utils.get_timestamp_in_ms()} + {"load_checkpoint_start_time": one_logger_utils.get_timestamp_in_ms()} ) - timers('load-checkpoint', log_level=0).start(barrier=True) + timers("load-checkpoint", log_level=0).start(barrier=True) args.iteration, args.num_floating_point_operations_so_far = load_checkpoint( model, @@ -1995,12 +2151,12 @@ def setup_model_and_optimizer( and getattr(args, "use_torch_fsdp2", False) and args.ckpt_format == "torch_dist", ) - timers('load-checkpoint').stop(barrier=True) - timers.log(['load-checkpoint']) + timers("load-checkpoint").stop(barrier=True) + timers.log(["load-checkpoint"]) one_logger and one_logger.log_metrics( { - 'load_checkpoint_finish_time': one_logger_utils.get_timestamp_in_ms(), - 'load_checkpoint_time': timers('load-checkpoint').active_time(), + "load_checkpoint_finish_time": one_logger_utils.get_timestamp_in_ms(), + "load_checkpoint_time": timers("load-checkpoint").active_time(), } ) else: @@ -2011,7 +2167,7 @@ def setup_model_and_optimizer( if ( args.iteration == 0 and len(unwrapped_model) == 1 - and hasattr(unwrapped_model[0], 'init_state_dict_from_bert') + and hasattr(unwrapped_model[0], "init_state_dict_from_bert") ): print_rank_0("Initializing ICT from pretrained BERT model") unwrapped_model[0].init_state_dict_from_bert() @@ -2052,7 +2208,16 @@ def dummy_train_step(data_iterator): batch = get_batch_on_this_cp_rank(batch) -def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=None): +def train_step( + forward_step_func, + data_iterator, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=None, +): """Single training step.""" args = get_args() timers = get_timers() @@ -2069,10 +2234,12 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch profile_cuda = getattr(args, "straggler_enable_gpu_profile", True) rerun_state_machine = get_rerun_state_machine() - save_dgrads_in_this_iteration = (args.save_dgrads_interval is not None and - (iteration + 1) % args.save_dgrads_interval == 0) - save_wgrads_in_this_iteration = (args.save_wgrads_interval is not None and - (iteration + 1) % args.save_wgrads_interval == 0) + save_dgrads_in_this_iteration = ( + args.save_dgrads_interval is not None and (iteration + 1) % args.save_dgrads_interval == 0 + ) + save_wgrads_in_this_iteration = ( + args.save_wgrads_interval is not None and (iteration + 1) % args.save_wgrads_interval == 0 + ) while rerun_state_machine.should_run_forward_backward(data_iterator): # Set grad to zero. for model_chunk in model: @@ -2158,11 +2325,16 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch return {}, True, should_checkpoint, should_exit, exit_code, None, None, 0 ########## FlagScale Begin ########## - if args.auto_skip_spiky_loss and (args.consumed_train_samples > args.lr_warmup_samples and args.curr_iteration > args.lr_warmup_iters): + if args.auto_skip_spiky_loss and ( + args.consumed_train_samples > args.lr_warmup_samples + and args.curr_iteration > args.lr_warmup_iters + ): spiky_loss_detector = get_spiky_loss_detector() loss_ = spiky_loss_detector.reduce_losses(losses_reduced) is_spiky_loss = spiky_loss_detector.is_spkiy_loss(loss_) - is_spiky_loss_tensor = torch.tensor(is_spiky_loss, dtype=torch.int, device=cur_platform.device_name()) + is_spiky_loss_tensor = torch.tensor( + is_spiky_loss, dtype=torch.int, device=cur_platform.device_name() + ) torch.distributed.all_reduce(is_spiky_loss_tensor, op=torch.distributed.ReduceOp.MAX) is_spiky_loss = is_spiky_loss_tensor.item() if is_spiky_loss > 0: @@ -2186,15 +2358,15 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch enabled=should_profile_straggler, profile_cuda=profile_cuda, ): - timers('optimizer', log_level=1).start(barrier=args.barrier_with_L1_time) + timers("optimizer", log_level=1).start(barrier=args.barrier_with_L1_time) update_successful, grad_norm, num_zeros_in_grad = optimizer.step() # get max attention logit for logging and run clip_qk() # Part of MuonClip Optimizer step log_max_attention_logit = 0 if args.qk_clip or args.log_max_attention_logit: log_max_attention_logit = clip_qk(model, log_max_only=not args.qk_clip) - - timers('optimizer').stop() + + timers("optimizer").stop() # when freezing sub-models we may have a mixture of successful and unsucessful ranks, # so we must gather across mp ranks @@ -2236,8 +2408,7 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch # over the total number of tokens across the global batch. val = torch.vstack(val).sum(dim=0) torch.distributed.all_reduce( - val, - group=mpu.get_data_parallel_group(with_context_parallel=True) + val, group=mpu.get_data_parallel_group(with_context_parallel=True) ) loss_reduced[key] = val[0] / val[1] elif val[0].numel() == 1: @@ -2256,7 +2427,16 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch num_zeros_in_grad, log_max_attention_logit, ) - return {}, skipped_iter, should_checkpoint, should_exit, exit_code, grad_norm, num_zeros_in_grad, log_max_attention_logit + return ( + {}, + skipped_iter, + should_checkpoint, + should_exit, + exit_code, + grad_norm, + num_zeros_in_grad, + log_max_attention_logit, + ) def training_log( @@ -2286,9 +2466,9 @@ def training_log( should_reset = not is_first_iteration # Advanced, skipped, and Nan iterations. - advanced_iters_key = 'advanced iterations' - skipped_iters_key = 'skipped iterations' - nan_iters_key = 'nan iterations' + advanced_iters_key = "advanced iterations" + skipped_iters_key = "skipped iterations" + nan_iters_key = "nan iterations" # Advanced iterations. if not skipped_iter: total_loss_dict[advanced_iters_key] = total_loss_dict.get(advanced_iters_key, 0) + 1 @@ -2302,49 +2482,55 @@ def training_log( for key in loss_dict: if not skipped_iter: total_loss_dict[key] = ( - total_loss_dict.get(key, torch.tensor([0.0], dtype=torch.float, device=cur_platform.device_name())) + total_loss_dict.get( + key, torch.tensor([0.0], dtype=torch.float, device=cur_platform.device_name()) + ) + loss_dict[key] ) else: value = loss_dict[key].float().sum().item() - is_nan = value == float('inf') or value == -float('inf') or value != value + is_nan = value == float("inf") or value == -float("inf") or value != value got_nan = got_nan or is_nan total_loss_dict[nan_iters_key] = total_loss_dict.get(nan_iters_key, 0) + int(got_nan) # Logging. timers_to_log = [] if args.timing_log_level >= 1: - timers_to_log.extend([ - 'forward-backward', - 'layernorm-grads-all-reduce', - 'embedding-grads-all-reduce', - 'all-grads-sync', - 'params-all-gather', - 'optimizer-copy-to-main-grad', - 'optimizer-unscale-and-check-inf', - 'optimizer-clip-main-grad', - 'optimizer-count-zeros', - 'optimizer-inner-step', - 'optimizer-copy-main-to-model-params', - 'optimizer', - ]) + timers_to_log.extend( + [ + "forward-backward", + "layernorm-grads-all-reduce", + "embedding-grads-all-reduce", + "all-grads-sync", + "params-all-gather", + "optimizer-copy-to-main-grad", + "optimizer-unscale-and-check-inf", + "optimizer-clip-main-grad", + "optimizer-count-zeros", + "optimizer-inner-step", + "optimizer-copy-main-to-model-params", + "optimizer", + ] + ) if args.timing_log_level >= 2: - timers_to_log.extend([ - 'batch-generator', - 'forward-compute', - 'backward-compute', - 'forward-recv', - 'forward-send', - 'backward-recv', - 'backward-send', - 'forward-send-forward-recv', - 'forward-send-backward-recv', - 'backward-send-forward-recv', - 'backward-send-backward-recv', - 'forward-backward-send-forward-backward-recv', - ]) + timers_to_log.extend( + [ + "batch-generator", + "forward-compute", + "backward-compute", + "forward-recv", + "forward-send", + "backward-recv", + "backward-send", + "forward-send-forward-recv", + "forward-send-backward-recv", + "backward-send-forward-recv", + "backward-send-backward-recv", + "forward-backward-send-forward-backward-recv", + ] + ) # Add timers from RL loop if needed. - if getattr(args, 'perform_rl_step', False): + if getattr(args, "perform_rl_step", False): timers_to_log.extend(RL_LOGGABLE_TIMER_NAMES) # Calculate batch size. @@ -2360,27 +2546,40 @@ def training_log( # Tensorboard values. if is_last_rank() and (iteration % args.tensorboard_log_interval == 0): if wandb_writer: - wandb_writer.log({'samples vs steps': args.consumed_train_samples}, iteration) + wandb_writer.log({"samples vs steps": args.consumed_train_samples}, iteration) if learning_rate is not None: if writer: - writer.add_scalar('learning-rate', learning_rate, iteration) - writer.add_scalar('learning-rate vs samples', learning_rate, args.consumed_train_samples) + writer.add_scalar("learning-rate", learning_rate, iteration) + writer.add_scalar( + "learning-rate vs samples", learning_rate, args.consumed_train_samples + ) if wandb_writer: - wandb_writer.log({'learning-rate': learning_rate}, iteration) - wandb_writer.log({'consumed-tokens': args.consumed_train_samples * args.seq_length / 1000. / 1000 / 1000}, iteration) + wandb_writer.log({"learning-rate": learning_rate}, iteration) + wandb_writer.log( + { + "consumed-tokens": args.consumed_train_samples + * args.seq_length + / 1000.0 + / 1000 + / 1000 + }, + iteration, + ) if args.decoupled_lr is not None: if writer: - writer.add_scalar('decoupled-learning-rate', args.decoupled_learning_rate, iteration) + writer.add_scalar( + "decoupled-learning-rate", args.decoupled_learning_rate, iteration + ) if args.skipped_train_samples > 0: if writer: - writer.add_scalar('skipped-train-samples', args.skipped_train_samples, iteration) + writer.add_scalar("skipped-train-samples", args.skipped_train_samples, iteration) if wandb_writer: - wandb_writer.log({'skipped-train-samples': args.skipped_train_samples}, iteration) + wandb_writer.log({"skipped-train-samples": args.skipped_train_samples}, iteration) if writer: - writer.add_scalar('batch-size', batch_size, iteration) - writer.add_scalar('batch-size vs samples', batch_size, args.consumed_train_samples) + writer.add_scalar("batch-size", batch_size, iteration) + writer.add_scalar("batch-size vs samples", batch_size, args.consumed_train_samples) if wandb_writer: - wandb_writer.log({'batch-size': batch_size}, iteration) + wandb_writer.log({"batch-size": batch_size}, iteration) # Log bins for packed mode if has_rl_utils and args.rl_use_sequence_packing: packing_metrics = rl_utils.get_sequence_packing_tensorboard_metrics(args) @@ -2392,47 +2591,55 @@ def training_log( for key in loss_dict: if writer: writer.add_scalar(key, loss_dict[key], iteration) - writer.add_scalar(key + ' vs samples', loss_dict[key], args.consumed_train_samples) + writer.add_scalar(key + " vs samples", loss_dict[key], args.consumed_train_samples) if wandb_writer: wandb_writer.log({key: loss_dict[key]}, iteration) if args.log_loss_scale_to_tensorboard: if writer: - writer.add_scalar('loss-scale', loss_scale, iteration) - writer.add_scalar('loss-scale vs samples', loss_scale, args.consumed_train_samples) + writer.add_scalar("loss-scale", loss_scale, iteration) + writer.add_scalar("loss-scale vs samples", loss_scale, args.consumed_train_samples) if wandb_writer: - wandb_writer.log({'loss-scale': loss_scale}, iteration) + wandb_writer.log({"loss-scale": loss_scale}, iteration) if args.log_world_size_to_tensorboard: if writer: - writer.add_scalar('world-size', args.world_size, iteration) - writer.add_scalar('world-size vs samples', args.world_size, args.consumed_train_samples) + writer.add_scalar("world-size", args.world_size, iteration) + writer.add_scalar( + "world-size vs samples", args.world_size, args.consumed_train_samples + ) if wandb_writer: - wandb_writer.log({'world-size': args.world_size}, iteration) + wandb_writer.log({"world-size": args.world_size}, iteration) if grad_norm is not None: if writer: - writer.add_scalar('grad-norm', grad_norm, iteration) - writer.add_scalar('grad-norm vs samples', grad_norm, args.consumed_train_samples) + writer.add_scalar("grad-norm", grad_norm, iteration) + writer.add_scalar("grad-norm vs samples", grad_norm, args.consumed_train_samples) if wandb_writer: - wandb_writer.log({'grad-norm': grad_norm}, iteration) + wandb_writer.log({"grad-norm": grad_norm}, iteration) if num_zeros_in_grad is not None: if writer: - writer.add_scalar('num-zeros', num_zeros_in_grad, iteration) + writer.add_scalar("num-zeros", num_zeros_in_grad, iteration) writer.add_scalar( - 'num-zeros vs samples', num_zeros_in_grad, args.consumed_train_samples + "num-zeros vs samples", num_zeros_in_grad, args.consumed_train_samples ) if wandb_writer: - wandb_writer.log({'num-zeros': num_zeros_in_grad}, iteration) + wandb_writer.log({"num-zeros": num_zeros_in_grad}, iteration) if params_norm is not None: if writer: - writer.add_scalar('params-norm', params_norm, iteration) - writer.add_scalar('params-norm vs samples', params_norm, args.consumed_train_samples) + writer.add_scalar("params-norm", params_norm, iteration) + writer.add_scalar( + "params-norm vs samples", params_norm, args.consumed_train_samples + ) if wandb_writer: - wandb_writer.log({'params-norm': params_norm}, iteration) + wandb_writer.log({"params-norm": params_norm}, iteration) if args.perform_rl_step: - grpo_collection_iteration = iteration // (args.grpo_iterations * ( ( args.grpo_samples_per_iteration )// args.global_batch_size )) + grpo_collection_iteration = iteration // ( + args.grpo_iterations * ((args.grpo_samples_per_iteration) // args.global_batch_size) + ) if writer: - writer.add_scalar('grpo_collection_iteration', grpo_collection_iteration, iteration) + writer.add_scalar("grpo_collection_iteration", grpo_collection_iteration, iteration) if wandb_writer: - wandb_writer.log({'grpo_collection_iteration': grpo_collection_iteration}, iteration) + wandb_writer.log( + {"grpo_collection_iteration": grpo_collection_iteration}, iteration + ) if args.log_memory_to_tensorboard: mem_stats = cur_platform.memory_stats() if writer: @@ -2445,7 +2652,9 @@ def training_log( writer.add_scalar( "mem-max-allocated-bytes", mem_stats["allocated_bytes.all.peak"], iteration ) - writer.add_scalar("mem-allocated-count", mem_stats["allocation.all.current"], iteration) + writer.add_scalar( + "mem-allocated-count", mem_stats["allocation.all.current"], iteration + ) if wandb_writer: wandb_writer.log( {"mem-reserved-bytes": mem_stats["reserved_bytes.all.current"]}, iteration @@ -2458,9 +2667,9 @@ def training_log( ) if args.log_max_attention_logit: if writer: - writer.add_scalar('max_attention_logit', max_attention_logit, iteration) + writer.add_scalar("max_attention_logit", max_attention_logit, iteration) if wandb_writer: - wandb_writer.log({'max_attention_logit': max_attention_logit}, iteration) + wandb_writer.log({"max_attention_logit": max_attention_logit}, iteration) # Log MoE metrics. if args.num_experts is not None: @@ -2479,8 +2688,10 @@ def training_log( from operator import itemgetter from megatron.core.ssm.mamba_hybrid_layer_allocation import ( - Symbols, get_hybrid_layer_counts, + Symbols, + get_hybrid_layer_counts, ) + layers = itemgetter(Symbols.MOE)(get_hybrid_layer_counts(args.hybrid_layer_pattern)) else: layers = args.num_layers @@ -2520,14 +2731,16 @@ def training_log( # Dump memory snapshot and print metrics to stdout. if iteration % args.log_interval == 0 or is_first_iteration: - if args.record_memory_history and (is_last_rank() or torch.distributed.get_backend() == 'fake'): + if args.record_memory_history and ( + is_last_rank() or torch.distributed.get_backend() == "fake" + ): snapshot = torch.cuda.memory._snapshot() from pickle import dump - with open(args.memory_snapshot_path, 'wb') as f: + with open(args.memory_snapshot_path, "wb") as f: dump(snapshot, f) - elapsed_time = timers('interval-time').elapsed(barrier=True, reset=should_reset) + elapsed_time = timers("interval-time").elapsed(barrier=True, reset=should_reset) elapsed_time_per_iteration = elapsed_time / total_iterations throughput = num_floating_point_operations(args, batch_size) / ( @@ -2541,61 +2754,63 @@ def training_log( # WandB according to the regular schedule. if args.log_timers_to_tensorboard and not is_first_iteration: if writer: - writer.add_scalar('iteration-time', elapsed_time_per_iteration, iteration) + writer.add_scalar("iteration-time", elapsed_time_per_iteration, iteration) if wandb_writer: - wandb_writer.log({'iteration-time': elapsed_time_per_iteration}, iteration) + wandb_writer.log({"iteration-time": elapsed_time_per_iteration}, iteration) log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')}]" - log_string += ' iteration {:8d}/{:8d} |'.format(iteration, args.train_iters) - log_string += ' consumed samples: {:12d} |'.format(args.consumed_train_samples) + log_string += " iteration {:8d}/{:8d} |".format(iteration, args.train_iters) + log_string += " consumed samples: {:12d} |".format(args.consumed_train_samples) if has_rl_utils and args.rl_use_sequence_packing: log_string += rl_utils.get_sequence_packing_log_info(args) if args.skipped_train_samples > 0: - log_string += ' skipped samples: {:12d} |'.format(args.skipped_train_samples) - log_string += ' elapsed time per iteration (ms): {:.1f} |'.format( + log_string += " skipped samples: {:12d} |".format(args.skipped_train_samples) + log_string += " elapsed time per iteration (ms): {:.1f} |".format( elapsed_time_per_iteration * 1000.0 ) if args.log_throughput: - log_string += f' throughput per GPU (TFLOP/s/GPU): {throughput:.1f} |' + log_string += f" throughput per GPU (TFLOP/s/GPU): {throughput:.1f} |" if args.log_timers_to_tensorboard: if writer: - writer.add_scalar('throughput', throughput, iteration) + writer.add_scalar("throughput", throughput, iteration) if wandb_writer: - wandb_writer.log({'throughput': throughput}, iteration) + wandb_writer.log({"throughput": throughput}, iteration) if args.log_energy: energy = (energy_monitor.lap() / total_iterations) / args.world_size power = energy / elapsed_time_per_iteration - log_string += f' energy per GPU (J/iter/GPU): {energy:.1f} |' - log_string += f' power per GPU (W/GPU): {power:.1f} |' + log_string += f" energy per GPU (J/iter/GPU): {energy:.1f} |" + log_string += f" power per GPU (W/GPU): {power:.1f} |" if writer: - writer.add_scalar('iter-energy/gpu', energy, iteration) - writer.add_scalar('power/gpu', power, iteration) + writer.add_scalar("iter-energy/gpu", energy, iteration) + writer.add_scalar("power/gpu", power, iteration) if wandb_writer: - wandb_writer.log({'iter-energy/gpu': energy}, iteration) - wandb_writer.log({'power/gpu': power}, iteration) + wandb_writer.log({"iter-energy/gpu": energy}, iteration) + wandb_writer.log({"power/gpu": power}, iteration) # Decoupled_learning_rate should be not None only on first and last pipeline stage. if learning_rate is not None: - log_string += f' learning rate: {learning_rate:.6E} |' - log_string += f' global batch size: {batch_size:5d} |' + log_string += f" learning rate: {learning_rate:.6E} |" + log_string += f" global batch size: {batch_size:5d} |" for key in total_loss_dict: if key not in [advanced_iters_key, skipped_iters_key, nan_iters_key]: avg = total_loss_dict[key].item() / float( max(1, total_loss_dict[advanced_iters_key]) ) if avg >= 0.0: - log_string += ' {}: {:.6E} |'.format(key, avg) + log_string += " {}: {:.6E} |".format(key, avg) if should_reset: - total_loss_dict[key] = torch.tensor([0.0], dtype=torch.float, device=cur_platform.device_name()) - log_string += f' loss scale: {loss_scale:.1f} |' + total_loss_dict[key] = torch.tensor( + [0.0], dtype=torch.float, device=cur_platform.device_name() + ) + log_string += f" loss scale: {loss_scale:.1f} |" if grad_norm is not None: - log_string += f' grad norm: {grad_norm:.3f} |' + log_string += f" grad norm: {grad_norm:.3f} |" if num_zeros_in_grad is not None: - log_string += f' num zeros: {num_zeros_in_grad} |' + log_string += f" num zeros: {num_zeros_in_grad} |" if params_norm is not None: - log_string += f' params norm: {params_norm:.3f} |' - log_string += ' number of skipped iterations: {:3d} |'.format( + log_string += f" params norm: {params_norm:.3f} |" + log_string += " number of skipped iterations: {:3d} |".format( total_loss_dict[skipped_iters_key] ) - log_string += ' number of nan iterations: {:3d} |'.format(total_loss_dict[nan_iters_key]) + log_string += " number of nan iterations: {:3d} |".format(total_loss_dict[nan_iters_key]) if should_reset: total_loss_dict[advanced_iters_key] = 0 total_loss_dict[skipped_iters_key] = 0 @@ -2607,19 +2822,26 @@ def training_log( if torch.distributed.get_rank() == 0: num_microbatches = get_num_microbatches() report_theoretical_memory(args, num_microbatches=num_microbatches, verbose=True) - report_memory(f'(after {iteration} iterations)') + report_memory(f"(after {iteration} iterations)") reported_memory_in_this_iteration = True loaded_iteration = max(get_loaded_iteration() or 0, 0) if iteration > (loaded_iteration + 1): # Make sure the memory after the second iteration is reported to include optimizer state memory. report_memory_flag = False - if args.log_memory_interval is not None and iteration % args.log_memory_interval == 0 and \ - not reported_memory_in_this_iteration: - report_memory(f'(after {iteration} iterations)') + if ( + args.log_memory_interval is not None + and iteration % args.log_memory_interval == 0 + and not reported_memory_in_this_iteration + ): + report_memory(f"(after {iteration} iterations)") # Write timers to wandb, don't reset the counts. if args.log_timers_to_tensorboard: - timers.write(timers_to_log, writer, iteration, normalizer=args.log_interval, reset=False) - timers.write(timers_to_log, wandb_writer, iteration, normalizer=args.log_interval, reset=False) + timers.write( + timers_to_log, writer, iteration, normalizer=args.log_interval, reset=False + ) + timers.write( + timers_to_log, wandb_writer, iteration, normalizer=args.log_interval, reset=False + ) # Log timers to stdout timers.log(timers_to_log, normalizer=args.log_interval, reset=should_reset) ########## FlagScale Begin ########## @@ -2629,11 +2851,11 @@ def training_log( if torch.distributed.get_rank() == 0: num_microbatches = get_num_microbatches() report_theoretical_memory(args, num_microbatches=num_microbatches, verbose=True) - report_memory(f'(after {iteration} iterations)') + report_memory(f"(after {iteration} iterations)") report_memory_flag = False else: num_microbatches = get_num_microbatches() - report_memory(f'(after {iteration} iterations)') + report_memory(f"(after {iteration} iterations)") report_memory_flag = False ########## FlagScale End ########## @@ -2663,7 +2885,7 @@ def compute_throughputs_and_append_to_progress_log(iteration, num_floating_point ) / (elapsed_time * 10**12 * args.world_size) tokens_so_far = args.consumed_train_samples * args.seq_length - saved_ckpt_prefix = 'Saving async checkpoint' if args.async_save else 'Saved checkpoint' + saved_ckpt_prefix = "Saving async checkpoint" if args.async_save else "Saved checkpoint" append_to_progress_log( f"{saved_ckpt_prefix}\tIteration: {iteration}\t" f"Job throughput: {job_throughput:.1f} TFLOP/s/GPU\t" @@ -2690,10 +2912,12 @@ def force_param_sync(model_chunks: list[DDP]) -> None: assert isinstance(model_chunk, DDP) model_chunk.start_param_sync(force_sync=True) + # Only report memory for first 3 checkpoint saves. num_checkpoints_memory_reported = 0 MAX_NUM_CHECKPOINTS_MEMORY_REPORTED = 3 + def save_checkpoint_and_time( iteration, model, @@ -2714,10 +2938,11 @@ def save_checkpoint_and_time( # Stop timer to get accurate train interval time and exclude checkpointing duration timers('interval-time').stop() - energy_monitor.pause() + if args.log_energy: + energy_monitor.pause() # Extra barrier is added to make sure all ranks report the max time. - timer_key = 'save-checkpoint-non-persistent' if non_persistent_ckpt else 'save-checkpoint' + timer_key = "save-checkpoint-non-persistent" if non_persistent_ckpt else "save-checkpoint" timers(timer_key, log_level=0).start(barrier=True) # Log E2E metrics before save-checkpoint @@ -2726,7 +2951,7 @@ def save_checkpoint_and_time( # that the async checkpoint worker process has enough GPU headroom for # D2H tensor transfers. for model_chunk in model: - if hasattr(model_chunk, 'free_overlap_buffers'): + if hasattr(model_chunk, "free_overlap_buffers"): model_chunk.free_overlap_buffers() cur_platform.empty_cache() @@ -2772,7 +2997,8 @@ def save_checkpoint_and_time( ) # Recover timing - energy_monitor.resume() + if args.log_energy: + energy_monitor.resume() timers('interval-time', log_level=0).start(barrier=True) @@ -2783,7 +3009,7 @@ def post_training_step_callbacks( iteration, prof, num_floating_point_operations_since_last_log_event, - nsys_nvtx_context = None, + nsys_nvtx_context=None, ): """Run all post-training-step functions (e.g., FT heartbeats, GC).""" args = get_args() @@ -2805,9 +3031,9 @@ def post_training_step_callbacks( ): if should_disable_forward_pre_hook(args): disable_forward_pre_hook(model) - assert check_param_hashes_across_dp_replicas( - model, cross_check=True - ), "Parameter hashes not matching across DP replicas" + assert check_param_hashes_across_dp_replicas(model, cross_check=True), ( + "Parameter hashes not matching across DP replicas" + ) torch.distributed.barrier() print_rank_0(f">>> Weight hashes match after {iteration} iterations...") if should_disable_forward_pre_hook(args): @@ -2821,8 +3047,7 @@ def post_training_step_callbacks( if ( args.profile and iteration == args.profile_step_end - and (len(args.profile_ranks) == 0 or - torch.distributed.get_rank() in args.profile_ranks) + and (len(args.profile_ranks) == 0 or torch.distributed.get_rank() in args.profile_ranks) ): if args.use_pytorch_profiler: assert prof is not None @@ -2873,7 +3098,7 @@ def checkpoint_and_decide_exit( checkpointing_context, train_data_iterator=train_data_iterator, ) - print_datetime('exiting program after receiving SIGTERM.') + print_datetime("exiting program after receiving SIGTERM.") return True @@ -2911,7 +3136,9 @@ def checkpoint_and_decide_exit( if args.exit_duration_in_mins: train_time = (time.time() - _TRAIN_START_TIME) / 60.0 done_cuda = torch.tensor( - [train_time > args.exit_duration_in_mins], dtype=torch.int, device=cur_platform.device_name() + [train_time > args.exit_duration_in_mins], + dtype=torch.int, + device=cur_platform.device_name(), ) torch.distributed.all_reduce(done_cuda, op=torch.distributed.ReduceOp.MAX) done = done_cuda.item() @@ -2926,17 +3153,13 @@ def checkpoint_and_decide_exit( checkpointing_context, train_data_iterator=train_data_iterator, ) - print_datetime(f'exiting program after {train_time} minutes') + print_datetime(f"exiting program after {train_time} minutes") return True # Exit based on iterations. - if ( - args.exit_interval - and iteration % args.exit_interval == 0 - ) or ( - args.phase_transition_iterations - and iteration in args.phase_transition_iterations + if (args.exit_interval and iteration % args.exit_interval == 0) or ( + args.phase_transition_iterations and iteration in args.phase_transition_iterations ): if args.save and not saved_checkpoint: save_checkpoint_and_time( @@ -2948,7 +3171,7 @@ def checkpoint_and_decide_exit( checkpointing_context, train_data_iterator=train_data_iterator, ) - print_datetime(f'exiting program at iteration {iteration}') + print_datetime(f"exiting program at iteration {iteration}") return True @@ -2967,7 +3190,7 @@ def train( checkpointing_context, non_loss_data_func, inference_model=None, - extra_valid_dataset_provider=None, ####### FlagScale Add ####### + extra_valid_dataset_provider=None, ####### FlagScale Add ####### ): """Training function: run train_step desired number of times, run validation, checkpoint.""" args = get_args() @@ -2993,29 +3216,31 @@ def train( args.load = None args.finetune = True load_checkpoint( - model, - None, # Don't load optimizer state - None, # Don't load scheduler state - checkpointing_context=checkpointing_context, - skip_load_to_model_and_opt=HAVE_FSDP2 - and getattr(args, "use_torch_fsdp2", False) - and args.ckpt_format == "torch_dist", - ) - ref_state_dict = {k: (v.cpu() if v is not None else v) for k, v in model[0].state_dict().items()} + model, + None, # Don't load optimizer state + None, # Don't load scheduler state + checkpointing_context=checkpointing_context, + skip_load_to_model_and_opt=HAVE_FSDP2 + and getattr(args, "use_torch_fsdp2", False) + and args.ckpt_format == "torch_dist", + ) + ref_state_dict = { + k: (v.cpu() if v is not None else v) for k, v in model[0].state_dict().items() + } # Reload RL training checkpoint weights args.load = load args.finetune = finetune print_rank_0("> Reloading RL training checkpoint...") load_checkpoint( - model, - None, - None, - checkpointing_context=checkpointing_context, - skip_load_to_model_and_opt=HAVE_FSDP2 - and getattr(args, "use_torch_fsdp2", False) - and args.ckpt_format == "torch_dist", - ) + model, + None, + None, + checkpointing_context=checkpointing_context, + skip_load_to_model_and_opt=HAVE_FSDP2 + and getattr(args, "use_torch_fsdp2", False) + and args.ckpt_format == "torch_dist", + ) args.no_load_optim = no_load_optim @@ -3024,8 +3249,9 @@ def train( print_rank_0("> Reinitializing microbatch calculator for GRPO training...") from megatron.core.num_microbatches_calculator import ( destroy_num_microbatches_calculator, - init_num_microbatches_calculator + init_num_microbatches_calculator, ) + # First destroy the existing calculator destroy_num_microbatches_calculator() # Then initialize with the correct perform_rl_step=True context @@ -3035,7 +3261,7 @@ def train( args.global_batch_size, args.micro_batch_size, mpu.get_data_parallel_world_size(), - args.decrease_batch_size_if_needed + args.decrease_batch_size_if_needed, ) print_rank_0(f"> GRPO training: num_microbatches set to {get_num_microbatches()}") @@ -3073,8 +3299,10 @@ def train( # Make sure rerun_state_machine has the right iteration loaded from checkpoint. rerun_state_machine = get_rerun_state_machine() if rerun_state_machine.current_iteration != iteration: - print_rank_0(f"Overwriting rerun_state_machine.current_iteration from " - f"{rerun_state_machine.current_iteration} to {iteration}...") + print_rank_0( + f"Overwriting rerun_state_machine.current_iteration from " + f"{rerun_state_machine.current_iteration} to {iteration}..." + ) rerun_state_machine.current_iteration = iteration # Track E2E metrics at the start of training. @@ -3097,8 +3325,8 @@ def train( config.timers = timers if isinstance(model[0], (megatron_FSDP, DDP)) and args.overlap_grad_reduce: assert config.no_sync_func is None, ( - 'When overlap_grad_reduce is True, config.no_sync_func must be None; ' - 'a custom no_sync_func is not supported when overlapping grad-reduce' + "When overlap_grad_reduce is True, config.no_sync_func must be None; " + "a custom no_sync_func is not supported when overlapping grad-reduce" ) config.no_sync_func = [model_chunk.no_sync for model_chunk in model] if len(model) == 1: @@ -3117,8 +3345,8 @@ def train( energy_monitor.setup() energy_monitor.resume() - timers('interval-time', log_level=0).start(barrier=True) - print_datetime('before the start of training step') + timers("interval-time", log_level=0).start(barrier=True) + print_datetime("before the start of training step") report_memory_flag = True perf_callback = initialize_perf_monitor(args) pre_hook_enabled = False @@ -3129,9 +3357,9 @@ def train( if args.manual_gc: # Disable the default garbage collector and perform the collection manually. # This is to align the timing of garbage collection across ranks. - assert ( - args.manual_gc_interval >= 0 - ), 'Manual garbage collection interval should be larger than or equal to 0' + assert args.manual_gc_interval >= 0, ( + "Manual garbage collection interval should be larger than or equal to 0" + ) gc.disable() gc.collect() @@ -3170,9 +3398,13 @@ def train( # Wrap forward_backward_func for Full iteration CUDA graph forward_backward_func = get_forward_backward_func() if args.cuda_graph_impl == "local" and CudaGraphScope.full_iteration in args.cuda_graph_scope: - forward_backward_func = FullCudaGraphWrapper(forward_backward_func, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps) + forward_backward_func = FullCudaGraphWrapper( + forward_backward_func, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps + ) if args.optimizer_cuda_graph: - optimizer.step = OptimizerCudaGraphWrapper(optimizer.step, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps) + optimizer.step = OptimizerCudaGraphWrapper( + optimizer.step, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps + ) def get_e2e_base_metrics(): """Get base metrics values for one-logger to calculate E2E tracking metrics.""" @@ -3180,44 +3412,47 @@ def get_e2e_base_metrics(): num_floating_point_operations_so_far - args.num_floating_point_operations_so_far ) return { - 'iteration': iteration, - 'train_duration': timers('interval-time').active_time(), - 'eval_duration': eval_duration, - 'eval_iterations': eval_iterations, - 'total_flops_since_current_train_start': num_floating_point_operations_since_current_train_start, - 'num_floating_point_operations_so_far': num_floating_point_operations_so_far, - 'consumed_train_samples': args.consumed_train_samples, - 'world_size': args.world_size, - 'seq_length': args.seq_length, + "iteration": iteration, + "train_duration": timers("interval-time").active_time(), + "eval_duration": eval_duration, + "eval_iterations": eval_iterations, + "total_flops_since_current_train_start": num_floating_point_operations_since_current_train_start, + "num_floating_point_operations_so_far": num_floating_point_operations_so_far, + "consumed_train_samples": args.consumed_train_samples, + "world_size": args.world_size, + "seq_length": args.seq_length, ########## FlagScale Begin ########## - 'extra_eval_duration': extra_eval_duration, - 'extra_eval_iterations': extra_eval_iterations, + "extra_eval_duration": extra_eval_duration, + "extra_eval_iterations": extra_eval_iterations, ########## FlagScale End ########## } # Cache into one-logger for callback. if one_logger: with one_logger.get_context_manager(): - one_logger.store_set('get_e2e_base_metrics', get_e2e_base_metrics) + one_logger.store_set("get_e2e_base_metrics", get_e2e_base_metrics) prof = None - nsys_nvtx_context = None # reference to context for nsys profiling, so it can be cleaned up + nsys_nvtx_context = None # reference to context for nsys profiling, so it can be cleaned up if ( args.profile - and (len(args.profile_ranks) == 0 or - torch.distributed.get_rank() in args.profile_ranks) + and (len(args.profile_ranks) == 0 or torch.distributed.get_rank() in args.profile_ranks) and args.use_pytorch_profiler ): if args.pytorch_profiler_collect_chakra: et_dir = Path(f"{args.tensorboard_dir}/../chakra") et_dir.mkdir(parents=True, exist_ok=True) - et = torch.profiler.ExecutionTraceObserver().register_callback(f"{et_dir}/rank-{torch.distributed.get_rank()}.json.gz") + et = torch.profiler.ExecutionTraceObserver().register_callback( + f"{et_dir}/rank-{torch.distributed.get_rank()}.json.gz" + ) else: et = None + def trace_handler(p): profile_dir = Path(f"{args.tensorboard_dir}/../torch_profile") profile_dir.mkdir(parents=True, exist_ok=True) p.export_chrome_trace(f"{profile_dir}/rank-{torch.distributed.get_rank()}.json.gz") + prof = torch.profiler.profile( schedule=torch.profiler.schedule( wait=max(args.profile_step_start - 1, 0), @@ -3245,9 +3480,9 @@ def trace_handler(p): pre_hook_enabled = False # Also, check weight hash across DP replicas to be very pedantic. if args.check_weight_hash_across_dp_replicas_interval is not None: - assert check_param_hashes_across_dp_replicas( - model, cross_check=True - ), "Parameter hashes not matching across DP replicas" + assert check_param_hashes_across_dp_replicas(model, cross_check=True), ( + "Parameter hashes not matching across DP replicas" + ) torch.distributed.barrier() print_rank_0(f">>> Weight hashes match after {iteration} iterations...") @@ -3264,9 +3499,9 @@ def trace_handler(p): # Run training iterations till done. buffered_rollouts = None while iteration < args.train_iters: - if (args.profile - and (len(args.profile_ranks) == 0 or - torch.distributed.get_rank() in args.profile_ranks)): + if args.profile and ( + len(args.profile_ranks) == 0 or torch.distributed.get_rank() in args.profile_ranks + ): if args.use_pytorch_profiler: prof.step() elif iteration == args.profile_step_start: @@ -3280,7 +3515,10 @@ def trace_handler(p): # Update the timeout for all process groups after initialization # We update the timeout after the first successful iteration, # which takes longer than others usually - if args.distributed_timeout_seconds_after_init is not None and iteration == start_iteration+1: + if ( + args.distributed_timeout_seconds_after_init is not None + and iteration == start_iteration + 1 + ): # TODO: some dynamic timeout setting is required # based on the iteration time considering interval-based steps (e.g. eval, checkpoint) # e.g. timeout for normal iterations vs timeout for iterations with checkpoint @@ -3295,8 +3533,11 @@ def trace_handler(p): update_num_microbatches(args.consumed_train_samples, consistency_check=False, verbose=True) # Skip automatic checkpoint on microbatch changes when sequence packing is active # as it intentionally reconfigures microbatches - if get_num_microbatches() != num_microbatches and iteration != 0 \ - and args.save_when_num_microbatches_change: ####### FlagScale Add ####### + if ( + get_num_microbatches() != num_microbatches + and iteration != 0 + and args.save_when_num_microbatches_change + ): ####### FlagScale Add ####### if args.rl_use_sequence_packing: print_rank_0( f"[Sequence Packing] Skipping automatic checkpoint at iteration {iteration} " @@ -3355,9 +3596,20 @@ def trace_handler(p): start_skip_iteration = 0 end_skip_iteration = 0 if args.skip_samples_range: - if args.consumed_train_samples + current_global_batch_size > args.skip_samples_range[0] and args.consumed_train_samples < args.skip_samples_range[1]: - num_skipped_iters = (args.skip_samples_range[1] - args.consumed_train_samples + current_global_batch_size - 1) // current_global_batch_size - args.skip_samples_range[1] = args.consumed_train_samples + num_skipped_iters * current_global_batch_size + if ( + args.consumed_train_samples + current_global_batch_size + > args.skip_samples_range[0] + and args.consumed_train_samples < args.skip_samples_range[1] + ): + num_skipped_iters = ( + args.skip_samples_range[1] + - args.consumed_train_samples + + current_global_batch_size + - 1 + ) // current_global_batch_size + args.skip_samples_range[1] = ( + args.consumed_train_samples + num_skipped_iters * current_global_batch_size + ) start_skip_iteration = iteration end_skip_iteration = iteration + num_skipped_iters else: @@ -3368,10 +3620,14 @@ def trace_handler(p): if mpu.is_pipeline_first_stage() or mpu.is_pipeline_last_stage(): for _ in range(get_num_microbatches()): _ = next(train_data_iterator) - args.consumed_train_samples += mpu.get_data_parallel_world_size() * \ - args.micro_batch_size * \ - get_num_microbatches() - update_num_microbatches(args.consumed_train_samples, consistency_check=True, verbose=True) + args.consumed_train_samples += ( + mpu.get_data_parallel_world_size() + * args.micro_batch_size + * get_num_microbatches() + ) + update_num_microbatches( + args.consumed_train_samples, consistency_check=True, verbose=True + ) iteration += 1 args.curr_iteration = iteration @@ -3389,7 +3645,11 @@ def trace_handler(p): cur_platform.empty_cache() with torch.no_grad(): train_data_iterator = rl_utils.get_grpo_data_iterator( - model, inference_model, optimizer, iteration, ref_state_dict, + model, + inference_model, + optimizer, + iteration, + ref_state_dict, grpo_iterations=args.grpo_iterations, grpo_prompts_per_step=args.grpo_prompts_per_step, grpo_group_size=args.grpo_group_size, @@ -3426,7 +3686,14 @@ def trace_handler(p): num_zeros_in_grad, max_attention_logit, ) = train_step( - forward_step_func, train_data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=iteration + forward_step_func, + train_data_iterator, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=iteration, ) ft_integration.on_training_step_end() if perf_callback is not None: @@ -3466,9 +3733,9 @@ def trace_handler(p): args.cuda_graph_impl == "transformer_engine" and args.cuda_graph_warmup_steps == 0 ): - assert ( - cuda_graph_helper.capture_finished() - ), "CUDA Graph capture should have been finished." + assert cuda_graph_helper.capture_finished(), ( + "CUDA Graph capture should have been finished." + ) cuda_graph_helper.cuda_graph_set_manual_hooks() iteration += 1 @@ -3477,7 +3744,7 @@ def trace_handler(p): if ( getattr(args, "fsdp_manual_registration", False) and getattr(args, "use_megatron_fsdp", False) - and iteration == start_iteration + 1 + and iteration == start_iteration + 1 ): for model_chunk in model: if isinstance(model_chunk, megatron_FSDP) and getattr( @@ -3554,21 +3821,23 @@ def trace_handler(p): if fs_straggler.should_report(iteration): report = fs_straggler.generate_report(step=iteration) if _is_global_rank_zero(): - _save_straggler_report(report, getattr(args, "straggler_log_dir", None), iteration) + _save_straggler_report( + report, getattr(args, "straggler_log_dir", None), iteration + ) # Evaluation. if args.eval_interval and iteration % args.eval_interval == 0 and args.do_valid: if args.log_energy: energy_monitor.pause() - timers('interval-time').stop() + timers("interval-time").stop() if should_disable_forward_pre_hook(args): disable_forward_pre_hook(model) pre_hook_enabled = False if args.manual_gc and args.manual_gc_eval: # Collect all objects. gc.collect() - prefix = f'iteration {iteration}' - timers('eval-time', log_level=0).start(barrier=True) + prefix = f"iteration {iteration}" + timers("eval-time", log_level=0).start(barrier=True) if args.perform_rl_step: rl_eval_model = model rl_training_model = None @@ -3591,15 +3860,24 @@ def trace_handler(p): training_model=rl_training_model, ) else: - evaluate_and_print_results(prefix, forward_step_func, - valid_data_iterator, model, - iteration, process_non_loss_data_func, - config, verbose=False, write_to_tensorboard=True, - non_loss_data_func=non_loss_data_func) - - eval_duration += timers('eval-time').elapsed() - eval_iterations += sum(args.eval_iters) if isinstance(args.eval_iters, list) else args.eval_iters - timers('eval-time').stop() + evaluate_and_print_results( + prefix, + forward_step_func, + valid_data_iterator, + model, + iteration, + process_non_loss_data_func, + config, + verbose=False, + write_to_tensorboard=True, + non_loss_data_func=non_loss_data_func, + ) + + eval_duration += timers("eval-time").elapsed() + eval_iterations += ( + sum(args.eval_iters) if isinstance(args.eval_iters, list) else args.eval_iters + ) + timers("eval-time").stop() one_logger_utils.track_e2e_metrics() if args.manual_gc and args.manual_gc_eval: @@ -3608,7 +3886,7 @@ def trace_handler(p): if should_disable_forward_pre_hook(args): enable_forward_pre_hook(model) pre_hook_enabled = True - timers('interval-time', log_level=0).start(barrier=True) + timers("interval-time", log_level=0).start(barrier=True) if args.log_energy: energy_monitor.resume() if args.num_experts is not None: @@ -3617,7 +3895,11 @@ def trace_handler(p): ######## FlagScale Begin ######## # Extra Evaluation ===================================================================== if args.extra_eval_interval and iteration % args.extra_eval_interval == 0: - from megatron.training.extra_valid import extra_evaluate_and_print_results, build_extra_valid_data_iterators + from megatron.training.extra_valid import ( + extra_evaluate_and_print_results, + build_extra_valid_data_iterators, + ) + # NOTE(zhaoyinglia): Must rebuild the dataloaders for extra validation here, # to guarantee extra validation start from extra_iter=0 every time, # but we don't need to rebuild the datasets. @@ -3630,14 +3912,13 @@ def trace_handler(p): elif args.use_dualpipev: extra_valid_data_iterator = [] for _ in range(2): - extra_iterators = build_extra_valid_data_iterators( - extra_valid_dataset_provider) + extra_iterators = build_extra_valid_data_iterators(extra_valid_dataset_provider) extra_valid_data_iterator.append(extra_iterators) else: - extra_valid_data_iterator = ( - build_extra_valid_data_iterators(extra_valid_dataset_provider) + extra_valid_data_iterator = build_extra_valid_data_iterators( + extra_valid_dataset_provider ) - timers('interval-time').stop() + timers("interval-time").stop() # do_extra_valid flag is used to indicate that we are doing extra validation # and is set in the build_extra_valid_data_iterators function if getattr(args, "do_extra_valid", False): @@ -3647,9 +3928,9 @@ def trace_handler(p): if args.manual_gc and args.manual_gc_eval: # Collect all objects. gc.collect() - prefix = 'iteration {}'.format(iteration) + prefix = "iteration {}".format(iteration) for extra_valid_index, extra_valid_data_itr in enumerate(extra_valid_data_iterator): - timers('extra-eval-time', log_level=0).start(barrier=True) + timers("extra-eval-time", log_level=0).start(barrier=True) extra_eval_iters = args.extra_eval_iters_list[extra_valid_index] extra_evaluate_and_print_results( extra_valid_index, @@ -3662,11 +3943,11 @@ def trace_handler(p): config, verbose=False, write_to_tensorboard=True, - non_loss_data_func=non_loss_data_func + non_loss_data_func=non_loss_data_func, ) - extra_eval_duration += timers('extra-eval-time').elapsed() + extra_eval_duration += timers("extra-eval-time").elapsed() extra_eval_iterations += extra_eval_iters - timers('extra-eval-time').stop() + timers("extra-eval-time").stop() one_logger_utils.track_e2e_metrics() if args.manual_gc and args.manual_gc_eval: @@ -3675,7 +3956,7 @@ def trace_handler(p): if should_disable_forward_pre_hook(args): enable_forward_pre_hook(model) pre_hook_enabled = True - timers('interval-time', log_level=0).start(barrier=True) + timers("interval-time", log_level=0).start(barrier=True) # ======================================================================================= ######## FlagScale End ######## # Miscellaneous post-training-step functions (e.g., FT heartbeats, GC). @@ -3759,13 +4040,13 @@ def evaluate( verbose=False, non_loss_data_func=None, eval_iters=None, - extra_valid_index=None, ####### FlagScale Add ####### + extra_valid_index=None, ####### FlagScale Add ####### ): """Evaluation.""" args = get_args() timers = get_timers() - timers('evaluate', log_level=0).start(barrier=True) + timers("evaluate", log_level=0).start(barrier=True) if args.vision_pretraining and args.vision_pretraining_type == "dino": from megatron.legacy.model.vision.knn_monitor import compute_feature_bank @@ -3788,7 +4069,9 @@ def evaluate( eval_num_microbatches = eval_batch_size // (args.micro_batch_size * args.data_parallel_size) forward_backward_func = get_forward_backward_func() if args.cuda_graph_impl == "local" and CudaGraphScope.full_iteration in args.cuda_graph_scope: - forward_backward_func = FullCudaGraphWrapper(forward_backward_func, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps) + forward_backward_func = FullCudaGraphWrapper( + forward_backward_func, cuda_graph_warmup_steps=args.cuda_graph_warmup_steps + ) if has_nvidia_modelopt: # [ModelOpt]: Pipeline-parallel Distillation stacks student and teacher tensors @@ -3803,8 +4086,9 @@ def evaluate( ######## FlagScale Begin ######## if extra_valid_index is not None: - assert getattr(args, "extra_eval_iters_list") is not None, \ + assert getattr(args, "extra_eval_iters_list") is not None, ( "extra_eval_iters_list must be provided if extra_valid_index is not None" + ) eval_iters = args.extra_eval_iters_list[extra_valid_index] elif eval_iters is None: eval_iters = args.eval_iters @@ -3813,11 +4097,11 @@ def evaluate( with torch.no_grad(): iteration = 0 if verbose: - print_rank_0(f'Evaluating on {eval_iters * eval_batch_size} samples') + print_rank_0(f"Evaluating on {eval_iters * eval_batch_size} samples") while iteration < eval_iters: iteration += 1 if verbose: - print_rank_0(f'Evaluating iter {iteration}/{eval_iters}') + print_rank_0(f"Evaluating iter {iteration}/{eval_iters}") # Don't care about timing during evaluation config.timers = None @@ -3844,7 +4128,9 @@ def evaluate( # Reduce across processes. for key in loss_dicts[0].keys(): if key not in total_loss_dict: - total_loss_dict[key] = torch.tensor([0.0, 0.0], dtype=torch.float, device=cur_platform.device()) + total_loss_dict[key] = torch.tensor( + [0.0, 0.0], dtype=torch.float, device=cur_platform.device() + ) val = [x[key].view(-1) for x in loss_dicts] if val[0].numel() == 2: @@ -3854,19 +4140,17 @@ def evaluate( val = val[:, 0] / val[:, 1].clamp(min=1) val = val.mean() torch.distributed.all_reduce( - val, - group=mpu.get_data_parallel_group(with_context_parallel=True) + val, group=mpu.get_data_parallel_group(with_context_parallel=True) ) val /= torch.distributed.get_world_size( group=mpu.get_data_parallel_group(with_context_parallel=True) ) total_loss_dict[key][0] += val total_loss_dict[key][1] += 1 - else : + else: val = torch.vstack(val).sum(dim=0) torch.distributed.all_reduce( - val, - group=mpu.get_data_parallel_group(with_context_parallel=True) + val, group=mpu.get_data_parallel_group(with_context_parallel=True) ) total_loss_dict[key] += val elif val[0].numel() == 1: @@ -3881,13 +4165,15 @@ def evaluate( if args.exit_duration_in_mins: train_time = (time.time() - _TRAIN_START_TIME) / 60.0 done_cuda = torch.tensor( - [train_time > args.exit_duration_in_mins], dtype=torch.int, device=cur_platform.device() + [train_time > args.exit_duration_in_mins], + dtype=torch.int, + device=cur_platform.device(), ) torch.distributed.all_reduce(done_cuda, op=torch.distributed.ReduceOp.MAX) done = done_cuda.item() if done: rerun_state_machine.set_mode(rerun_mode) - print_rank_0('Exiting during evaluation, timelimit reached') + print_rank_0("Exiting during evaluation, timelimit reached") return None, None, True collected_non_loss_data = None @@ -3914,8 +4200,8 @@ def evaluate( numerator, denominator = total_loss_dict[key] total_loss_dict[key] = numerator / denominator - timers('evaluate').stop() - timers.log(['evaluate']) + timers("evaluate").stop() + timers.log(["evaluate"]) rerun_state_machine.set_mode(rerun_mode) @@ -3955,9 +4241,13 @@ def evaluate_and_print_results( # with full validation we need to distribute eval_iters to all ranks if mpu.get_tensor_model_parallel_rank() == 0: - eval_iters = torch.tensor(args.eval_iters, dtype=torch.long, device=cur_platform.device_name()) + eval_iters = torch.tensor( + args.eval_iters, dtype=torch.long, device=cur_platform.device_name() + ) else: - eval_iters = torch.tensor([0] * len(eval_iters), dtype=torch.long, device=cur_platform.device_name()) + eval_iters = torch.tensor( + [0] * len(eval_iters), dtype=torch.long, device=cur_platform.device_name() + ) torch.distributed.broadcast(eval_iters, 0) eval_iters = eval_iters.tolist() args.eval_iters = eval_iters[0] if not args.multiple_validation_sets else eval_iters @@ -3983,35 +4273,40 @@ def evaluate_and_print_results( # Timelimit hit during evaluation if timelimit: return - string = f' validation{suffix} loss at {prefix} | ' + string = f" validation{suffix} loss at {prefix} | " for key in total_loss_dict: - string += '{} value: {:.6E} | '.format(key, total_loss_dict[key].item()) + string += "{} value: {:.6E} | ".format(key, total_loss_dict[key].item()) ppl = math.exp(min(20, total_loss_dict[key].item())) - string += '{} PPL: {:.6E} | '.format(key, ppl) + string += "{} PPL: {:.6E} | ".format(key, ppl) if writer: - writer.add_scalar('{} validation{}'.format(key, suffix), total_loss_dict[key].item(), iteration) writer.add_scalar( - '{} validation{} vs samples'.format(key, suffix), + "{} validation{}".format(key, suffix), total_loss_dict[key].item(), iteration + ) + writer.add_scalar( + "{} validation{} vs samples".format(key, suffix), total_loss_dict[key].item(), args.consumed_train_samples, ) if args.log_validation_ppl_to_tensorboard: - writer.add_scalar('{} validation{} ppl'.format(key, suffix), ppl, iteration) + writer.add_scalar("{} validation{} ppl".format(key, suffix), ppl, iteration) writer.add_scalar( - '{} validation{} ppl vs samples'.format(key, suffix), ppl, args.consumed_train_samples + "{} validation{} ppl vs samples".format(key, suffix), + ppl, + args.consumed_train_samples, ) if wandb_writer and is_last_rank(): wandb_writer.log( - {'{} validation{}'.format(key, suffix): total_loss_dict[key].item()}, iteration + {"{} validation{}".format(key, suffix): total_loss_dict[key].item()}, + iteration, ) if process_non_loss_data_func is not None and writer and is_last_rank(): process_non_loss_data_func(collected_non_loss_data, iteration, writer) length = len(string) + 1 - print_rank_last('-' * length) + print_rank_last("-" * length) print_rank_last(string) - print_rank_last('-' * length) + print_rank_last("-" * length) def cyclic_iter(iter): @@ -4043,7 +4338,11 @@ def get_train_valid_test_num_samples(): # Get train_samples in current phase. if args.phase_transition_iterations: - phase_transition_samples = [0] + [t * args.global_batch_size for t in args.phase_transition_iterations] + [args.train_samples] + phase_transition_samples = ( + [0] + + [t * args.global_batch_size for t in args.phase_transition_iterations] + + [args.train_samples] + ) current_sample = args.iteration * args.global_batch_size last_transition_sample = max(s for s in phase_transition_samples if s <= current_sample) next_transition_sample = min(s for s in phase_transition_samples if s > current_sample) @@ -4054,14 +4353,16 @@ def get_train_valid_test_num_samples(): return (train_samples_in_current_phase, eval_samples, test_samples) -def build_train_valid_test_datasets(build_train_valid_test_datasets_provider, train_valid_test_num_samples=None): +def build_train_valid_test_datasets( + build_train_valid_test_datasets_provider, train_valid_test_num_samples=None +): """Build pretraining datasets.""" if train_valid_test_num_samples is None: train_valid_test_num_samples = get_train_valid_test_num_samples() - print_rank_0(' > datasets target sizes (minimum size):') - print_rank_0(' train: {}'.format(train_valid_test_num_samples[0])) - print_rank_0(' validation: {}'.format(train_valid_test_num_samples[1])) - print_rank_0(' test: {}'.format(train_valid_test_num_samples[2])) + print_rank_0(" > datasets target sizes (minimum size):") + print_rank_0(" train: {}".format(train_valid_test_num_samples[0])) + print_rank_0(" validation: {}".format(train_valid_test_num_samples[1])) + print_rank_0(" test: {}".format(train_valid_test_num_samples[2])) return build_train_valid_test_datasets_provider(train_valid_test_num_samples) @@ -4072,13 +4373,13 @@ def build_train_valid_test_data_loaders(build_train_valid_test_datasets_provider (train_dataloader, valid_dataloaders, test_dataloader) = (None, None, None) - print_rank_0('> building train, validation, and test datasets ...') + print_rank_0("> building train, validation, and test datasets ...") # Backward compatibility, assume fixed batch size. if args.iteration > 0 and args.consumed_train_samples == 0: - assert ( - args.train_samples is None - ), 'Only backward compatiblity support for iteration-based training' + assert args.train_samples is None, ( + "Only backward compatiblity support for iteration-based training" + ) args.consumed_train_samples = args.iteration * args.global_batch_size if args.iteration > 0 and args.consumed_valid_samples == 0: @@ -4089,8 +4390,14 @@ def build_train_valid_test_data_loaders(build_train_valid_test_datasets_provider # Get consumed train samples in this phase. if args.phase_transition_iterations: - last_transition = max(iteration for iteration in (0, *args.phase_transition_iterations) if iteration <= args.iteration) - consumed_train_samples_in_current_phase = (args.iteration - last_transition) * args.global_batch_size + last_transition = max( + iteration + for iteration in (0, *args.phase_transition_iterations) + if iteration <= args.iteration + ) + consumed_train_samples_in_current_phase = ( + args.iteration - last_transition + ) * args.global_batch_size else: consumed_train_samples_in_current_phase = args.consumed_train_samples @@ -4099,7 +4406,6 @@ def build_train_valid_test_data_loaders(build_train_valid_test_datasets_provider # Construct the data pipeline if is_distributed or mpu.get_tensor_model_parallel_rank() == 0: - # Build datasets and dataloders. if args.perform_rl_step: # we don't need to build any dataloaders for RL training @@ -4107,17 +4413,21 @@ def build_train_valid_test_data_loaders(build_train_valid_test_datasets_provider valid_dataloaders = None test_dataloader = None do_train = (args.train_iters or 0) > 0 - do_valid = (args.full_validation or args.eval_iters > 0) - do_test = (args.full_validation or args.eval_iters > 0) + do_valid = args.full_validation or args.eval_iters > 0 + do_test = args.full_validation or args.eval_iters > 0 else: # Build datasets. - train_ds, valid_ds, test_ds = build_train_valid_test_datasets(build_train_valid_test_datasets_provider) + train_ds, valid_ds, test_ds = build_train_valid_test_datasets( + build_train_valid_test_datasets_provider + ) valid_ds = [valid_ds] if not isinstance(valid_ds, list) else valid_ds if args.skip_train: train_dataloader = None else: - train_dataloader = build_pretraining_data_loader(train_ds, consumed_train_samples_in_current_phase) + train_dataloader = build_pretraining_data_loader( + train_ds, consumed_train_samples_in_current_phase + ) valid_dataloaders = [] for valid_d in valid_ds: if args.skip_train or args.full_validation: @@ -4126,17 +4436,25 @@ def build_train_valid_test_data_loaders(build_train_valid_test_datasets_provider if args.multiple_validation_sets: # TODO(bnorick): for multiple validation sets without full validation, args.consumed_valid_samples is not # correct and needs to be calculated/set per validation set - raise NotImplementedError("--multiple-validation-sets currently requires --full-validation") - valid_dataloaders.append(build_pretraining_data_loader(valid_d, args.consumed_valid_samples)) + raise NotImplementedError( + "--multiple-validation-sets currently requires --full-validation" + ) + valid_dataloaders.append( + build_pretraining_data_loader(valid_d, args.consumed_valid_samples) + ) if not args.multiple_validation_sets: assert len(valid_dataloaders) == 1 test_dataloader = build_pretraining_data_loader(test_ds, 0) do_train = train_dataloader is not None and (args.skip_train or args.train_iters > 0) - do_valid = valid_dataloaders is not None and (args.full_validation or args.eval_iters > 0) + do_valid = valid_dataloaders is not None and ( + args.full_validation or args.eval_iters > 0 + ) do_test = test_dataloader is not None and (args.full_validation or args.eval_iters > 0) flags = torch.tensor( - [int(do_train), int(do_valid), int(do_test)], dtype=torch.long, device=cur_platform.device_name() + [int(do_train), int(do_valid), int(do_test)], + dtype=torch.long, + device=cur_platform.device_name(), ) else: flags = torch.tensor([0, 0, 0], dtype=torch.long, device=cur_platform.device_name()) @@ -4161,7 +4479,7 @@ def build_train_valid_test_data_iterators(build_train_valid_test_datasets_provid # Build iterators. dl_type = args.dataloader_type - assert dl_type in ['single', 'cyclic', 'external'] + assert dl_type in ["single", "cyclic", "external"] def _get_iterator(dataloader_type, dataloader): """Return dataset iterator.""" @@ -4190,7 +4508,7 @@ def _get_iterator(dataloader_type, dataloader): if args.full_validation: if args.multiple_validation_sets: if valid_dataloaders[0] is None: - args.eval_iters = [None]*len(valid_dataloaders) + args.eval_iters = [None] * len(valid_dataloaders) else: args.eval_iters = [len(dl) for dl in valid_dataloaders] else: @@ -4202,8 +4520,9 @@ def _get_iterator(dataloader_type, dataloader): else: valid_dl_type = "cyclic" if args.full_validation else dl_type print( - f"[VALID DATA LOADER LENGTHS] " - ", ".join(f"{idx}: {len(dl)}" for idx, dl in enumerate(valid_dataloaders)) + f"[VALID DATA LOADER LENGTHS] , ".join( + f"{idx}: {len(dl)}" for idx, dl in enumerate(valid_dataloaders) + ) ) valid_data_iterators = [ _get_iterator(valid_dl_type, dl) for dl in valid_dataloaders @@ -4227,6 +4546,6 @@ def should_disable_forward_pre_hook(args): """Block forward pre-hook for certain configurations.""" return ( not args.use_megatron_fsdp - and (args.use_distributed_optimizer or 'dist' in args.optimizer) + and (args.use_distributed_optimizer or "dist" in args.optimizer) and args.overlap_param_gather ) diff --git a/flagscale/train/megatron/training/utils.py b/flagscale/train/megatron/training/utils.py index 7c5d657300..0a9a01c25c 100644 --- a/flagscale/train/megatron/training/utils.py +++ b/flagscale/train/megatron/training/utils.py @@ -1,6 +1,7 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """General utilities.""" + import json import os import sys @@ -22,9 +23,9 @@ from apex.multi_tensor_apply import multi_tensor_applier except ImportError: warnings.warn( - f'Transformer Engine and Apex are not installed. ' - 'Falling back to local implementations of ' - 'multi_tensor_applier and multi_tensor_l2norm' + f"Transformer Engine and Apex are not installed. " + "Falling back to local implementations of " + "multi_tensor_applier and multi_tensor_l2norm" ) from megatron.core.utils import ( @@ -47,16 +48,18 @@ from megatron.plugin.utils import get_device_type_for_comm from megatron.plugin.platform import get_platform + cur_platform = get_platform() from megatron.plugin.decorators import overridable + def calc_params_l2_norm(model, force_create_fp32_copy=False): """Calculate l2 norm of parameters""" args = get_args() if not isinstance(model, list): model = [model] - if getattr(args, 'use_megatron_fsdp', False): + if getattr(args, "use_megatron_fsdp", False): # All Megatron FSDP parameters are expected to be PyTorch DTensor. # params_data is a dict of device_mesh -> list of local tensors. params = [] @@ -86,13 +89,13 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False): if not is_not_tp_duplicate: continue assert is_not_tp_duplicate - if not getattr(param, 'allreduce', True): + if not getattr(param, "allreduce", True): if not getattr(param, "is_engram_embedding", False): assert param_is_not_shared(param) param = to_local_if_dtensor(param) if args.bf16: - if not force_create_fp32_copy and hasattr(param, 'main_param'): - if getattr(param, 'main_param_sharded', False): + if not force_create_fp32_copy and hasattr(param, "main_param"): + if getattr(param, "main_param_sharded", False): if param.main_param is not None: sharded_params_data.append(param.main_param) else: @@ -112,8 +115,8 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False): if param_is_not_shared(param): param = to_local_if_dtensor(param) if args.bf16: - if not force_create_fp32_copy and hasattr(param, 'main_param'): - if getattr(param, 'main_param_sharded', False): + if not force_create_fp32_copy and hasattr(param, "main_param"): + if getattr(param, "main_param_sharded", False): if param.main_param is not None: sharded_params_data.append(param.main_param) else: @@ -129,7 +132,10 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False): dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device=cur_platform.device_name()) if len(params_data) > 0: norm, _ = multi_tensor_applier( - multi_tensor_l2norm, dummy_overflow_buf, [params_data], False # no per-parameter norm. + multi_tensor_l2norm, + dummy_overflow_buf, + [params_data], + False, # no per-parameter norm. ) norm_2 = norm * norm else: @@ -159,7 +165,7 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False): torch.distributed.all_reduce( sharded_norm_2, op=torch.distributed.ReduceOp.SUM, - group=mpu.get_data_parallel_group(with_context_parallel=True) + group=mpu.get_data_parallel_group(with_context_parallel=True), ) norm_2 += sharded_norm_2 @@ -178,7 +184,7 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False): # See details in https://gitlab-master.nvidia.com/ADLR/megatron-lm/-/issues/409 else: moe_norm_2 = torch.zeros_like(norm_2) - + # Add norm contribution from engram embedding. if len(engram_embedding_data) > 0: engram_embedding_norm, _ = multi_tensor_applier( @@ -201,18 +207,16 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False): original_norm_2 = norm_2.clone().detach() for mp_group in mp_groups: norm_2.copy_(original_norm_2) - torch.distributed.all_reduce( - norm_2, op=torch.distributed.ReduceOp.SUM, group=mp_group - ) + torch.distributed.all_reduce(norm_2, op=torch.distributed.ReduceOp.SUM, group=mp_group) if len(moe_params_data) > 0: emp_groups = mpu.get_expert_tensor_model_pipeline_parallel_group() comm_device = get_device_type_for_comm(emp_groups) if comm_device == "cpu": moe_norm_2 = moe_norm_2.cpu() - assert isinstance( - emp_groups, list - ), "emp_groups should be a list if mp_groups is a list" + assert isinstance(emp_groups, list), ( + "emp_groups should be a list if mp_groups is a list" + ) original_norm_2 = moe_norm_2.clone().detach() for emp_group in emp_groups: moe_norm_2.copy_(original_norm_2) @@ -223,14 +227,15 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False): assert len(engram_embedding_data) <= 0, "Engram embedding does not support hetero." ########## FlagScale End ########## else: # original code - # Reduce norm across model parallel groups (dense and expert). # Dense params should sum across all model-parallel GPUs (tensor + pipeline). dense_reduce_group = mpu.get_model_parallel_group() ranks_in_dense_reduce_group = torch.distributed.get_process_group_ranks(dense_reduce_group) # Expert params should sum across all model-parallel GPUs (expert + tensor + pipeline). expert_reduce_group = mpu.get_expert_tensor_model_pipeline_parallel_group() - ranks_in_expert_reduce_group = torch.distributed.get_process_group_ranks(expert_reduce_group) + ranks_in_expert_reduce_group = torch.distributed.get_process_group_ranks( + expert_reduce_group + ) # Engram params should sum across engram-embed-parallel GPUs.(Engram embedding and pipeline, for which has no engram module, the engram_module_initialized is False, and the param_norm is set to 0.) engram_mp_group = mpu.get_engram_model_parallel_group() @@ -249,7 +254,7 @@ def calc_params_l2_norm(model, force_create_fp32_copy=False): moe_norm_2, op=torch.distributed.ReduceOp.SUM, group=expert_reduce_group ) norm_2 += moe_norm_2 - # Reduce and add engram embedding norm if the group exists. + # Reduce and add engram embedding norm if the group exists. # Because engram_mp_group is different with other two groups in most cases, in order to reduce the impact of original code, allreduce and add independently here even if the engram embedding parallel group is the same as dense or expert group. if engram_mp_group is not None: torch.distributed.all_reduce( @@ -279,7 +284,10 @@ def calc_dtensor_params_l2_norm(params): norm = torch.zeros((1,), dtype=torch.float32, device=cur_platform.device_name()) else: norm, _ = multi_tensor_applier( - multi_tensor_l2norm, dummy_overflow_buf, [local_tensors], False # no per-parameter norm. + multi_tensor_l2norm, + dummy_overflow_buf, + [local_tensors], + False, # no per-parameter norm. ) norm_2 = norm * norm for pg, placement in zip( @@ -287,16 +295,12 @@ def calc_dtensor_params_l2_norm(params): dtensor_spec.placements, ): if placement.is_shard(): - torch.distributed.all_reduce( - norm_2, op=torch.distributed.ReduceOp.SUM, group=pg - ) + torch.distributed.all_reduce(norm_2, op=torch.distributed.ReduceOp.SUM, group=pg) elif placement.is_replicate(): # Replicated parameters are already summed across all ranks. pass else: - raise RuntimeError( - f"Unsupported placement {placement} for Megatron FSDP." - ) + raise RuntimeError(f"Unsupported placement {placement} for Megatron FSDP.") total_norm_2 += norm_2 return total_norm_2.item() ** 0.5 @@ -323,12 +327,16 @@ def reduce_max_stat_across_model_parallel_group(stat: float) -> float | None: stat = -1.0 model_parallel_groups = mpu.get_model_parallel_group() if not isinstance(model_parallel_groups, list): - stat = torch.tensor([stat], dtype=torch.float32, device=get_device_type_for_comm(model_parallel_groups)) + stat = torch.tensor( + [stat], dtype=torch.float32, device=get_device_type_for_comm(model_parallel_groups) + ) torch.distributed.all_reduce( stat, op=torch.distributed.ReduceOp.MAX, group=mpu.get_model_parallel_group() ) else: - stat = torch.tensor([stat], dtype=torch.float32, device=get_device_type_for_comm(model_parallel_groups[0])) + stat = torch.tensor( + [stat], dtype=torch.float32, device=get_device_type_for_comm(model_parallel_groups[0]) + ) for model_parallel_group in model_parallel_groups: torch.distributed.all_reduce( stat, op=torch.distributed.ReduceOp.MAX, group=model_parallel_group @@ -350,12 +358,16 @@ def logical_and_across_model_parallel_group(input: bool) -> bool: input = 0 model_parallel_groups = mpu.get_model_parallel_group() if not isinstance(model_parallel_groups, list): - input = torch.tensor([input], dtype=torch.int, device=get_device_type_for_comm(model_parallel_groups)) + input = torch.tensor( + [input], dtype=torch.int, device=get_device_type_for_comm(model_parallel_groups) + ) torch.distributed.all_reduce( input, op=torch.distributed.ReduceOp.MIN, group=mpu.get_model_parallel_group() ) else: - input = torch.tensor([input], dtype=torch.int, device=get_device_type_for_comm(model_parallel_groups[0])) + input = torch.tensor( + [input], dtype=torch.int, device=get_device_type_for_comm(model_parallel_groups[0]) + ) for model_parallel_group in model_parallel_groups: torch.distributed.all_reduce( input, op=torch.distributed.ReduceOp.MIN, group=model_parallel_group @@ -367,13 +379,15 @@ def report_memory(name): """Simple GPU memory report.""" args = get_args() mega_bytes = 1024.0 * 1024.0 - string = name + ' memory (MB)' - string += ' | allocated: {}'.format(cur_platform.memory_allocated() / mega_bytes) - string += ' | max allocated: {}'.format(cur_platform.max_memory_allocated() / mega_bytes) - string += ' | reserved: {}'.format(cur_platform.memory_reserved() / mega_bytes) - string += ' | max reserved: {}'.format(cur_platform.max_memory_reserved() / mega_bytes) + string = name + " memory (MB)" + string += " | allocated: {}".format(cur_platform.memory_allocated() / mega_bytes) + string += " | max allocated: {}".format(cur_platform.max_memory_allocated() / mega_bytes) + string += " | reserved: {}".format(cur_platform.memory_reserved() / mega_bytes) + string += " | max reserved: {}".format(cur_platform.max_memory_reserved() / mega_bytes) if args.log_device_memory_used: - string += ' | total device memory used: {}'.format(cur_platform.device_memory_used() / mega_bytes) + string += " | total device memory used: {}".format( + cur_platform.device_memory_used() / mega_bytes + ) if mpu.get_data_parallel_rank() == 0: print("[Rank {}] {}".format(torch.distributed.get_rank(), string), flush=True) @@ -382,18 +396,18 @@ def print_params_min_max_norm(optimizer, iteration): """Print min, max, and norm of all parameters.""" index = 0 rank = torch.distributed.get_rank() - string = 'iteration, rank, index, tensor-model-parallel, min, max, norm\n' + string = "iteration, rank, index, tensor-model-parallel, min, max, norm\n" optimizer_ = optimizer.optimizer for param_group in optimizer_.param_groups: - for param in param_group['params']: + for param in param_group["params"]: index += 1 min_ = param.data.min() max_ = param.data.max() norm = torch.linalg.norm(param.data) - string += '{:7d}, {:4d}, {:4d}, {:2d}, '.format( + string += "{:7d}, {:4d}, {:4d}, {:2d}, ".format( iteration, rank, index, int(param.tensor_model_parallel) ) - string += '{:.6E}, {:.6E}, {:.6E}\n'.format(min_, max_, norm) + string += "{:.6E}, {:.6E}, {:.6E}\n".format(min_, max_, norm) print(string, flush=True) @@ -415,13 +429,15 @@ def check_adlr_autoresume_termination(iteration, model, optimizer, opt_param_sch sys.exit(0) -def get_ltor_masks_and_position_ids(data, - eod_token, - pad_token, - reset_position_ids, - reset_attention_mask, - eod_mask_loss, - pad_mask_loss): +def get_ltor_masks_and_position_ids( + data, + eod_token, + pad_token, + reset_position_ids, + reset_attention_mask, + eod_mask_loss, + pad_mask_loss, +): """Build masks and position id for left to right model.""" # Extract batch size and sequence length. @@ -453,9 +469,10 @@ def get_ltor_masks_and_position_ids(data, if reset_position_ids or reset_attention_mask: # Loop through the batches: for b in range(micro_batch_size): - # Find indecies where EOD token is. - eod_index = position_ids[b, data[b] == eod_token] & position_ids[b, data[b] == pad_token] + eod_index = ( + position_ids[b, data[b] == eod_token] & position_ids[b, data[b] == pad_token] + ) # Detach indecies from positions if going to modify positions. if reset_position_ids: eod_index = eod_index.clone() @@ -511,13 +528,12 @@ def is_last_rank(): else: return torch.distributed.get_rank() == mpu.get_last_rank_when_using_pipeline() else: - return torch.distributed.get_rank() == ( - torch.distributed.get_world_size() - 1) + return torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1) def print_rank_last(message): """If distributed is initialized, print only on last rank.""" - if torch.distributed.is_initialized() and torch.distributed.get_backend() != 'fake': + if torch.distributed.is_initialized() and torch.distributed.get_backend() != "fake": if is_last_rank(): print(message, flush=True) else: @@ -538,9 +554,16 @@ def is_first_or_last_pipeline_stage(vp_stage, dualpipev_stage=None): ignore_virtual = False if dualpipev_stage is not None: ignore_dualpipev = False - return ( - mpu.is_pipeline_first_stage(ignore_virtual=ignore_virtual, vp_stage=vp_stage, ignore_dualpipev=ignore_dualpipev, dualpipev_stage=dualpipev_stage) - or mpu.is_pipeline_last_stage(ignore_virtual=ignore_virtual, vp_stage=vp_stage, ignore_dualpipev=ignore_dualpipev, dualpipev_stage=dualpipev_stage) + return mpu.is_pipeline_first_stage( + ignore_virtual=ignore_virtual, + vp_stage=vp_stage, + ignore_dualpipev=ignore_dualpipev, + dualpipev_stage=dualpipev_stage, + ) or mpu.is_pipeline_last_stage( + ignore_virtual=ignore_virtual, + vp_stage=vp_stage, + ignore_dualpipev=ignore_dualpipev, + dualpipev_stage=dualpipev_stage, ) @@ -559,8 +582,8 @@ def append_to_progress_log(string, barrier=True): if barrier: torch.distributed.barrier() if torch.distributed.get_rank() == 0: - with open_file(progress_log_filename, 'a') as f: - job_id = os.getenv('SLURM_JOB_ID', '') + with open_file(progress_log_filename, "a") as f: + job_id = os.getenv("SLURM_JOB_ID", "") num_gpus = args.world_size f.write( f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\tJob ID: {job_id}\t" @@ -584,14 +607,14 @@ def get_blend_and_blend_per_split(args): if use_data_path: if args.data_args_path is not None: assert args.data_path is None - with open_file(args.data_args_path, 'r') as f: + with open_file(args.data_args_path, "r") as f: blend = get_blend_from_list(f.read().split()) else: assert args.data_path is not None blend = get_blend_from_list(args.data_path) elif use_per_split_data_path: if args.per_split_data_args_path is not None: - with open_file(args.per_split_data_args_path, 'r') as f: + with open_file(args.per_split_data_args_path, "r") as f: per_split_data_args = json.load(f) # Each element in blend_per_split should be a list of files (and optional # weights), so split string if needed. @@ -617,7 +640,6 @@ def get_blend_and_blend_per_split(args): def get_batch_on_this_tp_rank(data_iterator, mtp_on_this_rank: bool = False): - args = get_args() def _broadcast(item): @@ -629,30 +651,31 @@ def _broadcast(item): ) if mpu.get_tensor_model_parallel_rank() == 0: - assert data_iterator is not None data = next(data_iterator) batch = { - 'tokens': data["tokens"].to(device=cur_platform.device(), non_blocking=True), - 'labels': data["labels"].to(device=cur_platform.device(), non_blocking=True), - 'loss_mask': data["loss_mask"].to(device=cur_platform.device(), non_blocking=True), - 'attention_mask': ( + "tokens": data["tokens"].to(device=cur_platform.device(), non_blocking=True), + "labels": data["labels"].to(device=cur_platform.device(), non_blocking=True), + "loss_mask": data["loss_mask"].to(device=cur_platform.device(), non_blocking=True), + "attention_mask": ( None if "attention_mask" not in data else data["attention_mask"].to(device=cur_platform.device(), non_blocking=True) ), - 'position_ids': data["position_ids"].to(device=cur_platform.device(), non_blocking=True), - 'cu_seqlens': ( + "position_ids": data["position_ids"].to( + device=cur_platform.device(), non_blocking=True + ), + "cu_seqlens": ( None if "cu_seqlens" not in data else data["cu_seqlens"].to(device=cur_platform.device(), non_blocking=True) ), - 'max_seqlen': ( + "max_seqlen": ( None if "max_seqlen" not in data else data["max_seqlen"].to(device=cur_platform.device(), non_blocking=True) ), - 'local_cp_size': ( + "local_cp_size": ( None if "local_cp_size" not in data else data["local_cp_size"].to(device=cur_platform.device(), non_blocking=True) @@ -675,48 +698,49 @@ def _broadcast_cu_seqlens(cu_seqlens): _broadcast(buf) if args.hybrid_context_parallel: - seq_len = torch.tensor(batch['tokens'].shape[0], dtype=torch.int32, device=cur_platform.device()) + seq_len = torch.tensor( + batch["tokens"].shape[0], dtype=torch.int32, device=cur_platform.device() + ) _broadcast(seq_len) - + if args.pipeline_model_parallel_size == 1 or mtp_on_this_rank: - _broadcast(batch['tokens']) - _broadcast(batch['labels']) - _broadcast(batch['loss_mask']) - _broadcast(batch['attention_mask']) - _broadcast(batch['position_ids']) - _broadcast_cu_seqlens(batch['cu_seqlens']) - _broadcast(batch['max_seqlen']) - _broadcast(batch['local_cp_size']) + _broadcast(batch["tokens"]) + _broadcast(batch["labels"]) + _broadcast(batch["loss_mask"]) + _broadcast(batch["attention_mask"]) + _broadcast(batch["position_ids"]) + _broadcast_cu_seqlens(batch["cu_seqlens"]) + _broadcast(batch["max_seqlen"]) + _broadcast(batch["local_cp_size"]) elif mpu.is_pipeline_first_stage(): - _broadcast(batch['tokens']) - _broadcast(batch['attention_mask']) - _broadcast(batch['position_ids']) - _broadcast_cu_seqlens(batch['cu_seqlens']) - _broadcast(batch['max_seqlen']) + _broadcast(batch["tokens"]) + _broadcast(batch["attention_mask"]) + _broadcast(batch["position_ids"]) + _broadcast_cu_seqlens(batch["cu_seqlens"]) + _broadcast(batch["max_seqlen"]) ######### FlagScale Begin ######## if mpu.get_dualpipev_pipeline_model_parallel_world_size() is not None: - _broadcast(batch['loss_mask']) - _broadcast(batch['labels']) + _broadcast(batch["loss_mask"]) + _broadcast(batch["labels"]) ######### FlagScale End ######## elif mpu.is_pipeline_last_stage(): # Multi-Token Prediction (MTP) layers need tokens and position_ids to calculate embedding. # Currently the Multi-Token Prediction (MTP) layers is fixed on the last stage, so we need # to broadcast tokens and position_ids to all of the tensor parallel ranks on the last stage. - _broadcast(batch['labels']) - _broadcast(batch['loss_mask']) - _broadcast(batch['attention_mask']) - + _broadcast(batch["labels"]) + _broadcast(batch["loss_mask"]) + _broadcast(batch["attention_mask"]) else: if args.hybrid_context_parallel: seq_len = torch.tensor(0, dtype=torch.int32, device=cur_platform.device()) _broadcast(seq_len) - shape = (seq_len.item()) + shape = seq_len.item() else: shape = (args.micro_batch_size, args.seq_length) - + tokens = torch.empty( shape, dtype=torch.int64, @@ -733,7 +757,11 @@ def _broadcast_cu_seqlens(cu_seqlens): device=cur_platform.device(), ) if args.create_attention_mask_in_dataloader: - shape_attention_mask = (args.micro_batch_size, 1, args.seq_length, args.seq_length) if not args.hybrid_context_parallel else (1, 1, shape[0], shape[0]) + shape_attention_mask = ( + (args.micro_batch_size, 1, args.seq_length, args.seq_length) + if not args.hybrid_context_parallel + else (1, 1, shape[0], shape[0]) + ) attention_mask = torch.empty( shape_attention_mask, dtype=torch.bool, @@ -756,12 +784,15 @@ def _broadcast_cu_seqlens(cu_seqlens): else: max_seqlen = None - - local_cp_size = torch.empty( - 1, - dtype=torch.int32, - device=cur_platform.device(), - ) if args.hybrid_context_parallel else None + local_cp_size = ( + torch.empty( + 1, + dtype=torch.int32, + device=cur_platform.device(), + ) + if args.hybrid_context_parallel + else None + ) def _broadcast_cu_seqlens(): dev = cur_platform.device() @@ -815,16 +846,15 @@ def _broadcast_cu_seqlens(): _broadcast(loss_mask) _broadcast(attention_mask) - batch = { - 'tokens': tokens, - 'labels': labels, - 'loss_mask': loss_mask, - 'attention_mask': attention_mask, - 'position_ids': position_ids, - 'cu_seqlens': cu_seqlens, - 'max_seqlen': max_seqlen, - 'local_cp_size': local_cp_size, + "tokens": tokens, + "labels": labels, + "loss_mask": loss_mask, + "attention_mask": attention_mask, + "position_ids": position_ids, + "cu_seqlens": cu_seqlens, + "max_seqlen": max_seqlen, + "local_cp_size": local_cp_size, } return batch @@ -842,7 +872,7 @@ def to_empty_if_meta_device(module: torch.nn.Module, *, device: torch.device, re accidently overwrite buffers with precomputed values during construction. Given the goal is to only materialize those tensors on meta device, this function checks the device first and only move the tensor to the destination if it is not on meta device. - + Args: module: The target module to apply this transformation. device: The desired device of the parameters @@ -857,9 +887,7 @@ def _empty_like_if_meta(tensor: torch.Tensor, *, device: torch.device): else: return tensor.to(device) - return module._apply( - lambda t: _empty_like_if_meta(t, device=device), recurse=recurse - ) + return module._apply(lambda t: _empty_like_if_meta(t, device=device), recurse=recurse) def get_nvtx_range(): @@ -891,18 +919,19 @@ def nvtx_range(msg, time=False, log_level=1): return nvtx_range except: + @contextmanager def dummy_range(msg, time=False, log_level=1): yield - return dummy_range + return dummy_range def has_nvrx_installed(): """Checks if nvidia-resiliency-ext is installed.""" try: import nvidia_resiliency_ext + return True except (ImportError, ModuleNotFoundError): return False - diff --git a/flagscale/train/megatron/training/wandb_utils.py b/flagscale/train/megatron/training/wandb_utils.py index 055e8ca1b6..02c904cb44 100644 --- a/flagscale/train/megatron/training/wandb_utils.py +++ b/flagscale/train/megatron/training/wandb_utils.py @@ -16,7 +16,9 @@ def _get_artifact_name_and_version(save_dir: Path, checkpoint_path: Path) -> Tup return save_dir.stem, checkpoint_path.stem -def on_save_checkpoint_success(checkpoint_path: str, tracker_filename: str, save_dir: str, iteration: int) -> None: +def on_save_checkpoint_success( + checkpoint_path: str, tracker_filename: str, save_dir: str, iteration: int +) -> None: """Function to be called after checkpointing succeeds and checkpoint is persisted for logging it as an artifact in W&B Args: @@ -30,7 +32,9 @@ def on_save_checkpoint_success(checkpoint_path: str, tracker_filename: str, save if wandb_writer: metadata = {"iteration": iteration} - artifact_name, artifact_version = _get_artifact_name_and_version(Path(save_dir), Path(checkpoint_path)) + artifact_name, artifact_version = _get_artifact_name_and_version( + Path(save_dir), Path(checkpoint_path) + ) artifact = wandb_writer.Artifact(artifact_name, type="model", metadata=metadata) # wandb's artifact.add_reference requires absolute paths checkpoint_path = str(Path(checkpoint_path).resolve()) @@ -51,10 +55,12 @@ def on_load_checkpoint_success(checkpoint_path: str, load_dir: str) -> None: """ wandb_writer = get_wandb_writer() - + if wandb_writer: try: - artifact_name, artifact_version = _get_artifact_name_and_version(Path(load_dir), Path(checkpoint_path)) + artifact_name, artifact_version = _get_artifact_name_and_version( + Path(load_dir), Path(checkpoint_path) + ) wandb_tracker_filename = _get_wandb_artifact_tracker_filename(load_dir) artifact_path = "" if wandb_tracker_filename.is_file(): diff --git a/flagscale/train/megatron/training/yaml_arguments.py b/flagscale/train/megatron/training/yaml_arguments.py index d73e5b3509..612e071cf9 100644 --- a/flagscale/train/megatron/training/yaml_arguments.py +++ b/flagscale/train/megatron/training/yaml_arguments.py @@ -22,27 +22,32 @@ # Taken from https://stackoverflow.com/questions/65414773/parse-environment-variable-from-yaml-with-pyyaml # Allows for yaml to use environment variables env_pattern = re.compile(r".*?\${(.*?)}.*?") + + def env_constructor(loader, node): value = loader.construct_scalar(node) for group in env_pattern.findall(value): assert os.environ.get(group) is not None, f"environment variable {group} in yaml not found" value = value.replace(f"${{{group}}}", os.environ.get(group)) return value + + yaml.add_implicit_resolver("!pathex", env_pattern) yaml.add_constructor("!pathex", env_constructor) str_dtype_to_torch = { - "float32" : torch.float32, - "float16" : torch.float16, - "bfloat16" : torch.bfloat16 + "float32": torch.float32, + "float16": torch.float16, + "bfloat16": torch.bfloat16, } from megatron.plugin.platform import get_platform + cur_platform = get_platform() + def validate_yaml(args, defaults={}): - # This is for legacy script env var setting if type(args.data_path) is str: # If no white space its a single path @@ -52,42 +57,63 @@ def validate_yaml(args, defaults={}): # Tensor model parallel size. args.model_parallel.tensor_model_parallel_size = min( - args.model_parallel.tensor_model_parallel_size, args.world_size) - assert args.world_size % args.model_parallel.tensor_model_parallel_size == 0, 'world size'\ - ' ({}) is not divisible by tensor model parallel size ({})'.format( - args.world_size, args.model_parallel.tensor_model_parallel_size) + args.model_parallel.tensor_model_parallel_size, args.world_size + ) + assert args.world_size % args.model_parallel.tensor_model_parallel_size == 0, ( + "world size ({}) is not divisible by tensor model parallel size ({})".format( + args.world_size, args.model_parallel.tensor_model_parallel_size + ) + ) # Pipeline model parallel size. args.model_parallel.pipeline_model_parallel_size = min( args.model_parallel.pipeline_model_parallel_size, - (args.world_size // args.model_parallel.tensor_model_parallel_size)) + (args.world_size // args.model_parallel.tensor_model_parallel_size), + ) args.model_parallel.transformer_pipeline_model_parallel_size = ( args.model_parallel.pipeline_model_parallel_size - 1 - if args.account_for_embedding_in_pipeline_split else - args.model_parallel.pipeline_model_parallel_size + if args.account_for_embedding_in_pipeline_split + else args.model_parallel.pipeline_model_parallel_size ) # Checks. - model_parallel_size = args.model_parallel.pipeline_model_parallel_size * \ - args.model_parallel.tensor_model_parallel_size - assert args.world_size % (model_parallel_size * args.model_parallel.context_parallel_size) == 0, \ - 'world size ({}) is not divisible by tensor parallel size ({}) times ' \ - 'pipeline parallel size ({}) times context parallel size ({})'.format( - args.world_size, args.model_parallel.tensor_model_parallel_size, - args.model_parallel.pipeline_model_parallel_size, args.model_parallel.context_parallel_size) - + model_parallel_size = ( + args.model_parallel.pipeline_model_parallel_size + * args.model_parallel.tensor_model_parallel_size + ) + assert ( + args.world_size % (model_parallel_size * args.model_parallel.context_parallel_size) == 0 + ), ( + "world size ({}) is not divisible by tensor parallel size ({}) times " + "pipeline parallel size ({}) times context parallel size ({})".format( + args.world_size, + args.model_parallel.tensor_model_parallel_size, + args.model_parallel.pipeline_model_parallel_size, + args.model_parallel.context_parallel_size, + ) + ) + # data_parallel_size is not in model parallel config - args.data_parallel_size = args.world_size // (model_parallel_size * args.model_parallel.context_parallel_size) + args.data_parallel_size = args.world_size // ( + model_parallel_size * args.model_parallel.context_parallel_size + ) if args.rank == 0: - print('using world size: {}, data-parallel size: {}, ' - 'context-parallel size: {}, ' - 'tensor-model-parallel size: {}, ' - 'pipeline-model-parallel size: {}'.format( - args.world_size, args.data_parallel_size, - args.model_parallel.context_parallel_size, - args.model_parallel.tensor_model_parallel_size, - args.model_parallel.pipeline_model_parallel_size), flush=True) + print( + "using world size: {}, data-parallel size: {}, " + "context-parallel size: {}, " + "tensor-model-parallel size: {}, " + "pipeline-model-parallel size: {}".format( + args.world_size, + args.data_parallel_size, + args.model_parallel.context_parallel_size, + args.model_parallel.tensor_model_parallel_size, + args.model_parallel.pipeline_model_parallel_size, + ), + flush=True, + ) if args.model_parallel.tp_comm_overlap: - assert args.model_parallel.sequence_parallel == True, 'Tensor parallel communication/GEMM overlap can happen only when sequence parallelism is enabled' + assert args.model_parallel.sequence_parallel == True, ( + "Tensor parallel communication/GEMM overlap can happen only when sequence parallelism is enabled" + ) # Set input defaults. for key in defaults: @@ -96,10 +122,11 @@ def validate_yaml(args, defaults={}): # ensuring the arg is set to None. if getattr(args, key, None) is not None: if args.rank == 0: - print('WARNING: overriding default arguments for {key}:{v} \ - with {key}:{v2}'.format(key=key, v=defaults[key], - v2=getattr(args, key)), - flush=True) + print( + "WARNING: overriding default arguments for {key}:{v} \ + with {key}:{v2}".format(key=key, v=defaults[key], v2=getattr(args, key)), + flush=True, + ) else: setattr(args, key, defaults[key]) @@ -109,35 +136,46 @@ def validate_yaml(args, defaults={}): if args.global_batch_size is None: args.global_batch_size = args.micro_batch_size * args.data_parallel_size if args.rank == 0: - print('setting global batch size to {}'.format( - args.global_batch_size), flush=True) + print("setting global batch size to {}".format(args.global_batch_size), flush=True) assert args.global_batch_size > 0 # num_layers_per_virtual_pipeline_stage is not insde model parallel for checkpointing if args.num_layers_per_virtual_pipeline_stage is not None: - assert args.model_parallel.pipeline_model_parallel_size > 2, \ - 'pipeline-model-parallel size should be greater than 2 with ' \ - 'interleaved schedule' - assert args.language_model.num_layers % args.model_parallel.transformer_pipeline_model_parallel_size == 0, \ - 'number of layers should be divisible by the pipeline parallel size' - num_layers_per_pipeline_stage = args.language_model.num_layers // args.model_parallel.transformer_pipeline_model_parallel_size - assert num_layers_per_pipeline_stage % args.num_layers_per_virtual_pipeline_stage == 0, \ - 'number of layers per pipeline stage must be divisible number of layers per virtual pipeline stage' - args.model_parallel.virtual_pipeline_model_parallel_size = num_layers_per_pipeline_stage // \ - args.num_layers_per_virtual_pipeline_stage + assert args.model_parallel.pipeline_model_parallel_size > 2, ( + "pipeline-model-parallel size should be greater than 2 with interleaved schedule" + ) + assert ( + args.language_model.num_layers + % args.model_parallel.transformer_pipeline_model_parallel_size + == 0 + ), "number of layers should be divisible by the pipeline parallel size" + num_layers_per_pipeline_stage = ( + args.language_model.num_layers + // args.model_parallel.transformer_pipeline_model_parallel_size + ) + assert num_layers_per_pipeline_stage % args.num_layers_per_virtual_pipeline_stage == 0, ( + "number of layers per pipeline stage must be divisible number of layers per virtual pipeline stage" + ) + args.model_parallel.virtual_pipeline_model_parallel_size = ( + num_layers_per_pipeline_stage // args.num_layers_per_virtual_pipeline_stage + ) else: args.model_parallel.virtual_pipeline_model_parallel_size = None # Overlap P2P communication is disabled if not using the interleaved schedule. args.model_parallel.overlap_p2p_comm = False if args.rank == 0: - print('WARNING: Setting args.overlap_p2p_comm to False since non-interleaved ' - 'schedule does not support overlapping p2p communication') + print( + "WARNING: Setting args.overlap_p2p_comm to False since non-interleaved " + "schedule does not support overlapping p2p communication" + ) if args.overlap_param_gather: - assert args.use_distributed_optimizer, \ - '--overlap-param-gather only supported with distributed optimizer' - assert args.overlap_grad_reduce, \ - '--overlap-grad-reduce should be turned on when using --overlap-param-gather' + assert args.use_distributed_optimizer, ( + "--overlap-param-gather only supported with distributed optimizer" + ) + assert args.overlap_grad_reduce, ( + "--overlap-grad-reduce should be turned on when using --overlap-param-gather" + ) # Parameters dtype. if args.model_parallel.fp16: @@ -151,15 +189,16 @@ def validate_yaml(args, defaults={}): if not args.accumulate_allreduce_grads_in_fp32: args.accumulate_allreduce_grads_in_fp32 = True if args.rank == 0: - print('accumulate and all-reduce gradients in fp32 for ' - 'bfloat16 data type.', flush=True) + print( + "accumulate and all-reduce gradients in fp32 for bfloat16 data type.", + flush=True, + ) if args.rank == 0: - print('using {} for parameters ...'.format(args.model_parallel.params_dtype), - flush=True) + print("using {} for parameters ...".format(args.model_parallel.params_dtype), flush=True) if args.dataloader_type is None: - args.dataloader_type = 'single' + args.dataloader_type = "single" # Consumed tokens. args.consumed_train_samples = 0 @@ -176,46 +215,44 @@ def validate_yaml(args, defaults={}): if args.train_iters: # If we use iteration-based training, make sure the # sample-based options are off. - assert args.train_samples is None, \ - 'expected iteration-based training' - assert args.lr_decay_samples is None, \ - 'expected iteration-based learning rate decay' - assert args.lr_warmup_samples == 0, \ - 'expected iteration-based learning rate warmup' - assert args.rampup_batch_size is None, \ - 'expected no batch-size rampup for iteration-based training' + assert args.train_samples is None, "expected iteration-based training" + assert args.lr_decay_samples is None, "expected iteration-based learning rate decay" + assert args.lr_warmup_samples == 0, "expected iteration-based learning rate warmup" + assert args.rampup_batch_size is None, ( + "expected no batch-size rampup for iteration-based training" + ) if args.lr_warmup_fraction is not None: - assert args.lr_warmup_iters == 0, \ - 'can only specify one of lr-warmup-fraction and lr-warmup-iters' + assert args.lr_warmup_iters == 0, ( + "can only specify one of lr-warmup-fraction and lr-warmup-iters" + ) # Sample-based training. if args.train_samples: # If we use sample-based training, make sure the # iteration-based options are off. - assert args.train_iters is None, \ - 'expected sample-based training' - assert args.lr_decay_iters is None, \ - 'expected sample-based learning rate decay' - assert args.lr_warmup_iters == 0, \ - 'expected sample-based learnig rate warmup' + assert args.train_iters is None, "expected sample-based training" + assert args.lr_decay_iters is None, "expected sample-based learning rate decay" + assert args.lr_warmup_iters == 0, "expected sample-based learnig rate warmup" if args.lr_warmup_fraction is not None: - assert args.lr_warmup_samples == 0, \ - 'can only specify one of lr-warmup-fraction ' \ - 'and lr-warmup-samples' + assert args.lr_warmup_samples == 0, ( + "can only specify one of lr-warmup-fraction and lr-warmup-samples" + ) # How to handle this better if args.language_model.num_layers is not None: - assert args.encoder_num_layers is None, \ - 'cannot have both num-layers and encoder-num-layers specified' + assert args.encoder_num_layers is None, ( + "cannot have both num-layers and encoder-num-layers specified" + ) args.encoder_num_layers = args.language_model.num_layers else: - assert args.encoder_num_layers is not None, \ - 'either num-layers or encoder-num-layers should be specified' + assert args.encoder_num_layers is not None, ( + "either num-layers or encoder-num-layers should be specified" + ) args.language_model.num_layers = args.encoder_num_layers # Check required arguments. # removed max_position_embeddings from reqs - required_args = ['num_layers', 'hidden_size', 'num_attention_heads'] + required_args = ["num_layers", "hidden_size", "num_attention_heads"] for req_arg in required_args: _check_arg_is_not_none(args.language_model, req_arg) @@ -227,15 +264,19 @@ def validate_yaml(args, defaults={}): # the same ballpark as the counterpart with 4*h size # we keep it a multiple of 64, which means the actual tensor size # will be a multiple of 64 / tp_size - args.language_model.ffn_hidden_size = int((4 * args.language_model.hidden_size * 2 / 3) / 64) * 64 + args.language_model.ffn_hidden_size = ( + int((4 * args.language_model.hidden_size * 2 / 3) / 64) * 64 + ) else: args.language_model.ffn_hidden_size = 4 * args.language_model.hidden_size if args.language_model.kv_channels is None: assert args.language_model.hidden_size % args.language_model.num_attention_heads == 0 - args.language_model.kv_channels = args.language_model.hidden_size // args.language_model.num_attention_heads + args.language_model.kv_channels = ( + args.language_model.hidden_size // args.language_model.num_attention_heads + ) - #TODO: Implement arguments for encoder-decoder + # TODO: Implement arguments for encoder-decoder if args.seq_length is not None: assert args.encoder_seq_length is None args.encoder_seq_length = args.seq_length @@ -253,17 +294,18 @@ def validate_yaml(args, defaults={}): assert args.save_interval is not None # Mixed precision checks. if args.fp16_lm_cross_entropy: - assert args.fp16, 'lm cross entropy in fp16 only support in fp16 mode.' + assert args.fp16, "lm cross entropy in fp16 only support in fp16 mode." if args.language_model.fp32_residual_connection: - assert args.model_parallel.fp16 or args.model_parallel.bf16, \ - 'residual connection in fp32 only supported when using fp16 or bf16.' + assert args.model_parallel.fp16 or args.model_parallel.bf16, ( + "residual connection in fp32 only supported when using fp16 or bf16." + ) if args.language_model.moe_grouped_gemm: - assert args.model_parallel.bf16, 'Currently GroupedGEMM for MoE only supports bf16 dtype.' + assert args.model_parallel.bf16, "Currently GroupedGEMM for MoE only supports bf16 dtype." dc = cur_platform.get_device_capability() assert dc[0] >= 8, "Unsupported compute capability for GroupedGEMM kernels." - if args.weight_decay_incr_style == 'constant': + if args.weight_decay_incr_style == "constant": assert args.start_weight_decay is None assert args.end_weight_decay is None args.start_weight_decay = args.weight_decay @@ -276,30 +318,33 @@ def validate_yaml(args, defaults={}): if not is_torch_min_version("1.11.0a0"): args.language_model.persist_layer_norm = False if args.rank == 0: - print('Persistent fused layer norm kernel is supported from ' - 'pytorch v1.11 (nvidia pytorch container paired with v1.11). ' - 'Defaulting to no_persist_layer_norm=True') + print( + "Persistent fused layer norm kernel is supported from " + "pytorch v1.11 (nvidia pytorch container paired with v1.11). " + "Defaulting to no_persist_layer_norm=True" + ) # Activation recomputing. if args.language_model.distribute_saved_activations: - assert args.model_parallel.tensor_model_parallel_size > 1, 'can distribute ' \ - 'recomputed activations only across tensor model ' \ - 'parallel groups' - assert args.language_model.recompute_granularity == 'full', \ - 'distributed recompute activations is only '\ - 'application to full recompute granularity' - assert args.language_model.recompute_method is not None, \ - 'for distributed recompute activations to work you '\ - 'need to use a recompute method ' - assert is_torch_min_version("1.10.0a0"), \ - 'distributed recompute activations are supported for pytorch ' \ - 'v1.10 and above (Nvidia Pytorch container >= 21.07). Current ' \ - f'pytorch version is v{get_torch_version()}.' - - if args.language_model.recompute_granularity == 'selective': - assert args.language_model.recompute_method is None, \ - 'recompute method is not yet supported for ' \ - 'selective recomputing granularity' + assert args.model_parallel.tensor_model_parallel_size > 1, ( + "can distribute recomputed activations only across tensor model parallel groups" + ) + assert args.language_model.recompute_granularity == "full", ( + "distributed recompute activations is only application to full recompute granularity" + ) + assert args.language_model.recompute_method is not None, ( + "for distributed recompute activations to work you need to use a recompute method " + ) + assert is_torch_min_version("1.10.0a0"), ( + "distributed recompute activations are supported for pytorch " + "v1.10 and above (Nvidia Pytorch container >= 21.07). Current " + f"pytorch version is v{get_torch_version()}." + ) + + if args.language_model.recompute_granularity == "selective": + assert args.language_model.recompute_method is None, ( + "recompute method is not yet supported for selective recomputing granularity" + ) # disable sequence parallelism when tp=1 # to avoid change in numerics when @@ -307,31 +352,38 @@ def validate_yaml(args, defaults={}): if args.model_parallel.tensor_model_parallel_size == 1: args.model_parallel.sequence_parallel = False - if os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') != "1": + if os.environ.get("CUDA_DEVICE_MAX_CONNECTIONS") != "1": if args.model_parallel.sequence_parallel: raise RuntimeError( "Using sequence parallelism requires setting the environment variable " - "CUDA_DEVICE_MAX_CONNECTIONS to 1") - + "CUDA_DEVICE_MAX_CONNECTIONS to 1" + ) + # MoE Spec check if args.language_model.num_moe_experts is not None: assert args.spec is None, "Model Spec must be None when using MoEs" if args.model_parallel.tensor_model_parallel_size > 1: - assert args.model_parallel.sequence_parallel, \ + assert args.model_parallel.sequence_parallel, ( "When using MoE and tensor parallelism, sequence parallelism must be used." + ) # Expert parallelism check - if args.model_parallel.expert_model_parallel_size > 1: - assert args.language_model.num_moe_experts is not None, "num_experts must be non None to use expert model parallelism" - assert args.language_model.num_moe_experts % args.model_parallel.expert_model_parallel_size == 0, \ - "Number of experts should be a multiple of expert model parallel_size." - assert not args.model_parallel.fp16, \ + if args.model_parallel.expert_model_parallel_size > 1: + assert args.language_model.num_moe_experts is not None, ( + "num_experts must be non None to use expert model parallelism" + ) + assert ( + args.language_model.num_moe_experts % args.model_parallel.expert_model_parallel_size + == 0 + ), "Number of experts should be a multiple of expert model parallel_size." + assert not args.model_parallel.fp16, ( "Expert parallelism is not supported with fp16 training." + ) # Print arguments. _print_args("arguments", args) - #TODO: Added as much of the global initialization requires the model parallel arguments + # TODO: Added as much of the global initialization requires the model parallel arguments args = SimpleNamespace(**args.__dict__, **args.model_parallel.__dict__) args = SimpleNamespace(**args.__dict__, **args.language_model.__dict__) # For GPT Layer spec in pretrain_gpt @@ -339,19 +391,19 @@ def validate_yaml(args, defaults={}): return args + def _print_args(title, args): """Print arguments.""" if args.rank == 0: - print(f'------------------------ {title} ------------------------', - flush=True) + print(f"------------------------ {title} ------------------------", flush=True) str_list = [] for arg in vars(args): - dots = '.' * (48 - len(arg)) - str_list.append(' {} {} {}'.format(arg, dots, getattr(args, arg))) + dots = "." * (48 - len(arg)) + str_list.append(" {} {} {}".format(arg, dots, getattr(args, arg))) for arg in sorted(str_list, key=lambda x: x.lower()): print(arg, flush=True) - print(f'-------------------- end of {title} ---------------------', - flush=True) + print(f"-------------------- end of {title} ---------------------", flush=True) + def core_config_from_args(args, dataclass=TransformerConfig): """Builds core config object from namespace args from given dataclass @@ -359,7 +411,7 @@ def core_config_from_args(args, dataclass=TransformerConfig): Raises exception if argument missing in args Args: - args(SimpleNamespace, optional): Namespace to pull argument values from + args(SimpleNamespace, optional): Namespace to pull argument values from dataclass (dataclass, optional): Core dataclass config to pull argument names from @@ -374,55 +426,63 @@ def core_config_from_args(args, dataclass=TransformerConfig): raise Exception(f"Missing argument {f.name} for {str(dataclass)} config") return kw_args + def _check_arg_is_not_none(args, arg): - assert getattr(args, arg) is not None, '{} argument is None'.format(arg) + assert getattr(args, arg) is not None, "{} argument is None".format(arg) -def core_transformer_config_from_yaml(args, transfomer_key = "language_model"): + +def core_transformer_config_from_yaml(args, transfomer_key="language_model"): # Combine transfomer config with model parallel args args = SimpleNamespace(**vars(getattr(args, transfomer_key)), **vars(args.model_parallel)) # Translate args to core transformer configuration - kw_args = core_config_from_args(args, TransformerConfig) - - # Hardcoded - kw_args['deallocate_pipeline_outputs'] = True - kw_args['pipeline_dtype'] = kw_args['params_dtype'] - kw_args['batch_p2p_comm'] = not args.overlap_p2p_comm - - assert args.activation_func in ["swiglu","squaredrelu","gelu"], f"{args.activation_func} is not a supported activation function" + kw_args = core_config_from_args(args, TransformerConfig) + + # Hardcoded + kw_args["deallocate_pipeline_outputs"] = True + kw_args["pipeline_dtype"] = kw_args["params_dtype"] + kw_args["batch_p2p_comm"] = not args.overlap_p2p_comm + + assert args.activation_func in ["swiglu", "squaredrelu", "gelu"], ( + f"{args.activation_func} is not a supported activation function" + ) if args.activation_func == "swiglu": - kw_args['activation_func'] = F.silu - kw_args['gated_linear_unit'] = True - kw_args['bias_activation_fusion'] = args.bias_swiglu_fusion + kw_args["activation_func"] = F.silu + kw_args["gated_linear_unit"] = True + kw_args["bias_activation_fusion"] = args.bias_swiglu_fusion elif args.activation_func == "squaredrelu": + def squared_relu(x): return torch.pow(F.relu(x), 2) - kw_args['activation_func'] = squared_relu + + kw_args["activation_func"] = squared_relu elif args.activation_func == "gelu": - kw_args['activation_func'] = F.gelu + kw_args["activation_func"] = F.gelu if args.add_bias_linear: - kw_args['bias_activation_fusion'] = False + kw_args["bias_activation_fusion"] = False else: - kw_args['bias_activation_fusion'] = args.bias_activation_fusion - + kw_args["bias_activation_fusion"] = args.bias_activation_fusion + if args.init_method == "xavier_uniform": - kw_args['init_method'] = torch.nn.init.xavier_uniform_ - kw_args['scaled_init_method'] = torch.nn.init.xavier_uniform_ + kw_args["init_method"] = torch.nn.init.xavier_uniform_ + kw_args["scaled_init_method"] = torch.nn.init.xavier_uniform_ if args.embedding_init_method == "xavier_uniform": - kw_args['embedding_init_method'] = torch.nn.init.xavier_uniform_ - + kw_args["embedding_init_method"] = torch.nn.init.xavier_uniform_ + # Return Transformer config. if getattr(args, "multi_latent_attention", False): return MLATransformerConfig(**kw_args) else: return TransformerConfig(**kw_args) + def load_yaml(yaml_path): print(f"warning using experimental yaml arguments feature, argparse arguments will be ignored") with open(yaml_path, "r") as f: config = yaml.safe_load(f) # Convert to nested namespace - config_namespace = json.loads(json.dumps(config), object_hook=lambda item: SimpleNamespace(**item)) + config_namespace = json.loads( + json.dumps(config), object_hook=lambda item: SimpleNamespace(**item) + ) # Add config location to namespace config_namespace.yaml_cfg = yaml_path return config_namespace - diff --git a/flagscale/train/perf_monitor/__init__.py b/flagscale/train/perf_monitor/__init__.py index bc36983289..2317b0212f 100644 --- a/flagscale/train/perf_monitor/__init__.py +++ b/flagscale/train/perf_monitor/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """FlagScale performance monitor utilities.""" from .hooks import ( diff --git a/flagscale/train/perf_monitor/flops_calculator.py b/flagscale/train/perf_monitor/flops_calculator.py index 77eb3d90ac..9cfcd84066 100644 --- a/flagscale/train/perf_monitor/flops_calculator.py +++ b/flagscale/train/perf_monitor/flops_calculator.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """FLOPS estimation helpers for performance monitoring.""" from __future__ import annotations diff --git a/flagscale/train/perf_monitor/hooks.py b/flagscale/train/perf_monitor/hooks.py index 94ec5e72ba..f2e35f9b59 100644 --- a/flagscale/train/perf_monitor/hooks.py +++ b/flagscale/train/perf_monitor/hooks.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Training-loop hooks for the performance monitor.""" from __future__ import annotations diff --git a/flagscale/train/perf_monitor/perf_logger.py b/flagscale/train/perf_monitor/perf_logger.py index e46ac542a4..084ad9addb 100644 --- a/flagscale/train/perf_monitor/perf_logger.py +++ b/flagscale/train/perf_monitor/perf_logger.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """File logger for performance monitor output.""" from __future__ import annotations diff --git a/flagscale/train/perf_monitor/perf_metrics.py b/flagscale/train/perf_monitor/perf_metrics.py index ce4e6a4b45..f1528a4c19 100644 --- a/flagscale/train/perf_monitor/perf_metrics.py +++ b/flagscale/train/perf_monitor/perf_metrics.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Performance metrics collection for training.""" from __future__ import annotations diff --git a/flagscale/train/train_config.py b/flagscale/train/train_config.py index 7c093e70f9..4d6c1f6dc1 100644 --- a/flagscale/train/train_config.py +++ b/flagscale/train/train_config.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """ Training configuration models using Pydantic. """ @@ -191,7 +205,9 @@ class DataConfig(BaseModel): model_config = {"extra": "allow", "arbitrary_types_allowed": True} dataset_type: str = "lerobot" - data_path: str | None = Field(default=None, description="Path to training dataset (unused when data_mix is set)") + data_path: str | None = Field( + default=None, description="Path to training dataset (unused when data_mix is set)" + ) tolerance_s: float = 0.0001 use_imagenet_stats: bool = True rename_map: dict[str, str] | None = None @@ -233,7 +249,9 @@ class ModelConfig(BaseModel): # Required fields to identify which model and checkpoint to use model_name: str = Field(..., description="Model name: 'pi0' or 'pi0.5'") # None when the policy loads pretrained sub-components (e.g. VLM) internally - checkpoint_dir: str | None = Field(default=None, description="Path to pretrained model checkpoint") + checkpoint_dir: str | None = Field( + default=None, description="Path to pretrained model checkpoint" + ) freeze: FreezeConfig | None = None optimizer: OptimizerConfig = Field(default_factory=OptimizerConfig) raw: DictConfig | None = Field(default=None, exclude=True) @@ -262,9 +280,7 @@ def validate_model_name(cls, v): def get_model_config_dict(self) -> dict[str, Any]: """Get all model-specific config fields (excluding train-level fields).""" - return self.model_dump( - exclude={"model_name", "checkpoint_dir", "freeze", "optimizer"} - ) + return self.model_dump(exclude={"model_name", "checkpoint_dir", "freeze", "optimizer"}) class TrainConfig(BaseModel): @@ -311,11 +327,13 @@ def from_pretrained(cls, pretrained_path: str | Path) -> "TrainConfig": def to_omegaconf(self) -> DictConfig: """Reconstruct the full OmegaConf config from stored raw DictConfigs.""" - return OmegaConf.create({ - "system": self.system.raw, - "model": self.model.raw, - "data": self.data.raw, - }) + return OmegaConf.create( + { + "system": self.system.raw, + "model": self.model.raw, + "data": self.data.raw, + } + ) class Config: # Allow arbitrary types for complex objects diff --git a/flagscale/train/train_gr00t_n1_5.py b/flagscale/train/train_gr00t_n1_5.py index bf14259ed4..0598d6377b 100644 --- a/flagscale/train/train_gr00t_n1_5.py +++ b/flagscale/train/train_gr00t_n1_5.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Mainly adopted from # https://github.com/huggingface/lerobot/blob/2b304eeb841ae6c371e3dd341bbbb9dd254b07cb/src/lerobot/scripts/lerobot_train.py @@ -16,7 +30,11 @@ import torch.distributed as dist from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy from torch.distributed.device_mesh import init_device_mesh -from torch.distributed.checkpoint.state_dict import get_model_state_dict, get_optimizer_state_dict, StateDictOptions +from torch.distributed.checkpoint.state_dict import ( + get_model_state_dict, + get_optimizer_state_dict, + StateDictOptions, +) from torch.optim import Optimizer from flagscale.logger import logger @@ -53,16 +71,20 @@ from flash_attn import flash_attn_func, flash_attn_varlen_func from flash_attn.bert_padding import pad_input, unpad_input import transformers.modeling_flash_attention_utils as _fa_utils + _fa_utils._flash_fn = flash_attn_func _fa_utils._flash_varlen_fn = flash_attn_varlen_func _fa_utils._pad_fn = pad_input _fa_utils._unpad_fn = unpad_input + def _patched_lazy_imports(implementation=None): return flash_attn_func, flash_attn_varlen_func, pad_input, unpad_input + _fa_utils._lazy_imports = _patched_lazy_imports except ImportError: pass + def set_seed(seed: int): random.seed(seed) np.random.seed(seed) @@ -71,9 +93,10 @@ def set_seed(seed: int): if get_platform().name() == "cuda": torch.backends.cudnn.enabled = True torch.backends.cudnn.benchmark = False - torch.backends.cudnn.deterministic = False + torch.backends.cudnn.deterministic = False torch.backends.cuda.matmul.allow_tf32 = False + def apply_fsdp2(policy, device_mesh): """Apply FSDP2 sharding to Gr00tN15. @@ -191,7 +214,6 @@ def _format_meter_val(meter: AverageMeter) -> str: return " ".join(display_list) - def make_policy( config: PreTrainedConfig, ds_meta: LeRobotDatasetMetadata | None = None, @@ -200,11 +222,7 @@ def make_policy( # Use == instead of `is` for FeatureType.ACTION comparison # because flagscale.FeatureType and lerobot.FeatureType are different enum classes - output_features = { - key: ft - for key, ft in features.items() - if ft.type == FeatureType.ACTION - } + output_features = {key: ft for key, ft in features.items() if ft.type == FeatureType.ACTION} input_features = {key: ft for key, ft in features.items() if key not in output_features} config.output_features = output_features @@ -383,7 +401,9 @@ def update_policy( optimizer.zero_grad() autocast_context = ( - torch.amp.autocast(get_platform().amp_device_type(), dtype=torch.bfloat16) if use_amp else nullcontext() + torch.amp.autocast(get_platform().amp_device_type(), dtype=torch.bfloat16) + if use_amp + else nullcontext() ) with autocast_context: @@ -408,7 +428,9 @@ def update_policy( policy.update() train_metrics.loss = loss.item() - train_metrics.grad_norm = grad_norm.full_tensor().item() if hasattr(grad_norm, 'full_tensor') else grad_norm.item() + train_metrics.grad_norm = ( + grad_norm.full_tensor().item() if hasattr(grad_norm, "full_tensor") else grad_norm.item() + ) train_metrics.lr = optimizer.param_groups[0]["lr"] train_metrics.update_s = time.perf_counter() - start_time @@ -419,7 +441,7 @@ def main(config: TrainConfig, seed: int): set_seed(seed) policy_config = PreTrainedConfig.from_train_config(config) - + local_rank = int(os.environ["LOCAL_RANK"]) get_platform().set_device(local_rank) dist.init_process_group(backend=get_platform().dist_backend()) @@ -517,7 +539,10 @@ def main(config: TrainConfig, seed: int): resume_from = config.system.checkpoint.resume_from if resume_from: step, _dl_state = load_training_state_fsdp2( - Path(resume_from), policy, optimizer, lr_scheduler, + Path(resume_from), + policy, + optimizer, + lr_scheduler, ) saved_rng = serialize_rng_state() for _ in range(step): diff --git a/flagscale/train/train_pi.py b/flagscale/train/train_pi.py index cf77b59d3e..68cbd1532a 100644 --- a/flagscale/train/train_pi.py +++ b/flagscale/train/train_pi.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Mainly adopted from # https://github.com/huggingface/lerobot/blob/2b304eeb841ae6c371e3dd341bbbb9dd254b07cb/src/lerobot/scripts/lerobot_train.py @@ -74,14 +88,16 @@ def set_seed(seed: int): if get_platform().name() == "cuda": torch.backends.cudnn.enabled = True torch.backends.cudnn.benchmark = False - torch.backends.cudnn.deterministic = False + torch.backends.cudnn.deterministic = False torch.backends.cuda.matmul.allow_tf32 = False def init_distributed(): local_rank = int(os.environ.get("LOCAL_RANK", 0)) get_platform().set_device(local_rank) - torch.distributed.init_process_group(backend=get_platform().dist_backend(), init_method="env://") + torch.distributed.init_process_group( + backend=get_platform().dist_backend(), init_method="env://" + ) return local_rank @@ -130,9 +146,7 @@ def make_dataset(cfg: DataConfig, policy_config): # TODO: (yupu) Remove hard-coded video backend video_backend = "torchcodec" - image_transforms = ( - ImageTransforms(cfg.image_transforms) if enable_image_transform else None - ) + image_transforms = ImageTransforms(cfg.image_transforms) if enable_image_transform else None # Leave the revision to None ds_meta = LeRobotDatasetMetadata(root=cfg.data_path, revision=None) delta_timestamps = resolve_delta_timestamps(policy_config, ds_meta) @@ -150,16 +164,12 @@ def make_dataset(cfg: DataConfig, policy_config): if cfg.use_imagenet_stats: for key in dataset.meta.camera_keys: for stats_type, stats in IMAGENET_STATS.items(): - dataset.meta.stats[key][stats_type] = torch.tensor( - stats, dtype=torch.float32 - ) + dataset.meta.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32) return dataset -def resolve_delta_timestamps( - cfg, ds_meta: LeRobotDatasetMetadata -) -> dict[str, list] | None: +def resolve_delta_timestamps(cfg, ds_meta: LeRobotDatasetMetadata) -> dict[str, list] | None: """Resolves delta_timestamps by reading from the 'delta_indices' properties of the PreTrainedConfig. Args: @@ -182,9 +192,7 @@ def resolve_delta_timestamps( if key == ACTION and cfg.action_delta_indices is not None: delta_timestamps[key] = [i / ds_meta.fps for i in cfg.action_delta_indices] if key.startswith(OBS_PREFIX) and cfg.observation_delta_indices is not None: - delta_timestamps[key] = [ - i / ds_meta.fps for i in cfg.observation_delta_indices - ] + delta_timestamps[key] = [i / ds_meta.fps for i in cfg.observation_delta_indices] if len(delta_timestamps) == 0: delta_timestamps = None @@ -247,9 +255,7 @@ def validate_visual_features_consistency( cfg (PreTrainedConfig): The model or policy configuration containing input_features and type. features (Dict[str, PolicyFeature]): A mapping of feature names to PolicyFeature objects. """ - expected_visuals = { - k for k, v in cfg.input_features.items() if v.type == FeatureType.VISUAL - } + expected_visuals = {k for k, v in cfg.input_features.items() if v.type == FeatureType.VISUAL} provided_visuals = {k for k, v in features.items() if v.type == FeatureType.VISUAL} if not provided_visuals.issubset(expected_visuals): raise_feature_mismatch_error(provided_visuals, expected_visuals) @@ -436,9 +442,13 @@ def update_policy( policy_model = policy.module if isinstance(policy, FSDP) else policy use_amp = getattr(policy_model.config, "use_amp", False) - autocast_context = torch.amp.autocast(get_platform().amp_device_type(), dtype=torch.bfloat16) if use_amp else nullcontext() + autocast_context = ( + torch.amp.autocast(get_platform().amp_device_type(), dtype=torch.bfloat16) + if use_amp + else nullcontext() + ) with autocast_context: - loss, _= policy.forward(batch) + loss, _ = policy.forward(batch) # TODO(rcadene): policy.unnormalize_outputs(out_dict) loss.backward() @@ -482,9 +492,7 @@ def main(config: TrainConfig, seed: int): model_name = config.model.model_name.lower() if model_name not in ["pi0", "pi0.5"]: - raise ValueError( - f"Invalid model_name: {model_name}. Must be 'pi0' or 'pi0.5'" - ) + raise ValueError(f"Invalid model_name: {model_name}. Must be 'pi0' or 'pi0.5'") # Load base config from checkpoint if model_name == "pi0.5": @@ -691,7 +699,10 @@ def main(config: TrainConfig, seed: int): if step % config.system.log_freq == 0 and is_main_process: logger.info(f"step: {step} loss: {train_tracker}") - if config.system.checkpoint.save_checkpoint and step % config.system.checkpoint.save_freq == 0: + if ( + config.system.checkpoint.save_checkpoint + and step % config.system.checkpoint.save_freq == 0 + ): dist.barrier() state_dict = get_model_state_dict( policy, @@ -743,7 +754,7 @@ def main(config: TrainConfig, seed: int): # Extract experiment config (seed, exp_dir, etc.) experiment_config = OmegaConf.to_container(config.experiment, resolve=True) - seed = experiment_config.get('seed', 42) + seed = experiment_config.get("seed", 42) logger.info("=" * 100) logger.info(f"Experiment: {experiment_config}") diff --git a/flagscale/train/train_qwen_gr00t.py b/flagscale/train/train_qwen_gr00t.py index 08d0771e10..ebf90670ee 100644 --- a/flagscale/train/train_qwen_gr00t.py +++ b/flagscale/train/train_qwen_gr00t.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Mainly adopted from # https://github.com/huggingface/lerobot/blob/2b304eeb841ae6c371e3dd341bbbb9dd254b07cb/src/lerobot/scripts/lerobot_train.py @@ -17,7 +31,11 @@ import torch.distributed as dist from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy from torch.distributed.device_mesh import init_device_mesh -from torch.distributed.checkpoint.state_dict import get_model_state_dict, get_optimizer_state_dict, StateDictOptions +from torch.distributed.checkpoint.state_dict import ( + get_model_state_dict, + get_optimizer_state_dict, + StateDictOptions, +) from torch.optim import Optimizer from flagscale.logger import logger @@ -97,7 +115,6 @@ def make_dataset(config: TrainConfig, policy_config: PreTrainedConfig): ds_meta = LeRobotDatasetMetadata(root=config.data.data_path, revision=None) delta_timestamps = _resolve_delta_timestamps(policy_config, ds_meta) - # torchcodec depends on NVIDIA NVDEC which is not available on all platforms (e.g. MUSA); # fall back to pyav for non-CUDA platforms. video_backend = "torchcodec" if get_platform().name() == "cuda" else "pyav" @@ -258,7 +275,6 @@ def _format_meter_val(meter: AverageMeter) -> str: return " ".join(display_list) - def make_pre_post_processors( policy, data_config, @@ -464,7 +480,9 @@ def update_policy( optimizer.zero_grad() autocast_context = ( - torch.amp.autocast(get_platform().amp_device_type(), dtype=torch.bfloat16) if use_amp else nullcontext() + torch.amp.autocast(get_platform().amp_device_type(), dtype=torch.bfloat16) + if use_amp + else nullcontext() ) with autocast_context: output = policy(batch, vlm_batch=vlm_batch) @@ -491,7 +509,9 @@ def update_policy( policy.update() train_metrics.loss = loss.item() - train_metrics.grad_norm = grad_norm.full_tensor().item() if hasattr(grad_norm, 'full_tensor') else grad_norm.item() + train_metrics.grad_norm = ( + grad_norm.full_tensor().item() if hasattr(grad_norm, "full_tensor") else grad_norm.item() + ) train_metrics.lr = optimizer.param_groups[0]["lr"] train_metrics.update_s = time.perf_counter() - start_time if "vlm_loss" in output and "vlm_loss" in train_metrics.metrics: @@ -575,10 +595,17 @@ def main(config: TrainConfig, seed: int): policy = make_policy(policy_config, ds_meta) dist.barrier() - dataset_stats = ds_meta.stats if not isinstance(dataset, LeRobotMixtureDataset) else dataset.merged_stats + dataset_stats = ( + ds_meta.stats + if not isinstance(dataset, LeRobotMixtureDataset) + else dataset.merged_stats + ) # Create processors - only provide dataset_stats if not resuming from saved processors preprocessor, postprocessor = make_pre_post_processors( - policy, config.data, dataset_stats=dataset_stats, device=device.type, + policy, + config.data, + dataset_stats=dataset_stats, + device=device.type, ) num_workers = 0 # config.system.num_workers @@ -604,7 +631,6 @@ def main(config: TrainConfig, seed: int): prefetch_factor=2 if num_workers > 0 else None, ) - dl_iter = cycle(dataloader) num_frames = dataset.num_frames if hasattr(dataset, "num_frames") else len(dataset) num_episodes = dataset.num_episodes if hasattr(dataset, "num_episodes") else 1 @@ -631,7 +657,10 @@ def main(config: TrainConfig, seed: int): resume_from = config.system.checkpoint.resume_from if resume_from: step, dl_state = load_training_state_fsdp2( - Path(resume_from), policy, optimizer, lr_scheduler, + Path(resume_from), + policy, + optimizer, + lr_scheduler, ) # Restore dataloader position via StatefulDistributedSampler if state # was saved; otherwise fall back to advancing the iterator manually. @@ -729,7 +758,11 @@ def main(config: TrainConfig, seed: int): checkpoint_dir = get_step_checkpoint_dir( output_dir, config.system.train_steps, step ) - dl_state = sampler.state_dict() if sampler is not None and hasattr(sampler, "state_dict") else None + dl_state = ( + sampler.state_dict() + if sampler is not None and hasattr(sampler, "state_dict") + else None + ) save_checkpoint( checkpoint_dir=checkpoint_dir, step=step, diff --git a/flagscale/train/utils/chunked_cross_entropy.py b/flagscale/train/utils/chunked_cross_entropy.py index c3a4f651c8..c0817655e8 100644 --- a/flagscale/train/utils/chunked_cross_entropy.py +++ b/flagscale/train/utils/chunked_cross_entropy.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import torch import torch.nn.functional as F diff --git a/flagscale/train/utils/optim_setup.py b/flagscale/train/utils/optim_setup.py index d8dfc02756..6a4fe1dd16 100644 --- a/flagscale/train/utils/optim_setup.py +++ b/flagscale/train/utils/optim_setup.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Optimizer setup utilities. Supports: diff --git a/flagscale/transformations/__init__.py b/flagscale/transformations/__init__.py index 837240c455..02e0461a66 100644 --- a/flagscale/transformations/__init__.py +++ b/flagscale/transformations/__init__.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from omegaconf import DictConfig from .log_io_transformation import LogIOTransformation diff --git a/flagscale/transformations/hook.py b/flagscale/transformations/hook.py index 9b73c320e3..5695513b51 100644 --- a/flagscale/transformations/hook.py +++ b/flagscale/transformations/hook.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Modified from # https://github.com/huggingface/diffusers/blob/4a7556eaecc9872dea50ce161301edfa6392693c/src/diffusers/hooks/hooks.py diff --git a/flagscale/transformations/log_io_transformation.py b/flagscale/transformations/log_io_transformation.py index 8f80a62a38..0e597e5346 100644 --- a/flagscale/transformations/log_io_transformation.py +++ b/flagscale/transformations/log_io_transformation.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from typing import Any import torch diff --git a/flagscale/transformations/state_scope_transformation.py b/flagscale/transformations/state_scope_transformation.py index 585dc84b1b..2cac277348 100644 --- a/flagscale/transformations/state_scope_transformation.py +++ b/flagscale/transformations/state_scope_transformation.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from typing import Any import torch.nn as nn diff --git a/flagscale/transformations/state_store.py b/flagscale/transformations/state_store.py index db5fd280d0..fb15a263f8 100644 --- a/flagscale/transformations/state_store.py +++ b/flagscale/transformations/state_store.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Modified from # https://github.com/huggingface/diffusers/blob/4a7556eaecc9872dea50ce161301edfa6392693c/src/diffusers/hooks/hooks.py diff --git a/flagscale/transformations/transformation.py b/flagscale/transformations/transformation.py index c869c0a61a..38a29d5161 100644 --- a/flagscale/transformations/transformation.py +++ b/flagscale/transformations/transformation.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import importlib from abc import ABC, abstractmethod from collections.abc import Iterable diff --git a/flagscale/utils.py b/flagscale/utils.py index 4197164f34..34d7abe1ad 100644 --- a/flagscale/utils.py +++ b/flagscale/utils.py @@ -1,3 +1,18 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + def flatten_dict_to_args(config_dict, ignore_keys=[]): args = [] for key, value in config_dict.items(): diff --git a/pyproject.toml b/pyproject.toml index d413d31f61..275ff7a9cf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + [build-system] requires = [ "setuptools>=77.0", diff --git a/requirements/ascend/base.txt b/requirements/ascend/base.txt index 3b6035360d..01039bd372 100644 --- a/requirements/ascend/base.txt +++ b/requirements/ascend/base.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Ascend platform-specific base dependencies -r ../common.txt diff --git a/requirements/ascend/serve.txt b/requirements/ascend/serve.txt index faa8d89419..96c5e73360 100644 --- a/requirements/ascend/serve.txt +++ b/requirements/ascend/serve.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Ascend serve-specific dependencies -r ./base.txt diff --git a/requirements/ascend/train.txt b/requirements/ascend/train.txt index ae046a9e32..1b30b5a515 100644 --- a/requirements/ascend/train.txt +++ b/requirements/ascend/train.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Ascend train-specific dependencies -r ./base.txt diff --git a/requirements/common.txt b/requirements/common.txt index c2beea110c..031f349459 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Common dependencies used by all tasks # Platform-agnostic packages including configuration management and experiment tracking diff --git a/requirements/cuda/all.txt b/requirements/cuda/all.txt index f9d526607e..ddb3e5fa05 100644 --- a/requirements/cuda/all.txt +++ b/requirements/cuda/all.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # All task dependencies combined # Includes train, inference, rl, hetero_train diff --git a/requirements/cuda/base.txt b/requirements/cuda/base.txt index ae5cd5f3c5..066009e3c3 100644 --- a/requirements/cuda/base.txt +++ b/requirements/cuda/base.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # CUDA platform-specific base dependencies -r ../common.txt diff --git a/requirements/cuda/hetero_train.txt b/requirements/cuda/hetero_train.txt index bac4885122..fc2d741387 100644 --- a/requirements/cuda/hetero_train.txt +++ b/requirements/cuda/hetero_train.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Heterogeneous training-specific dependencies # Currently extends training dependencies but separated for future flexibility diff --git a/requirements/cuda/inference.txt b/requirements/cuda/inference.txt index e1e1cc2cd2..840e60a1fe 100644 --- a/requirements/cuda/inference.txt +++ b/requirements/cuda/inference.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Inference-specific dependencies -r ./base.txt diff --git a/requirements/cuda/rl.txt b/requirements/cuda/rl.txt index a79865ebd5..5d192dd9b9 100644 --- a/requirements/cuda/rl.txt +++ b/requirements/cuda/rl.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Reinforcement Learning specific dependencies -r ./base.txt diff --git a/requirements/cuda/serve.txt b/requirements/cuda/serve.txt index 50bfd806e2..ee99610be7 100644 --- a/requirements/cuda/serve.txt +++ b/requirements/cuda/serve.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # serve-specific dependencies -r ./base.txt diff --git a/requirements/cuda/train.txt b/requirements/cuda/train.txt index d700b0e60b..f3b31ad8df 100644 --- a/requirements/cuda/train.txt +++ b/requirements/cuda/train.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Training-specific dependencies -r ./base.txt diff --git a/requirements/dev.txt b/requirements/dev.txt index a5a3110e64..53da11028e 100644 --- a/requirements/dev.txt +++ b/requirements/dev.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Development dependencies # Includes build tools, linting, testing utilities # Install with: ./tools/install/install.sh --dev ... diff --git a/requirements/musa/base.txt b/requirements/musa/base.txt index 11e10d5a25..46602d98ba 100644 --- a/requirements/musa/base.txt +++ b/requirements/musa/base.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # MUSA platform-specific base dependencies -r ../common.txt diff --git a/requirements/musa/train.txt b/requirements/musa/train.txt index 6fd66860c5..33680f2e4b 100644 --- a/requirements/musa/train.txt +++ b/requirements/musa/train.txt @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Training-specific dependencies -r ./base.txt diff --git a/run.py b/run.py index c3bf2a0645..5e3809bf5e 100644 --- a/run.py +++ b/run.py @@ -1,4 +1,19 @@ #!/usr/bin/env python + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Backward compatibility wrapper for flagscale.run. This file is kept for backward compatibility with existing scripts and documentation. diff --git a/setup.py b/setup.py index 5b10239287..b0f7685811 100644 --- a/setup.py +++ b/setup.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import platform import re diff --git a/tests/__init__.py b/tests/__init__.py index db5dae1f3c..3898f31428 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1 +1,15 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # This file makes tests/ a Python package diff --git a/tests/conftest.py b/tests/conftest.py index cbe266e8b8..ea75027363 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import pytest diff --git a/tests/functional_tests/benchmark/qwen3/conf/0_6b_metax.yaml b/tests/functional_tests/benchmark/qwen3/conf/0_6b_metax.yaml index 908aa872e7..84877cf063 100644 --- a/tests/functional_tests/benchmark/qwen3/conf/0_6b_metax.yaml +++ b/tests/functional_tests/benchmark/qwen3/conf/0_6b_metax.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: 0_6b_metax diff --git a/tests/functional_tests/benchmark/qwen3/conf/tp2_pp2.yaml b/tests/functional_tests/benchmark/qwen3/conf/tp2_pp2.yaml index e1148e89c6..9c0d3a4b2b 100644 --- a/tests/functional_tests/benchmark/qwen3/conf/tp2_pp2.yaml +++ b/tests/functional_tests/benchmark/qwen3/conf/tp2_pp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: tp2_pp2 diff --git a/tests/functional_tests/benchmark/qwen3/conf/train/0_6b_metax.yaml b/tests/functional_tests/benchmark/qwen3/conf/train/0_6b_metax.yaml index 71df858e65..c4f99a3912 100644 --- a/tests/functional_tests/benchmark/qwen3/conf/train/0_6b_metax.yaml +++ b/tests/functional_tests/benchmark/qwen3/conf/train/0_6b_metax.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - data diff --git a/tests/functional_tests/benchmark/qwen3/conf/train/data.yaml b/tests/functional_tests/benchmark/qwen3/conf/train/data.yaml index 8d9d4bd4ca..2430cc9a71 100644 --- a/tests/functional_tests/benchmark/qwen3/conf/train/data.yaml +++ b/tests/functional_tests/benchmark/qwen3/conf/train/data.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + data: data_path: /home/gitlab-runner/data/pile_wikipedia_demo/pile_wikipedia_demo split: 1 diff --git a/tests/functional_tests/benchmark/qwen3/conf/train/tp2_pp2.yaml b/tests/functional_tests/benchmark/qwen3/conf/train/tp2_pp2.yaml index 6e22e9bee5..cb78e6c010 100644 --- a/tests/functional_tests/benchmark/qwen3/conf/train/tp2_pp2.yaml +++ b/tests/functional_tests/benchmark/qwen3/conf/train/tp2_pp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - data diff --git a/tests/functional_tests/hetero_train/aquila/conf/dp2dp4_shared_embedding.yaml b/tests/functional_tests/hetero_train/aquila/conf/dp2dp4_shared_embedding.yaml index e5e31421b6..037e4785be 100644 --- a/tests/functional_tests/hetero_train/aquila/conf/dp2dp4_shared_embedding.yaml +++ b/tests/functional_tests/hetero_train/aquila/conf/dp2dp4_shared_embedding.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: dp2dp4_shared_embedding diff --git a/tests/functional_tests/hetero_train/aquila/conf/tp2dp1pp1_tp2dp2pp1_tp1dp2pp1.yaml b/tests/functional_tests/hetero_train/aquila/conf/tp2dp1pp1_tp2dp2pp1_tp1dp2pp1.yaml index 579f6cf28c..57ea645122 100644 --- a/tests/functional_tests/hetero_train/aquila/conf/tp2dp1pp1_tp2dp2pp1_tp1dp2pp1.yaml +++ b/tests/functional_tests/hetero_train/aquila/conf/tp2dp1pp1_tp2dp2pp1_tp1dp2pp1.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: tp2dp1pp1_tp2dp2pp1_tp1dp2pp1 diff --git a/tests/functional_tests/hetero_train/aquila/conf/tp2pp1_tp4pp1_tp2pp1.yaml b/tests/functional_tests/hetero_train/aquila/conf/tp2pp1_tp4pp1_tp2pp1.yaml index 8dbab54730..f99e9346c4 100644 --- a/tests/functional_tests/hetero_train/aquila/conf/tp2pp1_tp4pp1_tp2pp1.yaml +++ b/tests/functional_tests/hetero_train/aquila/conf/tp2pp1_tp4pp1_tp2pp1.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: tp2pp1_tp4pp1_tp2pp1 diff --git a/tests/functional_tests/hetero_train/aquila/conf/train/data.yaml b/tests/functional_tests/hetero_train/aquila/conf/train/data.yaml index b89276d50c..07c146b5d6 100644 --- a/tests/functional_tests/hetero_train/aquila/conf/train/data.yaml +++ b/tests/functional_tests/hetero_train/aquila/conf/train/data.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + data: data_path: /home/gitlab-runner/data/pile_wikipedia_demo/pile_wikipedia_demo split: 1 diff --git a/tests/functional_tests/hetero_train/aquila/conf/train/dp2dp4_shared_embedding.yaml b/tests/functional_tests/hetero_train/aquila/conf/train/dp2dp4_shared_embedding.yaml index 4726708ab5..1f12f2c3df 100644 --- a/tests/functional_tests/hetero_train/aquila/conf/train/dp2dp4_shared_embedding.yaml +++ b/tests/functional_tests/hetero_train/aquila/conf/train/dp2dp4_shared_embedding.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - data diff --git a/tests/functional_tests/hetero_train/aquila/conf/train/tp2dp1pp1_tp2dp2pp1_tp1dp2pp1.yaml b/tests/functional_tests/hetero_train/aquila/conf/train/tp2dp1pp1_tp2dp2pp1_tp1dp2pp1.yaml index 4c9753a0fd..a9fbb9b873 100644 --- a/tests/functional_tests/hetero_train/aquila/conf/train/tp2dp1pp1_tp2dp2pp1_tp1dp2pp1.yaml +++ b/tests/functional_tests/hetero_train/aquila/conf/train/tp2dp1pp1_tp2dp2pp1_tp1dp2pp1.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - data diff --git a/tests/functional_tests/hetero_train/aquila/conf/train/tp2pp1_tp4pp1_tp2pp1.yaml b/tests/functional_tests/hetero_train/aquila/conf/train/tp2pp1_tp4pp1_tp2pp1.yaml index 8a55935ba8..db5cbc4079 100644 --- a/tests/functional_tests/hetero_train/aquila/conf/train/tp2pp1_tp4pp1_tp2pp1.yaml +++ b/tests/functional_tests/hetero_train/aquila/conf/train/tp2pp1_tp4pp1_tp2pp1.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - data diff --git a/tests/functional_tests/inference/deepseek_r1_distill_qwen/conf/7b_tp2.yaml b/tests/functional_tests/inference/deepseek_r1_distill_qwen/conf/7b_tp2.yaml index c9ef51bd9b..6ab50c4966 100644 --- a/tests/functional_tests/inference/deepseek_r1_distill_qwen/conf/7b_tp2.yaml +++ b/tests/functional_tests/inference/deepseek_r1_distill_qwen/conf/7b_tp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: 7b_tp2 diff --git a/tests/functional_tests/inference/deepseek_r1_distill_qwen/conf/inference/7b_tp2.yaml b/tests/functional_tests/inference/deepseek_r1_distill_qwen/conf/inference/7b_tp2.yaml index ce82580d37..9169d14aa0 100644 --- a/tests/functional_tests/inference/deepseek_r1_distill_qwen/conf/inference/7b_tp2.yaml +++ b/tests/functional_tests/inference/deepseek_r1_distill_qwen/conf/inference/7b_tp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: model: /home/gitlab-runner/data/DeepSeek-R1-Distill-Qwen-7B tokenizer: /home/gitlab-runner/data/DeepSeek-R1-Distill-Qwen-7B diff --git a/tests/functional_tests/inference/deepseek_r1_distill_qwen/results_gold/7b_tp2 b/tests/functional_tests/inference/deepseek_r1_distill_qwen/results_gold/7b_tp2 index 2bf225d4e2..075af71b36 100644 --- a/tests/functional_tests/inference/deepseek_r1_distill_qwen/results_gold/7b_tp2 +++ b/tests/functional_tests/inference/deepseek_r1_distill_qwen/results_gold/7b_tp2 @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + ************************************************** output.prompt='The president of the United States' output.outputs[0].text=', George W. Bush, is known for his leadership in the 19' diff --git a/tests/functional_tests/inference/deepseek_r1_distill_qwen_flaggems/conf/7b_tp2.yaml b/tests/functional_tests/inference/deepseek_r1_distill_qwen_flaggems/conf/7b_tp2.yaml index e11a52e0c0..0ffabe251f 100644 --- a/tests/functional_tests/inference/deepseek_r1_distill_qwen_flaggems/conf/7b_tp2.yaml +++ b/tests/functional_tests/inference/deepseek_r1_distill_qwen_flaggems/conf/7b_tp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: 7b_tp2 diff --git a/tests/functional_tests/inference/deepseek_r1_distill_qwen_flaggems/conf/inference/7b_tp2.yaml b/tests/functional_tests/inference/deepseek_r1_distill_qwen_flaggems/conf/inference/7b_tp2.yaml index ce82580d37..9169d14aa0 100644 --- a/tests/functional_tests/inference/deepseek_r1_distill_qwen_flaggems/conf/inference/7b_tp2.yaml +++ b/tests/functional_tests/inference/deepseek_r1_distill_qwen_flaggems/conf/inference/7b_tp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: model: /home/gitlab-runner/data/DeepSeek-R1-Distill-Qwen-7B tokenizer: /home/gitlab-runner/data/DeepSeek-R1-Distill-Qwen-7B diff --git a/tests/functional_tests/inference/deepseek_r1_distill_qwen_flaggems/results_gold/7b_tp2 b/tests/functional_tests/inference/deepseek_r1_distill_qwen_flaggems/results_gold/7b_tp2 index 1b1d333737..7000a3ff68 100644 --- a/tests/functional_tests/inference/deepseek_r1_distill_qwen_flaggems/results_gold/7b_tp2 +++ b/tests/functional_tests/inference/deepseek_r1_distill_qwen_flaggems/results_gold/7b_tp2 @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + ************************************************** output.prompt='The president of the United States' output.outputs[0].text=', George W. Bush, has been given a 12-hour clock problem' diff --git a/tests/functional_tests/inference/qwen3/conf/4b_tp2.yaml b/tests/functional_tests/inference/qwen3/conf/4b_tp2.yaml index ea940956de..8de2ac2d1e 100644 --- a/tests/functional_tests/inference/qwen3/conf/4b_tp2.yaml +++ b/tests/functional_tests/inference/qwen3/conf/4b_tp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: 4b_tp2 diff --git a/tests/functional_tests/inference/qwen3/conf/4b_tp2_ascend.yaml b/tests/functional_tests/inference/qwen3/conf/4b_tp2_ascend.yaml index e13f8c4353..84e2adb1cc 100644 --- a/tests/functional_tests/inference/qwen3/conf/4b_tp2_ascend.yaml +++ b/tests/functional_tests/inference/qwen3/conf/4b_tp2_ascend.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: 4b_tp2_ascend diff --git a/tests/functional_tests/inference/qwen3/conf/inference/4b_tp2.yaml b/tests/functional_tests/inference/qwen3/conf/inference/4b_tp2.yaml index 7140169825..00e76dc7d7 100644 --- a/tests/functional_tests/inference/qwen3/conf/inference/4b_tp2.yaml +++ b/tests/functional_tests/inference/qwen3/conf/inference/4b_tp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: model: /home/gitlab-runner/data/Qwen3-4B tokenizer: /home/gitlab-runner/data/Qwen3-4B diff --git a/tests/functional_tests/inference/qwen3/conf/inference/4b_tp2_ascend.yaml b/tests/functional_tests/inference/qwen3/conf/inference/4b_tp2_ascend.yaml index 366e720810..74574fdaaf 100644 --- a/tests/functional_tests/inference/qwen3/conf/inference/4b_tp2_ascend.yaml +++ b/tests/functional_tests/inference/qwen3/conf/inference/4b_tp2_ascend.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: model: /home/gitlab-runner/data/Qwen3-4B tokenizer: /home/gitlab-runner/data/Qwen3-4B diff --git a/tests/functional_tests/inference/qwen3/results_gold/4b_tp2 b/tests/functional_tests/inference/qwen3/results_gold/4b_tp2 index 73a1700546..f428c1b542 100644 --- a/tests/functional_tests/inference/qwen3/results_gold/4b_tp2 +++ b/tests/functional_tests/inference/qwen3/results_gold/4b_tp2 @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + ************************************************** output.prompt='The first President of the United States' output.outputs[0].text=', George Washington, was born in 1732. The 4' diff --git a/tests/functional_tests/inference/qwen3/results_gold/4b_tp2_ascend b/tests/functional_tests/inference/qwen3/results_gold/4b_tp2_ascend index 73a1700546..f428c1b542 100644 --- a/tests/functional_tests/inference/qwen3/results_gold/4b_tp2_ascend +++ b/tests/functional_tests/inference/qwen3/results_gold/4b_tp2_ascend @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + ************************************************** output.prompt='The first President of the United States' output.outputs[0].text=', George Washington, was born in 1732. The 4' diff --git a/tests/functional_tests/inference/qwen3_flaggems/conf/4b_tp2.yaml b/tests/functional_tests/inference/qwen3_flaggems/conf/4b_tp2.yaml index 47fe827f84..3d7f1d928b 100644 --- a/tests/functional_tests/inference/qwen3_flaggems/conf/4b_tp2.yaml +++ b/tests/functional_tests/inference/qwen3_flaggems/conf/4b_tp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: 4b_tp2 diff --git a/tests/functional_tests/inference/qwen3_flaggems/conf/inference/4b_tp2.yaml b/tests/functional_tests/inference/qwen3_flaggems/conf/inference/4b_tp2.yaml index 7140169825..00e76dc7d7 100644 --- a/tests/functional_tests/inference/qwen3_flaggems/conf/inference/4b_tp2.yaml +++ b/tests/functional_tests/inference/qwen3_flaggems/conf/inference/4b_tp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: model: /home/gitlab-runner/data/Qwen3-4B tokenizer: /home/gitlab-runner/data/Qwen3-4B diff --git a/tests/functional_tests/inference/qwen3_flaggems/results_gold/4b_tp2 b/tests/functional_tests/inference/qwen3_flaggems/results_gold/4b_tp2 index ae870aba03..bfd81c02d1 100644 --- a/tests/functional_tests/inference/qwen3_flaggems/results_gold/4b_tp2 +++ b/tests/functional_tests/inference/qwen3_flaggems/results_gold/4b_tp2 @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + ************************************************** output.prompt='The first President of the United States' output.outputs[0].text=', George Washington, was born in 1732. The 4' diff --git a/tests/functional_tests/inference/robobrain2/conf/7b_tp2.yaml b/tests/functional_tests/inference/robobrain2/conf/7b_tp2.yaml index f281dcc56b..e4df658b73 100644 --- a/tests/functional_tests/inference/robobrain2/conf/7b_tp2.yaml +++ b/tests/functional_tests/inference/robobrain2/conf/7b_tp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: 7b_tp2 diff --git a/tests/functional_tests/inference/robobrain2/conf/inference/7b_tp2.yaml b/tests/functional_tests/inference/robobrain2/conf/inference/7b_tp2.yaml index d64c90a75a..0bfb76c47d 100644 --- a/tests/functional_tests/inference/robobrain2/conf/inference/7b_tp2.yaml +++ b/tests/functional_tests/inference/robobrain2/conf/inference/7b_tp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: model: /home/gitlab-runner/data/BAAI/RoboBrain2.0-7B tokenizer: /home/gitlab-runner/data/BAAI/RoboBrain2.0-7B diff --git a/tests/functional_tests/inference/robobrain2/results_gold/7b_tp2 b/tests/functional_tests/inference/robobrain2/results_gold/7b_tp2 index 9706168582..d13c7a4b0b 100644 --- a/tests/functional_tests/inference/robobrain2/results_gold/7b_tp2 +++ b/tests/functional_tests/inference/robobrain2/results_gold/7b_tp2 @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + ************************************************** output.prompt='What is shown in this image?' output.outputs[0].text='The image shows a large glass container filled with yellow' diff --git a/tests/functional_tests/inference/robobrain2_flaggems/conf/7b_tp2.yaml b/tests/functional_tests/inference/robobrain2_flaggems/conf/7b_tp2.yaml index 53bdc6a5c8..610d1a802c 100644 --- a/tests/functional_tests/inference/robobrain2_flaggems/conf/7b_tp2.yaml +++ b/tests/functional_tests/inference/robobrain2_flaggems/conf/7b_tp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - inference: 7b_tp2 diff --git a/tests/functional_tests/inference/robobrain2_flaggems/conf/inference/7b_tp2.yaml b/tests/functional_tests/inference/robobrain2_flaggems/conf/inference/7b_tp2.yaml index d64c90a75a..0bfb76c47d 100644 --- a/tests/functional_tests/inference/robobrain2_flaggems/conf/inference/7b_tp2.yaml +++ b/tests/functional_tests/inference/robobrain2_flaggems/conf/inference/7b_tp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + llm: model: /home/gitlab-runner/data/BAAI/RoboBrain2.0-7B tokenizer: /home/gitlab-runner/data/BAAI/RoboBrain2.0-7B diff --git a/tests/functional_tests/inference/robobrain2_flaggems/results_gold/7b_tp2 b/tests/functional_tests/inference/robobrain2_flaggems/results_gold/7b_tp2 index 9706168582..d13c7a4b0b 100644 --- a/tests/functional_tests/inference/robobrain2_flaggems/results_gold/7b_tp2 +++ b/tests/functional_tests/inference/robobrain2_flaggems/results_gold/7b_tp2 @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + ************************************************** output.prompt='What is shown in this image?' output.outputs[0].text='The image shows a large glass container filled with yellow' diff --git a/tests/functional_tests/rl/qwen2_5/conf/0_5b.yaml b/tests/functional_tests/rl/qwen2_5/conf/0_5b.yaml index 694e8a3707..c082385fcf 100644 --- a/tests/functional_tests/rl/qwen2_5/conf/0_5b.yaml +++ b/tests/functional_tests/rl/qwen2_5/conf/0_5b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - rl: 0_5b - _self_ diff --git a/tests/functional_tests/rl/qwen2_5/conf/rl/0_5b.yaml b/tests/functional_tests/rl/qwen2_5/conf/rl/0_5b.yaml index 9367024cfd..492d3ba1ba 100644 --- a/tests/functional_tests/rl/qwen2_5/conf/rl/0_5b.yaml +++ b/tests/functional_tests/rl/qwen2_5/conf/rl/0_5b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + #default verl config path config-path: config diff --git a/tests/functional_tests/serve/base/conf/multiple_model.yaml b/tests/functional_tests/serve/base/conf/multiple_model.yaml index 727794f93a..2b6ecf029a 100644 --- a/tests/functional_tests/serve/base/conf/multiple_model.yaml +++ b/tests/functional_tests/serve/base/conf/multiple_model.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: multiple_model diff --git a/tests/functional_tests/serve/base/conf/serve/multiple_model.yaml b/tests/functional_tests/serve/base/conf/serve/multiple_model.yaml index 25d86c6537..6f5910d5c4 100644 --- a/tests/functional_tests/serve/base/conf/serve/multiple_model.yaml +++ b/tests/functional_tests/serve/base/conf/serve/multiple_model.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: A module: ./tests/functional_tests/serve/utils/models/base_model.py name: ModelA diff --git a/tests/functional_tests/serve/qwen2_5/conf/0.5b.yaml b/tests/functional_tests/serve/qwen2_5/conf/0.5b.yaml index c23595fb65..e760d521e4 100644 --- a/tests/functional_tests/serve/qwen2_5/conf/0.5b.yaml +++ b/tests/functional_tests/serve/qwen2_5/conf/0.5b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 0.5b diff --git a/tests/functional_tests/serve/qwen2_5/conf/0.5b_multiple_instance.yaml b/tests/functional_tests/serve/qwen2_5/conf/0.5b_multiple_instance.yaml index e4f03ac2f5..b5b79d34c8 100644 --- a/tests/functional_tests/serve/qwen2_5/conf/0.5b_multiple_instance.yaml +++ b/tests/functional_tests/serve/qwen2_5/conf/0.5b_multiple_instance.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 0.5b_multiple_instance diff --git a/tests/functional_tests/serve/qwen2_5/conf/serve/0.5b.yaml b/tests/functional_tests/serve/qwen2_5/conf/serve/0.5b.yaml index 7a3c4eeeb9..2a931a1a14 100644 --- a/tests/functional_tests/serve/qwen2_5/conf/serve/0.5b.yaml +++ b/tests/functional_tests/serve/qwen2_5/conf/serve/0.5b.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine: vllm engine_args: diff --git a/tests/functional_tests/serve/qwen2_5/conf/serve/0.5b_multiple_instance.yaml b/tests/functional_tests/serve/qwen2_5/conf/serve/0.5b_multiple_instance.yaml index bed1aafc0e..9d5f2c79c2 100644 --- a/tests/functional_tests/serve/qwen2_5/conf/serve/0.5b_multiple_instance.yaml +++ b/tests/functional_tests/serve/qwen2_5/conf/serve/0.5b_multiple_instance.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine: vllm engine_args: diff --git a/tests/functional_tests/serve/qwen3/conf/4b_tp2_ascend.yaml b/tests/functional_tests/serve/qwen3/conf/4b_tp2_ascend.yaml index e5253150a0..9cb11e458a 100644 --- a/tests/functional_tests/serve/qwen3/conf/4b_tp2_ascend.yaml +++ b/tests/functional_tests/serve/qwen3/conf/4b_tp2_ascend.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - serve: 4b_tp2_ascend diff --git a/tests/functional_tests/serve/qwen3/conf/serve/4b_tp2_ascend.yaml b/tests/functional_tests/serve/qwen3/conf/serve/4b_tp2_ascend.yaml index ff478d3fc2..08dc7864c0 100644 --- a/tests/functional_tests/serve/qwen3/conf/serve/4b_tp2_ascend.yaml +++ b/tests/functional_tests/serve/qwen3/conf/serve/4b_tp2_ascend.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + - serve_id: vllm_model engine: vllm engine_args: diff --git a/tests/functional_tests/serve/utils/models/base_model.py b/tests/functional_tests/serve/utils/models/base_model.py index a44098f336..3e90914271 100644 --- a/tests/functional_tests/serve/utils/models/base_model.py +++ b/tests/functional_tests/serve/utils/models/base_model.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from util_models.util_model import fn diff --git a/tests/functional_tests/serve/utils/models/util_models/util_model.py b/tests/functional_tests/serve/utils/models/util_models/util_model.py index 025044b590..22f0052d61 100644 --- a/tests/functional_tests/serve/utils/models/util_models/util_model.py +++ b/tests/functional_tests/serve/utils/models/util_models/util_model.py @@ -1,3 +1,18 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + def fn(input_data): res = input_data + "__add_process_fn" return res diff --git a/tests/functional_tests/train/aquila/conf/tp2_pp2.yaml b/tests/functional_tests/train/aquila/conf/tp2_pp2.yaml index 76091741cb..d502fa11ad 100644 --- a/tests/functional_tests/train/aquila/conf/tp2_pp2.yaml +++ b/tests/functional_tests/train/aquila/conf/tp2_pp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: tp2_pp2 diff --git a/tests/functional_tests/train/aquila/conf/tp4_pp2.yaml b/tests/functional_tests/train/aquila/conf/tp4_pp2.yaml index 560aa9a06e..be1fdf2509 100644 --- a/tests/functional_tests/train/aquila/conf/tp4_pp2.yaml +++ b/tests/functional_tests/train/aquila/conf/tp4_pp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: tp4_pp2 diff --git a/tests/functional_tests/train/aquila/conf/train/data.yaml b/tests/functional_tests/train/aquila/conf/train/data.yaml index 81c70152f4..db6bc2d6a6 100644 --- a/tests/functional_tests/train/aquila/conf/train/data.yaml +++ b/tests/functional_tests/train/aquila/conf/train/data.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + data: data_path: /home/gitlab-runner/data/pile_wikipedia_demo/pile_wikipedia_demo # data_path: /share/project/ayl/test_data/pile_wikipedia_demo diff --git a/tests/functional_tests/train/aquila/conf/train/tp2_pp2.yaml b/tests/functional_tests/train/aquila/conf/train/tp2_pp2.yaml index 10bf1bdf41..d13acfcf00 100644 --- a/tests/functional_tests/train/aquila/conf/train/tp2_pp2.yaml +++ b/tests/functional_tests/train/aquila/conf/train/tp2_pp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + system: tensor_model_parallel_size: 2 pipeline_model_parallel_size: 2 diff --git a/tests/functional_tests/train/aquila/conf/train/tp4_pp2.yaml b/tests/functional_tests/train/aquila/conf/train/tp4_pp2.yaml index 7ac171632c..810e9530d2 100644 --- a/tests/functional_tests/train/aquila/conf/train/tp4_pp2.yaml +++ b/tests/functional_tests/train/aquila/conf/train/tp4_pp2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - data diff --git a/tests/functional_tests/train/deepseek/conf/tp1_pp2_ep2_v4.yaml b/tests/functional_tests/train/deepseek/conf/tp1_pp2_ep2_v4.yaml index fd65fb3273..01a80b0201 100644 --- a/tests/functional_tests/train/deepseek/conf/tp1_pp2_ep2_v4.yaml +++ b/tests/functional_tests/train/deepseek/conf/tp1_pp2_ep2_v4.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: tp1_pp2_ep2_v4 diff --git a/tests/functional_tests/train/deepseek/conf/tp2_pp2_ep2.yaml b/tests/functional_tests/train/deepseek/conf/tp2_pp2_ep2.yaml index 686853877b..90001db6aa 100644 --- a/tests/functional_tests/train/deepseek/conf/tp2_pp2_ep2.yaml +++ b/tests/functional_tests/train/deepseek/conf/tp2_pp2_ep2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: tp2_pp2_ep2 diff --git a/tests/functional_tests/train/deepseek/conf/tp2_pp2_ep2_engram.yaml b/tests/functional_tests/train/deepseek/conf/tp2_pp2_ep2_engram.yaml index 14ed97c215..db5449fccd 100644 --- a/tests/functional_tests/train/deepseek/conf/tp2_pp2_ep2_engram.yaml +++ b/tests/functional_tests/train/deepseek/conf/tp2_pp2_ep2_engram.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: tp2_pp2_ep2_engram diff --git a/tests/functional_tests/train/deepseek/conf/train/data.yaml b/tests/functional_tests/train/deepseek/conf/train/data.yaml index 8f32abc235..22dccb0f7e 100644 --- a/tests/functional_tests/train/deepseek/conf/train/data.yaml +++ b/tests/functional_tests/train/deepseek/conf/train/data.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + data: data_path: /home/gitlab-runner/data/pile_wikipedia_demo/pile_wikipedia_demo split: 1 diff --git a/tests/functional_tests/train/deepseek/conf/train/tp1_pp2_ep2_v4.yaml b/tests/functional_tests/train/deepseek/conf/train/tp1_pp2_ep2_v4.yaml index e670eb0462..50ac32303d 100644 --- a/tests/functional_tests/train/deepseek/conf/train/tp1_pp2_ep2_v4.yaml +++ b/tests/functional_tests/train/deepseek/conf/train/tp1_pp2_ep2_v4.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - data diff --git a/tests/functional_tests/train/deepseek/conf/train/tp2_pp2_ep2.yaml b/tests/functional_tests/train/deepseek/conf/train/tp2_pp2_ep2.yaml index d4466f8ba1..76cabdb43d 100644 --- a/tests/functional_tests/train/deepseek/conf/train/tp2_pp2_ep2.yaml +++ b/tests/functional_tests/train/deepseek/conf/train/tp2_pp2_ep2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - data diff --git a/tests/functional_tests/train/deepseek/conf/train/tp2_pp2_ep2_engram.yaml b/tests/functional_tests/train/deepseek/conf/train/tp2_pp2_ep2_engram.yaml index 3a4513895b..e8aee429d2 100644 --- a/tests/functional_tests/train/deepseek/conf/train/tp2_pp2_ep2_engram.yaml +++ b/tests/functional_tests/train/deepseek/conf/train/tp2_pp2_ep2_engram.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - data diff --git a/tests/functional_tests/train/mixtral/conf/tp2_pp1_ep2.yaml b/tests/functional_tests/train/mixtral/conf/tp2_pp1_ep2.yaml index a711648de0..1a5c0aab18 100644 --- a/tests/functional_tests/train/mixtral/conf/tp2_pp1_ep2.yaml +++ b/tests/functional_tests/train/mixtral/conf/tp2_pp1_ep2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: tp2_pp1_ep2 diff --git a/tests/functional_tests/train/mixtral/conf/tp4_pp1_ep2.yaml b/tests/functional_tests/train/mixtral/conf/tp4_pp1_ep2.yaml index a48e17e647..9ae3dcf2c6 100644 --- a/tests/functional_tests/train/mixtral/conf/tp4_pp1_ep2.yaml +++ b/tests/functional_tests/train/mixtral/conf/tp4_pp1_ep2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: tp4_pp1_ep2 diff --git a/tests/functional_tests/train/mixtral/conf/train/data.yaml b/tests/functional_tests/train/mixtral/conf/train/data.yaml index f9ee6d97dd..246c45e05d 100644 --- a/tests/functional_tests/train/mixtral/conf/train/data.yaml +++ b/tests/functional_tests/train/mixtral/conf/train/data.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + data: data_path: /home/gitlab-runner/data/llama_00_text_document/llama_00_text_document split: 1 diff --git a/tests/functional_tests/train/mixtral/conf/train/tp2_pp1_ep2.yaml b/tests/functional_tests/train/mixtral/conf/train/tp2_pp1_ep2.yaml index 57e0ef4a0d..ad01c2027c 100644 --- a/tests/functional_tests/train/mixtral/conf/train/tp2_pp1_ep2.yaml +++ b/tests/functional_tests/train/mixtral/conf/train/tp2_pp1_ep2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - data diff --git a/tests/functional_tests/train/mixtral/conf/train/tp4_pp1_ep2.yaml b/tests/functional_tests/train/mixtral/conf/train/tp4_pp1_ep2.yaml index 88abb2e190..d98c94e138 100644 --- a/tests/functional_tests/train/mixtral/conf/train/tp4_pp1_ep2.yaml +++ b/tests/functional_tests/train/mixtral/conf/train/tp4_pp1_ep2.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - data diff --git a/tests/functional_tests/train/qwen3/conf/0_6b_metax.yaml b/tests/functional_tests/train/qwen3/conf/0_6b_metax.yaml index c5ef0e02dd..69f80038c2 100644 --- a/tests/functional_tests/train/qwen3/conf/0_6b_metax.yaml +++ b/tests/functional_tests/train/qwen3/conf/0_6b_metax.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - _self_ - train: 0_6b_metax diff --git a/tests/functional_tests/train/qwen3/conf/train/0_6b_metax.yaml b/tests/functional_tests/train/qwen3/conf/train/0_6b_metax.yaml index cf9ad69ec7..2cab3203a3 100644 --- a/tests/functional_tests/train/qwen3/conf/train/0_6b_metax.yaml +++ b/tests/functional_tests/train/qwen3/conf/train/0_6b_metax.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + defaults: - data diff --git a/tests/functional_tests/train/qwen3/conf/train/data.yaml b/tests/functional_tests/train/qwen3/conf/train/data.yaml index 8d9d4bd4ca..2430cc9a71 100644 --- a/tests/functional_tests/train/qwen3/conf/train/data.yaml +++ b/tests/functional_tests/train/qwen3/conf/train/data.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + data: data_path: /home/gitlab-runner/data/pile_wikipedia_demo/pile_wikipedia_demo split: 1 diff --git a/tests/test_utils/config/platforms/ascend.yaml b/tests/test_utils/config/platforms/ascend.yaml index 6b5e3ab78a..2fb3f90e8e 100644 --- a/tests/test_utils/config/platforms/ascend.yaml +++ b/tests/test_utils/config/platforms/ascend.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Ascend NPU Platform Configuration # Flexible test selection mechanism for Huawei Ascend NPU environments # diff --git a/tests/test_utils/config/platforms/cuda.yaml b/tests/test_utils/config/platforms/cuda.yaml index 965336fd23..9a5dbc1dcf 100644 --- a/tests/test_utils/config/platforms/cuda.yaml +++ b/tests/test_utils/config/platforms/cuda.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # CUDA-Based Platform Configuration # Flexible test selection mechanism for NVIDIA GPU environments # diff --git a/tests/test_utils/config/platforms/metax.yaml b/tests/test_utils/config/platforms/metax.yaml index cfe883dadf..206c8131da 100644 --- a/tests/test_utils/config/platforms/metax.yaml +++ b/tests/test_utils/config/platforms/metax.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # MetaX C550 Platform Configuration # Flexible test selection mechanism for MetaX C550 GPU (MACA) environments # diff --git a/tests/test_utils/config/platforms/template.yaml b/tests/test_utils/config/platforms/template.yaml index 7b7dab5419..975410fc96 100644 --- a/tests/test_utils/config/platforms/template.yaml +++ b/tests/test_utils/config/platforms/template.yaml @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Platform Configuration Template # This is a TEMPLATE file - copy and customize it for your specific hardware platform # diff --git a/tests/test_utils/runners/check_env.sh b/tests/test_utils/runners/check_env.sh index 7c444d0f9d..dc4f056160 100644 --- a/tests/test_utils/runners/check_env.sh +++ b/tests/test_utils/runners/check_env.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # check_env.sh - Print a snapshot of the hardware/software environment. # Intended to run at the start of a CI job to aid debugging. set -eo pipefail diff --git a/tests/test_utils/runners/check_results.py b/tests/test_utils/runners/check_results.py index e7cd9f7550..5f3fcd88b7 100644 --- a/tests/test_utils/runners/check_results.py +++ b/tests/test_utils/runners/check_results.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import json import os import re diff --git a/tests/test_utils/runners/helpers.py b/tests/test_utils/runners/helpers.py index 04ea7a689a..0c1725c63b 100755 --- a/tests/test_utils/runners/helpers.py +++ b/tests/test_utils/runners/helpers.py @@ -1,4 +1,19 @@ #!/usr/bin/env python3 + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """ Helper functions for FlagScale test runners. Provides utilities for parsing test configurations and formatting output. diff --git a/tests/test_utils/runners/parse_benchmark_output.py b/tests/test_utils/runners/parse_benchmark_output.py index 3c88332b1c..4354b60129 100644 --- a/tests/test_utils/runners/parse_benchmark_output.py +++ b/tests/test_utils/runners/parse_benchmark_output.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Parse host_0_localhost.output log file and produce benchmark.json. Usage: diff --git a/tests/test_utils/runners/parse_config.py b/tests/test_utils/runners/parse_config.py index 9a2ac337f4..5bdb7d16ba 100644 --- a/tests/test_utils/runners/parse_config.py +++ b/tests/test_utils/runners/parse_config.py @@ -1,4 +1,19 @@ #!/usr/bin/env python3 + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import json import os diff --git a/tests/test_utils/runners/run_functional_tests.sh b/tests/test_utils/runners/run_functional_tests.sh index 4286b9e8e3..6f01fb8c1e 100755 --- a/tests/test_utils/runners/run_functional_tests.sh +++ b/tests/test_utils/runners/run_functional_tests.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Functional Test Runner set -eo pipefail diff --git a/tests/test_utils/runners/run_tests.sh b/tests/test_utils/runners/run_tests.sh index 19d8fd48be..64c4f0da2c 100755 --- a/tests/test_utils/runners/run_tests.sh +++ b/tests/test_utils/runners/run_tests.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # FlagScale Test Runner - Unified entry point for all tests set -eo pipefail diff --git a/tests/test_utils/runners/run_unit_tests.sh b/tests/test_utils/runners/run_unit_tests.sh index d2c36d3a8c..8feaf00a21 100755 --- a/tests/test_utils/runners/run_unit_tests.sh +++ b/tests/test_utils/runners/run_unit_tests.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Unit Test Runner set -eo pipefail diff --git a/tests/test_utils/runners/setup_training_test_env.sh b/tests/test_utils/runners/setup_training_test_env.sh index 2424f9674f..81cd6f71a4 100755 --- a/tests/test_utils/runners/setup_training_test_env.sh +++ b/tests/test_utils/runners/setup_training_test_env.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Prepare the Python/runtime environment for training tests. set -euo pipefail diff --git a/tests/test_utils/runners/utils.sh b/tests/test_utils/runners/utils.sh index 3b20370146..b086da4923 100755 --- a/tests/test_utils/runners/utils.sh +++ b/tests/test_utils/runners/utils.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Shared Utility Functions for FlagScale Test Runners # Logging diff --git a/tests/unit_tests/__init__.py b/tests/unit_tests/__init__.py index 9add4932b0..7232a9cce8 100644 --- a/tests/unit_tests/__init__.py +++ b/tests/unit_tests/__init__.py @@ -1 +1,15 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # This file makes tests/unit_tests/ a Python package diff --git a/tests/unit_tests/models/vla/test_protocols.py b/tests/unit_tests/models/vla/test_protocols.py index edcd3a51cd..29cf65bc9b 100644 --- a/tests/unit_tests/models/vla/test_protocols.py +++ b/tests/unit_tests/models/vla/test_protocols.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import unittest import torch diff --git a/tests/unit_tests/models/vla/test_qwen_gr00t_config.py b/tests/unit_tests/models/vla/test_qwen_gr00t_config.py index 295f0e4810..1fbcf8a612 100644 --- a/tests/unit_tests/models/vla/test_qwen_gr00t_config.py +++ b/tests/unit_tests/models/vla/test_qwen_gr00t_config.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Unit tests for QwenGr00tConfig and PreTrainedConfig.""" import json diff --git a/tests/unit_tests/models/vla/test_registry.py b/tests/unit_tests/models/vla/test_registry.py index 23d9c82583..daecb5859e 100644 --- a/tests/unit_tests/models/vla/test_registry.py +++ b/tests/unit_tests/models/vla/test_registry.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import unittest from flagscale.models.vla.registry import ( diff --git a/tests/unit_tests/models/vla/test_utils.py b/tests/unit_tests/models/vla/test_utils.py index 0ff384d481..bc92d8efa0 100644 --- a/tests/unit_tests/models/vla/test_utils.py +++ b/tests/unit_tests/models/vla/test_utils.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import unittest from flagscale.models.vla.utils import get_vlm_config diff --git a/tests/unit_tests/models/vla/vlm/test_qwenvl_backbone.py b/tests/unit_tests/models/vla/vlm/test_qwenvl_backbone.py index 82fee6575c..99ff7e5e8a 100644 --- a/tests/unit_tests/models/vla/vlm/test_qwenvl_backbone.py +++ b/tests/unit_tests/models/vla/vlm/test_qwenvl_backbone.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import unittest from flagscale.models.vla.registry import VLM_REGISTRY diff --git a/tests/unit_tests/models/vla/vlm/test_vlm_init.py b/tests/unit_tests/models/vla/vlm/test_vlm_init.py index 61087f33f0..f8eb9d6cbf 100644 --- a/tests/unit_tests/models/vla/vlm/test_vlm_init.py +++ b/tests/unit_tests/models/vla/vlm/test_vlm_init.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import unittest diff --git a/tests/unit_tests/runner/elastic/test_diagnostic.py b/tests/unit_tests/runner/elastic/test_diagnostic.py index 178947487d..5cda4a922d 100644 --- a/tests/unit_tests/runner/elastic/test_diagnostic.py +++ b/tests/unit_tests/runner/elastic/test_diagnostic.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import tempfile from unittest.mock import MagicMock, patch diff --git a/tests/unit_tests/runner/elastic/test_gpu_health_check.py b/tests/unit_tests/runner/elastic/test_gpu_health_check.py index 6aed8dc121..bb3b43b746 100644 --- a/tests/unit_tests/runner/elastic/test_gpu_health_check.py +++ b/tests/unit_tests/runner/elastic/test_gpu_health_check.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os from unittest.mock import MagicMock, patch diff --git a/tests/unit_tests/runner/elastic/test_log_collector.py b/tests/unit_tests/runner/elastic/test_log_collector.py index 898276f080..027ea25aff 100644 --- a/tests/unit_tests/runner/elastic/test_log_collector.py +++ b/tests/unit_tests/runner/elastic/test_log_collector.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import tempfile from unittest.mock import patch diff --git a/tests/unit_tests/runner/straggler/__init__.py b/tests/unit_tests/runner/straggler/__init__.py index 79fdc150be..d4f0bee0d3 100644 --- a/tests/unit_tests/runner/straggler/__init__.py +++ b/tests/unit_tests/runner/straggler/__init__.py @@ -1 +1,15 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Tests for FlagScale straggler utilities.""" diff --git a/tests/unit_tests/runner/straggler/test_config.py b/tests/unit_tests/runner/straggler/test_config.py index e5d9372730..07f40cda46 100644 --- a/tests/unit_tests/runner/straggler/test_config.py +++ b/tests/unit_tests/runner/straggler/test_config.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Unit tests for straggler configuration.""" from flagscale.runner.straggler.config import StragglerConfig diff --git a/tests/unit_tests/runner/straggler/test_detector.py b/tests/unit_tests/runner/straggler/test_detector.py index 25fcd1b249..754a9b15db 100644 --- a/tests/unit_tests/runner/straggler/test_detector.py +++ b/tests/unit_tests/runner/straggler/test_detector.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Unit tests for straggler detector.""" import json diff --git a/tests/unit_tests/runner/straggler/test_healthcheck.py b/tests/unit_tests/runner/straggler/test_healthcheck.py index 54fc3a08ca..b75a2aca9a 100644 --- a/tests/unit_tests/runner/straggler/test_healthcheck.py +++ b/tests/unit_tests/runner/straggler/test_healthcheck.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Unit tests for straggler health checks.""" import os diff --git a/tests/unit_tests/runner/straggler/test_report.py b/tests/unit_tests/runner/straggler/test_report.py index 02f947575b..5bc7caa790 100644 --- a/tests/unit_tests/runner/straggler/test_report.py +++ b/tests/unit_tests/runner/straggler/test_report.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Unit tests for straggler reports.""" import json diff --git a/tests/unit_tests/runner/straggler/test_section.py b/tests/unit_tests/runner/straggler/test_section.py index 5d25ac3a69..10e0c0448f 100644 --- a/tests/unit_tests/runner/straggler/test_section.py +++ b/tests/unit_tests/runner/straggler/test_section.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Unit tests for section profiling helpers.""" import time diff --git a/tests/unit_tests/runner/test_backend_megatron_perf.py b/tests/unit_tests/runner/test_backend_megatron_perf.py index 3e051accfe..b0226c628a 100644 --- a/tests/unit_tests/runner/test_backend_megatron_perf.py +++ b/tests/unit_tests/runner/test_backend_megatron_perf.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import sys import tempfile diff --git a/tests/unit_tests/runner/test_backend_native_train.py b/tests/unit_tests/runner/test_backend_native_train.py index 7dd6e3c18d..7ee027d001 100644 --- a/tests/unit_tests/runner/test_backend_native_train.py +++ b/tests/unit_tests/runner/test_backend_native_train.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import tempfile from unittest.mock import patch diff --git a/tests/unit_tests/runner/test_config_validation.py b/tests/unit_tests/runner/test_config_validation.py index a112ae6202..b4db870ee4 100644 --- a/tests/unit_tests/runner/test_config_validation.py +++ b/tests/unit_tests/runner/test_config_validation.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Tests for config error wrapping in flagscale/run.py and flagscale/cli.py. Verifies that OmegaConf errors are caught and re-raised with detailed context diff --git a/tests/unit_tests/runner/test_launcher_ssh_perf.py b/tests/unit_tests/runner/test_launcher_ssh_perf.py index a3ad7521fd..a645234256 100644 --- a/tests/unit_tests/runner/test_launcher_ssh_perf.py +++ b/tests/unit_tests/runner/test_launcher_ssh_perf.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from omegaconf import OmegaConf from flagscale.runner.launcher.launcher_ssh import _get_runner_cmd_train diff --git a/tests/unit_tests/runner/test_parse_hostfile.py b/tests/unit_tests/runner/test_parse_hostfile.py index 8f2e16b818..1979d1fc95 100644 --- a/tests/unit_tests/runner/test_parse_hostfile.py +++ b/tests/unit_tests/runner/test_parse_hostfile.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import pytest from flagscale.runner.utils import parse_hostfile diff --git a/tests/unit_tests/runner/test_path_utils.py b/tests/unit_tests/runner/test_path_utils.py index ed0125b8d4..a34ae73424 100644 --- a/tests/unit_tests/runner/test_path_utils.py +++ b/tests/unit_tests/runner/test_path_utils.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import pytest diff --git a/tests/unit_tests/runner/test_runner_train_straggler.py b/tests/unit_tests/runner/test_runner_train_straggler.py index e045ceac95..1363cffbdb 100644 --- a/tests/unit_tests/runner/test_runner_train_straggler.py +++ b/tests/unit_tests/runner/test_runner_train_straggler.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import sys import types diff --git a/tests/unit_tests/serve/processor/test_image_resize_processor.py b/tests/unit_tests/serve/processor/test_image_resize_processor.py index b8a8dfaeb7..5308d6031e 100644 --- a/tests/unit_tests/serve/processor/test_image_resize_processor.py +++ b/tests/unit_tests/serve/processor/test_image_resize_processor.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import numpy as np import pytest diff --git a/tests/unit_tests/test_basic.py b/tests/unit_tests/test_basic.py index 82959decdc..2f9c3326c4 100644 --- a/tests/unit_tests/test_basic.py +++ b/tests/unit_tests/test_basic.py @@ -1,2 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + def test_import(): pass diff --git a/tests/unit_tests/test_cli.py b/tests/unit_tests/test_cli.py index ee0f486bce..08099a44c8 100644 --- a/tests/unit_tests/test_cli.py +++ b/tests/unit_tests/test_cli.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import sys from pathlib import Path from unittest.mock import patch diff --git a/tests/unit_tests/test_init.py b/tests/unit_tests/test_init.py index b672bd1c05..7f032e1a41 100644 --- a/tests/unit_tests/test_init.py +++ b/tests/unit_tests/test_init.py @@ -1,3 +1,18 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + def test_get_version_from_importlib_metadata(mocker): """Test version retrieval from importlib.metadata""" mocker.patch("importlib.metadata.version", return_value="1.2.3") diff --git a/tests/unit_tests/test_parallel_context.py b/tests/unit_tests/test_parallel_context.py index 5240f3eff1..22b05e3793 100644 --- a/tests/unit_tests/test_parallel_context.py +++ b/tests/unit_tests/test_parallel_context.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import torch import megatron.training.global_vars as mcore_global_vars diff --git a/tests/unit_tests/test_run.py b/tests/unit_tests/test_run.py index fd81459c80..e92974a485 100644 --- a/tests/unit_tests/test_run.py +++ b/tests/unit_tests/test_run.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import pytest from omegaconf import OmegaConf diff --git a/tests/unit_tests/test_setup.py b/tests/unit_tests/test_setup.py index ba1ac6a009..59bf93d49d 100644 --- a/tests/unit_tests/test_setup.py +++ b/tests/unit_tests/test_setup.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import sys import unittest.mock diff --git a/tests/unit_tests/test_spiky_loss_detector.py b/tests/unit_tests/test_spiky_loss_detector.py index 929741024b..6678dfe4f7 100644 --- a/tests/unit_tests/test_spiky_loss_detector.py +++ b/tests/unit_tests/test_spiky_loss_detector.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import torch from megatron.training.spiky_loss import SpikyLossDetector diff --git a/tests/unit_tests/test_utilities.py b/tests/unit_tests/test_utilities.py index eb76795f30..a8cb8f290b 100644 --- a/tests/unit_tests/test_utilities.py +++ b/tests/unit_tests/test_utilities.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os from datetime import timedelta diff --git a/tests/unit_tests/train/megatron/test_train_robobrain_x0.py b/tests/unit_tests/train/megatron/test_train_robobrain_x0.py index 7620bd4622..9bbe7b43d4 100644 --- a/tests/unit_tests/train/megatron/test_train_robobrain_x0.py +++ b/tests/unit_tests/train/megatron/test_train_robobrain_x0.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """ Unit tests for train_robobrain_x0.py dataloader consistency. diff --git a/tests/unit_tests/train/test_perf_monitor.py b/tests/unit_tests/train/test_perf_monitor.py index b149f5484c..6978e97a8f 100644 --- a/tests/unit_tests/train/test_perf_monitor.py +++ b/tests/unit_tests/train/test_perf_monitor.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import json import time from pathlib import Path diff --git a/tests/unit_tests/train/test_train_config.py b/tests/unit_tests/train/test_train_config.py index 5cf7cd0fac..cb4f4e3562 100644 --- a/tests/unit_tests/train/test_train_config.py +++ b/tests/unit_tests/train/test_train_config.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import unittest from omegaconf import OmegaConf diff --git a/tests/unit_tests/train/utils/test_optim_setup.py b/tests/unit_tests/train/utils/test_optim_setup.py index 437f1a1f5a..0f23e93252 100644 --- a/tests/unit_tests/train/utils/test_optim_setup.py +++ b/tests/unit_tests/train/utils/test_optim_setup.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Unit tests for optimizer setup utilities.""" import unittest diff --git a/tests/unit_tests/train/utils/test_train_utils.py b/tests/unit_tests/train/utils/test_train_utils.py index 4cd5dd6f15..6796edd9e6 100644 --- a/tests/unit_tests/train/utils/test_train_utils.py +++ b/tests/unit_tests/train/utils/test_train_utils.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Non-FSDP2 tests adapted from # https://github.com/huggingface/lerobot/blob/2b304eeb841ae6c371e3dd341bbbb9dd254b07cb/tests/utils/test_train_utils.py diff --git a/tests/unit_tests/transformations/diffusion/test_taylorseer_transformation.py b/tests/unit_tests/transformations/diffusion/test_taylorseer_transformation.py index 1b4113aa43..33e6982f3f 100644 --- a/tests/unit_tests/transformations/diffusion/test_taylorseer_transformation.py +++ b/tests/unit_tests/transformations/diffusion/test_taylorseer_transformation.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import unittest import torch diff --git a/tests/unit_tests/transformations/diffusion/test_timestep_tracker_transformation.py b/tests/unit_tests/transformations/diffusion/test_timestep_tracker_transformation.py index 8158816852..f9fd56ea55 100644 --- a/tests/unit_tests/transformations/diffusion/test_timestep_tracker_transformation.py +++ b/tests/unit_tests/transformations/diffusion/test_timestep_tracker_transformation.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import unittest import torch diff --git a/tests/unit_tests/transformations/test_hook.py b/tests/unit_tests/transformations/test_hook.py index ed00f8db5e..3b665656a0 100644 --- a/tests/unit_tests/transformations/test_hook.py +++ b/tests/unit_tests/transformations/test_hook.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Modified from https://github.com/huggingface/diffusers/blob/main/tests/hooks/test_hooks.py import gc diff --git a/tests/unit_tests/transformations/test_state_scope_transform.py b/tests/unit_tests/transformations/test_state_scope_transform.py index af79a82354..a9ca9ba4a3 100644 --- a/tests/unit_tests/transformations/test_state_scope_transform.py +++ b/tests/unit_tests/transformations/test_state_scope_transform.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import unittest from typing import Any diff --git a/tests/unit_tests/transformations/test_state_store.py b/tests/unit_tests/transformations/test_state_store.py index 721e564aaa..2f54f2a9dd 100644 --- a/tests/unit_tests/transformations/test_state_store.py +++ b/tests/unit_tests/transformations/test_state_store.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import unittest from flagscale.transformations.state_store import StateStore diff --git a/tools/__init__.py b/tools/__init__.py index e69de29bb2..a491ed2726 100644 --- a/tools/__init__.py +++ b/tools/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tools/checkpoint/aquila/args.py b/tools/checkpoint/aquila/args.py index 08e22d824d..8a47b06dc2 100644 --- a/tools/checkpoint/aquila/args.py +++ b/tools/checkpoint/aquila/args.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import json import os diff --git a/tools/checkpoint/aquila/ckpt.py b/tools/checkpoint/aquila/ckpt.py index 754f831b9b..425ba6f0fc 100644 --- a/tools/checkpoint/aquila/ckpt.py +++ b/tools/checkpoint/aquila/ckpt.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import sys import torch diff --git a/tools/checkpoint/aquila/model.py b/tools/checkpoint/aquila/model.py index 558ba73d60..893b4f9feb 100644 --- a/tools/checkpoint/aquila/model.py +++ b/tools/checkpoint/aquila/model.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import time from megatron.core.enums import ModelType diff --git a/tools/checkpoint/convert.py b/tools/checkpoint/convert.py index 10fc55d1b2..4e4c3fe7f8 100644 --- a/tools/checkpoint/convert.py +++ b/tools/checkpoint/convert.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import copy import importlib diff --git a/tools/checkpoint/deepseek_v3/args.py b/tools/checkpoint/deepseek_v3/args.py index f0f2fc0468..7e0b4e8670 100644 --- a/tools/checkpoint/deepseek_v3/args.py +++ b/tools/checkpoint/deepseek_v3/args.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import json import os import sys diff --git a/tools/checkpoint/deepseek_v3/ckpt.py b/tools/checkpoint/deepseek_v3/ckpt.py index bc63843a11..45eb618477 100644 --- a/tools/checkpoint/deepseek_v3/ckpt.py +++ b/tools/checkpoint/deepseek_v3/ckpt.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import sys import torch diff --git a/tools/checkpoint/deepseek_v3/model.py b/tools/checkpoint/deepseek_v3/model.py index ffecc84233..567cf423a8 100644 --- a/tools/checkpoint/deepseek_v3/model.py +++ b/tools/checkpoint/deepseek_v3/model.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import time import torch diff --git a/tools/checkpoint/deepseek_v3/moonlight_deepseek/configuration_deepseek.py b/tools/checkpoint/deepseek_v3/moonlight_deepseek/configuration_deepseek.py index 30e20568b4..e07f6a1f1e 100644 --- a/tools/checkpoint/deepseek_v3/moonlight_deepseek/configuration_deepseek.py +++ b/tools/checkpoint/deepseek_v3/moonlight_deepseek/configuration_deepseek.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Copy from https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/configuration_deepseek.py from transformers.configuration_utils import PretrainedConfig diff --git a/tools/checkpoint/llama/args.py b/tools/checkpoint/llama/args.py index 873b0f72ca..a7cea8c9f0 100644 --- a/tools/checkpoint/llama/args.py +++ b/tools/checkpoint/llama/args.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import json import os diff --git a/tools/checkpoint/llama/ckpt.py b/tools/checkpoint/llama/ckpt.py index 754f831b9b..425ba6f0fc 100644 --- a/tools/checkpoint/llama/ckpt.py +++ b/tools/checkpoint/llama/ckpt.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import sys import torch diff --git a/tools/checkpoint/llama/model.py b/tools/checkpoint/llama/model.py index 0a9e02bf3c..3480651c4d 100644 --- a/tools/checkpoint/llama/model.py +++ b/tools/checkpoint/llama/model.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import sys sys.path.append("..") diff --git a/tools/checkpoint/loader_mcore.py b/tools/checkpoint/loader_mcore.py index dadd83b1d0..76805f9885 100644 --- a/tools/checkpoint/loader_mcore.py +++ b/tools/checkpoint/loader_mcore.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import sys import json @@ -10,30 +24,35 @@ def add_arguments(parser): - group = parser.add_argument_group(title='Megatron loader') + group = parser.add_argument_group(title="Megatron loader") - group.add_argument('--true-vocab-size', type=int, default=None, - help='original size of vocab, if specified will trim padding from embedding table.') - group.add_argument('--megatron-path', type=str, default=None, - help='Base directory of megatron repository') - group.add_argument('--position-embedding-type', - type=str, - default='learned_absolute', - choices=['learned_absolute', 'rope'], - help='Position embedding type.') + group.add_argument( + "--true-vocab-size", + type=int, + default=None, + help="original size of vocab, if specified will trim padding from embedding table.", + ) + group.add_argument( + "--megatron-path", type=str, default=None, help="Base directory of megatron repository" + ) + group.add_argument( + "--position-embedding-type", + type=str, + default="learned_absolute", + choices=["learned_absolute", "rope"], + help="Position embedding type.", + ) def _load_checkpoint(queue, args): - """ prepare import module """ # Search in directory above this root_path = os.path.abspath( - os.path.join(os.path.dirname(__file__), - os.path.pardir, - os.path.pardir)) + os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir) + ) sys.path.insert(0, root_path) sys.path.insert(0, os.path.join(root_path, "flagscale/train")) @@ -47,9 +66,11 @@ def _load_checkpoint(queue, args): from megatron.legacy.model import module from megatron.core import mpu from megatron.core.tensor_parallel.random import ( - get_cuda_rng_tracker, _DATA_PARALLEL_RNG_TRACKER_NAME, - _EXPERT_PARALLEL_RNG_TRACKER_NAME, _MODEL_PARALLEL_RNG_TRACKER_NAME - ) + get_cuda_rng_tracker, + _DATA_PARALLEL_RNG_TRACKER_NAME, + _EXPERT_PARALLEL_RNG_TRACKER_NAME, + _MODEL_PARALLEL_RNG_TRACKER_NAME, + ) from tools.checkpoint.utils import ( _ConverterFakeProcessGroup, get_expert_model_parallel_rank, @@ -59,7 +80,9 @@ def _load_checkpoint(queue, args): validate_mcore_parallel_size, ) except ModuleNotFoundError: - print("Unable to import Megatron, please specify the path to Megatron using --megatron-path. Exiting.") + print( + "Unable to import Megatron, please specify the path to Megatron using --megatron-path. Exiting." + ) queue.put("exit") exit(1) @@ -80,23 +103,26 @@ def queue_put(name, msg): # We want all arguments to come from us. sys.argv = [ - 'script.py', - '--no-masked-softmax-fusion', - '--no-bias-gelu-fusion', - '--no-bias-dropout-fusion', - '--use-cpu-initialization', - '--micro-batch-size', '1', - '--no-load-optim', - '--no-load-rng', - '--no-save-optim', - '--no-save-rng', - '--no-initialization', - '--mock-data', # To pass the "blend data checks" in arguments.py - '--transformer-impl', 'transformer_engine', - '--load', args.load_dir, - '--exit-on-missing-checkpoint', - '--use-mp-args-from-checkpoint-args', - '--no-one-logger', + "script.py", + "--no-masked-softmax-fusion", + "--no-bias-gelu-fusion", + "--no-bias-dropout-fusion", + "--use-cpu-initialization", + "--micro-batch-size", + "1", + "--no-load-optim", + "--no-load-rng", + "--no-save-optim", + "--no-save-rng", + "--no-initialization", + "--mock-data", # To pass the "blend data checks" in arguments.py + "--transformer-impl", + "transformer_engine", + "--load", + args.load_dir, + "--exit-on-missing-checkpoint", + "--use-mp-args-from-checkpoint-args", + "--no-one-logger", ] margs = parse_args() @@ -174,13 +200,13 @@ def _set_arg(arg_name): os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1" # Layernorm has bias; RMSNorm does not. - if hasattr(checkpoint_args, 'normalization'): + if hasattr(checkpoint_args, "normalization"): margs.norm_has_bias = checkpoint_args.normalization == "LayerNorm" else: # older models only supported LayerNorm margs.norm_has_bias = True - print("*"*20 + "validate loader arguments" + "*"*20) + print("*" * 20 + "validate loader arguments" + "*" * 20) margs = validate_args(margs) validate_mcore_parallel_size(margs) @@ -194,22 +220,22 @@ def check_for_arg(arg_name, default=None): queue.put("exit") exit(1) - check_for_arg('tensor_model_parallel_size') - check_for_arg('pipeline_model_parallel_size') - check_for_arg('expert_model_parallel_size') - check_for_arg('num_layers') - check_for_arg('hidden_size') - check_for_arg('seq_length') - check_for_arg('num_attention_heads') - check_for_arg('max_position_embeddings') - check_for_arg('position_embedding_type') - check_for_arg('tokenizer_type') - check_for_arg('iteration') - check_for_arg('bert_binary_head') - check_for_arg('params_dtype') - check_for_arg('swiglu', False) - check_for_arg('disable_bias_linear', not getattr(margs, "add_bias_linear", False)) - check_for_arg('add_qkv_bias', getattr(margs, "add_bias_linear_qkv", False)) + check_for_arg("tensor_model_parallel_size") + check_for_arg("pipeline_model_parallel_size") + check_for_arg("expert_model_parallel_size") + check_for_arg("num_layers") + check_for_arg("hidden_size") + check_for_arg("seq_length") + check_for_arg("num_attention_heads") + check_for_arg("max_position_embeddings") + check_for_arg("position_embedding_type") + check_for_arg("tokenizer_type") + check_for_arg("iteration") + check_for_arg("bert_binary_head") + check_for_arg("params_dtype") + check_for_arg("swiglu", False) + check_for_arg("disable_bias_linear", not getattr(margs, "add_bias_linear", False)) + check_for_arg("add_qkv_bias", getattr(margs, "add_bias_linear_qkv", False)) # Determine how to make our models. margs.model_type = model_plugin.model_type @@ -265,7 +291,9 @@ def check_for_arg(arg_name, default=None): edp_parallel_size = mcore_model_parallel_size // (etp_size * ep_size) fake_edp_group = _ConverterFakeProcessGroup(size=edp_parallel_size) fake_etp_ep_group = _ConverterFakeProcessGroup(size=etp_size * ep_size) - fake_tcp_group = _ConverterFakeProcessGroup(size=margs.tensor_model_parallel_size*margs.context_parallel_size) + fake_tcp_group = _ConverterFakeProcessGroup( + size=margs.tensor_model_parallel_size * margs.context_parallel_size + ) mpu._PIPELINE_MODEL_PARALLEL_GROUP = fake_pp_group mpu._CONTEXT_PARALLEL_GROUP = fake_cp_group mpu._DATA_PARALLEL_GROUP = fake_dp_group @@ -309,12 +337,13 @@ def check_for_arg(arg_name, default=None): md.previous_expert_parallel_size = margs.expert_model_parallel_size md.previous_expert_tensor_parallel_size = margs.expert_tensor_parallel_size md.previous_decoder_first_pipeline_num_layers = margs.decoder_first_pipeline_num_layers - md.true_vocab_size = args.true_vocab_size # true (non-padded) vocab size + md.true_vocab_size = args.true_vocab_size # true (non-padded) vocab size md.make_vocab_size_divisible_by = margs.make_vocab_size_divisible_by md.checkpoint_args = checkpoint_args consumed_train_samples = None consumed_valid_samples = None + def get_models(count, dtype): # for one pp stage nonlocal consumed_train_samples @@ -355,12 +384,12 @@ def get_models(count, dtype): load_checkpoint(model_, None, None) if consumed_train_samples is not None: - assert(margs.consumed_train_samples == consumed_train_samples) + assert margs.consumed_train_samples == consumed_train_samples else: consumed_train_samples = margs.consumed_train_samples if consumed_valid_samples is not None: - assert(margs.consumed_valid_samples == consumed_valid_samples) + assert margs.consumed_valid_samples == consumed_valid_samples else: consumed_valid_samples = margs.consumed_valid_samples @@ -377,7 +406,7 @@ def get_models(count, dtype): fake_pp_group = _ConverterFakeProcessGroup(rank=0, size=pp_size) mpu._PIPELINE_MODEL_PARALLEL_GROUP = fake_pp_group all_models = [get_models(mcore_model_parallel_size, margs.params_dtype)] - models = all_models[0][0] # pp0vpp0 + models = all_models[0][0] # pp0vpp0 md.consumed_train_samples = consumed_train_samples md.consumed_valid_samples = consumed_valid_samples @@ -409,7 +438,7 @@ def get_models(count, dtype): message = dict() margs.total_layer_num = total_layer_num - engram_layer_id = total_layer_num # get_global_layer_id + engram_layer_id = total_layer_num # get_global_layer_id if getattr(margs, "use_engram", False) and engram_layer_id in getattr( margs, "engram_layer_ids", [] ): diff --git a/tools/checkpoint/loader_transformers.py b/tools/checkpoint/loader_transformers.py index bc7e86b93b..6b5e1f5fe5 100644 --- a/tools/checkpoint/loader_transformers.py +++ b/tools/checkpoint/loader_transformers.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import sys import types @@ -7,37 +21,43 @@ def add_arguments(parser): - group = parser.add_argument_group(title='Transformers loader') - - group.add_argument('--true-vocab-size', type=int, default=None, - help='original size of vocab, if specified will trim padding from embedding table.') - group.add_argument('--megatron-path', type=str, default=None, - help='Base directory of megatron repository') - group.add_argument('--position-embedding-type', - type=str, - default='learned_absolute', - choices=['learned_absolute', 'rope'], - help='Position embedding type.') + group = parser.add_argument_group(title="Transformers loader") + + group.add_argument( + "--true-vocab-size", + type=int, + default=None, + help="original size of vocab, if specified will trim padding from embedding table.", + ) + group.add_argument( + "--megatron-path", type=str, default=None, help="Base directory of megatron repository" + ) + group.add_argument( + "--position-embedding-type", + type=str, + default="learned_absolute", + choices=["learned_absolute", "rope"], + help="Position embedding type.", + ) def _load_checkpoint(queue, args): - """ prepare import module """ try: import transformers - major, minor, _ = map(int, transformers.__version__.split('.')) + + major, minor, _ = map(int, transformers.__version__.split(".")) assert major >= 4 and minor >= 36 except: raise ImportError("transformers version >= 4.36.0 ") # Search in directory above this. root_path = os.path.abspath( - os.path.join(os.path.dirname(__file__), - os.path.pardir, - os.path.pardir)) + os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir) + ) sys.path.insert(0, root_path) sys.path.insert(0, os.path.join(root_path, "flagscale/train")) @@ -47,7 +67,9 @@ def _load_checkpoint(queue, args): try: from megatron.training.arguments import parse_args, validate_args except ModuleNotFoundError: - print("Unable to import Megatron, please specify the path to Megatron using --megatron-path. Exiting.") + print( + "Unable to import Megatron, please specify the path to Megatron using --megatron-path. Exiting." + ) queue.put("exit") exit(1) @@ -69,32 +91,39 @@ def queue_put(name, msg): # We want all arguments to come from us. sys.argv = [ - 'script.py', - '--no-masked-softmax-fusion', - '--no-bias-gelu-fusion', - '--no-bias-dropout-fusion', - '--use-cpu-initialization', - '--micro-batch-size', '1', - '--no-load-optim', - '--no-load-rng', - '--no-save-optim', - '--no-save-rng', - '--no-initialization', - '--mock-data', # To pass the "blend data checks" in arguments.py - '--use-mcore-models', - '--transformer-impl', 'transformer_engine', - '--load', args.load_dir, - '--exit-on-missing-checkpoint', - '--use-mp-args-from-checkpoint-args', - '--no-one-logger', + "script.py", + "--no-masked-softmax-fusion", + "--no-bias-gelu-fusion", + "--no-bias-dropout-fusion", + "--use-cpu-initialization", + "--micro-batch-size", + "1", + "--no-load-optim", + "--no-load-rng", + "--no-save-optim", + "--no-save-rng", + "--no-initialization", + "--mock-data", # To pass the "blend data checks" in arguments.py + "--use-mcore-models", + "--transformer-impl", + "transformer_engine", + "--load", + args.load_dir, + "--exit-on-missing-checkpoint", + "--use-mp-args-from-checkpoint-args", + "--no-one-logger", ] margs = parse_args() args_plugin.load_args_hf2mg(margs) - print("*"*20 + "validate loader arguments" + "*"*20) + print("*" * 20 + "validate loader arguments" + "*" * 20) margs = validate_args(margs) - margs.world_size = margs.tensor_model_parallel_size * margs.pipeline_model_parallel_size * margs.expert_model_parallel_size + margs.world_size = ( + margs.tensor_model_parallel_size + * margs.pipeline_model_parallel_size + * margs.expert_model_parallel_size + ) def check_for_arg(arg_name, default=None): if getattr(margs, arg_name, None) is None: @@ -106,23 +135,23 @@ def check_for_arg(arg_name, default=None): queue.put("exit") exit(1) - check_for_arg('tensor_model_parallel_size') - check_for_arg('pipeline_model_parallel_size') - check_for_arg('expert_model_parallel_size') - check_for_arg('num_layers') - check_for_arg('hidden_size') - check_for_arg('seq_length') - check_for_arg('num_attention_heads') - check_for_arg('max_position_embeddings') - check_for_arg('position_embedding_type') - check_for_arg('tokenizer_type') - check_for_arg('iteration') - check_for_arg('bert_binary_head') - check_for_arg('params_dtype') - check_for_arg('swiglu', False) - check_for_arg('disable_bias_linear', not getattr(margs, "add_bias_linear", False)) - check_for_arg('add_qkv_bias', getattr(margs, "add_bias_linear_qkv", False)) - check_for_arg('qk_layernorm', getattr(margs, "qk_layernorm", False)) + check_for_arg("tensor_model_parallel_size") + check_for_arg("pipeline_model_parallel_size") + check_for_arg("expert_model_parallel_size") + check_for_arg("num_layers") + check_for_arg("hidden_size") + check_for_arg("seq_length") + check_for_arg("num_attention_heads") + check_for_arg("max_position_embeddings") + check_for_arg("position_embedding_type") + check_for_arg("tokenizer_type") + check_for_arg("iteration") + check_for_arg("bert_binary_head") + check_for_arg("params_dtype") + check_for_arg("swiglu", False) + check_for_arg("disable_bias_linear", not getattr(margs, "add_bias_linear", False)) + check_for_arg("add_qkv_bias", getattr(margs, "add_bias_linear_qkv", False)) + check_for_arg("qk_layernorm", getattr(margs, "qk_layernorm", False)) """ use megatron args build object and init env @@ -169,7 +198,7 @@ def check_for_arg(arg_name, default=None): # Send embeddings. message = dict() - message["word embeddings"]= hf_model.model.embed_tokens.weight + message["word embeddings"] = hf_model.model.embed_tokens.weight queue_put("embeddings", message) # Send transformer layers diff --git a/tools/checkpoint/mistral/args.py b/tools/checkpoint/mistral/args.py index dd3de55349..6ed697e52e 100644 --- a/tools/checkpoint/mistral/args.py +++ b/tools/checkpoint/mistral/args.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import json import os diff --git a/tools/checkpoint/mistral/ckpt.py b/tools/checkpoint/mistral/ckpt.py index 5514640acf..e2b5a46337 100644 --- a/tools/checkpoint/mistral/ckpt.py +++ b/tools/checkpoint/mistral/ckpt.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import sys import torch diff --git a/tools/checkpoint/mistral/model.py b/tools/checkpoint/mistral/model.py index 0a9e02bf3c..3480651c4d 100644 --- a/tools/checkpoint/mistral/model.py +++ b/tools/checkpoint/mistral/model.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import sys sys.path.append("..") diff --git a/tools/checkpoint/mixtral/args.py b/tools/checkpoint/mixtral/args.py index 3404878fc4..c7053f872f 100644 --- a/tools/checkpoint/mixtral/args.py +++ b/tools/checkpoint/mixtral/args.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import json import os diff --git a/tools/checkpoint/mixtral/ckpt.py b/tools/checkpoint/mixtral/ckpt.py index 5f7b18e717..b019bc978f 100644 --- a/tools/checkpoint/mixtral/ckpt.py +++ b/tools/checkpoint/mixtral/ckpt.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import sys import torch diff --git a/tools/checkpoint/mixtral/model.py b/tools/checkpoint/mixtral/model.py index e6131c244c..05e90e30cc 100644 --- a/tools/checkpoint/mixtral/model.py +++ b/tools/checkpoint/mixtral/model.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import time from megatron.core.enums import ModelType diff --git a/tools/checkpoint/qwen2_5_vl/convert.md b/tools/checkpoint/qwen2_5_vl/convert.md index 79b08afce7..5508f752e8 100644 --- a/tools/checkpoint/qwen2_5_vl/convert.md +++ b/tools/checkpoint/qwen2_5_vl/convert.md @@ -1,3 +1,19 @@ + + ### 📎 Reference Mainly based on official [Pai-Megatron-Patch](https://github.com/alibaba/Pai-Megatron-Patch/tree/main/examples/qwen2_5_vl),with necessary modifications for integration into the current training framework. diff --git a/tools/checkpoint/qwen2_5_vl/hf2mcore_qwen2.5_vl_convertor.sh b/tools/checkpoint/qwen2_5_vl/hf2mcore_qwen2.5_vl_convertor.sh index e1b3599cce..c8c99a36d4 100644 --- a/tools/checkpoint/qwen2_5_vl/hf2mcore_qwen2.5_vl_convertor.sh +++ b/tools/checkpoint/qwen2_5_vl/hf2mcore_qwen2.5_vl_convertor.sh @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Copied from https://github.com/alibaba/Pai-Megatron-Patch/blob/8949a6647cbf6b39837ad3dd911fa4aa0726895b/toolkits/model_checkpoints_convertor/qwen/hf2mcore_qwen2.5_vl.py #!/bin/bash diff --git a/tools/checkpoint/qwen2_5_vl/run.sh b/tools/checkpoint/qwen2_5_vl/run.sh index 2c7aedfb3f..a80b63dc6e 100644 --- a/tools/checkpoint/qwen2_5_vl/run.sh +++ b/tools/checkpoint/qwen2_5_vl/run.sh @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + export PYTHONPATH=../../../:$PYTHONPATH bash hf2mcore_qwen2.5_vl_convertor.sh 7B \ /share/project/lizhiyu/data/Qwen2.5-VL-7B-Instruct \ diff --git a/tools/checkpoint/qwen2_5_vl/utils.py b/tools/checkpoint/qwen2_5_vl/utils.py index be7f97a5f3..715cbf397a 100644 --- a/tools/checkpoint/qwen2_5_vl/utils.py +++ b/tools/checkpoint/qwen2_5_vl/utils.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Copied from https://github.com/alibaba/Pai-Megatron-Patch/blob/8949a6647cbf6b39837ad3dd911fa4aa0726895b/toolkits/model_checkpoints_convertor/utils/__init__.py. import gc diff --git a/tools/checkpoint/qwen3/args.py b/tools/checkpoint/qwen3/args.py index 3a3cb4f74a..d67a1dd683 100644 --- a/tools/checkpoint/qwen3/args.py +++ b/tools/checkpoint/qwen3/args.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import json import os diff --git a/tools/checkpoint/qwen3/ckpt.py b/tools/checkpoint/qwen3/ckpt.py index 481718718b..e418d564b9 100644 --- a/tools/checkpoint/qwen3/ckpt.py +++ b/tools/checkpoint/qwen3/ckpt.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import sys sys.path.append("..") diff --git a/tools/checkpoint/qwen3/model.py b/tools/checkpoint/qwen3/model.py index ed481e9d82..51df0dc0b9 100644 --- a/tools/checkpoint/qwen3/model.py +++ b/tools/checkpoint/qwen3/model.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import time import torch diff --git a/tools/checkpoint/qwen35/convert_qwen35.py b/tools/checkpoint/qwen35/convert_qwen35.py index 94ef17a633..263966faee 100644 --- a/tools/checkpoint/qwen35/convert_qwen35.py +++ b/tools/checkpoint/qwen35/convert_qwen35.py @@ -1,4 +1,19 @@ #!/usr/bin/env python3 + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Unified Qwen3.5 HF <-> Megatron checkpoint converter. Supports both dense and MoE models. Model type is auto-detected from the diff --git a/tools/checkpoint/qwen35/qwen35/__init__.py b/tools/checkpoint/qwen35/qwen35/__init__.py index a99638de95..39c6401256 100644 --- a/tools/checkpoint/qwen35/qwen35/__init__.py +++ b/tools/checkpoint/qwen35/qwen35/__init__.py @@ -1 +1,15 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Qwen3.5 checkpoint conversion utilities.""" diff --git a/tools/checkpoint/qwen35/qwen35/attention.py b/tools/checkpoint/qwen35/qwen35/attention.py index 831f6f6edb..578bdd521f 100644 --- a/tools/checkpoint/qwen35/qwen35/attention.py +++ b/tools/checkpoint/qwen35/qwen35/attention.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Attention QKV merge/split helpers shared across directions and model types.""" import torch diff --git a/tools/checkpoint/qwen35/qwen35/config.py b/tools/checkpoint/qwen35/qwen35/config.py index 919a9321b4..45ff30237b 100644 --- a/tools/checkpoint/qwen35/qwen35/config.py +++ b/tools/checkpoint/qwen35/qwen35/config.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Unified configuration loader for Qwen3.5 checkpoint conversion.""" import json diff --git a/tools/checkpoint/qwen35/qwen35/constants.py b/tools/checkpoint/qwen35/qwen35/constants.py index 6254015efc..48fb735eae 100644 --- a/tools/checkpoint/qwen35/qwen35/constants.py +++ b/tools/checkpoint/qwen35/qwen35/constants.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Shared constants for Qwen3.5 checkpoint conversion.""" # Layer norm adjustment for zero-centered gamma. diff --git a/tools/checkpoint/qwen35/qwen35/converter.py b/tools/checkpoint/qwen35/qwen35/converter.py index ae91a2cea6..b04eb770da 100644 --- a/tools/checkpoint/qwen35/qwen35/converter.py +++ b/tools/checkpoint/qwen35/qwen35/converter.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Unified dense / MoE Qwen3.5 checkpoint converter.""" import os diff --git a/tools/checkpoint/qwen35/qwen35/gdn.py b/tools/checkpoint/qwen35/qwen35/gdn.py index 3e647ca3d8..76619c6460 100644 --- a/tools/checkpoint/qwen35/qwen35/gdn.py +++ b/tools/checkpoint/qwen35/qwen35/gdn.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """GDN (Gated DeltaNet) layer helpers shared across directions and model types.""" import torch diff --git a/tools/checkpoint/qwen35/qwen35/io.py b/tools/checkpoint/qwen35/qwen35/io.py index 795452e9ac..41cd9a9c64 100644 --- a/tools/checkpoint/qwen35/qwen35/io.py +++ b/tools/checkpoint/qwen35/qwen35/io.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Checkpoint I/O helpers for HF and Megatron formats.""" import json diff --git a/tools/checkpoint/qwen35/qwen35/mlp.py b/tools/checkpoint/qwen35/qwen35/mlp.py index f8346d022f..2a3c4b3880 100644 --- a/tools/checkpoint/qwen35/qwen35/mlp.py +++ b/tools/checkpoint/qwen35/qwen35/mlp.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """MLP conversion helpers (dense + MoE) shared across directions.""" import torch diff --git a/tools/checkpoint/qwen35/qwen35/mtp.py b/tools/checkpoint/qwen35/qwen35/mtp.py index 386f7feb44..a9496c02dd 100644 --- a/tools/checkpoint/qwen35/qwen35/mtp.py +++ b/tools/checkpoint/qwen35/qwen35/mtp.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """MTP (Multi-Token Prediction) conversion helpers (HF <-> Megatron).""" import torch diff --git a/tools/checkpoint/qwen35/qwen35/sharding.py b/tools/checkpoint/qwen35/qwen35/sharding.py index 9ecebdfae9..674d9e985e 100644 --- a/tools/checkpoint/qwen35/qwen35/sharding.py +++ b/tools/checkpoint/qwen35/qwen35/sharding.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """TP/PP/EP sharding helpers for Megatron checkpoints.""" import os diff --git a/tools/checkpoint/qwen35/qwen35/validation.py b/tools/checkpoint/qwen35/qwen35/validation.py index cdcea37189..1a83f8c7fd 100644 --- a/tools/checkpoint/qwen35/qwen35/validation.py +++ b/tools/checkpoint/qwen35/qwen35/validation.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Validation helpers for generated checkpoints.""" import os diff --git a/tools/checkpoint/qwen35/qwen35/vision.py b/tools/checkpoint/qwen35/qwen35/vision.py index 51737e2de3..77991cd027 100644 --- a/tools/checkpoint/qwen35/qwen35/vision.py +++ b/tools/checkpoint/qwen35/qwen35/vision.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Vision model conversion helpers (HF <-> Megatron).""" diff --git a/tools/checkpoint/qwen35/run.sh b/tools/checkpoint/qwen35/run.sh index 4a64a48f81..5c4f3549f1 100755 --- a/tools/checkpoint/qwen35/run.sh +++ b/tools/checkpoint/qwen35/run.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Unified checkpoint conversion entry point for Qwen3.5 # Dispatches to convert_qwen35.py with --direction. # diff --git a/tools/checkpoint/qwen3_engram/args.py b/tools/checkpoint/qwen3_engram/args.py index 5932159456..5a8f373113 100644 --- a/tools/checkpoint/qwen3_engram/args.py +++ b/tools/checkpoint/qwen3_engram/args.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import json import os diff --git a/tools/checkpoint/qwen3_engram/ckpt.py b/tools/checkpoint/qwen3_engram/ckpt.py index e48c76fbef..896ad833e4 100644 --- a/tools/checkpoint/qwen3_engram/ckpt.py +++ b/tools/checkpoint/qwen3_engram/ckpt.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import sys sys.path.append("..") diff --git a/tools/checkpoint/qwen3_engram/model.py b/tools/checkpoint/qwen3_engram/model.py index 64f9142ed8..a830d22e8b 100644 --- a/tools/checkpoint/qwen3_engram/model.py +++ b/tools/checkpoint/qwen3_engram/model.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import time import torch diff --git a/tools/checkpoint/qwen3_engram/modeling_hf/configuration_qwen3engram.py b/tools/checkpoint/qwen3_engram/modeling_hf/configuration_qwen3engram.py index 32b16fad75..105ebbaa8a 100644 --- a/tools/checkpoint/qwen3_engram/modeling_hf/configuration_qwen3engram.py +++ b/tools/checkpoint/qwen3_engram/modeling_hf/configuration_qwen3engram.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from transformers.configuration_utils import PretrainedConfig from transformers.modeling_rope_utils import rope_config_validation from transformers.utils import logging diff --git a/tools/checkpoint/qwen3_engram/modeling_hf/modeling_qwen3engram.py b/tools/checkpoint/qwen3_engram/modeling_hf/modeling_qwen3engram.py index c29cc81365..cabdcc14a9 100644 --- a/tools/checkpoint/qwen3_engram/modeling_hf/modeling_qwen3engram.py +++ b/tools/checkpoint/qwen3_engram/modeling_hf/modeling_qwen3engram.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from __future__ import annotations import math diff --git a/tools/checkpoint/qwen3_vl/convert.md b/tools/checkpoint/qwen3_vl/convert.md index 07b96d524b..50c572f68f 100644 --- a/tools/checkpoint/qwen3_vl/convert.md +++ b/tools/checkpoint/qwen3_vl/convert.md @@ -1,3 +1,19 @@ + + ### 📎 Reference Mainly based on official [Pai-Megatron-Patch](https://github.com/alibaba/Pai-Megatron-Patch/tree/main/examples/qwen3_vl),with necessary modifications for integration into the current training framework. diff --git a/tools/checkpoint/qwen3_vl/hf2mcore_qwen3_vl_convertor.sh b/tools/checkpoint/qwen3_vl/hf2mcore_qwen3_vl_convertor.sh index fa854786d3..469531693b 100644 --- a/tools/checkpoint/qwen3_vl/hf2mcore_qwen3_vl_convertor.sh +++ b/tools/checkpoint/qwen3_vl/hf2mcore_qwen3_vl_convertor.sh @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Copied from https://github.com/alibaba/Pai-Megatron-Patch/blob/8949a6647cbf6b39837ad3dd911fa4aa0726895b/toolkits/model_checkpoints_convertor/qwen/hf2mcore_qwen2.5_vl.py #!/bin/bash diff --git a/tools/checkpoint/qwen3_vl/run.sh b/tools/checkpoint/qwen3_vl/run.sh index 07bc0fb69d..9cb906b1cf 100644 --- a/tools/checkpoint/qwen3_vl/run.sh +++ b/tools/checkpoint/qwen3_vl/run.sh @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + export CUDA_VISIBLE_DEVICES=0 export PYTHONPATH=../../../:$PYTHONPATH # export MP_PP0_LAYERS=30 diff --git a/tools/checkpoint/run.sh b/tools/checkpoint/run.sh index 7b7bfd1e5a..a2fd026235 100644 --- a/tools/checkpoint/run.sh +++ b/tools/checkpoint/run.sh @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # sparse model change the parallel config python convert.py \ diff --git a/tools/checkpoint/rwkv/convert_ckpt.py b/tools/checkpoint/rwkv/convert_ckpt.py index 3a583324fc..ebae0ac500 100644 --- a/tools/checkpoint/rwkv/convert_ckpt.py +++ b/tools/checkpoint/rwkv/convert_ckpt.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import torch diff --git a/tools/checkpoint/rwkv/convert_ckpt.sh b/tools/checkpoint/rwkv/convert_ckpt.sh index 41b5aa126c..b46896985d 100644 --- a/tools/checkpoint/rwkv/convert_ckpt.sh +++ b/tools/checkpoint/rwkv/convert_ckpt.sh @@ -1,2 +1,16 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Usage: python convert_ckpt.py --direction ['ds2mg', 'mg2ds'] python convert_ckpt.py deepspeed_model.pt megatron_model.pth --direction ds2mg diff --git a/tools/checkpoint/saver_mcore.py b/tools/checkpoint/saver_mcore.py index 430c9d11c6..1e21445e1e 100644 --- a/tools/checkpoint/saver_mcore.py +++ b/tools/checkpoint/saver_mcore.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import sys import importlib @@ -6,47 +20,69 @@ def add_arguments(parser): - group = parser.add_argument_group(title='Megatron saver') - - group.add_argument('--megatron-path', type=str, default=None, - help='Base directory of Megatron repository') - - group.add_argument('--target-tensor-parallel-size', type=int, - help='Target tensor model parallel size, defaults to the tensor parallel size ' - 'in the input checkpoint if provided by the loader, otherwise to 1') - group.add_argument('--target-pipeline-parallel-size', type=int, - help='Target tensor model parallel size, default to the pipeline parall size ' - 'in the input checkpoint if provided by the loader, otherwise to 1') - group.add_argument('--target-decoder-first-pipeline-num-layers', type=int, - help='Target num layers of first pipeline stage ' - 'in the input checkpoint if provided by the loader, otherwise to None') - group.add_argument("--target-expert-parallel-size", type=int, - help='Target expert model parallel size, default to the expert parallel size ' - 'in the input checkpoint if provided by the loader, otherwise to 1.') - group.add_argument("--target-expert-tensor-parallel-size", type=int, - help='Target expert tensor parallel size, default to the expert tensor parallel size ' - 'in the input checkpoint if provided by the loader, otherwise to the target tensor parallel size.') - group.add_argument("--target-num-experts", type=int, default=None, - help='Target num of experts, default to the num_experts in the input checkpoint' - 'if provided by the loader, otherwise to None. NOTE: Do not support target_num_experts' - 'is not None and the num_experts is not equal to the provieded by the loader.') - group.add_argument("--target-params-dtype", type=str, default=None, - help='The dtype of the converted checkpoint. ' - 'Only used when converting a Transformers checkpoint to a Megatron checkpoint.') + group = parser.add_argument_group(title="Megatron saver") + + group.add_argument( + "--megatron-path", type=str, default=None, help="Base directory of Megatron repository" + ) + + group.add_argument( + "--target-tensor-parallel-size", + type=int, + help="Target tensor model parallel size, defaults to the tensor parallel size " + "in the input checkpoint if provided by the loader, otherwise to 1", + ) + group.add_argument( + "--target-pipeline-parallel-size", + type=int, + help="Target tensor model parallel size, default to the pipeline parall size " + "in the input checkpoint if provided by the loader, otherwise to 1", + ) + group.add_argument( + "--target-decoder-first-pipeline-num-layers", + type=int, + help="Target num layers of first pipeline stage " + "in the input checkpoint if provided by the loader, otherwise to None", + ) + group.add_argument( + "--target-expert-parallel-size", + type=int, + help="Target expert model parallel size, default to the expert parallel size " + "in the input checkpoint if provided by the loader, otherwise to 1.", + ) + group.add_argument( + "--target-expert-tensor-parallel-size", + type=int, + help="Target expert tensor parallel size, default to the expert tensor parallel size " + "in the input checkpoint if provided by the loader, otherwise to the target tensor parallel size.", + ) + group.add_argument( + "--target-num-experts", + type=int, + default=None, + help="Target num of experts, default to the num_experts in the input checkpoint" + "if provided by the loader, otherwise to None. NOTE: Do not support target_num_experts" + "is not None and the num_experts is not equal to the provieded by the loader.", + ) + group.add_argument( + "--target-params-dtype", + type=str, + default=None, + help="The dtype of the converted checkpoint. " + "Only used when converting a Transformers checkpoint to a Megatron checkpoint.", + ) group.add_argument("--build-model-with-initialization", action="store_true") def save_checkpoint(queue, args): - """ prepare import module """ # Search in directory above this root_path = os.path.abspath( - os.path.join(os.path.dirname(__file__), - os.path.pardir, - os.path.pardir)) + os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir) + ) sys.path.insert(0, root_path) sys.path.insert(0, os.path.join(root_path, "flagscale/train")) @@ -60,9 +96,11 @@ def save_checkpoint(queue, args): from megatron.training.tokenizer.tokenizer import vocab_size_with_padding from megatron.core import mpu from megatron.core.tensor_parallel.random import ( - get_cuda_rng_tracker, _DATA_PARALLEL_RNG_TRACKER_NAME, - _EXPERT_PARALLEL_RNG_TRACKER_NAME, _MODEL_PARALLEL_RNG_TRACKER_NAME - ) + get_cuda_rng_tracker, + _DATA_PARALLEL_RNG_TRACKER_NAME, + _EXPERT_PARALLEL_RNG_TRACKER_NAME, + _MODEL_PARALLEL_RNG_TRACKER_NAME, + ) from tools.checkpoint.utils import ( _ConverterFakeProcessGroup, get_expert_model_parallel_rank, @@ -72,7 +110,9 @@ def save_checkpoint(queue, args): validate_mcore_parallel_size, ) except ModuleNotFoundError: - print("Unable to import Megatron, please specify the path to Megatron using --megatron-path. Exiting.") + print( + "Unable to import Megatron, please specify the path to Megatron using --megatron-path. Exiting." + ) queue.put("exit") exit(1) @@ -110,115 +150,153 @@ def check_message(msg): md = queue_get() if args.target_tensor_parallel_size is None: - if hasattr(md, 'previous_tensor_parallel_size'): + if hasattr(md, "previous_tensor_parallel_size"): args.target_tensor_parallel_size = md.previous_tensor_parallel_size else: - print("loader did not provide a tensor parallel size and --target-tensor-parallel-size not provided on command line. " - "Default to 1.") + print( + "loader did not provide a tensor parallel size and --target-tensor-parallel-size not provided on command line. " + "Default to 1." + ) args.target_tensor_parallel_size = 1 if args.target_pipeline_parallel_size is None: - if hasattr(md, 'previous_pipeline_parallel_size'): + if hasattr(md, "previous_pipeline_parallel_size"): args.target_pipeline_parallel_size = md.previous_pipeline_parallel_size else: - print("loader did not provide a pipeline parallel size and --target-pipeline-parallel-size not provided on command line. " - "Default to 1.") + print( + "loader did not provide a pipeline parallel size and --target-pipeline-parallel-size not provided on command line. " + "Default to 1." + ) args.target_pipeline_parallel_size = 1 if args.target_pipeline_parallel_size > 1: if args.target_decoder_first_pipeline_num_layers is None: - if hasattr(md, 'previous_decoder_first_pipeline_num_layers'): - args.target_decoder_first_pipeline_num_layers = md.previous_decoder_first_pipeline_num_layers + if hasattr(md, "previous_decoder_first_pipeline_num_layers"): + args.target_decoder_first_pipeline_num_layers = ( + md.previous_decoder_first_pipeline_num_layers + ) else: - print("loader did not provide a pipeline parallel size and --target-decoder-first-pipeline-num-layers not provided on command line. " - "Default to None.") + print( + "loader did not provide a pipeline parallel size and --target-decoder-first-pipeline-num-layers not provided on command line. " + "Default to None." + ) args.target_decoder_first_pipeline_num_layers = None else: args.target_decoder_first_pipeline_num_layers = None if args.target_expert_parallel_size is None: - if hasattr(md, 'previous_expert_parallel_size'): + if hasattr(md, "previous_expert_parallel_size"): args.target_expert_parallel_size = md.previous_expert_parallel_size else: - print("loader did not provide a expert parallel size and --target-expert-parallel-size not provided on command line. " - "Default to 1.") + print( + "loader did not provide a expert parallel size and --target-expert-parallel-size not provided on command line. " + "Default to 1." + ) args.target_expert_parallel_size = 1 if args.target_expert_tensor_parallel_size is None: - if hasattr(md, 'previous_expert_tensor_parallel_size'): + if hasattr(md, "previous_expert_tensor_parallel_size"): args.target_expert_tensor_parallel_size = md.previous_expert_tensor_parallel_size else: args.target_expert_tensor_parallel_size = args.target_tensor_parallel_size if args.target_num_experts is None: - if hasattr(md, 'previous_num_experts'): + if hasattr(md, "previous_num_experts"): args.target_num_experts = md.previous_num_experts else: - print("loader did not provide a num experts and --target-num-experts not provided on command line. " - "Default to None.") + print( + "loader did not provide a num experts and --target-num-experts not provided on command line. " + "Default to None." + ) if args.target_num_experts is not None and md.previous_num_experts is not None: - assert args.target_num_experts >= md.previous_num_experts, \ + assert args.target_num_experts >= md.previous_num_experts, ( "target_num_experts should be greater than previous_num_experts" + ) if args.target_num_experts > md.previous_num_experts: - print(f"Warning: experts[{md.previous_num_experts}-{args.target_num_experts}] will be random initailized.") + print( + f"Warning: experts[{md.previous_num_experts}-{args.target_num_experts}] will be random initailized." + ) target_model_parallel_size = max( args.target_tensor_parallel_size, args.target_expert_parallel_size * args.target_expert_tensor_parallel_size, ) - os.environ["WORLD_SIZE"] = f'{target_model_parallel_size * args.target_pipeline_parallel_size}' + os.environ["WORLD_SIZE"] = f"{target_model_parallel_size * args.target_pipeline_parallel_size}" os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1" # We want all arguments to come from us sys.argv = [ - 'script.py', - '--num-layers', str(md.num_layers), - '--hidden-size', str(md.hidden_size), - '--seq-length', str(md.seq_length), - '--num-experts', str(getattr(md, "num_experts", 0)), - '--num-attention-heads', str(md.num_attention_heads), - '--max-position-embeddings', str(md.max_position_embeddings), - '--position-embedding-type', str(md.position_embedding_type), - '--tokenizer-type', str(md.tokenizer_type), - '--tensor-model-parallel-size', str(args.target_tensor_parallel_size), - '--pipeline-model-parallel-size', str(args.target_pipeline_parallel_size), - '--expert-model-parallel-size', str(args.target_expert_parallel_size), - '--expert-tensor-parallel-size', str(args.target_expert_tensor_parallel_size), - '--context-parallel-size', '1', - '--no-masked-softmax-fusion', - '--no-bias-gelu-fusion', - '--no-bias-dropout-fusion', - '--use-cpu-initialization', - '--transformer-impl', 'transformer_engine', - '--micro-batch-size', '1', - '--no-load-optim', - '--no-load-rng', - '--no-save-optim', - '--no-save-rng', - '--save-interval', '1', - '--save', args.save_dir, - '--ckpt-format', 'torch', # only 'torch' supported for conversion - '--no-one-logger', + "script.py", + "--num-layers", + str(md.num_layers), + "--hidden-size", + str(md.hidden_size), + "--seq-length", + str(md.seq_length), + "--num-experts", + str(getattr(md, "num_experts", 0)), + "--num-attention-heads", + str(md.num_attention_heads), + "--max-position-embeddings", + str(md.max_position_embeddings), + "--position-embedding-type", + str(md.position_embedding_type), + "--tokenizer-type", + str(md.tokenizer_type), + "--tensor-model-parallel-size", + str(args.target_tensor_parallel_size), + "--pipeline-model-parallel-size", + str(args.target_pipeline_parallel_size), + "--expert-model-parallel-size", + str(args.target_expert_parallel_size), + "--expert-tensor-parallel-size", + str(args.target_expert_tensor_parallel_size), + "--context-parallel-size", + "1", + "--no-masked-softmax-fusion", + "--no-bias-gelu-fusion", + "--no-bias-dropout-fusion", + "--use-cpu-initialization", + "--transformer-impl", + "transformer_engine", + "--micro-batch-size", + "1", + "--no-load-optim", + "--no-load-rng", + "--no-save-optim", + "--no-save-rng", + "--save-interval", + "1", + "--save", + args.save_dir, + "--ckpt-format", + "torch", # only 'torch' supported for conversion + "--no-one-logger", ] if args.target_decoder_first_pipeline_num_layers is not None: - sys.argv.extend(['--decoder-first-pipeline-num-layers', str(args.target_decoder_first_pipeline_num_layers)]) + sys.argv.extend( + [ + "--decoder-first-pipeline-num-layers", + str(args.target_decoder_first_pipeline_num_layers), + ] + ) if args.target_num_experts is not None: - sys.argv.extend(['--num-experts', str(args.target_num_experts)]) - sys.argv.append('--sequence-parallel') + sys.argv.extend(["--num-experts", str(args.target_num_experts)]) + sys.argv.append("--sequence-parallel") if not args.build_model_with_initialization: - sys.argv.append('--no-initialization') + sys.argv.append("--no-initialization") if md.make_vocab_size_divisible_by is not None: - sys.argv.extend(['--make-vocab-size-divisible-by', str(md.make_vocab_size_divisible_by)]) + sys.argv.extend(["--make-vocab-size-divisible-by", str(md.make_vocab_size_divisible_by)]) if md.output_layer: - sys.argv.append('--untie-embeddings-and-output-weights') + sys.argv.append("--untie-embeddings-and-output-weights") if not md.add_bias_linear: - sys.argv.append('--disable-bias-linear') + sys.argv.append("--disable-bias-linear") if md.add_qkv_bias: - sys.argv.append('--add-qkv-bias') + sys.argv.append("--add-qkv-bias") if md.qk_layernorm: - sys.argv.append('--qk-layernorm') + sys.argv.append("--qk-layernorm") if args.target_params_dtype is not None: assert args.target_params_dtype in ["fp32", "fp16", "bf16"] @@ -231,30 +309,61 @@ def check_message(msg): print(f"> convert params_dtype to {md.params_dtype}") if md.params_dtype == torch.float16: - sys.argv.append('--fp16') + sys.argv.append("--fp16") elif md.params_dtype == torch.bfloat16: - sys.argv.append('--bf16') + sys.argv.append("--bf16") margs = parse_args() - if hasattr (md, 'checkpoint_args'): + if hasattr(md, "checkpoint_args"): # These are arguments that we are either changing, or cause problems for validation if they are set # Note that some of these deal with T5 so will need to be changed if we support T5. - args_to_keep = ['tensor_model_parallel_size', 'pipeline_model_parallel_size', 'expert_model_parallel_size', - 'expert_tensor_parallel_size', 'world_size', 'params_dtype', - 'num_layers_per_virtual_pipeline_stage', 'virtual_pipeline_model_parallel_size', - 'masked_softmax_fusion', 'bias_gelu_fusion', 'bias_dropout_fusion', - 'sequence_parallel', - 'no_load_optim', 'no_load_rng', 'no_save_optim', 'no_save_rng', - 'vocab_file', 'tokenizer_model', - 'save_interval', 'save', 'load', 'use_mcore_models', 'num_experts', - 'perform_initialization', 'use_cpu_initialization', - 'recompute_granularity', 'recompute_num_layers', 'recompute_method', - 'encoder_num_layers', 'encoder_seq_length', - 'distribute_saved_activations', 'fp16', 'bf16', 'context_parallel_size', - 'train_iters', 'lr_decay_iters', 'lr_warmup_iters', 'lr_warmup_fraction', - 'start_weight_decay', 'end_weight_decay', - 'main_grads_dtype', 'main_params_dtype', 'exp_avg_dtype', 'exp_avg_sq_dtype', - 'ckpt_format', 'decoder_first_pipeline_num_layers', + args_to_keep = [ + "tensor_model_parallel_size", + "pipeline_model_parallel_size", + "expert_model_parallel_size", + "expert_tensor_parallel_size", + "world_size", + "params_dtype", + "num_layers_per_virtual_pipeline_stage", + "virtual_pipeline_model_parallel_size", + "masked_softmax_fusion", + "bias_gelu_fusion", + "bias_dropout_fusion", + "sequence_parallel", + "no_load_optim", + "no_load_rng", + "no_save_optim", + "no_save_rng", + "vocab_file", + "tokenizer_model", + "save_interval", + "save", + "load", + "use_mcore_models", + "num_experts", + "perform_initialization", + "use_cpu_initialization", + "recompute_granularity", + "recompute_num_layers", + "recompute_method", + "encoder_num_layers", + "encoder_seq_length", + "distribute_saved_activations", + "fp16", + "bf16", + "context_parallel_size", + "train_iters", + "lr_decay_iters", + "lr_warmup_iters", + "lr_warmup_fraction", + "start_weight_decay", + "end_weight_decay", + "main_grads_dtype", + "main_params_dtype", + "exp_avg_dtype", + "exp_avg_sq_dtype", + "ckpt_format", + "decoder_first_pipeline_num_layers", ] for arg, value in vars(md.checkpoint_args).items(): @@ -264,7 +373,9 @@ def check_message(msg): print(f"Checkpoint had argument {arg} but new arguments does not have this.") continue if getattr(margs, arg) != value: - print(f"Overwriting default {arg} value {getattr(margs, arg)} with value from checkpoint {value}.") + print( + f"Overwriting default {arg} value {getattr(margs, arg)} with value from checkpoint {value}." + ) setattr(margs, arg, value) margs.inference_batch_times_seqlen_threshold = -1 @@ -281,7 +392,7 @@ def check_message(msg): if margs.num_experts > 1 and args.target_tensor_parallel_size > 1: margs.sequence_parallel = True - print("*"*20 + "validate saver arguments" + "*"*20) + print("*" * 20 + "validate saver arguments" + "*" * 20) validate_args(margs) validate_mcore_parallel_size(margs) @@ -301,11 +412,13 @@ def check_message(msg): margs = get_args() # validate consumed_samples - if hasattr(md, 'consumed_train_samples'): + if hasattr(md, "consumed_train_samples"): margs.consumed_train_samples = md.consumed_train_samples margs.consumed_valid_samples = md.consumed_valid_samples - print(f"Setting consumed_train_samples to {margs.consumed_train_samples}" - f" and consumed_valid_samples to {margs.consumed_valid_samples}") + print( + f"Setting consumed_train_samples to {margs.consumed_train_samples}" + f" and consumed_valid_samples to {margs.consumed_valid_samples}" + ) else: print("consumed_train_samples not provided.") @@ -366,7 +479,9 @@ def check_message(msg): edp_parallel_size = mcore_model_parallel_size // (etp_size * ep_size) fake_edp_group = _ConverterFakeProcessGroup(size=edp_parallel_size) fake_etp_ep_group = _ConverterFakeProcessGroup(size=etp_size * ep_size) - fake_tcp_group = _ConverterFakeProcessGroup(size=margs.tensor_model_parallel_size*margs.context_parallel_size) + fake_tcp_group = _ConverterFakeProcessGroup( + size=margs.tensor_model_parallel_size * margs.context_parallel_size + ) mpu._PIPELINE_MODEL_PARALLEL_GROUP = fake_pp_group mpu._CONTEXT_PARALLEL_GROUP = fake_cp_group mpu._DATA_PARALLEL_GROUP = fake_dp_group @@ -425,9 +540,10 @@ def get_models(count, dtype): # process transformer layer total_layer_num = 0 for pp_rank in range(pp_size): - mpu.set_pipeline_model_parallel_rank(pp_rank) - fake_pp_group = _ConverterFakeProcessGroup(rank=pp_rank, size=margs.pipeline_model_parallel_size) + fake_pp_group = _ConverterFakeProcessGroup( + rank=pp_rank, size=margs.pipeline_model_parallel_size + ) mpu._PIPELINE_MODEL_PARALLEL_GROUP = fake_pp_group if pp_rank > 0: @@ -450,7 +566,9 @@ def get_models(count, dtype): if md.output_layer: msg = queue_get("output layer") - assert hasattr(models[0], 'output_layer'), "ERROR: got an output layer, but model does not have one" + assert hasattr(models[0], "output_layer"), ( + "ERROR: got an output layer, but model does not have one" + ) ckpt_plugin.set_output_layer_ckpt(msg, models, md, margs) if margs.mtp_num_layers: @@ -466,11 +584,18 @@ def get_models(count, dtype): tp_rank, ep_rank, _ = set_model_parallel_rank(tp_ep_rank) checkpoint_name = get_checkpoint_name(margs.save, md.iteration) print(f"megtron model is saving to {checkpoint_name} ...") - save_checkpoint(md.iteration, [model], None, None, - num_floating_point_operations_so_far=0, - pipeline_rank=pp_rank, pipeline_parallel=args.target_pipeline_parallel_size > 1, - expert_rank=ep_rank, expert_parallel=args.target_expert_parallel_size > 1, - tensor_rank=tp_rank) + save_checkpoint( + md.iteration, + [model], + None, + None, + num_floating_point_operations_so_far=0, + pipeline_rank=pp_rank, + pipeline_parallel=args.target_pipeline_parallel_size > 1, + expert_rank=ep_rank, + expert_parallel=args.target_expert_parallel_size > 1, + tensor_rank=tp_rank, + ) # release the uselese model parts models[tp_ep_rank] = None diff --git a/tools/checkpoint/saver_transformers.py b/tools/checkpoint/saver_transformers.py index 7407b1f020..11a6eaebcd 100644 --- a/tools/checkpoint/saver_transformers.py +++ b/tools/checkpoint/saver_transformers.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import os import sys import importlib @@ -6,47 +20,64 @@ def add_arguments(parser): - group = parser.add_argument_group(title='Transformers saver') - - group.add_argument('--megatron-path', type=str, default=None, - help='Base directory of Megatron repository') - - group.add_argument('--target-tensor-parallel-size', type=int, - help='Target tensor model parallel size, defaults to the tensor parallel size ' - 'in the input checkpoint if provided by the loader, otherwise to 1') - group.add_argument('--target-pipeline-parallel-size', type=int, - help='Target tensor model parallel size, default to the pipeline parall size ' - 'in the input checkpoint if provided by the loader, otherwise to 1') - group.add_argument("--target-expert-parallel-size", type=int, - help='Target expert model parallel size, default to the expert parallel size ' - 'in the input checkpoint if provided by the loader, otherwise to 1.') - group.add_argument("--target-num-experts", type=int, default=None, - help='Target num of experts, default to the num_experts in the input checkpoint' - 'if provided by the loader, otherwise to None. NOTE: Do not support target_num_experts' - 'is not None and the num_experts is not equal to the provieded by the loader.') - group.add_argument("--target-params-dtype", type=str, default=None, - help='The dtype of the converted checkpoint. ' - 'Only used when converting a Transformers checkpoint to a Megatron checkpoint.') + group = parser.add_argument_group(title="Transformers saver") + + group.add_argument( + "--megatron-path", type=str, default=None, help="Base directory of Megatron repository" + ) + + group.add_argument( + "--target-tensor-parallel-size", + type=int, + help="Target tensor model parallel size, defaults to the tensor parallel size " + "in the input checkpoint if provided by the loader, otherwise to 1", + ) + group.add_argument( + "--target-pipeline-parallel-size", + type=int, + help="Target tensor model parallel size, default to the pipeline parall size " + "in the input checkpoint if provided by the loader, otherwise to 1", + ) + group.add_argument( + "--target-expert-parallel-size", + type=int, + help="Target expert model parallel size, default to the expert parallel size " + "in the input checkpoint if provided by the loader, otherwise to 1.", + ) + group.add_argument( + "--target-num-experts", + type=int, + default=None, + help="Target num of experts, default to the num_experts in the input checkpoint" + "if provided by the loader, otherwise to None. NOTE: Do not support target_num_experts" + "is not None and the num_experts is not equal to the provieded by the loader.", + ) + group.add_argument( + "--target-params-dtype", + type=str, + default=None, + help="The dtype of the converted checkpoint. " + "Only used when converting a Transformers checkpoint to a Megatron checkpoint.", + ) def save_checkpoint(queue, args): - """ prepare import module """ try: import transformers - major, minor, _ = map(int, transformers.__version__.split('.')) + + major, minor, _ = map(int, transformers.__version__.split(".")) assert major >= 4 and minor >= 36 except: raise ImportError("transformers version >= 4.36.0 ") # Search in directory above this root_path = os.path.abspath( - os.path.join(os.path.dirname(__file__), - os.path.pardir, - os.path.pardir)) + os.path.join(os.path.dirname(__file__), os.path.pardir, os.path.pardir) + ) sys.path.insert(0, root_path) sys.path.insert(0, os.path.join(root_path, "flagscale/train")) @@ -57,7 +88,9 @@ def save_checkpoint(queue, args): from megatron.training.arguments import parse_args, validate_args from megatron.training.tokenizer.tokenizer import vocab_size_with_padding except ModuleNotFoundError: - print("Unable to import Megatron, please specify the path to Megatron using --megatron-path. Exiting.") + print( + "Unable to import Megatron, please specify the path to Megatron using --megatron-path. Exiting." + ) queue.put("exit") exit(1) @@ -101,86 +134,116 @@ def check_message(msg): # assert args.target_num_experts is None if args.target_tensor_parallel_size is None: - if hasattr(md, 'previous_tensor_parallel_size'): + if hasattr(md, "previous_tensor_parallel_size"): args.target_tensor_parallel_size = md.previous_tensor_parallel_size else: - print("loader did not provide a tensor parallel size and --target-tensor-parallel-size not provided on command line. " - "Default to 1.") + print( + "loader did not provide a tensor parallel size and --target-tensor-parallel-size not provided on command line. " + "Default to 1." + ) args.target_tensor_parallel_size = 1 if args.target_pipeline_parallel_size is None: - if hasattr(md, 'previous_pipeline_parallel_size'): + if hasattr(md, "previous_pipeline_parallel_size"): args.target_pipeline_parallel_size = md.previous_pipeline_parallel_size else: - print("loader did not provide a pipeline parallel size and --target-pipeline-parallel-size not provided on command line. " - "Default to 1.") + print( + "loader did not provide a pipeline parallel size and --target-pipeline-parallel-size not provided on command line. " + "Default to 1." + ) args.target_pipeline_parallel_size = 1 if args.target_expert_parallel_size is None: - if hasattr(md, 'previous_expert_parallel_size'): + if hasattr(md, "previous_expert_parallel_size"): args.target_expert_parallel_size = md.previous_expert_parallel_size else: - print("loader did not provide a expert parallel size and --target-expert-parallel-size not provided on command line. " - "Default to 1.") + print( + "loader did not provide a expert parallel size and --target-expert-parallel-size not provided on command line. " + "Default to 1." + ) args.target_expert_parallel_size = 1 if args.target_num_experts is None: - if hasattr(md, 'previous_num_experts'): + if hasattr(md, "previous_num_experts"): args.target_num_experts = md.previous_num_experts else: - print("loader did not provide a num experts and --target-num-experts not provided on command line. " - "Default to None.") + print( + "loader did not provide a num experts and --target-num-experts not provided on command line. " + "Default to None." + ) if args.target_num_experts is not None and md.previous_num_experts is not None: - assert args.target_num_experts >= md.previous_num_experts, \ + assert args.target_num_experts >= md.previous_num_experts, ( "target_num_experts should be greater than previous_num_experts" + ) if args.target_num_experts > md.previous_num_experts: - print(f"Warning: experts[{md.previous_num_experts}-{args.target_num_experts}] will be random initailized.") + print( + f"Warning: experts[{md.previous_num_experts}-{args.target_num_experts}] will be random initailized." + ) - os.environ["WORLD_SIZE"] = f'{args.target_tensor_parallel_size * args.target_pipeline_parallel_size * args.target_expert_parallel_size}' + os.environ["WORLD_SIZE"] = ( + f"{args.target_tensor_parallel_size * args.target_pipeline_parallel_size * args.target_expert_parallel_size}" + ) os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1" # We want all arguments to come from us sys.argv = [ - 'script.py', - '--num-layers', str(md.num_layers), - '--hidden-size', str(md.hidden_size), - '--seq-length', str(md.seq_length), - '--num-experts', str(getattr(md, "num_experts", 0)), - '--num-attention-heads', str(md.num_attention_heads), - '--max-position-embeddings', str(md.max_position_embeddings), - '--position-embedding-type', str(md.position_embedding_type), - '--tokenizer-type', str(md.tokenizer_type), - '--tensor-model-parallel-size', str(args.target_tensor_parallel_size), - '--pipeline-model-parallel-size', str(args.target_pipeline_parallel_size), - '--expert-model-parallel-size', str(args.target_expert_parallel_size), - '--context-parallel-size', '1', - '--no-masked-softmax-fusion', - '--no-bias-gelu-fusion', - '--no-bias-dropout-fusion', - '--use-cpu-initialization', - '--transformer-impl', 'transformer_engine', - '--micro-batch-size', '1', - '--no-load-optim', - '--no-load-rng', - '--no-save-optim', - '--no-save-rng', - '--no-initialization', - '--save-interval', '1', - '--save', args.save_dir, - '--ckpt-format', 'torch', # only 'torch' supported for conversion - '--no-one-logger', + "script.py", + "--num-layers", + str(md.num_layers), + "--hidden-size", + str(md.hidden_size), + "--seq-length", + str(md.seq_length), + "--num-experts", + str(getattr(md, "num_experts", 0)), + "--num-attention-heads", + str(md.num_attention_heads), + "--max-position-embeddings", + str(md.max_position_embeddings), + "--position-embedding-type", + str(md.position_embedding_type), + "--tokenizer-type", + str(md.tokenizer_type), + "--tensor-model-parallel-size", + str(args.target_tensor_parallel_size), + "--pipeline-model-parallel-size", + str(args.target_pipeline_parallel_size), + "--expert-model-parallel-size", + str(args.target_expert_parallel_size), + "--context-parallel-size", + "1", + "--no-masked-softmax-fusion", + "--no-bias-gelu-fusion", + "--no-bias-dropout-fusion", + "--use-cpu-initialization", + "--transformer-impl", + "transformer_engine", + "--micro-batch-size", + "1", + "--no-load-optim", + "--no-load-rng", + "--no-save-optim", + "--no-save-rng", + "--no-initialization", + "--save-interval", + "1", + "--save", + args.save_dir, + "--ckpt-format", + "torch", # only 'torch' supported for conversion + "--no-one-logger", ] if args.target_num_experts is not None: - sys.argv.extend(['--num-experts', str(args.target_num_experts)]) + sys.argv.extend(["--num-experts", str(args.target_num_experts)]) if md.make_vocab_size_divisible_by is not None: - sys.argv.extend(['--make-vocab-size-divisible-by', str(md.make_vocab_size_divisible_by)]) + sys.argv.extend(["--make-vocab-size-divisible-by", str(md.make_vocab_size_divisible_by)]) if md.output_layer: - sys.argv.append('--untie-embeddings-and-output-weights') + sys.argv.append("--untie-embeddings-and-output-weights") if not md.add_bias_linear: - sys.argv.append('--disable-bias-linear') + sys.argv.append("--disable-bias-linear") if md.add_qkv_bias: - sys.argv.append('--add-qkv-bias') + sys.argv.append("--add-qkv-bias") if args.target_params_dtype is not None: assert args.target_params_dtype in ["fp32", "fp16", "bf16"] @@ -193,29 +256,57 @@ def check_message(msg): print(f"> convert params_dtype to {md.params_dtype}") if md.params_dtype == torch.float16: - sys.argv.append('--fp16') + sys.argv.append("--fp16") elif md.params_dtype == torch.bfloat16: - sys.argv.append('--bf16') + sys.argv.append("--bf16") margs = parse_args() - if hasattr (md, 'checkpoint_args'): + if hasattr(md, "checkpoint_args"): # These are arguments that we are either changing, or cause problems for validation if they are set # Note that some of these deal with T5 so will need to be changed if we support T5. - args_to_keep = ['tensor_model_parallel_size', 'pipeline_model_parallel_size', 'expert_model_parallel_size', - 'expert_tensor_parallel_size', 'world_size', 'params_dtype', - 'num_layers_per_virtual_pipeline_stage', 'virtual_pipeline_model_parallel_size', - 'masked_softmax_fusion', 'bias_gelu_fusion', 'bias_dropout_fusion', - 'sequence_parallel', - 'no_load_optim', 'no_load_rng', 'no_save_optim', 'no_save_rng', - 'vocab_file', 'tokenizer_model', - 'save_interval', 'save', 'load', 'use_mcore_models', 'num_experts', - 'perform_initialization', 'use_cpu_initialization', - 'recompute_granularity', 'recompute_num_layers', 'recompute_method', - 'encoder_num_layers', 'encoder_seq_length', - 'distribute_saved_activations', 'fp16', 'bf16', 'context_parallel_size', - 'train_iters', 'lr_decay_iters', 'lr_warmup_iters', 'lr_warmup_fraction', - 'start_weight_decay', 'end_weight_decay', - 'ckpt_format', 'decoder_first_pipeline_num_layers', + args_to_keep = [ + "tensor_model_parallel_size", + "pipeline_model_parallel_size", + "expert_model_parallel_size", + "expert_tensor_parallel_size", + "world_size", + "params_dtype", + "num_layers_per_virtual_pipeline_stage", + "virtual_pipeline_model_parallel_size", + "masked_softmax_fusion", + "bias_gelu_fusion", + "bias_dropout_fusion", + "sequence_parallel", + "no_load_optim", + "no_load_rng", + "no_save_optim", + "no_save_rng", + "vocab_file", + "tokenizer_model", + "save_interval", + "save", + "load", + "use_mcore_models", + "num_experts", + "perform_initialization", + "use_cpu_initialization", + "recompute_granularity", + "recompute_num_layers", + "recompute_method", + "encoder_num_layers", + "encoder_seq_length", + "distribute_saved_activations", + "fp16", + "bf16", + "context_parallel_size", + "train_iters", + "lr_decay_iters", + "lr_warmup_iters", + "lr_warmup_fraction", + "start_weight_decay", + "end_weight_decay", + "ckpt_format", + "decoder_first_pipeline_num_layers", ] for arg, value in vars(md.checkpoint_args).items(): @@ -225,7 +316,9 @@ def check_message(msg): print(f"Checkpoint had argument {arg} but new arguments does not have this.") continue if getattr(margs, arg) != value: - print(f"Overwriting default {arg} value {getattr(margs, arg)} with value from checkpoint {value}.") + print( + f"Overwriting default {arg} value {getattr(margs, arg)} with value from checkpoint {value}." + ) setattr(margs, arg, value) margs.inference_batch_times_seqlen_threshold = -1 @@ -242,7 +335,7 @@ def check_message(msg): if margs.num_experts > 1 and args.target_tensor_parallel_size > 1: margs.sequence_parallel = True - print("*"*20 + "validate saver arguments" + "*"*20) + print("*" * 20 + "validate saver arguments" + "*" * 20) margs = validate_args(margs) # Use M-core models & unset loaded paths. diff --git a/tools/checkpoint/sfpt_ckpt/dcp_to_sfpt.py b/tools/checkpoint/sfpt_ckpt/dcp_to_sfpt.py index abc52c630b..6005b8393f 100644 --- a/tools/checkpoint/sfpt_ckpt/dcp_to_sfpt.py +++ b/tools/checkpoint/sfpt_ckpt/dcp_to_sfpt.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import os diff --git a/tools/checkpoint/sfpt_ckpt/sfpt_to_dcp.py b/tools/checkpoint/sfpt_ckpt/sfpt_to_dcp.py index 6d8767b0b5..af786123cb 100644 --- a/tools/checkpoint/sfpt_ckpt/sfpt_to_dcp.py +++ b/tools/checkpoint/sfpt_ckpt/sfpt_to_dcp.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import os from argparse import Namespace diff --git a/tools/checkpoint/utils.py b/tools/checkpoint/utils.py index ecd4f0e1e5..bf6028a429 100644 --- a/tools/checkpoint/utils.py +++ b/tools/checkpoint/utils.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from contextlib import contextmanager import psutil diff --git a/tools/codestyle/pre-commit.sh b/tools/codestyle/pre-commit.sh index 2c97826e47..434fda8db3 100644 --- a/tools/codestyle/pre-commit.sh +++ b/tools/codestyle/pre-commit.sh @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + set +x if ! [[ $(pre-commit --version) == *"4.2.0"* ]]; then diff --git a/tools/datasets/llava_onevision/make_llava_ov_wds.sh b/tools/datasets/llava_onevision/make_llava_ov_wds.sh index 3634324e6f..33a8d0bffb 100644 --- a/tools/datasets/llava_onevision/make_llava_ov_wds.sh +++ b/tools/datasets/llava_onevision/make_llava_ov_wds.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # set environment variables of the gpu cluster export GLOO_SOCKET_IFNAME=bond0 export NCCL_SOCKET_IFNAME=bond0 diff --git a/tools/datasets/qwenvl/build_llava_frame_dataset.py b/tools/datasets/qwenvl/build_llava_frame_dataset.py index b94ddbeea9..b16e9233ad 100644 --- a/tools/datasets/qwenvl/build_llava_frame_dataset.py +++ b/tools/datasets/qwenvl/build_llava_frame_dataset.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Copied from https://github.com/alibaba/Pai-Megatron-Patch/blob/8949a6647cbf6b39837ad3dd911fa4aa0726895b/toolkits/multimodal_data_preprocessing/build_llava_frame_dataset.py import json import multiprocessing.pool as mpp diff --git a/tools/datasets/qwenvl/convert_custom_dataset_to_wds_chatml_str.py b/tools/datasets/qwenvl/convert_custom_dataset_to_wds_chatml_str.py index 02d8b33c4e..e8e8ca2090 100644 --- a/tools/datasets/qwenvl/convert_custom_dataset_to_wds_chatml_str.py +++ b/tools/datasets/qwenvl/convert_custom_dataset_to_wds_chatml_str.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Adopted from https://github.com/alibaba/Pai-Megatron-Patch/blob/8949a6647cbf6b39837ad3dd911fa4aa0726895b/toolkits/multimodal_data_preprocessing/convert_custom_dataset_to_wds_chatml.py # We must store the path of vision data, not the real data. diff --git a/tools/datasets/qwenvl/convert_custom_dataset_to_wds_chatml_str_legacy.py b/tools/datasets/qwenvl/convert_custom_dataset_to_wds_chatml_str_legacy.py index 94f3caabe9..785134079c 100644 --- a/tools/datasets/qwenvl/convert_custom_dataset_to_wds_chatml_str_legacy.py +++ b/tools/datasets/qwenvl/convert_custom_dataset_to_wds_chatml_str_legacy.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Adopted from https://github.com/alibaba/Pai-Megatron-Patch/blob/8949a6647cbf6b39837ad3dd911fa4aa0726895b/toolkits/multimodal_data_preprocessing/convert_custom_dataset_to_wds_chatml.py # We must store the path of vision data, not the real data. diff --git a/tools/datasets/qwenvl/convert_llava_pretrain_to_wds.py b/tools/datasets/qwenvl/convert_llava_pretrain_to_wds.py index d4d9b27a1b..6c57363bb1 100644 --- a/tools/datasets/qwenvl/convert_llava_pretrain_to_wds.py +++ b/tools/datasets/qwenvl/convert_llava_pretrain_to_wds.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Copied from https://github.com/alibaba/Pai-Megatron-Patch/blob/8949a6647cbf6b39837ad3dd911fa4aa0726895b/toolkits/multimodal_data_preprocessing/convert_llava_pretrain_to_wds.py import json import os diff --git a/tools/datasets/qwenvl/data/__init__.py b/tools/datasets/qwenvl/data/__init__.py index e69de29bb2..a491ed2726 100644 --- a/tools/datasets/qwenvl/data/__init__.py +++ b/tools/datasets/qwenvl/data/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tools/datasets/qwenvl/data/dataset_helpers_robotics.py b/tools/datasets/qwenvl/data/dataset_helpers_robotics.py index 5edea80aa6..2d8e0dce6e 100644 --- a/tools/datasets/qwenvl/data/dataset_helpers_robotics.py +++ b/tools/datasets/qwenvl/data/dataset_helpers_robotics.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Adopted from https://github.com/alibaba/Pai-Megatron-Patch/blob/8949a6647cbf6b39837ad3dd911fa4aa0726895b/megatron_patch/data/energon/chatml.py. import logging diff --git a/tools/datasets/qwenvl/data/energon/chatml.py b/tools/datasets/qwenvl/data/energon/chatml.py index f566293881..8d2882a60e 100644 --- a/tools/datasets/qwenvl/data/energon/chatml.py +++ b/tools/datasets/qwenvl/data/energon/chatml.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Copied from https://github.com/alibaba/Pai-Megatron-Patch/blob/8949a6647cbf6b39837ad3dd911fa4aa0726895b/megatron_patch/data/energon/chatml.py. import pickle diff --git a/tools/datasets/qwenvl/dataset_preparation.md b/tools/datasets/qwenvl/dataset_preparation.md index ce08cf1c3f..529d36e5fb 100644 --- a/tools/datasets/qwenvl/dataset_preparation.md +++ b/tools/datasets/qwenvl/dataset_preparation.md @@ -1,3 +1,19 @@ + + # 📎 Reference Mainly based on [Pai-Megatron-Patch](https://github.com/alibaba/Pai-Megatron-Patch/tree/main/toolkits/multimodal_data_preprocessing/),with necessary modifications for integration into the current training framework. diff --git a/tools/datasets/qwenvl/replace_llava_image_key.py b/tools/datasets/qwenvl/replace_llava_image_key.py index 867977a93d..a33fe323b0 100644 --- a/tools/datasets/qwenvl/replace_llava_image_key.py +++ b/tools/datasets/qwenvl/replace_llava_image_key.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Copied from https://github.com/alibaba/Pai-Megatron-Patch/blob/8949a6647cbf6b39837ad3dd911fa4aa0726895b/toolkits/multimodal_data_preprocessing/replace_llava_image_key.py import json import os diff --git a/tools/datasets/qwenvl/run.sh b/tools/datasets/qwenvl/run.sh index df6cf378ea..00e2c803b1 100644 --- a/tools/datasets/qwenvl/run.sh +++ b/tools/datasets/qwenvl/run.sh @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + export PYTHONPATH=$PYTHONPATH:../../../ python convert_custom_dataset_to_wds_chatml_str.py \ --dataset-root=/share/project/lizhiyu/LLaMA-Factory/data/sample_dataset/text_only/ \ diff --git a/tools/datasets/vla/convert.py b/tools/datasets/vla/convert.py index 86f0b768e6..752ee375bd 100644 --- a/tools/datasets/vla/convert.py +++ b/tools/datasets/vla/convert.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Adopted from https://github.com/alibaba/Pai-Megatron-Patch/blob/8949a6647cbf6b39837ad3dd911fa4aa0726895b/toolkits/multimodal_data_preprocessing/convert_custom_dataset_to_wds_chatml.py import json diff --git a/tools/datasets/vla/data/dataset_helpers.py b/tools/datasets/vla/data/dataset_helpers.py index d9ca941444..411b4fe9a6 100644 --- a/tools/datasets/vla/data/dataset_helpers.py +++ b/tools/datasets/vla/data/dataset_helpers.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import json import logging diff --git a/tools/datasets/vla/data/energon/chatml.py b/tools/datasets/vla/data/energon/chatml.py index 973d0ce794..f4b512a932 100644 --- a/tools/datasets/vla/data/energon/chatml.py +++ b/tools/datasets/vla/data/energon/chatml.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Adopted from https://github.com/alibaba/Pai-Megatron-Patch/blob/8949a6647cbf6b39837ad3dd911fa4aa0726895b/megatron_patch/data/energon/chatml.py. import json diff --git a/tools/install/ascend/env.sh b/tools/install/ascend/env.sh index 2af6ff753f..b18ff25305 100755 --- a/tools/install/ascend/env.sh +++ b/tools/install/ascend/env.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ============================================================================= # FlagScale Ascend Environment Variables # ============================================================================= diff --git a/tools/install/ascend/install_base.sh b/tools/install/ascend/install_base.sh index 3228bd2791..3f8537863f 100755 --- a/tools/install/ascend/install_base.sh +++ b/tools/install/ascend/install_base.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Base phase (Ascend): requirements/ascend/base.txt SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" diff --git a/tools/install/ascend/install_serve.sh b/tools/install/ascend/install_serve.sh index 744a93ad82..c5aef53287 100755 --- a/tools/install/ascend/install_serve.sh +++ b/tools/install/ascend/install_serve.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Serve task (Ascend): requirements/ascend/serve.txt + source deps (FlagGems, vllm-plugin-FL) SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" diff --git a/tools/install/cuda/env.sh b/tools/install/cuda/env.sh index 7d934479b8..90da94cd17 100755 --- a/tools/install/cuda/env.sh +++ b/tools/install/cuda/env.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ============================================================================= # FlagScale CUDA Environment Variables # ============================================================================= diff --git a/tools/install/cuda/install_base.sh b/tools/install/cuda/install_base.sh index eedd46d8a6..a9d545325c 100755 --- a/tools/install/cuda/install_base.sh +++ b/tools/install/cuda/install_base.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Base phase (CUDA): apt packages + requirements/cuda/base.txt SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" diff --git a/tools/install/cuda/install_hetero_train.sh b/tools/install/cuda/install_hetero_train.sh index 78279f2562..4ceebe9f8d 100755 --- a/tools/install/cuda/install_hetero_train.sh +++ b/tools/install/cuda/install_hetero_train.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Source dependencies for hetero_train task (same as train) # Delegates to install_train.sh since both tasks use Megatron-LM-FL diff --git a/tools/install/cuda/install_inference.sh b/tools/install/cuda/install_inference.sh index 11fd74b76f..c2fc523b74 100755 --- a/tools/install/cuda/install_inference.sh +++ b/tools/install/cuda/install_inference.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Inference task (CUDA): requirements/cuda/inference.txt SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" diff --git a/tools/install/cuda/install_rl.sh b/tools/install/cuda/install_rl.sh index 4b981cc023..b597bd9bf9 100755 --- a/tools/install/cuda/install_rl.sh +++ b/tools/install/cuda/install_rl.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # RL task (CUDA): requirements/cuda/rl.txt SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" diff --git a/tools/install/cuda/install_serve.sh b/tools/install/cuda/install_serve.sh index 3f6726f20d..44aa605ff2 100755 --- a/tools/install/cuda/install_serve.sh +++ b/tools/install/cuda/install_serve.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Serve task (CUDA): requirements/cuda/serve.txt + source deps SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" diff --git a/tools/install/cuda/install_train.sh b/tools/install/cuda/install_train.sh index 2e0ab080a4..ab320edb88 100755 --- a/tools/install/cuda/install_train.sh +++ b/tools/install/cuda/install_train.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Train task (CUDA): requirements/cuda/train.txt + source deps SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" diff --git a/tools/install/install.sh b/tools/install/install.sh index a99da3a5e5..b89c12494d 100755 --- a/tools/install/install.sh +++ b/tools/install/install.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ============================================================================= # FlagScale Dependency Installation # ============================================================================= diff --git a/tools/install/install_dev.sh b/tools/install/install_dev.sh index 68b5a38c4e..08c7d34ce4 100755 --- a/tools/install/install_dev.sh +++ b/tools/install/install_dev.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Dev phase: requirements/dev.txt + dev tools (sccache) SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" diff --git a/tools/install/install_system.sh b/tools/install/install_system.sh index 86ad041e8e..7087ba1843 100755 --- a/tools/install/install_system.sh +++ b/tools/install/install_system.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ============================================================================= # FlagScale System Dependencies Installation # ============================================================================= diff --git a/tools/install/utils/load_platform_config.sh b/tools/install/utils/load_platform_config.sh index 41cff09cb9..83e6ebe636 100755 --- a/tools/install/utils/load_platform_config.sh +++ b/tools/install/utils/load_platform_config.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # Platform configuration loading script # Extracts platform configuration and groups tests by task # Usage: source load_platform_config.sh && load_platform_config diff --git a/tools/install/utils/parse_requirements.py b/tools/install/utils/parse_requirements.py index 0ea224505a..21f8e2cbe4 100644 --- a/tools/install/utils/parse_requirements.py +++ b/tools/install/utils/parse_requirements.py @@ -1,4 +1,19 @@ #!/usr/bin/env python3 + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Parse requirements files with per-package annotation support. Shared by setup.py (Python API) and tools/install shell scripts (CLI). diff --git a/tools/install/utils/pkg_utils.sh b/tools/install/utils/pkg_utils.sh index 544097cdc9..5d361201dc 100644 --- a/tools/install/utils/pkg_utils.sh +++ b/tools/install/utils/pkg_utils.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ============================================================================= # Package Manager Utilities # ============================================================================= diff --git a/tools/install/utils/pyenv_utils.sh b/tools/install/utils/pyenv_utils.sh index 3fe086bfbf..ebb8bba68c 100644 --- a/tools/install/utils/pyenv_utils.sh +++ b/tools/install/utils/pyenv_utils.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ============================================================================= # Python Environment Utilities # ============================================================================= diff --git a/tools/install/utils/retry_utils.sh b/tools/install/utils/retry_utils.sh index f0f66fdc25..fbbfd607dd 100755 --- a/tools/install/utils/retry_utils.sh +++ b/tools/install/utils/retry_utils.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ============================================================================= # Retry Utilities # ============================================================================= diff --git a/tools/install/utils/utils.sh b/tools/install/utils/utils.sh index e435dfa322..4986930b1c 100755 --- a/tools/install/utils/utils.sh +++ b/tools/install/utils/utils.sh @@ -1,4 +1,19 @@ #!/bin/bash + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + # ============================================================================= # Common Utilities # ============================================================================= diff --git a/tools/perf_monitor/perf_smoke.py b/tools/perf_monitor/perf_smoke.py index a60290fcef..7408e637be 100644 --- a/tools/perf_monitor/perf_smoke.py +++ b/tools/perf_monitor/perf_smoke.py @@ -1,4 +1,19 @@ #!/usr/bin/env python3 + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Standalone smoke test for the FlagScale performance monitor.""" from __future__ import annotations diff --git a/tools/straggler/gpu_burner.py b/tools/straggler/gpu_burner.py index 7f4026b7a9..5bcce48fbb 100644 --- a/tools/straggler/gpu_burner.py +++ b/tools/straggler/gpu_burner.py @@ -1,4 +1,19 @@ #!/usr/bin/env python3 + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Simple external GPU burner used to trigger straggler behavior.""" from __future__ import annotations diff --git a/tools/straggler/straggler_smoke.py b/tools/straggler/straggler_smoke.py index acacce4803..4e0d085904 100644 --- a/tools/straggler/straggler_smoke.py +++ b/tools/straggler/straggler_smoke.py @@ -1,4 +1,19 @@ #!/usr/bin/env python3 + +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Standalone smoke test for FlagScale straggler detection. This script does not depend on ``run.py`` or Megatron training startup. diff --git a/tools/visualize/2d_visualize.py b/tools/visualize/2d_visualize.py index be61854403..fc6f0b2cd3 100644 --- a/tools/visualize/2d_visualize.py +++ b/tools/visualize/2d_visualize.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import os diff --git a/tools/visualize/3d_visualize.py b/tools/visualize/3d_visualize.py index 672b3f8b47..8d2ce46067 100644 --- a/tools/visualize/3d_visualize.py +++ b/tools/visualize/3d_visualize.py @@ -1,3 +1,17 @@ +# Copyright 2026 FlagOS Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """ Usage: python pca_3d_visualizer.py latent1.npy latent2.npy ...