Merge branch 'main' into chcui/iris/optimize-nemotron-omni

Signed-off-by: adityavavreNVDA <avavre@nvidia.com>

Author: adityavavreNVDA

.dev.commit

@@ -1 +1 @@
-b3b1e2d793c596c2fa61076d001ab93f294e4e87
+2f2f8ebaef90f4d62a18a43808398223b1fc5809

.github/scripts/test_docker_dependency_layers.sh

@@ -252,9 +252,18 @@ if ! grep -Fq \
   echo "Bridge's lock metadata must preserve the MCore TE extra" >&2
   exit 1
 fi
-if grep -q 'transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@' pyproject.toml || \
-  grep -q '^name = "transformer-engine"$' pyproject.toml; then
-  echo "Bridge must inherit the TransformerEngine source and metadata from the selected MCore ref" >&2
+if grep -q 'transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@' pyproject.toml; then
+  echo "Bridge must inherit the TransformerEngine source from the selected MCore ref" >&2
+  exit 1
+fi
+
+if ! grep -Fq '    NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD=1 \' "$dockerfile"; then
+  echo "CI builds must disable TransformerEngine's recursive build-time submodule fetch" >&2
+  exit 1
+fi
+if grep -R -qE 'git submodule update.*--recursive|git submodule update --init --recursive' \
+  "$dockerfile" docker .github/actions; then
+  echo "CI build surfaces must not fetch recursive submodules at build time" >&2
   exit 1
 fi
 

.main.commit

@@ -1 +1 @@
-2a75ac12c54ba6a42b34024756820bf819a2e25f
+731b791469004f8fdcb896e65d610d6b1b0ebb32

3rdparty/Megatron-LM

@@ -55,6 +55,7 @@ ENV PATH="/opt/venv/bin:/usr/local/bin:$PATH" \
     UV_VERSION="0.7.2" \
     UV_HTTP_TIMEOUT=120 \
     NVTE_BUILD_NUM_PHILOX_ROUNDS=3 \
+    NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD=1 \
     HYBRID_EP_MULTINODE=1 \
     RDMA_CORE_HOME=/opt/rdma-core/build \
     LD_LIBRARY_PATH="/usr/local/cuda/lib64/:$LD_LIBRARY_PATH"
@@ -192,7 +193,7 @@ COPY --chown=1000:1000 . /opt/Megatron-Bridge
 # prepared environment, and uses /nemo_run for runtime outputs. Keep those paths owned by the
 # existing unprivileged base-image user instead of relying on a root final user.
 RUN install -d -o 1000 -g 1000 -m 0755 /nemo_run /home/ubuntu/.cache && \
-    chown 1000:1000 /opt/Megatron-Bridge
+    chown -R 1000:1000 /opt/Megatron-Bridge
 
 ENV HOME=/home/ubuntu
 USER 1000:1000

docker/Dockerfile.ci

@@ -13,10 +13,12 @@
 # limitations under the License.
 
 # Build layer variables
-# Supported mode: FW_DEP_BUILDER=base, FW_BASE_FINAL=fw_toolkit_builder
+# Supported mode: FW_DEP_BUILDER=base, FW_BASE_FINAL=fw_toolkit_builder,
+# VLLM_WHEEL_SRC=vllm_wheel_build|vllm_wheel_none
 # TRT-LLM stages are not present in this Dockerfile.
 ARG FW_DEP_BUILDER
 ARG FW_BASE_FINAL
+ARG VLLM_WHEEL_SRC=vllm_wheel_build
 
 ARG NEMO_FW_BASE_IMAGE
 FROM ${NEMO_FW_BASE_IMAGE} AS base
@@ -31,7 +33,7 @@ ENV NVIDIA_PRODUCT_NAME="NeMo Framework" \
 ##
 ##############################################################################
 
-FROM base AS vllm_wheel
+FROM base AS vllm_wheel_build
 
 ARG VLLM_VERSION=v0.14.1
 ARG MAX_JOBS=4
@@ -47,6 +49,13 @@ RUN mkdir -p /src/vllm && \
     pip install -r requirements/build/cuda.txt && \
     pip wheel --no-deps --no-build-isolation -v .
 
+# /src/vllm must exist for the install step's bind mount.
+FROM base AS vllm_wheel_none
+
+RUN mkdir -p /src/vllm
+
+FROM ${VLLM_WHEEL_SRC} AS vllm_wheel
+
 ##############################################################################
 ##
 ## Top layer of FW_BASE container
@@ -127,10 +136,14 @@ RUN --mount=type=bind,source=docker/common/install_nsys.sh,target=/opt/install_n
 # Copy Wheels
 RUN --mount=type=bind,from=vllm_wheel,source=/src/vllm/,target=/tmp/vllm/ \
     --mount=type=bind,source=docker/patches/vllm.patch,target=/opt/vllm.patch \
+    if ! compgen -G "/tmp/vllm/vllm*.whl" >/dev/null; then \
+    echo "No vLLM wheel produced; skipping vLLM install."; \
+    else \
     pip install /tmp/vllm/vllm*.whl && \
     pushd /usr/local/lib/python3.12/dist-packages/vllm && \
     patch -p1 < /opt/vllm.patch && \
-    popd
+    popd; \
+    fi
 
 ##############################################################################
 ##

docker/Dockerfile.fw_base

@@ -132,6 +132,7 @@ docker build \
 | `FW_BASE_FINAL` | Output stage. `trtllm_install` (with TRT-LLM) or `fw_toolkit_builder` (without) |
 | `UV_VERSION` | uv version to install |
 | `VLLM_VERSION` | vLLM git tag to build |
+| `VLLM_WHEEL_SRC` | Stage supplying the vLLM wheel. `vllm_wheel_build` (default) builds it from source; `vllm_wheel_none` skips both the build and the install |
 | `TRT_LLM_COMMIT` | TensorRT-LLM git commit or tag |
 | `TRT_LLM_VERSION` | TensorRT-LLM version string embedded as an image environment variable |
 | `TRT_VER` | TensorRT version for the TRT-LLM install scripts |

docker/README.md

@@ -76,24 +76,46 @@ Translate your RL framework config into Megatron Bridge's `ConfigContainer` for
 
 ```python
 import torch
+from megatron.bridge import AutoBridge
 from megatron.bridge.training.config import (
     ConfigContainer,
     TrainingConfig,
     OptimizerConfig,
     SchedulerConfig,
     DistributedDataParallelConfig,
     CheckpointConfig,
+    LoggerConfig,
     TokenizerConfig,
 )
 from nemo_rl.models.policy import PolicyConfig  # or your own policy cfg type
 
 # Example: map your RL config to Megatron config
 def build_megatron_config(rl_cfg: PolicyConfig, pretrained_ckpt_dir: str) -> ConfigContainer:
-    model_cfg = rl_cfg["megatron_cfg"].copy()
+    bridge = AutoBridge.from_hf_pretrained(rl_cfg["model_name"])
+    model_cfg = bridge.to_megatron_provider(load_weights=False)
+
+    # Keep framework-owned optimizer/scheduler/DDP mappings out of the model
+    # provider. Apply only provider fields, with Bridge validating each name.
+    provider_fields = (
+        "tensor_model_parallel_size",
+        "pipeline_model_parallel_size",
+        "context_parallel_size",
+        "expert_model_parallel_size",
+        "expert_tensor_parallel_size",
+        "sequence_parallel",
+        "recompute_granularity",
+        "recompute_method",
+        "recompute_num_layers",
+    )
+    model_overrides = {name: rl_cfg["megatron_cfg"][name] for name in provider_fields if name in rl_cfg["megatron_cfg"]}
+
     # Precision
-    dtype = rl_cfg["precision"]
-    model_cfg["bf16"] = dtype == "bfloat16"
-    model_cfg["fp16"] = dtype == "float16"
+    dtype = {
+        "float32": torch.float32,
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+    }[rl_cfg["precision"]]
+    model_cfg.apply_overrides_and_finalize(dtype=dtype, overrides=model_overrides)
 
     checkpoint = CheckpointConfig(
         save_interval=100,
@@ -130,17 +152,19 @@ def build_megatron_config(rl_cfg: PolicyConfig, pretrained_ckpt_dir: str) -> Con
         tokenizer_model=rl_cfg["model_name"],
     )
 
-    return ConfigContainer(
+    cfg = ConfigContainer(
         model=model_cfg,
         checkpoint=checkpoint,
-        logger=None,
+        logger=LoggerConfig(),
         train=train,
         optimizer=opt,
         ddp=ddp,
         scheduler=sch,
         dataset=None,
         tokenizer=tokenizer,
     )
+    cfg.validate()
+    return cfg
 ```
 
 Initialize Megatron-Core using a helper similar to `setup_megatron_model` from NeMo-RL:

docs/bridge-rl-integration.md

@@ -72,24 +72,46 @@ Translate your RL framework config into Megatron Bridge's `ConfigContainer` for
 
 ```python
 import torch
+from megatron.bridge import AutoBridge
 from megatron.bridge.training.config import (
     ConfigContainer,
     TrainingConfig,
     OptimizerConfig,
     SchedulerConfig,
     DistributedDataParallelConfig,
     CheckpointConfig,
+    LoggerConfig,
     TokenizerConfig,
 )
 from nemo_rl.models.policy import PolicyConfig  # or your own policy cfg type
 
 # Example: map your RL config to Megatron config
 def build_megatron_config(rl_cfg: PolicyConfig, pretrained_ckpt_dir: str) -> ConfigContainer:
-    model_cfg = rl_cfg["megatron_cfg"].copy()
+    bridge = AutoBridge.from_hf_pretrained(rl_cfg["model_name"])
+    model_cfg = bridge.to_megatron_provider(load_weights=False)
+
+    # Keep framework-owned optimizer/scheduler/DDP mappings out of the model
+    # provider. Apply only provider fields, with Bridge validating each name.
+    provider_fields = (
+        "tensor_model_parallel_size",
+        "pipeline_model_parallel_size",
+        "context_parallel_size",
+        "expert_model_parallel_size",
+        "expert_tensor_parallel_size",
+        "sequence_parallel",
+        "recompute_granularity",
+        "recompute_method",
+        "recompute_num_layers",
+    )
+    model_overrides = {name: rl_cfg["megatron_cfg"][name] for name in provider_fields if name in rl_cfg["megatron_cfg"]}
+
     # Precision
-    dtype = rl_cfg["precision"]
-    model_cfg["bf16"] = dtype == "bfloat16"
-    model_cfg["fp16"] = dtype == "float16"
+    dtype = {
+        "float32": torch.float32,
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+    }[rl_cfg["precision"]]
+    model_cfg.apply_overrides_and_finalize(dtype=dtype, overrides=model_overrides)
 
     checkpoint = CheckpointConfig(
         save_interval=100,
@@ -126,17 +148,19 @@ def build_megatron_config(rl_cfg: PolicyConfig, pretrained_ckpt_dir: str) -> Con
         tokenizer_model=rl_cfg["model_name"],
     )
 
-    return ConfigContainer(
+    cfg = ConfigContainer(
         model=model_cfg,
         checkpoint=checkpoint,
-        logger=None,
+        logger=LoggerConfig(),
         train=train,
         optimizer=opt,
         ddp=ddp,
         scheduler=sch,
         dataset=None,
         tokenizer=tokenizer,
     )
+    cfg.validate()
+    return cfg
 ```
 
 Initialize Megatron-Core using a helper similar to `setup_megatron_model` from NeMo-RL:

docs/fern/versions/0.4.2/pages/bridge-rl-integration.mdx

@@ -72,24 +72,46 @@ Translate your RL framework config into Megatron Bridge's `ConfigContainer` for
 
 ```python
 import torch
+from megatron.bridge import AutoBridge
 from megatron.bridge.training.config import (
     ConfigContainer,
     TrainingConfig,
     OptimizerConfig,
     SchedulerConfig,
     DistributedDataParallelConfig,
     CheckpointConfig,
+    LoggerConfig,
     TokenizerConfig,
 )
 from nemo_rl.models.policy import PolicyConfig  # or your own policy cfg type
 
 # Example: map your RL config to Megatron config
 def build_megatron_config(rl_cfg: PolicyConfig, pretrained_ckpt_dir: str) -> ConfigContainer:
-    model_cfg = rl_cfg["megatron_cfg"].copy()
+    bridge = AutoBridge.from_hf_pretrained(rl_cfg["model_name"])
+    model_cfg = bridge.to_megatron_provider(load_weights=False)
+
+    # Keep framework-owned optimizer/scheduler/DDP mappings out of the model
+    # provider. Apply only provider fields, with Bridge validating each name.
+    provider_fields = (
+        "tensor_model_parallel_size",
+        "pipeline_model_parallel_size",
+        "context_parallel_size",
+        "expert_model_parallel_size",
+        "expert_tensor_parallel_size",
+        "sequence_parallel",
+        "recompute_granularity",
+        "recompute_method",
+        "recompute_num_layers",
+    )
+    model_overrides = {name: rl_cfg["megatron_cfg"][name] for name in provider_fields if name in rl_cfg["megatron_cfg"]}
+
     # Precision
-    dtype = rl_cfg["precision"]
-    model_cfg["bf16"] = dtype == "bfloat16"
-    model_cfg["fp16"] = dtype == "float16"
+    dtype = {
+        "float32": torch.float32,
+        "bfloat16": torch.bfloat16,
+        "float16": torch.float16,
+    }[rl_cfg["precision"]]
+    model_cfg.apply_overrides_and_finalize(dtype=dtype, overrides=model_overrides)
 
     checkpoint = CheckpointConfig(
         save_interval=100,
@@ -126,17 +148,19 @@ def build_megatron_config(rl_cfg: PolicyConfig, pretrained_ckpt_dir: str) -> Con
         tokenizer_model=rl_cfg["model_name"],
     )
 
-    return ConfigContainer(
+    cfg = ConfigContainer(
         model=model_cfg,
         checkpoint=checkpoint,
-        logger=None,
+        logger=LoggerConfig(),
         train=train,
         optimizer=opt,
         ddp=ddp,
         scheduler=sch,
         dataset=None,
         tokenizer=tokenizer,
     )
+    cfg.validate()
+    return cfg
 ```
 
 Initialize Megatron-Core using a helper similar to `setup_megatron_model` from NeMo-RL: