TypeError: cannot pickle '_LazyModule' object
Author: lancekung
Pull request: no

@stas00 edit: please see https://github.com/huggingface/transformers/issues/12549#issuecomment-875287701 for the short reproduction script.

----------------

## Environment info

- `transformers` version: 4.9.0.dev0
- Platform: Linux with Nvidia P40
- Python version: 3.8.0
- PyTorch version (GPU?): 1.8.0
- Tensorflow version (GPU?):
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: yes

### Who can help

 @stas00 @patrickvonplaten, @LysandreJik

## Information

Model I am using (Bert, XLNet ...): GPT2

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [√] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [√] my own task or dataset: (give details below)

## To reproduce

I am running the minimal command:
```
python run_clm.py \
    --model_name_or_path /mycheckpoin/ \
    --train_file train.txt \
    --validation_file  eval.txt \
    --do_train \
    --do_eval \
    --output_dir ./models/ \
    --no_cuda False \
    --fp16 \
    --sharded_ddp simple \
    --num_train_epochs 3.0 \
    --disable_tqdm False \
    --save_steps 100 \
    --preprocessing_num_workers 32 \
    --per_device_train_batch_size 4 \
    --per_device_eval_batch_size 4
```

and I modified the following parts of the script ‘run_clm.py’, and the parameter rank passed in training_args.local_rank

```
def init_process(rank, size, fn, backend='gloo'):
    """ Initialize the distributed environment. """
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group(backend, rank=rank, world_size=size)
    fn(rank)

if __name__ == "__main__":
    # main()
    # size = int(os.environ['WORLD_SIZE'])
    size = int(torch.cuda.device_count())
    print(size)
    processes = []
    mp.set_start_method("spawn")
    for rank in range(size):
        p = mp.Process(target=init_process, args=(rank, main))
        p.start()
        processes.append(p)
    for p in processes:
        p.join()
```


the traceback informations are:

```
Process Process-2:
Traceback (most recent call last):
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/media/cfs/gonglixing/9Nctl/gpt_v2/run_clm_v3.py", line 511, in init_process
    fn(rank, size)
  File "/media/cfs/gonglixing/9Nctl/gpt_v2/run_clm_v3.py", line 367, in main
    tokenized_datasets = raw_datasets.map(
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/datasets/dataset_dict.py", line 471, in map
    {
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/datasets/dataset_dict.py", line 472, in <dictcomp>
    k: dataset.map(
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 1736, in map
    transformed_shards = [r.get() for r in results]
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 1736, in <listcomp>
    transformed_shards = [r.get() for r in results]
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/multiprocess/pool.py", line 771, in get
    raise self._value
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/multiprocess/pool.py", line 537, in _handle_tasks
    put(task)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/multiprocess/connection.py", line 209, in send
    self._send_bytes(_ForkingPickler.dumps(obj))
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/multiprocess/reduction.py", line 54, in dumps
    cls(buf, protocol, *args, **kwds).dump(obj)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/dill/_dill.py", line 498, in dump
    StockPickler.dump(self, obj)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 487, in dump
    self.save(obj)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 901, in save_tuple
    save(element)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/dill/_dill.py", line 990, in save_module_dict
    StockPickler.save_dict(pickler, obj)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 971, in save_dict
    self._batch_setitems(obj.items())
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 997, in _batch_setitems
    save(v)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/dill/_dill.py", line 1493, in save_function
    pickler.save_reduce(_create_function, (obj.__code__,
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 692, in save_reduce
    save(args)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 901, in save_tuple
    save(element)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/dill/_dill.py", line 990, in save_module_dict
    StockPickler.save_dict(pickler, obj)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 971, in save_dict
    self._batch_setitems(obj.items())
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 997, in _batch_setitems
    save(v)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 578, in save
    rv = reduce(self.proto)
TypeError: cannot pickle '_LazyModule' object
```

I run the following command based on the original script, it works well. The reason why I don't use this command is that our cluster doesn't support this way of passing parameters: "-m torch.distributed.launch --nproc_per_node=4 "



```
python -m torch.distributed.launch --nproc_per_node=4 run_clm.py \
    --model_name_or_path /mycheckpoin/ \
    --train_file train.txt \
    --validation_file  eval.txt \
    --do_train \
    --do_eval \
    --output_dir ./models/ \
    --no_cuda False \
    --fp16 \
    --sharded_ddp simple \
    --num_train_epochs 3.0 \
    --disable_tqdm False \
    --save_steps 100 \
    --preprocessing_num_workers 32 \
    --per_device_train_batch_size 4 \
    --per_device_eval_batch_size 4

```


## Expected behavior


Comments:
Could you please attach the final script you used or a branch that we can use to reproduce your code exactly? Thanks.

note: I took the liberty to edit your OP to use code formatting which is much easier to read. If possible use a similar approach in future reports. Thank you!


> Could you please attach the final script you used or a branch that we can use to reproduce your code exactly? Thanks.
> 
> note: I took the liberty to edit your OP to use code formatting which is much easier to read. If possible use a similar approach in future reports. Thank you!

this is my scripts, thanks very much!
[run_clm.py.zip](https://github.com/huggingface/transformers/files/6774180/run_clm.py.zip)


Thank you. The attached script fails for me. You also didn't supply the data, but I assume it doesn't matter. In the future please supply everything or adapt your runtime so that we could run it out of the box and not need to spend a lot of time to try to make things work.

```
python run_clm.py \
>     --model_name_or_path gpt2 \
>     --dataset_name wikitext \
>     --dataset_config_name wikitext-2-raw-v1 \
>     --do_train \
>     --do_eval \
>     --output_dir /tmp/test-clm
2021-07-06 21:18:15.064178: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0
2
2021-07-06 21:18:17.425481: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0
2021-07-06 21:18:17.425484: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library libcudart.so.11.0
Process Process-1:
Traceback (most recent call last):
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
TypeError: init_process() missing 1 required positional argument: 'fn'
Process Process-2:
Traceback (most recent call last):
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
TypeError: init_process() missing 1 required positional argument: 'fn'
```

same failure with distributed.

> Thank you. The attached script fails for me. You also didn't supply the data, but I assume it doesn't matter. In the future please supply everything or adapt your runtime so that we could run it out of the box and not need to spend a lot of time to try to make things work.

So sorry, it's my fault, I gave you the wrong version.
This is the right version.
[run_clm.py.zip](https://github.com/huggingface/transformers/files/6774314/run_clm.py.zip)


I'm able to reproduce the problem - great! 

Let's see what the culprit is.

So the trigger is: `--preprocessing_num_workers 32`

and the minimal reproduction cmd is:
```
python run_clm.py --model_name_or_path sshleifer/tiny-gpt2 --dataset_name wikitext \
--dataset_config_name wikitext-2-raw-v1 --do_train --do_eval --output_dir /tmp/test-clm \
--overwrite_output_dir --preprocessing_num_workers 32
```

It happens only with your version of the script. I tested with the one in `master` it works fine there.

The problem is unrelated to the change in https://github.com/huggingface/transformers/pull/11168 as you have discovered yourself, since your code removed my changes and you're just passing:
```
    def tokenize_function1(examples):
        return tokenizer(examples[text_column_name])
```

So need to look elsewhere for the cause.

From a quick look I suspect that perhaps this is an issue in `datasets` when `num_proc > 1`? Could you try to reduce the script to the bare minimum, so that it runs just:

```
    with training_args.main_process_first(desc="dataset map tokenization"):
        tokenized_datasets = raw_datasets.map(
            None,
            num_proc=5,
        )
```

inside the multi-proc modifications you made.

e.g. the above is enough to trigger the same error in the script so removing most of the code should 


OK, here is the minimal reproducible script. Totally unrelated to `transformers` it seems except for the import of `transformers`

```
import logging
import math
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

import datasets
from datasets import load_dataset

import transformers

import torch
import torch.distributed as dist
import torch.multiprocessing as mp

def main(rank, size):

    def tokenize_function(examples):
        return None

    raw_datasets = load_dataset("wikitext", "wikitext-2-raw-v1")
    tokenized_datasets = raw_datasets.map(
        tokenize_function,
        num_proc=32,
    )

def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()

def init_process(rank, size, fn, backend='gloo'):
    """ Initialize the distributed environment. """
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group(backend, rank=rank, world_size=size)
    fn(rank, size)


if __name__ == "__main__":
    # main()
    # size = int(os.environ['WORLD_SIZE'])
    size = int(torch.cuda.device_count())
    print(size)
    processes = []
    mp.set_start_method("spawn")
    for rank in range(size):
        p = mp.Process(target=init_process, args=(rank, size, main))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

```

this still fails with the same error.

```
python run_clm.py
2
Reusing dataset wikitext (/home/stas/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20)
Reusing dataset wikitext (/home/stas/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20)
Process Process-1:
Process Process-2:
Traceback (most recent call last):
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/mnt/nvme1/code/huggingface/users/lancekung/run_clm.py", line 60, in init_process
    fn(rank, size)
  File "/mnt/nvme1/code/huggingface/users/lancekung/run_clm.py", line 46, in main
    tokenized_datasets = raw_datasets.map(
  File "/mnt/nvme1/code/huggingface/datasets-master/src/datasets/dataset_dict.py", line 471, in map
    {
  File "/mnt/nvme1/code/huggingface/datasets-master/src/datasets/dataset_dict.py", line 472, in <dictcomp>
    k: dataset.map(
  File "/mnt/nvme1/code/huggingface/datasets-master/src/datasets/arrow_dataset.py", line 1736, in map
    transformed_shards = [r.get() for r in results]
  File "/mnt/nvme1/code/huggingface/datasets-master/src/datasets/arrow_dataset.py", line 1736, in <listcomp>
    transformed_shards = [r.get() for r in results]
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/multiprocess/pool.py", line 771, in get
    raise self._value
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/multiprocess/pool.py", line 537, in _handle_tasks
    put(task)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/multiprocess/connection.py", line 209, in send
    self._send_bytes(_ForkingPickler.dumps(obj))
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/multiprocess/reduction.py", line 54, in dumps
    cls(buf, protocol, *args, **kwds).dump(obj)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/dill/_dill.py", line 498, in dump
    StockPickler.dump(self, obj)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 487, in dump
    self.save(obj)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 901, in save_tuple
    save(element)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/dill/_dill.py", line 990, in save_module_dict
    StockPickler.save_dict(pickler, obj)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 971, in save_dict
    self._batch_setitems(obj.items())
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 997, in _batch_setitems
    save(v)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/dill/_dill.py", line 1493, in save_function
    pickler.save_reduce(_create_function, (obj.__code__,
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 692, in save_reduce
    save(args)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 901, in save_tuple
    save(element)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/dill/_dill.py", line 990, in save_module_dict
    StockPickler.save_dict(pickler, obj)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 971, in save_dict
    self._batch_setitems(obj.items())
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 997, in _batch_setitems
    save(v)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 578, in save
    rv = reduce(self.proto)
TypeError: cannot pickle '_LazyModule' object
Traceback (most recent call last):
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/mnt/nvme1/code/huggingface/users/lancekung/run_clm.py", line 60, in init_process
    fn(rank, size)
  File "/mnt/nvme1/code/huggingface/users/lancekung/run_clm.py", line 46, in main
    tokenized_datasets = raw_datasets.map(
  File "/mnt/nvme1/code/huggingface/datasets-master/src/datasets/dataset_dict.py", line 471, in map
    {
  File "/mnt/nvme1/code/huggingface/datasets-master/src/datasets/dataset_dict.py", line 472, in <dictcomp>
    k: dataset.map(
  File "/mnt/nvme1/code/huggingface/datasets-master/src/datasets/arrow_dataset.py", line 1736, in map
    transformed_shards = [r.get() for r in results]
  File "/mnt/nvme1/code/huggingface/datasets-master/src/datasets/arrow_dataset.py", line 1736, in <listcomp>
    transformed_shards = [r.get() for r in results]
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/multiprocess/pool.py", line 771, in get
    raise self._value
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/multiprocess/pool.py", line 537, in _handle_tasks
    put(task)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/multiprocess/connection.py", line 209, in send
    self._send_bytes(_ForkingPickler.dumps(obj))
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/multiprocess/reduction.py", line 54, in dumps
    cls(buf, protocol, *args, **kwds).dump(obj)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/dill/_dill.py", line 498, in dump
    StockPickler.dump(self, obj)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 487, in dump
    self.save(obj)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 901, in save_tuple
    save(element)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/dill/_dill.py", line 990, in save_module_dict
    StockPickler.save_dict(pickler, obj)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 971, in save_dict
    self._batch_setitems(obj.items())
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 997, in _batch_setitems
    save(v)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/dill/_dill.py", line 1493, in save_function
    pickler.save_reduce(_create_function, (obj.__code__,
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 692, in save_reduce
    save(args)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 901, in save_tuple
    save(element)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/dill/_dill.py", line 990, in save_module_dict
    StockPickler.save_dict(pickler, obj)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 971, in save_dict
    self._batch_setitems(obj.items())
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 997, in _batch_setitems
    save(v)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 578, in save
    rv = reduce(self.proto)
TypeError: cannot pickle '_LazyModule' object
```

But if you either:
* comment out `import transformers` 
* or set `num_proc=1` in `datasets.map` (instead of `n>1`)
all is good.

@lhoestq, @albertvillanova - does this ring any bells? Clearly `transformers` loads some module lazily and trips up `datasets` even though transformers isn't really used here directly. Thank you.



> OK, here is the minimal reproducible script. Totally unrelated to `transformers` it seems except for the import of `transformers`
> 
> ```
> import logging
> import math
> import os
> import sys
> from dataclasses import dataclass, field
> from typing import Optional
> 
> import datasets
> from datasets import load_dataset
> 
> import transformers
> 
> import torch
> import torch.distributed as dist
> import torch.multiprocessing as mp
> 
> def main(rank, size):
> 
>     def tokenize_function(examples):
>         return None
> 
>     raw_datasets = load_dataset("wikitext", "wikitext-2-raw-v1")
>     tokenized_datasets = raw_datasets.map(
>         tokenize_function,
>         num_proc=32,
>     )
> 
> def _mp_fn(index):
>     # For xla_spawn (TPUs)
>     main()
> 
> def init_process(rank, size, fn, backend='gloo'):
>     """ Initialize the distributed environment. """
>     os.environ['MASTER_ADDR'] = '127.0.0.1'
>     os.environ['MASTER_PORT'] = '29500'
>     dist.init_process_group(backend, rank=rank, world_size=size)
>     fn(rank, size)
> 
> 
> if __name__ == "__main__":
>     # main()
>     # size = int(os.environ['WORLD_SIZE'])
>     size = int(torch.cuda.device_count())
>     print(size)
>     processes = []
>     mp.set_start_method("spawn")
>     for rank in range(size):
>         p = mp.Process(target=init_process, args=(rank, size, main))
>         p.start()
>         processes.append(p)
> 
>     for p in processes:
>         p.join()
> ```
> 
> this still fails with the same error.
> 
> ```
> python run_clm.py
> 2
> Reusing dataset wikitext (/home/stas/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20)
> Reusing dataset wikitext (/home/stas/.cache/huggingface/datasets/wikitext/wikitext-2-raw-v1/1.0.0/aa5e094000ec7afeb74c3be92c88313cd6f132d564c7effd961c10fd47c76f20)
> Process Process-1:
> Process Process-2:
> Traceback (most recent call last):
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/multiprocessing/process.py", line 315, in _bootstrap
>     self.run()
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/multiprocessing/process.py", line 108, in run
>     self._target(*self._args, **self._kwargs)
>   File "/mnt/nvme1/code/huggingface/users/lancekung/run_clm.py", line 60, in init_process
>     fn(rank, size)
>   File "/mnt/nvme1/code/huggingface/users/lancekung/run_clm.py", line 46, in main
>     tokenized_datasets = raw_datasets.map(
>   File "/mnt/nvme1/code/huggingface/datasets-master/src/datasets/dataset_dict.py", line 471, in map
>     {
>   File "/mnt/nvme1/code/huggingface/datasets-master/src/datasets/dataset_dict.py", line 472, in <dictcomp>
>     k: dataset.map(
>   File "/mnt/nvme1/code/huggingface/datasets-master/src/datasets/arrow_dataset.py", line 1736, in map
>     transformed_shards = [r.get() for r in results]
>   File "/mnt/nvme1/code/huggingface/datasets-master/src/datasets/arrow_dataset.py", line 1736, in <listcomp>
>     transformed_shards = [r.get() for r in results]
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/multiprocess/pool.py", line 771, in get
>     raise self._value
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/multiprocess/pool.py", line 537, in _handle_tasks
>     put(task)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/multiprocess/connection.py", line 209, in send
>     self._send_bytes(_ForkingPickler.dumps(obj))
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/multiprocess/reduction.py", line 54, in dumps
>     cls(buf, protocol, *args, **kwds).dump(obj)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/dill/_dill.py", line 498, in dump
>     StockPickler.dump(self, obj)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 487, in dump
>     self.save(obj)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
>     f(self, obj)  # Call unbound method with explicit self
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 901, in save_tuple
>     save(element)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
>     f(self, obj)  # Call unbound method with explicit self
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/dill/_dill.py", line 990, in save_module_dict
>     StockPickler.save_dict(pickler, obj)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 971, in save_dict
>     self._batch_setitems(obj.items())
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 997, in _batch_setitems
>     save(v)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
>     f(self, obj)  # Call unbound method with explicit self
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/dill/_dill.py", line 1493, in save_function
>     pickler.save_reduce(_create_function, (obj.__code__,
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 692, in save_reduce
>     save(args)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
>     f(self, obj)  # Call unbound method with explicit self
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 901, in save_tuple
>     save(element)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
>     f(self, obj)  # Call unbound method with explicit self
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/dill/_dill.py", line 990, in save_module_dict
>     StockPickler.save_dict(pickler, obj)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 971, in save_dict
>     self._batch_setitems(obj.items())
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 997, in _batch_setitems
>     save(v)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 578, in save
>     rv = reduce(self.proto)
> TypeError: cannot pickle '_LazyModule' object
> Traceback (most recent call last):
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/multiprocessing/process.py", line 315, in _bootstrap
>     self.run()
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/multiprocessing/process.py", line 108, in run
>     self._target(*self._args, **self._kwargs)
>   File "/mnt/nvme1/code/huggingface/users/lancekung/run_clm.py", line 60, in init_process
>     fn(rank, size)
>   File "/mnt/nvme1/code/huggingface/users/lancekung/run_clm.py", line 46, in main
>     tokenized_datasets = raw_datasets.map(
>   File "/mnt/nvme1/code/huggingface/datasets-master/src/datasets/dataset_dict.py", line 471, in map
>     {
>   File "/mnt/nvme1/code/huggingface/datasets-master/src/datasets/dataset_dict.py", line 472, in <dictcomp>
>     k: dataset.map(
>   File "/mnt/nvme1/code/huggingface/datasets-master/src/datasets/arrow_dataset.py", line 1736, in map
>     transformed_shards = [r.get() for r in results]
>   File "/mnt/nvme1/code/huggingface/datasets-master/src/datasets/arrow_dataset.py", line 1736, in <listcomp>
>     transformed_shards = [r.get() for r in results]
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/multiprocess/pool.py", line 771, in get
>     raise self._value
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/multiprocess/pool.py", line 537, in _handle_tasks
>     put(task)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/multiprocess/connection.py", line 209, in send
>     self._send_bytes(_ForkingPickler.dumps(obj))
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/multiprocess/reduction.py", line 54, in dumps
>     cls(buf, protocol, *args, **kwds).dump(obj)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/dill/_dill.py", line 498, in dump
>     StockPickler.dump(self, obj)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 487, in dump
>     self.save(obj)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
>     f(self, obj)  # Call unbound method with explicit self
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 901, in save_tuple
>     save(element)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
>     f(self, obj)  # Call unbound method with explicit self
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/dill/_dill.py", line 990, in save_module_dict
>     StockPickler.save_dict(pickler, obj)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 971, in save_dict
>     self._batch_setitems(obj.items())
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 997, in _batch_setitems
>     save(v)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
>     f(self, obj)  # Call unbound method with explicit self
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/dill/_dill.py", line 1493, in save_function
>     pickler.save_reduce(_create_function, (obj.__code__,
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 692, in save_reduce
>     save(args)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
>     f(self, obj)  # Call unbound method with explicit self
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 901, in save_tuple
>     save(element)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 560, in save
>     f(self, obj)  # Call unbound method with explicit self
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/dill/_dill.py", line 990, in save_module_dict
>     StockPickler.save_dict(pickler, obj)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 971, in save_dict
>     self._batch_setitems(obj.items())
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 997, in _batch_setitems
>     save(v)
>   File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/pickle.py", line 578, in save
>     rv = reduce(self.proto)
> TypeError: cannot pickle '_LazyModule' object
> ```
> 
> But if you either:
> 
> * comment out `import transformers`
> * or set `num_proc=1` in `datasets.map` (instead of `n>1`)
>   all is good.
> 
> @lhoestq, @albertvillanova - does this ring any bells? Clearly `transformers` loads some module lazily and trips up `datasets` even though transformers isn't really used here directly. Thank you.

Thank you so much for your time, and hope other experts can give some tips about this problem.

Hi @stas00, thanks for pinging. 

I'm having a look and after a first search, I think you are right and the problem comes from the fact that `transformers` makes a lazy import when importing it. I guess this affects `datasets` here: https://github.com/huggingface/datasets/blob/master/src/datasets/utils/py_utils.py#L319 (PR: https://github.com/huggingface/datasets/pull/502), which is used by dumps to pickle objects in a multiprocessing setup.

cc: @lhoestq 

> Hi @stas00, thanks for pinging.
> 
> I'm having a look and after a first search, I think you are right and the problem comes from the fact that `transformers` makes a lazy import when importing it. I guess this affects `datasets` here: https://github.com/huggingface/datasets/blob/master/src/datasets/utils/py_utils.py#L319 (PR: [huggingface/datasets#502](https://github.com/huggingface/datasets/pull/502)), which is used by dumps to pickle objects in a multiprocessing setup.
> 
> cc: @lhoestq

hi albertvillanova, I removed import of transformers according to the following code, it still can't work.


`def _no_cache_fields(obj):
    try:
        if (
            "PreTrainedTokenizerBase" in [base_class.__name__ for base_class in type(obj).__mro__]
            and hasattr(obj, "cache")
            and isinstance(obj.cache, dict)
        )`


Note that we can easily make `_LazyModule` picklable. I can open a PR if needed to implement a `__reduce__` method for `_LazyModule`. It's the only object that prevents `transformers` from being picklable.

EDIT: here it is: https://github.com/huggingface/transformers/pull/12552

This is just a way to easily fix this issue, but I think we should definitely keep trying to figure out why it tried to pickle `transformers` in the first place. This might come from `dill` that pickles the globals of some environments when pickling any object

Linking to the new PR: https://github.com/huggingface/transformers/pull/12567


Should be closed by #12567, please let us know if the problem persists.

> Should be closed by #12567, please let us know if the problem persists.

Hi, a new problem has arisen
we can pickle "LazyModule" now, but can't  pickle <class 'types.AutoModelForCausalLM'>





Traceback (most recent call last):
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/media/cfs/gonglixing/9Nctl/gpt_v2/run_clm_v3.py", line 509, in init_process
    fn(rank, size)
  File "/media/cfs/gonglixing/9Nctl/gpt_v2/run_clm_v3.py", line 367, in main
    tokenized_datasets = raw_datasets.map(
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/datasets/dataset_dict.py", line 471, in map
    {
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/datasets/dataset_dict.py", line 472, in <dictcomp>
    k: dataset.map(
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 1736, in map
    transformed_shards = [r.get() for r in results]
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 1736, in <listcomp>
    transformed_shards = [r.get() for r in results]
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/multiprocess/pool.py", line 771, in get
    raise self._value
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/multiprocess/pool.py", line 537, in _handle_tasks
    put(task)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/multiprocess/connection.py", line 209, in send
    self._send_bytes(_ForkingPickler.dumps(obj))
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/multiprocess/reduction.py", line 54, in dumps
    cls(buf, protocol, *args, **kwds).dump(obj)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/dill/_dill.py",line 498, in dump
    StockPickler.dump(self, obj)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 487, in dump
    self.save(obj)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 901, in save_tuple
    save(element)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/dill/_dill.py",line 990, in save_module_dict
    StockPickler.save_dict(pickler, obj)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 971, in save_dict
    self._batch_setitems(obj.items())
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 997, in _batch_setitems
    save(v)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/dill/_dill.py",line 1493, in save_function
    pickler.save_reduce(_create_function, (obj.__code__,
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 692, in save_reduce
    save(args)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 901, in save_tuple
    save(element)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/dill/_dill.py",line 990, in save_module_dict
    StockPickler.save_dict(pickler, obj)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 971, in save_dict
    self._batch_setitems(obj.items())
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 997, in _batch_setitems
    save(v)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 560, in save
    f(self, obj)  # Call unbound method with explicit self
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/site-packages/dill/_dill.py",line 1439, in save_type
    StockPickler.save_global(pickler, obj, name=name)
  File "/usr/local/anaconda3/envs/py38/lib/python3.8/pickle.py", line 1070, in save_global
    raise PicklingError(
_pickle.PicklingError: Can't pickle <class 'types.AutoModelForCausalLM'>: it's notfound as types.AutoModelForCausalLM