transfer mrope_position_delta to device when first running (#11047)
This commit is contained in:
@@ -575,9 +575,15 @@ class ForwardBatch:
|
|||||||
device=model_runner.device,
|
device=model_runner.device,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
mrope_position_deltas = mm_input.mrope_position_delta.flatten().to(
|
if mm_input.mrope_position_delta.device.type != model_runner.device:
|
||||||
|
# transfer mrope_position_delta to device when the first running,
|
||||||
|
# avoiding successvie host-to-device data transfer
|
||||||
|
mm_input.mrope_position_delta = (
|
||||||
|
mm_input.mrope_position_delta.to(
|
||||||
model_runner.device, non_blocking=True
|
model_runner.device, non_blocking=True
|
||||||
)
|
)
|
||||||
|
)
|
||||||
|
mrope_position_deltas = mm_input.mrope_position_delta.flatten()
|
||||||
mrope_positions_list[batch_idx] = (
|
mrope_positions_list[batch_idx] = (
|
||||||
(mrope_position_deltas + self.seq_lens[batch_idx] - 1)
|
(mrope_position_deltas + self.seq_lens[batch_idx] - 1)
|
||||||
.unsqueeze(0)
|
.unsqueeze(0)
|
||||||
|
|||||||
Reference in New Issue
Block a user