Add block with Llama-like implementations #346
Conversation
olmo/model.py
Outdated
| @@ -309,10 +309,12 @@ def apply_rotary_pos_emb(self, pos_sin: torch.Tensor, pos_cos: torch.Tensor, t: | |||
| return out.to(t.dtype) | |||
|
|
|||
| def forward(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: | |||
| q_, k_ = q.float(), k.float() | |||
| q_, k_ = q.to(dtype=self.config.rope_precision), k.to(dtype=self.config.rope_precision) | |||
There was a problem hiding this comment.
These changes to rope should not affect the overall result if rope_precision_type = fp32 (the default).
| with torch.autocast(q.device.type, enabled=False): | ||
| query_len, key_len = q_.shape[-2], k_.shape[-2] # could be different if layer_past not None | ||
| pos_sin, pos_cos = self.get_rotary_embedding(key_len, q_.device) | ||
| pos_sin = pos_sin.type_as(q_) |
There was a problem hiding this comment.
Is this going to change any existing runs that use Rope?
There was a problem hiding this comment.
The rotary embeddings are fp32 by construction, so if rope_precision_type = fp32 (the default) then this shouldn't change the type.
olmo/model.py
Outdated
| ) -> torch.Tensor: | ||
| attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(q.size(-1)) | ||
|
|
||
| attn_bias = torch.zeros_like(attn_weights) |
There was a problem hiding this comment.
Yes. I can change this to re-use the attn_bias (Llama does this). I think there might be some code in place already that I can leverage
There was a problem hiding this comment.
The problem is that it changes attn_bias in place?
There was a problem hiding this comment.
I've updated the code to leverage existing get_causal_attention_bias. Now the bias should be re-used between calls (in the Llama implementation only). The moving around of get_causal_attention_bias to achieve this is not the cleanest.
olmo/model.py
Outdated
| attn_bias.masked_fill_(context_mask.logical_not(), torch.finfo(attn_bias.dtype).min) | ||
|
|
||
| if attn_mask is not None: | ||
| attn_bias += attn_mask.to(q.dtype) |
There was a problem hiding this comment.
Yes. I'm not sure how we can avoid doing this addition though. Llama does this addition every time it computes attention
|
Once this change is in, using Llama will require setting |
olmo/config.py
Outdated
| @@ -280,6 +286,11 @@ class ModelConfig(BaseConfig): | |||
| Use rotary positional embeddings (RoPE). Mutually exclusive with ``alibi``. | |||
| """ | |||
|
|
|||
| rope_precision_type: str = "fp32" | |||
There was a problem hiding this comment.
Consider making a StrEnum for the options here.
There was a problem hiding this comment.
Or never mind if you go with my other suggestion.
olmo/config.py
Outdated
| @@ -280,6 +286,11 @@ class ModelConfig(BaseConfig): | |||
| Use rotary positional embeddings (RoPE). Mutually exclusive with ``alibi``. | |||
| """ | |||
|
|
|||
| rope_precision_type: str = "fp32" | |||
| """ | |||
| Precision with which to apply RoPE (e.g. "amp_bf16", "amp_fp16", or "fp32"). | |||
There was a problem hiding this comment.
I don't think "amp_*" is meaningful here. It seems like there should really be two options:
- fp32, or
- whatever type q and k are
So maybe change this to a flag called rope_full_precision or something?
There was a problem hiding this comment.
I went with this suggestion instead of making it a StrEnum as in your other one.
When investigating why OLMo and Llama produce different results, we found a few different causes. This change adds a Llama block that we can use to address the following causes:
F.scaled_dot_product_attention.Also, OLMo always applies rotary embeddings in fp32, whereas Llama does it in the current type (which can be bf16). I have added a config that allows us to configure how rotary embeddings are applied.
Closes #345.