[BUG] Can't pickle local object 'instrument_w_nvtx.<locals>.wrapped_fn'

Describe the bug
A clear and concise description of what the bug is.
bug occurs while calling dataloader with multi num workers.
Here, ’trainer‘ is initialized from transformers. If I only debug with dataloader as follows, the code works.
‘’‘
for _, data in tqdm(enumerate(trainer.get_train_dataloader())):
print('dataloader: ', _, data.keys())
’‘’

However, if running trainer.train(), then the code throws the error as the title.

To Reproduce
Steps to reproduce the behavior:

1. Go to '...'
2. Click on '....'
3. Scroll down to '....'
4. See error

Expected behavior
A clear and concise description of what you expected to happen.

ds_report output
Please run ds_report to give us details about your setup.

Screenshots

System info (please complete the following information):

• OS: [e.g. Ubuntu 18.04]
• GPU count and types [e.g. two machines with x8 A100s each]
• Interconnects (if applicable) [e.g., two machines connected with 100 Gbps IB]
• Python version
• Any other relevant info about your setup

Launcher context
launching your experiment with the deepspeed launcher, MPI, or something else?

Docker context
Do not use docker.

Additional context
Add any other context about the problem here.

Start with the trainer.train() path and the instrument_w_nvtx.<locals>.wrapped_fn shown in the error, then reproduce it with the multi-worker dataloader described. Done means training no longer raises the pickle error; the issue names no file or test, so inspect the relevant DeepSpeed and PyTorch integration paths first.