gpt-oss sft.py script failing on default dataset loading, KeyError: 'text' #326
Description
root@dcc8ba51571c:/workspace/TensorRT-Model-Optimizer/examples/gpt-oss# HF_HOME=/mnt/data/tmp python3 sft.py --config configs/sft_lora.yaml --model_name_or_path openai/gpt-oss-20b --quant_cfg MXFP4_MLP_WEIGHT_ONLY_CFG --output_dir gpt-oss-20b-lora-qat
/usr/local/lib/python3.12/dist-packages/modelopt/torch/utils/import_utils.py:32: UserWarning: Failed to import transformer engine plugin due to: AttributeError("module 'transformer_engine' has no attribute 'pytorch'"). You may ignore this warning if you do not need this plugin.
warnings.warn(
/usr/local/lib/python3.12/dist-packages/modelopt/torch/utils/import_utils.py:32: UserWarning: Failed to import transformer_engine plugin due to: ImportError("cannot import name 'pytorch' from 'transformer_engine' (/usr/local/lib/python3.12/dist-packages/transformer_engine/init.py)"). You may ignore this warning if you do not need this plugin.
warnings.warn(
ModelOpt save/restore enabled for transformers library.
ModelOpt save/restore enabled for diffusers library.
ModelOpt save/restore enabled for peft library.
Loading checkpoint shards: 100%|██████████████████████████████████████████████████████| 3/3 [00:02<00:00, 1.23it/s]
ModelOpt save/restore enabled for transformers library.
ModelOpt save/restore enabled for diffusers library.
ModelOpt save/restore enabled for peft library.
/usr/local/lib/python3.12/dist-packages/peft/tuners/lora/layer.py:159: UserWarning: Unsupported layer type '<class 'transformers.models.gpt_oss.modeling_gpt_oss.GptOssExperts'>' encountered, proceed at your own risk.
warnings.warn(
Map (num_proc=12): 0%| | 0/900 [00:03<?, ? examples/s]
multiprocess.pool.RemoteTraceback:
"""
Traceback (most recent call last):
File "/usr/local/lib/python3.12/dist-packages/multiprocess/pool.py", line 125, in worker
result = (True, func(*args, **kwds))
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/datasets/utils/py_utils.py", line 678, in _write_generator_to_queue
for i, result in enumerate(func(**kwargs)):
File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_dataset.py", line 3458, in _map_single
batch = apply_function_on_filtered_inputs(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_dataset.py", line 3320, in apply_function_on_filtered_inputs
processed_inputs = function(*fn_args, *additional_args, **fn_kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/trl/trainer/sft_trainer.py", line 416, in tokenize
element[dataset_text_field] if formatting_func is None else formatting_func(element),
~~~~~~~^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/datasets/formatting/formatting.py", line 277, in getitem
value = self.data[key]
~~~~~~~~~^^^^^
KeyError: 'text'
"""
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/workspace/TensorRT-Model-Optimizer/examples/gpt-oss/sft.py", line 123, in
main(script_args, training_args, model_args, quant_args)
File "/workspace/TensorRT-Model-Optimizer/examples/gpt-oss/sft.py", line 100, in main
trainer = QATSFTTrainer(
^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/modelopt/torch/quantization/plugins/transformers_trainer.py", line 182, in init
super().init(*args, **kwargs)
File "/usr/local/lib/python3.12/dist-packages/modelopt/torch/opt/plugins/transformers.py", line 159, in init
super().init(*args, **kwargs)
File "/usr/local/lib/python3.12/dist-packages/transformers/utils/deprecation.py", line 172, in wrapped_func
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/trl/trainer/sft_trainer.py", line 265, in init
train_dataset = self._prepare_dataset(
^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/trl/trainer/sft_trainer.py", line 380, in _prepare_dataset
return self._prepare_non_packed_dataloader(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/trl/trainer/sft_trainer.py", line 455, in _prepare_non_packed_dataloader
tokenized_dataset = dataset.map(tokenize, **map_kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_dataset.py", line 560, in wrapper
out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_dataset.py", line 3147, in map
for rank, done, content in iflatmap_unordered(
File "/usr/local/lib/python3.12/dist-packages/datasets/utils/py_utils.py", line 718, in iflatmap_unordered
[async_result.get(timeout=0.05) for async_result in async_results]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/multiprocess/pool.py", line 774, in get
raise self._value
KeyError: 'text'