Merge branch 'main' into feat/odm-wo-categories

Author: dushyantbehl

build/nvcr.Dockerfile

@@ -47,8 +47,9 @@ RUN pip install --upgrade --force-reinstall torch torchaudio torchvision --index
 # Install main package + flash attention
 COPY . ${SOURCE_DIR}
 RUN cd ${SOURCE_DIR}
-RUN pip install --no-cache-dir ${SOURCE_DIR} && \
-    pip install --no-cache-dir ${SOURCE_DIR}[flash-attn]
+
+RUN pip install --no-cache-dir ${SOURCE_DIR}
+RUN pip install --user --no-build-isolation ${SOURCE_DIR}[flash-attn]
 
 # Optional extras
 RUN if [[ "${ENABLE_FMS_ACCELERATION}" == "true" ]]; then \

pyproject.toml

@@ -13,7 +13,7 @@ authors = [
   {name = "Sukriti Sharma", email = "[email protected]"},
   {name = "Anh Uong", email = "[email protected]"},
 ]
-license = "Apache-2.0"
+license = {text = "Apache-2.0"}
 readme = "README.md"
 requires-python = "~=3.9"
 keywords = ['fms-hf-tuning', 'python', 'tuning']
@@ -30,6 +30,7 @@ dependencies = [
 "accelerate>=1.9.0,<2.0.0",
 "transformers>=4.55.0,<=4.55.4",
 "torch>2.7.0,<2.9.0",
+"torchvision<0.24",
 "sentencepiece>=0.1.99,<0.3",
 "tokenizers<=0.22",
 "tqdm>=4.66.2,<5.0",

tuning/data/setup_dataprocessor.py

@@ -570,7 +570,7 @@ def setup_train_dataset_for_odm(
         reward_type=odm_config.odm.reward_type,
         auto_categorize_config=auto_categorize_config,
     )
-    train_args.accelerator_config = {"split_batches": True}
+    train_args.accelerator_config = {"dispatch_batches": False}
     return (True, train_dataset, data_collator)