updage docker image to nightly; add turbo fa env

Author: wenxie-amd

.github/workflows/ci.yaml

@@ -85,8 +85,8 @@ jobs:
           fi
       - name: Build and Push Docker Image
         run: |
-          echo "> Login to Docker Hub"
-          docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }}
+          echo "> Login to ROCm Docker Hub"
+          docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }}
 
           echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}"
           start_time=$(date +%s)
@@ -100,6 +100,9 @@ jobs:
           echo "> Docker tag image for Docker Hub"
           docker tag tasimage/primus:${{env.IMAGE_TAG}} docker.io/tasimage/primus:${{env.IMAGE_TAG}}
 
+          echo "> Login to Primus Docker Hub"
+          docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }}
+
           echo "> Docker push to Docker Hub"
           start_time=$(date +%s)
           docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}

.github/workflows/docker/Dockerfile

@@ -1,19 +1,13 @@
 # Base image
-FROM docker.io/rocm/megatron-lm:v25.9_gfx942
+# FROM docker.io/rocm/megatron-lm:v25.9_gfx942
+FROM docker.io/rocm/pyt-megatron-lm-jax-nightly-private:pytorch_rocm7.0_20251024
 
 # Specify the commit of Primus-Turbo when building: docker build --build-arg PRIMUS_TURBO_COMMIT=xxx .)
 ARG PRIMUS_TURBO_COMMIT
 
 # Install basic dependencies
 RUN apt-get update
 
-RUN pip3 install fbgemm-gpu && \
-    pip3 install torch==2.10.0.dev20251019+rocm7.0 \
-                 torchao==0.15.0.dev20251015+rocm7.0 \
-                 torchvision==0.25.0.dev20251020+rocm7.0 \
-                 --index-url https://download.pytorch.org/whl/nightly/rocm7.0 \
-                 --force-reinstall
-
 # Clone and install the Primus-Turbo
 WORKDIR /opt
 RUN mkdir -p /opt && cd /opt && \
@@ -24,7 +18,6 @@ RUN mkdir -p /opt && cd /opt && \
     pip3 install -r requirements.txt && \
     GPU_ARCHS="gfx942;gfx950" pip3 install --no-build-isolation .
 
-
 # Set the default working directory
 WORKDIR /opt
 

examples/run_local_pretrain.sh

@@ -73,6 +73,9 @@ done < <(env | grep "^PRIMUS_")
 while IFS='=' read -r name _; do
     ENV_ARGS+=("--env" "$name")
 done < <(env | grep "^NCCL_")
+while IFS='=' read -r name _; do
+    ENV_ARGS+=("--env" "$name")
+done < <(env | grep "^PRIMUS_TURBO_")
 ENV_ARGS+=("--env" "EXP")
 ENV_ARGS+=("--env" "HF_TOKEN")
 

examples/run_pretrain.sh

@@ -224,6 +224,9 @@ export NVTE_USE_OPTIMIZED_HIPIFIED_CAST_TRANSPOSE=0
 # Note: Disable v3 due to accuracy issues. Will fix after TE version 2.1.
 export NVTE_CK_USES_BWD_V3=${NVTE_CK_USES_BWD_V3:-0}
 
+# Note: Disable fp32 atomic due if you find any accuracy issue.
+export PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32=${PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32:0}
+
 # nvte debug envs
 export NVTE_DEBUG=0 # 0, 1
 export NVTE_DEBUG_LEVEL=0 # 0, 1, 2
@@ -241,6 +244,7 @@ LOG_INFO_RANK0 "NCCL_P2P_NET_CHUNKSIZE: $NCCL_P2P_NET_CHUNKSIZE"
 LOG_INFO_RANK0 "NVTE_CK_USES_BWD_V3: $NVTE_CK_USES_BWD_V3"
 LOG_INFO_RANK0 "NVTE_USE_CAST_TRANSPOSE_TRITON: $NVTE_USE_CAST_TRANSPOSE_TRITON"
 LOG_INFO_RANK0 "NVTE_USE_OPTIMIZED_HIPIFIED_CAST_TRANSPOSE: $NVTE_USE_OPTIMIZED_HIPIFIED_CAST_TRANSPOSE"
+LOG_INFO_RANK0 "PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32: $PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32"
 if [[ "$PATCH_TE_FLASH_ATTN" == "1" ]]; then
     LOG_INFO_RANK0 'Patching _flash_attn_max_version in attention.py...'
     sed -i 's/_flash_attn_max_version = PkgVersion(\".*\")/_flash_attn_max_version = PkgVersion(\"3.0.0.post1\")/' \