[AIR] HuggingFaceTrainer&Predictor implementation

doc/source/ray-air/getting-started.rst

@@ -60,6 +60,10 @@ Trainer
     :members:
     :show-inheritance:
 
+.. automodule:: ray.ml.train.integrations.huggingface
+    :members:
+    :show-inheritance:
+
 .. automodule:: ray.ml.train.integrations.sklearn
     :members:
     :show-inheritance:

python/ray/data/dataset.py

@@ -26,6 +26,7 @@
     import ray.util.sgd
     import torch
     import tensorflow as tf
+    import torch.utils.data
     from ray.data.dataset_pipeline import DatasetPipeline
     from ray.data.grouped_dataset import GroupedDataset
 
@@ -302,7 +303,8 @@ def transform(block: Block) -> Iterable[Block]:
                 ):
                     raise ValueError(
                         "The map batches UDF returned the value "
-                        f"{applied}, which is not allowed. "
+                        f"{applied} of type {type(applied)}, "
+                        "which is not allowed. "
                         "The return type must be either list, "
                         "pandas.DataFrame, or pyarrow.Table"
                     )
@@ -2072,6 +2074,7 @@ def to_torch(
         prefetch_blocks: int = 0,
         drop_last: bool = False,
         unsqueeze_label_tensor: bool = True,
+        unsqueeze_feature_tensors: bool = True,
     ) -> "torch.utils.data.IterableDataset":
         """Return a Torch IterableDataset over this dataset.
 
@@ -2145,6 +2148,10 @@ def to_torch(
                 be left as is, that is (N, ). In general, regression loss
                 functions expect an unsqueezed tensor, while classification
                 loss functions expect a squeezed one. Defaults to True.
+            unsqueeze_feature_tensors (bool): If set to True, the features tensors
+                will be unsqueezed (reshaped to (N, 1)) before being concatenated into
+                the final features tensor. Otherwise, they will be left as is, that is
+                (N, ). Defaults to True.
 
         Returns:
             A torch IterableDataset.
@@ -2196,10 +2203,13 @@ def make_generator():
                 drop_last=drop_last,
             ):
                 if label_column:
-                    label_vals = batch.pop(label_column).values
-                    label_tensor = torch.as_tensor(label_vals, dtype=label_column_dtype)
-                    if unsqueeze_label_tensor:
-                        label_tensor = label_tensor.view(-1, 1)
+                    label_tensor = convert_pandas_to_torch_tensor(
+                        batch,
+                        [label_column],
+                        label_column_dtype,
+                        unsqueeze=unsqueeze_label_tensor,
+                    )
+                    batch.pop(label_column)
                 else:
                     label_tensor = None
 
@@ -2211,6 +2221,7 @@ def make_generator():
                             feature_column_dtypes[key]
                             if isinstance(feature_column_dtypes, dict)
                             else feature_column_dtypes,
+                            unsqueeze=unsqueeze_feature_tensors,
                         )
                         for key in feature_columns
                     }
@@ -2219,6 +2230,7 @@ def make_generator():
                         batch,
                         columns=feature_columns,
                         column_dtypes=feature_column_dtypes,
+                        unsqueeze=unsqueeze_feature_tensors,
                     )
 
                 yield (features_tensor, label_tensor)

python/ray/ml/examples/huggingface_example.py

@@ -0,0 +1,94 @@
+# Based on
+# huggingface/notebooks/examples/language_modeling_from_scratch.ipynb
+
+from datasets import load_dataset
+from transformers import (
+    AutoTokenizer,
+    AutoConfig,
+    AutoModelForCausalLM,
+    Trainer,
+    TrainingArguments,
+)
+
+import ray
+import ray.data
+from ray.ml.train.integrations.huggingface import HuggingFaceTrainer
+
+model_checkpoint = "gpt2"
+tokenizer_checkpoint = "sgugger/gpt2-like-tokenizer"
+
+# block_size = tokenizer.model_max_length
+block_size = 128
+
+
+def get_dataset():
+    datasets = load_dataset("wikitext", "wikitext-2-raw-v1")
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_checkpoint)
+
+    def tokenize_function(examples):
+        return tokenizer(examples["text"])
+
+    tokenized_datasets = datasets.map(
+        tokenize_function, batched=True, num_proc=1, remove_columns=["text"]
+    )
+
+    def group_texts(examples):
+        # Concatenate all texts.
+        concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
+        total_length = len(concatenated_examples[list(examples.keys())[0]])
+        # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+        # customize this part to your needs.
+        total_length = (total_length // block_size) * block_size
+        # Split by chunks of max_len.
+        result = {
+            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+            for k, t in concatenated_examples.items()
+        }
+        result["labels"] = result["input_ids"].copy()
+        return result
+
+    lm_datasets = tokenized_datasets.map(
+        group_texts,
+        batched=True,
+        batch_size=1000,
+        num_proc=1,
+    )
+    return lm_datasets
+
+
+lm_dataset = get_dataset()
+ray_train = ray.data.from_arrow(lm_dataset["train"]._data.table)
+ray_validation = ray.data.from_arrow(lm_dataset["validation"]._data.table)
+
+
+def train_function(train_dataset, eval_dataset=None, **config):
+    model_config = AutoConfig.from_pretrained(model_checkpoint)
+    model = AutoModelForCausalLM.from_config(model_config)
+    print("Initializing TrainingArguments...")
+    training_args = TrainingArguments(
+        f"{model_checkpoint}-wikitext2",
+        evaluation_strategy="epoch",
+        num_train_epochs=2,
+        learning_rate=2e-5,
+        weight_decay=0.01,
+        disable_tqdm=True,
+        save_strategy="epoch",
+    )
+    print("Initializing Trainer...")
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+    )
+    print("Trainer initialized! Starting training...")
+    return trainer
+
+
+trainer = HuggingFaceTrainer(
+    trainer_init_per_worker=train_function,
+    scaling_config={"num_workers": 2, "use_gpu": False},
+    datasets={"train": ray_train.limit(16), "evaluation": ray_validation.limit(8)},
+)
+results = trainer.fit()
+print(results.metrics)

python/ray/ml/train/data_parallel_trainer.py

@@ -249,6 +249,9 @@ def _validate_train_loop_per_worker(
                 f"but it accepts {num_params} arguments instead."
             )
 
+    def _get_checkpoint_manager(self) -> TuneCheckpointManager:
+        return _DataParallelCheckpointManager()
+
     def training_loop(self) -> None:
         scaling_config_dataclass = ScalingConfigDataClass(**self.scaling_config)
 
@@ -271,7 +274,7 @@ def training_loop(self) -> None:
             max_retries=0,
         )
 
-        checkpoint_manager = _DataParallelCheckpointManager()
+        checkpoint_manager = self._get_checkpoint_manager()
         checkpoint_manager.on_init(preprocessor=self.preprocessor)
 
         # Start the remote actors.