Search code examples
pythonpytorchhuggingface-transformers

Why cant I set TrainingArguments.device in Huggingface?

Question

When I try to set the .device attribute to torch.device('cpu'), I get an error. How am I supposed to set device then?

Python Code

from transformers import TrainingArguments
from transformers import Trainer
import torch

training_args = TrainingArguments(
    output_dir="./some_local_dir",
    overwrite_output_dir=True,

    per_device_train_batch_size=4,
    dataloader_num_workers=2,

    max_steps=500,
    logging_steps=1,

    evaluation_strategy="steps",
    eval_steps=5
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    compute_metrics=compute_metrics,
)

training_args.device = torch.device('cpu')

Python Error

AttributeError                            Traceback (most recent call last)
<ipython-input-11-30a92c0570b8> in <cell line: 28>()
     26 )
     27 
---> 28 training_args.device = torch.device('cpu')

AttributeError: can't set attribute
Solution

  • From the docs of the TrainingArguments object doesn't have a settable device attribute.

    But interestingly device is initialized but non mutable:

    import torch
from transformers import TrainingArguments

args = TrainingArguments('./')

args.device  # [out]: device(type='cpu')

args.device = torch.device(type='cpu')

    [out]:

    ---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-12-dcb5ef23be68> in <cell line: 8>()
      6 
      7 
----> 8 args.device = torch.device(type='cpu')

AttributeError: can't set attribute

    From the code, it looks like the device is set after the initialization: https://github.com/huggingface/transformers/blob/main/src/transformers/training_args.py#L1113

    And the device is only setup when the TrainingArguments.device is called

    Perhaps you are referring to the huggingface model and you want to try training it on the CPU. To do that, usually the Trainer will automatically detect it if you accelerate https://huggingface.co/docs/accelerate/index

    By default, the TrainingArguments is already set to CPU if _n_gpu = -1

    But if you want to explicitly set the model to be used on CPU, try:

    model = model.to('cpu')
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    compute_metrics=compute_metrics,
)

    Or:

    trainer = Trainer(
    model=model.to('cpu'),
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    compute_metrics=compute_metrics,
    place_model_on_device=True
)