Merge branch 'main' into jwilber/docs-chatbot

Author: jwilber

3rdparty/Megatron-LM

@@ -29,6 +29,7 @@
 # @sichu2023 - Simon Chu
 # @skothenhill-nv - Steven Kothen-Hill
 # @trvachov - Timur Rvachov
+# @tshimko-nv - Tyler Shimko
 # @yzhang123 - Yang Zhang
 
 # TODO: make this a team of bionemo-core contributors
@@ -44,8 +45,8 @@ license_header @dorotat-nv @jstjohn @malcolmgreaves @trvachov
 #
 ## DOCUMENTATION
 #
-**.md @dorotat-nv @jstjohn @malcolmgreaves @pstjohn @trvachov @jwilber
-docs @dorotat-nv @jstjohn @malcolmgreaves @pstjohn @trvachov @jwilber
+**.md @dorotat-nv @jstjohn @malcolmgreaves @pstjohn @trvachov @jwilber @tshimko-nv
+docs @dorotat-nv @jstjohn @malcolmgreaves @pstjohn @trvachov @jwilber @tshimko-nv
 
 
 #

3rdparty/NeMo

@@ -19,7 +19,7 @@
 #   https://gitlab-master.nvidia.com/dl/JoC/nemo-ci/-/blob/main/.gitlab-ci.yml
 #  We should keep versions in our container up to date to ensure that we get the latest tested perf improvements and
 #   training loss curves from NeMo.
-ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:25.01-py3
+ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:25.04-py3
 
 FROM rust:1.86.0 AS rust-env
 
@@ -56,6 +56,13 @@ apt-get upgrade -qyy \
 rm -rf /tmp/* /var/tmp/*
 EOF
 
+
+## BUMP TE as a solution to the issue https://github.com/NVIDIA/bionemo-framework/issues/422. Drop this when pytorch images ship the fixed commit.
+ ARG TE_TAG=9d4e11eaa508383e35b510dc338e58b09c30be73
+ RUN PIP_CONSTRAINT= NVTE_FRAMEWORK=pytorch NVTE_WITH_USERBUFFERS=1 MPI_HOME=/usr/local/mpi \
+    pip --disable-pip-version-check --no-cache-dir install \
+    git+https://github.com/NVIDIA/TransformerEngine.git@${TE_TAG}
+
 # Install AWS CLI based on architecture
 RUN if [ "$TARGETARCH" = "arm64" ]; then \
       curl "https://awscli.amazonaws.com/awscli-exe-linux-aarch64.zip" -o "awscliv2.zip"; \
@@ -68,6 +75,7 @@ RUN if [ "$TARGETARCH" = "arm64" ]; then \
     ./aws/install && \
     rm -rf aws awscliv2.zip
 
+
 # Use a branch of causal_conv1d while the repository works on Blackwell support.
 ARG CAUSAL_CONV_TAG=52e06e3d5ca10af0c7eb94a520d768c48ef36f1f
 RUN CAUSAL_CONV1D_FORCE_BUILD=TRUE pip --disable-pip-version-check --no-cache-dir install git+https://github.com/trvachov/causal-conv1d.git@${CAUSAL_CONV_TAG}
@@ -123,6 +131,12 @@ fi
 ###############################################################################
 # /end ARM
 ###############################################################################
+# Fix the version of scikit-misc to 0.3.1 because newer versions of scikit-misc require numpy >= 2.0 to be built.
+# Since there are not pre-built wheels for arm64, we need to install this specific version.
+# Once bionemo is compatible with numpy >= 2.0, we can remove this.
+# Technically, this is only needed for the ARM build, but we apply to all architectures to avoid library version
+# divergence.
+RUN pip install scikit-misc==0.3.1
 
 # Mamba dependancy installation
 RUN pip --disable-pip-version-check --no-cache-dir install \
@@ -188,14 +202,16 @@ rm -rf nvidia-resiliency-ext/
 sed -i "/ngcsdk/d" ./sub-packages/bionemo-core/pyproject.toml
 # Remove llama-index because bionemo doesn't use it and it adds CVEs to container
 sed -i "/llama-index/d" ./3rdparty/NeMo/requirements/requirements_nlp.txt
+# Pin 'nvidia-modelopt' to 0.27.1 due to an API incompatibility of version 0.25.0
+sed -i -E "s|nvidia-modelopt\[torch\]>=[^,]+,<=([^ ;]+)|nvidia-modelopt[torch]==\1|" ./3rdparty/NeMo/requirements/requirements_nlp.txt
 uv pip install --no-build-isolation \
 ./3rdparty/*  \
 ./sub-packages/bionemo-* \
 -r /requirements-cve.txt \
 -r /requirements-test.txt
 
 # Install back ngcsdk, as a WAR for the protobuf version conflict with nemo_toolkit.
-uv pip install ngcsdk
+uv pip install ngcsdk==3.64.3  # Temporary fix for changed filename, see https://nvidia.slack.com/archives/C074Z808N05/p1746231345981209
 
 # Addressing security scan issue - CVE vulnerability https://github.com/advisories/GHSA-g4r7-86gm-pgqc The package is a
 # dependency of lm_eval from NeMo requirements_eval.txt. We also remove zstandard, another dependency of lm_eval, which
@@ -322,6 +338,8 @@ COPY ./docs ./docs
 COPY --from=rust-env /usr/local/cargo /usr/local/cargo
 COPY --from=rust-env /usr/local/rustup /usr/local/rustup
 
+# Fix a CRIT vuln: https://github.com/advisories/GHSA-vqfr-h8mv-ghfj
+RUN uv pip install h11==0.16.0
 
 # RUN rm -rf /usr/local/cargo /usr/local/rustup
 RUN chmod 777 -R /workspace/bionemo2/

CODEOWNERS

@@ -1 +1 @@
-2.5
+2.6

Dockerfile

@@ -13,12 +13,12 @@ script_args:
   data_path: /data/cellxgene_scdl
   model: geneformer
   variant: train
-  config_name: geneformer_config
+  config_name: 10M
   precision: [bf16-mixed]
   nodes: [2]
   gpus: 8
   batch_size: 32
-  max_steps: 37000
+  max_steps: 30000
   lr: 0.001
   val_check_interval: 500
   acc_grad: 1
@@ -38,10 +38,12 @@ script: |-
     --resume-if-exists \
     --log-every-n-steps 50 \
     --lr ${lr} \
+    --create-tflops-callback \
     --create-tensorboard-logger \
     --result-dir=${tensorboard_dir} \
     --wandb-project ${wandb_project_name} \
     --wandb-job-type=${pipeline_label} \
+    --wandb-group=${model}_${variant}_${config_name}__${target} \
     --cosine-rampup-frac 0.004331629559040111 \
     --cosine-hold-frac 0.021658147795200554 \
     --accumulate-grad-batches ${acc_grad} \

VERSION

@@ -3,6 +3,7 @@ time_limit: 1800
 key_segments:
   # Modify keys to be renamed (str) or excluded (False) from run identifier. By default, all args under script_args are included.
   data_path: False
+  dfpnl: False
 script_args:
   # All arguments referenced in the script string must be specified here.
   # Arguments not referenced in the script string must have the 'arg' field specified.
@@ -17,24 +18,28 @@ script_args:
   stop_steps: 200
   gpus: 8
   acc_grad: 1
+  dfpnl: ""
   products:
     - nodes: 1
       batch_size: 16
       pp: 1
       tp: 1
-#      FIXME (broken pp): https://github.com/NVIDIA/bionemo-framework/issues/784
-#    - nodes: 2
-#      batch_size: 16
-#      pp: 2
-#      tp: 1
+      dfpnl: ""
+    - nodes: 2
+      batch_size: 16
+      pp: 2
+      tp: 1
+      dfpnl: "--decoder-first-pipeline-num-layers=17"
     - nodes: 2
       batch_size: 16
       pp: 1
       tp: 2
+      dfpnl: ""
     - nodes: 2
       batch_size: 16
       pp: 1
       tp: 1
+      dfpnl: ""
 script: |-
   WANDB_API_KEY=$BIONEMO_WANDB_API_KEY ${variant}_${model} \
     --train-cluster-path=${data_path}/train_clusters.parquet \
@@ -51,6 +56,7 @@ script: |-
     --min-seq-length=1024 \
     --max-seq-length=1024 \
     --num-layers=33 \
+    ${dfpnl} \
     --hidden-size=1280 \
     --num-attention-heads=20 \
     --ffn-hidden-size=5120 \

ci/benchmarks/partial-conv/geneformer_pretrain.yaml

@@ -13,7 +13,7 @@ script_args:
   data_path: /data/cellxgene_scdl
   model: geneformer
   variant: train
-  config_name: geneformer_config
+  config_name: 10M
   precision: [bf16-mixed]
   gpus:  8
   max_steps: 1000
@@ -41,6 +41,7 @@ script: |-
     --micro-batch-size ${batch_size} \
     --resume-if-exists \
     --log-every-n-steps 50 \
+    --create-tflops-callback \
     --lr ${lr} \
     --create-tensorboard-logger \
     --result-dir=${tensorboard_dir} \

ci/benchmarks/perf/esm2_pretrain.yaml

@@ -25,11 +25,13 @@ usage() {
 Usage: $(basename "$0") [OPTIONS]
 
 Options:
-    --skip-docs    Skip running tests in the docs directory
-    --no-nbval     Skip jupyter notebook validation tests
-    --skip-slow    Skip tests marked as slow (@pytest.mark.slow)
-    --only-slow    Only run tests marked as slow (@pytest.mark.slow)
+    --skip-docs         Skip running tests in the docs directory
+    --no-nbval          Skip jupyter notebook validation tests
+    --skip-slow         Skip tests marked as slow (@pytest.mark.slow)
+    --only-slow         Only run tests marked as slow (@pytest.mark.slow)
     --allow-no-tests    Allow sub-packages with no found tests (for example no slow tests if --only-slow is set)
+    --ignore-files      Skip files from tests using glob patterns (comma-separated, no spaces).
+                            Example: --ignore-files docs/*.ipynb,src/specific_test.py
 
 Note: Documentation tests (docs/) are only run when notebook validation
       is enabled (--no-nbval not set) and docs are not skipped
@@ -54,6 +56,8 @@ NO_NBVAL=false
 SKIP_SLOW=false
 ONLY_SLOW=false
 ALLOW_NO_TESTS=false
+# TODO(@cspades): Ignore this Evo2 notebook test, which has a tendency to leave a 32GB orphaned process in GPU.
+declare -a IGNORE_FILES=("sub-packages/bionemo-evo2/examples/fine-tuning-tutorial.ipynb")
 error=false
 
 # Parse command line arguments
@@ -64,6 +68,10 @@ while (( $# > 0 )); do
         --skip-slow) SKIP_SLOW=true ;;
         --only-slow) ONLY_SLOW=true ;;
         --allow-no-tests) ALLOW_NO_TESTS=true ;;
+        --ignore-files)
+            shift
+            IFS=',' read -ra IGNORE_FILES <<< "$1"
+            ;;
         -h|--help) usage ;;
         *) echo "Unknown option: $1" >&2; usage 1 ;;
     esac
@@ -82,6 +90,10 @@ PYTEST_OPTIONS=(
     --cov-append
     --cov-report=xml:coverage.xml
 )
+# Add multiple file ignores if specified
+for ignore_file in "${IGNORE_FILES[@]}"; do
+    PYTEST_OPTIONS+=(--ignore-glob="$ignore_file")
+done
 [[ "$NO_NBVAL" != true ]] && PYTEST_OPTIONS+=(--nbval-lax)
 [[ "$SKIP_SLOW" == true ]] && PYTEST_OPTIONS+=(-m "not slow")
 [[ "$ONLY_SLOW" == true ]] && PYTEST_OPTIONS+=(-m "slow")

ci/benchmarks/perf/geneformer_pretrain.yaml

@@ -21,7 +21,45 @@ docker run --rm -it -p 8000:8000 \
 And then navigate to [`http://0.0.0.0:8000`](http://0.0.0.0:8000) on your local
 machine.
 
+## Sub-package Documentation
+
+When adding documentation for a new sub-package, ensure it is properly integrated into the documentation site by:
+
+1. Adding an entry to `docs/docs/user-guide/examples/SUMMARY.md` to include it in the Tutorials section
+2. Adding an entry to `docs/docs/user-guide/developer-guide/SUMMARY.md` to include it in the Developer Guide section
+
+This ensures the sub-package documentation is properly indexed and accessible through the navigation menu.
+
+The sub-package specific documentation itself must be placed alongside the sub-package code in the `sub-packages/bionemo-<sub-package-name>/` directory:
+
+- `README.md` - A root level file that describes the sub-package and how to use it.
+- `examples/` - A directory that contains documentation or examples specific to the sub-package, in the form of `.md` or `.ipynb` files.
+- `assets/` - A folder that contains any static assets used in any of the above files, e.g. `.png` files.
+
+When the docs are built, these documentation files will be fetched (via the [scripts/gen_ref_pages.py](./scripts/gen_ref_pages.py) script) for rendering in the main documentation site.
+
+- The `README.md` will be rendered as an individual page in the `User Guide -> Developer Guide -> <sub-package-name>/` section of the documentation site.
+- Every file in the `examples/` directory will be rendered as an individual page in the `User Guide -> Tutorials -> <sub-package-name>/` section of the documentation site.
+
+An example sub-package structure is shown below:
+
+```
+bionemo-<sub-package-name>/
+└── assets/
+    ├── example_1.png
+├── examples/
+│   ├── example_1.md
+│   └── example_2.ipynb
+├── src/
+├── tests/
+├── LICENSE
+├── pyproject.toml
+├── README.md
+├── VERSION
+```
+
 ## Hiding/collapsing `.ipynb` cells
+
 To remove cells from the rendered `mkdocs` html you can add a `remove-cell` tag to the cell. Note that `remove-output` is also an option to hide outputs but not the code cell. Unfortunately
 `remove-input` does not seem to be supported.