upgrade datasets to 4.4.1 (#3266)

* upgrade datasets

* cleanup pip cache earlier

* cleanup unused things from worker

* also cleanup sdist

.github/workflows/tests.yml

@@ -59,6 +59,10 @@ jobs:
     timeout-minutes: 20
 
     steps:
+      - name: cleanup node
+        run: |
+          sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc /opt/hostedtoolcache/CodeQL
+
       - name: Check out repository code
         uses: actions/checkout@v4
 
@@ -91,6 +95,10 @@ jobs:
           python scripts/cutcrossentropy_install.py | sh
           pip3 install -r requirements-dev.txt -r requirements-tests.txt
 
+      - name: cleanup pip cache
+        run: |
+          find "$(pip cache dir)/http-v2" -type f -mtime +14 -exec rm {} \;
+
       - name: Make sure PyTorch version wasn't clobbered
         run: |
           python -c "import torch; assert '${{ matrix.pytorch_version }}' in torch.__version__"
@@ -118,10 +126,6 @@ jobs:
           flags: unittests,pytorch-${{ matrix.pytorch_version }}
           fail_ci_if_error: false
 
-      - name: cleanup pip cache
-        run: |
-          find "$(pip cache dir)/http-v2" -type f -mtime +14 -exec rm {} \;
-
   pytest-sdist:
     name: PyTest from Source Dist
     runs-on: ubuntu-latest
@@ -134,6 +138,10 @@ jobs:
     timeout-minutes: 20
 
     steps:
+      - name: cleanup node
+        run: |
+          sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc /opt/hostedtoolcache/CodeQL
+
       - name: Check out repository code
         uses: actions/checkout@v4
 
@@ -167,6 +175,10 @@ jobs:
           python scripts/cutcrossentropy_install.py | sh
           pip3 install -r requirements-dev.txt -r requirements-tests.txt
 
+      - name: cleanup pip cache
+        run: |
+          find "$(pip cache dir)/http-v2" -type f -mtime +14 -exec rm {} \;
+
       - name: Make sure PyTorch version wasn't clobbered
         run: |
           python -c "import torch; assert '${{ matrix.pytorch_version }}' in torch.__version__"
@@ -184,10 +196,6 @@ jobs:
           pytest -v --durations=10 tests/monkeypatch/ --cov=axolotl --cov-append --cov-report=xml
           pytest -v --durations=10 tests/cli/
 
-      - name: cleanup pip cache
-        run: |
-          find "$(pip cache dir)/http-v2" -type f -mtime +14 -exec rm {} \;
-
   gate-skip-e2e:
     needs: [pre-commit, pytest, pytest-sdist]
     runs-on: ubuntu-latest

requirements.txt

@@ -15,7 +15,7 @@ peft>=0.17.1
 tokenizers>=0.22.1
 transformers==4.57.1
 accelerate==1.11.0
-datasets==4.3.0
+datasets==4.4.1
 deepspeed>=0.17.0
 trl==0.25.0
 hf_xet==1.2.0