Method repartition

SwissArmyTransformer/sat/model/transformer.py:115–118

(self)

            return HOOKS_DEFAULT['attention_forward'](self, hidden_states, mask, **kw_args)

    def repartition(self):
        world_size = get_model_parallel_world_size()
        self.num_attention_heads_per_partition = divide(self.num_attention_heads, world_size)
        self.hidden_size_per_partition = self.hidden_size_per_attention_head * self.num_attention_heads_per_partition

class CrossAttention(torch.nn.Module):
    """Parallel cross-attention layer for Transformer"""