|
|
|
@ -64,41 +64,56 @@ class TransformerConnectionHandler(ConnectionHandler):
|
|
|
|
|
async with self._allocate_caches(requested_backends, batch_size, max_length) as cache_handles:
|
|
|
|
|
assert len(cache_handles) == len(requested_backends)
|
|
|
|
|
while request.tensors: # iterate while user is willing to supply tensors
|
|
|
|
|
hidden_states = [deserialize_torch_tensor(tensor) for tensor in request.tensors]
|
|
|
|
|
length_increment = hidden_states[0].shape[1] # how many tokens are added this step (in each seq)
|
|
|
|
|
hidden_states, prompts, hypo_ids = [deserialize_torch_tensor(tensor) for tensor in request.tensors]
|
|
|
|
|
|
|
|
|
|
# Cast inputs to backend dtype
|
|
|
|
|
hidden_states = hidden_states.to(requested_backends[0].dtype)
|
|
|
|
|
assert hypo_ids.dtype == torch.int64, f"hypo ids must be int64, got {hypo_ids.dtype}"
|
|
|
|
|
|
|
|
|
|
# parse deep prompts (optional argument)
|
|
|
|
|
if prompts is None or is_dummy(prompts) or is_dummy(prompts):
|
|
|
|
|
prompts = [DUMMY] * len(requested_backends)
|
|
|
|
|
else:
|
|
|
|
|
prompts = [p.squeeze(0) for p in prompts.to(requested_backends[0].dtype).split(1, dim=0)]
|
|
|
|
|
|
|
|
|
|
if not (len(requested_backends) == len(prompts)):
|
|
|
|
|
raise ValueError(f"Received {len(prompts)} prompts for {len(requested_backends)} backends")
|
|
|
|
|
|
|
|
|
|
length_increment = hidden_states.shape[1] # how many tokens are added this step (in each seq)
|
|
|
|
|
if prefix_length + length_increment > max_length:
|
|
|
|
|
raise ValueError(
|
|
|
|
|
f"Maximum length exceeded: prefix {prefix_length} + current {length_increment}"
|
|
|
|
|
f" exceeds pre-allocated maximum {max_length}"
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
# Cast inputs to backend dtype
|
|
|
|
|
hidden_states = [tensor.to(requested_backends[0].dtype) for tensor in hidden_states]
|
|
|
|
|
|
|
|
|
|
# run request tensors through all requested modules, update caches
|
|
|
|
|
for backend, cache_handle in zip(requested_backends, cache_handles):
|
|
|
|
|
for backend, prompt, cache_handle in zip(requested_backends, prompts, cache_handles):
|
|
|
|
|
if not is_dummy(prompt):
|
|
|
|
|
hidden_states[:, : prompt.shape[1]] += prompt
|
|
|
|
|
|
|
|
|
|
cache_metadata[:, 0], cache_metadata[:, 1] = cache_handle, prefix_length
|
|
|
|
|
assert isinstance(
|
|
|
|
|
hidden_states, torch.Tensor
|
|
|
|
|
), f"hidden states must be tensor, got {type(hidden_states)}"
|
|
|
|
|
assert (
|
|
|
|
|
len(hidden_states) == 1 and hidden_states[0].ndim == 3
|
|
|
|
|
hidden_states.ndim == 3
|
|
|
|
|
), f"inputs to {type(backend)} must be a list with a single 3d tensor of hidden states"
|
|
|
|
|
|
|
|
|
|
hidden_states = await backend.inference_pool.submit_task(cache_metadata, *hidden_states)
|
|
|
|
|
assert isinstance(hidden_states, (list, tuple))
|
|
|
|
|
assert len(hidden_states) == 1 and hidden_states[0].ndim == 3
|
|
|
|
|
(hidden_states,) = await backend.inference_pool.submit_task(
|
|
|
|
|
cache_metadata, hidden_states, hypo_ids
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
# serialize and send last layer outputs
|
|
|
|
|
yield runtime_pb2.ExpertResponse(
|
|
|
|
|
tensors=[
|
|
|
|
|
serialize_torch_tensor(result.to(proto.dtype), proto.compression, allow_inplace=True)
|
|
|
|
|
for result, proto in zip(
|
|
|
|
|
hidden_states, nested_flatten(requested_backends[-1].outputs_schema)
|
|
|
|
|
(hidden_states,), nested_flatten(requested_backends[-1].outputs_schema)
|
|
|
|
|
)
|
|
|
|
|
]
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
# prepare for next step
|
|
|
|
|
prefix_length += hidden_states[0].shape[1]
|
|
|
|
|
prefix_length += hidden_states.shape[1]
|
|
|
|
|
request = await (anext(requests))
|
|
|
|
|
finally:
|
|
|
|
|
print("CLOSED RPC_INFERENCE")
|
|
|
|
@ -238,23 +253,20 @@ async def _rpc_forward(*flat_tensors: torch.Tensor, requested_backends: Sequence
|
|
|
|
|
:param requested_backends: a sequence of transformer blocks in the same order as they appear in forward pass
|
|
|
|
|
:returns: hidden states after the last layer [batch_size, seq_length, hid_size]
|
|
|
|
|
"""
|
|
|
|
|
hidden_states, *prompts = flat_tensors
|
|
|
|
|
hidden_states, prompts = flat_tensors
|
|
|
|
|
dtype = requested_backends[0].dtype
|
|
|
|
|
# check parse input tensors and cast dtypes
|
|
|
|
|
hidden_states = hidden_states.to(dtype)
|
|
|
|
|
assert hidden_states.ndim == 3
|
|
|
|
|
if not prompts or is_dummy(prompts[0]):
|
|
|
|
|
if prompts is None or is_dummy(prompts):
|
|
|
|
|
prompts = [DUMMY] * len(requested_backends)
|
|
|
|
|
pre_seq_len = 0
|
|
|
|
|
else:
|
|
|
|
|
prompts = [prompts[0].to(requested_backends[0].dtype)]
|
|
|
|
|
prompts = [p.squeeze(0) for p in prompts[0].split(1)]
|
|
|
|
|
pre_seq_len = prompts[0].shape[-2]
|
|
|
|
|
prompts = [p.squeeze(0) for p in prompts.to(requested_backends[0].dtype).split(1, dim=0)]
|
|
|
|
|
|
|
|
|
|
# Run a chain of requested backends
|
|
|
|
|
for backend, prompt in zip(requested_backends, prompts):
|
|
|
|
|
if not is_dummy(prompt):
|
|
|
|
|
hidden_states[:, :pre_seq_len] += prompt
|
|
|
|
|
hidden_states[:, : prompt.shape[1]] += prompt
|
|
|
|
|
(hidden_states,) = await backend.forward_pool.submit_task(hidden_states)
|
|
|
|
|
assert isinstance(hidden_states, torch.Tensor)
|
|
|
|
|
assert (
|
|
|
|
@ -268,18 +280,15 @@ async def _rpc_forward(*flat_tensors: torch.Tensor, requested_backends: Sequence
|
|
|
|
|
async def _rpc_backward(
|
|
|
|
|
*flat_tensors: torch.Tensor, requested_backends: Sequence[TransformerBackend]
|
|
|
|
|
) -> Union[torch.Tensor, Sequence[torch.Tensor]]:
|
|
|
|
|
inputs, grad_outputs, *prompts = flat_tensors
|
|
|
|
|
inputs, grad_outputs, prompts = flat_tensors
|
|
|
|
|
# Cast inputs & grad outputs to backend dtype
|
|
|
|
|
inputs = inputs.to(requested_backends[0].dtype)
|
|
|
|
|
grad_outputs = grad_outputs.to(requested_backends[-1].dtype)
|
|
|
|
|
|
|
|
|
|
if not prompts or is_dummy(prompts[0]):
|
|
|
|
|
if prompts is None or is_dummy(prompts):
|
|
|
|
|
prompts = [DUMMY] * len(requested_backends)
|
|
|
|
|
pre_seq_len = 0
|
|
|
|
|
else:
|
|
|
|
|
prompts = [prompts[0].to(requested_backends[0].dtype)]
|
|
|
|
|
prompts = [p.squeeze(0) for p in prompts[0].split(1)]
|
|
|
|
|
pre_seq_len = prompts[0].shape[-2]
|
|
|
|
|
prompts = [p.squeeze(0) for p in prompts.to(requested_backends[0].dtype).split(1, dim=0)]
|
|
|
|
|
|
|
|
|
|
# Run a forward chain to collect intermediate inputs
|
|
|
|
|
# Note that we do not forward for the last module since we do not need its output
|
|
|
|
@ -287,13 +296,13 @@ async def _rpc_backward(
|
|
|
|
|
for backend, prompt in zip(requested_backends[:-1], prompts[:-1]):
|
|
|
|
|
assert inputs.ndim == 3, f"inputs to {type(backend)} must be a single 3d tensor of hidden states"
|
|
|
|
|
if not is_dummy(prompt):
|
|
|
|
|
inputs[:, :pre_seq_len] += prompt
|
|
|
|
|
inputs[:, : prompt.shape[1]] += prompt
|
|
|
|
|
inter_inputs.append(inputs)
|
|
|
|
|
(inputs,) = await backend.forward_pool.submit_task(inputs)
|
|
|
|
|
assert isinstance(inputs, torch.Tensor)
|
|
|
|
|
|
|
|
|
|
if not is_dummy(prompts[-1]):
|
|
|
|
|
inputs[:, :pre_seq_len] += prompts[-1]
|
|
|
|
|
inputs[:, : prompts[-1].shape[1]] += prompts[-1]
|
|
|
|
|
inter_inputs.append(inputs)
|
|
|
|
|
|
|
|
|
|
assert len(inter_inputs) == len(prompts) == len(requested_backends), "internal shape error during backward"
|
|
|
|
@ -303,7 +312,7 @@ async def _rpc_backward(
|
|
|
|
|
(grad_outputs,) = await backend.backward_pool.submit_task(inp, grad_outputs)
|
|
|
|
|
assert isinstance(grad_outputs, torch.Tensor)
|
|
|
|
|
if not is_dummy(prompt):
|
|
|
|
|
grad_prompts_reversed.append(grad_outputs[:, :pre_seq_len].unsqueeze(0))
|
|
|
|
|
grad_prompts_reversed.append(grad_outputs[:, : prompt.shape[1]].unsqueeze(0))
|
|
|
|
|
|
|
|
|
|
grad_prompts = torch.cat(grad_prompts_reversed[::-1], dim=0) if grad_prompts_reversed else DUMMY
|
|
|
|
|
return [grad_outputs] if is_dummy(grad_prompts) else [grad_outputs, grad_prompts] # TODO un-duct-tape
|
|
|
|
|