Merge branch 'main' into vjawa/fix_memory_profiling_with_rmm

ci/test_gpu.sh

@@ -18,18 +18,15 @@ fi
 echo "Installing pytorch,transformers and pytest to the environment for crossfit tests..."
 mamba install \
   cuda-version=$CUDA_VERSION \
-  conda-forge::pytorch \
-  conda-forge::transformers \
-  conda-forge::pytest \
+  "pytorch>=2.*=*cuda*" \
+  transformers \
+  pytest \
+  sentence-transformers  \
+  sentencepiece \
   -c conda-forge \
-  --override-channels \
+  -c nvidia \
   --yes
 
-# Have to install sentence-transformers from pip
-# because conda-forge leads to a torch vision conflict
-# which leads to it being installed on CPUs
-pip3 install sentence-transformers sentencepiece
-
 # Install the crossfit package in editable mode with test dependencies
 pip3 install -e '.[test]'
 # Running tests

conda/environments/cuda_dev.yaml

@@ -20,13 +20,15 @@ dependencies:
   - pytest
   - pytest-benchmark
   - pytest-cov>=2
-  - pytorch>=2.0
+  - pytorch>=2.0=*cuda*
+  - tiktoken
   - transformers
   - tokenizers
   - torchvision
   - scikit-learn>=0.23.1
   - scipy
   - setuptools>=61.0.0
+  - sentencepiece
   - sentence-transformers
   - pip:
     - pytrec-eval

crossfit/op/base.py

@@ -26,7 +26,10 @@ def __init__(self, pre=None, cols=False, keep_cols=None):
         self.pre = pre
         self.cols = cols
         self.keep_cols = keep_cols or []
-        self.worker_name = getattr(self.get_worker(), "name", 0)
+
+    @property
+    def worker_name(self):
+        return getattr(self.get_worker(), "name", 0)
 
     def setup(self):
         pass

requirements/cuda12x.txt

@@ -11,3 +11,4 @@ curated-transformers>=1.0
 bitsandbytes>=0.30
 sentence-transformers>=2.0
 sentencepiece
+tiktoken