Archives/petals
2
0
Fork 0
mirror of https://github.com/bigscience-workshop/petals synced 2024-11-11 19:11:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add checks for forward() inputs on the client side (#123)

Browse Source
This commit is contained in:
justheuristic 2022-12-03 14:02:48 +03:00 committed by GitHub
parent 055f85b83e
commit 8491ed2bd3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/petals/client/remote_sequential.py
View File

@ -53,6 +53,8 @@ class RemoteSequential(nn.Module):
self.is_subsequence = self.sequence_manager.sequence_info.block_uids != block_uids self.is_subsequence = self.sequence_manager.sequence_info.block_uids != block_uids
def forward(self, inputs: torch.Tensor, prompts: torch.Tensor = DUMMY): def forward(self, inputs: torch.Tensor, prompts: torch.Tensor = DUMMY):
assert inputs.ndim == 3, "inputs must be a tensor of shape [batch_size, seq_length, hidden_size]"
assert inputs.shape[1] <= 2048, "The sequence length is capped at 2048 tokens in this version"
outputs = _RemoteSequentialAutogradFunction.apply(inputs, prompts, self.sequence_manager) outputs = _RemoteSequentialAutogradFunction.apply(inputs, prompts, self.sequence_manager)
return outputs return outputs