chore: bump mbridge and MLM

3rdparty/Megatron-Bridge-workspace/setup.py

@@ -27,7 +27,7 @@
 
 # Default dependencies from pyproject.toml
 CACHED_DEPENDENCIES = [
-    "transformers>=5.0.0,<=5.3.0",
+    "transformers>=5.8.1,<5.9.0",
     "peft>=0.18.1",
     "datasets>=2.20.0",
     "accelerate",
@@ -48,17 +48,16 @@
     "hydra-core>1.3,<=1.3.2",
     "megatron-core[dev,mlm]",
     "qwen-vl-utils",
-    # TODO(https://github.com/NVIDIA-NeMo/RL/issues/2111): upgrade to core_cu13 when we move to CUDA 13 base container
-    "transformer-engine[pytorch,core_cu13]",
-    "mamba-ssm",
     "nvidia-resiliency-ext",
-    "causal-conv1d",
     "flash-linear-attention",
     "timm",
     "open-clip-torch>=3.2.0",
     "mlflow>=3.9.0",
     "comet-ml>=3.50.0",
     "torch>=2.6.0",
+    "mistral-common>=1.10.0",
+    "flashinfer-cubin==0.6.8.post1",
+    "flashinfer-python==0.6.8.post1",
 ]
 
 # If the bridge source exists, compare cached dependencies with the submodule's pyproject

docker/Dockerfile

@@ -190,6 +190,13 @@ ENV NEMO_RL_VENV_DIR=/opt/ray_venvs
 ENV CUDNN_HOME=/opt/nemo_rl_venv/lib/python3.13/site-packages/nvidia/cudnn
 ENV LD_LIBRARY_PATH="/opt/nemo_rl_venv/lib/python3.13/site-packages/nvidia/cudnn/lib:${LD_LIBRARY_PATH}"
 # Verify with: python -c "import transformer_engine.pytorch as te; print(te.get_cudnn_version())"
+# Make z3 (bundled by tilelang's TVM) findable by the dynamic linker. z3-solver installs
+# libz3.so.4.15 into Python site-packages, which is not on the default linker search path.
+# Without this, fla's tilelang.is_available() raises OSError instead of ImportError,
+# crashing worker venvs that co-install flash-linear-attention (via megatron-bridge) + tilelang.
+RUN /opt/nemo_rl_venv/bin/python -c \
+    "import z3, os; print(os.path.join(os.path.dirname(z3.__file__), 'lib'))" \
+    > /etc/ld.so.conf.d/z3-solver.conf && ldconfig
 
 WORKDIR /opt/nemo-rl
 

nemo_rl/modelopt/models/policy/workers/megatron_quant_policy_worker.py

@@ -23,6 +23,7 @@
     has_modelopt_state,
     load_modelopt_state,
 )
+from megatron.bridge.utils.instantiate_utils import register_allowed_target_prefix
 from megatron.core.utils import unwrap_model
 from modelopt.torch.quantization.nn.modules.quant_module import QuantModule
 from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer
@@ -65,6 +66,10 @@ def __init__(self, config, *args, **kwargs):
         # setup_reference_model_state and runs before load_checkpoint to resume
         # quantizers on the model.
         self._model_import_post_wrap_hook = self._quantize
+
+        ## need to add nemo_rl to mbridge's allowed target prefixes in order to successfully
+        ## restore from the mbridge checkpoint.
+        register_allowed_target_prefix("nemo_rl.")
         self._transformer_layer_spec = quantization_layer_spec
         self._pre_load_checkpoint_hook = self._restore_modelopt_state_pre_load
         super().__init__(config, *args, **kwargs)

pyproject.toml

@@ -170,6 +170,8 @@ mcore = [
   "emerging-optimizers==0.2.0",
   "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480 ; platform_machine == 'x86_64'",
   "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@a48493600c4886c1b297aaa78db0e1ebc2d8dd6c ; platform_machine == 'aarch64'",
+  "mamba-ssm",
+  "causal-conv1d",
 ]
 modelopt = ["nvidia-modelopt"]
 nvrx = [
@@ -251,7 +253,7 @@ mamba-ssm = { git = "https://github.com/state-spaces/mamba.git", rev = "d68d16ed
 nv-grouped-gemm = { git = "https://github.com/fanshiqing/grouped_gemm", tag = "v1.1.4.post7" }
 emerging-optimizers = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git", rev = "v0.2.0" }
 nvidia-modelopt = { git = "https://github.com/NVIDIA/Model-Optimizer", rev = "905018803414702e414a86716484ed4115b37ba6" }
-nvidia-resiliency-ext = { git = "https://github.com/NVIDIA/nvidia-resiliency-ext.git", rev = "15a851565a4ce846c04431ecb0cf09903ab4837e" }
+nvidia-resiliency-ext = { git = "https://github.com/NVIDIA/nvidia-resiliency-ext.git", rev = "6c5f2a13c7688d92a7ac7ee6e464721eb8b7345d" }
 
 [tool.uv.workspace]
 members = [
@@ -457,7 +459,7 @@ name = "megatron-bridge"
 # Must stay in sync with 3rdparty/Megatron-Bridge-workspace/setup.py::CACHED_DEPENDENCIES.
 version = "0.0.0"
 requires-dist = [
-  "transformers>=5.0.0,<=5.3.0",
+  "transformers>=5.8.1,<5.9.0",
   "peft>=0.18.1",
   "datasets>=2.20.0",
   "accelerate",
@@ -479,25 +481,24 @@ requires-dist = [
   # megatron-core omitted here: it's a workspace member and already a direct dep of the mcore extra.
   # Non-workspace path deps cannot depend on workspace members (uv name-shadowing restriction).
   "qwen-vl-utils",
-  # TODO(https://github.com/NVIDIA-NeMo/RL/issues/2111): upgrade to core_cu13 when we move to CUDA 13
-  "transformer-engine[pytorch,core_cu13]",
-  "mamba-ssm",
   "nvidia-resiliency-ext",
-  "causal-conv1d",
   "flash-linear-attention",
   "timm",
   "open-clip-torch>=3.2.0",
   "mlflow>=3.9.0",
   "comet-ml>=3.50.0",
   "torch>=2.6.0",
+  "mistral-common>=1.10.0",
+  "flashinfer-cubin==0.6.8.post1",
+  "flashinfer-python==0.6.8.post1",
 ]
 
 # Override logsage metadata to remove numpy<=2.0.2 and pandas<=2.3.3 upper bounds
 # (numpy cap conflicts with onnx>=1.21.0rc4 CVE fix via nvidia-resiliency-ext)
 # Tracking: https://github.com/NVIDIA/nvidia-resiliency-ext/issues/301
 [[tool.uv.dependency-metadata]]
 name = "logsage"
-version = "0.1.5"
+version = "0.1.7"
 requires-dist = [
   "drain3>=0.9.11,<0.10.0",
   "langchain>=0.3.27,<0.4.0",