Skip to content

Commit 03dee88

Browse files
akoumpaakoumpa
committed
Merge branch 'main' into akoumpa/feat_custom_lora_backward
Conflicts resolved: - checkpoint/checkpointing.py: main moved _is_geq_torch_2_9 / CheckpointingConfig / SaveConsolidatedMode into checkpoint/config.py and added dtype-mapping normalization; kept main's import-based structure and preserved the nemotron-singlegpu-lora module NOTE. Dropped the now-redundant local _is_geq_torch_2_9 definition. - models/nemotron_v3/model.py: kept both independent method sets — this branch's gradient_checkpointing_enable/disable and main's PP-stage helpers (#2316). - examples/llm_finetune/gpt_oss/gpt_oss_20b_single_gpu.yaml: kept deleted (this branch removed it together with its test-config references; nothing in the merged tree references it — only the _peft variant remains). Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
2 parents 244f0ca + e2b1227 commit 03dee88

436 files changed

Lines changed: 31946 additions & 5127 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.agents/skills‎

Lines changed: 0 additions & 1 deletion
This file was deleted.

‎.claude/skills/skills‎

Lines changed: 0 additions & 1 deletion
This file was deleted.

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

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,12 @@ jobs:
5959
uses: astral-sh/setup-uv@v1
6060
with:
6161
version: 0.8.22
62+
- name: Free up disk space
63+
run: |
64+
sudo rm -rf /usr/local/lib/android || true
65+
sudo rm -rf /usr/local/.ghcup || true
66+
sudo rm -rf /usr/local/lib/node_modules || true
67+
pip cache purge || true
6268
- name: Install ruff
6369
env:
6470
UV_PROJECT_ENVIRONMENT: ./venv
@@ -69,12 +75,6 @@ jobs:
6975
export PATH="./bin/:$PATH"
7076
7177
uv sync --link-mode copy --locked --group linting
72-
- name: Free up disk space
73-
run: |
74-
sudo rm -rf /usr/local/lib/android || true
75-
sudo rm -rf /usr/local/.ghcup || true
76-
sudo rm -rf /usr/local/lib/node_modules || true
77-
pip cache purge || true
7878
- name: Run ruff
7979
run: |
8080
source ./venv/bin/activate
@@ -130,6 +130,12 @@ jobs:
130130
uses: astral-sh/setup-uv@v1
131131
with:
132132
version: 0.8.22
133+
- name: Free up disk space
134+
run: |
135+
sudo rm -rf /usr/local/lib/android || true
136+
sudo rm -rf /usr/local/.ghcup || true
137+
sudo rm -rf /usr/local/lib/node_modules || true
138+
pip cache purge || true
133139
- name: Install ty
134140
env:
135141
UV_PROJECT_ENVIRONMENT: ./venv

‎.github/workflows/install-test.yml‎

Lines changed: 29 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -11,10 +11,24 @@ concurrency:
1111
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-${{ github.event.label.name || 'main' }}-${{ github.event_name }}
1212
cancel-in-progress: true
1313

14+
env:
15+
TORCH_CU_INDEX: https://download.pytorch.org/whl/cu130
16+
CUDA_CONTAINER_IMAGE: nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04
17+
1418
jobs:
1519
pre-flight:
1620
uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.64.2
1721

22+
setup-install-image:
23+
runs-on: ubuntu-latest
24+
outputs:
25+
cuda_image: ${{ steps.cfg.outputs.cuda_image }}
26+
steps:
27+
- id: cfg
28+
shell: bash
29+
run: |
30+
echo "cuda_image=${CUDA_CONTAINER_IMAGE}" >> "$GITHUB_OUTPUT"
31+
1832
pip-test:
1933
runs-on: ${{ matrix.arch }}
2034
name: Pip - Python${{ matrix.python-version }} - ${{ matrix.arch == 'ubuntu-latest' && 'AMD64/Linux' || (matrix.arch == 'ubuntu-24.04-arm' && 'ARM64/Linux' || 'ARM64/Darwin') }} - No CUDA
@@ -60,12 +74,12 @@ jobs:
6074
ngc-cuda-test-uv:
6175
runs-on: linux-amd64-cpu16
6276
name: UV - Python${{ matrix.python-version }} - AMD64/Linux - NGC CUDA
63-
needs: [pre-flight]
77+
needs: [pre-flight, setup-install-image]
6478
if: |
6579
!(needs.pre-flight.outputs.docs_only == 'true'
6680
|| needs.pre-flight.outputs.is_deployment_workflow == 'true')
6781
container:
68-
image: nvcr.io/nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04
82+
image: ${{ needs.setup-install-image.outputs.cuda_image }}
6983
environment: nemo-ci
7084
strategy:
7185
fail-fast: false
@@ -142,9 +156,9 @@ jobs:
142156
runs-on: linux-amd64-cpu16
143157
name: Build CUDA wheelhouse - Python${{ matrix.python-version }} - AMD64/Linux - NGC CUDA
144158
container:
145-
image: nvcr.io/nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04
159+
image: ${{ needs.setup-install-image.outputs.cuda_image }}
146160
environment: nemo-ci
147-
needs: [pre-flight]
161+
needs: [pre-flight, setup-install-image]
148162
if: |
149163
!(needs.pre-flight.outputs.docs_only == 'true'
150164
|| needs.pre-flight.outputs.is_deployment_workflow == 'true')
@@ -183,7 +197,7 @@ jobs:
183197
. ./venv/bin/activate
184198
185199
pip install --upgrade pip
186-
pip install --index-url https://download.pytorch.org/whl/cu128 "torch<=2.8.0"
200+
pip install --index-url ${TORCH_CU_INDEX} "torch>=2.9.0"
187201
pip install numpy packaging psutil pybind11 setuptools wheel wheel_stub
188202
189203
mkdir -p "${WHEELHOUSE_DIR}"
@@ -192,7 +206,7 @@ jobs:
192206
mamba-ssm \
193207
nv-grouped-gemm \
194208
transformer-engine-torch \
195-
"transformer-engine[pytorch]<=2.11.0"
209+
"transformer-engine[pytorch]>=2.14.1"
196210
197211
- name: Verify CUDA wheelhouse contents
198212
shell: bash -x -e -u -o pipefail {0}
@@ -216,17 +230,21 @@ jobs:
216230
runs-on: linux-amd64-cpu16
217231
name: Pip - Python${{ matrix.python-version }}${{ matrix.extra-groups != '' && format('[{0}]', matrix.extra-groups) || '' }} - AMD64/Linux - NGC CUDA
218232
container:
219-
image: nvcr.io/nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04
233+
image: ${{ needs.setup-install-image.outputs.cuda_image }}
220234
environment: nemo-ci
221-
needs: [pre-flight, cuda-wheelhouse]
235+
needs: [pre-flight, setup-install-image, cuda-wheelhouse]
222236
if: |
223237
!(needs.pre-flight.outputs.docs_only == 'true'
224238
|| needs.pre-flight.outputs.is_deployment_workflow == 'true')
225239
strategy:
226240
fail-fast: false
227241
matrix:
228242
python-version: ["3.12"]
229-
extra-groups: ["", "cuda", "vlm", "fa", "all"]
243+
# TEMP (PR #2397): "fa" removed from the matrix. The [fa] install compiles
244+
# flash-attn from source on every run (~2h+) and dominates CI wall-time.
245+
# Re-enable once a prebuilt flash-attn wheel is cached in the CUDA wheelhouse.
246+
# cc @ko3n1g @thomasdhc (oliver könig, Dong Hyuk Chang)
247+
extra-groups: ["", "cuda", "vlm", "all"]
230248
env:
231249
EXTRA: ${{ matrix.extra-groups != '' && format('[{0}]', matrix.extra-groups) || '' }}
232250
WHEELHOUSE_DIR: /tmp/cuda-wheelhouse
@@ -260,7 +278,7 @@ jobs:
260278
PIP_ARGS=()
261279
262280
if [[ $EXTRA == *"fa"* || $EXTRA == *"cuda"* || $EXTRA == *"all"* ]]; then
263-
pip install --find-links "${PIP_FIND_LINKS}" --index-url https://download.pytorch.org/whl/cu128 "torch<=2.8.0"
281+
pip install --find-links "${PIP_FIND_LINKS}" --index-url ${TORCH_CU_INDEX} "torch>=2.9.0"
264282
pip install --find-links "${PIP_FIND_LINKS}" numpy packaging psutil pybind11 setuptools wheel wheel_stub
265283
PIP_ARGS=(--no-build-isolation)
266284
export TORCH_CUDA_ARCH_LIST="9.0 10.0 12.0"
@@ -335,7 +353,7 @@ jobs:
335353
python-binary: ./venv/bin/python
336354

337355
install-test-summary:
338-
needs: [pip-test, uv-test, ngc-cuda-test-uv, cuda-wheelhouse, ngc-cuda-test-pip, pre-flight]
356+
needs: [pip-test, uv-test, ngc-cuda-test-uv, cuda-wheelhouse, ngc-cuda-test-pip, pre-flight, setup-install-image]
339357
runs-on: ubuntu-latest
340358
name: Install test summary
341359
if: |

‎.github/workflows/release-freeze.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ on:
3535

3636
jobs:
3737
code-freeze:
38-
uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_code_freeze.yml@v0.86.0
38+
uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_code_freeze.yml@v1.4.2
3939
with:
4040
library-name: NeMo-Automodel
4141
python-package: nemo_automodel

‎.github/workflows/release.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ jobs:
8080
!cancelled() && !failure()
8181
&& !(needs.pre-flight.outputs.docs_only == 'true'
8282
|| needs.pre-flight.outputs.is_deployment_workflow == 'true')
83-
uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_release_library.yml@v1.4.0
83+
uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_release_library.yml@v1.4.3
8484
with:
8585
release-ref: ${{ inputs.release-ref || github.sha }}
8686
python-package: nemo_automodel

‎.github/workflows/uv-lock-generation.yml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ on:
1818
pull_request:
1919
paths:
2020
- 'pyproject.toml'
21+
- 'docker/common/uv-pytorch.toml'
2122

2223
jobs:
2324
check_uv_lock_and_update:

‎README.md‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
</div>
2222

2323
## 📣 News and Discussions
24+
- [06/04/2026][**Nemotron-3 Ultra**](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16) We now support finetuning NVIDIA's Nemotron 3 Ultra 550B A55B. Check out our [recipe](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/nemotron/nemotron_ultra_v3_hellaswag_peft.yaml) and [guide](https://github.com/NVIDIA-NeMo/Automodel/blob/main/docs/guides/llm/nemotron-3-ultra.md).
25+
- [06/03/2026][**Gemma 4 12B**](https://huggingface.co/google/gemma-4-12B) We now support finetuning the dense `google/gemma-4-12B` model. Check out our [recipe](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/gemma/gemma_4_12b_hellaswag.yaml).
2426
- [05/27/2026][**Step-3.7-Flash**](https://huggingface.co/stepfun-ai/Step-3.7-Flash) We added model coverage for Stepfun AI's 198B-A13B MoE vision-language model, targeting image/video agentic developer workflows with a 256k context language backbone and 1.8B ViT vision tower. See the [model coverage page](https://github.com/NVIDIA-NeMo/Automodel/blob/main/docs/model-coverage/vlm/stepfun-ai/step-3-7.md).
2527
- [05/19/2026][**Ling 2.0**](https://huggingface.co/collections/inclusionAI/ling-20) We now support finetuning the inclusionAI Ling 2.0 MoE family (`inclusionAI/Ling-mini-2.0`, `inclusionAI/Ling-flash-2.0`, and `inclusionAI/Ling-1T`), thanks to [@Hayden727](https://github.com/Hayden727). Check out our [recipes](https://github.com/NVIDIA-NeMo/Automodel/tree/main/examples/llm_finetune/ling).
2628
- [05/17/2026][**ERNIE 4.5**](https://huggingface.co/baidu) and [**MiMo-V2-Flash**](https://huggingface.co/XiaomiMiMo/MiMo-V2-Flash) We now support finetuning `baidu/ERNIE-4.5-0.3B-PT`, `baidu/ERNIE-4.5-21B-A3B-PT`, and `XiaomiMiMo/MiMo-V2-Flash`. Check out our ERNIE [dense recipe](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/ernie4_5/ernie4_5_0p3b_hellaswag.yaml), ERNIE [MoE recipe](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/ernie4_5/ernie4_5_21b_a3b_hellaswag.yaml), and MiMo [recipe](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/mimo_v2_flash/mimo_v2_flash_hellaswag.yaml).
@@ -362,7 +364,7 @@ YAML sketch:
362364
checkpoint:
363365
enabled: true
364366
checkpoint_dir: ./checkpoints
365-
save_consolidated: true
367+
save_consolidated: final
366368
model_save_format: safetensors
367369
``` -->
368370

‎docker/Dockerfile‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
ARG CUDA_IMAGE=nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04
15+
ARG CUDA_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04
1616
ARG PYTORCH_IMAGE=nvcr.io/nvidia/pytorch:26.04-py3
1717
ARG BASE_IMAGE=cuda
1818

@@ -62,8 +62,8 @@ RUN if [ -f /usr/local/bin/torchrun ]; then \
6262
FROM update_base_container AS automodel_dep
6363

6464
# Install TE
65-
ARG INSTALL_TE=False
66-
ARG TE_COMMIT=release_v2.11
65+
ARG INSTALL_TE=True
66+
ARG TE_COMMIT=release_v2.14
6767
RUN if [ "$INSTALL_TE" = "True" ]; then \
6868
git clone https://github.com/NVIDIA/TransformerEngine.git && \
6969
cd TransformerEngine && \

‎docker/common/uv-pytorch.lock‎

Lines changed: 22 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -74,24 +74,30 @@ constraints = [
7474
{ name = "wheel", specifier = ">=0.46.2" },
7575
]
7676
overrides = [
77+
{ name = "cuda-bindings", marker = "sys_platform == 'never'" },
78+
{ name = "cuda-pathfinder", marker = "sys_platform == 'never'" },
7779
{ name = "flash-attn", marker = "sys_platform == 'never'" },
78-
{ name = "nvidia-cublas-cu12", marker = "sys_platform == 'never'" },
79-
{ name = "nvidia-cuda-cupti-cu12", marker = "sys_platform == 'never'" },
80-
{ name = "nvidia-cuda-nvrtc-cu12", marker = "sys_platform == 'never'" },
81-
{ name = "nvidia-cuda-runtime-cu12", marker = "sys_platform == 'never'" },
82-
{ name = "nvidia-cudnn-cu12", marker = "sys_platform == 'never'" },
83-
{ name = "nvidia-cufft-cu12", marker = "sys_platform == 'never'" },
84-
{ name = "nvidia-cufile-cu12", marker = "sys_platform == 'never'" },
85-
{ name = "nvidia-curand-cu12", marker = "sys_platform == 'never'" },
86-
{ name = "nvidia-cusolver-cu12", marker = "sys_platform == 'never'" },
87-
{ name = "nvidia-cusparse-cu12", marker = "sys_platform == 'never'" },
88-
{ name = "nvidia-cusparselt-cu12", marker = "sys_platform == 'never'" },
89-
{ name = "nvidia-nccl-cu12", marker = "sys_platform == 'never'" },
80+
{ name = "nvidia-cublas", marker = "sys_platform == 'never'" },
81+
{ name = "nvidia-cuda-cupti", marker = "sys_platform == 'never'" },
82+
{ name = "nvidia-cuda-nvrtc", marker = "sys_platform == 'never'" },
83+
{ name = "nvidia-cuda-runtime", marker = "sys_platform == 'never'" },
84+
{ name = "nvidia-cudnn-cu13", marker = "sys_platform == 'never'" },
85+
{ name = "nvidia-cufft", marker = "sys_platform == 'never'" },
86+
{ name = "nvidia-cufile", marker = "sys_platform == 'never'" },
87+
{ name = "nvidia-curand", marker = "sys_platform == 'never'" },
88+
{ name = "nvidia-cusolver", marker = "sys_platform == 'never'" },
89+
{ name = "nvidia-cusparse", marker = "sys_platform == 'never'" },
90+
{ name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'never'" },
91+
{ name = "nvidia-nccl-cu13", marker = "sys_platform == 'never'" },
92+
{ name = "nvidia-nvjitlink", marker = "sys_platform == 'never'" },
93+
{ name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'never'" },
94+
{ name = "nvidia-nvtx", marker = "sys_platform == 'never'" },
9095
{ name = "torch", marker = "sys_platform == 'never'", index = "https://download.pytorch.org/whl/cpu" },
9196
{ name = "torchao", marker = "sys_platform == 'never'" },
9297
{ name = "torchcodec", marker = "sys_platform == 'never'" },
9398
{ name = "torchvision", marker = "sys_platform == 'never'", index = "https://download.pytorch.org/whl/cpu" },
9499
{ name = "transformer-engine", marker = "sys_platform == 'never'" },
100+
{ name = "transformer-engine-cu13", marker = "sys_platform == 'never'" },
95101
{ name = "transformer-engine-torch", marker = "sys_platform == 'never'" },
96102
{ name = "triton", marker = "sys_platform == 'never'" },
97103
]
@@ -3603,14 +3609,14 @@ requires-dist = [
36033609
{ name = "timm", marker = "extra == 'vlm'", specifier = "<=1.0.22" },
36043610
{ name = "torch", marker = "sys_platform != 'darwin' and sys_platform != 'linux'", specifier = ">=2.6.0", index = "https://download.pytorch.org/whl/cpu" },
36053611
{ name = "torch", marker = "sys_platform == 'darwin'", specifier = ">=2.6.0", index = "https://pypi.org/simple" },
3606-
{ name = "torch", marker = "sys_platform == 'linux'", specifier = ">=2.6.0", index = "https://download.pytorch.org/whl/cu129" },
3612+
{ name = "torch", marker = "sys_platform == 'linux'", specifier = ">=2.6.0", index = "https://download.pytorch.org/whl/cu130" },
36073613
{ name = "torchao" },
36083614
{ name = "torchcodec", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin' and extra == 'vlm'" },
36093615
{ name = "torchdata" },
36103616
{ name = "torchvision", marker = "sys_platform == 'darwin' and extra == 'diffusion'", index = "https://pypi.org/simple" },
36113617
{ name = "torchvision", marker = "sys_platform != 'darwin' and sys_platform != 'linux' and extra == 'diffusion'", index = "https://download.pytorch.org/whl/cpu" },
3612-
{ name = "torchvision", marker = "sys_platform == 'linux' and extra == 'diffusion'", index = "https://download.pytorch.org/whl/cu129" },
3613-
{ name = "transformer-engine", extras = ["pytorch"], marker = "extra == 'cuda'", specifier = "<=2.11.0" },
3618+
{ name = "torchvision", marker = "sys_platform == 'linux' and extra == 'diffusion'", index = "https://download.pytorch.org/whl/cu130" },
3619+
{ name = "transformer-engine", extras = ["pytorch"], marker = "extra == 'cuda'", specifier = ">=2.14.1" },
36143620
{ name = "transformers", specifier = "==5.8.1" },
36153621
{ name = "wandb", specifier = ">=0.26.1" },
36163622
]
@@ -3623,7 +3629,7 @@ build = [
36233629
{ name = "setuptools" },
36243630
{ name = "torch", marker = "sys_platform != 'darwin' and sys_platform != 'linux'", specifier = "<=2.10.0", index = "https://download.pytorch.org/whl/cpu" },
36253631
{ name = "torch", marker = "sys_platform == 'darwin'", specifier = "<=2.10.0", index = "https://pypi.org/simple" },
3626-
{ name = "torch", marker = "sys_platform == 'linux'", specifier = "<=2.10.0", index = "https://download.pytorch.org/whl/cu129" },
3632+
{ name = "torch", marker = "sys_platform == 'linux'", specifier = "<=2.10.0", index = "https://download.pytorch.org/whl/cu130" },
36273633
]
36283634
dev = [
36293635
{ name = "cut-cross-entropy", git = "https://github.com/apple/ml-cross-entropy.git?rev=87a86aba72cfd2f0d8abecaf81c13c4528ea07d8" },

0 commit comments

Comments
 (0)