refactor(runner): unify pp_proxy_tensors forward kwarg into one helper (#28382)
This commit is contained in:
@@ -3292,6 +3292,14 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
forward_batch_template=forward_batch,
|
||||
)
|
||||
|
||||
def _pp_kwargs(self, pp_proxy_tensors) -> dict:
|
||||
"""Build the pp_proxy_tensors forward kwarg, in one place.
|
||||
|
||||
Pipeline-parallel proxy tensors are threaded into model.forward only
|
||||
when the model accepts them (``support_pp``).
|
||||
"""
|
||||
return {"pp_proxy_tensors": pp_proxy_tensors} if self.support_pp else {}
|
||||
|
||||
def forward_decode(
|
||||
self,
|
||||
forward_batch: ForwardBatch,
|
||||
@@ -3315,9 +3323,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
else:
|
||||
self.attn_backend.init_forward_metadata(forward_batch)
|
||||
# FIXME: add pp_proxy_tensors arg to all models
|
||||
kwargs = {}
|
||||
if self.support_pp:
|
||||
kwargs["pp_proxy_tensors"] = pp_proxy_tensors
|
||||
kwargs = self._pp_kwargs(pp_proxy_tensors)
|
||||
|
||||
# Launch forward
|
||||
ctx = (
|
||||
@@ -3350,9 +3356,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput], bool
|
||||
]:
|
||||
# Setup extra arguments
|
||||
kwargs = {}
|
||||
if self.support_pp:
|
||||
kwargs["pp_proxy_tensors"] = pp_proxy_tensors
|
||||
kwargs = self._pp_kwargs(pp_proxy_tensors)
|
||||
if forward_batch.input_embeds is not None:
|
||||
kwargs["input_embeds"] = forward_batch.input_embeds.bfloat16()
|
||||
if (
|
||||
@@ -3454,9 +3458,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
else:
|
||||
self.attn_backend.forward_metadata = None
|
||||
|
||||
kwargs = {}
|
||||
if self.support_pp:
|
||||
kwargs["pp_proxy_tensors"] = pp_proxy_tensors
|
||||
kwargs = self._pp_kwargs(pp_proxy_tensors)
|
||||
ctx = (
|
||||
self.device_timer.wrap(metadata={"category": "idle"})
|
||||
if self.device_timer
|
||||
|
||||
Reference in New Issue
Block a user