szymon_dybczak
Esteemed Contributor III

Hi ,

The he problem is that TrainingArguments triggers a distributed training detection routine internally, which tries to inspect the environment for things like MPI, OpenMPI, or other distributed frameworks. In Databricks, this probing hangs because the cluster environment has partial distributed computing infrastructure (Spark) that responds to some of those checks but never completes them.

Try to set following environment variables:

import os
os.environ["WORLD_SIZE"] = "1"
os.environ["LOCAL_RANK"] = "-1"

from transformers import TrainingArguments

print("start")

args = TrainingArguments(
    output_dir="test",
    use_cpu=True              
)

print("end")