bigscience-workshop / bigscience-workshop/petals
Petals doesn't deal with server failure properly
- Dominant language
- Python
- Stars
- 10.6k
- Forks
- 650
- PR merge metrics
- No merged PRs in 30d
Description
Hi there, we'd like to report our findings on testing Petals' availability of fault tolerance.
We note that the current implementation of the method _step_ in the class __ServerInferenceSession_ from _inference_session.py_ contains the following content:
```
if self.history is None:
self.history = inputs
elif self.history.shape[1] == self._position:
self.history = torch.cat([self.history, inputs[:, -n_input_tokens:]], dim=1)
assert self.history.shape[1] == self._position + n_input_tokens, (
f"Broken input cache: span={self.span} shape={self.history.shape} "
f"position={self._position} n_input_tokens={n_input_tokens}"
)
```
where the attributes self.history and self.position are initialized as None and 0 respectively when an object of the class __ServerInferenceSession_ is created. The problem is, when a server fails, Petals replaces it with another server that serves the same blocks. However, the new server session is just initialized when joining the inference session, and its attribute _position_ is 0.
In the method __update_sequence_ of the class _InferenceSession_, the new server session's _history_ will be assigned the history of the failed server session:
`updated_sessions[0].history = self._server_sessions[server_idx].history
`
And during the inference, _n_input_tokens_ will always be 1. Thus, the assertion:
`assert self.history.shape[1] == self._position + n_input_tokens
`
is always likely to throw exceptions of "Broken input cache".
One possible solution is described as follow:
Delete the assert statement so that no exception will be thrown during the recovery process. Then, change the last few lines of code of method _step_ to:
```
self._position += n_input_tokens
s1 = outputs[0].shape[1]
if self.recover:
return outputs[0][0:1, s1 - 1:s1, :]
#return outputs[0][0][outputs[0].shape[1] - 1]
return outputs[0]
```
In which the _self.recover_ is a newly difined attribute of class __ServerInferenceSession_, representing whether or not this server session is to recover from a failed one, initialized as **False**, and will be set to True in the method __update_sequence_. This change is to tackle the problem that: when simply delete the assert statement, the returned value of _outputs[0]_ in the recoverd session will be a tensor of shape [1, (num of it's history inputs), 8192] instead of expected [1,1,8192].
By testing tens of examples, we believe this change work properly when dealing with server failures. The final outputs in case some server fail, are the same as the ones where no server fails.
Please check if there are such problems. Many thanks.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.