fix toml for mbridge

Author: chenyushuo

docs/sphinx_doc/source/tutorial/example_megatron.md

@@ -20,9 +20,10 @@ pip install -e ".[megatron]"
 # uv sync -extra megatron
 ```
 
-Then, install NVIDIA's Apex library for mixed-precision training:
+Then, install mbridge and NVIDIA's Apex library for mixed-precision training:
 
 ```bash
+pip install -U git+https://github.com/ISEEKYAN/mbridge.git
 pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation \
     --config-settings "--build-option=--cpp_ext" \
     --config-settings "--build-option=--cuda_ext" \

docs/sphinx_doc/source_zh/tutorial/example_megatron.md

@@ -20,9 +20,10 @@ pip install -e ".[megatron]"
 # uv sync -extra megatron
 ```
 
-另外还需要从源码安装 NVIDIA 的 Apex 库以支持混合精度训练：
+另外还需要从源码安装 mbridge 和 NVIDIA 的 Apex 库以支持混合精度训练：
 
 ```bash
+pip install -U git+https://github.com/ISEEKYAN/mbridge.git
 pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation \
     --config-settings "--build-option=--cpp_ext" \
     --config-settings "--build-option=--cuda_ext" \

pyproject.toml

@@ -81,7 +81,7 @@ megatron = [
     # "transformer_engine[pytorch]==2.8.0",
 
     # Install mbridge from main branch (unreleased version)
-    "mbridge @ git+https://github.com/ISEEKYAN/mbridge.git@20e9ffbbe72ae7b1df83bfe1bc3c11f7382f2612",
+    # "mbridge @ git+https://github.com/ISEEKYAN/mbridge.git@20e9ffbbe72ae7b1df83bfe1bc3c11f7382f2612",
 ]
 tinker = [
     "tinker; python_version >= '3.11'",

scripts/docker/Dockerfile.megatron

@@ -32,6 +32,7 @@ RUN pip install --upgrade pip \
     && pip install flash_attn==2.8.1 --no-build-isolation \
     && pip install -e .[megatron] \
     && pip install transformer_engine[pytorch]==2.8.0 --no-build-isolation --no-cache-dir \
+    && pip install -U git+https://github.com/ISEEKYAN/mbridge.git \
     && NVCC_APPEND_FLAGS="--threads 4" APEX_PARALLEL_BUILD=8 pip install -v \
         --disable-pip-version-check --no-cache-dir --no-build-isolation \
         --config-settings "--build-option=--cpp_ext" \