[train] support megatron-bridge for PT/SFT training (#10645)

docker/docker-cuda/Dockerfile.mbridge

@@ -0,0 +1,85 @@
+# LLaMA-Factory + Megatron Bridge (CUDA) runtime
+# Mirrors the verified host env: Python 3.12, PyTorch 2.12.1+cu126,
+# TransformerEngine 2.17, megatron-core 0.18, megatron-bridge 0.5.
+#
+# CUDA user-mode libs come from the PyTorch cu126 wheels (same as the host venv).
+# Layers are intentionally few (helps vfs / nested-docker disk usage).
+#
+# Build from repo root:
+#   bash docker/docker-cuda/build-megatron.sh
+
+ARG BASE_IMAGE=ubuntu:24.04
+FROM ${BASE_IMAGE}
+
+ARG PIP_INDEX=https://mirrors.aliyun.com/pypi/simple
+ARG PYPI_TRUSTED_HOST=mirrors.aliyun.com
+ARG TORCH_INDEX=https://download.pytorch.org/whl/cu126
+ARG APT_MIRROR=https://mirrors.tuna.tsinghua.edu.cn/ubuntu/
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PIP_ROOT_USER_ACTION=ignore \
+    PIP_DISABLE_PIP_VERSION_CHECK=1 \
+    PIP_BREAK_SYSTEM_PACKAGES=1 \
+    PIP_CONSTRAINT="" \
+    MAX_JOBS=8 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    VLLM_WORKER_MULTIPROC_METHOD=spawn \
+    DISABLE_VERSION_CHECK=1 \
+    USE_MEGATRON_BRIDGE=1 \
+    GRADIO_SERVER_PORT=7860 \
+    API_PORT=8000 \
+    http_proxy= \
+    https_proxy=
+
+SHELL ["/bin/bash", "-c"]
+WORKDIR /app
+
+# System deps + Ubuntu Python 3.12
+RUN if [ -f /etc/apt/sources.list.d/ubuntu.sources ]; then \
+      sed -i "s|http://archive.ubuntu.com/ubuntu|${APT_MIRROR}|g; s|http://security.ubuntu.com/ubuntu|${APT_MIRROR}|g" /etc/apt/sources.list.d/ubuntu.sources; \
+    elif [ -f /etc/apt/sources.list ]; then \
+      sed -i "s|http://archive.ubuntu.com/ubuntu/|${APT_MIRROR}|g; s|http://security.ubuntu.com/ubuntu/|${APT_MIRROR}|g" /etc/apt/sources.list; \
+    fi && \
+    apt-get update && apt-get install -y --no-install-recommends \
+        ca-certificates curl git vim wget \
+        build-essential ninja-build cmake libgomp1 zip unzip \
+        python3 python3-pip python3-dev python3-venv && \
+    ln -sf /usr/bin/python3 /usr/bin/python && \
+    rm -rf /var/lib/apt/lists/* && \
+    pip install --no-cache-dir --upgrade pip setuptools wheel packaging ninja pybind11 \
+        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX}
+
+COPY . /app
+
+# PyTorch + TE + Megatron Bridge + LLaMA-Factory in one layer
+RUN pip install --no-cache-dir \
+        torch==2.12.1 torchvision==0.27.1 torchaudio==2.11.0 \
+        --index-url ${TORCH_INDEX} && \
+    pip install --no-cache-dir --no-build-isolation \
+        "transformer-engine[pytorch]==2.17.0" \
+        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
+    pip install --no-cache-dir --no-build-isolation \
+        "megatron-bridge==0.5.0" \
+        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
+    pip install --no-cache-dir --no-build-isolation -e . \
+        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
+    pip install --no-cache-dir --no-build-isolation \
+        -r requirements/metrics.txt \
+        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
+    pip install --no-cache-dir --no-build-isolation \
+        "megatron-bridge==0.5.0" \
+        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
+    python - <<'PY'
+import torch
+import megatron.core
+import transformer_engine
+from megatron.bridge import AutoBridge  # noqa: F401
+import llamafactory
+print("torch", torch.__version__, "cuda", torch.version.cuda)
+print("te", transformer_engine.__version__)
+print("megatron-bridge import ok")
+PY
+
+EXPOSE 7860 8000
+CMD ["bash"]

docker/docker-cuda/README.md

@@ -104,6 +104,37 @@ sudo usermod -aG docker $USER
 # Log out and back in for changes to take effect
 ```
 
+## Megatron Bridge Image
+
+`Dockerfile.megatron` builds a CUDA runtime for LLaMA-Factory + [Megatron Bridge](https://docs.nvidia.com/nemo/megatron-bridge/latest/):
+
+| Component | Version |
+| --- | --- |
+| Base | `ubuntu:22.04` (Python 3.12) |
+| PyTorch | 2.12.1+cu126 (CUDA libs from wheels) |
+| TransformerEngine | 2.17.0 |
+| megatron-core | 0.18.x (via megatron-bridge) |
+| megatron-bridge | 0.5.0 |
+
+### Build
+
+From repo root:
+
+```bash
+docker build -f docker/docker-cuda/Dockerfile.megatron \
+  -t llamafactory-megatron-bridge:latest .
+```
+
+### Run training
+
+```bash
+docker run --rm -it --gpus all --ipc=host --shm-size=16g \
+  -e DISABLE_VERSION_CHECK=1 \
+  -e USE_MEGATRON_BRIDGE=1 \
+  -v "$PWD":/app -w /app \
+  llamafactory-megatron-bridge:latest
+```
+
 ## Additional Notes
 
 - The default image is built on Ubuntu 22.04 (x86_64), CUDA 12.4, Python 3.11, PyTorch 2.6.0, and Flash-attn 2.7.4