[BUG] AttributeError: Can't pickle local object 'main.<locals>.collate_fn' when running train_control_lora.sh

Hello! thanks for sharing your work!
when I run train_control_lora.sh(deepspeed version), I got error "AttributeError: Can't pickle local object 'main..collate_fn'"

related part of the logs is as :

  File "/opt/conda/envs/wanc/lib/python3.10/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/opt/conda/envs/wanc/lib/python3.10/multiprocessing/context.py", line 224, in _Popen
        main()return _default_context.get_context().Process._Popen(process_obj)

  File "/home1/jovyan/workspace/VideoX-Fun/scripts/wan2.1_fun/train_control_lora.py", line 1500, in main
  File "/opt/conda/envs/wanc/lib/python3.10/multiprocessing/context.py", line 288, in _Popen
    return Popen(process_obj)
  File "/opt/conda/envs/wanc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/opt/conda/envs/wanc/lib/python3.10/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/opt/conda/envs/wanc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/opt/conda/envs/wanc/lib/python3.10/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
AttributeError: Can't pickle local object 'main.<locals>.collate_fn'
    for step, batch in enumerate(train_dataloader):
  File "/opt/conda/envs/wanc/lib/python3.10/site-packages/accelerate/data_loader.py", line 564, in __iter__
    dataloader_iter = self.base_dataloader.__iter__()
  File "/opt/conda/envs/wanc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 434, in __iter__
    self._iterator = self._get_iterator()
  File "/opt/conda/envs/wanc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 387, in _get_iterator
    return _MultiProcessingDataLoaderIter(self)
  File "/opt/conda/envs/wanc/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1040, in __init__
    w.start()
  File "/opt/conda/envs/wanc/lib/python3.10/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/opt/conda/envs/wanc/lib/python3.10/multiprocessing/context.py", line 224, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
  File "/opt/conda/envs/wanc/lib/python3.10/multiprocessing/context.py", line 288, in _Popen
    return Popen(process_obj)
  File "/opt/conda/envs/wanc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/opt/conda/envs/wanc/lib/python3.10/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/opt/conda/envs/wanc/lib/python3.10/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/opt/conda/envs/wanc/lib/python3.10/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
AttributeError: Can't pickle local object 'main.<locals>.collate_fn'
Steps:   0%|                                                                                                                                                                                  | 0/102700 [00:05<?, ?it/s]

I tried to move collate_fn out of main funciton but many dependent problems occured..
My torch is 2.2.0+cu121/python3.10, 8xH100 GPU

I could run training without deepspeed but I had to reduce image_size to 512 to avoid OOM....
please help me!

Start with scripts/wan2.1_fun/train_control_lora.py around line 1500 and the train_control_lora.sh entry point. Inspect how the local collate_fn is passed to the DataLoader when the DeepSpeed path enables multiprocessing, then reproduce with the reported Python and PyTorch setup. Done means the DeepSpeed script starts training without the pickling error while retaining the intended data loading behavior.