diff --git a/drivers/vfio/pci/mlx5/cmd.c b/drivers/vfio/pci/mlx5/cmd.c
index a1b17cd688b9..ef1c141dc5e0 100644
--- a/drivers/vfio/pci/mlx5/cmd.c
+++ b/drivers/vfio/pci/mlx5/cmd.c
@@ -70,9 +70,18 @@ int mlx5vf_cmd_query_vhca_migration_state(struct mlx5vf_pci_core_device *mvdev,
 	 * Running both in parallel, might end-up with a failure in the
 	 * incremental query command on un-tracked vhca.
 	 */
-	if (query_flags & MLX5VF_QUERY_INC)
+	if (query_flags & MLX5VF_QUERY_INC) {
 		wait_event(mvdev->saving_migf->save_wait,
 			   !mvdev->saving_migf->save_cb_active);
+		if (mvdev->saving_migf->precopy_err) {
+			/* In case we had a PRE_COPY error, only query full image for final image */
+			if (!(query_flags & MLX5VF_QUERY_FINAL)) {
+				*state_size = 0;
+				return 0;
+			}
+			query_flags &= ~MLX5VF_QUERY_INC;
+		}
+	}
 	MLX5_SET(query_vhca_migration_state_in, in, opcode,
 		 MLX5_CMD_OP_QUERY_VHCA_MIGRATION_STATE);
 	MLX5_SET(query_vhca_migration_state_in, in, vhca_id, mvdev->vhca_id);
@@ -291,7 +300,10 @@ void mlx5vf_mig_file_cleanup_cb(struct work_struct *_work)
 
 	mutex_lock(&migf->lock);
 	if (async_data->status) {
-		migf->is_err = true;
+		if (async_data->status == MLX5_CMD_STAT_BAD_RES_STATE_ERR)
+			migf->precopy_err = true;
+		else
+			migf->is_err = true;
 		wake_up_interruptible(&migf->poll_wait);
 	}
 	mutex_unlock(&migf->lock);
@@ -328,6 +340,8 @@ static void mlx5vf_save_callback(int status, struct mlx5_async_work *context)
 	 * The error and the cleanup flows can't run from an
 	 * interrupt context
 	 */
+	if (status == -EREMOTEIO)
+		status = MLX5_GET(save_vhca_state_out, async_data->out, status);
 	async_data->status = status;
 	queue_work(migf->mvdev->cb_wq, &async_data->work);
 }
@@ -356,6 +370,18 @@ int mlx5vf_cmd_save_vhca_state(struct mlx5vf_pci_core_device *mvdev,
 	async_data = &migf->async_data;
 	async_data->sgt = (!track && inc) ? &migf->final_table.sgt :
 		&migf->table.sgt;
+	if (migf->precopy_err) {
+		/*
+		 * In case we had a PRE_COPY error, SAVE is triggered only for
+		 * the final image, read device full image.
+		 */
+		inc = false;
+		/*
+		 * Turn off precopy_err to let reader proceed only once this
+		 * SAVE call is completed, otherwise final state might be lost.
+		 */
+		migf->precopy_err = false;
+	}
 	err = dma_map_sgtable(mdev->device, async_data->sgt,
 			      DMA_FROM_DEVICE, 0);
 	if (err)
diff --git a/drivers/vfio/pci/mlx5/cmd.h b/drivers/vfio/pci/mlx5/cmd.h
index 03f3b5e99879..784670848a7c 100644
--- a/drivers/vfio/pci/mlx5/cmd.h
+++ b/drivers/vfio/pci/mlx5/cmd.h
@@ -32,6 +32,7 @@ struct mlx5_vf_migration_file {
 	struct mutex lock;
 	u8 disabled:1;
 	u8 is_err:1;
+	u8 precopy_err:1;
 	u8 save_cb_active:1;
 	u8 header_read:1;
 
@@ -134,6 +135,7 @@ struct mlx5vf_pci_core_device {
 
 enum {
 	MLX5VF_QUERY_INC = (1UL << 0),
+	MLX5VF_QUERY_FINAL = (1UL << 1),
 };
 
 int mlx5vf_cmd_suspend_vhca(struct mlx5vf_pci_core_device *mvdev, u16 op_mod);
diff --git a/drivers/vfio/pci/mlx5/main.c b/drivers/vfio/pci/mlx5/main.c
index 6cdd4fc93818..db2d0166a0f5 100644
--- a/drivers/vfio/pci/mlx5/main.c
+++ b/drivers/vfio/pci/mlx5/main.c
@@ -243,7 +243,8 @@ static ssize_t mlx5vf_save_read(struct file *filp, char __user *buf, size_t len,
 
 	if (!(filp->f_flags & O_NONBLOCK)) {
 		if (wait_event_interruptible(migf->poll_wait,
-			     (MIGF_HAS_DATA(migf) || migf->is_err)))
+			     (MIGF_HAS_DATA(migf) || migf->is_err ||
+			      migf->precopy_err)))
 			return -ERESTARTSYS;
 	}
 
@@ -351,7 +352,7 @@ static __poll_t mlx5vf_save_poll(struct file *filp,
 	mutex_lock(&migf->lock);
 	if (migf->disabled || migf->is_err)
 		pollflags = EPOLLIN | EPOLLRDNORM | EPOLLRDHUP;
-	else if (MIGF_HAS_DATA(migf))
+	else if (MIGF_HAS_DATA(migf) || migf->precopy_err)
 		pollflags = EPOLLIN | EPOLLRDNORM;
 	mutex_unlock(&migf->lock);
 
@@ -490,7 +491,7 @@ static int mlx5vf_pci_save_device_inc_data(struct mlx5vf_pci_core_device *mvdev)
 	int ret;
 
 	ret = mlx5vf_cmd_query_vhca_migration_state(mvdev, &length,
-						    MLX5VF_QUERY_INC);
+					MLX5VF_QUERY_INC | MLX5VF_QUERY_FINAL);
 	if (ret)
 		return ret;
 
