fix megatron build

Author: pan-x-c

scripts/docker/Dockerfile.megatron

@@ -31,7 +31,8 @@ RUN pip install --upgrade pip \
     && pip install -e .[mm,dev] \
     && pip install flash_attn==2.8.1 --no-build-isolation \
     && pip install -e .[megatron] \
-    && pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation \
+    && 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" \
         --config-settings "--build-option=--cuda_ext" \
         --resume-retries 20 git+https://github.com/NVIDIA/apex.git

scripts/docker/Dockerfile.uv

@@ -42,6 +42,7 @@ RUN . /opt/venv/bin/activate && \
 RUN . /opt/venv/bin/activate && \
     uv pip install -e.[flash_attn] && \
     uv pip install -e .[megatron] && \
+    NVCC_APPEND_FLAGS="--threads 4" APEX_PARALLEL_BUILD=8 \
     uv pip install -v --no-build-isolation \
     --config-settings="--build-option=--cpp_ext" \
     --config-settings="--build-option=--cuda_ext" \