Add mosaicml/llm-foundry Docker workflow (mosaicml#254)

.github/workflows/docker.yaml

@@ -0,0 +1,54 @@
+name: Docker
+on:
+  push:
+    branches:
+    - main
+  workflow_dispatch: {}
+jobs:
+  docker-build:
+    runs-on: ubuntu-latest
+    if: github.repository_owner == 'mosaicml'
+    strategy:
+      matrix:
+        include:
+        - name: '1.13.1_cu117'
+          base_image: mosaicml/pytorch:1.13.1_cu117-python3.10-ubuntu20.04
+        - name: '2.0.1_cu118'
+          base_image: mosaicml/pytorch:2.0.1_cu118-python3.10-ubuntu20.04
+
+    steps:
+    - name: Checkout
+      uses: actions/checkout@v3
+
+    - name: Setup QEMU
+      uses: docker/setup-qemu-action@v2
+
+    - name: Setup Docker Buildx
+      uses: docker/setup-buildx-action@v2
+
+    - name: Login to DockerHub
+      uses: docker/login-action@v2
+      with:
+        username: ${{ secrets.DOCKER_HUB_USERNAME }}
+        password: ${{ secrets.DOCKER_HUB_PASSWORD }}
+
+    - name: Calculate Docker Image Variables
+      run: |
+        set -euxo pipefail
+
+        ###################
+        # Calculate the tag
+        ###################
+        GIT_SHA=$(echo ${{ github.sha }} | cut -c1-7)
+        echo "IMAGE_TAG=${GIT_SHA}" >> ${GITHUB_ENV}
+
+    - name: Build and Push the Docker Image
+      uses: docker/build-push-action@v3
+      with:
+        context: .
+        tags: mosaicml/llm-foundry:${{ matrix.name }}-latest,
+          mosaicml/llm-foundry:${{ matrix.name }}-${{ env.IMAGE_TAG }}
+        push: true
+        cache-from: type=registry,ref=mosaicml/llm-foundry:${{ matrix.name }}-buildcache
+        cache-to: type=registry,ref=mosaicml/llm-foundry:${{ matrix.name }}-buildcache,mode=max
+        build-args: BASE_IMAGE=${{ matrix.base_image }}

Dockerfile

@@ -0,0 +1,12 @@
+# Copyright 2022 MosaicML LLM Foundry authors
+# SPDX-License-Identifier: Apache-2.0
+
+ARG BASE_IMAGE
+FROM $BASE_IMAGE
+
+
+# Install and uninstall foundry to cache foundry requirements
+RUN git clone -b main https://github.com/mosaicml/llm-foundry.git && \
+    pip install --no-cache-dir "./llm-foundry[gpu]" && \
+    pip uninstall -y llm-foundry && \
+    rm -rf llm-foundry

README.md

@@ -44,12 +44,12 @@ You'll find in this repo:
 MPT-7B is a GPT-style model, and the first in the MosaicML Foundation Series of models. Trained on 1T tokens of a MosaicML-curated dataset, MPT-7B is open-source, commercially usable, and equivalent to LLaMa 7B on evaluation metrics. The MPT architecture contains all the latest techniques on LLM modeling -- Flash Attention for efficiency, Alibi for context length extrapolation, and stability improvements to mitigate loss spikes. The base model and several variants, including a 64K context length fine-tuned model (!!) are all available:
 
 
-| Model              | Context Length | Download                                           | Demo                                                           | Commercial use? |
-|--------------------|----------------|----------------------------------------------------|----------------------------------------------------------------|-----------------|
-| MPT-7B             | 2048           | https://huggingface.co/mosaicml/mpt-7b             |                                                                | Yes             |
-| MPT-7B-Instruct    | 2048           | https://huggingface.co/mosaicml/mpt-7b-instruct    | [Demo](https://huggingface.co/spaces/mosaicml/mpt-7b-instruct) | Yes             |
-| MPT-7B-Chat        | 2048           | https://huggingface.co/mosaicml/mpt-7b-chat        | [Demo](https://huggingface.co/spaces/mosaicml/mpt-7b-chat)     | No              |
-| MPT-7B-StoryWriter | 65536          | https://huggingface.co/mosaicml/mpt-7b-storywriter |                                                                | Yes             |
+| Model              | Context Length | Download                                           | Demo                                                             | Commercial use? |
+|--------------------|----------------|----------------------------------------------------|------------------------------------------------------------------|-----------------|
+| MPT-7B             | 2048           | https://huggingface.co/mosaicml/mpt-7b             |                                                                  | Yes             |
+| MPT-7B-Instruct    | 2048           | https://huggingface.co/mosaicml/mpt-7b-instruct    | [Demo](https://huggingface.co/spaces/mosaicml/mpt-7b-instruct)   | Yes             |
+| MPT-7B-Chat        | 2048           | https://huggingface.co/mosaicml/mpt-7b-chat        | [Demo](https://huggingface.co/spaces/mosaicml/mpt-7b-chat)       | No              |
+| MPT-7B-StoryWriter | 65536          | https://huggingface.co/mosaicml/mpt-7b-storywriter | [Demo](https://huggingface.co/spaces/mosaicml/mpt-7b-storywriter)| Yes             |
 
 To try out these models locally, [follow the instructions](https://github.com/mosaicml/llm-foundry/tree/main/scripts/inference#interactive-generation-with-modelgenerate) in `scripts/inference/README.md` to prompt HF models using our [hf_generate.py](https://github.com/mosaicml/llm-foundry/blob/main/scripts/inference/hf_generate.py) or [hf_chat.py](https://github.com/mosaicml/llm-foundry/blob/main/scripts/inference/hf_chat.py) scripts.
 
@@ -77,17 +77,38 @@ Something missing? Contribute with a PR!
 
 
 
-# Prerequisites
-Here's what you need to get started with our LLM stack:
-* Use a Docker image with PyTorch 1.13+, e.g. [MosaicML's PyTorch base image](https://hub.docker.com/r/mosaicml/pytorch/tags)
-   * Recommended tag: `mosaicml/pytorch:1.13.1_cu117-python3.10-ubuntu20.04`
-   * This image comes pre-configured with the following dependencies:
-      * PyTorch Version: 1.13.1
-      * CUDA Version: 11.7
-      * Python Version: 3.10
-      * Ubuntu Version: 20.04
-      * FlashAttention kernels from [HazyResearch](https://github.com/HazyResearch/flash-attention)
-* Use a system with NVIDIA GPUs
+# Hardware and Software Requirements
+This codebase has been tested with PyTorch 1.13.1 and PyTorch 2.0.1 on systems with NVIDIA A100s and H100s.
+This codebase may also work on systems with other devices, such as consumer NVIDIA cards and AMD cards, but we are not actively testing these systems.
+If you have success/failure using LLM Foundry on other systems, please let us know in a Github issue and we will update the support matrix!
+
+## Supported hardware
+
+| Device                    | Torch Version    | Cuda Version | Status                        |
+|---------------------------|------------------|--------------|-------------------------------|
+| A100-40GB/80GB            | 1.13.1           | 11.7         | :white_check_mark: Supported  |
+| A100-40GB/80GB            | 2.0.1            | 11.7, 11.8   | :white_check_mark: Supported  |
+| H100-80GB                 | 1.13.1           | 11.7         | :x: Not Supported             |
+| H100-80GB                 | 2.0.1            | 11.8         | :white_check_mark: Supported  |
+| A10-24GB                  | 1.13.1           | 11.7         | :construction: In Progress    |
+| A10-24GB                  | 2.0.1            | 11.7, 11.8   | :construction: In Progress    |
+
+## MosaicML Docker Images
+We highly recommend using our prebuilt Docker images. You can find them here: https://hub.docker.com/orgs/mosaicml/repositories.
+
+The `mosaicml/pytorch` images are pinned to specific PyTorch and CUDA versions, and are stable and rarely updated.
+
+The `mosaicml/llm-foundry` images are built with new tags upon every commit to the `main` branch.
+You can select a specific commit hash such as `mosaicml/llm-foundry:1.13.1_cu117-f678575` or take the latest one using `mosaicml/llm-foundry:1.13.1_cu117-latest`.
+**Plese Note:** The `mosaicml/llm-foundry` images do not install the package itself, just the dependencies. You will still need to `pip install llm-foundry` either from PyPi or from source.
+
+| Docker Image                                                | Torch Version  | Cuda Version | LLM Foundry dependencies installed? |
+|-------------------------------------------------------------|----------------|--------------|-------------------------------------|
+| `mosaicml/pytorch:1.13.1_cu117-python3.10-ubuntu20.04`      | 1.13.1         | 11.7         | No                                  |
+| `mosaicml/pytorch:2.0.1_cu118-python3.10-ubuntu20.04`       | 2.0.1          | 11.8         | No                                  |
+| `mosaicml/llm-foundry:1.13.1_cu117-latest`                  | 1.13.1         | 11.7         | Yes                                 |
+| `mosaicml/llm-foundry:2.0.1_cu118-latest`                   | 2.0.1          | 11.8         | Yes                                 |
+
 
 # Installation
 

mcli/mcli-hf-eval.yaml

@@ -46,7 +46,7 @@ parameters:
   model:
     name: hf_causal_lm
     pretrained_model_name_or_path: ${model_name_or_path}
-    device: cpu
+    init_device: cpu
     pretrained: true
     use_auth_token: false