chiyuzhang94 (NONE) commented on 2020-10-17 19:21:59:
Hi @thomwolf @lhoestq ,

Thanks for your help. 
Finally, the preprocess is completed. I used 32 CPUs. The preprocess took 55 hours. The `dataset.arrow` size is 1.18TB. 
I saved the processed dataset and used it for RoBERTa pre-training. You can find my scripts [here](https://github.com/chiyuzhang94/transformers/blob/master/examples/language-modeling/run_language_modeling.py). 
I loaded the dataset by `load_from_disk()`.
`dataset = load_from_disk(file_path) `

But the job runs out of memory when it loads batch samples at the first step. Here is the error log. I can find that the dataset was loaded successfully. The job was terminated at the first step. I expect the `datasets` library can handle the memory issue effectively. But it still runs out of memory.  
```
I1017 00:22:08.109311 140737353971520 run_language_modeling.py:344]   Total optimization steps = 228650
I1017 00:22:08.109332 140737353971520 run_language_modeling.py:343]   Gradient Accumulation steps = 4
I1017 00:22:08.109332 140737353971520 run_language_modeling.py:343]   Gradient Accumulation steps = 4
I1017 00:22:08.109340 140737353971520 run_language_modeling.py:343]   Gradient Accumulation steps = 4
I1017 00:22:08.109347 140737353971520 run_language_modeling.py:345]   Warm-up steps = 10000
I1017 00:22:08.109370 140737353971520 run_language_modeling.py:344]   Total optimization steps = 228650
I1017 00:22:08.109369 140737353971520 run_language_modeling.py:344]   Total optimization steps = 228650
I1017 00:22:08.109378 140737353971520 run_language_modeling.py:344]   Total optimization steps = 228650
I1017 00:22:08.109405 140737353971520 run_language_modeling.py:345]   Warm-up steps = 10000
I1017 00:22:08.109406 140737353971520 run_language_modeling.py:345]   Warm-up steps = 10000
I1017 00:22:08.109414 140737353971520 run_language_modeling.py:345]   Warm-up steps = 10000
I1017 00:22:08.109496 140737353971520 run_language_modeling.py:359]   Continuing training from checkpoint, will skip to saved global_step
I1017 00:22:08.109497 140737353971520 run_language_modeling.py:359]   Continuing training from checkpoint, will skip to saved global_step
I1017 00:22:08.109499 140737353971520 run_language_modeling.py:359]   Continuing training from checkpoint, will skip to saved global_step
I1017 00:22:08.109500 140737353971520 run_language_modeling.py:359]   Continuing training from checkpoint, will skip to saved global_step
I1017 00:22:08.109534 140737353971520 run_language_modeling.py:360]   Continuing training from epoch 0
I1017 00:22:08.109534 140737353971520 run_language_modeling.py:360]   Continuing training from epoch 0
I1017 00:22:08.109535 140737353971520 run_language_modeling.py:360]   Continuing training from epoch 0
I1017 00:22:08.109537 140737353971520 run_language_modeling.py:360]   Continuing training from epoch 0
I1017 00:22:08.109573 140737353971520 run_language_modeling.py:361]   Continuing training from global step 1
I1017 00:22:08.109574 140737353971520 run_language_modeling.py:361]   Continuing training from global step 1
I1017 00:22:08.109574 140737353971520 run_language_modeling.py:361]   Continuing training from global step 1
I1017 00:22:08.109575 140737353971520 run_language_modeling.py:361]   Continuing training from global step 1
I1017 00:22:08.109614 140737353971520 run_language_modeling.py:362]   Will skip the first 1 steps in the first epoch
I1017 00:22:08.109613 140737353971520 run_language_modeling.py:362]   Will skip the first 1 steps in the first epoch
I1017 00:22:08.109613 140737353971520 run_language_modeling.py:362]   Will skip the first 1 steps in the first epoch
I1017 00:22:08.109620 140737353971520 run_language_modeling.py:362]   Will skip the first 1 steps in the first epoch
^MEpoch:   0%|          | 0/5 [00:00<?, ?it/s]
^MIteration:   0%|          | 0/182922 [00:00<?, ?it/s]^[[ATraceback (most recent call last):
  File "/project/evn_py36/lib/python3.6/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/project/evn_py36/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/project/evn_py36/lib/python3.6/site-packages/torch/distributed/launch.py", line 261, in <module>
    main()
  File "/project/evn_py36/lib/python3.6/site-packages/torch/distributed/launch.py", line 257, in main
    cmd=cmd)
subprocess.CalledProcessError: Command '['/project/evn_py36/bin/python3', '-u', '/scratch/roberta_og/run_language_modeling.py', '--local_rank=3', '--gradient_accumulation_steps', '4', '--train_data_file', '/tmp/pbs.1286839.pbsha.ib.sockeye/data/', '--output_dir', '/scratch/roberta_og/ckpt_model_hg/', '--model_type', 'roberta', '--logging_dir', '/scratch/roberta_og/runs/', '--mlm', '--fp16', '--data_cache_dir', '/tmp/pbs.1286839.pbsha.ib.sockeye/cache/', '--num_workers', '0', '--warmup_steps', '10000', '--lazy_loading', '--model_name_or_path', '/scratch/roberta_og/roberta/', '--config_name', '/scratch/roberta_og/roberta', '--tokenizer_name', '/scratch/roberta_og/roberta', '--do_train', '--block_size', '60', '--learning_rate', '5e-5', '--num_train_epochs', '5', '--save_total_limit', '5', '--per_gpu_train_batch_size', '256', '--seed', '42']' died with <Signals.SIGKILL: 9>.

```

Could you please investigate this OOM issue? Do you have any suggestions?

I run this job with 6 nodes, each node has 4 32-GB GPUs, 24 CPUs, and 192GB CPU memory. I use distributed training. 

Versions:
Python version 3.6.8
PyTorch version 1.6.0
TensorFlow version 2.3.0
Transformers: 3.0.1
datasets version: 1.1.2