Huggingface Trainer로 TPU Multicore 활용 학습 예제코드

from datasets import load_dataset
from transformers import AutoTokenizer, DataCollatorWithPadding
from transformers import TrainingArguments, AutoModelForSequenceClassification
from transformers import Trainer

def main(index):
    # 예제로 활용할 데이터셋
    raw_datasets = load_dataset('glue','mrpc')
    # 예제 활용 모델
    model_name = 'klue/roberta-large'
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    
    #raw_dataset은 huggingface dataset의 구조체로 해당 함수를 이용해 각 데이터셋 원소를
    #tokenize함
    def tokenize_fn(batch):
        return tokenizer(batch['sentence1'], batch['sentence2'], max_length=64, padding='max_length', truncation=True)

    tokenized_ds = raw_datasets.map(tokenize_fn, batched=True)

    # tpu 학습시 tpu_num_cores로 해당 tpu core 갯수를 설정해야함
    # per device batch size가 16이므로, global batch size는 128
    args = TrainingArguments('test-train',
                                per_device_train_batch_size=16
                                tpu_num_cores=8)
    #샘플 모델
    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=50)
    
    trainer = Trainer(
        model,
        args,
        train_dataset=tokenized_ds['train'],
        eval_dataset=tokenized_ds['validation'],
        tokenizer=tokenizer,
    )

    trainer.train()

if __name__ == '__main__':
    #xrt_tpu_config가 설정되어야 tpu 정상작동 보장
    import os
    os.environ['XRT_TPU_CONFIG']="localservice;0;localhost:51011"
    # xmp를 통해 위 main 함수를 multiprocessing으로 실행함
    # 각 프로세스별로 데이터셋, 토크나이저, 모델을 따로 만든 뒤
    # 최종 Trainer 단계에서 통합으로 동작함.
    import torch_xla.distributed.xla_multiprocessing as xmp
    xmp.spawn(main, args=(), nprocs=8,)