2D Rope into develop #1540
2D Rope into develop #1540
Conversation
- Add copyright attribution for rotate_half() and apply_rotary_pos_emb() functions - Rename apply_rotary_pos_emb_2d() to rotary_pos_emb_2d() for consistency - Rename config parameter use_2D_rope to rope_2D for better extensibility when supporting different RoPE variants in the future
clessig
left a comment
clessig
left a comment
There was a problem hiding this comment.
We can simplify the code if we change the syntax of the forward fct of the attention blocks slightly. This should probably be done consistently for all heads, even those without 2D rope.
| ) | ||
|
|
||
| return iter_start, iter_end | ||
| return iter_start, iter_end |
| def forward(self, tokens, coords=None, use_reentrant=False): | ||
| for block in self.ae_aggregation_blocks: | ||
| tokens = checkpoint(block, tokens, use_reentrant=use_reentrant) | ||
| if isinstance(block, MultiSelfAttentionHead | MultiSelfAttentionHeadLocal): |
There was a problem hiding this comment.
We could avoid the branches if we would follow the approach in the MLP:
There was a problem hiding this comment.
Or some compromise with *args and **kwargs
| def forward(self, tokens, coords=None, use_reentrant=False): | ||
| for block in self.ae_global_blocks: | ||
| tokens = checkpoint(block, tokens, use_reentrant=use_reentrant) | ||
| if isinstance(block, MultiSelfAttentionHead | MultiSelfAttentionHeadLocal): |
There was a problem hiding this comment.
Same as above: if-statement is not necessary
| for _b_idx, block in enumerate(self.fe_blocks): | ||
| if isinstance(block, torch.nn.modules.normalization.LayerNorm): | ||
| tokens = block(tokens) | ||
| elif isinstance(block, MultiSelfAttentionHead | MultiSelfAttentionHeadLocal): |
| ) | ||
| self.pe_global = torch.nn.Parameter(pe, requires_grad=False) | ||
|
|
||
| ### ROPE COORDS ### |
There was a problem hiding this comment.
standard comment:
# rope coords
| ks = self.lnorm_k(self.proj_heads_k(x).reshape(s)).to(self.dtype) | ||
| vs = self.proj_heads_v(x).reshape(s).to(self.dtype) | ||
|
|
||
| if coords is not None: |
There was a problem hiding this comment.
MultiSelfAttenionHead is a vanilla self attention block. The 2D rope positional encoding makes strong assumption of the structure of the qs and ks. At the very least, we need some documentation on what the option means and when it makes sense. Not sure if there's another good way to ensure proper use of this option (one could have a separate class but then one should make sure there's no code duplication).
3 participants
Description
This is @csjfwang 's implementation of 2D rope, coming from this orignal PR: #1445
Issue Number
Closes #1445
Closes #1109
Checklist before asking for review
./scripts/actions.sh lint./scripts/actions.sh unit-test./scripts/actions.sh integration-testlaunch-slurm.py --time 60