nccl error

Thanks for your work! I try to train bagel following the script:
```
export PYTHONPATH=$(pwd)

export NCCL_SHM_DISABLE=1

export NCCL_P2P_DISABLE=1

# Fine-tuning
torchrun \
--nnodes=1 \
--node_rank=0 \
--nproc_per_node=2 \
--master_addr=127.0.0.1 \
--master_port=29501 \
train/pretrain_unified_navit.py \
--dataset_config_file ./data/configs/example.yaml \
--model_path /inspire/hdd/ws-f4d69b29-e0a5-44e6-bd92-acf4de9990f0/public-project/zhengkaipeng-240108120123/weilai/models/BAGEL-7B-MoT \
--layer_module Qwen2MoTDecoderLayer \
--use_flex True \
--max_latent_size 64 \
--resume-from /inspire/hdd/ws-f4d69b29-e0a5-44e6-bd92-acf4de9990f0/public-project/zhengkaipeng-240108120123/weilai/models/BAGEL-7B-MoT \
--finetune_from_hf True \
--num_replicate=1 \
--num_shard=2 \
--auto_resume True \
--resume-model-only True \
--finetune-from-ema True \
--log_every 1 \
--lr 2e-5 \
--num_worker 1 \
--expected_num_tokens 10240 \
--max_num_tokens 11520 \
--max_num_tokens_per_sample 10240 \
--wandb_offline true
```

I use torch== 2.5.1, flash-attn==2.7.2.post1

But I met this error:
```
[rank0]: Traceback (most recent call last):
[rank0]: File "/inspire/hdd/ws-f4d69b29-e0a5-44e6-bd92-acf4de9990f0/public-project/zhengkaipeng-240108120123/weilai/codes/BAGEL/train/pretrain_unified_navit.py", line 701, in
[rank0]: main()
[rank0]: File "/inspire/hdd/ws-f4d69b29-e0a5-44e6-bd92-acf4de9990f0/public-project/zhengkaipeng-240108120123/weilai/codes/BAGEL/train/pretrain_unified_navit.py", line 624, in main
[rank0]: loss.backward()
[rank0]: File "/opt/conda/lib/python3.12/site-packages/torch/_tensor.py", line 581, in backward
[rank0]: torch.autograd.backward(
[rank0]: File "/opt/conda/lib/python3.12/site-packages/torch/autograd/__init__.py", line 347, in backward
[rank0]: _engine_run_backward(
[rank0]: File "/opt/conda/lib/python3.12/site-packages/torch/autograd/graph.py", line 825, in _engine_run_backward
[rank0]: return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/opt/conda/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]: return func(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/opt/conda/lib/python3.12/site-packages/torch/distributed/fsdp/_runtime_utils.py", line 755, in _post_backward_hook
[rank0]: _reduce_grad(state, handle)
[rank0]: File "/opt/conda/lib/python3.12/site-packages/torch/distributed/fsdp/_runtime_utils.py", line 861, in _reduce_grad
[rank0]: dist.all_reduce(new_sharded_grad, group=state._inter_node_pg)
[rank0]: File "/opt/conda/lib/python3.12/site-packages/torch/distributed/c10d_logger.py", line 83, in wrapper
[rank0]: return func(*args, **kwargs)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^
[rank0]: File "/opt/conda/lib/python3.12/site-packages/torch/distributed/distributed_c10d.py", line 2501, in all_reduce
[rank0]: work = group.allreduce([tensor], opts)
[rank0]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: torch.distributed.DistBackendError: NCCL error in: ../torch/csrc/distributed/c10d/NCCLUtils.hpp:317, unhandled system error (run with NCCL_DEBUG=INFO for details), NCCL version 2.21.5
[rank0]: ncclSystemError: System call (e.g. socket, malloc) or external library call failed or device error.
[rank0]: Last error:
[rank0]: Error while creating shared memory segment /dev/shm/nccl-OEIhdX (size 7340384)
```

How can I solve this? Thanks for your help!

The failing entry point is train/pretrain_unified_navit.py at loss.backward() on line 624; start by reproducing the provided torchrun command with NCCL_DEBUG=INFO. Check the reported /dev/shm shared-memory failure and the torch and flash-attn versions listed in the issue. Done means the same two-process training command completes backward without the NCCL error.

Written by the indexing model from the issue text.