How to run SFT on 2 L40 GPUs

BASE_MODEL=Qwen/Qwen2.5-3B-Instruct
DATA_DIR=datasets/blobs/50_shot/no_reasoning
PROJECT_NAME=SFT
EXPERIMENT_NAME=blobs-50shot-qwen2.5-3b-instruct-sft-no-reasoning
VLLM_ATTENTION_BACKEND=XFORMERS

CUDA_VISIBLE_DEVICE=0,1 torchrun --nproc_per_node=2 verl/trainer/fsdp_sft_trainer.py \
data.train_files=$DATA_DIR/train.parquet \
data.val_files=$DATA_DIR/test.parquet \
data.train_batch_size=64 \
data.micro_batch_size=8 \
data.max_length=2048 \
data.prompt_key=prompt \
+data.prompt_dict_keys=["content"] \
data.response_key=label \
model.partial_pretrain=$BASE_MODEL \
trainer.logger=['wandb'] \
+trainer.val_before_train=False \
trainer.default_hdfs_dir=null \
trainer.project_name=$PROJECT_NAME \
trainer.experiment_name=$EXPERIMENT_NAME \
trainer.total_epochs=15 2>&1 | tee verl_demo.log \
trainer.default_local_dir=checkpoints/$PROJECT_NAME/$EXPERIMENT_NAME

/var/lib/condor/execute/slot1/dir_3556967/liftr/verl/utils/dataset/sft_dataset.py:78: FutureWarning: Series.getitem treating keys as positions is deprecated. In a future version, integer 
keys will always be treated as labels (consistent with DataFrame behavior). To access a value by position, use ser.iloc[pos]
  ls = ls[0]
Flash Attention 2.0 only supports torch.float16 and torch.bfloat16 dtypes, but the current dype in Qwen2ForCausalLM is torch.float32. You should run training or inference using Automatic Mixe
d-Precision via the with torch.autocast(device_type='torch_device'): decorator, or load the model with the torch_dtype argument. Example: model = AutoModel.from_pretrained("openai/whispe r-tiny", attn_implementation="flash_attention_2", torch_dtype=torch.float16)
You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with model.to('cuda').
Loading checkpoint shards: 100%|██████████| 2/2 [00:02<00:00,  1.20s/it]
functools.partial(<function transformer_auto_wrap_policy at 0x14c0085007c0>, transformer_layer_cls={<class 'transformers.models.qwen2.modeling_qwen2.Qwen2DecoderLayer'>})
Loading checkpoint shards: 100%|██████████| 2/2 [00:02<00:00,  1.15s/it]
dsigpu4002:3691059:3691059 [0] NCCL INFO Bootstrap : Using bond0:128.105.68.87<0>
dsigpu4002:3691059:3691059 [0] NCCL INFO NET/Plugin : dlerror=libnccl-net.so: cannot open shared object file: No such file or directory No plugin found (libnccl-net.so), using internal implem
entation
dsigpu4002:3691059:3691059 [0] NCCL INFO cudaDriverVersion 12040
NCCL version 2.20.5+cuda12.4
dsigpu4002:3691060:3691060 [1] NCCL INFO cudaDriverVersion 12040
dsigpu4002:3691060:3691060 [1] NCCL INFO Bootstrap : Using bond0:128.105.68.87<0>
dsigpu4002:3691060:3691060 [1] NCCL INFO NET/Plugin : dlerror=libnccl-net.so: cannot open shared object file: No such file or directory No plugin found (libnccl-net.so), using internal implem
entation
dsigpu4002:3691059:3691145 [0] NCCL INFO NET/IB : Using [0]mlx5_bond_0:1/RoCE [RO]; OOB bond0:128.105.68.87<0>
dsigpu4002:3691059:3691145 [0] NCCL INFO Using non-device net plugin version 0
dsigpu4002:3691059:3691145 [0] NCCL INFO Using network IB
dsigpu4002:3691060:3691146 [1] NCCL INFO NET/IB : Using [0]mlx5_bond_0:1/RoCE [RO]; OOB bond0:128.105.68.87<0>
dsigpu4002:3691060:3691146 [1] NCCL INFO Using non-device net plugin version 0
dsigpu4002:3691060:3691146 [1] NCCL INFO Using network IB
dsigpu4002:3691060:3691146 [1] NCCL INFO comm 0xf5f3b80 rank 1 nranks 2 cudaDev 1 nvmlDev 1 busId d1000 commId 0x472e0f07d7c3c135 - Init START
dsigpu4002:3691059:3691145 [0] NCCL INFO comm 0x10b6e8f0 rank 0 nranks 2 cudaDev 0 nvmlDev 0 busId ce000 commId 0x472e0f07d7c3c135 - Init START
dsigpu4002:3691060:3691146 [1] NCCL INFO Setting affinity for GPU 1 to ffffff00,0000ffff,ff000000
dsigpu4002:3691059:3691145 [0] NCCL INFO Setting affinity for GPU 0 to ffffff00,0000ffff,ff000000
dsigpu4002:3691059:3691145 [0] NCCL INFO comm 0x10b6e8f0 rank 0 nRanks 2 nNodes 1 localRanks 2 localRank 0 MNNVL 0
dsigpu4002:3691060:3691146 [1] NCCL INFO comm 0xf5f3b80 rank 1 nRanks 2 nNodes 1 localRanks 2 localRank 1 MNNVL 0
dsigpu4002:3691060:3691146 [1] NCCL INFO Trees [0] -1/-1/-1->1->0 [1] 0/-1/-1->1->-1 [2] -1/-1/-1->1->0 [3] 0/-1/-1->1->-1
dsigpu4002:3691060:3691146 [1] NCCL INFO P2P Chunksize set to 131072
dsigpu4002:3691059:3691145 [0] NCCL INFO Channel 00/04 :    0   1
dsigpu4002:3691059:3691145 [0] NCCL INFO Channel 01/04 :    0   1
dsigpu4002:3691059:3691145 [0] NCCL INFO Channel 02/04 :    0   1
dsigpu4002:3691059:3691145 [0] NCCL INFO Channel 03/04 :    0   1
dsigpu4002:3691059:3691145 [0] NCCL INFO Trees [0] 1/-1/-1->0->-1 [1] -1/-1/-1->0->1 [2] 1/-1/-1->0->-1 [3] -1/-1/-1->0->1
dsigpu4002:3691059:3691145 [0] NCCL INFO P2P Chunksize set to 131072
dsigpu4002:3691059:3691145 [0] NCCL INFO Channel 00/0 : 0[0] -> 1[1] via P2P/CUMEM
dsigpu4002:3691060:3691146 [1] NCCL INFO Channel 00/0 : 1[1] -> 0[0] via P2P/CUMEM
dsigpu4002:3691059:3691145 [0] NCCL INFO Channel 01/0 : 0[0] -> 1[1] via P2P/CUMEM
dsigpu4002:3691060:3691146 [1] NCCL INFO Channel 01/0 : 1[1] -> 0[0] via P2P/CUMEM
dsigpu4002:3691059:3691145 [0] NCCL INFO Channel 02/0 : 0[0] -> 1[1] via P2P/CUMEM
dsigpu4002:3691060:3691146 [1] NCCL INFO Channel 02/0 : 1[1] -> 0[0] via P2P/CUMEM
dsigpu4002:3691059:3691145 [0] NCCL INFO Channel 03/0 : 0[0] -> 1[1] via P2P/CUMEM
dsigpu4002:3691060:3691146 [1] NCCL INFO Channel 03/0 : 1[1] -> 0[0] via P2P/CUMEM
dsigpu4002:3691060:3691146 [1] NCCL INFO Connected all rings
dsigpu4002:3691060:3691146 [1] NCCL INFO Connected all trees
dsigpu4002:3691059:3691145 [0] NCCL INFO Connected all rings
dsigpu4002:3691059:3691145 [0] NCCL INFO Connected all trees
dsigpu4002:3691060:3691146 [1] NCCL INFO threadThresholds 8/8/64 | 16/8/64 | 512 | 512
dsigpu4002:3691060:3691146 [1] NCCL INFO 4 coll channels, 0 collnet channels, 0 nvls channels, 4 p2p channels, 2 p2p channels per peer
dsigpu4002:3691059:3691145 [0] NCCL INFO threadThresholds 8/8/64 | 16/8/64 | 512 | 512
dsigpu4002:3691059:3691145 [0] NCCL INFO 4 coll channels, 0 collnet channels, 0 nvls channels, 4 p2p channels, 2 p2p channels per peer
dsigpu4002:3691060:3691146 [1] NCCL INFO comm 0xf5f3b80 rank 1 nranks 2 cudaDev 1 nvmlDev 1 busId d1000 commId 0x472e0f07d7c3c135 - Init COMPLETE
dsigpu4002:3691059:3691145 [0] NCCL INFO comm 0x10b6e8f0 rank 0 nranks 2 cudaDev 0 nvmlDev 0 busId ce000 commId 0x472e0f07d7c3c135 - Init COMPLETE