Merge branch 'main' into dfridman/deepseek-conversion-mtp-support

Author: ko3n1g

.github/workflows/_update_dependencies.yml

@@ -102,7 +102,7 @@ jobs:
         uses: peter-evans/create-pull-request@v6
         id: create-pull-request
         env:
-          title: "chore(beep boop 🤖): Bump `uv.lock` (${{ inputs.target-branch}}) (${{ needs.pre-flight.outputs.date }})"
+          title: "chore(beep boop 🤖): Bump `uv.lock` (${{ env.TARGET_BRANCH}}, mcore-${{ inputs.mcore-branch }}) (${{ needs.pre-flight.outputs.date }})"
         with:
           branch: ${{ env.SOURCE_BRANCH }}
           base: ${{ env.TARGET_BRANCH }}
@@ -115,6 +115,7 @@ jobs:
           commit-message: ${{ env.title }}
           signoff: true
           committer: "github-actions[bot] <github-actions[bot]@users.noreply.github.com>"
+          labels: needs-more-tests
 
       - name: Post /ok to test comment
         env:
@@ -197,8 +198,11 @@ jobs:
       - name: Revert submodule bump for dev branch
         if: ${{ inputs.mcore-branch == 'dev' }}
         run: |
-          git checkout 3rdparty/Megatron-LM
-          git checkout .main.commit
+          git config user.name "github-actions[bot]"
+          git config user.email "github-actions[bot]@users.noreply.github.com"
+          git fetch origin ${{ env.TARGET_BRANCH }}
+          git restore --source=origin/${{ env.TARGET_BRANCH }} -- 3rdparty/Megatron-LM uv.lock
+          git add 3rdparty/Megatron-LM uv.lock
           git commit -m "Revert submodule bump for dev branch"
           git push origin ${{ env.SOURCE_BRANCH }}
 

.github/workflows/build-test-publish-wheel.yml

@@ -34,12 +34,13 @@ concurrency:
 
 jobs:
   pre-flight:
-    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.73.2
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.80.1
     with:
       default_runner_prefix: ${{ vars.DEFAULT_RUNNER_PREFIX }}
       non_nvidia_runner_prefix: ${{ vars.NON_NVIDIA_RUNNER_PREFIX }}
       default_test_data_path: ${{ vars.DEFAULT_TEST_DATA_PATH }}
       non_nvidia_test_data_path: ${{ vars.NON_NVIDIA_TEST_DATA_PATH }}
+      sso_users_filename: ${{ vars.SSO_USERS_FILENAME }}
     secrets:
       NVIDIA_MANAGEMENT_ORG_PAT: ${{ secrets.NVIDIA_MANAGEMENT_ORG_PAT }}
 

.github/workflows/cicd-main.yml

@@ -61,12 +61,13 @@ env:
 
 jobs:
   pre-flight:
-    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.73.2
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.80.1
     with:
       default_runner_prefix: ${{ vars.DEFAULT_RUNNER_PREFIX }}
       non_nvidia_runner_prefix: ${{ vars.NON_NVIDIA_RUNNER_PREFIX }}
       default_test_data_path: ${{ vars.DEFAULT_TEST_DATA_PATH }}
       non_nvidia_test_data_path: ${{ vars.NON_NVIDIA_TEST_DATA_PATH }}
+      sso_users_filename: ${{ vars.SSO_USERS_FILENAME }}
     secrets:
       NVIDIA_MANAGEMENT_ORG_PAT: ${{ secrets.NVIDIA_MANAGEMENT_ORG_PAT }}
 
@@ -722,9 +723,10 @@ jobs:
           GH_TOKEN: ${{ github.token }}
           # A tier is "expected" if its if-condition (excluding success()/is_ci_workload/force_run_all) would be true.
           # Unexpected skips in an expected tier are failures.
-          EXPECT_L0: ${{ github.event.inputs.test_suite == '' || github.event.inputs.test_suite == 'all' || github.event.inputs.test_suite == 'functional-only' }}
-          EXPECT_L1: ${{ (github.event.inputs.test_suite == '' || github.event.inputs.test_suite == 'all' || github.event.inputs.test_suite == 'functional-only') && (github.ref == 'refs/heads/main' || github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' || needs.check-labels.outputs.needs_more_tests == 'true') }}
-          EXPECT_L2: ${{ (github.event.inputs.test_suite == '' || github.event.inputs.test_suite == 'all' || github.event.inputs.test_suite == 'functional-only') && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch') }}
+          # is_ci_workload must be true, otherwise docs-only/deployment runs have no tests and trigger false positives.
+          EXPECT_L0: ${{ needs.pre-flight.outputs.is_ci_workload == 'true' && (github.event.inputs.test_suite == '' || github.event.inputs.test_suite == 'all' || github.event.inputs.test_suite == 'functional-only') }}
+          EXPECT_L1: ${{ needs.pre-flight.outputs.is_ci_workload == 'true' && (github.event.inputs.test_suite == '' || github.event.inputs.test_suite == 'all' || github.event.inputs.test_suite == 'functional-only') && (github.ref == 'refs/heads/main' || github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' || needs.check-labels.outputs.needs_more_tests == 'true') }}
+          EXPECT_L2: ${{ needs.pre-flight.outputs.is_ci_workload == 'true' && (github.event.inputs.test_suite == '' || github.event.inputs.test_suite == 'all' || github.event.inputs.test_suite == 'functional-only') && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch') }}
         run: |
           FAILED_JOBS=$(gh run view $GITHUB_RUN_ID --json jobs --jq '[.jobs[] | select(.status == "completed" and .conclusion != "success" and .conclusion != "skipped")] | length') || echo 0
 
@@ -763,12 +765,13 @@ jobs:
     runs-on: ubuntu-latest
     needs: [Nemo_CICD_Test, pre-flight]
     if: |
-      (
+      always()
+      && !cancelled()
+      && needs.pre-flight.outputs.is_ci_workload == 'false'
+      && (
         needs.pre-flight.outputs.docs_only == 'true'
         || needs.pre-flight.outputs.is_deployment_workflow == 'true'
       )
-      && needs.pre-flight.outputs.is_ci_workload == 'false'
-      && !cancelled()
     steps:
       - name: Generate fake coverage report
         uses: actions/github-script@v6
@@ -788,10 +791,10 @@ jobs:
     runs-on: ubuntu-latest
     needs: [Nemo_CICD_Test, pre-flight]
     if: |
-      (
-        (needs.pre-flight.outputs.is_ci_workload == 'true' && !failure())
-        || success()
-      )
+      needs.Nemo_CICD_Test.result == 'success'
+      && needs.pre-flight.outputs.docs_only == 'false'
+      && needs.pre-flight.outputs.is_deployment_workflow == 'false'
+      && github.event.inputs.mcore_ref == ''
       && !cancelled()
     strategy:
       matrix:

.github/workflows/install-test.yml

@@ -57,10 +57,11 @@ concurrency:
 
 jobs:
   pre-flight:
-    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.73.2
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.80.1
     with:
       default_runner_prefix: ${{ vars.INSTALL_TEST_DEFAULT_RUNNER_PREFIX }}
       non_nvidia_runner_prefix: ${{ vars.INSTALL_TEST_NON_NVIDIA_RUNNER_PREFIX }}
+      sso_users_filename: ${{ vars.SSO_USERS_FILENAME }}
     secrets:
       NVIDIA_MANAGEMENT_ORG_PAT: ${{ secrets.NVIDIA_MANAGEMENT_ORG_PAT }}
 

.github/workflows/release.yaml

@@ -56,12 +56,13 @@ permissions:
 
 jobs:
   pre-flight:
-    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.73.2
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.80.1
     with:
       default_runner_prefix: ${{ vars.DEFAULT_RUNNER_PREFIX }}
       non_nvidia_runner_prefix: ${{ vars.NON_NVIDIA_RUNNER_PREFIX }}
       default_test_data_path: ${{ vars.DEFAULT_TEST_DATA_PATH }}
       non_nvidia_test_data_path: ${{ vars.NON_NVIDIA_TEST_DATA_PATH }}
+      sso_users_filename: ${{ vars.SSO_USERS_FILENAME }}
     secrets:
       NVIDIA_MANAGEMENT_ORG_PAT: ${{ secrets.NVIDIA_MANAGEMENT_ORG_PAT }}
 

.main.commit

@@ -1 +1 @@
-905c0e38655a61e603529b972ca9fdffd7daa667
+9ed8b0c4ac8e65d80fb35eecb3656c96b30df943

3rdparty/Megatron-LM

@@ -2,16 +2,16 @@
 
 ## Repository Overview
 
-- Megatron-Bridge (/opt/Megatron-Bridge/)
-- Megatron-LM (/opt/Megatron-Bridge/3rdparty/Megatron-LM/)
-- Evaluator (/opt/Evaluator/)
-- Export-Deploy (/opt/Export-Deploy/)
-- NeMo (/opt/NeMo/)
-- Run (/opt/Run/)
+- Megatron-Bridge (`/opt/Megatron-Bridge/`)
+- Megatron-LM (`/opt/Megatron-Bridge/3rdparty/Megatron-LM/`)
+- Evaluator (`/opt/Evaluator/`)
+- Export-Deploy (`/opt/Export-Deploy/`)
+- NeMo (`/opt/NeMo/`)
+- Run (`/opt/Run/`)
 
 ---
 
-## Installed Packged
+## Installed Packages
 
 Pip installed packages:
 
@@ -29,7 +29,7 @@ Uv virtualenv installed packages (/opt/venv/)
 
 Execute `uv pip list` to see full list of installed packages. Packages installed in /opt/venv take precedence over pip installed packages.
 
-See /opt/NeMo-FW/pyproject.toml for additonal uv configurations.
+See `/opt/NeMo-FW/pyproject.toml` for additonal uv configurations.
 
 ---
 
@@ -45,38 +45,55 @@ docker run -v <local-folder-path>:<container-folder-path> <container-image>
 
 #### 1 Overwrite existing directory
 
-1. Define the `<container-folder-path>` to correspond to the path defined in Repository Overview. (ie. /my-path/Megatron-Bridge/:/opt/Megatron-Bridge/)
+1. Define the `<container-folder-path>` to correspond to the path defined in Repository Overview. (ie. `/my-path/Megatron-Bridge/:/opt/Megatron-Bridge/`)
 
 2. `cd /opt/NeMo-FW` and run `uv sync --no-cache-dir --all-groups --inexact`
 
 3. Local development directory is synced to run in the container
 
 #### 2 Mount directory to a different path
 
-1. Modify `/opt/NeMo-FW/pyproject.toml` sections `tool.uv.sources` and `tool.uv; override-dependencies` to reflect new path. (ie. "megatron-bridge" = { path = `<container-folder-path>`, editable = true })
+1. Modify `/opt/NeMo-FW/pyproject.toml` sections `tool.uv.sources` and `tool.uv; override-dependencies` to reflect new path. (ie. `"megatron-bridge" = { path = <container-folder-path>, editable = true }`)
 
 2. `cd /opt/NeMo-FW` and run `uv sync --no-cache-dir --all-groups --inexact`
 
 3. Local development directory is synced to run in the container
 
 ### Installing packages inside the container
 
-Packages can be installed into the uv virtualenv (`/opt/venv/`) from two locations:
+All packages share a single uv virtualenv (`/opt/venv/`). The *location* you install
+from determines which resolution rules apply — use the guide below.
 
-#### From `/opt/NeMo-FW` (default)
+#### Which code needs this package?
 
-Use this when unsure, or when the dependency is not specific to Megatron:
+**Megatron-Bridge/Megatron-LM** (training, model architecture):
 
 ```bash
-cd /opt/NeMo-FW
+cd /opt/Megatron-Bridge
 uv pip install <package>
 ```
 
-#### From `/opt/Megatron-Bridge`
+Examples of Megatron-Bridge managed packages are TransformerEngine, nvidia-resiliency-ext, nvidia-modelopt. For a complete list, please visit `/opt/Megatron-Bridge/pyproject.toml`. Any of these dependencies can be re-installed via this directory.  
+
+It is unsafe to install packages that are not part of Megatron-Bridge from within this directory. For changing those dependencies, please visit the next section.  
 
-Use this only for Megatron-specific dependencies:
+**NeMo toolkit, Export-Deploy, Run, or Evaluator**:
 
 ```bash
-cd /opt/Megatron-Bridge
+cd /opt/NeMo-FW
 uv pip install <package>
 ```
+
+This is safe to use for general packages. This directory also prevents accidental re-install of heavy dependencies like trt-llm or vllm. MBridge-managed packages
+(TransformerEngine, nvidia-resiliency-ext, nvidia-modelopt, etc.) are protected
+and will not be overwritten. For reinstalling those, please visit the next section.
+
+**vllm, tensorrt-llm, or tensorrt**:
+
+These are built from source and baked into the container at build time (see
+`docker/Dockerfile.fw_base`). They cannot be managed via `uv pip install`.
+To change them, rebuild the container.
+
+**general note:**
+
+Running uv pip install outside any of the two directories above might lead to a re-install of torch, thus breaking all dependencies that have been compiled against the original torch version. By running uv pip install inside any of the two directories, we can avoid this unwanted side-effect.

docker/common/README.md

@@ -108,10 +108,14 @@ Note:
     - `gemma2_2b_pretrain_config`: Pre-training configuration for Gemma 2 2B
     - `gemma2_9b_pretrain_config`: Pre-training configuration for Gemma 2 9B
     - `gemma2_27b_pretrain_config`: Pre-training configuration for Gemma 2 27B
-  - **Finetuning:**
-    - `gemma2_2b_finetune_config`: Finetuning configuration for Gemma 2 2B with PEFT support (LoRA, DoRA)
-    - `gemma2_9b_finetune_config`: Finetuning configuration for Gemma 2 9B with PEFT support (LoRA, DoRA)
-    - `gemma2_27b_finetune_config`: Finetuning configuration for Gemma 2 27B with PEFT support (LoRA, DoRA)
+  - **SFT:**
+    - `gemma2_2b_sft_config`: Full SFT for Gemma 2 2B
+    - `gemma2_9b_sft_config`: Full SFT for Gemma 2 9B
+    - `gemma2_27b_sft_config`: Full SFT for Gemma 2 27B
+  - **PEFT** (LoRA, DoRA):
+    - `gemma2_2b_peft_config`: PEFT for Gemma 2 2B
+    - `gemma2_9b_peft_config`: PEFT for Gemma 2 9B
+    - `gemma2_27b_peft_config`: PEFT for Gemma 2 27B
 
 Before training, ensure the following environment variables are set:
 1. `SAVE_DIR`: checkpoint and log saving directory
@@ -164,21 +168,19 @@ config = gemma2_27b_pretrain_config(
 ```bash
 torchrun --nproc-per-node=8 run/run_recipe.py \
 --pretrained-checkpoint /models/gemma-2-2b \
---recipe gemma2_2b_finetune_config \
---peft none \
+--recipe gemma2_2b_sft_config \
 train.global_batch_size=64 \
 train.train_iters=1000 \
 checkpoint.save=$SAVE_DIR/gemma2_2b_finetune
 ```
 
 Or programmatically:
 ```python
-from megatron.bridge.recipes.gemma import gemma2_2b_finetune_config
+from megatron.bridge.recipes.gemma import gemma2_2b_sft_config
 
-config = gemma2_2b_finetune_config(
+config = gemma2_2b_sft_config(
     name="gemma2_2b_full_finetune",
     pretrained_checkpoint="/models/gemma-2-2b",
-    peft="none",
     train_iters=1000,
     global_batch_size=64,
 )
@@ -188,8 +190,7 @@ config = gemma2_2b_finetune_config(
 ```bash
 torchrun --nproc-per-node=8 run/run_recipe.py \
 --pretrained-checkpoint /models/gemma-2-9b \
---recipe gemma2_9b_finetune_config \
---peft none \
+--recipe gemma2_9b_sft_config \
 train.global_batch_size=64 \
 train.train_iters=1000 \
 checkpoint.save=$SAVE_DIR/gemma2_9b_finetune
@@ -199,8 +200,7 @@ checkpoint.save=$SAVE_DIR/gemma2_9b_finetune
 ```bash
 torchrun --nproc-per-node=16 run/run_recipe.py \
 --pretrained-checkpoint /models/gemma-2-27b \
---recipe gemma2_27b_finetune_config \
---peft none \
+--recipe gemma2_27b_sft_config \
 train.global_batch_size=64 \
 train.train_iters=1000 \
 checkpoint.save=$SAVE_DIR/gemma2_27b_finetune
@@ -212,50 +212,50 @@ checkpoint.save=$SAVE_DIR/gemma2_27b_finetune
 ```bash
 torchrun --nproc-per-node=8 run/run_recipe.py \
 --pretrained-checkpoint /models/gemma-2-2b \
---recipe gemma2_2b_finetune_config \
---peft lora \
+--recipe gemma2_2b_peft_config \
+--peft_scheme lora \
 train.global_batch_size=128 \
 checkpoint.save=$SAVE_DIR/gemma2_2b_lora
 ```
 
 PEFT options:
-- `--peft`: Set to `lora` for LoRA or `dora` for DoRA. Set to `none` for full finetuning.
+- `--peft_scheme`: Set to `lora` for LoRA or `dora` for DoRA. Use `gemma2_*_sft_config` for full finetuning.
 
 Or programmatically:
 ```python
-from megatron.bridge.recipes.gemma import gemma2_2b_finetune_config
+from megatron.bridge.recipes.gemma import gemma2_2b_peft_config
 
 # LoRA finetuning
-config = gemma2_2b_finetune_config(
+config = gemma2_2b_peft_config(
     name="gemma2_2b_lora_finetune",
     pretrained_checkpoint="/models/gemma-2-2b",
-    peft="lora",  # or "dora"
+    peft_scheme="lora",  # or "dora"
     train_iters=1000,
     global_batch_size=128,
 )
 ```
 
 #### Gemma 2 9B LoRA
 ```python
-from megatron.bridge.recipes.gemma import gemma2_9b_finetune_config
+from megatron.bridge.recipes.gemma import gemma2_9b_peft_config
 
-config = gemma2_9b_finetune_config(
+config = gemma2_9b_peft_config(
     name="gemma2_9b_lora_finetune",
     pretrained_checkpoint="/models/gemma-2-9b",
-    peft="lora",
+    peft_scheme="lora",
     train_iters=1000,
     global_batch_size=128,
 )
 ```
 
 #### Gemma 2 27B LoRA
 ```python
-from megatron.bridge.recipes.gemma import gemma2_27b_finetune_config
+from megatron.bridge.recipes.gemma import gemma2_27b_peft_config
 
-config = gemma2_27b_finetune_config(
+config = gemma2_27b_peft_config(
     name="gemma2_27b_lora_finetune",
     pretrained_checkpoint="/models/gemma-2-27b",
-    peft="lora",
+    peft_scheme="lora",
     train_iters=1000,
     global_batch_size=128,
 )