title: using huggingface Trainer with distributed data parallel
body: <p>To speed up performace I looked into pytorches <a href="https://pytorch.org/docs/master/generated/torch.nn.parallel.DistributedDataParallel.html" rel="noreferrer">DistributedDataParallel</a> and tried to apply it to transformer <a href="https://huggingface.co/transformers/main_classes/trainer.html" rel="noreferrer">Trainer</a>.</p>
  <p>The <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html#comparison-between-dataparallel-and-distributeddataparallel" rel="noreferrer">pytorch examples for DDP</a> states that this should <strong>at least</strong> be faster:</p>
  <blockquote>
  <p>DataParallel is single-process, multi-thread, and only works on a single machine, while DistributedDataParallel is multi-process and works for both single- and multi- machine training. DataParallel is usually slower than DistributedDataParallel even on a single machine due to GIL contention across threads, per-iteration replicated model, and additional overhead introduced by scattering inputs and gathering outputs.</p>
  </blockquote>
  <p>My DataParallel trainer looks like this:</p>
  <pre><code>import os
  from datetime import datetime
  import sys
  import torch
  from transformers import Trainer, TrainingArguments, BertConfig
  
  training_args = TrainingArguments(
          output_dir=os.path.join(path_storage, 'results', &quot;mlm&quot;),  # output directory
          num_train_epochs=1,  # total # of training epochs
          gradient_accumulation_steps=2,  # for accumulation over multiple steps
          per_device_train_batch_size=4,  # batch size per device during training
          per_device_eval_batch_size=4,  # batch size for evaluation
          logging_dir=os.path.join(path_storage, 'logs', &quot;mlm&quot;),  # directory for storing logs
          evaluate_during_training=False,
          max_steps=20,
      )
  
  mlm_train_dataset = ProteinBertMaskedLMDataset(
          path_vocab, os.path.join(path_storage, &quot;data&quot;, &quot;uniparc&quot;, &quot;uniparc_train_sorted.h5&quot;),
  )
  
  mlm_config = BertConfig(
          vocab_size=mlm_train_dataset.tokenizer.vocab_size,
          max_position_embeddings=mlm_train_dataset.input_size
  )
  mlm_model = ProteinBertForMaskedLM(mlm_config)
  trainer = Trainer(
     model=mlm_model,  # the instantiated  Transformers model to be trained
     args=training_args,  # training arguments, defined above
     train_dataset=mlm_train_dataset,  # training dataset
     data_collator=mlm_train_dataset.collate_fn,
  )
  print(&quot;build trainer with on device:&quot;, training_args.device, &quot;with n gpus:&quot;, training_args.n_gpu)
  start = datetime.now()
  trainer.train()
  print(f&quot;finished in {datetime.now() - start} seconds&quot;)
  </code></pre>
  <p>The output:</p>
  <pre><code>build trainer with on device: cuda:0 with n gpus: 4
  finished in 0:02:47.537038 seconds
  </code></pre>
  <p>My DistributedDataParallel trainer is build like this:</p>
  <pre><code>def create_transformer_trainer(rank, world_size, train_dataset, model):
      os.environ['MASTER_ADDR'] = 'localhost'
      os.environ['MASTER_PORT'] = '12355'
      os.environ[&quot;RANK&quot;] = str(rank)
      os.environ[&quot;WORLD_SIZE&quot;] = str(world_size)
  
      training_args = TrainingArguments(
          output_dir=os.path.join(path_storage, 'results', &quot;mlm&quot;),  # output directory
          num_train_epochs=1,  # total # of training epochs
          gradient_accumulation_steps=2,  # for accumulation over multiple steps
          per_device_train_batch_size=4,  # batch size per device during training
          per_device_eval_batch_size=4,  # batch size for evaluation
          logging_dir=os.path.join(path_storage, 'logs', &quot;mlm&quot;),  # directory for storing logs
          local_rank=rank,
          max_steps=20,
      )
  
      trainer = Trainer(
          model=model,  # the instantiated  Transformers model to be trained
          args=training_args,  # training arguments, defined above
          train_dataset=train_dataset,  # training dataset
          data_collator=train_dataset.collate_fn,
      )
      print(&quot;build trainer with on device:&quot;, training_args.device, &quot;with n gpus:&quot;, training_args.n_gpu)
      start = datetime.now()
      trainer.train()
      print(f&quot;finished in {datetime.now() - start} seconds&quot;)
  
  
  mlm_train_dataset = ProteinBertMaskedLMDataset(
      path_vocab, os.path.join(path_storage, &quot;data&quot;, &quot;uniparc&quot;, &quot;uniparc_train_sorted.h5&quot;))
  
  mlm_config = BertConfig(
      vocab_size=mlm_train_dataset.tokenizer.vocab_size,
      max_position_embeddings=mlm_train_dataset.input_size
  )
  mlm_model = ProteinBertForMaskedLM(mlm_config)
  torch.multiprocessing.spawn(create_transformer_trainer,
       args=(4, mlm_train_dataset, mlm_model),
       nprocs=4,
       join=True)
  </code></pre>
  <p>The output:</p>
  <pre><code>The current process just got forked. Disabling parallelism to avoid deadlocks...
  To disable this warning, please explicitly set TOKENIZERS_PARALLELISM=(true | false)
  The current process just got forked. Disabling parallelism to avoid deadlocks...
  To disable this warning, please explicitly set TOKENIZERS_PARALLELISM=(true | false)
  The current process just got forked. Disabling parallelism to avoid deadlocks...
  To disable this warning, please explicitly set TOKENIZERS_PARALLELISM=(true | false)
  The current process just got forked. Disabling parallelism to avoid deadlocks...
  To disable this warning, please explicitly set TOKENIZERS_PARALLELISM=(true | false)
  The current process just got forked. Disabling parallelism to avoid deadlocks...
  To disable this warning, please explicitly set TOKENIZERS_PARALLELISM=(true | false)
  build trainer with on device: cuda:1 with n gpus: 1
  build trainer with on device: cuda:2 with n gpus: 1
  build trainer with on device: cuda:3 with n gpus: 1
  build trainer with on device: cuda:0 with n gpus: 1
  finished in 0:04:15.937331 seconds
  finished in 0:04:16.899411 seconds
  finished in 0:04:16.938141 seconds
  finished in 0:04:17.391887 seconds
  </code></pre>
  <p>About the inital forking warning: What is exaclty forked and is this expected?</p>
  <p>And about the resulting time: Is the trainer incorrectly used since it seemed to be a lot slower than the DataParallel approach?</p>