feat(megatron): add initial megatron-lm image

Author: harubaru

.github/workflows/megatron.yml

@@ -0,0 +1,22 @@
+on:
+  workflow_dispatch:
+    inputs:
+      base-image:
+        description: 'Base image to use'
+        required: true
+  push:
+    paths:
+      - "megatron/**"
+      - ".github/workflows/megatron.yml"
+      - ".github/workflows/build.yml"
+
+
+jobs:
+  build:
+    uses: ./.github/workflows/build.yml
+    secrets: inherit
+    with:
+      image-name: megatron
+      folder: megatron
+      build-args: |
+        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'}}

megatron/Dockerfile

@@ -0,0 +1,8 @@
+ARG BASE_IMAGE
+FROM $BASE_IMAGE
+
+COPY requirements.txt /tmp/requirements.txt
+
+RUN python3 -m pip install -U --no-cache-dir \
+    -r /tmp/requirements.txt \
+    && rm /tmp/requirements.txt

megatron/requirements.txt

@@ -0,0 +1,4 @@
+pyyaml==6.0.2
+regex==2024.9.11
+transformers==4.45.2
+triton==3.0.0