Merge branch 'main' into rwkv-x-eagle-compile

RWKV · Feb 9, 2024 · 651247d · 651247d
2 parents 4c88076 + ba60880
commit 651247d
Show file tree

Hide file tree

Showing 59 changed files with 8,968 additions and 900 deletions.
diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml
@@ -1,4 +1,4 @@
-name: Docker Env Image (cuda-12-3)
+name: Docker Env Image (cuda-12-1)
 
 on:
   push:
@@ -21,7 +21,7 @@ env:
 
 jobs:
   build_env:
-    name: Docker Env Image (cuda-12-3)
+    name: Docker Env Image (cuda-12-1)
 
     runs-on: ubuntu-latest
     permissions:
@@ -103,110 +103,110 @@ jobs:
       
       # Build and push Docker image with Buildx (don't push on PR)
       # https://github.com/docker/build-push-action
-      - name: Build and push Docker image (env-cuda-12-3)
+      - name: Build and push Docker image (env-cuda-12-1)
         id: build-and-push
         uses: docker/build-push-action@v4
         with:
-          context: "{{defaultContext}}:docker/env-cuda-12-3"
+          context: "{{defaultContext}}:docker/env-cuda-12-1"
           push: ${{ github.event_name != 'pull_request' }} # Don't push on PR
-          tags: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME_LC }}:env-cuda-12-3
+          tags: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME_LC }}:env-cuda-12-1
           # tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
-          cache-from: type=gha,src=docker/env-cuda-12-3
+          cache-from: type=gha,src=docker/env-cuda-12-1
           cache-to: type=gha,mode=max
 
-  # build_runner:
-  #   name: Docker Env Image (github-worker-12-3)
+  build_runner:
+    name: Docker Env Image (github-worker-12-1)
 
-  #   needs: build_env
-  #   runs-on: ubuntu-latest
-  #   permissions:
-  #     contents: read
-  #     packages: write
-  #     # This is used to complete the identity challenge
-  #     # with sigstore/fulcio when running outside of PRs.
-  #     id-token: write
-
-  #   steps:
-  #     # Get and log the free space
-  #     - name: Get system free space (Before reclaim)
-  #       run: |
-  #         echo "Free space:"
-  #         df -h
+    needs: build_env
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+      # This is used to complete the identity challenge
+      # with sigstore/fulcio when running outside of PRs.
+      id-token: write
+
+    steps:
+      # Get and log the free space
+      - name: Get system free space (Before reclaim)
+        run: |
+          echo "Free space:"
+          df -h
       
-  #     # Due to the docker image being > available space on the runner
-  #     # we need to do some optimization, to create more space.
-  #     # https://github.com/marketplace/actions/disk-space-reclaimer
-  #     # https://stackoverflow.com/questions/76294509/github-actions-docker-service-container-25gb-cannot-be-loaded
-  #     - name: Maximize build space
-  #       uses: insightsengineering/disk-space-reclaimer@v1
-  #       with:
-  #         # this might remove tools that are actually needed,
-  #         # if set to "true" but frees about 6 GB
-  #         tools-cache: true
-
-  #         # all of these default to true, but feel free to set to
-  #         # "false" if necessary for your workflow
-  #         android: true
-  #         dotnet: true
-  #         haskell: true
-  #         large-packages: true
-  #         swap-storage: true
-  #         docker-images: true
+      # Due to the docker image being > available space on the runner
+      # we need to do some optimization, to create more space.
+      # https://github.com/marketplace/actions/disk-space-reclaimer
+      # https://stackoverflow.com/questions/76294509/github-actions-docker-service-container-25gb-cannot-be-loaded
+      - name: Maximize build space
+        uses: insightsengineering/disk-space-reclaimer@v1
+        with:
+          # this might remove tools that are actually needed,
+          # if set to "true" but frees about 6 GB
+          tools-cache: true
+
+          # all of these default to true, but feel free to set to
+          # "false" if necessary for your workflow
+          android: true
+          dotnet: true
+          haskell: true
+          large-packages: true
+          swap-storage: true
+          docker-images: true
 
-  #     # Get and log the free space
-  #     - name: Get system free space (After reclaim)
-  #       run: |
-  #         echo "Free space:"
-  #         df -h
-
-  #     - name: Checkout repository
-  #       uses: actions/checkout@v3
-
-  #     # Install the cosign tool except on PR
-  #     # https://github.com/sigstore/cosign-installer
-  #     - name: Install cosign
-  #       if: github.event_name != 'pull_request'
-  #       uses: sigstore/[email protected]
-  #       with:
-  #         cosign-release: 'v3.3.0'
-
-  #     # Workaround: https://github.com/docker/build-push-action/issues/461
-  #     - name: Setup Docker buildx
-  #       uses: docker/setup-buildx-action@v2
-
-  #     # Login against a Docker registry except on PR
-  #     # https://github.com/docker/login-action
-  #     - name: Log into registry ${{ env.REGISTRY }}
-  #       if: github.event_name != 'pull_request'
-  #       uses: docker/login-action@28218f9b04b4f3f62068d7b6ce6ca5b26e35336c
-  #       with:
-  #         registry: ${{ env.REGISTRY }}
-  #         username: ${{ github.actor }}
-  #         password: ${{ secrets.GITHUB_TOKEN }}
-
-  #     # Extract metadata (tags, labels) for Docker
-  #     # https://github.com/docker/metadata-action
-  #     - name: Extract Docker metadata
-  #       id: meta
-  #       uses: docker/metadata-action@98669ae865ea3cffbcbaa878cf57c20bbf1c6c38
-  #       with:
-  #         images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
+      # Get and log the free space
+      - name: Get system free space (After reclaim)
+        run: |
+          echo "Free space:"
+          df -h
+
+      - name: Checkout repository
+        uses: actions/checkout@v3
+
+      # Install the cosign tool except on PR
+      # https://github.com/sigstore/cosign-installer
+      - name: Install cosign
+        if: github.event_name != 'pull_request'
+        uses: sigstore/[email protected]
+        # with:
+        #   cosign-release: 'v2.2.0'
+
+      # Workaround: https://github.com/docker/build-push-action/issues/461
+      - name: Setup Docker buildx
+        uses: docker/setup-buildx-action@v2
+
+      # Login against a Docker registry except on PR
+      # https://github.com/docker/login-action
+      - name: Log into registry ${{ env.REGISTRY }}
+        if: github.event_name != 'pull_request'
+        uses: docker/login-action@28218f9b04b4f3f62068d7b6ce6ca5b26e35336c
+        with:
+          registry: ${{ env.REGISTRY }}
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      # Extract metadata (tags, labels) for Docker
+      # https://github.com/docker/metadata-action
+      - name: Extract Docker metadata
+        id: meta
+        uses: docker/metadata-action@98669ae865ea3cffbcbaa878cf57c20bbf1c6c38
+        with:
+          images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
 
-  #     - name: downcase IMAGE_NAME
-  #       run: |
-  #         echo "IMAGE_NAME_LC=${IMAGE_NAME,,}" >>${GITHUB_ENV}
+      - name: downcase IMAGE_NAME
+        run: |
+          echo "IMAGE_NAME_LC=${IMAGE_NAME,,}" >>${GITHUB_ENV}
       
-  #     # Build and push Docker image with Buildx (don't push on PR)
-  #     # https://github.com/docker/build-push-action
-  #     - name: Build and push Docker image (github-worker-cuda-12-3)
-  #       id: build-and-push
-  #       uses: docker/build-push-action@v4
-  #       with:
-  #         context: "{{defaultContext}}:docker/github-worker-cuda-12-3"
-  #         push: ${{ github.event_name != 'pull_request' }} # Don't push on PR
-  #         tags: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME_LC }}:github-worker-cuda-12-3
-  #         # tags: ${{ steps.meta.outputs.tags }}
-  #         labels: ${{ steps.meta.outputs.labels }}
-  #         cache-from: type=gha,src=docker/github-worker-cuda-12-3
-  #         cache-to: type=gha,mode=max
+      # Build and push Docker image with Buildx (don't push on PR)
+      # https://github.com/docker/build-push-action
+      - name: Build and push Docker image (github-worker-cuda-12-1)
+        id: build-and-push
+        uses: docker/build-push-action@v4
+        with:
+          context: "{{defaultContext}}:docker/github-worker-cuda-12-1"
+          push: ${{ github.event_name != 'pull_request' }} # Don't push on PR
+          tags: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME_LC }}:github-worker-cuda-12-1
+          # tags: ${{ steps.meta.outputs.tags }}
+          labels: ${{ steps.meta.outputs.labels }}
+          cache-from: type=gha,src=docker/github-worker-cuda-12-1
+          cache-to: type=gha,mode=max
diff --git a/RWKV-v5/config-example.yaml b/RWKV-v5/config-example.yaml
@@ -373,6 +373,10 @@ data:
   # source: "teven/enwiki_00k"   # Hugging face dataset
   # source: text                 # Text mode, used with source_data_dir
 
+  # Dataset split to use from HF dataset
+  # ---
+  # source_dataset_split: train
+
   # Additional source dataset params, used to grab subsets of the dataset
   # ---
   # source_dataset_params:
@@ -419,6 +423,7 @@ data:
 
   # Custom text column to use, useful for dataset with alternative training columns labels
   # This is checked before multi column merging, default is null (disabled)
+  # If set this takes priority
   # eg: 'code'
   # ---
   # custom_text_key: 'code'

diff --git a/RWKV-v5/datapack-example.yaml b/RWKV-v5/datapack-example.yaml
@@ -35,15 +35,20 @@ datapack:
 
   # Mixing mode to use, this is used to alternate between datasets
   #
-  # - batch  : Meaning one dataset worth per batch, partial batches are discarded
-  # - sample : Dataset is mixed on a per sample level
-  mixing_mode: "batch"
+  # - concat  : Keep It Simple Silly, lets just concat the datasets together
+  # - shuffle : Dataset is mixed on a per sample level
+  #
+  # (@TODO: Advance operations)
+  # - batch   : Meaning one dataset worth per batch, partial batches are discarded
+  mixing_mode: "shuffle"
 
   # # Mixing distribution to use
   # # - weighted : Dataset batches/mixture is distrbuted randomly, but weighted by dataset size
   # # - uniform  : Dataset batches/mixture is distrbuted randomly, but with uniform probability
   # distribution: "weighted"
 
+  # (@TODO: Advance operations)
+  #
   # Mixed batch percentage
   #
   # % of batches which will contain a mixture of records from multiple datasets

diff --git a/RWKV-v5/lightning_trainer.py b/RWKV-v5/lightning_trainer.py
@@ -23,7 +23,7 @@
 # Get the JIT / torch compile option flags default specific for lightning training mode
 # This enables torch compile by default
 RWKV_JIT_ON         = os.getenv("RWKV_JIT_ON", "1").lower() in ("1", "true", "yes")
-RWKV_TORCH_COMPILE  = os.getenv("RWKV_TORCH_COMPILE", f"1").lower() in ("1", "true", "yes")
+RWKV_TORCH_COMPILE  = os.getenv("RWKV_TORCH_COMPILE", f"0").lower() in ("1", "true", "yes")
 
 # Disable CUDA for RWKV
 RWKV_NO_CUDA        = os.getenv("RWKV_NO_CUDA", "0").lower() in ("1", "true", "yes")