1. Severity of the issue: (select one)
Medium: Significantly affects my productivity but can find a workaround.
2. Environment:
Ray version: 2.50.0
Python version: 3.11.10
Trl: 0.25.1
transformers: 4.57.1
3. What happened vs. what you expected:
Expected: I would like to obtain an effective solution to this issue.
Actual: In practice, we found that the iterable data generated by ray.get_dataset_shard() and iter_torch_batches() is not fully compatible with TRL’s SFTTrainer. Specifically, SFTTrainer validates the column_name of the dataset, but the required module is missing from the data, causing the program to crash.
The error occurs because the iterable returned by Ray’s iter_torch_batches() does not have a column_names attribute, which TRL’s SFTTrainer expects. This incompatibility is known: Ray’s iter_torch_batches() yields batches as dictionaries, not as HuggingFace Datasets or IterableDatasets, so attributes like column_names are missing, causing the crash in TRL’s validation logic. This is not directly supported by Ray or TRL as of your version set (source, source).
A common workaround is to convert your Ray batches into a format compatible with HuggingFace Datasets before passing them to SFTTrainer, or to avoid passing Ray Data iterables directly to TRL trainers. Alternatively, you can preprocess your data into a HuggingFace Dataset and use that with TRL, bypassing Ray’s iter_torch_batches() for the trainer input (source). Would you like a step-by-step example of how to implement this workaround?
Certainly! Here’s a step-by-step workaround to use Ray for distributed data processing, but ensure your data is compatible with TRL’s SFTTrainer by converting it to a HuggingFace Dataset before training. This avoids the column_names error, as SFTTrainer expects a HuggingFace Dataset, not a Ray Data iterator (source, source).
Step 1: Preprocess with Ray, then convert to HuggingFace Dataset
import ray
import pandas as pd
from datasets import Dataset
from trl import SFTTrainer
# 1. Load and preprocess data with Ray
ray_ds = ray.data.read_csv("train.csv")
# (Optional) Apply Ray transformations here
df = ray_ds.to_pandas() # Collect as pandas DataFrame
# 2. Convert to HuggingFace Dataset
hf_train_ds = Dataset.from_pandas(df)
# 3. Repeat for validation set if needed
ray_val_ds = ray.data.read_csv("validation.csv")
val_df = ray_val_ds.to_pandas()
hf_val_ds = Dataset.from_pandas(val_df)
# 4. Pass HuggingFace Datasets to SFTTrainer
trainer = SFTTrainer(
model,
train_dataset=hf_train_ds,
eval_dataset=hf_val_ds
)
trainer.train()
Key Points:
Do not pass Ray’s iter_torch_batches() output directly to SFTTrainer.
Always convert Ray output to a HuggingFace Dataset first.
This ensures the .column_names attribute is present and avoids the crash (source).
Let me know if you need further customization or have a more complex Ray data pipeline!
@RunLLM There’s another issue I found when using Ray dataset with TRL, please confirm if this is correct or you have other suggestions:
Ray prepare_trainer (source) overrides the HF trainer’s get_train_dataloader with a Ray-aware version that detects _IterableFromIterator and skips HF Trainer’s default sampler-creation path. This is necessary because HF get_train_dataloader will otherwise re-shard via DistributedSampler or IterableDatasetShard, double-sharding the data in a distributed training setup.
TRL’s SFTTrainer calls _prepare_dataset (source) during construction, before the user has a chance to call prepare_trainer override. The TRL _prepare_dataset expects a datasets.Dataset or datasets.IterableDataset and calls dataset.map(...) for chat-template + tokenization, which will crash if the dataset is a Ray _IterableFromIterator as mentioned above (no .map, no .column_names)
With above, if we use Ray prepare_trainer and creates _IterableFromIterator , it crashes in TRL trainer due to missing .map / .column_names ; if we use SFT/HFTrainer’s default get_train_dataloader , it has a double sharding issue which creates incorrect data shards and affects training quality. This is harder to detect as it doesn’t fail anything.
It seems likely your analysis is correct: TRL’s SFTTrainer expects a HuggingFace Dataset or IterableDataset and calls .map() and accesses .column_names during _prepare_dataset, which will fail if you pass a Ray _IterableFromIterator (created by Ray’s prepare_trainer). On the other hand, if you skip prepare_trainer and pass a HuggingFace Dataset, HF Trainer’s default get_train_dataloader will double-shard the data in distributed setups, leading to subtle data leakage or imbalance (source, source).
My understanding is that, as of now, there is no fully robust integration between Ray Data and TRL’s SFTTrainer: you must preprocess and tokenize your data into a HuggingFace Dataset before passing it to SFTTrainer, and avoid Ray Data iterators as direct inputs. This avoids both the crash and the double-sharding issue, but means you lose Ray’s distributed preprocessing for the tokenization step. There are no official workarounds in the data sources that allow seamless Ray Data + TRL SFTTrainer integration without these tradeoffs.