多gpu 的时候运行ppo_training.py报错，

[cqray1990]

if self.is_encoder_decoder:
            input_ids = input_kwargs["decoder_input_ids"]
            attention_mask = input_kwargs["decoder_attention_mask"]
        else:
            input_ids = input_kwargs["input_ids"]
            attention_mask = input_kwargs["attention_mask"]

        logprobs = logprobs_from_logits(logits[:, :-1, :], input_ids[:, 1:])

/lib/python3.10/site-packages/trl/core.py", line 139, in logprobs_from_logits logpy = torch.gather(logp, 2, labels.unsqueeze(2)).squeeze(-1) RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:7 and cuda:0! (when checking argument for argument index in method wrapper_CUDA_gather)

model device_map 设置为auto,奖励模型默认设置cuda:0上 model = AutoModelForCausalLMWithValueHead.from_pretrained( args.model_name_or_path, config=config, torch_dtype=torch_dtype, load_in_4bit=args.load_in_4bit, load_in_8bit=args.load_in_8bit, device_map=args.device_map, trust_remote_code=args.trust_remote_code, peft_config=peft_config if args.use_peft else None, )

奖励模型默认设置cuda:0上 reward_model = AutoModelForSequenceClassification.from_pretrained( args.reward_model_name_or_path, config=reward_config, load_in_8bit=args.load_in_8bit, trust_remote_code=args.trust_remote_code, ) reward_model.to(device)

model device_map 设置为auto,奖励模型默认设置cuda:0上,这样导致计算的时候不在统一gpu，大佬没遇到这种情况？多GPU的话一定会出现呢

[shibing624]

奖励模型设置cuda:0，ppo model set cuda:1