[transformer] Faster model parallel
@@ -907,23 +907,29 @@ class TransformerDecoder(nn.Module):
)
work_items = PipelineHelper.schedule_work_items(self.layers, chunks)
- new_incr_state = [{} for _ in chunks]
+ new_incr_state = {i: [] for i, _ in enumerate(self.layers)}
for chunk_idx, layer_nos, next_device in work_items:
s_tensor, s_enc_out, s_enc_mask, s_incr_state = chunks[chunk_idx]
for layer_no in layer_nos:
- s_tensor, new_incr_state[chunk_idx][layer_no] = self.layers[layer_no](
+ s_tensor, nis = self.layers[layer_no](
x=s_tensor,
encoder_output=s_enc_out,
encoder_mask=s_enc_mask,
incr_state=s_incr_state.get(layer_no),
)
- chunks[chunk_idx] = PipelineHelper.chunk_to(
- (s_tensor, s_enc_out, s_enc_mask, s_incr_state), next_device
+ new_incr_state[layer_no].append(nis)
+ # don't move incr state, it's always on the correct device
+ s_tensor, s_enc_out, s_enc_mask = PipelineHelper.chunk_to(
+ (s_tensor, s_enc_out, s_enc_mask), next_device
)
+ chunks[chunk_idx] = (s_tensor, s_enc_out, s_enc_mask, s_incr_state)
tensor_out = PipelineHelper.join([c[0] for c in chunks])
- new_incr_state = PipelineHelper.join(new_incr_state)
+ new_incr_state = {
+ layer_no: PipelineHelper.join(pieces)
+ for layer_no, pieces in new_incr_state.items()
+ }
return tensor_out, new_incr_state
