Wrong ouputs after implementing IO bindings despite both model producing same logits,past_key_states and encoder outputs.

[VikasOjha666]

class T5Encoder(Module):
def init(self,enc_session):
super().init()
self.enc_session=enc_session
self.main_input_name = "input_ids"

def encoder_output_shape(self,input_ids):
input_shape=input_ids.shape
batch_size=input_shape[0]
seq_len=input_shape[1]
return (batch_size,seq_len,768)](url)

def forward(self,input_ids,attention_mask,**kwargs):
io_bindings=self.enc_session.io_binding()
io_bindings.bind_input(name="input_ids",
device_type="cuda",
device_id= 0,
element_type=np.longlong,
shape=list(input_ids.shape),
buffer_ptr=input_ids.data_ptr())

io_bindings.bind_input(name="attention_mask",
                            device_type="cuda",
                            device_id= 0,
                            element_type=np.longlong,
                            shape=list(attention_mask.shape),
                            buffer_ptr=attention_mask.data_ptr())
out_shape=self.encoder_output_shape(input_ids)
out_tensor=torch.empty(out_shape, dtype=torch.float32, device="cuda")
io_bindings.bind_output("hidden_states", "cuda")
io_bindings.bind_output(
        name="hidden_states",
        device_type="cuda",
        device_id=0,
        element_type=np.float32,  # hard coded output type
        shape=out_shape,
        buffer_ptr=out_tensor.data_ptr(),
    )
self.enc_session.run_with_iobinding(io_bindings)
return BaseModelOutput(out_tensor)

class T5DecoderInit(torch.nn.Module):
def init(self, decoder_sess):
super().init()
self.decoder = decoder_sess

def forward(self, input_ids, encoder_attention_mask, encoder_hidden_states):
    input_ids_shape=input_ids.shape
    device = input_ids.device.type
    dec_init_io_binding = self.decoder.io_binding()
    dec_init_io_binding.bind_input(name="input_ids",
                            device_type="cuda",
                            device_id=0,
                            element_type=np.longlong,
                            shape=list(input_ids.shape),
                            buffer_ptr=input_ids.data_ptr())
    dec_init_io_binding.bind_input(name="encoder_attention_mask",
                            device_type="cuda",
                            device_id= 0,
                            element_type=np.longlong,
                            shape=list(encoder_attention_mask.shape),
                            buffer_ptr=encoder_attention_mask.data_ptr())
                          
    dec_init_io_binding.bind_input(name="encoder_hidden_states",
                            device_type="cuda",
                            device_id=0,
                            element_type=np.float32,
                            shape=list(encoder_hidden_states.shape),
                            buffer_ptr=encoder_hidden_states.data_ptr())
    dec_init_io_binding.bind_output("logits", device)
    dec_init_io_binding.bind_output("past_key_values", device)

    dec_init_logits_tensor=torch.empty((input_ids_shape[0],input_ids_shape[1],32102), dtype=torch.float32, device="cuda")
    dec_init_io_binding.bind_output(
        name="logits",
        device_type="cuda",
        device_id=0,
        element_type=np.float32,  # hard coded output type
        shape=(input_ids_shape[0],input_ids_shape[1],32102),
        buffer_ptr=dec_init_logits_tensor.data_ptr(),
    )

    for arg in self.decoder.get_outputs():
        dec_init_io_binding.bind_output(arg.name, device)

    self.decoder.run_with_iobinding(dec_init_io_binding)
    ort_output = dec_init_io_binding.get_outputs()
    logits = ort_output[0]

    list_pkv =tuple([torch.from_numpy(out.numpy()).cuda() for out in ort_output[1:]])
    out_past_key_values = tuple(
        list_pkv[i : i + 4] for i in range(0, len(list_pkv), 4)
    )

    return torch.from_numpy(logits.numpy()).cuda(),out_past_key_values
    # return ort_output

class T5Decoder(torch.nn.Module):
def init(self, decoder_sess):
super().init()
self.decoder = decoder_sess

def forward(self, input_ids, attention_mask, encoder_output, past_key_values):
    input_ids_shape=input_ids.shape
    device = input_ids.device.type
    dec_io_binding = self.decoder.io_binding()
    dec_io_binding.bind_input(name="input_ids",
                            device_type="cuda",
                            device_id=0,
                            element_type=np.longlong,
                            shape=list(input_ids.shape),
                            buffer_ptr=input_ids.data_ptr())
    dec_io_binding.bind_input(name="encoder_attention_mask",
                            device_type="cuda",
                            device_id=0,
                            element_type=np.longlong,
                            shape=list(attention_mask.shape),
                            buffer_ptr=attention_mask.data_ptr())
                          
    dec_io_binding.bind_input(name="encoder_hidden_states",
                            device_type="cuda",
                            device_id=0,
                            element_type=np.float32,
                            shape=list(encoder_output.shape),
                            buffer_ptr=encoder_output.data_ptr())

    flat_past_key_values = functools.reduce(operator.iconcat, past_key_values, [])

    past_key_values = [
        (f"pkv_{i}", pkv) for i, pkv in enumerate(flat_past_key_values)
    ]
    
    for pkv in past_key_values:
        pkl_val=OrtValue.ortvalue_from_numpy(pkv[1].cpu().numpy(), 'cuda')
        dec_io_binding.bind_ortvalue_input(pkv[0],pkl_val)

    for arg in self.decoder.get_outputs():
        dec_io_binding.bind_output(arg.name, device)
    logits_tensor=torch.empty((input_ids_shape[0],input_ids_shape[1],32102), dtype=torch.float32, device="cuda")
    dec_io_binding.bind_output(
        name="logits",
        device_type="cuda",
        device_id=0,
        element_type=np.float32,  # hard coded output type
        shape=(input_ids_shape[0],input_ids_shape[1],32102),
        buffer_ptr=logits_tensor.data_ptr(),
    )

    self.decoder.run_with_iobinding(dec_io_binding)
    ort_output = dec_io_binding.get_outputs()
    logits = ort_output[0]
    

    list_pkv = tuple(torch.from_numpy(x.numpy()).cuda() for x in ort_output[1:])

    # creates a tuple of tuples of shape 6x4 from the above tuple
    out_past_key_values = tuple(
        list_pkv[i : i + 4] for i in range(0, len(list_pkv), 4)
    )

    # values of logits are not directly accessible. The workaround implies creating a new Tensor from
    # the numpy representation. A direct way to forward the Tensor would increase speed.
    return torch.from_numpy(logits.numpy()).cuda(), out_past_key_values

Hi @Ki6an .I implemented the IO bindings version for GPU currently I am haven't binded the outputs of decoder and decoder init but planning to do it soon. Everything else stays the same as your code but I am noticing very strange behaviour that the generate function with your encoder and decoders produce correct outputs while for mine generate function is producing wrong output. I have also tried using torch.cuda.synchronise in some debugging cases to synchronise the cuda kernels but that even didn't effected. The output of my encoder,decoder,decoder init is same as yours except they are cuda tensors instead of CPU tensor.Please help.

[shiqingzhangCSU]

hi！
Did you solve the problem? Can you share your experience?