# Training on TPU does not work

**URL:** </static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ftraining-on-tpu-does-not-work%2F83878%3E%5C
**Category:** Beginners\
**Created:** [April 26, 2024, 11:47am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ftraining-on-tpu-does-not-work%2F83878 "2024-04-26T11:47:22Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![madslun](https://avatars.discourse-cdn.com/v4/letter/m/51bf81/32.png) [@madslun](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fmadslun)%5C
**Post date:** [April 26, 2024, 11:47am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ftraining-on-tpu-does-not-work%2F83878%2F1 "2024-04-26T11:47:23Z")

</div>

I am currently trying to train mBERT to do a binary classification on TPU.

I am connected to the TPU-VM through ssh. I have set up trainer and training args like this

```auto

def preprocess_text(df):
    length = 200
    tokenized_text = tokenizer(df['concat_text'], padding="max_length", max_length=length)
    tokenized_text["label"] = df['post']
    return tokenized_text

def prepare_dataset():
    train_data = load_dataset("csv", data_files=training_file) \
        .map(preprocess_text, batched=True)
    test_data = load_dataset("csv", data_files=test_file) \
        .map(preprocess_text, batched=True)

    return train_data, test_data

def get_training_args(pre_model):
    return TrainingArguments(
        output_dir="model_run",
        logging_dir="model_run_logs",
        learning_rate=9e-6,
        dataloader_num_workers=4,
        do_train=True,
        num_train_epochs=5,
        weight_decay=0.01,
        save_steps=5000,
        eval_steps=5000,
        evaluation_strategy="steps",
        save_strategy="steps",
        load_best_model_at_end=True,
        push_to_hub=False,
        save_total_limit=4,
        run_name=pre_model,
        auto_find_batch_size=True,
        tpu_num_cores=8,
        tpu_metrics_debug=True

    )

def run_training(pre_model):
    global tokenizer
    tokenizer = AutoTokenizer.from_pretrained(pre_model)
    data_collator = DefaultDataCollator()

    tokenized_train, tokenized_test = prepare_dataset()

    model = AutoModelForSequenceClassification.from_pretrained(
        pre_model, num_labels=2, id2label=id2label, label2id=label2id, trust_remote_code=True,
        problem_type="single_label_classification"
    ).to(device)

    trainer = Trainer(
        model=model,
        args=get_training_args(pre_model),
        train_dataset=tokenized_train["train"],
        eval_dataset=tokenized_test["train"],
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
        preprocess_logits_for_metrics=pre_process_logits,

    )

```

I am setting TPU as device by doing:

```auto
os.environ["PJRT_DEVICE"]="TPU"
device = xm.xla_device()

```

I then run by doing

```auto
curr_model = "bert-base-multilingual-cased"
pre_model = curr_model
run_training(curr_model)

```

Running this fine-tuning takes forever, so I seem to do something wrong. I read that to avoid TPU compiling for each step you must set max\_length for tokenizer. However, this does not seem to work as it still takes forever.

What am I doing wrong? Any help would be much appreciated.