Merge branch 'main' into dcp-hf-util

Author: dushyantbehl

README.md

@@ -8,7 +8,7 @@
   - [Advanced Data Processing](./docs/advanced-data-preprocessing.md#data-config)
   - [Guidelines on supported data formats](./docs/advanced-data-preprocessing.md#use-cases-supported-via-command-line-argument-training_data_path)
   - [Offline data processing](#offline-data-preprocessing)
-  - [Online data mixing](./docs/online-data-mixing.md)
+  - [Online data mixing](./docs/advanced-data-preprocessing.md#online-data-mixing-section)
 - [Additional Frameworks](#additional-frameworks)
   - [Inference](#inference)
   - [Validation](#validation)

build/Dockerfile

@@ -25,6 +25,7 @@ ARG ENABLE_MLFLOW=false
 ARG ENABLE_FMS_ACCELERATION=true
 ARG ENABLE_SCANNER=false
 ARG ENABLE_CLEARML=false
+ARG ENABLE_RECOMMENDER=true
 
 ## Base Layer ##################################################################
 FROM registry.access.redhat.com/ubi9/ubi:${BASE_UBI_IMAGE_TAG} AS base
@@ -188,6 +189,9 @@ RUN if [[ "${ENABLE_SCANNER}" == "true" ]]; then \
 RUN if [[ "${ENABLE_CLEARML}" == "true" ]]; then \
         python -m pip install --user "$(head bdist_name)[clearml]"; \
     fi
+RUN if [[ "${ENABLE_RECOMMENDER}" == "true" ]]; then \
+        python -m pip install --user "$(head bdist_name)[tuning-config-recommender]"; \
+    fi
 
     # Clean up the wheel module. It's only needed by flash-attn install
 RUN python -m pip uninstall wheel build -y && \

build/nvcr.Dockerfile

@@ -34,6 +34,7 @@ ARG ENABLE_MLFLOW=false
 ARG ENABLE_SCANNER=false
 ARG ENABLE_CLEARML=true
 ARG ENABLE_TRITON_KERNELS=true
+ARG ENABLE_RECOMMENDER=true
 
 # Ensures to always build mamba_ssm from source
 ENV PIP_NO_BINARY=mamba-ssm,mamba_ssm
@@ -76,6 +77,9 @@ RUN if [[ "${ENABLE_MLFLOW}" == "true" ]]; then \
 RUN if [[ "${ENABLE_SCANNER}" == "true" ]]; then \
         pip install --no-cache-dir ${SOURCE_DIR}[scanner-dev]; \
     fi
+RUN if [[ "${ENABLE_RECOMMENDER}" == "true" ]]; then \
+        python -m pip install --user "$(head bdist_name)[tuning-config-recommender]"; \
+    fi
 
 # cleanup
 RUN rm -rf /root/.cache /tmp/* /opt/pytorch

docs/advanced-data-preprocessing.md

@@ -162,7 +162,8 @@ Each data handler has:
     - `sampling` (optional, float): The sampling ratio (0.0 to 1.0) with which to sample a dataset in case of interleaving.
     - `split` (optional, dict[str: float]): Defines how to split the dataset into training and validation sets. Requires both `train` and `validation` keys.
     - `data_handlers` (optional, list): A list of data handler configurations which preprocess the dataset.
-
+    - `dataset_split_name` (optional, str): Name of the dataset split. This is useful for loading HuggingFace datasets with split names that are different from the standard (eg: `train_sft` instead of `train`). If no `dataset_split_name` is provided, `train` is used.
+    - `shuffle` (optional, bool): If the dataset should be shuffled while splitting into train and validation split. Defaults to `True`. Use caution when using this field and only use when the dataset is already shuffled.
 
 We do provide some sample `data_configs` here, [predefined_data_configs](../tests/artifacts/predefined_data_configs/).
 

pyproject.toml

@@ -38,12 +38,12 @@ dependencies = [
 "peft>=0.18.0,< 0.19.0",
 "datasets>=4.0.0,<5.0.0",
 "simpleeval>=0.9.13,<2.0",
-"pillow>=11.0.0,<12.0",
+"pillow>=12.1.1",
 "kernels<=0.9.0",
 ]
 
 [project.optional-dependencies]
-dev = ["wheel>=0.42.0,<1.0", "packaging>=23.2,<25", "ninja>=1.11.1.1,<2.0", "scikit-learn>=1.0, <2.0", "boto3>=1.34, <2.0"]
+dev = ["wheel>=0.42.0,<1.0", "packaging>=23.2,<25", "ninja>=1.11.1.1,<2.0", "scikit-learn>=1.0, <2.0", "boto3>=1.34, <2.0", "hf_transfer>=0.1.9"]
 flash-attn = ["flash-attn>=2.8.3"]
 aim = ["aim>=3.19.0,<4.0"]
 mlflow = ["mlflow"]
@@ -60,6 +60,7 @@ fms-accel-all = [
     "fms-acceleration-moe",
     "fms-acceleration-odm"
 ]
+tuning-config-recommender=["tuning-config-recommender>=0.1.5"]
 
 [tool.setuptools.packages.find]
 exclude = ["tests", "tests.*"]

tests/artifacts/predefined_data_configs/__init__.py

@@ -86,3 +86,6 @@
 DATA_CONFIG_SKIP_LARGE_COLUMNS_HANDLER = os.path.join(
     PREDEFINED_DATA_CONFIGS, "skip_large_columns_data_handler_template.yaml"
 )
+DATA_CONFIG_CUSTOM_SPLIT_NAME = os.path.join(
+    PREDEFINED_DATA_CONFIGS, "dataset_with_custom_split.yaml"
+)

tests/artifacts/predefined_data_configs/dataset_with_custom_split.yaml

@@ -0,0 +1,34 @@
+dataprocessor:
+    type: default
+    sampling_stopping_strategy: first_exhausted
+    seed: 66
+datasets:
+  - name: dataset_split_custom_split_name
+    split:
+      train: 0.8
+      validation: 0.2
+    sampling: 0.5
+    dataset_split_name: "train_sft"
+    data_paths:
+      - "FILE_PATH"
+    data_handlers:
+      - name: tokenize_and_apply_chat_template_with_masking
+        arguments:
+          remove_columns: all
+          batched: false
+          fn_kwargs:
+            formatted_text_column_name: "formatted_chat_data"
+            conversation_column: "messages"
+  - name: dataset_wo_split_custom_split_name
+    sampling: 0.5
+    dataset_split_name: "train_sft"
+    data_paths:
+      - "FILE_PATH"
+    data_handlers:
+      - name: tokenize_and_apply_chat_template_with_masking
+        arguments:
+          remove_columns: all
+          batched: false
+          fn_kwargs:
+            formatted_text_column_name: "formatted_chat_data"
+            conversation_column: "messages"

tests/artifacts/testdata/__init__.py

@@ -86,6 +86,7 @@
 IMAGE_DATASET = os.path.join(JSONL_DATA_DIR, "image_dataset.jsonl")
 EMPTY_DATA = os.path.join(JSON_DATA_DIR, "empty_data.json")
 MALFORMATTED_DATA = os.path.join(JSON_DATA_DIR, "malformatted_data.json")
+CHAT_DATA_HF_HOSTED_CUSTOM_SPLIT = "rom7/test-OpenHermes-2.5-H4"
 
 # Other constants
 CUSTOM_TOKENIZER_TINYLLAMA = os.path.join(

tests/test_sft_trainer.py

@@ -42,6 +42,7 @@
 from tests.artifacts.language_models import MAYKEYE_TINY_LLAMA_CACHED, TINYMIXTRAL_MOE
 from tests.artifacts.predefined_data_configs import (
     CHAT_TEMPLATE_JINJA,
+    DATA_CONFIG_CUSTOM_SPLIT_NAME,
     DATA_CONFIG_DUPLICATE_COLUMNS,
     DATA_CONFIG_INVALID_BASE64_CHAT_TEMPLATE,
     DATA_CONFIG_MULTIPLE_DATASETS_ODM_YAML,
@@ -61,6 +62,7 @@
     GRANITE_3_1_B_CHAT_TEMPLATE,
 )
 from tests.artifacts.testdata import (
+    CHAT_DATA_HF_HOSTED_CUSTOM_SPLIT,
     CHAT_DATA_MULTI_TURN,
     CHAT_DATA_MULTI_TURN_CONVERSATIONS,
     CHAT_DATA_MULTI_TURN_GRANITE_3_1B,
@@ -949,16 +951,16 @@ def test_run_causallm_lora_add_special_tokens():
             ["lm_head"],
             ["embed_tokens"],
             marks=pytest.mark.skipif(
-                version.parse(peft.__version__) <= version.parse("0.18.0"),
-                reason="Not released in PEFT <= 0.18.0",
+                version.parse(peft.__version__) <= version.parse("0.18.1"),
+                reason="Not released in PEFT <= 0.18.1",
             ),
         ),
         pytest.param(
             ["embed_tokens", "lm_head"],
             ["embed_tokens"],
             marks=pytest.mark.skipif(
-                version.parse(peft.__version__) <= version.parse("0.18.0"),
-                reason="Not released in PEFT <= 0.18.0",
+                version.parse(peft.__version__) <= version.parse("0.18.1"),
+                reason="Not released in PEFT <= 0.18.1",
             ),
         ),
     ],
@@ -1010,8 +1012,8 @@ def test_run_causallm_lora_tied_weights_in_modules_to_save(modules_to_save, expe
     ],
 )
 @pytest.mark.skipif(
-    version.parse(peft.__version__) <= version.parse("0.18.0"),
-    reason="Not released in PEFT <= 0.18.0",
+    version.parse(peft.__version__) <= version.parse("0.18.1"),
+    reason="Not released in PEFT <= 0.18.1",
 )
 def test_run_causallm_lora_tied_weights_in_target_modules(target_modules, expected):
     """Check if a model with tied weights in target_modules is correctly trained"""
@@ -1916,6 +1918,61 @@ def test_run_moe_ft_with_save_model_dir(dataset_path):
         assert os.path.exists(os.path.join(save_model_dir))
 
 
+@pytest.mark.parametrize(
+    "datafiles, dataconfigfile",
+    [
+        (
+            [CHAT_DATA_HF_HOSTED_CUSTOM_SPLIT, CHAT_DATA_HF_HOSTED_CUSTOM_SPLIT],
+            DATA_CONFIG_CUSTOM_SPLIT_NAME,
+        )
+    ],
+)
+def test_run_chat_style_ft_using_custom_split_name(datafiles, dataconfigfile):
+    """Check if we can select custom split for a dataset."""
+    with tempfile.TemporaryDirectory() as tempdir:
+        data_args = copy.deepcopy(DATA_ARGS)
+        data_args.training_data_path = None
+        data_args.response_template = None
+        data_args.dataset_text_field = None
+        data_args.chat_template = CHAT_TEMPLATE_JINJA
+
+        model_args = copy.deepcopy(MODEL_ARGS)
+        model_args.model_name_or_path = TINYMIXTRAL_MOE
+        model_args.tokenizer_name_or_path = TINYMIXTRAL_MOE
+
+        train_args = copy.deepcopy(TRAIN_ARGS)
+        train_args.output_dir = tempdir
+
+        with tempfile.NamedTemporaryFile(
+            "w", delete=False, suffix=".yaml"
+        ) as temp_yaml_file:
+            with open(dataconfigfile, "r", encoding="utf-8") as f:
+                data = yaml.safe_load(f)
+            datasets = data["datasets"]
+            for i, d in enumerate(datasets):
+                d["data_paths"] = [datafiles[i]]
+            yaml.dump(data, temp_yaml_file)
+            data_args.data_config_path = temp_yaml_file.name
+
+        sft_trainer.train(model_args, data_args, train_args)
+
+        # validate the configs
+        _validate_training(tempdir)
+        checkpoint_path = _get_checkpoint_path(tempdir)
+
+        # Load the model
+        loaded_model = TunedCausalLM.load(checkpoint_path, MODEL_NAME)
+
+        # Run inference on the text
+        output_inference = loaded_model.run(
+            '<|user|>\nProvide two rhyming words for the word "love"\n\
+            <nopace></s><|assistant|>',
+            max_new_tokens=50,
+        )
+        assert len(output_inference) > 0
+        assert 'Provide two rhyming words for the word "love"' in output_inference
+
+
 ############################# Helper functions #############################
 def _test_run_causallm_ft(training_args, model_args, data_args, tempdir):
     train_args = copy.deepcopy(training_args)

tox.ini

@@ -5,7 +5,7 @@ envlist = py, lint, fmt
 description = run unit tests
 deps =
     pytest>=7
-    .[aim,mlflow,clearml,scanner-dev]
+    .[mlflow,clearml,scanner-dev]
 commands =
     pytest {posargs:tests}
 
@@ -56,9 +56,9 @@ commands =
 
 [testenv:accel]
 description = run all unit tests including requring GPU support
-deps = 
+deps =
     pytest>=7
-    .[aim,mlflow,clearml,scanner-dev,fms-accel-all]
+    .[mlflow,clearml,scanner-dev,fms-accel-all]
 setenv =
     CUDA_VISIBLE_DEVICES=0
 commands_pre = 
@@ -74,9 +74,9 @@ commands =
 
 [testenv:gpu]
 description = run all unit tests including requring GPU support
-deps = 
+deps =
     pytest>=7
-    .[aim,mlflow,clearml,scanner-dev,fms-accel-all]
+    .[mlflow,clearml,scanner-dev,fms-accel-all]
 setenv =
     CUDA_VISIBLE_DEVICES=0
 commands_pre =