feat(megatron): clone megatron source into image

coreweave · Oct 8, 2024 · 14465e0 · 14465e0
1 parent 02e94df
commit 14465e0
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/.github/workflows/megatron.yml b/.github/workflows/megatron.yml
@@ -4,6 +4,9 @@ on:
       base-image:
         description: 'Base image to use'
         required: true
+      commit:
+        description: 'Commit of Megatron to include'
+        required: true
   push:
     paths:
       - "megatron/**"
@@ -13,12 +16,11 @@ on:
 
 jobs:
   build:
-    env:
-      DEFAULT_IMAGE: ghcr.io/coreweave/ml-containers/torch-extras:bfe03aa-nccl-cuda12.4.1-ubuntu22.04-nccl2.21.5-1-torch2.4.0-vision0.19.0-audio2.4.0
     uses: ./.github/workflows/build.yml
     secrets: inherit
     with:
       image-name: megatron
       folder: megatron
       build-args: |
-        BASE_IMAGE=${{ inputs.base-image || env.DEFAULT_IMAGE }}
+        BASE_IMAGE=${{ inputs.base-image || 'ghcr.io/coreweave/ml-containers/torch-extras:bfe03aa-nccl-cuda12.4.1-ubuntu22.04-nccl2.21.5-1-torch2.4.0-vision0.19.0-audio2.4.0'}}
+        COMMIT=${{ inputs.commit || 'main'}}
diff --git a/megatron/Dockerfile b/megatron/Dockerfile
@@ -6,3 +6,8 @@ COPY requirements.txt /tmp/requirements.txt
 RUN python3 -m pip install -U --no-cache-dir \
     -r /tmp/requirements.txt \
     && rm /tmp/requirements.txt
+
+ARG COMMIT
+RUN git clone https://github.com/NVIDIA/megatron-lm && \
+    cd megatron-lm && \
+    git checkout ${COMMIT}